xref: /linux/tools/testing/selftests/bpf/progs/cgroup_hierarchical_stats.c (revision fcab107abe1ab5be9dbe874baa722372da8f4f73)
1 // SPDX-License-Identifier: GPL-2.0-only
2 /*
3  * Copyright 2022 Google LLC.
4  */
5 #include "vmlinux.h"
6 #include <bpf/bpf_helpers.h>
7 #include <bpf/bpf_tracing.h>
8 #include <bpf/bpf_core_read.h>
9 
10 char _license[] SEC("license") = "GPL";
11 
12 struct percpu_attach_counter {
13 	/* Previous percpu state, to figure out if we have new updates */
14 	__u64 prev;
15 	/* Current percpu state */
16 	__u64 state;
17 };
18 
19 struct attach_counter {
20 	/* State propagated through children, pending aggregation */
21 	__u64 pending;
22 	/* Total state, including all cpus and all children */
23 	__u64 state;
24 };
25 
26 struct {
27 	__uint(type, BPF_MAP_TYPE_PERCPU_HASH);
28 	__uint(max_entries, 1024);
29 	__type(key, __u64);
30 	__type(value, struct percpu_attach_counter);
31 } percpu_attach_counters SEC(".maps");
32 
33 struct {
34 	__uint(type, BPF_MAP_TYPE_HASH);
35 	__uint(max_entries, 1024);
36 	__type(key, __u64);
37 	__type(value, struct attach_counter);
38 } attach_counters SEC(".maps");
39 
40 extern void css_rstat_updated(
41 		struct cgroup_subsys_state *css, int cpu) __ksym;
42 extern void css_rstat_flush(struct cgroup_subsys_state *css) __ksym;
43 
44 static uint64_t cgroup_id(struct cgroup *cgrp)
45 {
46 	return cgrp->kn->id;
47 }
48 
49 static int create_percpu_attach_counter(__u64 cg_id, __u64 state)
50 {
51 	struct percpu_attach_counter pcpu_init = {.state = state, .prev = 0};
52 
53 	return bpf_map_update_elem(&percpu_attach_counters, &cg_id,
54 				   &pcpu_init, BPF_NOEXIST);
55 }
56 
57 static int create_attach_counter(__u64 cg_id, __u64 state, __u64 pending)
58 {
59 	struct attach_counter init = {.state = state, .pending = pending};
60 
61 	return bpf_map_update_elem(&attach_counters, &cg_id,
62 				   &init, BPF_NOEXIST);
63 }
64 
65 SEC("fentry/cgroup_attach_task")
66 int BPF_PROG(counter, struct cgroup *dst_cgrp, struct task_struct *leader,
67 	     bool threadgroup)
68 {
69 	__u64 cg_id = cgroup_id(dst_cgrp);
70 	struct percpu_attach_counter *pcpu_counter = bpf_map_lookup_elem(
71 			&percpu_attach_counters,
72 			&cg_id);
73 
74 	if (pcpu_counter)
75 		pcpu_counter->state += 1;
76 	else if (create_percpu_attach_counter(cg_id, 1))
77 		return 0;
78 
79 	css_rstat_updated(&dst_cgrp->self, bpf_get_smp_processor_id());
80 	return 0;
81 }
82 
83 SEC("fentry/bpf_rstat_flush")
84 int BPF_PROG(flusher, struct cgroup *cgrp, struct cgroup *parent, int cpu)
85 {
86 	struct percpu_attach_counter *pcpu_counter;
87 	struct attach_counter *total_counter, *parent_counter;
88 	__u64 cg_id = cgroup_id(cgrp);
89 	__u64 parent_cg_id = parent ? cgroup_id(parent) : 0;
90 	__u64 state;
91 	__u64 delta = 0;
92 
93 	/* Add CPU changes on this level since the last flush */
94 	pcpu_counter = bpf_map_lookup_percpu_elem(&percpu_attach_counters,
95 						  &cg_id, cpu);
96 	if (pcpu_counter) {
97 		state = pcpu_counter->state;
98 		delta += state - pcpu_counter->prev;
99 		pcpu_counter->prev = state;
100 	}
101 
102 	total_counter = bpf_map_lookup_elem(&attach_counters, &cg_id);
103 	if (!total_counter) {
104 		if (create_attach_counter(cg_id, delta, 0))
105 			return 0;
106 		goto update_parent;
107 	}
108 
109 	/* Collect pending stats from subtree */
110 	if (total_counter->pending) {
111 		delta += total_counter->pending;
112 		total_counter->pending = 0;
113 	}
114 
115 	/* Propagate changes to this cgroup's total */
116 	total_counter->state += delta;
117 
118 update_parent:
119 	/* Skip if there are no changes to propagate, or no parent */
120 	if (!delta || !parent_cg_id)
121 		return 0;
122 
123 	/* Propagate changes to cgroup's parent */
124 	parent_counter = bpf_map_lookup_elem(&attach_counters,
125 					     &parent_cg_id);
126 	if (parent_counter)
127 		parent_counter->pending += delta;
128 	else
129 		create_attach_counter(parent_cg_id, 0, delta);
130 	return 0;
131 }
132 
133 SEC("iter.s/cgroup")
134 int BPF_PROG(dumper, struct bpf_iter_meta *meta, struct cgroup *cgrp)
135 {
136 	struct seq_file *seq = meta->seq;
137 	struct attach_counter *total_counter;
138 	__u64 cg_id = cgrp ? cgroup_id(cgrp) : 0;
139 
140 	/* Do nothing for the terminal call */
141 	if (!cg_id)
142 		return 1;
143 
144 	/* Flush the stats to make sure we get the most updated numbers */
145 	css_rstat_flush(&cgrp->self);
146 
147 	total_counter = bpf_map_lookup_elem(&attach_counters, &cg_id);
148 	if (!total_counter) {
149 		BPF_SEQ_PRINTF(seq, "cg_id: %llu, attach_counter: 0\n",
150 			       cg_id);
151 	} else {
152 		BPF_SEQ_PRINTF(seq, "cg_id: %llu, attach_counter: %llu\n",
153 			       cg_id, total_counter->state);
154 	}
155 	return 0;
156 }
157