aiops/agent/main.go

// Package main 是 AIOps Agent 的入口。
// Agent 定时采集 CPU、内存、磁盘、端口、进程 TOP、Swap、FD、上下文切换等指标，
// 并通过 HTTP 上报至 AIOps 后端。
//
// 优化要点（1000 台服务器规模）：
//   - 启动随机 Jitter 避免惊群效应
//   - 连接池复用 HTTP Keep-Alive
//   - 进程 TOP 限制为 3（可配置）
//   - Agent 自检心跳（序列号丢包检测）
//   - 端口变化率采集（增量检测，减少 payload）
//
// 署名 Logan.Li
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"ezops/agent/collector"
	"ezops/agent/model"
	"ezops/agent/reporter"
	"ezops/agent/task"
)

var (
	version   = "dev"
	buildTime = "unknown"
)

func main() {
	server := flag.String("server", "http://localhost:8080", "AIOps 后端上报地址")
	interval := flag.Int("interval", 10, "采集与上报间隔，单位：秒")
	token := flag.String("token", "", "Bearer 鉴权令牌")
	hostnameFlag := flag.String("hostname", "", "上报主机名，默认使用系统主机名")
	topN := flag.Int("top-n", 3, "进程 TOP N 数量")
	flag.Parse()

	logger := log.New(os.Stdout, "[aiops-agent] ", log.LstdFlags|log.Lmicroseconds)

	hostname := *hostnameFlag
	if hostname == "" {
		if h, err := os.Hostname(); err == nil {
			hostname = h
		} else {
			hostname = "unknown"
		}
	}

	if *server == "" {
		logger.Fatal("参数 -server 不能为空")
	}
	if *interval < 5 {
		logger.Fatalf("参数 -interval 必须 >= 5，当前为 %d（1000 台规模推荐 30~60s）", *interval)
	}

	logger.Printf("AIOps Agent v%s (built %s) go=%s", version, buildTime, runtime.Version())
	logger.Printf("配置: server=%s interval=%ds top-n=%d hostname=%q", *server, *interval, *topN, hostname)

	// ═══ 初始化采集器 ═══
	cpuCollector := collector.NewCPUCollector()
	memCollector := collector.NewMemCollector()
	diskCollector := collector.NewDiskCollector()
	portCollector := collector.NewPortCollector()
	netCollector := collector.NewNetworkCollector()
	diskIOCollector := collector.NewDiskIOCollector()
	procTopCollector := collector.NewProcessTopCollector()
	swapCollector := collector.NewSwapCollector()
	rep := reporter.New(*server, *token, logger)

	// 系统静态信息（仅启动时一次）
	sysInfo := collector.CollectSystem()
	logger.Printf("系统信息: OS=%q Kernel=%q", sysInfo.OSVersion, sysInfo.Kernel)

	// 探测本机 IP
	localIP := detectLocalIP()
	logger.Printf("探测本机 IP: %s", localIP)

	sysReportCount := 0 // 前 3 次上报携带系统信息

	// 任务执行器
	taskExecutor := task.New(logger)
	taskClient := task.NewClient(*server, *token, hostname, logger)

	// ═══ 启动随机 Jitter：避免 1000 台同时上报 ═══
	startupJitter := time.Duration(rand.Intn(*interval)) * time.Second
	logger.Printf("启动 Jitter: %v（防止惊群效应）", startupJitter)

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	// Agent 自检序列号
	var reportSeq uint64

	// collectAndReport 执行单次采集并上报
	collectAndReport := func() {
		report := &model.Report{
			Hostname:     hostname,
			IP:           localIP,
			Timestamp:    time.Now().UTC().Format(time.RFC3339),
			AgentVersion: version,
		}

		// 系统信息（前 3 次）
		if sysReportCount < 3 && (sysInfo.OSVersion != "" || sysInfo.Kernel != "") {
			report.System = &model.SystemStat{
				OSVersion:    sysInfo.OSVersion,
				Kernel:       sysInfo.Kernel,
				Uptime:       sysInfo.UptimeSeconds,
			}
			sysReportCount++
		}

		// CPU
		if stat, err := cpuCollector.Collect(); err != nil {
			logger.Printf("CPU 采集异常: %v", err)
		} else {
			report.CPU = stat
		}

		// 内存
		if stat, err := memCollector.Collect(); err != nil {
			logger.Printf("内存采集异常: %v", err)
		} else {
			report.Memory = stat
		}

		// Swap（仅当系统配置了 swap 才上报）
		if swap, err := swapCollector.Collect(); err != nil {
			logger.Printf("Swap 采集异常: %v", err)
		} else if swap != nil {
			report.Swap = swap
		}

		// 磁盘
		if disks, err := diskCollector.Collect(); err != nil {
			logger.Printf("磁盘采集异常: %v", err)
		} else {
			report.Disks = disks
		}

		// 监听端口
		if ports, err := portCollector.Collect(); err != nil {
			logger.Printf("端口采集异常: %v", err)
		} else {
			report.Ports = ports
		}

		// 网络
		if netStat, err := netCollector.Collect(); err != nil {
			logger.Printf("网络采集异常: %v", err)
		} else {
			report.Network = netStat
		}

		// 磁盘 IO
		if dioStat, err := diskIOCollector.Collect(); err != nil {
			logger.Printf("磁盘IO采集异常: %v", err)
		} else {
			report.DiskIO = dioStat
		}

		// 进程 TOP N
		if procs, err := procTopCollector.Collect(*topN); err != nil {
			logger.Printf("进程TOP采集异常: %v", err)
		} else {
			report.ProcessTop = procs
		}

		// 系统 FD 句柄
		if fdStat, err := collectFD(); err != nil {
			logger.Printf("FD采集异常: %v", err)
		} else {
			report.FD = fdStat
		}

		// 上下文切换
		if csStat, err := collectContextSwitches(); err != nil {
			// 静默跳过（某些内核/容器环境不可用）
		} else {
			report.ContextSwitches = csStat
		}

		// Agent 自检心跳
		reportSeq++
		var memStats runtime.MemStats
		runtime.ReadMemStats(&memStats)
		report.Heartbeat = model.AgentHeartbeat{
			Version:   version,
			UptimeSec: sysInfo.UptimeSeconds,
			SelfCPU:   0, // TODO: self-CPU via /proc/self/stat
			SelfRSSMB: float64(memStats.Alloc) / 1024.0 / 1024.0,
			ReportSeq: reportSeq,
		}

		// 上报
		if err := rep.Send(ctx, report); err != nil {
			logger.Printf("上报失败: %v", err)
			return
		}
		logger.Printf("上报成功: cpu=%.1f%% mem=%.1f%% swap=%.1f%% disks=%d ports=%d procs=%d fd=%d/%d cs=%.0f/s",
			report.CPU.UsagePercent,
			report.Memory.UsagePercent,
			func() float64 {
				if report.Swap != nil {
					return report.Swap.UsagePercent
				}
				return 0
			}(),
			len(report.Disks),
			len(report.Ports),
			len(report.ProcessTop),
			func() int {
				if report.FD != nil {
					return report.FD.Allocated
				}
				return 0
			}(),
			func() int64 {
				if report.FD != nil {
					return report.FD.Max
				}
				return 0
			}(),
			func() float64 {
				if report.ContextSwitches != nil {
					return report.ContextSwitches.CTXPerSec
				}
				return 0
			}(),
		)
	}

	// pollAndExecute 任务轮询
	pollAndExecute := func() {
		tasks, err := taskClient.Poll(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			logger.Printf("任务轮询失败: %v", err)
			return
		}
		if len(tasks) == 0 {
			return
		}
		logger.Printf("拉取到 %d 个待执行任务", len(tasks))
		for _, t := range tasks {
			if ctx.Err() != nil {
				return
			}
			result := taskExecutor.ExecuteTask(ctx, t)
			if err := taskClient.Submit(ctx, result); err != nil {
				if ctx.Err() != nil {
					return
				}
				logger.Printf("任务 %s 结果回传失败: %v", t.TaskID, err)
			} else {
				logger.Printf("任务 %s 结果回传成功 status=%s", t.TaskID, result.Status)
			}
		}
	}

	// ═══ 任务轮询协程（5s 间隔，不受 -interval 影响）═══
	const taskPollInterval = 5 * time.Second
	go func() {
		pollAndExecute()
		ticker := time.NewTicker(taskPollInterval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				pollAndExecute()
			case <-ctx.Done():
				return
			}
		}
	}()

	// ═══ 主循环：启动先 sleep jitter，再按 interval 采集 ═══
	time.Sleep(startupJitter)
	collectAndReport()
	ticker := time.NewTicker(time.Duration(*interval) * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			collectAndReport()
		case <-ctx.Done():
			logger.Println("收到退出信号，Agent 正常关闭")
			return
		}
	}
}

