[PATCH 18/19] io-wq: only do io_uringlet_end() at the first schedule time

[Hao Xu <hao.xu@linux.dev>]

From: Hao Xu <howeyxu@tencent.com>

A request may block multiple times during its life cycle. We should
only do io_uringlet_end() at the first time since this function may
modify ctx->submit_state info and for the non-first time, the task
already lost the control of submitting sqes. Allowing it to do so will
damage the submission state.

Signed-off-by: Hao Xu <howeyxu@tencent.com>
---
 io_uring/io-wq.c    | 14 +++++++++++---
 io_uring/io_uring.c |  2 +-
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/io_uring/io-wq.c b/io_uring/io-wq.c
index 9fcaeea7a478..f845b7daced8 100644
--- a/io_uring/io-wq.c
+++ b/io_uring/io-wq.c
@@ -400,16 +400,24 @@ static void io_wqe_dec_running(struct io_worker *worker)
 
 	if (io_wq_is_uringlet(wq)) {
 		bool activated;
+		bool first_block;
 
 		if (!io_worker_test_submit(worker))
 			return;
 
-		io_uringlet_end(wq->private);
-		io_worker_set_scheduled(worker);
 		raw_spin_lock(&wq->lock);
-		wq->owner = IO_WQ_OWNER_TRANSMIT;
+		first_block = (wq->owner == worker ? true : false);
 		raw_spin_unlock(&wq->lock);
 
+		io_worker_set_scheduled(worker);
+
+		if (first_block) {
+			io_uringlet_end(wq->private);
+			raw_spin_lock(&wq->lock);
+			wq->owner = IO_WQ_OWNER_TRANSMIT;
+			raw_spin_unlock(&wq->lock);
+		}
+
 		raw_spin_lock(&wqe->lock);
 		rcu_read_lock();
 		activated = io_wqe_activate_free_worker(wqe, acct);
diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index a48e34f63845..7ebc83b3a33f 100644
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -2171,7 +2171,7 @@ int io_submit_sqes_let(struct io_wq_work *work)
 
 	io_get_task_refs(entries);
 	io_submit_state_start(&ctx->submit_state, entries);
-	ctx->submit_state->need_plug = false;
+	ctx->submit_state.need_plug = false;
 	do {
 		const struct io_uring_sqe *sqe;
 		struct io_kiocb *req;
-- 
2.25.1