1 // SPDX-License-Identifier: GPL-3.0
2 // Copyright (c) 2022 Nicolas Sterchele
3 #include "vmlinux.h"
4 #include <bpf/bpf_helpers.h>
5 #include <bpf/bpf_core_read.h>
6 #include <bpf/bpf_tracing.h>
7 #include "wakeuptime.h"
8 #include "maps.bpf.h"
9
10 #define PF_KTHREAD 0x00200000 /* kernel thread */
11
12 const volatile pid_t targ_pid = 0;
13 const volatile __u64 max_block_ns = -1;
14 const volatile __u64 min_block_ns = 1;
15 const volatile bool user_threads_only = false;
16
17 struct {
18 __uint(type, BPF_MAP_TYPE_HASH);
19 __uint(max_entries, MAX_ENTRIES);
20 __type(key, struct key_t);
21 __type(value, u64);
22 } counts SEC(".maps");
23
24 struct {
25 __uint(type, BPF_MAP_TYPE_HASH);
26 __uint(max_entries, MAX_ENTRIES);
27 __type(key, u32);
28 __type(value, u64);
29 } start SEC(".maps");
30
31 struct {
32 __uint(type, BPF_MAP_TYPE_STACK_TRACE);
33 __uint(key_size, sizeof(u32));
34 } stackmap SEC(".maps");
35
offcpu_sched_switch(struct task_struct * prev)36 static int offcpu_sched_switch(struct task_struct *prev)
37 {
38 u64 pid_tgid = bpf_get_current_pid_tgid();
39 u32 pid = pid_tgid >> 32;
40 u32 tid = (u32)pid_tgid;
41 u64 ts;
42
43 if (targ_pid && targ_pid != pid)
44 return 0;
45
46 if (user_threads_only && prev->flags & PF_KTHREAD)
47 return 0;
48
49 ts = bpf_ktime_get_ns();
50 bpf_map_update_elem(&start, &tid, &ts, BPF_ANY);
51 return 0;
52 }
53
wakeup(void * ctx,struct task_struct * p)54 static int wakeup(void *ctx, struct task_struct *p)
55 {
56 u32 pid = p->tgid;
57 u32 tid = p->pid;
58 u64 delta, *count_key, *tsp;
59 static const u64 zero;
60 struct key_t key = {};
61
62 if (targ_pid && targ_pid != pid)
63 return 0;
64 tsp = bpf_map_lookup_elem(&start, &tid);
65 if (tsp == 0)
66 return 0;
67 bpf_map_delete_elem(&start, &tid);
68
69 delta = bpf_ktime_get_ns() - *tsp;
70 if ((delta < min_block_ns) || (delta > max_block_ns))
71 return 0;
72
73 key.w_k_stack_id = bpf_get_stackid(ctx, &stackmap, 0);
74 bpf_probe_read_kernel(&key.target, sizeof(key.target), p->comm);
75 bpf_get_current_comm(&key.waker, sizeof(key.waker));
76
77 count_key = bpf_map_lookup_or_try_init(&counts, &key, &zero);
78 if (count_key)
79 __atomic_add_fetch(count_key, delta, __ATOMIC_RELAXED);
80
81 return 0;
82 }
83
84
85 SEC("tp_btf/sched_switch")
BPF_PROG(sched_switch,bool preempt,struct task_struct * prev,struct task_struct * next)86 int BPF_PROG(sched_switch, bool preempt, struct task_struct *prev, struct task_struct *next)
87 {
88 return offcpu_sched_switch(prev);
89 }
90
91 SEC("tp_btf/sched_wakeup")
BPF_PROG(sched_wakeup,struct task_struct * p)92 int BPF_PROG(sched_wakeup, struct task_struct *p)
93 {
94 return wakeup(ctx, p);
95 }
96
97 char LICENSE[] SEC("license") = "GPL";
98