Katyusha's blog

1
package collector
2

3
import (
4
  "bufio"
5
  "fmt"
6
  "os"
7
  "strconv"
8
  "strings"
9
  "time"
10

11
  "github.com/cilium/ebpf/link"
12
  "github.com/cilium/ebpf/rlimit"
13
)
14

15
// memStat 与 mem.bpf.c 中的 struct mem_stat 二进制布局一致。
16
type memStat struct {
17
  DirectReclaimCount uint64
18
  DirectReclaimNs    uint64
19
  Comm               [16]byte
20
}
21

22
type fault struct{ maj, min uint64 }
23

24
// MemProc 是单个进程在窗口内的内存压力贡献。
25
type MemProc struct {
26
  Pid                uint32
27
  Comm               string
28
  DirectReclaimCount uint64  // 窗口内直接回收次数
29
  DirectReclaimMs    float64 // 窗口内直接回收耗时(ms)
30
  MajFlt             uint64  // 窗口内 major fault 增量
31
  MinFlt             uint64  // 窗口内 minor fault 增量
32
}
33

34
// MemSnapshot 是一个窗口的系统内存状态与 per-process 压力贡献。
35
type MemSnapshot struct {
36
  MemAvailablePct float64
37
  MemTotalKB      uint64
38
  MemAvailableKB  uint64
39
  KswapdWakes     uint64 // 窗口内增量
40
  Procs           []MemProc
41
}
42

43
// MemCollector 加载内存场景的 eBPF 程序并汇总系统内存压力。
44
type MemCollector struct {
45
  objs       memObjects
46
  links      []link.Link
47
  prev       map[uint32]memStat
48
  prevFault  map[uint32]fault
49
  prevKswapd uint64
50
}
51

52
// NewMemCollector 加载字节码、挂载 vmscan tracepoint。
53
func NewMemCollector() (*MemCollector, error) {
54
  if err := rlimit.RemoveMemlock(); err != nil {
55
    return nil, fmt.Errorf("remove memlock: %w", err)
56
  }
57
  c := &MemCollector{
58
    prev:      make(map[uint32]memStat),
59
    prevFault: make(map[uint32]fault),
60
  }
61
  if err := loadMemObjects(&c.objs, nil); err != nil {
62
    return nil, fmt.Errorf("load bpf objects: %w", err)
63
  }
64

65
  begin, err := link.Tracepoint("vmscan", "mm_vmscan_direct_reclaim_begin", c.objs.HandleDirectBegin, nil)
66
  if err != nil {
67
    c.Close()
68
    return nil, fmt.Errorf("attach direct_reclaim_begin: %w", err)
69
  }
70
  c.links = append(c.links, begin)
71

72
  end, err := link.Tracepoint("vmscan", "mm_vmscan_direct_reclaim_end", c.objs.HandleDirectEnd, nil)
73
  if err != nil {
74
    c.Close()
75
    return nil, fmt.Errorf("attach direct_reclaim_end: %w", err)
76
  }
77
  c.links = append(c.links, end)
78

79
  kw, err := link.Tracepoint("vmscan", "mm_vmscan_kswapd_wake", c.objs.HandleKswapdWake, nil)
80
  if err != nil {
81
    c.Close()
82
    return nil, fmt.Errorf("attach kswapd_wake: %w", err)
83
  }
84
  c.links = append(c.links, kw)
85
  return c, nil
86
}
87

88
// Close 卸载探针并释放资源。
89
func (c *MemCollector) Close() {
90
  for _, l := range c.links {
91
    _ = l.Close()
92
  }
93
  c.objs.Close()
94
}
95

96
// Poll 汇总一个窗口的系统内存状态与进程压力贡献。
97
func (c *MemCollector) Poll(interval time.Duration) (MemSnapshot, error) {
98
  total, avail := readMemInfo()
99
  snap := MemSnapshot{MemTotalKB: total, MemAvailableKB: avail}
100
  if total > 0 {
101
    snap.MemAvailablePct = float64(avail) / float64(total) * 100
102
  }
103

104
  // kswapd 唤醒(全局)增量
105
  var kw uint64
106
  if err := c.objs.KswapdWakes.Lookup(uint32(0), &kw); err == nil {
107
    snap.KswapdWakes = kw - c.prevKswapd
108
    c.prevKswapd = kw
109
  }
110

111
  cur := make(map[uint32]memStat)
112
  var key uint32
113
  var val memStat
114
  it := c.objs.MemStats.Iterate()
115
  for it.Next(&key, &val) {
116
    cur[key] = val
117
  }
118
  if err := it.Err(); err != nil {
119
    return snap, fmt.Errorf("iterate mem_stats: %w", err)
120
  }
121

122
  for pid, v := range cur {
123
    p := c.prev[pid]
124
    dCount := v.DirectReclaimCount - p.DirectReclaimCount
125
    dNs := v.DirectReclaimNs - p.DirectReclaimNs
126
    maj, min := readProcFaults(pid)
127
    pf := c.prevFault[pid]
128
    proc := MemProc{
129
      Pid:                pid,
130
      Comm:               commToString(v.Comm),
131
      DirectReclaimCount: dCount,
132
      DirectReclaimMs:    float64(dNs) / 1e6,
133
      MajFlt:             saturatingSub(maj, pf.maj),
134
      MinFlt:             saturatingSub(min, pf.min),
135
    }
136
    c.prevFault[pid] = fault{maj: maj, min: min}
137
    if proc.DirectReclaimCount > 0 || proc.MajFlt > 0 {
138
      snap.Procs = append(snap.Procs, proc)
139
    }
140
  }
141
  c.prev = cur
142
  return snap, nil
143
}
144

145
func saturatingSub(a, b uint64) uint64 {
146
  if a < b {
147
    return 0
148
  }
149
  return a - b
150
}
151

152
// readMemInfo 返回 /proc/meminfo 的 MemTotal、MemAvailable（kB）。
153
func readMemInfo() (total, avail uint64) {
154
  f, err := os.Open("/proc/meminfo")
155
  if err != nil {
156
    return 0, 0
157
  }
158
  defer f.Close()
159
  sc := bufio.NewScanner(f)
160
  for sc.Scan() {
161
    fields := strings.Fields(sc.Text())
162
    if len(fields) < 2 {
163
      continue
164
    }
165
    v, _ := strconv.ParseUint(fields[1], 10, 64)
166
    switch fields[0] {
167
    case "MemTotal:":
168
      total = v
169
    case "MemAvailable:":
170
      avail = v
171
    }
172
  }
173
  return total, avail
174
}
175

176
// readProcFaults 读取 /proc/<pid>/stat 的 minflt(字段10)、majflt(字段12)。
177
func readProcFaults(pid uint32) (maj, min uint64) {
178
  b, err := os.ReadFile(fmt.Sprintf("/proc/%d/stat", pid))
179
  if err != nil {
180
    return 0, 0
181
  }
182
  s := string(b)
183
  // comm 可能含空格/括号，跳到最后一个 ')' 之后再按空格切分。
184
  rp := strings.LastIndexByte(s, ')')
185
  if rp < 0 || rp+2 >= len(s) {
186
    return 0, 0
187
  }
188
  fields := strings.Fields(s[rp+2:]) // fields[0]=state(字段3)
189
  // 字段10(minflt) -> idx 7；字段12(majflt) -> idx 9
190
  if len(fields) > 9 {
191
    min, _ = strconv.ParseUint(fields[7], 10, 64)
192
    maj, _ = strconv.ParseUint(fields[9], 10, 64)
193
  }
194
  return maj, min
195
}