[PATCH net-next 1/3] net: netmem: add net_iov_area freelist helpers

[Stanislav Fomichev <sdf.kernel@gmail.com>]

io_uring zero-copy receive and devmem both maintain a bounded LIFO for
net_iovs in a contiguous area. Store the freelist in struct net_iov_area
and provide common push and pop helpers.

Leave synchronization to area owners. Keep devmem's area adjacent to its
lock. Use u32 indices and counts, which halves devmem's freelist storage
on 64-bit systems. Reject devmem areas with more than U32_MAX entries
before narrowing the count. With 4 KiB chunks, the limit is almost 16 TiB.

Signed-off-by: Stanislav Fomichev <sdf@fomichev.me>
---
 include/net/netmem.h | 28 ++++++++++++++++++++++++++-
 io_uring/zcrx.c      | 25 +++++++++---------------
 io_uring/zcrx.h      |  4 ----
 net/core/devmem.c    | 46 ++++++++++++++++++++++----------------------
 net/core/devmem.h    |  7 +++----
 5 files changed, 62 insertions(+), 48 deletions(-)

diff --git a/include/net/netmem.h b/include/net/netmem.h
index bccacd21b6c3..da885d95ea63 100644
--- a/include/net/netmem.h
+++ b/include/net/netmem.h
@@ -101,10 +101,15 @@ struct net_iov {
 struct net_iov_area {
 	/* Array of net_iovs for this area. */
 	struct net_iov *niovs;
-	size_t num_niovs;
+
+	/* Stack of free net_iov indices. */
+	u32 *freelist;
 
 	/* Offset into the dma-buf where this chunk starts.  */
 	unsigned long base_virtual;
+
+	u32 num_niovs;
+	u32 free_count;
 };
 
 static inline struct net_iov_area *net_iov_owner(const struct net_iov *niov)
@@ -117,6 +122,27 @@ static inline unsigned int net_iov_idx(const struct net_iov *niov)
 	return niov - net_iov_owner(niov)->niovs;
 }
 
