[PATCH RFC 14/20] mempool: introduce mempool_use_prealloc_only

[andrey.konovalov@linux.dev]

From: Andrey Konovalov <andreyknvl@google.com>

Introduce a new mempool_use_prealloc_only API that tells the mempool to
only use the elements preallocated during the mempool's creation and to
not attempt allocating new ones.

This API is required to test the KASAN poisoning/unpoisoning functinality
in KASAN tests, but it might be also useful on its own.

Signed-off-by: Andrey Konovalov <andreyknvl@google.com>
---
 include/linux/mempool.h |  2 ++
 mm/mempool.c            | 27 ++++++++++++++++++++++++---
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/include/linux/mempool.h b/include/linux/mempool.h
index 4aae6c06c5f2..822adf1e7567 100644
--- a/include/linux/mempool.h
+++ b/include/linux/mempool.h
@@ -18,6 +18,7 @@ typedef struct mempool_s {
 	int min_nr;		/* nr of elements at *elements */
 	int curr_nr;		/* Current nr of elements at *elements */
 	void **elements;
+	bool use_prealloc_only;	/* Use only preallocated elements */
 
 	void *pool_data;
 	mempool_alloc_t *alloc;
@@ -48,6 +49,7 @@ extern mempool_t *mempool_create_node(int min_nr, mempool_alloc_t *alloc_fn,
 			mempool_free_t *free_fn, void *pool_data,
 			gfp_t gfp_mask, int nid);
 
+extern void mempool_use_prealloc_only(mempool_t *pool);
 extern int mempool_resize(mempool_t *pool, int new_min_nr);
 extern void mempool_destroy(mempool_t *pool);
 extern void *mempool_alloc(mempool_t *pool, gfp_t gfp_mask) __malloc;
diff --git a/mm/mempool.c b/mm/mempool.c
index f67ca6753332..59f7fcd355b3 100644
--- a/mm/mempool.c
+++ b/mm/mempool.c
@@ -365,6 +365,20 @@ int mempool_resize(mempool_t *pool, int new_min_nr)
 }
 EXPORT_SYMBOL(mempool_resize);
 
+/**
+ * mempool_use_prealloc_only - mark a pool to only use preallocated elements
+ * @pool:      pointer to the memory pool that should be marked
+ *
+ * This function should only be called right after the pool creation via
+ * mempool_init() or mempool_create() and must not be called concurrently with
+ * mempool_alloc().
+ */
+void mempool_use_prealloc_only(mempool_t *pool)
+{
+	pool->use_prealloc_only = true;
+}
+EXPORT_SYMBOL(mempool_use_prealloc_only);
+
 /**
  * mempool_alloc - allocate an element from a specific memory pool
  * @pool:      pointer to the memory pool which was allocated via
@@ -397,9 +411,11 @@ void *mempool_alloc(mempool_t *pool, gfp_t gfp_mask)
 
 repeat_alloc:
 
-	element = pool->alloc(gfp_temp, pool->pool_data);
-	if (likely(element != NULL))
-		return element;
+	if (!pool->use_prealloc_only) {
+		element = pool->alloc(gfp_temp, pool->pool_data);
+		if (likely(element != NULL))
+			return element;
+	}
 
 	spin_lock_irqsave(&pool->lock, flags);
 	if (likely(pool->curr_nr)) {
@@ -415,6 +431,11 @@ void *mempool_alloc(mempool_t *pool, gfp_t gfp_mask)
 		return element;
 	}
 
+	if (pool->use_prealloc_only) {
+		spin_unlock_irqrestore(&pool->lock, flags);
+		return NULL;
+	}
+
 	/*
 	 * We use gfp mask w/o direct reclaim or IO for the first round.  If
 	 * alloc failed with that and @pool was empty, retry immediately.
-- 
2.25.1