[PATCH] io_uring: clear IO_RING_F_DRAIN_ACTIVE on drain completion flush

[Wei Gao <wegao@suse.com>]

When running a test that repeatedly submits requests followed by a
drained request across iterations on the same ring (such as liburing's
test/multicqes_drain.t), the test occasionally deadlocks and times out
waiting for completions (-ETIME).

Test mechanism (multicqes_drain.t):
The test executes the following sequence in each iteration on the same
ring:
1. Submits 2 poll requests on pipes:
   - sqe[0]: multishot poll on pipe1 (user_data = 10 / 0xa)
   - sqe[1]: oneshot poll on pipe2 (user_data = 20 / 0x14)
2. Triggers events via trigger_event():
   - Writes "foo" to pipe1 twice, sleeps 1ms, and reads "foo" back.
     Generates 2 CQEs for multishot poll #10 (IORING_CQE_F_MORE).
   - Writes "foo" to pipe2 once, sleeps 1ms, and reads "foo" back.
     Generates 1 CQE for oneshot poll #20.
3. Submits poll_remove and a drained NOP:
   - sqe[0]: poll_remove #10 (user_data = 30 / 0x1e)
   - sqe[1]: NOP with IOSQE_IO_DRAIN (user_data = 40 / 0x28)
   Produces 1 CQE for poll_remove (ret=0), 1 CQE for multishot
   cancel (ret=-ECANCELED), and 1 CQE for the drained NOP.
4. Harvests all 6 expected CQEs with io_uring_wait_cqe_timeout().
5. Loops back to step 1 on the same ring for the next iteration.

Why the failure occurs:
Prior to commit 8fb7aee05591 ("io_uring: drain based on allocates reqs"),
io_queue_deferred() checked CQ sequence numbers:
(seq + ctx->cq_extra != ctx->cached_cq_tail). When the prior drain request
posted its CQE, sequence numbers matched immediately. Even if
IO_RING_F_DRAIN_ACTIVE was not cleared upon completion, the next
iteration's non-drain request dequeued right away, emptied
ctx->defer_list, and cleared the flag via io_drain_req():

    if (!drain && list_empty(&ctx->defer_list))
        ctx->int_flags &= ~IO_RING_F_DRAIN_ACTIVE;

However, commit 8fb7aee05591 ("io_uring: drain based on allocates reqs")
changed the defer check:

-   if ((drain_seen || first) && io_drain_defer_seq(de->req, de->seq))
-       break;
+   if ((drain_seen || first) && ctx->nr_req_allocated != ctx->nr_drained)
+       return;

This change severed the self-healing mechanism:
1. Originally, the check verified whether all CQEs had been posted
   (seq + cq_extra == cached_cq_tail, e.g. 6 == 6). Even if the io-wq
   worker thread was delayed in freeing request memory, cached_cq_tail had
   already advanced, allowing the request to dequeue and self-heal.
2. The commit changed the check to whether slab memory has been freed
   (nr_req_allocated != nr_drained). Drained requests run via io-wq
   (due to forced async) and have an asynchronous refcount teardown: the
   CQE is reaped by userspace, but the io-wq worker thread drops its
   reference slightly later via io_wq_free_work() -> io_free_req().
3. If iteration N+1 submits before io_wq_free_work() runs, the allocation
   counter still accounts for the prior NOP (allocated=2, drained=1).
4. Seeing 2 != 1, io_queue_deferred() aborts early via return. This
   completely bypasses list_del_init() and the self-healing clearance in
   io_drain_req(), stranding the new poll requests on ctx->defer_list.
5. While polls are stranded on ctx->defer_list, userspace writes to
   and reads from the pipe. When polls are finally released later, the
   pipe data is already gone. The multishot poll permanently misses the
   event, producing fewer CQEs and causing a timeout (-ETIME).

Trace comparison (Passing vs Failing iteration):

