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

优化结构

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
cmd/wire_gen.go generated
View File

@ -113,14 +113,14 @@ func wireApp(confServer *conf.Server, runtime *conf.Runtime, logger *slog.Logger
taskRepo := task.NewTaskRepo(dataData)
registry := app.TaskRegistry(catalog)
taskUsecase := task2.NewTaskUsecaseWithRegistry(taskRepo, registry)
mediaRepo := system.NewMediaRepo(dataData)
mediaUsecase := system2.NewMediaUsecase(mediaRepo, reloadable, runtimeSettings)
taskExecutor := worker.NewTaskExecutorWithRegistry(taskUsecase, mediaUsecase, runtime, registry)
taskExecutor := worker.NewTaskExecutorWithRegistry(taskUsecase, registry)
taskScheduler := worker.NewTaskScheduler(taskUsecase, authorityUsecase, taskExecutor, logger)
taskRuntime := worker.NewTaskRuntime(taskScheduler)
taskApplicationUsecase := task2.NewTaskApplicationUsecase(taskUsecase, taskRuntime)
v13 := task3.NewTaskService(taskApplicationUsecase)
handlerTask := handler.NewTask(v13)
mediaRepo := system.NewMediaRepo(dataData)
mediaUsecase := system2.NewMediaUsecase(mediaRepo, reloadable, runtimeSettings)
v14 := system3.NewMediaService(mediaUsecase, runtimeSettings)
media := handler.NewMedia(v14)
auditQueryRepo := system.NewAuditRepo(dataData)

85
internal/data/config_watch.go
View File

@ -33,43 +33,77 @@ func (d *Data) watchConfig() func() {
return func() {}
}
done := make(chan struct{})
finished := make(chan struct{})
var once sync.Once
go func() {
defer close(finished)
var timer *time.Timer
var timerC <-chan time.Time
events := watcher.Events
watchErrors := watcher.Errors
for {
select {
case event, ok := <-watcher.Events:
if !ok || filepath.Clean(event.Name) != absolute || event.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Rename) == 0 {
case event, ok := <-events:
if !ok {
events = nil
continue
}
if filepath.Clean(event.Name) != absolute || event.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Rename) == 0 {
continue
}
if timer != nil {
timer.Stop()
if !timer.Stop() {
select {
case <-timer.C:
default:
}
}
}
timer = time.AfterFunc(100*time.Millisecond, func() {
next, loadErr := readBootstrap(absolute)
if loadErr != nil {
logger.Error("reload changed config", "mod", "system", "error", loadErr)
return
}
if next.Data == nil || next.Admin == nil {
logger.Error("reload changed config: data and admin configuration are required", "mod", "system")
return
}
if current := d.runtime.Admin(); current != nil {
next.Admin.Storage = current.Storage
next.Admin.Email = current.Email
}
next.Admin.ConfigPath = absolute
d.runtime.Replace(next.Data, next.Admin)
logger.Info("config file changed", "mod", "system", "path", absolute)
})
case watchErr, ok := <-watcher.Errors:
if ok {
logger.Error("config watcher error", "mod", "system", "error", watchErr)
if timer == nil {
timer = time.NewTimer(100 * time.Millisecond)
} else {
timer.Reset(100 * time.Millisecond)
}
timerC = timer.C
case <-timerC:
// The debounce callback runs on this goroutine so shutdown can
// wait for it to finish without racing runtime replacement.
timerC = nil
select {
case <-done:
continue
default:
}
next, loadErr := readBootstrap(absolute)
if loadErr != nil {
logger.Error("reload changed config", "mod", "system", "error", loadErr)
continue
}
if next.Data == nil || next.Admin == nil {
logger.Error("reload changed config: data and admin configuration are required", "mod", "system")
continue
}
if current := d.runtime.Admin(); current != nil {
next.Admin.Storage = current.Storage
next.Admin.Email = current.Email
}
next.Admin.ConfigPath = absolute
d.runtime.Replace(next.Data, next.Admin)
logger.Info("config file changed", "mod", "system", "path", absolute)
case watchErr, ok := <-watchErrors:
if !ok {
watchErrors = nil
continue
}
logger.Error("config watcher error", "mod", "system", "error", watchErr)
case <-done:
if timer != nil {
timer.Stop()
if !timer.Stop() {
select {
case <-timer.C:
default:
}
}
}
return
}
@ -79,6 +113,7 @@ func (d *Data) watchConfig() func() {
once.Do(func() {
close(done)
_ = watcher.Close()
<-finished
})
}
}

59
internal/data/config_watch_test.go Normal file
View File

@ -0,0 +1,59 @@
package data
import (
"os"
"path/filepath"
"sync"
"testing"
"time"
"kra/internal/conf"
)
func TestConfigWatcherStopWaitsForDebouncedReload(t *testing.T) {
root := t.TempDir()
configPath := filepath.Join(root, "config.yaml")
config := []byte("data:\n database:\n driver: sqlite\n path: " + filepath.ToSlash(root) + "\nadmin: {}\n")
if err := os.WriteFile(configPath, config, 0o600); err != nil {
t.Fatal(err)
}
runtime := conf.NewRuntime(&conf.Data{}, &conf.AdminBackend{ConfigPath: configPath})
started := make(chan struct{})
release := make(chan struct{})
var startOnce sync.Once
stopListener := runtime.Subscribe(func(*conf.Data, *conf.AdminBackend) {
startOnce.Do(func() { close(started) })
<-release
})
defer stopListener()
data := &Data{runtime: runtime}
stop := data.watchConfig()
// A write event schedules the 100ms debounce reload.
if err := os.WriteFile(configPath, config, 0o600); err != nil {
t.Fatal(err)
}
select {
case <-started:
case <-time.After(3 * time.Second):
stop()
t.Fatal("debounced reload did not start")
}
stopped := make(chan struct{})
go func() {
stop()
close(stopped)
}()
select {
case <-stopped:
t.Fatal("watcher stop returned before reload callback finished")
case <-time.After(100 * time.Millisecond):
}
close(release)
select {
case <-stopped:
case <-time.After(time.Second):
t.Fatal("watcher stop did not wait for callback completion")
}
}

16
internal/initialize/compatibility.go
View File

@ -274,12 +274,22 @@ func parseCompatibleDuration(raw string) (time.Duration, error) {
return value, nil
}
if index := strings.Index(raw, "d"); index >= 0 {
days, _ := strconv.Atoi(raw[:index])
value := time.Duration(days) * 24 * time.Hour
remainder, err := time.ParseDuration(raw[index+1:])
if index == 0 {
return 0, strconv.ErrSyntax
}
days, err := strconv.Atoi(raw[:index])
if err != nil {
return 0, err
}
value := time.Duration(days) * 24 * time.Hour
remainderText := strings.TrimSpace(raw[index+1:])
if remainderText == "" {
return value, nil
}
remainder, err := time.ParseDuration(remainderText)
if err != nil {
return 0, err
}
return value + remainder, nil
}
nanoseconds, err := strconv.ParseInt(raw, 10, 64)

38
internal/initialize/compatibility_test.go Normal file
View File

@ -0,0 +1,38 @@
package initialize
import (
"testing"
"time"
)
func TestParseCompatibleDuration(t *testing.T) {
tests := []struct {
name string
raw string
want time.Duration
wantErr bool
}{
{name: "days only", raw: "3d", want: 72 * time.Hour},
{name: "days and duration", raw: "3d4h", want: 76 * time.Hour},
{name: "missing day count", raw: "xd", wantErr: true},
{name: "invalid remainder", raw: "3dgarbage", wantErr: true},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
got, err := parseCompatibleDuration(test.raw)
if test.wantErr {
if err == nil {
t.Fatalf("parseCompatibleDuration(%q) accepted malformed input", test.raw)
}
return
}
if err != nil {
t.Fatalf("parseCompatibleDuration(%q) error = %v", test.raw, err)
}
if got != test.want {
t.Fatalf("parseCompatibleDuration(%q) = %s, want %s", test.raw, got, test.want)
}
})
}
}

