[PATCH 2/3] mm/page_ext: support early allocation for SPARSEMEM

[Muchun Song <songmuchun@bytedance.com>]

Moving KMSAN's per-page metadata into page_ext requires page extensions
to be available when kmsan_init_shadow() runs, before memblock is freed.
FLATMEM page extensions are available by then, but SPARSEMEM normally
allocates them later from the page or vmalloc allocators.

Add need_early so clients can request SPARSEMEM section backing from
memblock during page_ext_init_early(). Use page_ext_init() as the common
late path for client callbacks.

KMSAN starts requesting early page extensions in the following patch.

Signed-off-by: Muchun Song <songmuchun@bytedance.com>
---
 include/linux/page_ext.h | 26 +++--------
 mm/mm_init.c             | 11 +----
 mm/page_ext.c            | 99 ++++++++++++++++++++++++++++++----------
 3 files changed, 82 insertions(+), 54 deletions(-)

diff --git a/include/linux/page_ext.h b/include/linux/page_ext.h
index 79c53ec45dfa..3b66720ddb18 100644
--- a/include/linux/page_ext.h
+++ b/include/linux/page_ext.h
@@ -16,6 +16,8 @@ struct pglist_data;
  * @size: The size of the client data within page_ext.
  * @need: Function that returns true if client requires page_ext.
  * @init: (optional) Called to initialize client once page_exts are allocated.
+ * @need_early: True if the client requires page_ext during early boot,
+ *              before normal memory allocation is available.
  * @need_shared_flags: True when client is using shared page_ext->flags
  *                     field.
  *
@@ -27,6 +29,7 @@ struct page_ext_operations {
 	size_t size;
 	bool (*need)(void);
 	void (*init)(void);
+	bool need_early;
 	bool need_shared_flags;
 };
 
@@ -61,15 +64,10 @@ static inline bool early_page_ext_enabled(void)
 	return early_page_ext;
 }
 
-#ifdef CONFIG_SPARSEMEM
-static inline void page_ext_init_flatmem(void)
-{
-}
-extern void page_ext_init(void);
-static inline void page_ext_init_flatmem_late(void)
-{
-}
+void page_ext_init_early(void);
+void page_ext_init(void);
 
+#ifdef CONFIG_SPARSEMEM
 static inline bool page_ext_iter_next_fast_possible(unsigned long next_pfn)
 {
 	/*
@@ -79,12 +77,6 @@ static inline bool page_ext_iter_next_fast_possible(unsigned long next_pfn)
 	return next_pfn % PAGES_PER_SECTION;
 }
 #else
-extern void page_ext_init_flatmem(void);
-extern void page_ext_init_flatmem_late(void);
-static inline void page_ext_init(void)
-{
-}
-
 static inline bool page_ext_iter_next_fast_possible(unsigned long next_pfn)
 {
 	return true;
@@ -205,11 +197,7 @@ static inline void page_ext_init(void)
 {
 }
 
-static inline void page_ext_init_flatmem_late(void)
-{
-}
-
-static inline void page_ext_init_flatmem(void)
+static inline void page_ext_init_early(void)
 {
 }
 
diff --git a/mm/mm_init.c b/mm/mm_init.c
index 2367af05570b..35cf799971c3 100644
--- a/mm/mm_init.c
+++ b/mm/mm_init.c
@@ -2707,11 +2707,7 @@ void __init mm_core_init(void)
 	build_all_zonelists(NULL);
 	page_alloc_init_cpuhp();
 	alloc_tag_sec_init();
-	/*
-	 * page_ext requires contiguous pages,
-	 * bigger than MAX_PAGE_ORDER unless SPARSEMEM.
-	 */
-	page_ext_init_flatmem();
+	page_ext_init_early();
 	mem_debugging_and_hardening_init();
 	kfence_alloc_pool_and_metadata();
 	report_meminit();
@@ -2727,11 +2723,6 @@ void __init mm_core_init(void)
 	memblock_free_all();
 	mem_init();
 	kmem_cache_init();
-	/*
-	 * page_owner must be initialized after buddy is ready, and also after
-	 * slab is ready so that stack_depot_init() works properly
-	 */
-	page_ext_init_flatmem_late();
 	kmemleak_init();
 	ptlock_cache_init();
 	pgtable_cache_init();
diff --git a/mm/page_ext.c b/mm/page_ext.c
index b679a8c1f7d7..8578f52a32db 100644
--- a/mm/page_ext.c
+++ b/mm/page_ext.c
@@ -12,6 +12,7 @@
 #include <linux/rcupdate.h>
 #include <linux/pgalloc_tag.h>
 #include <linux/iommu-debug-pagealloc.h>
+#include <linux/slab.h>
 
 /*
  * struct page extension
@@ -135,7 +136,12 @@ static bool __init invoke_need_callbacks(void)
 		if (page_ext_ops[i]->need()) {
 			page_ext_ops[i]->offset = page_ext_size;
 			page_ext_size += page_ext_ops[i]->size;
-			need = true;
+			/*
+			 * page_ext requires contiguous pages, bigger than
+			 * MAX_PAGE_ORDER unless SPARSEMEM.
+			 */
+			need |= !IS_ENABLED(CONFIG_SPARSEMEM) ||
+				page_ext_ops[i]->need_early;
 		}
 	}
 
@@ -159,7 +165,7 @@ static inline struct page_ext *get_entry(void *base, unsigned long index)
 }
 
 #ifndef CONFIG_SPARSEMEM
