xref: /linux/mm/bpf_memcontrol.c (revision 5904db9891f80f84283648121e2d8c8a506296a8)
1 // SPDX-License-Identifier: GPL-2.0-or-later
2 /*
3  * Memory Controller-related BPF kfuncs and auxiliary code
4  *
5  * Author: Roman Gushchin <roman.gushchin@linux.dev>
6  */
7 
8 #include <linux/memcontrol.h>
9 #include <linux/bpf.h>
10 
11 __bpf_kfunc_start_defs();
12 
13 /**
14  * bpf_get_mem_cgroup - Get a reference to a memory cgroup
15  * @css: pointer to the css structure
16  *
17  * It's fine to pass a css which belongs to any cgroup controller,
18  * e.g. unified hierarchy's main css.
19  *
20  * Implements KF_ACQUIRE semantics.
21  *
22  * Return: A pointer to a mem_cgroup structure after bumping
23  * the corresponding css's reference counter.
24  */
25 __bpf_kfunc struct mem_cgroup *
26 bpf_get_mem_cgroup(struct cgroup_subsys_state *css)
27 {
28 	struct mem_cgroup *memcg = NULL;
29 	bool rcu_unlock = false;
30 
31 	if (mem_cgroup_disabled() || !root_mem_cgroup)
32 		return NULL;
33 
34 	if (root_mem_cgroup->css.ss != css->ss) {
35 		struct cgroup *cgroup = css->cgroup;
36 		int ssid = root_mem_cgroup->css.ss->id;
37 
38 		rcu_read_lock();
39 		rcu_unlock = true;
40 		css = rcu_dereference_raw(cgroup->subsys[ssid]);
41 	}
42 
43 	if (css && css_tryget(css))
44 		memcg = container_of(css, struct mem_cgroup, css);
45 
46 	if (rcu_unlock)
47 		rcu_read_unlock();
48 
49 	return memcg;
50 }
51 
52 /**
53  * bpf_put_mem_cgroup - Put a reference to a memory cgroup
54  * @memcg: memory cgroup to release
55  *
56  * Releases a previously acquired memcg reference.
57  * Implements KF_RELEASE semantics.
58  */
59 __bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg)
60 {
61 	css_put(&memcg->css);
62 }
63 
64 __bpf_kfunc_end_defs();
65 
66 BTF_KFUNCS_START(bpf_memcontrol_kfuncs)
67 BTF_ID_FLAGS(func, bpf_get_mem_cgroup, KF_ACQUIRE | KF_RET_NULL | KF_RCU)
68 BTF_ID_FLAGS(func, bpf_put_mem_cgroup, KF_RELEASE)
69 
70 BTF_KFUNCS_END(bpf_memcontrol_kfuncs)
71 
72 static const struct btf_kfunc_id_set bpf_memcontrol_kfunc_set = {
73 	.owner          = THIS_MODULE,
74 	.set            = &bpf_memcontrol_kfuncs,
75 };
76 
77 static int __init bpf_memcontrol_init(void)
78 {
79 	int err;
80 
81 	err = register_btf_kfunc_id_set(BPF_PROG_TYPE_UNSPEC,
82 					&bpf_memcontrol_kfunc_set);
83 	if (err)
84 		pr_warn("error while registering bpf memcontrol kfuncs: %d", err);
85 
86 	return err;
87 }
88 late_initcall(bpf_memcontrol_init);
89