+static inline struct net_iov *net_iov_area_pop(struct net_iov_area *area)
+{
+	u32 idx;
+
+	if (unlikely(!area->free_count))
+		return NULL;
+
+	idx = area->freelist[--area->free_count];
+	return &area->niovs[idx];
+}
+
+static inline void net_iov_area_push(struct net_iov_area *area,
+				     struct net_iov *niov)
+{
+	if (WARN_ON_ONCE(net_iov_owner(niov) != area ||
+			 area->free_count >= area->num_niovs))
+		return;
+
+	area->freelist[area->free_count++] = net_iov_idx(niov);
+}
+
 /* Initialize a niov: stamp the owning area, the memory provider type.
  */
 static inline void net_iov_init(struct net_iov *niov,
diff --git a/io_uring/zcrx.c b/io_uring/zcrx.c
index 86d580d4410d..fa6061127b9d 100644
--- a/io_uring/zcrx.c
+++ b/io_uring/zcrx.c
@@ -462,7 +462,7 @@ static void io_zcrx_free_area(struct io_zcrx_ifq *ifq,
 		io_unaccount_mem(ifq->user, ifq->mm_account,
 				 area->mem.account_pages);
 
-	kvfree(area->freelist);
+	kvfree(area->nia.freelist);
 	kvfree(area->nia.niovs);
 	kvfree(area->user_refs);
 	kfree(area);
@@ -548,9 +548,10 @@ static int __zcrx_create_area(struct io_zcrx_ifq *ifq,
 	if (!area->nia.niovs)
 		goto err;
 
-	area->freelist = kvmalloc_array(nr_iovs, sizeof(area->freelist[0]),
-					GFP_KERNEL_ACCOUNT | __GFP_ZERO);
-	if (!area->freelist)
+	area->nia.freelist = kvmalloc_array(nr_iovs,
+					    sizeof(area->nia.freelist[0]),
+					    GFP_KERNEL_ACCOUNT | __GFP_ZERO);
+	if (!area->nia.freelist)
 		goto err;
 
 	area->user_refs = kvmalloc_objs(area->user_refs[0], nr_iovs,
@@ -562,7 +563,7 @@ static int __zcrx_create_area(struct io_zcrx_ifq *ifq,
 		struct net_iov *niov = &area->nia.niovs[i];
 
 		net_iov_init(niov, &area->nia, NET_IOV_IOURING);
-		area->freelist[i] = i;
+		area->nia.freelist[i] = i;
 		atomic_set(&area->user_refs[i], 0);
 	}
 
@@ -572,7 +573,7 @@ static int __zcrx_create_area(struct io_zcrx_ifq *ifq,
 			goto err;
 	}
 
-	area->free_count = nr_iovs;
+	area->nia.free_count = nr_iovs;
 	/* we're only supporting one area per ifq for now */
 	area->area_id = zcrx_next_area_id(ifq);
 	area_reg->rq_area_token = zcrx_area_id_to_token(area->area_id);
@@ -689,22 +690,14 @@ static void io_zcrx_return_niov_freelist(struct net_iov *niov)
 	struct io_zcrx_ifq *ifq = area->ifq;
 
 	guard(spinlock_bh)(&ifq->alloc_lock);
-	if (WARN_ON_ONCE(area->free_count >= area->nia.num_niovs))
-		return;
-	area->freelist[area->free_count++] = net_iov_idx(niov);
+	net_iov_area_push(&area->nia, niov);
 }
 
 static struct net_iov *zcrx_get_free_niov(struct io_zcrx_area *area)
 {
-	unsigned niov_idx;
-
 	lockdep_assert_held(&area->ifq->alloc_lock);
 
-	if (unlikely(!area->free_count))
-		return NULL;
-
-	niov_idx = area->freelist[--area->free_count];
-	return &area->nia.niovs[niov_idx];
+	return net_iov_area_pop(&area->nia);
 }
 
 static void io_zcrx_return_niov(struct net_iov *niov)
diff --git a/io_uring/zcrx.h b/io_uring/zcrx.h
index d4a54b4e17fd..a923291defde 100644
--- a/io_uring/zcrx.h
+++ b/io_uring/zcrx.h
@@ -36,10 +36,6 @@ struct io_zcrx_area {
 	bool			is_mapped;
 	u16			area_id;
 
-	/* freelist */
-	u32			free_count;
-	u32			*freelist;
-
 	struct io_zcrx_mem	mem;
 };
 
diff --git a/net/core/devmem.c b/net/core/devmem.c
index a9d86b5a5588..c1c1872b88de 100644
--- a/net/core/devmem.c
+++ b/net/core/devmem.c
@@ -42,10 +42,10 @@ void __net_devmem_dmabuf_binding_free(struct work_struct *wq)
 {
 	struct net_devmem_dmabuf_binding *binding = container_of(wq, typeof(*binding), unbind_w);
 
-	if (binding->freelist)
-		WARN(binding->free_count != binding->area.num_niovs,
-		     "destroying dmabuf binding with outstanding net_iovs: total=%zu, free=%zu",
-		     binding->area.num_niovs, binding->free_count);
+	if (binding->area.freelist)
+		WARN(binding->area.free_count != binding->area.num_niovs,
+		     "destroying dmabuf binding with outstanding net_iovs: total=%u, free=%u",
+		     binding->area.num_niovs, binding->area.free_count);
 
 	kvfree(binding->area.niovs);
 	dma_buf_unmap_attachment_unlocked(binding->attachment, binding->sgt,
@@ -54,7 +54,7 @@ void __net_devmem_dmabuf_binding_free(struct work_struct *wq)
 	dma_buf_put(binding->dmabuf);
 	xa_destroy(&binding->bound_rxqs);
 	percpu_ref_exit(&binding->ref);
-	kvfree(binding->freelist);
+	kvfree(binding->area.freelist);
 	kvfree(binding->tx_vec);
 	kfree(binding);
 }
@@ -67,12 +67,10 @@ net_devmem_alloc_dmabuf_bulk(struct net_devmem_dmabuf_binding *binding,
 
 	spin_lock_bh(&binding->freelist_lock);
 
-	count = min_t(size_t, count, binding->free_count);
-	for (i = 0; i < count; i++) {
-		struct net_iov *niov = binding->freelist[--binding->free_count];
-
-		netmems[i] = net_iov_to_netmem(niov);
-	}
+	count = min(count, binding->area.free_count);
+	for (i = 0; i < count; i++)
+		netmems[i] =
+			net_iov_to_netmem(net_iov_area_pop(&binding->area));
 
 	spin_unlock_bh(&binding->freelist_lock);
 
@@ -84,12 +82,7 @@ void net_devmem_free_dmabuf(struct net_iov *niov)
 	struct net_devmem_dmabuf_binding *binding = net_devmem_iov_binding(niov);
 
 	spin_lock_bh(&binding->freelist_lock);
-	if (WARN_ON_ONCE(binding->free_count >= binding->area.num_niovs)) {
-		spin_unlock_bh(&binding->freelist_lock);
-		return;
-	}
-
-	binding->freelist[binding->free_count++] = niov;
+	net_iov_area_push(&binding->area, niov);
 	spin_unlock_bh(&binding->freelist_lock);
 }
 
@@ -228,6 +221,12 @@ net_devmem_bind_dmabuf(struct net_device *dev, void *vdev,
 		goto err_unmap;
 	}
 
+	if ((dmabuf->size >> niov_shift) > U32_MAX) {
+		err = -E2BIG;
+		NL_SET_ERR_MSG(extack, "dmabuf contains too many net_iovs");
+		goto err_unmap;
+	}
+
 	binding->area.base_virtual = 0;
 	binding->area.num_niovs = dmabuf->size >> niov_shift;
 	if (direction == DMA_TO_DEVICE) {
@@ -239,10 +238,11 @@ net_devmem_bind_dmabuf(struct net_device *dev, void *vdev,
 		}
 	} else {
 		spin_lock_init(&binding->freelist_lock);
-		binding->freelist = kvmalloc_array(binding->area.num_niovs,
-						   sizeof(binding->freelist[0]),
-						   GFP_KERNEL);
-		if (!binding->freelist) {
+		binding->area.freelist =
+			kvmalloc_array(binding->area.num_niovs,
+				       sizeof(binding->area.freelist[0]),
+				       GFP_KERNEL);
+		if (!binding->area.freelist) {
 			err = -ENOMEM;
 			goto err_unmap;
 		}
@@ -279,7 +279,7 @@ net_devmem_bind_dmabuf(struct net_device *dev, void *vdev,
 			if (direction == DMA_TO_DEVICE)
 				binding->tx_vec[niov_idx] = niov;
 			else
-				binding->freelist[binding->free_count++] = niov;
+				net_iov_area_push(&binding->area, niov);
 			dma_addr += niov_size;
 		}
 	}
@@ -297,7 +297,7 @@ net_devmem_bind_dmabuf(struct net_device *dev, void *vdev,
 err_free_niovs:
 	kvfree(binding->area.niovs);
 err_free_freelist:
-	kvfree(binding->freelist);
+	kvfree(binding->area.freelist);
 	kvfree(binding->tx_vec);
 err_unmap:
 	dma_buf_unmap_attachment_unlocked(binding->attachment, binding->sgt,
diff --git a/net/core/devmem.h b/net/core/devmem.h
index f71d9a2bcb3d..5edb763b89b2 100644
--- a/net/core/devmem.h
+++ b/net/core/devmem.h
@@ -17,6 +17,9 @@ struct netlink_ext_ack;
 
 struct net_devmem_dmabuf_binding {
 	struct net_iov_area area;
+	/* Protects area.freelist and area.free_count. */
+	spinlock_t freelist_lock;
+
 	struct dma_buf *dmabuf;
 	struct dma_buf_attachment *attachment;
 	struct sg_table *sgt;
@@ -57,10 +60,6 @@ struct net_devmem_dmabuf_binding {
 	/* rxq's this binding is active on. */
 	struct xarray bound_rxqs;
 
-	spinlock_t freelist_lock ____cacheline_aligned_in_smp;
-	size_t free_count;
-	struct net_iov **freelist;
-
 	/* ID of this binding. Globally unique to all bindings currently
 	 * active.
 	 */
-- 
2.53.0-Meta