[PATCH 2/3] mm: support fallible mempool_alloc_bulk()

[Eric Biggers <ebiggers@kernel.org>]

To fix a deadlock, blk-crypto-fallback needs to be able to make fallible
mempool_alloc_bulk() allocations.  But mempool_alloc_bulk() hardcodes
GFP_KERNEL and infinite retries, which differs from mempool_alloc()
which supports fallible allocations via its gfp_mask argument.

Therefore, add a gfp_mask argument to mempool_alloc_bulk().  As with
mempool_alloc(), the presence of __GFP_DIRECT_RECLAIM in the mask
selects between the fallible and infallible modes.

For now it just provides all-or-nothing semantics and returns a bool,
similar to kmem_cache_alloc_bulk().

Signed-off-by: Eric Biggers <ebiggers@kernel.org>
---
 block/blk-crypto-fallback.c |  6 ++----
 include/linux/mempool.h     |  4 ++--
 mm/mempool.c                | 42 +++++++++++++++++++++++++------------
 3 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/block/blk-crypto-fallback.c b/block/blk-crypto-fallback.c
index 2a5c52ab74b4..bda913c39381 100644
--- a/block/blk-crypto-fallback.c
+++ b/block/blk-crypto-fallback.c
@@ -172,7 +172,6 @@ static void blk_crypto_fallback_encrypt_endio(struct bio *enc_bio)
 static struct bio *blk_crypto_alloc_enc_bio(struct bio *bio_src,
 		unsigned int nr_segs, struct page ***pages_ret)
 {
-	unsigned int memflags = memalloc_noio_save();
 	unsigned int nr_allocated;
 	struct page **pages;
 	struct bio *bio;
@@ -206,12 +205,11 @@ static struct bio *blk_crypto_alloc_enc_bio(struct bio *bio_src,
 	 * any non-zero slot already contains a valid allocation.
 	 */
 	memset(pages, 0, sizeof(struct page *) * nr_segs);
-	nr_allocated = alloc_pages_bulk(GFP_KERNEL, nr_segs, pages);
+	nr_allocated = alloc_pages_bulk(GFP_NOIO, nr_segs, pages);
 	if (nr_allocated < nr_segs)
 		mempool_alloc_bulk(blk_crypto_bounce_page_pool,
 				(void **)pages + nr_allocated,
-				nr_segs - nr_allocated);
-	memalloc_noio_restore(memflags);
+				nr_segs - nr_allocated, GFP_NOIO);
 	*pages_ret = pages;
 	return bio;
 }
diff --git a/include/linux/mempool.h b/include/linux/mempool.h
index a0fa6d43e0dc..f7898cd1512b 100644
--- a/include/linux/mempool.h
+++ b/include/linux/mempool.h
@@ -65,8 +65,8 @@ void mempool_destroy(struct mempool *pool);
 void *mempool_alloc_noprof(struct mempool *pool, gfp_t gfp_mask) __malloc;
 #define mempool_alloc(...)						\
 	alloc_hooks(mempool_alloc_noprof(__VA_ARGS__))
-int mempool_alloc_bulk_noprof(struct mempool *pool, void **elem,
-		unsigned int count);
+bool mempool_alloc_bulk_noprof(struct mempool *pool, void **elem,
+		unsigned int count, gfp_t gfp_mask);
 #define mempool_alloc_bulk(...)						\
 	alloc_hooks(mempool_alloc_bulk_noprof(__VA_ARGS__))
 
diff --git a/mm/mempool.c b/mm/mempool.c
index d454bc9f39e9..d741e5f62554 100644
--- a/mm/mempool.c
+++ b/mm/mempool.c
@@ -473,25 +473,28 @@ static inline gfp_t mempool_adjust_gfp(gfp_t *gfp_mask)
 /**
  * mempool_alloc_bulk - allocate multiple elements from a memory pool
  * @pool:	pointer to the memory pool
- * @elems:	partially or fully populated elements array
- * @count:	number of entries in @elem that need to be allocated
+ * @elems:	pointer to array into which the element pointers will be stored
+ * @count:	number of elements to allocate
+ * @gfp_mask:	GFP_* flags.  %__GFP_ZERO is not supported.  If this mask
+ *		includes %__GFP_DIRECT_RECLAIM, then the allocation is retried
+ *		indefinitely until it succeeds and the return value is always
+ *		%true.  If the mask doesn't include %__GFP_DIRECT_RECLAIM, then
+ *		failure is allowed and %false can be returned.
  *
  * Allocate @count elements into @elems.  This is done by first calling into the
  * alloc_fn supplied at pool initialization time, and dipping into the reserved
- * pool when alloc_fn fails to allocate an element.
- *
- * On return all @count elements in @elems will be populated.
+ * pool to atomically allocate the remaining elements if alloc_fn fails.
  *
- * Return: Always 0.  If it wasn't for %$#^$ alloc tags, it would return void.
+ * Return: %true if the allocation succeeded, or %false if it failed.
  */
-int mempool_alloc_bulk_noprof(struct mempool *pool, void **elems,
-		unsigned int count)
+bool mempool_alloc_bulk_noprof(struct mempool *pool, void **elems,
+		unsigned int count, gfp_t gfp_mask)
 {
-	gfp_t gfp_mask = GFP_KERNEL;
 	gfp_t gfp_temp = mempool_adjust_gfp(&gfp_mask);
 	unsigned int allocated = 0;
 
 	VM_WARN_ON_ONCE(count > pool->min_nr);
+	VM_WARN_ON_ONCE(gfp_mask & __GFP_ZERO);
 	might_alloc(gfp_mask);
 
 	/*
@@ -516,13 +519,26 @@ int mempool_alloc_bulk_noprof(struct mempool *pool, void **elems,
 		allocated++;
 	}
 
-	return 0;
+	return true;
 
 use_pool:
+	/* Try to atomically allocate the remaining elements from the pool. */
 	if (mempool_alloc_from_pool(pool, elems, count, allocated, gfp_temp))
-		return 0;
-	gfp_temp = gfp_mask;
-	goto repeat_alloc;
+		return true;
+	/* Retry if this was just the opportunistic first pass. */
+	if (gfp_temp != gfp_mask) {
+		gfp_temp = gfp_mask;
+		goto repeat_alloc;
+	}
+	/* Retry indefinitely if __GFP_DIRECT_RECLAIM is set. */
+	if (gfp_mask & __GFP_DIRECT_RECLAIM)
+		goto repeat_alloc;
+	/* On failure, roll back any successful allocations from ->alloc(). */
+	while (allocated--) {
+		pool->free(elems[allocated], pool->pool_data);
+		elems[allocated] = NULL;
+	}
+	return false;
 }
 EXPORT_SYMBOL_GPL(mempool_alloc_bulk_noprof);
 
-- 
2.55.0