Katyusha's blog

1
package collector
2

3
import (
4
  "fmt"
5
  "time"
6

7
  "github.com/cilium/ebpf/link"
8
  "github.com/cilium/ebpf/rlimit"
9
)
10

11
// taskStat 与 cpu.bpf.c 中的 struct task_stat 二进制布局一致（无填充）。
12
type taskStat struct {
13
  RunNs  uint64
14
  RunqNs uint64
15
  Ctx    uint64
16
  Comm   [16]byte
17
}
18

19
// Sample 是单个线程在一个采样窗口内的派生指标。
20
type Sample struct {
21
  Pid        uint32  // 内核 pid 字段，实际为线程 id(tid)
22
  Comm       string  // 线程名
23
  CPUUtil    float64 // 单核 CPU 占用率（1.0 ≈ 占满一个核）
24
  CtxPerMin  float64 // 每分钟上下文切换次数
25
  RunqWaitUs float64 // 平均运行队列等待时间(微秒)
26
}
27

28
// CPUCollector 加载 CPU 场景的 eBPF 程序并按窗口读取聚合数据。
29
type CPUCollector struct {
30
  objs  cpuObjects
31
  links []link.Link
32
  prev  map[uint32]taskStat
33
}
34

35
// NewCPUCollector 加载字节码、挂载 tracepoint。
36
func NewCPUCollector() (*CPUCollector, error) {
37
  if err := rlimit.RemoveMemlock(); err != nil {
38
    return nil, fmt.Errorf("remove memlock: %w", err)
39
  }
40
  c := &CPUCollector{prev: make(map[uint32]taskStat)}
41
  if err := loadCpuObjects(&c.objs, nil); err != nil {
42
    return nil, fmt.Errorf("load bpf objects: %w", err)
43
  }
44
  sw, err := link.Tracepoint("sched", "sched_switch", c.objs.HandleSwitch, nil)
45
  if err != nil {
46
    c.Close()
47
    return nil, fmt.Errorf("attach sched_switch: %w", err)
48
  }
49
  c.links = append(c.links, sw)
50

51
  wk, err := link.Tracepoint("sched", "sched_wakeup", c.objs.HandleWakeup, nil)
52
  if err != nil {
53
    c.Close()
54
    return nil, fmt.Errorf("attach sched_wakeup: %w", err)
55
  }
56
  c.links = append(c.links, wk)
57
  return c, nil
58
}
59

60
// Close 卸载探针并释放资源。
61
func (c *CPUCollector) Close() {
62
  for _, l := range c.links {
63
    _ = l.Close()
64
  }
65
  c.objs.Close()
66
}
67

68
// Poll 读取 stats map，计算自上次调用以来的差分指标。
69
func (c *CPUCollector) Poll(interval time.Duration) ([]Sample, error) {
70
  cur := make(map[uint32]taskStat)
71
  var key uint32
72
  var val taskStat
73
  it := c.objs.Stats.Iterate()
74
  for it.Next(&key, &val) {
75
    cur[key] = val
76
  }
77
  if err := it.Err(); err != nil {
78
    return nil, fmt.Errorf("iterate stats: %w", err)
79
  }
80

81
  intervalNs := float64(interval.Nanoseconds())
82
  intervalMin := interval.Minutes()
83
  var samples []Sample
84
  for tid, v := range cur {
85
    var dRun, dRunq, dCtx uint64
86
    if p, ok := c.prev[tid]; ok {
87
      dRun = v.RunNs - p.RunNs
88
      dRunq = v.RunqNs - p.RunqNs
89
      dCtx = v.Ctx - p.Ctx
90
    } else {
91
      dRun, dRunq, dCtx = v.RunNs, v.RunqNs, v.Ctx
92
    }
93
    if dRun == 0 && dCtx == 0 {
94
      continue
95
    }
96
    s := Sample{
97
      Pid:     tid,
98
      Comm:    commToString(v.Comm),
99
      CPUUtil: float64(dRun) / intervalNs,
100
    }
101
    if intervalMin > 0 {
102
      s.CtxPerMin = float64(dCtx) / intervalMin
103
    }
104
    if dCtx > 0 {
105
      s.RunqWaitUs = float64(dRunq) / float64(dCtx) / 1000.0
106
    }
107
    samples = append(samples, s)
108
  }
109
  c.prev = cur
110
  return samples, nil
111
}
112

113
func commToString(b [16]byte) string {
114
  n := 0
115
  for n < len(b) && b[n] != 0 {
116
    n++
117
  }
118
  return string(b[:n])
119
}