xref: /linux/mm/bpf_memcontrol.c (revision 5c7db3239c9fbe3c62cb0d89b64959ea23af2de9)
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_root_mem_cgroup - Returns a pointer to the root memory cgroup
15  *
16  * The function has KF_ACQUIRE semantics, even though the root memory
17  * cgroup is never destroyed after being created and doesn't require
18  * reference counting. And it's perfectly safe to pass it to
19  * bpf_put_mem_cgroup()
20  *
21  * Return: A pointer to the root memory cgroup.
22  */
23 __bpf_kfunc struct mem_cgroup *bpf_get_root_mem_cgroup(void)
24 {
25 	if (mem_cgroup_disabled())
26 		return NULL;
27 
28 	/* css_get() is not needed */
29 	return root_mem_cgroup;
30 }
31 
32 /**
33  * bpf_get_mem_cgroup - Get a reference to a memory cgroup
34  * @css: pointer to the css structure
35  *
36  * It's fine to pass a css which belongs to any cgroup controller,
37  * e.g. unified hierarchy's main css.
38  *
39  * Implements KF_ACQUIRE semantics.
40  *
41  * Return: A pointer to a mem_cgroup structure after bumping
42  * the corresponding css's reference counter.
43  */
44 __bpf_kfunc struct mem_cgroup *
45 bpf_get_mem_cgroup(struct cgroup_subsys_state *css)
46 {
47 	struct mem_cgroup *memcg = NULL;
48 	bool rcu_unlock = false;
49 
50 	if (mem_cgroup_disabled() || !root_mem_cgroup)
51 		return NULL;
52 
53 	if (root_mem_cgroup->css.ss != css->ss) {
54 		struct cgroup *cgroup = css->cgroup;
55 		int ssid = root_mem_cgroup->css.ss->id;
56 
57 		rcu_read_lock();
58 		rcu_unlock = true;
59 		css = rcu_dereference_raw(cgroup->subsys[ssid]);
60 	}
61 
62 	if (css && css_tryget(css))
63 		memcg = container_of(css, struct mem_cgroup, css);
64 
65 	if (rcu_unlock)
66 		rcu_read_unlock();
67 
68 	return memcg;
69 }
70 
71 /**
72  * bpf_put_mem_cgroup - Put a reference to a memory cgroup
73  * @memcg: memory cgroup to release
74  *
75  * Releases a previously acquired memcg reference.
76  * Implements KF_RELEASE semantics.
77  */
78 __bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg)
79 {
80 	css_put(&memcg->css);
81 }
82 
83 __bpf_kfunc_end_defs();
84 
85 BTF_KFUNCS_START(bpf_memcontrol_kfuncs)
86 BTF_ID_FLAGS(func, bpf_get_root_mem_cgroup, KF_ACQUIRE | KF_RET_NULL)
87 BTF_ID_FLAGS(func, bpf_get_mem_cgroup, KF_ACQUIRE | KF_RET_NULL | KF_RCU)
88 BTF_ID_FLAGS(func, bpf_put_mem_cgroup, KF_RELEASE)
89 
90 BTF_KFUNCS_END(bpf_memcontrol_kfuncs)
91 
92 static const struct btf_kfunc_id_set bpf_memcontrol_kfunc_set = {
93 	.owner          = THIS_MODULE,
94 	.set            = &bpf_memcontrol_kfuncs,
95 };
96 
97 static int __init bpf_memcontrol_init(void)
98 {
99 	int err;
100 
101 	err = register_btf_kfunc_id_set(BPF_PROG_TYPE_UNSPEC,
102 					&bpf_memcontrol_kfunc_set);
103 	if (err)
104 		pr_warn("error while registering bpf memcontrol kfuncs: %d", err);
105 
106 	return err;
107 }
108 late_initcall(bpf_memcontrol_init);
109