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 cgroup_rstat_updated(struct cgroup *cgrp, int cpu) __ksym; |
41 | extern void cgroup_rstat_flush(struct cgroup *cgrp) __ksym; |
42 | |
43 | static uint64_t cgroup_id(struct cgroup *cgrp) |
44 | { |
45 | return cgrp->kn->id; |
46 | } |
47 | |
48 | static int create_percpu_attach_counter(__u64 cg_id, __u64 state) |
49 | { |
50 | struct percpu_attach_counter pcpu_init = {.state = state, .prev = 0}; |
51 | |
52 | return bpf_map_update_elem(&percpu_attach_counters, &cg_id, |
53 | &pcpu_init, BPF_NOEXIST); |
54 | } |
55 | |
56 | static int create_attach_counter(__u64 cg_id, __u64 state, __u64 pending) |
57 | { |
58 | struct attach_counter init = {.state = state, .pending = pending}; |
59 | |
60 | return bpf_map_update_elem(&attach_counters, &cg_id, |
61 | &init, BPF_NOEXIST); |
62 | } |
63 | |
64 | SEC("fentry/cgroup_attach_task" ) |
65 | int BPF_PROG(counter, struct cgroup *dst_cgrp, struct task_struct *leader, |
66 | bool threadgroup) |
67 | { |
68 | __u64 cg_id = cgroup_id(dst_cgrp); |
69 | struct percpu_attach_counter *pcpu_counter = bpf_map_lookup_elem( |
70 | &percpu_attach_counters, |
71 | &cg_id); |
72 | |
73 | if (pcpu_counter) |
74 | pcpu_counter->state += 1; |
75 | else if (create_percpu_attach_counter(cg_id, 1)) |
76 | return 0; |
77 | |
78 | cgroup_rstat_updated(dst_cgrp, bpf_get_smp_processor_id()); |
79 | return 0; |
80 | } |
81 | |
82 | SEC("fentry/bpf_rstat_flush" ) |
83 | int BPF_PROG(flusher, struct cgroup *cgrp, struct cgroup *parent, int cpu) |
84 | { |
85 | struct percpu_attach_counter *pcpu_counter; |
86 | struct attach_counter *total_counter, *parent_counter; |
87 | __u64 cg_id = cgroup_id(cgrp); |
88 | __u64 parent_cg_id = parent ? cgroup_id(parent) : 0; |
89 | __u64 state; |
90 | __u64 delta = 0; |
91 | |
92 | /* Add CPU changes on this level since the last flush */ |
93 | pcpu_counter = bpf_map_lookup_percpu_elem(&percpu_attach_counters, |
94 | &cg_id, cpu); |
95 | if (pcpu_counter) { |
96 | state = pcpu_counter->state; |
97 | delta += state - pcpu_counter->prev; |
98 | pcpu_counter->prev = state; |
99 | } |
100 | |
101 | total_counter = bpf_map_lookup_elem(&attach_counters, &cg_id); |
102 | if (!total_counter) { |
103 | if (create_attach_counter(cg_id, delta, 0)) |
104 | return 0; |
105 | goto update_parent; |
106 | } |
107 | |
108 | /* Collect pending stats from subtree */ |
109 | if (total_counter->pending) { |
110 | delta += total_counter->pending; |
111 | total_counter->pending = 0; |
112 | } |
113 | |
114 | /* Propagate changes to this cgroup's total */ |
115 | total_counter->state += delta; |
116 | |
117 | update_parent: |
118 | /* Skip if there are no changes to propagate, or no parent */ |
119 | if (!delta || !parent_cg_id) |
120 | return 0; |
121 | |
122 | /* Propagate changes to cgroup's parent */ |
123 | parent_counter = bpf_map_lookup_elem(&attach_counters, |
124 | &parent_cg_id); |
125 | if (parent_counter) |
126 | parent_counter->pending += delta; |
127 | else |
128 | create_attach_counter(parent_cg_id, 0, delta); |
129 | return 0; |
130 | } |
131 | |
132 | SEC("iter.s/cgroup" ) |
133 | int BPF_PROG(dumper, struct bpf_iter_meta *meta, struct cgroup *cgrp) |
134 | { |
135 | struct seq_file *seq = meta->seq; |
136 | struct attach_counter *total_counter; |
137 | __u64 cg_id = cgrp ? cgroup_id(cgrp) : 0; |
138 | |
139 | /* Do nothing for the terminal call */ |
140 | if (!cg_id) |
141 | return 1; |
142 | |
143 | /* Flush the stats to make sure we get the most updated numbers */ |
144 | cgroup_rstat_flush(cgrp); |
145 | |
146 | total_counter = bpf_map_lookup_elem(&attach_counters, &cg_id); |
147 | if (!total_counter) { |
148 | BPF_SEQ_PRINTF(seq, "cg_id: %llu, attach_counter: 0\n" , |
149 | cg_id); |
150 | } else { |
151 | BPF_SEQ_PRINTF(seq, "cg_id: %llu, attach_counter: %llu\n" , |
152 | cg_id, total_counter->state); |
153 | } |
154 | return 0; |
155 | } |
156 | |