1 // SPDX-License-Identifier: GPL-2.0
2 // Copyright (c) 2019 Facebook
3 #include "vmlinux.h"
4 #include <bpf/bpf_core_read.h>
5 #include <bpf/bpf_helpers.h>
6 #include "runqslower.h"
7 
8 #define TASK_RUNNING	0
9 
10 const volatile __u64 min_us = 0;
11 const volatile pid_t targ_pid = 0;
12 const volatile pid_t targ_tgid = 0;
13 
14 // Dummy instance to get skeleton to generate definition for `struct event`
15 struct event _event = {0};
16 
17 // Kernel 5.14 changed the state field to __state
18 struct task_struct___pre_5_14 {
19 	long int state;
20 };
21 
22 struct {
23 	__uint(type, BPF_MAP_TYPE_HASH);
24 	__uint(max_entries, 10240);
25 	__type(key, u32);
26 	__type(value, u64);
27 } start SEC(".maps");
28 
29 struct {
30 	__uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
31 	__uint(key_size, sizeof(u32));
32 	__uint(value_size, sizeof(u32));
33 } events SEC(".maps");
34 
35 /* record enqueue timestamp */
36 static __always_inline
trace_enqueue(u32 tgid,u32 pid)37 int trace_enqueue(u32 tgid, u32 pid)
38 {
39 	u64 ts;
40 
41 	if (!pid)
42 		return 0;
43 	if (targ_tgid && targ_tgid != tgid)
44 		return 0;
45 	if (targ_pid && targ_pid != pid)
46 		return 0;
47 
48 	ts = bpf_ktime_get_ns();
49 	bpf_map_update_elem(&start, &pid, &ts, 0);
50 	return 0;
51 }
52 
53 SEC("tp_btf/sched_wakeup")
handle__sched_wakeup(u64 * ctx)54 int handle__sched_wakeup(u64 *ctx)
55 {
56 	/* TP_PROTO(struct task_struct *p) */
57 	struct task_struct *p = (void *)ctx[0];
58 
59 	return trace_enqueue(p->tgid, p->pid);
60 }
61 
62 SEC("tp_btf/sched_wakeup_new")
handle__sched_wakeup_new(u64 * ctx)63 int handle__sched_wakeup_new(u64 *ctx)
64 {
65 	/* TP_PROTO(struct task_struct *p) */
66 	struct task_struct *p = (void *)ctx[0];
67 
68 	return trace_enqueue(p->tgid, p->pid);
69 }
70 
get_task_state(struct task_struct * t)71 static inline long get_task_state(struct task_struct *t)
72 {
73 	if (bpf_core_field_exists(t->__state))
74 		return t->__state;
75 
76 	return ((struct task_struct___pre_5_14*)t)->state;
77 }
78 
79 SEC("tp_btf/sched_switch")
handle__sched_switch(u64 * ctx)80 int handle__sched_switch(u64 *ctx)
81 {
82 	/* TP_PROTO(bool preempt, struct task_struct *prev,
83 	 *	    struct task_struct *next)
84 	 */
85 	struct task_struct *prev = (struct task_struct *)ctx[1];
86 	struct task_struct *next = (struct task_struct *)ctx[2];
87 	struct event event = {};
88 	u64 *tsp, delta_us;
89 	long state = get_task_state(prev);
90 	u32 pid;
91 
92 	/* ivcsw: treat like an enqueue event and store timestamp */
93 	if (state == TASK_RUNNING)
94 		trace_enqueue(prev->tgid, prev->pid);
95 
96 	pid = next->pid;
97 
98 	/* fetch timestamp and calculate delta */
99 	tsp = bpf_map_lookup_elem(&start, &pid);
100 	if (!tsp)
101 		return 0;   /* missed enqueue */
102 
103 	delta_us = (bpf_ktime_get_ns() - *tsp) / 1000;
104 	if (min_us && delta_us <= min_us)
105 		return 0;
106 
107 	event.pid = pid;
108 	event.delta_us = delta_us;
109 	bpf_probe_read_kernel_str(&event.task, sizeof(event.task), next->comm);
110 
111 	/* output */
112 	bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU,
113 			      &event, sizeof(event));
114 
115 	bpf_map_delete_elem(&start, &pid);
116 	return 0;
117 }
118 
119 char LICENSE[] SEC("license") = "GPL";
120