91 lines
3.2 KiB
Go
91 lines
3.2 KiB
Go
package runs
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log/slog"
|
|
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
"github.com/riverqueue/river"
|
|
"github.com/riverqueue/river/riverdriver/riverpgxv5"
|
|
|
|
"textmachine/platform/internal/pgstore"
|
|
)
|
|
|
|
// 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}
|
|
}
|
|
|
|
// spawnWorker runs SpawnArgs jobs.
|
|
type spawnWorker struct {
|
|
river.WorkerDefaults[SpawnArgs]
|
|
svc *Service
|
|
}
|
|
|
|
func (w *spawnWorker) Work(ctx context.Context, job *river.Job[SpawnArgs]) error {
|
|
return w.svc.Spawn(ctx, job.Args.RunID)
|
|
}
|
|
|
|
// Queue is the River client, wired to this service.
|
|
type Queue struct {
|
|
client *river.Client[pgstore.Tx]
|
|
}
|
|
|
|
// NewQueue builds the queue and its worker pool.
|
|
//
|
|
// Concurrency is deliberately small: a 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 NewQueue(pool *pgxpool.Pool, svc *Service, log *slog.Logger, workers int) (*Queue, error) {
|
|
if workers <= 0 {
|
|
workers = 4
|
|
}
|
|
w := river.NewWorkers()
|
|
if err := river.AddWorkerSafely(w, &spawnWorker{svc: svc}); err != nil {
|
|
return nil, fmt.Errorf("runs: register 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}},
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("runs: 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("runs: enqueue %s: %w", runID, 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) }
|