From 40cbfe3db2c0954fb3c0eb5f3832f08f049aefd7 Mon Sep 17 00:00:00 2001
From: Jens Axboe <axboe@kernel.dk>
Date: Fri, 21 Oct 2022 13:55:37 -0700
Subject: [PATCH] src/queue: don't wait twice if looping in _io_uring_get_cqe()

If the caller asks for > 1 events with a timeout and we get just 1 event,
then we end up looping around and issuing an IORING_ENTER_GETEVENTS
with the timeout set twice. This results in twice the timeout that the
caller asked for.

Add a ->has_ts argument to struct get_data and don't re-enter if we've
already looped and it is set.

Also be a bit saner in not overwriting the return value, if set, and
return the initial value from io_uring_enter().

Signed-off-by: Jens Axboe <axboe@kernel.dk>
---
 src/queue.c | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/queue.c b/src/queue.c
index d29d422b..c06bcc35 100644
--- a/src/queue.c
+++ b/src/queue.c
@@ -55,6 +55,7 @@ struct get_data {
 	unsigned wait_nr;
 	unsigned get_flags;
 	int sz;
+	int has_ts;
 	void *arg;
 };
 
@@ -64,7 +65,7 @@ static int _io_uring_get_cqe(struct io_uring *ring,
 {
 	struct io_uring_cqe *cqe = NULL;
 	bool looped = false;
-	int err;
+	int err = 0;
 
 	do {
 		bool need_enter = false;
@@ -72,9 +73,12 @@ static int _io_uring_get_cqe(struct io_uring *ring,
 		unsigned nr_available;
 		int ret;
 
-		err = __io_uring_peek_cqe(ring, &cqe, &nr_available);
-		if (err)
+		ret = __io_uring_peek_cqe(ring, &cqe, &nr_available);
+		if (ret) {
+			if (!err)
+				err = ret;
 			break;
+		}
 		if (!cqe && !data->wait_nr && !data->submit) {
 			/*
 			 * If we already looped once, we already entererd
@@ -82,7 +86,8 @@ static int _io_uring_get_cqe(struct io_uring *ring,
 			 * wait for, don't keep retrying.
 			 */
 			if (looped || !cq_ring_needs_enter(ring)) {
-				err = -EAGAIN;
+				if (!err)
+					err = -EAGAIN;
 				break;
 			}
 			need_enter = true;
@@ -95,6 +100,13 @@ static int _io_uring_get_cqe(struct io_uring *ring,
 			need_enter = true;
 		if (!need_enter)
 			break;
+		if (looped && data->has_ts) {
+			struct io_uring_getevents_arg *arg = data->arg;
+
+			if (!cqe && arg->ts && !err)
+				err = -ETIME;
+			break;
+		}
 
 		if (ring->int_flags & INT_FLAG_REG_RING)
 			flags |= IORING_ENTER_REGISTERED_RING;
@@ -102,14 +114,18 @@ static int _io_uring_get_cqe(struct io_uring *ring,
 					    data->wait_nr, flags, data->arg,
 					    data->sz);
 		if (ret < 0) {
-			err = ret;
+			if (!err)
+				err = ret;
 			break;
 		}
 
 		data->submit -= ret;
 		if (cqe)
 			break;
-		looped = true;
+		if (!looped) {
+			looped = true;
+			err = ret;
+		}
 	} while (1);
 
 	*cqe_ptr = cqe;
@@ -233,6 +249,7 @@ static int io_uring_wait_cqes_new(struct io_uring *ring,
 		.wait_nr	= wait_nr,
 		.get_flags	= IORING_ENTER_EXT_ARG,
 		.sz		= sizeof(arg),
+		.has_ts		= ts != NULL,
 		.arg		= &arg
 	};
 
@@ -317,6 +334,7 @@ int io_uring_submit_and_wait_timeout(struct io_uring *ring,
 				.wait_nr	= wait_nr,
 				.get_flags	= IORING_ENTER_EXT_ARG,
 				.sz		= sizeof(arg),
+				.has_ts		= ts != NULL,
 				.arg		= &arg
 			};