Re: [PATCH 4/6] io_uring: switch normal task_work to a mpscq

On 6/15/26 12:47 PM, Jens Axboe wrote:
> On 6/15/26 12:33 PM, Caleb Sander Mateos wrote:
>> On Sat, Jun 13, 2026 at 5:08?AM Jens Axboe <[email protected]> wrote:
>>>
>>> On 6/12/26 8:26 PM, Caleb Sander Mateos wrote:
>>>> On Fri, Jun 12, 2026 at 12:37?PM Jens Axboe <[email protected]> wrote:
>>>>>
>>>>> On 6/12/26 12:59 PM, Caleb Sander Mateos wrote:
>>>>>>> @@ -236,10 +262,14 @@ void io_req_normal_work_add(struct io_kiocb *req)
>>>>>>>                 return;
>>>>>>>         }
>>>>>>>
>>>>>>> +       /* task_work must only be added once */
>>>>>>> +       if (test_and_set_bit(0, &tctx->tw_pending))
>>>>>>> +               return;
>>>>>>
>>>>>> Is tw_pending necessary? How come the task_work_add() exclusivity
>>>>>> isn't already provided by the mpscq_push() check above?
>>>>>
>>>>> It is, because the transition from empty -> not-empty no longer works
>>>>> for that, as the mpscq emtpies one-by-one rather than with a delete-all
>>>>> kind of primitive.
>>>>
>>>> Sorry, I'm still not following why the empty check doesn't suffice.
>>>> It's true that mpscq elements can be removed from the head one at a
>>>> time, but mpscq_push() will continue to return false until the
>>>> consumer pops all the elements and successfully sets tail back to
>>>> &stub. mpscq_push() will return true once when tail transitions away
>>>> from &stub, and then not again until the task work runs and sets tail
>>>> back to &stub.
>>>
>>> Let's say the task_work is currently running, a producer is adding more.
>>> It finds queue empty, re-adds the task_work. That part is fine, we can
>>> add the task_work while it's running as it has been detached already.
>>> The task_work keeps running and also prunes this new item. Producer adds
>>> another one, finds the queue empty, re-adds task_work. This one is not
>>> OK, the task_work was already re-added when it previously found it
>>> empty. Boom.
>>
>> Ah right, I forgot that mpscq_pop() can both return a popped node and
>> set the tail back to &stub. Maybe it would make sense for it to return
>> whether the queue has been marked empty and break out of
>> tctx_task_work_run() in that case instead of relying on a separate
>> call to mpscq_empty()? The atomic RMW for tw_pending every time the
>> queue transitions between empty and non-empty seems like it could be
>> quite expensive.
> 
> We could tweak it like that. I didn't look too closely as this is the
> !DEFER case and hence a lot less interesting, but if you want to send a
> patch my way I'd be happy to stage it on top.

I took a look, and yes I think it actually comes out nicer this way.
Good suggestion! It also helps cap the number of task_work items run,
which is a nice side effect. What do you think?

Needs a commit message obviously.

commit 572a1fb6d0f25b706ff044fcf141827f49db2ec0
Author: Jens Axboe <[email protected]>
Date:   Mon Jun 15 13:43:16 2026 -0600

    io_uring: get rid of tw_pending for !DEFER task work
    
    Suggested-by: Caleb Sander Mateos <[email protected]>
    Signed-off-by: Jens Axboe <[email protected]>

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 6415a3353ee0..87151a5b62c1 100644
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -149,8 +149,6 @@ struct io_uring_task {
 
 	struct { /* task_work */
 		struct mpscq		task_list;
-		/* BIT(0) guards adding tw only once */
-		unsigned long		tw_pending;
 		struct callback_head	task_work;
 	} ____cacheline_aligned_in_smp;
 };
diff --git a/io_uring/mpscq.h b/io_uring/mpscq.h
index c801384c6a0a..f910526766fd 100644
--- a/io_uring/mpscq.h
+++ b/io_uring/mpscq.h
@@ -122,4 +122,13 @@ static inline struct llist_node *mpscq_pop(struct mpscq *q,
 	return NULL;
 }
 
+/*
+ * Returns true if the most recent mpscq_pop() that returned a node also
+ * emptied the queue. Consumer must be serialized.
+ */
+static inline bool mpscq_pop_emptied(struct mpscq *q, struct llist_node *head)
+{
+	return head == &q->stub;
+}
+
 #endif /* IOU_MPSCQ_H */
diff --git a/io_uring/tw.c b/io_uring/tw.c
index e74372233f40..f2ce806b01a1 100644
--- a/io_uring/tw.c
+++ b/io_uring/tw.c
@@ -34,10 +34,6 @@ void io_tctx_fallback_work(struct work_struct *work)
 						  fallback_work);
 	unsigned int count = 0;
 
-	/* see tctx_task_work() - a set bit must always have a run coming */
-	clear_bit(0, &tctx->tw_pending);
-	smp_mb__after_atomic();
-
 	/*
 	 * Run the entries directly. We're in PF_KTHRED context, hence
 	 * io_should_terminate_tw() is true and they will be marked as
@@ -101,6 +97,13 @@ void tctx_task_work_run(struct io_uring_task *tctx, unsigned int max_entries,
 				io_poll_task_func, io_req_rw_complete,
 				(struct io_tw_req){req}, ts);
 		(*count)++;
+		/*
+		 * Break if most recent pop emptied the queue. This helps
+		 * bound task_work run, and also protects the regular
+		 * task_work addition.
+		 */
+		if (mpscq_pop_emptied(&tctx->task_list, tctx->task_head))
+			break;
 		if (unlikely(need_resched())) {
 			ctx_flush_and_put(ctx, ts);
 			ctx = NULL;
@@ -127,8 +130,6 @@ void tctx_task_work(struct callback_head *cb)
 	unsigned int count = 0;
 
 	tctx = container_of(cb, struct io_uring_task, task_work);
-	clear_bit(0, &tctx->tw_pending);
-	smp_mb__after_atomic();
 	tctx_task_work_run(tctx, UINT_MAX, &count);
 }
 
@@ -206,7 +207,7 @@ void io_req_normal_work_add(struct io_kiocb *req)
 	struct io_uring_task *tctx = req->tctx;
 	struct io_ring_ctx *ctx = req->ctx;
 
-	/* task_work already pending, we're done */
+	/* tw run already pending, nothing else to do */
 	if (!mpscq_push(&tctx->task_list, &req->io_task_work.node))
 		return;
 
@@ -223,10 +224,6 @@ void io_req_normal_work_add(struct io_kiocb *req)
 		return;
 	}
 
-	/* task_work must only be added once */
-	if (test_and_set_bit(0, &tctx->tw_pending))
-		return;
-
 	if (likely(!task_work_add(tctx->task, &tctx->task_work, ctx->notify_method)))
 		return;
 

-- 
Jens Axboe