io_uring: introduce callback driven main loop

The io_uring_enter() has a fixed order of execution: it submits
requests, waits for completions, and returns to the user. Allow to
optionally replace it with a custom loop driven by a callback called
loop_step. The basic requirements to the callback is that it should be
able to submit requests, wait for completions, parse them and repeat.
Most of the communication including parameter passing can be implemented
via shared memory.

The callback should return IOU_LOOP_CONTINUE to continue execution or
IOU_LOOP_STOP to return to the user space. Note that the kernel may
decide to prematurely terminate it as well, e.g. in case the process was
signalled or killed.

The hook takes a structure with parameters. It can be used to ask the
kernel to wait for CQEs by setting cq_wait_idx to the CQE index it wants
to wait for. Spurious wake ups are possible and even likely, the callback
is expected to handle it. There will be more parameters in the future
like timeout.

It can be used with kernel callbacks, for example, as a slow path
deprecation mechanism overwiting SQEs and emulating the wanted
behaviour, however it's more useful together with BPF programs
implemented in following patches.

Note that keeping it separately from the normal io_uring wait loop
makes things much simpler and cleaner. It keeps it in one place instead
of spreading a bunch of checks in different places including disabling
the submission path. It holds the lock by default, which is a better fit
for BPF synchronisation and the loop execution model. It nicely avoids
existing quirks like forced wake ups on timeout request completion. And
it should be easier to implement new features.

Signed-off-by: Pavel Begunkov <asml.silence@gmail.com>
Link: https://patch.msgid.link/a2d369aa1c9dd23ad7edac9220cffc563abcaed6.1772109579.git.asml.silence@gmail.com
Signed-off-by: Jens Axboe <axboe@kernel.dk>

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 4dbd7083dd540e0a92c6cb33e5df4fce37598993..344b634b8989886c67405c26fc15df61dbcba086 100644 (file)
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -41,6 +41,8 @@ enum io_uring_cmd_flags {
        IO_URING_F_COMPAT               = (1 << 12),
 };
 
+struct iou_loop_params;
+
 struct io_wq_work_node {
        struct io_wq_work_node *next;
 };
