supervisor.go

package backstage

import (
	"context"
	"log/slog"
	"sync"
	"time"
)

const defaultDrainTimeout = 30 * time.Second

// Supervisor manages a set of named queues and scheduled tasks. Start it with
// [Supervisor.Serve], which blocks until the context is cancelled and then
// drains all in-flight workers before returning.
type Supervisor struct {
	name         string
	log          *slog.Logger
	drainTimeout time.Duration

	mu     sync.Mutex
	queues map[string]*queue
	tasks  []*scheduledTask
}

// New creates a new Supervisor. Options are applied in order; the last
// [WithLogger] call wins. If no logger is set, [slog.Default] is used.
func New(name string, opts ...Option) *Supervisor {
	s := &Supervisor{
		name:         name,
		drainTimeout: defaultDrainTimeout,
		queues:       make(map[string]*queue),
	}

	for _, opt := range opts {
		opt(s)
	}

	// Set default logger if WithLogger was not called.
	if s.log == nil {
		s.log = newLogger(name, nil)
	}

	return s
}

// RegisterQueue registers a named worker pool. Must be called before [Supervisor.Serve].
// When cfg.Store is nil, an in-memory buffered channel of depth cfg.Buffer is used.
func (s *Supervisor) RegisterQueue(name string, cfg QueueConfig) {
	if cfg.Workers < 1 {
		cfg.Workers = 1
	}

	store := cfg.Store
	if store == nil {
		buf := cfg.Buffer
		if buf < 1 {
			buf = 32
		}
		store = newMemStore(buf)
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	q := newQueue(name, cfg.Workers, store, s.log)
	s.queues[name] = q

	s.log.Info("queue registered",
		slog.String("queue", name),
		slog.Int("workers", cfg.Workers),
	)
}

// Dispatch enqueues a job onto the named queue without blocking.
// Returns [ErrQueueNotFound] if the queue has not been registered.
// Returns [ErrQueueFull] if the buffer is full and the job cannot be accepted.
func (s *Supervisor) Dispatch(queueName string, job Job) error {
	s.mu.Lock()
	q, ok := s.queues[queueName]
	s.mu.Unlock()

	if !ok {
		s.log.Error("dispatch to unknown queue",
			slog.String("queue", queueName),
			slog.String("job", job.Name),
		)
		return ErrQueueNotFound
	}

	if err := q.dispatch(job); err != nil {
		s.log.Warn("job dispatch failed",
			slog.String("queue", queueName),
			slog.String("job", job.Name),
			slog.Int("queue_depth", q.depth()),
			slog.Any("error", err),
		)
		return err
	}

	return nil
}

// Schedule registers a task that runs on the given schedule directly in its
// own goroutine (not dispatched to any queue).
func (s *Supervisor) Schedule(schedule Schedule, job Job) {
	s.addTask(&scheduledTask{job: job, schedule: schedule})
}

// ScheduleOnQueue registers a task that runs on the given schedule and is
// dispatched onto the named queue when it fires.
func (s *Supervisor) ScheduleOnQueue(queueName string, schedule Schedule, job Job) {
	s.addTask(&scheduledTask{job: job, schedule: schedule, queue: queueName})
}

func (s *Supervisor) addTask(t *scheduledTask) {
	s.mu.Lock()
	defer s.mu.Unlock()

	t.nextRun = t.schedule.Next(time.Now())
	s.tasks = append(s.tasks, t)

	s.log.Info("task scheduled",
		slog.String("task", t.job.Name),
		slog.String("queue", t.queue),
		slog.Time("next_at", t.nextRun),
	)
}

// Serve starts all registered queues and the scheduler loop, then blocks until
// ctx is cancelled. On cancellation it waits up to DrainTimeout for in-flight
// workers to finish before returning nil.
func (s *Supervisor) Serve(ctx context.Context) error {
	s.mu.Lock()
	queueNames := make([]string, 0, len(s.queues))
	for name := range s.queues {
		queueNames = append(queueNames, name)
	}
	taskNames := make([]string, 0, len(s.tasks))
	for _, t := range s.tasks {
		taskNames = append(taskNames, t.job.Name)
	}
	s.mu.Unlock()

	s.log.Info("supervisor starting",
		slog.Any("queues", queueNames),
		slog.Any("tasks", taskNames),
	)

	// workerCtx is cancelled on shutdown to signal all workers to stop.
	workerCtx, cancelWorkers := context.WithCancel(ctx)
	defer cancelWorkers()

	// wg tracks all running worker goroutines across all queues.
	var wg sync.WaitGroup

	s.mu.Lock()
	for _, q := range s.queues {
		q.start(workerCtx, &wg)
	}
	s.mu.Unlock()

	go s.schedulerLoop(workerCtx)

	// Block until the parent context is cancelled.
	<-ctx.Done()

	s.log.Info("supervisor shutting down",
		slog.String("reason", ctx.Err().Error()),
		slog.Duration("drain_timeout", s.drainTimeout),
	)

	cancelWorkers()

	// Wait for all workers to finish, bounded by DrainTimeout.
	drained := make(chan struct{})
	go func() {
		wg.Wait()
		close(drained)
	}()

	select {
	case <-drained:
		s.log.Info("supervisor stopped: all workers finished")
	case <-time.After(s.drainTimeout):
		s.log.Warn("supervisor stopped: drain timeout exceeded",
			slog.Duration("drain_timeout", s.drainTimeout),
		)
	}

	return nil
}

// --------------------------------------------------------------------------
// Scheduler loop
// --------------------------------------------------------------------------

func (s *Supervisor) schedulerLoop(ctx context.Context) {
	s.log.Debug("scheduler loop started")
	defer s.log.Debug("scheduler loop stopped")

	for {
		s.mu.Lock()
		next := s.earliestFire()
		s.mu.Unlock()

		var timer *time.Timer
		if next.IsZero() {
			// No tasks registered yet; check back in an hour.
			timer = time.NewTimer(time.Hour)
		} else {
			d := time.Until(next)
			if d < 0 {
				d = 0
			}
			timer = time.NewTimer(d)
			s.log.Debug("scheduler waiting",
				slog.Time("next_fire", next),
				slog.Duration("in", d),
			)
		}

		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case now := <-timer.C:
			s.fireDue(ctx, now)
		}
	}
}

