[PATCH 4/8] io_uring: move SQPOLL post-wakeup ring need wakeup flag into wake handler

[Jens Axboe <axboe@xxxxxxxxx>]

We need to decouple the clearing on wakeup from the the inline schedule,
as that is going to be required for handling multiple rings in one
thread.

Wrap our wakeup handler so we can clear it when we get the wakeup, by
definition that is when we no longer need the flag set.

Signed-off-by: Jens Axboe <axboe@xxxxxxxxx>
---
 fs/io_uring.c | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/fs/io_uring.c b/fs/io_uring.c
index 76f02db37ffc..95c81e0395d9 100644
--- a/fs/io_uring.c
+++ b/fs/io_uring.c
@@ -6613,6 +6613,18 @@ static inline void io_ring_clear_wakeup_flag(struct io_ring_ctx *ctx)
 	spin_unlock_irq(&ctx->completion_lock);
 }
 
+static int io_sq_wake_function(struct wait_queue_entry *wqe, unsigned mode,
+			       int sync, void *key)
+{
+	struct io_ring_ctx *ctx = container_of(wqe, struct io_ring_ctx, sqo_wait_entry);
+	int ret;
+
+	ret = autoremove_wake_function(wqe, mode, sync, key);
+	if (ret)
+		io_ring_clear_wakeup_flag(ctx);
+	return ret;
+}
+
 static int io_sq_thread(void *data)
 {
 	struct io_ring_ctx *ctx = data;
@@ -6621,6 +6633,7 @@ static int io_sq_thread(void *data)
 	int ret = 0;
 
 	init_wait(&ctx->sqo_wait_entry);
+	ctx->sqo_wait_entry.func = io_sq_wake_function;
 
 	complete(&ctx->sq_thread_comp);
 
@@ -6707,7 +6720,6 @@ static int io_sq_thread(void *data)
 				schedule();
 				finish_wait(ctx->sqo_wait, &ctx->sqo_wait_entry);
 
-				io_ring_clear_wakeup_flag(ctx);
 				ret = 0;
 				continue;
 			}
-- 
2.28.0