Katyusha's blog

1
// SPDX-License-Identifier: GPL-2.0
2
//
3
// 场景①：CPU 异常占用 / 调度延迟 采集探针 (CO-RE, libbpf)
4
//
5
// 通过 sched_switch / sched_wakeup 两个稳定 ABI tracepoint，按线程(tid)聚合：
6
//   - run_ns ：累计在 CPU 上运行的时间   -> 计算 CPU 占用率
7
//   - runq_ns：累计在运行队列等待的时间   -> 计算调度等待时间
8
//   - ctx    ：被切出 CPU 的次数          -> 计算上下文切换频次
9
// 这些正是赛题"CPU 异常场景"列出的关键证据点。
10
// 内核态只做聚合，用户态按窗口读取并做差分，开销极低（无 per-event 上送）。
11

12
#include "vmlinux.h"
13
#include <bpf/bpf_helpers.h>
14

15
char LICENSE[] SEC("license") = "GPL";
16

17
#define TASK_COMM_LEN 16
18

19
struct task_stat {
20
  __u64 run_ns;   // 累计在 CPU 上运行的纳秒数
21
  __u64 runq_ns;  // 累计运行队列等待纳秒数
22
  __u64 ctx;      // 上下文切换(被切出)次数
23
  char  comm[TASK_COMM_LEN];
24
};
25

26
// tid -> 本次上 CPU 的时间戳
27
struct {
28
  __uint(type, BPF_MAP_TYPE_HASH);
29
  __uint(max_entries, 16384);
30
  __type(key, __u32);
31
  __type(value, __u64);
32
} oncpu_start SEC(".maps");
33

34
// tid -> 被唤醒(入队)的时间戳
35
struct {
36
  __uint(type, BPF_MAP_TYPE_HASH);
37
  __uint(max_entries, 16384);
38
  __type(key, __u32);
39
  __type(value, __u64);
40
} wakeup_ts SEC(".maps");
41

42
// tid -> 累计统计
43
struct {
44
  __uint(type, BPF_MAP_TYPE_HASH);
45
  __uint(max_entries, 16384);
46
  __type(key, __u32);
47
  __type(value, struct task_stat);
48
} stats SEC(".maps");
49

50
// tracepoint 上下文布局（稳定 ABI，直接读取，无需 CO-RE 重定位）
51
struct sched_switch_tp {
52
  __u64 pad;
53
  char  prev_comm[TASK_COMM_LEN];
54
  __s32 prev_pid;
55
  __s32 prev_prio;
56
  __s64 prev_state;
57
  char  next_comm[TASK_COMM_LEN];
58
  __s32 next_pid;
59
  __s32 next_prio;
60
};
61

62
struct sched_wakeup_tp {
63
  __u64 pad;
64
  char  comm[TASK_COMM_LEN];
65
  __s32 pid;
66
  __s32 prio;
67
  __s32 target_cpu;
68
};
69

70
static __always_inline struct task_stat *get_stat(__u32 tid)
71
{
72
  struct task_stat *st = bpf_map_lookup_elem(&stats, &tid);
73
  if (st)
74
    return st;
75
  struct task_stat zero = {};
76
  bpf_map_update_elem(&stats, &tid, &zero, BPF_NOEXIST);
77
  return bpf_map_lookup_elem(&stats, &tid);
78
}
79

80
SEC("tracepoint/sched/sched_wakeup")
81
int handle_wakeup(struct sched_wakeup_tp *ctx)
82
{
83
  __u32 tid = (__u32)ctx->pid;
84
  if (tid == 0)
85
    return 0;
86
  __u64 now = bpf_ktime_get_ns();
87
  bpf_map_update_elem(&wakeup_ts, &tid, &now, BPF_ANY);
88
  return 0;
89
}
90

91
SEC("tracepoint/sched/sched_switch")
92
int handle_switch(struct sched_switch_tp *ctx)
93
{
94
  __u64 now = bpf_ktime_get_ns();
95
  __u32 prev = (__u32)ctx->prev_pid;
96
  __u32 next = (__u32)ctx->next_pid;
97

98
  // prev 被切出 CPU：累计运行时间与切换次数
99
  if (prev != 0) {
100
    __u64 *ts = bpf_map_lookup_elem(&oncpu_start, &prev);
101
    struct task_stat *st = get_stat(prev);
102
    if (st) {
103
      if (ts && now > *ts)
104
        st->run_ns += now - *ts;
105
      st->ctx += 1;
106
      bpf_probe_read_kernel_str(&st->comm, sizeof(st->comm),
107
              ctx->prev_comm);
108
    }
109
  }
110

111
  // next 上 CPU：记录起点；若有唤醒时间戳则累计运行队列等待
112
  if (next != 0) {
113
    bpf_map_update_elem(&oncpu_start, &next, &now, BPF_ANY);
114
    __u64 *wts = bpf_map_lookup_elem(&wakeup_ts, &next);
115
    if (wts && now > *wts) {
116
      struct task_stat *st = get_stat(next);
117
      if (st)
118
        st->runq_ns += now - *wts;
119
      bpf_map_delete_elem(&wakeup_ts, &next);
120
    }
121
  }
122
  return 0;
123
}