1 // SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause)
2 // Copyright (c) 2021 Facebook
3 #include "vmlinux.h"
4 #include <bpf/bpf_helpers.h>
5 #include <bpf/bpf_tracing.h>
6 #include "bperf_u.h"
7
8 #define MAX_ENTRIES 102400
9
10 struct {
11 __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
12 __uint(key_size, sizeof(__u32));
13 __uint(value_size, sizeof(struct bpf_perf_event_value));
14 __uint(max_entries, 1);
15 } diff_readings SEC(".maps");
16
17 struct {
18 __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
19 __uint(key_size, sizeof(__u32));
20 __uint(value_size, sizeof(struct bpf_perf_event_value));
21 __uint(max_entries, 1);
22 } accum_readings SEC(".maps");
23
24 struct {
25 __uint(type, BPF_MAP_TYPE_HASH);
26 __uint(key_size, sizeof(__u32));
27 __uint(value_size, sizeof(struct bperf_filter_value));
28 __uint(max_entries, MAX_ENTRIES);
29 __uint(map_flags, BPF_F_NO_PREALLOC);
30 } filter SEC(".maps");
31
32 enum bperf_filter_type type = 0;
33 int enabled = 0;
34 int inherit;
35
36 SEC("fexit/XXX")
BPF_PROG(fexit_XXX)37 int BPF_PROG(fexit_XXX)
38 {
39 struct bpf_perf_event_value *diff_val, *accum_val;
40 __u32 filter_key, zero = 0;
41 __u32 accum_key;
42 struct bperf_filter_value *fval;
43
44 if (!enabled)
45 return 0;
46
47 switch (type) {
48 case BPERF_FILTER_GLOBAL:
49 accum_key = zero;
50 goto do_add;
51 case BPERF_FILTER_CPU:
52 filter_key = bpf_get_smp_processor_id();
53 break;
54 case BPERF_FILTER_PID:
55 filter_key = bpf_get_current_pid_tgid() & 0xffffffff;
56 break;
57 case BPERF_FILTER_TGID:
58 /* Use pid as the filter_key to exclude new task counts
59 * when inherit is disabled. Don't worry about the existing
60 * children in TGID losing their counts, bpf_counter has
61 * already added them to the filter map via perf_thread_map
62 * before this bpf prog runs.
63 */
64 filter_key = inherit ?
65 bpf_get_current_pid_tgid() >> 32 :
66 bpf_get_current_pid_tgid() & 0xffffffff;
67 break;
68 default:
69 return 0;
70 }
71
72 fval = bpf_map_lookup_elem(&filter, &filter_key);
73 if (!fval)
74 return 0;
75
76 accum_key = fval->accum_key;
77 if (fval->exited)
78 bpf_map_delete_elem(&filter, &filter_key);
79
80 do_add:
81 diff_val = bpf_map_lookup_elem(&diff_readings, &zero);
82 if (!diff_val)
83 return 0;
84
85 accum_val = bpf_map_lookup_elem(&accum_readings, &accum_key);
86 if (!accum_val)
87 return 0;
88
89 accum_val->counter += diff_val->counter;
90 accum_val->enabled += diff_val->enabled;
91 accum_val->running += diff_val->running;
92
93 return 0;
94 }
95
96 /* The program is only used for PID or TGID filter types. */
97 SEC("tp_btf/task_newtask")
BPF_PROG(on_newtask,struct task_struct * task,__u64 clone_flags)98 int BPF_PROG(on_newtask, struct task_struct *task, __u64 clone_flags)
99 {
100 __u32 parent_key, child_key;
101 struct bperf_filter_value *parent_fval;
102 struct bperf_filter_value child_fval = { 0 };
103
104 if (!enabled)
105 return 0;
106
107 switch (type) {
108 case BPERF_FILTER_PID:
109 parent_key = bpf_get_current_pid_tgid() & 0xffffffff;
110 child_key = task->pid;
111 break;
112 case BPERF_FILTER_TGID:
113 parent_key = bpf_get_current_pid_tgid() >> 32;
114 child_key = task->tgid;
115 if (child_key == parent_key)
116 return 0;
117 break;
118 default:
119 return 0;
120 }
121
122 /* Check if the current task is one of the target tasks to be counted */
123 parent_fval = bpf_map_lookup_elem(&filter, &parent_key);
124 if (!parent_fval)
125 return 0;
126
127 /* Start counting for the new task by adding it into filter map,
128 * inherit the accum key of its parent task so that they can be
129 * counted together.
130 */
131 child_fval.accum_key = parent_fval->accum_key;
132 child_fval.exited = 0;
133 bpf_map_update_elem(&filter, &child_key, &child_fval, BPF_NOEXIST);
134
135 return 0;
136 }
137
138 /* The program is only used for PID or TGID filter types. */
139 SEC("tp_btf/sched_process_exit")
BPF_PROG(on_exittask,struct task_struct * task)140 int BPF_PROG(on_exittask, struct task_struct *task)
141 {
142 __u32 pid;
143 struct bperf_filter_value *fval;
144
145 if (!enabled)
146 return 0;
147
148 /* Stop counting for this task by removing it from filter map.
149 * For TGID type, if the pid can be found in the map, it means that
150 * this pid belongs to the leader task. After the task exits, the
151 * tgid of its child tasks (if any) will be 1, so the pid can be
152 * safely removed.
153 */
154 pid = task->pid;
155 fval = bpf_map_lookup_elem(&filter, &pid);
156 if (fval)
157 fval->exited = 1;
158
159 return 0;
160 }
161
162 char LICENSE[] SEC("license") = "Dual BSD/GPL";
163