[v3] mm/memory_hotplug: Fix remove_memory() lockdep splat

diff --git a/Documentation/core-api/memory-hotplug.rst b/Documentation/core-api/memory-hotplug.rst
index de7467e48067..637b467378b7 100644
--- a/Documentation/core-api/memory-hotplug.rst
+++ b/Documentation/core-api/memory-hotplug.rst
@@ -90,12 +90,13 @@  Locking Internals
 =================
 
 When adding/removing memory that uses memory block devices (i.e. ordinary RAM),
-the device_hotplug_lock should be held to:
+the device_hotplug_lock is held to:
 
 - synchronize against online/offline requests (e.g. via sysfs). This way, memory
   block devices can only be accessed (.online/.state attributes) by user
-  space once memory has been fully added. And when removing memory, we
-  know nobody is in critical sections.
+  space once memory has been fully added. And when removing memory, the
+  memory block device is invalidated (mem->section count set to 0) under the
+  lock to abort any in-flight online requests.
 - synchronize against CPU hotplug and similar (e.g. relevant for ACPI and PPC)
 
 Especially, there is a possible lock inversion that is avoided using
@@ -112,7 +113,13 @@  can result in a lock inversion.
 
 onlining/offlining of memory should be done via device_online()/
 device_offline() - to make sure it is properly synchronized to actions
-via sysfs. Holding device_hotplug_lock is advised (to e.g. protect online_type)
+via sysfs. Holding device_hotplug_lock is required to prevent online racing
+removal. The device_hotplug_lock and memblock invalidation allows
+remove_memory_block_devices() to run outside of mem_hotplug_lock to avoid lock
+dependency conflicts with memblock-sysfs teardown. The add_memory() path
+performs create_memory_block_devices() under mem_hotplug_lock so that if it
+fails it can perform an arch_remove_memory() cleanup. There are no known lock
+dependency problems with memblock-sysfs setup.
 
 When adding/removing/onlining/offlining memory or adding/removing
 heterogeneous/device memory, we should always hold the mem_hotplug_lock in
diff --git a/drivers/base/core.c b/drivers/base/core.c
index 42a672456432..5d5036370c92 100644
--- a/drivers/base/core.c
+++ b/drivers/base/core.c
@@ -1146,6 +1146,11 @@  void unlock_device_hotplug(void)
 	mutex_unlock(&device_hotplug_lock);
 }
 
+void assert_held_device_hotplug(void)
+{
+	lockdep_assert_held(&device_hotplug_lock);
+}
+
 int lock_device_hotplug_sysfs(void)
 {
 	if (mutex_trylock(&device_hotplug_lock))
diff --git a/drivers/base/memory.c b/drivers/base/memory.c
index 799b43191dea..91c6fbd2383e 100644
--- a/drivers/base/memory.c
+++ b/drivers/base/memory.c
@@ -280,6 +280,10 @@  static int memory_subsys_online(struct device *dev)
 	if (mem->state == MEM_ONLINE)
 		return 0;
 
+	/* online lost the race with hot-unplug, abort */
+	if (!mem->section_count)
+		return -ENXIO;
+
 	/*
 	 * If we are called from state_store(), online_type will be
 	 * set >= 0 Otherwise we were called from the device online
@@ -736,8 +740,6 @@  int create_memory_block_devices(unsigned long start, unsigned long size)
  * Remove memory block devices for the given memory area. Start and size
  * have to be aligned to memory block granularity. Memory block devices
  * have to be offline.
- *
- * Called under device_hotplug_lock.
  */
 void remove_memory_block_devices(unsigned long start, unsigned long size)
 {
@@ -746,6 +748,8 @@  void remove_memory_block_devices(unsigned long start, unsigned long size)
 	struct memory_block *mem;
 	unsigned long block_id;
 
+	assert_held_device_hotplug();
+
 	if (WARN_ON_ONCE(!IS_ALIGNED(start, memory_block_size_bytes()) ||
 			 !IS_ALIGNED(size, memory_block_size_bytes())))
 		return;
diff --git a/include/linux/device.h b/include/linux/device.h
index 96ff76731e93..e042da3b1953 100644
--- a/include/linux/device.h
+++ b/include/linux/device.h
@@ -1553,6 +1553,7 @@  static inline bool device_supports_offline(struct device *dev)
 extern void lock_device_hotplug(void);
 extern void unlock_device_hotplug(void);
 extern int lock_device_hotplug_sysfs(void);
+extern void assert_held_device_hotlpug(void);
 extern int device_offline(struct device *dev);
 extern int device_online(struct device *dev);
 extern void set_primary_fwnode(struct device *dev, struct fwnode_handle *fwnode);
diff --git a/mm/memory_hotplug.c b/mm/memory_hotplug.c
index 55ac23ef11c1..0158cd4cca48 100644
--- a/mm/memory_hotplug.c
+++ b/mm/memory_hotplug.c
@@ -1763,8 +1763,6 @@  static int __ref try_remove_memory(int nid, u64 start, u64 size)
 
 	BUG_ON(check_hotplug_memory_range(start, size));
 
-	mem_hotplug_begin();
-
 	/*
 	 * All memory blocks must be offlined before removing memory.  Check
 	 * whether all memory blocks in question are offline and return error
@@ -1777,9 +1775,15 @@  static int __ref try_remove_memory(int nid, u64 start, u64 size)
 	/* remove memmap entry */
 	firmware_map_remove(start, start + size, "System RAM");
 
-	/* remove memory block devices before removing memory */
+	/*
+	 * Remove memory block devices before removing memory and before
+	 * mem_hotplug_begin() (see Documentation/core-api/memory-hotplug.rst
+	 * "Locking Internals").
+	 */
 	remove_memory_block_devices(start, size);
 
+	mem_hotplug_begin();
+
 	arch_remove_memory(nid, start, size, NULL);
 	memblock_free(start, size);
 	memblock_remove(start, size);