kra-new/internal/worker/task_runtime_test.go

package worker

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	taskbiz "kra/internal/biz/task"
)

type workerTaskRepo struct {
	mu      sync.Mutex
	items   []*taskbiz.TimedTask
	listErr error
	logs    chan *taskbiz.TimedTaskLog
}

func (r *workerTaskRepo) CreateTask(context.Context, *taskbiz.TimedTask) error { return nil }
func (r *workerTaskRepo) UpdateTask(context.Context, *taskbiz.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) (*taskbiz.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, *taskbiz.TimedTask) ([]*taskbiz.TimedTask, int64, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.listErr != nil {
		return nil, 0, r.listErr
	}
	items := make([]*taskbiz.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 *taskbiz.TimedTaskLog) error {
	if r.logs != nil {
		copy := *value
		r.logs <- &copy
	}
	return nil
}
func (r *workerTaskRepo) ListTaskLogs(context.Context, int, int, uint, string) ([]*taskbiz.TimedTaskLog, int64, error) {
	return nil, 0, nil
}

func newTestTaskScheduler(repo *workerTaskRepo) *TaskScheduler {
	tasks := taskbiz.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 := &taskbiz.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 = []*taskbiz.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(&taskbiz.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
	taskbiz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
		calls.Add(1)
		started <- struct{}{}
		<-release
		return nil
	})

	task := &taskbiz.TimedTask{ID: 9, Name: "single", ExecutorType: taskbiz.TaskExecutorMethod, MethodName: methodName}
	repo := &workerTaskRepo{items: []*taskbiz.TimedTask{task}, logs: make(chan *taskbiz.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
	taskbiz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
		calls.Add(1)
		started <- struct{}{}
		<-release
		return nil
	})

	task := &taskbiz.TimedTask{ID: 10, Name: "shared-gate", ExecutorType: taskbiz.TaskExecutorMethod, MethodName: methodName}
	repo := &workerTaskRepo{logs: make(chan *taskbiz.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()")
	}
}