270 lines
6.4 KiB
Go
270 lines
6.4 KiB
Go
package scheduler
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"os/exec"
|
||
"runtime"
|
||
"strings"
|
||
"time"
|
||
|
||
"ppgo_job/dao"
|
||
"ppgo_job/model/entity"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
// TaskResult 任务执行结果
|
||
type TaskResult struct {
|
||
Output string
|
||
Error string
|
||
IsOk bool
|
||
IsTimeout bool
|
||
}
|
||
|
||
// executeTask 执行任务主逻辑
|
||
func (s *schedulerService) executeTask(ctx context.Context, task *entity.Task) {
|
||
startTime := time.Now()
|
||
g.Log().Infof(ctx, "开始执行任务 [%d] %s", task.Id, task.TaskName)
|
||
|
||
var finalResult *TaskResult
|
||
serverName := "本地服务器"
|
||
var serverId int
|
||
|
||
if task.TaskType == "http" {
|
||
// HTTP 任务:直接从调度器所在机器发起请求
|
||
result := s.execHTTP(ctx, task)
|
||
finalResult = result
|
||
} else {
|
||
// Shell 任务:按原有逻辑执行(本地/远程)
|
||
serverIds := strings.Split(task.ServerIds, ",")
|
||
for _, sid := range serverIds {
|
||
sid = strings.TrimSpace(sid)
|
||
if sid == "" {
|
||
continue
|
||
}
|
||
|
||
var result *TaskResult
|
||
if sid == "0" {
|
||
serverName = "本地服务器"
|
||
serverId = 0
|
||
result = s.execLocal(ctx, task.Command, task.Timeout)
|
||
} else {
|
||
if srv, _ := dao.TaskServer.GetById(ctx, parseInt(sid)); srv != nil {
|
||
serverName = srv.ServerName
|
||
serverId = srv.Id
|
||
}
|
||
result = s.execRemote(ctx, task.Command, task.Timeout, parseInt(sid))
|
||
}
|
||
|
||
if finalResult == nil || !result.IsOk {
|
||
finalResult = result
|
||
}
|
||
}
|
||
}
|
||
|
||
if finalResult == nil {
|
||
finalResult = &TaskResult{IsOk: true}
|
||
}
|
||
|
||
elapsed := time.Since(startTime).Milliseconds()
|
||
|
||
// 记录日志
|
||
log := &entity.TaskLog{
|
||
TaskId: task.Id,
|
||
ServerId: serverId,
|
||
ServerName: serverName,
|
||
Output: finalResult.Output,
|
||
Error: finalResult.Error,
|
||
ProcessTime: int(elapsed),
|
||
CreateTime: time.Now().Unix(),
|
||
}
|
||
|
||
if finalResult.IsTimeout {
|
||
log.Status = -2
|
||
} else if !finalResult.IsOk {
|
||
log.Status = -1
|
||
} else {
|
||
log.Status = 0
|
||
}
|
||
|
||
_, _ = dao.TaskLog.Insert(ctx, log)
|
||
|
||
// 更新任务统计
|
||
_ = dao.Task.UpdateFields(ctx, task.Id, g.Map{"prev_time": startTime.Unix()})
|
||
_, _ = g.DB().Exec(ctx, "UPDATE pp_task SET execute_times = COALESCE(execute_times, 0) + 1 WHERE id = ?", task.Id)
|
||
|
||
// 失败时通知
|
||
if log.Status < 0 && task.IsNotify == 1 {
|
||
s.sendNotify(ctx, task, log)
|
||
}
|
||
|
||
g.Log().Infof(ctx, "任务 [%d] %s 执行完成, 耗时 %dms", task.Id, task.TaskName, elapsed)
|
||
}
|
||
|
||
// execLocal 本地执行命令
|
||
func (s *schedulerService) execLocal(ctx context.Context, command string, timeoutSec int) *TaskResult {
|
||
var cmd *exec.Cmd
|
||
if runtime.GOOS == "windows" {
|
||
cmd = exec.Command("CMD", "/C", command)
|
||
} else {
|
||
cmd = exec.Command("sh", "-c", command)
|
||
}
|
||
|
||
result := &TaskResult{}
|
||
done := make(chan error, 1)
|
||
|
||
go func() {
|
||
out, err := cmd.CombinedOutput()
|
||
result.Output = string(out)
|
||
done <- err
|
||
}()
|
||
|
||
timeout := time.Duration(timeoutSec) * time.Second
|
||
if timeoutSec <= 0 {
|
||
timeout = 24 * time.Hour
|
||
}
|
||
|
||
select {
|
||
case err := <-done:
|
||
result.IsOk = (err == nil)
|
||
if err != nil {
|
||
result.Error = err.Error()
|
||
}
|
||
case <-time.After(timeout):
|
||
_ = cmd.Process.Kill()
|
||
result.IsTimeout = true
|
||
result.Error = fmt.Sprintf("任务执行超过 %d 秒,已强制终止", timeoutSec)
|
||
}
|
||
|
||
return result
|
||
}
|
||
|
||
// execRemote 远程执行(SSH)
|
||
func (s *schedulerService) execRemote(ctx context.Context, command string, timeoutSec int, serverId int) *TaskResult {
|
||
if serverId <= 0 {
|
||
return s.execLocal(ctx, command, timeoutSec)
|
||
}
|
||
|
||
server, err := dao.TaskServer.GetById(ctx, serverId)
|
||
if err != nil || server == nil {
|
||
return &TaskResult{IsOk: false, Error: fmt.Sprintf("服务器不存在: %d", serverId)}
|
||
}
|
||
|
||
g.Log().Infof(ctx, "远程执行 -> %s@%s:%d", server.ServerAccount, server.ServerIp, server.Port)
|
||
|
||
sshClient := &SSHClient{}
|
||
connectErr := sshClient.Connect(server.ServerIp, server.Port, server.ServerAccount, server.Password, server.PrivateKeySrc, server.Type)
|
||
if connectErr != nil {
|
||
return &TaskResult{IsOk: false, Error: "SSH连接失败: " + connectErr.Error()}
|
||
}
|
||
defer sshClient.Close()
|
||
|
||
result := &TaskResult{}
|
||
stdout, stderr, runErr := sshClient.Exec(command, timeoutSec)
|
||
result.Output = stdout
|
||
result.Error = stderr
|
||
|
||
if runErr != nil {
|
||
if strings.Contains(runErr.Error(), "timeout") {
|
||
result.IsTimeout = true
|
||
} else {
|
||
result.IsOk = false
|
||
}
|
||
if stderr == "" {
|
||
result.Error = runErr.Error()
|
||
}
|
||
} else {
|
||
result.IsOk = true
|
||
}
|
||
|
||
return result
|
||
}
|
||
|
||
// execHTTP 执行 HTTP 请求任务
|
||
func (s *schedulerService) execHTTP(ctx context.Context, task *entity.Task) *TaskResult {
|
||
result := &TaskResult{}
|
||
client := &http.Client{Timeout: time.Duration(task.Timeout) * time.Second}
|
||
if task.Timeout <= 0 {
|
||
client.Timeout = 30 * time.Second
|
||
}
|
||
|
||
method := strings.ToUpper(task.Method)
|
||
if method == "" {
|
||
method = "GET"
|
||
}
|
||
|
||
var bodyReader io.Reader
|
||
if task.Body != "" && (method == "POST" || method == "PUT" || method == "PATCH") {
|
||
bodyReader = bytes.NewBufferString(task.Body)
|
||
}
|
||
|
||
req, err := http.NewRequest(method, task.Url, bodyReader)
|
||
if err != nil {
|
||
result.Error = "创建请求失败: " + err.Error()
|
||
return result
|
||
}
|
||
|
||
// 解析并设置自定义 Header
|
||
if task.Headers != "" {
|
||
var headers map[string]string
|
||
if err := json.Unmarshal([]byte(task.Headers), &headers); err == nil {
|
||
for k, v := range headers {
|
||
req.Header.Set(k, v)
|
||
}
|
||
} else {
|
||
result.Error = "Headers 格式错误(需为 JSON 对象)"
|
||
return result
|
||
}
|
||
}
|
||
|
||
// 默认 Content-Type
|
||
if req.Header.Get("Content-Type") == "" && bodyReader != nil {
|
||
req.Header.Set("Content-Type", "application/json")
|
||
}
|
||
|
||
g.Log().Infof(ctx, "HTTP请求 -> %s %s", method, task.Url)
|
||
resp, err := client.Do(req)
|
||
if err != nil {
|
||
result.Error = "请求失败: " + err.Error()
|
||
return result
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||
bodyStr := string(bodyBytes)
|
||
|
||
// 截断过长的输出
|
||
if len(bodyStr) > 10000 {
|
||
bodyStr = bodyStr[:10000] + "\n... (truncated)"
|
||
}
|
||
|
||
statusLine := fmt.Sprintf("[HTTP %d]", resp.StatusCode)
|
||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||
result.IsOk = true
|
||
result.Output = statusLine + " " + bodyStr
|
||
} else {
|
||
result.Error = statusLine + " " + bodyStr
|
||
}
|
||
|
||
g.Log().Infof(ctx, "HTTP响应 -> %s %s -> %s", method, task.Url, statusLine)
|
||
return result
|
||
}
|
||
|
||
// parseInt 简单字符串转整数
|
||
func parseInt(s string) int {
|
||
id := 0
|
||
for _, c := range s {
|
||
if c >= '0' && c <= '9' {
|
||
id = id*10 + int(c-'0')
|
||
} else {
|
||
break
|
||
}
|
||
}
|
||
return id
|
||
}
|