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 
41 // stats
42 __s64 total;
43 __s64 count;
44 __s64 max;
45 __s64 min;
46 
47 const volatile int has_cpu = 0;
48 const volatile int has_task = 0;
49 const volatile int use_nsec = 0;
50 const volatile unsigned int bucket_range;
51 const volatile unsigned int min_latency;
52 const volatile unsigned int max_latency;
53 const volatile unsigned int bucket_num = NUM_BUCKET;
54 
55 SEC("kprobe/func")
BPF_PROG(func_begin)56 int BPF_PROG(func_begin)
57 {
58 	__u64 key, now;
59 
60 	if (!enabled)
61 		return 0;
62 
63 	key = bpf_get_current_pid_tgid();
64 
65 	if (has_cpu) {
66 		__u32 cpu = bpf_get_smp_processor_id();
67 		__u8 *ok;
68 
69 		ok = bpf_map_lookup_elem(&cpu_filter, &cpu);
70 		if (!ok)
71 			return 0;
72 	}
73 
74 	if (has_task) {
75 		__u32 pid = key & 0xffffffff;
76 		__u8 *ok;
77 
78 		ok = bpf_map_lookup_elem(&task_filter, &pid);
79 		if (!ok)
80 			return 0;
81 	}
82 
83 	now = bpf_ktime_get_ns();
84 
85 	// overwrite timestamp for nested functions
86 	bpf_map_update_elem(&functime, &key, &now, BPF_ANY);
87 	return 0;
88 }
89 
90 SEC("kretprobe/func")
BPF_PROG(func_end)91 int BPF_PROG(func_end)
92 {
93 	__u64 tid;
94 	__u64 *start;
95 	__u64 cmp_base = use_nsec ? 1 : 1000;
96 
97 	if (!enabled)
98 		return 0;
99 
100 	tid = bpf_get_current_pid_tgid();
101 
102 	start = bpf_map_lookup_elem(&functime, &tid);
103 	if (start) {
104 		__s64 delta = bpf_ktime_get_ns() - *start;
105 		__u64 val = delta;
106 		__u32 key = 0;
107 		__u64 *hist;
108 
109 		bpf_map_delete_elem(&functime, &tid);
110 
111 		if (delta < 0)
112 			return 0;
113 
114 		if (bucket_range != 0) {
115 			val = delta / cmp_base;
116 
117 			if (min_latency > 0) {
118 				if (val > min_latency)
119 					val -= min_latency;
120 				else
121 					goto do_lookup;
122 			}
123 
124 			// Less than 1 unit (ms or ns), or, in the future,
125 			// than the min latency desired.
126 			if (val > 0) { // 1st entry: [ 1 unit .. bucket_range units )
127 				key = val / bucket_range + 1;
128 				if (key >= bucket_num)
129 					key = bucket_num - 1;
130 			}
131 
132 			goto do_lookup;
133 		}
134 		// calculate index using delta
135 		for (key = 0; key < (bucket_num - 1); key++) {
136 			if (delta < (cmp_base << key))
137 				break;
138 		}
139 
140 do_lookup:
141 		hist = bpf_map_lookup_elem(&latency, &key);
142 		if (!hist)
143 			return 0;
144 
145 		__sync_fetch_and_add(hist, 1);
146 
147 		__sync_fetch_and_add(&total, delta); // always in nsec
148 		__sync_fetch_and_add(&count, 1);
149 
150 		if (delta > max)
151 			max = delta;
152 		if (delta < min)
153 			min = delta;
154 	}
155 
156 	return 0;
157 }
158