[RFC PATCH 1/2] io_uring: don't let a preempted task_work add hold back the wake up

[Austin Schuh <austin.linux@gmail.com>]

Only one local task_work add per wait cycle wakes the waiter: the one
that takes ->cq_wait_nr's credit.  Every later add sees the credit gone
and returns.  The credit is taken after the entry is queued and the
wake up comes after that, with nothing held in between, so an adder
preempted in that window leaves the waiter asleep with work queued.  A
lazy add preempted before it counts itself down does the same: the add
that would have reached zero stops one short.

The waiter stays asleep for as long as anything of higher priority runs
on the adder's CPU, whatever the waiter's own priority.  Nothing boosts
the adder, since it holds no lock.  This bites IORING_OP_MSG_RING into
a DEFER_TASKRUN ring in particular, where the adds come from other
tasks of any priority: a real-time sender's message waits for a
lower-priority sender that took the credit and was preempted.  On
PREEMPT_RT, where threaded interrupt handlers preempt everything below
them, that is routine.

Traced on a Jetson Orin Nano (6.18.39-rt), with a SCHED_FIFO 60 sender
posting to a SCHED_FIFO 60 receiver every 500us while lower-priority
processes post to the same ring from CPUs loaded with SCHED_FIFO 50
bursts: of the 17 wakeups over 100us late in 60s, 15 were delivered by
a lower-priority sender whose MSG_RING had been in flight since before
the real-time sender's returned, for 12 to 479us.

On for-7.4/io_uring with PREEMPT_RT, before and after this patch, the
real-time sender on CPU 4 posts a MSG_RING to the real-time receiver's
DEFER_TASKRUN ring on CPU 5 every 500us.  On CPUs 0 and 1,
lower-priority processes post to the same ring every 200us, alongside
SCHED_FIFO 50 threads running 20us to 1ms bursts and normal-priority
interrupt and system load.  Ten minutes a run, 1.2 million messages
each, timed from the send to the receiver handling it:

                              late by    late by
                     runs   over 100us  over 1ms   longest
  FIFO 40, before      3     889-1139   341-417   8.6-11.8ms
  FIFO 40, after       3        15-21         0   0.69-0.94ms
  FIFO 60, before      1         1058       438     8.3ms
  FIFO 60, after       1            0         0      71us

The three runs at FIFO 40 are that setup as is, plus a thread in each
real-time process mapping and unmapping memory on CPU 0 or 1, and with
no lower-priority processes starting or stopping.  What remains at FIFO
40 is the sender's own MSG_RING being slow, 44 to 60 sends over 100us a
run, before and after.  Traced on 6.18, its request allocation waits on
the slab node's list_lock, held by a lower-priority process that
inheritance only lifts to 40, under the FIFO 50 bursts.  At 60, nothing
is left.

Keep preemption disabled from queueing the entry to the wake up, as
mpscq_push() already does for the push itself.  That covers a few
atomics and try_to_wake_up().  eventfd_signal() takes a spinlock_t,
which sleeps on PREEMPT_RT, so signal the eventfd and wake ring pollers
after that instead of before.

Fixes: 8751d15426a3 ("io_uring: reduce scheduling due to tw")
Assisted-by: LLM
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
---
 io_uring/tw.c | 80 +++++++++++++++++++++++++++++++++------------------
 1 file changed, 52 insertions(+), 28 deletions(-)

diff --git a/io_uring/tw.c b/io_uring/tw.c
index 0b5066f97224..cfb591508662 100644
--- a/io_uring/tw.c
+++ b/io_uring/tw.c
@@ -179,37 +179,14 @@ static void io_local_work_wake_poll(struct io_ring_ctx *ctx, unsigned int flags)
 	call_rcu_hurry(&ctx->poll_wq_rcu, io_poll_wq_wake_rcu);
 }
 
