Re: [PATCH] io_uring: clear IO_RING_F_DRAIN_ACTIVE on drain completion flush

[Jens Axboe <axboe@kernel.dk>]

On 10/18/26 1:01 AM, Wei Gao wrote:
> 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.

Didn't quite feel right to me, below is with some LLM assistance
to try and get to the bottom of it.

Thanks for digging into this. I don't think it's just a race with a
slow io-wq free, though. DRAIN_ACTIVE is still set when the next request gets initialized, so it's forced async and goes through the defer list
every time, even if the drain request is freed in time.een read.

Basing the clear off any REQ_F_IO_DRAIN request in the flush is not
enough:

- A drain request marked REQ_F_REISSUE counts as has_drain, but
  io_free_batch_list() punts it back to io-wq in that same flush. With
  an empty defer list, we clear DRAIN_ACTIVE while the drain is still
  pending.

- A second IOSQE_IO_DRAIN request that fails init after io_init_drain()
  has run gets completed via io_req_defer_failed() right away. If that
  happens while an earlier drain is still running in io-wq, we end the
  drain, and the next request is issued before the earlier drain
  completes.

It also only works if the drain request completes through the flush.
Without DEFER_TASKRUN, io-wq posts the CQE directly from
io_req_complete_post(), so the application can see the CQE and submit
before anything clears the flag.

A drain request dispatched from the defer list always runs alone, so
we can just remember which one it is and end the drain when its CQE is
flushed, and have io-wq complete drain requests via task_work so that
happens under ->uring_lock. Something like the below. Can you give it
a spin with your reproducer?

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 7c5d99a8f75d..2e5a786f2a2a 100644
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -524,6 +524,8 @@ struct io_ring_ctx {
 	atomic_t			poll_wq_rcu_pending;
 	struct list_head		defer_list;
 	unsigned			nr_drained;
+	/* drain request dispatched from ->defer_list, until it completes */
+	struct io_kiocb			*drain_req;
 
 	/* protected by ->completion_lock */
 	unsigned			nr_req_allocated;
diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index e6569c204917..8ffe0c9a4e29 100644
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -480,6 +480,9 @@ static __cold noinline void io_queue_deferred(struct io_ring_ctx *ctx)
 
 	lockdep_assert_held(&ctx->uring_lock);
 	__io_req_caches_free(ctx);
+	/* nothing outside of ->defer_list is alive */
+	if (ctx->nr_req_allocated == ctx->nr_drained)
+		ctx->drain_req = NULL;
 
 	while (!list_empty(&ctx->defer_list)) {
 		struct io_defer_entry *de = list_first_entry(&ctx->defer_list,
@@ -491,6 +494,9 @@ static __cold noinline void io_queue_deferred(struct io_ring_ctx *ctx)
 
 		list_del_init(&de->list);
 		ctx->nr_drained -= io_linked_nr(de->req);
+		/* it runs alone, track its completion unless it's a link */
+		if ((de->req->flags & REQ_F_IO_DRAIN) && !de->req->link)
+			ctx->drain_req = de->req;
 		io_req_task_queue(de->req);
 		kfree(de);
 		first = false;
@@ -980,9 +986,11 @@ static void io_req_complete_post(struct io_kiocb *req, unsigned issue_flags)
 
 	/*
 	 * Handle special CQ sync cases via task_work. DEFER_TASKRUN requires
-	 * the submitter task context, IOPOLL protects with uring_lock.
+	 * the submitter task context, IOPOLL protects with uring_lock. A drain
+	 * request ends the drain under uring_lock as its CQE is posted.
 	 */
-	if ((ctx->int_flags & IO_RING_F_LOCKLESS_CQ) || (req->flags & REQ_F_REISSUE)) {
+	if ((ctx->int_flags & IO_RING_F_LOCKLESS_CQ) ||
+	    (req->flags & (REQ_F_REISSUE | REQ_F_IO_DRAIN))) {
 defer_complete:
 		req->io_task_work.func = io_req_task_complete;
 		io_req_task_work_add(req);
@@ -1185,11 +1193,33 @@ static void io_free_batch_list(struct io_ring_ctx *ctx,
 	} while (node);
 }
 
+/*
+ * The dispatched drain request ran alone, and anything submitted after it was
+ * deferred. Once it completed, new requests may run if nothing is waiting.
+ */
+static __cold bool io_drain_req_completed(struct io_ring_ctx *ctx)
+{
+	struct io_wq_work_node *node;
+
+	__wq_list_for_each(node, &ctx->submit_state.compl_reqs) {
+		struct io_kiocb *req = container_of(node, struct io_kiocb,
+					    comp_list);
+
+		/* a requeued request hasn't completed */
+		if (req == ctx->drain_req && !(req->flags & REQ_F_REISSUE)) {
+			ctx->drain_req = NULL;
+			return true;
+		}
+	}
+	return false;
+}
+
 void __io_submit_flush_completions(struct io_ring_ctx *ctx)
 	__must_hold(&ctx->uring_lock)
 {
 	struct io_submit_state *state = &ctx->submit_state;
 	struct io_wq_work_node *node;
+	bool drain_done = false;
 
 	/*
 	 * Drop the files before the CQEs are posted, so they're released by
@@ -1224,13 +1254,19 @@ void __io_submit_flush_completions(struct io_ring_ctx *ctx)
 	}
 	__io_cq_unlock_post(ctx);
 
+	if (unlikely(ctx->int_flags & IO_RING_F_DRAIN_ACTIVE) && ctx->drain_req)
+		drain_done = io_drain_req_completed(ctx);
+
 	if (!wq_list_empty(&state->compl_reqs)) {
 		io_free_batch_list(ctx, state->compl_reqs.first);
 		INIT_WQ_LIST(&state->compl_reqs);
 	}
 
-	if (unlikely(ctx->int_flags & IO_RING_F_DRAIN_ACTIVE))
+	if (unlikely(ctx->int_flags & IO_RING_F_DRAIN_ACTIVE)) {
 		io_queue_deferred(ctx);
+		if (drain_done && list_empty(&ctx->defer_list))
+			ctx->int_flags &= ~IO_RING_F_DRAIN_ACTIVE;
+	}
 
 	ctx->submit_state.cq_flush = false;
 }

-- 
Jens Axboe