aiops/agent/collector/network.go

// Package collector 提供服务器各项系统指标的采集能力。
// 本文件实现网络连接数、网卡流量等网络指标采集。
//
// 署名 Hermes (extended from Logan.Li pattern)
package collector

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

	"ezops/agent/model"
)

// netSample 保存上一次的网络采样数据，用于计算速率
type netSample struct {
	time     time.Time
	rxBytes  uint64
	txBytes  uint64
}

// NetworkCollector 网络指标采集器
type NetworkCollector struct {
	mu   sync.Mutex
	prev *netSample
}

// NewNetworkCollector 创建网络采集器
func NewNetworkCollector() *NetworkCollector {
	return &NetworkCollector{}
}

// Collect 采集网络指标（并发连接数、TIME_WAIT、网卡 rx/tx B/s）
func (c *NetworkCollector) Collect() (model.NetworkStat, error) {
	stat := model.NetworkStat{}

	// 1. 采集 TCP 连接统计
	if tcpStat, err := readSockStat(); err == nil {
		stat.TCPEstablished = tcpStat.established
		stat.TCPTimeWait = tcpStat.timeWait
		stat.TCPTotal = tcpStat.total
	}
	stat.TCPCloseWait = countCloseWait()

	// 2. 采集网卡流量（取所有非 lo 网卡合计）
	curRx, curTx, nicName, err := readNetDevTotal()
	if err != nil {
		return stat, fmt.Errorf("读取网卡流量失败: %w", err)
	}
	stat.NicName = nicName

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

	now := time.Now()
	cur := &netSample{time: now, rxBytes: curRx, txBytes: curTx}

	if c.prev != nil && now.After(c.prev.time) {
		elapsed := now.Sub(c.prev.time).Seconds()
		if elapsed > 0 && cur.rxBytes >= c.prev.rxBytes && cur.txBytes >= c.prev.txBytes {
			stat.NetRxBPS = float64(cur.rxBytes-c.prev.rxBytes) / elapsed
			stat.NetTxBPS = float64(cur.txBytes-c.prev.txBytes) / elapsed
		}
	}
	c.prev = cur

	return stat, nil
}

// tcpSockStat 解析 /proc/net/sockstat 或 /proc/net/tcp 的结果
type tcpSockStat struct {
	established int
	timeWait    int
	total       int
}

// readSockStat 从 /proc/net/sockstat 读取 TCP 连接统计
func readSockStat() (tcpSockStat, error) {
	var result tcpSockStat

	f, err := os.Open("/proc/net/sockstat")
	if err != nil {
		// Fallback: try /proc/net/sockstat6
		return result, err
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "TCP:") {
			continue
		}
		// 格式: "TCP: inuse 123 orphan 0 tw 45 alloc 200 mem 30"
		fields := strings.Fields(line)
		for i, f := range fields {
			switch f {
			case "inuse":
				if i+1 < len(fields) {
					result.established, _ = strconv.Atoi(fields[i+1])
				}
			case "tw":
				if i+1 < len(fields) {
					result.timeWait, _ = strconv.Atoi(fields[i+1])
				}
			case "alloc":
				if i+1 < len(fields) {
					result.total, _ = strconv.Atoi(fields[i+1])
				}
			}
		}
		result.total = result.established // total = established count (inuse)
		break
	}

	return result, scanner.Err()
}

// readNetDevTotal 读取 /proc/net/dev，取第一个非 lo 且活跃的物理网卡总流量
func readNetDevTotal() (rxBytes, txBytes uint64, nicName string, err error) {
	f, err := os.Open("/proc/net/dev")
	if err != nil {
		return 0, 0, "", err
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	// Skip header lines (first 2 lines)
	lineNum := 0
	for scanner.Scan() {
		lineNum++
		if lineNum <= 2 {
			continue
		}
		line := strings.TrimSpace(scanner.Text())
		parts := strings.SplitN(line, ":", 2)
		if len(parts) != 2 {
			continue
		}
		name := strings.TrimSpace(parts[0])
		// Skip loopback and virtual interfaces
		if name == "lo" || strings.HasPrefix(name, "veth") || strings.HasPrefix(name, "docker") ||
			strings.HasPrefix(name, "br-") || strings.HasPrefix(name, "virbr") {
			continue
		}

		fields := strings.Fields(parts[1])
		if len(fields) < 10 {
			continue
		}

		// fields[0]=rx_bytes, fields[8]=tx_bytes
		rb, _ := strconv.ParseUint(fields[0], 10, 64)
		tb, _ := strconv.ParseUint(fields[8], 10, 64)

		if rb > 0 || tb > 0 {
			nicName = name
			rxBytes += rb
			txBytes += tb
		}
	}

	if nicName == "" {
		return 0, 0, "eth0", nil
	}

	return rxBytes, txBytes, nicName, scanner.Err()
}

// countCloseWait 从 /proc/net/tcp 统计 CLOSE_WAIT 连接数
func countCloseWait() int {
	count := 0
	for _, path := range []string{"/proc/net/tcp", "/proc/net/tcp6"} {
		data, err := os.ReadFile(path)
		if err != nil {
			continue
		}
		for _, line := range strings.Split(string(data), "\n") {
			fields := strings.Fields(line)
			if len(fields) >= 4 && fields[3] == "08" {
				count++
			}
		}
	}
	return count
}