[PATCH 4/4] io_uring: flush task work before waiting for ring exit

[Jens Axboe <axboe@kernel.dk>]

We could have triggered task work when we removed poll completions.
Be sure to check if we're called off the exit path, as the exit task
work could have already been queued. If it is, then the poll flush
will have queued to the io-wq helper thread, so there's no need to run
the work.

Ensure we do so before flushing the CQ ring overflow, in case the poll
flush puts us into overflow mode.

Cc: Oleg Nesterov <oleg@redhat.com>
Cc: Peter Zijlstra <peterz@infradead.org>
Signed-off-by: Jens Axboe <axboe@kernel.dk>
---
 fs/io_uring.c | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/fs/io_uring.c b/fs/io_uring.c
index 7fb51c383e51..4e760b7cd772 100644
--- a/fs/io_uring.c
+++ b/fs/io_uring.c
@@ -7293,10 +7293,15 @@ static void io_ring_ctx_wait_and_kill(struct io_ring_ctx *ctx)
 		io_wq_cancel_all(ctx->io_wq);
 
 	io_iopoll_reap_events(ctx);
+	idr_for_each(&ctx->personality_idr, io_remove_personalities, ctx);
+
+	if (current->task_works != &task_work_exited)
+		task_work_run();
+
 	/* if we failed setting up the ctx, we might not have any rings */
 	if (ctx->rings)
 		io_cqring_overflow_flush(ctx, true);
-	idr_for_each(&ctx->personality_idr, io_remove_personalities, ctx);
+
 	wait_for_completion(&ctx->completions[0]);
 	io_ring_ctx_free(ctx);
 }
-- 
2.26.0