[PATCH liburing 0/3] zcrx: add support for notifications and statistic

[PATCH liburing 0/3] zcrx: add support for notifications and statistic

[Clément Léger]

Add support for io_uring notifications and statistics. This allows to
receive CQE notifications upon out of buffers or when copy fallback is
used to copy packets. An associated stat structure is also supported and
allows to track copy fallback statistics without requiring the user to
handle notifications CQEs.

Clément Léger (3):
  Update uapi headers to add ZCRX notification
  test/zcrx: add ZCRX notification/stats tests
  examples/zcrx: add notification support

 examples/zcrx.c                       | 109 ++++++++
 src/include/liburing/io_uring.h       |  34 ++-
 src/include/liburing/io_uring/query.h |  12 +
 test/zcrx.c                           | 376 ++++++++++++++++++++++++++
 4 files changed, 530 insertions(+), 1 deletion(-)

-- 
Clément Léger

[PATCH liburing 1/3] Update uapi headers to add ZCRX notification

[Clément Léger]

Extend the ZCRX uapi with notification infrastructure: notification
types (no-buffers, copy-fallback), stats structure, notification
descriptor for registration and arm control operation.

A new IORING_REGISTER_QUERY op allows userspace to discover the
notification stats memory layout. Additionnal flags are as well required
to register the statistics structure and enable notification upon
registration.

Signed-off-by: Clément Léger <cleger@meta.com>
---
 src/include/liburing/io_uring.h       | 34 ++++++++++++++++++++++++++-
 src/include/liburing/io_uring/query.h | 12 ++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/src/include/liburing/io_uring.h b/src/include/liburing/io_uring.h
index 983aa265..a479d0b5 100644
--- a/src/include/liburing/io_uring.h
+++ b/src/include/liburing/io_uring.h
@@ -1068,6 +1068,30 @@ enum zcrx_features {
 	 * value in struct io_uring_zcrx_ifq_reg::rx_buf_len.
 	 */
 	ZCRX_FEATURE_RX_PAGE_SIZE	= 1 << 0,
+	ZCRX_FEATURE_NOTIFICATION	= 1 << 1,
+};
+
+enum zcrx_notification_type {
+	ZCRX_NOTIF_NO_BUFFERS = 1 << 0,
+	ZCRX_NOTIF_COPY = 1 << 1
+};
+
+enum zcrx_notification_desc_flags {
+	/* If set, assume that stats_offset hold a correct offset to  */
+	ZCRX_NOTIF_DESC_FLAG_STATS = 1 << 0,
+};
+
+struct io_uring_zcrx_notif_stats {
+	__u64	copy_count;	/* cumulative copy-fallback CQEs */
+	__u64	copy_bytes;	/* cumulative bytes copied */
+};
+
+struct zcrx_notification_desc {
+	__u64	user_data;
+	__u32	type_mask;
+	__u32	flags; /* see enum zcrx_notification_desc_flags */
+	__u64	stats_offset; /* offset from the beginning of refill ring region for stats */
+	__u64	__resv2[9];
 };
 
 /*
@@ -1085,12 +1109,14 @@ struct io_uring_zcrx_ifq_reg {
 	struct io_uring_zcrx_offsets offsets;
 	__u32	zcrx_id;
 	__u32	rx_buf_len;
-	__u64	__resv[3];
+	__u64	notif_desc; /* see struct zcrx_notification_desc */
+	__u64	__resv[2];
 };
 
 enum zcrx_ctrl_op {
 	ZCRX_CTRL_FLUSH_RQ,
 	ZCRX_CTRL_EXPORT,
+	ZCRX_CTRL_ARM_NOTIFICATION,
 
 	__ZCRX_CTRL_LAST,
 };
@@ -1104,6 +1130,11 @@ struct zcrx_ctrl_export {
 	__u32 		__resv1[11];
 };
 
