[RFC PATCH 3/3] io_uring/rsrc: prefill the node cache when a file table is registered empty

[Uzair Beg <uzairbeg11@gmail.com>]

Registering a sparse fixed file table allocates no nodes at
registration time; each node is allocated later, on the install path,
where MSG_RING SEND_FD pays for it. Bare-metal measurement of the
4,096-slot first fill shows the cost is not the allocator call
(bulk refill was neutral) nor fresh slab pages (priming the slab was
neutral), but the per-object SLUB allocation path itself. The only way
to take it off the install path is to not allocate there.

When a sparse table of N slots is registered, grow the per-ring node
cache to min(N, IO_ALLOC_CACHE_PREFILL_MAX) and bulk-fill it, so the
subsequent installs hit the cache. Prefill is best-effort: on any
failure the cache is left in a valid state (a successfully grown
pointer array is retained) and registration proceeds unchanged.
Non-sparse registrations are untouched, since they allocate every node
inline anyway.

On the reported 4,096-slot first fill this is 9.8% faster than
unpatched. Because the enlarged cache also retains nodes released by
FILES_UPDATE, a same-ring remove-and-refill of 4,096 files is 18.8%
faster. The cost is moved to registration rather than removed: a
one-shot register-then-fill is unchanged overall, and a program that
registers many slots and installs few pays for nodes it never uses.
Whether that trade is acceptable, or should be behind a registration
flag, is the question this patch is intended to raise.

The stash loop from the bulk refill path is factored into a helper so
both callers share it.

Reported-by: Chengfeng Lin <lin2530632123@gmail.com>
Closes: https://lore.kernel.org/io-uring/CANGjgdmt0FQ=offsdfn+wEaDxbOFoAa6bi92X_vEo4S6aCZ56A@mail.gmail.com/
Tested-by: Chengfeng Lin <lin2530632123@gmail.com>
Co-developed-by: Chengfeng Lin <lin2530632123@gmail.com>
Signed-off-by: Chengfeng Lin <lin2530632123@gmail.com>
Signed-off-by: Uzair Beg <uzairbeg11@gmail.com>
---
 io_uring/alloc_cache.c | 58 ++++++++++++++++++++++++++++++++++--------
 io_uring/alloc_cache.h |  2 ++
 io_uring/rsrc.c        |  4 +++
 3 files changed, 54 insertions(+), 10 deletions(-)

diff --git a/io_uring/alloc_cache.c b/io_uring/alloc_cache.c
index cba0e6c5d66..2c6e09313d2 100644
--- a/io_uring/alloc_cache.c
+++ b/io_uring/alloc_cache.c
@@ -38,6 +38,22 @@ bool io_alloc_cache_init(struct io_alloc_cache *cache,
 	return false;
 }
 
+static void io_cache_stash(struct io_alloc_cache *cache, void **slot,
+			   unsigned int nr)
+{
+	unsigned int i;
+
+	for (i = 0; i < nr; i++) {
+		if (cache->init_clear)
+			memset(slot[i], 0, cache->init_clear);
+		if (unlikely(!kasan_mempool_poison_object(slot[i])))
+			break;
+		cache->nr_cached++;
+	}
+	for (; i < nr; i++)
+		kmem_cache_free(cache->slab, slot[i]);
+}
+
 void *io_cache_alloc_new(struct io_alloc_cache *cache, gfp_t gfp)
 {
 	void *obj;
@@ -45,7 +61,7 @@ void *io_cache_alloc_new(struct io_alloc_cache *cache, gfp_t gfp)
 	if (cache->slab) {
 		unsigned int room = cache->max_cached - cache->nr_cached;
 		void **slot = &cache->entries[cache->nr_cached];
-		unsigned int batch, got, i;
+		unsigned int batch, got;
 
 		if (unlikely(!room))
 			return kmem_cache_alloc(cache->slab, gfp);
@@ -57,15 +73,7 @@ void *io_cache_alloc_new(struct io_alloc_cache *cache, gfp_t gfp)
 
 		/* return one object, stash the rest in the cache */
 		obj = slot[got - 1];
-		for (i = 0; i < got - 1; i++) {
-			if (cache->init_clear)
-				memset(slot[i], 0, cache->init_clear);
-			if (unlikely(!kasan_mempool_poison_object(slot[i])))
-				break;
-			cache->nr_cached++;
-		}
-		for (; i < got - 1; i++)
-			kmem_cache_free(cache->slab, slot[i]);
+		io_cache_stash(cache, slot, got - 1);
 	} else {
 		obj = kmalloc(cache->elem_size, gfp);
 	}
@@ -73,3 +81,33 @@ void *io_cache_alloc_new(struct io_alloc_cache *cache, gfp_t gfp)
 		memset(obj, 0, cache->init_clear);
 	return obj;
 }
+
+void io_alloc_cache_prefill(struct io_alloc_cache *cache, unsigned int nr)
+{
+	gfp_t gfp = GFP_KERNEL | __GFP_NOWARN;
+	unsigned int got;
+	void **entries;
+
+	if (!cache->slab || !cache->entries)
+		return;
+
+	nr = min_t(unsigned int, nr, IO_ALLOC_CACHE_PREFILL_MAX);
+	if (nr <= cache->nr_cached)
+		return;
+
+	if (nr > cache->max_cached) {
+		entries = kvmalloc_array(nr, sizeof(void *), gfp);
+		if (!entries)
+			return;
+		memcpy(entries, cache->entries,
+		       cache->nr_cached * sizeof(void *));
+		kvfree(cache->entries);
+		cache->entries = entries;
+		cache->max_cached = nr;
+	}
+
+	got = kmem_cache_alloc_bulk(cache->slab, gfp, nr - cache->nr_cached,
+				    &cache->entries[cache->nr_cached]);
+	if (got)
+		io_cache_stash(cache, &cache->entries[cache->nr_cached], got);
+}
diff --git a/io_uring/alloc_cache.h b/io_uring/alloc_cache.h
index 82d552c7517..ca6af52dd7d 100644
--- a/io_uring/alloc_cache.h
+++ b/io_uring/alloc_cache.h
@@ -8,6 +8,7 @@
  */
 #define IO_ALLOC_CACHE_MAX	128
 #define IO_ALLOC_CACHE_REFILL	32
+#define IO_ALLOC_CACHE_PREFILL_MAX	4096
 
 void io_alloc_cache_free(struct io_alloc_cache *cache,
 			 void (*free)(const void *));
@@ -16,6 +17,7 @@ bool io_alloc_cache_init(struct io_alloc_cache *cache,
 			 unsigned int init_bytes);
 
 void *io_cache_alloc_new(struct io_alloc_cache *cache, gfp_t gfp);
+void io_alloc_cache_prefill(struct io_alloc_cache *cache, unsigned int nr);
 
 static inline bool io_alloc_cache_put(struct io_alloc_cache *cache,
 				      void *entry)
diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index 6413682ebe4..b7f78b6b514 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -560,6 +560,10 @@ int io_sqe_files_register(struct io_ring_ctx *ctx, void __user *arg,
 	if (!io_alloc_file_tables(ctx, &ctx->file_table, nr_args))
 		return -ENOMEM;
 
+	/* sparse table: nodes are installed later, so cache them now */
+	if (!fds)
+		io_alloc_cache_prefill(&ctx->node_cache, nr_args);
+
 	for (i = 0; i < nr_args; i++) {
 		struct io_rsrc_node *node;
 		u64 tag = 0;
-- 
2.43.0