Add support for IORING_REGISTER_PBUF_STATUSpbuf-status

Basic support for querying the provided ring buffer head, and a helper
that returns the number of available entries. Includes a basic test
case as well.

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

5 files changed, 222 insertions, 0 deletions

diff --git a/src/include/liburing.h b/src/include/liburing.h
index e66c75c6..b3f82302 100644
--- a/src/include/liburing.h
+++ b/src/include/liburing.h
@@ -239,6 +239,7 @@ int io_uring_close_ring_fd(struct io_uring *ring);
 int io_uring_register_buf_ring(struct io_uring *ring,
 			       struct io_uring_buf_reg *reg, unsigned int flags);
 int io_uring_unregister_buf_ring(struct io_uring *ring, int bgid);
+int io_uring_buf_ring_head(struct io_uring *ring, int buf_group, unsigned *head);
 int io_uring_register_sync_cancel(struct io_uring *ring,
 				 struct io_uring_sync_cancel_reg *reg);
 
@@ -1488,6 +1489,20 @@ IOURINGINLINE void io_uring_buf_ring_cq_advance(struct io_uring *ring,
 	__io_uring_buf_ring_cq_advance(ring, br, count, count);
 }
 
+IOURINGINLINE int io_uring_buf_ring_available(struct io_uring *ring,
+					      struct io_uring_buf_ring *br,
+					      unsigned short bgid)
+{
+	unsigned head;
+	int ret;
+
+	ret = io_uring_buf_ring_head(ring, bgid, &head);
+	if (ret)
+		return ret;
+
+	return br->tail - head;
+}
+
 #ifndef LIBURING_INTERNAL
 IOURINGINLINE struct io_uring_sqe *io_uring_get_sqe(struct io_uring *ring)
 {
diff --git a/src/include/liburing/io_uring.h b/src/include/liburing/io_uring.h
index dfeb5fa3..2a473d1b 100644
--- a/src/include/liburing/io_uring.h
+++ b/src/include/liburing/io_uring.h
@@ -562,6 +562,9 @@ enum {
 	/* register a range of fixed file slots for automatic slot allocation */
 	IORING_REGISTER_FILE_ALLOC_RANGE	= 25,
 
+	/* return status information for a buffer group */
+	IORING_REGISTER_PBUF_STATUS		= 26,
+
 	/* this goes last */
 	IORING_REGISTER_LAST,
 
@@ -688,6 +691,13 @@ struct io_uring_buf_reg {
 	__u64	resv[3];
 };
 
+/* argument for IORING_REGISTER_PBUF_STATUS */
+struct io_uring_buf_status {
+	__u32	buf_group;	/* input */
+	__u32	head;		/* output */
+	__u32	resv[8];
+};
+
 /*
  * io_uring_restriction->opcode values
  */
diff --git a/src/register.c b/src/register.c
index bf010b0d..782736d9 100644
--- a/src/register.c
+++ b/src/register.c
@@ -326,6 +326,20 @@ int io_uring_unregister_buf_ring(struct io_uring *ring, int bgid)
 	return do_register(ring, IORING_UNREGISTER_PBUF_RING, &reg, 1);
 }
 
+int io_uring_buf_ring_head(struct io_uring *ring, int buf_group, unsigned *head)
+{
+	struct io_uring_buf_status buf_status = {
+		.buf_group	= buf_group,
+	};
+	int ret;
+
+	ret = do_register(ring, IORING_REGISTER_PBUF_STATUS, &buf_status, 1);
+	if (ret)
+		return ret;
+	*head = buf_status.head;
+	return 0;
+}
+
 int io_uring_register_sync_cancel(struct io_uring *ring,
 				  struct io_uring_sync_cancel_reg *reg)
 {
diff --git a/test/Makefile b/test/Makefile
index 25abd542..9ccc36ee 100644
--- a/test/Makefile
+++ b/test/Makefile
@@ -155,6 +155,7 @@ test_srcs := \
 	register-restrictions.c \
 	rename.c \
 	ringbuf-read.c \
+	ringbuf-status.c \
 	ring-leak2.c \
 	ring-leak.c \
 	rsrc_tags.c \
diff --git a/test/ringbuf-status.c b/test/ringbuf-status.c
new file mode 100644
index 00000000..b36216e5
--- /dev/null
+++ b/test/ringbuf-status.c
@@ -0,0 +1,182 @@
+/* SPDX-License-Identifier: MIT */
+/*
+ * Description: test reading provided ring buf head
+ *
+ */
+#include <errno.h>
+#include <stdio.h>
+#include <unistd.h>
+#include <stdlib.h>
+#include <string.h>
+#include <fcntl.h>
+
+#include "liburing.h"
+#include "helpers.h"
+
+#define BUF_SIZE	32
+#define NR_BUFS		8
+#define FSIZE		(BUF_SIZE * NR_BUFS)
+
+#define BR_MASK		(NR_BUFS - 1)
+#define BGID		1
+
+static int no_buf_ring;
+static int no_buf_ring_status;
+
+static int test(int invalid)
+{
+	struct io_uring_sqe *sqe;
+	struct io_uring_cqe *cqe;
+	struct io_uring ring;
+	struct io_uring_buf_ring *br;
+	int ret, i, fds[2];
+	unsigned head;
+	char *buf;
+	void *ptr;
+	char output[16];
+
+	memset(output, 0x55, sizeof(output));
+
+	ret = io_uring_queue_init(NR_BUFS, &ring, 0);
+	if (ret) {
+		fprintf(stderr, "ring setup failed: %d\n", ret);
+		return 1;
+	}
+
+	if (pipe(fds) < 0) {
+		perror("pipe");
+		return T_EXIT_FAIL;
+	}
+
+	if (posix_memalign((void **) &buf, 4096, FSIZE))
+		return 1;
+
+	br = io_uring_setup_buf_ring(&ring, NR_BUFS, BGID, 0, &ret);
+	if (!br) {
+		if (ret == -EINVAL) {
+			no_buf_ring = 1;
+			return 0;
+		}
+		fprintf(stderr, "Buffer ring register failed %d\n", ret);
+		return 1;
+	}
+
+	ptr = buf;
+	for (i = 0; i < NR_BUFS; i++) {
+		io_uring_buf_ring_add(br, ptr, BUF_SIZE, i + 1, BR_MASK, i);
+		ptr += BUF_SIZE;
+	}
+	io_uring_buf_ring_advance(br, NR_BUFS);
+
+	/* head should be zero at this point */
+	head = 1;
+	if (!invalid)
+		ret = io_uring_buf_ring_head(&ring, BGID, &head);
+	else
+		ret = io_uring_buf_ring_head(&ring, BGID + 10, &head);
+	if (ret) {
+		if (ret == -EINVAL) {
+			no_buf_ring_status = 1;
+			return T_EXIT_SKIP;
+		}
+		if (invalid && ret == -ENOENT)
+			return T_EXIT_PASS;
+		fprintf(stderr, "buf_ring_head: %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+	if (invalid) {
+		fprintf(stderr, "lookup of bad group id succeeded\n");
+		return T_EXIT_FAIL;
+	}
+	if (head != 0) {
+		fprintf(stderr, "bad head %d\n", head);
+		return T_EXIT_FAIL;
+	}
+
+	ret = io_uring_buf_ring_available(&ring, br, BGID);
+	if (ret != NR_BUFS) {
+		fprintf(stderr, "ring available %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	sqe = io_uring_get_sqe(&ring);
+	io_uring_prep_read(sqe, fds[0], NULL, BUF_SIZE, i * BUF_SIZE);
+	sqe->buf_group = BGID;
+	sqe->flags |= IOSQE_BUFFER_SELECT;
+	sqe->user_data = 1;
+
+	ret = io_uring_submit(&ring);
+	if (ret != 1) {
+		fprintf(stderr, "submit: %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	/* head should still be zero at this point, no buffers consumed */
+	head = 1;
+	ret = io_uring_buf_ring_head(&ring, BGID, &head);
+	if (head != 0) {
+		fprintf(stderr, "bad head after submit %d\n", head);
+		return T_EXIT_FAIL;
+	}
+
+	ret = write(fds[1], output, sizeof(output));
+	if (ret != sizeof(output)) {
+		fprintf(stderr, "pipe buffer write %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	ret = io_uring_wait_cqe(&ring, &cqe);
+	if (ret) {
+		fprintf(stderr, "wait cqe failed %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+	if (cqe->res != sizeof(output)) {
+		fprintf(stderr, "cqe res %d\n", cqe->res);
+		return T_EXIT_FAIL;
+	}
+	if (!(cqe->flags & IORING_CQE_F_BUFFER)) {
+		fprintf(stderr, "no buffer selected\n");
+		return T_EXIT_FAIL;
+	}
+	io_uring_cqe_seen(&ring, cqe);
+
+	/* head should now be one, we consumed a buffer */
+	ret = io_uring_buf_ring_head(&ring, BGID, &head);
+	if (head != 1) {
+		fprintf(stderr, "bad head after cqe %d\n", head);
+		return T_EXIT_FAIL;
+	}
+
+	ret = io_uring_buf_ring_available(&ring, br, BGID);
+	if (ret != NR_BUFS - 1) {
+		fprintf(stderr, "ring available %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	close(fds[0]);
+	close(fds[1]);
+	free(buf);
+	io_uring_queue_exit(&ring);
+	return T_EXIT_PASS;
+}
+
+int main(int argc, char *argv[])
+{
+	int ret;
+
+	ret = test(0);
+	if (ret == T_EXIT_FAIL) {
+		fprintf(stderr, "test 0 failed\n");
+		return T_EXIT_FAIL;
+	}
+	if (no_buf_ring || no_buf_ring_status)
+		return T_EXIT_SKIP;
+
+	ret = test(1);
+	if (ret == T_EXIT_FAIL) {
+		fprintf(stderr, "test 1 failed\n");
+		return T_EXIT_FAIL;
+	}
+
+	return T_EXIT_PASS;
+}