141 lines
4.6 KiB
Go
141 lines
4.6 KiB
Go
// Package task 实现 AIOps Agent 的任务拉取与执行能力。
|
||
// Agent 通过轮询接口从后端获取待执行任务（命令/脚本），本地执行后回传结果。
|
||
//
|
||
// 署名 Logan.Li
|
||
package task
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"fmt"
|
||
"log"
|
||
"os/exec"
|
||
"time"
|
||
)
|
||
|
||
// 任务类型常量
|
||
const (
|
||
TypeCommand = "command" // 单条 shell 命令
|
||
TypeScript = "script" // 多行 bash 脚本
|
||
)
|
||
|
||
// Task 表示从后端拉取到的待执行任务。
|
||
type Task struct {
|
||
TaskID string `json:"task_id"` // 任务唯一标识，回传结果时使用
|
||
Type string `json:"type"` // 任务类型：command / script
|
||
Command string `json:"command"` // 命令内容（command 时为单条命令，script 时为脚本正文）
|
||
Timeout int `json:"timeout"` // 执行超时，单位：秒；<=0 表示不限制
|
||
}
|
||
|
||
// Result 表示任务执行结果，用于回传后端。
|
||
type Result struct {
|
||
TaskID string `json:"task_id"` // 对应的任务 ID
|
||
ExitCode int `json:"exit_code"` // 进程退出码；超时/启动失败时为 -1
|
||
Result string `json:"result"` // 合并后的 stdout+stderr 输出
|
||
Status string `json:"status"` // 执行状态：success / failed / timeout
|
||
}
|
||
|
||
// Executor 负责本地执行任务并产出结果。
|
||
type Executor struct {
|
||
logger *log.Logger
|
||
}
|
||
|
||
// New 创建 Executor 实例。
|
||
func New(logger *log.Logger) *Executor {
|
||
return &Executor{logger: logger}
|
||
}
|
||
|
||
// ExecuteTask 执行单个任务，返回执行结果。
|
||
// 根据任务类型选择执行方式：
|
||
// - command: 通过 "sh -c <command>" 执行单条命令
|
||
// - script: 通过 "bash" 以标准输入喂入脚本正文执行
|
||
//
|
||
// 超时控制通过 context 实现；网络/信号导致的取消不会 panic。
|
||
// 超时或执行失败均会返回 Result，不返回 error，便于调用方统一回传。
|
||
func (e *Executor) ExecuteTask(ctx context.Context, t Task) Result {
|
||
res := Result{TaskID: t.TaskID}
|
||
|
||
// 构造子进程：根据任务类型选择解释器
|
||
var cmd *exec.Cmd
|
||
switch t.Type {
|
||
case TypeCommand:
|
||
// 单条命令：交给 /bin/sh -c 解释执行
|
||
cmd = exec.CommandContext(ctx, "sh", "-c", t.Command)
|
||
case TypeScript:
|
||
// 脚本：通过 bash 标准输入执行，便于支持多行与 heredoc 等语法
|
||
cmd = exec.CommandContext(ctx, "bash")
|
||
cmd.Stdin = bytes.NewReader([]byte(t.Command))
|
||
default:
|
||
// 未知任务类型，直接标记失败，避免误执行
|
||
res.ExitCode = -1
|
||
res.Status = "failed"
|
||
res.Result = fmt.Sprintf("未知的任务类型: %q", t.Type)
|
||
e.logger.Printf("任务 %s 类型非法: %q", t.TaskID, t.Type)
|
||
return res
|
||
}
|
||
|
||
// 合并 stdout 与 stderr，输出顺序与终端一致，便于排查问题
|
||
var buf bytes.Buffer
|
||
cmd.Stdout = &buf
|
||
cmd.Stderr = &buf
|
||
|
||
// 设置执行超时上下文（若任务指定了有效 timeout）
|
||
execCtx, cancel := e.withTimeout(ctx, t.Timeout)
|
||
defer cancel()
|
||
|
||
startTime := time.Now()
|
||
err := cmd.Run()
|
||
elapsed := time.Since(startTime)
|
||
|
||
// 优先判定超时：Run 返回的 deadline exceeded 与 ctx.Err() 一起识别
|
||
if execCtx.Err() == context.DeadlineExceeded {
|
||
res.ExitCode = -1
|
||
res.Status = "timeout"
|
||
res.Result = fmt.Sprintf("任务执行超时（%ds），已输出:\n%s", t.Timeout, buf.String())
|
||
e.logger.Printf("任务 %s 执行超时（%ds）", t.TaskID, t.Timeout)
|
||
return res
|
||
}
|
||
|
||
// 正常退出码提取
|
||
res.ExitCode = exitCode(err)
|
||
res.Result = buf.String()
|
||
|
||
// 根据退出码判定状态
|
||
if res.ExitCode == 0 {
|
||
res.Status = "success"
|
||
e.logger.Printf("任务 %s 执行成功 exit=0 耗时=%s 输出=%dB",
|
||
t.TaskID, elapsed, buf.Len())
|
||
} else {
|
||
res.Status = "failed"
|
||
e.logger.Printf("任务 %s 执行失败 exit=%d 耗时=%s 输出=%dB",
|
||
t.TaskID, res.ExitCode, elapsed, buf.Len())
|
||
}
|
||
|
||
return res
|
||
}
|
||
|
||
// withTimeout 在父上下文基础上叠加超时控制。
|
||
// 当 timeout <= 0 时直接返回父上下文，不做超时限制。
|
||
// 返回的取消函数保证安全可调用（即使未创建新的超时上下文）。
|
||
func (e *Executor) withTimeout(parent context.Context, timeout int) (context.Context, context.CancelFunc) {
|
||
if timeout <= 0 {
|
||
// 无超时要求：返回空取消函数，保持调用方代码统一
|
||
return parent, func() {}
|
||
}
|
||
return context.WithTimeout(parent, time.Duration(timeout)*time.Second)
|
||
}
|
||
|
||
// exitCode 从 exec.Run 返回的错误中提取退出码。
|
||
// 成功时返回 0；无法识别的错误（如进程未启动）返回 -1。
|
||
func exitCode(err error) int {
|
||
if err == nil {
|
||
return 0
|
||
}
|
||
// 尝试断言为 *exec.ExitError 以获取 ExitCode
|
||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||
return exitErr.ExitCode()
|
||
}
|
||
// 其他类型错误（如二进制不存在）统一记为 -1
|
||
return -1
|
||
}
|