// maskToken 对令牌做脱敏处理
func maskToken(token string) string {
	if token == "" {
		return "(empty)"
	}
	if len(token) <= 4 {
		return "****"
	}
	return token[:4] + "****"
}

// detectLocalIP 探测本机主 IP
func detectLocalIP() string {
	conn, err := net.Dial("udp", "8.8.8.8:53")
	if err == nil {
		defer conn.Close()
		if addr, ok := conn.LocalAddr().(*net.UDPAddr); ok && addr.IP != nil {
			return addr.IP.String()
		}
	}
	ifaces, err := net.Interfaces()
	if err == nil {
		for _, iface := range ifaces {
			if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 {
				continue
			}
			addrs, _ := iface.Addrs()
			for _, addr := range addrs {
				if ipnet, ok := addr.(*net.IPNet); ok && ipnet.IP.To4() != nil {
					return ipnet.IP.String()
				}
			}
		}
	}
	return "unknown"
}

// ══════════════════════════════════════════════
// 补充采集器（内联简单实现，不单独建文件）
// ══════════════════════════════════════════════

// collectFD 读取 /proc/sys/fs/file-nr 和 /proc/sys/fs/file-max
func collectFD() (*model.SocketFDStat, error) {
	alloc, err := os.ReadFile("/proc/sys/fs/file-nr")
	if err != nil {
		return nil, err
	}
	var allocated, _unused1, _unused2 int
	fmt.Sscanf(string(alloc), "%d %d %d", &allocated, &_unused1, &_unused2)

	max, err := os.ReadFile("/proc/sys/fs/file-max")
	if err != nil {
		return nil, err
	}
	var maxFD int64
	fmt.Sscanf(string(max), "%d", &maxFD)

	stat := &model.SocketFDStat{
		Allocated: allocated,
		Max:       maxFD,
	}
	if maxFD > 0 {
		stat.UsagePct = float64(allocated) / float64(maxFD) * 100.0
	}
	return stat, nil
}

// collectContextSwitches 读取 /proc/stat 中的 ctxt 和 intr 行
// 需要两次采样才能计算速率；这里使用简单的累计值（首次返回 0）
func collectContextSwitches() (*model.ContextSwitchStat, error) {
	data, err := os.ReadFile("/proc/stat")
	if err != nil {
		return nil, err
	}
	stat := &model.ContextSwitchStat{}
	lines := string(data)
	for _, line := range splitLines(lines) {
		if hasPrefix(line, "ctxt ") {
			fmt.Sscanf(line, "ctxt %f", &stat.CTXPerSec)
		}
		if hasPrefix(line, "intr ") {
			fmt.Sscanf(line, "intr %f", &stat.IRQPerSec)
		}
	}
	// NOTE: 这里返回的是累计值而非速率；后端可用两次上报的差值计算速率
	return stat, nil
}

func splitLines(s string) []string {
	var result []string
	start := 0
	for i := 0; i < len(s); i++ {
		if s[i] == '\n' {
			result = append(result, s[start:i])
			start = i + 1
		}
	}
	if start < len(s) {
		result = append(result, s[start:])
	}
	return result
}

func hasPrefix(s, prefix string) bool {
	return len(s) >= len(prefix) && s[:len(prefix)] == prefix
}