1 // SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause)
2 // Copyright (c) 2021 Google
3 #include "vmlinux.h"
4 #include <bpf/bpf_helpers.h>
5 #include <bpf/bpf_tracing.h>
6 
7 // This should be in sync with "util/ftrace.h"
8 #define NUM_BUCKET  22
9 
10 struct {
11 	__uint(type, BPF_MAP_TYPE_HASH);
12 	__uint(key_size, sizeof(__u64));
13 	__uint(value_size, sizeof(__u64));
14 	__uint(max_entries, 10000);
15 } functime SEC(".maps");
16 
17 struct {
18 	__uint(type, BPF_MAP_TYPE_HASH);
19 	__uint(key_size, sizeof(__u32));
20 	__uint(value_size, sizeof(__u8));
21 	__uint(max_entries, 1);
22 } cpu_filter SEC(".maps");
23 
24 struct {
25 	__uint(type, BPF_MAP_TYPE_HASH);
26 	__uint(key_size, sizeof(__u32));
27 	__uint(value_size, sizeof(__u8));
28 	__uint(max_entries, 1);
29 } task_filter SEC(".maps");
30 
31 struct {
32 	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
33 	__uint(key_size, sizeof(__u32));
34 	__uint(value_size, sizeof(__u64));
35 	__uint(max_entries, NUM_BUCKET);
36 } latency SEC(".maps");
37 
38 
39 int enabled = 0;
40 int has_cpu = 0;
41 int has_task = 0;
42 int use_nsec = 0;
43 
44 SEC("kprobe/func")
BPF_PROG(func_begin)45 int BPF_PROG(func_begin)
46 {
47 	__u64 key, now;
48 
49 	if (!enabled)
50 		return 0;
51 
52 	key = bpf_get_current_pid_tgid();
53 
54 	if (has_cpu) {
55 		__u32 cpu = bpf_get_smp_processor_id();
56 		__u8 *ok;
57 
58 		ok = bpf_map_lookup_elem(&cpu_filter, &cpu);
59 		if (!ok)
60 			return 0;
61 	}
62 
63 	if (has_task) {
64 		__u32 pid = key & 0xffffffff;
65 		__u8 *ok;
66 
67 		ok = bpf_map_lookup_elem(&task_filter, &pid);
68 		if (!ok)
69 			return 0;
70 	}
71 
72 	now = bpf_ktime_get_ns();
73 
74 	// overwrite timestamp for nested functions
75 	bpf_map_update_elem(&functime, &key, &now, BPF_ANY);
76 	return 0;
77 }
78 
79 SEC("kretprobe/func")
BPF_PROG(func_end)80 int BPF_PROG(func_end)
81 {
82 	__u64 tid;
83 	__u64 *start;
84 	__u64 cmp_base = use_nsec ? 1 : 1000;
85 
86 	if (!enabled)
87 		return 0;
88 
89 	tid = bpf_get_current_pid_tgid();
90 
91 	start = bpf_map_lookup_elem(&functime, &tid);
92 	if (start) {
93 		__s64 delta = bpf_ktime_get_ns() - *start;
94 		__u32 key;
95 		__u64 *hist;
96 
97 		bpf_map_delete_elem(&functime, &tid);
98 
99 		if (delta < 0)
100 			return 0;
101 
102 		// calculate index using delta
103 		for (key = 0; key < (NUM_BUCKET - 1); key++) {
104 			if (delta < (cmp_base << key))
105 				break;
106 		}
107 
108 		hist = bpf_map_lookup_elem(&latency, &key);
109 		if (!hist)
110 			return 0;
111 
112 		*hist += 1;
113 	}
114 
115 	return 0;
116 }
117