13
internal/integration/storage/qiniu_storage.go
View File

@ -22,7 +22,6 @@ type qiniuStorage struct {
upload *qstorage.FormUploader
manager *qstorage.BucketManager
mac *qbox.Mac
token string
}
func newQiniuStorage(config *conf.AdminBackend_Qiniu) (system.FileStorage, error) {
@ -43,10 +42,14 @@ func newQiniuStorage(config *conf.AdminBackend_Qiniu) (system.FileStorage, error
cfg.Zone = &qstorage.ZoneXinjiapo
}
mac := qbox.NewMac(config.AccessKey, config.SecretKey)
policy := qstorage.PutPolicy{Scope: config.Bucket}
token := policy.UploadToken(mac)
return &qiniuStorage{config: config, upload: qstorage.NewFormUploader(&cfg), manager: qstorage.NewBucketManager(mac, &cfg), mac: mac, token: token}, nil
return &qiniuStorage{config: config, upload: qstorage.NewFormUploader(&cfg), manager: qstorage.NewBucketManager(mac, &cfg), mac: mac}, nil
}
func (s *qiniuStorage) uploadToken() string {
policy := qstorage.PutPolicy{Scope: s.config.Bucket, Expires: 3600}
return policy.UploadToken(s.mac)
}
func (s *qiniuStorage) file(key string, size int64) *system.StoredFile {
return &system.StoredFile{Name: path.Base(key), Path: key, URL: strings.TrimSuffix(s.config.BaseUrl, "/") + "/" + key, Size: size}
}
@ -68,7 +71,7 @@ func (s *qiniuStorage) Put(ctx context.Context, name string, reader io.Reader) (
}
defer temporary.Close()
ret := qstorage.PutRet{}
if err = s.upload.Put(ctx, &ret, s.token, strings.TrimPrefix(name, "/"), temporary, size, &qstorage.PutExtra{}); err != nil {
if err = s.upload.Put(ctx, &ret, s.uploadToken(), strings.TrimPrefix(name, "/"), temporary, size, &qstorage.PutExtra{}); err != nil {
return nil, err
}
return s.file(ret.Key, size), nil

65
internal/integration/storage/reloadable.go
View File

@ -2,6 +2,7 @@ package storage
import (
"context"
"errors"
"io"
"kra/internal/biz/system"
"sync"
@ -31,38 +32,66 @@ func NewReloadable(config *conf.AdminBackend) (*Reloadable, error) {
}
func (s *Reloadable) Replace(current system.FileStorage) {
if s == nil {
return
}
s.mu.Lock()
s.current = current
s.mu.Unlock()
}
func (s *Reloadable) Put(ctx context.Context, name string, reader io.Reader) (*system.StoredFile, error) {
func (s *Reloadable) active() (system.FileStorage, error) {
if s == nil {
return nil, errors.New("file storage is unavailable")
}
s.mu.RLock()
defer s.mu.RUnlock()
return s.current.Put(ctx, name, reader)
current := s.current
s.mu.RUnlock()
if current == nil {
return nil, errors.New("file storage is unavailable")
}
return current, nil
}
func (s *Reloadable) Put(ctx context.Context, name string, reader io.Reader) (*system.StoredFile, error) {
current, err := s.active()
if err != nil {
return nil, err
}
return current.Put(ctx, name, reader)
}
func (s *Reloadable) Open(ctx context.Context, name string) (io.ReadCloser, error) {
s.mu.RLock()
defer s.mu.RUnlock()
return s.current.Open(ctx, name)
current, err := s.active()
if err != nil {
return nil, err
}
return current.Open(ctx, name)
}
func (s *Reloadable) Delete(ctx context.Context, name string) error {
s.mu.RLock()
defer s.mu.RUnlock()
return s.current.Delete(ctx, name)
current, err := s.active()
if err != nil {
return err
}
return current.Delete(ctx, name)
}
func (s *Reloadable) Compose(ctx context.Context, names []string, destination string) (*system.StoredFile, string, error) {
s.mu.RLock()
defer s.mu.RUnlock()
return s.current.Compose(ctx, names, destination)
current, err := s.active()
if err != nil {
return nil, "", err
}
return current.Compose(ctx, names, destination)
}
func (s *Reloadable) DeletePrefix(ctx context.Context, prefix string) error {
s.mu.RLock()
defer s.mu.RUnlock()
return s.current.DeletePrefix(ctx, prefix)
current, err := s.active()
if err != nil {
return err
}
return current.DeletePrefix(ctx, prefix)
}
func (s *Reloadable) List(ctx context.Context, prefix, cursor string, limit int) ([]*system.StoredFile, string, bool, error) {
s.mu.RLock()
defer s.mu.RUnlock()
return s.current.List(ctx, prefix, cursor, limit)
current, err := s.active()
if err != nil {
return nil, "", false, err
}
return current.List(ctx, prefix, cursor, limit)
}

96
internal/integration/storage/reloadable_test.go Normal file
View File

@ -0,0 +1,96 @@
package storage
import (
"context"
"errors"
"io"
"testing"
"time"
"kra/internal/biz/system"
)
type blockingStorage struct {
started chan struct{}
release chan struct{}
}
func (s *blockingStorage) Put(context.Context, string, io.Reader) (*system.StoredFile, error) {
close(s.started)
<-s.release
return &system.StoredFile{Name: "uploaded"}, nil
}
func (*blockingStorage) Open(context.Context, string) (io.ReadCloser, error) {
return nil, errors.New("not implemented")
}
func (*blockingStorage) Delete(context.Context, string) error { return nil }
func (*blockingStorage) Compose(context.Context, []string, string) (*system.StoredFile, string, error) {
return nil, "", errors.New("not implemented")
}
func (*blockingStorage) DeletePrefix(context.Context, string) error { return nil }
func (*blockingStorage) List(context.Context, string, string, int) ([]*system.StoredFile, string, bool, error) {
return nil, "", false, errors.New("not implemented")
}
type emptyStorage struct{}
func (*emptyStorage) Put(context.Context, string, io.Reader) (*system.StoredFile, error) {
return &system.StoredFile{Name: "replacement"}, nil
}
func (*emptyStorage) Open(context.Context, string) (io.ReadCloser, error) {
return nil, errors.New("not implemented")
}
func (*emptyStorage) Delete(context.Context, string) error { return nil }
func (*emptyStorage) Compose(context.Context, []string, string) (*system.StoredFile, string, error) {
return nil, "", errors.New("not implemented")
}
func (*emptyStorage) DeletePrefix(context.Context, string) error { return nil }
func (*emptyStorage) List(context.Context, string, string, int) ([]*system.StoredFile, string, bool, error) {
return nil, "", false, errors.New("not implemented")
}
func TestReloadableReplaceDoesNotWaitForInFlightOperation(t *testing.T) {
old := &blockingStorage{started: make(chan struct{}), release: make(chan struct{})}
reloadable := &Reloadable{current: old}
putDone := make(chan error, 1)
go func() {
_, err := reloadable.Put(context.Background(), "file", nil)
putDone <- err
}()
select {
case <-old.started:
case <-time.After(time.Second):
t.Fatal("storage operation did not start")
}
replaced := make(chan struct{})
go func() {
reloadable.Replace(&emptyStorage{})
close(replaced)
}()
select {
case <-replaced:
case <-time.After(200 * time.Millisecond):
t.Fatal("Replace waited for an in-flight storage operation")
}
close(old.release)
select {
case err := <-putDone:
if err != nil {
t.Fatalf("Put() error = %v", err)
}
case <-time.After(time.Second):
t.Fatal("storage operation did not finish")
}
}
func TestReloadableNilStorageReturnsError(t *testing.T) {
var reloadable *Reloadable
if _, err := reloadable.Put(context.Background(), "file", nil); err == nil {
t.Fatal("nil Reloadable.Put() succeeded")
}
if _, err := (&Reloadable{}).Open(context.Background(), "file"); err == nil {
t.Fatal("empty Reloadable.Open() succeeded")
}
}

66
internal/worker/task_executor.go
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:

11
internal/worker/task_executor_test.go
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)
}
}

18
internal/worker/task_registry.go
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)
}

25
internal/worker/task_registry_test.go Normal file
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)
}
}
}

22
internal/worker/task_runtime_test.go
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)
}
}

134
internal/worker/task_scheduler.go
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("任务正在执行或调度器正在停止")
}

19
internal/worker/task_scheduler_test.go
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)
}
}