[RFC 2/2] io_uring/loop: introduce wait timeouts

The basic waiting functionality should be able to time out of CQE
waiting. Implement timeouts for the loop logic and let BPF to control
it. The BPF API consists of two flags it can set in struct
iou_loop_params, IOU_LOOP_TIMEOUT[_ABS], which instruct the loop to arm
a timeout while waiting using the value passed in struct
iou_loop_params::timeout_ns as either a relative of absolute timeout. It
can also find out if the last waiting timed out by checking for the
IOU_LOOP_WAIT_TIMED_OUT flag.

Signed-off-by: Pavel Begunkov <[email protected]>
---
 io_uring/bpf-ops.c |  6 ++++++
 io_uring/loop.c    | 21 ++++++++++++++++++++-
 io_uring/loop.h    | 14 ++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/io_uring/bpf-ops.c b/io_uring/bpf-ops.c
index 5a50f0675fe5..ee99f675a43c 100644
--- a/io_uring/bpf-ops.c
+++ b/io_uring/bpf-ops.c
@@ -95,6 +95,12 @@ static int bpf_io_btf_struct_access(struct bpf_verifier_log *log,
 	if (t == loop_params_type) {
 		if (off + size <= offsetofend(struct iou_loop_params, cq_wait_idx))
 			return SCALAR_VALUE;
+		if (off >= offsetof(struct iou_loop_params, timeout_ns) &&
+		    off + size <= offsetofend(struct iou_loop_params, timeout_ns))
+			return SCALAR_VALUE;
+		if (off >= offsetof(struct iou_loop_params, flags) &&
+		    off + size <= offsetofend(struct iou_loop_params, flags))
+			return SCALAR_VALUE;
 	}
 
 	return -EACCES;
diff --git a/io_uring/loop.c b/io_uring/loop.c
index affaee440dc3..de0d326a3f81 100644
--- a/io_uring/loop.c
+++ b/io_uring/loop.c
@@ -5,6 +5,7 @@
 
 struct io_loop_state {
 	struct iou_loop_params lp;
+	__u32 cur_flags;
 };
 
 static inline int io_loop_nr_cqes(const struct io_ring_ctx *ctx,
@@ -40,7 +41,20 @@ static void io_loop_wait(struct io_ring_ctx *ctx, struct io_loop_state *ls,
 	}
 
 	mutex_unlock(&ctx->uring_lock);
-	schedule();
+
+	if (ls->cur_flags & IOU_LOOP_TIMEOUT) {
+		ktime_t timeout = ns_to_ktime(lp->timeout_ns);
+		enum hrtimer_mode mode = HRTIMER_MODE_REL;
+
+		if (ls->cur_flags & IOU_LOOP_TIMEOUT_ABS)
+			mode = HRTIMER_MODE_ABS;
+
+		if (!schedule_hrtimeout_range_clock(&timeout, 0, mode,ctx->clockid))
+			ls->lp.flags |= IOU_LOOP_WAIT_TIMED_OUT;
+	} else {
+		schedule();
+	}
+
 	io_loop_wait_finish(ctx);
 	mutex_lock(&ctx->uring_lock);
 }
@@ -60,6 +74,11 @@ static int __io_run_loop(struct io_ring_ctx *ctx)
 			break;
 		if (step_res != IOU_LOOP_CONTINUE)
 			return -EINVAL;
+		if (ls.lp.flags & ~IOU_LOOP_VALID_FLAGS)
+			return -EINVAL;
+
+		ls.cur_flags = ls.lp.flags;
+		ls.lp.flags = 0;
 
 		nr_wait = io_loop_nr_cqes(ctx, &ls.lp);
 		if (nr_wait > 0)
diff --git a/io_uring/loop.h b/io_uring/loop.h
index 4dd4fb3aefef..555210721a84 100644
--- a/io_uring/loop.h
+++ b/io_uring/loop.h
@@ -4,12 +4,26 @@
 
 #include <linux/io_uring_types.h>
 
+enum iou_loop_flags {
+	IOU_LOOP_TIMEOUT			= 1U << 0,
+	IOU_LOOP_TIMEOUT_ABS			= 1U << 1,
+	/* If set, the last loop waiting was interrupted by a timeout */
+	IOU_LOOP_WAIT_TIMED_OUT			= 1U << 2,
+};
+
+#define IOU_LOOP_VALID_FLAGS	(IOU_LOOP_TIMEOUT |\
+				 IOU_LOOP_TIMEOUT_ABS |\
+				 IOU_LOOP_WAIT_TIMED_OUT)
+
 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;
+	/* see IOU_LOOP_* flags */
+	__u32			flags;
+	__u64			timeout_ns;
 };
 
 enum {
-- 
2.54.0