Re: [PATCH bpf-next v1 2/6] mm: introduce BPF kfuncs to deal with memcg pointers

[Roman Gushchin <roman.gushchin@linux.dev>]

Shakeel Butt <shakeel.butt@linux.dev> writes:

> On Thu, Dec 18, 2025 at 05:57:46PM -0800, Roman Gushchin wrote:
>> To effectively operate with memory cgroups in BPF there is a need
>> to convert css pointers to memcg pointers. A simple container_of
>> cast which is used in the kernel code can't be used in BPF because
>> from the verifier's point of view that's a out-of-bounds memory access.
>> 
>> Introduce helper get/put kfuncs which can be used to get
>> a refcounted memcg pointer from the css pointer:
>>   - bpf_get_mem_cgroup,
>>   - bpf_put_mem_cgroup.
>> 
>> bpf_get_mem_cgroup() can take both memcg's css and the corresponding
>> cgroup's "self" css. It allows it to be used with the existing cgroup
>> iterator which iterates over cgroup tree, not memcg tree.
>> 
>> Signed-off-by: Roman Gushchin <roman.gushchin@linux.dev>
>> ---
>>  mm/Makefile         |  3 ++
>>  mm/bpf_memcontrol.c | 88 +++++++++++++++++++++++++++++++++++++++++++++
>
> Let's add this file to MAINTAINERS file.

Will do. I planned to create a new entry for mm-related bpf files
as part of the bpf oom patchset.

>
>>  2 files changed, 91 insertions(+)
>>  create mode 100644 mm/bpf_memcontrol.c
>> 
>> diff --git a/mm/Makefile b/mm/Makefile
>> index 9175f8cc6565..79c39a98ff83 100644
>> --- a/mm/Makefile
>> +++ b/mm/Makefile
>> @@ -106,6 +106,9 @@ obj-$(CONFIG_MEMCG) += memcontrol.o vmpressure.o
>>  ifdef CONFIG_SWAP
>>  obj-$(CONFIG_MEMCG) += swap_cgroup.o
>>  endif
>> +ifdef CONFIG_BPF_SYSCALL
>> +obj-$(CONFIG_MEMCG) += bpf_memcontrol.o
>> +endif
>>  obj-$(CONFIG_CGROUP_HUGETLB) += hugetlb_cgroup.o
>>  obj-$(CONFIG_GUP_TEST) += gup_test.o
>>  obj-$(CONFIG_DMAPOOL_TEST) += dmapool_test.o
>> diff --git a/mm/bpf_memcontrol.c b/mm/bpf_memcontrol.c
>> new file mode 100644
>> index 000000000000..8aa842b56817
>> --- /dev/null
>> +++ b/mm/bpf_memcontrol.c
>> @@ -0,0 +1,88 @@
>> +// SPDX-License-Identifier: GPL-2.0-or-later
>> +/*
>> + * Memory Controller-related BPF kfuncs and auxiliary code
>> + *
>> + * Author: Roman Gushchin <roman.gushchin@linux.dev>
>> + */
>> +
>> +#include <linux/memcontrol.h>
>> +#include <linux/bpf.h>
>> +
>> +__bpf_kfunc_start_defs();
>> +
>> +/**
>> + * bpf_get_mem_cgroup - Get a reference to a memory cgroup
>> + * @css: pointer to the css structure
>> + *
>> + * Returns a pointer to a mem_cgroup structure after bumping
>> + * the corresponding css's reference counter.
>> + *
>> + * It's fine to pass a css which belongs to any cgroup controller,
>> + * e.g. unified hierarchy's main css.
>> + *
>> + * Implements KF_ACQUIRE semantics.
>> + */
>> +__bpf_kfunc struct mem_cgroup *
>> +bpf_get_mem_cgroup(struct cgroup_subsys_state *css)
>> +{
>> +	struct mem_cgroup *memcg = NULL;
>> +	bool rcu_unlock = false;
>> +
>> +	if (!root_mem_cgroup)
>> +		return NULL;
>
> Should we also handle mem_cgroup_disabled() here?

Good point, will add in v2. Same with bpf_get_root_mem_cgroup() patch.

>
>> +
>> +	if (root_mem_cgroup->css.ss != css->ss) {
>> +		struct cgroup *cgroup = css->cgroup;
>> +		int ssid = root_mem_cgroup->css.ss->id;
>> +
>> +		rcu_read_lock();
>> +		rcu_unlock = true;
>> +		css = rcu_dereference_raw(cgroup->subsys[ssid]);
>> +	}
>> +
>> +	if (css && css_tryget(css))
>> +		memcg = container_of(css, struct mem_cgroup, css);
>> +
>> +	if (rcu_unlock)
>> +		rcu_read_unlock();
>
> Any reason to handle rcu lock like this? Why not just take the rcu read
> lock irrespective? It is cheap.

Idk, it's cheap but not entirely free and I think the code is still
perfectly readable.

>
>> +
>> +	return memcg;
>> +}
>> +
>> +/**
>> + * bpf_put_mem_cgroup - Put a reference to a memory cgroup
>> + * @memcg: memory cgroup to release
>> + *
>> + * Releases a previously acquired memcg reference.
>> + * Implements KF_RELEASE semantics.
>> + */
>> +__bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg)
>> +{
>> +	css_put(&memcg->css);
>
> Should we NULL check memcg here? bpf_get_mem_cgroup() can return NULL.

No, the verifier ensures it's a valid memcg pointer. No need for an
additional check.

>
>> +}
>> +
>> +__bpf_kfunc_end_defs();
>> +
>> +BTF_KFUNCS_START(bpf_memcontrol_kfuncs)
>> +BTF_ID_FLAGS(func, bpf_get_mem_cgroup, KF_TRUSTED_ARGS | KF_ACQUIRE | KF_RET_NULL | KF_RCU)
>> +BTF_ID_FLAGS(func, bpf_put_mem_cgroup, KF_TRUSTED_ARGS | KF_RELEASE)
>
> Will the verifier enforce that bpf_put_mem_cgroup() can not be called
> with NULL?

Yep.

Thanks for the review!