admin
e6956aa001
- React frontend with route-level code splitting - Backend rebranded from Baihu to TaskPool - DB brand migration script and local compatibility
845 lines
21 KiB
Go
845 lines
21 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/url"
|
||
"os"
|
||
"runtime"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/engigu/taskpool/internal/constant"
|
||
"github.com/engigu/taskpool/internal/executor"
|
||
"github.com/engigu/taskpool/internal/logger"
|
||
"github.com/engigu/taskpool/internal/utils"
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
// WebSocket 消息类型
|
||
const (
|
||
WSTypeHeartbeat = constant.WSTypeHeartbeat
|
||
WSTypeHeartbeatAck = constant.WSTypeHeartbeatAck
|
||
WSTypeTasks = constant.WSTypeTasks
|
||
WSTypeTaskResult = constant.WSTypeTaskResult
|
||
WSTypeUpdate = constant.WSTypeUpdate
|
||
WSTypeConnected = constant.WSTypeConnected
|
||
WSTypeDisabled = constant.WSTypeDisabled
|
||
WSTypeEnabled = constant.WSTypeEnabled
|
||
WSTypeFetchTasks = constant.WSTypeFetchTasks
|
||
WSTypeTaskLog = constant.WSTypeTaskLog
|
||
WSTypeExecute = constant.WSTypeExecute
|
||
WSTypeTaskHeartbeat = constant.WSTypeTaskHeartbeat
|
||
WSTypeStop = constant.WSTypeStop
|
||
)
|
||
|
||
type WSMessage struct {
|
||
Type string `json:"type"`
|
||
Data json.RawMessage `json:"data,omitempty"`
|
||
}
|
||
|
||
type AgentTask struct {
|
||
ID string `json:"id"`
|
||
Name string `json:"name"`
|
||
Command string `json:"command"`
|
||
PreCommand string `json:"pre_command"`
|
||
PostCommand string `json:"post_command"`
|
||
Schedule string `json:"schedule"`
|
||
Cron string `json:"cron"`
|
||
Timeout int `json:"timeout"`
|
||
WorkDir string `json:"work_dir"`
|
||
Envs string `json:"envs"`
|
||
Languages []map[string]string `json:"languages"`
|
||
RandomRange int `json:"random_range"`
|
||
Secrets []string `json:"secrets"`
|
||
Enabled bool `json:"enabled"`
|
||
}
|
||
|
||
func (t *AgentTask) GetID() string {
|
||
return t.ID
|
||
}
|
||
|
||
func (t *AgentTask) GetName() string {
|
||
return t.Name
|
||
}
|
||
|
||
func (t *AgentTask) GetCommand() string {
|
||
return t.Command
|
||
}
|
||
|
||
func (t *AgentTask) GetPreCommand() string {
|
||
return t.PreCommand
|
||
}
|
||
|
||
func (t *AgentTask) GetPostCommand() string {
|
||
return t.PostCommand
|
||
}
|
||
|
||
func (t *AgentTask) GetTimeout() int {
|
||
return t.Timeout
|
||
}
|
||
|
||
func (t *AgentTask) GetWorkDir() string {
|
||
return t.WorkDir
|
||
}
|
||
|
||
func (t *AgentTask) GetEnvs() string {
|
||
return t.Envs
|
||
}
|
||
|
||
func (t *AgentTask) GetEnvVars() []string {
|
||
return nil
|
||
}
|
||
|
||
func (t *AgentTask) GetSecrets() []string {
|
||
return t.Secrets
|
||
}
|
||
|
||
func (t *AgentTask) GetLanguages() []map[string]string {
|
||
return t.Languages
|
||
}
|
||
|
||
func (t *AgentTask) GetUseMise() bool {
|
||
return false
|
||
}
|
||
|
||
func (t *AgentTask) UseMise() bool {
|
||
return false
|
||
}
|
||
|
||
func (t *AgentTask) GetSchedule() string {
|
||
if t.Schedule != "" {
|
||
return t.Schedule
|
||
}
|
||
return t.Cron
|
||
}
|
||
|
||
func (t *AgentTask) GetRandomRange() int {
|
||
return t.RandomRange
|
||
}
|
||
|
||
type TaskResult struct {
|
||
TaskID string `json:"task_id"`
|
||
LogID string `json:"log_id"`
|
||
AgentID string `json:"agent_id"` // 仅用于 HTTP 上报时后端补充
|
||
Command string `json:"command"`
|
||
Output string `json:"output"`
|
||
Error string `json:"error"`
|
||
Status string `json:"status"`
|
||
Duration int64 `json:"duration"`
|
||
ExitCode int `json:"exit_code"`
|
||
StartTime int64 `json:"start_time"`
|
||
EndTime int64 `json:"end_time"`
|
||
}
|
||
|
||
type Agent struct {
|
||
config *Config
|
||
configFile string
|
||
machineID string
|
||
scheduler *executor.Scheduler
|
||
cronManager *executor.CronManager
|
||
tasks map[string]*AgentTask // 本地任务缓存，用于执行 lookup
|
||
lastTaskCount int
|
||
mu sync.RWMutex
|
||
client *http.Client
|
||
wsConn *websocket.Conn
|
||
wsMu sync.Mutex
|
||
stopCh chan struct{}
|
||
wsStopCh chan struct{} // 用于停止当前 WebSocket 相关的 goroutine
|
||
taskLogs map[string][]string // 记录最近的日志行，用于失败显示
|
||
logMu sync.Mutex // taskLogs 的锁
|
||
schedulerStarted bool // 调度器是否已经启动
|
||
}
|
||
|
||
func NewAgent(config *Config, configFile string) *Agent {
|
||
a := &Agent{
|
||
config: config,
|
||
configFile: configFile,
|
||
machineID: utils.GenerateMachineID(),
|
||
tasks: make(map[string]*AgentTask),
|
||
client: &http.Client{Timeout: 30 * time.Second},
|
||
stopCh: make(chan struct{}),
|
||
lastTaskCount: -1,
|
||
taskLogs: make(map[string][]string),
|
||
}
|
||
|
||
// 初始化调度器
|
||
handler := &AgentHandler{agent: a}
|
||
schedCfg := executor.SchedulerConfig{
|
||
WorkerCount: runtime.NumCPU(),
|
||
QueueSize: 100,
|
||
RateInterval: 100 * time.Millisecond,
|
||
Verbose: true,
|
||
}
|
||
a.scheduler = executor.NewScheduler(schedCfg, handler)
|
||
a.scheduler.SetLogger(logger.NewSchedulerLogger())
|
||
a.cronManager = executor.NewCronManager(a.scheduler)
|
||
a.cronManager.SetLogger(logger.NewSchedulerLogger())
|
||
|
||
return a
|
||
}
|
||
|
||
// AgentHandler 实现 executor.SchedulerEventHandler
|
||
type AgentHandler struct {
|
||
agent *Agent
|
||
}
|
||
|
||
func (h *AgentHandler) OnTaskScheduled(req *executor.ExecutionRequest) {}
|
||
|
||
func (h *AgentHandler) OnTaskExecuting(req *executor.ExecutionRequest) (io.Writer, io.Writer, error) {
|
||
if req.LogID != "" {
|
||
writer := &RealTimeLogWriter{agent: h.agent, logID: req.LogID}
|
||
return writer, writer, nil
|
||
}
|
||
return nil, nil, nil
|
||
}
|
||
|
||
func (h *AgentHandler) OnTaskHeartbeat(req *executor.ExecutionRequest, duration int64) {
|
||
if req.LogID != "" {
|
||
h.agent.sendWSMessage(WSTypeTaskHeartbeat, map[string]interface{}{
|
||
"log_id": req.LogID,
|
||
"duration": duration,
|
||
})
|
||
}
|
||
|
||
// 每分钟打印一次任务还在运行的日志，提升长任务的存在感
|
||
if duration >= 60000 && (duration/60000 > (duration-3000)/60000) {
|
||
logger.Infof("[Scheduler] 任务 #%s 仍在运行中... (已耗时: %v)",
|
||
req.TaskID, (time.Duration(duration) * time.Millisecond).Round(time.Second))
|
||
}
|
||
}
|
||
|
||
func (h *AgentHandler) OnTaskStarted(req *executor.ExecutionRequest) {}
|
||
|
||
func (h *AgentHandler) OnTaskCompleted(req *executor.ExecutionRequest, result *executor.ExecutionResult) {
|
||
h.agent.sendTaskResult(&TaskResult{
|
||
TaskID: req.TaskID,
|
||
LogID: result.LogID,
|
||
Command: req.Command,
|
||
Output: result.Output,
|
||
Error: result.Error,
|
||
Status: result.Status,
|
||
Duration: result.Duration,
|
||
ExitCode: result.ExitCode,
|
||
StartTime: result.StartTime.Unix(),
|
||
EndTime: result.EndTime.Unix(),
|
||
})
|
||
|
||
if result.Status == constant.TaskStatusFailed {
|
||
h.agent.printLastLogs(result.LogID)
|
||
}
|
||
h.agent.clearTaskLog(result.LogID)
|
||
}
|
||
|
||
func (h *AgentHandler) OnTaskFailed(req *executor.ExecutionRequest, err error) {
|
||
errMsg := fmt.Sprintf("任务执行失败: %v", err)
|
||
// 先发送日志，确保服务端能收到错误信息
|
||
h.agent.sendWSMessage(WSTypeTaskLog, map[string]interface{}{
|
||
"log_id": req.LogID,
|
||
"content": errMsg,
|
||
})
|
||
|
||
h.agent.sendTaskResult(&TaskResult{
|
||
TaskID: req.TaskID,
|
||
LogID: req.LogID,
|
||
Command: req.Command,
|
||
Output: "",
|
||
Error: err.Error(),
|
||
Status: constant.TaskStatusFailed,
|
||
Duration: 0,
|
||
ExitCode: 1,
|
||
StartTime: time.Now().Unix(),
|
||
EndTime: time.Now().Unix(),
|
||
})
|
||
|
||
h.agent.printLastLogs(req.LogID)
|
||
h.agent.clearTaskLog(req.LogID)
|
||
}
|
||
|
||
func (h *AgentHandler) OnCronNextRun(req *executor.ExecutionRequest, nextRun time.Time) {}
|
||
|
||
func (a *Agent) Start() error {
|
||
if a.config.Token == "" {
|
||
return fmt.Errorf("缺少令牌，请在配置文件中设置 token")
|
||
}
|
||
|
||
logger.Infof("机器识别码: %s", a.machineID[:16]+"...")
|
||
// 调度器暂不在此启动，等待 WebSocket 连接成功并获取到调度配置后再启动
|
||
go a.wsLoop()
|
||
|
||
logger.Info("Agent 已启动 (时区: Asia/Shanghai, 模式: WebSocket)")
|
||
return nil
|
||
}
|
||
|
||
func (a *Agent) Stop() {
|
||
close(a.stopCh)
|
||
a.closeWS()
|
||
|
||
a.mu.Lock()
|
||
started := a.schedulerStarted
|
||
a.schedulerStarted = false
|
||
a.mu.Unlock()
|
||
|
||
if started {
|
||
a.cronManager.Stop()
|
||
a.scheduler.Stop()
|
||
}
|
||
logger.Info("Agent 已停止")
|
||
}
|
||
|
||
// wsLoop WebSocket 连接循环
|
||
func (a *Agent) wsLoop() {
|
||
for {
|
||
select {
|
||
case <-a.stopCh:
|
||
return
|
||
default:
|
||
}
|
||
|
||
if err := a.connectWS(); err != nil {
|
||
logger.Warnf("WebSocket 连接失败: %v，5秒后重试...", err)
|
||
time.Sleep(5 * time.Second)
|
||
continue
|
||
}
|
||
|
||
a.readWS()
|
||
|
||
logger.Warn("WebSocket 连接断开，5秒后重连...")
|
||
time.Sleep(5 * time.Second)
|
||
}
|
||
}
|
||
|
||
func (a *Agent) connectWS() error {
|
||
serverURL := a.config.ServerURL
|
||
wsURL := strings.Replace(serverURL, "http://", "ws://", 1)
|
||
wsURL = strings.Replace(wsURL, "https://", "wss://", 1)
|
||
wsURL = fmt.Sprintf("%s/api/agent/ws?token=%s&machine_id=%s", wsURL, url.QueryEscape(a.config.Token), url.QueryEscape(a.machineID))
|
||
|
||
logger.Infof("正在连接 WebSocket: %s", wsURL)
|
||
logger.Infof("Token: %s..., MachineID: %s...", a.config.Token[:8], a.machineID[:16])
|
||
|
||
dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second}
|
||
conn, resp, err := dialer.Dial(wsURL, nil)
|
||
if err != nil {
|
||
if resp != nil {
|
||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||
logger.Errorf("WebSocket 握手失败: HTTP %d, Body: %s", resp.StatusCode, string(bodyBytes))
|
||
resp.Body.Close()
|
||
} else {
|
||
logger.Errorf("WebSocket 连接失败: %v", err)
|
||
}
|
||
return err
|
||
}
|
||
|
||
a.wsMu.Lock()
|
||
a.wsConn = conn
|
||
a.wsStopCh = make(chan struct{})
|
||
a.wsMu.Unlock()
|
||
|
||
logger.Info("WebSocket 已连接")
|
||
a.sendHeartbeat()
|
||
go a.heartbeatLoop()
|
||
|
||
return nil
|
||
}
|
||
|
||
func (a *Agent) closeWS() {
|
||
a.wsMu.Lock()
|
||
defer a.wsMu.Unlock()
|
||
if a.wsStopCh != nil {
|
||
close(a.wsStopCh)
|
||
a.wsStopCh = nil
|
||
}
|
||
if a.wsConn != nil {
|
||
a.wsConn.Close()
|
||
a.wsConn = nil
|
||
}
|
||
}
|
||
|
||
func (a *Agent) readWS() {
|
||
defer func() {
|
||
logger.Info("readWS 退出，准备关闭连接")
|
||
a.closeWS()
|
||
}()
|
||
|
||
for {
|
||
a.wsMu.Lock()
|
||
conn := a.wsConn
|
||
a.wsMu.Unlock()
|
||
|
||
if conn == nil {
|
||
logger.Warn("readWS: wsConn 为 nil")
|
||
return
|
||
}
|
||
|
||
_, message, err := conn.ReadMessage()
|
||
if err != nil {
|
||
logger.Warnf("WebSocket 读取错误: %v", err)
|
||
return
|
||
}
|
||
|
||
var msg WSMessage
|
||
if err := json.Unmarshal(message, &msg); err != nil {
|
||
continue
|
||
}
|
||
|
||
a.handleWSMessage(&msg)
|
||
}
|
||
}
|
||
|
||
func (a *Agent) handleWSMessage(msg *WSMessage) {
|
||
switch msg.Type {
|
||
case WSTypeConnected:
|
||
a.handleConnected(msg.Data)
|
||
case WSTypeHeartbeatAck:
|
||
a.handleHeartbeatAck(msg.Data)
|
||
case WSTypeTasks:
|
||
a.handleTasks(msg.Data)
|
||
case WSTypeUpdate:
|
||
logger.Info("收到更新指令，开始更新...")
|
||
go a.selfUpdate()
|
||
case WSTypeDisabled:
|
||
logger.Warn("Agent 已被禁用，清空所有任务")
|
||
a.clearAllTasks()
|
||
case WSTypeEnabled:
|
||
logger.Info("Agent 已被启用，主动拉取任务")
|
||
a.fetchTasks()
|
||
case WSTypeExecute:
|
||
a.handleExecute(msg.Data)
|
||
case WSTypeStop:
|
||
a.handleStop(msg.Data)
|
||
}
|
||
}
|
||
|
||
func (a *Agent) fetchTasks() {
|
||
logger.Info("正在从服务器拉取任务列表...")
|
||
if err := a.sendWSMessage(WSTypeFetchTasks, map[string]interface{}{}); err != nil {
|
||
logger.Warnf("请求任务列表失败: %v", err)
|
||
}
|
||
}
|
||
|
||
func (a *Agent) handleConnected(data json.RawMessage) {
|
||
var resp struct {
|
||
AgentID string `json:"agent_id"`
|
||
Name string `json:"name"`
|
||
IsNewAgent bool `json:"is_new_agent"`
|
||
MachineID string `json:"machine_id"`
|
||
SchedulerConfig map[string]interface{} `json:"scheduler_config"`
|
||
}
|
||
json.Unmarshal(data, &resp)
|
||
|
||
if resp.IsNewAgent {
|
||
logger.Infof("注册成功: Agent #%s, 机器码: %s", resp.AgentID, a.machineID[:16]+"...")
|
||
} else {
|
||
logger.Infof("连接成功: Agent #%s (已存在), 机器码: %s", resp.AgentID, a.machineID[:16]+"...")
|
||
}
|
||
|
||
// 更新调度器配置
|
||
if resp.SchedulerConfig != nil {
|
||
a.updateSchedulerConfig(resp.SchedulerConfig)
|
||
}
|
||
|
||
a.fetchTasks()
|
||
}
|
||
|
||
func (a *Agent) updateSchedulerConfig(config map[string]interface{}) {
|
||
// 获取当前配置作为基础
|
||
currentCfg := a.scheduler.GetConfig()
|
||
newCfg := currentCfg
|
||
|
||
// 更新配置项
|
||
if val, ok := config["worker_count"]; ok {
|
||
if v, ok := val.(float64); ok { // JSON 数字解析为 float64
|
||
newCfg.WorkerCount = int(v)
|
||
}
|
||
}
|
||
if val, ok := config["queue_size"]; ok {
|
||
if v, ok := val.(float64); ok {
|
||
newCfg.QueueSize = int(v)
|
||
}
|
||
}
|
||
if val, ok := config["rate_interval"]; ok {
|
||
if v, ok := val.(float64); ok {
|
||
newCfg.RateInterval = time.Duration(v) * time.Millisecond
|
||
}
|
||
}
|
||
if val, ok := config["strict_queue"]; ok {
|
||
if v, ok := val.(bool); ok {
|
||
newCfg.StrictQueue = v
|
||
}
|
||
}
|
||
|
||
a.mu.Lock()
|
||
started := a.schedulerStarted
|
||
a.schedulerStarted = true
|
||
a.mu.Unlock()
|
||
|
||
if !started {
|
||
logger.Infof("首次连接成功，启动调度器配置: workers=%d, queue=%d, rate=%v, strict=%t",
|
||
newCfg.WorkerCount, newCfg.QueueSize, newCfg.RateInterval, newCfg.StrictQueue)
|
||
// 用下发的最新配置加载并启动调度器与计划任务管理器
|
||
a.scheduler.Reload(newCfg)
|
||
a.cronManager.Start()
|
||
} else if newCfg != currentCfg {
|
||
logger.Infof("收到调度配置更新: workers=%d, queue=%d, rate=%v, strict=%t",
|
||
newCfg.WorkerCount, newCfg.QueueSize, newCfg.RateInterval, newCfg.StrictQueue)
|
||
a.scheduler.Reload(newCfg)
|
||
} else {
|
||
logger.Infof("当前调度配置未改变: workers=%d, queue=%d, rate=%v, strict=%t",
|
||
newCfg.WorkerCount, newCfg.QueueSize, newCfg.RateInterval, newCfg.StrictQueue)
|
||
}
|
||
}
|
||
|
||
func (a *Agent) handleHeartbeatAck(data json.RawMessage) {
|
||
var resp struct {
|
||
AgentID string `json:"agent_id"`
|
||
Name string `json:"name"`
|
||
NeedUpdate bool `json:"need_update"`
|
||
ForceUpdate bool `json:"force_update"`
|
||
LatestVersion string `json:"latest_version"`
|
||
}
|
||
json.Unmarshal(data, &resp)
|
||
|
||
if resp.NeedUpdate && (a.config.AutoUpdate || resp.ForceUpdate) {
|
||
logger.Infof("发现新版本 %s，开始更新...", resp.LatestVersion)
|
||
go a.selfUpdate()
|
||
}
|
||
}
|
||
|
||
func (a *Agent) handleTasks(data json.RawMessage) {
|
||
var resp struct {
|
||
Tasks []AgentTask `json:"tasks"`
|
||
}
|
||
json.Unmarshal(data, &resp)
|
||
|
||
newCount := len(resp.Tasks)
|
||
if newCount != a.lastTaskCount || newCount == 0 {
|
||
logger.Infof("任务列表同步成功: 共获取到 %d 个任务", newCount)
|
||
a.lastTaskCount = newCount
|
||
}
|
||
|
||
a.updateTasks(resp.Tasks)
|
||
}
|
||
|
||
func (a *Agent) handleExecute(data json.RawMessage) {
|
||
var req struct {
|
||
TaskID string `json:"task_id"`
|
||
LogID string `json:"log_id"`
|
||
Envs string `json:"envs"`
|
||
Secrets []string `json:"secrets"`
|
||
Command string `json:"command"`
|
||
PreCommand string `json:"pre_command"`
|
||
PostCommand string `json:"post_command"`
|
||
}
|
||
if err := json.Unmarshal(data, &req); err != nil {
|
||
logger.Errorf("解析立即执行请求失败: %v", err)
|
||
return
|
||
}
|
||
|
||
// 查找任务
|
||
a.mu.RLock()
|
||
task, exists := a.tasks[req.TaskID]
|
||
a.mu.RUnlock()
|
||
|
||
if !exists {
|
||
logger.Warnf("任务 #%s 不存在，无法执行", req.TaskID)
|
||
return
|
||
}
|
||
|
||
// 准备执行请求
|
||
// 如果消息中携带了环境变量或指令，则优先使用（确保即时生效）
|
||
envs := task.Envs
|
||
if req.Envs != "" {
|
||
envs = req.Envs
|
||
}
|
||
|
||
command := task.Command
|
||
if req.Command != "" {
|
||
command = req.Command
|
||
}
|
||
|
||
preCommand := task.PreCommand
|
||
if req.PreCommand != "" {
|
||
preCommand = req.PreCommand
|
||
}
|
||
|
||
postCommand := task.PostCommand
|
||
if req.PostCommand != "" {
|
||
postCommand = req.PostCommand
|
||
}
|
||
|
||
execReq := &executor.ExecutionRequest{
|
||
TaskID: task.ID,
|
||
LogID: req.LogID,
|
||
Name: task.Name,
|
||
Command: command,
|
||
PreCommand: preCommand,
|
||
PostCommand: postCommand,
|
||
WorkDir: task.WorkDir,
|
||
Envs: executor.ParseEnvVars(envs),
|
||
Secrets: req.Secrets,
|
||
Timeout: task.Timeout,
|
||
Languages: task.Languages,
|
||
UseMise: task.UseMise(),
|
||
Type: executor.TaskTypeManual,
|
||
}
|
||
|
||
// 立即执行任务（加入队列）
|
||
a.scheduler.EnqueueOrExecute(execReq)
|
||
}
|
||
|
||
func (a *Agent) handleStop(data json.RawMessage) {
|
||
var req struct {
|
||
LogID string `json:"log_id"`
|
||
}
|
||
if err := json.Unmarshal(data, &req); err != nil {
|
||
logger.Errorf("解析停止请求失败: %v", err)
|
||
return
|
||
}
|
||
|
||
logger.Infof("[Agent] 收到停止指令 LogID: %s", req.LogID)
|
||
if a.scheduler.StopLog(req.LogID) {
|
||
logger.Infof("[Agent] 任务执行 #%s 已成功停止", req.LogID)
|
||
} else {
|
||
logger.Warnf("[Agent] 任务执行 #%s 停止失败（可能已完成或不在运行队列中）", req.LogID)
|
||
}
|
||
}
|
||
|
||
// RealTimeLogWriter 实时日志写入器，通过 WebSocket 发送日志
|
||
type RealTimeLogWriter struct {
|
||
agent *Agent
|
||
logID string
|
||
}
|
||
|
||
func (w *RealTimeLogWriter) Write(p []byte) (n int, err error) {
|
||
if len(p) == 0 {
|
||
return 0, nil
|
||
}
|
||
|
||
// 记录到本地缓存，用于失败时显示
|
||
w.agent.addTaskLog(w.logID, p)
|
||
|
||
// 构造消息
|
||
msg := map[string]interface{}{
|
||
"log_id": w.logID,
|
||
"content": string(p),
|
||
}
|
||
|
||
// 发送消息
|
||
if err := w.agent.sendWSMessage(WSTypeTaskLog, msg); err != nil {
|
||
// 如果发送失败，不阻塞程序执行，只记录日志
|
||
return len(p), nil
|
||
}
|
||
|
||
return len(p), nil
|
||
}
|
||
|
||
func (a *Agent) sendWSMessage(msgType string, data interface{}) error {
|
||
a.wsMu.Lock()
|
||
defer a.wsMu.Unlock()
|
||
|
||
if a.wsConn == nil {
|
||
return fmt.Errorf("WebSocket 未连接")
|
||
}
|
||
|
||
dataBytes, _ := json.Marshal(data)
|
||
msg := WSMessage{Type: msgType, Data: dataBytes}
|
||
msgBytes, _ := json.Marshal(msg)
|
||
|
||
a.wsConn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
||
if err := a.wsConn.WriteMessage(websocket.TextMessage, msgBytes); err != nil {
|
||
logger.Warnf("发送消息失败 (%s): %v", msgType, err)
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (a *Agent) heartbeatLoop() {
|
||
ticker := time.NewTicker(time.Duration(a.config.Interval) * time.Second)
|
||
defer ticker.Stop()
|
||
|
||
a.wsMu.Lock()
|
||
wsStopCh := a.wsStopCh
|
||
a.wsMu.Unlock()
|
||
|
||
if wsStopCh == nil {
|
||
return
|
||
}
|
||
|
||
for {
|
||
select {
|
||
case <-a.stopCh:
|
||
return
|
||
case <-wsStopCh:
|
||
return
|
||
case <-ticker.C:
|
||
a.wsMu.Lock()
|
||
conn := a.wsConn
|
||
a.wsMu.Unlock()
|
||
if conn == nil {
|
||
return
|
||
}
|
||
a.sendHeartbeat()
|
||
}
|
||
}
|
||
}
|
||
|
||
func (a *Agent) sendHeartbeat() {
|
||
hostname, _ := os.Hostname()
|
||
data := map[string]interface{}{
|
||
"version": Version,
|
||
"build_time": BuildTime,
|
||
"hostname": hostname,
|
||
"os": runtime.GOOS,
|
||
"arch": runtime.GOARCH,
|
||
"auto_update": a.config.AutoUpdate,
|
||
}
|
||
if err := a.sendWSMessage(WSTypeHeartbeat, data); err != nil {
|
||
logger.Warnf("发送心跳失败: %v", err)
|
||
}
|
||
}
|
||
|
||
func (a *Agent) sendTaskResult(result *TaskResult) {
|
||
if err := a.sendWSMessage(WSTypeTaskResult, result); err != nil {
|
||
logger.Warnf("发送任务结果失败: %v，尝试 HTTP 上报", err)
|
||
a.reportResultHTTP(result)
|
||
}
|
||
}
|
||
|
||
func (a *Agent) reportResultHTTP(result *TaskResult) error {
|
||
resp, err := a.doRequest("POST", "/api/agent/report", result)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer resp.Body.Close()
|
||
return nil
|
||
}
|
||
|
||
func (a *Agent) updateTasks(tasks []AgentTask) {
|
||
a.mu.Lock()
|
||
defer a.mu.Unlock()
|
||
|
||
newTasks := make(map[string]*AgentTask)
|
||
for i := range tasks {
|
||
newTasks[tasks[i].ID] = &tasks[i]
|
||
}
|
||
|
||
// 1. 移除不再存在的任务
|
||
for id := range a.tasks {
|
||
if _, exists := newTasks[id]; !exists {
|
||
a.cronManager.RemoveTask(id)
|
||
delete(a.tasks, id)
|
||
logger.Infof("移除调度任务 #%s", id)
|
||
}
|
||
}
|
||
|
||
// 2. 添加或更新任务
|
||
for id, task := range newTasks {
|
||
oldTask, exists := a.tasks[id]
|
||
if !exists || oldTask.Schedule != task.Schedule || oldTask.Command != task.Command ||
|
||
oldTask.PreCommand != task.PreCommand || oldTask.PostCommand != task.PostCommand ||
|
||
oldTask.Enabled != task.Enabled || oldTask.Timeout != task.Timeout ||
|
||
oldTask.WorkDir != task.WorkDir || oldTask.Envs != task.Envs ||
|
||
oldTask.RandomRange != task.RandomRange {
|
||
if task.Enabled {
|
||
err := a.cronManager.AddTask(task)
|
||
if err != nil {
|
||
logger.Errorf("添加调度任务 #%s 失败: %v", id, err)
|
||
continue
|
||
}
|
||
logger.Infof("已添加调度任务 #%s %s (%s)", id, task.Name, task.GetSchedule())
|
||
} else {
|
||
a.cronManager.RemoveTask(id)
|
||
logger.Infof("调度任务 #%s 已禁用", id)
|
||
}
|
||
a.tasks[id] = task
|
||
}
|
||
}
|
||
}
|
||
|
||
func (a *Agent) clearAllTasks() {
|
||
a.mu.Lock()
|
||
defer a.mu.Unlock()
|
||
|
||
for id := range a.tasks {
|
||
a.cronManager.RemoveTask(id)
|
||
logger.Infof("移除任务 #%s", id)
|
||
}
|
||
|
||
a.tasks = make(map[string]*AgentTask)
|
||
a.lastTaskCount = 0
|
||
logger.Info("所有任务已清空")
|
||
}
|
||
|
||
func (a *Agent) addTaskLog(logID string, p []byte) {
|
||
if logID == "" {
|
||
return
|
||
}
|
||
a.logMu.Lock()
|
||
defer a.logMu.Unlock()
|
||
|
||
content := string(p)
|
||
lines := strings.Split(strings.TrimSuffix(content, "\n"), "\n")
|
||
|
||
a.taskLogs[logID] = append(a.taskLogs[logID], lines...)
|
||
if len(a.taskLogs[logID]) > 50 {
|
||
a.taskLogs[logID] = a.taskLogs[logID][len(a.taskLogs[logID])-50:]
|
||
}
|
||
}
|
||
|
||
func (a *Agent) printLastLogs(logID string) {
|
||
if logID == "" {
|
||
return
|
||
}
|
||
a.logMu.Lock()
|
||
lines, ok := a.taskLogs[logID]
|
||
a.logMu.Unlock()
|
||
|
||
if !ok || len(lines) == 0 {
|
||
return
|
||
}
|
||
|
||
logger.Errorf("--- 任务 #%s 失败日志预览 (最近 %d 行) ---", logID, len(lines))
|
||
for _, line := range lines {
|
||
fmt.Println(" " + line)
|
||
}
|
||
logger.Errorf("--- 任务 #%s 结束 ---", logID)
|
||
}
|
||
|
||
func (a *Agent) clearTaskLog(logID string) {
|
||
if logID == "" {
|
||
return
|
||
}
|
||
a.logMu.Lock()
|
||
defer a.logMu.Unlock()
|
||
delete(a.taskLogs, logID)
|
||
}
|
||
|
||
// executeTask 已被 AgentHandler.OnTaskCompleted 代替，此处删除旧实现
|
||
|
||
func (a *Agent) doRequest(method, path string, body interface{}) (*http.Response, error) {
|
||
var bodyReader io.Reader
|
||
if body != nil {
|
||
data, err := json.Marshal(body)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
bodyReader = bytes.NewReader(data)
|
||
}
|
||
|
||
req, err := http.NewRequest(method, a.config.ServerURL+path, bodyReader)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
req.Header.Set("Authorization", "Bearer "+a.config.Token)
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("X-Machine-ID", a.machineID)
|
||
|
||
return a.client.Do(req)
|
||
}
|