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
cgroup_id(struct cgroup * cgrp)43 static uint64_t cgroup_id(struct cgroup *cgrp)
44 {
45 return cgrp->kn->id;
46 }
47
create_percpu_attach_counter(__u64 cg_id,__u64 state)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
create_attach_counter(__u64 cg_id,__u64 state,__u64 pending)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")
BPF_PROG(counter,struct cgroup * dst_cgrp,struct task_struct * leader,bool threadgroup)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")
BPF_PROG(flusher,struct cgroup * cgrp,struct cgroup * parent,int cpu)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")
BPF_PROG(dumper,struct bpf_iter_meta * meta,struct cgroup * cgrp)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