[PATCH] fixup! memblock: add KHO support for reserve_mem

[Changyuan Lyu <changyuanl@xxxxxxxxxx>]

This patch includes the suggested changes from
https://lore.kernel.org/lkml/aAeaJ2iqkrv_ffhT@xxxxxxxxxx/ and can be
squashed with "memblock: add KHO support for reserve_mem".

Fixes: 2e257a656639 ("memblock: add KHO support for reserve_mem")
Suggested-by: Mike Rapoport (Microsoft) <rppt@xxxxxxxxxx>
Signed-off-by: Changyuan Lyu <changyuanl@xxxxxxxxxx>
---
 mm/memblock.c | 69 +++++++++++++++++++++++++--------------------------
 1 file changed, 34 insertions(+), 35 deletions(-)

diff --git a/mm/memblock.c b/mm/memblock.c
index 3571a859f2fe1..d38a72f07ea0c 100644
--- a/mm/memblock.c
+++ b/mm/memblock.c
@@ -2490,15 +2490,6 @@ static int reserve_mem_kho_finalize(struct kho_serialization *ser)
 {
 	int err = 0, i;

-	if (!reserved_mem_count)
-		return NOTIFY_DONE;
-
-	if (IS_ERR(kho_fdt)) {
-		err = PTR_ERR(kho_fdt);
-		pr_err("memblock FDT was not prepared successfully: %d\n", err);
-		return notifier_from_errno(err);
-	}
-
 	for (i = 0; i < reserved_mem_count; i++) {
 		struct reserve_mem_table *map = &reserved_mem_table[i];

@@ -2528,19 +2519,14 @@ static struct notifier_block reserve_mem_kho_nb = {
 	.notifier_call = reserve_mem_kho_notifier,
 };

-static void __init prepare_kho_fdt(void)
+static int __init prepare_kho_fdt(void)
 {
 	int err = 0, i;
 	void *fdt;

-	if (!reserved_mem_count)
-		return;
-
 	kho_fdt = alloc_page(GFP_KERNEL);
-	if (!kho_fdt) {
-		kho_fdt = ERR_PTR(-ENOMEM);
-		return;
-	}
+	if (!kho_fdt)
+		return -ENOMEM;

 	fdt = page_to_virt(kho_fdt);

@@ -2565,18 +2551,30 @@ static void __init prepare_kho_fdt(void)
 	if (err) {
 		pr_err("failed to prepare memblock FDT for KHO: %d\n", err);
 		put_page(kho_fdt);
-		kho_fdt = ERR_PTR(-EINVAL);
+		kho_fdt = NULL;
 	}
+
+	return err;
 }

 static int __init reserve_mem_init(void)
 {
-	if (!kho_is_enabled())
+	int err;
+
+	if (!kho_is_enabled() || !reserved_mem_count)
 		return 0;

-	prepare_kho_fdt();
+	err = prepare_kho_fdt();
+	if (err)
+		return err;
+
+	err = register_kho_notifier(&reserve_mem_kho_nb);
+	if (err) {
+		put_page(kho_fdt);
+		kho_fdt = NULL;
+	}

-	return register_kho_notifier(&reserve_mem_kho_nb);
+	return err;
 }
 late_initcall(reserve_mem_init);

@@ -2586,33 +2584,38 @@ static void *__init reserve_mem_kho_retrieve_fdt(void)
 {
 	phys_addr_t fdt_phys;
 	struct folio *fdt_folio;
-	void *fdt;
 	int err;

+	if (kho_fdt_in)
+		return kho_fdt_in;
+
 	err = kho_retrieve_subtree(MEMBLOCK_KHO_FDT, &fdt_phys);
 	if (err) {
 		if (err != -ENOENT)
 			pr_warn("failed to retrieve FDT '%s' from KHO: %d\n",
 				MEMBLOCK_KHO_FDT, err);
-		return ERR_PTR(err);
+		goto out;
 	}

 	fdt_folio = kho_restore_folio(fdt_phys);
 	if (!fdt_folio) {
 		pr_warn("failed to restore memblock KHO FDT (0x%llx)\n", fdt_phys);
-		return ERR_PTR(-EFAULT);
+		err = -EFAULT;
+		goto out;
 	}

-	fdt = page_to_virt(folio_page(fdt_folio, 0));
+	kho_fdt_in = folio_address(fdt_folio);

-	err = fdt_node_check_compatible(fdt, 0, MEMBLOCK_KHO_NODE_COMPATIBLE);
+	err = fdt_node_check_compatible(kho_fdt_in, 0, MEMBLOCK_KHO_NODE_COMPATIBLE);
 	if (err) {
 		pr_warn("FDT '%s' is incompatible with '%s': %d\n",
 			MEMBLOCK_KHO_FDT, MEMBLOCK_KHO_NODE_COMPATIBLE, err);
-		return ERR_PTR(-EINVAL);
+		err = -EFAULT;
 	}
-
-	return fdt;
+out:
+	if (err)
+		kho_fdt_in = ERR_PTR(err);
+	return kho_fdt_in;
 }

 static bool __init reserve_mem_kho_revive(const char *name, phys_addr_t size,
@@ -2622,14 +2625,10 @@ static bool __init reserve_mem_kho_revive(const char *name, phys_addr_t size,
 	const phys_addr_t *p_start, *p_size;
 	const void *fdt;

-	if (!kho_fdt_in)
-		kho_fdt_in = reserve_mem_kho_retrieve_fdt();
-
-	if (IS_ERR(kho_fdt_in))
+	fdt = reserve_mem_kho_retrieve_fdt();
+	if (IS_ERR(fdt))
 		return false;

-	fdt = kho_fdt_in;
-
 	offset = fdt_subnode_offset(fdt, 0, name);
 	if (offset < 0) {
 		pr_warn("FDT '%s' has no child '%s': %d\n",
--
2.49.0.805.g082f7c87e0-goog