agentic-coding-workflow/internal/jobs/module.go

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, known := scheduleEntry(cfg, name)
		if !known {
			log.Warn("jobs.module.unknown_runner", "name", name)
		}
		sch.Add(name, entry.interval, entry.enabled, fn)
	}

	handlerTimeout := handlerTimeoutFor(cfg.JobReaper)
	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,
			HandlerTimeout: handlerTimeout,
			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
	}
	timer := time.NewTimer(grace)
	defer timer.Stop()
	select {
	case <-done:
		m.log.Info("jobs.module.stopped")
	case <-timer.C:
		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) }

// handlerTimeoutFor derives the per-job handler timeout from the reaper's lease
// timeout, leaving a margin below it so a healthy long handler finishes (or is
// canceled) before the reaper re-leases the job and double-runs it. Returns 0
// (no timeout) when reaping is disabled or no positive lease timeout is set.
func handlerTimeoutFor(rc JobReaperConfig) time.Duration {
	if !rc.Enabled || rc.LeaseTimeout <= 0 {
		return 0
	}
	// 90% of the lease window leaves headroom for the write-back to land before
	// the reaper's cutoff.
	return rc.LeaseTimeout - rc.LeaseTimeout/10
}

type schedEntry struct {
	interval time.Duration
	enabled  bool
}

func scheduleEntry(cfg Config, name string) (schedEntry, bool) {
	switch name {
	case "workspace_fetch":
		return schedEntry{cfg.WorkspaceFetch.Interval, cfg.WorkspaceFetch.Enabled}, true
	case "worktree_prune":
		return schedEntry{cfg.WorktreePrune.Interval, cfg.WorktreePrune.Enabled}, true
	case "attachment_cleanup":
		return schedEntry{cfg.AttachmentCleanup.Interval, cfg.AttachmentCleanup.Enabled}, true
	case "job_reaper":
		return schedEntry{cfg.JobReaper.Interval, cfg.JobReaper.Enabled}, true
	case "job_purge":
		return schedEntry{cfg.JobPurge.Interval, cfg.JobPurge.Enabled}, true
	case "acp_events_purge":
		return schedEntry{cfg.AcpEventsPurge.Interval, cfg.AcpEventsPurge.Enabled}, true
	case "mcp_tokens_purge":
		return schedEntry{cfg.MCPTokensPurge.Interval, cfg.MCPTokensPurge.Enabled}, true
	}
	return schedEntry{0, false}, false
}