Test support for timed waittimeout

Signed-off-by: Jens Axboe <axboe@kernel.dk>

3 files changed, 94 insertions, 9 deletions

diff --git a/src/include/liburing.h b/src/include/liburing.h
index 338d81d..d0c3769 100644
--- a/src/include/liburing.h
+++ b/src/include/liburing.h
@@ -74,6 +74,8 @@ extern int io_uring_peek_cqe(struct io_uring *ring,
 	struct io_uring_cqe **cqe_ptr);
 extern int io_uring_wait_cqe(struct io_uring *ring,
 	struct io_uring_cqe **cqe_ptr);
+extern int io_uring_wait_cqe_timeout(struct io_uring *ring,
+	struct io_uring_cqe **cqe_ptr, struct timespec *ts);
 extern int io_uring_submit(struct io_uring *ring);
 extern int io_uring_submit_and_wait(struct io_uring *ring, unsigned wait_nr);
 extern struct io_uring_sqe *io_uring_get_sqe(struct io_uring *ring);
@@ -88,6 +90,8 @@ extern int io_uring_unregister_files(struct io_uring *ring);
 extern int io_uring_register_eventfd(struct io_uring *ring, int fd);
 extern int io_uring_unregister_eventfd(struct io_uring *ring);
 
+#define LIBURING_TIMEOUT	((__u64) -1)
+
 #define io_uring_for_each_cqe(ring, head, cqe)					\
 	/*									\
 	 * io_uring_smp_load_acquire() enforces the order of tail		\
diff --git a/src/queue.c b/src/queue.c
index 203b534..41b93a5 100644
--- a/src/queue.c
+++ b/src/queue.c
@@ -12,25 +12,35 @@
 #include "liburing/barrier.h"
 
 static int __io_uring_get_cqe(struct io_uring *ring,
-			      struct io_uring_cqe **cqe_ptr, int wait)
+			      struct io_uring_cqe **cqe_ptr, unsigned submit,
+			      int wait)
 {
+	int ret, err = 0;
 	unsigned head;
-	int ret;
 
 	do {
 		io_uring_for_each_cqe(ring, head, *cqe_ptr)
 			break;
-		if (*cqe_ptr)
+		if (*cqe_ptr) {
+			if ((*cqe_ptr)->user_data == LIBURING_TIMEOUT) {
+				if ((*cqe_ptr)->res < 0)
+					err = (*cqe_ptr)->res;
+				io_uring_cq_advance(ring, 1);
+				if (!err)
+					continue;
+				*cqe_ptr = NULL;
+			}
 			break;
+		}
 		if (!wait)
 			return -EAGAIN;
-		ret = io_uring_enter(ring->ring_fd, 0, 1,
-					IORING_ENTER_GETEVENTS, NULL);
+		ret = io_uring_enter(ring->ring_fd, submit, 1,
+				IORING_ENTER_GETEVENTS, NULL);
 		if (ret < 0)
 			return -errno;
 	} while (1);
 
-	return 0;
+	return err;
 }
 
 /*
@@ -39,7 +49,7 @@ static int __io_uring_get_cqe(struct io_uring *ring,
  */
 int io_uring_peek_cqe(struct io_uring *ring, struct io_uring_cqe **cqe_ptr)
 {
-	return __io_uring_get_cqe(ring, cqe_ptr, 0);
+	return __io_uring_get_cqe(ring, cqe_ptr, 0, 0);
 }
 
 /*
@@ -48,7 +58,21 @@ int io_uring_peek_cqe(struct io_uring *ring, struct io_uring_cqe **cqe_ptr)
  */
 int io_uring_wait_cqe(struct io_uring *ring, struct io_uring_cqe **cqe_ptr)
 {
-	return __io_uring_get_cqe(ring, cqe_ptr, 1);
+	return __io_uring_get_cqe(ring, cqe_ptr, 0, 1);
+}
+
+int io_uring_wait_cqe_timeout(struct io_uring *ring,
+		struct io_uring_cqe **cqe_ptr,
+		struct timespec *ts)
+{
+	struct io_uring_sqe *sqe;
+
+	sqe = io_uring_get_sqe(ring);
+	io_uring_prep_timeout(sqe, ts, 1);
+	sqe->user_data = LIBURING_TIMEOUT;
+	io_uring_submit(ring);
+
+	return __io_uring_get_cqe(ring, cqe_ptr, 1, 1);
 }
 
 /*
@@ -115,7 +139,7 @@ static int __io_uring_submit(struct io_uring *ring, unsigned wait_nr)
 		}
 
 		ret = io_uring_enter(ring->ring_fd, submitted, wait_nr, flags,
-					NULL);
+				NULL);
 		if (ret < 0)
 			return -errno;
 	} else
diff --git a/test/timeout.c b/test/timeout.c
index 641f4ac..00581e0 100644
--- a/test/timeout.c
+++ b/test/timeout.c
@@ -177,6 +177,57 @@ err:
 	return 1;
 }
 
+static int test_single_timeout_wait(struct io_uring *ring)
+{
+	struct io_uring_cqe *cqe;
+	struct io_uring_sqe *sqe;
+	struct timespec ts;
+	int i, ret;
+
+	sqe = io_uring_get_sqe(ring);
+	io_uring_prep_nop(sqe);
+	io_uring_sqe_set_data(sqe, (void *) 1);
+
+	sqe = io_uring_get_sqe(ring);
+	io_uring_prep_nop(sqe);
+	io_uring_sqe_set_data(sqe, (void *) 1);
+
+	ret = io_uring_submit(ring);
+	if (ret <= 0) {
+		printf("sqe submit failed: %d\n", ret);
+		goto err;
+	}
+
+	ts.tv_sec = 1;
+	ts.tv_nsec = 0;
+
+	i = 0;
+	do {
+		ret = io_uring_wait_cqe_timeout(ring, &cqe, &ts);
+		if (ret == -ETIME)
+			break;
+		if (ret < 0) {
+			printf("wait timeout failed: %d\n", ret);
+			goto err;
+		}
+
+		if (cqe->res < 0) {
+			printf("res: %d\n", cqe->res);
+			goto err;
+		}
+		io_uring_cqe_seen(ring, cqe);
+		i++;
+	} while (1);
+
+	if (i != 2) {
+		printf("got %d completions\n", i);
+		goto err;
+	}
+	return 0;
+err:
+	return 1;
+}
+
 /*
  * Test single timeout waking us up
  */
@@ -293,6 +344,12 @@ int main(int argc, char *argv[])
 		return ret;
 	}
 
+	ret = test_single_timeout_wait(&ring);
+	if (ret) {
+		printf("test_single_timeout_wait failed\n");
+		return ret;
+	}
+
 	/*
 	 * this test must go last, it kills the ring
 	 */