Vellum/internal/jobs/jobs_test.go

package jobs

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

func testDB(t *testing.T) *sql.DB {
	t.Helper()
	db, err := core.OpenDB(filepath.Join(t.TempDir(), "jobs.db"))
	if err != nil {
		t.Fatalf("OpenDB: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })
	if err := core.Migrate(context.Background(), db, migrations.FS, discard()); err != nil {
		t.Fatalf("Migrate: %v", err)
	}
	return db
}

func discard() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) }

func stateOf(t *testing.T, db *sql.DB, id int64) (string, int, string) {
	t.Helper()
	var state, errMsg string
	var attempts int
	if err := db.QueryRow(
		`SELECT state, attempts, error FROM jobs WHERE id = ?`, id).Scan(&state, &attempts, &errMsg); err != nil {
		t.Fatal(err)
	}
	return state, attempts, errMsg
}

func TestRunOnceExecutesHandler(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())

	var got json.RawMessage
	q.Register("greet", func(_ context.Context, payload json.RawMessage) error {
		got = payload
		return nil
	})

	if err := q.Enqueue(context.Background(), "greet", map[string]string{"name": "ada"}); err != nil {
		t.Fatal(err)
	}

	ran, err := q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if !ran {
		t.Fatal("RunOnce reported nothing to do")
	}
	if string(got) != `{"name":"ada"}` {
		t.Errorf("payload = %s", got)
	}

	// Nothing left to run.
	ran, err = q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if ran {
		t.Error("RunOnce ran a job that was already done")
	}

	var state string
	var n int
	if err := db.QueryRow(`SELECT state, count(*) FROM jobs GROUP BY state`).Scan(&state, &n); err != nil {
		t.Fatal(err)
	}
	if state != StateDone || n != 1 {
		t.Errorf("state=%s count=%d, want done/1", state, n)
	}
}

func TestUnknownKindFailsImmediately(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())

	if err := q.Enqueue(context.Background(), "nobody-handles-this", nil); err != nil {
		t.Fatal(err)
	}
	if _, err := q.RunOnce(context.Background()); err != nil {
		t.Fatal(err)
	}

	var state, errMsg string
	if err := db.QueryRow(`SELECT state, error FROM jobs`).Scan(&state, &errMsg); err != nil {
		t.Fatal(err)
	}
	if state != StateFailed {
		t.Errorf("state = %s, want failed (a missing handler must not retry)", state)
	}
	if errMsg == "" {
		t.Error("failure reason was not recorded")
	}
}

func TestFailureBacksOffThenFails(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())
	q.maxAttempts = 2
	q.baseBackoff = 30 * 1000 * 1000 * 1000 // 30s

	boom := errors.New("transient")
	q.Register("flaky", func(context.Context, json.RawMessage) error { return boom })

	if err := q.Enqueue(context.Background(), "flaky", nil); err != nil {
		t.Fatal(err)
	}

	// First run fails and reschedules into the future.
	if _, err := q.RunOnce(context.Background()); err != nil {
		t.Fatal(err)
	}
	state, attempts, errMsg := stateOf(t, db, 1)
	if state != StateQueued || attempts != 1 {
		t.Fatalf("after first failure: state=%s attempts=%d", state, attempts)
	}
	if errMsg == "" {
		t.Error("failure reason was not recorded")
	}

	// Not due yet, so nothing runs.
	ran, err := q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if ran {
		t.Error("a backed-off job ran before its run_after")
	}

	// Force it due and run again: attempts is now at the ceiling.
	if _, err := db.Exec(`UPDATE jobs SET run_after = '1970-01-01T00:00:00Z'`); err != nil {
		t.Fatal(err)
	}
	if _, err := q.RunOnce(context.Background()); err != nil {
		t.Fatal(err)
	}
	state, attempts, _ = stateOf(t, db, 1)
	if state != StateFailed {
		t.Errorf("after exhausting attempts: state=%s, want failed", state)
	}
	if attempts != 2 {
		t.Errorf("attempts = %d, want 2", attempts)
	}
}

func TestEnqueueAfterDelays(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())
	q.Register("later", func(context.Context, json.RawMessage) error { return nil })

	if err := q.EnqueueAfter(context.Background(), "later", nil, 60*1000*1000*1000); err != nil {
		t.Fatal(err)
	}
	ran, err := q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if ran {
		t.Error("a delayed job ran immediately")
	}
}

