1// SPDX-License-Identifier: GPL-2.0-only
2/*
3 * This test makes sure BPF stats collection using rstat works correctly.
4 * The test uses 3 BPF progs:
5 * (a) counter: This BPF prog is invoked every time we attach a process to a
6 *              cgroup and locklessly increments a percpu counter.
7 *              The program then calls cgroup_rstat_updated() to inform rstat
8 *              of an update on the (cpu, cgroup) pair.
9 *
10 * (b) flusher: This BPF prog is invoked when an rstat flush is ongoing, it
11 *              aggregates all percpu counters to a total counter, and also
12 *              propagates the changes to the ancestor cgroups.
13 *
14 * (c) dumper: This BPF prog is a cgroup_iter. It is used to output the total
15 *             counter of a cgroup through reading a file in userspace.
16 *
17 * The test sets up a cgroup hierarchy, and the above programs. It spawns a few
18 * processes in the leaf cgroups and makes sure all the counters are aggregated
19 * correctly.
20 *
21 * Copyright 2022 Google LLC.
22 */
23#include <asm-generic/errno.h>
24#include <errno.h>
25#include <sys/types.h>
26#include <sys/mount.h>
27#include <sys/stat.h>
28#include <unistd.h>
29
30#include <test_progs.h>
31#include <bpf/libbpf.h>
32#include <bpf/bpf.h>
33
34#include "cgroup_helpers.h"
35#include "cgroup_hierarchical_stats.skel.h"
36
37#define PAGE_SIZE 4096
38#define MB(x) (x << 20)
39
40#define PROCESSES_PER_CGROUP 3
41
42#define BPFFS_ROOT "/sys/fs/bpf/"
43#define BPFFS_ATTACH_COUNTERS BPFFS_ROOT "attach_counters/"
44
45#define CG_ROOT_NAME "root"
46#define CG_ROOT_ID 1
47
48#define CGROUP_PATH(p, n) {.path = p"/"n, .name = n}
49
50static struct {
51	const char *path, *name;
52	unsigned long long id;
53	int fd;
54} cgroups[] = {
55	CGROUP_PATH("/", "test"),
56	CGROUP_PATH("/test", "child1"),
57	CGROUP_PATH("/test", "child2"),
58	CGROUP_PATH("/test/child1", "child1_1"),
59	CGROUP_PATH("/test/child1", "child1_2"),
60	CGROUP_PATH("/test/child2", "child2_1"),
61	CGROUP_PATH("/test/child2", "child2_2"),
62};
63
64#define N_CGROUPS ARRAY_SIZE(cgroups)
65#define N_NON_LEAF_CGROUPS 3
66
67static int root_cgroup_fd;
68static bool mounted_bpffs;
69
70/* reads file at 'path' to 'buf', returns 0 on success. */
71static int read_from_file(const char *path, char *buf, size_t size)
72{
73	int fd, len;
74
75	fd = open(path, O_RDONLY);
76	if (fd < 0)
77		return fd;
78
79	len = read(fd, buf, size);
80	close(fd);
81	if (len < 0)
82		return len;
83
84	buf[len] = 0;
85	return 0;
86}
87
88/* mounts bpffs and mkdir for reading stats, returns 0 on success. */
89static int setup_bpffs(void)
90{
91	int err;
92
93	/* Mount bpffs */
94	err = mount("bpf", BPFFS_ROOT, "bpf", 0, NULL);
95	mounted_bpffs = !err;
96	if (ASSERT_FALSE(err && errno != EBUSY, "mount"))
97		return err;
98
99	/* Create a directory to contain stat files in bpffs */
100	err = mkdir(BPFFS_ATTACH_COUNTERS, 0755);
101	if (!ASSERT_OK(err, "mkdir"))
102		return err;
103
104	return 0;
105}
106
107static void cleanup_bpffs(void)
108{
109	/* Remove created directory in bpffs */
110	ASSERT_OK(rmdir(BPFFS_ATTACH_COUNTERS), "rmdir "BPFFS_ATTACH_COUNTERS);
111
112	/* Unmount bpffs, if it wasn't already mounted when we started */
113	if (mounted_bpffs)
114		return;
115
116	ASSERT_OK(umount(BPFFS_ROOT), "unmount bpffs");
117}
118
119/* sets up cgroups, returns 0 on success. */
120static int setup_cgroups(void)
121{
122	int i, fd, err;
123
124	err = setup_cgroup_environment();
125	if (!ASSERT_OK(err, "setup_cgroup_environment"))
126		return err;
127
128	root_cgroup_fd = get_root_cgroup();
129	if (!ASSERT_GE(root_cgroup_fd, 0, "get_root_cgroup"))
130		return root_cgroup_fd;
131
132	for (i = 0; i < N_CGROUPS; i++) {
133		fd = create_and_get_cgroup(cgroups[i].path);
134		if (!ASSERT_GE(fd, 0, "create_and_get_cgroup"))
135			return fd;
136
137		cgroups[i].fd = fd;
138		cgroups[i].id = get_cgroup_id(cgroups[i].path);
139	}
140	return 0;
141}
142
143static void cleanup_cgroups(void)
144{
145	close(root_cgroup_fd);
146	for (int i = 0; i < N_CGROUPS; i++)
147		close(cgroups[i].fd);
148	cleanup_cgroup_environment();
149}
150
151/* Sets up cgroup hiearchary, returns 0 on success. */
152static int setup_hierarchy(void)
153{
154	return setup_bpffs() || setup_cgroups();
155}
156
157static void destroy_hierarchy(void)
158{
159	cleanup_cgroups();
160	cleanup_bpffs();
161}
162
163static int attach_processes(void)
164{
165	int i, j, status;
166
167	/* In every leaf cgroup, attach 3 processes */
168	for (i = N_NON_LEAF_CGROUPS; i < N_CGROUPS; i++) {
169		for (j = 0; j < PROCESSES_PER_CGROUP; j++) {
170			pid_t pid;
171
172			/* Create child and attach to cgroup */
173			pid = fork();
174			if (pid == 0) {
175				if (join_parent_cgroup(cgroups[i].path))
176					exit(EACCES);
177				exit(0);
178			}
179
180			/* Cleanup child */
181			waitpid(pid, &status, 0);
182			if (!ASSERT_TRUE(WIFEXITED(status), "child process exited"))
183				return 1;
184			if (!ASSERT_EQ(WEXITSTATUS(status), 0,
185				       "child process exit code"))
186				return 1;
187		}
188	}
189	return 0;
190}
191
192static unsigned long long
193get_attach_counter(unsigned long long cgroup_id, const char *file_name)
194{
195	unsigned long long attach_counter = 0, id = 0;
196	static char buf[128], path[128];
197
198	/* For every cgroup, read the file generated by cgroup_iter */
199	snprintf(path, 128, "%s%s", BPFFS_ATTACH_COUNTERS, file_name);
200	if (!ASSERT_OK(read_from_file(path, buf, 128), "read cgroup_iter"))
201		return 0;
202
203	/* Check the output file formatting */
204	ASSERT_EQ(sscanf(buf, "cg_id: %llu, attach_counter: %llu\n",
205			 &id, &attach_counter), 2, "output format");
206
207	/* Check that the cgroup_id is displayed correctly */
208	ASSERT_EQ(id, cgroup_id, "cgroup_id");
209	/* Check that the counter is non-zero */
210	ASSERT_GT(attach_counter, 0, "attach counter non-zero");
211	return attach_counter;
212}
213
214static void check_attach_counters(void)
215{
216	unsigned long long attach_counters[N_CGROUPS], root_attach_counter;
217	int i;
218
219	for (i = 0; i < N_CGROUPS; i++)
220		attach_counters[i] = get_attach_counter(cgroups[i].id,
221							cgroups[i].name);
222
223	/* Read stats for root too */
224	root_attach_counter = get_attach_counter(CG_ROOT_ID, CG_ROOT_NAME);
225
226	/* Check that all leafs cgroups have an attach counter of 3 */
227	for (i = N_NON_LEAF_CGROUPS; i < N_CGROUPS; i++)
228		ASSERT_EQ(attach_counters[i], PROCESSES_PER_CGROUP,
229			  "leaf cgroup attach counter");
230
231	/* Check that child1 == child1_1 + child1_2 */
232	ASSERT_EQ(attach_counters[1], attach_counters[3] + attach_counters[4],
233		  "child1_counter");
234	/* Check that child2 == child2_1 + child2_2 */
235	ASSERT_EQ(attach_counters[2], attach_counters[5] + attach_counters[6],
236		  "child2_counter");
237	/* Check that test == child1 + child2 */
238	ASSERT_EQ(attach_counters[0], attach_counters[1] + attach_counters[2],
239		  "test_counter");
240	/* Check that root >= test */
241	ASSERT_GE(root_attach_counter, attach_counters[1], "root_counter");
242}
243
244/* Creates iter link and pins in bpffs, returns 0 on success, -errno on failure.
245 */
246static int setup_cgroup_iter(struct cgroup_hierarchical_stats *obj,
247			     int cgroup_fd, const char *file_name)
248{
249	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
250	union bpf_iter_link_info linfo = {};
251	struct bpf_link *link;
252	static char path[128];
253	int err;
254
255	/*
256	 * Create an iter link, parameterized by cgroup_fd. We only want to
257	 * traverse one cgroup, so set the traversal order to "self".
258	 */
259	linfo.cgroup.cgroup_fd = cgroup_fd;
260	linfo.cgroup.order = BPF_CGROUP_ITER_SELF_ONLY;
261	opts.link_info = &linfo;
262	opts.link_info_len = sizeof(linfo);
263	link = bpf_program__attach_iter(obj->progs.dumper, &opts);
264	if (!ASSERT_OK_PTR(link, "attach_iter"))
265		return -EFAULT;
266
267	/* Pin the link to a bpffs file */
268	snprintf(path, 128, "%s%s", BPFFS_ATTACH_COUNTERS, file_name);
269	err = bpf_link__pin(link, path);
270	ASSERT_OK(err, "pin cgroup_iter");
271
272	/* Remove the link, leaving only the ref held by the pinned file */
273	bpf_link__destroy(link);
274	return err;
275}
276
277/* Sets up programs for collecting stats, returns 0 on success. */
278static int setup_progs(struct cgroup_hierarchical_stats **skel)
279{
280	int i, err;
281
282	*skel = cgroup_hierarchical_stats__open_and_load();
283	if (!ASSERT_OK_PTR(*skel, "open_and_load"))
284		return 1;
285
286	/* Attach cgroup_iter program that will dump the stats to cgroups */
287	for (i = 0; i < N_CGROUPS; i++) {
288		err = setup_cgroup_iter(*skel, cgroups[i].fd, cgroups[i].name);
289		if (!ASSERT_OK(err, "setup_cgroup_iter"))
290			return err;
291	}
292
293	/* Also dump stats for root */
294	err = setup_cgroup_iter(*skel, root_cgroup_fd, CG_ROOT_NAME);
295	if (!ASSERT_OK(err, "setup_cgroup_iter"))
296		return err;
297
298	bpf_program__set_autoattach((*skel)->progs.dumper, false);
299	err = cgroup_hierarchical_stats__attach(*skel);
300	if (!ASSERT_OK(err, "attach"))
301		return err;
302
303	return 0;
304}
305
306static void destroy_progs(struct cgroup_hierarchical_stats *skel)
307{
308	static char path[128];
309	int i;
310
311	for (i = 0; i < N_CGROUPS; i++) {
312		/* Delete files in bpffs that cgroup_iters are pinned in */
313		snprintf(path, 128, "%s%s", BPFFS_ATTACH_COUNTERS,
314			 cgroups[i].name);
315		ASSERT_OK(remove(path), "remove cgroup_iter pin");
316	}
317
318	/* Delete root file in bpffs */
319	snprintf(path, 128, "%s%s", BPFFS_ATTACH_COUNTERS, CG_ROOT_NAME);
320	ASSERT_OK(remove(path), "remove cgroup_iter root pin");
321	cgroup_hierarchical_stats__destroy(skel);
322}
323
324void test_cgroup_hierarchical_stats(void)
325{
326	struct cgroup_hierarchical_stats *skel = NULL;
327
328	if (setup_hierarchy())
329		goto hierarchy_cleanup;
330	if (setup_progs(&skel))
331		goto cleanup;
332	if (attach_processes())
333		goto cleanup;
334	check_attach_counters();
335cleanup:
336	destroy_progs(skel);
337hierarchy_cleanup:
338	destroy_hierarchy();
339}
340