worker.go

package pgqueue

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	sq "github.com/Masterminds/squirrel"
)

// Worker errors.
var (
	ErrWorkerAlreadyStarted = errors.New("pgqueue: worker already started")
	ErrWorkerStopped        = errors.New("pgqueue: worker already stopped (single-use)")
	ErrWorkerNotRunning     = errors.New("pgqueue: worker not running")
)

// Spec defaults.
const (
	defaultPollInterval   = 5 * time.Second
	defaultMaxRetries     = int16(3)
	defaultRetryDelay     = 30 * time.Second
	defaultLeaseDuration  = 5 * time.Minute
	defaultFinalizeBuffer = 10 * time.Second
	defaultReapInterval   = 1 * time.Minute
	defaultMonitorInterval = 5 * time.Minute
)

type workerState uint32

const (
	workerStateIdle     workerState = 0
	workerStateStarting workerState = 1
	workerStateRunning  workerState = 2
	workerStateStopping workerState = 3
	workerStateStopped  workerState = 4
)

// runnerEntry is a type-erased handler entry created by Register/RegisterTicker.
type runnerEntry struct {
	queue          string
	pollInterval   time.Duration
	retryDelay     time.Duration
	maxRetries     int16
	leaseDuration  time.Duration
	finalizeBuffer time.Duration
	isTicker       bool
	tickerEvery    time.Duration
	ensureTicker   func(ctx context.Context) error
	run            func(ctx context.Context, task *Task)
}

// Worker polls for tasks and dispatches them to registered handlers.
type Worker struct {
	queue *Queue
	log   *slog.Logger

	mu      sync.Mutex
	state   atomic.Uint32
	entries []runnerEntry
	queues  map[string]struct{}

	reapInterval time.Duration

	workerCtx    context.Context
	workerCancel context.CancelFunc
	pollCtx      context.Context
	pollCancel   context.CancelFunc
	wg           sync.WaitGroup
	done         chan struct{}
}

// NewWorker creates a Worker bound to a Queue.
func NewWorker(queue *Queue, opts ...WorkerOption) *Worker {
	w := &Worker{
		queue:        queue,
		log:          queue.log,
		queues:       make(map[string]struct{}),
		reapInterval: defaultReapInterval,
	}
	for _, o := range opts {
		o(w)
	}
	return w
}

// WorkerOption configures a Worker.
type WorkerOption func(*Worker)

// WithWorkerLogger overrides the logger inherited from Queue.
func WithWorkerLogger(l *slog.Logger) WorkerOption {
	return func(w *Worker) { w.log = l }
}

// WithReapInterval overrides the default reaper interval (1m).
// Panics if d is not positive.
func WithReapInterval(d time.Duration) WorkerOption {
	if d <= 0 {
		panic("pgqueue: WithReapInterval requires a positive duration")
	}
	return func(w *Worker) { w.reapInterval = d }
}

// IsRunning returns true if the worker is in the running state.
func (w *Worker) IsRunning() bool {
	return workerState(w.state.Load()) == workerStateRunning
}

