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

["Clément Léger" <cleger@meta.com>]

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