aiops/agent/collector/diskio.go

// Package collector — 磁盘 IO 指标采集器
// 通过读取 /proc/diskstats 获取磁盘读写 BPS 和 IOPS。
//
// 署名 Hermes
package collector

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"ezops/agent/model"
)

// diskIOSample 保存上一次采样，用于计算速率
type diskIOSample struct {
	time       time.Time
	readSects  uint64
	writeSects uint64
	readOps    uint64
	writeOps   uint64
}

// DiskIOCollector 磁盘 IO 采集器
type DiskIOCollector struct {
	mu   sync.Mutex
	prev *diskIOSample
}

// NewDiskIOCollector 创建磁盘 IO 采集器
func NewDiskIOCollector() *DiskIOCollector {
	return &DiskIOCollector{}
}

// Collect 采集磁盘 IO 指标
func (c *DiskIOCollector) Collect() (model.DiskIOStat, error) {
	stat := model.DiskIOStat{}

	rSects, wSects, rOps, wOps, err := readDiskStats()
	if err != nil {
		return stat, fmt.Errorf("读取磁盘 IO 失败: %w", err)
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	now := time.Now()
	cur := &diskIOSample{
		time:       now,
		readSects:  rSects,
		writeSects: wSects,
		readOps:    rOps,
		writeOps:   wOps,
	}

	if c.prev != nil && now.After(c.prev.time) {
		elapsed := now.Sub(c.prev.time).Seconds()
		if elapsed > 0 {
			sectorSize := uint64(512) // 标准扇区 512 bytes
			stat.ReadBPS = float64((cur.readSects-c.prev.readSects)*sectorSize) / elapsed
			stat.WriteBPS = float64((cur.writeSects-c.prev.writeSects)*sectorSize) / elapsed
			stat.ReadIOPS = float64(cur.readOps-c.prev.readOps) / elapsed
			stat.WriteIOPS = float64(cur.writeOps-c.prev.writeOps) / elapsed
		}
	}
	c.prev = cur

	return stat, nil
}

// readDiskStats 读取 /proc/diskstats，汇总所有物理磁盘的读写数据
// fields[2]=设备名, fields[3]=读完成次数, fields[5]=读扇区数, fields[7]=写完成次数, fields[9]=写扇区数
func readDiskStats() (rSects, wSects, rOps, wOps uint64, err error) {
	f, err := os.Open("/proc/diskstats")
	if err != nil {
		return 0, 0, 0, 0, err
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		fields := strings.Fields(line)
		if len(fields) < 14 {
			continue
		}

		dev := fields[2]
		// 只采集物理磁盘 (sdX, nvmeXn1, vdX)，跳过分区和 loop 设备
		if strings.HasPrefix(dev, "loop") ||
			(strings.HasPrefix(dev, "sd") && len(dev) > 3) || // sda1, sda2 等分区
			(strings.HasPrefix(dev, "nvme") && strings.Contains(dev, "p")) || // nvme0n1p1
			(strings.HasPrefix(dev, "ram")) ||
			strings.HasPrefix(dev, "dm-") {
			continue
		}

		// field 3 = 读完成次数 (rd_ios)
		// field 5 = 读扇区数 (rd_sectors)
		// field 7 = 写完成次数 (wr_ios)
		// field 9 = 写扇区数 (wr_sectors)
		if v, e := strconv.ParseUint(fields[3], 10, 64); e == nil {
			rOps += v
		}
		if v, e := strconv.ParseUint(fields[5], 10, 64); e == nil {
			rSects += v
		}
		if v, e := strconv.ParseUint(fields[7], 10, 64); e == nil {
			wOps += v
		}
		if v, e := strconv.ParseUint(fields[9], 10, 64); e == nil {
			wSects += v
		}
	}

	return rSects, wSects, rOps, wOps, scanner.Err()
}