yvan/kra-oa
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

优化结构

Browse Source
This commit is contained in:
Yvan
2026-08-23 07:48:02 +08:00
parent 40632c5756
commit 7a9a738c6a
15 changed files with 590 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/worker/task_executor.go
+53 -13
View File
@@ -14,27 +14,19 @@ import (
"syscall"
"time"
"kra/internal/biz/system"
taskbiz "kra/internal/biz/task"
"kra/internal/conf"
)
type TaskExecutor struct {
tasks *taskbiz.TaskUsecase
media *system.MediaUsecase
runtime *conf.Runtime
methods taskbiz.TaskMethodRegistry
}
func NewTaskExecutor(tasks *taskbiz.TaskUsecase, media *system.MediaUsecase, runtime *conf.Runtime) *TaskExecutor {
return NewTaskExecutorWithRegistry(tasks, media, runtime, taskbiz.DefaultTaskMethodRegistry())
}
func NewTaskExecutorWithRegistry(tasks *taskbiz.TaskUsecase, media *system.MediaUsecase, runtime *conf.Runtime, methods taskbiz.TaskMethodRegistry) *TaskExecutor {
func NewTaskExecutorWithRegistry(tasks *taskbiz.TaskUsecase, methods taskbiz.TaskMethodRegistry) *TaskExecutor {
if methods == nil {
methods = taskbiz.DefaultTaskMethodRegistry()
}
return &TaskExecutor{tasks: tasks, media: media, runtime: runtime, methods: methods}
return &TaskExecutor{tasks: tasks, methods: methods}
}
func privateIP(ip net.IP) bool {
@@ -63,6 +55,12 @@ func taskHTTPClient(allowPrivate bool) *http.Client {
}
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)
@@ -70,6 +68,9 @@ func (e *TaskExecutor) runHTTP(ctx context.Context, task *taskbiz.TimedTask) (st
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
@@ -83,6 +84,9 @@ func (e *TaskExecutor) runHTTP(ctx context.Context, task *taskbiz.TimedTask) (st
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)
@@ -92,8 +96,11 @@ func (e *TaskExecutor) runHTTP(ctx context.Context, task *taskbiz.TimedTask) (st
return "", err
}
defer response.Body.Close()
body, _ := io.ReadAll(io.LimitReader(response.Body, 1<<20))
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)
}
@@ -110,7 +117,24 @@ func truncateTaskText(value string) string {
return value[:limit] + "...(截断)"
}
func (e *TaskExecutor) recordTaskLog(ctx context.Context, log *taskbiz.TimedTaskLog) {
defer func() { _ = recover() }()
if e == nil || e.tasks == nil || log == nil {
return
}
_ = e.tasks.RecordTaskLog(ctx, log)
}
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 taskbiz.TaskMethodFunc
var ok bool
if e.methods != nil {
@@ -151,7 +175,11 @@ func (e *TaskExecutor) Run(ctx context.Context, task *taskbiz.TimedTask, trigger
ctx = context.Background()
}
started := time.Now()
log = &taskbiz.TimedTaskLog{TaskID: task.ID, TaskName: task.Name, TriggerType: trigger, StartedAt: started, Status: "success"}
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"
@@ -162,8 +190,20 @@ func (e *TaskExecutor) Run(ctx context.Context, task *taskbiz.TimedTask, trigger
log.DurationMS = log.FinishedAt.Sub(started).Milliseconds()
logCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
defer cancel()
_ = e.tasks.RecordTaskLog(logCtx, log)
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:
internal/worker/task_executor_test.go
+11
View File
@@ -95,3 +95,14 @@ func TestRunMethodHonorsDeadlineWhenMethodIgnoresContext(t *testing.T) {
}
close(release)
}
func TestRunHandlesNilTask(t *testing.T) {
executor := &TaskExecutor{}
log := executor.Run(context.Background(), nil, "manual")
if log == nil {
t.Fatal("Run() returned nil log")
}
if log.Status != "fail" || log.ErrorMsg != "任务不能为空" {
t.Fatalf("Run() log = %#v, want a failed nil-task result", log)
}
}
internal/worker/task_registry.go
+17 -1
View File
@@ -3,6 +3,7 @@ package worker
import (
"context"
"encoding/json"
"errors"
"kra/internal/biz/system"
taskbiz "kra/internal/biz/task"
@@ -31,6 +32,12 @@ func (methods *TaskMethods) RegisterTasks(registry *platformtask.Registry) {
}
registry.Register(platformtask.Method{
Name: taskbiz.TaskMethodClearDB, Description: "清理数据库过期日志(操作记录/JWT黑名单/定时任务执行日志)", Run: func(ctx context.Context, _ json.RawMessage) error {
if ctx == nil {
ctx = context.Background()
}
if methods.maintenance == nil || methods.tasks == nil {
return errors.New("ClearDB 任务依赖未就绪")
}
if err := methods.maintenance.CleanupExpired(ctx); err != nil {
return err
}
@@ -39,8 +46,17 @@ func (methods *TaskMethods) RegisterTasks(registry *platformtask.Registry) {
})
registry.Register(platformtask.Method{
Name: taskbiz.TaskMethodUploads, Description: "清理过期大文件上传会话", Run: func(ctx context.Context, _ json.RawMessage) error {
if ctx == nil {
ctx = context.Background()
}
if methods.media == nil {
return errors.New("CleanStaleUploads 任务依赖未就绪")
}
ttl := 24
config := methods.runtime.Admin()
var config *conf.AdminBackend
if methods.runtime != nil {
config = methods.runtime.Admin()
}
if config != nil && config.Media != nil && config.Media.SessionTtl > 0 {
ttl = int(config.Media.SessionTtl)
}
internal/worker/task_registry_test.go
+25
View File
@@ -0,0 +1,25 @@
package worker
import (
"context"
"encoding/json"
"testing"
taskbiz "kra/internal/biz/task"
platformtask "kra/pkg/task"
)
func TestTaskMethodsReportUnavailableDependencies(t *testing.T) {
registry := platformtask.NewRegistry()
(&TaskMethods{}).RegisterTasks(registry)
for _, name := range []string{taskbiz.TaskMethodClearDB, taskbiz.TaskMethodUploads} {
method, ok := registry.Lookup(name)
if !ok {
t.Fatalf("task method %q was not registered", name)
}
if err := method(context.Background(), json.RawMessage(`{}`)); err == nil {
t.Fatalf("task method %q returned nil with missing dependencies", name)
}
}
}
internal/worker/task_runtime_test.go
+22
View File
@@ -308,3 +308,25 @@ func TestStartReturnsAfterStop(t *testing.T) {
t.Fatal("Start() did not return after Stop()")
}
}
func TestStopBeforeStartPreventsSchedulerStartup(t *testing.T) {
scheduler := newTestTaskScheduler(&workerTaskRepo{})
if err := scheduler.Stop(context.Background()); err != nil {
t.Fatalf("Stop() error = %v", err)
}
started := make(chan error, 1)
go func() {
started <- scheduler.Start(context.Background())
}()
select {
case err := <-started:
if err != nil {
t.Fatalf("Start() error = %v", err)
}
case <-time.After(time.Second):
t.Fatal("Start() did not return after a pre-start Stop()")
}
if got := len(scheduler.standard.Entries()) + len(scheduler.seconds.Entries()); got != 0 {
t.Fatalf("pre-start Stop() left cron entries: %d", got)
}
}
internal/worker/task_scheduler.go
+120 -14
View File
@@ -23,6 +23,8 @@ type TaskScheduler struct {
seconds *cron.Cron
mu sync.Mutex
entries map[uint]scheduledEntry
lifecycleMu sync.Mutex
started bool
ctxMu sync.RWMutex
runContext context.Context
cancel context.CancelFunc
@@ -48,31 +50,49 @@ func NewTaskScheduler(tasks *taskbiz.TaskUsecase, authorities *system.AuthorityU
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) taskbiz.TaskRuntime { return scheduler }
func NewTaskRuntime(scheduler *TaskScheduler) taskbiz.TaskRuntime {
if scheduler == nil {
return nil
}
return scheduler
}
func NewTaskReloader(scheduler *TaskScheduler) system.TaskReloader { return scheduler }
func NewTaskReloader(scheduler *TaskScheduler) system.TaskReloader {
if scheduler == nil {
return nil
}
return scheduler
}
func (s *TaskScheduler) Start(ctx context.Context) error {
if ctx == nil {
ctx = context.Background()
}
runContext, cancel := context.WithCancel(ctx)
s.ctxMu.Lock()
s.runContext, s.cancel = runContext, cancel
s.ctxMu.Unlock()
logger := s.logger
if logger == nil {
logger = slog.Default()
}
s.lifecycleMu.Lock()
s.runMu.Lock()
if s.stopping {
if s.started || s.stopping {
s.runMu.Unlock()
s.lifecycleMu.Unlock()
cancel()
return nil
}
s.started = true
s.stopping = false
s.runMu.Unlock()
s.ctxMu.Lock()
s.runContext, s.cancel = runContext, cancel
s.ctxMu.Unlock()
s.standard.Start()
s.seconds.Start()
s.lifecycleMu.Unlock()
if err := s.Reload(runContext); err != nil {
s.logger.WarnContext(runContext, "timed task table is not ready", "error", err)
logger.WarnContext(runContext, "timed task table is not ready", "error", err)
}
s.runMu.Unlock()
// Stop cancels runContext. Waiting on the child context is important: the
// application lifecycle passes a long-lived parent context to Start and
// invokes Stop separately during graceful shutdown.
@@ -89,8 +109,10 @@ func (s *TaskScheduler) Stop(ctx context.Context) error {
if ctx == nil {
ctx = context.Background()
}
s.lifecycleMu.Lock()
s.runMu.Lock()
s.stopping = true
s.started = false
idle := s.idle
if idle == nil {
idle = make(chan struct{})
@@ -110,6 +132,7 @@ func (s *TaskScheduler) Stop(ctx context.Context) error {
s.ctxMu.Lock()
cancel := s.cancel
s.ctxMu.Unlock()
s.lifecycleMu.Unlock()
if cancel != nil {
cancel()
}
@@ -142,6 +165,12 @@ func (s *TaskScheduler) removeLocked(id uint) {
}
func (s *TaskScheduler) Reload(ctx context.Context) error {
if ctx == nil {
ctx = context.Background()
}
if s.tasks == nil {
return errors.New("任务用例未初始化")
}
s.mu.Lock()
defer s.mu.Unlock()
@@ -149,12 +178,18 @@ func (s *TaskScheduler) Reload(ctx context.Context) error {
if err != nil {
return err
}
if err := ctx.Err(); err != nil {
return err
}
type preparedTask struct {
task *taskbiz.TimedTask
schedule cron.Schedule
}
prepared := make([]preparedTask, 0, len(items))
for _, task := range items {
if task == nil {
return errors.New("任务列表包含空任务")
}
if task.Enabled {
schedule, parseErr := parseTaskSchedule(task)
if parseErr != nil {
@@ -164,13 +199,27 @@ func (s *TaskScheduler) Reload(ctx context.Context) error {
}
}
s.runMu.Lock()
if err := ctx.Err(); err != nil {
s.runMu.Unlock()
return err
}
if s.stopping {
s.runMu.Unlock()
return errors.New("任务调度器正在停止")
}
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))
s.runMu.Unlock()
logger := s.logger
if logger == nil {
logger = slog.Default()
}
logger.InfoContext(ctx, "定时任务重载完成", "task_count", len(items))
return nil
}
@@ -216,9 +265,16 @@ func (s *TaskScheduler) finishRun(id uint) {
}
func (s *TaskScheduler) dispatch(task *taskbiz.TimedTask, trigger string, async bool) bool {
if task == nil {
return false
}
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)
logger := s.logger
if logger == nil {
logger = slog.Default()
}
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() {
@@ -234,25 +290,45 @@ func (s *TaskScheduler) dispatch(task *taskbiz.TimedTask, trigger string, async
}
func (s *TaskScheduler) run(ctx context.Context, task *taskbiz.TimedTask, trigger string) {
if ctx == nil {
ctx = context.Background()
}
logger := s.logger
if logger == nil {
logger = slog.Default()
}
var taskID uint
var taskName string
if task != nil {
taskID, taskName = task.ID, task.Name
}
if s.executor == nil {
logger.Error("timed task execution skipped because executor is unavailable", "task_id", taskID, "task_name", taskName, "trigger_type", trigger)
return
}
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...)
logger.Info("timed task finished", attributes...)
} else {
s.logger.Error("timed task finished", attributes...)
logger.Error("timed task finished", attributes...)
}
if log.Status != "success" {
if s.authorities == nil {
logger.Warn("timed task alert skipped because authority service is unavailable", attributes...)
return
}
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)
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)})
s.PublishToUsers(ids, map[string]any{"taskId": log.TaskID, "name": log.TaskName, "error": log.ErrorMsg, "time": time.Now().Format(time.RFC3339)})
}
}
@@ -270,6 +346,12 @@ func (s *TaskScheduler) Schedule(task *taskbiz.TimedTask) error {
}
s.mu.Lock()
defer s.mu.Unlock()
s.runMu.Lock()
stopping := s.stopping
s.runMu.Unlock()
if stopping {
return errors.New("任务调度器正在停止")
}
s.removeLocked(task.ID)
if !task.Enabled {
return nil
@@ -306,12 +388,21 @@ func (s *TaskScheduler) scheduleLocked(task *taskbiz.TimedTask, schedule cron.Sc
}
func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error {
if ctx == nil {
ctx = context.Background()
}
if s.tasks == nil {
return errors.New("任务用例未初始化")
}
s.mu.Lock()
defer s.mu.Unlock()
task, err := s.tasks.FindTask(ctx, id)
if err != nil {
return err
}
if task == nil {
return errors.New("任务不存在")
}
var schedule cron.Schedule
if task.Enabled {
schedule, err = parseTaskSchedule(task)
@@ -319,6 +410,12 @@ func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error {
return err
}
}
s.runMu.Lock()
stopping := s.stopping
s.runMu.Unlock()
if stopping {
return errors.New("任务调度器正在停止")
}
s.removeLocked(task.ID)
if task.Enabled {
s.scheduleLocked(cloneTimedTask(task), schedule)
@@ -327,10 +424,19 @@ func (s *TaskScheduler) ScheduleID(ctx context.Context, id uint) error {
}
func (s *TaskScheduler) TriggerID(ctx context.Context, id uint) error {
if ctx == nil {
ctx = context.Background()
}
if s.tasks == nil {
return errors.New("任务用例未初始化")
}
task, err := s.tasks.FindTask(ctx, id)
if err != nil {
return err
}
if task == nil {
return errors.New("任务不存在")
}
if !s.Trigger(task) {
return errors.New("任务正在执行或调度器正在停止")
}
internal/worker/task_scheduler_test.go
+18 -1
View File
@@ -1,6 +1,9 @@
package worker
import "testing"
import (
"context"
"testing"
)
func TestCloseSubscribersClosesEveryAlertStream(t *testing.T) {
scheduler := &TaskScheduler{subscribers: map[uint]map[chan []byte]struct{}{}}
@@ -48,3 +51,17 @@ func TestSubscribeCapsConnectionsPerUser(t *testing.T) {
}
scheduler.closeSubscribers()
}
func TestSubscribeAfterStopReturnsClosedStream(t *testing.T) {
scheduler := newTestTaskScheduler(&workerTaskRepo{})
if err := scheduler.Stop(context.Background()); err != nil {
t.Fatalf("Stop() error = %v", err)
}
ch := scheduler.Subscribe(9)
if _, open := <-ch; open {
t.Fatal("subscription remains open after scheduler shutdown")
}
if len(scheduler.subscribers) != 0 {
t.Fatalf("subscription was added after shutdown: %#v", scheduler.subscribers)
}
}