// earliestFire returns the soonest nextRun across all tasks.
// Caller must hold s.mu.
func (s *Supervisor) earliestFire() time.Time {
	var earliest time.Time
	for _, t := range s.tasks {
		if earliest.IsZero() || t.nextRun.Before(earliest) {
			earliest = t.nextRun
		}
	}
	return earliest
}

// fireDue runs all tasks whose nextRun is at or before now, then reschedules them.
func (s *Supervisor) fireDue(ctx context.Context, now time.Time) {
	s.mu.Lock()
	var due []*scheduledTask
	for _, t := range s.tasks {
		if !t.nextRun.After(now) {
			due = append(due, t)
			t.nextRun = t.schedule.Next(now)
		}
	}
	s.mu.Unlock()

	for _, t := range due {
		t := t
		s.log.Debug("task fired",
			slog.String("task", t.job.Name),
			slog.String("queue", t.queue),
			slog.Time("fired_at", now),
			slog.Time("next_at", t.nextRun),
		)

		if t.queue != "" {
			if err := s.Dispatch(t.queue, t.job); err != nil {
				s.log.Warn("scheduled task dispatch failed",
					slog.String("task", t.job.Name),
					slog.String("queue", t.queue),
					slog.Any("error", err),
				)
			}
			continue
		}

		// Run directly in a goroutine — not bound to any queue.
		go s.runDirect(ctx, t.job)
	}
}

// runDirect executes a job in its own goroutine, retrying up to job.Retries
// times on failure. A panic aborts the job immediately without retrying.
func (s *Supervisor) runDirect(ctx context.Context, job Job) {
	log := s.log
	if job.ID != "" {
		log = s.log.With(slog.String("job_id", job.ID))
	}

	start := time.Now()
	log.Debug("direct task started", slog.String("task", job.Name))

	defer func() {
		if r := recover(); r != nil {
			log.Error("direct task panicked",
				slog.String("task", job.Name),
				slog.Any("panic", r),
			)
		}
	}()

	var err error
	for attempt := 0; attempt <= job.Retries; attempt++ {
		if attempt > 0 {
			if !sleepRetry(ctx, job.RetryDelay) {
				return
			}
			log.Debug("direct task retrying",
				slog.String("task", job.Name),
				slog.Int("attempt", attempt+1),
				slog.Int("of", job.Retries+1),
			)
		}
		err = job.Run(ctx)
		if err == nil {
			break
		}
	}

	dur := time.Since(start)

	if err != nil {
		log.Warn("direct task finished with error",
			slog.String("task", job.Name),
			slog.Duration("duration", dur),
			slog.Any("error", err),
		)
		return
	}

	log.Debug("direct task finished",
		slog.String("task", job.Name),
		slog.Duration("duration", dur),
	)
}