// Register registers a one-shot job handler. Must be called before Start.
func Register[P any](w *Worker, spec JobSpec[P], handler JobHandler[P]) error {
	spec = applyJobDefaults(spec)
	if err := validateJobSpec(spec); err != nil {
		return err
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	if s := workerState(w.state.Load()); s != workerStateIdle {
		return fmt.Errorf("pgqueue: cannot register after worker has started (state=%d)", s)
	}
	if _, exists := w.queues[spec.Queue]; exists {
		return fmt.Errorf("pgqueue: queue %q already registered", spec.Queue)
	}
	w.queues[spec.Queue] = struct{}{}

	w.entries = append(w.entries, runnerEntry{
		queue:          spec.Queue,
		pollInterval:   spec.PollInterval,
		retryDelay:     spec.RetryDelay,
		maxRetries:     spec.MaxRetries,
		leaseDuration:  spec.LeaseDuration,
		finalizeBuffer: spec.FinalizeBuffer,
		isTicker:       false,
		run:            makeJobRunner(w, spec, handler),
	})
	return nil
}

// RegisterTicker registers a recurring ticker handler. Must be called before Start.
func RegisterTicker[P any](w *Worker, spec TickerSpec[P], handler TickerHandler[P]) error {
	spec = applyTickerDefaults(spec)
	if err := validateTickerSpec(spec); err != nil {
		return err
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	if s := workerState(w.state.Load()); s != workerStateIdle {
		return fmt.Errorf("pgqueue: cannot register after worker has started (state=%d)", s)
	}
	if _, exists := w.queues[spec.Queue]; exists {
		return fmt.Errorf("pgqueue: queue %q already registered", spec.Queue)
	}
	w.queues[spec.Queue] = struct{}{}

	w.entries = append(w.entries, runnerEntry{
		queue:          spec.Queue,
		pollInterval:   spec.PollInterval,
		retryDelay:     spec.RetryDelay,
		maxRetries:     spec.MaxRetries,
		leaseDuration:  spec.LeaseDuration,
		finalizeBuffer: spec.FinalizeBuffer,
		isTicker:       true,
		tickerEvery:    spec.Every,
		ensureTicker:   makeTickerEnsurer(w, spec),
		run:            makeTickerRunner(w, spec, handler),
	})
	return nil
}

// Start runs preflight checks, then starts poll loops, reaper, and quarantine
// monitor. Blocks until all goroutines exit. Workers are single-use.
func (w *Worker) Start(ctx context.Context) error {
	w.mu.Lock()
	if s := workerState(w.state.Load()); s != workerStateIdle {
		w.mu.Unlock()
		if s == workerStateStopped {
			return ErrWorkerStopped
		}
		return ErrWorkerAlreadyStarted
	}
	w.state.Store(uint32(workerStateStarting))
	w.mu.Unlock()

	// Preflight — all fallible work before flipping to running.
	preflightCtx, preflightCancel := context.WithCancel(ctx)
	defer preflightCancel()

	for _, entry := range w.entries {
		if entry.ensureTicker != nil {
			if err := entry.ensureTicker(preflightCtx); err != nil {
				w.state.Store(uint32(workerStateIdle))
				return fmt.Errorf("pgqueue: ensure ticker %q: %w", entry.queue, err)
			}
		}
	}

	w.warnOrphanedQueues(preflightCtx)

	// Check if parent context was cancelled during preflight.
	if err := preflightCtx.Err(); err != nil {
		w.state.Store(uint32(workerStateIdle))
		return err
	}

	// Preflight passed — transition to running.
	w.mu.Lock()
	w.state.Store(uint32(workerStateRunning))
	w.workerCtx, w.workerCancel = context.WithCancel(ctx)
	w.pollCtx, w.pollCancel = context.WithCancel(w.workerCtx)
	w.done = make(chan struct{})
	w.mu.Unlock()

	for _, entry := range w.entries {
		w.wg.Add(1)
		go func(r runnerEntry) {
			defer w.wg.Done()
			w.runLoop(r)
		}(entry)
	}

	w.wg.Add(1)
	go func() {
		defer w.wg.Done()
		w.reapLoop()
	}()

	w.wg.Add(1)
	go func() {
		defer w.wg.Done()
		w.monitorQuarantineLoop()
	}()

	w.wg.Wait()
	w.state.Store(uint32(workerStateStopped))
	close(w.done)
	return nil
}

// Stop cancels polling and waits for in-flight work to drain.
// Returns nil on clean drain, context.DeadlineExceeded on timeout.
func (w *Worker) Stop(drainCtx context.Context) error {
	w.mu.Lock()
	s := workerState(w.state.Load())
	if s == workerStateStopping || s == workerStateStopped {
		w.mu.Unlock()
		return nil
	}
	if s != workerStateRunning {
		w.mu.Unlock()
		return ErrWorkerNotRunning
	}
	w.state.Store(uint32(workerStateStopping))
	w.mu.Unlock()

	w.pollCancel()

	select {
	case <-w.done:
		return nil
	case <-drainCtx.Done():
		w.workerCancel()
		return context.DeadlineExceeded
	}
}

// --- Poll loop ---

func (w *Worker) runLoop(r runnerEntry) {
	ticker := time.NewTicker(r.pollInterval)
	defer ticker.Stop()

	for {
		select {
		case <-w.pollCtx.Done():
			return
		case <-ticker.C:
			w.poll(r)
		}
	}
}

func (w *Worker) poll(r runnerEntry) {
	now := time.Now().UTC()
	claimToken, err := generateToken()
	if err != nil {
		w.log.Error("generate claim token failed", "queue", r.queue, "error", err)
		return
	}
	leaseUntil := now.Add(r.leaseDuration)

	claimed, err := w.queue.tasks.ClaimForUpdate(w.pollCtx,
		sq.And{
			sq.Eq{"status": TaskStatusPending, "queue": r.queue},
			sq.LtOrEq{"run_at": now},
		},
		[]string{"run_at ASC"},
		1,
		func(task *Task) error {
			task.Status = TaskStatusRunning
			task.LastRanAt = now
			task.ClaimedAt = &now
			task.LeaseUntil = &leaseUntil
			task.ClaimToken = &claimToken
			task.Try++
			return nil
		},
	)
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			w.log.Error("claim failed", "queue", r.queue, "error", err)
		}
		return
	}
	if len(claimed) == 0 {
		return
	}
	task := claimed[0]

	// Execute with cooperative deadline.
	deadline := leaseUntil.Add(-r.finalizeBuffer)
	handlerCtx, handlerCancel := context.WithDeadline(w.workerCtx, deadline)
	r.run(handlerCtx, task)
	handlerCancel()

	// Finalize with fencing.
	finalizeCtx, finalizeCancel := context.WithTimeout(w.workerCtx, 30*time.Second)
	rowsAffected, err := w.queue.finalizeTask(finalizeCtx, task, claimToken)
	finalizeCancel()

	if err != nil {
		w.log.Error("finalize failed",
			"queue", r.queue, "task_id", task.ID, "error", err)
	} else if rowsAffected == 0 {
		w.log.Warn("stale finalize (task was reaped)",
			"queue", r.queue, "task_id", task.ID, "claim_token", claimToken)
	}
}

