On Tue, Aug 26, 2025 at 03:19:40PM +0800, Yafang Shao wrote: > We will utilize this new kfunc bpf_mm_get_mem_cgroup() to retrieve the > associated mem_cgroup from the given @mm. The obtained mem_cgroup must > be released by calling bpf_put_mem_cgroup() as a paired operation. > > Signed-off-by: Yafang Shao <laoar.shao@xxxxxxxxx> > --- > mm/bpf_thp.c | 51 ++++++++++++++++++++++++++++++++++++++++++++++++++- > 1 file changed, 50 insertions(+), 1 deletion(-) > > diff --git a/mm/bpf_thp.c b/mm/bpf_thp.c > index fbff3b1bb988..b757e8f425fd 100644 > --- a/mm/bpf_thp.c > +++ b/mm/bpf_thp.c > @@ -175,10 +175,59 @@ static struct bpf_struct_ops bpf_bpf_thp_ops = { > .name = "bpf_thp_ops", > }; > > +__bpf_kfunc_start_defs(); > + > +/** > + * bpf_mm_get_mem_cgroup - Get the memory cgroup associated with a mm_struct. > + * @mm: The mm_struct to query > + * > + * The obtained mem_cgroup must be released by calling bpf_put_mem_cgroup(). > + * > + * Return: The associated mem_cgroup on success, or NULL on failure. Note that > + * this function depends on CONFIG_MEMCG being enabled - it will always return > + * NULL if CONFIG_MEMCG is not configured. > + */ > +__bpf_kfunc struct mem_cgroup *bpf_mm_get_mem_cgroup(struct mm_struct *mm) > +{ > + return get_mem_cgroup_from_mm(mm); > +} > + > +/** > + * bpf_put_mem_cgroup - Release a memory cgroup obtained from bpf_mm_get_mem_cgroup() > + * @memcg: The memory cgroup to release > + */ > +__bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg) > +{ > +#ifdef CONFIG_MEMCG > + if (!memcg) > + return; > + css_put(&memcg->css); > +#endif Just use mem_cgroup_put() here. > +} > + > +__bpf_kfunc_end_defs(); > + > +BTF_KFUNCS_START(bpf_thp_ids) > +BTF_ID_FLAGS(func, bpf_mm_get_mem_cgroup, KF_TRUSTED_ARGS | KF_ACQUIRE | KF_RET_NULL) > +BTF_ID_FLAGS(func, bpf_put_mem_cgroup, KF_RELEASE) > +BTF_KFUNCS_END(bpf_thp_ids) > + > +static const struct btf_kfunc_id_set bpf_thp_set = { > + .owner = THIS_MODULE, > + .set = &bpf_thp_ids, > +}; > + > static int __init bpf_thp_ops_init(void) > { > - int err = register_bpf_struct_ops(&bpf_bpf_thp_ops, bpf_thp_ops); > + int err; > + > + err = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &bpf_thp_set); > + if (err) { > + pr_err("bpf_thp: Failed to register kfunc sets (%d)\n", err); > + return err; > + } > > + err = register_bpf_struct_ops(&bpf_bpf_thp_ops, bpf_thp_ops); > if (err) > pr_err("bpf_thp: Failed to register struct_ops (%d)\n", err); > return err; > -- > 2.47.3 >