Vellum/internal/jobs/jobs.go

// Package jobs is a table-backed work queue.
//
// It lives in the same SQLite file as everything else: no Redis, no second
// process, and the queue is included in the same backup. SQLite's single
// writer makes the claim below atomic without any extra locking.
package jobs

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"math/rand/v2"
	"sync"
	"time"

	"github.com/mithral/vellum/internal/core"
)

// States a job can be in.
const (
	StateQueued  = "queued"
	StateRunning = "running"
	StateDone    = "done"
	StateFailed  = "failed"
)

// Handler processes one job payload.
type Handler func(ctx context.Context, payload json.RawMessage) error

// Queue claims and runs jobs.
//
// Handlers registered with Register now run CONCURRENTLY (see Run), so a
// handler must not assume it is the only thing touching the database: use
// transactions or per-row writes rather than read-modify-write across calls.
type Queue struct {
	db          *sql.DB
	log         *slog.Logger
	mu          sync.RWMutex
	handlers    map[string]Handler
	maxAttempts int
	pollEvery   time.Duration
	baseBackoff time.Duration

	// Workers is how many jobs Run processes at once. A long job (a library
	// scan) therefore no longer blocks every other job kind behind it.
	Workers int
}

// New builds a Queue with sensible defaults.
func New(db *sql.DB, log *slog.Logger) *Queue {
	return &Queue{
		db:          db,
		log:         log,
		handlers:    map[string]Handler{},
		maxAttempts: 5,
		pollEvery:   time.Second,
		baseBackoff: 30 * time.Second,
		Workers:     4,
	}
}

// Register installs the handler for a job kind.
func (q *Queue) Register(kind string, h Handler) {
	q.mu.Lock()
	defer q.mu.Unlock()
	q.handlers[kind] = h
}

// handler looks up a registered handler under a read lock, so Register is safe
// to call while workers are running.
func (q *Queue) handler(kind string) (Handler, bool) {
	q.mu.RLock()
	defer q.mu.RUnlock()
	h, ok := q.handlers[kind]
	return h, ok
}

// Enqueue adds a job to run as soon as possible.
func (q *Queue) Enqueue(ctx context.Context, kind string, payload any) error {
	return q.EnqueueAfter(ctx, kind, payload, 0)
}

// EnqueueRef adds a job that refers to a specific entity (a library, an item),
// so its state and result can be looked up by that id later.
func (q *Queue) EnqueueRef(ctx context.Context, kind string, refID int64, payload any) (int64, error) {
	body := "{}"
	if payload != nil {
		raw, err := json.Marshal(payload)
		if err != nil {
			return 0, fmt.Errorf("jobs: marshal payload: %w", err)
		}
		body = string(raw)
	}
	var id int64
	err := q.db.QueryRowContext(ctx, `
		INSERT INTO jobs (kind, ref_id, payload, state, run_after)
		VALUES (?, ?, ?, ?, ?)
		RETURNING id`,
		kind, refID, body, StateQueued, core.Now()).Scan(&id)
	if err != nil {
		return 0, fmt.Errorf("jobs: enqueue %s: %w", kind, err)
	}
	return id, nil
}

// EnqueueAfter adds a job that will not run before delay has elapsed.
func (q *Queue) EnqueueAfter(ctx context.Context, kind string, payload any, delay time.Duration) error {
	body := "{}"
	if payload != nil {
		raw, err := json.Marshal(payload)
		if err != nil {
			return fmt.Errorf("jobs: marshal payload: %w", err)
		}
		body = string(raw)
	}
	_, err := q.db.ExecContext(ctx, `
		INSERT INTO jobs (kind, payload, state, run_after)
		VALUES (?, ?, ?, ?)`,
		kind, body, StateQueued, core.Stamp(time.Now().Add(delay)))
	if err != nil {
		return fmt.Errorf("jobs: enqueue %s: %w", kind, err)
	}
	return nil
}