// --- Operational goroutines ---

func (w *Worker) reapLoop() {
	ticker := time.NewTicker(w.reapInterval)
	defer ticker.Stop()

	for {
		select {
		case <-w.pollCtx.Done():
			return
		case <-ticker.C:
			w.reap()
		}
	}
}

func (w *Worker) reap() {
	now := time.Now().UTC()
	tag, err := w.queue.tasks.Query.Exec(w.workerCtx,
		psql.Update(w.queue.tasks.Name).
			Set("status", TaskStatusPending).
			Set("run_at", now).
			Set("claimed_at", nil).
			Set("lease_until", nil).
			Set("claim_token", nil).
			Where(sq.And{
				sq.Eq{"status": TaskStatusRunning},
				sq.Lt{"lease_until": now},
			}),
	)
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			w.log.Error("reap failed", "error", err)
		}
		return
	}
	if tag.RowsAffected() > 0 {
		w.log.Info("reaped expired leases", "count", tag.RowsAffected())
	}
}

func (w *Worker) monitorQuarantineLoop() {
	ticker := time.NewTicker(defaultMonitorInterval)
	defer ticker.Stop()

	for {
		select {
		case <-w.pollCtx.Done():
			return
		case <-ticker.C:
			w.monitorQuarantine()
		}
	}
}

func (w *Worker) monitorQuarantine() {
	count, err := w.queue.tasks.Count(w.workerCtx, sq.Eq{"status": TaskStatusDisabled})
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			w.log.Error("quarantine monitor failed", "error", err)
		}
		return
	}
	if count > 0 {
		w.log.Warn("tasks in quarantine (disabled)", "count", count)
	}
}

func (w *Worker) warnOrphanedQueues(ctx context.Context) {
	// Use a grouped query to avoid loading all tasks into memory.
	query := psql.Select("queue", "COUNT(*) as cnt").
		From(w.queue.tasks.Name).
		Where(sq.Eq{"status": []TaskStatus{TaskStatusPending, TaskStatusRunning, TaskStatusDisabled}}).
		GroupBy("queue")

	rows, err := w.queue.tasks.Query.QueryRows(ctx, query)
	if err != nil {
		w.log.Warn("orphaned queue check failed", "error", err)
		return
	}
	defer rows.Close()

	for rows.Next() {
		var queue string
		var count int64
		if err := rows.Scan(&queue, &count); err != nil {
			w.log.Warn("orphaned queue check: scan failed", "error", err)
			return
		}
		if _, registered := w.queues[queue]; !registered {
			w.log.Warn("orphaned queue: tasks with no registered handler",
				"queue", queue, "count", count)
		}
	}
}

// --- Type-erased handler adapters ---

func clearLease(t *Task) {
	t.ClaimedAt = nil
	t.LeaseUntil = nil
	t.ClaimToken = nil
}

