1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
3 #include <test_progs.h>
4 #include <bpf/libbpf.h>
5 #include <bpf/btf.h>
6 #include <fcntl.h>
7 #include <sys/mman.h>
8 #include <unistd.h>
9 #include "cgroup_helpers.h"
10 #include "cgroup_iter_memcg.h"
11 #include "cgroup_iter_memcg.skel.h"
12
13 /*
14 * memcg stats are cached per-cpu and only become visible once the periodic
15 * flusher runs (FLUSH_TIME, 2s), or once pending updates cross
16 * MEMCG_CHARGE_BATCH * num_online_cpus(). That threshold grows with the CPU
17 * count, so on a large machine a single pass does not reach it and
18 * bpf_mem_cgroup_flush_stats() returns without flushing anything. Retry for
19 * long enough to cover a flusher cycle.
20 */
21 #define MEMCG_STAT_RETRIES 16
22 #define MEMCG_STAT_RETRY_DELAY_US (250 * 1000)
23
read_stats(struct bpf_link * link)24 static int read_stats(struct bpf_link *link)
25 {
26 int fd, ret = 0;
27 ssize_t bytes;
28
29 fd = bpf_iter_create(bpf_link__fd(link));
30 if (!ASSERT_OK_FD(fd, "bpf_iter_create"))
31 return 1;
32
33 /*
34 * Invoke iter program by reading from its fd. We're not expecting any
35 * data to be written by the bpf program so the result should be zero.
36 * Results will be read directly through the custom data section
37 * accessible through skel->data_query.memcg_query.
38 */
39 bytes = read(fd, NULL, 0);
40 if (!ASSERT_EQ(bytes, 0, "read fd"))
41 ret = 1;
42
43 close(fd);
44 return ret;
45 }
46
test_anon(struct bpf_link * link,struct memcg_query * memcg_query)47 static void test_anon(struct bpf_link *link, struct memcg_query *memcg_query)
48 {
49 int retries = 0;
50 void *map;
51 size_t len;
52
53 len = sysconf(_SC_PAGESIZE) * 1024;
54
55 retry:
56 /*
57 * Increase memcg anon usage by mapping and writing
58 * to a new anon region.
59 */
60 map = mmap(NULL, len, PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
61 if (!ASSERT_NEQ(map, MAP_FAILED, "mmap anon"))
62 return;
63
64 memset(map, 1, len);
65
66 if (!ASSERT_OK(read_stats(link), "read stats"))
67 goto cleanup;
68
69 if (!memcg_query->nr_anon_mapped && ++retries < MEMCG_STAT_RETRIES) {
70 usleep(MEMCG_STAT_RETRY_DELAY_US);
71 munmap(map, len);
72 goto retry;
73 }
74
75 ASSERT_GT(memcg_query->nr_anon_mapped, 0, "final anon mapped val");
76
77 cleanup:
78 munmap(map, len);
79 }
80
test_file(struct bpf_link * link,struct memcg_query * memcg_query)81 static void test_file(struct bpf_link *link, struct memcg_query *memcg_query)
82 {
83 int retries = 0;
84 void *map;
85 size_t len;
86 char *path;
87 int fd;
88
89 len = sysconf(_SC_PAGESIZE) * 1024;
90 path = "/tmp/test_cgroup_iter_memcg";
91
92 /*
93 * Increase memcg file usage by creating and writing
94 * to a mapped file.
95 */
96 fd = open(path, O_CREAT | O_RDWR, 0644);
97 if (!ASSERT_OK_FD(fd, "open fd"))
98 return;
99 retry:
100 if (!ASSERT_OK(ftruncate(fd, len), "ftruncate"))
101 goto cleanup_fd;
102
103 map = mmap(NULL, len, PROT_WRITE, MAP_SHARED, fd, 0);
104 if (!ASSERT_NEQ(map, MAP_FAILED, "mmap file"))
105 goto cleanup_fd;
106
107 memset(map, 1, len);
108
109 if (!ASSERT_OK(read_stats(link), "read stats"))
110 goto cleanup_map;
111
112 if ((!memcg_query->nr_file_pages || !memcg_query->nr_file_mapped) &&
113 ++retries < MEMCG_STAT_RETRIES) {
114 usleep(MEMCG_STAT_RETRY_DELAY_US);
115 munmap(map, len);
116 goto retry;
117 }
118
119 ASSERT_GT(memcg_query->nr_file_pages, 0, "final file value");
120 ASSERT_GT(memcg_query->nr_file_mapped, 0, "final file mapped value");
121
122 cleanup_map:
123 munmap(map, len);
124 cleanup_fd:
125 close(fd);
126 unlink(path);
127 }
128
test_shmem(struct bpf_link * link,struct memcg_query * memcg_query)129 static void test_shmem(struct bpf_link *link, struct memcg_query *memcg_query)
130 {
131 int retries = 0;
132 size_t len;
133 int fd;
134
135 len = sysconf(_SC_PAGESIZE) * 1024;
136
137 /*
138 * Increase memcg shmem usage by creating and writing
139 * to a memfd backed by shmem/tmpfs.
140 */
141 fd = memfd_create("tmp_shmem", 0);
142 if (!ASSERT_OK_FD(fd, "memfd_create"))
143 return;
144
145 retry:
146 if (!ASSERT_OK(fallocate(fd, 0, 0, len), "fallocate"))
147 goto cleanup;
148
149 if (!ASSERT_OK(read_stats(link), "read stats"))
150 goto cleanup;
151
152 if (!memcg_query->nr_shmem && ++retries < MEMCG_STAT_RETRIES) {
153 usleep(MEMCG_STAT_RETRY_DELAY_US);
154 goto retry;
155 }
156
157 ASSERT_GT(memcg_query->nr_shmem, 0, "final shmem value");
158
159 cleanup:
160 close(fd);
161 }
162
test_pgfault(struct bpf_link * link,struct memcg_query * memcg_query)163 static void test_pgfault(struct bpf_link *link, struct memcg_query *memcg_query)
164 {
165 int retries = 0;
166 void *map;
167 size_t len;
168
169 len = sysconf(_SC_PAGESIZE) * 1024;
170
171 retry:
172 /* Create region to use for triggering a page fault. */
173 map = mmap(NULL, len, PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
174 if (!ASSERT_NEQ(map, MAP_FAILED, "mmap anon"))
175 return;
176
177 /* Trigger page fault. */
178 memset(map, 1, len);
179
180 if (!ASSERT_OK(read_stats(link), "read stats"))
181 goto cleanup;
182
183 if (!memcg_query->pgfault && ++retries < MEMCG_STAT_RETRIES) {
184 usleep(MEMCG_STAT_RETRY_DELAY_US);
185 munmap(map, len);
186 goto retry;
187 }
188
189 ASSERT_GT(memcg_query->pgfault, 0, "final pgfault val");
190
191 cleanup:
192 munmap(map, len);
193 }
194
test_cgroup_iter_memcg(void)195 void test_cgroup_iter_memcg(void)
196 {
197 char *cgroup_rel_path = "/cgroup_iter_memcg_test";
198 struct cgroup_iter_memcg *skel;
199 struct bpf_link *link;
200 int cgroup_fd;
201
202 cgroup_fd = cgroup_setup_and_join(cgroup_rel_path);
203 if (!ASSERT_OK_FD(cgroup_fd, "cgroup_setup_and_join"))
204 return;
205
206 skel = cgroup_iter_memcg__open_and_load();
207 if (!ASSERT_OK_PTR(skel, "cgroup_iter_memcg__open_and_load"))
208 goto cleanup_cgroup_fd;
209
210 DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
211 union bpf_iter_link_info linfo = {
212 .cgroup.cgroup_fd = cgroup_fd,
213 .cgroup.order = BPF_CGROUP_ITER_SELF_ONLY,
214 };
215 opts.link_info = &linfo;
216 opts.link_info_len = sizeof(linfo);
217
218 link = bpf_program__attach_iter(skel->progs.cgroup_memcg_query, &opts);
219 if (!ASSERT_OK_PTR(link, "bpf_program__attach_iter"))
220 goto cleanup_skel;
221
222 if (test__start_subtest("cgroup_iter_memcg__anon"))
223 test_anon(link, &skel->data_query->memcg_query);
224 if (test__start_subtest("cgroup_iter_memcg__shmem"))
225 test_shmem(link, &skel->data_query->memcg_query);
226 if (test__start_subtest("cgroup_iter_memcg__file"))
227 test_file(link, &skel->data_query->memcg_query);
228 if (test__start_subtest("cgroup_iter_memcg__pgfault"))
229 test_pgfault(link, &skel->data_query->memcg_query);
230
231 bpf_link__destroy(link);
232 cleanup_skel:
233 cgroup_iter_memcg__destroy(skel);
234 cleanup_cgroup_fd:
235 close(cgroup_fd);
236 cleanup_cgroup_environment();
237 }
238