[PATCH v3 liburing 7/7] add accept with overflow test

[Dylan Yudaken <dylany@fb.com>]

add test to exercise the overflow codepath for multishot accept.
this doesn't actually fail previously, but does at least exerceise
the codepath and ensure that some invariants hold wrt flags and
IORING_CQE_F_MORE.

Signed-off-by: Dylan Yudaken <dylany@fb.com>
---
 test/accept.c | 129 +++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 101 insertions(+), 28 deletions(-)

diff --git a/test/accept.c b/test/accept.c
index 77e3ebc..0463173 100644
--- a/test/accept.c
+++ b/test/accept.c
@@ -24,6 +24,9 @@
 #include "liburing.h"
 
 #define MAX_FDS 32
+#define NOP_USER_DATA (1LLU << 50)
+#define INITIAL_USER_DATA 1000
+
 static int no_accept;
 static int no_accept_multi;
 
@@ -39,6 +42,7 @@ struct accept_test_args {
 	bool queue_accept_before_connect;
 	bool multishot;
 	int extra_loops;
+	bool overflow;
 };
 
 static void close_fds(int fds[], int nr)
@@ -86,6 +90,24 @@ static void queue_recv(struct io_uring *ring, int fd, bool fixed)
 		sqe->flags |= IOSQE_FIXED_FILE;
 }
 
