[PATCH 3/5] io_uring/io-wq: allow put_work return next work

[Pavel Begunkov <asml.silence@xxxxxxxxx>]

Formerly work->func() was returning next work to exeucute. Make put_work
do the same. As put_work() is the last thing happening with work during
issuing, it have all info needed to deduce the next job.

Signed-off-by: Pavel Begunkov <asml.silence@xxxxxxxxx>
---
 fs/io-wq.c    | 17 +++++------------
 fs/io-wq.h    |  2 +-
 fs/io_uring.c | 29 ++++++++++++++++++++++++++---
 3 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/fs/io-wq.c b/fs/io-wq.c
index a830eddaffbe..8bdda5e23dcd 100644
--- a/fs/io-wq.c
+++ b/fs/io-wq.c
@@ -443,7 +443,7 @@ static void io_wq_switch_creds(struct io_worker *worker,
 static void io_worker_handle_work(struct io_worker *worker)
 	__releases(wqe->lock)
 {
-	struct io_wq_work *work, *old_work = NULL, *put_work = NULL;
+	struct io_wq_work *work, *old_work = NULL;
 	struct io_wqe *wqe = worker->wqe;
 	struct io_wq *wq = wqe->wq;
 
@@ -464,8 +464,6 @@ static void io_worker_handle_work(struct io_worker *worker)
 			wqe->flags |= IO_WQE_FLAG_STALLED;
 
 		spin_unlock_irq(&wqe->lock);
-		if (put_work && wq->put_work)
-			wq->put_work(old_work);
 		if (!work)
 			break;
 next:
@@ -497,10 +495,8 @@ static void io_worker_handle_work(struct io_worker *worker)
 		if (test_bit(IO_WQ_BIT_CANCEL, &wq->state))
 			work->flags |= IO_WQ_WORK_CANCEL;
 
-		if (wq->get_work && !(work->flags & IO_WQ_WORK_INTERNAL)) {
-			put_work = work;
+		if (wq->get_work && !(work->flags & IO_WQ_WORK_INTERNAL))
 			wq->get_work(work);
-		}
 
 		old_work = work;
 		work->func(work);
@@ -509,6 +505,9 @@ static void io_worker_handle_work(struct io_worker *worker)
 		worker->cur_work = NULL;
 		spin_unlock_irq(&worker->lock);
 
+		if (wq->put_work && !(work->flags & IO_WQ_WORK_INTERNAL))
+			wq->put_work(&work);
+
 		spin_lock_irq(&wqe->lock);
 
 		if (hash != -1U) {
@@ -517,12 +516,6 @@ static void io_worker_handle_work(struct io_worker *worker)
 		}
 		if (work && work != old_work) {
 			spin_unlock_irq(&wqe->lock);
-
-			if (put_work && wq->put_work) {
-				wq->put_work(put_work);
-				put_work = NULL;
-			}
-
 			/* dependent work not hashed */
 			hash = -1U;
 			goto next;
diff --git a/fs/io-wq.h b/fs/io-wq.h
index 508615af4552..f1d717e9acc1 100644
--- a/fs/io-wq.h
+++ b/fs/io-wq.h
@@ -83,7 +83,7 @@ struct io_wq_work {
 	} while (0)						\
 
 typedef void (get_work_fn)(struct io_wq_work *);
-typedef void (put_work_fn)(struct io_wq_work *);
+typedef void (put_work_fn)(struct io_wq_work **);
 
 struct io_wq_data {
 	struct user_struct *user;
diff --git a/fs/io_uring.c b/fs/io_uring.c
index c49ed2846f85..9b220044b608 100644
--- a/fs/io_uring.c
+++ b/fs/io_uring.c
@@ -5939,11 +5939,34 @@ static int io_sqe_files_update(struct io_ring_ctx *ctx, void __user *arg,
 	return __io_sqe_files_update(ctx, &up, nr_args);
 }
 
-static void io_put_work(struct io_wq_work *work)
+static void io_link_work_cb(struct io_wq_work *work)
 {
-	struct io_kiocb *req = container_of(work, struct io_kiocb, work);
+	struct io_kiocb *link = work->data;
 
-	io_put_req(req);
+	io_queue_linked_timeout(link);
+	io_wq_submit_work(work);
+}
+
+static void io_wq_assign_next(struct io_wq_work **workptr, struct io_kiocb *req)
+{
+	struct io_kiocb *link;
+
+	io_prep_next_work(req, &link);
+	*workptr = &req->work;
+	if (link) {
+		req->work.func = io_link_work_cb;
+		req->work.data = link;
+	}
+}
+
+static void io_put_work(struct io_wq_work **workptr)
+{
+	struct io_kiocb *req = container_of(*workptr, struct io_kiocb, work);
+	struct io_kiocb *nxt = NULL;
+
+	io_put_req_find_next(req, &nxt);
+	if (nxt)
+		io_wq_assign_next(workptr, nxt);
 }
 
 static void io_get_work(struct io_wq_work *work)
-- 
2.24.0