Katyusha's blog

1
package collector
2

3
import (
4
  "bufio"
5
  "fmt"
6
  "os"
7
  "time"
8

9
  "github.com/cilium/ebpf/link"
10
  "github.com/cilium/ebpf/rlimit"
11
)
12

13
const ioNSlots = 32
14

15
// ioStat 与 block.bpf.c 中的 struct dev_stat 二进制布局一致。
16
type ioStat struct {
17
  Count      uint64
18
  TotalLatNs uint64
19
  MaxLatNs   uint64
20
  Bytes      uint64
21
  Inflight   int64
22
  Slots      [ioNSlots]uint64
23
}
24

25
// IOSample 是单个块设备在一个窗口内的派生 I/O 指标。
26
type IOSample struct {
27
  Dev            uint32
28
  DevName        string  // "8:0 sda"
29
  IOPS           float64
30
  AvgLatMs       float64
31
  P99LatMs       float64
32
  MaxLatMs       float64
33
  ThroughputMBps float64
34
  QueueDepth     int64
35
}
36

37
// IOCollector 加载 I/O 场景的 eBPF 程序并按窗口读取块层时延数据。
38
type IOCollector struct {
39
  objs    blockObjects
40
  links   []link.Link
41
  prev    map[uint32]ioStat
42
  devName map[uint32]string
43
}
44

45
// NewIOCollector 加载字节码、挂载块层 tracepoint、载入设备名表。
46
func NewIOCollector() (*IOCollector, error) {
47
  if err := rlimit.RemoveMemlock(); err != nil {
48
    return nil, fmt.Errorf("remove memlock: %w", err)
49
  }
50
  c := &IOCollector{
51
    prev:    make(map[uint32]ioStat),
52
    devName: loadPartitions(),
53
  }
54
  if err := loadBlockObjects(&c.objs, nil); err != nil {
55
    return nil, fmt.Errorf("load bpf objects: %w", err)
56
  }
57
  is, err := link.Tracepoint("block", "block_rq_issue", c.objs.HandleIssue, nil)
58
  if err != nil {
59
    c.Close()
60
    return nil, fmt.Errorf("attach block_rq_issue: %w", err)
61
  }
62
  c.links = append(c.links, is)
63

64
  cp, err := link.Tracepoint("block", "block_rq_complete", c.objs.HandleComplete, nil)
65
  if err != nil {
66
    c.Close()
67
    return nil, fmt.Errorf("attach block_rq_complete: %w", err)
68
  }
69
  c.links = append(c.links, cp)
70
  return c, nil
71
}
72

73
// Close 卸载探针并释放资源。
74
func (c *IOCollector) Close() {
75
  for _, l := range c.links {
76
    _ = l.Close()
77
  }
78
  c.objs.Close()
79
}
80

81
// Poll 读取 dev_stats，计算自上次调用以来的差分与 P99。
82
func (c *IOCollector) Poll(interval time.Duration) ([]IOSample, error) {
83
  cur := make(map[uint32]ioStat)
84
  var key uint32
85
  var val ioStat
86
  it := c.objs.DevStats.Iterate()
87
  for it.Next(&key, &val) {
88
    cur[key] = val
89
  }
90
  if err := it.Err(); err != nil {
91
    return nil, fmt.Errorf("iterate dev_stats: %w", err)
92
  }
93

94
  secs := interval.Seconds()
95
  var samples []IOSample
96
  for dev, v := range cur {
97
    p := c.prev[dev]
98
    dCount := v.Count - p.Count
99
    dLat := v.TotalLatNs - p.TotalLatNs
100
    dBytes := v.Bytes - p.Bytes
101
    if dCount == 0 {
102
      continue
103
    }
104
    s := IOSample{
105
      Dev:            dev,
106
      DevName:        c.devString(dev),
107
      MaxLatMs:       float64(v.MaxLatNs) / 1e6,
108
      QueueDepth:     v.Inflight,
109
      ThroughputMBps: float64(dBytes) / secs / 1e6,
110
    }
111
    if secs > 0 {
112
      s.IOPS = float64(dCount) / secs
113
    }
114
    s.AvgLatMs = float64(dLat) / float64(dCount) / 1e6
115
    s.P99LatMs = p99FromSlots(v.Slots, p.Slots, dCount)
116
    samples = append(samples, s)
117
  }
118
  c.prev = cur
119
  return samples, nil
120
}
121

122
// p99FromSlots 由 log2(ns) 直方图差分估算 P99（取该槽位上界 2^slot ns）。
123
func p99FromSlots(cur, prev [ioNSlots]uint64, total uint64) float64 {
124
  if total == 0 {
125
    return 0
126
  }
127
  target := float64(total) * 0.99
128
  var cum float64
129
  for slot := 0; slot < ioNSlots; slot++ {
130
    cum += float64(cur[slot] - prev[slot])
131
    if cum >= target {
132
      return float64(uint64(1)<<uint(slot)) / 1e6 // ns -> ms
133
    }
134
  }
135
  return float64(uint64(1)<<uint(ioNSlots-1)) / 1e6
136
}
137

138
func (c *IOCollector) devString(dev uint32) string {
139
  maj, min := dev>>20, dev&0xfffff
140
  if name, ok := c.devName[dev]; ok {
141
    return fmt.Sprintf("%d:%d %s", maj, min, name)
142
  }
143
  return fmt.Sprintf("%d:%d", maj, min)
144
}
145

146
// loadPartitions 从 /proc/partitions 构建 dev_t -> 设备名 映射。
147
func loadPartitions() map[uint32]string {
148
  m := make(map[uint32]string)
149
  f, err := os.Open("/proc/partitions")
150
  if err != nil {
151
    return m
152
  }
153
  defer f.Close()
154
  sc := bufio.NewScanner(f)
155
  for sc.Scan() {
156
    var maj, min uint32
157
    var blocks uint64
158
    var name string
159
    if n, _ := fmt.Sscanf(sc.Text(), "%d %d %d %s", &maj, &min, &blocks, &name); n == 4 {
160
      m[(maj<<20)|(min&0xfffff)] = name
161
    }
162
  }
163
  return m
164
}