[PATCH 2/2] io_uring: reap iopoll events before exiting io wq

[Ming Lei <ming.lei@redhat.com>]

io_uring delays reaping iopoll events into exit work, which is scheduled
in io_uring_release(). But io wq is exited inside do_exit(), and wq code
path may share resource with iopoll code path, such as request. If iopoll
events aren't reaped, io wq exit can't be done, and cause IO hang.

The issue can be triggered when terminating 't/io_uring -n4 /dev/nullb0' with
default null_blk parameters.

Fix it by reaping iopoll events in io_uring_clean_tctx() and moving io wq
exit into workqueue context.

Closes: https://lore.kernel.org/linux-block/3893581.1691785261@warthog.procyon.org.uk/
Reported-by: David Howells <dhowells@redhat.com>
Signed-off-by: Ming Lei <ming.lei@redhat.com>
---
 io_uring/io_uring.c |  2 +-
 io_uring/io_uring.h |  1 +
 io_uring/tctx.c     | 60 +++++++++++++++++++++++++++++++++++++++------
 3 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index c4adb44f1aa4..ab7844c3380c 100644
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -3214,7 +3214,7 @@ static __cold bool io_uring_try_cancel_iowq(struct io_ring_ctx *ctx)
 	return ret;
 }
 
-static bool iopoll_reap_events(struct io_ring_ctx *ctx, bool reap_all)
+bool iopoll_reap_events(struct io_ring_ctx *ctx, bool reap_all)
 {
 	bool reapped = false;
 
diff --git a/io_uring/io_uring.h b/io_uring/io_uring.h
index 547c30582fb8..f1556666a064 100644
--- a/io_uring/io_uring.h
+++ b/io_uring/io_uring.h
@@ -63,6 +63,7 @@ void io_req_task_queue_fail(struct io_kiocb *req, int ret);
 void io_req_task_submit(struct io_kiocb *req, struct io_tw_state *ts);
 void tctx_task_work(struct callback_head *cb);
 __cold void io_uring_cancel_generic(bool cancel_all, struct io_sq_data *sqd);
+bool iopoll_reap_events(struct io_ring_ctx *ctx, bool reap_all);
 int io_uring_alloc_task_context(struct task_struct *task,
 				struct io_ring_ctx *ctx);
 
diff --git a/io_uring/tctx.c b/io_uring/tctx.c
index c043fe93a3f2..582b9149bab1 100644
--- a/io_uring/tctx.c
+++ b/io_uring/tctx.c
@@ -6,6 +6,7 @@
 #include <linux/slab.h>
 #include <linux/nospec.h>
 #include <linux/io_uring.h>
+#include <linux/delay.h>
 
 #include <uapi/linux/io_uring.h>
 
@@ -175,24 +176,67 @@ __cold void io_uring_del_tctx_node(unsigned long index)
 	kfree(node);
 }
 
+struct wq_exit_work {
+	struct work_struct work;
+	struct io_wq *wq;
+	bool done;
+};
+
+static void io_uring_wq_exit_work(struct work_struct *work)
+{
+	struct wq_exit_work *wq_work =
+                  container_of(work, struct wq_exit_work, work);
+	struct io_wq *wq = wq_work->wq;
+
+	/*
+	 * Must be after io_uring_del_tctx_node() (removes nodes under
+	 * uring_lock) to avoid race with io_uring_try_cancel_iowq().
+	 */
+	io_wq_put_and_exit(wq);
+	wq_work->done = true;
+}
+
 __cold void io_uring_clean_tctx(struct io_uring_task *tctx)
 {
 	struct io_wq *wq = tctx->io_wq;
+	struct wq_exit_work work = {
+		.wq = wq,
+		.done = true,
+	};
 	struct io_tctx_node *node;
 	unsigned long index;
 
-	xa_for_each(&tctx->xa, index, node) {
-		io_uring_del_tctx_node(index);
-		cond_resched();
-	}
+	/*
+	 * io_wq may depend on reaping iopoll events because pending
+	 * requests in io_wq may share resource with polled requests,
+	 * meantime new polled IO may be submitted from io_wq after
+	 * getting resource.
+	 *
+	 * So io_wq has to be exited from workqueue context for avoiding
+	 * IO hang.
+	 */
 	if (wq) {
+		work.done = false;
+		INIT_WORK(&work.work, io_uring_wq_exit_work);
+		queue_work(system_unbound_wq, &work.work);
+	}
+
+	while (!work.done) {
+		xa_for_each(&tctx->xa, index, node)
+			iopoll_reap_events(node->ctx, true);
+
 		/*
-		 * Must be after io_uring_del_tctx_node() (removes nodes under
-		 * uring_lock) to avoid race with io_uring_try_cancel_iowq().
+		 * Wait a little while and reap again since new polled
+		 * IO may get resource from io_wq and be submitted.
 		 */
-		io_wq_put_and_exit(wq);
-		tctx->io_wq = NULL;
+		msleep(10);
+	}
+
+	xa_for_each(&tctx->xa, index, node) {
+		io_uring_del_tctx_node(index);
+		cond_resched();
 	}
+	tctx->io_wq = NULL;
 }
 
 void io_uring_unreg_ringfd(void)
-- 
2.40.1