[PATCH 08/11] io_uring: allocate a uarg for freeing zero copy skbs

[David Wei <dw@davidwei.uk>]

From: David Wei <davidhwei@meta.com>

As ZC skbs are marked as zero copy, they will bypass the default skb
frag destructor. This patch adds a static uarg that is attached to ZC
bufs and a callback that returns them to the freelist of a ZC pool.

Signed-off-by: David Wei <davidhwei@meta.com>
Co-developed-by: Jonathan Lemon <jonathan.lemon@gmail.com>
---
 include/linux/io_uring.h  |  7 +++++++
 include/linux/netdevice.h |  1 +
 io_uring/zc_rx.c          | 44 +++++++++++++++++++++++++++++++++++++++
 io_uring/zc_rx.h          |  2 ++
 4 files changed, 54 insertions(+)

diff --git a/include/linux/io_uring.h b/include/linux/io_uring.h
index 61eae25a8f1d..e2a4f92df814 100644
--- a/include/linux/io_uring.h
+++ b/include/linux/io_uring.h
@@ -62,6 +62,8 @@ const char *io_uring_get_opcode(u8 opcode);
 
 struct io_zc_rx_ifq;
 struct io_zc_rx_buf *io_zc_rx_get_buf(struct io_zc_rx_ifq *ifq);
+struct io_zc_rx_buf *io_zc_rx_buf_from_page(struct io_zc_rx_ifq *ifq,
+					    struct page *page);
 void io_zc_rx_put_buf(struct io_zc_rx_ifq *ifq, struct io_zc_rx_buf *buf);
 static inline dma_addr_t io_zc_rx_buf_dma(struct io_zc_rx_buf *buf)
 {
@@ -123,6 +125,11 @@ static inline struct io_zc_rx_buf *io_zc_rx_get_buf(struct io_zc_rx_ifq *ifq)
 {
 	return NULL;
 }
+static inline struct io_zc_rx_buf *io_zc_rx_buf_from_page(struct io_zc_rx_ifq *ifq,
+							  struct page *page)
+{
+	return NULL;
+}
 void io_zc_rx_put_buf(struct io_zc_rx_ifq *ifq, struct io_zc_rx_buf *buf)
 {
 }
diff --git a/include/linux/netdevice.h b/include/linux/netdevice.h
index a20a5c847916..bf133cbee721 100644
--- a/include/linux/netdevice.h
+++ b/include/linux/netdevice.h
@@ -1043,6 +1043,7 @@ struct netdev_bpf {
 		struct {
 			struct io_zc_rx_ifq *ifq;
 			u16 queue_id;
+			struct ubuf_info *uarg;
 		} zc_rx;
 	};
 };
diff --git a/io_uring/zc_rx.c b/io_uring/zc_rx.c
index 14bc063f1c6c..b8dd699e2777 100644
--- a/io_uring/zc_rx.c
+++ b/io_uring/zc_rx.c
@@ -44,6 +44,11 @@ static u64 mk_page_info(u16 pool_id, u32 pgid)
 	return (u64)0xface << 48 | (u64)pool_id << 32 | (u64)pgid;
 }
 
+static bool is_zc_rx_page(struct page *page)
+{
+	return PagePrivate(page) && ((page_private(page) >> 48) == 0xface);
+}
+
 typedef int (*bpf_op_t)(struct net_device *dev, struct netdev_bpf *bpf);
 
 static int __io_queue_mgmt(struct net_device *dev, struct io_zc_rx_ifq *ifq,
@@ -59,6 +64,7 @@ static int __io_queue_mgmt(struct net_device *dev, struct io_zc_rx_ifq *ifq,
 	cmd.command = XDP_SETUP_ZC_RX;
 	cmd.zc_rx.ifq = ifq;
 	cmd.zc_rx.queue_id = queue_id;
+	cmd.zc_rx.uarg = &ifq->uarg;
 
 	return ndo_bpf(dev, &cmd);
 }
@@ -73,6 +79,26 @@ static int io_close_zc_rxq(struct io_zc_rx_ifq *ifq)
 	return __io_queue_mgmt(ifq->dev, NULL, ifq->if_rxq_id);
 }
 
+static void io_zc_rx_skb_free(struct sk_buff *skb, struct ubuf_info *uarg,
+			      bool success)
+{
+	struct skb_shared_info *shinfo = skb_shinfo(skb);
+	struct io_zc_rx_ifq *ifq;
+	struct io_zc_rx_buf *buf;
+	struct page *page;
+	int i;
+
+	ifq = container_of(uarg, struct io_zc_rx_ifq, uarg);
+	for (i = 0; i < shinfo->nr_frags; i++) {
+		page = skb_frag_page(&shinfo->frags[i]);
+		buf = io_zc_rx_buf_from_page(ifq, page);
+		if (likely(buf))
+			io_zc_rx_put_buf(ifq, buf);
+		else
+			__skb_frag_unref(&shinfo->frags[i], skb->pp_recycle);
+	}
+}
+
 static int io_zc_rx_map_buf(struct device *dev, struct page *page, u16 pool_id,
 			    u32 pgid, struct io_zc_rx_buf *buf)
 {
@@ -268,6 +294,8 @@ int io_register_zc_rx_ifq(struct io_ring_ctx *ctx,
 	if (ret)
 		goto err;
 
+	ifq->uarg.callback = io_zc_rx_skb_free;
+	ifq->uarg.flags = SKBFL_ALL_ZEROCOPY | SKBFL_FIXED_FRAG;
 	ifq->rq_entries = reg.rq_entries;
 	ifq->cq_entries = reg.cq_entries;
 	ifq->cached_rq_head = 0;
@@ -407,3 +435,19 @@ void io_zc_rx_put_buf(struct io_zc_rx_ifq *ifq, struct io_zc_rx_buf *buf)
 	io_zc_rx_recycle_buf(pool, buf);
 }
 EXPORT_SYMBOL(io_zc_rx_put_buf);
+
+struct io_zc_rx_buf *io_zc_rx_buf_from_page(struct io_zc_rx_ifq *ifq,
+					    struct page *page)
+{
+	struct io_zc_rx_pool *pool;
+	int pgid;
+
+	if (!is_zc_rx_page(page))
+		return NULL;
+
+	pool = ifq->pool;
+	pgid = page_private(page) & 0xffffffff;
+
+	return &pool->bufs[pgid];
+}
+EXPORT_SYMBOL(io_zc_rx_buf_from_page);
diff --git a/io_uring/zc_rx.h b/io_uring/zc_rx.h
index b063a3c81ccb..ee7e36295f91 100644
--- a/io_uring/zc_rx.h
+++ b/io_uring/zc_rx.h
@@ -3,6 +3,7 @@
 #define IOU_ZC_RX_H
 
 #include <linux/io_uring_types.h>
+#include <linux/skbuff.h>
 
 struct io_zc_rx_ifq {
 	struct io_ring_ctx	*ctx;
@@ -10,6 +11,7 @@ struct io_zc_rx_ifq {
 	struct io_rbuf_ring	*ring;
 	struct io_uring_rbuf_rqe *rqes;
 	struct io_uring_rbuf_cqe *cqes;
+	struct ubuf_info	uarg;
 	u32			rq_entries, cq_entries;
 	u32			cached_rq_head;
 	u32			cached_cq_tail;
-- 
2.39.3