func makeJobRunner[P any](w *Worker, spec JobSpec[P], handler JobHandler[P]) func(context.Context, *Task) {
	return func(ctx context.Context, task *Task) {
		var p P
		if err := task.DecodePayload(&p); err != nil {
			w.log.Error("decode payload",
				"queue", spec.Queue, "task_id", task.ID, "try", task.Try,
				"stage", "decode", "error", err)
			clearLease(task)
			task.Status = TaskStatusFailed
			return
		}
		job := &Job[P]{
			Task:    TaskInfo{ID: task.ID, Queue: task.Queue, Try: task.Try, RunAt: task.RunAt, CreatedAt: task.CreatedAt},
			Payload: p,
		}
		result := handler.RunTask(ctx, job)
		if result.action == resultSkip {
			result.action = resultFail
		}
		clearLease(task)
		applyJobResult(w.log, task, result, spec)
	}
}

func applyJobResult[P any](log *slog.Logger, t *Task, r Result, spec JobSpec[P]) {
	switch r.action {
	case resultDone:
		t.Status = TaskStatusCompleted
	case resultRetry:
		if r.err != nil {
			log.Error("task retry",
				"queue", spec.Queue, "task_id", t.ID, "try", t.Try,
				"stage", "handler", "error", r.err)
		}
		if t.Try >= spec.MaxRetries {
			t.Status = TaskStatusFailed
		} else {
			t.Status = TaskStatusPending
			t.RunAt = time.Now().UTC().Add(time.Duration(t.Try) * spec.RetryDelay)
		}
	case resultFail:
		if r.err != nil {
			log.Error("task failed",
				"queue", spec.Queue, "task_id", t.ID, "try", t.Try,
				"stage", "handler", "error", r.err)
		}
		t.Status = TaskStatusFailed
	}
}

func makeTickerEnsurer[P any](w *Worker, spec TickerSpec[P]) func(context.Context) error {
	return func(ctx context.Context) error {
		payload, err := json.Marshal(spec.InitialPayload)
		if err != nil {
			return fmt.Errorf("marshal initial payload: %w", err)
		}
		hash := spec.Key
		task := &Task{
			Queue:   spec.Queue,
			Status:  TaskStatusPending,
			RunAt:   time.Now().UTC(),
			Payload: payload,
			Hash:    &hash,
		}
		return w.queue.tasks.Upsert(ctx, []string{"queue", "hash"}, nil, task)
	}
}

func makeTickerRunner[P any](w *Worker, spec TickerSpec[P], handler TickerHandler[P]) func(context.Context, *Task) {
	return func(ctx context.Context, task *Task) {
		var p P
		if err := task.DecodePayload(&p); err != nil {
			w.log.Error("decode payload",
				"queue", spec.Queue, "task_id", task.ID, "try", task.Try,
				"stage", "decode", "error", err)
			clearLease(task)
			if task.Try >= spec.MaxRetries {
				w.log.Error("ticker quarantined (repeated codec failures)",
					"queue", spec.Queue, "task_id", task.ID, "tries", task.Try)
				task.Status = TaskStatusDisabled
			} else {
				task.Status = TaskStatusPending
				task.RunAt = time.Now().UTC().Add(time.Duration(task.Try) * spec.RetryDelay)
			}
			return
		}
		job := &Job[P]{
			Task:    TaskInfo{ID: task.ID, Queue: task.Queue, Try: task.Try, RunAt: task.RunAt, CreatedAt: task.CreatedAt},
			Payload: p,
		}
		result := handler.RunTick(ctx, job)
		clearLease(task)
		applyTickerResult(w.log, task, job, result, spec)
	}
}

