This flag is set unless we can be sure the caller isn't in an atomic
context.

The allocator will soon start needing to call set_direct_map_* APIs
which cannot be called with IRQs off. It will need to do this even
before direct reclaim is possible.

Despite the fact that, in principle, ALLOC_NOBLOCK is distinct from
__GFP_DIRECT_RECLAIM, in order to avoid introducing a GFP flag, just
infer the former based on whether the caller set the latter. This means
that, in practice, ALLOC_NOBLOCK is just !__GFP_DIRECT_RECLAIM, except
that it is not influenced by gfp_allowed_mask. This requires some rather
ugly plumbing to get it into the slowpath without clobbering it.

Call it ALLOC_NOBLOCK in order to try and mitigate confusion vs the
recently-removed ALLOC_NON_BLOCK, which meant something different.

Signed-off-by: Brendan Jackman <jackmanb@google.com>
---
 mm/page_alloc.c | 23 +++++++++++++++++++----
 mm/page_alloc.h |  1 +
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/mm/page_alloc.c b/mm/page_alloc.c
index fb522a09f2e60..ecdd3780192e8 100644
--- a/mm/page_alloc.c
+++ b/mm/page_alloc.c
@@ -5286,6 +5286,19 @@ __alloc_pages_slowpath(gfp_t gfp_mask, unsigned int order,
 	return page;
 }
 
+/* Set alloc flags before applying gfp_allowed_mask. */
+static inline unsigned int init_alloc_flags(gfp_t gfp_mask, unsigned int alloc_flags)
+{
+	/*
+	 * If the caller allowed __GFP_DIRECT_RECLAIM, they can't be atomic.
+	 * Note this is a separate determination from whether direct reclaim is
+	 * actually allowed, it must happen before applying gfp_allowed_mask.
+	 */
+	if (!(gfp_mask & __GFP_DIRECT_RECLAIM))
+		alloc_flags |= ALLOC_NOBLOCK;
+	return alloc_flags;
+}
+
 static inline bool prepare_alloc_pages(gfp_t gfp_mask, unsigned int order,
 		int preferred_nid, nodemask_t *nodemask,
 		struct alloc_context *ac, gfp_t *alloc_gfp,
@@ -5364,8 +5377,10 @@ unsigned long alloc_pages_bulk_noprof(gfp_t gfp, int preferred_nid,
 	struct zoneref *z;
 	struct per_cpu_pages *pcp;
 	struct list_head *pcp_list;
-	struct alloc_context ac;
-	unsigned int alloc_flags = ALLOC_WMARK_LOW;
+	struct alloc_context ac = {
+		.alloc_flags = init_alloc_flags(gfp, ALLOC_DEFAULT),
+	};
+	unsigned int alloc_flags = ac.alloc_flags | ALLOC_WMARK_LOW;
 	int nr_populated = 0, nr_account = 0;
 
 	/*
@@ -5584,9 +5599,9 @@ struct page *__alloc_frozen_pages_noprof(gfp_t gfp, unsigned int order,
 	struct page *page;
 	gfp_t alloc_gfp; /* The gfp_t that was actually used for allocation */
 	struct alloc_context ac = {
-		.alloc_flags = alloc_flags,
+		.alloc_flags = init_alloc_flags(gfp, alloc_flags),
 	};
-	unsigned int fastpath_alloc_flags = alloc_flags;
+	unsigned int fastpath_alloc_flags = ac.alloc_flags;
 
 	/* Other flags could be supported later if needed. */
 	if (WARN_ON(alloc_flags & ~(ALLOC_NOLOCK | ALLOC_NO_CODETAG)))
diff --git a/mm/page_alloc.h b/mm/page_alloc.h
index 2307b173a459f..02db12c9a1dc2 100644
--- a/mm/page_alloc.h
+++ b/mm/page_alloc.h
@@ -75,6 +75,7 @@
  */
 #define ALLOC_UNMAPPED	       0x2000
 #endif
+#define ALLOC_NOBLOCK	       0x4000 /* Caller may be atomic */
 
 /* Flags that allow allocations below the min watermark. */
 #define ALLOC_RESERVES (ALLOC_HARDER|ALLOC_MIN_RESERVE|ALLOC_HIGHATOMIC|ALLOC_OOM)

-- 
2.54.0