aiops/agent/collector/health_probe.go

// Package collector 的健康探针采集器。
// 在后端下发的探针配置驱动下，本地执行 HTTP GET / TCP 连接 / 命令执行，
// 并将结果返回供上报。
//
// 署名 Logan.Li
package collector

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"os/exec"
	"strings"
	"time"

	"ezops/agent/model"
)

// HealthProbeCollector 健康探针执行器
type HealthProbeCollector struct {
	httpClient *http.Client
}

// NewHealthProbeCollector 创建采集器实例
func NewHealthProbeCollector() *HealthProbeCollector {
	return &HealthProbeCollector{
		httpClient: &http.Client{
			Timeout: 10 * time.Second,
			Transport: &http.Transport{
				DialContext: (&net.Dialer{
					Timeout: 5 * time.Second,
				}).DialContext,
			},
		},
	}
}

// Execute 执行所有探针配置，返回结果列表。
// parentCtx 用于传播取消信号——探针执行尽量尊重主循环的生命周期。
func (h *HealthProbeCollector) Execute(parentCtx context.Context, configs []model.ProbeConfig) []model.ProbeResult {
	results := make([]model.ProbeResult, 0, len(configs))

	for _, cfg := range configs {
		// 检查父上下文是否已取消
		if err := parentCtx.Err(); err != nil {
			break
		}

		result := model.ProbeResult{
			ProbeID:   cfg.ID,
			Name:      cfg.Name,
			ProbeType: cfg.ProbeType,
		}

		start := time.Now()

		switch cfg.ProbeType {
		case "http_get":
			result = h.execHTTP(parentCtx, cfg, start)
		case "tcp_socket":
			result = h.execTCP(cfg, start)
		case "exec_command":
			result = h.execCmd(parentCtx, cfg, start)
		default:
			result.Status = "down"
			result.Message = fmt.Sprintf("unknown probe type: %s", cfg.ProbeType)
		}

		results = append(results, result)
	}

	return results
}

// execHTTP 执行 HTTP 探测，接收父 context
func (h *HealthProbeCollector) execHTTP(parentCtx context.Context, cfg model.ProbeConfig, start time.Time) model.ProbeResult {
	result := model.ProbeResult{
		ProbeID:   cfg.ID,
		Name:      cfg.Name,
		ProbeType: cfg.ProbeType,
	}

	timeout := time.Duration(cfg.Timeout) * time.Second
	if timeout <= 0 {
		timeout = 10 * time.Second
	}

	ctx, cancel := context.WithTimeout(parentCtx, timeout)
	defer cancel()

	method := cfg.HTTPMethod
	if method == "" {
		method = "GET"
	}

	req, err := http.NewRequestWithContext(ctx, method, cfg.HTTPURL, nil)
	if err != nil {
		result.Status = "down"
		result.Message = fmt.Sprintf("构造请求失败: %v", err)
		result.LatencyMs = int(time.Since(start).Milliseconds())
		return result
	}
	req.Header.Set("User-Agent", "aiops-agent-health-probe/2.1")

	resp, err := h.httpClient.Do(req)
	latencyMs := int(time.Since(start).Milliseconds())
	result.LatencyMs = latencyMs

	if err != nil {
		result.Status = "down"
		result.Message = fmt.Sprintf("HTTP请求失败: %v", err)
		return result
	}
	defer resp.Body.Close()

	expectCode := cfg.HTTPExpectCode
	if expectCode == 0 {
		expectCode = 200
	}
	result.StatusCode = resp.StatusCode

	if resp.StatusCode != expectCode {
		result.Status = "down"
		result.Message = fmt.Sprintf("状态码不匹配: got %d, want %d", resp.StatusCode, expectCode)
		return result
	}

	// 可选：检查响应体
	if cfg.HTTPExpectBody != "" {
		bodyBytes, err := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
		if err != nil {
			result.Status = "down"
			result.Message = fmt.Sprintf("读取响应体失败: %v", err)
			return result
		}
		if !strings.Contains(string(bodyBytes), cfg.HTTPExpectBody) {
			result.Status = "down"
			result.Message = fmt.Sprintf("响应体不包含期望字符串: %q", cfg.HTTPExpectBody)
			return result
		}
	}

	result.Status = "up"
	result.Message = fmt.Sprintf("HTTP %d OK (%dms)", resp.StatusCode, latencyMs)
	return result
}

// execTCP 执行 TCP 连接探测
func (h *HealthProbeCollector) execTCP(cfg model.ProbeConfig, start time.Time) model.ProbeResult {
	result := model.ProbeResult{
		ProbeID:   cfg.ID,
		Name:      cfg.Name,
		ProbeType: cfg.ProbeType,
	}

	timeout := time.Duration(cfg.Timeout) * time.Second
	if timeout <= 0 {
		timeout = 10 * time.Second
	}

	addr := fmt.Sprintf("%s:%d", cfg.TCPHost, cfg.TCPPort)
	conn, err := net.DialTimeout("tcp", addr, timeout)
	latencyMs := int(time.Since(start).Milliseconds())
	result.LatencyMs = latencyMs

	if err != nil {
		result.Status = "down"
		result.Message = fmt.Sprintf("TCP连接失败 %s: %v", addr, err)
		return result
	}
	conn.Close()

	result.Status = "up"
	result.Message = fmt.Sprintf("TCP %s connected (%dms)", addr, latencyMs)
	return result
}

// execCmd 执行命令探测（exit code 0 = up），接收父 context
func (h *HealthProbeCollector) execCmd(parentCtx context.Context, cfg model.ProbeConfig, start time.Time) model.ProbeResult {
	result := model.ProbeResult{
		ProbeID:   cfg.ID,
		Name:      cfg.Name,
		ProbeType: cfg.ProbeType,
	}

	timeout := time.Duration(cfg.Timeout) * time.Second
	if timeout <= 0 {
		timeout = 10 * time.Second
	}

	ctx, cancel := context.WithTimeout(parentCtx, timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, "sh", "-c", cfg.ExecCommand)
	output, err := cmd.CombinedOutput()
	latencyMs := int(time.Since(start).Milliseconds())
	result.LatencyMs = latencyMs
	result.StatusCode = cmd.ProcessState.ExitCode()

	if err != nil {
		result.Status = "down"
		// 截断输出避免过大
		outMsg := strings.TrimSpace(string(output))
		if len(outMsg) > 200 {
			outMsg = outMsg[:200] + "..."
		}
		result.Message = fmt.Sprintf("命令执行失败(exit=%d): %s", cmd.ProcessState.ExitCode(), outMsg)
		return result
	}

	result.Status = "up"
	result.Message = fmt.Sprintf("命令执行成功 (%dms)", latencyMs)
	return result
}