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

On Mon, Nov 06, 2023 at 09:10PM +0100, andrey.konovalov@linux.dev wrote:
> 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 */

This increases the struct size from 56 to 64 bytes (64 bit arch).
mempool_t is embedded in lots of other larger structs, and this may
result in some unwanted bloat.

Is there a way to achieve the same thing without adding a new bool to
the mempool struct?

It seems a little excessive only for the purpose of the tests.

On Wed, Nov 22, 2023 at 6:21 PM Marco Elver <elver@google.com> wrote:
>
> On Mon, Nov 06, 2023 at 09:10PM +0100, andrey.konovalov@linux.dev wrote:
> > 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 */
>
> This increases the struct size from 56 to 64 bytes (64 bit arch).
> mempool_t is embedded in lots of other larger structs, and this may
> result in some unwanted bloat.
>
> Is there a way to achieve the same thing without adding a new bool to
> the mempool struct?

We could split out the part of mempool_alloc that uses preallocated
elements without what waiting part and expose it in another API
function named something like mempool_alloc_preallocated. Would that
be better?

On Thu, 23 Nov 2023 at 19:06, Andrey Konovalov <andreyknvl@gmail.com> wrote:
>
> On Wed, Nov 22, 2023 at 6:21 PM Marco Elver <elver@google.com> wrote:
> >
> > On Mon, Nov 06, 2023 at 09:10PM +0100, andrey.konovalov@linux.dev wrote:
> > > 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 */
> >
> > This increases the struct size from 56 to 64 bytes (64 bit arch).
> > mempool_t is embedded in lots of other larger structs, and this may
> > result in some unwanted bloat.
> >
> > Is there a way to achieve the same thing without adding a new bool to
> > the mempool struct?
>
> We could split out the part of mempool_alloc that uses preallocated
> elements without what waiting part and expose it in another API
> function named something like mempool_alloc_preallocated. Would that
> be better?

Yes, that might be better. As long as other users of mempool (esp if
KASAN is disabled) are unaffected then it should be fine.