textmachine/platform/internal/jobs/jobs.go

// Package jobs is the platform's queue: one River client, and the job kinds the control plane hands
// to itself.
//
// It knows nothing about what a job MEANS. Both workers below call one method of one service and
// return, because the work itself belongs to the package that owns the object — the run lifecycle to
// `runs`, the intake to `books` — and a queue that also owned the work would be a second place where
// a run can be started.
//
// One client for both kinds, deliberately: River runs its own maintenance (scheduler, rescuer,
// cleaner) per client, and a second client in one process would be a second set of those against the
// same tables for no gain.
package jobs

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/riverqueue/river"
	"github.com/riverqueue/river/riverdriver/riverpgxv5"

	"textmachine/platform/internal/pgstore"
)

// Spawner starts the engine for a run that has already been admitted.
type Spawner interface {
	Spawn(ctx context.Context, runID string) error
}

// Parser turns a received book file into a chapter tree.
type Parser interface {
	Parse(ctx context.Context, bookID string) error
}

// SpawnArgs is one queued permission to start a run.
//
// The job carries an id and nothing else. Everything about the run — the book, the ceiling, the
// binary it is pinned to — is in Postgres, and a job that carried its own copy would be a second
// answer that goes stale the moment the run is restarted with a smaller budget.
type SpawnArgs struct {
	RunID string `json:"run_id"`
}

// Kind is River's name for this job type.
func (SpawnArgs) Kind() string { return "tm_spawn_run" }

// InsertOpts pins the queue-level policy.
//
// MaxAttempts is 1 on purpose, and it is the opposite of the usual queue reflex. A retry here does
// not repeat lost work: the run row already exists and holds the account's money, and the thing that
// would be repeated is spawning an engine. What recovers a run whose spawn failed is the reconciler,
// which reads the world instead of assuming the job's view of it — and which is the ONLY component
// that can tell "the unit never started" from "the unit is running and this platform was restarted".
func (SpawnArgs) InsertOpts() river.InsertOpts {
	return river.InsertOpts{MaxAttempts: 1, Queue: river.QueueDefault}
}

// ParseArgs is one queued book waiting to be cut into chapters.
type ParseArgs struct {
	BookID string `json:"book_id"`
}

// Kind is River's name for this job type.
func (ParseArgs) Kind() string { return "tm_parse_book" }

// InsertOpts pins the queue-level policy.
//
// MaxAttempts is 1 for the same reason as the spawn's, arrived at from the other direction: a parse
// is free and idempotent, so a retry would be harmless — but it would also be a SECOND recovery
// mechanism next to the intake sweep, which already re-claims a parse whose process is gone and
// which is the only one that can bound how many times a broken host re-runs the engine.
func (ParseArgs) InsertOpts() river.InsertOpts {
	return river.InsertOpts{MaxAttempts: 1, Queue: river.QueueDefault}
}

type spawnWorker struct {
	river.WorkerDefaults[SpawnArgs]
	svc Spawner
}

func (w *spawnWorker) Work(ctx context.Context, job *river.Job[SpawnArgs]) error {
	return w.svc.Spawn(ctx, job.Args.RunID)
}

type parseWorker struct {
	river.WorkerDefaults[ParseArgs]
	svc Parser
}

func (w *parseWorker) Work(ctx context.Context, job *river.Job[ParseArgs]) error {
	return w.svc.Parse(ctx, job.Args.BookID)
}

// JobTimeout bounds ONE job. Generous because the work is an engine call over a whole book, and
// finite because a worker that never returns is a worker slot that never comes back.
//
// Exported because it is half of a pair: the intake's claim grace must OUTLIVE it, or a backstop
// sweep takes the claim off a parse this queue is still running. The other half is written as this
// constant plus a margin.
const JobTimeout = 15 * time.Minute

// Queue is the River client, wired to the services that do the work.
type Queue struct {
	client *river.Client[pgstore.Tx]
}

// New builds the queue and its worker pool.
//
// Concurrency is deliberately small: a spawn worker's whole job is to create a transient unit, and
// the runs themselves are bounded by their own cgroups and by the one-live-run-per-book index, not
// by how many workers exist.
func New(pool *pgxpool.Pool, spawner Spawner, parser Parser, log *slog.Logger, workers int) (*Queue, error) {
	if workers <= 0 {
		workers = 4
	}
	w := river.NewWorkers()
	if spawner != nil {
		if err := river.AddWorkerSafely(w, &spawnWorker{svc: spawner}); err != nil {
			return nil, fmt.Errorf("jobs: register spawn worker: %w", err)
		}
	}
	if parser != nil {
		if err := river.AddWorkerSafely(w, &parseWorker{svc: parser}); err != nil {
			return nil, fmt.Errorf("jobs: register parse worker: %w", err)
		}
	}
	c, err := river.NewClient(riverpgxv5.New(pool), &river.Config{
		Logger:  log,
		Workers: w,
		Queues:  map[string]river.QueueConfig{river.QueueDefault: {MaxWorkers: workers}},
		// CHOSEN, not inherited. River's own default is one minute, and both jobs here call the engine
		// against a whole book: a parse ingests and cuts a source of up to the intake's limit, and a
		// spawn reads the book's meter first (seconds of CPU on a 23 MB book, unified backlog row 100).
		// A minute would kill those, and a killed process is indistinguishable from a host that cannot
		// run the engine — so the intake would count a deployment fault against a book whose only sin
		// was being large.
		JobTimeout: JobTimeout,
	})
	if err != nil {
		return nil, fmt.Errorf("jobs: river client: %w", err)
	}
	return &Queue{client: c}, nil
}

// EnqueueRun inserts the job in the CALLER's transaction, so the run row, its hold and its queue
// entry commit together or not at all.
func (q *Queue) EnqueueRun(ctx context.Context, tx pgstore.Tx, runID string) error {
	if _, err := q.client.InsertTx(ctx, tx, SpawnArgs{RunID: runID}, nil); err != nil {
		return fmt.Errorf("jobs: enqueue run %s: %w", runID, err)
	}
	return nil
}

// EnqueueRunNow inserts the job on its own, outside any transaction.
//
// The caller is a RESUME: the run it names has already been re-opened and its money already held, in
// a transaction that is committed by the time this runs. A failed insert therefore costs one sweep
// interval and nothing else — the reconciler finds an attempt with no unit and spawns it — which is
// exactly what a lost queue entry costs on the admission path too.
func (q *Queue) EnqueueRunNow(ctx context.Context, runID string) error {
	if _, err := q.client.Insert(ctx, SpawnArgs{RunID: runID}, nil); err != nil {
		return fmt.Errorf("jobs: enqueue run %s: %w", runID, err)
	}
	return nil
}

// EnqueueParse inserts the job in the CALLER's transaction, so a book that says it is being parsed
// and the job that parses it commit together.
func (q *Queue) EnqueueParse(ctx context.Context, tx pgstore.Tx, bookID string) error {
	if _, err := q.client.InsertTx(ctx, tx, ParseArgs{BookID: bookID}, nil); err != nil {
		return fmt.Errorf("jobs: enqueue parse %s: %w", bookID, err)
	}
	return nil
}

// Start begins working jobs.
func (q *Queue) Start(ctx context.Context) error { return q.client.Start(ctx) }

// Stop drains the workers. The runs themselves are untouched: they are transient units, not
// children, and outliving this process is what they are for.
func (q *Queue) Stop(ctx context.Context) error { return q.client.Stop(ctx) }