Kprobes registered for tracing:
  p:drain_check io_queue_deferred+0xd7 \
      allocated=+0x464:u32 drained=+0x460:u32
  p:drain_blocked io_queue_deferred+0xeb \
      allocated=+0x464:u32 drained=+0x460:u32
  p:drain_dequeued io_queue_deferred+0x88 drained=+0x460:u32
  p:free_work io_wq_free_work work=%di
  p:flush_compl __io_submit_flush_completions \
      allocated=+0x464:u32 int_flags=+0x4:u32

[Passing iteration - iou-wrk frees slab before next submission]:
1. Prior drained NOP completes, slab memory freed in time (1 == 1):
   400.411558: flush_compl: allocated=1 int_flags=208
   400.411558: io_uring_complete: user_data 0x28, result 0, cflags 0x0
2. New poll requests enter io_drain_req(); check matches (1 == 1):
   400.411563: drain_check: allocated=1 drained=1
   400.411564: drain_dequeued: drained=1
   400.411566: free_work: work freed
3. Polls arm on pipes immediately before userspace writes; events caught:
   400.412638: io_uring_complete: user_data 0xa, result 1, cflags 0x2
   400.412638: flush_compl: allocated=9 int_flags=192
   400.413703: io_uring_complete: user_data 0xa, result 1, cflags 0x2
   400.413703: flush_compl: allocated=9 int_flags=192
   400.414768: flush_compl: allocated=9 int_flags=192
   400.414768: io_uring_complete: user_data 0x14, result 1, cflags 0x0
4. poll_remove and drained NOP submitted; drained NOP waits:
   400.414769: drain_check: allocated=3 drained=1
   400.414770: drain_blocked: allocated=3 drained=1
   400.414770: flush_compl: allocated=3 int_flags=208
   400.414770: io_uring_complete: user_data 0x1e, result 0, cflags 0x0
   400.414770: drain_check: allocated=2 drained=1
   400.414770: drain_blocked: allocated=2 drained=1
   400.414771: flush_compl: allocated=2 int_flags=208
   400.414771: io_uring_complete: user_data 0xa, result -125, cflags 0x0
5. All prior requests done (1 == 1); drained NOP dequeues and completes:
   400.414771: drain_check: allocated=1 drained=1
   400.414771: drain_dequeued: drained=1
   400.414773: free_work: work freed
   400.414774: flush_compl: allocated=1 int_flags=208
   400.414774: io_uring_complete: user_data 0x28, result 0, cflags 0x0
   => All 6 CQEs reaped successfully (PASS).

[Failing iteration - iou-wrk delayed by 15 us -> FAIL]:
1. Iteration N NOP completes, but iou-wrk free_work has not run yet:
   400.424650: flush_compl: allocated=1 int_flags=208
   400.424650: io_uring_complete: user_data 0x28, result 0, cflags 0x0
2. 10 us later, Iteration N+1 submits sqe[0] (POLL_ADD) and sqe[1].
   Because DRAIN_ACTIVE was not cleared, requests enter io_drain_req().
   Because slab memory for the prior NOP is not freed yet, counters
   mismatch (2 != 1):
   400.424660: drain_check: allocated=2 drained=1
   400.424661: drain_blocked: allocated=2 drained=1
   400.424662: drain_check: allocated=3 drained=2
   400.424662: drain_blocked: allocated=3 drained=2
3. 5 us later, iou-wrk finally frees old NOP (polls already blocked):
   400.424665: free_work: work freed
4. While polls are trapped on defer_list, userspace writes and reads pipe.
   When polls are finally released 1ms later, the pipe is already empty:
   400.425726: flush_compl: allocated=3 int_flags=208
   400.425727: drain_check: allocated=2 drained=2
   400.425727: drain_dequeued: drained=2
   400.425727: drain_dequeued: drained=1