func applyTickerResult[P any](log *slog.Logger, t *Task, job *Job[P], r Result, spec TickerSpec[P]) {
	switch r.action {
	case resultDone:
		if encoded, err := json.Marshal(job.Payload); err != nil {
			log.Error("encode payload",
				"queue", spec.Queue, "task_id", t.ID, "stage", "encode", "error", err)
			if t.Try >= spec.MaxRetries {
				log.Error("ticker quarantined (repeated encode failures)",
					"queue", spec.Queue, "task_id", t.ID, "tries", t.Try)
				t.Status = TaskStatusDisabled
			} else {
				t.Status = TaskStatusPending
				t.RunAt = time.Now().UTC().Add(time.Duration(t.Try) * spec.RetryDelay)
			}
		} else {
			t.Payload = encoded
			t.Try = 0
			t.RunAt = time.Now().UTC().Add(spec.Every)
			t.Status = TaskStatusPending
		}
	case resultRetry:
		if r.err != nil {
			log.Error("ticker retry",
				"queue", spec.Queue, "task_id", t.ID, "try", t.Try,
				"stage", "handler", "error", r.err)
		}
		if t.Try >= spec.MaxRetries {
			log.Warn("ticker exhausted retries, rescheduling at next interval",
				"queue", spec.Queue, "task_id", t.ID, "tries", t.Try)
			t.Try = 0
			t.RunAt = time.Now().UTC().Add(spec.Every)
		} else {
			t.RunAt = time.Now().UTC().Add(time.Duration(t.Try) * spec.RetryDelay)
		}
		t.Status = TaskStatusPending
	case resultSkip:
		if r.err != nil {
			log.Warn("ticker run skipped",
				"queue", spec.Queue, "task_id", t.ID,
				"stage", "handler", "error", r.err)
		}
		if encoded, err := json.Marshal(job.Payload); err != nil {
			log.Error("encode payload on skip",
				"queue", spec.Queue, "task_id", t.ID, "stage", "encode", "error", err)
			if t.Try >= spec.MaxRetries {
				log.Error("ticker quarantined (encode failure on skip)",
					"queue", spec.Queue, "task_id", t.ID, "tries", t.Try)
				t.Status = TaskStatusDisabled
			} else {
				t.Status = TaskStatusPending
				t.RunAt = time.Now().UTC().Add(time.Duration(t.Try) * spec.RetryDelay)
			}
		} else {
			t.Payload = encoded
			t.Try = 0
			t.RunAt = time.Now().UTC().Add(spec.Every)
			t.Status = TaskStatusPending
		}
	case resultFail:
		if r.err != nil {
			log.Error("ticker permanently failed",
				"queue", spec.Queue, "task_id", t.ID,
				"stage", "handler", "error", r.err)
		}
		t.Status = TaskStatusFailed
	}
}

// --- Finalize ---

func (q *Queue) finalizeTask(ctx context.Context, task *Task, claimToken string) (int64, error) {
	query := psql.Update(q.tasks.Name).
		Set("status", task.Status).
		Set("try", task.Try).
		Set("run_at", task.RunAt).
		Set("last_ran_at", task.LastRanAt).
		Set("payload", task.Payload).
		Set("claimed_at", nil).
		Set("lease_until", nil).
		Set("claim_token", nil).
		Where(sq.Eq{"id": task.ID, "claim_token": claimToken})

	tag, err := q.tasks.Query.Exec(ctx, query)
	if err != nil {
		return 0, fmt.Errorf("pgqueue: finalize task: %w", err)
	}
	return tag.RowsAffected(), nil
}

// --- Defaults ---

func applyJobDefaults[P any](s JobSpec[P]) JobSpec[P] {
	if s.PollInterval == 0 {
		s.PollInterval = defaultPollInterval
	}
	if s.MaxRetries < 0 {
		s.MaxRetries = defaultMaxRetries
	}
	if s.RetryDelay == 0 {
		s.RetryDelay = defaultRetryDelay
	}
	if s.LeaseDuration == 0 {
		s.LeaseDuration = defaultLeaseDuration
	}
	if s.FinalizeBuffer == 0 {
		s.FinalizeBuffer = defaultFinalizeBuffer
	}
	return s
}

func applyTickerDefaults[P any](s TickerSpec[P]) TickerSpec[P] {
	if s.PollInterval == 0 {
		s.PollInterval = max(min(s.Every/2, 30*time.Second), time.Millisecond)
	}
	if s.MaxRetries < 0 {
		s.MaxRetries = defaultMaxRetries
	}
	if s.RetryDelay == 0 {
		s.RetryDelay = defaultRetryDelay
	}
	if s.LeaseDuration == 0 {
		s.LeaseDuration = defaultLeaseDuration
	}
	if s.FinalizeBuffer == 0 {
		s.FinalizeBuffer = defaultFinalizeBuffer
	}
	return s
}

// --- Helpers ---

func generateToken() (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("pgqueue: generate claim token: %w", err)
	}
	return hex.EncodeToString(b), nil
}