// Run processes jobs until the context is cancelled.
//
// It starts q.Workers goroutines, each of which claims and runs one job at a
// time. That keeps a long-running job (a library scan) from blocking every
// other job kind. Run only returns once every worker has stopped, which happens
// when ctx is cancelled.
func (q *Queue) Run(ctx context.Context) {
	workers := q.Workers
	if workers < 1 {
		workers = 1
	}

	var wg sync.WaitGroup
	wg.Add(workers)
	for i := 0; i < workers; i++ {
		go func() {
			defer wg.Done()
			q.worker(ctx)
		}()
	}
	wg.Wait()
}

// worker drains due jobs until ctx is cancelled. When the queue is empty it
// waits pollEvery plus a small jitter before claiming again, so idle workers do
// not busy-spin (and do not all wake on the same tick).
func (q *Queue) worker(ctx context.Context) {
	for {
		if ctx.Err() != nil {
			return
		}

		ran, err := q.RunOnce(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			q.log.Error("job queue", "err", err)
		}
		if ran {
			continue
		}

		select {
		case <-ctx.Done():
			return
		case <-time.After(q.idleWait()):
		}
	}
}

// idleWait is pollEvery with up to a quarter of that added as jitter.
func (q *Queue) idleWait() time.Duration {
	jitter := q.pollEvery / 4
	if jitter <= 0 {
		return q.pollEvery
	}
	return q.pollEvery + time.Duration(rand.Int64N(int64(jitter)))
}

// RunOnce claims and runs at most one job. It reports whether it ran anything,
// which is what makes the queue testable without a ticker.
func (q *Queue) RunOnce(ctx context.Context) (bool, error) {
	job, ok, err := q.claim(ctx)
	if err != nil || !ok {
		return false, err
	}

	handler, exists := q.handler(job.Kind)
	if !exists {
		// A missing handler is a programming error, not a transient fault.
		// Retrying it just delays the inevitable and clutters the queue.
		q.failPermanently(ctx, job, fmt.Errorf("no handler registered for %q", job.Kind))
		return true, nil
	}

	q.runHandler(ctx, job, handler)
	return true, nil
}

// runHandler invokes a handler and records its outcome. A panic is recovered so
// it cannot kill the worker goroutine or leave the job stuck in 'running'; the
// job is marked failed instead.
func (q *Queue) runHandler(ctx context.Context, job claimedJob, handler Handler) {
	defer func() {
		if r := recover(); r != nil {
			q.log.Error("job panicked", "kind", job.Kind, "id", job.ID, "panic", r)
			q.failPermanently(ctx, job, fmt.Errorf("job handler panicked: %v", r))
		}
	}()

	if err := handler(withJobID(ctx, job.ID), job.Payload); err != nil {
		q.log.Warn("job failed", "kind", job.Kind, "id", job.ID, "attempt", job.Attempts, "err", err)
		q.finish(ctx, job, err)
		return
	}
	q.finish(ctx, job, nil)
}

