优化结构
This commit is contained in:
Yvan
@@ -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()")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user