1 // SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause)
2 // Copyright (c) 2022 Google
3 #include "vmlinux.h"
4 #include <bpf/bpf_helpers.h>
5 #include <bpf/bpf_tracing.h>
6 #include <bpf/bpf_core_read.h>
7 
8 /* maximum stack trace depth */
9 #define MAX_STACKS   8
10 
11 /* default buffer size */
12 #define MAX_ENTRIES  10240
13 
14 struct contention_key {
15 	__s32 stack_id;
16 };
17 
18 struct contention_data {
19 	__u64 total_time;
20 	__u64 min_time;
21 	__u64 max_time;
22 	__u32 count;
23 	__u32 flags;
24 };
25 
26 struct tstamp_data {
27 	__u64 timestamp;
28 	__u64 lock;
29 	__u32 flags;
30 	__s32 stack_id;
31 };
32 
33 /* callstack storage  */
34 struct {
35 	__uint(type, BPF_MAP_TYPE_STACK_TRACE);
36 	__uint(key_size, sizeof(__u32));
37 	__uint(value_size, MAX_STACKS * sizeof(__u64));
38 	__uint(max_entries, MAX_ENTRIES);
39 } stacks SEC(".maps");
40 
41 /* maintain timestamp at the beginning of contention */
42 struct {
43 	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
44 	__uint(map_flags, BPF_F_NO_PREALLOC);
45 	__type(key, int);
46 	__type(value, struct tstamp_data);
47 } tstamp SEC(".maps");
48 
49 /* actual lock contention statistics */
50 struct {
51 	__uint(type, BPF_MAP_TYPE_HASH);
52 	__uint(key_size, sizeof(struct contention_key));
53 	__uint(value_size, sizeof(struct contention_data));
54 	__uint(max_entries, MAX_ENTRIES);
55 } lock_stat SEC(".maps");
56 
57 struct {
58 	__uint(type, BPF_MAP_TYPE_HASH);
59 	__uint(key_size, sizeof(__u32));
60 	__uint(value_size, sizeof(__u8));
61 	__uint(max_entries, 1);
62 } cpu_filter SEC(".maps");
63 
64 struct {
65 	__uint(type, BPF_MAP_TYPE_HASH);
66 	__uint(key_size, sizeof(__u32));
67 	__uint(value_size, sizeof(__u8));
68 	__uint(max_entries, 1);
69 } task_filter SEC(".maps");
70 
71 /* control flags */
72 int enabled;
73 int has_cpu;
74 int has_task;
75 int stack_skip;
76 
77 /* error stat */
78 int lost;
79 
80 static inline int can_record(void)
81 {
82 	if (has_cpu) {
83 		__u32 cpu = bpf_get_smp_processor_id();
84 		__u8 *ok;
85 
86 		ok = bpf_map_lookup_elem(&cpu_filter, &cpu);
87 		if (!ok)
88 			return 0;
89 	}
90 
91 	if (has_task) {
92 		__u8 *ok;
93 		__u32 pid = bpf_get_current_pid_tgid();
94 
95 		ok = bpf_map_lookup_elem(&task_filter, &pid);
96 		if (!ok)
97 			return 0;
98 	}
99 
100 	return 1;
101 }
102 
103 SEC("tp_btf/contention_begin")
104 int contention_begin(u64 *ctx)
105 {
106 	struct task_struct *curr;
107 	struct tstamp_data *pelem;
108 
109 	if (!enabled || !can_record())
110 		return 0;
111 
112 	curr = bpf_get_current_task_btf();
113 	pelem = bpf_task_storage_get(&tstamp, curr, NULL,
114 				     BPF_LOCAL_STORAGE_GET_F_CREATE);
115 	if (!pelem || pelem->lock)
116 		return 0;
117 
118 	pelem->timestamp = bpf_ktime_get_ns();
119 	pelem->lock = (__u64)ctx[0];
120 	pelem->flags = (__u32)ctx[1];
121 	pelem->stack_id = bpf_get_stackid(ctx, &stacks, BPF_F_FAST_STACK_CMP | stack_skip);
122 
123 	if (pelem->stack_id < 0)
124 		lost++;
125 	return 0;
126 }
127 
128 SEC("tp_btf/contention_end")
129 int contention_end(u64 *ctx)
130 {
131 	struct task_struct *curr;
132 	struct tstamp_data *pelem;
133 	struct contention_key key;
134 	struct contention_data *data;
135 	__u64 duration;
136 
137 	if (!enabled)
138 		return 0;
139 
140 	curr = bpf_get_current_task_btf();
141 	pelem = bpf_task_storage_get(&tstamp, curr, NULL, 0);
142 	if (!pelem || pelem->lock != ctx[0])
143 		return 0;
144 
145 	duration = bpf_ktime_get_ns() - pelem->timestamp;
146 
147 	key.stack_id = pelem->stack_id;
148 	data = bpf_map_lookup_elem(&lock_stat, &key);
149 	if (!data) {
150 		struct contention_data first = {
151 			.total_time = duration,
152 			.max_time = duration,
153 			.min_time = duration,
154 			.count = 1,
155 			.flags = pelem->flags,
156 		};
157 
158 		bpf_map_update_elem(&lock_stat, &key, &first, BPF_NOEXIST);
159 		pelem->lock = 0;
160 		return 0;
161 	}
162 
163 	__sync_fetch_and_add(&data->total_time, duration);
164 	__sync_fetch_and_add(&data->count, 1);
165 
166 	/* FIXME: need atomic operations */
167 	if (data->max_time < duration)
168 		data->max_time = duration;
169 	if (data->min_time > duration)
170 		data->min_time = duration;
171 
172 	pelem->lock = 0;
173 	return 0;
174 }
175 
176 char LICENSE[] SEC("license") = "Dual BSD/GPL";
177