[PATCH v2] drivers/base/memory: set mem->altmap after successful device registration

[Georgi Djakov <georgi.djakov@oss.qualcomm.com>]

If __add_memory_block() fails at xa_store() (under memory pressure
for example), device_unregister() is called, which eventually
triggers memory_block_release() with mem->altmap still set, causing
a WARN_ON(mem->altmap). This was triggered by modifying virtio-mem
driver.

Fix this by delaying the assignment of mem->altmap until after
__add_memory_block() has succeeded.

Fixes: 1a8c64e11043 ("mm/memory_hotplug: embed vmem_altmap details in memory block")
Signed-off-by: Georgi Djakov <georgi.djakov@oss.qualcomm.com>
---

v1: https://patch.msgid.link/20260513130900.2204499-1-georgi.djakov@oss.qualcomm.com

 drivers/base/memory.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/drivers/base/memory.c b/drivers/base/memory.c
index 11d57cfa8d72..78e9600e6d45 100644
--- a/drivers/base/memory.c
+++ b/drivers/base/memory.c
@@ -807,7 +807,6 @@ static int add_memory_block(unsigned long block_id, int nid, unsigned long state
 	mem->start_section_nr = block_id * sections_per_block;
 	mem->state = state;
 	mem->nid = nid;
-	mem->altmap = altmap;
 	INIT_LIST_HEAD(&mem->group_next);
 
 #ifndef CONFIG_NUMA
@@ -825,6 +824,8 @@ static int add_memory_block(unsigned long block_id, int nid, unsigned long state
 	if (ret)
 		return ret;
 
+	mem->altmap = altmap;
+
 	if (group) {
 		mem->group = group;
 		list_add(&mem->group_next, &group->memory_blocks);