bperf_follower.bpf.c

  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
 10struct {
 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
 17struct {
 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
 24struct {
 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
 32enum bperf_filter_type type = 0;
 33int enabled = 0;
 34int inherit;
 35
 36SEC("fexit/XXX")
 37int 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
 80do_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. */
 97SEC("tp_btf/task_newtask")
 98int 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. */
139SEC("tp_btf/sched_process_exit")
140int 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
162char LICENSE[] SEC("license") = "Dual BSD/GPL";