Re: [PATCH] drivers/base/memory: make memory block get/put explicit

[Muchun Song <muchun.song@linux.dev>]

> On May 11, 2026, at 21:23, Muchun Song <muchun.song@linux.dev> wrote:
> 
> 
> 
>> On May 11, 2026, at 20:22, David Hildenbrand (Arm) <david@kernel.org> wrote:
>> 
>> On 5/11/26 13:18, Muchun Song wrote:
>>> Rename the memory block lookup helper to make the acquired reference
>>> explicit, add memory_block_put() to wrap put_device(), and collapse the
>>> redundant section-number wrapper into a single block-id based lookup
>>> interface.
>>> 
>>> This makes it clearer to callers that a successful lookup holds a
>>> reference that must be dropped, reducing the chance of forgetting the
>>> matching put and leaking the memory block device reference.
>> 
>> Better mention some of the other changes here, like removing find_memory_block().
> 
> Will do.
> 
>> 
>> [...]
>> 
>>>   unlock_device_hotplug();
>>> diff --git a/include/linux/memory.h b/include/linux/memory.h
>>> index 5bb5599c6b2b..29edef1f975c 100644
>>> --- a/include/linux/memory.h
>>> +++ b/include/linux/memory.h
>>> @@ -158,7 +158,11 @@ int create_memory_block_devices(unsigned long start, unsigned long size,
>>> void remove_memory_block_devices(unsigned long start, unsigned long size);
>>> extern void memory_dev_init(void);
>>> extern int memory_notify(enum memory_block_state state, void *v);
>>> -extern struct memory_block *find_memory_block(unsigned long section_nr);
>>> +extern struct memory_block *memory_block_get(unsigned long block_id);
>> 
>> While at it, please drop the "extern".
> 
> OK.
> 
>> 
>>> +static inline void memory_block_put(struct memory_block *mem)
>>> +{
>>> +    put_device(&mem->dev);
>>> +}
>>> typedef int (*walk_memory_blocks_func_t)(struct memory_block *, void *);
>>> extern int walk_memory_blocks(unsigned long start, unsigned long size,
>>>                 void *arg, walk_memory_blocks_func_t func);
>>> @@ -171,7 +175,6 @@ struct memory_group *memory_group_find_by_id(int mgid);
>>> typedef int (*walk_memory_groups_func_t)(struct memory_group *, void *);
>>> int walk_dynamic_memory_groups(int nid, walk_memory_groups_func_t func,
>>>                  struct memory_group *excluded, void *arg);
>>> -struct memory_block *find_memory_block_by_id(unsigned long block_id);
>>> #define hotplug_memory_notifier(fn, pri) ({        \
>>>   static __meminitdata struct notifier_block fn##_mem_nb =\
>>>       { .notifier_call = fn, .priority = pri };\
>>> diff --git a/mm/memory_hotplug.c b/mm/memory_hotplug.c
>>> index 462d8dcd636d..890c6453e887 100644
>>> --- a/mm/memory_hotplug.c
>>> +++ b/mm/memory_hotplug.c
>>> @@ -1417,14 +1417,13 @@ static void remove_memory_blocks_and_altmaps(u64 start, u64 size)
>>>       struct vmem_altmap *altmap = NULL;
>>>       struct memory_block *mem;
>>> 
>>> -        mem = find_memory_block(pfn_to_section_nr(PFN_DOWN(cur_start)));
>>> +        mem = memory_block_get(phys_to_block_id(cur_start));
>>>       if (WARN_ON_ONCE(!mem))
>>>           continue;
>>> 
>>>       altmap = mem->altmap;
>>>       mem->altmap = NULL;
>>> -        /* drop the ref. we got via find_memory_block() */
>>> -        put_device(&mem->dev);
>>> +        memory_block_put(mem);
>> 
>> Would guards come in handy here?
> 
> You mean to introduce something like:
> 
> scoped_guard(memory_block, id) {
> }
> 
> Right? If yes, I will give it a try.

Hi David,

Did I get that right?

5 files changed, 26 insertions(+), 35 deletions(-)

diff --git a/drivers/base/memory.c b/drivers/base/memory.c
index 5b5d41089e81..eba48320fc30 100644
--- a/drivers/base/memory.c
+++ b/drivers/base/memory.c
@@ -997,7 +997,6 @@ int walk_memory_blocks(unsigned long start, unsigned long size,
 {
 	const unsigned long start_block_id = phys_to_block_id(start);
 	const unsigned long end_block_id = phys_to_block_id(start + size - 1);
-	struct memory_block *mem;
 	unsigned long block_id;
 	int ret = 0;

@@ -1005,12 +1004,11 @@ int walk_memory_blocks(unsigned long start, unsigned long size,
 		return 0;

 	for (block_id = start_block_id; block_id <= end_block_id; block_id++) {
-		mem = memory_block_get(block_id);
+		CLASS(memory_block_get, mem)(block_id);
 		if (!mem)
 			continue;

 		ret = func(mem, arg);
-		memory_block_put(mem);
 		if (ret)
 			break;
 	}
@@ -1218,23 +1216,19 @@ int walk_dynamic_memory_groups(int nid, walk_memory_groups_func_t func,
 void memblk_nr_poison_inc(unsigned long pfn)
 {
 	const unsigned long block_id = pfn_to_block_id(pfn);
-	struct memory_block *mem = memory_block_get(block_id);
+	CLASS(memory_block_get, mem)(block_id);

-	if (mem) {
+	if (mem)
 		atomic_long_inc(&mem->nr_hwpoison);
-		memory_block_put(mem);
-	}
 }

 void memblk_nr_poison_sub(unsigned long pfn, long i)
 {
 	const unsigned long block_id = pfn_to_block_id(pfn);
-	struct memory_block *mem = memory_block_get(block_id);
+	CLASS(memory_block_get, mem)(block_id);

-	if (mem) {
+	if (mem)
 		atomic_long_sub(i, &mem->nr_hwpoison);
-		memory_block_put(mem);
-	}
 }

 static unsigned long memblk_nr_poison(struct memory_block *mem)
diff --git a/drivers/base/node.c b/drivers/base/node.c
index b3333ca92090..32a5431cbb60 100644
--- a/drivers/base/node.c
+++ b/drivers/base/node.c
@@ -845,15 +845,12 @@ static void register_memory_blocks_under_nodes(void)
 			continue;

 		for (block_id = start_block_id; block_id <= end_block_id; block_id++) {
-			struct memory_block *mem;
-
-			mem = memory_block_get(block_id);
+			CLASS(memory_block_get, mem)(block_id);
 			if (!mem)
 				continue;

 			memory_block_add_nid_early(mem, nid);
 			do_register_memory_block_under_node(nid, mem);
-			memory_block_put(mem);
 		}

 	}
diff --git a/drivers/s390/char/sclp_mem.c b/drivers/s390/char/sclp_mem.c
index 6df1926d4c62..c9a1a7adcab6 100644
--- a/drivers/s390/char/sclp_mem.c
+++ b/drivers/s390/char/sclp_mem.c
@@ -237,13 +237,12 @@ static ssize_t sclp_config_mem_store(struct kobject *kobj, struct kobj_attribute
 	} else {
 		if (!sclp_mem->config)
 			goto out_unlock;
-		mem = memory_block_get(phys_to_block_id(addr));
-		if (mem->state != MEM_OFFLINE) {
-			memory_block_put(mem);
-			rc = -EBUSY;
-			goto out_unlock;
+		scoped_class(memory_block_get, mem, phys_to_block_id(addr)) {
+			if (mem->state != MEM_OFFLINE) {
+				rc = -EBUSY;
+				goto out_unlock;
+			}
 		}
-		memory_block_put(mem);
 		sclp_mem_change_state(addr, block_size, 0);
 		__remove_memory(addr, block_size);
 #ifdef CONFIG_KASAN
@@ -293,12 +292,12 @@ static ssize_t sclp_memmap_on_memory_store(struct kobject *kobj, struct kobj_att
 		return rc;
 	block_size = memory_block_size_bytes();
 	sclp_mem = container_of(kobj, struct sclp_mem, kobj);
-	mem = memory_block_get(phys_to_block_id(sclp_mem->id * block_size));
-	if (!mem) {
-		WRITE_ONCE(sclp_mem->memmap_on_memory, value);
-	} else {
-		memory_block_put(mem);
-		rc = -EBUSY;
+	scoped_class(memory_block_get, mem,
+		     phys_to_block_id(sclp_mem->id * block_size)) {
+		if (!mem)
+			WRITE_ONCE(sclp_mem->memmap_on_memory, value);
+		else
+			rc = -EBUSY;
 	}
 	unlock_device_hotplug();
 	return rc ? rc : count;
diff --git a/include/linux/memory.h b/include/linux/memory.h
index 463dc02f6cff..fb67041941a8 100644
--- a/include/linux/memory.h
+++ b/include/linux/memory.h
@@ -18,6 +18,7 @@

 #include <linux/node.h>
 #include <linux/compiler.h>
+#include <linux/cleanup.h>
 #include <linux/mutex.h>
 #include <linux/memory_hotplug.h>

@@ -163,6 +164,8 @@ static inline void memory_block_put(struct memory_block *mem)
 {
 	put_device(&mem->dev);
 }
+DEFINE_CLASS(memory_block_get, struct memory_block *, if (_T) memory_block_put(_T),
+	     memory_block_get(block_id), unsigned long block_id)
 typedef int (*walk_memory_blocks_func_t)(struct memory_block *, void *);
 extern int walk_memory_blocks(unsigned long start, unsigned long size,
 			      void *arg, walk_memory_blocks_func_t func);
diff --git a/mm/memory_hotplug.c b/mm/memory_hotplug.c
index 890c6453e887..3f89ffb286d6 100644
--- a/mm/memory_hotplug.c
+++ b/mm/memory_hotplug.c
@@ -1415,15 +1415,13 @@ static void remove_memory_blocks_and_altmaps(u64 start, u64 size)
 	for (cur_start = start; cur_start < start + size;
 	     cur_start += memblock_size) {
 		struct vmem_altmap *altmap = NULL;
-		struct memory_block *mem;
-
-		mem = memory_block_get(phys_to_block_id(cur_start));
-		if (WARN_ON_ONCE(!mem))
-			continue;
+		scoped_class(memory_block_get, mem, phys_to_block_id(cur_start)) {
+			if (WARN_ON_ONCE(!mem))
+				continue;

-		altmap = mem->altmap;
-		mem->altmap = NULL;
-		memory_block_put(mem);
+			altmap = mem->altmap;
+			mem->altmap = NULL;
+		}

 		remove_memory_block_devices(cur_start, memblock_size);

Thanks,
Muchun

> 
>> 
>> In general
>> 
>> Acked-by: David Hildenbrand (Arm) <david@kernel.org>
> 
> Thanks.
> 
> Muchun
> 
>> 
>> --
>> Cheers,
>> 
>> David