aiops/agent/main.go

// Package main 是 AIOps Agent 的入口。
// Agent 定时采集 CPU、内存、磁盘及监听端口指标，并通过 HTTP 上报至 AIOps 后端。
//
// 以下变量可在编译时通过 -ldflags "-X main.version=... -X main.buildTime=..." 注入。
//
// 署名 Logan.Li
package main

import (
	"context"
	"flag"
	"log"
	"os"
	"os/signal"
	"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 后端上报地址（含 scheme 与端口）")
	interval := flag.Int("interval", 10, "采集与上报间隔，单位：秒")
	token := flag.String("token", "", "Bearer 鉴权令牌")
	hostnameFlag := flag.String("hostname", "", "上报主机名，默认使用系统主机名")
	flag.Parse()

	// 统一日志格式：带微秒时间戳，便于排查问题
	logger := log.New(os.Stdout, "[aiops-agent] ", log.LstdFlags|log.Lmicroseconds)

	// 解析主机名：优先使用 -hostname，否则回退到 os.Hostname()
	hostname := *hostnameFlag
	if hostname == "" {
		if h, err := os.Hostname(); err == nil {
			hostname = h
		} else {
			hostname = "unknown"
		}
	}

	// 参数校验
	if *server == "" {
		logger.Fatal("参数 -server 不能为空")
	}
	if *interval < 1 {
		logger.Fatalf("参数 -interval 必须 >= 1，当前为 %d", *interval)
	}

	logger.Printf("AIOps Agent 启动: version=%s buildTime=%s", version, buildTime)
	logger.Printf("配置: server=%s interval=%ds hostname=%q token=%s",
		*server, *interval, hostname, maskToken(*token))

	// 初始化各采集器与上报器
	cpuCollector := collector.NewCPUCollector()
	memCollector := collector.NewMemCollector()
	diskCollector := collector.NewDiskCollector()
	portCollector := collector.NewPortCollector()
	netCollector := collector.NewNetworkCollector()
	diskIOCollector := collector.NewDiskIOCollector()
	rep := reporter.New(*server, *token, logger)

	// 初始化任务执行器与任务客户端（与 reporter 复用相同的 server/token/hostname）
	taskExecutor := task.New(logger)
	taskClient := task.NewClient(*server, *token, hostname, logger)

	// 监听 SIGINT / SIGTERM，支持优雅退出
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	ticker := time.NewTicker(time.Duration(*interval) * time.Second)
	defer ticker.Stop()

	// collectAndReport 执行单次采集并上报；各项采集独立进行，单项失败不影响整体上报
	collectAndReport := func() {
		report := &model.Report{
			Hostname:  hostname,
			Timestamp: time.Now().UTC().Format(time.RFC3339),
		}

		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
		}

		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
		}

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

		if err := rep.Send(ctx, report); err != nil {
			logger.Printf("上报失败: %v", err)
			return
		}
		logger.Printf("上报成功: cpu=%.2f%% mem=%.2f%% disks=%d ports=%d net=%dconn/%dwait netIO=%.0f/%.0fB/s diskIO=%.0f/%.0fB/s",
			report.CPU.UsagePercent, report.Memory.UsagePercent, len(report.Disks), len(report.Ports),
			report.Network.TCPEstablished, report.Network.TCPTimeWait,
			report.Network.NetRxBPS, report.Network.NetTxBPS,
			report.DiskIO.ReadBPS, report.DiskIO.WriteBPS)
	}

	// pollAndExecute 执行单次任务轮询：拉取待执行任务 → 逐个执行 → 回传结果。
	// 网络错误只记录日志，不 panic，下一个周期继续重试。
	pollAndExecute := func() {
		tasks, err := taskClient.Poll(ctx)
		if err != nil {
			// 上下文已取消（Agent 退出中）属于正常情况，不输出告警
			if ctx.Err() != nil {
				return
			}
			logger.Printf("任务轮询失败: %v", err)
			return
		}
		if len(tasks) == 0 {
			return // 暂无任务，静默返回
		}
		logger.Printf("拉取到 %d 个待执行任务", len(tasks))

		for _, t := range tasks {
			// Agent 退出时跳过后续任务，避免退出延迟
			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)
			}
		}
	}

	// 启动任务轮询协程：每 5 秒轮询一次，与指标采集并行运行。
	// 使用独立的 5 秒固定间隔，不受 -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
			}
		}
	}()

	// 主协程：启动后立即执行一次指标采集，再按 ticker 周期执行
	collectAndReport()
	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] + "****"
}