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"
	"time"
)

// 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.
type Queue struct {
	db          *sql.DB
	log         *slog.Logger
	handlers    map[string]Handler
	maxAttempts int
	pollEvery   time.Duration
	baseBackoff time.Duration
}

// 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,
	}
}

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

// 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, time.Now().UTC().Format(time.RFC3339Nano)).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, time.Now().UTC().Add(delay).Format(time.RFC3339Nano))
	if err != nil {
		return fmt.Errorf("jobs: enqueue %s: %w", kind, err)
	}
	return nil
}

// Run processes jobs until the context is cancelled.
func (q *Queue) Run(ctx context.Context) {
	ticker := time.NewTicker(q.pollEvery)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			// Drain everything currently due before waiting again.
			for {
				ran, err := q.RunOnce(ctx)
				if err != nil {
					if ctx.Err() != nil {
						return
					}
					q.log.Error("job queue", "err", err)
					break
				}
				if !ran {
					break
				}
			}
		}
	}
}

// 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.handlers[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
	}

	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 true, nil
	}
	q.finish(ctx, job, nil)
	return true, 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, time.Now().UTC().Format(time.RFC3339Nano),
	).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 := time.Now().UTC().Add(backoff).Format(time.RFC3339Nano)
	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()
}