@@ -361,6 +363,9 @@ struct io_ring_ctx {
                struct io_alloc_cache   rw_cache;
                struct io_alloc_cache   cmd_cache;
 
+               int (*loop_step)(struct io_ring_ctx *ctx,
+                                struct iou_loop_params *);
+
                /*
                 * Any cancelable uring_cmd is added to this list in
                 * ->uring_cmd() by io_uring_cmd_insert_cancelable()

diff --git a/io_uring/Makefile b/io_uring/Makefile
index 931f9156132a0dea83ccc40968821a69e319418d..1c1f47de32a4cbb666346631bd21a819f7229eea 100644 (file)
--- a/io_uring/Makefile
+++ b/io_uring/Makefile
@@ -14,7 +14,7 @@ obj-$(CONFIG_IO_URING)                += io_uring.o opdef.o kbuf.o rsrc.o notif.o \
                                        advise.o openclose.o statx.o timeout.o \
                                        cancel.o waitid.o register.o \
                                        truncate.o memmap.o alloc_cache.o \
-                                       query.o
+                                       query.o loop.o
 
 obj-$(CONFIG_IO_URING_ZCRX)    += zcrx.o
 obj-$(CONFIG_IO_WQ)            += io-wq.o

diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index 74cd62b44d94baf7a9be86383bdb2e003f870838..960d36c49ffe426c69b17fa3a3ba80206710056f 100644 (file)
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -95,6 +95,7 @@
 #include "eventfd.h"
 #include "wait.h"
 #include "bpf_filter.h"
+#include "loop.h"
 
 #define SQE_COMMON_FLAGS (IOSQE_FIXED_FILE | IOSQE_IO_LINK | \
                          IOSQE_IO_HARDLINK | IOSQE_ASYNC)
@@ -588,6 +589,11 @@ void io_cqring_do_overflow_flush(struct io_ring_ctx *ctx)
        mutex_unlock(&ctx->uring_lock);
 }
 
+void io_cqring_overflow_flush_locked(struct io_ring_ctx *ctx)
+{
+       __io_cqring_overflow_flush(ctx, false);
+}
+
 /* must to be called somewhat shortly after putting a request */
 static inline void io_put_task(struct io_kiocb *req)
 {
@@ -2571,6 +2577,11 @@ SYSCALL_DEFINE6(io_uring_enter, unsigned int, fd, u32, to_submit,
        if (unlikely(smp_load_acquire(&ctx->flags) & IORING_SETUP_R_DISABLED))
                goto out;
 
+       if (io_has_loop_ops(ctx)) {
+               ret = io_run_loop(ctx);
+               goto out;
+       }
+
        /*
         * For SQ polling, the thread will do all submissions and completions.
         * Just return the requested submit count, and wake the thread if

diff --git a/io_uring/loop.c b/io_uring/loop.c
new file mode 100644 (file)
index 0000000..31843cc
--- /dev/null
+++ b/io_uring/loop.c
@@ -0,0 +1,91 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#include "io_uring.h"
+#include "wait.h"
+#include "loop.h"
+
+static inline int io_loop_nr_cqes(const struct io_ring_ctx *ctx,
+                                 const struct iou_loop_params *lp)
+{
+       return lp->cq_wait_idx - READ_ONCE(ctx->rings->cq.tail);
+}
+
+static inline void io_loop_wait_start(struct io_ring_ctx *ctx, unsigned nr_wait)
+{
+       atomic_set(&ctx->cq_wait_nr, nr_wait);
+       set_current_state(TASK_INTERRUPTIBLE);
+}
+
+static inline void io_loop_wait_finish(struct io_ring_ctx *ctx)
+{
+       __set_current_state(TASK_RUNNING);
+       atomic_set(&ctx->cq_wait_nr, IO_CQ_WAKE_INIT);
+}
+
+static void io_loop_wait(struct io_ring_ctx *ctx, struct iou_loop_params *lp,
+                        unsigned nr_wait)
+{
+       io_loop_wait_start(ctx, nr_wait);
+
+       if (unlikely(io_local_work_pending(ctx) ||
+                    io_loop_nr_cqes(ctx, lp) <= 0) ||
+                    READ_ONCE(ctx->check_cq)) {
+               io_loop_wait_finish(ctx);
+               return;
+       }
+
+       mutex_unlock(&ctx->uring_lock);
+       schedule();
+       io_loop_wait_finish(ctx);
+       mutex_lock(&ctx->uring_lock);
+}
+
+static int __io_run_loop(struct io_ring_ctx *ctx)
+{
+       struct iou_loop_params lp = {};
+
+       while (true) {
+               int nr_wait, step_res;
+
+               if (unlikely(!ctx->loop_step))
+                       return -EFAULT;
+
+               step_res = ctx->loop_step(ctx, &lp);
+               if (step_res == IOU_LOOP_STOP)
+                       break;
+               if (step_res != IOU_LOOP_CONTINUE)
+                       return -EINVAL;
+
+               nr_wait = io_loop_nr_cqes(ctx, &lp);
+               if (nr_wait > 0)
+                       io_loop_wait(ctx, &lp, nr_wait);
+               else
+                       nr_wait = 0;
+
+               if (task_work_pending(current)) {
+                       mutex_unlock(&ctx->uring_lock);
+                       io_run_task_work();
+                       mutex_lock(&ctx->uring_lock);
+               }
+               if (unlikely(task_sigpending(current)))
+                       return -EINTR;
+               io_run_local_work_locked(ctx, nr_wait);
+
+               if (READ_ONCE(ctx->check_cq) & BIT(IO_CHECK_CQ_OVERFLOW_BIT))
+                       io_cqring_overflow_flush_locked(ctx);
+       }
+
+       return 0;
+}
+
+int io_run_loop(struct io_ring_ctx *ctx)
+{
+       int ret;
+
+       if (!io_allowed_run_tw(ctx))
+               return -EEXIST;
+
+       mutex_lock(&ctx->uring_lock);
+       ret = __io_run_loop(ctx);
+       mutex_unlock(&ctx->uring_lock);
+       return ret;
+}

diff --git a/io_uring/loop.h b/io_uring/loop.h
new file mode 100644 (file)
index 0000000..d7718b9
--- /dev/null
+++ b/io_uring/loop.h
@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: GPL-2.0
+#ifndef IOU_LOOP_H
+#define IOU_LOOP_H
+
+#include <linux/io_uring_types.h>
+
+struct iou_loop_params {
+       /*
+        * The CQE index to wait for. Only serves as a hint and can still be
+        * woken up earlier.
+        */
+       __u32                   cq_wait_idx;
+};
+
+enum {
+       IOU_LOOP_CONTINUE = 0,
+       IOU_LOOP_STOP,
+};
+
+static inline bool io_has_loop_ops(struct io_ring_ctx *ctx)
+{
+       return data_race(ctx->loop_step);
+}
+
+int io_run_loop(struct io_ring_ctx *ctx);
+
+#endif

diff --git a/io_uring/wait.h b/io_uring/wait.h
index 5e236f74e1aff8baaf2146dd69ae166254ccd5eb..037e512dd80c1d7087380c087ae692cdbbb4f6c8 100644 (file)
--- a/io_uring/wait.h
+++ b/io_uring/wait.h
@@ -25,6 +25,7 @@ int io_cqring_wait(struct io_ring_ctx *ctx, int min_events, u32 flags,
                   struct ext_arg *ext_arg);
 int io_run_task_work_sig(struct io_ring_ctx *ctx);
 void io_cqring_do_overflow_flush(struct io_ring_ctx *ctx);
+void io_cqring_overflow_flush_locked(struct io_ring_ctx *ctx);
 
 static inline unsigned int __io_cqring_events(struct io_ring_ctx *ctx)
 {