295 lines
8.4 KiB
Go
295 lines
8.4 KiB
Go
package worker
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
"unicode/utf8"
|
|
|
|
taskbiz "kra/internal/biz/task"
|
|
platformtask "kra/pkg/task"
|
|
)
|
|
|
|
type TaskExecutor struct {
|
|
tasks *taskbiz.TaskUsecase
|
|
methods taskbiz.TaskMethodRegistry
|
|
// orphans tracks method goroutines that outlived their execution timeout.
|
|
// A method cannot be killed, so the task stays claimed until it returns.
|
|
orphanMu sync.Mutex
|
|
orphans map[uint]struct{}
|
|
}
|
|
|
|
func NewTaskExecutorWithRegistry(tasks *taskbiz.TaskUsecase, methods taskbiz.TaskMethodRegistry) *TaskExecutor {
|
|
return &TaskExecutor{tasks: tasks, methods: methods}
|
|
}
|
|
|
|
func privateIP(ip net.IP) bool {
|
|
return ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalMulticast() || ip.IsLinkLocalUnicast() || ip.IsUnspecified()
|
|
}
|
|
|
|
// Task HTTP clients are built once and shared: a per-request client would
|
|
// discard every pooled connection and never release its transport.
|
|
var (
|
|
taskHTTPClientPublic = newTaskHTTPClient(false)
|
|
taskHTTPClientPrivate = newTaskHTTPClient(true)
|
|
)
|
|
|
|
func taskHTTPClient(allowPrivate bool) *http.Client {
|
|
if allowPrivate {
|
|
return taskHTTPClientPrivate
|
|
}
|
|
return taskHTTPClientPublic
|
|
}
|
|
|
|
func newTaskHTTPClient(allowPrivate bool) *http.Client {
|
|
dialer := &net.Dialer{Timeout: 10 * time.Second, Control: func(_ string, address string, _ syscall.RawConn) error {
|
|
if allowPrivate {
|
|
return nil
|
|
}
|
|
host, _, err := net.SplitHostPort(address)
|
|
if err != nil {
|
|
return fmt.Errorf("解析拨号地址失败: %w", err)
|
|
}
|
|
ip := net.ParseIP(host)
|
|
if ip == nil {
|
|
return fmt.Errorf("非法拨号 IP: %s", host)
|
|
}
|
|
if privateIP(ip) {
|
|
return fmt.Errorf("目标解析为内网/环回/链路本地地址, 已被 SSRF 防护拒绝(可在任务上开启\"允许内网\"豁免): %s", ip)
|
|
}
|
|
return nil
|
|
}}
|
|
transport := &http.Transport{Proxy: nil, DialContext: dialer.DialContext, MaxIdleConnsPerHost: 4, IdleConnTimeout: 90 * time.Second}
|
|
return &http.Client{Timeout: 30 * time.Second, Transport: transport}
|
|
}
|
|
|
|
func (e *TaskExecutor) runHTTP(ctx context.Context, task *taskbiz.TimedTask) (string, error) {
|
|
if task == nil {
|
|
return "", errors.New("任务不能为空")
|
|
}
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
parsed, err := url.Parse(task.HTTPURL)
|
|
if err != nil {
|
|
return "", fmt.Errorf("URL 非法: %w", err)
|
|
}
|
|
if parsed.Scheme != "http" && parsed.Scheme != "https" {
|
|
return "", fmt.Errorf("仅允许 http/https, 实际为 %q", parsed.Scheme)
|
|
}
|
|
if parsed.Host == "" {
|
|
return "", errors.New("URL 缺少主机名")
|
|
}
|
|
method := strings.ToUpper(strings.TrimSpace(task.HTTPMethod))
|
|
if method == "" {
|
|
method = http.MethodGet
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, method, task.HTTPURL, bytes.NewBufferString(task.HTTPBody))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
headers := map[string]string{}
|
|
if len(task.HTTPHeader) > 0 {
|
|
if err = json.Unmarshal(task.HTTPHeader, &headers); err != nil {
|
|
return "", fmt.Errorf("http_header 必须是 JSON 对象: %w", err)
|
|
}
|
|
if headers == nil {
|
|
return "", errors.New("http_header 必须是 JSON 对象")
|
|
}
|
|
}
|
|
for key, value := range headers {
|
|
request.Header.Set(key, value)
|
|
}
|
|
response, err := taskHTTPClient(task.HTTPAllowPrivate).Do(request)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer response.Body.Close()
|
|
body, readErr := io.ReadAll(io.LimitReader(response.Body, 1<<20))
|
|
output := fmt.Sprintf("HTTP %d: %s", response.StatusCode, string(body))
|
|
if readErr != nil {
|
|
return output, fmt.Errorf("读取 HTTP 响应失败: %w", readErr)
|
|
}
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
return output, fmt.Errorf("非 2xx 响应: %d", response.StatusCode)
|
|
}
|
|
return output, nil
|
|
}
|
|
|
|
var (
|
|
errTaskTimeout = errors.New("任务执行超时")
|
|
errTaskOrphaned = errors.New("上一次执行超时后仍在运行, 本次执行已跳过")
|
|
)
|
|
|
|
func truncateTaskText(value string) string {
|
|
const limit = 4000
|
|
if len(value) <= limit {
|
|
return value
|
|
}
|
|
// Cut on a rune boundary: slicing raw bytes would split multi-byte text
|
|
// (Chinese output in particular) into an invalid UTF-8 sequence.
|
|
cut := limit
|
|
for cut > 0 && !utf8.RuneStart(value[cut]) {
|
|
cut--
|
|
}
|
|
return value[:cut] + "...(截断)"
|
|
}
|
|
|
|
func (e *TaskExecutor) recordTaskLog(ctx context.Context, log *taskbiz.TimedTaskLog) {
|
|
if e == nil || e.tasks == nil || log == nil {
|
|
return
|
|
}
|
|
// This runs from Run's deferred block, so a panic here would replace the
|
|
// task's own result with a log-writing panic. Recover, but report it: a
|
|
// silently swallowed panic hides the fact that no execution log was stored.
|
|
defer func() {
|
|
if recovered := recover(); recovered != nil {
|
|
slog.Default().Error("recording the timed task log panicked", "mod", "timedTask", "task_id", log.TaskID, "task_name", log.TaskName, "panic", recovered)
|
|
}
|
|
}()
|
|
if err := e.tasks.RecordTaskLog(ctx, log); err != nil {
|
|
slog.Default().Error("recording the timed task log failed", "mod", "timedTask", "task_id", log.TaskID, "task_name", log.TaskName, "error", err)
|
|
}
|
|
}
|
|
|
|
// methodOrphaned reports whether a previous invocation of the task's method is
|
|
// still running after its execution timeout elapsed.
|
|
func (e *TaskExecutor) methodOrphaned(id uint) bool {
|
|
e.orphanMu.Lock()
|
|
defer e.orphanMu.Unlock()
|
|
_, orphaned := e.orphans[id]
|
|
return orphaned
|
|
}
|
|
|
|
// markMethodOrphan keeps the task claimed until the abandoned goroutine
|
|
// returns. Without it the scheduler's running-set is cleared as soon as the
|
|
// timeout fires, so the next trigger would run the same method concurrently.
|
|
func (e *TaskExecutor) markMethodOrphan(id uint, done <-chan error) {
|
|
e.orphanMu.Lock()
|
|
if e.orphans == nil {
|
|
e.orphans = map[uint]struct{}{}
|
|
}
|
|
e.orphans[id] = struct{}{}
|
|
e.orphanMu.Unlock()
|
|
go func() {
|
|
<-done
|
|
e.orphanMu.Lock()
|
|
delete(e.orphans, id)
|
|
e.orphanMu.Unlock()
|
|
}()
|
|
}
|
|
|
|
func (e *TaskExecutor) runMethod(ctx context.Context, task *taskbiz.TimedTask) error {
|
|
if task == nil {
|
|
return errors.New("任务不能为空")
|
|
}
|
|
if e == nil {
|
|
return errors.New("任务执行器未初始化")
|
|
}
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
var method platformtask.MethodFunc
|
|
var ok bool
|
|
if e.methods != nil {
|
|
method, ok = e.methods.Lookup(task.MethodName)
|
|
}
|
|
if !ok {
|
|
return fmt.Errorf("方法 %s 未注册(需通过 platform/task.Registry 注册)", task.MethodName)
|
|
}
|
|
if e.methodOrphaned(task.ID) {
|
|
return errTaskOrphaned
|
|
}
|
|
runCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
|
|
defer cancel()
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
defer func() {
|
|
if recovered := recover(); recovered != nil {
|
|
done <- fmt.Errorf("panic: %v", recovered)
|
|
}
|
|
}()
|
|
done <- method(runCtx, json.RawMessage(task.Params))
|
|
}()
|
|
select {
|
|
case err := <-done:
|
|
if errors.Is(err, context.DeadlineExceeded) {
|
|
return errTaskTimeout
|
|
}
|
|
return err
|
|
case <-runCtx.Done():
|
|
e.markMethodOrphan(task.ID, done)
|
|
if errors.Is(runCtx.Err(), context.DeadlineExceeded) {
|
|
return errTaskTimeout
|
|
}
|
|
return runCtx.Err()
|
|
}
|
|
}
|
|
|
|
func (e *TaskExecutor) Run(ctx context.Context, task *taskbiz.TimedTask, trigger string) (log *taskbiz.TimedTaskLog) {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
started := time.Now()
|
|
log = &taskbiz.TimedTaskLog{TriggerType: trigger, StartedAt: started, Status: "success"}
|
|
if task != nil {
|
|
log.TaskID = task.ID
|
|
log.TaskName = task.Name
|
|
}
|
|
defer func() {
|
|
if recovered := recover(); recovered != nil {
|
|
log.Status = "fail"
|
|
log.ErrorMsg = truncateTaskText(fmt.Sprint(recovered))
|
|
}
|
|
log.Output = truncateTaskText(log.Output)
|
|
log.FinishedAt = time.Now()
|
|
log.DurationMS = log.FinishedAt.Sub(started).Milliseconds()
|
|
logCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
|
|
defer cancel()
|
|
if task != nil {
|
|
e.recordTaskLog(logCtx, log)
|
|
}
|
|
}()
|
|
if task == nil {
|
|
log.Status = "fail"
|
|
log.ErrorMsg = "任务不能为空"
|
|
return log
|
|
}
|
|
if e == nil {
|
|
log.Status = "fail"
|
|
log.ErrorMsg = "任务执行器未初始化"
|
|
return log
|
|
}
|
|
var err error
|
|
switch task.ExecutorType {
|
|
case taskbiz.TaskExecutorMethod:
|
|
err = e.runMethod(ctx, task)
|
|
case taskbiz.TaskExecutorHTTP:
|
|
runCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
|
log.Output, err = e.runHTTP(runCtx, task)
|
|
cancel()
|
|
default:
|
|
err = fmt.Errorf("未知执行器类型: %s", task.ExecutorType)
|
|
}
|
|
if err != nil {
|
|
if errors.Is(err, errTaskTimeout) || errors.Is(err, context.DeadlineExceeded) {
|
|
log.Status = "timeout"
|
|
} else {
|
|
log.Status = "fail"
|
|
}
|
|
log.ErrorMsg = truncateTaskText(err.Error())
|
|
}
|
|
return log
|
|
}
|