-void __init page_ext_init_flatmem_late(void)
+void __init page_ext_init(void)
 {
 	invoke_init_callbacks();
 }
@@ -217,9 +223,8 @@ static int __init alloc_node_page_ext(int nid)
 	return 0;
 }
 
-void __init page_ext_init_flatmem(void)
+void __init page_ext_init_early(void)
 {
-
 	int nid, fail;
 
 	if (!invoke_need_callbacks())
@@ -262,11 +267,19 @@ static struct page_ext *lookup_page_ext(const struct page *page)
 	return get_entry(page_ext, pfn);
 }
 
-static void *__meminit alloc_page_ext(size_t size, int nid)
+static __ref void *alloc_page_ext(size_t size, int nid)
 {
 	gfp_t flags = GFP_KERNEL | __GFP_ZERO | __GFP_NOWARN;
 	void *addr = NULL;
 
+	if (!slab_is_available()) {
+		addr = memblock_alloc_try_nid(size, PAGE_SIZE, 0,
+					      MEMBLOCK_ALLOC_ACCESSIBLE, nid);
+		if (addr)
+			memmap_boot_pages_add(size / PAGE_SIZE);
+		return addr;
+	}
+
 	addr = alloc_pages_exact_nid(nid, size, flags);
 	if (addr)
 		kmemleak_alloc(addr, size, 1, flags);
@@ -274,7 +287,7 @@ static void *__meminit alloc_page_ext(size_t size, int nid)
 		addr = vzalloc_node(size, nid);
 
 	if (addr)
-		memmap_pages_add(DIV_ROUND_UP(size, PAGE_SIZE));
+		memmap_pages_add(size / PAGE_SIZE);
 
 	return addr;
 }
@@ -290,7 +303,7 @@ static int __meminit init_section_page_ext(unsigned long pfn, int nid)
 	if (section->page_ext)
 		return 0;
 
-	table_size = page_ext_size * PAGES_PER_SECTION;
+	table_size = PAGE_ALIGN(page_ext_size * PAGES_PER_SECTION);
 	base = alloc_page_ext(table_size, nid);
 
 	/*
@@ -317,20 +330,27 @@ static int __meminit init_section_page_ext(unsigned long pfn, int nid)
 
 static void free_page_ext(void *addr)
 {
-	size_t table_size;
+	size_t table_size = PAGE_ALIGN(page_ext_size * PAGES_PER_SECTION);
+	long nr_pages = table_size / PAGE_SIZE;
 	struct page *page;
 
-	table_size = page_ext_size * PAGES_PER_SECTION;
-	memmap_pages_add(-1L * (DIV_ROUND_UP(table_size, PAGE_SIZE)));
-
 	if (is_vmalloc_addr(addr)) {
+		memmap_pages_add(-nr_pages);
 		vfree(addr);
-	} else {
-		page = virt_to_page(addr);
-		BUG_ON(PageReserved(page));
-		kmemleak_free(addr);
-		free_pages_exact(addr, table_size);
+		return;
+	}
+
+	page = virt_to_page(addr);
+	if (PageReserved(page)) {
+		memmap_boot_pages_add(-nr_pages);
+		kmemleak_free_part_phys(__pa(addr), table_size);
+		free_reserved_area(addr, addr + table_size, -1, "page_ext");
+		return;
 	}
+
+	memmap_pages_add(-nr_pages);
+	kmemleak_free(addr);
+	free_pages_exact(addr, table_size);
 }
 
 static void __free_page_ext(unsigned long pfn)
@@ -444,14 +464,28 @@ static int __meminit page_ext_callback(struct notifier_block *self,
 	return notifier_from_errno(ret);
 }
 
-void __init page_ext_init(void)
+static int __init page_ext_init_sections_early(void)
+{
+	unsigned long start, end, pfn;
+	int i, nid;
+
+	for_each_mem_pfn_range(i, MAX_NUMNODES, &start, &end, &nid) {
+		for (pfn = start; pfn < end; pfn = SECTION_ALIGN_UP(pfn + 1)) {
+			if (!pfn_valid(pfn))
+				continue;
+			if (init_section_page_ext(pfn, nid))
+				return -ENOMEM;
+		}
+	}
+
+	return 0;
+}
+
+static int __init page_ext_init_sections(void)
 {
 	unsigned long pfn;
 	int nid;
 
-	if (!invoke_need_callbacks())
-		return;
-
 	for_each_node_state(nid, N_MEMORY) {
 		unsigned long start_pfn, end_pfn;
 
@@ -476,17 +510,32 @@ void __init page_ext_init(void)
 			if (pfn_to_nid(pfn) != nid)
 				continue;
 			if (init_section_page_ext(pfn, nid))
-				goto oom;
+				return -ENOMEM;
 			cond_resched();
 		}
 	}
+
+	return 0;
+}
+
+void __init page_ext_init_early(void)
+{
+	if (!invoke_need_callbacks())
+		return;
+	if (page_ext_init_sections_early())
+		panic("Out of memory");
+}
+
+void __init page_ext_init(void)
+{
+	if (!page_ext_size)
+		return;
+	if (page_ext_init_sections())
+		panic("Out of memory");
+
 	hotplug_memory_notifier(page_ext_callback, DEFAULT_CALLBACK_PRI);
 	pr_info("allocated %ld bytes of page_ext\n", total_usage);
 	invoke_init_callbacks();
-	return;
-
-oom:
-	panic("Out of memory");
 }
 
 #endif
-- 
2.54.0