+static void queue_accept_multishot(struct io_uring *ring, int fd,
+				   int idx, bool fixed)
+{
+	struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
+	int ret;
+
+	if (fixed)
+		io_uring_prep_multishot_accept_direct(sqe, fd,
+						NULL, NULL,
+						0);
+	else
+		io_uring_prep_multishot_accept(sqe, fd, NULL, NULL, 0);
+
+	io_uring_sqe_set_data64(sqe, idx);
+	ret = io_uring_submit(ring);
+	assert(ret != -1);
+}
+
 static void queue_accept_conn(struct io_uring *ring, int fd,
 			      struct accept_test_args args)
 {
@@ -93,40 +115,51 @@ static void queue_accept_conn(struct io_uring *ring, int fd,
 	int ret;
 	int fixed_idx = args.fixed ? 0 : -1;
 	int count = 1 + args.extra_loops;
-	bool multishot = args.multishot;
+
+	if (args.multishot) {
+		queue_accept_multishot(ring, fd, INITIAL_USER_DATA, args.fixed);
+		return;
+	}
 
 	while (count--) {
 		sqe = io_uring_get_sqe(ring);
 		if (fixed_idx < 0) {
-			if (!multishot)
-				io_uring_prep_accept(sqe, fd, NULL, NULL, 0);
-			else
-				io_uring_prep_multishot_accept(sqe, fd, NULL,
-							       NULL, 0);
+			io_uring_prep_accept(sqe, fd, NULL, NULL, 0);
 		} else {
-			if (!multishot)
-				io_uring_prep_accept_direct(sqe, fd, NULL, NULL,
-							    0, fixed_idx);
-			else
-				io_uring_prep_multishot_accept_direct(sqe, fd,
-								      NULL, NULL,
-								      0);
+			io_uring_prep_accept_direct(sqe, fd, NULL, NULL,
+						    0, fixed_idx);
 		}
-
 		ret = io_uring_submit(ring);
 		assert(ret != -1);
 	}
 }
 
-static int accept_conn(struct io_uring *ring, int fixed_idx, bool multishot)
+static int accept_conn(struct io_uring *ring, int fixed_idx, int *multishot, int fd)
 {
-	struct io_uring_cqe *cqe;
+	struct io_uring_cqe *pcqe;
+	struct io_uring_cqe cqe;
 	int ret;
 
-	ret = io_uring_wait_cqe(ring, &cqe);
-	assert(!ret);
-	ret = cqe->res;
-	io_uring_cqe_seen(ring, cqe);
+	do {
+		ret = io_uring_wait_cqe(ring, &pcqe);
+		assert(!ret);
+		cqe = *pcqe;
+		io_uring_cqe_seen(ring, pcqe);
+	} while (cqe.user_data == NOP_USER_DATA);
+
+	if (*multishot) {
+		if (!(cqe.flags & IORING_CQE_F_MORE)) {
+			(*multishot)++;
+			queue_accept_multishot(ring, fd, *multishot, fixed_idx == 0);
+		} else {
+			if (cqe.user_data != *multishot) {
+				fprintf(stderr, "received multishot after told done!\n");
+				return -ECANCELED;
+			}
+		}
+	}
+
+	ret = cqe.res;
 
 	if (fixed_idx >= 0) {
 		if (ret > 0) {
@@ -203,6 +236,32 @@ static int set_client_fd(struct sockaddr_in *addr)
 	return fd;
 }
 
+static void cause_overflow(struct io_uring *ring)
+{
+	int i, ret;
+
+	for (i = 0; i < *ring->cq.kring_entries; i++) {
+		struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
+
+		io_uring_prep_nop(sqe);
+		io_uring_sqe_set_data64(sqe, NOP_USER_DATA);
+		ret = io_uring_submit(ring);
+		assert(ret != -1);
+	}
+
+}
+
+static void clear_overflow(struct io_uring *ring)
+{
+	struct io_uring_cqe *cqe;
+
+	while (!io_uring_peek_cqe(ring, &cqe)) {
+		if (cqe->user_data != NOP_USER_DATA)
+			break;
+		io_uring_cqe_seen(ring, cqe);
+	}
+}
+
 static int test_loop(struct io_uring *ring,
 		     struct accept_test_args args,
 		     int recv_s0,
@@ -215,15 +274,22 @@ static int test_loop(struct io_uring *ring,
 	bool multishot = args.multishot;
 	uint32_t multishot_mask = 0;
 	int nr_fds = multishot ? MAX_FDS : 1;
+	int multishot_idx = multishot ? INITIAL_USER_DATA : 0;
 
-	for (i = 0; i < nr_fds; i++)
+	if (args.overflow)
+		cause_overflow(ring);
+
+	for (i = 0; i < nr_fds; i++) {
 		c_fd[i] = set_client_fd(addr);
+		if (args.overflow && i == nr_fds / 2)
+			clear_overflow(ring);
+	}
 
 	if (!args.queue_accept_before_connect)
 		queue_accept_conn(ring, recv_s0, args);
 
 	for (i = 0; i < nr_fds; i++) {
-		s_fd[i] = accept_conn(ring, fixed ? 0 : -1, multishot);
+		s_fd[i] = accept_conn(ring, fixed ? 0 : -1, &multishot_idx, recv_s0);
 		if (s_fd[i] == -EINVAL) {
 			if (args.accept_should_error)
 				goto out;
@@ -527,14 +593,15 @@ static int test_accept(int count, bool before)
 	return ret;
 }
 
-static int test_multishot_accept(int count, bool before)
+static int test_multishot_accept(int count, bool before, bool overflow)
 {
 	struct io_uring m_io_uring;
 	int ret;
 	struct accept_test_args args = {
 		.queue_accept_before_connect = before,
 		.multishot = true,
-		.extra_loops = count - 1
+		.extra_loops = count - 1,
+		.overflow = overflow
 	};
 
 	if (no_accept_multi)
@@ -779,15 +846,21 @@ int main(int argc, char *argv[])
 		return ret;
 	}
 
-	ret = test_multishot_accept(1, false);
+	ret = test_multishot_accept(1, true, true);
+	if (ret) {
+		fprintf(stderr, "test_multishot_accept(1, false, true) failed\n");
+		return ret;
+	}
+
+	ret = test_multishot_accept(1, false, false);
 	if (ret) {
-		fprintf(stderr, "test_multishot_accept(1, false) failed\n");
+		fprintf(stderr, "test_multishot_accept(1, false, false) failed\n");
 		return ret;
 	}
 
-	ret = test_multishot_accept(1, true);
+	ret = test_multishot_accept(1, true, false);
 	if (ret) {
-		fprintf(stderr, "test_multishot_accept(1, true) failed\n");
+		fprintf(stderr, "test_multishot_accept(1, true, false) failed\n");
 		return ret;
 	}
 
-- 
2.30.2