func TestStats(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())
	q.Register("noop", func(context.Context, json.RawMessage) error { return nil })

	for i := 0; i < 3; i++ {
		if err := q.Enqueue(context.Background(), "noop", nil); err != nil {
			t.Fatal(err)
		}
	}
	if _, err := q.RunOnce(context.Background()); err != nil {
		t.Fatal(err)
	}

	stats, err := q.Stats(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if stats[StateDone] != 1 || stats[StateQueued] != 2 {
		t.Errorf("stats = %v, want done:1 queued:2", stats)
	}
}

func TestRecoverStaleRequeuesRunning(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())

	for _, kind := range []string{"a", "b"} {
		if _, err := db.Exec(
			`INSERT INTO jobs (kind, state, run_after) VALUES (?, ?, '1970-01-01T00:00:00Z')`,
			kind, StateRunning); err != nil {
			t.Fatal(err)
		}
	}
	if _, err := db.Exec(`INSERT INTO jobs (kind, state) VALUES ('c', ?)`, StateQueued); err != nil {
		t.Fatal(err)
	}
	if _, err := db.Exec(`INSERT INTO jobs (kind, state) VALUES ('d', ?)`, StateDone); err != nil {
		t.Fatal(err)
	}

	n, err := q.RecoverStale(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if n != 2 {
		t.Fatalf("RecoverStale moved %d rows, want 2", n)
	}

	var running int
	if err := db.QueryRow(`SELECT count(*) FROM jobs WHERE state = ?`, StateRunning).Scan(&running); err != nil {
		t.Fatal(err)
	}
	if running != 0 {
		t.Errorf("still %d rows in %s", running, StateRunning)
	}

	// Requeued rows must be due now, not stuck in the future. Scope to the
	// recovered rows: rows inserted via SQLite's strftime default use a
	// different fractional-second format and are not what this test checks.
	var future int
	if err := db.QueryRow(`SELECT count(*) FROM jobs WHERE state = ? AND kind IN ('a', 'b') AND run_after > ?`,
		StateQueued, core.Now()).Scan(&future); err != nil {
		t.Fatal(err)
	}
	if future != 0 {
		t.Errorf("%d requeued rows have a future run_after", future)
	}

	// A queued row that was already fine and a done row are left alone.
	var done int
	if err := db.QueryRow(`SELECT count(*) FROM jobs WHERE state = ?`, StateDone).Scan(&done); err != nil {
		t.Fatal(err)
	}
	if done != 1 {
		t.Errorf("done rows = %d, want 1", done)
	}
}

func TestRecoverStaleEmpty(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())

	n, err := q.RecoverStale(context.Background())
	if err != nil {
		t.Fatalf("RecoverStale on empty table: %v", err)
	}
	if n != 0 {
		t.Errorf("RecoverStale on empty table moved %d rows, want 0", n)
	}
}

func TestPanicFailsJobWithoutCrashing(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())
	q.Register("boom", func(context.Context, json.RawMessage) error {
		panic("kaboom")
	})

	if err := q.Enqueue(context.Background(), "boom", nil); err != nil {
		t.Fatal(err)
	}

	ran, err := q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if !ran {
		t.Fatal("RunOnce reported nothing to do")
	}

	state, _, errMsg := stateOf(t, db, 1)
	if state != StateFailed {
		t.Errorf("state = %s, want failed after a panic", state)
	}
	if errMsg == "" {
		t.Error("panic reason was not recorded")
	}
}

func TestRunUsesConcurrentWorkers(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())
	q.Workers = 4
	q.pollEvery = 5 * time.Millisecond

	const n = 8
	var running, peak int32
	var wg sync.WaitGroup
	wg.Add(n)
	q.Register("work", func(context.Context, json.RawMessage) error {
		cur := atomic.AddInt32(&running, 1)
		for {
			p := atomic.LoadInt32(&peak)
			if cur <= p || atomic.CompareAndSwapInt32(&peak, p, cur) {
				break
			}
		}
		time.Sleep(20 * time.Millisecond)
		atomic.AddInt32(&running, -1)
		wg.Done()
		return nil
	})

	for i := 0; i < n; i++ {
		if err := q.Enqueue(context.Background(), "work", nil); err != nil {
			t.Fatal(err)
		}
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan struct{})
	go func() { q.Run(ctx); close(done) }()

	allDone := make(chan struct{})
	go func() { wg.Wait(); close(allDone) }()
	select {
	case <-allDone:
	case <-time.After(10 * time.Second):
		t.Fatal("jobs did not all run")
	}

	cancel()
	select {
	case <-done:
	case <-time.After(10 * time.Second):
		t.Fatal("Run did not return after its context was cancelled")
	}

	if got := atomic.LoadInt32(&peak); got < 2 {
		t.Errorf("peak concurrency = %d, want >= 2 (workers should not be serialised)", got)
	}
}