优化结构
This commit is contained in:
Yvan
@@ -1,174 +0,0 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"kra/internal/biz"
|
||||
"kra/internal/conf"
|
||||
)
|
||||
|
||||
type TaskExecutor struct {
|
||||
tasks *biz.TaskUsecase
|
||||
media *biz.MediaUsecase
|
||||
runtime *conf.Runtime
|
||||
}
|
||||
|
||||
func NewTaskExecutor(tasks *biz.TaskUsecase, media *biz.MediaUsecase, runtime *conf.Runtime) *TaskExecutor {
|
||||
executor := &TaskExecutor{tasks: tasks, media: media, runtime: runtime}
|
||||
registerTaskMethods(executor)
|
||||
return executor
|
||||
}
|
||||
|
||||
func privateIP(ip net.IP) bool {
|
||||
return ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalMulticast() || ip.IsLinkLocalUnicast() || ip.IsUnspecified()
|
||||
}
|
||||
|
||||
func taskHTTPClient(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
|
||||
}}
|
||||
return &http.Client{Timeout: 30 * time.Second, Transport: &http.Transport{Proxy: nil, DialContext: dialer.DialContext}}
|
||||
}
|
||||
|
||||
func (e *TaskExecutor) runHTTP(ctx context.Context, task *biz.TimedTask) (string, error) {
|
||||
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)
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
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, _ := io.ReadAll(io.LimitReader(response.Body, 1<<20))
|
||||
output := fmt.Sprintf("HTTP %d: %s", response.StatusCode, string(body))
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return output, fmt.Errorf("非 2xx 响应: %d", response.StatusCode)
|
||||
}
|
||||
return output, nil
|
||||
}
|
||||
|
||||
var errTaskTimeout = errors.New("任务执行超时")
|
||||
|
||||
func truncateTaskText(value string) string {
|
||||
const limit = 4000
|
||||
if len(value) <= limit {
|
||||
return value
|
||||
}
|
||||
return value[:limit] + "...(截断)"
|
||||
}
|
||||
|
||||
func (e *TaskExecutor) runMethod(ctx context.Context, task *biz.TimedTask) error {
|
||||
method, ok := biz.TaskMethodByName(task.MethodName)
|
||||
if !ok {
|
||||
return fmt.Errorf("方法 %s 未注册(需在 internal/worker/task_registry.go 经 biz.RegisterTaskMethod 注册)", task.MethodName)
|
||||
}
|
||||
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():
|
||||
if errors.Is(runCtx.Err(), context.DeadlineExceeded) {
|
||||
return errTaskTimeout
|
||||
}
|
||||
return runCtx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (e *TaskExecutor) Run(ctx context.Context, task *biz.TimedTask, trigger string) (log *biz.TimedTaskLog) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
started := time.Now()
|
||||
log = &biz.TimedTaskLog{TaskID: task.ID, TaskName: task.Name, TriggerType: trigger, StartedAt: started, Status: "success"}
|
||||
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()
|
||||
_ = e.tasks.RecordTaskLog(logCtx, log)
|
||||
}()
|
||||
var err error
|
||||
switch task.ExecutorType {
|
||||
case biz.TaskExecutorMethod:
|
||||
err = e.runMethod(ctx, task)
|
||||
case biz.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
|
||||
}
|
||||
@@ -1,97 +0,0 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"kra/internal/biz"
|
||||
)
|
||||
|
||||
func TestPrivateIP(t *testing.T) {
|
||||
for _, value := range []string{"127.0.0.1", "10.0.0.1", "172.16.0.1", "192.168.1.1", "169.254.1.1", "::1"} {
|
||||
if !privateIP(net.ParseIP(value)) {
|
||||
t.Fatalf("expected %s to be private", value)
|
||||
}
|
||||
}
|
||||
if privateIP(net.ParseIP("8.8.8.8")) {
|
||||
t.Fatal("public IP was classified as private")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMethodReportsParentDeadlineAsTimeout(t *testing.T) {
|
||||
const methodName = "worker-test-parent-deadline"
|
||||
biz.RegisterTaskMethod(methodName, "test", func(ctx context.Context, _ json.RawMessage) error {
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
})
|
||||
executor := &TaskExecutor{}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := executor.runMethod(ctx, &biz.TimedTask{MethodName: methodName}); !errors.Is(err, errTaskTimeout) {
|
||||
t.Fatalf("runMethod() error = %v, want errTaskTimeout", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMethodConvertsPanicToError(t *testing.T) {
|
||||
const methodName = "worker-test-panic"
|
||||
biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
|
||||
panic("boom")
|
||||
})
|
||||
executor := &TaskExecutor{}
|
||||
err := executor.runMethod(context.Background(), &biz.TimedTask{MethodName: methodName})
|
||||
if err == nil || !strings.Contains(err.Error(), "panic: boom") {
|
||||
t.Fatalf("runMethod() error = %v, want recovered panic", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMethodUsesParentContext(t *testing.T) {
|
||||
const methodName = "worker-test-parent-context"
|
||||
started := make(chan struct{})
|
||||
biz.RegisterTaskMethod(methodName, "test", func(ctx context.Context, _ json.RawMessage) error {
|
||||
close(started)
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
})
|
||||
executor := &TaskExecutor{}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- executor.runMethod(ctx, &biz.TimedTask{MethodName: methodName})
|
||||
}()
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("registered method did not start")
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case err := <-done:
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("runMethod() error = %v, want context.Canceled", err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("runMethod() did not stop after parent cancellation")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunMethodHonorsDeadlineWhenMethodIgnoresContext(t *testing.T) {
|
||||
const methodName = "worker-test-ignores-context"
|
||||
release := make(chan struct{})
|
||||
biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
|
||||
<-release
|
||||
return nil
|
||||
})
|
||||
executor := &TaskExecutor{}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := executor.runMethod(ctx, &biz.TimedTask{MethodName: methodName}); !errors.Is(err, errTaskTimeout) {
|
||||
close(release)
|
||||
t.Fatalf("runMethod() error = %v, want errTaskTimeout", err)
|
||||
}
|
||||
close(release)
|
||||
}
|
||||
@@ -1,32 +0,0 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
|
||||
"kra/internal/biz"
|
||||
)
|
||||
|
||||
// registerTaskMethods defines the in-process methods available to the task
|
||||
// management page. Add application-specific registrations here.
|
||||
func registerTaskMethods(executor *TaskExecutor) {
|
||||
biz.RegisterTaskMethod(
|
||||
biz.TaskMethodClearDB,
|
||||
"清理数据库过期日志(操作记录/JWT黑名单/定时任务执行日志)",
|
||||
func(ctx context.Context, _ json.RawMessage) error {
|
||||
return executor.tasks.CleanupLogs(ctx)
|
||||
},
|
||||
)
|
||||
biz.RegisterTaskMethod(
|
||||
biz.TaskMethodUploads,
|
||||
"清理过期大文件上传会话",
|
||||
func(ctx context.Context, _ json.RawMessage) error {
|
||||
ttl := 24
|
||||
config := executor.runtime.Admin()
|
||||
if config != nil && config.Media != nil && config.Media.SessionTtl > 0 {
|
||||
ttl = int(config.Media.SessionTtl)
|
||||
}
|
||||
return executor.media.CleanupStale(ctx, ttl)
|
||||
},
|
||||
)
|
||||
}
|
||||
@@ -1,274 +0,0 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"kra/internal/biz"
|
||||
)
|
||||
|
||||
type workerTaskRepo struct {
|
||||
mu sync.Mutex
|
||||
items []*biz.TimedTask
|
||||
listErr error
|
||||
logs chan *biz.TimedTaskLog
|
||||
}
|
||||
|
||||
func (r *workerTaskRepo) CreateTask(context.Context, *biz.TimedTask) error { return nil }
|
||||
func (r *workerTaskRepo) UpdateTask(context.Context, *biz.TimedTask) error { return nil }
|
||||
func (r *workerTaskRepo) DeleteTask(context.Context, uint) error { return nil }
|
||||
func (r *workerTaskRepo) ToggleTask(context.Context, uint, bool) error { return nil }
|
||||
func (r *workerTaskRepo) CleanupLogs(context.Context) error { return nil }
|
||||
func (r *workerTaskRepo) TaskNameExists(context.Context, string, uint) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
func (r *workerTaskRepo) FindTask(_ context.Context, id uint) (*biz.TimedTask, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, item := range r.items {
|
||||
if item.ID == id {
|
||||
return cloneTimedTask(item), nil
|
||||
}
|
||||
}
|
||||
return nil, errors.New("task not found")
|
||||
}
|
||||
func (r *workerTaskRepo) ListTasks(context.Context, int, int, *biz.TimedTask) ([]*biz.TimedTask, int64, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.listErr != nil {
|
||||
return nil, 0, r.listErr
|
||||
}
|
||||
items := make([]*biz.TimedTask, 0, len(r.items))
|
||||
for _, item := range r.items {
|
||||
items = append(items, cloneTimedTask(item))
|
||||
}
|
||||
return items, int64(len(items)), nil
|
||||
}
|
||||
func (r *workerTaskRepo) RecordTaskLog(_ context.Context, value *biz.TimedTaskLog) error {
|
||||
if r.logs != nil {
|
||||
copy := *value
|
||||
r.logs <- ©
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (r *workerTaskRepo) ListTaskLogs(context.Context, int, int, uint, string) ([]*biz.TimedTaskLog, int64, error) {
|
||||
return nil, 0, nil
|
||||
}
|
||||
|
||||
func newTestTaskScheduler(repo *workerTaskRepo) *TaskScheduler {
|
||||
tasks := biz.NewTaskUsecase(repo)
|
||||
executor := &TaskExecutor{tasks: tasks}
|
||||
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
return NewTaskScheduler(tasks, nil, executor, logger)
|
||||
}
|
||||
|
||||
func TestReloadPreservesExistingScheduleOnPreparationFailure(t *testing.T) {
|
||||
repo := &workerTaskRepo{}
|
||||
scheduler := newTestTaskScheduler(repo)
|
||||
old := &biz.TimedTask{ID: 1, Name: "old", Spec: "0 0 * * *", Enabled: true}
|
||||
if err := scheduler.Schedule(old); err != nil {
|
||||
t.Fatalf("schedule old task: %v", err)
|
||||
}
|
||||
|
||||
t.Run("query failure", func(t *testing.T) {
|
||||
repo.mu.Lock()
|
||||
repo.listErr = errors.New("database unavailable")
|
||||
repo.mu.Unlock()
|
||||
if err := scheduler.Reload(context.Background()); err == nil {
|
||||
t.Fatal("Reload() error = nil, want query error")
|
||||
}
|
||||
assertOnlyScheduledTask(t, scheduler, old.ID)
|
||||
})
|
||||
|
||||
t.Run("invalid cron", func(t *testing.T) {
|
||||
repo.mu.Lock()
|
||||
repo.listErr = nil
|
||||
repo.items = []*biz.TimedTask{{ID: 2, Name: "invalid", Spec: "bad cron", Enabled: true}}
|
||||
repo.mu.Unlock()
|
||||
if err := scheduler.Reload(context.Background()); err == nil {
|
||||
t.Fatal("Reload() error = nil, want cron error")
|
||||
}
|
||||
assertOnlyScheduledTask(t, scheduler, old.ID)
|
||||
})
|
||||
}
|
||||
|
||||
func assertOnlyScheduledTask(t *testing.T, scheduler *TaskScheduler, id uint) {
|
||||
t.Helper()
|
||||
scheduler.mu.Lock()
|
||||
defer scheduler.mu.Unlock()
|
||||
if len(scheduler.entries) != 1 {
|
||||
t.Fatalf("scheduled task count = %d, want 1", len(scheduler.entries))
|
||||
}
|
||||
if _, ok := scheduler.entries[id]; !ok {
|
||||
t.Fatalf("scheduled task %d was replaced after failed reload", id)
|
||||
}
|
||||
if got := len(scheduler.standard.Entries()) + len(scheduler.seconds.Entries()); got != 1 {
|
||||
t.Fatalf("cron entry count = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentScheduleKeepsSingleEntry(t *testing.T) {
|
||||
scheduler := newTestTaskScheduler(&workerTaskRepo{})
|
||||
const workers = 32
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if err := scheduler.Schedule(&biz.TimedTask{ID: 7, Name: "same", Spec: "0 0 * * *", Enabled: true}); err != nil {
|
||||
t.Errorf("Schedule() error = %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
assertOnlyScheduledTask(t, scheduler, 7)
|
||||
}
|
||||
|
||||
func TestManualTriggerSkipsOverlappingExecution(t *testing.T) {
|
||||
const methodName = "worker-test-no-overlap"
|
||||
started := make(chan struct{}, 2)
|
||||
release := make(chan struct{})
|
||||
var calls atomic.Int32
|
||||
biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
|
||||
calls.Add(1)
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
return nil
|
||||
})
|
||||
|
||||
task := &biz.TimedTask{ID: 9, Name: "single", ExecutorType: biz.TaskExecutorMethod, MethodName: methodName}
|
||||
repo := &workerTaskRepo{items: []*biz.TimedTask{task}, logs: make(chan *biz.TimedTaskLog, 2)}
|
||||
scheduler := newTestTaskScheduler(repo)
|
||||
if err := scheduler.TriggerID(context.Background(), task.ID); err != nil {
|
||||
t.Fatalf("first manual trigger failed: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first task execution did not start")
|
||||
}
|
||||
|
||||
if err := scheduler.TriggerID(context.Background(), task.ID); err == nil {
|
||||
t.Fatal("overlapping manual trigger returned success")
|
||||
}
|
||||
if got := calls.Load(); got != 1 {
|
||||
t.Fatalf("execution count while first run is active = %d, want 1", got)
|
||||
}
|
||||
close(release)
|
||||
select {
|
||||
case <-repo.logs:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first task execution did not finish")
|
||||
}
|
||||
waitForSchedulerIdle(t, scheduler)
|
||||
|
||||
if err := scheduler.TriggerID(context.Background(), task.ID); err != nil {
|
||||
t.Fatalf("manual trigger failed after the previous run finished: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("task could not run again after the first execution finished")
|
||||
}
|
||||
select {
|
||||
case <-repo.logs:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("second task execution did not finish")
|
||||
}
|
||||
waitForSchedulerIdle(t, scheduler)
|
||||
if got := calls.Load(); got != 2 {
|
||||
t.Fatalf("total execution count = %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutomaticAndManualTriggersShareOverlapGate(t *testing.T) {
|
||||
const methodName = "worker-test-auto-manual-overlap"
|
||||
started := make(chan struct{}, 1)
|
||||
release := make(chan struct{})
|
||||
var calls atomic.Int32
|
||||
biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
|
||||
calls.Add(1)
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
return nil
|
||||
})
|
||||
|
||||
task := &biz.TimedTask{ID: 10, Name: "shared-gate", ExecutorType: biz.TaskExecutorMethod, MethodName: methodName}
|
||||
repo := &workerTaskRepo{logs: make(chan *biz.TimedTaskLog, 1)}
|
||||
scheduler := newTestTaskScheduler(repo)
|
||||
if !scheduler.dispatch(task, "auto", true) {
|
||||
t.Fatal("automatic dispatch was rejected while idle")
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("automatic task execution did not start")
|
||||
}
|
||||
if scheduler.Trigger(task) {
|
||||
t.Fatal("manual trigger bypassed an active automatic execution")
|
||||
}
|
||||
if got := calls.Load(); got != 1 {
|
||||
t.Fatalf("execution count while automatic run is active = %d, want 1", got)
|
||||
}
|
||||
close(release)
|
||||
select {
|
||||
case <-repo.logs:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("automatic task execution did not finish")
|
||||
}
|
||||
waitForSchedulerIdle(t, scheduler)
|
||||
}
|
||||
|
||||
func waitForSchedulerIdle(t *testing.T, scheduler *TaskScheduler) {
|
||||
t.Helper()
|
||||
scheduler.runMu.Lock()
|
||||
idle := scheduler.idle
|
||||
scheduler.runMu.Unlock()
|
||||
select {
|
||||
case <-idle:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("scheduler did not become idle")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopCancelsAndWaitsForActiveRuns(t *testing.T) {
|
||||
scheduler := newTestTaskScheduler(&workerTaskRepo{})
|
||||
runCtx, cancel := context.WithCancel(context.Background())
|
||||
scheduler.ctxMu.Lock()
|
||||
scheduler.runContext = runCtx
|
||||
scheduler.cancel = cancel
|
||||
scheduler.ctxMu.Unlock()
|
||||
|
||||
taskCtx, ok := scheduler.beginRun(12)
|
||||
if !ok {
|
||||
t.Fatal("beginRun() rejected an idle scheduler")
|
||||
}
|
||||
finished := make(chan struct{})
|
||||
go func() {
|
||||
<-taskCtx.Done()
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
close(finished)
|
||||
scheduler.finishRun(12)
|
||||
}()
|
||||
|
||||
stopCtx, stopCancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer stopCancel()
|
||||
if err := scheduler.Stop(stopCtx); err != nil {
|
||||
t.Fatalf("Stop() error = %v", err)
|
||||
}
|
||||
select {
|
||||
case <-finished:
|
||||
default:
|
||||
t.Fatal("Stop() returned before the active run finished")
|
||||
}
|
||||
if _, ok := scheduler.beginRun(13); ok {
|
||||
t.Fatal("scheduler accepted a new run after Stop()")
|
||||
}
|
||||
}
|
||||
@@ -1,386 +0,0 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/robfig/cron/v3"
|
||||
"kra/internal/biz"
|
||||
)
|
||||
|
||||
type TaskScheduler struct {
|
||||
tasks *biz.TaskUsecase
|
||||
authorities *biz.AuthorityUsecase
|
||||
executor *TaskExecutor
|
||||
logger *slog.Logger
|
||||
standard *cron.Cron
|
||||
seconds *cron.Cron
|
||||
mu sync.Mutex
|
||||
entries map[uint]scheduledEntry
|
||||
ctxMu sync.RWMutex
|
||||
runContext context.Context
|
||||
cancel context.CancelFunc
|
||||
runMu sync.Mutex
|
||||
running map[uint]struct{}
|
||||
idle chan struct{}
|
||||
stopping bool
|
||||
subMu sync.RWMutex
|
||||
subscribers map[uint]map[chan []byte]struct{}
|
||||
}
|
||||
|
||||
type scheduledEntry struct {
|
||||
seconds bool
|
||||
entry cron.EntryID
|
||||
}
|
||||
|
||||
func NewTaskScheduler(tasks *biz.TaskUsecase, authorities *biz.AuthorityUsecase, executor *TaskExecutor, logger *slog.Logger) *TaskScheduler {
|
||||
idle := make(chan struct{})
|
||||
close(idle)
|
||||
return &TaskScheduler{tasks: tasks, authorities: authorities, executor: executor, logger: logger.With("mod", "timedTask"), standard: cron.New(), seconds: cron.New(cron.WithSeconds()), entries: map[uint]scheduledEntry{}, running: map[uint]struct{}{}, idle: idle, subscribers: map[uint]map[chan []byte]struct{}{}}
|
||||
}
|
||||
|
||||
func NewTaskRuntime(scheduler *TaskScheduler) biz.TaskRuntime { return scheduler }
|
||||
|
||||
func (s *TaskScheduler) Start(ctx context.Context) error {
|
||||
runContext, cancel := context.WithCancel(ctx)
|
||||
s.ctxMu.Lock()
|
||||
s.runContext, s.cancel = runContext, cancel
|
||||
s.ctxMu.Unlock()
|
||||
s.runMu.Lock()
|
||||
s.stopping = false
|
||||
s.runMu.Unlock()
|
||||
s.standard.Start()
|
||||
s.seconds.Start()
|
||||
if err := s.Reload(ctx); err != nil {
|
||||
s.logger.WarnContext(ctx, "timed task table is not ready", "error", err)
|
||||
}
|
||||
<-ctx.Done()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Stop(ctx context.Context) error {
|
||||
s.closeSubscribers()
|
||||
s.runMu.Lock()
|
||||
s.stopping = true
|
||||
idle := s.idle
|
||||
if idle == nil {
|
||||
idle = make(chan struct{})
|
||||
close(idle)
|
||||
}
|
||||
s.runMu.Unlock()
|
||||
s.ctxMu.Lock()
|
||||
cancel := s.cancel
|
||||
s.ctxMu.Unlock()
|
||||
if cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
standardDone, secondsDone := s.standard.Stop().Done(), s.seconds.Stop().Done()
|
||||
for _, done := range []<-chan struct{}{standardDone, secondsDone, idle} {
|
||||
select {
|
||||
case <-done:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Remove(id uint) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.removeLocked(id)
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) removeLocked(id uint) {
|
||||
if old, ok := s.entries[id]; ok {
|
||||
if old.seconds {
|
||||
s.seconds.Remove(old.entry)
|
||||
} else {
|
||||
s.standard.Remove(old.entry)
|
||||
}
|
||||
delete(s.entries, id)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Reload(ctx context.Context) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
items, _, err := s.tasks.ListTasks(ctx, 0, 0, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
type preparedTask struct {
|
||||
task *biz.TimedTask
|
||||
schedule cron.Schedule
|
||||
}
|
||||
prepared := make([]preparedTask, 0, len(items))
|
||||
for _, task := range items {
|
||||
if task.Enabled {
|
||||
schedule, parseErr := parseTaskSchedule(task)
|
||||
if parseErr != nil {
|
||||
return fmt.Errorf("定时任务 %d 的 cron 表达式非法: %w", task.ID, parseErr)
|
||||
}
|
||||
prepared = append(prepared, preparedTask{task: cloneTimedTask(task), schedule: schedule})
|
||||
}
|
||||
}
|
||||
|
||||
for id := range s.entries {
|
||||
s.removeLocked(id)
|
||||
}
|
||||
for _, item := range prepared {
|
||||
s.scheduleLocked(item.task, item.schedule)
|
||||
}
|
||||
s.logger.InfoContext(ctx, "定时任务重载完成", "task_count", len(items))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) executionContext() context.Context {
|
||||
s.ctxMu.RLock()
|
||||
defer s.ctxMu.RUnlock()
|
||||
if s.runContext != nil {
|
||||
return s.runContext
|
||||
}
|
||||
return context.Background()
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) beginRun(id uint) (context.Context, bool) {
|
||||
ctx := s.executionContext()
|
||||
s.runMu.Lock()
|
||||
defer s.runMu.Unlock()
|
||||
if s.stopping || ctx.Err() != nil {
|
||||
return nil, false
|
||||
}
|
||||
if s.running == nil {
|
||||
s.running = map[uint]struct{}{}
|
||||
}
|
||||
if _, exists := s.running[id]; exists {
|
||||
return nil, false
|
||||
}
|
||||
if len(s.running) == 0 {
|
||||
s.idle = make(chan struct{})
|
||||
}
|
||||
s.running[id] = struct{}{}
|
||||
return ctx, true
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) finishRun(id uint) {
|
||||
s.runMu.Lock()
|
||||
defer s.runMu.Unlock()
|
||||
if _, exists := s.running[id]; !exists {
|
||||
return
|
||||
}
|
||||
delete(s.running, id)
|
||||
if len(s.running) == 0 && s.idle != nil {
|
||||
close(s.idle)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) dispatch(task *biz.TimedTask, trigger string, async bool) bool {
|
||||
ctx, ok := s.beginRun(task.ID)
|
||||
if !ok {
|
||||
s.logger.Warn("timed task skipped because it is already running or the scheduler is stopping", "task_id", task.ID, "task_name", task.Name, "trigger_type", trigger)
|
||||
return false
|
||||
}
|
||||
run := func() {
|
||||
defer s.finishRun(task.ID)
|
||||
s.run(ctx, task, trigger)
|
||||
}
|
||||
if async {
|
||||
go run()
|
||||
return true
|
||||
}
|
||||
run()
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) run(ctx context.Context, task *biz.TimedTask, trigger string) {
|
||||
log := s.executor.Run(ctx, task, trigger)
|
||||
attributes := []any{"task_id", log.TaskID, "task_name", log.TaskName, "trigger_type", log.TriggerType, "status", log.Status, "duration_ms", log.DurationMS, "started_at", log.StartedAt, "finished_at", log.FinishedAt}
|
||||
if log.ErrorMsg != "" {
|
||||
attributes = append(attributes, "error", log.ErrorMsg)
|
||||
}
|
||||
if log.Status == "success" {
|
||||
s.logger.Info("timed task finished", attributes...)
|
||||
} else {
|
||||
s.logger.Error("timed task finished", attributes...)
|
||||
}
|
||||
if log.Status != "success" {
|
||||
alertCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
|
||||
defer cancel()
|
||||
ids, err := s.authorities.AuthorityUserIDs(alertCtx, 888)
|
||||
if err != nil {
|
||||
s.logger.Error("query timed task alert recipients failed", "error", err)
|
||||
return
|
||||
}
|
||||
s.PublishToUsers(ids, map[string]any{"taskId": task.ID, "name": task.Name, "error": log.ErrorMsg, "time": time.Now().Format(time.RFC3339)})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Schedule(task *biz.TimedTask) error {
|
||||
if task == nil {
|
||||
return errors.New("定时任务不能为空")
|
||||
}
|
||||
var schedule cron.Schedule
|
||||
var err error
|
||||
if task.Enabled {
|
||||
schedule, err = parseTaskSchedule(task)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.removeLocked(task.ID)
|
||||
if !task.Enabled {
|
||||
return nil
|
||||
}
|
||||
s.scheduleLocked(cloneTimedTask(task), schedule)
|
||||
return nil
|
||||
}
|
||||
|
||||
func parseTaskSchedule(task *biz.TimedTask) (cron.Schedule, error) {
|
||||
if task.WithSeconds {
|
||||
return cron.NewParser(cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor).Parse(task.Spec)
|
||||
}
|
||||
return cron.ParseStandard(task.Spec)
|
||||
}
|
||||
|
||||
func cloneTimedTask(task *biz.TimedTask) *biz.TimedTask {
|
||||
copy := *task
|
||||
copy.Params = append([]byte(nil), task.Params...)
|
||||
copy.HTTPHeader = append([]byte(nil), task.HTTPHeader...)
|
||||
return ©
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) scheduleLocked(task *biz.TimedTask, schedule cron.Schedule) {
|
||||
run := func() {
|
||||
s.dispatch(task, "auto", false)
|
||||
}
|
||||
if task.WithSeconds {
|
||||
id := s.seconds.Schedule(schedule, cron.FuncJob(run))
|
||||
s.entries[task.ID] = scheduledEntry{seconds: true, entry: id}
|
||||
} else {
|
||||
id := s.standard.Schedule(schedule, cron.FuncJob(run))
|
||||
s.entries[task.ID] = scheduledEntry{entry: id}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
task, err := s.tasks.FindTask(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var schedule cron.Schedule
|
||||
if task.Enabled {
|
||||
schedule, err = parseTaskSchedule(task)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
s.removeLocked(task.ID)
|
||||
if task.Enabled {
|
||||
s.scheduleLocked(cloneTimedTask(task), schedule)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error {
|
||||
task, err := s.tasks.FindTask(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !s.Trigger(task) {
|
||||
return errors.New("任务正在执行或调度器正在停止")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) NextRuns() map[uint]time.Time {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := map[uint]time.Time{}
|
||||
for id, item := range s.entries {
|
||||
if item.seconds {
|
||||
out[id] = s.seconds.Entry(item.entry).Next
|
||||
} else {
|
||||
out[id] = s.standard.Entry(item.entry).Next
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Trigger(task *biz.TimedTask) bool {
|
||||
if task == nil {
|
||||
return false
|
||||
}
|
||||
return s.dispatch(cloneTimedTask(task), "manual", true)
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) Subscribe(userID uint) chan []byte {
|
||||
ch := make(chan []byte, 16)
|
||||
s.subMu.Lock()
|
||||
if s.subscribers[userID] == nil {
|
||||
s.subscribers[userID] = map[chan []byte]struct{}{}
|
||||
}
|
||||
if len(s.subscribers[userID]) >= 10 {
|
||||
for old := range s.subscribers[userID] {
|
||||
delete(s.subscribers[userID], old)
|
||||
close(old)
|
||||
break
|
||||
}
|
||||
}
|
||||
s.subscribers[userID][ch] = struct{}{}
|
||||
s.subMu.Unlock()
|
||||
return ch
|
||||
}
|
||||
func (s *TaskScheduler) Unsubscribe(userID uint, ch chan []byte) {
|
||||
s.subMu.Lock()
|
||||
if subscribers := s.subscribers[userID]; subscribers != nil {
|
||||
if _, ok := subscribers[ch]; !ok {
|
||||
s.subMu.Unlock()
|
||||
return
|
||||
}
|
||||
delete(subscribers, ch)
|
||||
if len(subscribers) == 0 {
|
||||
delete(s.subscribers, userID)
|
||||
}
|
||||
close(ch)
|
||||
}
|
||||
s.subMu.Unlock()
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) closeSubscribers() {
|
||||
s.subMu.Lock()
|
||||
defer s.subMu.Unlock()
|
||||
for userID, subscribers := range s.subscribers {
|
||||
for ch := range subscribers {
|
||||
close(ch)
|
||||
}
|
||||
delete(s.subscribers, userID)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *TaskScheduler) PublishToUsers(userIDs []uint, value any) {
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
raw = []byte(fmt.Sprint(value))
|
||||
}
|
||||
s.subMu.RLock()
|
||||
defer s.subMu.RUnlock()
|
||||
for _, userID := range userIDs {
|
||||
for ch := range s.subscribers[userID] {
|
||||
select {
|
||||
case ch <- raw:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,50 +0,0 @@
|
||||
package worker
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestCloseSubscribersClosesEveryAlertStream(t *testing.T) {
|
||||
scheduler := &TaskScheduler{subscribers: map[uint]map[chan []byte]struct{}{}}
|
||||
first := scheduler.Subscribe(1)
|
||||
second := scheduler.Subscribe(1)
|
||||
third := scheduler.Subscribe(2)
|
||||
|
||||
scheduler.closeSubscribers()
|
||||
|
||||
for _, ch := range []chan []byte{first, second, third} {
|
||||
if _, open := <-ch; open {
|
||||
t.Fatal("subscriber channel remains open after scheduler shutdown")
|
||||
}
|
||||
}
|
||||
if len(scheduler.subscribers) != 0 {
|
||||
t.Fatalf("subscriber registry was not cleared: %#v", scheduler.subscribers)
|
||||
}
|
||||
|
||||
// Handlers defer Unsubscribe. It must remain safe after shutdown has already
|
||||
// closed and removed every channel.
|
||||
scheduler.Unsubscribe(1, first)
|
||||
}
|
||||
|
||||
func TestSubscribeCapsConnectionsPerUser(t *testing.T) {
|
||||
scheduler := &TaskScheduler{subscribers: map[uint]map[chan []byte]struct{}{}}
|
||||
channels := make([]chan []byte, 0, 11)
|
||||
for i := 0; i < 11; i++ {
|
||||
channels = append(channels, scheduler.Subscribe(7))
|
||||
}
|
||||
if got := len(scheduler.subscribers[7]); got != 10 {
|
||||
t.Fatalf("subscriber count = %d, want 10", got)
|
||||
}
|
||||
closed := 0
|
||||
for _, ch := range channels {
|
||||
select {
|
||||
case _, open := <-ch:
|
||||
if !open {
|
||||
closed++
|
||||
}
|
||||
default:
|
||||
}
|
||||
}
|
||||
if closed != 1 {
|
||||
t.Fatalf("evicted subscriber count = %d, want 1", closed)
|
||||
}
|
||||
scheduler.closeSubscribers()
|
||||
}
|
||||
@@ -1,6 +0,0 @@
|
||||
package worker
|
||||
|
||||
import "github.com/google/wire"
|
||||
|
||||
// ProviderSet contains background task runtime providers.
|
||||
var ProviderSet = wire.NewSet(NewTaskExecutor, NewTaskScheduler, NewTaskRuntime)
|
||||
Reference in New Issue
Block a user