-void io_req_local_work_add(struct io_kiocb *req, unsigned flags)
+/*
+ * Issue the wake up of the waiter's cycle if this add is the one that takes
+ * ->cq_wait_nr's credit, see io_cq_wait_arm().
+ */
+static void io_local_work_wake(struct io_ring_ctx *ctx, unsigned int flags)
 {
-	struct io_ring_ctx *ctx = req->ctx;
 	int nr_wait;
 
-	/* pairs with synchronize_rcu() in io_ring_exit_work() */
-	guard(rcu)();
-
-	/*
-	 * We don't know how many requests there are in the link and whether
-	 * they can even be queued lazily, fall back to non-lazy.
-	 */
-	if (req->flags & IO_REQ_LINK_FLAGS)
-		flags &= ~IOU_F_TWQ_LAZY_WAKE;
-
-	/*
-	 * The xchg() in mpscq_push() implies a full barrier, which pairs with
-	 * the barrier in io_cq_wait_arm() on the io_cqring_wait() side. This
-	 * ensures that either we see the updated ->cq_wait_nr, or waiters going
-	 * to sleep will observe the work added to the list, which is similar to
-	 * the wait/wake task state sync.
-	 */
-	if (mpscq_push(&ctx->work_list, &req->io_task_work.node)) {
-		io_ctx_mark_taskrun(ctx);
-		if (data_race(ctx->int_flags) & IO_RING_F_HAS_EVFD)
-			io_eventfd_signal(ctx, false, flags & IOU_F_TWQ_IN_WAKE);
-		/* whoever sets the flag wakes ->poll_wq afterwards */
-		if (data_race(ctx->int_flags) & IO_RING_F_POLL_ACTIVATED)
-			io_local_work_wake_poll(ctx, flags);
-	}
-
 	/*
 	 * No one is waiting (IO_CQ_WAKE_INIT), or this cycle's wake up has
 	 * already been issued (zero or negative, see below).
@@ -239,6 +216,53 @@ void io_req_local_work_add(struct io_kiocb *req, unsigned flags)
 	wake_up_state(ctx->submitter_task, TASK_INTERRUPTIBLE);
 }
 
+void io_req_local_work_add(struct io_kiocb *req, unsigned flags)
+{
+	struct io_ring_ctx *ctx = req->ctx;
+	bool first;
+
+	/* pairs with synchronize_rcu() in io_ring_exit_work() */
+	guard(rcu)();
+
+	/*
+	 * We don't know how many requests there are in the link and whether
+	 * they can even be queued lazily, fall back to non-lazy.
+	 */
+	if (req->flags & IO_REQ_LINK_FLAGS)
+		flags &= ~IOU_F_TWQ_LAZY_WAKE;
+
+	/*
+	 * Only the add that takes the credit wakes the waiter, later ones
+	 * return early, and a lazy add that hasn't counted itself down yet
+	 * holds back the one that would. Neither may be preempted between
+	 * queueing its entry and the wake up: the waiter, whatever its
+	 * priority, sleeps with the entry queued for as long as anything of
+	 * higher priority runs on the adder's CPU, and nothing boosts the
+	 * adder, as it holds no lock.
+	 *
+	 * The xchg() in mpscq_push() implies a full barrier, which pairs with
+	 * the barrier in io_cq_wait_arm() on the io_cqring_wait() side. This
+	 * ensures that either we see the updated ->cq_wait_nr, or waiters going
+	 * to sleep will observe the work added to the list, which is similar to
+	 * the wait/wake task state sync.
+	 */
+	preempt_disable();
+	first = mpscq_push(&ctx->work_list, &req->io_task_work.node);
+	if (first)
+		io_ctx_mark_taskrun(ctx);
+	io_local_work_wake(ctx, flags);
+	preempt_enable();
+
+	/* eventfd_signal() takes a spinlock_t, not allowed above on PREEMPT_RT */
+	if (first) {
+		if (data_race(ctx->int_flags) & IO_RING_F_HAS_EVFD)
+			io_eventfd_signal(ctx, false, flags & IOU_F_TWQ_IN_WAKE);
+		/* whoever sets the flag wakes ->poll_wq afterwards */
+		if (data_race(ctx->int_flags) & IO_RING_F_POLL_ACTIVATED)
+			io_local_work_wake_poll(ctx, flags);
+	}
+}
+
 void io_req_normal_work_add(struct io_kiocb *req)
 {
 	struct io_uring_task *tctx = req->tctx;
-- 
2.47.3