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"
	"errors"
	"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
}

// Exporter builds one requested reader's copy of a book.
type Exporter interface {
	Build(ctx context.Context, exportID 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}
}

// ExportArgs is one queued export waiting to be built.
type ExportArgs struct {
	ExportID string `json:"export_id"`
}

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

// InsertOpts pins the queue-level policy.
//
// MaxAttempts is 1, the third time this package refuses the usual queue reflex and the first time
// the reason is a USER-VISIBLE promise rather than money. A retry would build a second file for a
// row that already carries a verdict, and the poll the canon requires to END would have been ended
// by the first attempt already. What recovers a build whose worker is gone is the export sweep,
// which reads the row's age instead of assuming the job's view of it.
func (ExportArgs) 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 {
	err := w.svc.Parse(ctx, job.Args.BookID)
	if errors.Is(err, ErrTryAgainLater) {
		// The pass established NOTHING and said so. Snoozing rather than failing, and the difference
		// matters because MaxAttempts is 1: a returned error consumes this job outright, and the book
		// would then wait out the intake sweep's whole grace over a host that was busy for a moment.
		// A snooze does not increment the attempt (river.JobSnooze), so it is the SAME recovery
		// mechanism coming back — not the second one this queue's policy refuses.
		return river.JobSnooze(RetryDelay)
	}
	return err
}

// ErrTryAgainLater is a pass that did nothing and wants the job back rather than spent.
//
// Declared HERE rather than by the service, because what it selects is a QUEUE policy and this is the
// package that owns one. A service wraps its own reason in it (books.giveBack); this package decides
// what the queue does about it, and the two cannot drift into different opinions about whether the
// job is finished.
var ErrTryAgainLater = errors.New("jobs: this pass established nothing and the job should come back")

// RetryDelay is how long a job that established nothing waits before it is offered again.
//
// Shorter than a cut, because what it usually waits out is one: the ordinary producer of
// ErrTryAgainLater is a host that has as many books under the engine as it will take, and a slot
// frees when one of them finishes. Long enough that a saturated host is not re-asked in a tight loop.
const RetryDelay = 30 * time.Second

type exportWorker struct {
	river.WorkerDefaults[ExportArgs]
	svc Exporter
}

func (w *exportWorker) Work(ctx context.Context, job *river.Job[ExportArgs]) error {
	return w.svc.Build(ctx, job.Args.ExportID)
}

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

// DefaultWorkers is how many jobs this queue runs at once when a deployment does not say.
//
// The number a host is sized for, and therefore the ONE place it is written: the intake's cap on
// concurrent engine cuts is the same figure said for every way of starting one rather than only for
// the way that goes through here (books.DefaultMaxCuts), and two literals coupled by prose is the
// drift this package has already paid for elsewhere.
const DefaultWorkers = 4

// 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, exporter Exporter, log *slog.Logger, workers int) (*Queue, error) {
	if workers <= 0 {
		workers = DefaultWorkers
	}
	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)
		}
	}
	if exporter != nil {
		if err := river.AddWorkerSafely(w, &exportWorker{svc: exporter}); err != nil {
			return nil, fmt.Errorf("jobs: register export 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
}

// EnqueueExport inserts the job in the CALLER's transaction, so the export row and the job that
// builds it commit together — a row with no job is a poll the stale sweep has to end, and a job
// with no row is a worker that reads nothing.
func (q *Queue) EnqueueExport(ctx context.Context, tx pgstore.Tx, exportID string) error {
	if _, err := q.client.InsertTx(ctx, tx, ExportArgs{ExportID: exportID}, nil); err != nil {
		return fmt.Errorf("jobs: enqueue export %s: %w", exportID, 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) }