diff --git a/internal/jobs/module.go b/internal/jobs/module.go new file mode 100644 index 0000000..5fc7b65 --- /dev/null +++ b/internal/jobs/module.go @@ -0,0 +1,129 @@ +package jobs + +import ( + "context" + "log/slog" + "strconv" + "sync" + "time" + + "github.com/yan1h/agent-coding-workflow/internal/infra/clock" +) + +// RunnerFunc is the signature of a periodic runner's tick. +type RunnerFunc func(ctx context.Context) + +// ModuleDeps groups the externally-supplied collaborators. +type ModuleDeps struct { + Repo Repository + Handlers map[JobType]Handler + RunnerFns map[string]RunnerFunc // name → tick fn (e.g. "workspace_fetch") + Clock clock.Clock + Log *slog.Logger + WorkerBackoff []time.Duration + OnDeadLetter func(ctx context.Context, job Job, err error) +} + +// Module wires the scheduler, worker pool and runners. Owners construct it via +// NewModule + Start, then call Stop on shutdown. +type Module struct { + cfg Config + scheduler *Scheduler + workers []*Worker + log *slog.Logger + wg sync.WaitGroup +} + +// NewModule constructs but does not start the module. +func NewModule(cfg Config, deps ModuleDeps) *Module { + clk := deps.Clock + if clk == nil { + clk = clock.Real() + } + log := deps.Log + if log == nil { + log = slog.Default() + } + + sch := NewScheduler(log) + for name, fn := range deps.RunnerFns { + entry := scheduleEntry(cfg, name) + sch.Add(name, entry.interval, entry.enabled, fn) + } + + workers := make([]*Worker, cfg.Workers) + for i := 0; i < cfg.Workers; i++ { + workers[i] = NewWorker(deps.Repo, deps.Handlers, WorkerOptions{ + ID: workerID(i), + PollInterval: 200 * time.Millisecond, + Backoff: deps.WorkerBackoff, + Clock: clk, + Log: log, + OnDeadLetter: deps.OnDeadLetter, + }) + } + return &Module{cfg: cfg, scheduler: sch, workers: workers, log: log} +} + +// Start kicks off the scheduler and worker goroutines. They exit when ctx is +// canceled. Caller must call Stop to wait for clean shutdown. +func (m *Module) Start(ctx context.Context) { + if !m.cfg.Enabled { + m.log.Warn("jobs.module.disabled") + return + } + m.scheduler.Start(ctx) + for _, w := range m.workers { + m.wg.Add(1) + go func(w *Worker) { + defer m.wg.Done() + w.Run(ctx) + }(w) + } + m.log.Info("jobs.module.started", "workers", len(m.workers)) +} + +// Stop waits up to ShutdownGrace for goroutines to exit. +func (m *Module) Stop(ctx context.Context) { + done := make(chan struct{}) + go func() { + m.scheduler.Wait() + m.wg.Wait() + close(done) + }() + grace := m.cfg.ShutdownGrace + if grace <= 0 { + grace = 30 * time.Second + } + select { + case <-done: + m.log.Info("jobs.module.stopped") + case <-time.After(grace): + m.log.Warn("jobs.module.shutdown_timeout", "grace_ms", grace.Milliseconds()) + case <-ctx.Done(): + m.log.Warn("jobs.module.shutdown_ctx_canceled") + } +} + +func workerID(i int) string { return "worker#" + strconv.Itoa(i) } + +type schedEntry struct { + interval time.Duration + enabled bool +} + +func scheduleEntry(cfg Config, name string) schedEntry { + switch name { + case "workspace_fetch": + return schedEntry{cfg.WorkspaceFetch.Interval, cfg.WorkspaceFetch.Enabled} + case "worktree_prune": + return schedEntry{cfg.WorktreePrune.Interval, cfg.WorktreePrune.Enabled} + case "attachment_cleanup": + return schedEntry{cfg.AttachmentCleanup.Interval, cfg.AttachmentCleanup.Enabled} + case "job_reaper": + return schedEntry{cfg.JobReaper.Interval, cfg.JobReaper.Enabled} + case "job_purge": + return schedEntry{cfg.JobPurge.Interval, cfg.JobPurge.Enabled} + } + return schedEntry{0, false} +} diff --git a/internal/jobs/module_test.go b/internal/jobs/module_test.go new file mode 100644 index 0000000..d5b6b14 --- /dev/null +++ b/internal/jobs/module_test.go @@ -0,0 +1,97 @@ +package jobs_test + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/yan1h/agent-coding-workflow/internal/jobs" +) + +func TestModule_DisabledIsNoOp(t *testing.T) { + t.Parallel() + var called atomic.Int32 + m := jobs.NewModule( + jobs.Config{Enabled: false, Workers: 1, ShutdownGrace: time.Second, + WorkspaceFetch: jobs.RunnerConfig{Enabled: true, Interval: 10 * time.Millisecond}}, + jobs.ModuleDeps{ + Repo: nil, + Handlers: map[jobs.JobType]jobs.Handler{}, + RunnerFns: map[string]jobs.RunnerFunc{ + "workspace_fetch": func(_ context.Context) { called.Add(1) }, + }, + }) + require.NotNil(t, m) + ctx, cancel := context.WithCancel(context.Background()) + m.Start(ctx) + time.Sleep(50 * time.Millisecond) + cancel() + m.Stop(context.Background()) + assert.Zero(t, called.Load(), "disabled module must not run any runner") +} + +func TestModule_EnabledRunsRegisteredRunner(t *testing.T) { + t.Parallel() + var ticks atomic.Int32 + m := jobs.NewModule( + jobs.Config{Enabled: true, Workers: 0, ShutdownGrace: time.Second, + WorkspaceFetch: jobs.RunnerConfig{Enabled: true, Interval: 10 * time.Millisecond}}, + jobs.ModuleDeps{ + Repo: nil, + Handlers: map[jobs.JobType]jobs.Handler{}, + RunnerFns: map[string]jobs.RunnerFunc{ + "workspace_fetch": func(_ context.Context) { ticks.Add(1) }, + }, + }) + ctx, cancel := context.WithCancel(context.Background()) + m.Start(ctx) + require.Eventually(t, func() bool { return ticks.Load() >= 1 }, 200*time.Millisecond, 5*time.Millisecond) + cancel() + m.Stop(context.Background()) + assert.GreaterOrEqual(t, ticks.Load(), int32(1)) +} + +func TestModule_UnknownRunnerNameDisabled(t *testing.T) { + t.Parallel() + var called atomic.Int32 + m := jobs.NewModule( + jobs.Config{Enabled: true, Workers: 0, ShutdownGrace: time.Second}, + jobs.ModuleDeps{ + Repo: nil, + Handlers: map[jobs.JobType]jobs.Handler{}, + RunnerFns: map[string]jobs.RunnerFunc{ + "made_up_runner": func(_ context.Context) { called.Add(1) }, + }, + }) + ctx, cancel := context.WithCancel(context.Background()) + m.Start(ctx) + time.Sleep(50 * time.Millisecond) + cancel() + m.Stop(context.Background()) + assert.Zero(t, called.Load(), "unknown runner name should default to disabled") +} + +func TestModule_StopRespectsShutdownGrace(t *testing.T) { + t.Parallel() + m := jobs.NewModule( + jobs.Config{Enabled: true, Workers: 0, ShutdownGrace: 50 * time.Millisecond, + WorkspaceFetch: jobs.RunnerConfig{Enabled: true, Interval: 10 * time.Millisecond}}, + jobs.ModuleDeps{ + Repo: nil, + Handlers: map[jobs.JobType]jobs.Handler{}, + RunnerFns: map[string]jobs.RunnerFunc{ + "workspace_fetch": func(_ context.Context) {}, + }, + }) + ctx, cancel := context.WithCancel(context.Background()) + m.Start(ctx) + cancel() + start := time.Now() + m.Stop(context.Background()) + elapsed := time.Since(start) + assert.Less(t, elapsed, 500*time.Millisecond, "Stop should return promptly") +} diff --git a/internal/jobs/worker.go b/internal/jobs/worker.go index 6217b20..45d2a8d 100644 --- a/internal/jobs/worker.go +++ b/internal/jobs/worker.go @@ -2,6 +2,7 @@ package jobs import ( "context" + "errors" "fmt" "log/slog" "sync" @@ -19,6 +20,12 @@ type WorkerOptions struct { Backoff []time.Duration Clock clock.Clock Log *slog.Logger + // OnDeadLetter, if non-nil, is invoked once after a job is moved to dead-letter + // (either because IsPermanent(err) or attempts exhausted). The hook runs in + // the worker goroutine; panics are recovered and logged but do not stop the + // worker. Use this for fallback admin notifications when external delivery + // fails permanently. + OnDeadLetter func(ctx context.Context, job Job, err error) } // Worker pulls leased jobs from the repository, dispatches them to handlers, @@ -88,6 +95,7 @@ func (w *Worker) execute(ctx context.Context, job Job) { w.opts.Log.Error("job.dead_letter_write_failed", "id", job.ID, "err", derr.Error()) } w.opts.Log.Error("job.dead_letter", "id", job.ID, "type", job.Type, "reason", "unknown_type") + w.fireOnDeadLetter(ctx, job, errors.New("unknown job type: "+string(job.Type))) return } @@ -110,6 +118,7 @@ func (w *Worker) execute(ctx context.Context, job Job) { } w.opts.Log.Error("job.dead_letter", "id", job.ID, "type", job.Type, "attempts", job.Attempts, "err", msg) + w.fireOnDeadLetter(ctx, job, err) return } @@ -148,6 +157,19 @@ func backoffFor(seq []time.Duration, attempt int) time.Duration { return seq[idx] } +// fireOnDeadLetter invokes opts.OnDeadLetter, if set, recovering any panics. +func (w *Worker) fireOnDeadLetter(ctx context.Context, job Job, err error) { + if w.opts.OnDeadLetter == nil { + return + } + defer func() { + if r := recover(); r != nil { + w.opts.Log.Error("job.dead_letter_hook_panic", "id", job.ID, "panic", r) + } + }() + w.opts.OnDeadLetter(ctx, job, err) +} + func truncate(s string, n int) string { if len(s) <= n { return s diff --git a/internal/jobs/worker_test.go b/internal/jobs/worker_test.go index de19863..0aa1a94 100644 --- a/internal/jobs/worker_test.go +++ b/internal/jobs/worker_test.go @@ -10,6 +10,7 @@ import ( "testing" "time" + "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -166,3 +167,33 @@ func TestWorker_UnknownTypeDeadLetters(t *testing.T) { cancel() w.Wait() } + +func TestWorker_DeadLetterHookCalled(t *testing.T) { + t.Parallel() + ctx := context.Background() + repo, _ := setupRepo(t) + enqueued, err := repo.Enqueue(ctx, jobs.TypeWebhookDeliver, json.RawMessage(`{}`), time.Now(), 5) + require.NoError(t, err) + + h := &recordingHandler{typ: jobs.TypeWebhookDeliver, wantErr: jobs.Permanent(errors.New("4xx"))} + var called atomic.Int32 + var seenID atomic.Value // uuid.UUID + w := jobs.NewWorker(repo, map[jobs.JobType]jobs.Handler{jobs.TypeWebhookDeliver: h}, + jobs.WorkerOptions{ + ID: "t", PollInterval: 5 * time.Millisecond, Backoff: []time.Duration{time.Second}, + Clock: clock.Real(), Log: discardLogger(), + OnDeadLetter: func(_ context.Context, j jobs.Job, _ error) { + seenID.Store(j.ID) + called.Add(1) + }, + }) + cctx, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + defer cancel() + go w.Run(cctx) + require.Eventually(t, func() bool { return called.Load() >= 1 }, 500*time.Millisecond, 10*time.Millisecond) + cancel() + w.Wait() + if id, ok := seenID.Load().(uuid.UUID); ok { + assert.Equal(t, enqueued.ID, id) + } +}