// RecoverStale returns every job stuck in StateRunning to StateQueued so it is
// retried, and reports how many rows it moved.
//
// This is safe because the server is a single process and RecoverStale is meant
// to be called once at startup: at that point nothing can legitimately be
// mid-job, so any 'running' row is left over from a process that died (or was
// restarted) before it could finish. Requeuing them costs at most one repeated
// execution, which is the usual at-least-once guarantee of a work queue. It
// relies only on columns that already exist, and returns (0, nil) on an empty
// table.
func (q *Queue) RecoverStale(ctx context.Context) (int, error) {
	res, err := q.db.ExecContext(ctx, `
		UPDATE jobs SET state = ?, run_after = ?
		WHERE state = ?`,
		StateQueued, core.Now(), StateRunning)
	if err != nil {
		return 0, fmt.Errorf("jobs: recover stale: %w", err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("jobs: recover stale: %w", err)
	}
	return int(n), nil
}

type claimedJob struct {
	ID       int64
	Kind     string
	Payload  json.RawMessage
	Attempts int
}

// claim atomically takes the oldest due job.
//
// The UPDATE ... RETURNING form means a job cannot be claimed twice even with
// several workers: the subselect and the state change are one statement.
func (q *Queue) claim(ctx context.Context) (claimedJob, bool, error) {
	var (
		job     claimedJob
		payload string
	)
	err := q.db.QueryRowContext(ctx, `
		UPDATE jobs SET state = ?, attempts = attempts + 1
		WHERE id = (
			SELECT id FROM jobs
			WHERE state = ? AND run_after <= ?
			ORDER BY run_after, id
			LIMIT 1
		)
		RETURNING id, kind, payload, attempts`,
		StateRunning, StateQueued, core.Now(),
	).Scan(&job.ID, &job.Kind, &payload, &job.Attempts)
	if errors.Is(err, sql.ErrNoRows) {
		return claimedJob{}, false, nil
	}
	if err != nil {
		return claimedJob{}, false, fmt.Errorf("jobs: claim: %w", err)
	}
	job.Payload = json.RawMessage(payload)
	return job, true, nil
}

// finish marks a job done, or reschedules it with exponential backoff until it
// exhausts its attempts.
func (q *Queue) finish(ctx context.Context, job claimedJob, runErr error) {
	if runErr == nil {
		if _, err := q.db.ExecContext(ctx,
			`UPDATE jobs SET state = ?, error = '' WHERE id = ?`, StateDone, job.ID); err != nil {
			q.log.Error("jobs: mark done", "id", job.ID, "err", err)
		}
		return
	}

	if job.Attempts >= q.maxAttempts {
		if _, err := q.db.ExecContext(ctx,
			`UPDATE jobs SET state = ?, error = ? WHERE id = ?`,
			StateFailed, runErr.Error(), job.ID); err != nil {
			q.log.Error("jobs: mark failed", "id", job.ID, "err", err)
		}
		return
	}

	backoff := q.baseBackoff * time.Duration(1<<(job.Attempts-1))
	runAfter := core.Stamp(time.Now().Add(backoff))
	if _, err := q.db.ExecContext(ctx,
		`UPDATE jobs SET state = ?, error = ?, run_after = ? WHERE id = ?`,
		StateQueued, runErr.Error(), runAfter, job.ID); err != nil {
		q.log.Error("jobs: reschedule", "id", job.ID, "err", err)
	}
}

// failPermanently marks a job failed with no further attempts.
func (q *Queue) failPermanently(ctx context.Context, job claimedJob, err error) {
	if _, err := q.db.ExecContext(ctx,
		`UPDATE jobs SET state = ?, error = ? WHERE id = ?`,
		StateFailed, err.Error(), job.ID); err != nil {
		q.log.Error("jobs: mark failed", "id", job.ID, "err", err)
	}
}

// jobCtxKey carries the running job's id to its handler, so a handler can
// attach a result without the queue needing to know the result's shape.
type jobCtxKey struct{}

func withJobID(ctx context.Context, id int64) context.Context {
	return context.WithValue(ctx, jobCtxKey{}, id)
}

// JobID returns the id of the job currently being handled.
func JobID(ctx context.Context) (int64, bool) {
	id, ok := ctx.Value(jobCtxKey{}).(int64)
	return id, ok
}

// SetResult stores a job's result so it can be read back later. Best-effort:
// a handler that produced a result should not fail because recording it did.
func SetResult(ctx context.Context, db *sql.DB, jobID int64, v any) error {
	body, err := json.Marshal(v)
	if err != nil {
		return fmt.Errorf("jobs: marshal result: %w", err)
	}
	_, err = db.ExecContext(ctx, `UPDATE jobs SET result = ? WHERE id = ?`, string(body), jobID)
	return err
}

// Stats reports job counts by state, for the health endpoint and debugging.
func (q *Queue) Stats(ctx context.Context) (map[string]int, error) {
	rows, err := q.db.QueryContext(ctx, `SELECT state, count(*) FROM jobs GROUP BY state`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := map[string]int{}
	for rows.Next() {
		var state string
		var n int
		if err := rows.Scan(&state, &n); err != nil {
			return nil, err
		}
		out[state] = n
	}
	return out, rows.Err()
}