io_uring: set TASK_RUNNING right after schedule

Instead of constantly watching that the state of the task is running
before executing tw or taking locks in io_cqring_wait(), switch it back
to TASK_RUNNING immediately.

Signed-off-by: Pavel Begunkov <asml.silence@gmail.com>
Link: https://lore.kernel.org/r/246dddee247d89fd52023f785ed17cc34962a008.1672916894.git.asml.silence@gmail.com
Signed-off-by: Jens Axboe <axboe@kernel.dk>

diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index 4cb9cce23c9084a366b139636363c6aa884acf9b..2ec011f0ba7d11c1c89ecb1860932895f0651322 100644 (file)
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -2544,6 +2544,7 @@ static int io_cqring_wait(struct io_ring_ctx *ctx, int min_events,
                ret = io_cqring_wait_schedule(ctx, &iowq, &timeout);
                if (ret < 0)
                        break;
+               __set_current_state(TASK_RUNNING);
                /*
                 * Run task_work after scheduling and before io_should_wake().
                 * If we got woken because of task_work being processed, run it
@@ -2556,10 +2557,8 @@ static int io_cqring_wait(struct io_ring_ctx *ctx, int min_events,
                check_cq = READ_ONCE(ctx->check_cq);
                if (unlikely(check_cq)) {
                        /* let the caller flush overflows, retry */
-                       if (check_cq & BIT(IO_CHECK_CQ_OVERFLOW_BIT)) {
-                               finish_wait(&ctx->cq_wait, &iowq.wq);
+                       if (check_cq & BIT(IO_CHECK_CQ_OVERFLOW_BIT))
                                io_cqring_do_overflow_flush(ctx);
-                       }
                        if (check_cq & BIT(IO_CHECK_CQ_DROPPED_BIT)) {
                                ret = -EBADR;
                                break;