+struct zcrx_ctrl_arm_notif {
+	__u32		type_mask;
+	__u32		__resv[11];
+};
+
 struct zcrx_ctrl {
 	__u32	zcrx_id;
 	__u32	op; /* see enum zcrx_ctrl_op */
@@ -1112,6 +1143,7 @@ struct zcrx_ctrl {
 	union {
 		struct zcrx_ctrl_export		zc_export;
 		struct zcrx_ctrl_flush_rq	zc_flush;
+		struct zcrx_ctrl_arm_notif	zc_arm_notif;
 	};
 };
 
diff --git a/src/include/liburing/io_uring/query.h b/src/include/liburing/io_uring/query.h
index 0b624817..ad5efd99 100644
--- a/src/include/liburing/io_uring/query.h
+++ b/src/include/liburing/io_uring/query.h
@@ -20,6 +20,7 @@ enum {
 	IO_URING_QUERY_OPCODES			= 0,
 	IO_URING_QUERY_ZCRX			= 1,
 	IO_URING_QUERY_SCQ			= 2,
+	IO_URING_QUERY_ZCRX_NOTIF		= 3,
 
 	__IO_URING_QUERY_MAX,
 };
@@ -59,6 +60,17 @@ struct io_uring_query_zcrx {
 	__u64 __resv2;
 };
 
+struct io_uring_query_zcrx_notif {
+	/* Bitmask of supported ZCRX_NOTIF_* flags*/
+	__u32 notif_flags;
+	/* Size of io_uring_zcrx_notif_stats */
+	__u32 notif_stats_size;
+	/* Required alignment for the stats buffer within the region (ie stats_offset) */
+	__u32 notif_stats_off_alignment;
+	__u32 resv1;
+	__u64 __resv2[10];
+};
+
 struct io_uring_query_scq {
 	/* The SQ/CQ rings header size */
 	__u64 hdr_size;
-- 
2.52.0

[PATCH liburing 2/3] test/zcrx: add ZCRX notification/stats tests

[Clément Léger]

Enhance zcrx testsutie so that it covers notification registration
validation, stats region placement, and arm operation error paths.

Signed-off-by: Clément Léger <cleger@meta.com>
---
 test/zcrx.c | 376 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 376 insertions(+)

diff --git a/test/zcrx.c b/test/zcrx.c
index a0e8a863..b13eab43 100644
--- a/test/zcrx.c
+++ b/test/zcrx.c
@@ -811,6 +811,364 @@ static int test_zcrx_clone(void)
 	return 0;
 }
 
+static int try_register_ifq_with_notif(struct zcrx_notification_desc *notif,
+				       void *area)
+{
+	struct io_uring_region_desc rq_region = {
+		.size = get_rq_size(RQ_ENTRIES),
+		.user_addr = uring_ptr_to_u64(def_rq_mem),
+		.flags = IORING_MEM_REGION_TYPE_USER,
+	};
+	struct io_uring_zcrx_area_reg area_reg = {
+		.addr = uring_ptr_to_u64(area),
+		.len = AREA_SZ,
+		.flags = 0,
+	};
+	struct io_uring_zcrx_ifq_reg reg = {
+		.flags = ZCRX_REG_NODEV,
+		.rq_entries = RQ_ENTRIES,
+		.area_ptr = uring_ptr_to_u64(&area_reg),
+		.region_ptr = uring_ptr_to_u64(&rq_region),
+		.notif_desc = uring_ptr_to_u64(notif),
+	};
+
+	return try_register_zcrx(&reg);
+}
+
+static int test_notif_registration(void *area)
+{
+	struct zcrx_notification_desc notif;
+	int ret;
+
+	/* Valid registration with no-buffers notification */
+	memset(&notif, 0, sizeof(notif));
+	notif.user_data = 42;
+	notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+
+	ret = try_register_ifq_with_notif(&notif, area);
+	if (ret) {
+		fprintf(stderr, "valid notif registration failed: %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	/* Valid registration with both notification types */
+	memset(&notif, 0, sizeof(notif));
+	notif.user_data = 42;
+	notif.type_mask = ZCRX_NOTIF_NO_BUFFERS |
+			  ZCRX_NOTIF_COPY;
+
+	ret = try_register_ifq_with_notif(&notif, area);
+	if (ret) {
+		fprintf(stderr, "valid notif both-types registration failed: %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	/* Invalid type_mask (bit beyond valid range) */
+	memset(&notif, 0, sizeof(notif));
+	notif.user_data = 42;
+	notif.type_mask = (1 << 31);
+
+	ret = try_register_ifq_with_notif(&notif, area);
+	if (ret != -EINVAL) {
+		fprintf(stderr, "invalid notif type_mask should fail with -EINVAL, got %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	/* Non-zero reserved field should fail */
+	memset(&notif, 0, sizeof(notif));
+	notif.user_data = 42;
+	notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+	notif.__resv2[0] = 1;
+
+	ret = try_register_ifq_with_notif(&notif, area);
+	if (ret != -EINVAL) {
+		fprintf(stderr, "non-zero resv2 should fail with -EINVAL, got %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	/* Invalid notif flags */
+	memset(&notif, 0, sizeof(notif));
+	notif.user_data = 42;
+	notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+	notif.flags = ~ZCRX_NOTIF_DESC_FLAG_STATS;
+
+	ret = try_register_ifq_with_notif(&notif, area);
+	if (ret != -EINVAL) {
+		fprintf(stderr, "invalid notif flags should fail with -EINVAL, got %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	return T_EXIT_PASS;
+}
+
+static int test_notif_stats(void *area)
+{
+	struct io_uring_query_zcrx zcrx_query = {};
+	struct io_uring_query_hdr zcrx_hdr = {
+		.size = sizeof(zcrx_query),
+		.query_data = uring_ptr_to_u64(&zcrx_query),
+		.query_op = IO_URING_QUERY_ZCRX,
+	};
+	struct io_uring_query_zcrx_notif notif_query = {};
+	struct io_uring_query_hdr notif_hdr = {
+		.size = sizeof(notif_query),
+		.query_data = uring_ptr_to_u64(&notif_query),
+		.query_op = IO_URING_QUERY_ZCRX_NOTIF,
+	};
+	struct zcrx_notification_desc notif;
+	size_t ring_size, rqes_end, stats_off, hdr_size;
+	int ret;
+	int tret = T_EXIT_FAIL;
+	void *ring_mem;
+
+	/*
+	 * Query the kernel for the actual refill ring header size and stats
+	 * struct layout so we can build a correctly-sized user region.
+	 */
+	ret = io_uring_register(-1, IORING_REGISTER_QUERY, &zcrx_hdr, 0);
+	if (ret < 0 || zcrx_hdr.result < 0) {
+		fprintf(stderr, "zcrx query not supported: %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	ret = io_uring_register(-1, IORING_REGISTER_QUERY, &notif_hdr, 0);
+	if (ret < 0 || notif_hdr.result < 0) {
+		fprintf(stderr, "zcrx notif query not supported: %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	if (!notif_query.notif_stats_size) {
+		fprintf(stderr, "notif stats not supported\n");
+		return T_EXIT_FAIL;
+	}
+
+	hdr_size = T_ALIGN_UP((__u64)zcrx_query.rq_hdr_size,
+			     zcrx_query.rq_hdr_alignment);
+	rqes_end = hdr_size + sizeof(struct io_uring_zcrx_rqe) * RQ_ENTRIES;
+	stats_off = T_ALIGN_UP(rqes_end, notif_query.notif_stats_off_alignment);
+	ring_size = stats_off + notif_query.notif_stats_size;
+	ring_size = T_ALIGN_UP(ring_size, page_size);
+
+	ring_mem = mmap(NULL, ring_size, PROT_READ | PROT_WRITE,
+			MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
+	if (ring_mem == MAP_FAILED) {
+		perror("mmap ring_mem");
+		return T_EXIT_FAIL;
+	}
+
+	struct io_uring_region_desc region = {
+		.user_addr = uring_ptr_to_u64(ring_mem),
+		.size = ring_size,
+		.flags = IORING_MEM_REGION_TYPE_USER,
+	};
+	struct io_uring_zcrx_area_reg area_reg = {
+		.addr = uring_ptr_to_u64(area),
+		.len = AREA_SZ,
+		.flags = 0,
+	};
+
+	/* Valid stats offset */
+	memset(&notif, 0, sizeof(notif));
+	notif.user_data = 42;
+	notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+	notif.flags = ZCRX_NOTIF_DESC_FLAG_STATS;
+	notif.stats_offset = stats_off;
+
+	struct io_uring_zcrx_ifq_reg reg = {
+		.flags = ZCRX_REG_NODEV,
+		.rq_entries = RQ_ENTRIES,
+		.area_ptr = uring_ptr_to_u64(&area_reg),
+		.region_ptr = uring_ptr_to_u64(&region),
+		.notif_desc = uring_ptr_to_u64(&notif),
+	};
+
+	/*
+	 * Pre-fill the stats area with non-zero values to verify the kernel
+	 * zeroes it on registration.
+	 */
+	memset((char *)ring_mem + stats_off, 0xff,
+	       sizeof(struct io_uring_zcrx_notif_stats));
+
+	ret = try_register_zcrx(&reg);
+	if (ret) {
+		fprintf(stderr, "notif stats registration failed: %d\n", ret);
+		goto unmap;
+	}
+
+	/* Verify the stats area was zeroed by the kernel */
+	struct io_uring_zcrx_notif_stats *stats =
+		(struct io_uring_zcrx_notif_stats *)((char *)ring_mem + stats_off);
+	if (stats->copy_count || stats->copy_bytes) {
+		fprintf(stderr, "stats not zeroed after registration\n");
+		goto unmap;
+	}
+
+	/* Misaligned stats offset should fail */
+	memset(&notif, 0, sizeof(notif));
+	notif.user_data = 42;
+	notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+	notif.flags = ZCRX_NOTIF_DESC_FLAG_STATS;
+	notif.stats_offset = stats_off + 1;
+
+	region.user_addr = uring_ptr_to_u64(ring_mem);
+	reg.region_ptr = uring_ptr_to_u64(&region);
+	reg.notif_desc = uring_ptr_to_u64(&notif);
+
+	ret = try_register_zcrx(&reg);
+	if (ret != -EINVAL) {
+		fprintf(stderr, "misaligned stats_offset should fail, got %d\n", ret);
+		goto unmap;
+	}
+
+	/* Stats offset overlapping with rqes should fail */
+	notif.stats_offset = sizeof(struct io_uring_zcrx_rqe);
+
+	ret = try_register_zcrx(&reg);
+	if (ret != -ERANGE) {
+		fprintf(stderr, "overlapping stats_offset should fail, got %d\n", ret);
+		goto unmap;
+	}
+
+	/* Stats offset beyond region should fail */
+	notif.stats_offset = ring_size;
+
+	ret = try_register_zcrx(&reg);
+	if (ret != -ERANGE) {
+		fprintf(stderr, "out-of-bounds stats_offset should fail, got %d\n", ret);
+		goto unmap;
+	}
+
+	tret = T_EXIT_PASS;
+unmap:
+	munmap(ring_mem, ring_size);
+	return tret;
+}
+
+static int test_arm_notification(void *area)
+{
+	struct zcrx_notification_desc notif;
+	struct io_uring ring;
+	struct zcrx_ctrl ctrl;
+	__u32 zcrx_id;
+	int ret;
+
+	struct io_uring_zcrx_area_reg area_reg = {
+		.addr = uring_ptr_to_u64(area),
+		.len = AREA_SZ,
+		.flags = 0,
+	};
+	struct io_uring_region_desc rq_region = {
+		.size = get_rq_size(RQ_ENTRIES),
+		.user_addr = uring_ptr_to_u64(def_rq_mem),
+		.flags = IORING_MEM_REGION_TYPE_USER,
+	};
+	struct io_uring_zcrx_ifq_reg reg = {
+		.flags = ZCRX_REG_NODEV,
+		.rq_entries = RQ_ENTRIES,
+		.area_ptr = uring_ptr_to_u64(&area_reg),
+		.region_ptr = uring_ptr_to_u64(&rq_region),
+	};
+
+	/* First, register without notifications and check arm fails */
+	ret = t_create_ring(128, &ring, RING_FLAGS);
+	if (ret != T_SETUP_OK) {
+		fprintf(stderr, "ring create failed: %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	ret = io_uring_register_ifq(&ring, &reg);
+	if (ret) {
+		fprintf(stderr, "ifq register failed: %d\n", ret);
+		io_uring_queue_exit(&ring);
+		return T_EXIT_FAIL;
+	}
+	zcrx_id = reg.zcrx_id;
+
+	memset(&ctrl, 0, sizeof(ctrl));
+	ctrl.zcrx_id = zcrx_id;
+	ctrl.op = ZCRX_CTRL_ARM_NOTIFICATION;
+	ctrl.zc_arm_notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+
+	ret = io_uring_register(ring.ring_fd, IORING_REGISTER_ZCRX_CTRL,
+				&ctrl, 0);
+	if (ret != -EINVAL) {
+		fprintf(stderr, "arm without prior notif should fail, got %d\n", ret);
+		io_uring_queue_exit(&ring);
+		return T_EXIT_FAIL;
+	}
+	io_uring_queue_exit(&ring);
+
+	/* Now register with notifications and try arm with invalid mask */
+	memset(&notif, 0, sizeof(notif));
+	notif.user_data = 42;
+	notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+
+	reg.notif_desc = uring_ptr_to_u64(&notif);
+
+	ret = t_create_ring(128, &ring, RING_FLAGS);
+	if (ret != T_SETUP_OK) {
+		fprintf(stderr, "ring create failed: %d\n", ret);
+		return T_EXIT_FAIL;
+	}
+
+	ret = io_uring_register_ifq(&ring, &reg);
+	if (ret) {
+		fprintf(stderr, "notif ifq register failed: %d\n", ret);
+		io_uring_queue_exit(&ring);
+		return T_EXIT_FAIL;
+	}
+	zcrx_id = reg.zcrx_id;
+
+	/* Arm with invalid type (bit 31) should fail */
+	memset(&ctrl, 0, sizeof(ctrl));
+	ctrl.zcrx_id = zcrx_id;
+	ctrl.op = ZCRX_CTRL_ARM_NOTIFICATION;
+	ctrl.zc_arm_notif.type_mask = (1 << 31);
+
+	ret = io_uring_register(ring.ring_fd, IORING_REGISTER_ZCRX_CTRL,
+				&ctrl, 0);
+	if (ret != -EINVAL) {
+		fprintf(stderr, "arm with invalid mask should fail, got %d\n", ret);
+		io_uring_queue_exit(&ring);
+		return T_EXIT_FAIL;
+	}
+
+	/*
+	 * Arm for a type that hasn't fired yet should fail because the kernel
+	 * only allows rearming notifications that have already been delivered.
+	 */
+	memset(&ctrl, 0, sizeof(ctrl));
+	ctrl.zcrx_id = zcrx_id;
+	ctrl.op = ZCRX_CTRL_ARM_NOTIFICATION;
+	ctrl.zc_arm_notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+
+	ret = io_uring_register(ring.ring_fd, IORING_REGISTER_ZCRX_CTRL,
+				&ctrl, 0);
+	if (ret != -EINVAL) {
+		fprintf(stderr, "arm for unfired notif should fail, got %d\n", ret);
+		io_uring_queue_exit(&ring);
+		return T_EXIT_FAIL;
+	}
+
+	/* Non-zero reserved in arm_notif should fail */
+	memset(&ctrl, 0, sizeof(ctrl));
+	ctrl.zcrx_id = zcrx_id;
+	ctrl.op = ZCRX_CTRL_ARM_NOTIFICATION;
+	ctrl.zc_arm_notif.type_mask = ZCRX_NOTIF_NO_BUFFERS;
+	ctrl.zc_arm_notif.__resv[0] = 1;
+
+	ret = io_uring_register(ring.ring_fd, IORING_REGISTER_ZCRX_CTRL,
+				&ctrl, 0);
+	if (ret != -EINVAL) {
+		fprintf(stderr, "arm with non-zero resv should fail, got %d\n", ret);
+		io_uring_queue_exit(&ring);
+		return T_EXIT_FAIL;
+	}
+
+	io_uring_queue_exit(&ring);
+	return T_EXIT_PASS;
+}
+
 static int test_rq_flush(void)
 {
 	struct t_executor ctx;
@@ -1112,6 +1470,24 @@ static int run_tests(void)
 		return T_EXIT_FAIL;
 	}
 
+	ret = test_notif_registration(def_area_mem);
+	if (ret) {
+		fprintf(stderr, "test_notif_registration failed\n");
+		return ret;
+	}
+
+	ret = test_notif_stats(def_area_mem);
+	if (ret) {
+		fprintf(stderr, "test_notif_stats failed\n");
+		return ret;
+	}
+
+	ret = test_arm_notification(def_area_mem);
+	if (ret) {
+		fprintf(stderr, "test_arm_notification failed\n");
+		return ret;
+	}
+
 	ret = test_recv();
 	if (ret) {
 		fprintf(stderr, "test_recv() failed %i\n", ret);
-- 
2.52.0

[PATCH liburing 3/3] examples/zcrx: add notification support

[Clément Léger]

Add support to register notifications and display stats when
notifications are received. This allows to provide an example of how to
register and configure notifications and stats support. If an error is
encountered while running, then the notification type will be display
along stat counter if associated to the notification (ie copy fallback).

Notification support is probe before zcrx setup and enabled by default
if detected.

Signed-off-by: Clément Léger <cleger@meta.com>
---
 examples/zcrx.c | 109 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 109 insertions(+)

diff --git a/examples/zcrx.c b/examples/zcrx.c
index b55b07ce..537c62ad 100644
--- a/examples/zcrx.c
+++ b/examples/zcrx.c
@@ -91,8 +91,11 @@ static size_t cfg_size = 0;
 static unsigned cfg_affinity_mode = AFFINITY_MODE_NONE;
 static unsigned cfg_rq_alloc_mode = RQ_ALLOC_USER;
 static unsigned cfg_area_type = AREA_TYPE_NORMAL;
+static bool cfg_notif = false;
 static struct sockaddr_in6 cfg_addr;
 
+#define NOTIF_USER_DATA		UINT64_MAX
+
 static long page_size;
 
 static void *area_ptr;
@@ -102,6 +105,10 @@ static struct io_uring_zcrx_rq rq_ring;
 static unsigned long area_token;
 static bool stop;
 static __u32 zcrx_id;
+static __u32 notif_stats_size;
+static __u32 notif_stats_alignment;
+static __u32 notif_stats_offset;
+static struct io_uring_zcrx_notif_stats *notif_stats;
 
 static int dmabuf_fd;
 static int memfd;
@@ -164,6 +171,13 @@ static inline size_t get_refill_ring_size(unsigned int rq_entries)
 	ring_size = rq_entries * sizeof(struct io_uring_zcrx_rqe);
 	/* add space for the header (head/tail/etc.) */
 	ring_size += page_size;
+
+	if (cfg_notif && notif_stats_size) {
+		notif_stats_offset = T_ALIGN_UP(ring_size,
+						notif_stats_alignment);
+		ring_size = notif_stats_offset + notif_stats_size;
+	}
+
 	return T_ALIGN_UP(ring_size, page_size);
 }
 
@@ -236,9 +250,48 @@ static void zcrx_populate_area(struct io_uring_zcrx_area_reg *area_reg)
 	area_reg->flags = 0;
 }
 
+static void zcrx_probe_notif(void)
+{
+	struct io_uring_query_zcrx zcrx_query;
+	struct io_uring_query_hdr zcrx_hdr = {
+		.size = sizeof(zcrx_query),
+		.query_data = uring_ptr_to_u64(&zcrx_query),
+		.query_op = IO_URING_QUERY_ZCRX,
+	};
+	struct io_uring_query_zcrx_notif notif_query;
+	struct io_uring_query_hdr notif_hdr = {
+		.size = sizeof(notif_query),
+		.query_data = uring_ptr_to_u64(&notif_query),
+		.query_op = IO_URING_QUERY_ZCRX_NOTIF,
+	};
+	int ret;
+
+	ret = io_uring_register(-1, IORING_REGISTER_QUERY, &zcrx_hdr, 0);
+	if (ret < 0 || zcrx_hdr.result < 0) {
+		printf("Failed to query ZCRX\n");
+		return;
+	}
+
+	if (!(zcrx_query.features & ZCRX_FEATURE_NOTIFICATION)) {
+		printf("Kernel doesn't support ZCRX notifications\n");
+		return;
+	}
+
+	ret = io_uring_register(-1, IORING_REGISTER_QUERY, &notif_hdr, 0);
+	if (ret < 0 || notif_hdr.result < 0) {
+		printf("Failed to query ZCRX stats parameters\n");
+		return;
+	}
+
+	notif_stats_size = notif_query.notif_stats_size;
+	notif_stats_alignment = notif_query.notif_stats_off_alignment;
+	cfg_notif = true;
+}
+
 static void setup_zcrx(struct io_uring *ring)
 {
 	struct io_uring_zcrx_area_reg area_reg;
+	struct zcrx_notification_desc notif;
 	unsigned int ifindex;
 	unsigned int rq_entries = cfg_rq_entries;
 	unsigned rq_flags = 0;
@@ -276,6 +329,18 @@ static void setup_zcrx(struct io_uring *ring)
 		.region_ptr = uring_ptr_to_u64(&region_reg),
 	};
 
+	if (cfg_notif) {
+		memset(&notif, 0, sizeof(notif));
+		notif.user_data = NOTIF_USER_DATA;
+		notif.type_mask = ZCRX_NOTIF_NO_BUFFERS |
+				  ZCRX_NOTIF_COPY;
+		if (notif_stats_size) {
+			notif.flags = ZCRX_NOTIF_DESC_FLAG_STATS;
+			notif.stats_offset = notif_stats_offset;
+		}
+		reg.notif_desc = uring_ptr_to_u64(&notif);
+	}
+
 	ret = io_uring_register_ifq(ring, &reg);
 	if (ret)
 		t_error(1, 0, "io_uring_register_ifq(): %d", ret);
@@ -297,6 +362,10 @@ static void setup_zcrx(struct io_uring *ring)
 
 	zcrx_id = reg.zcrx_id;
 	area_token = area_reg.rq_area_token;
+
+	if (cfg_notif && notif_stats_size)
+		notif_stats = (struct io_uring_zcrx_notif_stats *)
+				((char *)ring_ptr + notif_stats_offset);
 }
 
 static void add_accept(struct io_uring *ring, int sockfd)
@@ -473,6 +542,40 @@ static void process_recvzc(struct io_uring *ring,
 	return_buffer(&rq_ring, cqe);
 }
 
+static void rearm_notification(struct io_uring *ring, __u32 type_mask)
+{
+	struct zcrx_ctrl ctrl = {
+		.zcrx_id = zcrx_id,
+		.op = ZCRX_CTRL_ARM_NOTIFICATION,
+	};
+	int ret;
+
+	ctrl.zc_arm_notif.type_mask = type_mask;
+	ret = io_uring_register(ring->ring_fd, IORING_REGISTER_ZCRX_CTRL,
+				&ctrl, 0);
+	if (ret < 0)
+		fprintf(stderr, "arm notification failed: %d\n", ret);
+}
+
+static void process_notification(struct io_uring *ring,
+				 struct io_uring_cqe *cqe)
+{
+	__u32 type_mask = cqe->res;
+
+	if (type_mask & ZCRX_NOTIF_NO_BUFFERS)
+		printf("Notification: no buffers available\n");
+	if (type_mask & ZCRX_NOTIF_COPY) {
+		printf("Notification: copy fallback\n");
+
+		if (notif_stats)
+			printf("stats: copy_count=%llu copy_bytes=%llu\n",
+			       IO_URING_READ_ONCE(notif_stats->copy_count),
+			       IO_URING_READ_ONCE(notif_stats->copy_bytes));
+	}
+
+	rearm_notification(ring, type_mask);
+}
+
 static void server_loop(struct io_uring *ring)
 {
 	struct io_uring_cqe *cqe;
@@ -484,6 +587,11 @@ static void server_loop(struct io_uring *ring)
 		t_error(1, ret, "io_uring_submit_and_wait failed\n");
 
 	io_uring_for_each_cqe(ring, head, cqe) {
+		if (cqe->user_data == NOTIF_USER_DATA) {
+			process_notification(ring, cqe);
+			count++;
+			continue;
+		}
 		switch (cqe->user_data & REQ_TYPE_MASK) {
 		case REQ_TYPE_ACCEPT:
 			process_accept(ring, cqe);
@@ -534,6 +642,7 @@ static void run_server(void)
 	if (ret)
 		t_error(1, ret, "ring init failed");
 
+	zcrx_probe_notif();
 	setup_zcrx(&ring);
 	add_accept(&ring, listen_fd);
 
-- 
2.52.0