5. Pipe event permanently missed; 1.0s later, wait_cqe times out (-ETIME):
   401.428202: flush_compl: allocated=3 int_flags=208
   401.428203: io_uring_complete: user_data 0xa, result -125, cflags 0x0
   401.428207: drain_check: allocated=2 drained=2
   401.428208: drain_dequeued: drained=2
   401.428210: drain_blocked: allocated=2 drained=1
   401.428213: flush_compl: allocated=2 int_flags=208
   401.428213: io_uring_complete: user_data 0x1e, result -14, cflags 0x0
   401.428213: io_uring_complete: user_data 0x28, result -125, cflags 0x0
   => wait_cqe timed out: ret=-62 (-ETIME)

Fix:
Make io_queue_deferred() return a bool indicating whether any drain
request was active or deferred. In __io_submit_flush_completions(), track
whether any completing request had REQ_F_IO_DRAIN. Clear
IO_RING_F_DRAIN_ACTIVE only when a drained request has completed
(has_drain), no drain was active or deferred (!drain_dispatched), and
ctx->defer_list is empty. This prevents premature clearance if an
unrelated completion flushes while a drain request is still executing
in io-wq.

All 10,000 iterations of reproduce_drain and liburing multicqes_drain.t
pass reliably with this fix.

Fixes: 8fb7aee05591 ("io_uring: drain based on allocates reqs")
Signed-off-by: Wei Gao <wegao@suse.com>
---
 io_uring/io_uring.c | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index cf5218d100d9..4f7b054e6f97 100644
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -445,7 +445,7 @@ unsigned io_linked_nr(struct io_kiocb *req)
 	return nr;
 }
 
-static __cold noinline void io_queue_deferred(struct io_ring_ctx *ctx)
+static __cold noinline bool io_queue_deferred(struct io_ring_ctx *ctx)
 {
 	bool drain_seen = false, first = true;
 
@@ -458,7 +458,7 @@ static __cold noinline void io_queue_deferred(struct io_ring_ctx *ctx)
 
 		drain_seen |= de->req->flags & REQ_F_IO_DRAIN;
 		if ((drain_seen || first) && ctx->nr_req_allocated != ctx->nr_drained)
-			return;
+			return true;
 
 		list_del_init(&de->list);
 		ctx->nr_drained -= io_linked_nr(de->req);
@@ -466,6 +466,8 @@ static __cold noinline void io_queue_deferred(struct io_ring_ctx *ctx)
 		kfree(de);
 		first = false;
 	}
+
+	return drain_seen;
 }
 
 void __io_commit_cqring_flush(struct io_ring_ctx *ctx)
@@ -1131,12 +1133,16 @@ void __io_submit_flush_completions(struct io_ring_ctx *ctx)
 {
 	struct io_submit_state *state = &ctx->submit_state;
 	struct io_wq_work_node *node;
+	bool has_drain = false;
 
 	__io_cq_lock(ctx);
 	__wq_list_for_each(node, &state->compl_reqs) {
 		struct io_kiocb *req = container_of(node, struct io_kiocb,
 					    comp_list);
 
+		if (req->flags & REQ_F_IO_DRAIN)
+			has_drain = true;
+
 		/*
 		 * Requests marked with REQUEUE should not post a CQE, they
 		 * will go through the io-wq retry machinery and post one
@@ -1157,8 +1163,12 @@ void __io_submit_flush_completions(struct io_ring_ctx *ctx)
 		INIT_WQ_LIST(&state->compl_reqs);
 	}
 
-	if (unlikely(ctx->int_flags & IO_RING_F_DRAIN_ACTIVE))
-		io_queue_deferred(ctx);
+	if (unlikely(ctx->int_flags & IO_RING_F_DRAIN_ACTIVE)) {
+		bool drain_dispatched = io_queue_deferred(ctx);
+
+		if (has_drain && !drain_dispatched && list_empty(&ctx->defer_list))
+			ctx->int_flags &= ~IO_RING_F_DRAIN_ACTIVE;
+	}
 
 	ctx->submit_state.cq_flush = false;
 }
-- 
2.55.0