[PATCH] io_uring/tw: avoid local task work starvation

[pgray <hi@pgray.dev>]

IO_LOCAL_TW_DEFAULT_MAX (20) is too small for massive backlogs

this patch fixes by:
* not counting passive completions against budget
* add a ceiling of 512: IO_LOCAL_TW_COMPL_MAX
* clearing IORING_SQ_TASKRUN when local queue empty

Link: https://github.com/axboe/liburing/issues/1645
Fixes: f46b9cdb22f7 ("io_uring: limit local tw done")

Signed-off-by: pgray <hi@pgray.dev>
---
 io_uring/tw.c | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/io_uring/tw.c b/io_uring/tw.c
index f573bcc3af6a..236ed01fcfaa 100644
--- a/io_uring/tw.c
+++ b/io_uring/tw.c
@@ -16,6 +16,8 @@
 #include "wait.h"
 #include "mpscq.h"
 
+#define IO_LOCAL_TW_COMPL_MAX		512
+
 static void ctx_flush_and_put(struct io_ring_ctx *ctx, io_tw_token_t tw)
 {
 	if (!ctx)
@@ -287,17 +289,23 @@ static int __io_run_local_work_loop(struct io_ring_ctx *ctx,
 				    int events)
 {
 	int ret = 0;
+	int nr_events = 0;
 
-	while (ret < events) {
+	while (nr_events < events && ret < IO_LOCAL_TW_COMPL_MAX) {
 		struct llist_node *node = mpscq_pop(&ctx->work_list, &ctx->work_head);
 		struct io_kiocb *req;
 
 		if (!node)
 			break;
 		req = container_of(node, struct io_kiocb, io_task_work.node);
-		INDIRECT_CALL_2(req->io_task_work.func,
-				io_poll_task_func, io_req_rw_complete,
-				(struct io_tw_req){req}, tw);
+		if (req->io_task_work.func == io_req_task_complete) {
+			io_req_task_complete((struct io_tw_req){req}, tw);
+		} else {
+			INDIRECT_CALL_2(req->io_task_work.func,
+					io_poll_task_func, io_req_rw_complete,
+					(struct io_tw_req){req}, tw);
+			nr_events++;
+		}
 		ret++;
 	}
 
@@ -342,8 +350,11 @@ int io_run_local_work_locked(struct io_ring_ctx *ctx, int min_events)
 {
 	struct io_tw_state ts = {};
 
-	if (!io_local_work_pending(ctx))
+	if (!io_local_work_pending(ctx)) {
+		if (ctx->flags & IORING_SETUP_TASKRUN_FLAG)
+			atomic_andnot(IORING_SQ_TASKRUN, &ctx->rings->sq_flags);
 		return 0;
+	}
 	return __io_run_local_work(ctx, ts, min_events,
 					max(IO_LOCAL_TW_DEFAULT_MAX, min_events));
 }
-- 
2.55.0