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") 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") 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") 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