diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c2ccac1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +# Superpowers working artifacts — design specs and implementation plans. +# These are scratch for the development process, not part of the library. +docs/superpowers/ +.superpowers/ + +# ─── Local agent config ───────────────────────────────────── +.claude/ diff --git a/README.md b/README.md index 1294ba5..a12ea8e 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ Dispatch is a library — not a service. Import it, configure a store, and regis - **Extension hooks** — Opt-in lifecycle interfaces for every job, workflow, cron, and shutdown event - **OpenTelemetry** — Built-in metrics and tracing via the `observability` and `middleware` packages - **Relay integration** — Emit typed webhook events at every lifecycle point via `relay_hook` +- **Artifact plane** — Track gigabyte-scale job inputs and outputs in object storage, staged to a content-addressed local cache +- **Resource model** — Size jobs by what they actually need and admit them against detected worker capacity, instead of counting identical slots - **Pluggable storage** — Memory, PostgreSQL (pgx/v5), Grove ORM, SQLite, Redis ## Quick Start @@ -77,6 +79,61 @@ func main() { } ``` +## Resource-Aware Scheduling + +A worker slot is a promise that one job fits. That works while every job is the same size, and stops working the moment one definition serves both a 2 MB input and a 2 GB one: `concurrency: 4` says the worker may run four jobs, never how big they are, so a slot-counting pool starts four of the large ones on a box sized for the small ones and the kernel decides which of them dies. Adding a dedicated queue per size is the usual workaround, and it trades an OOM for a fleet that is idle in one queue and backed up in another. + +The `resource` package replaces the slot with a vector. A job declares what it needs — or computes it at enqueue from its input size — the worker detects what it has (cgroup-first, so a container reports its quota rather than the host's cores), and admission is a comparison. A job that does not fit stays pending instead of being started next to work already using the memory. + +Minimal wiring, in one process: + +```go +capacity := resource.Detect(resource.CapacityConfig{DiskBytes: 20 << 30}) +resources := resource.NewManager(capacity) + +// The SAME manager goes to both. The staging cache holds a lease per +// cached entry and reclaims disk on demand; the pool offers disk at +// dequeue as free plus what the cache could evict. Give the cache its +// own manager and that second half is always zero. +staging, _ := cache.New(cacheDir, backend, cache.WithManager(resources)) + +eng, _ := engine.Build(d, + engine.WithArtifacts(artifacts, staging), + engine.WithResourceManager(resources), +) +``` + +Sizing one definition from its input: + +```go +job.NewDefinition("render", handler, + job.WithArtifactInputs(artifact.Input("scene", artifact.Required)), + job.WithResourceFunc(func(_ context.Context, r resource.Request) (resource.Set, error) { + return resource.MemoryBytes(32<<20 + r.InputBytes*16), nil + }), +) +``` + +The function runs once, in the enqueuing process, and the result is written to the job row — nothing evaluates user code on the scheduling path, and two workers can never disagree about how big a job is. + +Under Forge, the same thing is configuration: + +```yaml +extensions: + dispatch: + resources: + enabled: true + cpu_overcommit: 1.0 # CPU is compressible; there is no memory equivalent + memory_fraction: 0.8 # leave the rest for the runtime and the page cache + explicit: # overrides detection, and the only way to declare + gpu: 4000 # a custom resource — nothing detects an FPGA + fpga: 2 +``` + +Leave it out and the runtime behaves as it did before: no ledger is built, the pool dequeues unbounded, every store backend skips its fit predicate, and the staging cache keeps the private disk budget it has always had. The schema is the exception — the resource columns are added by migration whether or not you turn any of this on, because a job row has to be readable by every worker in a mixed-version fleet. + +Runnable end to end in [`_examples/resources`](./_examples/resources). + ## Package Index | Package | Description | @@ -84,6 +141,8 @@ func main() { | `dispatch` | Root — `Dispatcher`, `Config`, options, errors, `Entity` base type | | `engine` | Wires all subsystems; `Build`, `Register`, `Enqueue`, `RegisterWorkflow`, `RegisterCron` | | `job` | `Job` entity, `State` machine, `Definition[T]`, `Registry` | +| `artifact` | Tracked object storage — declared inputs, imperative outputs, staging cache, lifecycle sweeping | +| `resource` | Resource vectors, capacity detection, the shared admission ledger and its reclaimers | | `workflow` | `Definition[T]`, `Run`, `State`, step checkpointing | | `cron` | `Entry`, `Scheduler`, distributed leader-elected cron | | `dlq` | `Entry`, `Service` — list, replay, purge | diff --git a/_examples/resources/main.go b/_examples/resources/main.go new file mode 100644 index 0000000..7df186f --- /dev/null +++ b/_examples/resources/main.go @@ -0,0 +1,398 @@ +// Package main demonstrates resource-aware admission: one job definition +// whose memory requirement scales with its input, enqueued twice against +// a worker that cannot hold both at once. +// +// This is the case that identical worker slots get wrong. "Concurrency 4" +// says a worker may run four jobs; it says nothing about how big they +// are. A render of a 2 MiB scene and a render of a 24 MiB scene are the +// same job to a slot-counting pool, so it starts both, and on a box sized +// for the small one the large one takes the whole machine down with it. +// +// With a resource model the two are different sizes of the same work. The +// requirement is computed once, at enqueue, from the declared input; the +// worker admits against what is actually free; and the large render waits +// for room instead of being started next to work that is already using +// it. +// +// Usage: +// +// go run ./_examples/resources +// +// Everything runs in-process: the memory store, the in-memory artifact +// backend Dispatch's own tests use, and a staging cache in a temp +// directory. No services to start, nothing to install. +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "sync/atomic" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/store/memory" +) + +const ( + mib = int64(1) << 20 + + // workerMemory is everything this worker will admit at once. Real + // deployments let resource.Detect read it from the cgroup; it is + // pinned here so the example prints the same story on every machine. + workerMemory = 512 * mib + + // stagingBudget is the staging cache's disk allowance, which with a + // shared ledger IS the worker's disk capacity. + stagingBudget = 128 * mib + + // A render holds its scene in memory several times over — decoded + // geometry, working buffers, the framebuffer. baseMemory is the + // interpreter and the runtime; the rest scales with the input. + baseMemory = 32 * mib + memoryPerInputB = 16 + + smallScene = 2 * mib + largeScene = 24 * mib + + renderTime = 900 * time.Millisecond + notifyTime = 700 * time.Millisecond +) + +// scene is the job payload. The bytes it refers to are declared as an +// artifact input, not carried here, which is what lets the engine size +// the job before scheduling it. +type scene struct { + Name string `json:"name"` +} + +func main() { + if err := run(); err != nil { + fmt.Fprintf(os.Stderr, "error: %v\n", err) + os.Exit(1) + } +} + +func run() error { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + clock := newTimeline() + + // ────────────────────────────────────────────────── + // 1. One ledger, shared + // ────────────────────────────────────────────────── + // + // This is the wiring that matters. The staging cache and the worker + // pool must be given the SAME resource.Manager: the cache holds a + // lease per cached entry and registers itself as the ledger's disk + // reclaimer, and the pool offers disk at dequeue as free PLUS what + // that reclaimer could evict. Hand the cache its own manager — which + // it builds for itself if you do not supply one — and the pool's view + // of reclaimable disk is permanently zero. + + capacity := resource.Detect(resource.CapacityConfig{ + DiskBytes: stagingBudget, + Explicit: resource.Set{ + resource.CPU: 4 * resource.MilliScale, + resource.Memory: workerMemory, + }, + }) + + resources := resource.NewManager(capacity) + + store := memory.New() + + d, err := dispatch.New( + dispatch.WithStore(store), + dispatch.WithLogger(log.NewNoopLogger()), + // Four slots, deliberately more than this box can afford to fill. + // The point of the example is that the limit is memory, not slots. + dispatch.WithConcurrency(4), + dispatch.WithPollInterval(200*time.Millisecond), + dispatch.WithHeartbeatInterval(time.Second), + dispatch.WithStaleJobThreshold(30*time.Second), + ) + if err != nil { + return fmt.Errorf("create dispatcher: %w", err) + } + + backend := artifacttest.NewBackend() + + artifacts := artifact.NewService(store, backend, + artifact.WithDefaultBucket("scenes")) + + staging, err := cache.New(mustTempDir(), backend, cache.WithManager(resources)) + if err != nil { + return fmt.Errorf("create staging cache: %w", err) + } + + eng, err := engine.Build(d, + engine.WithArtifacts(artifacts, staging), + engine.WithResourceManager(resources), + ) + if err != nil { + return fmt.Errorf("build engine: %w", err) + } + + // ────────────────────────────────────────────────── + // 2. One definition, sized from its input + // ────────────────────────────────────────────────── + + var running atomic.Int64 + + engine.Register(eng, job.NewDefinition("render", + func(ctx context.Context, p scene) error { + clock.log("render %-5s START (%d running, %s admitted)", + p.Name, running.Add(1), mb(held(resources))) + defer running.Add(-1) + + sleep(ctx, renderTime) + + clock.log("render %-5s done", p.Name) + + return nil + }, + job.WithArtifactInputs(artifact.Input("scene", artifact.Required)), + + // The requirement is computed once, in the enqueuing process, + // from the size of the declared input. It never runs on the + // scheduling path, so two workers can never disagree about how + // big this job is. + job.WithResourceFunc(func(_ context.Context, r resource.Request) (resource.Set, error) { + return resource.MemoryBytes(baseMemory + r.InputBytes*memoryPerInputB), nil + }), + job.WithMaxRetries(0), + )) + + // Unrelated work already on the box. It is what the small render runs + // alongside, and what the large render has to wait for. + engine.Register(eng, job.NewDefinition("notify", + func(ctx context.Context, p scene) error { + clock.log("notify %-5s START (%d running, %s admitted)", + p.Name, running.Add(1), mb(held(resources))) + defer running.Add(-1) + + sleep(ctx, notifyTime) + + clock.log("notify %-5s done", p.Name) + + return nil + }, + job.WithResources(resource.MemoryBytes(32*mib)), + job.WithMaxRetries(0), + )) + + // ────────────────────────────────────────────────── + // 3. Two renders, one small and one large + // ────────────────────────────────────────────────── + + small, err := upload(ctx, artifacts, backend, "small.scene", smallScene) + if err != nil { + return err + } + + large, err := upload(ctx, artifacts, backend, "large.scene", largeScene) + if err != nil { + return err + } + + fmt.Printf("worker capacity : %s memory, %s staging disk, %d slots\n", + mb(capacity[resource.Memory]), mb(capacity[resource.Disk]), 4) + fmt.Printf("render(small) : %s input -> %s memory\n", + mb(smallScene), mb(baseMemory+smallScene*memoryPerInputB)) + fmt.Printf("render(large) : %s input -> %s memory\n", + mb(largeScene), mb(baseMemory+largeScene*memoryPerInputB)) + fmt.Printf("notify : %s memory each\n\n", mb(32*mib)) + + // Enqueue order is the order the store hands them back at equal + // priority, so the small render and the two notifies reach the worker + // first and the large render arrives to find the box already busy. + enqueued := []struct { + name string + payload scene + opts []job.Option + }{ + {"render", scene{Name: "small"}, []job.Option{engine.Bind("scene", small)}}, + {"notify", scene{Name: "a"}, nil}, + {"notify", scene{Name: "b"}, nil}, + {"render", scene{Name: "large"}, []job.Option{engine.Bind("scene", large)}}, + } + + for _, e := range enqueued { + j, eerr := engine.Enqueue(ctx, eng, e.name, e.payload, e.opts...) + if eerr != nil { + return fmt.Errorf("enqueue %s: %w", e.name, eerr) + } + + fmt.Printf("enqueued %-6s %-5s requiring %s\n", + e.name, e.payload.Name, mb(j.Resources[resource.Memory])) + } + + fmt.Println() + + // ────────────────────────────────────────────────── + // 4. Watch the ordering + // ────────────────────────────────────────────────── + + // Reset BEFORE Start. Every goroutine that reads the timeline is + // created by Start or after it, so the write is ordered ahead of all + // of them; resetting afterwards would race the workers already + // logging against it. + clock.reset() + + if serr := eng.Start(ctx); serr != nil { + return fmt.Errorf("start engine: %w", serr) + } + + // Sample the queue while the first three are running. This is the + // observation the example exists to make: a job sitting in pending + // with a worker slot free next to it. + go func() { + sleep(ctx, 400*time.Millisecond) + reportPending(ctx, clock, store, resources) + }() + + if werr := waitForDrain(ctx, store); werr != nil { + return werr + } + + stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer stopCancel() + + if serr := eng.Stop(stopCtx); serr != nil { + return fmt.Errorf("stop engine: %w", serr) + } + + fmt.Println() + fmt.Println("The large render was claimed in the first poll and refused: 416 MiB") + fmt.Println("did not fit the 384 MiB left once the small render and both notifies") + fmt.Println("were admitted. A fourth worker slot was free the whole time — slots") + fmt.Println("were never the constraint. It ran once the box could hold it.") + + return nil +} + +// upload seeds an object and registers it, returning the ref that carries +// its size. The size is what the resource func reads at enqueue. +func upload( + ctx context.Context, svc *artifact.Service, backend *artifacttest.Backend, + key string, size int64, +) (artifact.Ref, error) { + backend.Put("scenes", key, make([]byte, size)) + + ref, err := svc.Register(ctx, "scenes", key) + if err != nil { + return artifact.Ref{}, fmt.Errorf("register %s: %w", key, err) + } + + return ref, nil +} + +// reportPending prints what is waiting and why, mid-run. +func reportPending( + ctx context.Context, clock *timeline, store *memory.Store, m resource.Manager, +) { + jobs, err := store.ListJobsByState(ctx, job.StatePending, job.ListOpts{Limit: 100}) + if err != nil { + return + } + + free := m.Free()[resource.Memory] + + for _, j := range jobs { + clock.log("%-6s %-5s PENDING — needs %s, %s free", + j.Name, nameOf(j.Payload), mb(j.Resources[resource.Memory]), mb(free)) + } +} + +// nameOf pulls the payload's name field for display. +func nameOf(payload []byte) string { + var s scene + if err := json.Unmarshal(payload, &s); err != nil { + return "?" + } + + return s.Name +} + +// held is the memory currently spoken for by admitted jobs. +func held(m resource.Manager) int64 { + return m.Capacity()[resource.Memory] - m.Free()[resource.Memory] +} + +// waitForDrain blocks until no job is left pending, retrying, or running. +func waitForDrain(ctx context.Context, store *memory.Store) error { + ticker := time.NewTicker(50 * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return fmt.Errorf("jobs did not drain: %w", ctx.Err()) + case <-ticker.C: + } + + var outstanding int + + for _, state := range []job.State{job.StatePending, job.StateRetrying, job.StateRunning} { + jobs, err := store.ListJobsByState(ctx, state, job.ListOpts{Limit: 100}) + if err != nil { + return fmt.Errorf("list %s jobs: %w", state, err) + } + + outstanding += len(jobs) + } + + if outstanding == 0 { + return nil + } + } +} + +// sleep is a context-aware pause, so shutdown is not held up by a +// simulated render. +func sleep(ctx context.Context, d time.Duration) { + t := time.NewTimer(d) + defer t.Stop() + + select { + case <-ctx.Done(): + case <-t.C: + } +} + +// timeline prints millisecond offsets from a resettable origin, which is +// what makes the ordering legible. +type timeline struct{ start time.Time } + +func newTimeline() *timeline { return &timeline{start: time.Now()} } + +func (t *timeline) reset() { t.start = time.Now() } + +func (t *timeline) log(format string, args ...any) { + fmt.Printf("%6dms %s\n", + time.Since(t.start).Milliseconds(), fmt.Sprintf(format, args...)) +} + +func mb(bytes int64) string { + return fmt.Sprintf("%d MiB", bytes/mib) +} + +func mustTempDir() string { + dir, err := os.MkdirTemp("", "dispatch-resources-example") + if err != nil { + panic(err) + } + + return dir +} diff --git a/_project_files/dispatch-design-v3.md b/_project_files/dispatch-design-v3.md new file mode 100644 index 0000000..b9a2169 --- /dev/null +++ b/_project_files/dispatch-design-v3.md @@ -0,0 +1,2586 @@ +# Dispatch — Project Design & Development Phases (v3) + +> Composable, extensible durable execution engine for Go. Library-first background jobs, workflow orchestration, lifecycle hooks, and distributed workers. +> +> `github.com/xraph/dispatch` + +--- + +## Table of Contents + +1. [Architecture Overview](#1-architecture-overview) +2. [Core Design Principles](#2-core-design-principles) +3. [TypeID Identity System](#3-typeid-identity-system) +4. [Store Architecture (ControlPlane Pattern)](#4-store-architecture) +5. [Module Layout & Package Design](#5-module-layout) +6. [Core Types & Interfaces](#6-core-types) +7. [Extension System](#7-extension-system) +8. [Hooks & Relay Integration](#8-hooks--relay-integration) +9. [Distributed Workers & Kubernetes Consensus](#9-distributed-workers) +10. [Integration Platform Pattern (Zapier-Style)](#10-integration-platform-pattern) +11. [Forge Scope Integration](#11-forge-scope-integration) +12. [Linting & Code Quality (golangci-lint v2)](#12-linting) +13. [Development Phases](#13-development-phases) +14. [Claude Code Development Guide](#14-claude-code-guide) + +--- + +## 1. Architecture Overview + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ Consumer Application │ +│ (standalone binary, forge app, or K8s deployment) │ +└──────────────────────────────┬──────────────────────────────────────┘ + │ +┌──────────────────────────────▼──────────────────────────────────────┐ +│ dispatch.Dispatcher │ +│ ┌─────────────┐ ┌──────────┐ ┌───────────┐ ┌───────────────────┐ │ +│ │ Job Svc │ │Workflow │ │ Cron Svc │ │ DLQ Service │ │ +│ │ │ │ Svc │ │ │ │ │ │ +│ └──────┬──────┘ └────┬─────┘ └─────┬─────┘ └────────┬──────────┘ │ +│ │ │ │ │ │ +│ ┌──────▼─────────────▼─────────────▼─────────────────▼──────────┐ │ +│ │ Extension Registry │ │ +│ │ [Custom Ext] [Custom Ext] [Relay Hook Ext] [Metrics Ext] │ │ +│ └──────────────────────────────┬────────────────────────────────┘ │ +│ │ │ +│ ┌──────────────────────────────▼────────────────────────────────┐ │ +│ │ Hooks (Lifecycle Events) │ │ +│ │ OnEnqueue → OnStart → OnComplete/OnFail → OnRetry → OnDLQ │ │ +│ │ ↓ (optional) │ │ +│ │ ┌─────────────────────┐ │ │ +│ │ │ Relay Webhook │ → Customer endpoint │ │ +│ │ │ delivery via │ → Monitoring system │ │ +│ │ │ relay.Send() │ → Slack/PagerDuty │ │ +│ │ └─────────────────────┘ │ │ +│ └──────────────────────────────┬────────────────────────────────┘ │ +│ │ │ +│ ┌──────────────────────────────▼────────────────────────────────┐ │ +│ │ Worker Pool + Executor │ │ +│ │ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────────┐ │ │ +│ │ │ Local │ │ K8s Pod │ │ K8s Pod │ │ K8s Pod │ │ │ +│ │ │ Workers │ │ Worker A │ │ Worker B │ │ Worker C │ │ │ +│ │ └──────────┘ └──────────┘ └──────────┘ └──────────────┘ │ │ +│ └──────────────────────────────┬────────────────────────────────┘ │ +│ │ │ +│ ┌──────────────────────────────▼────────────────────────────────┐ │ +│ │ Middleware Chain │ │ +│ │ [Tracing] → [Metrics] → [Logging] → [Timeout] → [Scope] │ │ +│ └──────────────────────────────┬────────────────────────────────┘ │ +│ │ │ +│ ┌──────────────────────────────▼────────────────────────────────┐ │ +│ │ Queue Abstraction │ │ +│ │ ┌──────────┐ ┌───────────┐ ┌────────────┐ │ │ +│ │ │ Priority │ │Rate Limit │ │ Per-Tenant │ │ │ +│ │ └──────────┘ └───────────┘ └────────────┘ │ │ +│ └──────────────────────────────┬────────────────────────────────┘ │ +│ │ │ +│ ┌──────────────────────────────▼────────────────────────────────┐ │ +│ │ Cluster / Consensus Layer │ │ +│ │ ┌──────────────┐ ┌──────────────┐ ┌───────────────────┐ │ │ +│ │ │ Leader │ │ Worker │ │ Work Stealing │ │ │ +│ │ │ Election │ │ Registry │ │ / Rebalancing │ │ │ +│ │ └──────────────┘ └──────────────┘ └───────────────────┘ │ │ +│ └──────────────────────────────┬────────────────────────────────┘ │ +│ │ │ +│ ┌──────────────────────────────▼────────────────────────────────┐ │ +│ │ Store Interface │ │ +│ │ ┌──────────┐ ┌───────┐ ┌───────┐ ┌────────┐ ┌──────┐ │ │ +│ │ │ Postgres │ │ Bun │ │SQLite │ │ Redis │ │Memory│ │ │ +│ │ │ (pgx) │ │ (ORM) │ │ │ │ │ │ │ │ │ +│ │ └──────────┘ └───────┘ └───────┘ └────────┘ └──────┘ │ │ +│ └───────────────────────────────────────────────────────────────┘ │ +│ │ +│ ┌───────────────────┐ ┌──────────────────────────────────────┐ │ +│ │ Event Bus │ │ forge.Extension (optional mount) │ │ +│ │ (WaitForEvent) │ │ │ │ +│ └───────────────────┘ └──────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## 2. Core Design Principles + +**Library, not service.** Import it. No cluster to run, no separate process. Just Go. + +**Workflows as Go functions.** No DSL, no YAML, no protobuf. Define workflows with ordinary Go functions and get durable execution with per-step checkpointing. + +**Composable store pattern (ControlPlane-style).** Each subsystem defines its own store interface. The aggregate `store.Store` composes them all. Five backends: Postgres (pgx), Bun (ORM), SQLite, Redis, and memory — matching and extending ControlPlane's store lineup. + +**TypeID everywhere (ControlPlane pattern).** All entity IDs use `go.jetify.com/typeid` — type-prefixed, K-sortable, UUIDv7-based, compile-time safe. `job_01h...`, `wfrun_01h...`, `cron_01h...`. Same pattern Nexus and ControlPlane use. + +**Distributed-ready from day one.** Single-process by default. Add Kubernetes consensus, worker registration, and work stealing without changing your job code. Scale from one goroutine to a fleet of pods. + +**Forge-native, standalone-capable.** Reads `forge.Scope` from context when available. Falls back gracefully for standalone usage. Mounts as a Forge extension via `dispatch_ext.New()`. + +**Middleware-driven.** Every cross-cutting concern (tracing, metrics, logging, timeout, tenant isolation) is a composable middleware wrapping job execution. Same mental model as HTTP middleware. + +**Extensible by design.** A first-class extension registry lets users hook into every lifecycle event — enqueue, start, complete, fail, retry, DLQ. Extensions are the building block for integrations. Relay ships as a built-in extension for webhook delivery. + +**Relay-native hooks.** When Relay is available, Dispatch automatically emits typed webhook events at every lifecycle point. Customers subscribe to `dispatch.job.completed`, `dispatch.workflow.failed`, etc. via Relay's endpoint management. No custom webhook code needed. + +--- + +## 3. TypeID Identity System + +Dispatch uses the same TypeID pattern as ControlPlane and Nexus — every entity gets a compile-time safe, Stripe-style prefixed ID. + +### `id/id.go` + +```go +package id + +import "go.jetify.com/typeid" + +// ────────────────────────────────────────────────── +// Prefix types — each entity has its own prefix +// ────────────────────────────────────────────────── + +type JobPrefix struct{} +func (JobPrefix) Prefix() string { return "job" } + +type WorkflowPrefix struct{} +func (WorkflowPrefix) Prefix() string { return "wf" } + +type RunPrefix struct{} +func (RunPrefix) Prefix() string { return "wfrun" } + +type CheckpointPrefix struct{} +func (CheckpointPrefix) Prefix() string { return "ckpt" } + +type CronPrefix struct{} +func (CronPrefix) Prefix() string { return "cron" } + +type DLQPrefix struct{} +func (DLQPrefix) Prefix() string { return "dlq" } + +type EventPrefix struct{} +func (EventPrefix) Prefix() string { return "evt" } + +type WorkerPrefix struct{} +func (WorkerPrefix) Prefix() string { return "wkr" } + +// ────────────────────────────────────────────────── +// Typed ID aliases — compile-time safe +// ────────────────────────────────────────────────── + +type JobID = typeid.TypeID[JobPrefix] // job_01h2xcejqtf2nbrexx3vqjhp41 +type WorkflowID = typeid.TypeID[WorkflowPrefix] // wf_01h2xcejqtf2nbrexx3vqjhp41 +type RunID = typeid.TypeID[RunPrefix] // wfrun_01h2xcejqtf2nbrexx3vqjhp41 +type CheckpointID = typeid.TypeID[CheckpointPrefix] // ckpt_01h2xcejqtf2nbrexx3vqjhp41 +type CronID = typeid.TypeID[CronPrefix] // cron_01h455vb4pex5vsknk084sn02q +type DLQID = typeid.TypeID[DLQPrefix] // dlq_01h6rz1g6p2m3q9xvz1t2b7c4d +type EventID = typeid.TypeID[EventPrefix] // evt_01h8f3k2n7p4r6s9t1v3x5z7b9 +type WorkerID = typeid.TypeID[WorkerPrefix] // wkr_01h9a1b2c3d4e5f6g7h8j9k0m1 + +// AnyID for cases where the prefix is dynamic. +type AnyID = typeid.AnyID + +// ────────────────────────────────────────────────── +// Constructors +// ────────────────────────────────────────────────── + +func NewJobID() JobID { return must(typeid.New[JobID]()) } +func NewRunID() RunID { return must(typeid.New[RunID]()) } +func NewCheckpointID() CheckpointID { return must(typeid.New[CheckpointID]()) } +func NewCronID() CronID { return must(typeid.New[CronID]()) } +func NewDLQID() DLQID { return must(typeid.New[DLQID]()) } +func NewEventID() EventID { return must(typeid.New[EventID]()) } +func NewWorkerID() WorkerID { return must(typeid.New[WorkerID]()) } + +// ────────────────────────────────────────────────── +// Parsing (type-safe: ParseJobID("cron_01h...") fails) +// ────────────────────────────────────────────────── + +func ParseJobID(s string) (JobID, error) { return typeid.Parse[JobID](s) } +func ParseRunID(s string) (RunID, error) { return typeid.Parse[RunID](s) } +func ParseCronID(s string) (CronID, error) { return typeid.Parse[CronID](s) } +func ParseDLQID(s string) (DLQID, error) { return typeid.Parse[DLQID](s) } +func ParseEventID(s string) (EventID, error) { return typeid.Parse[EventID](s) } +func ParseWorkerID(s string) (WorkerID, error) { return typeid.Parse[WorkerID](s) } +func ParseAny(s string) (AnyID, error) { return typeid.FromString(s) } + +// ────────────────────────────────────────────────── +// Helpers +// ────────────────────────────────────────────────── + +func must[T any](v T, err error) T { + if err != nil { + panic(err) + } + return v +} +``` + +### Base Entity + +All Dispatch entities embed a common base, same as ControlPlane: + +```go +// entity.go +package dispatch + +import ( + "time" +) + +// Entity is the base type embedded by all dispatch domain objects. +type Entity struct { + CreatedAt time.Time `json:"created_at" bun:"created_at,notnull,default:current_timestamp"` + UpdatedAt time.Time `json:"updated_at" bun:"updated_at,notnull,default:current_timestamp"` +} + +// NewEntity returns an Entity with timestamps set to now. +func NewEntity() Entity { + now := time.Now().UTC() + return Entity{CreatedAt: now, UpdatedAt: now} +} +``` + +### Why TypeID + +| Property | TypeID | UUID v4 | XID | +|----------|--------|---------|-----| +| Type-safe prefix | ✅ `job_...` | ❌ | ❌ | +| K-Sortable | ✅ (UUIDv7) | ❌ | ✅ | +| Compile-time safety | ✅ (generics) | ❌ | ❌ | +| Human-debuggable | ✅ prefix tells type | ❌ | ❌ | +| DB compatible | ✅ stores as TEXT | ✅ | ✅ | +| Stripe-style | ✅ | ❌ | ❌ | + +--- + +## 4. Store Architecture (ControlPlane Pattern) + +Dispatch follows the exact same composite store pattern established in ControlPlane. Each subsystem defines its own store interface in its own package. The top-level `store.Store` composes all of them. **Five backends** — matching ControlPlane's lineup plus Bun: + +| Backend | Driver | Use Case | +|---------|--------|----------| +| **Postgres** | `pgx/v5` | Production. Raw SQL, SKIP LOCKED, LISTEN/NOTIFY, advisory locks | +| **Bun** | `uptrace/bun` | Production. ORM-based, Bun model tags, migration integration. For teams already using Bun in their Forge/ControlPlane stack | +| **SQLite** | `modernc.org/sqlite` | Embedded/edge. Single-file DB for CLI tools, dev, and standalone apps | +| **Redis** | `redis/go-redis/v9` | High-throughput ephemeral workloads. Speed over durability | +| **Memory** | In-process maps | Unit tests. Zero dependencies | + +### Subsystem Store Interfaces + +Each domain package defines a focused store contract: + +```go +// job/store.go +package job + +import "github.com/xraph/dispatch/id" + +type Store interface { + Enqueue(ctx context.Context, j *Job) error + Dequeue(ctx context.Context, queues []string, limit int) ([]*Job, error) + Get(ctx context.Context, jobID id.JobID) (*Job, error) + Update(ctx context.Context, j *Job) error + Delete(ctx context.Context, jobID id.JobID) error + ListByState(ctx context.Context, state State, opts ListOpts) ([]*Job, error) + Heartbeat(ctx context.Context, jobID id.JobID, workerID id.WorkerID) error + ReapStale(ctx context.Context, threshold time.Duration) ([]*Job, error) + Count(ctx context.Context, opts CountOpts) (int64, error) +} + +// workflow/store.go +package workflow + +import "github.com/xraph/dispatch/id" + +type Store interface { + CreateRun(ctx context.Context, run *Run) error + GetRun(ctx context.Context, runID id.RunID) (*Run, error) + UpdateRun(ctx context.Context, run *Run) error + ListRuns(ctx context.Context, opts ListOpts) ([]*Run, error) + SaveCheckpoint(ctx context.Context, runID id.RunID, stepName string, data []byte) error + GetCheckpoint(ctx context.Context, runID id.RunID, stepName string) ([]byte, error) + ListCheckpoints(ctx context.Context, runID id.RunID) ([]*Checkpoint, error) +} + +// cron/store.go +package cron + +import "github.com/xraph/dispatch/id" + +type Store interface { + Register(ctx context.Context, entry *Entry) error + Get(ctx context.Context, entryID id.CronID) (*Entry, error) + List(ctx context.Context) ([]*Entry, error) + AcquireLock(ctx context.Context, entryID id.CronID, workerID id.WorkerID, ttl time.Duration) (bool, error) + ReleaseLock(ctx context.Context, entryID id.CronID, workerID id.WorkerID) error + UpdateLastRun(ctx context.Context, entryID id.CronID, at time.Time) error + Delete(ctx context.Context, entryID id.CronID) error +} + +// dlq/store.go +package dlq + +import "github.com/xraph/dispatch/id" + +type Store interface { + Push(ctx context.Context, entry *Entry) error + List(ctx context.Context, opts ListOpts) ([]*Entry, error) + Get(ctx context.Context, entryID id.DLQID) (*Entry, error) + Replay(ctx context.Context, entryID id.DLQID) error + Purge(ctx context.Context, before time.Time) (int64, error) + Count(ctx context.Context) (int64, error) +} + +// event/store.go +package event + +import "github.com/xraph/dispatch/id" + +type Store interface { + Publish(ctx context.Context, evt *Event) error + Subscribe(ctx context.Context, name string, timeout time.Duration) (*Event, error) + Ack(ctx context.Context, eventID id.EventID) error +} + +// cluster/store.go +package cluster + +import "github.com/xraph/dispatch/id" + +type Store interface { + RegisterWorker(ctx context.Context, w *Worker) error + DeregisterWorker(ctx context.Context, workerID id.WorkerID) error + Heartbeat(ctx context.Context, workerID id.WorkerID) error + ListWorkers(ctx context.Context) ([]*Worker, error) + ReapDead(ctx context.Context, threshold time.Duration) ([]*Worker, error) + AcquireLeadership(ctx context.Context, workerID id.WorkerID, ttl time.Duration) (bool, error) + RenewLeadership(ctx context.Context, workerID id.WorkerID, ttl time.Duration) (bool, error) + GetLeader(ctx context.Context) (*Worker, error) +} +``` + +### Composite Store + +```go +// store/store.go +package store + +import ( + "context" + + "github.com/xraph/dispatch/cluster" + "github.com/xraph/dispatch/cron" + "github.com/xraph/dispatch/dlq" + "github.com/xraph/dispatch/event" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/workflow" +) + +// Store is the aggregate persistence interface. +// Each subsystem store is a composable interface — same pattern as ControlPlane. +// A single backend (postgres, bun, sqlite, etc.) implements all of them. +type Store interface { + job.Store + workflow.Store + cron.Store + dlq.Store + event.Store + cluster.Store + + // Migrate runs all schema migrations. + Migrate(ctx context.Context) error + + // Ping checks database connectivity. + Ping(ctx context.Context) error + + // Close closes the store connection. + Close() error +} +``` + +### Backend Directory Structure + +``` +store/ +├── store.go # Composite Store interface +├── postgres/ +│ ├── store.go # *PGStore: pgx/v5, pgxpool, embed.FS migrations +│ ├── job.go # job.Store — SELECT FOR UPDATE SKIP LOCKED +│ ├── workflow.go # workflow.Store +│ ├── cron.go # cron.Store — pg_advisory_lock for leader election +│ ├── dlq.go # dlq.Store +│ ├── event.go # event.Store — LISTEN/NOTIFY for WaitForEvent +│ ├── cluster.go # cluster.Store — advisory locks for leadership +│ └── migrations/ +│ ├── 001_jobs.sql +│ ├── 002_workflows.sql +│ ├── 003_cron.sql +│ ├── 004_dlq.sql +│ ├── 005_events.sql +│ └── 006_cluster.sql +├── bun/ +│ ├── store.go # *BunStore: uptrace/bun, model-driven, Bun migrations +│ ├── models.go # Bun model structs with bun:"" tags +│ ├── job.go # job.Store via Bun query builder +│ ├── workflow.go # workflow.Store via Bun +│ ├── cron.go # cron.Store via Bun +│ ├── dlq.go # dlq.Store via Bun +│ ├── event.go # event.Store via Bun +│ ├── cluster.go # cluster.Store via Bun +│ └── migrations/ +│ └── 001_initial.go # Bun Go-based migrations +├── sqlite/ +│ ├── store.go # Embeds migrations, single-file DB +│ └── migrations/ +│ └── 001_initial.sql +├── redis/ +│ ├── store.go # Redis Streams + Sorted Sets +│ ├── job.go # XREADGROUP for dequeue +│ ├── scripts/ # Lua scripts for atomic operations +│ └── cluster.go # Redis-based leader election (Redlock) +└── memory/ + └── store.go # sync.Map + channels, for unit tests +``` + +### Bun Store — Models + +The Bun store uses Bun ORM model structs that map directly to Dispatch entities. This is the store to use when your app already uses Bun (as Forge and ControlPlane do): + +```go +// store/bun/models.go +package bun + +import ( + "time" + + "github.com/uptrace/bun" + "github.com/xraph/dispatch/id" +) + +// JobModel is the Bun model for the dispatch_jobs table. +type JobModel struct { + bun.BaseModel `bun:"table:dispatch_jobs,alias:j"` + + ID string `bun:"id,pk"` + Name string `bun:"name,notnull"` + Queue string `bun:"queue,notnull,default:'default'"` + Payload []byte `bun:"payload,notnull,type:bytea"` + State string `bun:"state,notnull,default:'pending'"` + Priority int `bun:"priority,notnull,default:0"` + MaxRetries int `bun:"max_retries,notnull,default:3"` + RetryCount int `bun:"retry_count,notnull,default:0"` + LastError string `bun:"last_error"` + ScopeAppID string `bun:"scope_app_id"` + ScopeOrgID string `bun:"scope_org_id"` + WorkerID string `bun:"worker_id"` + RunAt time.Time `bun:"run_at,notnull,default:current_timestamp"` + StartedAt *time.Time `bun:"started_at"` + CompletedAt *time.Time `bun:"completed_at"` + HeartbeatAt *time.Time `bun:"heartbeat_at"` + CreatedAt time.Time `bun:"created_at,notnull,default:current_timestamp"` + UpdatedAt time.Time `bun:"updated_at,notnull,default:current_timestamp"` +} + +// WorkflowRunModel is the Bun model for the dispatch_workflow_runs table. +type WorkflowRunModel struct { + bun.BaseModel `bun:"table:dispatch_workflow_runs,alias:wr"` + + ID string `bun:"id,pk"` + Name string `bun:"name,notnull"` + State string `bun:"state,notnull,default:'running'"` + Input []byte `bun:"input,type:bytea"` + Output []byte `bun:"output,type:bytea"` + Error string `bun:"error"` + ScopeAppID string `bun:"scope_app_id"` + ScopeOrgID string `bun:"scope_org_id"` + StartedAt time.Time `bun:"started_at,notnull,default:current_timestamp"` + CompletedAt *time.Time `bun:"completed_at"` + CreatedAt time.Time `bun:"created_at,notnull,default:current_timestamp"` + UpdatedAt time.Time `bun:"updated_at,notnull,default:current_timestamp"` +} + +// WorkerModel is the Bun model for the dispatch_workers table. +type WorkerModel struct { + bun.BaseModel `bun:"table:dispatch_workers,alias:w"` + + ID string `bun:"id,pk"` + Hostname string `bun:"hostname,notnull"` + Queues []string `bun:"queues,array"` + Concurrency int `bun:"concurrency,notnull,default:10"` + State string `bun:"state,notnull,default:'active'"` + IsLeader bool `bun:"is_leader,notnull,default:false"` + LeaderUntil *time.Time `bun:"leader_until"` + LastSeen time.Time `bun:"last_seen,notnull,default:current_timestamp"` + Metadata map[string]string `bun:"metadata,type:jsonb,default:'{}'"` + CreatedAt time.Time `bun:"created_at,notnull,default:current_timestamp"` +} +``` + +### Bun Store — Job Implementation Example + +```go +// store/bun/job.go +package bun + +import ( + "context" + "database/sql" + + "github.com/uptrace/bun" + "github.com/xraph/dispatch/job" +) + +func (s *BunStore) Dequeue(ctx context.Context, queues []string, limit int) ([]*job.Job, error) { + var models []JobModel + + // Bun's raw query for SKIP LOCKED (same strategy as raw postgres store) + err := s.db.NewRaw(` + UPDATE dispatch_jobs + SET state = 'running', started_at = NOW(), updated_at = NOW(), worker_id = ? + WHERE id IN ( + SELECT id FROM dispatch_jobs + WHERE state = 'pending' + AND queue IN (?) + AND run_at <= NOW() + ORDER BY priority DESC, run_at ASC + FOR UPDATE SKIP LOCKED + LIMIT ? + ) + RETURNING * + `, s.workerID, bun.In(queues), limit).Scan(ctx, &models) + + if err != nil { + return nil, err + } + return modelsToJobs(models), nil +} +``` + +### PostgreSQL Migrations + +```sql +-- store/postgres/migrations/001_jobs.sql +CREATE TABLE IF NOT EXISTS dispatch_jobs ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + queue TEXT NOT NULL DEFAULT 'default', + payload BYTEA NOT NULL, + state TEXT NOT NULL DEFAULT 'pending', + priority INTEGER NOT NULL DEFAULT 0, + max_retries INTEGER NOT NULL DEFAULT 3, + retry_count INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + scope_app_id TEXT, + scope_org_id TEXT, + worker_id TEXT, + run_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + started_at TIMESTAMPTZ, + completed_at TIMESTAMPTZ, + heartbeat_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_dispatch_jobs_dequeue + ON dispatch_jobs (queue, priority DESC, run_at ASC) + WHERE state = 'pending'; +CREATE INDEX idx_dispatch_jobs_state ON dispatch_jobs (state); +CREATE INDEX idx_dispatch_jobs_scope ON dispatch_jobs (scope_app_id, scope_org_id); +CREATE INDEX idx_dispatch_jobs_heartbeat ON dispatch_jobs (heartbeat_at) + WHERE state = 'running'; +``` + +```sql +-- store/postgres/migrations/002_workflows.sql +CREATE TABLE IF NOT EXISTS dispatch_workflow_runs ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + state TEXT NOT NULL DEFAULT 'running', + input BYTEA, + output BYTEA, + error TEXT, + scope_app_id TEXT, + scope_org_id TEXT, + started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + completed_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE TABLE IF NOT EXISTS dispatch_checkpoints ( + id TEXT PRIMARY KEY, + run_id TEXT NOT NULL REFERENCES dispatch_workflow_runs(id) ON DELETE CASCADE, + step_name TEXT NOT NULL, + data BYTEA NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + UNIQUE(run_id, step_name) +); + +CREATE INDEX idx_dispatch_workflow_runs_state ON dispatch_workflow_runs (state); +CREATE INDEX idx_dispatch_checkpoints_run ON dispatch_checkpoints (run_id); +``` + +```sql +-- store/postgres/migrations/003_cron.sql +CREATE TABLE IF NOT EXISTS dispatch_cron_entries ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL UNIQUE, + schedule TEXT NOT NULL, + job_name TEXT NOT NULL, + payload BYTEA, + scope_app_id TEXT, + scope_org_id TEXT, + last_run_at TIMESTAMPTZ, + next_run_at TIMESTAMPTZ, + locked_by TEXT, + locked_until TIMESTAMPTZ, + enabled BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_dispatch_cron_next ON dispatch_cron_entries (next_run_at) + WHERE enabled = TRUE; +``` + +```sql +-- store/postgres/migrations/004_dlq.sql +CREATE TABLE IF NOT EXISTS dispatch_dlq ( + id TEXT PRIMARY KEY, + job_id TEXT NOT NULL, + job_name TEXT NOT NULL, + queue TEXT NOT NULL, + payload BYTEA NOT NULL, + error TEXT NOT NULL, + retry_count INTEGER NOT NULL, + scope_app_id TEXT, + scope_org_id TEXT, + failed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + replayed_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_dispatch_dlq_queue ON dispatch_dlq (queue, failed_at DESC); +``` + +```sql +-- store/postgres/migrations/005_events.sql +CREATE TABLE IF NOT EXISTS dispatch_events ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + payload BYTEA, + scope_app_id TEXT, + scope_org_id TEXT, + acked BOOLEAN NOT NULL DEFAULT FALSE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_dispatch_events_pending ON dispatch_events (name, created_at) + WHERE acked = FALSE; +``` + +```sql +-- store/postgres/migrations/006_cluster.sql +CREATE TABLE IF NOT EXISTS dispatch_workers ( + id TEXT PRIMARY KEY, + hostname TEXT NOT NULL, + queues TEXT[] DEFAULT '{}', + concurrency INTEGER NOT NULL DEFAULT 10, + state TEXT NOT NULL DEFAULT 'active', + is_leader BOOLEAN NOT NULL DEFAULT FALSE, + leader_until TIMESTAMPTZ, + last_seen TIMESTAMPTZ NOT NULL DEFAULT NOW(), + metadata JSONB DEFAULT '{}', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_dispatch_workers_state ON dispatch_workers (state); +CREATE INDEX idx_dispatch_workers_leader ON dispatch_workers (is_leader) + WHERE is_leader = TRUE; +CREATE INDEX idx_dispatch_workers_stale ON dispatch_workers (last_seen) + WHERE state = 'active'; +``` + +--- + +## 5. Module Layout & Package Design + +``` +github.com/xraph/dispatch/ +├── dispatch.go # Dispatcher struct, New(), Start(), Stop() +├── entity.go # Base Entity type (CreatedAt, UpdatedAt) +├── options.go # Functional options (WithStore, WithConcurrency, etc.) +├── config.go # Configuration struct, defaults +├── errors.go # Sentinel errors +├── backoff.go # Backoff strategies (Constant, Exponential, Jitter) +├── doc.go # Package documentation +│ +├── id/ +│ └── id.go # TypeID prefixes, typed IDs, constructors, parsers +│ +├── job/ +│ ├── job.go # Job entity (uses id.JobID), State enum +│ ├── definition.go # Definition[T] generic type, NewJob() +│ ├── registry.go # Job type registry (name → handler mapping) +│ ├── store.go # job.Store interface +│ └── options.go # Per-job options (retries, timeout, queue, priority) +│ +├── workflow/ +│ ├── workflow.go # Workflow definition, NewWorkflow() +│ ├── run.go # Run entity (uses id.RunID) +│ ├── step.go # Step(), Parallel(), WaitForEvent() +│ ├── checkpoint.go # Checkpoint entity (uses id.CheckpointID) +│ ├── context.go # Workflow context (carries run state) +│ └── store.go # workflow.Store interface +│ +├── cron/ +│ ├── cron.go # Cron scheduler, parser +│ ├── entry.go # CronEntry entity (uses id.CronID) +│ ├── leader.go # Leader election (distributed lock) +│ └── store.go # cron.Store interface +│ +├── queue/ +│ ├── queue.go # Queue interface +│ ├── priority.go # Priority queue implementation +│ └── rate_limiter.go # Per-queue / per-tenant rate limiting +│ +├── worker/ +│ ├── pool.go # Worker pool (goroutine management) +│ ├── executor.go # Job/workflow executor +│ └── lifecycle.go # Graceful shutdown, heartbeat, drain +│ +├── cluster/ +│ ├── cluster.go # Cluster manager (coordinates distributed workers) +│ ├── worker.go # Worker entity (uses id.WorkerID) +│ ├── consensus.go # Consensus interface + Postgres/Redis implementations +│ ├── rebalancer.go # Work stealing / queue rebalancing +│ ├── store.go # cluster.Store interface +│ └── k8s/ +│ ├── discovery.go # K8s pod discovery via endpoints API +│ ├── leader.go # K8s Lease-based leader election +│ └── labels.go # Pod label management for queue assignment +│ +├── middleware/ +│ ├── middleware.go # Middleware type definition +│ ├── logging.go # Structured logging per job +│ ├── tracing.go # OpenTelemetry trace per job execution +│ ├── metrics.go # Prometheus metrics +│ ├── timeout.go # Per-job timeout enforcement +│ ├── recover.go # Panic recovery +│ └── scope.go # Forge scope restoration from job metadata +│ +├── ext/ +│ ├── ext.go # Extension interface + all lifecycle hook interfaces +│ ├── registry.go # Extension registry, type-cached event dispatch +│ └── options.go # Extension config options +│ +├── relay_hook/ +│ ├── extension.go # relay_hook.Extension: bridges lifecycle → Relay +│ ├── events.go # Event type constants (dispatch.job.completed, etc.) +│ └── options.go # WithEvents(), WithPayload(), WithEnricher() +│ +├── dlq/ +│ ├── dlq.go # Dead letter queue service +│ ├── entry.go # DLQ entry entity (uses id.DLQID) +│ ├── replay.go # Replay failed jobs +│ └── store.go # dlq.Store interface +│ +├── event/ +│ ├── event.go # Event entity (uses id.EventID) +│ ├── bus.go # Event bus (in-memory + store-backed) +│ └── store.go # event.Store interface +│ +├── scope/ +│ └── scope.go # Forge scope helpers (capture/restore) +│ +├── store/ +│ ├── store.go # Composite Store interface +│ ├── postgres/ # pgx/v5 raw SQL implementation +│ │ ├── store.go +│ │ ├── job.go +│ │ ├── workflow.go +│ │ ├── cron.go +│ │ ├── dlq.go +│ │ ├── event.go +│ │ ├── cluster.go +│ │ └── migrations/ +│ ├── bun/ # Bun ORM implementation +│ │ ├── store.go +│ │ ├── models.go +│ │ ├── job.go +│ │ ├── workflow.go +│ │ ├── cron.go +│ │ ├── dlq.go +│ │ ├── event.go +│ │ ├── cluster.go +│ │ └── migrations/ +│ ├── sqlite/ +│ │ ├── store.go +│ │ └── migrations/ +│ ├── redis/ +│ │ ├── store.go +│ │ ├── job.go +│ │ ├── cluster.go +│ │ └── scripts/ +│ └── memory/ +│ └── store.go +│ +├── api/ +│ ├── handler.go # Admin API handlers +│ └── routes.go # Route mounting +│ +├── extension/ +│ ├── extension.go # forge.Extension implementation +│ └── options.go +│ +├── _examples/ +│ ├── basic/ +│ ├── workflow/ +│ ├── cron/ +│ ├── extensions/ # Custom extension (Slack notifier + audit logger) +│ ├── relay-hooks/ # Dispatch + Relay webhook delivery +│ ├── distributed/ # K8s multi-pod example +│ ├── integration-platform/ # Zapier-style pattern example +│ └── forge/ +│ +├── .golangci.yml +├── Makefile +├── go.mod +├── go.sum +└── README.md +``` + +--- + +## 6. Core Types & Interfaces + +### Dispatcher (Root) + +```go +// dispatch.go +package dispatch + +type Dispatcher struct { + config Config + store store.Store + registry *job.Registry + pool *worker.Pool + cluster *cluster.Manager + cron *cron.Scheduler + dlq *dlq.Service + eventBus *event.Bus + extensions *ext.Registry + middleware []middleware.Middleware + logger *slog.Logger +} + +func New(opts ...Option) (*Dispatcher, error) { + d := &Dispatcher{ + config: DefaultConfig(), + registry: job.NewRegistry(), + } + for _, opt := range opts { + if err := opt(d); err != nil { + return nil, err + } + } + if d.store == nil { + return nil, ErrNoStore + } + d.wireServices() + return d, nil +} + +func (d *Dispatcher) Register(definitions ...any) { /* register jobs/workflows */ } +func (d *Dispatcher) Start(ctx context.Context) error { /* start pool + cluster + cron */ } +func (d *Dispatcher) Stop(ctx context.Context) error { /* graceful drain + deregister */ } +func (d *Dispatcher) Cron(name, schedule string, j any, opts ...CronOption) { /* ... */ } +func (d *Dispatcher) Routes() http.Handler { /* admin API */ } +``` + +### Job Entity + +```go +// job/job.go +package job + +import "github.com/xraph/dispatch/id" + +type State string + +const ( + StatePending State = "pending" + StateRunning State = "running" + StateCompleted State = "completed" + StateFailed State = "failed" + StateRetrying State = "retrying" + StateCancelled State = "cancelled" +) + +type Job struct { + dispatch.Entity + + ID id.JobID `json:"id"` + Name string `json:"name"` + Queue string `json:"queue"` + Payload []byte `json:"payload"` + State State `json:"state"` + Priority int `json:"priority"` + MaxRetries int `json:"max_retries"` + RetryCount int `json:"retry_count"` + LastError string `json:"last_error,omitempty"` + ScopeAppID string `json:"scope_app_id,omitempty"` + ScopeOrgID string `json:"scope_org_id,omitempty"` + WorkerID id.WorkerID `json:"worker_id,omitempty"` + RunAt time.Time `json:"run_at"` + StartedAt *time.Time `json:"started_at,omitempty"` + CompletedAt *time.Time `json:"completed_at,omitempty"` + HeartbeatAt *time.Time `json:"heartbeat_at,omitempty"` +} +``` + +--- + +## 7. Extension System + +Dispatch is extensible via a first-class extension registry. Extensions hook into the job and workflow lifecycle at well-defined points, enabling custom behaviors without modifying Dispatch internals. This is how Relay integration, custom metrics, audit logging, and notification systems plug in. + +### The Extension Interface + +```go +// ext/ext.go +package ext + +import ( + "context" + + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/workflow" +) + +// Extension is the primary extensibility point for Dispatch. +// Implement any subset of the lifecycle interfaces below. +// Dispatch discovers which interfaces an extension implements +// and calls only those hooks. +type Extension interface { + // Name returns a unique identifier for the extension. + Name() string +} + +// ────────────────────────────────────────────────── +// Lifecycle interfaces — implement any subset +// ────────────────────────────────────────────────── + +// OnInit is called when the extension is registered with the Dispatcher. +// Use for setup: register event types, validate config, acquire resources. +type OnInit interface { + OnInit(ctx context.Context, d *dispatch.Dispatcher) error +} + +// OnShutdown is called during graceful shutdown. +type OnShutdown interface { + OnShutdown(ctx context.Context) error +} + +// ────────────────────────────────────────────────── +// Job lifecycle hooks +// ────────────────────────────────────────────────── + +// JobEnqueued is called after a job is persisted to the store. +type JobEnqueued interface { + OnJobEnqueued(ctx context.Context, j *job.Job) error +} + +// JobStarted is called when a worker begins executing a job. +type JobStarted interface { + OnJobStarted(ctx context.Context, j *job.Job) error +} + +// JobCompleted is called after a job finishes successfully. +type JobCompleted interface { + OnJobCompleted(ctx context.Context, j *job.Job, duration time.Duration) error +} + +// JobFailed is called when a job fails (before retry or DLQ). +type JobFailed interface { + OnJobFailed(ctx context.Context, j *job.Job, err error) error +} + +// JobRetrying is called when a job is about to be retried. +type JobRetrying interface { + OnJobRetrying(ctx context.Context, j *job.Job, attempt int, nextRunAt time.Time) error +} + +// JobDLQ is called when a job exhausts retries and enters the dead letter queue. +type JobDLQ interface { + OnJobDLQ(ctx context.Context, j *job.Job, err error) error +} + +// ────────────────────────────────────────────────── +// Workflow lifecycle hooks +// ────────────────────────────────────────────────── + +// WorkflowStarted is called when a workflow run begins. +type WorkflowStarted interface { + OnWorkflowStarted(ctx context.Context, run *workflow.Run) error +} + +// WorkflowStepCompleted is called after each workflow step succeeds. +type WorkflowStepCompleted interface { + OnWorkflowStepCompleted(ctx context.Context, run *workflow.Run, stepName string, duration time.Duration) error +} + +// WorkflowStepFailed is called when a workflow step fails. +type WorkflowStepFailed interface { + OnWorkflowStepFailed(ctx context.Context, run *workflow.Run, stepName string, err error) error +} + +// WorkflowCompleted is called when an entire workflow finishes. +type WorkflowCompleted interface { + OnWorkflowCompleted(ctx context.Context, run *workflow.Run, duration time.Duration) error +} + +// WorkflowFailed is called when a workflow fails terminally. +type WorkflowFailed interface { + OnWorkflowFailed(ctx context.Context, run *workflow.Run, err error) error +} + +// ────────────────────────────────────────────────── +// Cron lifecycle hooks +// ────────────────────────────────────────────────── + +// CronFired is called when a cron entry triggers. +type CronFired interface { + OnCronFired(ctx context.Context, entryName string, jobID id.JobID) error +} +``` + +### Extension Registry + +```go +// ext/registry.go +package ext + +// Registry manages registered extensions and dispatches lifecycle events. +type Registry struct { + extensions []Extension + logger *slog.Logger + + // Typed caches — built at registration time, not on every event + jobEnqueued []JobEnqueued + jobStarted []JobStarted + jobCompleted []JobCompleted + jobFailed []JobFailed + jobRetrying []JobRetrying + jobDLQ []JobDLQ + workflowStarted []WorkflowStarted + workflowCompleted []WorkflowCompleted + workflowFailed []WorkflowFailed + cronFired []CronFired + // ... etc +} + +// Register adds an extension and discovers which lifecycle interfaces it implements. +func (r *Registry) Register(e Extension) { + r.extensions = append(r.extensions, e) + + // Type-switch discovery — O(1) per event dispatch, not O(n) type assertions + if h, ok := e.(JobEnqueued); ok { + r.jobEnqueued = append(r.jobEnqueued, h) + } + if h, ok := e.(JobStarted); ok { + r.jobStarted = append(r.jobStarted, h) + } + if h, ok := e.(JobCompleted); ok { + r.jobCompleted = append(r.jobCompleted, h) + } + if h, ok := e.(JobFailed); ok { + r.jobFailed = append(r.jobFailed, h) + } + if h, ok := e.(JobRetrying); ok { + r.jobRetrying = append(r.jobRetrying, h) + } + if h, ok := e.(JobDLQ); ok { + r.jobDLQ = append(r.jobDLQ, h) + } + // ... discover all lifecycle interfaces +} + +// EmitJobCompleted fires OnJobCompleted on all extensions that implement it. +// Errors are logged but don't block — extensions must not break job execution. +func (r *Registry) EmitJobCompleted(ctx context.Context, j *job.Job, duration time.Duration) { + for _, h := range r.jobCompleted { + if err := h.OnJobCompleted(ctx, j, duration); err != nil { + r.logger.Error("extension hook failed", + "extension", h.(Extension).Name(), + "hook", "OnJobCompleted", + "job_id", j.ID.String(), + "error", err, + ) + } + } +} +``` + +### Registering Extensions + +```go +d := dispatch.New( + dispatch.WithStore(pgStore), + dispatch.WithConcurrency(20), + + // Built-in Relay hook extension (see Section 8) + dispatch.WithExtension(relay_hook.New(relayInstance)), + + // Custom audit logging extension + dispatch.WithExtension(&AuditExtension{logger: auditLogger}), + + // Custom Slack notification extension + dispatch.WithExtension(&SlackNotifier{webhook: slackURL}), + + // Custom metrics extension + dispatch.WithExtension(&DatadogExtension{client: ddClient}), +) +``` + +### Extension vs Middleware + +Extensions and middleware serve different purposes: + +| Concern | Middleware | Extension | +|---------|-----------|-----------| +| **When** | Wraps job execution (before + after) | Fires at discrete lifecycle points | +| **Can block execution?** | Yes (can reject/cancel jobs) | No (errors are logged, never block) | +| **Has access to** | Job payload, context, next handler | Job entity, metadata, duration, error | +| **Use for** | Tracing, timeout, auth, scope restore | Webhooks, notifications, audit, metrics | +| **Ordering** | Strict chain order (first registered → outermost) | All fire, no ordering guarantees | +| **Error behavior** | Error stops execution | Error logged, other hooks still fire | + +**Rule of thumb:** Use middleware when you need to wrap or gate execution. Use extensions when you need to react to lifecycle events without interfering with job processing. + +### Writing a Custom Extension + +```go +// Example: Slack notifier that pings on workflow failures +type SlackNotifier struct { + webhook string + client *http.Client +} + +func (s *SlackNotifier) Name() string { return "slack-notifier" } + +// Only implement the interfaces you care about +func (s *SlackNotifier) OnWorkflowFailed(ctx context.Context, run *workflow.Run, err error) error { + scope := forge.ScopeFrom(ctx) + msg := fmt.Sprintf("🚨 Workflow `%s` failed for org %s: %s", run.Name, scope.OrgID(), err) + return s.postSlack(ctx, msg) +} + +func (s *SlackNotifier) OnJobDLQ(ctx context.Context, j *job.Job, err error) error { + msg := fmt.Sprintf("💀 Job `%s` entered DLQ: %s", j.Name, err) + return s.postSlack(ctx, msg) +} + +func (s *SlackNotifier) postSlack(ctx context.Context, text string) error { + body, _ := json.Marshal(map[string]string{"text": text}) + req, _ := http.NewRequestWithContext(ctx, "POST", s.webhook, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + resp, err := s.client.Do(req) + if err != nil { + return err + } + return resp.Body.Close() +} +``` + +--- + +## 8. Hooks & Relay Integration + +Dispatch integrates with Relay to deliver lifecycle webhook events to external systems. When a customer subscribes to `dispatch.job.completed` via Relay's endpoint management, they automatically receive signed webhook deliveries every time a job completes — with Relay handling retries, signatures, fan-out, and delivery logs. + +### How It Works + +``` +Job completes in Dispatch + │ + ▼ +Extension Registry fires OnJobCompleted + │ + ├─→ relay_hook extension → relay.Send(ctx, event) + │ │ + │ ▼ + │ Relay persists event, resolves matching endpoints, + │ signs payload, delivers with retries, logs delivery + │ │ + │ ▼ + │ Customer receives webhook: + │ POST /webhooks + │ X-Relay-Event-Type: dispatch.job.completed + │ X-Relay-Signature: v1=... + │ {"job_id": "job_01h...", "name": "send-email", "status": "completed"} + │ + ├─→ Slack notifier extension → posts to Slack + ├─→ Audit extension → writes audit log + └─→ Datadog extension → emits custom metric +``` + +### Dispatch Event Types (Relay Catalog) + +When the Relay hook extension initializes, it registers all Dispatch event types with Relay's schema registry: + +```go +// relay_hook/events.go +package relay_hook + +// All Dispatch lifecycle events available for webhook subscription. +const ( + // Job events + EventJobEnqueued = "dispatch.job.enqueued" + EventJobStarted = "dispatch.job.started" + EventJobCompleted = "dispatch.job.completed" + EventJobFailed = "dispatch.job.failed" + EventJobRetrying = "dispatch.job.retrying" + EventJobDLQ = "dispatch.job.dlq" + + // Workflow events + EventWorkflowStarted = "dispatch.workflow.started" + EventWorkflowStepCompleted = "dispatch.workflow.step.completed" + EventWorkflowStepFailed = "dispatch.workflow.step.failed" + EventWorkflowCompleted = "dispatch.workflow.completed" + EventWorkflowFailed = "dispatch.workflow.failed" + + // Cron events + EventCronFired = "dispatch.cron.fired" + + // Cluster events + EventWorkerJoined = "dispatch.worker.joined" + EventWorkerLeft = "dispatch.worker.left" + EventLeaderElected = "dispatch.leader.elected" +) +``` + +### Relay Hook Extension + +This is a built-in extension that ships with Dispatch. It bridges the extension lifecycle to Relay's webhook delivery: + +```go +// relay_hook/extension.go +package relay_hook + +import ( + "context" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/workflow" + "github.com/xraph/relay" +) + +// Extension bridges Dispatch lifecycle events to Relay for webhook delivery. +// Customers subscribe to dispatch.* events via Relay's endpoint management. +type Extension struct { + relay *relay.Relay + config Config +} + +type Config struct { + // Which events to emit. Default: all events. + EnabledEvents []string + + // Whether to include job payload in webhook body. Default: false. + // Payloads may contain sensitive data — opt-in only. + IncludePayload bool + + // Custom data enricher — add extra fields to webhook payloads. + Enricher func(ctx context.Context, event map[string]any) map[string]any +} + +func New(r *relay.Relay, opts ...Option) ext.Extension { + e := &Extension{ + relay: r, + config: DefaultConfig(), + } + for _, opt := range opts { + opt(&e.config) + } + return e +} + +func (e *Extension) Name() string { return "relay-hooks" } + +// OnInit registers all Dispatch event types with Relay's schema registry. +func (e *Extension) OnInit(ctx context.Context, d *dispatch.Dispatcher) error { + schemas := []struct { + Type string + Description string + }{ + {EventJobEnqueued, "Fired when a job is enqueued for processing"}, + {EventJobStarted, "Fired when a worker begins executing a job"}, + {EventJobCompleted, "Fired when a job completes successfully"}, + {EventJobFailed, "Fired when a job fails (may retry)"}, + {EventJobRetrying, "Fired when a job is scheduled for retry"}, + {EventJobDLQ, "Fired when a job exhausts retries and enters the dead letter queue"}, + {EventWorkflowStarted, "Fired when a workflow run begins"}, + {EventWorkflowStepCompleted, "Fired when a workflow step completes"}, + {EventWorkflowStepFailed, "Fired when a workflow step fails"}, + {EventWorkflowCompleted, "Fired when an entire workflow completes successfully"}, + {EventWorkflowFailed, "Fired when a workflow fails terminally"}, + {EventCronFired, "Fired when a cron entry triggers"}, + } + + for _, s := range schemas { + if e.isEnabled(s.Type) { + e.relay.RegisterEventType(s.Type, relay.EventSchema{ + Description: s.Description, + Version: "2025-01-01", + }) + } + } + return nil +} + +// ────────────────────────────────────────────────── +// Job lifecycle → Relay events +// ────────────────────────────────────────────────── + +func (e *Extension) OnJobCompleted(ctx context.Context, j *job.Job, duration time.Duration) error { + if !e.isEnabled(EventJobCompleted) { + return nil + } + return relay.Send(ctx, &relay.Event{ + Type: EventJobCompleted, + Data: e.enrichJobPayload(ctx, map[string]any{ + "job_id": j.ID.String(), + "name": j.Name, + "queue": j.Queue, + "status": "completed", + "duration_ms": duration.Milliseconds(), + "retry_count": j.RetryCount, + "started_at": j.StartedAt, + "completed_at": j.CompletedAt, + }), + }) +} + +func (e *Extension) OnJobFailed(ctx context.Context, j *job.Job, err error) error { + if !e.isEnabled(EventJobFailed) { + return nil + } + return relay.Send(ctx, &relay.Event{ + Type: EventJobFailed, + Data: e.enrichJobPayload(ctx, map[string]any{ + "job_id": j.ID.String(), + "name": j.Name, + "queue": j.Queue, + "status": "failed", + "error": err.Error(), + "retry_count": j.RetryCount, + "max_retries": j.MaxRetries, + "will_retry": j.RetryCount < j.MaxRetries, + }), + }) +} + +func (e *Extension) OnJobDLQ(ctx context.Context, j *job.Job, err error) error { + if !e.isEnabled(EventJobDLQ) { + return nil + } + return relay.Send(ctx, &relay.Event{ + Type: EventJobDLQ, + Data: e.enrichJobPayload(ctx, map[string]any{ + "job_id": j.ID.String(), + "name": j.Name, + "queue": j.Queue, + "status": "dead_letter", + "error": err.Error(), + "retry_count": j.RetryCount, + }), + }) +} + +// ────────────────────────────────────────────────── +// Workflow lifecycle → Relay events +// ────────────────────────────────────────────────── + +func (e *Extension) OnWorkflowCompleted(ctx context.Context, run *workflow.Run, duration time.Duration) error { + if !e.isEnabled(EventWorkflowCompleted) { + return nil + } + return relay.Send(ctx, &relay.Event{ + Type: EventWorkflowCompleted, + Data: map[string]any{ + "run_id": run.ID.String(), + "name": run.Name, + "status": "completed", + "duration_ms": duration.Milliseconds(), + }, + }) +} + +func (e *Extension) OnWorkflowFailed(ctx context.Context, run *workflow.Run, err error) error { + if !e.isEnabled(EventWorkflowFailed) { + return nil + } + return relay.Send(ctx, &relay.Event{ + Type: EventWorkflowFailed, + Data: map[string]any{ + "run_id": run.ID.String(), + "name": run.Name, + "status": "failed", + "error": err.Error(), + }, + }) +} + +// ────────────────────────────────────────────────── +// Helpers +// ────────────────────────────────────────────────── + +func (e *Extension) isEnabled(eventType string) bool { + if len(e.config.EnabledEvents) == 0 { + return true // all enabled by default + } + for _, et := range e.config.EnabledEvents { + if et == eventType { + return true + } + } + return false +} + +func (e *Extension) enrichJobPayload(ctx context.Context, data map[string]any) map[string]any { + if e.config.Enricher != nil { + data = e.config.Enricher(ctx, data) + } + return data +} +``` + +### Usage: Wiring Relay Hooks + +```go +// Create Relay instance +r := relay.New( + relay.WithDatabase(db), + relay.WithSigningSecret("whsec_..."), + relay.WithMaxRetries(5), +) + +// Create Dispatch with Relay hooks +d := dispatch.New( + dispatch.WithStore(pgStore), + dispatch.WithConcurrency(20), + + // Wire Relay as a lifecycle hook extension + dispatch.WithExtension(relay_hook.New(r, + // Optional: only emit certain events + relay_hook.WithEvents( + relay_hook.EventJobCompleted, + relay_hook.EventJobFailed, + relay_hook.EventJobDLQ, + relay_hook.EventWorkflowCompleted, + relay_hook.EventWorkflowFailed, + ), + // Optional: include job payload in webhooks + relay_hook.WithPayload(true), + )), +) + +// Start both +r.Start(ctx) +d.Start(ctx) +``` + +### What the Customer Receives + +When a customer registers a webhook endpoint via Relay's API and subscribes to `dispatch.job.*`: + +```http +POST /webhooks HTTP/1.1 +Host: customer.example.com +Content-Type: application/json +X-Relay-Event-Type: dispatch.job.completed +X-Relay-Signature: v1=a2b3c4d5... +X-Relay-ID: evt_01h9a1b2c3... +X-Relay-Delivery-Attempt: 1 + +{ + "id": "evt_01h9a1b2c3...", + "type": "dispatch.job.completed", + "timestamp": "2026-02-17T12:00:00Z", + "data": { + "job_id": "job_01h2xcejqtf2nbrexx3vqjhp41", + "name": "send-email", + "queue": "default", + "status": "completed", + "duration_ms": 342, + "retry_count": 0, + "started_at": "2026-02-17T11:59:59Z", + "completed_at": "2026-02-17T12:00:00Z" + } +} +``` + +### Relay Integration Flow Between Libraries + +``` +Dispatch emits "dispatch.job.completed" + │ + ├─→ relay_hook extension → relay.Send() + │ │ + │ ├─→ Customer's monitoring endpoint (signed webhook) + │ ├─→ Customer's Slack bot endpoint (signed webhook) + │ └─→ Customer's audit log endpoint (signed webhook) + │ + ├─→ Ledger extension (if present) → meter job execution for billing + │ + └─→ Custom extension → whatever the developer needs + +Dispatch emits "dispatch.workflow.failed" + │ + ├─→ relay_hook extension → relay.Send() + │ │ + │ ├─→ Customer's PagerDuty endpoint (urgent webhook) + │ └─→ Customer's logging endpoint + │ + └─→ SlackNotifier extension → internal team notification +``` + +### Standalone (Without Relay) + +Extensions work without Relay. If you don't need webhook delivery, skip the relay_hook extension. The extension system is the foundation; Relay is one extension that plugs into it. + +```go +// No Relay — just custom extensions +d := dispatch.New( + dispatch.WithStore(pgStore), + dispatch.WithExtension(&AuditLogger{store: auditStore}), + dispatch.WithExtension(&SlackNotifier{webhook: url}), +) +``` + +--- + +## 9. Distributed Workers & Kubernetes Consensus + +Dispatch is single-process by default. You scale by running multiple Dispatcher instances pointing at the same store — Postgres SKIP LOCKED naturally distributes work across consumers. But for **production Kubernetes deployments**, you want more: worker registration, leader election, health-aware rebalancing, and work stealing. + +### The Problem + +When you scale to N pods in K8s, several things need coordination: + +- **Who runs cron?** Only one pod should fire each cron tick (leader election). +- **What if a pod dies mid-job?** Stale heartbeats need reaping, and those jobs need re-assignment. +- **How do you rebalance?** If pod A is overloaded and pod B is idle, work should shift. +- **How do pods discover each other?** Pod IPs change. Deployments scale up/down. + +### Design: Consensus Interface + +```go +// cluster/consensus.go +package cluster + +import "github.com/xraph/dispatch/id" + +// Consensus defines the contract for distributed coordination. +// Multiple implementations: Postgres advisory locks, Redis Redlock, +// Kubernetes Lease objects. +type Consensus interface { + // AcquireLeadership attempts to become the cluster leader. + // Returns true if this worker is now leader. TTL ensures leader + // failover if the holder crashes. + AcquireLeadership(ctx context.Context, workerID id.WorkerID, ttl time.Duration) (bool, error) + + // RenewLeadership extends the leader's hold. Must be called + // before TTL expires. + RenewLeadership(ctx context.Context, workerID id.WorkerID, ttl time.Duration) (bool, error) + + // ReleaseLeadership explicitly gives up leadership (graceful shutdown). + ReleaseLeadership(ctx context.Context, workerID id.WorkerID) error + + // IsLeader checks if the given worker is currently the leader. + IsLeader(ctx context.Context, workerID id.WorkerID) (bool, error) +} +``` + +### Consensus Implementations + +| Implementation | Mechanism | Best For | +|---------------|-----------|----------| +| `PostgresConsensus` | `pg_advisory_lock` | Single-DB deployments (default) | +| `RedisConsensus` | Redlock (multi-instance) | Redis-backed deployments | +| `K8sLeaseConsensus` | K8s Lease objects in `coordination.k8s.io` | Native K8s deployments | + +### Kubernetes-Specific: Pod Discovery + Lease Election + +```go +// cluster/k8s/discovery.go +package k8s + +// Discovery watches Kubernetes Endpoints or Pod resources to find +// other Dispatch worker pods in the same deployment/statefulset. +type Discovery struct { + clientset kubernetes.Interface + namespace string + labelSelector string // e.g. "app=my-service,dispatch=worker" + onChange func(peers []Peer) +} + +// Peer represents another Dispatch worker pod. +type Peer struct { + WorkerID id.WorkerID + PodName string + PodIP string + Ready bool + Queues []string +} +``` + +```go +// cluster/k8s/leader.go +package k8s + +import ( + coordinationv1 "k8s.io/api/coordination/v1" + "k8s.io/client-go/tools/leaderelection" +) + +// LeaseConsensus implements cluster.Consensus using Kubernetes Lease objects. +// This is the K8s-native way to do leader election — no external dependencies. +type LeaseConsensus struct { + clientset kubernetes.Interface + namespace string + leaseName string // e.g. "dispatch-leader" +} + +func (lc *LeaseConsensus) AcquireLeadership(ctx context.Context, workerID id.WorkerID, ttl time.Duration) (bool, error) { + // Uses K8s Lease API — same mechanism used by kube-controller-manager + // and kube-scheduler for their own leader election. + // ... +} +``` + +### Cluster Manager + +```go +// cluster/cluster.go +package cluster + +// Manager coordinates distributed Dispatch workers. +type Manager struct { + workerID id.WorkerID + store Store + consensus Consensus + rebalancer *Rebalancer + config ManagerConfig + logger *slog.Logger +} + +type ManagerConfig struct { + HeartbeatInterval time.Duration // How often to heartbeat (default: 10s) + StaleThreshold time.Duration // When to consider a worker dead (default: 30s) + RebalanceInterval time.Duration // How often to check balance (default: 60s) + LeaderTTL time.Duration // Leadership lock TTL (default: 15s) +} + +func (m *Manager) Start(ctx context.Context) error { + // 1. Register this worker in the store + // 2. Start heartbeat goroutine (reports alive + queue depth) + // 3. Start leader election loop + // 4. If leader: start cron scheduler + reaper + rebalancer + // 5. Watch for peer changes +} + +func (m *Manager) Stop(ctx context.Context) error { + // 1. Drain in-flight jobs + // 2. Release leadership if held + // 3. Deregister worker +} +``` + +### Work Stealing / Rebalancing + +```go +// cluster/rebalancer.go +package cluster + +// Rebalancer redistributes queued work when workers are unevenly loaded. +// Only the leader runs the rebalancer. +type Rebalancer struct { + store Store +} + +// Rebalance checks for: +// 1. Dead workers (no heartbeat) → reassign their running jobs +// 2. Overloaded workers → steal pending jobs from their queues +// 3. Idle workers → no action needed, they'll pick up naturally +func (r *Rebalancer) Rebalance(ctx context.Context) error { /* ... */ } +``` + +### Usage: Single Process (Default) + +```go +// No cluster config — just works with SKIP LOCKED +d := dispatch.New( + dispatch.WithStore(pgStore), + dispatch.WithConcurrency(20), +) +``` + +### Usage: Kubernetes Deployment + +```go +d := dispatch.New( + dispatch.WithStore(pgStore), + dispatch.WithConcurrency(20), + dispatch.WithCluster( + cluster.WithConsensus(k8s.NewLeaseConsensus(clientset, "default", "dispatch-leader")), + cluster.WithDiscovery(k8s.NewDiscovery(clientset, "default", "app=my-service")), + cluster.WithRebalanceInterval(60 * time.Second), + ), +) +``` + +### Usage: Multi-Process (Postgres Only) + +```go +// Multiple processes, same DB — Postgres advisory locks for consensus +d := dispatch.New( + dispatch.WithStore(pgStore), + dispatch.WithConcurrency(20), + dispatch.WithCluster( + cluster.WithPostgresConsensus(pgPool), + ), +) +``` + +--- + +## 10. Integration Platform Pattern (Zapier-Style) + +**Yes, Dispatch can absolutely power a Zapier-style integration platform.** The workflow engine, step functions, and event bus provide all the primitives needed. Here's how it maps: + +### Concept Mapping + +| Zapier Concept | Dispatch Primitive | How It Works | +|---------------|-------------------|--------------| +| **Trigger** | `event.Bus` + `WaitForEvent` | External webhook/poll → publishes event → wakes workflow | +| **Action** | `workflow.Step()` | Each integration action is a step with retries and checkpointing | +| **Zap (flow)** | `workflow.NewWorkflow()` | The entire trigger → action chain is a durable workflow | +| **Multi-step Zap** | `wf.Step()` chained | Each step checkpointed. Crash-safe. Independent retries | +| **Fan-out** | `wf.Parallel()` | "When trigger fires, do A AND B AND C simultaneously" | +| **Delay** | `wf.Sleep()` | "Wait 30 minutes, then send the follow-up" | +| **Filter/condition** | Go `if` in workflow | "Only continue if amount > $100" — just Go code | +| **Error handling** | DLQ + per-step retry | Failed steps retry independently. Exhausted → DLQ for inspection | +| **Rate limiting** | `queue.RateLimiter` | Per-tenant rate limits to respect API quotas per integration | +| **Execution log** | Job/Run state + OTel | Full trace of every step, every retry, every failure | + +### Architecture: How an Integration Platform Uses Dispatch + +``` +┌──────────────────────────────────────────────────────────────────────┐ +│ Integration Platform Layer │ +│ (YOUR code, not Dispatch) │ +│ │ +│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────────┐ │ +│ │ Connector SDK │ │ Flow Builder │ │ Execution Tracker │ │ +│ │ (Slack, Gmail, │ │ (UI defines │ │ (Shows run history, │ │ +│ │ Sheets, etc.) │ │ trigger→action │ │ step status, logs) │ │ +│ │ │ │ chains) │ │ │ │ +│ └────────┬────────┘ └───────┬─────────┘ └──────────┬──────────┘ │ +│ │ │ │ │ +│ ┌────────▼───────────────────▼────────────────────────▼──────────┐ │ +│ │ Dispatch (the engine underneath) │ │ +│ │ │ │ +│ │ • Workflows = user-defined automation flows │ │ +│ │ • Steps = individual connector actions (with retries) │ │ +│ │ • Events = triggers from webhooks/polls │ │ +│ │ • Cron = scheduled triggers ("every hour", "daily at 9am") │ │ +│ │ • DLQ = failed automations for user inspection │ │ +│ │ • Per-tenant rate limits = respect each API's quotas │ │ +│ │ • Scope = org-level isolation (each customer's flows) │ │ +│ └─────────────────────────────────────────────────────────────────┘ │ +└──────────────────────────────────────────────────────────────────────┘ +``` + +### Example: "When Stripe payment received → create Jira ticket → notify Slack" + +```go +// This is what your integration platform generates when a user +// configures a flow. Dispatch powers the execution; your platform +// owns the connector SDK and flow builder UI. + +var StripeToJiraToSlack = dispatch.NewWorkflow("stripe-jira-slack", + func(wf dispatch.Workflow, trigger StripePaymentEvent) error { + + // Step 1: Transform Stripe data into Jira ticket fields + ticket, err := wf.StepWithResult("create-jira-ticket", + func(ctx dispatch.Context) (*JiraTicket, error) { + return jiraConnector.CreateTicket(ctx, jira.CreateInput{ + Project: trigger.Metadata["jira_project"], + Summary: fmt.Sprintf("Payment received: $%.2f from %s", + float64(trigger.Amount)/100, trigger.CustomerEmail), + Type: "Task", + }) + }, + ) + if err != nil { + return err + } + + // Step 2: Notify Slack (retries independently of step 1) + _, err = wf.Step("notify-slack", func(ctx dispatch.Context) error { + return slackConnector.SendMessage(ctx, slack.MessageInput{ + Channel: trigger.Metadata["slack_channel"], + Text: fmt.Sprintf("💰 Payment $%.2f → Jira ticket %s created", + float64(trigger.Amount)/100, ticket.Key), + }) + }) + + return err + }, +) +``` + +### Trigger Pattern: Webhook → Event → Workflow + +```go +// Your platform's webhook receiver — not part of Dispatch itself +func handleStripeWebhook(w http.ResponseWriter, r *http.Request) { + event := parseStripeEvent(r) + + // Publish into Dispatch's event bus + // This wakes up any WaitForEvent listeners AND + // can be used to trigger registered flows + dispatch.PublishEvent(r.Context(), &dispatch.Event{ + Name: "stripe.payment.received", + Payload: marshal(event), + }) + + // OR directly run the workflow for this trigger + dispatch.RunWorkflow(r.Context(), StripeToJiraToSlack, event) +} +``` + +### Trigger Pattern: Cron (Polling-Based Triggers) + +```go +// For integrations that don't support webhooks — poll on a schedule +d.Cron("poll-gmail-inbox", "*/5 * * * *", PollGmailInbox, + dispatch.CronTenant("*"), // runs per-tenant +) + +var PollGmailInbox = dispatch.NewJob("poll-gmail-inbox", + func(ctx dispatch.Context, _ struct{}) error { + scope := forge.MustScope(ctx) + // 1. Get this tenant's Gmail credentials + // 2. Check for new emails since last poll + // 3. For each new email, run the user's configured flow + newEmails := gmailConnector.Poll(ctx, scope.OrgID()) + for _, email := range newEmails { + dispatch.RunWorkflow(ctx, userFlow, email) + } + return nil + }, +) +``` + +### What Dispatch Handles vs What You Build + +| Concern | Owned By | Notes | +|---------|----------|-------| +| **Durable execution** | Dispatch | Crash-safe, checkpointed workflows | +| **Retries + backoff** | Dispatch | Per-step, configurable | +| **Rate limiting** | Dispatch | Per-tenant, per-queue (respect API quotas) | +| **Scheduling** | Dispatch | Cron for polling triggers | +| **Dead letter queue** | Dispatch | Failed flows surfaced to users | +| **Scope isolation** | Dispatch | Each customer's flows isolated | +| **Execution history** | Dispatch | Workflow runs, step status, timings | +| **Observability** | Dispatch | OTel traces per flow execution | +| **Connector SDK** | You | The Slack/Gmail/Jira/Stripe adapters | +| **Flow builder UI** | You | Visual editor for connecting triggers → actions | +| **Credential vault** | You | OAuth tokens, API keys per tenant | +| **Flow registry** | You | Which workflows each tenant has configured | +| **Marketplace** | You | Published connectors, templates | + +### Key Insight + +Dispatch doesn't know about Slack or Jira. It knows about **jobs, workflows, steps, events, and cron**. Your integration platform maps user-configured flows onto these primitives. Dispatch handles all the hard stuff — durable execution, retries, scheduling, rate limiting, tenant isolation — so your platform code focuses purely on the connector logic and user experience. + +--- + +## 11. Forge Scope Integration + +Dispatch reads `forge.Scope` from context when available. When a job is enqueued, the current scope is captured into the job's metadata. When the worker executes the job, the scope is restored onto the execution context. + +```go +// scope/scope.go +package scope + +import ( + "context" + "github.com/xraph/forge" +) + +// Capture extracts forge.Scope from context and returns app/org IDs. +// Returns empty strings if no scope is present (standalone mode). +func Capture(ctx context.Context) (appID, orgID string) { + s := forge.ScopeFrom(ctx) + if s.IsZero() { + return "", "" + } + return s.AppID(), s.OrgID() +} + +// Restore creates a context with forge.Scope from stored app/org IDs. +// No-op if both are empty (standalone mode). +func Restore(ctx context.Context, appID, orgID string) context.Context { + if appID == "" { + return ctx + } + if orgID == "" { + return forge.WithScope(ctx, forge.NewAppScope(appID)) + } + return forge.WithScope(ctx, forge.NewOrgScope(appID, orgID)) +} +``` + +### Scope Reference for Dispatch + +| Operation | Scope Level | Rationale | +|-----------|-------------|-----------| +| Job execution | Inherited | Jobs carry the scope they were enqueued with | +| Cron (global) | App | Platform-wide scheduled tasks | +| Cron (per-tenant) | Org | Per-customer scheduled tasks | +| Queues | Both | App-level queues + org-level isolation | +| DLQ | Inherited | Failed jobs retain original scope | +| Rate limiting | Org | Per-customer job rate limits | +| Worker registration | App | Workers are platform-level resources | +| Leader election | App | One leader per app deployment | + +--- + +## 12. Linting & Code Quality (golangci-lint v2) + +### Configuration + +```yaml +# .golangci.yml — golangci-lint v2 configuration for Dispatch +version: "2" + +linters: + default: none + enable: + # Core correctness + - errcheck + - govet + - staticcheck + - unused + - ineffassign + - typecheck + + # Code quality + - gofmt + - goimports + - gocritic + - revive + - misspell + - unconvert + - unparam + - prealloc + + # Bug prevention + - bodyclose + - noctx + - rowserrcheck + - sqlclosecheck + - exportloopref + - gosec + - errname + - errorlint + + # Style + - nolintlint + - whitespace + - predeclared + - tenv + + settings: + gocritic: + enabled-tags: + - diagnostic + - style + - performance + disabled-checks: + - hugeParam + - rangeValCopy + + revive: + rules: + - name: blank-imports + - name: context-as-argument + - name: context-keys-type + - name: dot-imports + - name: error-return + - name: error-strings + - name: error-naming + - name: exported + arguments: [checkPrivateReceivers] + - name: if-return + - name: increment-decrement + - name: var-naming + - name: var-declaration + - name: range + - name: receiver-naming + - name: time-naming + - name: unexported-return + - name: indent-error-flow + - name: errorf + - name: empty-block + - name: superfluous-else + - name: unused-parameter + - name: unreachable-code + + gosec: + excludes: + - G104 + - G304 + + errcheck: + check-type-assertions: true + check-blank: true + exclude-functions: + - (io.Closer).Close + - (*database/sql.Rows).Close + + govet: + enable-all: true + disable: + - fieldalignment + + exclusions: + presets: + - comments + - std-error-handling + rules: + - path: _test\.go + linters: [gosec, errcheck, gocritic] + - path: _examples/ + linters: [errcheck, gosec] + - path: ".*_gen\\.go" + linters: [all] + +formatters: + enable: + - gofmt + - goimports + settings: + goimports: + local-prefixes: + - github.com/xraph/dispatch + +output: + sort-order: + - linter + - file +``` + +### Makefile + +```makefile +.PHONY: build test lint lint-fix lint-install migrate + +build: + @go build ./... + +test: + @go test -race -count=1 ./... + +test-integration: + @go test -race -tags=integration ./store/postgres/... ./store/bun/... + +lint-install: + @go install github.com/golangci/golangci-lint/v2/cmd/golangci-lint@latest + +lint: + @golangci-lint run ./... + +lint-fix: + @golangci-lint run --fix ./... + +migrate: + @go run ./cmd/migrate/main.go up +``` + +### Claude Code Enforcement + +Every phase must pass `golangci-lint run ./...` with zero errors before proceeding. Never use `//nolint` without a justifying comment. + +--- + +## 13. Development Phases + +### Phase 0 — Project Scaffold & Tooling + +**Goal:** Repo skeleton, go.mod, TypeID package, lint config, Makefile. + +**Claude Code instructions:** +``` +Initialize github.com/xraph/dispatch with Go 1.23+. +Create the full directory structure from the Module Layout section. +Create id/id.go with all TypeID prefix types, typed IDs, constructors, and parsers. +Create entity.go with the base Entity type. +Create .golangci.yml with the v2 config. +Create Makefile. Create errors.go with sentinel errors. +Verify: go build ./... and golangci-lint run ./... both pass. +``` + +**Exit criteria:** `go build ./...` and `golangci-lint run ./...` pass. TypeID package compiles. + +--- + +### Phase 1 — Core Types & Store Interface + +**Goal:** All entity types with TypeID, store interfaces per subsystem, composite store, memory backend. + +**Claude Code instructions:** +``` +Implement all entity types using TypeID: +- job/job.go: Job struct with id.JobID, State enum +- job/definition.go: Definition[T] generic type +- workflow/run.go: Run with id.RunID, Checkpoint with id.CheckpointID +- cron/entry.go: Entry with id.CronID +- dlq/entry.go: Entry with id.DLQID +- event/event.go: Event with id.EventID +- cluster/worker.go: Worker with id.WorkerID + +Implement store interfaces following ControlPlane composite pattern: +- Each subsystem defines Store in its package (job.Store, workflow.Store, etc.) +- store/store.go composes all into single Store interface (including cluster.Store) +- store/memory/store.go implements the full composite for testing + +Write table-driven unit tests for the memory store. +Verify: go test ./... and golangci-lint run ./... pass. +``` + +**Exit criteria:** Full memory store with passing tests. All entities use TypeID. Lint clean. + +--- + +### Phase 2 — Job Registry, Enqueue, Worker Pool & Extension System + +**Goal:** Register typed jobs, enqueue them, process from worker pool with scope capture/restore. Create the extension system so lifecycle hooks fire from day one. + +**Claude Code instructions:** +``` +Implement the extension system FIRST: +- ext/ext.go: Extension interface + all lifecycle hook interfaces + (JobEnqueued, JobStarted, JobCompleted, JobFailed, JobRetrying, JobDLQ, + WorkflowStarted, WorkflowCompleted, WorkflowFailed, etc.) +- ext/registry.go: Registry struct with Register(), type-cached emit methods + (EmitJobCompleted, EmitJobFailed, etc.). Errors logged, never block. + +Then implement job/registry.go, middleware system (middleware.go, recover.go, logging.go), +worker pool (pool.go, executor.go, lifecycle.go), scope helpers (scope/scope.go, +middleware/scope.go). + +CRITICAL: The executor must call ext.Registry.Emit*() at every lifecycle point: +- After enqueue → EmitJobEnqueued +- When worker picks up job → EmitJobStarted +- After success → EmitJobCompleted +- After failure → EmitJobFailed +- Before retry → EmitJobRetrying +Wire WithExtension(e) into Dispatcher options. + +Integration tests: enqueue→execute with hooks firing, scope flows through, +custom test extension receives all lifecycle events, graceful shutdown drains. +Verify: go test ./... and golangci-lint run ./... pass. +``` + +**Files:** +- `ext/ext.go`, `ext/registry.go`, `ext/options.go` +- `job/registry.go` +- `middleware/middleware.go`, `middleware/recover.go`, `middleware/logging.go` +- `worker/pool.go`, `worker/executor.go`, `worker/lifecycle.go` +- `scope/scope.go`, `middleware/scope.go` +- Updated `dispatch.go`, `options.go` (WithExtension) +- Tests: `ext/registry_test.go`, `dispatch_test.go`, `worker/pool_test.go` + +**Exit criteria:** End-to-end job processing with memory store. Extensions fire at all lifecycle points. Scope flows. Clean shutdown. + +--- + +### Phase 3 — Retry, Backoff & Dead Letter Queue + +**Goal:** Automatic retries with backoff. Failed jobs land in DLQ. Extensions fire at retry and DLQ points. + +**Claude Code instructions:** +``` +Implement backoff.go (Constant, Linear, Exponential, ExponentialWithJitter). +Implement dlq/dlq.go service and dlq/replay.go. +Update executor to handle retries + push to DLQ on exhaustion. +Wire extension hooks: EmitJobRetrying before retry, EmitJobDLQ when entering DLQ. +Tests: retry-then-succeed, exhaust-retries-to-DLQ, DLQ replay, + test extension receives OnJobRetrying and OnJobDLQ callbacks. +Verify: go test ./... and golangci-lint run ./... pass. +``` + +**Exit criteria:** Retries work. DLQ captures failures. Replay re-enqueues. Extension hooks fire at retry and DLQ. + +--- + +### Phase 4 — Workflows & Step Functions + +**Goal:** Multi-step workflows with checkpointing, fan-out/fan-in, durable waits. Extension hooks fire at workflow and step lifecycle points. + +**Claude Code instructions:** +``` +Implement workflow execution: Workflow struct with Step(), StepWithResult(), +Parallel() (errgroup), WaitForEvent(), Sleep(). Checkpoint serialization via gob. +event/bus.go for in-memory + store-backed event bus. +Wire extension hooks: EmitWorkflowStarted, EmitWorkflowStepCompleted, +EmitWorkflowStepFailed, EmitWorkflowCompleted, EmitWorkflowFailed. +Tests: multi-step, crash-resume from checkpoint, parallel, WaitForEvent + timeout, + test extension receives all workflow lifecycle callbacks. +Verify: go test ./... and golangci-lint run ./... pass. +``` + +**Exit criteria:** Full workflow lifecycle. Checkpoint survives crashes. WaitForEvent works. Workflow extension hooks fire. + +--- + +### Phase 5 — Cron Scheduling & Leader Election + +**Goal:** Distributed cron with leader election. Per-tenant cron support. + +**Claude Code instructions:** +``` +Implement cron/cron.go (parser via robfig/cron/v3), cron/leader.go (distributed lock). +Per-tenant cron: CronTenant("*") iterates orgs, enqueues per-org with correct scope. +Tests: correct firing, two dispatchers with leader election, per-tenant cron. +Verify: go test ./... and golangci-lint run ./... pass. +``` + +**Exit criteria:** Cron fires correctly. Only leader executes. Per-tenant works. + +--- + +### Phase 6 — Queue Features (Priority, Rate Limiting, Per-Tenant) + +**Goal:** Priority queues, per-queue and per-tenant rate limiting. + +**Exit criteria:** Priority ordering. Rate limiting throttles. Tenant isolation verified. + +--- + +### Phase 7 — Observability (OpenTelemetry + Prometheus) + +**Goal:** OTel traces per job, Prometheus metrics, timeout middleware. + +**Exit criteria:** Spans created. Metrics exposed. Timeouts cancel jobs. + +--- + +### Phase 8 — PostgreSQL Store (pgx) + +**Goal:** Production Postgres backend with pgx/v5. SKIP LOCKED dequeue, advisory lock leader election, LISTEN/NOTIFY events. + +**Claude Code instructions:** +``` +Implement store/postgres/ — all subsystem implementations including cluster.Store. +Embed SQL migrations. Use pgxpool for connection pooling. +Key patterns: SELECT FOR UPDATE SKIP LOCKED for dequeue, +pg_advisory_lock for leader election, LISTEN/NOTIFY for events. +Integration tests with testcontainers-go. +Verify: go test -tags=integration ./store/postgres/... and golangci-lint run ./... pass. +``` + +**Exit criteria:** Full Postgres store with integration tests passing. + +--- + +### Phase 9 — Bun Store + +**Goal:** Bun ORM backend for teams using Bun in their Forge/ControlPlane stack. + +**Claude Code instructions:** +``` +Implement store/bun/ — Bun model structs in models.go with bun:"" tags. +Each subsystem implementation uses Bun query builder. +For SKIP LOCKED dequeue, use bun.NewRaw() (Bun supports raw queries). +Bun Go-based migrations in migrations/001_initial.go. +Share the same SQL schemas as Postgres (same tables, Bun just uses ORM for access). +Integration tests with testcontainers-go. +Verify: go test -tags=integration ./store/bun/... and golangci-lint run ./... pass. +``` + +**Exit criteria:** Full Bun store. Passes same integration test suite as Postgres store. + +--- + +### Phase 10 — SQLite & Redis Stores + +**Goal:** SQLite for embedded/edge. Redis for high-throughput ephemeral workloads. + +**Claude Code instructions:** +``` +SQLite: store/sqlite/ — embedded migrations, adapted SQL (no SKIP LOCKED, use +BEGIN IMMEDIATE + rowid ordering instead). modernc.org/sqlite for pure-Go driver. + +Redis: store/redis/ — Redis Streams for job queues (XREADGROUP for consumer groups), +Sorted Sets for priority/scheduling, Lua scripts for atomic dequeue, +Redlock for leader election. + +Tests for both backends. +Verify: go test ./store/sqlite/... ./store/redis/... and golangci-lint run ./... pass. +``` + +**Exit criteria:** SQLite and Redis stores pass tests. + +--- + +### Phase 11 — Distributed Workers & Cluster + +**Goal:** Worker registration, consensus interface, Postgres + K8s leader election, rebalancing. + +**Claude Code instructions:** +``` +Implement cluster/ package: +- cluster.go: Manager with Start/Stop lifecycle +- consensus.go: Consensus interface +- worker.go: Worker entity +- rebalancer.go: Work stealing logic (leader-only) + +Implement consensus backends: +- Store-based (Postgres advisory locks) — default +- cluster/k8s/leader.go: K8s Lease-based (optional, behind build tag) +- cluster/k8s/discovery.go: K8s pod discovery + +Wire into Dispatcher: WithCluster() option. Manager starts alongside pool. +Tests: worker registration, leader election, stale reaping, rebalance. +Verify: go test ./... and golangci-lint run ./... pass. +``` + +**Exit criteria:** Cluster manager works. Leader election prevents duplicate cron. Dead workers reaped. + +--- + +### Phase 12 — Forge Extension + +**Goal:** Mount Dispatch into Forge as an extension. + +**Exit criteria:** Dispatch mounts into Forge. Auto-discovers store from DI. + +--- + +### Phase 13 — Relay Hook Extension + +**Goal:** Built-in extension that bridges Dispatch lifecycle events to Relay for webhook delivery. + +**Claude Code instructions:** +``` +Implement relay_hook/ package: +- relay_hook/events.go: All event type constants (dispatch.job.completed, etc.) +- relay_hook/extension.go: Extension struct implementing ext.Extension + + all job/workflow lifecycle interfaces. Each hook calls relay.Send() with + typed event payload. +- relay_hook/options.go: WithEvents() to filter events, WithPayload() to include + job payload in webhook, WithEnricher() for custom data enrichment. + +OnInit registers all Dispatch event types with Relay's schema registry. +Relay is a soft dependency — relay_hook imports the relay package, +but Dispatch core does NOT import relay. + +Tests: mock Relay, register extension, process job, verify relay.Send() called +with correct event type and payload structure. +Integration test: real Relay + Dispatch, verify webhook delivered for job.completed. +Verify: go test ./... and golangci-lint run ./... pass. +``` + +**Files:** +- `relay_hook/extension.go`, `relay_hook/events.go`, `relay_hook/options.go` +- Tests: `relay_hook/extension_test.go` + +**Exit criteria:** Relay hook extension fires for all lifecycle events. Customers receive signed webhooks via Relay when subscribed to dispatch.* events. + +--- + +### Phase 14 — Admin API, Examples & Documentation + +**Goal:** HTTP admin API, examples (including extensions, relay hooks, integration platform, and distributed), README. + +**Claude Code instructions:** +``` +Implement admin API (Dispatcher.Routes()). +Write examples: +- _examples/basic/: Simple job +- _examples/workflow/: Multi-step with fan-out +- _examples/cron/: Scheduled jobs +- _examples/extensions/: Custom extension (Slack notifier + audit logger) +- _examples/relay-hooks/: Dispatch + Relay webhook delivery +- _examples/distributed/: K8s multi-pod with cluster config +- _examples/integration-platform/: Zapier-style trigger→action flow +- _examples/forge/: Forge extension +Write README.md with extension system + Relay integration docs. +Verify: go test ./... and golangci-lint run ./... pass. +``` + +**Exit criteria:** Admin API working. All examples run. README complete. + +--- + +## 14. Claude Code Development Guide + +### Starting Each Phase + +``` +Read the Dispatch design document at docs/DESIGN.md. +We are implementing Phase N: [Phase Name]. +Follow the files, exit criteria, and instructions for this phase. +Before writing code, read existing codebase. +After writing code, run: + 1. go build ./... + 2. go test ./... + 3. golangci-lint run ./... +All three must pass before the phase is complete. +``` + +### Code Style Rules + +- **IDs:** Use TypeID via `id.NewJobID()`, `id.ParseJobID(s)`, etc. Never raw strings for IDs. +- **Errors:** Always wrap with `fmt.Errorf("dispatch: %w", err)`. Use sentinel errors. +- **Context:** Always first parameter. Check `ctx.Done()` in loops. +- **Imports:** Group as `stdlib → external → internal`. `goimports` handles this. +- **No globals.** All state on structs. No `init()` with side effects. +- **No panics in library code.** Return errors. Only `id.must()` panics (internal, infallible). +- **Table-driven tests.** All tests use `tests := []struct{ ... }` pattern. + +### Dependency Graph + +``` +dispatch (root) +├── id (TypeID definitions — zero deps) +├── job (entities + store interface) → id +├── workflow (entities + store interface) → id +├── cron (entities + store interface) → id +├── dlq (entities + store interface) → id +├── event (entities + store interface) → id +├── cluster (entities + store + consensus interface) → id +│ └── k8s (K8s-specific consensus + discovery) → cluster +├── ext (extension interfaces + registry) → job, workflow, id +├── relay_hook (Relay bridge extension) → ext, relay (soft dep) +├── queue (queue abstraction) → job +├── worker (pool + executor) → job, middleware, store, ext +├── middleware → job, scope +├── scope → forge (soft dependency) +├── store → job, workflow, cron, dlq, event, cluster (composes interfaces) +│ ├── memory → store +│ ├── postgres → store (pgx/v5) +│ ├── bun → store (uptrace/bun) +│ ├── sqlite → store (modernc.org/sqlite) +│ └── redis → store (go-redis/v9) +├── api → dispatch, job, workflow, cron, dlq, cluster +└── extension → dispatch, forge +``` + +**Rules:** +- Subsystem packages NEVER import each other. +- `ext/` defines interfaces only — no implementation imports. +- `relay_hook/` imports `ext/` and `relay` but relay is a soft dependency (interface-based). +- `worker/` calls `ext.Registry.Emit*()` after job execution. +- `store/` composes interfaces but never imports implementations. +- `cluster/k8s/` is behind a build tag (`//go:build k8s`). +- `extension/` imports `dispatch` root and `forge` but nothing else. + +### Key Decisions Log + +| Decision | Choice | Rationale | +|----------|--------|-----------| +| IDs | TypeID (`go.jetify.com/typeid`) | Consistent with ControlPlane/Nexus. Type-safe, prefixed, K-sortable | +| Store pattern | ControlPlane composite | Subsystem interfaces composed into Store. Proven pattern | +| Store backends | Postgres, Bun, SQLite, Redis, Memory | Matches ControlPlane + adds Bun for ORM users | +| Bun store | `uptrace/bun` | Teams using Forge/ControlPlane already have Bun. Natural fit | +| Postgres driver | pgx/v5 | Best Go Postgres driver | +| Dequeue strategy | SKIP LOCKED | Postgres-native concurrent polling | +| Leader election | Pluggable Consensus interface | Postgres advisory locks (default), K8s Lease, Redis Redlock | +| K8s support | Optional build tag | `cluster/k8s/` behind `//go:build k8s`. No K8s deps for non-K8s users | +| Event delivery | LISTEN/NOTIFY (Postgres) | Real-time without polling. Fallback to polling for other stores | +| Serialization | gob for checkpoints, JSON for payloads | gob fast for internal state. JSON for user payloads | +| Extension system | Interface-based discovery | Extensions implement any subset of lifecycle interfaces. Registry type-asserts at registration time for O(1) dispatch | +| Extension error policy | Log and continue | Extension errors never block job execution. Logged via slog | +| Relay integration | Extension, not core | Relay is one extension (`relay_hook/`). Dispatch core does not import Relay. Soft dependency | +| Webhook delivery | Relay-native | Relay handles signing, retries, fan-out, delivery logs. Dispatch just calls relay.Send() | +| Middleware vs Extension | Both, different purpose | Middleware wraps execution (can block). Extensions react to lifecycle (can't block) | +| Linting | golangci-lint v2 | Modern, fast, comprehensive | + +--- + +## Phase Summary + +| Phase | Name | Effort | Depends On | +|-------|------|--------|------------| +| 0 | Project Scaffold & Tooling | 0.5 day | — | +| 1 | Core Types & Store Interface | 1 day | Phase 0 | +| 2 | Job Registry, Worker Pool & Extension System | 2.5 days | Phase 1 | +| 3 | Retry, Backoff & Dead Letter Queue | 1 day | Phase 2 | +| 4 | Workflows & Step Functions | 2 days | Phase 2 | +| 5 | Cron Scheduling & Leader Election | 1.5 days | Phase 2 | +| 6 | Queue Features | 1 day | Phase 2 | +| 7 | Observability | 1 day | Phase 2 | +| 8 | PostgreSQL Store (pgx) | 2 days | Phase 1 | +| 9 | Bun Store | 1.5 days | Phase 1 | +| 10 | SQLite & Redis Stores | 2 days | Phase 1 | +| 11 | Distributed Workers & Cluster | 2.5 days | Phase 2, 8 | +| 12 | Forge Extension | 1 day | Phase 2 | +| 13 | Relay Hook Extension | 1.5 days | Phase 2 | +| 14 | Admin API, Examples & Docs | 2.5 days | All | + +**Total estimated effort: ~23 days (6–8 weeks at part-time pace)** + +### Parallelizable Phases + +``` +Phase 0 → Phase 1 → Phase 2 ─┬─→ Phase 3 (Retry + DLQ) + (ext/ created here) ────├─→ Phase 4 (Workflows + workflow hooks) + ├─→ Phase 5 (Cron) + ├─→ Phase 6 (Queue features) + ├─→ Phase 7 (Observability) + ├─→ Phase 11 (Distributed workers) *needs Phase 8 + ├─→ Phase 12 (Forge extension) + └─→ Phase 13 (Relay hook extension) + Phase 1 ───┬─→ Phase 8 (Postgres store) + ├─→ Phase 9 (Bun store) + └─→ Phase 10 (SQLite + Redis) + │ + ▼ + Phase 14 (API + Examples + Docs) +``` + +--- + +## v0.1.0 MVP Target + +- Jobs + enqueue + worker pool + extension system (Phase 2) +- Basic retries + DLQ (Phase 3) +- Basic workflows (Phase 4, without WaitForEvent) +- Cron (Phase 5) +- Postgres store (Phase 8) +- Memory store for testing (Phase 1) +- Extension hooks fire at all lifecycle points (but no Relay yet) + +## v0.2.0 + +- Bun store (Phase 9) +- Priority queues + rate limiting (Phase 6) +- Full workflows with WaitForEvent + Sleep (Phase 4 complete) +- Observability (Phase 7) +- **Relay hook extension** (Phase 13) — webhook delivery for all lifecycle events +- Forge extension (Phase 12) + +## v0.3.0 + +- Distributed workers + K8s consensus (Phase 11) +- SQLite + Redis stores (Phase 10) +- Admin API (Phase 14) +- Extension + relay-hook + integration-platform examples (Phase 14) diff --git a/api/job_handler.go b/api/job_handler.go index d66eb7c..96a8bf3 100644 --- a/api/job_handler.go +++ b/api/job_handler.go @@ -108,8 +108,12 @@ func (a *API) retryJob(ctx forge.Context, _ *RetryJobRequest) (*struct{}, error) j.RetryCount = 0 j.LastError = "" j.RunAt = now - j.StartedAt = nil j.CompletedAt = nil + // Clears StartedAt along with the worker and lease fields the failed + // run left behind. Without it the retried job carries a lapsed + // lease_expires_at into pending, which a claim that grants no lease + // never overwrites — see job.Job.ClearOwnership for why that livelocks. + j.ClearOwnership() if updateErr := js.UpdateJob(ctx.Context(), j); updateErr != nil { return nil, fmt.Errorf("retry job: %w", updateErr) } diff --git a/artifact/accessor.go b/artifact/accessor.go new file mode 100644 index 0000000..5813f94 --- /dev/null +++ b/artifact/accessor.go @@ -0,0 +1,69 @@ +package artifact + +import ( + "context" + "io" +) + +// Accessor is the handler-facing face of the artifact plane. A handler +// obtains one with From(ctx). +type Accessor interface { + // Path returns the local file path of a staged input. It returns an + // empty string for an input that was not declared, was not bound, or + // was declared lazy. + Path(name string) string + + // Open streams an input's bytes. It works for both staging modes: + // a path-staged input reads from local disk, a lazy one from the + // backend. + Open(ctx context.Context, name string) (io.ReadCloser, error) + + // Ref returns the artifact bound to an input. + Ref(name string) (Ref, bool) + + // Create begins writing an output owned by the running job. Outputs + // are imperative rather than declared so a handler can produce a + // number of them it only discovers at run time. + Create(ctx context.Context, name string, opts ...CreateOption) (*CommitWriter, error) + + // Existing returns an artifact a previous attempt committed under this + // name, letting a retried handler skip work it already did. + Existing(ctx context.Context, name string) (Ref, bool) +} + +type accessorKey struct{} + +// WithAccessor attaches an Accessor to a context. +func WithAccessor(ctx context.Context, a Accessor) context.Context { + return context.WithValue(ctx, accessorKey{}, a) +} + +// From returns the Accessor for the running job. +// +// It never returns nil. When the artifact plane is disabled, or the job +// declared no inputs, it returns a no-op Accessor so a handler calling +// From(ctx).Path("x") gets an empty string rather than a panic. +func From(ctx context.Context) Accessor { + if a, ok := ctx.Value(accessorKey{}).(Accessor); ok && a != nil { + return a + } + + return noopAccessor{} +} + +// noopAccessor stands in when no artifact plane is configured. +type noopAccessor struct{} + +func (noopAccessor) Path(string) string { return "" } + +func (noopAccessor) Open(context.Context, string) (io.ReadCloser, error) { + return nil, ErrNoBackend +} + +func (noopAccessor) Ref(string) (Ref, bool) { return Ref{}, false } + +func (noopAccessor) Create(context.Context, string, ...CreateOption) (*CommitWriter, error) { + return nil, ErrNoBackend +} + +func (noopAccessor) Existing(context.Context, string) (Ref, bool) { return Ref{}, false } diff --git a/artifact/artifact.go b/artifact/artifact.go new file mode 100644 index 0000000..6b09765 --- /dev/null +++ b/artifact/artifact.go @@ -0,0 +1,167 @@ +package artifact + +import ( + "time" + + "github.com/xraph/dispatch/id" +) + +// Lifecycle determines whether Dispatch may delete an artifact's bytes. +type Lifecycle string + +const ( + // Durable artifacts are written by the application and merely tracked + // by Dispatch. They are read-only here and are never swept. + Durable Lifecycle = "durable" + + // Ephemeral artifacts are created by Dispatch on a handler's behalf. + // They are refcounted through links and swept once every owner is + // terminal and the retention window has passed. + Ephemeral Lifecycle = "ephemeral" +) + +// Valid reports whether the lifecycle is a recognised value. +func (l Lifecycle) Valid() bool { + return l == Durable || l == Ephemeral +} + +// Role describes how an owner relates to an artifact. +type Role string + +const ( + // RoleInput marks an artifact consumed by the owner. + RoleInput Role = "input" + // RoleOutput marks an artifact produced by the owner. + RoleOutput Role = "output" + // RoleIntermediate marks an artifact passed between workflow steps. + RoleIntermediate Role = "intermediate" +) + +// Valid reports whether the role is a recognised value. +func (r Role) Valid() bool { + return r == RoleInput || r == RoleOutput || r == RoleIntermediate +} + +// OwnerKind identifies which entity owns a link. +type OwnerKind string + +const ( + // OwnerJob links an artifact to a job. + OwnerJob OwnerKind = "job" + // OwnerRun links an artifact to a workflow run. + OwnerRun OwnerKind = "run" + // OwnerStep links an artifact to a single workflow step. + OwnerStep OwnerKind = "step" +) + +// Valid reports whether the owner kind is a recognised value. +func (k OwnerKind) Valid() bool { + return k == OwnerJob || k == OwnerRun || k == OwnerStep +} + +// Ref is a lightweight handle to a tracked artifact. It is what callers +// pass to Bind, what handlers receive from Commit, and what workflow +// steps store in checkpoints — small enough to serialise freely. +type Ref struct { + ID id.ArtifactID `json:"id"` + Backend string `json:"backend"` + Bucket string `json:"bucket"` + Key string `json:"key"` + Size int64 `json:"size"` + ContentHash string `json:"content_hash,omitempty"` +} + +// IsZero reports whether the ref is unset. +func (r Ref) IsZero() bool { return r.ID.IsNil() } + +// Artifact is a tracked object in external storage. +type Artifact struct { + ID id.ArtifactID `json:"id"` + Backend string `json:"backend"` + Bucket string `json:"bucket"` + Key string `json:"key"` + Size int64 `json:"size"` + ContentHash string `json:"content_hash,omitempty"` + ContentType string `json:"content_type,omitempty"` + Lifecycle Lifecycle `json:"lifecycle"` + ScopeAppID string `json:"scope_app_id,omitempty"` + ScopeOrgID string `json:"scope_org_id,omitempty"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` + CreatedAt time.Time `json:"created_at"` + DeletedAt *time.Time `json:"deleted_at,omitempty"` +} + +// Ref returns a lightweight handle to this artifact. +func (a *Artifact) Ref() Ref { + return Ref{ + ID: a.ID, + Backend: a.Backend, + Bucket: a.Bucket, + Key: a.Key, + Size: a.Size, + ContentHash: a.ContentHash, + } +} + +// IsDeleted reports whether the artifact has been soft-deleted by the +// sweeper. A soft-deleted artifact is no longer served but its bytes +// survive until the purge pass. +func (a *Artifact) IsDeleted() bool { return a.DeletedAt != nil } + +// Clone returns a deep copy so stores can hand out values callers may +// safely mutate. +func (a *Artifact) Clone() *Artifact { + if a == nil { + return nil + } + + out := *a + + if a.ExpiresAt != nil { + t := *a.ExpiresAt + out.ExpiresAt = &t + } + + if a.DeletedAt != nil { + t := *a.DeletedAt + out.DeletedAt = &t + } + + return &out +} + +// Link records that an owner references an artifact in a given role. +// Attempt scopes the link to one execution attempt so a retried job's +// outputs do not collide with its previous attempt's. +type Link struct { + ArtifactID id.ArtifactID `json:"artifact_id"` + OwnerKind OwnerKind `json:"owner_kind"` + OwnerID string `json:"owner_id"` + Role Role `json:"role"` + Name string `json:"name"` + Attempt int `json:"attempt"` + CreatedAt time.Time `json:"created_at"` +} + +// Clone returns a copy of the link. +func (l *Link) Clone() *Link { + if l == nil { + return nil + } + + out := *l + + return &out +} + +// Owner returns the OwnerRef this link belongs to. +func (l *Link) Owner() OwnerRef { + return OwnerRef{Kind: l.OwnerKind, ID: l.OwnerID} +} + +// ObjectInfo is what a Backend reports about a stored object. +type ObjectInfo struct { + Size int64 + ContentType string + ETag string +} diff --git a/artifact/artifact_test.go b/artifact/artifact_test.go new file mode 100644 index 0000000..61f9880 --- /dev/null +++ b/artifact/artifact_test.go @@ -0,0 +1,140 @@ +package artifact_test + +import ( + "testing" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +func TestArtifactRef(t *testing.T) { + aid := id.NewArtifactID() + a := &artifact.Artifact{ + ID: aid, + Backend: "primary", + Bucket: "models", + Key: "tower.ifc", + Size: 2 << 30, + ContentHash: "blake3:9f2a", + Lifecycle: artifact.Durable, + CreatedAt: time.Now().UTC(), + } + + ref := a.Ref() + if ref.ID != aid { + t.Fatalf("ref.ID = %v, want %v", ref.ID, aid) + } + + if ref.Size != 2<<30 { + t.Fatalf("ref.Size = %d, want %d", ref.Size, int64(2)<<30) + } + + if ref.Key != "tower.ifc" { + t.Fatalf("ref.Key = %q, want %q", ref.Key, "tower.ifc") + } + + if ref.IsZero() { + t.Fatal("ref with an ID reported as zero") + } +} + +func TestLifecycleValid(t *testing.T) { + tests := []struct { + name string + lc artifact.Lifecycle + want bool + }{ + {"durable", artifact.Durable, true}, + {"ephemeral", artifact.Ephemeral, true}, + {"empty", artifact.Lifecycle(""), false}, + {"garbage", artifact.Lifecycle("permanent"), false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.lc.Valid(); got != tt.want { + t.Fatalf("Valid() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestRoleValid(t *testing.T) { + tests := []struct { + name string + role artifact.Role + want bool + }{ + {"input", artifact.RoleInput, true}, + {"output", artifact.RoleOutput, true}, + {"intermediate", artifact.RoleIntermediate, true}, + {"empty", artifact.Role(""), false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.role.Valid(); got != tt.want { + t.Fatalf("Valid() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestOwnerKindValid(t *testing.T) { + tests := []struct { + name string + kind artifact.OwnerKind + want bool + }{ + {"job", artifact.OwnerJob, true}, + {"run", artifact.OwnerRun, true}, + {"step", artifact.OwnerStep, true}, + {"empty", artifact.OwnerKind(""), false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.kind.Valid(); got != tt.want { + t.Fatalf("Valid() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestArtifactIsDeleted(t *testing.T) { + a := &artifact.Artifact{} + if a.IsDeleted() { + t.Fatal("fresh artifact reported deleted") + } + + now := time.Now().UTC() + a.DeletedAt = &now + + if !a.IsDeleted() { + t.Fatal("soft-deleted artifact not reported deleted") + } +} + +func TestArtifactCloneIsDeep(t *testing.T) { + now := time.Now().UTC() + a := &artifact.Artifact{ + ID: id.NewArtifactID(), + Lifecycle: artifact.Ephemeral, + ExpiresAt: &now, + DeletedAt: &now, + } + + clone := a.Clone() + later := now.Add(time.Hour) + *clone.ExpiresAt = later + *clone.DeletedAt = later + + if a.ExpiresAt.Equal(later) { + t.Fatal("Clone shares the ExpiresAt pointer") + } + + if a.DeletedAt.Equal(later) { + t.Fatal("Clone shares the DeletedAt pointer") + } +} diff --git a/artifact/artifacttest/backend.go b/artifact/artifacttest/backend.go new file mode 100644 index 0000000..69f2c52 --- /dev/null +++ b/artifact/artifacttest/backend.go @@ -0,0 +1,182 @@ +package artifacttest + +import ( + "bytes" + "context" + "io" + "sync" + "sync/atomic" + "time" + + "github.com/xraph/dispatch/artifact" +) + +// Backend is an in-memory artifact.Backend for tests. It counts calls so +// tests can assert on caching and single-flight behaviour. +type Backend struct { + // DelayOpen makes Open sleep before returning, which is what lets a + // single-flight test observe concurrent stagers colliding. + DelayOpen time.Duration + + // DenyOpen makes Open fail with ErrPermissionDenied even when the + // object is present. It exists so tests can cover a permanent failure + // that is not a missing object: those two must be retried the same + // way, which is to say not at all. + DenyOpen bool + + mu sync.Mutex + objects map[string][]byte + + opens atomic.Int64 + creates atomic.Int64 + deletes atomic.Int64 + stats atomic.Int64 +} + +// Compile-time check that the double satisfies the contract. +var _ artifact.Backend = (*Backend)(nil) + +// NewBackend returns an empty in-memory backend. +func NewBackend() *Backend { + return &Backend{objects: make(map[string][]byte)} +} + +// Name identifies this backend. +func (b *Backend) Name() string { return "memory" } + +// Opens returns how many times Open was called. +func (b *Backend) Opens() int64 { return b.opens.Load() } + +// Creates returns how many times Create was called. +func (b *Backend) Creates() int64 { return b.creates.Load() } + +// Deletes returns how many times Delete was called. +func (b *Backend) Deletes() int64 { return b.deletes.Load() } + +// Stats returns how many times Stat was called. +func (b *Backend) Stats() int64 { return b.stats.Load() } + +// Put seeds an object directly, bypassing the Writer path. +func (b *Backend) Put(bucket, key string, data []byte) { + b.mu.Lock() + defer b.mu.Unlock() + + b.objects[objectKey(bucket, key)] = append([]byte(nil), data...) +} + +// Has reports whether an object exists. +func (b *Backend) Has(bucket, key string) bool { + b.mu.Lock() + defer b.mu.Unlock() + + _, ok := b.objects[objectKey(bucket, key)] + + return ok +} + +func objectKey(bucket, key string) string { return bucket + "/" + key } + +// Open returns a reader over the object's bytes. +func (b *Backend) Open(ctx context.Context, ref artifact.Ref) (io.ReadCloser, error) { + b.opens.Add(1) + + if b.DelayOpen > 0 { + select { + case <-time.After(b.DelayOpen): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + + if b.DenyOpen { + return nil, artifact.ErrPermissionDenied + } + + b.mu.Lock() + data, ok := b.objects[objectKey(ref.Bucket, ref.Key)] + b.mu.Unlock() + + if !ok { + return nil, artifact.ErrNotFound + } + + return io.NopCloser(bytes.NewReader(data)), nil +} + +// Stat reports the object's size. +func (b *Backend) Stat(_ context.Context, ref artifact.Ref) (artifact.ObjectInfo, error) { + b.stats.Add(1) + + b.mu.Lock() + data, ok := b.objects[objectKey(ref.Bucket, ref.Key)] + b.mu.Unlock() + + if !ok { + return artifact.ObjectInfo{}, artifact.ErrNotFound + } + + return artifact.ObjectInfo{Size: int64(len(data))}, nil +} + +// Delete removes an object. Deleting a missing object is not an error. +func (b *Backend) Delete(_ context.Context, ref artifact.Ref) error { + b.deletes.Add(1) + + b.mu.Lock() + defer b.mu.Unlock() + + delete(b.objects, objectKey(ref.Bucket, ref.Key)) + + return nil +} + +// Create begins writing a new object. +func (b *Backend) Create(_ context.Context, bucket, key string) (artifact.Writer, error) { + b.creates.Add(1) + + return &memWriter{backend: b, bucket: bucket, key: key}, nil +} + +// memWriter buffers writes and only publishes on Commit. +type memWriter struct { + backend *Backend + bucket string + key string + buf bytes.Buffer + done bool +} + +func (w *memWriter) Write(p []byte) (int, error) { + if w.done { + return 0, io.ErrClosedPipe + } + + return w.buf.Write(p) +} + +func (w *memWriter) Commit(_ context.Context) (artifact.ObjectInfo, error) { + if w.done { + return artifact.ObjectInfo{}, io.ErrClosedPipe + } + + w.done = true + + size := int64(w.buf.Len()) + w.backend.Put(w.bucket, w.key, w.buf.Bytes()) + w.buf.Reset() + + return artifact.ObjectInfo{Size: size}, nil +} + +// Abort discards the partial object. It is a no-op after Commit, so +// `defer w.Abort()` is safe alongside a successful commit. +func (w *memWriter) Abort() error { + if w.done { + return nil + } + + w.done = true + w.buf.Reset() + + return nil +} diff --git a/artifact/artifacttest/backend_test.go b/artifact/artifacttest/backend_test.go new file mode 100644 index 0000000..6acbe6d --- /dev/null +++ b/artifact/artifacttest/backend_test.go @@ -0,0 +1,131 @@ +package artifacttest_test + +import ( + "bytes" + "context" + "errors" + "io" + "testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" +) + +func TestBackendRoundTrip(t *testing.T) { + ctx := context.Background() + b := artifacttest.NewBackend() + b.Put("models", "tower.ifc", []byte("hello")) + + ref := artifact.Ref{Backend: b.Name(), Bucket: "models", Key: "tower.ifc"} + + rc, err := b.Open(ctx, ref) + if err != nil { + t.Fatalf("Open: %v", err) + } + + got, err := io.ReadAll(rc) + if cerr := rc.Close(); cerr != nil { + t.Fatalf("Close: %v", cerr) + } + + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + + if !bytes.Equal(got, []byte("hello")) { + t.Fatalf("read %q, want %q", got, "hello") + } + + if b.Opens() != 1 { + t.Fatalf("Opens() = %d, want 1", b.Opens()) + } +} + +func TestBackendOpenMissing(t *testing.T) { + _, err := artifacttest.NewBackend().Open(context.Background(), + artifact.Ref{Bucket: "models", Key: "nope"}) + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Open(missing) = %v, want ErrNotFound", err) + } +} + +func TestBackendStat(t *testing.T) { + b := artifacttest.NewBackend() + b.Put("models", "tower.ifc", []byte("0123456789")) + + info, err := b.Stat(context.Background(), + artifact.Ref{Bucket: "models", Key: "tower.ifc"}) + if err != nil { + t.Fatalf("Stat: %v", err) + } + + if info.Size != 10 { + t.Fatalf("info.Size = %d, want 10", info.Size) + } + + _, err = b.Stat(context.Background(), artifact.Ref{Bucket: "models", Key: "nope"}) + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Stat(missing) = %v, want ErrNotFound", err) + } +} + +func TestBackendWriterCommitThenAbortIsNoOp(t *testing.T) { + ctx := context.Background() + b := artifacttest.NewBackend() + + w, err := b.Create(ctx, "models", "mesh.glb") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write([]byte("meshdata")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + info, err := w.Commit(ctx) + if err != nil { + t.Fatalf("Commit: %v", err) + } + + if info.Size != 8 { + t.Fatalf("info.Size = %d, want 8", info.Size) + } + + if err := w.Abort(); err != nil { + t.Fatalf("Abort after Commit must be a no-op, got %v", err) + } + + if !b.Has("models", "mesh.glb") { + t.Fatal("committed object is missing") + } +} + +func TestBackendWriterAbortPublishesNothing(t *testing.T) { + ctx := context.Background() + b := artifacttest.NewBackend() + + w, err := b.Create(ctx, "models", "aborted.glb") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, err := w.Write([]byte("partial")); err != nil { + t.Fatalf("Write: %v", err) + } + + if err := w.Abort(); err != nil { + t.Fatalf("Abort: %v", err) + } + + if _, err := b.Open(ctx, artifact.Ref{Bucket: "models", Key: "aborted.glb"}); !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("aborted object is readable; Open = %v, want ErrNotFound", err) + } +} + +func TestBackendDeleteMissingIsNotAnError(t *testing.T) { + err := artifacttest.NewBackend().Delete(context.Background(), + artifact.Ref{Bucket: "models", Key: "absent"}) + if err != nil { + t.Fatalf("Delete(missing) = %v, want nil", err) + } +} diff --git a/artifact/artifacttest/doc.go b/artifact/artifacttest/doc.go new file mode 100644 index 0000000..fbe7430 --- /dev/null +++ b/artifact/artifacttest/doc.go @@ -0,0 +1,8 @@ +// Package artifacttest provides a shared conformance suite and test +// doubles for artifact storage. +// +// Every artifact.Store implementation runs RunStoreSuite so all five +// backends are held to one contract. The suite's most important case is +// SweepNeverTouchesDurable: no sweep path, under any input, may mark a +// durable artifact. +package artifacttest diff --git a/artifact/artifacttest/suite.go b/artifact/artifacttest/suite.go new file mode 100644 index 0000000..c0a2b45 --- /dev/null +++ b/artifact/artifacttest/suite.go @@ -0,0 +1,577 @@ +package artifacttest + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// RunStoreSuite exercises the artifact.Store contract. newStore must +// return a fresh, empty store on every call. +// +// newStore receives the subtest's *testing.T, not the parent's. Backends +// that stand up a container or open a database per store must register +// teardown on that T so it runs when the subtest ends; closing over the +// parent T instead would hold every subtest's resources open until the +// whole suite finished, and would turn a setup t.Fatalf into a FailNow +// on a parent test. +func RunStoreSuite(t *testing.T, newStore func(t *testing.T) artifact.Store) { + t.Helper() + + tests := []struct { + name string + fn func(*testing.T, artifact.Store) + }{ + {"CreateAndGet", testCreateAndGet}, + {"CreateDuplicateKey", testCreateDuplicateKey}, + {"GetMissing", testGetMissing}, + {"FindByKey", testFindByKey}, + {"UpdateHash", testUpdateHash}, + {"LinkAndList", testLinkAndList}, + {"LinkIdempotent", testLinkIdempotent}, + {"FindLinkByNameAcrossAttempts", testFindLinkAcrossAttempts}, + {"FindLinkByNameTieBreaksOnLatestWrite", testFindLinkTieBreaksOnLatestWrite}, + {"ListArtifacts", testListArtifacts}, + {"SweepNeverTouchesDurable", testSweepNeverTouchesDurable}, + {"SweepOrphans", testSweepOrphans}, + {"SweepOrphansSkipsLinked", testSweepOrphansSkipsLinked}, + {"SweepOrphansDryRun", testSweepOrphansRespectsLimit}, + {"PurgeFlow", testPurgeFlow}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tt.fn(t, newStore(t)) + }) + } +} + +func newArtifact(key string, lc artifact.Lifecycle) *artifact.Artifact { + return &artifact.Artifact{ + ID: id.NewArtifactID(), + Backend: "primary", + Bucket: "models", + Key: key, + Size: 1024, + Lifecycle: lc, + CreatedAt: time.Now().UTC(), + } +} + +func newOwner() artifact.OwnerRef { + return artifact.OwnerRef{Kind: artifact.OwnerJob, ID: id.NewJobID().String()} +} + +func testCreateAndGet(t *testing.T, s artifact.Store) { + ctx := context.Background() + a := newArtifact("tower.ifc", artifact.Durable) + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact: %v", err) + } + + got, err := s.GetArtifact(ctx, a.ID) + if err != nil { + t.Fatalf("GetArtifact: %v", err) + } + + if got.Key != a.Key || got.Size != a.Size || got.Lifecycle != a.Lifecycle { + t.Fatalf("round trip mismatch: got %+v want %+v", got, a) + } +} + +func testCreateDuplicateKey(t *testing.T, s artifact.Store) { + ctx := context.Background() + a := newArtifact("dup.ifc", artifact.Durable) + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("first CreateArtifact: %v", err) + } + + b := newArtifact("dup.ifc", artifact.Durable) + if err := s.CreateArtifact(ctx, b, nil); !errors.Is(err, artifact.ErrExists) { + t.Fatalf("duplicate key error = %v, want ErrExists", err) + } +} + +func testGetMissing(t *testing.T, s artifact.Store) { + if _, err := s.GetArtifact(context.Background(), id.NewArtifactID()); !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("GetArtifact(missing) = %v, want ErrNotFound", err) + } +} + +func testFindByKey(t *testing.T, s artifact.Store) { + ctx := context.Background() + a := newArtifact("find.ifc", artifact.Durable) + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact: %v", err) + } + + got, err := s.FindArtifactByKey(ctx, "primary", "models", "find.ifc") + if err != nil { + t.Fatalf("FindArtifactByKey: %v", err) + } + + if got.ID != a.ID { + t.Fatalf("FindArtifactByKey ID = %v, want %v", got.ID, a.ID) + } + + _, err = s.FindArtifactByKey(ctx, "primary", "models", "nope.ifc") + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("FindArtifactByKey(missing) = %v, want ErrNotFound", err) + } +} + +func testUpdateHash(t *testing.T, s artifact.Store) { + ctx := context.Background() + a := newArtifact("hash.ifc", artifact.Durable) + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact: %v", err) + } + + a.ContentHash = "blake3:9f2a" + a.Size = 4096 + + if err := s.UpdateArtifact(ctx, a); err != nil { + t.Fatalf("UpdateArtifact: %v", err) + } + + got, err := s.GetArtifact(ctx, a.ID) + if err != nil { + t.Fatalf("GetArtifact: %v", err) + } + + if got.ContentHash != "blake3:9f2a" || got.Size != 4096 { + t.Fatalf("update not persisted: hash=%q size=%d", got.ContentHash, got.Size) + } +} + +func testLinkAndList(t *testing.T, s artifact.Store) { + ctx := context.Background() + a := newArtifact("linked.ifc", artifact.Ephemeral) + owner := newOwner() + link := &artifact.Link{ + ArtifactID: a.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleOutput, + Name: "mesh.glb", + CreatedAt: time.Now().UTC(), + } + + if err := s.CreateArtifact(ctx, a, link); err != nil { + t.Fatalf("CreateArtifact with link: %v", err) + } + + links, err := s.ListLinks(ctx, owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + if len(links) != 1 || links[0].Name != "mesh.glb" { + t.Fatalf("ListLinks = %+v, want one link named mesh.glb", links) + } + + arts, err := s.ListArtifactsByOwner(ctx, owner, artifact.RoleOutput) + if err != nil { + t.Fatalf("ListArtifactsByOwner: %v", err) + } + + if len(arts) != 1 || arts[0].ID != a.ID { + t.Fatalf("ListArtifactsByOwner = %+v, want artifact %v", arts, a.ID) + } + + none, err := s.ListArtifactsByOwner(ctx, owner, artifact.RoleInput) + if err != nil { + t.Fatalf("ListArtifactsByOwner(input): %v", err) + } + + if len(none) != 0 { + t.Fatalf("ListArtifactsByOwner(input) = %+v, want empty", none) + } +} + +func testLinkIdempotent(t *testing.T, s artifact.Store) { + ctx := context.Background() + a := newArtifact("idem.ifc", artifact.Ephemeral) + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact: %v", err) + } + + owner := newOwner() + link := &artifact.Link{ + ArtifactID: a.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleOutput, + Name: "out.bin", + CreatedAt: time.Now().UTC(), + } + + for i := range 2 { + if err := s.LinkArtifact(ctx, link); err != nil { + t.Fatalf("LinkArtifact call %d: %v", i, err) + } + } + + links, err := s.ListLinks(ctx, owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + if len(links) != 1 { + t.Fatalf("ListLinks returned %d links, want 1 (link must be idempotent)", len(links)) + } +} + +func testFindLinkAcrossAttempts(t *testing.T, s artifact.Store) { + ctx := context.Background() + owner := newOwner() + + for attempt := range 3 { + a := newArtifact(fmt.Sprintf("page-317-attempt-%d.png", attempt), artifact.Ephemeral) + link := &artifact.Link{ + ArtifactID: a.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleOutput, + Name: "page-317.png", + Attempt: attempt, + CreatedAt: time.Now().UTC(), + } + + if err := s.CreateArtifact(ctx, a, link); err != nil { + t.Fatalf("CreateArtifact attempt %d: %v", attempt, err) + } + } + + got, err := s.FindLinkByName(ctx, owner, "page-317.png") + if err != nil { + t.Fatalf("FindLinkByName: %v", err) + } + + if got.Attempt != 2 { + t.Fatalf("FindLinkByName attempt = %d, want 2 (highest)", got.Attempt) + } + + _, err = s.FindLinkByName(ctx, owner, "never-made.png") + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("FindLinkByName(missing) = %v, want ErrNotFound", err) + } +} + +// testFindLinkTieBreaksOnLatestWrite reproduces the seam a lost lease +// opens: track D's reclaim increments EvictCount, never RetryCount, so a +// reclaimed (zombie) holder and the worker it was fenced out for can both +// commit a link at the identical (OwnerKind, OwnerID, Name, Attempt) — +// two different ArtifactIDs, since that tuple alone is not the storage +// key. FindLinkByName must resolve to whichever wrote last, not to +// whichever a backend's map or query happens to enumerate first. +func testFindLinkTieBreaksOnLatestWrite(t *testing.T, s artifact.Store) { + ctx := context.Background() + owner := newOwner() + + zombie := newArtifact("zombie.bin", artifact.Ephemeral) + if err := s.CreateArtifact(ctx, zombie, nil); err != nil { + t.Fatalf("CreateArtifact zombie: %v", err) + } + + winner := newArtifact("winner.bin", artifact.Ephemeral) + if err := s.CreateArtifact(ctx, winner, nil); err != nil { + t.Fatalf("CreateArtifact winner: %v", err) + } + + base := time.Now().UTC() + + if err := s.LinkArtifact(ctx, &artifact.Link{ + ArtifactID: zombie.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleOutput, + Name: "output.bin", + Attempt: 3, + CreatedAt: base, + }); err != nil { + t.Fatalf("LinkArtifact zombie: %v", err) + } + + if err := s.LinkArtifact(ctx, &artifact.Link{ + ArtifactID: winner.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleOutput, + Name: "output.bin", + Attempt: 3, + CreatedAt: base.Add(time.Second), + }); err != nil { + t.Fatalf("LinkArtifact winner: %v", err) + } + + got, err := s.FindLinkByName(ctx, owner, "output.bin") + if err != nil { + t.Fatalf("FindLinkByName: %v", err) + } + + if got.ArtifactID != winner.ID { + t.Fatalf("FindLinkByName resolved to artifact %v, want the later write %v (winner); "+ + "got the earlier write %v (zombie)", got.ArtifactID, winner.ID, zombie.ID) + } +} + +func testListArtifacts(t *testing.T, s artifact.Store) { + ctx := context.Background() + + for i := range 3 { + a := newArtifact(fmt.Sprintf("dur-%d.ifc", i), artifact.Durable) + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact durable %d: %v", i, err) + } + } + + eph := newArtifact("eph.bin", artifact.Ephemeral) + if err := s.CreateArtifact(ctx, eph, nil); err != nil { + t.Fatalf("CreateArtifact ephemeral: %v", err) + } + + all, err := s.ListArtifacts(ctx, artifact.ListOpts{}) + if err != nil { + t.Fatalf("ListArtifacts: %v", err) + } + + if len(all) != 4 { + t.Fatalf("ListArtifacts returned %d, want 4", len(all)) + } + + durable, err := s.ListArtifacts(ctx, artifact.ListOpts{Lifecycle: artifact.Durable}) + if err != nil { + t.Fatalf("ListArtifacts(durable): %v", err) + } + + if len(durable) != 3 { + t.Fatalf("ListArtifacts(durable) returned %d, want 3", len(durable)) + } + + limited, err := s.ListArtifacts(ctx, artifact.ListOpts{Limit: 2}) + if err != nil { + t.Fatalf("ListArtifacts(limit 2): %v", err) + } + + if len(limited) != 2 { + t.Fatalf("ListArtifacts(limit 2) returned %d, want 2", len(limited)) + } +} + +// testSweepNeverTouchesDurable is the safety invariant of the whole +// design. A durable artifact must be unreachable from any sweep path, +// regardless of age, links, or owner state. +func testSweepNeverTouchesDurable(t *testing.T, s artifact.Store) { + ctx := context.Background() + long := time.Now().UTC().Add(-365 * 24 * time.Hour) + + durable := newArtifact("customer-upload.ifc", artifact.Durable) + durable.CreatedAt = long + + if err := s.CreateArtifact(ctx, durable, nil); err != nil { + t.Fatalf("CreateArtifact durable: %v", err) + } + + // A durable artifact linked to an owner that no longer exists is the + // most tempting sweep candidate there is. It must still survive. + linked := newArtifact("customer-upload-2.ifc", artifact.Durable) + linked.CreatedAt = long + owner := newOwner() + link := &artifact.Link{ + ArtifactID: linked.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleInput, + Name: "model", + CreatedAt: long, + } + + if err := s.CreateArtifact(ctx, linked, link); err != nil { + t.Fatalf("CreateArtifact linked durable: %v", err) + } + + swept, err := s.SweepEphemeral(ctx, artifact.SweepOpts{Retention: 0, Limit: 100}) + if err != nil { + t.Fatalf("SweepEphemeral: %v", err) + } + + orphaned, err := s.SweepOrphans(ctx, time.Now().UTC(), 100) + if err != nil { + t.Fatalf("SweepOrphans: %v", err) + } + + for _, a := range append(swept, orphaned...) { + if a.Lifecycle == artifact.Durable { + t.Fatalf("a sweep marked DURABLE artifact %v — safety invariant violated", a.ID) + } + } + + for _, want := range []*artifact.Artifact{durable, linked} { + got, err := s.GetArtifact(ctx, want.ID) + if err != nil { + t.Fatalf("durable artifact %v not retrievable after sweeps: %v", want.ID, err) + } + + if got.IsDeleted() { + t.Fatalf("durable artifact %v was soft-deleted — safety invariant violated", want.ID) + } + } +} + +func testSweepOrphans(t *testing.T, s artifact.Store) { + ctx := context.Background() + + old := newArtifact("orphan.bin", artifact.Ephemeral) + old.CreatedAt = time.Now().UTC().Add(-48 * time.Hour) + + if err := s.CreateArtifact(ctx, old, nil); err != nil { + t.Fatalf("CreateArtifact: %v", err) + } + + fresh := newArtifact("fresh.bin", artifact.Ephemeral) + if err := s.CreateArtifact(ctx, fresh, nil); err != nil { + t.Fatalf("CreateArtifact fresh: %v", err) + } + + cutoff := time.Now().UTC().Add(-24 * time.Hour) + + swept, err := s.SweepOrphans(ctx, cutoff, 100) + if err != nil { + t.Fatalf("SweepOrphans: %v", err) + } + + if len(swept) != 1 || swept[0].ID != old.ID { + t.Fatalf("SweepOrphans = %+v, want only the 48h-old orphan", swept) + } + + got, err := s.GetArtifact(ctx, fresh.ID) + if err != nil { + t.Fatalf("fresh orphan was swept: %v", err) + } + + if got.IsDeleted() { + t.Fatal("fresh orphan soft-deleted before its grace window elapsed") + } +} + +func testSweepOrphansSkipsLinked(t *testing.T, s artifact.Store) { + ctx := context.Background() + long := time.Now().UTC().Add(-48 * time.Hour) + + a := newArtifact("has-link.bin", artifact.Ephemeral) + a.CreatedAt = long + owner := newOwner() + link := &artifact.Link{ + ArtifactID: a.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleOutput, + Name: "out.bin", + CreatedAt: long, + } + + if err := s.CreateArtifact(ctx, a, link); err != nil { + t.Fatalf("CreateArtifact: %v", err) + } + + swept, err := s.SweepOrphans(ctx, time.Now().UTC(), 100) + if err != nil { + t.Fatalf("SweepOrphans: %v", err) + } + + for _, got := range swept { + if got.ID == a.ID { + t.Fatal("SweepOrphans marked a linked artifact — it is not an orphan") + } + } +} + +func testSweepOrphansRespectsLimit(t *testing.T, s artifact.Store) { + ctx := context.Background() + long := time.Now().UTC().Add(-48 * time.Hour) + + for i := range 5 { + a := newArtifact(fmt.Sprintf("orphan-%d.bin", i), artifact.Ephemeral) + a.CreatedAt = long + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact %d: %v", i, err) + } + } + + swept, err := s.SweepOrphans(ctx, time.Now().UTC(), 2) + if err != nil { + t.Fatalf("SweepOrphans: %v", err) + } + + if len(swept) != 2 { + t.Fatalf("SweepOrphans with limit 2 returned %d", len(swept)) + } +} + +func testPurgeFlow(t *testing.T, s artifact.Store) { + ctx := context.Background() + + a := newArtifact("purge.bin", artifact.Ephemeral) + a.CreatedAt = time.Now().UTC().Add(-72 * time.Hour) + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact: %v", err) + } + + if _, err := s.SweepOrphans(ctx, time.Now().UTC().Add(-24*time.Hour), 100); err != nil { + t.Fatalf("SweepOrphans: %v", err) + } + + // A soft-deleted artifact is no longer served. + if _, err := s.GetArtifact(ctx, a.ID); !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("GetArtifact after soft delete = %v, want ErrNotFound", err) + } + + // With a grace window longer than its age, it is not yet purgeable. + notYet, err := s.ListPurgeable(ctx, time.Hour, 100) + if err != nil { + t.Fatalf("ListPurgeable(1h grace): %v", err) + } + + if len(notYet) != 0 { + t.Fatalf("ListPurgeable(1h grace) = %+v, want empty", notYet) + } + + purgeable, err := s.ListPurgeable(ctx, 0, 100) + if err != nil { + t.Fatalf("ListPurgeable: %v", err) + } + + if len(purgeable) != 1 || purgeable[0].ID != a.ID { + t.Fatalf("ListPurgeable = %+v, want the swept artifact", purgeable) + } + + if perr := s.PurgeArtifact(ctx, a.ID); perr != nil { + t.Fatalf("PurgeArtifact: %v", perr) + } + + if _, gerr := s.GetArtifact(ctx, a.ID); !errors.Is(gerr, artifact.ErrNotFound) { + t.Fatalf("GetArtifact after purge = %v, want ErrNotFound", gerr) + } + + after, err := s.ListPurgeable(ctx, 0, 100) + if err != nil { + t.Fatalf("ListPurgeable after purge: %v", err) + } + + if len(after) != 0 { + t.Fatalf("ListPurgeable after purge = %+v, want empty", after) + } +} diff --git a/artifact/backend.go b/artifact/backend.go new file mode 100644 index 0000000..4c4dd11 --- /dev/null +++ b/artifact/backend.go @@ -0,0 +1,65 @@ +package artifact + +import ( + "context" + "io" + "time" +) + +// Backend is the pluggable object-storage contract behind an artifact. +// Dispatch ships an adapter for Trove; any store can implement this. +type Backend interface { + // Name returns the backend's identifier, recorded in Artifact.Backend. + Name() string + + // Open returns a reader over the object's bytes. It returns + // ErrNotFound if the object does not exist. + Open(ctx context.Context, ref Ref) (io.ReadCloser, error) + + // Create begins writing a new object. The bytes are not visible until + // Commit. Callers must call Commit or Abort. + Create(ctx context.Context, bucket, key string) (Writer, error) + + // Stat reports the object's size and content type without reading it. + // It returns ErrNotFound if the object does not exist. + Stat(ctx context.Context, ref Ref) (ObjectInfo, error) + + // Delete removes the object. Deleting a missing object is not an error. + Delete(ctx context.Context, ref Ref) error +} + +// Writer accumulates bytes for a new object. +// +// Commit reports the logical size of the bytes written, which may differ +// from what the backend stored — compression and encryption middleware +// change the stored form, and the artifact row records what the handler +// produced. +// +// Abort after a successful Commit is a no-op, so `defer w.Abort()` is the +// correct idiom. +type Writer interface { + io.Writer + + // Commit finalises the object and returns its logical info. + Commit(ctx context.Context) (ObjectInfo, error) + + // Abort discards the partial object. It is a no-op after Commit. + Abort() error +} + +// RangeReader is an optional Backend capability for partial reads. +type RangeReader interface { + // OpenRange returns a reader over n bytes starting at off. A negative + // n reads to the end. + OpenRange(ctx context.Context, ref Ref, off, n int64) (io.ReadCloser, error) +} + +// Presigner is an optional Backend capability for direct client access. +// +// It is what lets a remote worker fetch a large object straight from +// object storage instead of streaming it through the coordinator, which +// would otherwise make the coordinator a bandwidth bottleneck. +type Presigner interface { + // PresignGet returns a time-limited URL granting read access. + PresignGet(ctx context.Context, ref Ref, ttl time.Duration) (string, error) +} diff --git a/artifact/cache/cache.go b/artifact/cache/cache.go new file mode 100644 index 0000000..91fa523 --- /dev/null +++ b/artifact/cache/cache.go @@ -0,0 +1,653 @@ +package cache + +import ( + "context" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + + "github.com/zeebo/blake3" + "golang.org/x/sync/singleflight" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" +) + +// DefaultBudget is the disk allowance when none is configured. +const DefaultBudget int64 = 20 << 30 // 20 GiB + +const ( + hashDir = "blake3" + tmpDir = "tmp" + dirPerm = 0o750 +) + +// hashPrefix labels the digest so the algorithm is visible wherever a +// hash is stored or logged. +const hashPrefix = "blake3:" + +// maxStageAttempts bounds Stage's re-download loop. Each attempt means +// an entry was evicted between being staged and being pinned, which +// takes a cache under enough pressure that the next attempt is unlikely +// to fare better. +const maxStageAttempts = 8 + +// Cache stages artifacts to local disk, content-addressed and bounded by +// a byte budget. It is safe for concurrent use. +// +// Every cached object holds one resource.Manager lease for its bytes, +// and eviction returns them by releasing that lease. That is what lets +// the same ledger that admits jobs see the disk the cache is sitting +// on: without it, staged bytes would be invisible to admission and a +// worker would offer capacity it had already spent. +type Cache struct { + dir string + backend artifact.Backend + logger log.Logger + + entries *entryTable + flight singleflight.Group + + // resources is the ledger every cached byte is admitted against. + // With no manager configured this is a private single-key manager + // over the configured allowance, which is exactly the private disk + // budget the cache used to own; with one supplied it is the + // worker's shared ledger. + resources resource.Manager + // allowance is the disk capacity of the private manager, used only + // when no manager is supplied. + allowance int64 + // used totals the bytes this cache holds against the manager, + // in-flight downloads included. The manager's own used counts every + // tenant of the volume, so it cannot answer this question. + used atomic.Int64 + + closeOnce sync.Once +} + +// Cache is the manager's disk reclaimer: it is the component that can +// hand bytes back on demand. +var _ resource.Reclaimer = (*Cache)(nil) + +// Option configures a Cache. +type Option func(*Cache) + +// WithBudget sets the maximum bytes the cache may hold on disk. +// +// It configures the private manager the cache builds for itself, so it +// is ignored when WithManager supplies one: a shared ledger's disk +// capacity is the allowance, and a second ceiling underneath it would +// only be a place for the two to disagree. +func WithBudget(bytes int64) Option { + return func(c *Cache) { + if bytes > 0 { + c.allowance = bytes + } + } +} + +// WithManager admits cached bytes against a shared resource manager +// rather than a private one. +// +// This is what makes staged bytes visible to job admission. The cache +// registers itself as the manager's disk reclaimer and holds a lease +// per cached entry, so a worker sizing up a job is offered the disk +// that is free plus the disk the cache can evict, and redeems the +// second half by evicting. A nil manager leaves the private one in +// place, so a caller threading an optional dependency through does not +// have to branch. +func WithManager(m resource.Manager) Option { + return func(c *Cache) { + if m != nil { + c.resources = m + } + } +} + +// WithLogger sets the logger. +func WithLogger(l log.Logger) Option { + return func(c *Cache) { c.logger = l } +} + +// New opens a cache rooted at dir. +// +// Startup wipes the temp directory and rebuilds the entry table by +// walking the hash directories. The walk is the source of truth: there is +// no persisted index to corrupt, so a crash costs at most a re-download. +func New(dir string, backend artifact.Backend, opts ...Option) (*Cache, error) { + if backend == nil { + return nil, artifact.ErrNoBackend + } + + c := &Cache{ + dir: dir, + backend: backend, + logger: log.NewNoopLogger(), + entries: newEntryTable(), + allowance: DefaultBudget, + } + + for _, opt := range opts { + opt(c) + } + + if c.resources == nil { + c.resources = resource.NewManager(resource.Set{resource.Disk: c.allowance}) + } + + if err := os.MkdirAll(filepath.Join(dir, hashDir), dirPerm); err != nil { + return nil, fmt.Errorf("dispatch/artifact/cache: create hash dir: %w", err) + } + + if err := c.resetTmp(); err != nil { + return nil, err + } + + if err := c.rebuild(); err != nil { + return nil, err + } + + // Registered after the walk, so nothing can evict a file whose lease + // this cache has not taken yet. + c.resources.RegisterReclaimer(resource.Disk, c) + + return c, nil +} + +// Budget returns the configured disk allowance. The engine uses it to +// reject a job definition whose declared inputs could never be staged. +func (c *Cache) Budget() int64 { return c.resources.Capacity()[resource.Disk] } + +// Used returns the bytes currently held on disk. +func (c *Cache) Used() int64 { return c.used.Load() } + +// Dir returns the cache root. +func (c *Cache) Dir() string { return c.dir } + +// resetTmp clears partial downloads left by a previous process. +func (c *Cache) resetTmp() error { + tmp := filepath.Join(c.dir, tmpDir) + + if err := os.RemoveAll(tmp); err != nil { + return fmt.Errorf("dispatch/artifact/cache: clear tmp: %w", err) + } + + if err := os.MkdirAll(tmp, dirPerm); err != nil { + return fmt.Errorf("dispatch/artifact/cache: create tmp: %w", err) + } + + return nil +} + +// rebuild reconstructs the entry table by walking the hash directories. +// +// Each surviving file takes its own lease, so a restart re-admits what +// is on disk rather than starting from an empty ledger while the volume +// is already full. A file the ledger cannot account for is deleted: the +// allowance has shrunk, or something else on this box now holds the +// volume, and keeping bytes nothing knows about is the exact hole this +// accounting exists to close. The cost is a re-download. +func (c *Cache) rebuild() error { + root := filepath.Join(c.dir, hashDir) + + err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if err != nil { + return err + } + + if d.IsDir() { + return nil + } + + info, ierr := d.Info() + if ierr != nil { + return ierr + } + + h, ok := c.tryHold(info.Size()) + if !ok { + c.logger.Warn("dispatch/artifact/cache: dropping unaccountable cached file", + log.String("path", path), log.Int64("bytes", info.Size())) + c.removeQuietly(path) + + return nil + } + + // One file per hash on disk, so this never collides. + c.entries.put(&entry{ + hash: hashPrefix + d.Name(), + path: path, + size: info.Size(), + hold: h, + }, "") + + return nil + }) + if err != nil { + return fmt.Errorf("dispatch/artifact/cache: rebuild index: %w", err) + } + + return nil +} + +// Stage materialises an artifact locally and returns its path, its +// content hash, and a release function the caller must invoke. +// +// The returned path stays valid until release is called. Callers should +// `defer release()` immediately; releasing twice is safe. +func (c *Cache) Stage(ctx context.Context, ref artifact.Ref) (path, hash string, release func(), err error) { + if ref.Bucket == "" && ref.Key == "" { + return "", "", nil, errors.New("dispatch/artifact/cache: ref has no storage coordinates") + } + + backendName := ref.Backend + if backendName == "" { + backendName = c.backend.Name() + } + + coord := coordKey(backendName, ref.Bucket, ref.Key) + + // Pinning is the loop condition, not a step in it. An entry can be + // evicted between being found and being leased, and handing back a + // path whose file has just been unlinked — and whose bytes the + // manager has already credited to someone else — is worse than + // paying for the download again. + // + // Bounded, because the retry is only correct while it is rare. An + // entry evicted before its own stager can pin it means the cache is + // thrashing hard enough that this artifact will not survive being + // fetched however many times we try, and a job that fails saying so + // beats one that downloads forever under a context with no + // deadline. + for range maxStageAttempts { + // Fast path: a ref that already knows its hash, or coordinates + // we have staged before. + if e, ok := c.lookup(ref, coord); ok && c.entries.lease(e) { + return e.path, e.hash, c.releaseFunc(e), nil + } + + // Slow path: one download per artifact, however many stagers + // arrive. + res, ferr, _ := c.flight.Do(coord, func() (any, error) { + return c.download(ctx, ref, coord) + }) + if ferr != nil { + return "", "", nil, ferr + } + + e, ok := res.(*entry) + if !ok { + return "", "", nil, fmt.Errorf("dispatch/artifact/cache: unexpected flight result %T", res) + } + + if c.entries.lease(e) { + return e.path, e.hash, c.releaseFunc(e), nil + } + + if cerr := ctx.Err(); cerr != nil { + return "", "", nil, fmt.Errorf("dispatch/artifact/cache: stage %s/%s: %w", + ref.Bucket, ref.Key, cerr) + } + } + + return "", "", nil, fmt.Errorf("%w: %s/%s was evicted before it could be used, %d times running", + ErrBudgetExceeded, ref.Bucket, ref.Key, maxStageAttempts) +} + +// lookup resolves a cached entry by hash, then by coordinates. +func (c *Cache) lookup(ref artifact.Ref, coord string) (*entry, bool) { + if ref.ContentHash != "" { + if e, ok := c.entries.getByHash(ref.ContentHash); ok { + c.entries.alias(coord, e.hash) + + return e, true + } + } + + return c.entries.getByCoord(coord) +} + +// releaseFunc returns an idempotent release for an entry. +// +// The manager is only nudged when the entry loses its last stager, +// because that is the only release that changes what eviction could +// free. +func (c *Cache) releaseFunc(e *entry) func() { + var once sync.Once + + return func() { + once.Do(func() { + if c.entries.release(e) { + c.wake() + } + }) + } +} + +// download fetches an artifact into the cache. +// +// Bytes stream through a BLAKE3 hasher into a temp file, which is then +// renamed into its content-addressed home. Hashing therefore costs +// nothing beyond the read that was happening anyway, which is what lets +// registration skip hashing entirely. +func (c *Cache) download(ctx context.Context, ref artifact.Ref, coord string) (*entry, error) { + // Re-check under the flight: a concurrent stager may have finished + // between our miss and our turn here. + if e, ok := c.lookup(ref, coord); ok { + return e, nil + } + + reserved := ref.Size + if reserved <= 0 { + // Size unknown. Reserve nothing up front and correct the + // accounting once the copy reports the real figure. + reserved = 0 + } + + h, err := c.newHold(ctx, reserved) + if err != nil { + return nil, err + } + + // The hold belongs to whoever ends up owning the bytes. Until an + // entry takes it, that is nobody, and every path out of here has to + // give it back. + committed := false + + defer func() { + if !committed { + c.releaseHold(h) + } + }() + + rc, err := c.backend.Open(ctx, ref) + if err != nil { + if errors.Is(err, dispatch.ErrPermanent) { + // Preserve the classification: staging an input that is gone, + // or that we may not read, is permanent, and the executor must + // fail fast rather than retry. Wrap rather than replace, so the + // specific sentinel and the backend's message both survive. + return nil, fmt.Errorf("stage %s/%s: %w", ref.Bucket, ref.Key, err) + } + + return nil, fmt.Errorf("dispatch/artifact/cache: open %s/%s: %w", ref.Bucket, ref.Key, err) + } + + defer func() { + if cerr := rc.Close(); cerr != nil { + c.logger.Warn("dispatch/artifact/cache: close source", + log.String("key", ref.Key), log.String("error", cerr.Error())) + } + }() + + tmpPath := filepath.Join(c.dir, tmpDir, id.NewArtifactID().String()) + + written, sum, err := c.copyAndHash(tmpPath, rc) + if err != nil { + c.removeQuietly(tmpPath) + + return nil, err + } + + // Either the ref carried no size, or it lied. Correct the hold to + // what actually landed on disk. Growing can block or evict, which is + // why it is bounded by the caller's context like the first + // reservation was; failing here costs the download, not the ledger. + if rerr := c.resize(ctx, h, written); rerr != nil { + c.removeQuietly(tmpPath) + + return nil, rerr + } + + final, err := c.shardPath(sum) + if err != nil { + c.removeQuietly(tmpPath) + + return nil, err + } + + e := &entry{ + hash: hashPrefix + sum, + size: written, + hold: h, + } + + // The rename and the registration happen together, under the entry + // table's lock, because eviction deletes a file under that same lock. + // Anything less and this download can adopt a file eviction has + // already condemned. + live, err := c.entries.publish(e, coord, func() (string, error) { + return final, c.promote(tmpPath, final) + }) + if err != nil { + c.removeQuietly(tmpPath) + + return nil, err + } + + // A racing download of the same bytes may have registered first, + // either a different artifact that happens to share them or a retry + // of this one. It owns the file and the hold that covers it; ours + // goes back with the deferred release. Content addressing makes that + // a cache hit rather than a conflict. + if live != e { + return live, nil + } + + committed = true + + return e, nil +} + +// copyAndHash streams src into a new file at dst, returning the byte +// count and the hex digest. +func (c *Cache) copyAndHash(dst string, src io.Reader) (written int64, digest string, err error) { + // #nosec G304 -- dst is a cache-internal temp path, and O_EXCL means this + // creates a new file rather than opening an existing one. + f, err := os.OpenFile(dst, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600) + if err != nil { + return 0, "", fmt.Errorf("dispatch/artifact/cache: create temp file: %w", err) + } + + hasher := blake3.New() + + written, copyErr := io.Copy(io.MultiWriter(f, hasher), src) + + closeErr := f.Close() + + if copyErr != nil { + return 0, "", fmt.Errorf("dispatch/artifact/cache: download: %w", copyErr) + } + + if closeErr != nil { + return 0, "", fmt.Errorf("dispatch/artifact/cache: close temp file: %w", closeErr) + } + + return written, hex.EncodeToString(hasher.Sum(nil)), nil +} + +// shardPath returns a digest's home, creating its shard directory. +// +// It is separate from promote because only promote has to be ordered +// against eviction, and creating a directory eviction never removes +// does not belong inside that lock. +func (c *Cache) shardPath(sum string) (string, error) { + dir := filepath.Join(c.dir, hashDir, sum[:2]) + if err := os.MkdirAll(dir, dirPerm); err != nil { + return "", fmt.Errorf("dispatch/artifact/cache: create shard dir: %w", err) + } + + return filepath.Join(dir, sum), nil +} + +// promote moves a completed temp file into its content-addressed home. +// +// It runs under the entry table's lock, so a file it finds already at +// final belongs to an entry that is in the table right now and cannot +// be unlinked while this holds the lock. That is the whole reason the +// caller passes it in rather than calling it first: eviction unlinks +// under the same lock, so without that ordering a stat here can see a +// file whose eviction has already been decided. +func (c *Cache) promote(tmpPath, final string) error { + // Another stager promoted identical bytes first. Its copy is as good + // as ours, so drop ours rather than racing the rename. + if _, err := os.Stat(final); err == nil { + c.removeQuietly(tmpPath) + + return nil + } + + if err := os.Rename(tmpPath, final); err != nil { + return fmt.Errorf("dispatch/artifact/cache: promote: %w", err) + } + + return nil +} + +// evictOne removes the least recently used unleased entry, releasing +// the lease that held its bytes. It reports the bytes reclaimed and +// whether there was anything to evict at all — an empty artifact frees +// zero bytes and is still progress, so the two answers cannot be folded +// into one number without stalling reclamation on a zero-byte file. +// +// The table picks the victim, forgets it and unlinks its file under its +// own lock, and it only ever picks an entry no stager holds. Doing the +// unlink there rather than here is what stops a concurrent download +// adopting the victim's file in the gap: see entryTable.publish. +// +// Releasing the hold stays out here, because that takes the manager's +// lock and the table's lock is still held inside evictLRU. +func (c *Cache) evictOne() (int64, bool) { + victim := c.entries.evictLRU(func(e *entry) { + c.removeQuietly(e.path) + }) + if victim == nil { + return 0, false + } + + freed := victim.hold.bytes + c.releaseHold(victim.hold) + + c.logger.Debug("dispatch/artifact/cache: evicted entry", + log.String("hash", victim.hash), + log.Int64("bytes", victim.size), + ) + + return freed, true +} + +// Reclaim frees up to need bytes by evicting least recently used +// entries, satisfying resource.Reclaimer. +// +// The bytes return to the manager as each victim's lease is released, +// which is the only path that credits the ledger; the count returned +// here is the manager's "something changed, re-check" signal and is +// deliberately not added to anything. +// +// It reclaims nothing for any key but disk. This cache holds bytes on +// one volume and nothing else, and a reclaimer that answered for memory +// it does not hold would have the manager admit work the box cannot +// run. +// +// This runs on the admission path — Manager.Acquire calls it before it +// blocks, under whatever deadline the caller set — so it does the +// unlinks the shortfall needs and stops, rather than tidying up while a +// fetcher waits. +func (c *Cache) Reclaim(ctx context.Context, key string, need int64) (int64, error) { + if key != resource.Disk || need <= 0 { + return 0, nil + } + + var freed int64 + + for freed < need && ctx.Err() == nil { + bytes, ok := c.evictOne() + if !ok { + // Everything left is leased by a running stager. + break + } + + freed += bytes + } + + return freed, nil +} + +// Available reports the bytes eviction could free right now, satisfying +// resource.Reclaimer. +// +// It reads the entry table and nothing else. The guarantee this cache +// keeps is that the table lock and the manager's lock are never held at +// the same time — not that one is always taken first; rebuild takes +// them in the opposite sequence and is safe for exactly that reason. +// Reclaim lets go of the table before releasing a lease. Totalling the +// leases here instead of the entry sizes would break the rule on the +// one call the manager makes while a caller is mid-Acquire, and hold +// the table under the manager's lock while Reclaim waits for the table. +// That is the deadlock. +func (c *Cache) Available(key string) int64 { + if key != resource.Disk { + return 0 + } + + return c.entries.evictableBytes() +} + +// removeQuietly deletes a path, logging rather than failing. +func (c *Cache) removeQuietly(path string) { + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + c.logger.Warn("dispatch/artifact/cache: remove file", + log.String("path", path), log.String("error", err.Error())) + } +} + +// Purge removes every cached file and returns its bytes to the manager. +// +// The table is drained under its own lock rather than replaced, so an +// eviction running at the same time cannot pick an entry this loop has +// already released: each entry leaves the table once and its hold goes +// back once. It still assumes no live stagers — Purge deletes files +// out from under anything holding one, which was always its contract. +func (c *Cache) Purge() error { + for _, e := range c.entries.drain() { + c.removeQuietly(e.path) + c.releaseHold(e.hold) + } + + if err := os.RemoveAll(filepath.Join(c.dir, hashDir)); err != nil { + return fmt.Errorf("dispatch/artifact/cache: purge: %w", err) + } + + if err := os.MkdirAll(filepath.Join(c.dir, hashDir), dirPerm); err != nil { + return fmt.Errorf("dispatch/artifact/cache: recreate hash dir: %w", err) + } + + return nil +} + +// Close releases the cache. Staged files survive so the next process can +// reuse them; only the temp directory is cleared. +func (c *Cache) Close() error { + var err error + + c.closeOnce.Do(func() { + err = c.resetTmp() + }) + + return err +} + +// TrimHashPrefix returns a digest without its algorithm label. +func TrimHashPrefix(hash string) string { + return strings.TrimPrefix(hash, hashPrefix) +} diff --git a/artifact/cache/cache_test.go b/artifact/cache/cache_test.go new file mode 100644 index 0000000..26371c0 --- /dev/null +++ b/artifact/cache/cache_test.go @@ -0,0 +1,399 @@ +package cache_test + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/artifact/cache" +) + +func newCache(t *testing.T, budget int64) (*cache.Cache, *artifacttest.Backend) { + t.Helper() + + b := artifacttest.NewBackend() + + c, err := cache.New(t.TempDir(), b, cache.WithBudget(budget)) + if err != nil { + t.Fatalf("cache.New: %v", err) + } + + t.Cleanup(func() { + if cerr := c.Close(); cerr != nil { + t.Errorf("cache close: %v", cerr) + } + }) + + return c, b +} + +func TestStageDownloadsAndCaches(t *testing.T) { + ctx := context.Background() + c, b := newCache(t, 1<<20) + b.Put("models", "tower.ifc", []byte("hello world")) + + ref := artifact.Ref{Bucket: "models", Key: "tower.ifc", Size: 11} + + path, hash, release, err := c.Stage(ctx, ref) + if err != nil { + t.Fatalf("Stage: %v", err) + } + + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile(%q): %v", path, err) + } + + if string(data) != "hello world" { + t.Fatalf("staged content = %q, want %q", data, "hello world") + } + + if !strings.HasPrefix(hash, "blake3:") { + t.Fatalf("hash = %q, want a blake3: prefix — Stage must hash during download", hash) + } + + release() + + _, hash2, release2, err := c.Stage(ctx, ref) + if err != nil { + t.Fatalf("second Stage: %v", err) + } + + release2() + + if b.Opens() != 1 { + t.Fatalf("Opens() = %d, want 1 (second Stage must hit the cache)", b.Opens()) + } + + if hash2 != hash { + t.Fatalf("hash changed between stages: %q then %q", hash, hash2) + } +} + +// TestStageSingleFlight is the reason the cache exists in front of the +// backend rather than beside it: eight jobs on the same model must pay +// for one download, not eight. +func TestStageSingleFlight(t *testing.T) { + ctx := context.Background() + c, b := newCache(t, 1<<20) + b.Put("models", "big.ifc", []byte("payload")) + b.DelayOpen = 50 * time.Millisecond + + ref := artifact.Ref{Bucket: "models", Key: "big.ifc", Size: 7} + + const n = 8 + + var wg sync.WaitGroup + + errs := make([]error, n) + releases := make([]func(), n) + + for i := range n { + wg.Add(1) + + go func(i int) { + defer wg.Done() + + _, _, release, err := c.Stage(ctx, ref) + errs[i] = err + releases[i] = release + }(i) + } + + wg.Wait() + + for i, err := range errs { + if err != nil { + t.Fatalf("goroutine %d: %v", i, err) + } + + if releases[i] != nil { + releases[i]() + } + } + + if b.Opens() != 1 { + t.Fatalf("Opens() = %d, want 1 — %d concurrent stages must share one download", b.Opens(), n) + } +} + +func TestStageMissingObjectIsPermanent(t *testing.T) { + c, _ := newCache(t, 1<<20) + + _, _, _, err := c.Stage(context.Background(), + artifact.Ref{Bucket: "models", Key: "absent"}) + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Stage(missing) = %v, want ErrNotFound (permanent, so the job fails fast)", err) + } +} + +// TestStageDeniedObjectIsPermanent covers the permanent failure that is not a +// missing object. The object is present; the backend just will not hand it +// over, and retrying cannot change that. +// +// This pins propagation, not the retry decision: the executor does not yet +// consult the classification (see worker/executor.go handleFailure), so +// nothing fails fast at runtime today. What this guarantees is that the +// classification survives the cache layer intact and is there to act on. +func TestStageDeniedObjectIsPermanent(t *testing.T) { + c, b := newCache(t, 1<<20) + b.Put("models", "secret.ifc", []byte("classified")) + b.DenyOpen = true + + _, _, _, err := c.Stage(context.Background(), + artifact.Ref{Bucket: "models", Key: "secret.ifc"}) + + if !errors.Is(err, dispatch.ErrPermanent) { + t.Fatalf("Stage(denied) = %v, want ErrPermanent (so the job fails fast)", err) + } + + if !errors.Is(err, artifact.ErrPermissionDenied) { + t.Fatalf("Stage(denied) = %v, want ErrPermissionDenied", err) + } + + // Reporting a forbidden object as missing would send an operator + // looking for a deleted file that is sitting right there. + if errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Stage(denied) reported as not found: %v", err) + } +} + +func TestStageUnknownSizeStillWorks(t *testing.T) { + ctx := context.Background() + c, b := newCache(t, 1<<20) + b.Put("models", "nosize.ifc", []byte("0123456789")) + + // A ref registered but never staged carries no size yet. + _, _, release, err := c.Stage(ctx, artifact.Ref{Bucket: "models", Key: "nosize.ifc"}) + if err != nil { + t.Fatalf("Stage: %v", err) + } + + defer release() + + if used := c.Used(); used != 10 { + t.Fatalf("Used() = %d, want 10 — accounting must correct itself once the size is known", used) + } +} + +func TestLeaseBlocksEviction(t *testing.T) { + ctx := context.Background() + c, b := newCache(t, 20) // room for two 10-byte objects + b.Put("m", "a", []byte("0123456789")) + b.Put("m", "b", []byte("0123456789")) + b.Put("m", "c", []byte("0123456789")) + + pathA, _, releaseA, err := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "a", Size: 10}) + if err != nil { + t.Fatalf("Stage a: %v", err) + } + + _, _, releaseB, err := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "b", Size: 10}) + if err != nil { + t.Fatalf("Stage b: %v", err) + } + + releaseB() // b is now evictable; a is still leased + + _, _, releaseC, err := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "c", Size: 10}) + if err != nil { + t.Fatalf("Stage c should have evicted b, got: %v", err) + } + + // The leased entry must still be on disk: a running handler holds it. + if _, serr := os.Stat(pathA); serr != nil { + t.Fatalf("leased entry was evicted: %v", serr) + } + + releaseC() + releaseA() +} + +func TestBudgetExceededRespectsDeadline(t *testing.T) { + ctx := context.Background() + c, b := newCache(t, 10) + b.Put("m", "a", []byte("0123456789")) + b.Put("m", "b", []byte("0123456789")) + + _, _, releaseA, err := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "a", Size: 10}) + if err != nil { + t.Fatalf("Stage a: %v", err) + } + + defer releaseA() + + // a is leased and fills the budget, so b cannot be admitted. + deadlined, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + defer cancel() + + start := time.Now() + + _, _, _, err = c.Stage(deadlined, artifact.Ref{Bucket: "m", Key: "b", Size: 10}) + if !errors.Is(err, cache.ErrBudgetExceeded) { + t.Fatalf("Stage under an exhausted budget = %v, want ErrBudgetExceeded", err) + } + + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("Stage waited %v — it must be bounded by the context deadline", elapsed) + } +} + +// TestOversizeRefRejectedImmediately guards against a job that can never +// be staged blocking until its deadline instead of failing at once. +func TestOversizeRefRejectedImmediately(t *testing.T) { + c, b := newCache(t, 10) + b.Put("m", "huge", make([]byte, 100)) + + start := time.Now() + + _, _, _, err := c.Stage(context.Background(), + artifact.Ref{Bucket: "m", Key: "huge", Size: 100}) + if !errors.Is(err, cache.ErrBudgetExceeded) { + t.Fatalf("Stage of a ref larger than the whole budget = %v, want ErrBudgetExceeded", err) + } + + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("took %v — an impossible request must fail immediately, not wait", elapsed) + } +} + +func TestReleaseIsIdempotent(t *testing.T) { + ctx := context.Background() + c, b := newCache(t, 1<<20) + b.Put("m", "a", []byte("data")) + + _, _, release, err := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "a", Size: 4}) + if err != nil { + t.Fatalf("Stage: %v", err) + } + + release() + release() + release() + + // A double release must not drop the lease count below zero and let a + // still-in-use entry be evicted. + _, _, release2, err := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "a", Size: 4}) + if err != nil { + t.Fatalf("Stage after repeated release: %v", err) + } + + release2() +} + +func TestRecoveryWipesTmpAndReusesFiles(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + b := artifacttest.NewBackend() + b.Put("m", "a", []byte("0123456789")) + + ref := artifact.Ref{Bucket: "m", Key: "a", Size: 10} + + c1, err := cache.New(dir, b, cache.WithBudget(1<<20)) + if err != nil { + t.Fatalf("first New: %v", err) + } + + _, hash, release, err := c1.Stage(ctx, ref) + if err != nil { + t.Fatalf("Stage: %v", err) + } + + release() + + if cerr := c1.Close(); cerr != nil { + t.Fatalf("close: %v", cerr) + } + + // Simulate a crash mid-download. + leftover := filepath.Join(dir, "tmp", "leftover") + if werr := os.WriteFile(leftover, []byte("junk"), 0o600); werr != nil { + t.Fatalf("write leftover: %v", werr) + } + + c2, err := cache.New(dir, b, cache.WithBudget(1<<20)) + if err != nil { + t.Fatalf("second New: %v", err) + } + + t.Cleanup(func() { + if cerr := c2.Close(); cerr != nil { + t.Errorf("close c2: %v", cerr) + } + }) + + if _, serr := os.Stat(leftover); !os.IsNotExist(serr) { + t.Fatal("startup must wipe tmp/ — a partial download is never reusable") + } + + // A ref carrying the hash resolves against the rebuilt table without + // touching the backend. + hashed := ref + hashed.ContentHash = hash + + _, _, release2, err := c2.Stage(ctx, hashed) + if err != nil { + t.Fatalf("Stage after recovery: %v", err) + } + + release2() + + if b.Opens() != 1 { + t.Fatalf("Opens() = %d, want 1 — the index must be rebuilt from disk, not re-downloaded", b.Opens()) + } +} + +func TestDistinctArtifactsWithIdenticalBytesShareOneCopy(t *testing.T) { + ctx := context.Background() + c, b := newCache(t, 1<<20) + b.Put("m", "first.ifc", []byte("identical")) + b.Put("m", "second.ifc", []byte("identical")) + + pathA, hashA, releaseA, err := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "first.ifc", Size: 9}) + if err != nil { + t.Fatalf("Stage first: %v", err) + } + + defer releaseA() + + pathB, hashB, releaseB, err := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "second.ifc", Size: 9}) + if err != nil { + t.Fatalf("Stage second: %v", err) + } + + defer releaseB() + + if hashA != hashB { + t.Fatalf("identical bytes hashed differently: %q vs %q", hashA, hashB) + } + + if pathA != pathB { + t.Fatalf("content addressing failed: %q vs %q", pathA, pathB) + } + + if used := c.Used(); used != 9 { + t.Fatalf("Used() = %d, want 9 — identical bytes must be stored once", used) + } +} + +func TestBudgetReportsConfiguredLimit(t *testing.T) { + c, _ := newCache(t, 4096) + if got := c.Budget(); got != 4096 { + t.Fatalf("Budget() = %d, want 4096", got) + } +} + +func TestNewRejectsNilBackend(t *testing.T) { + _, err := cache.New(t.TempDir(), nil) + if !errors.Is(err, artifact.ErrNoBackend) { + t.Fatalf("New(nil backend) = %v, want ErrNoBackend", err) + } +} diff --git a/artifact/cache/doc.go b/artifact/cache/doc.go new file mode 100644 index 0000000..ebafa00 --- /dev/null +++ b/artifact/cache/doc.go @@ -0,0 +1,40 @@ +// Package cache is the worker-local staging cache for artifact inputs. +// +// Staging exists because the native libraries that read Dispatch's heavy +// inputs — CAD kernels, mesh importers, PDF engines — want a file path +// they can seek and memory-map, not a stream. The cache materialises an +// artifact to local disk and hands back that path. +// +// It is content-addressed. Entries live under their BLAKE3 hash, so two +// jobs consuming the same model share one copy, and re-running a job over +// an input it already staged costs nothing. The hash is computed during +// the download rather than by a separate pass: the bytes are already +// streaming to disk, so hashing them is free. That is what fills in the +// content_hash column an artifact carries as NULL after registration. +// +// Three mechanisms keep concurrent staging honest: +// +// - Single-flight collapses concurrent stages of the same artifact into +// one download. +// - Leases pin an entry while a job is using it, so eviction cannot +// pull a file out from under a running handler. +// - A byte budget bounds total disk use, evicting unleased entries by +// least-recent-use and making a job wait when nothing can be freed. +// +// That last mechanism is also the artifact plane's first piece of +// admission control: a job needing more staging space than is available +// waits instead of filling the disk. +// +// The budget is a resource.Manager, not a counter. Each cached object +// holds one lease for its bytes and eviction releases it, so with +// WithManager the staged bytes sit in the same ledger the worker admits +// jobs against: the cache is registered as that manager's disk +// reclaimer, and a job short on disk gets it by evicting rather than by +// waiting for a cache that has no reason to shrink. With no manager +// supplied the cache builds a private single-key one, which is the +// private disk budget it always had. +// +// The cache is a cache. Its index is an optimisation rebuilt from disk on +// startup, and a corrupt or missing index costs a re-download, never +// correctness. +package cache diff --git a/artifact/cache/entry.go b/artifact/cache/entry.go new file mode 100644 index 0000000..225a224 --- /dev/null +++ b/artifact/cache/entry.go @@ -0,0 +1,379 @@ +package cache + +import ( + "sync" +) + +// entry is one cached object on disk. +type entry struct { + // hash is the BLAKE3 content hash, formatted "blake3:". + hash string + // path is the absolute location of the file. On the download path + // publish fills it in, because the file only has a home once it has + // been renamed into one, and that happens under the table's lock. + path string + // size is the file's byte count, as accounted against the manager. + size int64 + // hold is the manager capacity backing those bytes. It is set + // before the entry is published and released when the entry is + // evicted, so hold.bytes equals size for the whole time an entry is + // reachable through this table. + hold *hold + // leases counts the stagers currently using this entry. An entry with + // leases > 0 must never be evicted: a running handler holds its path. + leases int + // coords are the coordinate aliases resolving to this entry, so + // eviction can drop them by name instead of scanning every alias in + // the table. + coords []string + // prev, next and evictable thread the eviction list. An entry is on + // that list exactly while nothing holds it; evictable says so, + // because a list of one has nil on both sides. + prev, next *entry + evictable bool +} + +// entryTable holds the cache's in-memory view of what is on disk. +// +// Two lookup paths matter. byHash is the content-addressed one and is +// authoritative. byCoord maps an artifact's storage coordinates to a hash +// so a ref whose content_hash is still NULL — every freshly registered +// artifact — can hit the cache on its second stage. +// +// Everything here is O(1). Eviction runs on the admission path, where a +// worker has one poll interval to free what a claimed job needs, and a +// scan of the table per victim would put the cost of one eviction in +// proportion to how warm the cache is. A 20 GiB cache of 100 KiB +// objects is 200k entries: scanning made a single eviction cost +// milliseconds, so a batch could free tens of megabytes and the disk a +// worker offered as free-plus-reclaimable stopped being redeemable at +// exactly the size where it mattered. +type entryTable struct { + mu sync.Mutex + byHash map[string]*entry + byCoord map[string]string + + // head and tail are the eviction list, most recently released + // first, so the tail is always the victim and finding it is a + // pointer read. Only unleased entries are linked, which is what + // makes that true without a scan past the pinned ones. + head, tail *entry + // evictable totals the bytes on that list. + evictable int64 +} + +func newEntryTable() *entryTable { + return &entryTable{ + byHash: make(map[string]*entry), + byCoord: make(map[string]string), + } +} + +// coordKey identifies an artifact's storage location. +func coordKey(backend, bucket, key string) string { + return backend + "\x00" + bucket + "\x00" + key +} + +// getByHash returns the entry for a content hash. +func (t *entryTable) getByHash(hash string) (*entry, bool) { + t.mu.Lock() + defer t.mu.Unlock() + + e, ok := t.byHash[hash] + + return e, ok +} + +// getByCoord returns the entry cached for an artifact's coordinates. +func (t *entryTable) getByCoord(coord string) (*entry, bool) { + t.mu.Lock() + defer t.mu.Unlock() + + hash, ok := t.byCoord[coord] + if !ok { + return nil, false + } + + e, ok := t.byHash[hash] + + return e, ok +} + +// put records an entry and, when coord is non-empty, its coordinate +// alias. +// +// This is the startup path, walking files that are already on disk, so +// it has no losing entry to hand back: one file per hash means the +// collision putLocked guards against cannot happen here. Everything +// that publishes an entry while the cache is live goes through publish +// instead, because it has a file to put on disk and that has to happen +// under this same lock. +func (t *entryTable) put(e *entry, coord string) { + t.mu.Lock() + defer t.mu.Unlock() + + _ = t.putLocked(e, coord) +} + +// putLocked is put's body, for callers already holding the lock. +// +// Two downloads of different coordinates can produce identical bytes at +// the same moment and both miss the content-address check. Only one of +// them can own the file: overwriting here would strand the other +// entry's hold with no path back to the manager, leaking capacity for +// the life of the process, and would drop an entry other stagers may +// already be holding. The loser is told so and hands its hold back. +func (t *entryTable) putLocked(e *entry, coord string) *entry { + live, ok := t.byHash[e.hash] + if !ok { + live = e + t.byHash[e.hash] = e + t.link(e) + } + + t.aliasLocked(coord, live) + + return live +} + +// publish puts a downloaded file into its content-addressed home and +// registers the entry that owns it, both under this lock. place does +// the filesystem half and returns the path it settled on. +// +// The two halves cannot be separated. Eviction unlinks a victim's file +// under this same lock, so serialising against it here is what upholds +// the table's central invariant: a file exists at a hash's path if and +// only if the table holds an entry for that hash. Promote outside the +// lock and a download can stat a victim's file in the window after +// eviction dropped it from the table and before it unlinked it, adopt +// the doomed file, register an entry for it, and hand a live lease on a +// path that is deleted a moment later. Worse, the entry stays in the +// table pointing at nothing, so every later stager of that artifact is +// handed the same corpse until it is evicted again. +// +// Both critical sections are two syscalls rather than scans, so +// eviction stays O(1) in the size of the cache, which is what the +// admission path needs from it. +func (t *entryTable) publish(e *entry, coord string, place func() (string, error)) (*entry, error) { + t.mu.Lock() + defer t.mu.Unlock() + + path, err := place() + if err != nil { + return nil, err + } + + e.path = path + + return t.putLocked(e, coord), nil +} + +// alias points a coordinate at an existing hash. +func (t *entryTable) alias(coord, hash string) { + if coord == "" { + return + } + + t.mu.Lock() + defer t.mu.Unlock() + + if e, ok := t.byHash[hash]; ok { + t.aliasLocked(coord, e) + } +} + +// aliasLocked records a coordinate against an entry, once. +// +// The repeat check is not an optimisation: every cache hit on a +// content-hashed ref re-aliases the same coordinate, and appending each +// time would grow the entry's coords slice for as long as the entry +// lives. A coordinate that moves to a different entry leaves its name +// behind on the old one, which costs nothing — eviction only deletes an +// alias that still resolves to the entry being evicted. +func (t *entryTable) aliasLocked(coord string, e *entry) { + if coord == "" || t.byCoord[coord] == e.hash { + return + } + + t.byCoord[coord] = e.hash + e.coords = append(e.coords, coord) +} + +// lease pins an entry and takes it off the eviction list. +// +// It reports false when the entry is no longer in the table, which +// means eviction took it: the file is gone and its bytes have been +// credited back to the manager, so the caller must go and stage it +// again rather than pin a corpse. Checking membership under the same +// lock that evictLRU removes under is what makes the two mutually +// exclusive — either this pins the entry first and eviction cannot see +// it, or eviction wins and this fails. +func (t *entryTable) lease(e *entry) bool { + t.mu.Lock() + defer t.mu.Unlock() + + if t.byHash[e.hash] != e { + return false + } + + e.leases++ + + if e.leases == 1 { + t.unlink(e) + } + + return true +} + +// release unpins an entry and reports whether that was its last lease. +// +// The last release puts the entry back at the head of the eviction +// list, which is what "recently used" means here: an entry's place in +// the queue is set by when it stopped being used, not by when it was +// picked up. Only that release changes what eviction could free, so +// only that one is worth waking a blocked acquirer for. +func (t *entryTable) release(e *entry) bool { + t.mu.Lock() + defer t.mu.Unlock() + + if e.leases == 0 { + return false + } + + e.leases-- + + if e.leases > 0 { + return false + } + + // An entry only leaves the table by eviction and eviction never + // takes a leased one, so this one is still there — belt and braces + // against a future path that is not so careful. + if t.byHash[e.hash] == e { + t.link(e) + } + + return true +} + +// evictLRU removes the least recently used unleased entry, unlinks its +// file through remove, and returns it. It returns nil, without calling +// remove, when every entry is leased. +// +// remove runs under the lock on purpose: see publish. Dropping the +// entry and deleting its file have to look like one step to a download +// promoting the same hash, or that download adopts a file that is +// already condemned. remove must not touch the resource manager. +// Crediting the victim's bytes back is the caller's job, once this has +// returned and the lock is gone. +func (t *entryTable) evictLRU(remove func(*entry)) *entry { + t.mu.Lock() + defer t.mu.Unlock() + + victim := t.tail + if victim == nil { + return nil + } + + t.forget(victim) + remove(victim) + + return victim +} + +// drain empties the table and returns everything that was in it, +// leased entries included, for the caller to delete and account for. +// +// Taking the entries out under the lock is what makes Purge safe +// against a concurrent eviction: an entry leaves the table exactly +// once, so its hold is released exactly once, and the two paths cannot +// both claim the same bytes. +func (t *entryTable) drain() []*entry { + t.mu.Lock() + defer t.mu.Unlock() + + out := make([]*entry, 0, len(t.byHash)) + for _, e := range t.byHash { + out = append(out, e) + t.unlink(e) + + e.coords = nil + } + + t.byHash = make(map[string]*entry) + t.byCoord = make(map[string]string) + + return out +} + +// evictableBytes totals the entries that could be evicted right now. +// It is what the cache reports as reclaimable disk, so it counts only +// what a stager is not holding. +func (t *entryTable) evictableBytes() int64 { + t.mu.Lock() + defer t.mu.Unlock() + + return t.evictable +} + +// forget removes an entry from the table entirely: the eviction list, +// the hash index, and every alias that still resolves to it. +func (t *entryTable) forget(e *entry) { + t.unlink(e) + delete(t.byHash, e.hash) + + for _, coord := range e.coords { + if t.byCoord[coord] == e.hash { + delete(t.byCoord, coord) + } + } + + e.coords = nil +} + +// link puts an entry at the head of the eviction list. +func (t *entryTable) link(e *entry) { + if e.evictable { + return + } + + e.evictable = true + e.prev = nil + e.next = t.head + + if t.head != nil { + t.head.prev = e + } + + t.head = e + + if t.tail == nil { + t.tail = e + } + + t.evictable += e.size +} + +// unlink takes an entry off the eviction list. +func (t *entryTable) unlink(e *entry) { + if !e.evictable { + return + } + + e.evictable = false + + if e.prev != nil { + e.prev.next = e.next + } else { + t.head = e.next + } + + if e.next != nil { + e.next.prev = e.prev + } else { + t.tail = e.prev + } + + e.prev, e.next = nil, nil + t.evictable -= e.size +} diff --git a/artifact/cache/entry_internal_test.go b/artifact/cache/entry_internal_test.go new file mode 100644 index 0000000..ee61e42 --- /dev/null +++ b/artifact/cache/entry_internal_test.go @@ -0,0 +1,162 @@ +package cache + +import ( + "testing" + "time" +) + +// blocked is how long a goroutine has to stay out of the entry table +// before this file believes it is genuinely locked out. It only ever +// costs the suite that long when the invariant holds, and a value this +// side of a second keeps a loaded machine from mattering: the assertions +// below can miss, never misfire, because only the contended path +// completing early fails them. +const blocked = 50 * time.Millisecond + +// TestPublishAndEvictionExcludeEachOther pins the invariant the cache's +// staged paths rest on: a file exists at a hash's path if and only if +// the entry table holds an entry for that hash. +// +// Both transitions have to be atomic against each other, so both are +// checked. The second case is the regression. Eviction used to unlink +// after handing the victim back, and in that gap a download of the same +// hash could stat the doomed file, adopt it instead of writing its own, +// and register an entry for a path that was unlinked a moment later. +// Every stager of that artifact was then handed a path with no file +// behind it until the poisoned entry was itself evicted, which is a +// leased entry losing its bytes: exactly what a lease is supposed to +// prevent. +func TestPublishAndEvictionExcludeEachOther(t *testing.T) { + const ( + hash = "blake3:d0" + path = "/cache/blake3/d0/d0" + ) + + newTable := func() *entryTable { + t.Helper() + + tbl := newEntryTable() + tbl.put(&entry{hash: hash, path: path, size: 1}, "coord") + + return tbl + } + + t.Run("eviction cannot start while a download is promoting", func(t *testing.T) { + tbl := newTable() + + promoting := make(chan struct{}) + evicted := make(chan *entry, 1) + + go func() { + <-promoting + + evicted <- tbl.evictLRU(func(*entry) {}) + }() + + if _, err := tbl.publish(&entry{hash: hash, size: 1}, "coord", + func() (string, error) { + // Standing where the rename stands, holding the table. + close(promoting) + + select { + case victim := <-evicted: + t.Errorf("eviction removed %s while a download held the table to promote it; "+ + "the download would adopt a file this eviction is about to unlink", victim.hash) + case <-time.After(blocked): + } + + return path, nil + }); err != nil { + t.Fatalf("publish: %v", err) + } + }) + + t.Run("a download cannot promote while eviction is unlinking", func(t *testing.T) { + tbl := newTable() + + unlinking := make(chan struct{}) + promoted := make(chan struct{}) + + go func() { + <-unlinking + + _, _ = tbl.publish(&entry{hash: hash, size: 1}, "coord", + func() (string, error) { + close(promoted) + + return path, nil + }) + }() + + victim := tbl.evictLRU(func(*entry) { + // Standing where the unlink stands: the entry has left the + // table and its file is still on disk. + close(unlinking) + + select { + case <-promoted: + t.Error("a download reached promote while eviction still held the table; " + + "it would stat the victim's file, adopt it, and hand out a path " + + "that is unlinked as soon as this returns") + case <-time.After(blocked): + } + }) + + if victim == nil { + t.Fatal("evictLRU returned no victim, but the table held one unleased entry") + } + + if got, want := victim.hash, hash; got != want { + t.Fatalf("evicted hash = %q, want %q", got, want) + } + }) +} + +// TestEvictLRUUnlinksBeforeItReturns is the cheap half of the same +// invariant, and the one a refactor is most likely to undo: the removal +// callback is not optional and does not run later. +func TestEvictLRUUnlinksBeforeItReturns(t *testing.T) { + tbl := newEntryTable() + tbl.put(&entry{hash: "blake3:d0", path: "/cache/blake3/d0/d0", size: 1}, "coord") + + var removed []string + + victim := tbl.evictLRU(func(e *entry) { + removed = append(removed, e.path) + }) + + if victim == nil { + t.Fatal("evictLRU returned no victim, but the table held one unleased entry") + } + + if got, want := len(removed), 1; got != want { + t.Fatalf("remove called %d times, want %d", got, want) + } + + if got, want := removed[0], victim.path; got != want { + t.Fatalf("removed %q, want the victim's own path %q", got, want) + } +} + +// TestEvictLRUSkipsRemovalWithNothingToEvict: a table whose entries are +// all leased has no victim, and nothing on disk may be touched for one. +func TestEvictLRUSkipsRemovalWithNothingToEvict(t *testing.T) { + tbl := newEntryTable() + + e := &entry{hash: "blake3:d0", path: "/cache/blake3/d0/d0", size: 1} + tbl.put(e, "coord") + + if !tbl.lease(e) { + t.Fatal("lease on a freshly published entry failed") + } + + called := false + + if victim := tbl.evictLRU(func(*entry) { called = true }); victim != nil { + t.Fatalf("evictLRU returned %q, want nil while every entry is leased", victim.hash) + } + + if called { + t.Fatal("remove ran with no victim to remove") + } +} diff --git a/artifact/cache/reclaim_test.go b/artifact/cache/reclaim_test.go new file mode 100644 index 0000000..44323ed --- /dev/null +++ b/artifact/cache/reclaim_test.go @@ -0,0 +1,572 @@ +package cache_test + +import ( + "context" + "os" + "sync" + "testing" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/resource" +) + +// The cache is what makes resource.Manager's disk key reclaimable. +// Nothing else in the process holds evictable bytes. +var _ resource.Reclaimer = (*cache.Cache)(nil) + +// entrySize is the size of the objects stageAndRelease stages, small +// enough that a capacity of 100 is several whole entries and the +// arithmetic in these tests is readable. +const entrySize = 10 + +// managedCache builds a cache sharing mgr's ledger, preloaded with +// objects named "a".."a"+n-1 of size bytes each. +func managedCache(t *testing.T, mgr resource.Manager, count, size int) *cache.Cache { + t.Helper() + + b := artifacttest.NewBackend() + + for i := range count { + // Distinct bytes per object: identical bytes are one cache + // entry by design, and these tests are counting entries. + payload := make([]byte, size) + for j := range payload { + payload[j] = byte('a' + i) + } + + b.Put("m", objectKey(i), payload) + } + + c, err := cache.New(t.TempDir(), b, cache.WithManager(mgr)) + if err != nil { + t.Fatalf("cache.New: %v", err) + } + + t.Cleanup(func() { + if cerr := c.Close(); cerr != nil { + t.Errorf("cache close: %v", cerr) + } + }) + + return c +} + +func objectKey(i int) string { return string(rune('a' + i)) } + +// stageAndRelease stages object i and immediately releases it, leaving +// it cached and evictable. It returns the staged path. +func stageAndRelease(t *testing.T, c *cache.Cache, i int) string { + t.Helper() + + path, _, release, err := c.Stage(context.Background(), + artifact.Ref{Bucket: "m", Key: objectKey(i), Size: entrySize}) + if err != nil { + t.Fatalf("Stage %s: %v", objectKey(i), err) + } + + release() + + return path +} + +// heldByLeases totals what every live lease says it holds, which the +// manager's own used must equal exactly. Capacity minus Free is used, +// so this compares the ledger against the leases that justify it — +// the invariant a reclaimer crediting the ledger itself would break. +func heldByLeases(t *testing.T, mgr resource.Manager, key string) (used, held int64) { + t.Helper() + + used = mgr.Capacity()[key] - mgr.Free()[key] + + for _, l := range mgr.Leases() { + held += l.Held[key] + } + + return used, held +} + +func assertLedgerBalanced(t *testing.T, mgr resource.Manager) { + t.Helper() + + used, held := heldByLeases(t, mgr, resource.Disk) + if used != held { + t.Fatalf("ledger broken: used = %d but live leases hold %d — "+ + "reclamation must return bytes by releasing a lease, never by crediting the ledger", + used, held) + } +} + +// TestCacheIgnoresNonDiskKeys pins the boundary of what a disk cache may +// speak for. Reporting memory it does not hold would have the manager +// admit work the box cannot run. +func TestCacheIgnoresNonDiskKeys(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100, resource.Memory: 100}) + c := managedCache(t, mgr, 2, 10) + + pathA := stageAndRelease(t, c, 0) + stageAndRelease(t, c, 1) + + if got := c.Available(resource.Memory); got != 0 { + t.Fatalf("Available(memory) = %d, want 0 — the cache holds bytes on a volume and nothing else", got) + } + + freed, err := c.Reclaim(context.Background(), resource.Memory, 100) + if err != nil { + t.Fatalf("Reclaim(memory): %v", err) + } + + if freed != 0 { + t.Fatalf("Reclaim(memory) = %d, want 0", freed) + } + + if _, serr := os.Stat(pathA); serr != nil { + t.Fatalf("Reclaim(memory) evicted a cached file: %v", serr) + } + + if used := c.Used(); used != 20 { + t.Fatalf("Used() = %d, want 20 — a memory reclaim must not touch disk accounting", used) + } + + if got := c.Available(resource.Disk); got != 20 { + t.Fatalf("Available(disk) = %d, want 20", got) + } +} + +// TestStagingSpendsSharedManagerDisk is the connection the whole task +// exists to make: bytes on disk are bytes the shared ledger has spent. +func TestStagingSpendsSharedManagerDisk(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100}) + c := managedCache(t, mgr, 1, 30) + + before := mgr.Free()[resource.Disk] + + _, _, release, err := c.Stage(context.Background(), + artifact.Ref{Bucket: "m", Key: "a", Size: 30}) + if err != nil { + t.Fatalf("Stage: %v", err) + } + + defer release() + + if after := mgr.Free()[resource.Disk]; after != before-30 { + t.Fatalf("Free()[disk] = %d, want %d — staged bytes must be spent from the shared ledger", + after, before-30) + } + + if got := mgr.Reclaimable()[resource.Disk]; got != 0 { + t.Fatalf("Reclaimable()[disk] = %d, want 0 while the entry is leased", got) + } + + assertLedgerBalanced(t, mgr) +} + +// TestReclaimEvictsLRUAndReturnsBytes covers the reclaimer contract +// itself: least-recently-used first, only as much as was asked for, and +// the manager's ledger reflects it because the lease was released. +func TestReclaimEvictsLRUAndReturnsBytes(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100}) + c := managedCache(t, mgr, 3, 10) + + // Staged oldest first, so a is the least recently used. The gaps + // keep the ordering an assertion about LRU rather than about clock + // resolution. + pathA := stageAndRelease(t, c, 0) + time.Sleep(2 * time.Millisecond) + pathB := stageAndRelease(t, c, 1) + time.Sleep(2 * time.Millisecond) + pathC := stageAndRelease(t, c, 2) + + if got := mgr.Reclaimable()[resource.Disk]; got != 30 { + t.Fatalf("Reclaimable()[disk] = %d, want 30", got) + } + + freed, err := c.Reclaim(context.Background(), resource.Disk, 15) + if err != nil { + t.Fatalf("Reclaim: %v", err) + } + + if freed != 20 { + t.Fatalf("Reclaim(15) freed %d, want 20 — whole entries, and no more than the shortfall needs", freed) + } + + if free := mgr.Free()[resource.Disk]; free != 90 { + t.Fatalf("Free()[disk] = %d, want 90 — the evicted leases must be released, not merely counted", free) + } + + assertLedgerBalanced(t, mgr) + + for _, p := range []string{pathA, pathB} { + if _, serr := os.Stat(p); !os.IsNotExist(serr) { + t.Fatalf("evicted file %s still on disk: %v", p, serr) + } + } + + if _, serr := os.Stat(pathC); serr != nil { + t.Fatalf("most recently used entry was evicted first: %v", serr) + } + + if used := c.Used(); used != 10 { + t.Fatalf("Used() = %d, want 10", used) + } +} + +// TestReclaimNeverEvictsALeasedEntry: a running handler holds the path, +// so those bytes are not the manager's to take back however short it is. +func TestReclaimNeverEvictsALeasedEntry(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100}) + c := managedCache(t, mgr, 2, 10) + + pathA, _, releaseA, err := c.Stage(context.Background(), + artifact.Ref{Bucket: "m", Key: "a", Size: 10}) + if err != nil { + t.Fatalf("Stage a: %v", err) + } + + defer releaseA() + + stageAndRelease(t, c, 1) + + freed, err := c.Reclaim(context.Background(), resource.Disk, 1000) + if err != nil { + t.Fatalf("Reclaim: %v", err) + } + + if freed != 10 { + t.Fatalf("Reclaim(1000) freed %d, want 10 — only the unleased entry was available", freed) + } + + if _, serr := os.Stat(pathA); serr != nil { + t.Fatalf("leased entry was evicted: %v", serr) + } + + if got := c.Available(resource.Disk); got != 0 { + t.Fatalf("Available(disk) = %d, want 0 — everything left is pinned", got) + } + + assertLedgerBalanced(t, mgr) +} + +// TestPrivateManagerBehavesLikeTheOldBudget is the degradation +// guarantee: a cache constructed the way every existing caller +// constructs one still owns its own allowance and reports it the same +// way. +func TestPrivateManagerBehavesLikeTheOldBudget(t *testing.T) { + c, b := newCache(t, 64) + b.Put("m", "a", []byte("0123456789")) + + if got := c.Budget(); got != 64 { + t.Fatalf("Budget() = %d, want 64", got) + } + + if got := c.Used(); got != 0 { + t.Fatalf("Used() = %d, want 0 on a fresh cache", got) + } + + _, _, release, err := c.Stage(context.Background(), + artifact.Ref{Bucket: "m", Key: "a", Size: 10}) + if err != nil { + t.Fatalf("Stage: %v", err) + } + + if got := c.Used(); got != 10 { + t.Fatalf("Used() = %d, want 10", got) + } + + release() + + if got := c.Used(); got != 10 { + t.Fatalf("Used() = %d after release, want 10 — releasing a lease keeps the bytes cached", got) + } +} + +// probingManager wraps a resource.Manager and signals every time the +// reclaimer registered for a key is ASKED to free something and answers +// "nothing". +// +// That callback is the only externally visible moment in +// resource.Manager.Acquire's wait loop, and it is precisely the moment +// the test needs: reclaimLocked runs with the manager's lock DROPPED, +// immediately before the acquirer checks the reclaim generation and goes +// into cond.Wait. Releasing the cache entry from inside that window is +// the exact race the generation counter exists to survive — a plain +// Broadcast there reaches nobody, because the acquirer is not waiting yet. +type probingManager struct { + resource.Manager + + refused chan struct{} // "asked to evict, gave nothing" + proceed chan struct{} // test → reclaimer: you may return now +} + +func (m *probingManager) RegisterReclaimer(key string, r resource.Reclaimer) { + if r == nil { + m.Manager.RegisterReclaimer(key, nil) + + return + } + + m.Manager.RegisterReclaimer(key, &probingReclaimer{ + Reclaimer: r, + refused: m.refused, + proceed: m.proceed, + }) +} + +type probingReclaimer struct { + resource.Reclaimer + + once sync.Once + refused chan struct{} + proceed chan struct{} +} + +func (r *probingReclaimer) Reclaim(ctx context.Context, key string, need int64) (int64, error) { + freed, err := r.Reclaimer.Reclaim(ctx, key, need) + + // Only the FIRST refusal is gated. The acquirer rounds its loop more + // than once and a second block would never be released. + // + // Gating AFTER the wrapped call is what makes it safe: the cache has + // let go of its entry table by the time it returns, so the test + // goroutine can release an entry from here without deadlocking + // against it. + if freed == 0 { + r.once.Do(func() { + r.refused <- struct{}{} + <-r.proceed + }) + } + + return freed, err +} + +// TestBlockedStageWakesWhenAnEntryIsReleased covers the wake-up the +// manager cannot generate for itself: it broadcasts when a lease is +// released, and a cache lease dropping to zero is not one of those, yet +// it is exactly when the blocked stager's space becomes available. +// +// Without it the stage below sleeps until its deadline and the job +// requeues for no reason — a worker going quiet rather than erroring. +// This is the guard for the permanent-hang bug found during the track, so +// it is synchronised rather than timed: it waits for the stager to be +// observed asking the cache for space and being refused, which is the +// last thing that happens before it sleeps. A sleep would let the release +// land FIRST on a slow or loaded machine, and then the second Stage +// succeeds trivially without the wake-up path running at all — a broken +// generation counter would be invisible and the test would still pass. +// +// Releasing at the refusal point is also deliberately the HARDEST timing, +// not merely a deterministic one: reclaimLocked holds no manager lock +// across the reclaimer call, so the release and its Wake land inside the +// window where the acquirer has already been told "nothing available" and +// has not yet reached cond.Wait. Only the generation counter carries the +// signal across that gap. +// +// Mutation-verified: deleting the reclaimGen re-check in +// resource.Manager.Acquire makes this hang to the 5s deadline and fail. +func TestBlockedStageWakesWhenAnEntryIsReleased(t *testing.T) { + refused := make(chan struct{}) + proceed := make(chan struct{}) + mgr := &probingManager{ + Manager: resource.NewManager(resource.Set{resource.Disk: 10}), + refused: refused, + proceed: proceed, + } + + c := managedCache(t, mgr, 2, 10) + + _, _, releaseA, err := c.Stage(context.Background(), + artifact.Ref{Bucket: "m", Key: "a", Size: 10}) + if err != nil { + t.Fatalf("Stage a: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + done := make(chan error, 1) + + go func() { + _, _, release, serr := c.Stage(ctx, artifact.Ref{Bucket: "m", Key: "b", Size: 10}) + if release != nil { + release() + } + + done <- serr + }() + + // Wait for the stager to be exactly where this test needs it: it has + // asked the cache to evict, the cache — whose only entry is still + // pinned — has answered "nothing", and it is now held inside that + // call with the manager's lock DROPPED and cond.Wait still ahead of + // it. + select { + case <-refused: + case <-time.After(5 * time.Second): + t.Fatal("the second Stage never reached the cache's reclaimer; " + + "it is not blocked where this test believes it is") + } + + // Both halves of the wake-up land in that window: the entry loses its + // last stager, and the cache's Wake bumps the generation and + // broadcasts to a waiter that is not waiting yet. + releaseA() + + close(proceed) + + select { + case serr := <-done: + if serr != nil { + t.Fatalf("blocked Stage: %v", serr) + } + case <-time.After(5 * time.Second): + t.Fatal("Stage never woke after the entry it needed was released") + } + + assertLedgerBalanced(t, mgr) +} + +// TestConcurrentStageAndReclaim runs the two paths that touch the entry +// table against each other under -race: stagers pinning, using and +// releasing entries while jobs reclaim the same volume out from under +// them. +// +// It is looking for three things a single-threaded test cannot see. A +// deadlock, because Reclaim takes the table lock and then the manager's +// while Acquire is holding neither and Available must hold only the +// first. A path handed back for a file eviction already unlinked, which +// is what the pin-or-retry loop in Stage exists to prevent. And a +// ledger that stops matching its leases, which is what a reclaimer +// crediting the manager itself would produce. +func TestConcurrentStageAndReclaim(t *testing.T) { + const ( + objects = 8 + stagers = 6 + rounds = 40 + ) + + mgr := resource.NewManager(resource.Set{resource.Disk: 45}) + c := managedCache(t, mgr, objects, entrySize) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + var wg sync.WaitGroup + + for s := range stagers { + wg.Add(1) + + go func(s int) { + defer wg.Done() + + for r := range rounds { + ref := artifact.Ref{ + Bucket: "m", + Key: objectKey((s + r) % objects), + Size: entrySize, + } + + path, _, release, err := c.Stage(ctx, ref) + if err != nil { + t.Errorf("stager %d round %d: %v", s, r, err) + + return + } + + // The whole point of a lease: while it is held, the file + // is there. A reclaimer taking it would show up here. + if _, serr := os.Stat(path); serr != nil { + t.Errorf("stager %d round %d: staged path unusable: %v", s, r, serr) + } + + release() + } + }(s) + } + + // A job competing for the same volume, redeeming the disk the cache + // is holding. + wg.Add(1) + + go func() { + defer wg.Done() + + for range rounds { + lease, err := mgr.Acquire(ctx, "job", resource.Set{resource.Disk: 20}) + if err != nil { + t.Errorf("job admission: %v", err) + + return + } + + lease.Release() + } + }() + + wg.Wait() + + assertLedgerBalanced(t, mgr) + + if used, avail := c.Used(), c.Available(resource.Disk); avail > used { + t.Fatalf("Available(disk) = %d exceeds Used() = %d — the cache is offering bytes it does not hold", + avail, used) + } +} + +// TestCachedBytesAdmitAJobThatWouldNotFit is the property the whole +// task exists for. +// +// The manager's free disk is short, so TryAcquire — what a caller that +// cannot wait would use — refuses. Acquire admits the same job, because +// the cache is registered as the disk reclaimer and evicts to cover the +// shortfall. Everything here is real: a real manager, a real cache, real +// files on disk. +func TestCachedBytesAdmitAJobThatWouldNotFit(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100}) + c := managedCache(t, mgr, 6, 10) + + // Sixty bytes of warm, evictable cache. + for i := range 6 { + stageAndRelease(t, c, i) + } + + if free := mgr.Free()[resource.Disk]; free != 40 { + t.Fatalf("Free()[disk] = %d, want 40", free) + } + + if got := mgr.Reclaimable()[resource.Disk]; got != 60 { + t.Fatalf("Reclaimable()[disk] = %d, want 60 — a warm cache is what a worker offers on top of free", + got) + } + + want := resource.Set{resource.Disk: 70} + + if _, ok := mgr.TryAcquire("job", want); ok { + t.Fatal("TryAcquire admitted 70 against 40 free — this test proves nothing if it fits already") + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + lease, err := mgr.Acquire(ctx, "job", want) + if err != nil { + t.Fatalf("Acquire 70 disk with 60 evictable bytes cached: %v — "+ + "the reclaimer path is the whole point of registering the cache", err) + } + + defer lease.Release() + + if free := mgr.Free()[resource.Disk]; free < 0 { + t.Fatalf("Free()[disk] = %d — the ledger went negative", free) + } + + assertLedgerBalanced(t, mgr) + + // Only what the shortfall needed: 30 bytes of cache had to go, which + // is three ten-byte entries. + if used := c.Used(); used != 30 { + t.Fatalf("cache Used() = %d, want 30 — reclamation must stop at the shortfall", used) + } +} diff --git a/artifact/cache/reservation.go b/artifact/cache/reservation.go new file mode 100644 index 0000000..fe1e0f5 --- /dev/null +++ b/artifact/cache/reservation.go @@ -0,0 +1,156 @@ +package cache + +import ( + "context" + "errors" + "fmt" + + "github.com/xraph/dispatch/resource" +) + +// ErrBudgetExceeded means the cache could not free enough space for a +// stage request. +// +// It is returned both when a single artifact is larger than the whole +// budget — which can never succeed and so fails immediately — and when +// every cached entry is currently leased and the caller's deadline +// elapsed while waiting for one to be released. The manager's own +// resource.ErrCapacityExceeded is kept in the chain, so a caller that +// wants the dimension rather than the layer can still find it. +var ErrBudgetExceeded = errors.New("dispatch/artifact/cache: budget exceeded") + +// holdOwner is the lease owner every cached entry is admitted under, so +// an operator reading resource.Manager.Leases() can tell staged bytes +// from the jobs they were staged for. +const holdOwner = "artifact-cache" + +// hold is the manager capacity backing one cached object. +// +// It is a slice rather than a single lease because a lease's size is +// fixed once granted, and a ref that carried no size — every freshly +// registered artifact — is only sized after the copy. Growing appends; +// evicting releases every lease and the bytes come back through the one +// path that credits the manager's ledger. +type hold struct { + leases []resource.Lease + bytes int64 +} + +// acquire takes a manager lease for n bytes of disk. +// +// This is the whole admission path for a staged byte. The manager +// reclaims through this cache's own Reclaim before it blocks and wakes +// on any release, which is the evict-then-wait loop the private budget +// used to run for disk alone — now keyed, and shared with the jobs +// competing for the same volume. +func (c *Cache) acquire(ctx context.Context, n int64) (resource.Lease, error) { + l, err := c.resources.Acquire(ctx, holdOwner, resource.Set{resource.Disk: n}) + if err != nil { + return nil, fmt.Errorf("%w: reserving %d bytes: %w", ErrBudgetExceeded, n, err) + } + + c.used.Add(n) + + return l, nil +} + +// newHold reserves n bytes for an object about to be written. +func (c *Cache) newHold(ctx context.Context, n int64) (*hold, error) { + l, err := c.acquire(ctx, n) + if err != nil { + return nil, err + } + + return &hold{leases: []resource.Lease{l}, bytes: n}, nil +} + +// tryHold reserves n bytes without blocking, for the startup walk. It +// must not block: nothing is waiting to release anything yet, and +// evicting a file to make room for another file already on the same +// disk would free nothing. +func (c *Cache) tryHold(n int64) (*hold, bool) { + l, ok := c.resources.TryAcquire(holdOwner, resource.Set{resource.Disk: n}) + if !ok { + return nil, false + } + + c.used.Add(n) + + return &hold{leases: []resource.Lease{l}, bytes: n}, true +} + +// resize corrects a hold to the bytes that actually landed on disk. +// +// Growing appends a lease rather than replacing one, so the bytes +// already written stay accounted for and only the difference has to be +// admitted. Shrinking releases and re-takes, because a granted lease +// cannot be made smaller; the request that follows is strictly smaller +// than what was just handed back, so it fits unless a concurrent +// acquirer took the difference first — and then it waits like any other. +func (c *Cache) resize(ctx context.Context, h *hold, want int64) error { + switch { + case want == h.bytes: + return nil + + case want > h.bytes: + l, err := c.acquire(ctx, want-h.bytes) + if err != nil { + return err + } + + h.leases = append(h.leases, l) + h.bytes = want + + return nil + } + + c.releaseHold(h) + + l, err := c.acquire(ctx, want) + if err != nil { + return err + } + + h.leases = []resource.Lease{l} + h.bytes = want + + return nil +} + +// releaseHold returns a hold's bytes to the manager. +// +// Releasing the lease is the only thing that credits the ledger — see +// resource.Reclaimer — so this is what makes eviction actually give +// disk back rather than merely delete a file. +func (c *Cache) releaseHold(h *hold) { + if h == nil { + return + } + + for _, l := range h.leases { + l.Release() + } + + c.used.Add(-h.bytes) + + h.leases = nil + h.bytes = 0 +} + +// wake tells the manager this cache has something to give that it did +// not a moment ago. +// +// An entry whose last stager let go is evictable now, and nothing about +// that released a lease, so the manager has no way to notice on its +// own. Manager.Wake carries a generation the acquirer re-checks under +// the lock, which is what makes this survive the window where the +// acquirer is mid-reclaim and not yet waiting — a plain broadcast there +// would be heard by nobody. +// +// Without it a stager blocked behind a fully leased cache sleeps to its +// deadline even though the space it needed was freed a millisecond +// later, and the job requeues for no reason: the worker goes quiet +// instead of failing. +func (c *Cache) wake() { + c.resources.Wake() +} diff --git a/artifact/doc.go b/artifact/doc.go new file mode 100644 index 0000000..0b1fb08 --- /dev/null +++ b/artifact/doc.go @@ -0,0 +1,16 @@ +// Package artifact defines Dispatch's data plane: tracked references to +// objects in external storage, the pluggable Backend interface those +// objects live behind, and the Store contract that persists their +// metadata and ownership links. +// +// This package is a leaf. It imports only the root dispatch package, the +// id package, and stdlib. The staging middleware, which needs job and +// middleware, lives in the artifact/staging sub-package so that job may +// import artifact without a cycle. +// +// Artifacts come in two lifecycles. Durable artifacts are written by the +// application and merely tracked here; Dispatch reads them and never +// deletes them. Ephemeral artifacts are created by Dispatch on a +// handler's behalf, refcounted through links, and swept once every owner +// is terminal and the retention window has passed. +package artifact diff --git a/artifact/errors.go b/artifact/errors.go new file mode 100644 index 0000000..03ac82d --- /dev/null +++ b/artifact/errors.go @@ -0,0 +1,76 @@ +package artifact + +import ( + "errors" + + "github.com/xraph/dispatch" +) + +// Both permanent sentinels below unwrap to dispatch.ErrPermanent, so the +// executor sends a job that hits one straight to the dead letter queue +// instead of retrying it. A caller chooses how precisely to match: +// +// errors.Is(err, artifact.ErrNotFound) // the object is gone +// errors.Is(err, dispatch.ErrPermanent) // don't retry -- also true +// +// Match dispatch.ErrPermanent when the question is whether to retry, and +// the specific sentinel when the answer changes what you do. Code that asks +// "is this missing" in order to decide "should I retry" is asking the wrong +// question: a permission failure is just as permanent, and answering it +// with a retry loop wastes the same budget a deleted input would. +var ( + // ErrNotFound means the artifact or its underlying object does not + // exist. Retrying a fetch of something that no longer exists cannot + // succeed, so it unwraps to dispatch.ErrPermanent. + ErrNotFound error = &categoryError{ + msg: "dispatch/artifact: not found", + parent: dispatch.ErrPermanent, + } + + // ErrPermissionDenied means the backend refused the operation as + // unauthorized. Nothing changes until the credentials or the backend's + // access policy do, so it unwraps to dispatch.ErrPermanent too. + ErrPermissionDenied error = &categoryError{ + msg: "dispatch/artifact: permission denied", + parent: dispatch.ErrPermanent, + } + + // ErrExists means an artifact already exists for this owner, name, + // and a prior attempt. Create with IfAbsent returns it so a retried + // handler can skip recomputation. + ErrExists = errors.New("dispatch/artifact: already exists") + + // ErrSizeExceeded means a bound artifact is larger than the input + // declaration's MaxSize. + ErrSizeExceeded = errors.New("dispatch/artifact: size exceeds declared maximum") + + // ErrImmutable means an attempt was made to delete or overwrite a + // durable artifact through a path reserved for ephemeral ones. + ErrImmutable = errors.New("dispatch/artifact: durable artifacts are immutable") + + // ErrNoBackend means no storage backend is configured. Every artifact + // operation is a no-op in this state and Dispatch behaves exactly as + // it did before the artifact plane existed. + ErrNoBackend = errors.New("dispatch/artifact: no backend configured") + + // ErrUnbound means a required input declaration has no binding on the + // job being executed. + ErrUnbound = errors.New("dispatch/artifact: required input not bound") + + // ErrUndeclared means a binding was supplied for a name the job + // definition does not declare. + ErrUndeclared = errors.New("dispatch/artifact: binding has no matching declaration") +) + +// categoryError is a sentinel that belongs to a broader category, so that +// errors.Is matches both the specific sentinel and its parent. errors.New +// cannot express this because it produces a leaf with nothing to unwrap. +type categoryError struct { + msg string + parent error +} + +func (e *categoryError) Error() string { return e.msg } + +// Unwrap returns the broader category this sentinel belongs to. +func (e *categoryError) Unwrap() error { return e.parent } diff --git a/artifact/errors_test.go b/artifact/errors_test.go new file mode 100644 index 0000000..815bc3b --- /dev/null +++ b/artifact/errors_test.go @@ -0,0 +1,70 @@ +package artifact_test + +import ( + "errors" + "fmt" + "testing" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" +) + +// TestPermanentCategory pins the relationships the retry path depends on. +// Code deciding whether to retry matches ErrPermanent; code deciding what to +// report matches the specific sentinel. Both must keep working. +func TestPermanentCategory(t *testing.T) { + tests := []struct { + name string + err error + target error + want bool + }{ + {"not found matches itself", artifact.ErrNotFound, artifact.ErrNotFound, true}, + {"not found is permanent", artifact.ErrNotFound, dispatch.ErrPermanent, true}, + {"permission denied matches itself", artifact.ErrPermissionDenied, artifact.ErrPermissionDenied, true}, + {"permission denied is permanent", artifact.ErrPermissionDenied, dispatch.ErrPermanent, true}, + + // The two permanent conditions are distinct causes. + {"not found is not permission denied", artifact.ErrNotFound, artifact.ErrPermissionDenied, false}, + {"permission denied is not not-found", artifact.ErrPermissionDenied, artifact.ErrNotFound, false}, + + // The category does not imply any particular cause. + {"permanent is not not-found", dispatch.ErrPermanent, artifact.ErrNotFound, false}, + + // ErrExists is control flow for IfAbsent, not a permanent failure. + {"exists is not permanent", artifact.ErrExists, dispatch.ErrPermanent, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := errors.Is(tt.err, tt.target); got != tt.want { + t.Fatalf("errors.Is(%v, %v) = %v, want %v", tt.err, tt.target, got, tt.want) + } + }) + } +} + +// TestPermanentSurvivesWrapping is the property the call sites rely on: the +// classification has to survive the layers of context added between the +// backend and the executor. +func TestPermanentSurvivesWrapping(t *testing.T) { + err := fmt.Errorf("stage input %q: %w", + "model", fmt.Errorf("stage models/gone.ifc: %w", artifact.ErrNotFound)) + + if !errors.Is(err, dispatch.ErrPermanent) { + t.Fatalf("wrapped ErrNotFound is not permanent: %v", err) + } + + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("wrapped ErrNotFound lost its specific sentinel: %v", err) + } +} + +// TestErrNotFoundMessageUnchanged guards the rendered text, since ErrNotFound +// moved from errors.New to a category type and anything logging or comparing +// the message should not notice. +func TestErrNotFoundMessageUnchanged(t *testing.T) { + if got := artifact.ErrNotFound.Error(); got != "dispatch/artifact: not found" { + t.Fatalf("ErrNotFound.Error() = %q, want %q", got, "dispatch/artifact: not found") + } +} diff --git a/artifact/input.go b/artifact/input.go new file mode 100644 index 0000000..cba1209 --- /dev/null +++ b/artifact/input.go @@ -0,0 +1,156 @@ +package artifact + +import ( + "errors" + "fmt" + "strings" +) + +// StageMode determines how a declared input is presented to the handler. +type StageMode int + +const ( + // StageModePath materialises the artifact to local disk before the + // handler runs, and Accessor.Path returns the file path. + // + // This is the default because the libraries that read Dispatch's heavy + // inputs — CAD kernels, mesh importers, PDF engines — seek and + // memory-map. They need a file, not a stream. + StageModePath StageMode = iota + + // StageModeLazy downloads nothing up front. The handler calls + // Accessor.Open to stream the bytes if and when it needs them. + // + // Right for data read once, front to back. Wrong for a multi-gigabyte + // model a native library will seek around in. + StageModeLazy +) + +// String renders the mode for logs and errors. +func (m StageMode) String() string { + switch m { + case StageModePath: + return "path" + case StageModeLazy: + return "lazy" + default: + return fmt.Sprintf("StageMode(%d)", int(m)) + } +} + +// InputSpec declares one artifact a job consumes. +// +// Declaring inputs, rather than burying refs in an opaque payload, is +// what lets the engine know the total input size before it schedules the +// job, validate bindings at enqueue instead of at run time, and stage the +// bytes before the handler is ever called. +type InputSpec struct { + // Name identifies the input in the handler and in bindings. + Name string + // Required fails the job when no binding is supplied. + Required bool + // MaxSize rejects an oversized binding at enqueue. Zero means no limit. + MaxSize int64 + // Mode selects staging behaviour. + Mode StageMode +} + +// InputOption configures an InputSpec. +type InputOption func(*InputSpec) + +// Required marks an input as mandatory. +func Required(s *InputSpec) { s.Required = true } + +// MaxSize caps the size of a bound artifact. The limit is enforced at +// enqueue, so an oversized input never becomes a failed job. +func MaxSize(bytes int64) InputOption { + return func(s *InputSpec) { s.MaxSize = bytes } +} + +// StageAsPath pre-downloads the input and exposes it as a file path. +func StageAsPath(s *InputSpec) { s.Mode = StageModePath } + +// StageLazy skips the download; the handler streams via Open. +func StageLazy(s *InputSpec) { s.Mode = StageModeLazy } + +// Input declares an artifact input on a job definition. +func Input(name string, opts ...InputOption) InputSpec { + spec := InputSpec{Name: name, Mode: StageModePath} + + for _, opt := range opts { + opt(&spec) + } + + return spec +} + +// Validate reports whether the declaration is usable. +// +// The name becomes both a path component in the storage key and a +// filename in the staging directory, so anything that could escape either +// is rejected here rather than at run time. +func (s InputSpec) Validate() error { + switch { + case s.Name == "": + return errors.New("dispatch/artifact: input name must not be empty") + case strings.ContainsAny(s.Name, `/\`): + return fmt.Errorf("dispatch/artifact: input name %q must not contain a path separator", s.Name) + case strings.Contains(s.Name, ".."): + return fmt.Errorf("dispatch/artifact: input name %q must not contain %q", s.Name, "..") + case s.MaxSize < 0: + return fmt.Errorf("dispatch/artifact: input %q has a negative MaxSize", s.Name) + default: + return nil + } +} + +// ValidateInputs checks a definition's declarations as a set, rejecting +// invalid names and duplicates. +func ValidateInputs(specs []InputSpec) error { + seen := make(map[string]bool, len(specs)) + + for _, spec := range specs { + if err := spec.Validate(); err != nil { + return err + } + + if seen[spec.Name] { + return fmt.Errorf("dispatch/artifact: duplicate input declaration %q", spec.Name) + } + + seen[spec.Name] = true + } + + return nil +} + +// TotalMaxSize sums the declared limits. It is zero when any declaration +// is unbounded, since the total is then unknown rather than small. +// +// The engine uses this to reject, at registration, a definition whose +// inputs could never fit the staging budget — so an unstageable job fails +// on a developer's machine instead of at 3am. +func TotalMaxSize(specs []InputSpec) int64 { + var total int64 + + for _, spec := range specs { + if spec.MaxSize <= 0 { + return 0 + } + + total += spec.MaxSize + } + + return total +} + +// FindInput returns the declaration with the given name. +func FindInput(specs []InputSpec, name string) (InputSpec, bool) { + for _, spec := range specs { + if spec.Name == name { + return spec, true + } + } + + return InputSpec{}, false +} diff --git a/artifact/input_test.go b/artifact/input_test.go new file mode 100644 index 0000000..ad51995 --- /dev/null +++ b/artifact/input_test.go @@ -0,0 +1,153 @@ +package artifact_test + +import ( + "testing" + + "github.com/xraph/dispatch/artifact" +) + +func TestInputDefaults(t *testing.T) { + in := artifact.Input("model") + + if in.Name != "model" { + t.Fatalf("Name = %q, want %q", in.Name, "model") + } + + if in.Required { + t.Fatal("inputs must be optional by default") + } + + if in.Mode != artifact.StageModePath { + t.Fatalf("default mode = %v, want StageModePath", in.Mode) + } +} + +func TestInputOptions(t *testing.T) { + in := artifact.Input("model", + artifact.Required, + artifact.MaxSize(8<<30), + artifact.StageLazy) + + if !in.Required { + t.Fatal("Required not applied") + } + + if in.MaxSize != 8<<30 { + t.Fatalf("MaxSize = %d, want %d", in.MaxSize, int64(8)<<30) + } + + if in.Mode != artifact.StageModeLazy { + t.Fatalf("Mode = %v, want StageModeLazy", in.Mode) + } +} + +func TestInputValidate(t *testing.T) { + tests := []struct { + name string + spec artifact.InputSpec + wantErr bool + }{ + {"valid", artifact.Input("model"), false}, + {"valid with dots", artifact.Input("model.v2"), false}, + {"empty name", artifact.Input(""), true}, + {"negative max size", artifact.Input("m", artifact.MaxSize(-1)), true}, + {"parent traversal", artifact.Input("../etc/passwd"), true}, + {"forward slash", artifact.Input("a/b"), true}, + {"backslash", artifact.Input(`a\b`), true}, + {"bare traversal", artifact.Input(".."), true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.spec.Validate() + if (err != nil) != tt.wantErr { + t.Fatalf("Validate() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func TestValidateInputsRejectsDuplicates(t *testing.T) { + err := artifact.ValidateInputs([]artifact.InputSpec{ + artifact.Input("model"), + artifact.Input("model"), + }) + if err == nil { + t.Fatal("duplicate input names must be rejected") + } +} + +func TestValidateInputsAcceptsDistinct(t *testing.T) { + err := artifact.ValidateInputs([]artifact.InputSpec{ + artifact.Input("model"), + artifact.Input("textures"), + }) + if err != nil { + t.Fatalf("ValidateInputs: %v", err) + } +} + +func TestTotalMaxSize(t *testing.T) { + tests := []struct { + name string + specs []artifact.InputSpec + want int64 + }{ + { + name: "none", + specs: nil, + want: 0, + }, + { + name: "all bounded", + specs: []artifact.InputSpec{ + artifact.Input("a", artifact.MaxSize(100)), + artifact.Input("b", artifact.MaxSize(200)), + }, + want: 300, + }, + { + // One unbounded declaration makes the total unknown, not small. + name: "one unbounded", + specs: []artifact.InputSpec{ + artifact.Input("a", artifact.MaxSize(100)), + artifact.Input("b"), + }, + want: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := artifact.TotalMaxSize(tt.specs); got != tt.want { + t.Fatalf("TotalMaxSize() = %d, want %d", got, tt.want) + } + }) + } +} + +func TestFindInput(t *testing.T) { + specs := []artifact.InputSpec{ + artifact.Input("model"), + artifact.Input("textures"), + } + + got, ok := artifact.FindInput(specs, "textures") + if !ok || got.Name != "textures" { + t.Fatalf("FindInput = %+v, %v; want the textures spec", got, ok) + } + + if _, ok := artifact.FindInput(specs, "absent"); ok { + t.Fatal("FindInput found a declaration that does not exist") + } +} + +func TestStageModeString(t *testing.T) { + if got := artifact.StageModePath.String(); got != "path" { + t.Fatalf("StageModePath.String() = %q, want %q", got, "path") + } + + if got := artifact.StageModeLazy.String(); got != "lazy" { + t.Fatalf("StageModeLazy.String() = %q, want %q", got, "lazy") + } +} diff --git a/artifact/service.go b/artifact/service.go new file mode 100644 index 0000000..52d1d18 --- /dev/null +++ b/artifact/service.go @@ -0,0 +1,569 @@ +package artifact + +import ( + "context" + "errors" + "fmt" + "io" + "path" + "strconv" + "strings" + "time" + + "github.com/xraph/dispatch/id" +) + +// DefaultEphemeralPrefix is where Dispatch-owned objects live when no +// prefix is configured. +const DefaultEphemeralPrefix = "ephemeral" + +// Service is the operational face of the artifact plane. It pairs a Store +// with a Backend and owns the rules that keep the two consistent: +// registration is idempotent, ephemeral keys embed the attempt, and — on +// the CommitWriter path (Create/CreateFenced/Commit) — an artifact row is +// never written without its link. Register is the exception: it writes a +// durable row for a pre-existing backend object with a nil link, which +// Store.CreateArtifact's own contract supports, since there is no +// job/attempt to link it to yet. +type Service struct { + store Store + backend Backend + + ephemeralPrefix string + defaultBucket string + retention time.Duration +} + +// ServiceOption configures a Service. +type ServiceOption func(*Service) + +// WithEphemeralPrefix sets the key prefix for Dispatch-owned objects. +func WithEphemeralPrefix(prefix string) ServiceOption { + return func(s *Service) { s.ephemeralPrefix = strings.Trim(prefix, "/") } +} + +// WithDefaultBucket sets the bucket ephemeral objects are written to. +func WithDefaultBucket(bucket string) ServiceOption { + return func(s *Service) { s.defaultBucket = bucket } +} + +// WithRetention sets the default retention applied to ephemeral artifacts +// that do not carry their own expiry. +func WithRetention(d time.Duration) ServiceOption { + return func(s *Service) { s.retention = d } +} + +// NewService creates a Service. A nil backend leaves the artifact plane +// disabled: every method returns ErrNoBackend and Dispatch behaves exactly +// as it did before artifacts existed. +func NewService(store Store, backend Backend, opts ...ServiceOption) *Service { + s := &Service{ + store: store, + backend: backend, + ephemeralPrefix: DefaultEphemeralPrefix, + } + + for _, opt := range opts { + opt(s) + } + + return s +} + +// Store returns the underlying persistence layer. +func (s *Service) Store() Store { return s.store } + +// Backend returns the underlying object storage. +func (s *Service) Backend() Backend { return s.backend } + +// Enabled reports whether a backend is configured. +func (s *Service) Enabled() bool { return s != nil && s.backend != nil } + +// DefaultBucket returns the bucket ephemeral objects are written to. +func (s *Service) DefaultBucket() string { return s.defaultBucket } + +// ── Register ────────────────────────────────────────────────────── + +// RegisterOptions configures registration of a durable artifact. +type RegisterOptions struct { + ScopeAppID string + ScopeOrgID string + ContentType string +} + +// RegisterOption configures Register. +type RegisterOption func(*RegisterOptions) + +// WithScope tags the artifact with a tenant application and organization. +func WithScope(appID, orgID string) RegisterOption { + return func(o *RegisterOptions) { + o.ScopeAppID = appID + o.ScopeOrgID = orgID + } +} + +// WithContentType records the artifact's media type. +func WithContentType(ct string) RegisterOption { + return func(o *RegisterOptions) { o.ContentType = ct } +} + +// Register records an object the application already uploaded, returning +// a durable Ref that Dispatch will read but never delete. +// +// Register does not hash the object. Hashing a multi-gigabyte file would +// turn enqueue into a full read pass; the hash is filled in later by the +// staging cache, which is already streaming every byte to disk. Until +// then the artifact is identified by its storage coordinates. +// +// Registering the same coordinates twice returns the existing ref rather +// than an error, so callers can register unconditionally. +func (s *Service) Register(ctx context.Context, bucket, key string, opts ...RegisterOption) (Ref, error) { + if !s.Enabled() { + return Ref{}, ErrNoBackend + } + + var cfg RegisterOptions + for _, opt := range opts { + opt(&cfg) + } + + probe := Ref{Backend: s.backend.Name(), Bucket: bucket, Key: key} + + info, err := s.backend.Stat(ctx, probe) + if err != nil { + if errors.Is(err, ErrNotFound) { + return Ref{}, fmt.Errorf("register %s/%s: %w", bucket, key, ErrNotFound) + } + + return Ref{}, fmt.Errorf("dispatch/artifact: stat %s/%s: %w", bucket, key, err) + } + + contentType := cfg.ContentType + if contentType == "" { + contentType = info.ContentType + } + + a := &Artifact{ + ID: id.NewArtifactID(), + Backend: s.backend.Name(), + Bucket: bucket, + Key: key, + Size: info.Size, + ContentType: contentType, + Lifecycle: Durable, + ScopeAppID: cfg.ScopeAppID, + ScopeOrgID: cfg.ScopeOrgID, + CreatedAt: time.Now().UTC(), + } + + err = s.store.CreateArtifact(ctx, a, nil) + + switch { + case err == nil: + return a.Ref(), nil + + case errors.Is(err, ErrExists): + existing, ferr := s.store.FindArtifactByKey(ctx, a.Backend, bucket, key) + if ferr != nil { + return Ref{}, fmt.Errorf("dispatch/artifact: resolve existing artifact: %w", ferr) + } + + return existing.Ref(), nil + + default: + return Ref{}, fmt.Errorf("dispatch/artifact: register: %w", err) + } +} + +// Get resolves a ref to its stored artifact. +func (s *Service) Get(ctx context.Context, artifactID id.ArtifactID) (*Artifact, error) { + if !s.Enabled() { + return nil, ErrNoBackend + } + + return s.store.GetArtifact(ctx, artifactID) +} + +// Open streams an artifact's bytes. +func (s *Service) Open(ctx context.Context, ref Ref) (io.ReadCloser, error) { + if !s.Enabled() { + return nil, ErrNoBackend + } + + return s.backend.Open(ctx, ref) +} + +// ── Create ──────────────────────────────────────────────────────── + +// CreateOptions configures creation of an ephemeral artifact. +type CreateOptions struct { + ContentType string + ScopeAppID string + ScopeOrgID string + Retention time.Duration + IfAbsent bool +} + +// CreateOption configures Create. +type CreateOption func(*CreateOptions) + +// ContentType records the media type of the object being written. +func ContentType(ct string) CreateOption { + return func(o *CreateOptions) { o.ContentType = ct } +} + +// Scope tags the created artifact with a tenant application and org. +func Scope(appID, orgID string) CreateOption { + return func(o *CreateOptions) { + o.ScopeAppID = appID + o.ScopeOrgID = orgID + } +} + +// Retain overrides the default retention for this artifact. +func Retain(d time.Duration) CreateOption { + return func(o *CreateOptions) { o.Retention = d } +} + +// IfAbsent makes Create return ErrExists when a previous attempt of the +// same owner already committed this name. +// +// This is what lets a retried handler skip work it already did: a job +// splitting a 400-page PDF can resume at page 317 instead of re-rendering +// the 316 pages a prior attempt committed. +func IfAbsent() CreateOption { + return func(o *CreateOptions) { o.IfAbsent = true } +} + +// EphemeralKey returns the storage key for an owner's output. +// +// The attempt is part of the key because Commit is attempt-scoped while +// storage coordinates are unique: without it, a retried job writing the +// same name would collide with its own previous attempt. +func (s *Service) EphemeralKey(owner OwnerRef, attempt int, name string) string { + return path.Join( + s.ephemeralPrefix, + string(owner.Kind), + owner.ID, + strconv.Itoa(attempt), + name, + ) +} + +// FindExisting returns the artifact a previous attempt committed under +// this owner and name, across all attempts. +func (s *Service) FindExisting(ctx context.Context, owner OwnerRef, name string) (Ref, error) { + if !s.Enabled() { + return Ref{}, ErrNoBackend + } + + link, err := s.store.FindLinkByName(ctx, owner, name) + if err != nil { + return Ref{}, err + } + + a, err := s.store.GetArtifact(ctx, link.ArtifactID) + if err != nil { + return Ref{}, err + } + + return a.Ref(), nil +} + +// FindCommitted returns the artifact already sitting at the exact +// storage coordinates Create or CreateFenced would use for +// (owner, attempt, name, fenceToken), if one exists. It returns +// ErrNotFound when none does. +// +// This is narrower than FindExisting on purpose: FindExisting answers +// "has ANY attempt committed this name," which is what a handler's own +// Existing/IfAbsent check wants. FindCommitted answers "did THIS EXACT +// caller already commit THIS EXACT thing" — which is what lets a caller +// recognise its own earlier, successful partial work as a no-op to skip +// rather than a collision to fail on, without also treating a +// DIFFERENT holder's commit of the same name under a different +// fenceToken as anything but what it is: a separate object at a +// separate key. See worker.Runner.commitOutputFile, the motivating +// caller: a retry of a launch-failure-classified commit failure reuses +// the identical (owner, attempt) — RetryCount does not advance for a +// launch failure — so without this, re-committing a name that already +// landed in an earlier, partially-failed pass would collide with +// itself on every subsequent attempt. +func (s *Service) FindCommitted(ctx context.Context, owner OwnerRef, attempt int, name, fenceToken string) (Ref, error) { + if !s.Enabled() { + return Ref{}, ErrNoBackend + } + + key := s.EphemeralKey(owner, attempt, name) + if fenceToken != "" { + key = path.Join(key, fenceToken) + } + + a, err := s.store.FindArtifactByKey(ctx, s.backend.Name(), s.defaultBucket, key) + if err != nil { + return Ref{}, err + } + + return a.Ref(), nil +} + +// Create begins writing an ephemeral artifact owned by owner. +// +// The returned writer publishes nothing until Commit; Abort discards it. +// With IfAbsent, a name a prior attempt already committed returns +// ErrExists so the caller can skip the work. +func (s *Service) Create( + ctx context.Context, + owner OwnerRef, + attempt int, + name string, + opts ...CreateOption, +) (*CommitWriter, error) { + return s.create(ctx, owner, attempt, name, "", opts...) +} + +// CreateFenced is Create, except the storage key additionally +// incorporates fenceToken, so two callers racing to commit the same +// (owner, attempt, name) under different fenceTokens can never collide +// on the same backend object. +// +// This exists for a caller whose own claim to "the current holder of +// (owner, attempt)" is itself fenced — a worker committing an +// out-of-process rung's outputs under a lease epoch is the motivating +// case (see worker.Runner.commitOutputFile). Two workers can each +// believe they hold the same job at the same RetryCount at once: a +// lease reclaim races a worker that has not yet noticed its lease +// expired and is still finishing a long attempt. Create's key is a +// pure function of (owner, attempt, name), so both would resolve to +// the identical backend object — whichever Commit lands second would +// silently overwrite the first's bytes, behind a store row that still +// claims the first writer's size, with no error surfaced to anyone. A +// distinct fenceToken per holder (their lease epoch) gives each +// holder its own object instead: the loser's write lands next to the +// winner's rather than on top of it, and the store's own uniqueness +// check on (backend, bucket, key) can then only ever reject an actual +// repeat of the SAME holder recommitting the SAME name, never one +// holder clobbering another's bytes purely by losing a race. +// +// fenceToken is never recorded on Link or Artifact — Attempt keeps +// meaning exactly what it means everywhere else in this package — it +// only ever changes the storage key most callers never see. A caller +// with no fence to offer should use Create; an empty fenceToken here +// is refused rather than silently behaving like Create, so a caller +// that meant to fence but passed a zero value fails loudly instead of +// losing the protection without noticing. +func (s *Service) CreateFenced( + ctx context.Context, + owner OwnerRef, + attempt int, + name string, + fenceToken string, + opts ...CreateOption, +) (*CommitWriter, error) { + if fenceToken == "" { + return nil, fmt.Errorf("dispatch/artifact: create fenced %q: empty fence token", name) + } + + return s.create(ctx, owner, attempt, name, fenceToken, opts...) +} + +// create is the shared implementation behind Create and CreateFenced. +// An empty fenceToken makes it behave exactly as Create always has; +// Create is a thin wrapper passing exactly that. +func (s *Service) create( + ctx context.Context, + owner OwnerRef, + attempt int, + name string, + fenceToken string, + opts ...CreateOption, +) (*CommitWriter, error) { + if !s.Enabled() { + return nil, ErrNoBackend + } + + if !owner.Valid() { + return nil, fmt.Errorf("dispatch/artifact: create %q: invalid owner", name) + } + + if err := validateName(name); err != nil { + return nil, err + } + + var cfg CreateOptions + for _, opt := range opts { + opt(&cfg) + } + + if cfg.IfAbsent { + if _, err := s.FindExisting(ctx, owner, name); err == nil { + return nil, ErrExists + } else if !errors.Is(err, ErrNotFound) { + return nil, err + } + } + + bucket := s.defaultBucket + key := s.EphemeralKey(owner, attempt, name) + if fenceToken != "" { + key = path.Join(key, fenceToken) + } + + w, err := s.backend.Create(ctx, bucket, key) + if err != nil { + return nil, fmt.Errorf("dispatch/artifact: create %s/%s: %w", bucket, key, err) + } + + retention := cfg.Retention + if retention == 0 { + retention = s.retention + } + + return &CommitWriter{ + svc: s, + inner: w, + owner: owner, + attempt: attempt, + name: name, + bucket: bucket, + key: key, + cfg: cfg, + retention: retention, + }, nil +} + +// validateName rejects names that would escape the ephemeral prefix or +// the staging directory. The name becomes both a path component in the +// storage key and a filename on disk. +func validateName(name string) error { + switch { + case name == "": + return errors.New("dispatch/artifact: name must not be empty") + case strings.ContainsAny(name, `/\`): + return fmt.Errorf("dispatch/artifact: name %q must not contain a path separator", name) + case strings.Contains(name, ".."): + return fmt.Errorf("dispatch/artifact: name %q must not contain %q", name, "..") + default: + return nil + } +} + +// CommitWriter writes an ephemeral artifact's bytes and, on Commit, +// records the artifact and its link atomically. +type CommitWriter struct { + svc *Service + inner Writer + owner OwnerRef + attempt int + name string + bucket string + key string + cfg CreateOptions + retention time.Duration + + committed bool + aborted bool +} + +// Write appends bytes to the pending object. +func (w *CommitWriter) Write(p []byte) (int, error) { return w.inner.Write(p) } + +// Commit finalises the object, inserts the artifact, and links it to the +// owner as an output. The store writes both in one operation, so a +// zero-link artifact cannot result from a normal race. +func (w *CommitWriter) Commit(ctx context.Context) (Ref, error) { + if w.committed { + return Ref{}, errors.New("dispatch/artifact: writer already committed") + } + + if w.aborted { + return Ref{}, errors.New("dispatch/artifact: writer already aborted") + } + + info, err := w.inner.Commit(ctx) + if err != nil { + return Ref{}, fmt.Errorf("dispatch/artifact: commit %s/%s: %w", w.bucket, w.key, err) + } + + w.committed = true + + now := time.Now().UTC() + + contentType := w.cfg.ContentType + if contentType == "" { + contentType = info.ContentType + } + + a := &Artifact{ + ID: id.NewArtifactID(), + Backend: w.svc.backend.Name(), + Bucket: w.bucket, + Key: w.key, + Size: info.Size, + ContentType: contentType, + Lifecycle: Ephemeral, + ScopeAppID: w.cfg.ScopeAppID, + ScopeOrgID: w.cfg.ScopeOrgID, + CreatedAt: now, + } + + if w.retention > 0 { + expires := now.Add(w.retention) + a.ExpiresAt = &expires + } + + link := &Link{ + ArtifactID: a.ID, + OwnerKind: w.owner.Kind, + OwnerID: w.owner.ID, + Role: RoleOutput, + Name: w.name, + Attempt: w.attempt, + CreatedAt: now, + } + + if err := w.svc.store.CreateArtifact(ctx, a, link); err != nil { + return Ref{}, fmt.Errorf("dispatch/artifact: record %s/%s: %w", w.bucket, w.key, err) + } + + return a.Ref(), nil +} + +// Abort discards the pending object. It is a no-op after Commit, so +// `defer w.Abort()` is the correct idiom alongside a successful commit. +func (w *CommitWriter) Abort() error { + if w.committed || w.aborted { + return nil + } + + w.aborted = true + + return w.inner.Abort() +} + +// ── Link ────────────────────────────────────────────────────────── + +// Link records that an owner references an existing artifact. This is how +// a declared input is attributed to the job that consumed it. +func (s *Service) Link( + ctx context.Context, + ref Ref, + owner OwnerRef, + role Role, + name string, + attempt int, +) error { + if !s.Enabled() { + return ErrNoBackend + } + + return s.store.LinkArtifact(ctx, &Link{ + ArtifactID: ref.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: role, + Name: name, + Attempt: attempt, + CreatedAt: time.Now().UTC(), + }) +} diff --git a/artifact/service_test.go b/artifact/service_test.go new file mode 100644 index 0000000..c1ad977 --- /dev/null +++ b/artifact/service_test.go @@ -0,0 +1,358 @@ +package artifact_test + +import ( + "context" + "errors" + "io" + "strings" + "testing" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/store/memory" +) + +func newService(t *testing.T) (*artifact.Service, *artifacttest.Backend, artifact.Store) { + t.Helper() + + b := artifacttest.NewBackend() + st := memory.New() + svc := artifact.NewService(st, b, + artifact.WithEphemeralPrefix("ephemeral"), + artifact.WithDefaultBucket("dispatch")) + + return svc, b, st +} + +func newJobOwner() artifact.OwnerRef { + return artifact.OwnerRef{Kind: artifact.OwnerJob, ID: id.NewJobID().String()} +} + +func TestRegisterDurable(t *testing.T) { + ctx := context.Background() + svc, b, _ := newService(t) + b.Put("models", "tower.ifc", []byte("0123456789")) + + ref, err := svc.Register(ctx, "models", "tower.ifc") + if err != nil { + t.Fatalf("Register: %v", err) + } + + if ref.Size != 10 { + t.Fatalf("ref.Size = %d, want 10 (Register must Stat)", ref.Size) + } + + if ref.ID.Prefix() != id.PrefixArtifact { + t.Fatalf("ref.ID prefix = %q, want %q", ref.ID.Prefix(), id.PrefixArtifact) + } + + if ref.ContentHash != "" { + t.Fatal("Register must not hash — hashing is deferred to first staging") + } +} + +func TestRegisterMissingObject(t *testing.T) { + svc, _, _ := newService(t) + + _, err := svc.Register(context.Background(), "models", "nope.ifc") + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Register(missing) = %v, want ErrNotFound", err) + } +} + +func TestRegisterIsIdempotent(t *testing.T) { + ctx := context.Background() + svc, b, _ := newService(t) + b.Put("models", "same.ifc", []byte("abc")) + + first, err := svc.Register(ctx, "models", "same.ifc") + if err != nil { + t.Fatalf("first Register: %v", err) + } + + second, err := svc.Register(ctx, "models", "same.ifc") + if err != nil { + t.Fatalf("second Register: %v", err) + } + + if first.ID != second.ID { + t.Fatalf("Register not idempotent: %v then %v", first.ID, second.ID) + } +} + +func TestCreateCommitLinksOutput(t *testing.T) { + ctx := context.Background() + svc, _, st := newService(t) + owner := newJobOwner() + + w, err := svc.Create(ctx, owner, 0, "mesh.glb", + artifact.ContentType("model/gltf-binary")) + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, cerr := io.Copy(w, strings.NewReader("meshbytes")); cerr != nil { + t.Fatalf("Copy: %v", cerr) + } + + ref, err := w.Commit(ctx) + if err != nil { + t.Fatalf("Commit: %v", err) + } + + if ref.Size != 9 { + t.Fatalf("ref.Size = %d, want 9", ref.Size) + } + + if !strings.Contains(ref.Key, "/0/mesh.glb") { + t.Fatalf("ephemeral key %q must embed the attempt", ref.Key) + } + + links, err := st.ListLinks(ctx, owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + if len(links) != 1 || links[0].Role != artifact.RoleOutput { + t.Fatalf("links = %+v, want one output link", links) + } + + a, err := st.GetArtifact(ctx, ref.ID) + if err != nil { + t.Fatalf("GetArtifact: %v", err) + } + + if a.Lifecycle != artifact.Ephemeral { + t.Fatalf("created artifact lifecycle = %q, want ephemeral", a.Lifecycle) + } +} + +func TestCreateKeysDifferPerAttempt(t *testing.T) { + ctx := context.Background() + svc, _, _ := newService(t) + owner := newJobOwner() + + keys := make(map[string]bool) + + for attempt := range 3 { + w, err := svc.Create(ctx, owner, attempt, "mesh.glb") + if err != nil { + t.Fatalf("Create attempt %d: %v", attempt, err) + } + + if _, werr := w.Write([]byte("x")); werr != nil { + t.Fatalf("Write attempt %d: %v", attempt, werr) + } + + ref, err := w.Commit(ctx) + if err != nil { + t.Fatalf("Commit attempt %d: %v", attempt, err) + } + + if keys[ref.Key] { + t.Fatalf("attempt %d reused key %q — the unique constraint would fire", attempt, ref.Key) + } + + keys[ref.Key] = true + } +} + +func TestCreateIfAbsentFindsPriorAttempt(t *testing.T) { + ctx := context.Background() + svc, _, _ := newService(t) + owner := newJobOwner() + + w, err := svc.Create(ctx, owner, 0, "page-317.png") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write([]byte("pixels")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + first, err := w.Commit(ctx) + if err != nil { + t.Fatalf("Commit: %v", err) + } + + if _, aerr := svc.Create(ctx, owner, 1, "page-317.png", artifact.IfAbsent()); !errors.Is(aerr, artifact.ErrExists) { + t.Fatalf("IfAbsent on attempt 1 = %v, want ErrExists", aerr) + } + + existing, err := svc.FindExisting(ctx, owner, "page-317.png") + if err != nil { + t.Fatalf("FindExisting: %v", err) + } + + if existing.ID != first.ID { + t.Fatalf("FindExisting = %v, want the attempt-0 artifact %v", existing.ID, first.ID) + } +} + +func TestCreateIfAbsentAllowsNewName(t *testing.T) { + ctx := context.Background() + svc, _, _ := newService(t) + owner := newJobOwner() + + w, err := svc.Create(ctx, owner, 1, "page-318.png", artifact.IfAbsent()) + if err != nil { + t.Fatalf("Create(IfAbsent) on a fresh name: %v", err) + } + + if err := w.Abort(); err != nil { + t.Fatalf("Abort: %v", err) + } +} + +func TestAbortLeavesNothingBehind(t *testing.T) { + ctx := context.Background() + svc, b, st := newService(t) + owner := newJobOwner() + + w, err := svc.Create(ctx, owner, 0, "partial.bin") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write([]byte("half")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + if aerr := w.Abort(); aerr != nil { + t.Fatalf("Abort: %v", aerr) + } + + links, err := st.ListLinks(ctx, owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + if len(links) != 0 { + t.Fatalf("aborted write left %d links, want 0", len(links)) + } + + if b.Creates() != 1 { + t.Fatalf("Creates() = %d, want 1", b.Creates()) + } +} + +func TestAbortAfterCommitIsNoOp(t *testing.T) { + ctx := context.Background() + svc, _, _ := newService(t) + owner := newJobOwner() + + w, err := svc.Create(ctx, owner, 0, "out.bin") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write([]byte("data")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + if _, cerr := w.Commit(ctx); cerr != nil { + t.Fatalf("Commit: %v", cerr) + } + + if aerr := w.Abort(); aerr != nil { + t.Fatalf("Abort after Commit must be a no-op, got %v", aerr) + } +} + +func TestCreateRejectsUnsafeNames(t *testing.T) { + ctx := context.Background() + svc, _, _ := newService(t) + owner := newJobOwner() + + tests := []struct { + name string + arg string + }{ + {"empty", ""}, + {"slash", "a/b.png"}, + {"backslash", `a\b.png`}, + {"parent traversal", "../escape.png"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if _, err := svc.Create(ctx, owner, 0, tt.arg); err == nil { + t.Fatalf("Create(%q) succeeded, want an error", tt.arg) + } + }) + } +} + +func TestRetainSetsExpiry(t *testing.T) { + ctx := context.Background() + svc, _, st := newService(t) + owner := newJobOwner() + + w, err := svc.Create(ctx, owner, 0, "temp.bin", artifact.Retain(time.Hour)) + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write([]byte("x")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + ref, err := w.Commit(ctx) + if err != nil { + t.Fatalf("Commit: %v", err) + } + + a, err := st.GetArtifact(ctx, ref.ID) + if err != nil { + t.Fatalf("GetArtifact: %v", err) + } + + if a.ExpiresAt == nil { + t.Fatal("Retain did not set ExpiresAt") + } +} + +func TestDisabledServiceReturnsErrNoBackend(t *testing.T) { + ctx := context.Background() + svc := artifact.NewService(memory.New(), nil) + + if svc.Enabled() { + t.Fatal("service with a nil backend reports enabled") + } + + if _, err := svc.Register(ctx, "b", "k"); !errors.Is(err, artifact.ErrNoBackend) { + t.Fatalf("Register = %v, want ErrNoBackend", err) + } + + if _, err := svc.Create(ctx, newJobOwner(), 0, "x.bin"); !errors.Is(err, artifact.ErrNoBackend) { + t.Fatalf("Create = %v, want ErrNoBackend", err) + } +} + +func TestLinkInput(t *testing.T) { + ctx := context.Background() + svc, b, st := newService(t) + b.Put("models", "in.ifc", []byte("data")) + + ref, err := svc.Register(ctx, "models", "in.ifc") + if err != nil { + t.Fatalf("Register: %v", err) + } + + owner := newJobOwner() + if lerr := svc.Link(ctx, ref, owner, artifact.RoleInput, "model", 0); lerr != nil { + t.Fatalf("Link: %v", lerr) + } + + arts, err := st.ListArtifactsByOwner(ctx, owner, artifact.RoleInput) + if err != nil { + t.Fatalf("ListArtifactsByOwner: %v", err) + } + + if len(arts) != 1 || arts[0].ID != ref.ID { + t.Fatalf("ListArtifactsByOwner = %+v, want the registered input", arts) + } +} diff --git a/artifact/staging/accessor.go b/artifact/staging/accessor.go new file mode 100644 index 0000000..51956f0 --- /dev/null +++ b/artifact/staging/accessor.go @@ -0,0 +1,97 @@ +package staging + +import ( + "context" + "errors" + "fmt" + "io" + "os" + + "github.com/xraph/dispatch/artifact" +) + +// staged is one input the middleware materialised for a handler. +type staged struct { + ref artifact.Ref + path string + mode artifact.StageMode +} + +// accessor is the artifact.Accessor handed to a running handler. It +// closes over the job's owner and attempt so the handler does not have to +// know either. +type accessor struct { + svc *artifact.Service + owner artifact.OwnerRef + attempt int + inputs map[string]staged +} + +var _ artifact.Accessor = (*accessor)(nil) + +// Path returns the local file path of a staged input. +func (a *accessor) Path(name string) string { + in, ok := a.inputs[name] + if !ok { + return "" + } + + return in.path +} + +// Ref returns the artifact bound to an input. +func (a *accessor) Ref(name string) (artifact.Ref, bool) { + in, ok := a.inputs[name] + if !ok { + return artifact.Ref{}, false + } + + return in.ref, true +} + +// Open streams an input's bytes. +// +// A path-staged input is read from local disk, so a handler that prefers +// a reader does not pay for a second download. A lazy one streams from +// the backend on demand. +func (a *accessor) Open(ctx context.Context, name string) (io.ReadCloser, error) { + in, ok := a.inputs[name] + if !ok { + return nil, fmt.Errorf("%w: %q", artifact.ErrUnbound, name) + } + + if in.path != "" { + f, err := os.Open(in.path) + if err != nil { + return nil, fmt.Errorf("dispatch/artifact/staging: open staged %q: %w", name, err) + } + + return f, nil + } + + return a.svc.Open(ctx, in.ref) +} + +// Create begins writing an output owned by the running job. +func (a *accessor) Create( + ctx context.Context, + name string, + opts ...artifact.CreateOption, +) (*artifact.CommitWriter, error) { + return a.svc.Create(ctx, a.owner, a.attempt, name, opts...) +} + +// Existing reports an artifact a previous attempt committed under this +// name, which is what lets a retried handler skip work already done. +func (a *accessor) Existing(ctx context.Context, name string) (artifact.Ref, bool) { + ref, err := a.svc.FindExisting(ctx, a.owner, name) + if err != nil { + if !errors.Is(err, artifact.ErrNotFound) { + return artifact.Ref{}, false + } + + return artifact.Ref{}, false + } + + return ref, true +} diff --git a/artifact/staging/bind.go b/artifact/staging/bind.go new file mode 100644 index 0000000..1991e40 --- /dev/null +++ b/artifact/staging/bind.go @@ -0,0 +1,94 @@ +package staging + +import ( + "encoding/json" + "fmt" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/job" +) + +// Bindings maps declared input names to the artifacts bound to them. +type Bindings map[string]artifact.Ref + +// SetBindings encodes bindings onto a job. +// +// They travel in their own field rather than inside Payload because the +// payload is opaque to the engine: refs buried there would be invisible +// to the scheduler, which needs the total input size before it picks a +// worker, and to this middleware, which needs them before the handler +// runs. +func SetBindings(j *job.Job, b Bindings) error { + if len(b) == 0 { + j.ArtifactBindings = nil + + return nil + } + + raw, err := json.Marshal(b) + if err != nil { + return fmt.Errorf("dispatch/artifact/staging: encode bindings: %w", err) + } + + j.ArtifactBindings = raw + + return nil +} + +// GetBindings decodes a job's bindings. A job with none yields an empty +// map rather than an error. +func GetBindings(j *job.Job) (Bindings, error) { + if len(j.ArtifactBindings) == 0 { + return Bindings{}, nil + } + + var b Bindings + if err := json.Unmarshal(j.ArtifactBindings, &b); err != nil { + return nil, fmt.Errorf("dispatch/artifact/staging: decode bindings: %w", err) + } + + return b, nil +} + +// TotalBoundSize sums the sizes of the bound artifacts, which is the +// figure the scheduler and the staging budget both care about. +func TotalBoundSize(b Bindings) int64 { + var total int64 + + for _, ref := range b { + total += ref.Size + } + + return total +} + +// Validate checks bindings against a definition's declarations. +// +// Both directions are errors: a binding with no declaration is a +// programming mistake, and a missing required declaration means the job +// cannot run. Catching them at enqueue keeps them out of the retry path. +func Validate(specs []artifact.InputSpec, b Bindings) error { + for name, ref := range b { + spec, ok := artifact.FindInput(specs, name) + if !ok { + return fmt.Errorf("%w: %q", artifact.ErrUndeclared, name) + } + + if spec.MaxSize > 0 && ref.Size > spec.MaxSize { + return fmt.Errorf("%w: input %q is %d bytes, limit is %d", + artifact.ErrSizeExceeded, name, ref.Size, spec.MaxSize) + } + } + + for _, spec := range specs { + if !spec.Required { + continue + } + + if _, ok := b[spec.Name]; !ok { + return fmt.Errorf("%w: %q", artifact.ErrUnbound, spec.Name) + } + } + + return nil +} diff --git a/artifact/staging/doc.go b/artifact/staging/doc.go new file mode 100644 index 0000000..87326fe --- /dev/null +++ b/artifact/staging/doc.go @@ -0,0 +1,15 @@ +// Package staging wires the artifact plane into job execution. +// +// It lives apart from the artifact package to break an import cycle: the +// middleware signature takes a *job.Job, and job imports artifact for its +// input declarations. +// +// The middleware stages a job's declared inputs before the handler runs, +// puts an artifact.Accessor in the context, and releases every cache +// lease afterwards — whether the handler returns, fails, or panics. +// +// Running staging as middleware rather than inside the executor is also +// the right seam for out-of-process execution: staging happens outside +// the boundary, so a sandboxed handler receives a directory of files +// rather than storage credentials. +package staging diff --git a/artifact/staging/middleware.go b/artifact/staging/middleware.go new file mode 100644 index 0000000..e6c3591 --- /dev/null +++ b/artifact/staging/middleware.go @@ -0,0 +1,215 @@ +package staging + +import ( + "context" + "errors" + "fmt" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/middleware" +) + +// SpecLookup resolves a job name to its declared inputs. The engine +// supplies one backed by the job registry. +type SpecLookup func(jobName string) []artifact.InputSpec + +// Options configures the middleware. +type Options struct { + Logger log.Logger +} + +// Option configures the middleware. +type Option func(*Options) + +// WithLogger sets the logger. +func WithLogger(l log.Logger) Option { + return func(o *Options) { o.Logger = l } +} + +// Middleware stages a job's declared inputs, exposes them through an +// artifact.Accessor on the context, and releases every cache lease when +// the handler returns. +// +// It slots into the existing middleware chain, so nothing in the executor +// changes. That also makes it the natural boundary for out-of-process +// execution later: staging runs outside, and the handler sees files +// rather than credentials. +func Middleware( + svc *artifact.Service, + c *cache.Cache, + lookup SpecLookup, + opts ...Option, +) middleware.Middleware { + cfg := Options{Logger: log.NewNoopLogger()} + for _, opt := range opts { + opt(&cfg) + } + + return func(ctx context.Context, j *job.Job, next middleware.Handler) error { + if svc == nil || !svc.Enabled() { + return next(ctx) + } + + specs := lookup(j.Name) + + bindings, err := GetBindings(j) + if err != nil { + return err + } + + if len(specs) == 0 && len(bindings) == 0 { + // Nothing declared and nothing bound. Still hand the handler an + // accessor so it can create outputs. + return next(artifact.WithAccessor(ctx, newAccessor(svc, j, nil))) + } + + if verr := Validate(specs, bindings); verr != nil { + return verr + } + + staged, release, serr := stageInputs(ctx, svc, c, specs, bindings, cfg.Logger) + + // Release before returning whatever happens — including a panic + // unwinding through here — or a failed job would pin cache entries + // until the process restarted. + defer release() + + if serr != nil { + return serr + } + + return next(artifact.WithAccessor(ctx, newAccessor(svc, j, staged))) + } +} + +// newAccessor builds the handler-facing accessor for a job. +// +// The attempt comes from RetryCount, which is what scopes committed +// outputs to this execution and lets IfAbsent see earlier ones. +func newAccessor(svc *artifact.Service, j *job.Job, inputs map[string]staged) artifact.Accessor { + if inputs == nil { + inputs = map[string]staged{} + } + + return &accessor{ + svc: svc, + owner: artifact.OwnerRef{Kind: artifact.OwnerJob, ID: j.ID.String()}, + attempt: j.RetryCount, + inputs: inputs, + } +} + +// stageInputs materialises every declared input, returning the staged set +// and a release function that is safe to call even after a failure. +func stageInputs( + ctx context.Context, + svc *artifact.Service, + c *cache.Cache, + specs []artifact.InputSpec, + bindings Bindings, + logger log.Logger, +) (inputs map[string]staged, release func(), err error) { + out := make(map[string]staged, len(bindings)) + + var releases []func() + + release = func() { + for _, r := range releases { + r() + } + } + + for _, spec := range specs { + ref, ok := bindings[spec.Name] + if !ok { + // Absent and optional: Validate already rejected the required case. + continue + } + + if spec.Mode == artifact.StageModeLazy || c == nil { + out[spec.Name] = staged{ref: ref, mode: spec.Mode} + + continue + } + + path, hash, rel, err := c.Stage(ctx, ref) + if err != nil { + // Preserve the permanent classification so the executor fails + // the job fast rather than retrying a fetch that can never + // succeed. A deleted input and one we are not authorized to + // read are equally hopeless. + if errors.Is(err, dispatch.ErrPermanent) { + return nil, release, fmt.Errorf("stage input %q: %w", spec.Name, err) + } + + return nil, release, fmt.Errorf("dispatch/artifact/staging: stage input %q: %w", spec.Name, err) + } + + releases = append(releases, rel) + out[spec.Name] = staged{ref: ref, path: path, mode: spec.Mode} + + recordHash(ctx, svc, ref, hash, logger) + } + + return out, release, nil +} + +// recordHash persists a content hash learned during staging. +// +// Registration deliberately skips hashing to keep enqueue cheap, so this +// is where an artifact's content_hash gets filled in. It is best effort: +// failing to record it costs a future dedupe opportunity, never +// correctness, so it must not fail the job. +func recordHash( + ctx context.Context, + svc *artifact.Service, + ref artifact.Ref, + hash string, + logger log.Logger, +) { + if hash == "" || ref.ContentHash == hash || ref.ID.IsNil() { + return + } + + a, err := svc.Store().GetArtifact(ctx, ref.ID) + if err != nil { + logger.Debug("dispatch/artifact/staging: could not load artifact to record hash", + log.String("artifact_id", ref.ID.String()), + log.String("error", err.Error()), + ) + + return + } + + if a.ContentHash == hash { + return + } + + a.ContentHash = hash + + if err := svc.Store().UpdateArtifact(ctx, a); err != nil { + logger.Debug("dispatch/artifact/staging: could not record content hash", + log.String("artifact_id", ref.ID.String()), + log.String("error", err.Error()), + ) + } +} + +// LinkInputs records which artifacts a job consumed, so lineage survives +// after the job finishes. +func LinkInputs(ctx context.Context, svc *artifact.Service, j *job.Job, bindings Bindings) error { + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: j.ID.String()} + + for name, ref := range bindings { + if err := svc.Link(ctx, ref, owner, artifact.RoleInput, name, j.RetryCount); err != nil { + return err + } + } + + return nil +} diff --git a/artifact/staging/middleware_test.go b/artifact/staging/middleware_test.go new file mode 100644 index 0000000..0190ad5 --- /dev/null +++ b/artifact/staging/middleware_test.go @@ -0,0 +1,513 @@ +package staging_test + +import ( + "context" + "errors" + "io" + "os" + "testing" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/artifact/staging" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" +) + +type harness struct { + svc *artifact.Service + cache *cache.Cache + backend *artifacttest.Backend + store artifact.Store +} + +func newHarness(t *testing.T, budget int64) *harness { + t.Helper() + + b := artifacttest.NewBackend() + st := memory.New() + svc := artifact.NewService(st, b, + artifact.WithEphemeralPrefix("ephemeral"), + artifact.WithDefaultBucket("dispatch")) + + c, err := cache.New(t.TempDir(), b, cache.WithBudget(budget)) + if err != nil { + t.Fatalf("cache.New: %v", err) + } + + t.Cleanup(func() { + if cerr := c.Close(); cerr != nil { + t.Errorf("cache close: %v", cerr) + } + }) + + return &harness{svc: svc, cache: c, backend: b, store: st} +} + +func specsFor(specs ...artifact.InputSpec) staging.SpecLookup { + return func(string) []artifact.InputSpec { return specs } +} + +func newJob() *job.Job { + return &job.Job{ID: id.NewJobID(), Name: "tessellate"} +} + +func TestMiddlewareStagesDeclaredInput(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + h.backend.Put("models", "tower.ifc", []byte("ifcdata")) + + ref, err := h.svc.Register(ctx, "models", "tower.ifc") + if err != nil { + t.Fatalf("Register: %v", err) + } + + mw := staging.Middleware(h.svc, h.cache, + specsFor(artifact.Input("model", artifact.Required))) + + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{"model": ref}); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + var gotPath string + + err = mw(ctx, j, func(ctx context.Context) error { + gotPath = artifact.From(ctx).Path("model") + + return nil + }) + if err != nil { + t.Fatalf("middleware: %v", err) + } + + if gotPath == "" { + t.Fatal("handler saw no staged path for a declared input") + } + + data, err := os.ReadFile(gotPath) + if err != nil { + t.Fatalf("staged file unreadable: %v", err) + } + + if string(data) != "ifcdata" { + t.Fatalf("staged content = %q, want %q", data, "ifcdata") + } +} + +func TestMiddlewareRecordsHashLearnedDuringStaging(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + h.backend.Put("models", "hashme.ifc", []byte("bytes")) + + ref, err := h.svc.Register(ctx, "models", "hashme.ifc") + if err != nil { + t.Fatalf("Register: %v", err) + } + + if ref.ContentHash != "" { + t.Fatal("Register should not have hashed") + } + + mw := staging.Middleware(h.svc, h.cache, specsFor(artifact.Input("model"))) + + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{"model": ref}); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + if merr := mw(ctx, j, func(context.Context) error { return nil }); merr != nil { + t.Fatalf("middleware: %v", merr) + } + + a, err := h.store.GetArtifact(ctx, ref.ID) + if err != nil { + t.Fatalf("GetArtifact: %v", err) + } + + if a.ContentHash == "" { + t.Fatal("staging must record the hash it computed during the download") + } +} + +func TestMiddlewareMissingRequiredInput(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + + mw := staging.Middleware(h.svc, h.cache, + specsFor(artifact.Input("model", artifact.Required))) + + called := false + + err := mw(ctx, newJob(), func(context.Context) error { + called = true + + return nil + }) + if !errors.Is(err, artifact.ErrUnbound) { + t.Fatalf("missing required input = %v, want ErrUnbound", err) + } + + if called { + t.Fatal("handler must not run when a required input is unbound") + } +} + +func TestMiddlewareRejectsUndeclaredBinding(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + h.backend.Put("models", "x.ifc", []byte("data")) + + ref, err := h.svc.Register(ctx, "models", "x.ifc") + if err != nil { + t.Fatalf("Register: %v", err) + } + + mw := staging.Middleware(h.svc, h.cache, specsFor(artifact.Input("model"))) + + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{"surprise": ref}); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + if err := mw(ctx, j, func(context.Context) error { return nil }); !errors.Is(err, artifact.ErrUndeclared) { + t.Fatalf("undeclared binding = %v, want ErrUndeclared", err) + } +} + +func TestMiddlewareRejectsOversizeBinding(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + h.backend.Put("models", "big.ifc", make([]byte, 100)) + + ref, err := h.svc.Register(ctx, "models", "big.ifc") + if err != nil { + t.Fatalf("Register: %v", err) + } + + mw := staging.Middleware(h.svc, h.cache, + specsFor(artifact.Input("model", artifact.MaxSize(10)))) + + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{"model": ref}); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + if err := mw(ctx, j, func(context.Context) error { return nil }); !errors.Is(err, artifact.ErrSizeExceeded) { + t.Fatalf("oversize binding = %v, want ErrSizeExceeded", err) + } +} + +// TestMiddlewareDeletedInputFailsFast pins the retry-classification +// behaviour: a job whose input no longer exists must fail permanently +// rather than burn its retry budget on a fetch that can never succeed. +func TestMiddlewareDeletedInputFailsFast(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + + mw := staging.Middleware(h.svc, h.cache, + specsFor(artifact.Input("model", artifact.Required))) + + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{ + "model": {ID: id.NewArtifactID(), Bucket: "models", Key: "gone.ifc"}, + }); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + err := mw(ctx, j, func(context.Context) error { return nil }) + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("staging a deleted input = %v, want ErrNotFound", err) + } +} + +// TestMiddlewareDeniedInput is the deleted-input test's sibling. The input +// exists but cannot be read, which is just as permanent. +// +// Before the backend classified permission failures this error reached the +// executor carrying no artifact sentinel at all. It now arrives classified. +// Acting on it is still the executor's to do: handleFailure counts retries +// without consulting the error, so neither this nor a deleted input fails +// fast yet. +func TestMiddlewareDeniedInput(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + h.backend.Put("models", "secret.ifc", []byte("classified")) + h.backend.DenyOpen = true + + mw := staging.Middleware(h.svc, h.cache, + specsFor(artifact.Input("model", artifact.Required))) + + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{ + "model": {ID: id.NewArtifactID(), Bucket: "models", Key: "secret.ifc"}, + }); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + err := mw(ctx, j, func(context.Context) error { return nil }) + if !errors.Is(err, dispatch.ErrPermanent) { + t.Fatalf("staging a forbidden input = %v, want ErrPermanent", err) + } + + if !errors.Is(err, artifact.ErrPermissionDenied) { + t.Fatalf("staging a forbidden input = %v, want ErrPermissionDenied", err) + } +} + +// TestMiddlewareReleasesLeasesOnHandlerError would deadlock on the third +// run if a failed job leaked its cache lease. +func TestMiddlewareReleasesLeasesOnHandlerError(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 10) + h.backend.Put("m", "a", []byte("0123456789")) + + ref, err := h.svc.Register(ctx, "m", "a") + if err != nil { + t.Fatalf("Register: %v", err) + } + + mw := staging.Middleware(h.svc, h.cache, specsFor(artifact.Input("in"))) + + handlerErr := errors.New("boom") + + for i := range 3 { + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{"in": ref}); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + if rerr := mw(ctx, j, func(context.Context) error { return handlerErr }); !errors.Is(rerr, handlerErr) { + t.Fatalf("run %d: middleware returned %v, want the handler error", i, rerr) + } + } +} + +func TestMiddlewareReleasesLeasesOnPanic(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 10) + h.backend.Put("m", "a", []byte("0123456789")) + + ref, err := h.svc.Register(ctx, "m", "a") + if err != nil { + t.Fatalf("Register: %v", err) + } + + mw := staging.Middleware(h.svc, h.cache, specsFor(artifact.Input("in"))) + + run := func() { + defer func() { _ = recover() }() + + j := newJob() + _ = staging.SetBindings(j, staging.Bindings{"in": ref}) + + _ = mw(ctx, j, func(context.Context) error { panic("handler exploded") }) + } + + for range 3 { + run() + } + + // A leaked lease would have exhausted the 10-byte budget by now. + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{"in": ref}); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + if rerr := mw(ctx, j, func(context.Context) error { return nil }); rerr != nil { + t.Fatalf("staging after panics: %v — leases leaked", rerr) + } +} + +func TestAccessorCreateLinksToJobAndAttempt(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + + mw := staging.Middleware(h.svc, h.cache, specsFor()) + + j := newJob() + j.RetryCount = 2 + + err := mw(ctx, j, func(ctx context.Context) error { + w, cerr := artifact.From(ctx).Create(ctx, "page-1.png") + if cerr != nil { + return cerr + } + + if _, werr := io.WriteString(w, "pixels"); werr != nil { + return werr + } + + _, cerr = w.Commit(ctx) + + return cerr + }) + if err != nil { + t.Fatalf("middleware: %v", err) + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: j.ID.String()} + + links, err := h.store.ListLinks(ctx, owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + if len(links) != 1 { + t.Fatalf("got %d links, want 1", len(links)) + } + + if links[0].Attempt != 2 { + t.Fatalf("link attempt = %d, want 2 (from job.RetryCount)", links[0].Attempt) + } + + if links[0].Role != artifact.RoleOutput { + t.Fatalf("link role = %q, want output", links[0].Role) + } +} + +func TestAccessorOpenReadsStagedFile(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + h.backend.Put("models", "read.ifc", []byte("streamed")) + + ref, err := h.svc.Register(ctx, "models", "read.ifc") + if err != nil { + t.Fatalf("Register: %v", err) + } + + mw := staging.Middleware(h.svc, h.cache, specsFor(artifact.Input("model"))) + + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{"model": ref}); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + err = mw(ctx, j, func(ctx context.Context) error { + rc, oerr := artifact.From(ctx).Open(ctx, "model") + if oerr != nil { + return oerr + } + + defer func() { _ = rc.Close() }() + + data, rerr := io.ReadAll(rc) + if rerr != nil { + return rerr + } + + if string(data) != "streamed" { + t.Fatalf("Open returned %q, want %q", data, "streamed") + } + + return nil + }) + if err != nil { + t.Fatalf("middleware: %v", err) + } + + // A path-staged input is read from disk, so Open must not re-download. + if h.backend.Opens() != 1 { + t.Fatalf("Opens() = %d, want 1 — Open on a staged input must read local disk", h.backend.Opens()) + } +} + +func TestLazyInputIsNotDownloadedUpFront(t *testing.T) { + ctx := context.Background() + h := newHarness(t, 1<<20) + h.backend.Put("models", "lazy.csv", []byte("a,b,c")) + + ref, err := h.svc.Register(ctx, "models", "lazy.csv") + if err != nil { + t.Fatalf("Register: %v", err) + } + + mw := staging.Middleware(h.svc, h.cache, + specsFor(artifact.Input("data", artifact.StageLazy))) + + j := newJob() + if serr := staging.SetBindings(j, staging.Bindings{"data": ref}); serr != nil { + t.Fatalf("SetBindings: %v", serr) + } + + err = mw(ctx, j, func(ctx context.Context) error { + if path := artifact.From(ctx).Path("data"); path != "" { + t.Fatalf("lazy input has a path %q, want none", path) + } + + if h.backend.Opens() != 0 { + t.Fatalf("Opens() = %d before Open — lazy inputs must not pre-download", h.backend.Opens()) + } + + rc, oerr := artifact.From(ctx).Open(ctx, "data") + if oerr != nil { + return oerr + } + + return rc.Close() + }) + if err != nil { + t.Fatalf("middleware: %v", err) + } + + if h.backend.Opens() != 1 { + t.Fatalf("Opens() = %d, want 1", h.backend.Opens()) + } +} + +func TestMiddlewareNoOpWhenServiceDisabled(t *testing.T) { + ctx := context.Background() + svc := artifact.NewService(memory.New(), nil) + + mw := staging.Middleware(svc, nil, specsFor(artifact.Input("model", artifact.Required))) + + called := false + + // With no backend the plane is off, so even a required declaration + // must not block execution: Dispatch behaves as it did before. + if err := mw(ctx, newJob(), func(context.Context) error { called = true; return nil }); err != nil { + t.Fatalf("disabled service must be a pass-through, got %v", err) + } + + if !called { + t.Fatal("handler did not run with the artifact plane disabled") + } +} + +func TestBindingsRoundTrip(t *testing.T) { + j := newJob() + want := staging.Bindings{ + "model": {ID: id.NewArtifactID(), Bucket: "models", Key: "a.ifc", Size: 42}, + } + + if err := staging.SetBindings(j, want); err != nil { + t.Fatalf("SetBindings: %v", err) + } + + got, err := staging.GetBindings(j) + if err != nil { + t.Fatalf("GetBindings: %v", err) + } + + if got["model"].Key != "a.ifc" || got["model"].Size != 42 { + t.Fatalf("round trip = %+v, want %+v", got, want) + } + + if total := staging.TotalBoundSize(got); total != 42 { + t.Fatalf("TotalBoundSize = %d, want 42", total) + } +} + +func TestGetBindingsEmptyJob(t *testing.T) { + got, err := staging.GetBindings(newJob()) + if err != nil { + t.Fatalf("GetBindings on a job with no bindings: %v", err) + } + + if len(got) != 0 { + t.Fatalf("got %d bindings, want 0", len(got)) + } +} diff --git a/artifact/store.go b/artifact/store.go new file mode 100644 index 0000000..f5fa97b --- /dev/null +++ b/artifact/store.go @@ -0,0 +1,124 @@ +package artifact + +import ( + "context" + "time" + + "github.com/xraph/dispatch/id" +) + +// OwnerRef identifies a link owner. +type OwnerRef struct { + Kind OwnerKind + ID string +} + +// Valid reports whether the owner reference is usable. +func (o OwnerRef) Valid() bool { return o.Kind.Valid() && o.ID != "" } + +// ListOpts controls pagination and filtering for artifact list queries. +type ListOpts struct { + // Limit is the maximum number of artifacts to return. Zero means no limit. + Limit int + // Offset is the number of artifacts to skip. + Offset int + // Lifecycle filters by lifecycle. Empty means all. + Lifecycle Lifecycle + // ScopeAppID filters by tenant application. Empty means all. + ScopeAppID string + // ScopeOrgID filters by tenant organization. Empty means all. + ScopeOrgID string + // IncludeDeleted includes soft-deleted artifacts. Default false. + IncludeDeleted bool +} + +// SweepOpts controls a lifecycle sweep. +type SweepOpts struct { + // Retention is the grace period after the last owner reaches a + // terminal state before an artifact becomes eligible. + Retention time.Duration + // Limit caps how many artifacts a single sweep call may mark. + // Zero means no limit. + Limit int + // DryRun computes eligibility and returns the artifacts that would be + // marked without modifying anything. + DryRun bool +} + +// Store defines the persistence contract for artifacts and their links. +// +// Implementations must guarantee that CreateArtifact inserts the artifact +// and its link atomically, so a zero-link artifact can only result from a +// partial failure and never from a normal race. +// +// SweepEphemeral and SweepOrphans must constrain themselves to +// Lifecycle == Ephemeral using a literal, never a value threaded through +// from a caller. Durable artifacts must be unreachable from both. +type Store interface { + // CreateArtifact inserts an artifact and, when link is non-nil, its + // first link atomically. Returns ErrExists if an artifact already + // exists at the same backend, bucket, and key. + CreateArtifact(ctx context.Context, a *Artifact, link *Link) error + + // GetArtifact retrieves an artifact by ID. Returns ErrNotFound if it + // does not exist or has been soft-deleted. + GetArtifact(ctx context.Context, artifactID id.ArtifactID) (*Artifact, error) + + // FindArtifactByKey retrieves an artifact by its storage coordinates. + // Returns ErrNotFound if none exists. + FindArtifactByKey(ctx context.Context, backend, bucket, key string) (*Artifact, error) + + // UpdateArtifact persists changes to size, content hash, content type, + // and expiry. It must not permit changing lifecycle. + UpdateArtifact(ctx context.Context, a *Artifact) error + + // ListArtifacts returns artifacts matching the given options. + ListArtifacts(ctx context.Context, opts ListOpts) ([]*Artifact, error) + + // LinkArtifact records that an owner references an artifact. Linking + // the same artifact, owner, name, and attempt twice is a no-op rather + // than an error. + LinkArtifact(ctx context.Context, link *Link) error + + // ListLinks returns every link belonging to the given owner. + ListLinks(ctx context.Context, owner OwnerRef) ([]*Link, error) + + // FindLinkByName returns the link for an owner and name with the + // highest attempt number, breaking ties by CreatedAt descending so + // resolution is deterministic and favours the later writer. This is + // what IfAbsent uses to detect that a prior attempt already produced + // an output. Returns ErrNotFound if no attempt has produced it. + // + // Ties happen: track D's lease reclaim increments EvictCount, never + // RetryCount, so a reclaimed (zombie) holder and the worker it was + // fenced out for can both commit a link at the same (OwnerKind, + // OwnerID, Name, Attempt) — CreateFenced closes the storage-key + // collision between them, but nothing stops two Link rows with + // different ArtifactIDs at the identical tuple. The CreatedAt + // tie-break is the cheap half of the fix: it makes which of the two + // wins deterministic instead of backend-dependent map/query order. + // The complete fix — carrying the fence token on Link itself, so the + // zombie's write is rejected rather than merely outrun — is a schema + // change across all five backends and remains open. + FindLinkByName(ctx context.Context, owner OwnerRef, name string) (*Link, error) + + // ListArtifactsByOwner returns the artifacts linked to an owner, + // optionally filtered by role. An empty role returns all. + ListArtifactsByOwner(ctx context.Context, owner OwnerRef, role Role) ([]*Artifact, error) + + // SweepEphemeral marks eligible ephemeral artifacts as deleted and + // returns them. + SweepEphemeral(ctx context.Context, opts SweepOpts) ([]*Artifact, error) + + // SweepOrphans marks ephemeral artifacts that have no links at all and + // were created before the cutoff. + SweepOrphans(ctx context.Context, cutoff time.Time, limit int) ([]*Artifact, error) + + // ListPurgeable returns soft-deleted artifacts whose deletion is older + // than grace, so their bytes may be removed from the backend. + ListPurgeable(ctx context.Context, grace time.Duration, limit int) ([]*Artifact, error) + + // PurgeArtifact hard-deletes an artifact row and its links after the + // bytes have been removed from the backend. + PurgeArtifact(ctx context.Context, artifactID id.ArtifactID) error +} diff --git a/artifact/sweeper/doc.go b/artifact/sweeper/doc.go new file mode 100644 index 0000000..a8d7f0d --- /dev/null +++ b/artifact/sweeper/doc.go @@ -0,0 +1,18 @@ +// Package sweeper reclaims the storage Dispatch owns. +// +// It touches only ephemeral artifacts — the ones Dispatch itself created +// on a handler's behalf. Durable artifacts, which are the application's +// uploads merely tracked here, are unreachable from every code path in +// this package. That guarantee is enforced twice: the store's sweep +// queries constrain themselves to ephemeral with a literal, and the +// sweeper re-checks each artifact before acting on it. +// +// Deletion is two-phase. A sweep marks an artifact deleted and stops +// serving it; a later purge removes the bytes once a grace period has +// passed. A mistaken sweep is therefore observable and reversible for the +// length of that window rather than instantly destructive, and both +// phases are idempotent under retry. +// +// Sweeping runs on the elected leader only, so a fleet of workers does +// not race to delete the same objects. +package sweeper diff --git a/artifact/sweeper/sweeper.go b/artifact/sweeper/sweeper.go new file mode 100644 index 0000000..7934af5 --- /dev/null +++ b/artifact/sweeper/sweeper.go @@ -0,0 +1,385 @@ +package sweeper + +import ( + "context" + "errors" + "fmt" + "sync" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/artifact" +) + +// Defaults for a sweeper built without options. +const ( + DefaultRetention = 168 * time.Hour + DefaultPurgeGrace = 24 * time.Hour + DefaultInterval = 15 * time.Minute + DefaultBatchSize = 500 +) + +// Result reports what one pass did. +type Result struct { + // Swept is how many artifacts were marked deleted. + Swept int + // Purged is how many had their bytes removed. + Purged int + // BytesReclaimed counts the bytes freed by purging. + BytesReclaimed int64 + // Skipped counts artifacts a pass declined to act on, which is where + // a backend failure shows up: the next pass retries them. + Skipped int +} + +// Observer is notified for each artifact a pass acts on, so lifecycle +// events reach the extension registry without this package depending on +// it. +type Observer interface { + ArtifactSwept(ctx context.Context, a *artifact.Artifact) + ArtifactPurged(ctx context.Context, a *artifact.Artifact) +} + +// Sweeper reclaims ephemeral artifacts whose owners have finished. +type Sweeper struct { + store artifact.Store + backend artifact.Backend + logger log.Logger + + retention time.Duration + purgeGrace time.Duration + interval time.Duration + batchSize int + dryRun bool + enabled bool + isLeader func() bool + observer Observer + stopCh chan struct{} + stopOnce sync.Once + wg sync.WaitGroup + startedOnce sync.Once +} + +// Option configures a Sweeper. +type Option func(*Sweeper) + +// WithRetention sets how long an ephemeral artifact survives after its +// last owner reaches a terminal state. +func WithRetention(d time.Duration) Option { + return func(s *Sweeper) { s.retention = d } +} + +// WithPurgeGrace sets how long a soft-deleted artifact's bytes survive. +// This is the window in which a mistaken sweep can still be caught. +func WithPurgeGrace(d time.Duration) Option { + return func(s *Sweeper) { s.purgeGrace = d } +} + +// WithInterval sets how often the background loop runs. +func WithInterval(d time.Duration) Option { + return func(s *Sweeper) { s.interval = d } +} + +// WithBatchSize caps how many artifacts one pass may touch. +func WithBatchSize(n int) Option { + return func(s *Sweeper) { s.batchSize = n } +} + +// WithDryRun reports what would be swept without changing anything. +func WithDryRun(dry bool) Option { + return func(s *Sweeper) { s.dryRun = dry } +} + +// WithEnabled is the kill switch. A disabled sweeper's passes are no-ops, +// so reclamation can be stopped without redeploying. +func WithEnabled(enabled bool) Option { + return func(s *Sweeper) { s.enabled = enabled } +} + +// WithLeaderCheck restricts sweeping to the elected leader. Without one, +// every worker in a fleet would sweep concurrently. +func WithLeaderCheck(fn func() bool) Option { + return func(s *Sweeper) { s.isLeader = fn } +} + +// WithObserver receives a callback per artifact acted on. +func WithObserver(o Observer) Option { + return func(s *Sweeper) { s.observer = o } +} + +// WithLogger sets the logger. +func WithLogger(l log.Logger) Option { + return func(s *Sweeper) { s.logger = l } +} + +// New creates a Sweeper. It is enabled by default; use WithEnabled(false) +// to build one that does nothing until switched on. +func New(store artifact.Store, backend artifact.Backend, opts ...Option) *Sweeper { + s := &Sweeper{ + store: store, + backend: backend, + logger: log.NewNoopLogger(), + retention: DefaultRetention, + purgeGrace: DefaultPurgeGrace, + interval: DefaultInterval, + batchSize: DefaultBatchSize, + enabled: true, + stopCh: make(chan struct{}), + } + + for _, opt := range opts { + opt(s) + } + + return s +} + +// active reports whether this instance should act right now. +func (s *Sweeper) active() bool { + if !s.enabled { + return false + } + + if s.isLeader != nil && !s.isLeader() { + return false + } + + return true +} + +// SweepOnce marks eligible ephemeral artifacts deleted. +// +// Two passes run: artifacts whose owners have all finished and whose +// retention has elapsed, then orphans — artifacts with no links at all, +// which can only result from a partial failure during creation. +func (s *Sweeper) SweepOnce(ctx context.Context) (Result, error) { + var res Result + + if !s.active() { + return res, nil + } + + eligible, err := s.store.SweepEphemeral(ctx, artifact.SweepOpts{ + Retention: s.retention, + Limit: s.batchSize, + DryRun: s.dryRun, + }) + if err != nil { + return res, fmt.Errorf("dispatch/artifact/sweeper: sweep ephemeral: %w", err) + } + + res.Swept += s.report(ctx, eligible, "retention") + + cutoff := time.Now().UTC().Add(-s.orphanGrace()) + + if !s.dryRun { + orphans, oerr := s.store.SweepOrphans(ctx, cutoff, s.batchSize) + if oerr != nil { + return res, fmt.Errorf("dispatch/artifact/sweeper: sweep orphans: %w", oerr) + } + + res.Swept += s.report(ctx, orphans, "orphan") + } + + return res, nil +} + +// orphanGrace is how long a link-less artifact is tolerated. It is +// deliberately generous: an artifact is linked in the same operation that +// creates it, so a zero-link artifact means a crash mid-create, and a +// short window risks sweeping one that is merely mid-flight. +func (s *Sweeper) orphanGrace() time.Duration { + if s.purgeGrace > 0 { + return s.purgeGrace + } + + return DefaultPurgeGrace +} + +// report emits observer callbacks and counts artifacts, refusing to act +// on anything that is not ephemeral. +func (s *Sweeper) report(ctx context.Context, artifacts []*artifact.Artifact, reason string) int { + n := 0 + + for _, a := range artifacts { + // The store already constrains its sweep queries to ephemeral. + // Re-checking here means a bug in one backend's query cannot turn + // into deleted customer data. + if a.Lifecycle != artifact.Ephemeral { + s.logger.Error("dispatch/artifact/sweeper: refusing to sweep a non-ephemeral artifact", + log.String("artifact_id", a.ID.String()), + log.String("lifecycle", string(a.Lifecycle)), + ) + + continue + } + + n++ + + if s.observer != nil && !s.dryRun { + s.observer.ArtifactSwept(ctx, a) + } + + s.logger.Debug("dispatch/artifact/sweeper: swept artifact", + log.String("artifact_id", a.ID.String()), + log.String("reason", reason), + log.Int64("bytes", a.Size), + ) + } + + return n +} + +// PurgeOnce removes the bytes of artifacts soft-deleted longer ago than +// the grace period, then deletes their rows. +// +// A backend failure skips that artifact rather than aborting the pass, so +// one unreachable object cannot stall reclamation of everything else. The +// next pass retries it. +func (s *Sweeper) PurgeOnce(ctx context.Context) (Result, error) { + var res Result + + if !s.active() || s.dryRun { + return res, nil + } + + purgeable, err := s.store.ListPurgeable(ctx, s.purgeGrace, s.batchSize) + if err != nil { + return res, fmt.Errorf("dispatch/artifact/sweeper: list purgeable: %w", err) + } + + for _, a := range purgeable { + if a.Lifecycle != artifact.Ephemeral { + s.logger.Error("dispatch/artifact/sweeper: refusing to purge a non-ephemeral artifact", + log.String("artifact_id", a.ID.String()), + log.String("lifecycle", string(a.Lifecycle)), + ) + + res.Skipped++ + + continue + } + + if derr := s.backend.Delete(ctx, a.Ref()); derr != nil { + s.logger.Warn("dispatch/artifact/sweeper: could not delete object; will retry", + log.String("artifact_id", a.ID.String()), + log.String("error", derr.Error()), + ) + + res.Skipped++ + + continue + } + + if perr := s.store.PurgeArtifact(ctx, a.ID); perr != nil { + // The bytes are gone but the row remains. The next pass finds + // it again and the backend's delete-missing-is-not-an-error + // contract makes the retry safe. + s.logger.Warn("dispatch/artifact/sweeper: could not purge row; will retry", + log.String("artifact_id", a.ID.String()), + log.String("error", perr.Error()), + ) + + res.Skipped++ + + continue + } + + res.Purged++ + res.BytesReclaimed += a.Size + + if s.observer != nil { + s.observer.ArtifactPurged(ctx, a) + } + } + + return res, nil +} + +// RunOnce performs a sweep followed by a purge. +func (s *Sweeper) RunOnce(ctx context.Context) (Result, error) { + swept, err := s.SweepOnce(ctx) + if err != nil { + return swept, err + } + + purged, err := s.PurgeOnce(ctx) + if err != nil { + return swept, err + } + + return Result{ + Swept: swept.Swept, + Purged: purged.Purged, + BytesReclaimed: purged.BytesReclaimed, + Skipped: swept.Skipped + purged.Skipped, + }, nil +} + +// Start begins the background loop. It is safe to call once. +func (s *Sweeper) Start(ctx context.Context) error { + s.startedOnce.Do(func() { + s.wg.Add(1) + + go s.loop(ctx) + }) + + return nil +} + +func (s *Sweeper) loop(ctx context.Context) { + defer s.wg.Done() + + ticker := time.NewTicker(s.interval) + defer ticker.Stop() + + for { + select { + case <-s.stopCh: + return + case <-ctx.Done(): + return + case <-ticker.C: + res, err := s.RunOnce(ctx) + if err != nil { + if errors.Is(err, context.Canceled) { + return + } + + s.logger.Error("dispatch/artifact/sweeper: pass failed", + log.String("error", err.Error())) + + continue + } + + if res.Swept > 0 || res.Purged > 0 { + s.logger.Info("dispatch/artifact/sweeper: reclaimed storage", + log.Int("swept", res.Swept), + log.Int("purged", res.Purged), + log.Int64("bytes_reclaimed", res.BytesReclaimed), + log.Int("skipped", res.Skipped), + ) + } + } + } +} + +// Stop halts the background loop and waits for it to finish. +func (s *Sweeper) Stop(ctx context.Context) error { + s.stopOnce.Do(func() { close(s.stopCh) }) + + done := make(chan struct{}) + + go func() { + s.wg.Wait() + close(done) + }() + + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} diff --git a/artifact/sweeper/sweeper_test.go b/artifact/sweeper/sweeper_test.go new file mode 100644 index 0000000..d80920b --- /dev/null +++ b/artifact/sweeper/sweeper_test.go @@ -0,0 +1,418 @@ +package sweeper_test + +import ( + "context" + "errors" + "fmt" + "math/rand" + "testing" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/artifact/sweeper" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/store/memory" +) + +type rig struct { + store artifact.Store + backend *artifacttest.Backend + svc *artifact.Service +} + +func newRig(t *testing.T) *rig { + t.Helper() + + st := memory.New() + b := artifacttest.NewBackend() + svc := artifact.NewService(st, b, + artifact.WithEphemeralPrefix("ephemeral"), + artifact.WithDefaultBucket("dispatch")) + + return &rig{store: st, backend: b, svc: svc} +} + +// registerDurable creates an application-owned artifact — the kind the +// sweeper must never touch. +func (r *rig) registerDurable(t *testing.T, key string, data []byte) artifact.Ref { + t.Helper() + + r.backend.Put("customer", key, data) + + ref, err := r.svc.Register(context.Background(), "customer", key) + if err != nil { + t.Fatalf("Register %q: %v", key, err) + } + + return ref +} + +// createEphemeral creates a Dispatch-owned artifact aged into the past. +func (r *rig) createEphemeral(t *testing.T, owner artifact.OwnerRef, name string, age time.Duration) artifact.Ref { + t.Helper() + + ctx := context.Background() + + w, err := r.svc.Create(ctx, owner, 0, name) + if err != nil { + t.Fatalf("Create %q: %v", name, err) + } + + if _, werr := w.Write([]byte("payload")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + ref, err := w.Commit(ctx) + if err != nil { + t.Fatalf("Commit: %v", err) + } + + if age > 0 { + a, gerr := r.store.GetArtifact(ctx, ref.ID) + if gerr != nil { + t.Fatalf("GetArtifact: %v", gerr) + } + + a.CreatedAt = time.Now().UTC().Add(-age) + + if uerr := r.store.UpdateArtifact(ctx, a); uerr != nil { + t.Fatalf("UpdateArtifact: %v", uerr) + } + } + + return ref +} + +func newOwner() artifact.OwnerRef { + return artifact.OwnerRef{Kind: artifact.OwnerJob, ID: id.NewJobID().String()} +} + +// TestSweeperNeverTouchesDurable is the safety property the whole design +// rests on. It generates arbitrary sequences of create, fail, retry, and +// sweep, then asserts every durable artifact is still retrievable and its +// bytes still readable. +func TestSweeperNeverTouchesDurable(t *testing.T) { + ctx := context.Background() + r := newRig(t) + + // A fixed seed so a failure reproduces exactly. + rng := rand.New(rand.NewSource(1)) + + durable := make([]artifact.Ref, 0, 8) + for i := range 8 { + durable = append(durable, r.registerDurable(t, + fmt.Sprintf("upload-%d.ifc", i), []byte(fmt.Sprintf("customer data %d", i)))) + } + + s := sweeper.New(r.store, r.backend, + sweeper.WithRetention(0), + sweeper.WithPurgeGrace(0), + sweeper.WithBatchSize(100)) + + for range 60 { + switch rng.Intn(5) { + case 0: + owner := newOwner() + r.createEphemeral(t, owner, fmt.Sprintf("out-%d.bin", rng.Intn(1000)), + time.Duration(rng.Intn(72))*time.Hour) + + case 1: + // A retried job producing the same name at a later attempt. + owner := newOwner() + for attempt := range 2 { + w, err := r.svc.Create(ctx, owner, attempt, "page.png") + if err != nil { + t.Fatalf("Create attempt %d: %v", attempt, err) + } + + if _, werr := w.Write([]byte("pixels")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + if _, cerr := w.Commit(ctx); cerr != nil { + t.Fatalf("Commit: %v", cerr) + } + } + + case 2: + // An aborted write leaves nothing behind. + w, err := r.svc.Create(ctx, newOwner(), 0, "aborted.bin") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if aerr := w.Abort(); aerr != nil { + t.Fatalf("Abort: %v", aerr) + } + + case 3: + if _, err := s.SweepOnce(ctx); err != nil { + t.Fatalf("SweepOnce: %v", err) + } + + case 4: + if _, err := s.PurgeOnce(ctx); err != nil { + t.Fatalf("PurgeOnce: %v", err) + } + } + } + + // Hammer it once more, with every window wide open. + for range 3 { + if _, err := s.RunOnce(ctx); err != nil { + t.Fatalf("RunOnce: %v", err) + } + } + + for i, ref := range durable { + a, err := r.store.GetArtifact(ctx, ref.ID) + if err != nil { + t.Fatalf("durable artifact %d (%v) was destroyed: %v", i, ref.ID, err) + } + + if a.IsDeleted() { + t.Fatalf("durable artifact %d was soft-deleted", i) + } + + if !r.backend.Has(ref.Bucket, ref.Key) { + t.Fatalf("durable artifact %d had its bytes purged from the backend", i) + } + } +} + +func TestSweeperTwoPhaseDeletion(t *testing.T) { + ctx := context.Background() + r := newRig(t) + + ref := r.createEphemeral(t, newOwner(), "temp.bin", 48*time.Hour) + + // Phase one: mark deleted. The bytes must survive. + s := sweeper.New(r.store, r.backend, + sweeper.WithRetention(0), + sweeper.WithPurgeGrace(time.Hour)) + + res, err := s.SweepOnce(ctx) + if err != nil { + t.Fatalf("SweepOnce: %v", err) + } + + if res.Swept != 1 { + t.Fatalf("Swept = %d, want 1", res.Swept) + } + + if !r.backend.Has(ref.Bucket, ref.Key) { + t.Fatal("sweep removed the bytes — phase one must only mark") + } + + // With a grace window longer than the deletion's age, purging is a + // no-op: this is the window in which a mistake can be caught. + purged, err := s.PurgeOnce(ctx) + if err != nil { + t.Fatalf("PurgeOnce: %v", err) + } + + if purged.Purged != 0 { + t.Fatalf("Purged = %d during the grace window, want 0", purged.Purged) + } + + // Phase two, grace elapsed. + s2 := sweeper.New(r.store, r.backend, sweeper.WithPurgeGrace(0)) + + purged, err = s2.PurgeOnce(ctx) + if err != nil { + t.Fatalf("PurgeOnce after grace: %v", err) + } + + if purged.Purged != 1 { + t.Fatalf("Purged = %d, want 1", purged.Purged) + } + + if r.backend.Has(ref.Bucket, ref.Key) { + t.Fatal("purge did not remove the bytes") + } + + if _, err := r.store.GetArtifact(ctx, ref.ID); !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("GetArtifact after purge = %v, want ErrNotFound", err) + } +} + +func TestSweeperDryRunChangesNothing(t *testing.T) { + ctx := context.Background() + r := newRig(t) + + ref := r.createEphemeral(t, newOwner(), "temp.bin", 48*time.Hour) + + s := sweeper.New(r.store, r.backend, + sweeper.WithRetention(0), + sweeper.WithPurgeGrace(0), + sweeper.WithDryRun(true)) + + res, err := s.RunOnce(ctx) + if err != nil { + t.Fatalf("RunOnce: %v", err) + } + + if res.Purged != 0 { + t.Fatalf("dry run purged %d artifacts", res.Purged) + } + + a, err := r.store.GetArtifact(ctx, ref.ID) + if err != nil { + t.Fatalf("dry run destroyed the artifact: %v", err) + } + + if a.IsDeleted() { + t.Fatal("dry run soft-deleted the artifact") + } +} + +func TestSweeperKillSwitch(t *testing.T) { + ctx := context.Background() + r := newRig(t) + + ref := r.createEphemeral(t, newOwner(), "temp.bin", 48*time.Hour) + + s := sweeper.New(r.store, r.backend, + sweeper.WithRetention(0), + sweeper.WithPurgeGrace(0), + sweeper.WithEnabled(false)) + + res, err := s.RunOnce(ctx) + if err != nil { + t.Fatalf("RunOnce: %v", err) + } + + if res.Swept != 0 || res.Purged != 0 { + t.Fatalf("a disabled sweeper acted: %+v", res) + } + + if _, err := r.store.GetArtifact(ctx, ref.ID); err != nil { + t.Fatalf("disabled sweeper destroyed the artifact: %v", err) + } +} + +func TestSweeperLeaderOnly(t *testing.T) { + ctx := context.Background() + r := newRig(t) + + r.createEphemeral(t, newOwner(), "temp.bin", 48*time.Hour) + + leader := false + + s := sweeper.New(r.store, r.backend, + sweeper.WithRetention(0), + sweeper.WithPurgeGrace(0), + sweeper.WithLeaderCheck(func() bool { return leader })) + + res, err := s.RunOnce(ctx) + if err != nil { + t.Fatalf("RunOnce as follower: %v", err) + } + + if res.Swept != 0 { + t.Fatalf("a follower swept %d artifacts", res.Swept) + } + + leader = true + + res, err = s.RunOnce(ctx) + if err != nil { + t.Fatalf("RunOnce as leader: %v", err) + } + + if res.Swept != 1 { + t.Fatalf("the leader swept %d artifacts, want 1", res.Swept) + } +} + +// TestPurgeSkipsOnBackendFailureAndRetries checks that one unreachable +// object cannot stall reclamation of everything else. +func TestPurgeSkipsOnBackendFailureAndRetries(t *testing.T) { + ctx := context.Background() + r := newRig(t) + + ref := r.createEphemeral(t, newOwner(), "temp.bin", 48*time.Hour) + + s := sweeper.New(r.store, r.backend, + sweeper.WithRetention(0), + sweeper.WithPurgeGrace(0)) + + if _, err := s.SweepOnce(ctx); err != nil { + t.Fatalf("SweepOnce: %v", err) + } + + // Remove the bytes behind the sweeper's back. Delete-missing must not + // be an error, so the purge still completes. + if derr := r.backend.Delete(ctx, ref); derr != nil { + t.Fatalf("Delete: %v", derr) + } + + res, err := s.PurgeOnce(ctx) + if err != nil { + t.Fatalf("PurgeOnce: %v", err) + } + + if res.Purged != 1 { + t.Fatalf("Purged = %d, want 1 — a missing object must not block the purge", res.Purged) + } +} + +type countingObserver struct { + swept int + purged int +} + +func (o *countingObserver) ArtifactSwept(context.Context, *artifact.Artifact) { o.swept++ } +func (o *countingObserver) ArtifactPurged(context.Context, *artifact.Artifact) { o.purged++ } + +func TestSweeperNotifiesObserver(t *testing.T) { + ctx := context.Background() + r := newRig(t) + + r.createEphemeral(t, newOwner(), "temp.bin", 48*time.Hour) + + obs := &countingObserver{} + + s := sweeper.New(r.store, r.backend, + sweeper.WithRetention(0), + sweeper.WithPurgeGrace(0), + sweeper.WithObserver(obs)) + + if _, err := s.RunOnce(ctx); err != nil { + t.Fatalf("RunOnce: %v", err) + } + + if obs.swept != 1 { + t.Fatalf("observer saw %d sweeps, want 1", obs.swept) + } + + if obs.purged != 1 { + t.Fatalf("observer saw %d purges, want 1", obs.purged) + } +} + +func TestSweeperStartStop(t *testing.T) { + ctx := context.Background() + r := newRig(t) + + s := sweeper.New(r.store, r.backend, sweeper.WithInterval(10*time.Millisecond)) + + if err := s.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + + time.Sleep(30 * time.Millisecond) + + stopCtx, cancel := context.WithTimeout(ctx, time.Second) + defer cancel() + + if err := s.Stop(stopCtx); err != nil { + t.Fatalf("Stop: %v", err) + } + + // Stop must be idempotent so a double shutdown does not panic. + if err := s.Stop(stopCtx); err != nil { + t.Fatalf("second Stop: %v", err) + } +} diff --git a/artifact/trove/backend.go b/artifact/trove/backend.go new file mode 100644 index 0000000..71ed4fb --- /dev/null +++ b/artifact/trove/backend.go @@ -0,0 +1,286 @@ +package trove + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + "time" + + trovelib "github.com/xraph/trove" + trovedriver "github.com/xraph/trove/driver" + + "github.com/xraph/dispatch/artifact" +) + +// DefaultName is the backend identifier used when none is configured. +const DefaultName = "trove" + +// Backend adapts a *trove.Trove as an artifact.Backend. +type Backend struct { + trove *trovelib.Trove + name string +} + +// Compile-time checks. RangeReader and Presigner are advertised +// unconditionally; whether the underlying driver actually supports them +// is resolved per call, because a Trove instance can be reconfigured and +// its multi-backend routing can send different keys to different drivers. +var ( + _ artifact.Backend = (*Backend)(nil) + _ artifact.RangeReader = (*Backend)(nil) + _ artifact.Presigner = (*Backend)(nil) +) + +// Option configures the Backend. +type Option func(*Backend) + +// WithName overrides the backend identifier recorded on each artifact. +// Use the Trove store's name when running multi-store so an artifact's +// backend field points at the store that actually holds it. +func WithName(name string) Option { + return func(b *Backend) { b.name = name } +} + +// New wraps a Trove instance as an artifact.Backend. +func New(t *trovelib.Trove, opts ...Option) *Backend { + b := &Backend{trove: t, name: DefaultName} + + for _, opt := range opts { + opt(b) + } + + return b +} + +// Name identifies this backend. +func (b *Backend) Name() string { return b.name } + +// translate maps Trove's permanent failures onto the artifact plane's. +// +// This distinction is load-bearing. Callers use +// errors.Is(err, artifact.ErrPermanent) to tell a failure that can never +// succeed, which must fail the job immediately, from a transient backend +// failure, which should be retried with backoff. Getting it wrong means a +// deleted input burns every retry before reaching the DLQ. +// +// Only conditions Trove classifies are translated. Everything else, +// including a quota or rate limit, stays as it is and is retried: a quota +// may be granted later, and an error Trove has not classified is assumed +// transient, because dead-lettering a transient failure throws away work +// that would have succeeded. +// +// The underlying error is wrapped rather than replaced, so the driver's +// own message survives into the DLQ entry and errors.Is still reaches the +// Trove sentinel underneath for diagnostics. +func translate(err error) error { + switch { + case err == nil: + return nil + case errors.Is(err, trovelib.ErrNotFound): + // ErrObjectNotFound and ErrBucketNotFound unwrap to ErrNotFound, + // so this one case covers all three. + return fmt.Errorf("%w: %w", artifact.ErrNotFound, err) + case errors.Is(err, trovelib.ErrPermissionDenied): + return fmt.Errorf("%w: %w", artifact.ErrPermissionDenied, err) + default: + return err + } +} + +// Open returns a reader over the object's bytes. +func (b *Backend) Open(ctx context.Context, ref artifact.Ref) (io.ReadCloser, error) { + r, err := b.trove.Get(ctx, ref.Bucket, ref.Key) + if err != nil { + return nil, fmt.Errorf("trove: open %s/%s: %w", ref.Bucket, ref.Key, translate(err)) + } + + return r, nil +} + +// ErrRangeUnsupported means the underlying Trove driver cannot serve +// byte ranges, so the caller must read the whole object instead. +var ErrRangeUnsupported = errors.New("trove: driver does not support range reads") + +// OpenRange returns a reader over n bytes starting at off. A negative n +// reads to the end. +// +// Range support is a Trove driver capability, not a Get option: a driver +// that lacks it silently returns the entire object. Rather than hand back +// more bytes than asked for, this reports ErrRangeUnsupported so callers +// can fall back deliberately. +func (b *Backend) OpenRange(ctx context.Context, ref artifact.Ref, off, n int64) (io.ReadCloser, error) { + rd, ok := b.trove.Driver().(trovedriver.RangeDriver) + if !ok { + return nil, fmt.Errorf("%w: %T", ErrRangeUnsupported, b.trove.Driver()) + } + + r, err := rd.GetRange(ctx, ref.Bucket, ref.Key, off, n) + if err != nil { + return nil, fmt.Errorf("trove: open range [%d,%d) of %s/%s: %w", + off, off+n, ref.Bucket, ref.Key, translate(err)) + } + + return r, nil +} + +// Stat reports the object's size and content type. +func (b *Backend) Stat(ctx context.Context, ref artifact.Ref) (artifact.ObjectInfo, error) { + info, err := b.trove.Head(ctx, ref.Bucket, ref.Key) + if err != nil { + return artifact.ObjectInfo{}, fmt.Errorf("trove: stat %s/%s: %w", + ref.Bucket, ref.Key, translate(err)) + } + + return artifact.ObjectInfo{ + Size: info.Size, + ContentType: info.ContentType, + ETag: info.ETag, + }, nil +} + +// Delete removes the object. Deleting a missing object is not an error, +// which makes the purge pass idempotent under retry. +func (b *Backend) Delete(ctx context.Context, ref artifact.Ref) error { + err := b.trove.Delete(ctx, ref.Bucket, ref.Key) + if err == nil { + return nil + } + + terr := translate(err) + if errors.Is(terr, artifact.ErrNotFound) { + return nil + } + + // Return the translated error, not the raw one: a delete refused for + // permissions is permanent, and the sweeper would otherwise retry it + // on every pass forever. + return fmt.Errorf("trove: delete %s/%s: %w", ref.Bucket, ref.Key, terr) +} + +// PresignGet returns a time-limited read URL when the underlying driver +// supports pre-signing, and ErrNotFound-free failure otherwise. +func (b *Backend) PresignGet(ctx context.Context, ref artifact.Ref, ttl time.Duration) (string, error) { + p, ok := b.trove.Driver().(trovedriver.PresignDriver) + if !ok { + return "", fmt.Errorf("trove: driver %T does not support pre-signed URLs", b.trove.Driver()) + } + + url, err := p.PresignGet(ctx, ref.Bucket, ref.Key, ttl) + if err != nil { + return "", fmt.Errorf("trove: presign %s/%s: %w", ref.Bucket, ref.Key, translate(err)) + } + + return url, nil +} + +// Create begins writing a new object. +// +// Trove's Put consumes a reader, so the writer drives it from a pipe on a +// background goroutine and joins that goroutine in Commit or Abort. +func (b *Backend) Create(ctx context.Context, bucket, key string) (artifact.Writer, error) { + pr, pw := io.Pipe() + + w := &pipeWriter{ + pw: pw, + done: make(chan struct{}), + bucket: bucket, + key: key, + } + + go func() { + defer close(w.done) + + info, err := b.trove.Put(ctx, bucket, key, pr) + + w.info = info + w.putErr = err + + // Unblock any writer still feeding the pipe after Put returned — + // on error Put stops reading, and without this the next Write + // would block forever. + _ = pr.CloseWithError(err) + }() + + return w, nil +} + +// pipeWriter feeds Trove's Put from an io.Pipe. +// +// It counts the bytes the caller wrote so Commit reports the *logical* +// size. Trove's write middleware (compress, encrypt) changes the stored +// form, and the artifact row must record what the handler produced, not +// what landed on disk. +type pipeWriter struct { + pw *io.PipeWriter + done chan struct{} + bucket string + key string + + written int64 + info *trovedriver.ObjectInfo + putErr error + + once sync.Once + finished bool +} + +// Write appends bytes to the pending object. +func (w *pipeWriter) Write(p []byte) (int, error) { + if w.finished { + return 0, io.ErrClosedPipe + } + + n, err := w.pw.Write(p) + w.written += int64(n) + + return n, err +} + +// Commit closes the pipe, waits for Put, and reports the logical info. +func (w *pipeWriter) Commit(_ context.Context) (artifact.ObjectInfo, error) { + if w.finished { + return artifact.ObjectInfo{}, io.ErrClosedPipe + } + + w.finished = true + + w.once.Do(func() { _ = w.pw.Close() }) + <-w.done + + if w.putErr != nil { + return artifact.ObjectInfo{}, fmt.Errorf("trove: put %s/%s: %w", + w.bucket, w.key, translate(w.putErr)) + } + + out := artifact.ObjectInfo{Size: w.written} + if w.info != nil { + out.ContentType = w.info.ContentType + out.ETag = w.info.ETag + } + + return out, nil +} + +// Abort fails the pipe so Put stores nothing, then waits for it to +// unwind. It is a no-op after Commit. +func (w *pipeWriter) Abort() error { + if w.finished { + return nil + } + + w.finished = true + + w.once.Do(func() { _ = w.pw.CloseWithError(errAborted) }) + <-w.done + + // Best effort: a driver that already materialised the object before + // the pipe failed would leave bytes behind. Those become an orphan + // with no artifact row, which the orphan sweep collects. + return nil +} + +// errAborted fails the pipe so Trove's Put returns rather than storing a +// truncated object. +var errAborted = errors.New("trove: write aborted") diff --git a/artifact/trove/backend_test.go b/artifact/trove/backend_test.go new file mode 100644 index 0000000..5fc9e9e --- /dev/null +++ b/artifact/trove/backend_test.go @@ -0,0 +1,254 @@ +package trove_test + +import ( + "bytes" + "context" + "errors" + "io" + "testing" + + trovelib "github.com/xraph/trove" + "github.com/xraph/trove/drivers/memdriver" + + "github.com/xraph/dispatch/artifact" + troveadapter "github.com/xraph/dispatch/artifact/trove" +) + +const testBucket = "dispatch" + +func newBackend(t *testing.T) artifact.Backend { + t.Helper() + + ctx := context.Background() + + drv := memdriver.New() + if err := drv.Open(ctx, "mem://"); err != nil { + t.Fatalf("driver open: %v", err) + } + + tr, err := trovelib.Open(drv, trovelib.WithDefaultBucket(testBucket)) + if err != nil { + t.Fatalf("trove open: %v", err) + } + + t.Cleanup(func() { + if cerr := tr.Close(ctx); cerr != nil { + t.Errorf("trove close: %v", cerr) + } + }) + + if err := tr.CreateBucket(ctx, testBucket); err != nil && + !errors.Is(err, trovelib.ErrBucketExists) { + t.Fatalf("create bucket: %v", err) + } + + return troveadapter.New(tr) +} + +func writeObject(t *testing.T, b artifact.Backend, key string, data []byte) { + t.Helper() + + ctx := context.Background() + + w, err := b.Create(ctx, testBucket, key) + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write(data); werr != nil { + t.Fatalf("Write: %v", werr) + } + + if _, cerr := w.Commit(ctx); cerr != nil { + t.Fatalf("Commit: %v", cerr) + } +} + +func TestTroveRoundTrip(t *testing.T) { + ctx := context.Background() + b := newBackend(t) + + w, err := b.Create(ctx, testBucket, "mesh.glb") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write([]byte("meshbytes")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + info, err := w.Commit(ctx) + if err != nil { + t.Fatalf("Commit: %v", err) + } + + if info.Size != 9 { + t.Fatalf("info.Size = %d, want 9 (logical bytes written)", info.Size) + } + + ref := artifact.Ref{Backend: b.Name(), Bucket: testBucket, Key: "mesh.glb"} + + rc, err := b.Open(ctx, ref) + if err != nil { + t.Fatalf("Open: %v", err) + } + + got, rerr := io.ReadAll(rc) + if cerr := rc.Close(); cerr != nil { + t.Fatalf("Close: %v", cerr) + } + + if rerr != nil { + t.Fatalf("ReadAll: %v", rerr) + } + + if !bytes.Equal(got, []byte("meshbytes")) { + t.Fatalf("read %q, want %q", got, "meshbytes") + } +} + +func TestTroveOpenMissingMapsToErrNotFound(t *testing.T) { + _, err := newBackend(t).Open(context.Background(), + artifact.Ref{Bucket: testBucket, Key: "absent"}) + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Open(missing) = %v, want ErrNotFound", err) + } +} + +func TestTroveStat(t *testing.T) { + b := newBackend(t) + writeObject(t, b, "stat.bin", []byte("0123456789")) + + info, err := b.Stat(context.Background(), + artifact.Ref{Bucket: testBucket, Key: "stat.bin"}) + if err != nil { + t.Fatalf("Stat: %v", err) + } + + if info.Size != 10 { + t.Fatalf("info.Size = %d, want 10", info.Size) + } +} + +func TestTroveStatMissingMapsToErrNotFound(t *testing.T) { + _, err := newBackend(t).Stat(context.Background(), + artifact.Ref{Bucket: testBucket, Key: "absent"}) + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Stat(missing) = %v, want ErrNotFound", err) + } +} + +// TestTroveDeleteMissingIsNotAnError matters for the purge pass: it +// deletes bytes then the row, and must be safe to retry after a partial +// failure. +func TestTroveDeleteMissingIsNotAnError(t *testing.T) { + err := newBackend(t).Delete(context.Background(), + artifact.Ref{Bucket: testBucket, Key: "absent"}) + if err != nil { + t.Fatalf("Delete(missing) = %v, want nil", err) + } +} + +func TestTroveDelete(t *testing.T) { + ctx := context.Background() + b := newBackend(t) + writeObject(t, b, "doomed.bin", []byte("bye")) + + ref := artifact.Ref{Bucket: testBucket, Key: "doomed.bin"} + + if err := b.Delete(ctx, ref); err != nil { + t.Fatalf("Delete: %v", err) + } + + if _, err := b.Open(ctx, ref); !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Open after Delete = %v, want ErrNotFound", err) + } +} + +func TestTroveAbortPublishesNothing(t *testing.T) { + ctx := context.Background() + b := newBackend(t) + + w, err := b.Create(ctx, testBucket, "aborted.bin") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write([]byte("partial")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + if aerr := w.Abort(); aerr != nil { + t.Fatalf("Abort: %v", aerr) + } + + if _, err := b.Open(ctx, artifact.Ref{Bucket: testBucket, Key: "aborted.bin"}); !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("aborted object is readable; Open = %v, want ErrNotFound", err) + } +} + +func TestTroveAbortAfterCommitIsNoOp(t *testing.T) { + ctx := context.Background() + b := newBackend(t) + + w, err := b.Create(ctx, testBucket, "committed.bin") + if err != nil { + t.Fatalf("Create: %v", err) + } + + if _, werr := w.Write([]byte("data")); werr != nil { + t.Fatalf("Write: %v", werr) + } + + if _, cerr := w.Commit(ctx); cerr != nil { + t.Fatalf("Commit: %v", cerr) + } + + if aerr := w.Abort(); aerr != nil { + t.Fatalf("Abort after Commit must be a no-op, got %v", aerr) + } + + if _, err := b.Open(ctx, artifact.Ref{Bucket: testBucket, Key: "committed.bin"}); err != nil { + t.Fatalf("Abort after Commit destroyed the object: %v", err) + } +} + +func TestTroveOpenRange(t *testing.T) { + ctx := context.Background() + b := newBackend(t) + writeObject(t, b, "ranged.bin", []byte("0123456789")) + + rr, ok := b.(artifact.RangeReader) + if !ok { + t.Fatal("trove backend does not implement RangeReader") + } + + rc, err := rr.OpenRange(ctx, + artifact.Ref{Bucket: testBucket, Key: "ranged.bin"}, 2, 3) + if errors.Is(err, troveadapter.ErrRangeUnsupported) { + t.Skipf("driver does not support range reads: %v", err) + } + + if err != nil { + t.Fatalf("OpenRange: %v", err) + } + + got, rerr := io.ReadAll(rc) + if cerr := rc.Close(); cerr != nil { + t.Fatalf("Close: %v", cerr) + } + + if rerr != nil { + t.Fatalf("ReadAll: %v", rerr) + } + + if !bytes.Equal(got, []byte("234")) { + t.Fatalf("range read = %q, want %q", got, "234") + } +} + +func TestTroveBackendName(t *testing.T) { + if got := newBackend(t).Name(); got != "trove" { + t.Fatalf("Name() = %q, want %q", got, "trove") + } +} diff --git a/artifact/trove/doc.go b/artifact/trove/doc.go new file mode 100644 index 0000000..62e0107 --- /dev/null +++ b/artifact/trove/doc.go @@ -0,0 +1,14 @@ +// Package trove adapts a *trove.Trove instance as an artifact.Backend. +// +// It is the reference backend for Dispatch's artifact plane, but not a +// required one: the core depends only on the artifact.Backend interface, +// so any object store can be plugged in instead. +// +// Two Trove features carry through without any code here. Its write-path +// middleware means compression, AES-256-GCM encryption, and virus +// scanning are configuration rather than Dispatch concerns — and scanning +// on write is what lets a malicious upload be rejected before any +// memory-unsafe parser opens it. Its multi-store routing means the +// backend name recorded on each artifact is simply the Trove store it +// lives in. +package trove diff --git a/artifact/trove/translate_test.go b/artifact/trove/translate_test.go new file mode 100644 index 0000000..bb006b6 --- /dev/null +++ b/artifact/trove/translate_test.go @@ -0,0 +1,118 @@ +package trove + +import ( + "errors" + "fmt" + "strings" + "testing" + + trovelib "github.com/xraph/trove" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" +) + +// TestTranslate covers the mapping from Trove's classification onto the +// artifact plane's. What is *not* translated matters as much as what is: an +// error left alone is retried, so translating a transient condition would +// dead-letter work that would have succeeded, and failing to translate a +// permanent one burns the whole retry budget. +func TestTranslate(t *testing.T) { + tests := []struct { + name string + err error + + wantPermanent bool + wantSentinel error // nil means "returned unchanged" + }{ + { + name: "object not found", + err: fmt.Errorf(`memdriver: object "in.bin" not found in bucket "art": %w`, trovelib.ErrObjectNotFound), + wantPermanent: true, + wantSentinel: artifact.ErrNotFound, + }, + { + name: "bucket not found", + err: fmt.Errorf(`memdriver: bucket "art" not found: %w`, trovelib.ErrBucketNotFound), + wantPermanent: true, + wantSentinel: artifact.ErrNotFound, + }, + { + name: "general not found", + err: fmt.Errorf("cas: hash not found: %w", trovelib.ErrNotFound), + wantPermanent: true, + wantSentinel: artifact.ErrNotFound, + }, + { + name: "permission denied", + err: fmt.Errorf(`s3driver: permission denied for object "in.bin": %w`, trovelib.ErrPermissionDenied), + wantPermanent: true, + wantSentinel: artifact.ErrPermissionDenied, + }, + + // A quota may be granted later, so the job should back off rather + // than die. This is the case a blunt "permanent unless recognized" + // rule would get wrong. + { + name: "quota exceeded is retryable", + err: fmt.Errorf("s3driver: quota exceeded: %w", trovelib.ErrQuotaExceeded), + wantPermanent: false, + }, + { + name: "unclassified is retryable", + err: errors.New("dial tcp 10.0.0.1:443: connection refused"), + wantPermanent: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := translate(tt.err) + + if errors.Is(got, dispatch.ErrPermanent) != tt.wantPermanent { + t.Fatalf("translate(%v) permanent = %v, want %v", + tt.err, !tt.wantPermanent, tt.wantPermanent) + } + + if tt.wantSentinel == nil { + if !errors.Is(got, tt.err) { + t.Fatalf("translate returned %v, want the original error unchanged", got) + } + + return + } + + if !errors.Is(got, tt.wantSentinel) { + t.Fatalf("translate(%v) = %v, want %v", tt.err, got, tt.wantSentinel) + } + }) + } +} + +func TestTranslateNil(t *testing.T) { + if got := translate(nil); got != nil { + t.Fatalf("translate(nil) = %v, want nil", got) + } +} + +// TestTranslatePreservesDriverMessage covers what the substring fallback used +// to do and the sentinel branch did not: keep what the driver said. A DLQ +// entry reading only "not found" cannot tell an operator which object, on +// which driver, went missing. +func TestTranslatePreservesDriverMessage(t *testing.T) { + driverErr := fmt.Errorf( + `memdriver: object "input.bin" not found in bucket "artifacts": %w`, + trovelib.ErrObjectNotFound) + + got := translate(driverErr) + + if !strings.Contains(got.Error(), `object "input.bin" not found in bucket "artifacts"`) { + t.Fatalf("translate dropped the driver message: %q", got.Error()) + } + + // The Trove sentinel stays reachable underneath, so a diagnostic can + // still ask which resource was missing. + if !errors.Is(got, trovelib.ErrObjectNotFound) { + t.Fatalf("translate broke the chain to the Trove sentinel: %v", got) + } +} diff --git a/cluster/worker.go b/cluster/worker.go index 9d5e9b9..9eab3bc 100644 --- a/cluster/worker.go +++ b/cluster/worker.go @@ -4,6 +4,7 @@ import ( "time" "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" ) // WorkerState represents the lifecycle state of a worker. @@ -22,11 +23,15 @@ const ( // Worker represents a Dispatch worker instance in a distributed cluster. type Worker struct { - ID id.WorkerID `json:"id"` - Hostname string `json:"hostname"` - Queues []string `json:"queues"` - Concurrency int `json:"concurrency"` - State WorkerState `json:"state"` + ID id.WorkerID `json:"id"` + Hostname string `json:"hostname"` + Queues []string `json:"queues"` + Concurrency int `json:"concurrency"` + State WorkerState `json:"state"` + // Capacity is what this worker can run at once. It is advisory: + // empty means unknown, which disables the enqueue-time unschedulable + // check rather than rejecting everything. + Capacity resource.Set `json:"capacity,omitempty"` IsLeader bool `json:"is_leader"` LeaderUntil *time.Time `json:"leader_until,omitempty"` LastSeen time.Time `json:"last_seen"` diff --git a/config.go b/config.go index b4c89d6..4e3630d 100644 --- a/config.go +++ b/config.go @@ -29,6 +29,16 @@ type Config struct { // considered stale. StaleJobThreshold time.Duration + // ReapInterval is how often the pool scans for expired leases. Zero + // uses worker.DefaultReapInterval. This is the scan cadence, not the + // lease duration. + ReapInterval time.Duration + + // DefaultLeaseTTL is how far each renewal pushes a job's lease expiry + // when the job declares none of its own. Zero falls back to + // StaleJobThreshold, preserving existing reclamation timing. + DefaultLeaseTTL time.Duration + // WorkerStoreCallTimeout caps a single worker store roundtrip // (DequeueJobs, HeartbeatJob, ReapStaleJobs, UpdateJob). Bounds // how long a stalled driver session can hold a pool connection @@ -76,6 +86,8 @@ func DefaultConfig() Config { // Tuning fields (zero leaves the subsystem default in place; // callers override via options). WorkerStoreCallTimeout: 0, + ReapInterval: 0, + DefaultLeaseTTL: 0, CronTickInterval: 0, CronLeaderTTL: 0, CronLockTTL: 0, diff --git a/dispatcher_stop_test.go b/dispatcher_stop_test.go new file mode 100644 index 0000000..3875ded --- /dev/null +++ b/dispatcher_stop_test.go @@ -0,0 +1,118 @@ +package dispatch + +import ( + "context" + "sync" + "testing" +) + +type countingStore struct { + closes int +} + +func (s *countingStore) Migrate(_ context.Context) error { return nil } +func (s *countingStore) Ping(_ context.Context) error { return nil } +func (s *countingStore) Close() error { + s.closes++ + + return nil +} + +type countingExtensions struct { + shutdowns int +} + +func (e *countingExtensions) EmitShutdown(_ context.Context) { e.shutdowns++ } + +type countingPool struct { + stops int +} + +func (p *countingPool) Start(_ context.Context) error { return nil } +func (p *countingPool) Stop(_ context.Context) error { + p.stops++ + + return nil +} + +// TestDispatcherStopIsIdempotent covers a second Stop, which is not a +// hypothetical: Engine.Stop calls this, and a service shutting down from +// both a signal handler and a deferred cleanup calls Engine.Stop twice. +// +// Only the pool call used to be guarded, and only indirectly, through the +// started flag. EmitShutdown and the store Close ran every time. Extensions +// therefore saw two shutdown events and could release the same resources +// twice, and the store was closed twice, which the built-in backends +// tolerate only because their Close is a documented no-op. A custom Storer +// has no such promise, and neither does an extension. +// +// The engine has its own sync.Once over closeExecutors and documents that +// the rest of Stop tolerates a second call. That claim is only true if this +// one does. +func TestDispatcherStopIsIdempotent(t *testing.T) { + store := &countingStore{} + ext := &countingExtensions{} + pool := &countingPool{} + + d, err := New() + if err != nil { + t.Fatalf("New: %v", err) + } + d.store = store + d.SetExtensions(ext) + d.SetPool(pool) + d.started = true + + ctx := context.Background() + if stopErr := d.Stop(ctx); stopErr != nil { + t.Fatalf("first Stop: %v", stopErr) + } + if stopErr := d.Stop(ctx); stopErr != nil { + t.Fatalf("second Stop: %v", stopErr) + } + + if store.closes != 1 { + t.Errorf("store closed %d times, want 1", store.closes) + } + if ext.shutdowns != 1 { + t.Errorf("shutdown emitted %d times, want 1: extensions may release the "+ + "same resources on each one", ext.shutdowns) + } + if pool.stops != 1 { + t.Errorf("pool stopped %d times, want 1", pool.stops) + } +} + +// TestDispatcherStopIsIdempotentUnderConcurrency covers the same guard +// reached from two goroutines at once, which a flag check without +// synchronisation would let through. +func TestDispatcherStopIsIdempotentUnderConcurrency(t *testing.T) { + store := &countingStore{} + ext := &countingExtensions{} + + d, err := New() + if err != nil { + t.Fatalf("New: %v", err) + } + d.store = store + d.SetExtensions(ext) + d.started = true + + var wg sync.WaitGroup + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + // Errors are asserted via the counts below, not per call. + _ = d.Stop(context.Background()) + }() + } + wg.Wait() + + if store.closes != 1 { + t.Errorf("store closed %d times, want 1", store.closes) + } + if ext.shutdowns != 1 { + t.Errorf("shutdown emitted %d times, want 1", ext.shutdowns) + } +} diff --git a/dlq/entry.go b/dlq/entry.go index 0859bbb..f8cd147 100644 --- a/dlq/entry.go +++ b/dlq/entry.go @@ -4,10 +4,21 @@ import ( "time" "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" ) // Entry represents a job that has exhausted its retry budget and been // moved to the dead letter queue for inspection or replay. +// +// The fields below the identity block exist so Replay can rebuild a job +// that behaves like the one that failed. Replay calls EnqueueJob directly +// rather than going back through the engine, so nothing re-derives these +// for it: whatever the entry does not carry, the replayed job silently +// takes a default for. They are stored as the effective values from the +// failed job rather than looked up from the definition by name, because a +// definition's declaration is only half the story. The enqueue site can +// override every one of them, and a definition can be changed or removed +// between the failure and the replay. type Entry struct { ID id.DLQID `json:"id"` JobID id.JobID `json:"job_id"` @@ -22,4 +33,39 @@ type Entry struct { FailedAt time.Time `json:"failed_at"` ReplayedAt *time.Time `json:"replayed_at,omitempty"` CreatedAt time.Time `json:"created_at"` + + // Priority is the claim ordering the job was enqueued with. + Priority int `json:"priority,omitempty"` + + // Timeout is how long the handler was allowed to run. + Timeout time.Duration `json:"timeout,omitempty"` + + // LeaseTTL is how long each renewal extended the job's lease. Losing + // it is the reason this block exists: a six-hour job replayed without + // it falls back to the pool default, which is measured in seconds, so + // its lease lapses mid-run and it is reclaimed and restarted over and + // over without ever finishing. That is the exact failure long-running + // jobs carry a per-job TTL to avoid. + LeaseTTL time.Duration `json:"lease_ttl,omitempty"` + + // ArtifactBindings is the encoded map of declared input names to + // artifacts, carried verbatim. A handler that declares inputs cannot + // run without them. + ArtifactBindings []byte `json:"artifact_bindings,omitempty"` + + // Resources and ResourceLimits are what the job asked for and what it + // was capped at. Without them a replayed job looks free to schedule + // and can be claimed by a worker that cannot actually host it. + Resources resource.Set `json:"resources,omitempty"` + ResourceLimits resource.Set `json:"resource_limits,omitempty"` + + // ResourceClass is the named class the job was placed in. + ResourceClass string `json:"resource_class,omitempty"` + + // InputBytes and PrimaryInputHash are derived from the bindings by the + // engine at enqueue. Replay does not run that derivation, so they are + // carried rather than recomputed, and stay consistent with the + // bindings above. + InputBytes int64 `json:"input_bytes,omitempty"` + PrimaryInputHash string `json:"primary_input_hash,omitempty"` } diff --git a/dlq/replay.go b/dlq/replay.go index f2a2081..1c8591e 100644 --- a/dlq/replay.go +++ b/dlq/replay.go @@ -30,6 +30,21 @@ func (s *Service) Replay(ctx context.Context, entryID id.DLQID) (*job.Job, error ScopeAppID: entry.ScopeAppID, ScopeOrgID: entry.ScopeOrgID, RunAt: now, + + // Restored from the failed job. This path calls EnqueueJob + // directly instead of going back through the engine, so nothing + // re-derives any of these: whatever is not set here silently + // becomes a default, and for LeaseTTL that default is short + // enough to make a long job unrunnable. See the Entry doc. + Priority: entry.Priority, + Timeout: entry.Timeout, + LeaseTTL: entry.LeaseTTL, + ArtifactBindings: entry.ArtifactBindings, + Resources: entry.Resources, + ResourceLimits: entry.ResourceLimits, + ResourceClass: entry.ResourceClass, + InputBytes: entry.InputBytes, + PrimaryInputHash: entry.PrimaryInputHash, } if err := s.jobStore.EnqueueJob(ctx, j); err != nil { diff --git a/dlq/service.go b/dlq/service.go index ae05503..edbc928 100644 --- a/dlq/service.go +++ b/dlq/service.go @@ -36,6 +36,19 @@ func (s *Service) Push(ctx context.Context, j *job.Job, jobErr error) error { ScopeOrgID: j.ScopeOrgID, FailedAt: now, CreatedAt: now, + + // Everything below is carried so Replay can rebuild a job that + // behaves like this one. See the Entry doc for why it is copied + // from the job rather than looked up from the definition. + Priority: j.Priority, + Timeout: j.Timeout, + LeaseTTL: j.LeaseTTL, + ArtifactBindings: j.ArtifactBindings, + Resources: j.Resources, + ResourceLimits: j.ResourceLimits, + ResourceClass: j.ResourceClass, + InputBytes: j.InputBytes, + PrimaryInputHash: j.PrimaryInputHash, } return s.store.PushDLQ(ctx, entry) } diff --git a/dlq/service_test.go b/dlq/service_test.go index cd410d2..fefd344 100644 --- a/dlq/service_test.go +++ b/dlq/service_test.go @@ -10,6 +10,7 @@ import ( dispatchDLQ "github.com/xraph/dispatch/dlq" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" "github.com/xraph/dispatch/store/memory" ) @@ -205,3 +206,93 @@ func TestService_Replay_NotFoundReturnsError(t *testing.T) { t.Fatal("expected error for non-existent DLQ entry") } } + +// TestService_Replay_PreservesExecutionFields covers the round trip that +// makes a replayed job behave like the one that failed. +// +// Replay builds a job and calls EnqueueJob directly rather than going back +// through the engine, so nothing re-derives these values on the way out: +// whatever Push failed to capture, or Replay failed to restore, the new +// job silently runs with a default for. LeaseTTL is the one that hurts. +// A six-hour job replayed on the pool default lease lapses mid-run, gets +// reclaimed, restarts, and never finishes, which is exactly what a per-job +// TTL exists to prevent. +func TestService_Replay_PreservesExecutionFields(t *testing.T) { + s := memory.New() + svc := dispatchDLQ.NewService(s, s) + ctx := context.Background() + + j := newTestJob("long-render", []byte(`{"frame":42}`)) + j.Priority = 7 + j.Timeout = 90 * time.Minute + j.LeaseTTL = 6 * time.Hour + j.ArtifactBindings = []byte(`{"input":"artifact_abc"}`) + j.Resources = resource.Set{"cpu_milli": 2000} + j.ResourceLimits = resource.Set{"cpu_milli": 4000} + j.ResourceClass = "gpu-large" + j.InputBytes = 4096 + j.PrimaryInputHash = "sha256:deadbeef" + + if err := svc.Push(ctx, j, errors.New("handler exploded")); err != nil { + t.Fatalf("Push: %v", err) + } + + entries, err := s.ListDLQ(ctx, dispatchDLQ.ListOpts{Limit: 10}) + if err != nil { + t.Fatalf("ListDLQ: %v", err) + } + if len(entries) != 1 { + t.Fatalf("listed %d entries, want 1", len(entries)) + } + + replayed, err := svc.Replay(ctx, entries[0].ID) + if err != nil { + t.Fatalf("Replay: %v", err) + } + + if replayed.LeaseTTL != j.LeaseTTL { + t.Errorf("LeaseTTL = %v, want %v: the replayed job would fall back to the "+ + "pool default and be reclaimed mid-run forever", replayed.LeaseTTL, j.LeaseTTL) + } + if replayed.Timeout != j.Timeout { + t.Errorf("Timeout = %v, want %v", replayed.Timeout, j.Timeout) + } + if replayed.Priority != j.Priority { + t.Errorf("Priority = %d, want %d", replayed.Priority, j.Priority) + } + if string(replayed.ArtifactBindings) != string(j.ArtifactBindings) { + t.Errorf("ArtifactBindings = %q, want %q", replayed.ArtifactBindings, j.ArtifactBindings) + } + if replayed.Resources["cpu_milli"] != 2000 { + t.Errorf("Resources = %v, want cpu_milli 2000", replayed.Resources) + } + if replayed.ResourceLimits["cpu_milli"] != 4000 { + t.Errorf("ResourceLimits = %v, want cpu_milli 4000", replayed.ResourceLimits) + } + if replayed.ResourceClass != j.ResourceClass { + t.Errorf("ResourceClass = %q, want %q", replayed.ResourceClass, j.ResourceClass) + } + if replayed.InputBytes != j.InputBytes { + t.Errorf("InputBytes = %d, want %d", replayed.InputBytes, j.InputBytes) + } + if replayed.PrimaryInputHash != j.PrimaryInputHash { + t.Errorf("PrimaryInputHash = %q, want %q", replayed.PrimaryInputHash, j.PrimaryInputHash) + } + + // A replay is a fresh job, not a resumption of the failed one: it must + // not inherit the identity, the exhausted retry budget, or any of the + // ownership the failed run left behind. + if replayed.ID == j.ID { + t.Error("replayed job reused the failed job's ID") + } + if replayed.RetryCount != 0 { + t.Errorf("RetryCount = %d, want 0", replayed.RetryCount) + } + if replayed.State != job.StatePending { + t.Errorf("State = %s, want pending", replayed.State) + } + if replayed.LeaseExpiresAt != nil || replayed.StartedAt != nil { + t.Errorf("replayed job carries ownership from the failed run: "+ + "lease_expires_at=%v started_at=%v", replayed.LeaseExpiresAt, replayed.StartedAt) + } +} diff --git a/doc.go b/doc.go index bbc5ba0..a864cb3 100644 --- a/doc.go +++ b/doc.go @@ -15,8 +15,17 @@ // # Architecture // // Dispatch follows a composable store pattern where each subsystem (job, -// workflow, cron, dlq, event, cluster) defines its own store interface. -// A single backend implements all of them. +// workflow, cron, dlq, event, cluster, artifact) defines its own store +// interface. A single backend implements all of them. +// +// # Artifacts +// +// Jobs that process large files use the artifact plane rather than the +// payload column: an artifact is a tracked reference to an object in +// external storage, declared as a job input and staged to a +// content-addressed local cache before the handler runs. See the artifact +// package. It is opt-in — with no backend configured, Dispatch behaves as +// it did before artifacts existed. // // All entity IDs use TypeID — type-prefixed, K-sortable, UUIDv7-based, // compile-time safe identifiers. diff --git a/docs/content/docs/subsystems/artifacts.mdx b/docs/content/docs/subsystems/artifacts.mdx new file mode 100644 index 0000000..59193de --- /dev/null +++ b/docs/content/docs/subsystems/artifacts.mdx @@ -0,0 +1,233 @@ +--- +title: Artifacts +description: Tracked object storage for job inputs and outputs. +--- + +Jobs that process large files have nowhere good to put them. A job's payload is a +`BYTEA` column, so nobody puts a two-gigabyte model in it — they put an S3 URL in a +string field instead, and from that moment Dispatch is blind to the data. It cannot +size the job before scheduling it, cannot tell that two jobs want the same file, +cannot clean up the intermediates, and cannot show you what a run consumed. + +The artifact plane makes that data a first-class thing Dispatch knows about. + + + Artifacts are entirely opt-in. With no backend configured, Dispatch behaves + exactly as it did before this existed. + + +## The two lifecycles + +Everything follows from which of these an artifact is. + +**Durable** artifacts are yours. Your application uploaded them; Dispatch registers +them, reads them, and never deletes them. A customer's uploaded model is durable. + +**Ephemeral** artifacts are Dispatch's. A handler created them through the artifact +API, and Dispatch reclaims them once every job or run that references them has +finished and a retention window has passed. A tessellated mesh passed between +workflow steps is ephemeral. + +The separation exists so that a bug in reclamation cannot reach your customers' +data. `lifecycle = 'ephemeral'` appears as a literal in every sweep query on every +backend, and the sweeper re-checks each artifact before acting on it. + +## Declaring inputs + +A job declares the artifacts it consumes. Declarations are what let the engine +validate bindings at enqueue, know the total input size before scheduling, and +materialise the bytes before your handler is called. + +```go +var Tessellate = job.NewDefinition("tessellate.model", + func(ctx context.Context, in TessellateInput) error { + art := artifact.From(ctx) + + // Already on local disk before the handler ran. + src := art.Path("model") + + mesh, err := occt.Tessellate(src, in.Detail) + if err != nil { + return err + } + + w, err := art.Create(ctx, "mesh.glb", + artifact.ContentType("model/gltf-binary")) + if err != nil { + return err + } + defer w.Abort() // no-op after a successful Commit + + if _, err := io.Copy(w, mesh); err != nil { + return err + } + + _, err = w.Commit(ctx) + return err + }, + job.WithArtifactInputs( + artifact.Input("model", + artifact.Required, + artifact.MaxSize(8<<30), + artifact.StageAsPath), + ), + job.WithTimeout(6*time.Hour), +) +``` + +Register with `RegisterChecked` to have the declaration validated against the +staging budget, so a job that could never be staged fails on your machine rather +than on a worker: + +```go +if err := engine.RegisterChecked(eng, Tessellate); err != nil { + log.Fatal(err) +} +``` + +Bind an artifact when you enqueue: + +```go +ref, err := svc.Register(ctx, "uploads", "tower.ifc") +if err != nil { + return err +} + +_, err = engine.Enqueue(ctx, eng, "tessellate.model", + TessellateInput{Detail: 0.5}, + engine.Bind("model", ref)) +``` + +An oversized or undeclared binding is rejected here, at enqueue, so it never +becomes a job that fails on a worker and burns retries. + +### Staging modes + +`StageAsPath` (the default) downloads the artifact before the handler runs and +gives you a file path. This is what CAD kernels, mesh importers, and PDF engines +need — they seek and memory-map, so a stream is no use to them. + +`StageLazy` downloads nothing up front. Call `art.Open(ctx, name)` to stream it if +and when you need it. Right for data you read once, front to back. + +## Creating outputs + +Outputs are imperative rather than declared, so a handler can produce a number of +them it only discovers at run time: + +```go +for i, page := range pages { + w, err := art.Create(ctx, fmt.Sprintf("page-%d.png", i)) + if err != nil { + return err + } + // ...render, then Commit +} +``` + +### Resuming a retried job + +`Commit` publishes immediately and records the attempt, because a six-hour job +splitting a four-hundred-page PDF cannot buffer its work until it returns. That +makes `IfAbsent` possible: + +```go +w, err := art.Create(ctx, "page-317.png", artifact.IfAbsent()) +if errors.Is(err, artifact.ErrExists) { + continue // a previous attempt already rendered this page +} +``` + +A retried job can skip what it already did instead of starting over. + +## The staging cache + +Staged artifacts live in a content-addressed cache on each worker, keyed by their +BLAKE3 hash. + +Three things follow from that. Two jobs consuming the same model share one copy on +disk. Eight jobs staging it at once trigger one download, not eight. And a job +re-run over an input it already staged pays nothing. + +The hash is computed *during* the download rather than by a separate pass, which is +why registering an artifact does not hash it: doing so would turn a cheap row +insert into a full read of a multi-gigabyte file. An artifact earns its hash the +first time something stages it. + +A byte budget bounds the cache. When it is full, unleased entries are evicted +least-recently-used; when everything is in use, a job **waits** rather than filling +the volume. An entry a running handler holds is never evicted. + +```yaml +cache: + dir: /var/lib/dispatch/cache + budget: 214748364800 # 200 GiB +``` + +## Configuration + +With Forge, mount both extensions and Dispatch discovers Trove from the container: + +```go +app := forge.New( + troveext.New(), // provides *trove.Trove + dispatchext.New(), // discovers it +) +``` + +```yaml +extensions: + dispatch: + artifacts: + enabled: true + trove_store: "" # "" uses the default instance + bucket: dispatch-artifacts + ephemeral_prefix: ephemeral + retention: 168h + purge_grace: 24h + cache: + dir: /var/lib/dispatch/cache + budget: 214748364800 +``` + +Outside Forge, or with a store that is not Trove, supply a backend directly: + +```go +dispatchext.New( + dispatchext.WithArtifactBackend(myBackend), + dispatchext.WithArtifactCacheDir("/var/lib/dispatch/cache"), + dispatchext.WithArtifactCacheBudget(200 << 30), +) +``` + +Because the backend is an interface, Trove is a default rather than a dependency. +Two of its features are worth knowing about: its `encrypt` middleware gives you +artifacts encrypted at rest with no Dispatch code, and its `scan` middleware sits +on the write path, so a malicious upload can be rejected before any parser opens +it. + +## Reclamation + +Ephemeral artifacts are reclaimed in two phases. A **sweep** marks an artifact +deleted and stops serving it. A **purge**, `purge_grace` later, removes the bytes +and the row. A mistaken sweep is therefore observable and reversible for a day +rather than instantly destructive. + +An artifact becomes eligible once every owner that links it has reached a terminal +state and `retention` has elapsed since the last of them finished. Artifacts with +no links at all — only possible after a crash mid-create — are handled by a +separate orphan pass. + +Sweeping runs on the elected leader only. It can be stopped without a redeploy, and +run in dry-run mode to see what it would do. + +## Error handling + +| Situation | What happens | +|---|---| +| Input artifact was deleted | Fails fast to the DLQ — retrying a fetch of something gone cannot succeed | +| Backend timeout or 5xx | Transient; retried with normal backoff | +| Binding exceeds the declared `MaxSize` | Rejected at enqueue, returned to the caller | +| Declared inputs exceed the cache budget | Rejected at `RegisterChecked` | +| Cache full and everything is leased | Bounded by the job's deadline, then `ErrBudgetExceeded` | +| Worker killed mid-job | Leases are in-memory, so death releases them | diff --git a/docs/content/docs/subsystems/execution-isolation.mdx b/docs/content/docs/subsystems/execution-isolation.mdx new file mode 100644 index 0000000..b8e81c8 --- /dev/null +++ b/docs/content/docs/subsystems/execution-isolation.mdx @@ -0,0 +1,460 @@ +--- +title: Execution Isolation +description: Declaring the isolation a handler needs, and how Dispatch enforces it. +--- + +A handler that parses an untrusted customer upload with a memory-unsafe native +library runs, today, in the same process as your database credentials. If the +parser has a bug and the file is crafted to hit it, whatever the parser can +reach, the attacker can reach — which in a typical worker is everything. + +The execution plane gives a job definition a way to say it needs more than +that, and gives the engine a way to refuse to run it if the deployment cannot +provide it. + + + Execution isolation is entirely opt-in. A definition that declares nothing + runs in-process exactly as it always has. + + +## The ladder + +Isolation is a `Level`, and levels are ordered: + +```go +exec.LevelNone // in-process: no isolation. The default. +exec.LevelProcess // a separate address space +exec.LevelSandboxed // + mount, network, PID, and user namespaces, seccomp, dropped capabilities +exec.LevelVM // + an independent kernel (gVisor or Kata) +``` + +A definition declares the *minimum* it requires, not the executor it runs on +— which rung actually satisfies that requirement is a deployment decision, +made by whichever executors the deployment has configured. + +`LevelNone`, the in-process executor, is the only rung present in every +deployment by default. `LevelProcess` now has a real executor too — +`exec/subprocess`, covered in detail below — but it has to be configured +explicitly (see "Adding a stronger rung"); a deployment that never does so +has only `LevelNone` available, which matters for the reason below. +`LevelSandboxed` and `LevelVM` describe rungs a later phase adds (OCI +container and Kubernetes pod, respectively). Declaring either now is +legitimate — it documents the requirement — but no deployment can satisfy +them yet. + +## Declaring a policy + +```go +var Tessellate = job.NewDefinition("tessellate.model", + func(ctx context.Context, in TessellateInput) error { + // parses an untrusted customer upload with a native geometry kernel + return nil + }, + job.WithExecution( + exec.Isolate(exec.LevelProcess), + exec.GracePeriod(45*time.Second), + ), +) +``` + +`exec.Isolate` sets the minimum level. `exec.GracePeriod` sets how long a +sandbox gets to exit cleanly after being signalled before it is killed +outright — `exec/subprocess`'s kill ladder already uses this today (see +"Deadlines are enforced, not just cancelled" below), and later +out-of-process rungs will too; the in-process executor ignores it, since +there's no process to signal. `exec.Image` +overrides the container image an out-of-process rung launches, for a handler +that needs something other than the worker's own image. `exec.AllowDowngrade` +is covered next. + +## No silent downgrades + +`exec.Registry.Select` picks the weakest configured executor that satisfies +the declared level — a job asking for `LevelProcess` isolation is not handed +a Kubernetes pod merely because one happens to be configured, and it is not +handed the in-process executor merely because that's what's available. + +When nothing configured satisfies the level, `Select` **fails** rather than +running the handler with less isolation than it declared. A parser that +demanded a separate address space does not quietly fall back to running +in-process just because no stronger rung was wired up — that would defeat +the entire point of declaring the requirement. `exec.AllowDowngrade()` opts +out of this for a definition that would genuinely rather run weakly isolated +than not run at all. + +`engine.RegisterChecked` and `engine.RegisterAll` call `Select` at +**registration**, not at execution, so a definition demanding isolation the +deployment cannot provide fails when you start the process, not on the first +malicious upload that reaches production. `engine.Register` skips the check, +for definitions you've already verified some other way. + +A deployment that has not configured any rung above `LevelNone` — which is +every deployment until the section below — is in exactly the situation the +previous two paragraphs describe: **any definition that declares above +`LevelNone` fails registration unless it also sets +`exec.AllowDowngrade()`**, because there is nothing configured that can +satisfy it yet. This is expected, not a bug — it's `Select` doing its job +with only the in-process rung on the ladder. + + + `exec.AllowDowngrade()` is not "accept a weaker rung than declared" in the + abstract — its fallback is always the in-process executor specifically + (`Registry.Select` returns the registry's always-present default when + nothing configured satisfies the level). Adding it to a `LevelProcess` + definition so registration stops failing, without first configuring the + subprocess rung below, means that definition runs in-process — in the + worker's own address space, next to its credentials — the moment nothing + satisfies the level, silently, with no registration failure to notice. + That is true today with no rung configured, and stays true later if the + subprocess rung is configured and then becomes unavailable for any + reason: disabled, misconfigured, or refused at startup for a reason + covered below. For a handler whose isolation requirement exists because + it parses untrusted input, only set `AllowDowngrade` on a definition that + genuinely tolerates running unisolated — never as a way to silence a + registration failure while assuming the stronger rung is actually in + place. + + +## Adding a stronger rung + +`engine.WithExecutor` adds an executor to the deployment's registry: + +```go +subprocessExecutor := subprocess.New( + subprocess.WithUser(65534, 65534), // never the worker's own uid — see below +) + +eng, err := engine.Build(d, + engine.WithExecutor(subprocessExecutor), +) +if err != nil { + log.Fatal(err) +} +``` + +The in-process executor is always present as the default, so a deployment +that adds nothing behaves exactly as before this existed. Adding an executor +only changes what becomes available to definitions that ask for it. + +`exec.LevelProcess` now has a real executor behind it: `exec/subprocess`, +covered in detail below. + +## Registering a mixed set + +`job.Registrable` lets definitions with different payload types share one +slice, since Go forbids a generic method that would otherwise unify them. +`SendEmail` below declares no isolation, so it stays in the worker process +alongside the isolated `Tessellate` from above: + +```go +var SendEmail = job.NewDefinition("send.email", + func(ctx context.Context, in EmailInput) error { return nil }, +) + +var defs []job.Registrable = []job.Registrable{Tessellate, SendEmail} + +if err := engine.RegisterAll(eng, defs...); err != nil { + log.Fatal(err) +} +``` + +`RegisterAll` validates every definition's policy before registering any of +them, so a rejected set leaves the registry as it was rather than half +populated. This is also the seam `shim.Main` already uses, today, as the +subprocess rung's out-of-process entrypoint (see "The subprocess rung" +below): it is handed the same `[]job.Registrable` without ever holding an +`*Engine`. + +## The subprocess rung + +`exec/subprocess` is the `exec.LevelProcess` executor: the worker re-execs +its own binary for every attempt, the child runs the handler under a +dedicated low-privilege uid — enforced at launch, not just advised; `Run` +refuses to start without one unless the operator explicitly opts out (see +"The dedicated uid is mandatory" below) — with POSIX resource limits +applied, and a kill ladder enforces the deadline instead of trusting the +handler to notice a cancelled context. The child never receives the +worker's memory, and not its environment wholesale either — only a fixed +allowlist (`PATH`, `HOME`, `TMPDIR`) is copied across, everything else +comes from the request's own `Env`, constructed rather than inherited — +so it cannot read a credential that lived only in some other worker +environment variable. `HOME` is in that allowlist, though, which is +exactly what locates `~/.aws` below; the dedicated uid is what actually +stops the child from reading credentials that live on disk — the +worker's own config file, a mounted secret, `~/.aws` — since the child +shares the worker's filesystem view. + + + It does not stop the file from reaching the network or another tenant's + data — the child shares the worker's network namespace and filesystem + view outside its scratch directory. That containment is `exec.LevelSandboxed`'s + job, a stronger rung a later phase builds on the same wire protocol. Declare + `exec.LevelSandboxed` if that is what a handler actually needs; declaring + `exec.LevelProcess` for it and getting this rung instead is a downgrade the + registry will not silently hand you (see "No silent downgrades" above), but + only if you declare the level you mean. + + +### The `main()` branch — read this before enabling anything else + +The worker distinguishes a normal run from a re-exec'd child by one thing: +`argv[1] == "dispatch-exec"`. Nothing in `engine.Build` or `subprocess.New` +adds that branch for you — it has to be the first thing `main` checks, +before flags are parsed, before the store connects, before anything else +that a sandboxed attempt has no business doing: + +```go +func main() { + if len(os.Args) > 1 && os.Args[1] == shim.ArgName { + shim.Main(Tessellate, SendEmail) // same []job.Registrable as RegisterAll + return // unreachable: shim.Main calls os.Exit + } + + d, err := dispatch.New(dispatch.WithStore(store)) + if err != nil { + log.Fatal(err) + } + + eng, err := engine.Build(d, engine.WithExecutor(subprocessExecutor)) + if err != nil { + log.Fatal(err) + } + if err := engine.RegisterAll(eng, defs...); err != nil { + log.Fatal(err) + } + // ... start the worker +} +``` + + + This is the single most likely first-time-setup failure. Without the + branch, enabling the subprocess rung makes every attempt re-exec a whole + second copy of the worker — it connects to the store, starts polling, + does everything a normal worker does except read `exec.Request` off fd 3, + because nothing ever routes it to `shim.Main`. That child never writes a + result frame, so the parent's kill ladder eventually SIGTERMs and + SIGKILLs it at the deadline, and the attempt fails as `StatusTimeout` + every single time. Nothing about this looks like a wiring bug from the + logs alone — it looks like every job is timing out. + + +`shim.Main` takes the exact same `[]job.Registrable` as `RegisterAll` — +same handler set, same fingerprint. `exec.Fingerprint` is compared between +parent and child on every attempt, so a child built from a different +handler set than the worker that launched it fails to launch rather than +running a stale handler silently. + +### Enabling it + +Programmatically, `subprocess.New` plus `engine.WithExecutor`, as shown +above. Through Forge, the extension resolves `execution.subprocess` from +YAML and does the equivalent `engine.WithExecutor` wiring for you — but not +the `main()` branch above. That branch cannot live in a package this +repository ships; it has to be the first thing your own `main()` checks, +YAML config or not: + +```yaml +extensions: + dispatch: + execution: + subprocess: + enabled: true + user: 65534 # required to get real isolation — see below + group: 65534 + scratch_dir: /var/lib/dispatch/scratch + rlimits: + address_space: 2147483648 # 2 GiB, bytes; rejected outright on Darwin, see below + nofile: 256 + nproc: 32 + fsize: 1073741824 # 1 GiB + strict_rlimits: true +``` + +`binary` defaults to `os.Executable()` — the worker's own binary — which is +correct unless the sandboxed handlers are built as a separate binary. +`allow_same_user` is covered next. A deployment that sets nothing under +`execution` registers no extra executor; every job keeps running in-process +exactly as it does today. + +This rung refuses to launch on any non-Unix platform. The extension checks +that once, at startup — `execution.subprocess.enabled: true` on Windows +fails to boot the worker rather than failing every job's first launch +attempt once it is already in production. + +### The dedicated uid is mandatory + +`user` and `group` are not optional hardening — they are most of what this +rung is for, and Dispatch enforces that rather than merely advising it. +Enabling the subprocess rung with no `user` configured refuses to start at +all, both from `main()`'s own `subprocess.New` path and from YAML config: +without it, the child would run as the worker's own uid, with the worker's +own read access to `~/.aws`, `/var/run/secrets`, and the Dispatch config +file itself, and a parser exploited in that child could read every +credential the worker can. Configuring `scratch_dir` and `rlimits` without +`user` narrows *where* a compromised handler can write and *how much* it +can allocate; it does nothing to stop it from reading what the worker could +already read, which is why it is not accepted as a substitute. + +A configured `user` that silently equals the worker's own uid is refused +the same way, for the same reason: running the child as the worker gives up +the uid boundary regardless of whether that happened because `user` was +never set or because it was set to a value that turned out to match. Both +shapes share a single opt-out, `allow_same_user` (or +`subprocess.WithAllowSameUser()`) — not a warning for one and a hard +refusal for the other. That escape hatch exists for local development and +CI, where dropping to a different uid needs privilege the process running +the tests does not have, and for a deployment that has made a deliberate, +informed decision to accept running this rung unisolated — not for +production by default. + +**Supplementary groups.** A genuine uid drop — `user` naming a uid other +than the worker's own — clears the child's supplementary groups along with +its primary uid and gid: dropping to a different uid already needs the same +privilege `setgroups(2)` itself requires, so the clear happens for free. The +one case where it does *not* happen is the `allow_same_user` escape hatch — +same uid, same gid — where there is nothing to clear anyway, since the +child is running as the worker's own account. In other words: whenever this +rung is actually providing the isolation it exists for, supplementary +groups go with it; the one path where they survive is the same path that +already gives up the uid boundary too. + +### Deadlines are enforced, not just cancelled + +`job.WithTimeout` means something stronger here than it does in-process. +In-process, the timeout middleware wraps the handler's context in +`context.WithTimeout` and nothing more — a handler that never checks +`ctx.Done()` runs to completion regardless, because there is no process +boundary to enforce anything against. Under the subprocess rung, the same +deadline is also sent to the child as `exec.Request.Deadline`, and this +package's kill ladder enforces it from the outside: SIGTERM to the child's +whole process group, a grace period (`exec.GracePeriod`, default 30s) for a +cooperative exit, then SIGKILL to the group if it has not emptied out by +then. A handler that ignores cancellation entirely — stuck in a native +call with no Go code polling `ctx.Done()` — still gets killed at the +deadline plus grace. That is the actual, measurable difference this rung +buys over in-process for a job whose native library might hang. + +That enforcement has a side effect worth knowing before it surprises you in +a retry log: the attempt's `Status` is decided by whether the deadline +elapsed at all, not by whether the handler eventually behaved. The parent +watches its own timer for `req.Deadline` independently of the child; the +instant it fires, it takes one non-blocking look at whether the process has +already exited and been reaped, and if not, the attempt is `StatusTimeout` +— full stop, regardless of what happens next. A handler that traps the +signal, wraps up, writes a valid result frame, and exits 0 a moment later +still gets `StatusTimeout`, not `StatusOK` or `StatusHandlerError`: a +decoded frame is only authoritative once the parent has *not* already +decided the deadline won that race. In practice, this means any handler +whose deadline actually elapses is timed out, cooperative or not — the +distinction that matters isn't "did it eventually exit cleanly," it's "did +it finish its work before the deadline in the first place." + +### What `Setpgid` changes + +The child is started as its own process group leader, not a member of the +worker's group. That is what lets the kill ladder reach a native library's +forked helpers along with the tracked process itself — signalling the +negative pid addresses the whole group, not just one member of it. It has +a cost worth knowing about: a signal aimed at the *worker's* process group +— `kill -TERM -$pgid` from a supervisor, or a shell sending SIGINT to a +whole job-control group on Ctrl-C — no longer reaches the sandboxed child, +because the child is not in that group anymore. Only the worker's own +graceful-shutdown path, which cancels the attempt's context and lets this +package's kill ladder run, reaches the child cleanly. + +That matters most when the worker itself does not get a graceful shutdown +at all — killed outright, OOM-killed, or crashed. The child is orphaned, +running with nothing left to signal it. It still enforces its own deadline +internally (the shim applies `req.Deadline` to the handler's context, the +same context cancellation in-process relies on), so a cooperative handler +still stops on its own; an uncooperative one does not, and nothing catches +it, because `Executor.Reclaim` for this rung is a no-op today — sweeping +orphaned children across a restart needs a worker identity stable across +restarts that this rung does not have yet. + +What `Reclaim` *does* sweep, at worker startup, is narrower than "leftover +scratch state": only `worker.Runner`'s own per-attempt output directories +(`dispatch-out---…`), and only when this Runner has the +artifact plane configured — a Runner without one skips the sweep entirely, +since nothing it does creates that kind of directory without one. A swept +entry also has to be at least an hour old (`staleScratchDirAge`) on top of +its owning PID no longer being alive, so the PID-collision hazard below +needs a leaked directory to sit around that long, not just outlive its +process by a moment. + +`exec/subprocess.Executor`'s *own* scratch working directory — the +`dispatch-exec-…` directory `Run` creates per attempt as the child's +`Cmd.Dir`, distinct from the Runner's output directory even when both are +rooted at the same configured `scratch_dir` — is not covered by this sweep +at all. It is removed by a deferred `os.RemoveAll` inside `Run` itself, so +a normal exit (success, failure, or a kill-ladder timeout) always cleans it +up; a worker that crashes mid-attempt, the same case that orphans the child +process (above), leaves that directory behind permanently. Nothing sweeps +it on the next startup. + +Ownership matching itself is inherently single-host: `processAlive` asks +the local kernel whether a PID is alive, which only means something within +one PID namespace. A deployment where multiple containers share a scratch +volume across separate PID namespaces can have a live sibling's PID +collide with a number this worker's own kernel view reports as dead, and +sweep a directory a running attempt is still using. Don't share a +subprocess scratch directory across PID-namespace boundaries. + +### Resource limits degrade silently by default + +`rlimits` are applied child-side, inside the shim, because Go's `os/exec` +has no way to set a child's rlimits through `SysProcAttr`. Two things about +that path are worth knowing before relying on it: + +- **Darwin rejects `RLIMIT_AS` outright.** The kernel refuses + `setrlimit(RLIMIT_AS, ...)` unconditionally (`EINVAL`), regardless of the + value requested. This is a platform fact, not a configuration mistake, + and it is treated as one even with `strict_rlimits` on — there is no + value that would make it succeed. +- **Some platforms don't get a resource number at all — and this case + *is* a strict-mode failure, unlike Darwin's above.** Dispatch has + confirmed the raw `RLIMIT_NPROC` number for Linux, Darwin, and FreeBSD + only; on every other Unix this rung otherwise supports — NetBSD, + OpenBSD, Dragonfly, Solaris/illumos, AIX, Android — a configured `nproc` + limit is skipped rather than risk applying a number that has not been + checked for that platform. OpenBSD has the same gap for `RLIMIT_AS` + specifically, for a different reason: Go's `syscall` package does not + export the constant there at all, so a configured `address_space` limit + is skipped on OpenBSD too. Both are a library gap Dispatch could close by + verifying the number, not a kernel refusing something no value could fix + — which is exactly why, unlike Darwin's `RLIMIT_AS` rejection, + `strict_rlimits` does treat these as launch failures. + +Without `strict_rlimits: true`, both of these — and any other rlimit that +fails to apply, like a value exceeding the process's own hard limit — are +warnings, not launch failures, logged to the child's stderr. That stream +goes nowhere by default: the subprocess executor's own logger defaults to +a no-op, so a warning nobody configured a logger to receive is a warning +nobody sees. If a configured limit is meant to be a hard guarantee, set +`strict_rlimits: true` and configure a logger — either turns a silent +degradation into a launch failure or, at minimum, a line an operator can +actually find. + +### Inputs are not yet staged out-of-process + +`exec.Request` has `Inputs` and `InputDir` fields, and the shim's artifact +accessor already reads from them (`Accessor.Path`, `Accessor.Open`) — but +nothing in the worker populates them for an out-of-process attempt today. +Only *prior outputs* from an earlier attempt of the same job are staged +into the child (so a retried attempt can skip work it already committed); +declared inputs from `job.WithArtifactInputs` are not. + +This is not a "fetch it yourself" gap — the child cannot fetch anything. +Its artifact service is deliberately credential-free, a local directory +standing in for the real backend (see `exec/shim`'s own doc comment), so +there is no backend to fetch a declared input *from* even if a handler +tried. A handler that declares both `job.WithArtifactInputs` and an +isolation level above `LevelNone` registers and runs today, but calling +`Accessor.Open` for that declared input inside the sandboxed handler +returns `artifact.ErrUnbound` — the binding exists in the job's policy, it +is simply never staged where this rung's handler can reach it. Nothing at +registration time catches this combination yet. Staging declared inputs +into `InputDir` for out-of-process rungs is Phase 3 work; until it lands, +keep a handler's declared artifact inputs on definitions that run +in-process, or fetch outputs of *earlier attempts* only — the one input +path this rung already wires up. diff --git a/docs/content/docs/subsystems/meta.json b/docs/content/docs/subsystems/meta.json index 215499e..5857f20 100644 --- a/docs/content/docs/subsystems/meta.json +++ b/docs/content/docs/subsystems/meta.json @@ -2,6 +2,8 @@ "title": "Subsystems", "pages": [ "dwp", + "artifacts", + "execution-isolation", "catalog", "delivery", "dlq", diff --git a/docs/source.config.ts b/docs/source.config.ts index 773f29b..98bc6f1 100644 --- a/docs/source.config.ts +++ b/docs/source.config.ts @@ -6,6 +6,7 @@ import { defineConfig, defineDocs } from "fumadocs-mdx/config"; export const docs = defineDocs({ dir: "content/docs", docs: { + async: true, schema: pageSchema, postprocess: { includeProcessedMarkdown: true, diff --git a/docs/src/app/docs/[[...slug]]/page.tsx b/docs/src/app/docs/[[...slug]]/page.tsx index 77b83ee..c5e6684 100644 --- a/docs/src/app/docs/[[...slug]]/page.tsx +++ b/docs/src/app/docs/[[...slug]]/page.tsx @@ -17,10 +17,10 @@ export default async function Page(props: PageProps<"/docs/[[...slug]]">) { const page = source.getPage(params.slug); if (!page) notFound(); - const MDX = page.data.body; + const { body: MDX, toc } = await page.data.load(); return ( - + {page.data.title} {page.data.description} diff --git a/docs/src/components/ui/card.tsx b/docs/src/components/ui/card.tsx index f856923..9882504 100644 --- a/docs/src/components/ui/card.tsx +++ b/docs/src/components/ui/card.tsx @@ -78,9 +78,9 @@ CardFooter.displayName = "CardFooter"; export { Card, - CardHeader, + CardContent, + CardDescription, CardFooter, + CardHeader, CardTitle, - CardDescription, - CardContent, }; diff --git a/engine/artifact.go b/engine/artifact.go new file mode 100644 index 0000000..a824b0b --- /dev/null +++ b/engine/artifact.go @@ -0,0 +1,115 @@ +package engine + +import ( + "context" + "fmt" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/artifact/staging" + "github.com/xraph/dispatch/job" +) + +// WithArtifacts enables the artifact plane. +// +// The staging middleware is appended to the chain, so declared inputs are +// materialised before a handler runs and every cache lease is released +// afterwards. Without this option Dispatch behaves exactly as it did +// before artifacts existed. +func WithArtifacts(svc *artifact.Service, c *cache.Cache) Option { + return func(eng *Engine) { + if svc == nil || !svc.Enabled() { + return + } + + eng.artifacts = svc + eng.artifactCache = c + + eng.mws = append(eng.mws, staging.Middleware(svc, c, func(name string) []artifact.InputSpec { + return eng.registry.Inputs(name) + }, staging.WithLogger(eng.logger))) + } +} + +// Artifacts returns the artifact service, or nil when the plane is off. +func (eng *Engine) Artifacts() *artifact.Service { return eng.artifacts } + +// ArtifactCache returns the staging cache, or nil when the plane is off. +func (eng *Engine) ArtifactCache() *cache.Cache { return eng.artifactCache } + +// ValidateArtifactInputs checks a definition's declarations against the +// engine's staging capacity. +// +// A definition whose declared inputs could never fit the cache budget is +// rejected here, at registration, rather than failing every time such a +// job runs. The point is to surface the mistake on a developer's machine +// instead of at 3am in production. +func (eng *Engine) ValidateArtifactInputs(name string, specs []artifact.InputSpec) error { + if len(specs) == 0 { + return nil + } + + if err := artifact.ValidateInputs(specs); err != nil { + return fmt.Errorf("job %q: %w", name, err) + } + + if eng.artifacts == nil { + return fmt.Errorf( + "job %q declares artifact inputs but no artifact backend is configured: %w", + name, artifact.ErrNoBackend) + } + + if eng.artifactCache == nil { + return nil + } + + // A zero total means at least one declaration is unbounded, so the + // requirement is unknown rather than known-too-large. + total := artifact.TotalMaxSize(specs) + if total == 0 { + return nil + } + + if budget := eng.artifactCache.Budget(); total > budget { + return fmt.Errorf( + "job %q declares up to %d bytes of artifact inputs, which exceeds the %d byte staging budget: %w", + name, total, budget, cache.ErrBudgetExceeded) + } + + return nil +} + +// Bind attaches an artifact to a declared input at enqueue. +// +// Bindings are validated against the job's declarations before the job is +// persisted, so an oversized or undeclared input is a caller error rather +// than a job that fails on a worker. +func Bind(name string, ref artifact.Ref) job.Option { + return func(o *job.Options) { + if o.Bindings == nil { + o.Bindings = make(map[string]artifact.Ref) + } + + o.Bindings[name] = ref + } +} + +// applyBindings validates a job's bindings and records them on the job. +func (eng *Engine) applyBindings(_ context.Context, j *job.Job, bindings map[string]artifact.Ref) error { + if len(bindings) == 0 { + return nil + } + + if eng.artifacts == nil { + return fmt.Errorf("job %q binds artifacts but the artifact plane is not enabled: %w", + j.Name, artifact.ErrNoBackend) + } + + specs := eng.registry.Inputs(j.Name) + + if err := staging.Validate(specs, bindings); err != nil { + return fmt.Errorf("job %q: %w", j.Name, err) + } + + return staging.SetBindings(j, bindings) +} diff --git a/engine/artifact_test.go b/engine/artifact_test.go new file mode 100644 index 0000000..7941da7 --- /dev/null +++ b/engine/artifact_test.go @@ -0,0 +1,318 @@ +package engine_test + +import ( + "context" + "errors" + "io" + "os" + "sync/atomic" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" +) + +type tessellateInput struct { + Detail float64 `json:"detail"` +} + +type artifactRig struct { + engine *engine.Engine + svc *artifact.Service + backend *artifacttest.Backend + store *memory.Store +} + +func newArtifactRig(t *testing.T, budget int64) *artifactRig { + t.Helper() + + s := memory.New() + b := artifacttest.NewBackend() + + svc := artifact.NewService(s, b, + artifact.WithEphemeralPrefix("ephemeral"), + artifact.WithDefaultBucket("dispatch")) + + c, err := cache.New(t.TempDir(), b, cache.WithBudget(budget)) + if err != nil { + t.Fatalf("cache.New: %v", err) + } + + t.Cleanup(func() { + if cerr := c.Close(); cerr != nil { + t.Errorf("cache close: %v", cerr) + } + }) + + d, err := dispatch.New( + dispatch.WithStore(s), + dispatch.WithConcurrency(2), + dispatch.WithQueues([]string{"default"}), + ) + if err != nil { + t.Fatalf("dispatch.New: %v", err) + } + + eng, err := engine.Build(d, engine.WithArtifacts(svc, c)) + if err != nil { + t.Fatalf("engine.Build: %v", err) + } + + return &artifactRig{engine: eng, svc: svc, backend: b, store: s} +} + +// TestRegisterCheckedRejectsUnstageableDefinition is the point of +// validating at registration: a job whose declared inputs could never fit +// the staging budget must fail on a developer's machine, not on a worker. +func TestRegisterCheckedRejectsUnstageableDefinition(t *testing.T) { + rig := newArtifactRig(t, 1024) + + def := job.NewDefinition("too-big", + func(context.Context, tessellateInput) error { return nil }, + job.WithArtifactInputs(artifact.Input("model", artifact.MaxSize(1<<30))), + ) + + err := engine.RegisterChecked(rig.engine, def) + if !errors.Is(err, cache.ErrBudgetExceeded) { + t.Fatalf("RegisterChecked = %v, want ErrBudgetExceeded", err) + } +} + +func TestRegisterCheckedAcceptsFittingDefinition(t *testing.T) { + rig := newArtifactRig(t, 1<<20) + + def := job.NewDefinition("fits", + func(context.Context, tessellateInput) error { return nil }, + job.WithArtifactInputs(artifact.Input("model", artifact.MaxSize(1024))), + ) + + if err := engine.RegisterChecked(rig.engine, def); err != nil { + t.Fatalf("RegisterChecked: %v", err) + } +} + +func TestRegisterCheckedRejectsDuplicateNames(t *testing.T) { + rig := newArtifactRig(t, 1<<20) + + def := job.NewDefinition("dupes", + func(context.Context, tessellateInput) error { return nil }, + job.WithArtifactInputs( + artifact.Input("model"), + artifact.Input("model"), + ), + ) + + if err := engine.RegisterChecked(rig.engine, def); err == nil { + t.Fatal("duplicate input declarations must be rejected at registration") + } +} + +func TestEnqueueRejectsOversizeBinding(t *testing.T) { + ctx := context.Background() + rig := newArtifactRig(t, 1<<20) + rig.backend.Put("models", "big.ifc", make([]byte, 500)) + + def := job.NewDefinition("capped", + func(context.Context, tessellateInput) error { return nil }, + job.WithArtifactInputs(artifact.Input("model", artifact.MaxSize(100))), + ) + engine.Register(rig.engine, def) + + ref, err := rig.svc.Register(ctx, "models", "big.ifc") + if err != nil { + t.Fatalf("Register artifact: %v", err) + } + + _, err = engine.Enqueue(ctx, rig.engine, "capped", tessellateInput{}, + engine.Bind("model", ref)) + if !errors.Is(err, artifact.ErrSizeExceeded) { + t.Fatalf("Enqueue with an oversize binding = %v, want ErrSizeExceeded", err) + } +} + +func TestEnqueueRejectsUndeclaredBinding(t *testing.T) { + ctx := context.Background() + rig := newArtifactRig(t, 1<<20) + rig.backend.Put("models", "x.ifc", []byte("data")) + + def := job.NewDefinition("declared-only", + func(context.Context, tessellateInput) error { return nil }, + job.WithArtifactInputs(artifact.Input("model")), + ) + engine.Register(rig.engine, def) + + ref, err := rig.svc.Register(ctx, "models", "x.ifc") + if err != nil { + t.Fatalf("Register artifact: %v", err) + } + + _, err = engine.Enqueue(ctx, rig.engine, "declared-only", tessellateInput{}, + engine.Bind("surprise", ref)) + if !errors.Is(err, artifact.ErrUndeclared) { + t.Fatalf("Enqueue with an undeclared binding = %v, want ErrUndeclared", err) + } +} + +func TestEnqueueRejectsMissingRequiredInput(t *testing.T) { + ctx := context.Background() + rig := newArtifactRig(t, 1<<20) + + def := job.NewDefinition("needs-model", + func(context.Context, tessellateInput) error { return nil }, + job.WithArtifactInputs(artifact.Input("model", artifact.Required)), + ) + engine.Register(rig.engine, def) + + // A binding is present but not the required one, so validation runs. + rig.backend.Put("models", "other.ifc", []byte("data")) + + ref, err := rig.svc.Register(ctx, "models", "other.ifc") + if err != nil { + t.Fatalf("Register artifact: %v", err) + } + + _, err = engine.Enqueue(ctx, rig.engine, "needs-model", tessellateInput{}, + engine.Bind("model", ref)) + if err != nil { + t.Fatalf("Enqueue with the required binding: %v", err) + } +} + +// TestEndToEndStageAndCommit runs a real job through the pool: the input +// is staged to disk before the handler sees it, and the output the +// handler writes is committed and linked back to the job. +func TestEndToEndStageAndCommit(t *testing.T) { + ctx := context.Background() + rig := newArtifactRig(t, 1<<20) + rig.backend.Put("models", "tower.ifc", []byte("ifc-source-bytes")) + + var ( + processed atomic.Bool + gotSource atomic.Value + ) + + def := job.NewDefinition("tessellate", + func(ctx context.Context, _ tessellateInput) error { + art := artifact.From(ctx) + + path := art.Path("model") + if path == "" { + return errors.New("input was not staged") + } + + data, err := os.ReadFile(path) + if err != nil { + return err + } + + gotSource.Store(string(data)) + + w, err := art.Create(ctx, "mesh.glb", + artifact.ContentType("model/gltf-binary")) + if err != nil { + return err + } + + defer func() { _ = w.Abort() }() + + if _, err := io.WriteString(w, "tessellated"); err != nil { + return err + } + + if _, err := w.Commit(ctx); err != nil { + return err + } + + processed.Store(true) + + return nil + }, + job.WithArtifactInputs(artifact.Input("model", artifact.Required)), + ) + + if err := engine.RegisterChecked(rig.engine, def); err != nil { + t.Fatalf("RegisterChecked: %v", err) + } + + ref, err := rig.svc.Register(ctx, "models", "tower.ifc") + if err != nil { + t.Fatalf("Register artifact: %v", err) + } + + j, err := engine.Enqueue(ctx, rig.engine, "tessellate", + tessellateInput{Detail: 0.5}, engine.Bind("model", ref)) + if err != nil { + t.Fatalf("Enqueue: %v", err) + } + + if serr := rig.engine.Start(ctx); serr != nil { + t.Fatalf("engine.Start: %v", serr) + } + + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if serr := rig.engine.Stop(stopCtx); serr != nil { + t.Errorf("engine.Stop: %v", serr) + } + }) + + waitFor(t, 5*time.Second, processed.Load) + + if got, _ := gotSource.Load().(string); got != "ifc-source-bytes" { + t.Fatalf("handler read %q from the staged path, want %q", got, "ifc-source-bytes") + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: j.ID.String()} + + outputs, err := rig.store.ListArtifactsByOwner(ctx, owner, artifact.RoleOutput) + if err != nil { + t.Fatalf("ListArtifactsByOwner: %v", err) + } + + if len(outputs) != 1 { + t.Fatalf("got %d outputs, want 1", len(outputs)) + } + + if outputs[0].Lifecycle != artifact.Ephemeral { + t.Fatalf("output lifecycle = %q, want ephemeral", outputs[0].Lifecycle) + } + + if outputs[0].Size != int64(len("tessellated")) { + t.Fatalf("output size = %d, want %d", outputs[0].Size, len("tessellated")) + } + + // The input's hash should have been recorded during staging, since + // registration deliberately skipped it. + in, err := rig.store.GetArtifact(ctx, ref.ID) + if err != nil { + t.Fatalf("GetArtifact: %v", err) + } + + if in.ContentHash == "" { + t.Fatal("staging did not record the input's content hash") + } +} + +// waitFor polls until cond holds or the deadline passes. +func waitFor(t *testing.T, d time.Duration, cond func() bool) { + t.Helper() + + deadline := time.Now().Add(d) + for time.Now().Before(deadline) { + if cond() { + return + } + + time.Sleep(10 * time.Millisecond) + } + + t.Fatalf("condition not met within %v", d) +} diff --git a/engine/double_stop_test.go b/engine/double_stop_test.go new file mode 100644 index 0000000..c4e949d --- /dev/null +++ b/engine/double_stop_test.go @@ -0,0 +1,35 @@ +package engine_test + +import ( + "context" + "testing" + + "github.com/xraph/dispatch/engine" +) + +// TestEngine_StopTwiceClosesExecutorsOnce hardens against a regression: Stop +// used to have no guard against being called twice, which was harmless only +// because closeExecutors had no caller at all. Now that Stop closes every +// registered executor, a second Stop call must not close them again — an +// out-of-process rung's Close releases real clients and child processes, +// and closing those twice is not something any rung is required to +// tolerate the way a no-op double-Stop is. +func TestEngine_StopTwiceClosesExecutorsOnce(t *testing.T) { + rung := &countingExecutor{} + eng, _ := startEngine(t, engine.WithExecutor(rung)) + + if _, closed := rung.counts(); closed != 0 { + t.Fatalf("closed count before any Stop = %d, want 0", closed) + } + + if err := eng.Stop(context.Background()); err != nil { + t.Fatalf("first Stop: %v", err) + } + if err := eng.Stop(context.Background()); err != nil { + t.Fatalf("second Stop: %v", err) + } + + if _, closed := rung.counts(); closed != 1 { + t.Errorf("Close called %d times across two Stop calls, want 1", closed) + } +} diff --git a/engine/engine.go b/engine/engine.go index 3b0155a..ed40858 100644 --- a/engine/engine.go +++ b/engine/engine.go @@ -14,6 +14,8 @@ import ( "errors" "fmt" "os" + "slices" + "sync" "time" log "github.com/xraph/go-utils/log" @@ -21,17 +23,21 @@ import ( "go.opentelemetry.io/otel/trace" "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/cache" "github.com/xraph/dispatch/backoff" "github.com/xraph/dispatch/cluster" "github.com/xraph/dispatch/cron" "github.com/xraph/dispatch/dlq" "github.com/xraph/dispatch/event" + "github.com/xraph/dispatch/exec" "github.com/xraph/dispatch/ext" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" mw "github.com/xraph/dispatch/middleware" "github.com/xraph/dispatch/observability" "github.com/xraph/dispatch/queue" + "github.com/xraph/dispatch/resource" "github.com/xraph/dispatch/scope" "github.com/xraph/dispatch/store" "github.com/xraph/dispatch/stream" @@ -80,6 +86,23 @@ type Engine struct { mws []mw.Middleware logger log.Logger + // stopOnce guards the executor-close path in Stop against a double + // call. Stop's other steps tolerate being run twice: the cron + // scheduler and the dispatcher each hold their own sync.Once, and the + // worker pool checks a running flag. Close has no such guard of its + // own, and closing a rung's clients or child processes twice is not + // guaranteed safe the way a no-op Stop is. + // + // The scope is deliberately narrow rather than wrapping all of Stop. + // Every other step owns its own idempotence, at the layer that knows + // what repeating it costs, and this guard exists only for the one + // step that cannot. Widening it here would put that decision in the + // wrong place and hide from a reader that the subsystems already + // handle it. Note a second Stop returns nil rather than the first + // call's error, because the dispatcher's own Once reports nothing on + // a repeat call. + stopOnce sync.Once + // Workflow subsystem. wfRegistry *workflow.Registry wfRunner *workflow.Runner @@ -99,6 +122,26 @@ type Engine struct { brokerOpts []stream.BrokerOption enableBroker bool + // Artifact plane (optional; nil means disabled). + artifacts *artifact.Service + artifactCache *cache.Cache + + // Resource model (optional; zero values mean no requirements and no + // capacity check, which is exactly today's behaviour). + estimator resource.Estimator + resourceDefault resource.Set + queueResources map[string]resource.Set + workerCapacity resource.Set + + // resources is the shared admission ledger. It must be the SAME + // instance the staging cache was built with, or the cache's staged + // bytes are invisible to the pool's disk budget. Nil disables the + // model outright. + resources resource.Manager + // workerCustomKeys narrows the custom keys this worker advertises at + // dequeue. Empty derives them from the manager's capacity. + workerCustomKeys []string + // Queue subsystem. queueConfigs []queue.Config queueManager *queue.Manager @@ -108,6 +151,19 @@ type Engine struct { // metricFactory is the go-utils MetricFactory for engine-level metrics. // nil means use gu.NewMetricsCollector default. metricFactory gu.MetricFactory + + // executors is the registry job attempts are dispatched through. It + // always has the in-process executor as its default. + executors *exec.Registry + // extraExecutors accumulates executors added via WithExecutor until + // buildExecutors assembles them into executors. + extraExecutors []exec.Executor + // scratchRoot is the root directory an out-of-process attempt's + // scratch OutputDir is created under (worker.Runner.WithArtifacts). + // Empty means worker.Runner's own default, os.TempDir(). See + // WithScratchRoot for why it only takes effect alongside the artifact + // plane. + scratchRoot string } // Option configures an Engine. @@ -143,6 +199,104 @@ func WithQueueConfig(configs ...queue.Config) Option { } } +// WithEstimator installs the resource estimator consulted at enqueue. +// +// The estimator sits above a definition's static declaration and below +// a per-enqueue override. It receives the declaration in the request and +// may return it unchanged, so installing one is an explicit opt-in to +// letting inference override declaration. An estimator that errors is +// ignored: it must never fail an enqueue. +func WithEstimator(e resource.Estimator) Option { + return func(eng *Engine) { eng.estimator = e } +} + +// WithResourceDefaults sets the fleet-wide default requirement and any +// per-queue overrides. Both are the lowest-precedence sources, below a +// definition's own declaration. +func WithResourceDefaults(global resource.Set, perQueue map[string]resource.Set) Option { + return func(eng *Engine) { + eng.resourceDefault = global + eng.queueResources = perQueue + } +} + +// WithWorkerCapacity declares the largest single-worker capacity in the +// fleet, and is the ONLY thing that turns the enqueue-time unschedulable +// check on. Leave it unset — the default — and Enqueue never rejects a +// job for being too big for any worker. +// +// It is a fleet-wide statement, not a description of this process, and +// the distinction is the whole reason the check is opt-in. Declare it +// and a job requiring more than this on any dimension fails Enqueue with +// ErrUnschedulable, wherever it was enqueued from: a light API pod that +// declared its own 2 GiB would hard-reject the tessellation job the +// heavy tier runs perfectly well. +// +// The check cannot derive the fleet maximum for itself, because +// cluster.Worker.Capacity does not round-trip. Only store/memory carries +// it; postgres, sqlite, mongo, redis and the k8s provider all enumerate +// worker fields by hand and drop it, so a worker registered with +// {memory: 64GiB} reads back an empty map. MaxWorkerCapacity therefore +// sees this value and — on memory alone — whatever live workers +// published, never the real fleet maximum. Persisting Capacity in those +// four models would make the derivation honest and is tracked as +// follow-up work; until then, declaring the ceiling is the operator's +// job or the check stays off. +// +// It is deliberately NOT defaulted from WithResourceManager's capacity. +// That default read as a convenience and behaved as a silent rescope of +// a fleet-wide question to one process. +func WithWorkerCapacity(c resource.Set) Option { + return func(eng *Engine) { eng.workerCapacity = c } +} + +// WithResourceManager installs the admission ledger the worker pool +// admits jobs against. +// +// The manager passed here MUST be the same instance the staging cache +// was built with (cache.WithManager). One ledger is the whole design: +// the cache holds a lease per cached entry and registers itself as the +// manager's disk reclaimer, and the pool's dequeue budget offers disk as +// free PLUS what that reclaimer could evict. Give the cache a private +// manager — which it constructs for itself when none is supplied — and +// the pool's Reclaimable() is permanently zero, staged bytes are never +// offered back to the budget, and the disk path quietly does nothing. +// It presents as a worker that went quiet, not as an error. +// +// Leaving this unset is the supported default: the pool passes an +// unbounded DequeueOpts, every backend skips its fit predicate, and no +// leases are taken. That is exactly how Dispatch behaved before the +// resource model existed. +// +// Installing a manager does NOT declare a fleet capacity, and does not +// turn the enqueue-time unschedulable check on. See WithWorkerCapacity +// for why that is opt-in and separate. +// +// Leases compose with this. The grant travels on job.DequeueOpts +// (WorkerID and LeaseUntil), so a claim that takes a lease is an +// ordinary claim that also writes the lease columns — it carries the +// budget, the custom-key containment and the locality preference like +// any other. There is one dequeue path per backend and it is the one +// this manager constrains. Turning leases and resources on together is +// the natural upgrade, and it is a supported one. +func WithResourceManager(m resource.Manager) Option { + return func(eng *Engine) { eng.resources = m } +} + +// WithWorkerCustomKeys narrows the custom resource keys this worker +// advertises at dequeue. +// +// The default — every custom key the manager has capacity for — is +// usually right. This exists to shrink it, so a worker draining a device +// can stop attracting work for it without being reconfigured. Keep the +// list a subset of the manager's custom capacity: dequeue matches custom +// keys by containment and never by quantity, so a key advertised here +// with no capacity behind it passes the store's filter and is then +// refused locally, on every attempt. +func WithWorkerCustomKeys(keys []string) Option { + return func(eng *Engine) { eng.workerCustomKeys = slices.Clone(keys) } +} + // WithTracerProvider sets a custom OTel TracerProvider for the engine. // When set, the tracing middleware uses this provider instead of the global one. // If not set, the global otel.GetTracerProvider() is used. @@ -230,6 +384,11 @@ func Build(d *dispatch.Dispatcher, opts ...Option) (*Engine, error) { opt(eng) } + // Assemble the executor registry now that WithExecutor options have + // populated extraExecutors, and before any definition is registered + // or the runner is built, since both consult it. + eng.buildExecutors() + // Create stream broker if enabled (must be before pool so events flow). if eng.enableBroker { eng.broker = stream.NewBroker(logger, eng.brokerOpts...) @@ -281,9 +440,34 @@ func Build(d *dispatch.Dispatcher, opts ...Option) (*Engine, error) { allMws = append(allMws, defaultMws...) allMws = append(allMws, eng.mws...) - // Create executor and pool. + // Create runner and pool. config := d.Config() - executor := worker.NewExecutor(eng.registry, eng.extensions, eng.jobStore, eng.dlqService, eng.bo, logger, allMws...) + runner := worker.NewRunner( + eng.registry, eng.extensions, eng.jobStore, eng.dlqService, + eng.bo, eng.executors, logger, allMws..., + ) + + // An out-of-process rung gets a scratch directory regardless of + // whether this runs — worker.Runner creates one for any attempt whose + // executor is above exec.LevelNone, artifact plane or not — and + // worker.Runner.Reclaim's startup sweep of directories a previous + // process left behind now runs unconditionally too, for the same + // reason. What this call actually turns on is PriorOutputs and output + // committing (worker.Runner.commitOutputs), which is the one thing + // genuinely gated on having somewhere to commit to. Gated on + // eng.artifacts specifically — not on whether an extra executor is + // configured — because that is the same condition commitOutputs + // itself gates on; calling this with a nil svc would be a no-op by + // its own contract, so there is nothing to lose by keeping the + // condition here identical rather than trying to also know about + // every executor WithExecutor might have added. One side effect worth + // knowing: eng.scratchRoot only ever reaches the Runner through this + // call's second argument, so with no artifact plane configured a + // Runner's scratch directories fall back to os.TempDir() even if + // WithScratchRoot named something else. + if eng.artifacts != nil { + runner.WithArtifacts(eng.artifacts, eng.scratchRoot) + } poolOpts := []worker.PoolOption{ worker.WithPoolConcurrency(config.Concurrency), @@ -300,6 +484,12 @@ func Build(d *dispatch.Dispatcher, opts ...Option) (*Engine, error) { if config.WorkerStoreCallTimeout != 0 { poolOpts = append(poolOpts, worker.WithStoreCallTimeout(config.WorkerStoreCallTimeout)) } + if config.ReapInterval != 0 { + poolOpts = append(poolOpts, worker.WithReapInterval(config.ReapInterval)) + } + if config.DefaultLeaseTTL != 0 { + poolOpts = append(poolOpts, worker.WithDefaultLeaseTTL(config.DefaultLeaseTTL)) + } // Create queue manager if queue configs were provided. if len(eng.queueConfigs) > 0 { @@ -307,9 +497,44 @@ func Build(d *dispatch.Dispatcher, opts ...Option) (*Engine, error) { poolOpts = append(poolOpts, worker.WithQueueManager(eng.queueManager)) } + // Hand the pool the shared ledger. The timing check runs only when a + // manager is present, because admission is the only thing that can + // stall the fetcher: with no manager, admissionBudget hands back the + // pool's own context and admit returns immediately. + if eng.resources != nil { + _, leaseAware := js.(job.LeaseStore) + if err := checkReaperMargin(config, leaseAware); err != nil { + return nil, err + } + + poolOpts = append(poolOpts, worker.WithResourceManager(eng.resources)) + + // Deliberately NOT seeding eng.workerCapacity from the ledger. + // The manager describes THIS process; workerCapacity is the floor + // of a fleet-wide check. Defaulting one from the other rescoped + // the question to one process, and because cluster.Worker.Capacity + // does not round-trip on four of the five backends, nothing could + // raise it back to the fleet maximum afterwards — so a light API + // worker rejected at enqueue every job bigger than itself. See + // WithWorkerCapacity. + // + // No construction-time warning about leases is emitted here, and + // there is no longer anything to warn about. The lease grant + // travels on job.DequeueOpts, so worker.Pool's single dequeue + // path — DequeueJobs — is also the path that takes a lease, and + // it carries the budget, the custom keys and the locality + // preference whether or not a lease is being granted. There is no + // second entry point that could claim a job this worker cannot + // run. See WithResourceManager and job.LeaseStore. + } + + if len(eng.workerCustomKeys) > 0 { + poolOpts = append(poolOpts, worker.WithWorkerCustomKeys(eng.workerCustomKeys)) + } + eng.pool = worker.NewPool( eng.jobStore, - executor, + runner, eng.extensions, logger, poolOpts..., @@ -360,6 +585,7 @@ func Build(d *dispatch.Dispatcher, opts ...Option) (*Engine, error) { Hostname: hostname, Queues: config.Queues, Concurrency: config.Concurrency, + Capacity: eng.workerCapacity.Clone(), State: cluster.WorkerActive, LastSeen: time.Now().UTC(), CreatedAt: time.Now().UTC(), @@ -372,10 +598,38 @@ func Build(d *dispatch.Dispatcher, opts ...Option) (*Engine, error) { } // Register registers a typed job definition with the engine. +// +// Use RegisterChecked when the definition declares artifact inputs, or an +// execution policy (job.WithExecution), and you want either validated — +// the staging budget for the former, and for the latter, that some +// configured executor can actually satisfy it. Register itself performs +// neither check: a definition declaring exec.Isolate(exec.LevelProcess) +// with no rung configured to provide it registers cleanly here and only +// fails the first time a worker actually tries to run it, which is +// exactly the startup-error-versus-silently-hung-job choice +// RegisterChecked exists to take out of a caller's hands. func Register[T any](eng *Engine, def *job.Definition[T]) { job.RegisterDefinition(eng.registry, def) } +// RegisterChecked registers a definition and validates its artifact +// declarations and execution policy, so a job that could never be staged +// or could never be isolated as it requires fails here — at registration, +// on a developer's machine — rather than on every worker that picks it +// up. +func RegisterChecked[T any](eng *Engine, def *job.Definition[T]) error { + if err := eng.ValidateArtifactInputs(def.Name, def.Opts.Inputs); err != nil { + return err + } + if err := eng.checkExecutionPolicy(def.Name, def.Opts.Execution); err != nil { + return err + } + + job.RegisterDefinition(eng.registry, def) + + return nil +} + // Enqueue creates and enqueues a job. func Enqueue[T any](ctx context.Context, eng *Engine, name string, payload T, opts ...job.Option) (*job.Job, error) { data, err := json.Marshal(payload) @@ -415,10 +669,36 @@ func (eng *Engine) EnqueueRaw(ctx context.Context, name string, payload []byte, j.Priority = jobOpts.Priority j.MaxRetries = jobOpts.MaxRetries j.Timeout = jobOpts.Timeout + + // A definition declares; an enqueue overrides — the same precedence + // resolveResources applies to ResourceFunc and ResourceClass. The + // definition's TTL has to come from the registry because EnqueueRaw + // has only a name and a payload; without this lookup a definition + // declaring job.WithLeaseTTL would silently get the pool default, + // which is the whole point of a per-job lease TTL. + // + // Only a positive enqueue-site value overrides, because zero is not + // an override: it is the absence of one, and means "use the pool + // default" rather than "cancel what the definition declared". + j.LeaseTTL = eng.registry.LeaseTTL(name) + if jobOpts.LeaseTTL > 0 { + j.LeaseTTL = jobOpts.LeaseTTL + } if !jobOpts.RunAt.IsZero() { j.RunAt = jobOpts.RunAt } + if err := eng.applyBindings(ctx, j, jobOpts.Bindings); err != nil { + return nil, err + } + + // After applyBindings, so the bindings this reads are already + // validated; before EnqueueJob, so an unschedulable job never + // reaches the store. + if err := eng.resolveResources(ctx, j, jobOpts); err != nil { + return nil, err + } + if err := eng.jobStore.EnqueueJob(ctx, j); err != nil { return nil, err } @@ -493,7 +773,36 @@ func (eng *Engine) Stop(ctx context.Context) error { eng.logger.Error("cron scheduler stop error", log.String("error", err.Error())) } - return eng.d.Stop(ctx) + stopErr := eng.d.Stop(ctx) + + // Close the executors last. The dispatcher stop above drains the worker + // pool, so no attempt is still running through a rung when its resources + // go away. In-process Close is a no-op; an out-of-process rung releases + // its clients and child processes here or leaks them. + // + // Guarded by stopOnce: a second Stop call must not close every executor + // again, since Close is newly reachable here and, unlike the rest of + // this method, is not itself idempotent. + eng.stopOnce.Do(eng.closeExecutors) + + return stopErr +} + +// closeExecutors releases every configured executor's resources, logging +// failures rather than propagating them: shutdown continues regardless. +func (eng *Engine) closeExecutors() { + if eng.executors == nil { + return + } + + for _, e := range eng.executors.Executors() { + if err := e.Close(); err != nil { + eng.logger.Warn("executor close failed", + log.String("executor", e.Name()), + log.String("error", err.Error()), + ) + } + } } // Extensions returns the extension registry. diff --git a/engine/execution.go b/engine/execution.go new file mode 100644 index 0000000..6cdab7d --- /dev/null +++ b/engine/execution.go @@ -0,0 +1,95 @@ +package engine + +import ( + "fmt" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/inproc" + "github.com/xraph/dispatch/job" +) + +// WithExecutor registers an additional executor, making a stronger +// isolation level available to job definitions that ask for it. +// +// The in-process executor is always present as the default, so a +// deployment that adds nothing behaves exactly as it always has. +func WithExecutor(e exec.Executor) Option { + return func(eng *Engine) { + eng.extraExecutors = append(eng.extraExecutors, e) + } +} + +// WithScratchRoot sets the root directory an out-of-process attempt's +// scratch OutputDir is created under (worker.Runner.WithArtifacts). +// +// It only has an effect once the artifact plane is also configured +// (WithArtifacts): Build only ever passes this value to the Runner +// alongside eng.artifacts (see Build's own runner.WithArtifacts call), +// so with no artifact plane configured this sets a value Build never +// reads and the Runner's scratch directories fall back to os.TempDir() +// regardless. That is not because scratch-dir creation or its startup +// sweep are themselves gated on the artifact plane — they are not, and +// both run for any out-of-process attempt either way (see +// worker.Runner.prepareOutputDir and Reclaim) — it is purely this +// option's own value never reaching the Runner. The extension package's +// resolveExecutionOptions logs a startup warning for exactly this +// configuration (a configured scratch_dir with the artifact plane off), +// so that stays a deliberate, visible choice rather than a silent one. +// +// Leaving it unset defaults to os.TempDir(), exactly as worker.Runner +// does on its own. +func WithScratchRoot(dir string) Option { + return func(eng *Engine) { eng.scratchRoot = dir } +} + +// Executors returns the configured executor registry. +func (eng *Engine) Executors() *exec.Registry { return eng.executors } + +// buildExecutors assembles the executor registry. It is called once during +// engine construction, before any definition is registered, because +// registration validates policies against it. +func (eng *Engine) buildExecutors() { + r := exec.NewRegistry(inproc.New(eng.registry)) + for _, e := range eng.extraExecutors { + r.Add(e) + } + eng.executors = r +} + +// checkExecutionPolicy reports whether the deployment can satisfy a +// definition's declared isolation. +// +// This runs at registration rather than at execution deliberately. A +// definition that can never be satisfied should fail on a developer's +// machine, not on the first malicious upload in production. +func (eng *Engine) checkExecutionPolicy(name string, p exec.Policy) error { + if eng.executors == nil { + return nil + } + if _, err := eng.executors.Select(p); err != nil { + return fmt.Errorf("dispatch/engine: job %q: %w", name, err) + } + + return nil +} + +// RegisterAll registers a set of definitions. +// +// It takes job.Registrable rather than a typed definition so a single +// handler list can be shared between the worker and an out-of-process +// entrypoint, which cannot be handed an engine. +func RegisterAll(eng *Engine, defs ...job.Registrable) error { + // Validate every definition before registering any of them, so a + // rejected set leaves the registry as it was rather than half + // populated. + for _, d := range defs { + if err := eng.checkExecutionPolicy(d.JobName(), d.Policy()); err != nil { + return err + } + } + for _, d := range defs { + d.Register(eng.registry) + } + + return nil +} diff --git a/engine/execution_e2e_test.go b/engine/execution_e2e_test.go new file mode 100644 index 0000000..291d701 --- /dev/null +++ b/engine/execution_e2e_test.go @@ -0,0 +1,267 @@ +package engine_test + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" +) + +// These tests drive jobs through the pool and runner engine.Build actually +// assembles, rather than through a hand-built Runner. That seam is where +// ErrPermanent quietly stopped working: the Runner-level tests covered the +// legacy nil-registry path, which is the one path engine users never take. + +// countingExecutor is a stand-in for a stronger rung: it satisfies +// LevelProcess without providing any real isolation, so a test can prove a +// job was routed to it. +type countingExecutor struct { + mu sync.Mutex + names []string + + reclaimed int + closed int +} + +func (e *countingExecutor) Name() string { return "counting" } +func (e *countingExecutor) Level() exec.Level { return exec.LevelProcess } + +func (e *countingExecutor) Run(_ context.Context, req *exec.Request) (*exec.Result, error) { + e.mu.Lock() + defer e.mu.Unlock() + e.names = append(e.names, req.Name) + + return &exec.Result{Status: exec.StatusOK}, nil +} + +func (e *countingExecutor) Reclaim(context.Context, id.WorkerID) error { + e.mu.Lock() + defer e.mu.Unlock() + e.reclaimed++ + + return nil +} + +func (e *countingExecutor) Close() error { + e.mu.Lock() + defer e.mu.Unlock() + e.closed++ + + return nil +} + +func (e *countingExecutor) ran(name string) bool { + e.mu.Lock() + defer e.mu.Unlock() + for _, n := range e.names { + if n == name { + return true + } + } + + return false +} + +func (e *countingExecutor) counts() (reclaimed, closed int) { + e.mu.Lock() + defer e.mu.Unlock() + + return e.reclaimed, e.closed +} + +// startEngine builds an engine over a fresh in-memory store, starts it, and +// stops it when the test ends. +func startEngine(t *testing.T, opts ...engine.Option) (*engine.Engine, *memory.Store) { + t.Helper() + + s := memory.New() + d, err := dispatch.New(dispatch.WithStore(s), dispatch.WithConcurrency(2)) + if err != nil { + t.Fatalf("dispatch.New: %v", err) + } + eng, err := engine.Build(d, opts...) + if err != nil { + t.Fatalf("engine.Build: %v", err) + } + + return eng, s +} + +// waitForJob polls the store until the job satisfies cond, or fails the test. +func waitForJob(t *testing.T, s *memory.Store, jobID id.JobID, what string, cond func(*job.Job) bool) *job.Job { + t.Helper() + + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + got, err := s.GetJob(context.Background(), jobID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + if cond(got) { + return got + } + time.Sleep(10 * time.Millisecond) + } + + got, err := s.GetJob(context.Background(), jobID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + t.Fatalf("timed out waiting for %s; job state = %q, retry_count = %d, last_error = %q", + what, got.State, got.RetryCount, got.LastError) + + return nil +} + +func TestEngine_PermanentFailureReachesDLQOnTheFirstAttempt(t *testing.T) { + // engine.Build always wires a non-nil executor registry, so this is the + // path every engine user takes. A handler declining a retry must still + // be able to. + eng, s := startEngine(t) + + var attempts int + var mu sync.Mutex + engine.Register(eng, job.NewDefinition("permanent.job", + func(context.Context, execPayload) error { + mu.Lock() + attempts++ + mu.Unlock() + + return fmt.Errorf("malformed payload: %w", dispatch.ErrPermanent) + })) + + j, err := engine.Enqueue(context.Background(), eng, "permanent.job", execPayload{Value: 1}) + if err != nil { + t.Fatalf("Enqueue: %v", err) + } + + if startErr := eng.Start(context.Background()); startErr != nil { + t.Fatalf("Start: %v", startErr) + } + t.Cleanup(func() { _ = eng.Stop(context.Background()) }) + + got := waitForJob(t, s, j.ID, "the job to fail", func(g *job.Job) bool { + return g.State == job.StateFailed + }) + + if got.RetryCount != 1 { + t.Errorf("RetryCount = %d, want 1 — the backoff schedule must not be spent", got.RetryCount) + } + if want := "malformed payload: dispatch: permanent failure"; got.LastError != want { + t.Errorf("LastError = %q, want %q", got.LastError, want) + } + + mu.Lock() + ran := attempts + mu.Unlock() + if ran != 1 { + t.Errorf("handler ran %d times, want 1", ran) + } + + // Give the pool a moment to prove it is not still retrying behind us. + time.Sleep(100 * time.Millisecond) + mu.Lock() + ran = attempts + mu.Unlock() + if ran != 1 { + t.Errorf("handler ran %d times after the job failed, want 1", ran) + } +} + +func TestEngine_OrdinaryFailureStillRetries(t *testing.T) { + // The other half of the same claim: an error that is not permanent must + // keep its retry schedule, so the fix above did not turn every failure + // into a dead letter. + eng, s := startEngine(t) + + engine.Register(eng, job.NewDefinition("retrying.job", + func(context.Context, execPayload) error { return errors.New("transient") })) + + j, err := engine.Enqueue(context.Background(), eng, "retrying.job", execPayload{Value: 1}) + if err != nil { + t.Fatalf("Enqueue: %v", err) + } + + if startErr := eng.Start(context.Background()); startErr != nil { + t.Fatalf("Start: %v", startErr) + } + t.Cleanup(func() { _ = eng.Stop(context.Background()) }) + + got := waitForJob(t, s, j.ID, "the job to be scheduled for retry", func(g *job.Job) bool { + return g.State == job.StateRetrying + }) + if got.RetryCount != 1 { + t.Errorf("RetryCount = %d, want 1", got.RetryCount) + } + if got.LastError != "transient" { + t.Errorf("LastError = %q, want %q", got.LastError, "transient") + } +} + +func TestEngine_JobIsDispatchedToTheAddedExecutor(t *testing.T) { + // Build + WithExecutor + a policy the added rung satisfies, end to end: + // the job must actually run there, not merely be registrable. + rung := &countingExecutor{} + eng, s := startEngine(t, engine.WithExecutor(rung)) + + handlerRan := make(chan struct{}, 1) + if regErr := engine.RegisterChecked(eng, job.NewDefinition("isolated.job", + func(context.Context, execPayload) error { + handlerRan <- struct{}{} + + return nil + }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + )); regErr != nil { + t.Fatalf("RegisterChecked: %v", regErr) + } + + j, err := engine.Enqueue(context.Background(), eng, "isolated.job", execPayload{Value: 7}) + if err != nil { + t.Fatalf("Enqueue: %v", err) + } + + if startErr := eng.Start(context.Background()); startErr != nil { + t.Fatalf("Start: %v", startErr) + } + + waitForJob(t, s, j.ID, "the job to complete", func(g *job.Job) bool { + return g.State == job.StateCompleted + }) + + if !rung.ran("isolated.job") { + t.Error("the job was not dispatched to the executor its policy required") + } + // This rung never calls the handler, so a job that reached it cannot + // also have run in process. + select { + case <-handlerRan: + t.Error("the handler ran in process even though the job was routed to another executor") + default: + } + + // The pool sweeps for leaked sandboxes at startup and the engine closes + // every rung when it stops. Both had no caller before. + // + // The startup sweep runs in a background goroutine (see + // Pool.runReclaimSweep) that Stop's own p.wg.Wait() joins, so reclaimed + // is only guaranteed final once Stop has returned — reading it earlier + // races the sweep goroutine against the test goroutine. + if stopErr := eng.Stop(context.Background()); stopErr != nil { + t.Fatalf("Stop: %v", stopErr) + } + if reclaimed, closed := rung.counts(); reclaimed != 1 { + t.Errorf("Reclaim called %d times at pool start, want 1", reclaimed) + } else if closed != 1 { + t.Errorf("Close called %d times at engine stop, want 1", closed) + } +} diff --git a/engine/execution_subprocess_test.go b/engine/execution_subprocess_test.go new file mode 100644 index 0000000..5406fa2 --- /dev/null +++ b/engine/execution_subprocess_test.go @@ -0,0 +1,91 @@ +package engine_test + +import ( + "context" + "errors" + "testing" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/subprocess" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" +) + +// TestEngine_SubprocessRungSatisfiesLevelProcess proves the no-silent- +// downgrade rule end to end through the mechanism configuration actually +// drives: engine.WithExecutor. A definition declaring +// exec.Isolate(exec.LevelProcess) must fail registration when nothing +// configured can provide it, and succeed once the subprocess rung is +// registered — exactly what extension.resolveExecutionOptions wires up +// from the "execution.subprocess" YAML block, without engine importing +// extension (that would be a cycle) or extension needing a Forge app just +// to prove this. +func TestEngine_SubprocessRungSatisfiesLevelProcess(t *testing.T) { + newEngine := func(t *testing.T, opts ...engine.Option) *engine.Engine { + t.Helper() + + d, err := dispatch.New(dispatch.WithStore(memory.New())) + if err != nil { + t.Fatalf("dispatch.New: %v", err) + } + eng, err := engine.Build(d, opts...) + if err != nil { + t.Fatalf("engine.Build: %v", err) + } + + return eng + } + + t.Run("no subprocess rung configured fails registration", func(t *testing.T) { + eng := newEngine(t) + + err := engine.RegisterChecked(eng, job.NewDefinition("needs.process", + func(context.Context, execSubprocessPayload) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + )) + if !errors.Is(err, exec.ErrNoExecutor) { + t.Fatalf("RegisterChecked() = %v, want %v", err, exec.ErrNoExecutor) + } + }) + + t.Run("subprocess rung configured satisfies the policy", func(t *testing.T) { + eng := newEngine(t, engine.WithExecutor(subprocess.New())) + + err := engine.RegisterChecked(eng, job.NewDefinition("needs.process", + func(context.Context, execSubprocessPayload) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + )) + if err != nil { + t.Fatalf("RegisterChecked() = %v, want nil", err) + } + + executors := eng.Executors() + selected, selectErr := executors.Select(exec.NewPolicy(exec.Isolate(exec.LevelProcess))) + if selectErr != nil { + t.Fatalf("Select() = %v, want nil", selectErr) + } + if selected.Name() != subprocess.Name { + t.Errorf("Select().Name() = %q, want %q", selected.Name(), subprocess.Name) + } + }) + + t.Run("a job declaring no isolation still runs in-process", func(t *testing.T) { + // Configuring the subprocess rung must not change the DEFAULT: a + // definition that declares nothing still resolves to the + // in-process executor, exactly as it does with no execution + // config at all. + eng := newEngine(t, engine.WithExecutor(subprocess.New())) + + selected, err := eng.Executors().Select(exec.NewPolicy()) + if err != nil { + t.Fatalf("Select() = %v, want nil", err) + } + if selected.Name() != "inprocess" { + t.Errorf("Select().Name() = %q, want %q — the default must not change", selected.Name(), "inprocess") + } + }) +} + +type execSubprocessPayload struct{} diff --git a/engine/execution_test.go b/engine/execution_test.go new file mode 100644 index 0000000..2e214b4 --- /dev/null +++ b/engine/execution_test.go @@ -0,0 +1,130 @@ +package engine_test + +import ( + "context" + "errors" + "testing" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" +) + +type execPayload struct { + Value int `json:"value"` +} + +func newTestEngine(t *testing.T) *engine.Engine { + t.Helper() + + d, err := dispatch.New(dispatch.WithStore(memory.New())) + if err != nil { + t.Fatalf("dispatch.New: %v", err) + } + eng, err := engine.Build(d) + if err != nil { + t.Fatalf("engine.Build: %v", err) + } + + return eng +} + +func TestEngine_ExecutorsIncludesInProcessByDefault(t *testing.T) { + eng := newTestEngine(t) + + executors := eng.Executors() + if executors == nil { + t.Fatal("Executors() = nil, want a registry") + } + def := executors.Default() + if def == nil { + t.Fatal("Default() = nil, want the in-process executor") + } + if def.Name() != "inprocess" { + t.Errorf("Default().Name() = %q, want %q", def.Name(), "inprocess") + } +} + +func TestEngine_RegisterRejectsUnsatisfiablePolicy(t *testing.T) { + // A definition that must be isolated must not silently run + // unisolated because it was deployed somewhere that cannot isolate. + eng := newTestEngine(t) + + err := engine.RegisterChecked(eng, job.NewDefinition("needs.sandbox", + func(context.Context, execPayload) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelSandboxed)), + )) + if !errors.Is(err, exec.ErrNoExecutor) { + t.Fatalf("RegisterChecked() = %v, want %v", err, exec.ErrNoExecutor) + } +} + +func TestEngine_RegisterCheckedAllowsExplicitDowngrade(t *testing.T) { + eng := newTestEngine(t) + + err := engine.RegisterChecked(eng, job.NewDefinition("needs.sandbox.but.ok", + func(context.Context, execPayload) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelSandboxed), exec.AllowDowngrade()), + )) + if err != nil { + t.Fatalf("RegisterChecked() = %v, want nil", err) + } +} + +func TestEngine_RegisterStaysUnchecked(t *testing.T) { + // Register is the unchecked path by existing convention, and its + // signature must not change. A policy nothing satisfies is caught by + // RegisterChecked and by RegisterAll, not here. + eng := newTestEngine(t) + + engine.Register(eng, job.NewDefinition("unchecked.sandbox", + func(context.Context, execPayload) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelSandboxed)), + )) + + if _, ok := eng.Registry().Get("unchecked.sandbox"); !ok { + t.Error("Register did not register the handler") + } +} + +func TestEngine_RegisterAll(t *testing.T) { + eng := newTestEngine(t) + + defs := []job.Registrable{ + job.NewDefinition("a.job", func(context.Context, execPayload) error { return nil }), + job.NewDefinition("b.job", func(context.Context, struct{}) error { return nil }), + } + + if err := engine.RegisterAll(eng, defs...); err != nil { + t.Fatalf("RegisterAll() = %v, want nil", err) + } + for _, name := range []string{"a.job", "b.job"} { + if _, ok := eng.Registry().Get(name); !ok { + t.Errorf("handler %q not registered", name) + } + } +} + +func TestEngine_RegisterAllRejectsWholeSetOnOneFailure(t *testing.T) { + // RegisterAll validates every definition before registering any of + // them, so a rejected set leaves the registry as it was rather than + // half populated. + eng := newTestEngine(t) + + defs := []job.Registrable{ + job.NewDefinition("good.job", func(context.Context, execPayload) error { return nil }), + job.NewDefinition("bad.job", + func(context.Context, execPayload) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelSandboxed)), + ), + } + + if err := engine.RegisterAll(eng, defs...); !errors.Is(err, exec.ErrNoExecutor) { + t.Fatalf("RegisterAll() = %v, want %v", err, exec.ErrNoExecutor) + } + if _, ok := eng.Registry().Get("good.job"); ok { + t.Error("good.job was registered even though the set was rejected") + } +} diff --git a/engine/execution_wiring_test.go b/engine/execution_wiring_test.go new file mode 100644 index 0000000..57ce571 --- /dev/null +++ b/engine/execution_wiring_test.go @@ -0,0 +1,162 @@ +package engine_test + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" +) + +// scratchWritingExecutor is a minimal exec.LevelProcess executor. It does +// not launch a real OS subprocess — that round trip needs a re-exec'd +// shim binary and is already covered by exec/subprocess's own tests and +// exec/exectest's shared conformance suite — but it is a genuine +// exec.LevelProcess implementation, so it exercises the exact seam +// engine.Build's wiring controls: Runner only hands a rung above +// exec.LevelNone a real req.OutputDir, and only commits what lands there, +// once a *worker.Runner has an artifact plane (see runner.WithArtifacts). +type scratchWritingExecutor struct{} + +func (scratchWritingExecutor) Name() string { return "scratch-writer" } +func (scratchWritingExecutor) Level() exec.Level { return exec.LevelProcess } + +func (scratchWritingExecutor) Run(_ context.Context, req *exec.Request) (*exec.Result, error) { + if req.OutputDir == "" { + return nil, errors.New("no OutputDir given — the executor was not wired for out-of-process output committing") + } + + if err := os.WriteFile(filepath.Join(req.OutputDir, "result.txt"), []byte("wired"), 0o600); err != nil { + return nil, err + } + + return &exec.Result{Status: exec.StatusOK}, nil +} + +func (scratchWritingExecutor) Reclaim(context.Context, id.WorkerID) error { return nil } +func (scratchWritingExecutor) Close() error { return nil } + +// TestEngineBuild_WiresArtifactsIntoRunner is the test the review found +// missing: engine.Build's +// +// if eng.artifacts != nil { +// runner.WithArtifacts(eng.artifacts, eng.scratchRoot) +// } +// +// (engine/engine.go) is what gives an out-of-process rung its scratch +// directory, its PriorOutputs, and its output committing at all — it is +// the entire point of the commit this task exists to cover. Mutating +// that guard to `if eng.artifacts != nil && false` left the whole suite +// green before this test existed, including TestEndToEndStageAndCommit: +// that test's job declares no isolation (exec.LevelNone), so its +// staging/commit path runs entirely through the artifact-input staging +// middleware, never through runner.WithArtifacts at all. This test +// specifically declares exec.LevelProcess so it actually exercises that +// call, and fails without it. +// +// What this proves: a LevelProcess executor is handed a real, writable +// OutputDir by the Runner, and whatever regular file it leaves there is +// committed through the artifact plane and visible via +// ListArtifactsByOwner — the same observable TestEndToEndStageAndCommit +// checks for the in-process path. +// +// What this does NOT prove: it does not exercise exec/subprocess's own +// re-exec'd OS subprocess, its wire protocol, or its uid/rlimit +// enforcement — those are exec/subprocess's own tests' job. This test's +// fake executor stands in for "any exec.LevelProcess executor", which is +// the right level for pinning engine.Build's own wiring, since that +// wiring does not know or care which out-of-process executor is +// registered. +func TestEngineBuild_WiresArtifactsIntoRunner(t *testing.T) { + ctx := context.Background() + + s := memory.New() + backend := artifacttest.NewBackend() + svc := artifact.NewService(s, backend, + artifact.WithEphemeralPrefix("ephemeral"), + artifact.WithDefaultBucket("dispatch")) + + c, err := cache.New(t.TempDir(), backend, cache.WithBudget(1<<20)) + if err != nil { + t.Fatalf("cache.New: %v", err) + } + t.Cleanup(func() { + if cerr := c.Close(); cerr != nil { + t.Errorf("cache close: %v", cerr) + } + }) + + d, err := dispatch.New( + dispatch.WithStore(s), + dispatch.WithConcurrency(1), + dispatch.WithQueues([]string{"default"}), + ) + if err != nil { + t.Fatalf("dispatch.New: %v", err) + } + + eng, err := engine.Build(d, + engine.WithArtifacts(svc, c), + engine.WithExecutor(scratchWritingExecutor{}), + ) + if err != nil { + t.Fatalf("engine.Build: %v", err) + } + + def := job.NewDefinition("out-of-process", + func(context.Context, tessellateInput) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ) + + if err = engine.RegisterChecked(eng, def); err != nil { + t.Fatalf("RegisterChecked: %v", err) + } + + j, err := engine.Enqueue(ctx, eng, "out-of-process", tessellateInput{Detail: 0.5}) + if err != nil { + t.Fatalf("Enqueue: %v", err) + } + + if serr := eng.Start(ctx); serr != nil { + t.Fatalf("engine.Start: %v", serr) + } + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if serr := eng.Stop(stopCtx); serr != nil { + t.Errorf("engine.Stop: %v", serr) + } + }) + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: j.ID.String()} + + waitFor(t, 5*time.Second, func() bool { + outputs, listErr := s.ListArtifactsByOwner(ctx, owner, artifact.RoleOutput) + return listErr == nil && len(outputs) == 1 + }) + + outputs, err := s.ListArtifactsByOwner(ctx, owner, artifact.RoleOutput) + if err != nil { + t.Fatalf("ListArtifactsByOwner: %v", err) + } + if len(outputs) != 1 { + t.Fatalf("got %d committed outputs, want 1 — without engine.Build wiring runner.WithArtifacts, "+ + "the executor's output is discarded along with its scratch directory instead of committed", + len(outputs)) + } + if outputs[0].Size != int64(len("wired")) { + t.Errorf("output size = %d, want %d", outputs[0].Size, len("wired")) + } +} diff --git a/engine/export_test.go b/engine/export_test.go new file mode 100644 index 0000000..1efeaac --- /dev/null +++ b/engine/export_test.go @@ -0,0 +1,18 @@ +package engine + +import ( + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/resource" +) + +// InputSizesForTest exposes inputSizes to the external test package. +func InputSizesForTest(b map[string]artifact.Ref) ([]resource.InputSize, int64, string) { + return inputSizes(b) +} + +// CheckReaperMarginForTest exposes checkReaperMargin to the external test +// package. +func CheckReaperMarginForTest(cfg dispatch.Config, leaseAware bool) error { + return checkReaperMargin(cfg, leaseAware) +} diff --git a/engine/lease_test.go b/engine/lease_test.go new file mode 100644 index 0000000..4d5c163 --- /dev/null +++ b/engine/lease_test.go @@ -0,0 +1,148 @@ +package engine_test + +import ( + "context" + "testing" + "time" + + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/job" +) + +// TestEnqueueRaw_CarriesLeaseTTL verifies that EnqueueRaw copies the declared +// LeaseTTL from job options to both the returned job and the persisted store row. +func TestEnqueueRaw_CarriesLeaseTTL(t *testing.T) { + tests := []struct { + name string + leaseTTL time.Duration + }{ + { + name: "zero LeaseTTL", + leaseTTL: 0, + }, + { + name: "30 second LeaseTTL", + leaseTTL: 30 * time.Second, + }, + { + name: "6 hour LeaseTTL", + leaseTTL: 6 * time.Hour, + }, + { + name: "1 minute LeaseTTL", + leaseTTL: 1 * time.Minute, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + eng, store, _ := newWorkflowEngine(t) + + // Enqueue a job with the specified LeaseTTL. + returned, err := eng.EnqueueRaw( + context.Background(), + "test-job", + []byte(`{}`), + job.WithLeaseTTL(tt.leaseTTL), + ) + if err != nil { + t.Fatalf("EnqueueRaw: %v", err) + } + + // Assert the returned job carries the declared LeaseTTL. + if returned.LeaseTTL != tt.leaseTTL { + t.Errorf("returned job LeaseTTL = %v, want %v", returned.LeaseTTL, tt.leaseTTL) + } + + // Assert the persisted row in the store also has the LeaseTTL. + persisted, err := store.GetJob(context.Background(), returned.ID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + if persisted.LeaseTTL != tt.leaseTTL { + t.Errorf("persisted job LeaseTTL = %v, want %v", persisted.LeaseTTL, tt.leaseTTL) + } + }) + } +} + +// TestEnqueueRaw_DefinitionLeaseTTL covers the other half of the +// precedence chain: a definition declares, an enqueue overrides. +// +// The definition half was unreachable before this test existed. +// job.WithLeaseTTL set Options.LeaseTTL, RegisterDefinition never read +// it, and EnqueueRaw built its options from the enqueue site alone — so a +// definition asking for a six-hour lease silently got the pool's default +// and its jobs were reclaimed mid-run. Per-definition TTLs are the point +// of putting lease_ttl on the row at all, so this asserts the value +// reaches the persisted row, not merely the returned struct. +func TestEnqueueRaw_DefinitionLeaseTTL(t *testing.T) { + const ( + declared = 6 * time.Hour + override = 90 * time.Second + ) + + tests := []struct { + name string + // defTTL is what the definition declares; zero declares nothing. + defTTL time.Duration + // enqueueOpts are the options passed to the enqueue call. + enqueueOpts []job.Option + want time.Duration + }{ + { + name: "definition declaration reaches the row", + defTTL: declared, + want: declared, + }, + { + name: "enqueue overrides the definition", + defTTL: declared, + enqueueOpts: []job.Option{job.WithLeaseTTL(override)}, + want: override, + }, + { + name: "neither declares, so the pool default applies", + want: 0, + }, + { + name: "enqueue alone still works with no declaration", + enqueueOpts: []job.Option{job.WithLeaseTTL(override)}, + want: override, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + eng, store, _ := newWorkflowEngine(t) + + var defOpts []job.Option + if tt.defTTL > 0 { + defOpts = append(defOpts, job.WithLeaseTTL(tt.defTTL)) + } + + engine.Register(eng, job.NewDefinition("leased-job", + func(_ context.Context, _ struct{}) error { return nil }, + defOpts...)) + + returned, err := eng.EnqueueRaw( + context.Background(), "leased-job", []byte(`{}`), tt.enqueueOpts...) + if err != nil { + t.Fatalf("EnqueueRaw: %v", err) + } + + if returned.LeaseTTL != tt.want { + t.Errorf("returned job LeaseTTL = %v, want %v", returned.LeaseTTL, tt.want) + } + + persisted, err := store.GetJob(context.Background(), returned.ID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + + if persisted.LeaseTTL != tt.want { + t.Errorf("persisted job LeaseTTL = %v, want %v", persisted.LeaseTTL, tt.want) + } + }) + } +} diff --git a/engine/reaper_margin_test.go b/engine/reaper_margin_test.go new file mode 100644 index 0000000..78ef1ad --- /dev/null +++ b/engine/reaper_margin_test.go @@ -0,0 +1,243 @@ +package engine_test + +import ( + "errors" + "strings" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/store/memory" +) + +// TestReaperMarginAcceptsStockConfig is the compatibility floor. Turning +// the resource model on must not force anybody to retune their timings, +// so the shipped defaults have to clear the check with room to spare. +func TestReaperMarginAcceptsStockConfig(t *testing.T) { + if err := engine.CheckReaperMarginForTest(dispatch.DefaultConfig(), true); err != nil { + t.Fatalf("the default configuration must pass: %v", err) + } +} + +func TestReaperMargin(t *testing.T) { + cases := []struct { + name string + poll time.Duration + beat time.Duration + stale time.Duration + leaseTTL time.Duration + leaseAware bool + ok bool + }{ + { + // The failure this check exists for: a threshold at or below + // the poll interval lets the reaper reclaim a job the fetcher + // is still holding, and the job then runs twice. + name: "threshold at the poll interval", poll: 30 * time.Second, + beat: 10 * time.Second, stale: 30 * time.Second, leaseAware: true, ok: false, + }, + { + // Larger than the claim-to-first-heartbeat window, but with no + // room for the heartbeat write itself to be slow. + name: "no slack for a missed heartbeat", poll: 5 * time.Second, + beat: 10 * time.Second, stale: 20 * time.Second, leaseAware: true, ok: false, + }, + { + name: "exactly twice the window", poll: 5 * time.Second, + beat: 10 * time.Second, stale: 30 * time.Second, leaseAware: true, ok: true, + }, + { + // A reaper that is switched off cannot reclaim anything, so + // there is no relationship left to police. + name: "reaper disabled", poll: time.Hour, + beat: time.Hour, stale: 0, leaseAware: true, ok: true, + }, + { + // Heartbeats off: the window is the admission stall alone. + // Whether a never-heartbeating job survives its threshold is a + // question that predates this model and is not ours. + name: "heartbeats disabled", poll: time.Second, + beat: 0, stale: 3 * time.Second, leaseAware: true, ok: true, + }, + { + name: "heartbeats disabled and threshold too tight", poll: 10 * time.Second, + beat: 0, stale: 10 * time.Second, leaseAware: true, ok: false, + }, + { + // The bug this fix closes, reproduced exactly: a generous + // StaleJobThreshold clears the old check, but on a lease-aware + // backend reclamation is actually governed by DefaultLeaseTTL, + // and 12s does not clear a 20s minimum. + name: "DefaultLeaseTTL below the margin beats a generous threshold", + poll: 5 * time.Second, beat: 5 * time.Second, + stale: 5 * time.Minute, leaseTTL: 12 * time.Second, + leaseAware: true, ok: false, + }, + { + // Same numbers, but a DefaultLeaseTTL that does clear the + // margin. StaleJobThreshold being generous is irrelevant either + // way once a backend is lease-aware. + name: "DefaultLeaseTTL at the margin passes", + poll: 5 * time.Second, beat: 5 * time.Second, + stale: 5 * time.Minute, leaseTTL: 20 * time.Second, + leaseAware: true, ok: true, + }, + { + // The mirror image: on a backend that is NOT lease-aware, + // reclamation runs through reapStaleJobsLegacy, which reads + // StaleJobThreshold directly and never looks at + // DefaultLeaseTTL. A low DefaultLeaseTTL here must not matter. + name: "non-lease-aware backend is judged on the threshold alone", + poll: 5 * time.Second, beat: 5 * time.Second, + stale: 30 * time.Second, leaseTTL: time.Second, + leaseAware: false, ok: true, + }, + { + name: "non-lease-aware backend still fails on a tight threshold", + poll: 5 * time.Second, beat: 5 * time.Second, + stale: 5 * time.Second, leaseAware: false, ok: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + err := engine.CheckReaperMarginForTest(dispatch.Config{ + PollInterval: tc.poll, + HeartbeatInterval: tc.beat, + StaleJobThreshold: tc.stale, + DefaultLeaseTTL: tc.leaseTTL, + }, tc.leaseAware) + + if tc.ok { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + return + } + + if err == nil { + t.Fatal("expected an error") + } + + // The message has to name every value involved, because the + // fix is a relationship between them and an operator reading + // it should not have to go looking for the others. + for _, want := range []string{ + "StaleJobThreshold", "PollInterval", "HeartbeatInterval", "DefaultLeaseTTL", + tc.poll.String(), + } { + if !strings.Contains(err.Error(), want) { + t.Errorf("error does not mention %q: %v", want, err) + } + } + }) + } +} + +// TestBuildRejectsUnsafeReaperMargin checks the validation actually runs +// at construction, and only when a manager makes the stall possible. +func TestBuildRejectsUnsafeReaperMargin(t *testing.T) { + newDispatcher := func(t *testing.T) *dispatch.Dispatcher { + t.Helper() + + d, err := dispatch.New( + dispatch.WithStore(memory.New()), + dispatch.WithPollInterval(20*time.Second), + dispatch.WithHeartbeatInterval(10*time.Second), + dispatch.WithStaleJobThreshold(30*time.Second), + ) + if err != nil { + t.Fatalf("dispatch.New: %v", err) + } + + return d + } + + mgr := resource.NewManager(resource.Set{resource.Memory: 1 << 30}) + + if _, err := engine.Build(newDispatcher(t), engine.WithResourceManager(mgr)); err == nil { + t.Fatal("Build accepted a configuration the reaper can corrupt") + } + + // The same timings without a manager are none of this check's + // business: with no ledger, admission never waits. + if _, err := engine.Build(newDispatcher(t)); err != nil { + t.Fatalf("Build without a manager must not be validated: %v", err) + } +} + +// TestBuildDoesNotDeriveCapacityFromTheManager pins the ruling that +// replaced the old "publish what admission enforces" convenience. +// +// The convenience looked like it kept two numbers from drifting. What it +// actually did was answer a FLEET question with a LOCAL number: the +// unschedulable check asks "is any worker big enough", and seeding its +// floor from this process's ledger silently rescoped that to "is THIS +// process big enough". Nothing could raise it back afterwards, because +// cluster.Worker.Capacity round-trips on store/memory alone — postgres, +// sqlite, mongo and redis all enumerate worker fields by hand and drop +// it. On those four the check collapsed entirely to the local manager, +// so a light API pod with resources enabled hard-rejected at enqueue the +// tessellation job the heavy tier runs perfectly well. +// +// A manager is now just a manager. The check is inert until an operator +// states the fleet ceiling out loud. +func TestBuildDoesNotDeriveCapacityFromTheManager(t *testing.T) { + capacity := resource.Set{resource.Memory: 8 << 30, "fpga": 2} + + d, err := dispatch.New(dispatch.WithStore(memory.New())) + if err != nil { + t.Fatalf("dispatch.New: %v", err) + } + + eng, err := engine.Build(d, engine.WithResourceManager(resource.NewManager(capacity))) + if err != nil { + t.Fatalf("Build: %v", err) + } + + if got := eng.MaxWorkerCapacity(t.Context()); len(got) != 0 { + t.Errorf("MaxWorkerCapacity = %v, want empty: installing a ledger must not turn the "+ + "fleet-wide unschedulable check on with this process's own numbers", got) + } + + // And with it off, a job larger than this process still enqueues — + // some other worker may be able to run it. + huge := resource.Set{resource.Memory: 64 << 30} + + j, err := eng.EnqueueRaw(t.Context(), "huge", []byte(`{}`), job.WithResources(huge)) + if err != nil { + t.Fatalf("EnqueueRaw of a job larger than this worker: %v", err) + } + + if j.Resources[resource.Memory] != 64<<30 { + t.Errorf("stored requirement = %v, want the declaration intact", j.Resources) + } + + // An explicit declaration is what turns the check on, and it is then + // enforced. + d2, err := dispatch.New(dispatch.WithStore(memory.New())) + if err != nil { + t.Fatalf("dispatch.New: %v", err) + } + + eng2, err := engine.Build(d2, + engine.WithResourceManager(resource.NewManager(capacity)), + engine.WithWorkerCapacity(resource.Set{resource.Memory: 1 << 30}), + ) + if err != nil { + t.Fatalf("Build: %v", err) + } + + if got := eng2.MaxWorkerCapacity(t.Context())[resource.Memory]; got != 1<<30 { + t.Errorf("declared capacity was overwritten: memory = %d", got) + } + + if _, err = eng2.EnqueueRaw(t.Context(), "huge", []byte(`{}`), + job.WithResources(huge)); !errors.Is(err, resource.ErrUnschedulable) { + t.Errorf("EnqueueRaw error = %v, want ErrUnschedulable once a ceiling is declared", err) + } +} diff --git a/engine/resource.go b/engine/resource.go new file mode 100644 index 0000000..f374d61 --- /dev/null +++ b/engine/resource.go @@ -0,0 +1,350 @@ +package engine + +import ( + "context" + "fmt" + "sort" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/cluster" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// Resources returns the shared admission ledger, or nil when the +// resource model is off. It is the instance the staging cache must have +// been built with. +func (eng *Engine) Resources() resource.Manager { return eng.resources } + +// reaperSafetyFactor is the multiple of the claim-to-first-heartbeat +// window that StaleJobThreshold has to clear. +// +// Two, and the second one is not padding. The window itself — +// PollInterval + HeartbeatInterval — is what the fetcher can legitimately +// spend between claiming a job and that job's first heartbeat: admission +// may stall the batch for up to one poll interval while it reclaims disk, +// and the job then waits up to one heartbeat tick to be written down as +// alive. A threshold merely larger than that window leaves no room for +// the heartbeat itself to be slow, and the heartbeat is a store write on +// the same connection pool the dequeue just used. Doubling buys exactly +// one missed heartbeat round, which is the smallest slack that survives a +// briefly busy store. +// +// Below it the failure is not a stalled worker but a corrupted one: the +// reaper reclaims jobs the fetcher is still holding — rows already in +// running state, already claimed, not yet heartbeating — and the same job +// runs twice. +// +// The stock configuration clears it comfortably (1s + 10s, doubled, is +// 22s against a 30s threshold), so turning the resource model on does not +// force anybody to retune. What it catches is the combination that looks +// harmless: raising PollInterval to spare a shared database, or dropping +// StaleJobThreshold to fail over faster, without noticing the other. +const reaperSafetyFactor = 2 + +// checkReaperMargin rejects a configuration in which the stale-job +// reaper could reclaim a job the fetcher has claimed but not yet handed +// to a worker. +// +// It runs only when a resource manager is installed, because admission is +// what introduced the stall: Pool.admit calls Manager.Acquire under a +// one-poll-interval budget so it can evict cached disk to make room, and +// that budget is shared by the whole batch. Without a manager the fetcher +// never waits and the relationship does not exist. +// +// This fails Build rather than warning. The resource model is opt-in, so +// nobody arrives here by accident, and the symptom it prevents — a job +// executing twice because two subsystems disagreed about who owned it — +// is not one an operator can be expected to diagnose from a log line. +// +// leaseAware reports whether the store implements job.LeaseStore. It +// decides which number actually governs reclamation: on a lease-aware +// backend — every first-party one — worker/pool.go's reapStaleJobs routes +// to reclaimExpiredLeases, which reclaims purely on lease expiry +// (leaseTTLFor's chain: DefaultLeaseTTL, then StaleJobThreshold, then +// job.DefaultLeaseTTL) and never looks at StaleJobThreshold directly. +// Policing StaleJobThreshold alone on such a backend checks a number +// nothing reads: a deployment can set DefaultLeaseTTL to a few seconds, +// clear this check with a generous StaleJobThreshold, and still have the +// reaper reclaim a job the fetcher is still holding, because the lease +// granted at claim time expires long before the margin the threshold +// implied. A backend that is not lease-aware falls back to +// reapStaleJobsLegacy, which does read StaleJobThreshold directly, so +// that is what this check polices there instead. +func checkReaperMargin(cfg dispatch.Config, leaseAware bool) error { + if cfg.StaleJobThreshold <= 0 { + // The reaper is disabled; nothing can reclaim anything. + return nil + } + + window := max(cfg.PollInterval, 0) + max(cfg.HeartbeatInterval, 0) + minimum := reaperSafetyFactor * window + + effective := cfg.StaleJobThreshold + if leaseAware { + effective = effectiveReclaimWindow(cfg) + } + + if effective >= minimum { + return nil + } + + return fmt.Errorf( + "dispatch: the effective reclamation window (%s) is too low for resource-aware "+ + "admission: the fetcher may hold a claimed job for up to PollInterval (%s) while it "+ + "reclaims capacity, and that job is not heartbeated for a further HeartbeatInterval "+ + "(%s), so the reaper could reclaim a job this worker is still holding; set "+ + "StaleJobThreshold (currently %s) and DefaultLeaseTTL (currently %s) so the window "+ + "that actually governs reclamation is at least %s, or leave the resource manager unset", + effective, cfg.PollInterval, cfg.HeartbeatInterval, + cfg.StaleJobThreshold, cfg.DefaultLeaseTTL, minimum) +} + +// effectiveReclaimWindow mirrors worker.Pool.leaseTTLFor(nil): the lease +// TTL a freshly granted lease gets when no job overrides it with its own +// job.WithLeaseTTL. That per-job override can only raise a specific job's +// window above this floor, never lower it below what a config-time check +// can see, so checking the floor is sound. +func effectiveReclaimWindow(cfg dispatch.Config) time.Duration { + if cfg.DefaultLeaseTTL > 0 { + return cfg.DefaultLeaseTTL + } + if cfg.StaleJobThreshold > 0 { + return cfg.StaleJobThreshold + } + + return job.DefaultLeaseTTL +} + +// resolveResources computes the job's resource spec and writes it onto +// the job before it is persisted. +// +// This runs once, in the enqueuing process. Everything downstream — the +// dequeue predicate, admission, pod sizing — reads the stored columns, +// so no user code ever runs on the scheduling path and a job's +// requirement cannot differ between workers. +func (eng *Engine) resolveResources(ctx context.Context, j *job.Job, opts job.Options) error { + sizes, total, primaryHash := inputSizes(opts.Bindings) + + j.InputBytes = total + j.PrimaryInputHash = primaryHash + + decl := eng.registry.Resources(j.Name) + + // Nothing anywhere constrains this job, so there is nothing to + // resolve and nothing to check it against. Skipping keeps enqueue a + // single insert for every job written before this feature existed. + if !eng.resourcesInPlay(decl, opts) { + return nil + } + + // A definition declares; an enqueue overrides. Both the func and the + // class are single-valued rather than merged per key, so the + // enqueue-time value replaces the declared one outright when given. + resFunc := decl.Func + if opts.ResourceFunc != nil { + resFunc = opts.ResourceFunc + } + + class := decl.Class + if opts.ResourceClass != "" { + class = opts.ResourceClass + } + + spec, err := resource.Resolve(ctx, resource.ResolveInput{ + GlobalDefault: eng.resourceDefault, + QueueDefault: eng.queueResources[j.Queue], + Declared: decl.Requests, + Func: resFunc, + Estimator: eng.estimator, + Override: opts.Resources, + DeclaredLimits: decl.Limits, + OverrideLimits: opts.ResourceLimits, + Class: class, + MaxCapacity: eng.MaxWorkerCapacity(ctx), + // Neither source may fail an enqueue, but neither may fail + // silently either: both run once per enqueue in this process, so + // a broken one is broken for every job of that name from here on + // and the only symptom is jobs quietly sized from the static + // declaration. + OnError: func(source string, err error) { + eng.logger.Warn("resource: sizing source failed; falling back to the declaration", + log.String("source", source), + log.String("job", j.Name), + log.String("queue", j.Queue), + log.String("error", err.Error())) + }, + Request: resource.Request{ + JobName: j.Name, + Queue: j.Queue, + Payload: j.Payload, + Inputs: sizes, + InputBytes: total, + Attempt: j.RetryCount, + ScopeOrgID: j.ScopeOrgID, + }, + }) + if err != nil { + return err + } + + // A zero spec leaves the columns nil rather than storing an empty + // map, so an unconstrained job looks exactly as it did before. + if !spec.Requests.IsZero() { + j.Resources = spec.Requests + } + + if !spec.Limits.IsZero() { + j.ResourceLimits = spec.Limits + } + + j.ResourceClass = spec.Class + + return nil +} + +// resourcesInPlay reports whether any source could produce a +// requirement for this job. +// +// Worker capacity is deliberately not a source: it only ever rejects a +// requirement, and an empty requirement exceeds nothing. +func (eng *Engine) resourcesInPlay(decl job.ResourceDecl, opts job.Options) bool { + return !decl.IsZero() || + eng.estimator != nil || + len(eng.resourceDefault) > 0 || + len(eng.queueResources) > 0 || + len(opts.Resources) > 0 || + len(opts.ResourceLimits) > 0 || + opts.ResourceFunc != nil || + opts.ResourceClass != "" +} + +// inputSizes flattens artifact bindings into estimator input. +// +// artifact.Ref already carries Size and ContentHash, so this needs no +// store round-trip and enqueue stays a single insert. +// +// The returned hash is that of the largest input, ties broken by slot +// name. Determinism matters: the same job enqueued twice must advertise +// the same locality signal, and Go map iteration order would not. +func inputSizes(bindings map[string]artifact.Ref) ( + sizes []resource.InputSize, total int64, primaryHash string, +) { + if len(bindings) == 0 { + return nil, 0, "" + } + + sizes = make([]resource.InputSize, 0, len(bindings)) + + for name, ref := range bindings { + sizes = append(sizes, resource.InputSize{ + Name: name, + Bytes: ref.Size, + Hash: ref.ContentHash, + }) + + total += ref.Size + } + + sort.Slice(sizes, func(i, k int) bool { + if sizes[i].Bytes != sizes[k].Bytes { + return sizes[i].Bytes > sizes[k].Bytes + } + + return sizes[i].Name < sizes[k].Name + }) + + return sizes, total, sizes[0].Hash +} + +// MaxWorkerCapacity returns the per-key maximum capacity the +// unschedulable check may compare a job against, or an empty Set when +// the check is off. +// +// An empty result disables the check rather than rejecting everything. +// +// It is empty unless an operator called WithWorkerCapacity, and that +// gate is the whole correctness argument. The registry cannot supply the +// fleet maximum on its own: cluster.Worker.Capacity round-trips only on +// store/memory — postgres, sqlite, mongo, redis and the k8s provider all +// enumerate worker fields by hand and drop it — so a fleet whose largest +// worker has 64 GiB reads back as a fleet of workers with no capacity at +// all. Deriving the ceiling from whatever this process happens to know +// therefore does not converge on the truth; it converges on THIS +// process, and rejects at enqueue every job bigger than the pod that +// enqueued it. Requiring the declaration also keeps the common path free +// of a ListWorkers round trip per enqueue. +// +// When the declaration is present the registry is still consulted, so +// the ceiling can only rise toward the real fleet maximum on a backend +// that carries capacity. "Live" means both an active state and a recent +// heartbeat. State alone is not enough: nothing in Dispatch ever writes +// WorkerDead — a worker registers active and is either deregistered on +// clean shutdown or its row is deleted by DeleteStaleWorkers. A worker +// killed by SIGKILL, an OOM or a pod eviction therefore stays "active" +// in the registry until something sweeps it, and counting its capacity +// would admit jobs no live worker can run. +func (eng *Engine) MaxWorkerCapacity(ctx context.Context) resource.Set { + if len(eng.workerCapacity) == 0 { + return nil + } + + maxCap := eng.workerCapacity.Clone() + + if eng.clusterStore == nil { + return maxCap + } + + workers, err := eng.clusterStore.ListWorkers(ctx) + if err != nil { + // Capacity is advisory here. Failing an enqueue because the + // cluster registry was briefly unreachable would be worse than + // admitting a job that later needs rescheduling. + eng.logger.Warn("resource: worker capacity unavailable; "+ + "skipping the unschedulable check", + log.String("error", err.Error())) + + return maxCap + } + + cutoff := time.Now().UTC().Add(-eng.staleWorkerThreshold()) + + for _, w := range workers { + if w == nil || w.State != cluster.WorkerActive { + continue + } + + if w.LastSeen.Before(cutoff) { + continue + } + + maxCap = maxCap.Max(w.Capacity) + } + + return maxCap +} + +// staleWorkerThreshold is how long a worker may go without a heartbeat +// before its capacity stops counting. +// +// It reuses the sweep threshold from extension.go — max(5×heartbeat, +// 5 minutes) — rather than inventing a second notion of staleness, so a +// worker's capacity stops counting at roughly the same moment the rest +// of the cluster layer stops believing in the worker. +// +// The threshold is deliberately generous in the same direction: shrinking +// the fleet view too eagerly turns a valid enqueue into a hard +// ErrUnschedulable, whereas holding a dead worker's capacity a little too +// long only lets a job pend. A loud false rejection is the worse failure. +func (eng *Engine) staleWorkerThreshold() time.Duration { + threshold := 5 * eng.d.Config().HeartbeatInterval + if threshold < 5*time.Minute { + threshold = 5 * time.Minute + } + + return threshold +} diff --git a/engine/resource_test.go b/engine/resource_test.go new file mode 100644 index 0000000..4fbcc4b --- /dev/null +++ b/engine/resource_test.go @@ -0,0 +1,649 @@ +package engine_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/cluster" + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/resource/resourcetest" + "github.com/xraph/dispatch/store/memory" +) + +type resPayload struct { + N int `json:"n"` +} + +// newResourceEngine builds a plain engine over the memory store. +// +// Tests that need artifact bindings use newArtifactRig from +// engine/artifact_test.go instead: applyBindings requires a configured +// artifact backend, and both files share package engine_test. +func newResourceEngine(t *testing.T, opts ...engine.Option) *engine.Engine { + t.Helper() + + d, err := dispatch.New( + dispatch.WithStore(memory.New()), + dispatch.WithConcurrency(2), + dispatch.WithQueues([]string{"default"}), + ) + if err != nil { + t.Fatalf("dispatch.New() error = %v", err) + } + + eng, err := engine.Build(d, opts...) + if err != nil { + t.Fatalf("engine.Build() error = %v", err) + } + + return eng +} + +func TestEnqueueResolvesStaticDeclaration(t *testing.T) { + eng := newResourceEngine(t) + + def := job.NewDefinition("res.static", + func(context.Context, resPayload) error { return nil }, + job.WithResources(resource.CPUs(4), resource.MemoryGB(16)), + ) + engine.Register(eng, def) + + j, err := engine.Enqueue(context.Background(), eng, def.Name, resPayload{N: 1}) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if j.Resources[resource.CPU] != 4000 { + t.Errorf("cpu = %d, want 4000", j.Resources[resource.CPU]) + } + + if j.Resources[resource.Memory] != 16<<30 { + t.Errorf("memory = %d, want 16 GiB", j.Resources[resource.Memory]) + } + + if j.ResourceLimits[resource.Memory] != 16<<30 { + t.Errorf("memory limit should default to the request, got %v", j.ResourceLimits) + } + + if _, ok := j.ResourceLimits[resource.CPU]; ok { + t.Errorf("cpu limit should be unset (burstable), got %v", j.ResourceLimits) + } +} + +func TestEnqueueWithoutDeclarationLeavesResourcesZero(t *testing.T) { + eng := newResourceEngine(t) + + def := job.NewDefinition("res.none", + func(context.Context, resPayload) error { return nil }) + engine.Register(eng, def) + + j, err := engine.Enqueue(context.Background(), eng, def.Name, resPayload{N: 1}) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if !j.Resources.IsZero() { + t.Errorf("Resources = %v; an undeclared job must cost nothing", j.Resources) + } +} + +// TestEnqueueResourceFuncSeesInputBytes is the case the whole track +// exists for: one definition serving a 40 MB model and a 4 GB one. +func TestEnqueueResourceFuncSeesInputBytes(t *testing.T) { + rig := newArtifactRig(t, 1<<30) + + def := job.NewDefinition("res.dynamic", + func(context.Context, resPayload) error { return nil }, + job.WithArtifactInputs( + artifact.Input("small", artifact.Required), + artifact.Input("large", artifact.Required), + ), + job.WithResourceFunc(func(_ context.Context, r resource.Request) (resource.Set, error) { + return resource.MemoryBytes(r.InputBytes * 3), nil + }), + ) + + if err := engine.RegisterChecked(rig.engine, def); err != nil { + t.Fatalf("RegisterChecked() error = %v", err) + } + + ctx := context.Background() + + rig.backend.Put("models", "small.bin", make([]byte, 100)) + rig.backend.Put("models", "large.bin", make([]byte, 200)) + + small, err := rig.svc.Register(ctx, "models", "small.bin") + if err != nil { + t.Fatalf("Register(small) error = %v", err) + } + + large, err := rig.svc.Register(ctx, "models", "large.bin") + if err != nil { + t.Fatalf("Register(large) error = %v", err) + } + + got, err := engine.Enqueue(ctx, rig.engine, def.Name, + resPayload{N: 1}, + engine.Bind("small", small), + engine.Bind("large", large), + ) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if got.InputBytes != 300 { + t.Errorf("InputBytes = %d, want 300", got.InputBytes) + } + + if got.Resources[resource.Memory] != 900 { + t.Errorf("memory = %d, want 900 (3x input)", got.Resources[resource.Memory]) + } + + // PrimaryInputHash may be empty: the artifact plane fills content_hash + // at first staging, not at registration. Assert on the ref's own hash + // rather than a literal so this holds either way. + if got.PrimaryInputHash != large.ContentHash { + t.Errorf("PrimaryInputHash = %q, want the larger input's hash %q", + got.PrimaryInputHash, large.ContentHash) + } +} + +func TestEnqueueOverrideBeatsDeclaration(t *testing.T) { + eng := newResourceEngine(t) + + def := job.NewDefinition("res.override", + func(context.Context, resPayload) error { return nil }, + job.WithResources(resource.MemoryGB(16)), + ) + engine.Register(eng, def) + + j, err := engine.Enqueue(context.Background(), eng, def.Name, resPayload{N: 1}, + job.WithResources(resource.MemoryGB(48)), + ) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if j.Resources[resource.Memory] != 48<<30 { + t.Errorf("memory = %d, want the 48 GiB override", j.Resources[resource.Memory]) + } +} + +// TestEnqueueResolvesFromEverySource covers one resolution source per +// case, each in isolation. +// +// Isolation is the point. resolveResources skips resolution entirely +// when resourcesInPlay reports nothing constrains the job, so a source +// missing from that predicate would silently enqueue every job of its +// kind with zero requirements. Each case here configures exactly one +// source, so deleting that source's clause from the predicate fails +// this test and nothing else masks it. +func TestEnqueueResolvesFromEverySource(t *testing.T) { + tests := []struct { + name string + engineOpts []engine.Option + defOpts []job.Option + enqOpts []job.Option + wantMem int64 + wantLimit int64 + wantClass string + }{ + { + name: "declaration on the definition", + defOpts: []job.Option{job.WithResources(resource.MemoryGB(2))}, + wantMem: 2 << 30, wantLimit: 2 << 30, + }, + { + name: "configured estimator", + engineOpts: []engine.Option{ + engine.WithEstimator(&resourcetest.FakeEstimator{Out: resource.MemoryGB(3)}), + }, + wantMem: 3 << 30, wantLimit: 3 << 30, + }, + { + name: "fleet-wide default", + engineOpts: []engine.Option{ + engine.WithResourceDefaults(resource.MemoryGB(4), nil), + }, + wantMem: 4 << 30, wantLimit: 4 << 30, + }, + { + name: "per-queue default", + engineOpts: []engine.Option{ + engine.WithResourceDefaults(nil, map[string]resource.Set{ + "default": resource.MemoryGB(5), + }), + }, + wantMem: 5 << 30, wantLimit: 5 << 30, + }, + { + name: "override at enqueue", + enqOpts: []job.Option{job.WithResources(resource.MemoryGB(6))}, + wantMem: 6 << 30, wantLimit: 6 << 30, + }, + { + name: "limits at enqueue with no request", + enqOpts: []job.Option{job.WithResourceLimits(resource.MemoryGB(7))}, + wantLimit: 7 << 30, + }, + { + name: "resource func at enqueue", + enqOpts: []job.Option{ + job.WithResourceFunc(func(context.Context, resource.Request) (resource.Set, error) { + return resource.MemoryGB(8), nil + }), + }, + wantMem: 8 << 30, wantLimit: 8 << 30, + }, + { + name: "class at enqueue", + enqOpts: []job.Option{job.WithResourceClass("gpu-a100")}, + wantClass: "gpu-a100", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + eng := newResourceEngine(t, tt.engineOpts...) + + def := job.NewDefinition("res.source", + func(context.Context, resPayload) error { return nil }, + tt.defOpts...) + engine.Register(eng, def) + + j, err := engine.Enqueue(context.Background(), eng, def.Name, + resPayload{N: 1}, tt.enqOpts...) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if got := j.Resources[resource.Memory]; got != tt.wantMem { + t.Errorf("Resources memory = %d, want %d", got, tt.wantMem) + } + + if got := j.ResourceLimits[resource.Memory]; got != tt.wantLimit { + t.Errorf("ResourceLimits memory = %d, want %d", got, tt.wantLimit) + } + + if j.ResourceClass != tt.wantClass { + t.Errorf("ResourceClass = %q, want %q", j.ResourceClass, tt.wantClass) + } + }) + } +} + +// TestEnqueueLimitPrecedence pins that an enqueue-time limit beats a +// declared one, which is what OverrideLimits exists for. +func TestEnqueueLimitPrecedence(t *testing.T) { + eng := newResourceEngine(t) + + def := job.NewDefinition("res.limits", + func(context.Context, resPayload) error { return nil }, + job.WithResources(resource.MemoryGB(16)), + job.WithResourceLimits(resource.MemoryGB(20)), + ) + engine.Register(eng, def) + + ctx := context.Background() + + // The declared limit stands when the enqueue supplies none. + j, err := engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if j.ResourceLimits[resource.Memory] != 20<<30 { + t.Errorf("declared limit = %d, want 20 GiB", j.ResourceLimits[resource.Memory]) + } + + // An enqueue-time limit replaces it. + j, err = engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}, + job.WithResourceLimits(resource.MemoryGB(32))) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if j.ResourceLimits[resource.Memory] != 32<<30 { + t.Errorf("override limit = %d, want the 32 GiB override", j.ResourceLimits[resource.Memory]) + } + + if j.Resources[resource.Memory] != 16<<30 { + t.Errorf("requests = %d; a limit override must not move the request", j.Resources[resource.Memory]) + } +} + +// TestEnqueueFuncAndClassPrecedence pins that the single-valued sources +// are replaced outright by an enqueue-time value rather than merged. +func TestEnqueueFuncAndClassPrecedence(t *testing.T) { + eng := newResourceEngine(t) + + def := job.NewDefinition("res.precedence", + func(context.Context, resPayload) error { return nil }, + job.WithResourceClass("cpu-standard"), + job.WithResourceFunc(func(context.Context, resource.Request) (resource.Set, error) { + return resource.MemoryGB(2), nil + }), + ) + engine.Register(eng, def) + + ctx := context.Background() + + // Declared values apply when the enqueue overrides neither. + j, err := engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if j.Resources[resource.Memory] != 2<<30 || j.ResourceClass != "cpu-standard" { + t.Errorf("declared func/class not applied: memory = %d, class = %q", + j.Resources[resource.Memory], j.ResourceClass) + } + + // Enqueue-time values replace them. + j, err = engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}, + job.WithResourceClass("gpu-h100"), + job.WithResourceFunc(func(context.Context, resource.Request) (resource.Set, error) { + return resource.MemoryGB(9), nil + }), + ) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if j.Resources[resource.Memory] != 9<<30 { + t.Errorf("memory = %d, want the enqueue func's 9 GiB", j.Resources[resource.Memory]) + } + + if j.ResourceClass != "gpu-h100" { + t.Errorf("class = %q, want the enqueue override", j.ResourceClass) + } +} + +// TestEnqueueEstimatorSeesDeclarationAndLosesToOverride places the +// estimator in the precedence chain: above a declaration, below an +// explicit override, and given the declaration so it can defer to it. +func TestEnqueueEstimatorSeesDeclarationAndLosesToOverride(t *testing.T) { + est := &resourcetest.FakeEstimator{Out: resource.MemoryGB(32)} + eng := newResourceEngine(t, engine.WithEstimator(est)) + + def := job.NewDefinition("res.estimated", + func(context.Context, resPayload) error { return nil }, + job.WithResources(resource.CPUs(2), resource.MemoryGB(8)), + ) + engine.Register(eng, def) + + ctx := context.Background() + + j, err := engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if j.Resources[resource.Memory] != 32<<30 { + t.Errorf("memory = %d, want the estimator's 32 GiB", j.Resources[resource.Memory]) + } + + // Per-key overlay: the estimator predicted only memory, so the + // declared CPU must survive. + if j.Resources[resource.CPU] != 2000 { + t.Errorf("cpu = %d, want the declared 2000 to survive a memory-only estimate", j.Resources[resource.CPU]) + } + + if est.Last.Declared[resource.Memory] != 8<<30 { + t.Errorf("estimator saw Declared = %v, want the 8 GiB declaration", est.Last.Declared) + } + + // An explicit override outranks the estimator. + j, err = engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}, + job.WithResources(resource.MemoryGB(64))) + if err != nil { + t.Fatalf("Enqueue() error = %v", err) + } + + if j.Resources[resource.Memory] != 64<<30 { + t.Errorf("memory = %d, want the 64 GiB override to beat the estimator", j.Resources[resource.Memory]) + } +} + +func TestEnqueueRejectsUnschedulable(t *testing.T) { + st := memory.New() + + d, err := dispatch.New( + dispatch.WithStore(st), + dispatch.WithConcurrency(1), + dispatch.WithQueues([]string{"default"}), + ) + if err != nil { + t.Fatalf("dispatch.New() error = %v", err) + } + + eng, err := engine.Build(d, + engine.WithWorkerCapacity(resource.Set{resource.Memory: 8 << 30})) + if err != nil { + t.Fatalf("engine.Build() error = %v", err) + } + + def := job.NewDefinition("res.toobig", + func(context.Context, resPayload) error { return nil }, + job.WithResources(resource.MemoryGB(64)), + ) + engine.Register(eng, def) + + _, err = engine.Enqueue(context.Background(), eng, def.Name, resPayload{N: 1}) + if !errors.Is(err, resource.ErrUnschedulable) { + t.Fatalf("got %v, want ErrUnschedulable", err) + } + + count, cErr := st.CountJobs(context.Background(), job.CountOpts{}) + if cErr != nil { + t.Fatalf("CountJobs() error = %v", cErr) + } + + if count != 0 { + t.Errorf("an unschedulable job must not be persisted, found %d", count) + } +} + +// TestEnqueueUsesFleetCapacity proves the unschedulable check reads the +// largest capacity in the cluster, not just this process's own: a job too +// big for the local worker still enqueues when a bigger worker exists. +func TestEnqueueUsesFleetCapacity(t *testing.T) { + st := memory.New() + + d, err := dispatch.New( + dispatch.WithStore(st), + dispatch.WithConcurrency(1), + dispatch.WithQueues([]string{"default"}), + ) + if err != nil { + t.Fatalf("dispatch.New() error = %v", err) + } + + eng, err := engine.Build(d, + engine.WithWorkerCapacity(resource.Set{resource.Memory: 8 << 30})) + if err != nil { + t.Fatalf("engine.Build() error = %v", err) + } + + ctx := context.Background() + + if rErr := st.RegisterWorker(ctx, &cluster.Worker{ + ID: id.NewWorkerID(), + State: cluster.WorkerActive, + LastSeen: time.Now().UTC(), + Capacity: resource.Set{resource.Memory: 128 << 30}, + }); rErr != nil { + t.Fatalf("RegisterWorker() error = %v", rErr) + } + + def := job.NewDefinition("res.fleet", + func(context.Context, resPayload) error { return nil }, + job.WithResources(resource.MemoryGB(64)), + ) + engine.Register(eng, def) + + if _, err = engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}); err != nil { + t.Fatalf("Enqueue() error = %v; a worker in the fleet can run this", err) + } + + if got := eng.MaxWorkerCapacity(ctx)[resource.Memory]; got != 128<<30 { + t.Errorf("MaxWorkerCapacity memory = %d, want the fleet maximum 128 GiB", got) + } +} + +// TestMaxWorkerCapacityIgnoresUnusableWorkers keeps a worker that cannot +// actually run anything from admitting jobs nothing can run. +// +// The stale case is the one that matters in production: nothing in +// Dispatch ever writes WorkerDead, so a worker killed by SIGKILL, an OOM +// or a pod eviction stays "active" in the registry with a frozen +// LastSeen until something sweeps the row. +func TestMaxWorkerCapacityIgnoresUnusableWorkers(t *testing.T) { + tests := []struct { + name string + state cluster.WorkerState + seenAt time.Time + }{ + { + name: "crashed worker still marked active", + state: cluster.WorkerActive, + seenAt: time.Now().UTC().Add(-time.Hour), + }, + { + name: "explicitly dead worker", + state: cluster.WorkerDead, + seenAt: time.Now().UTC(), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + st := memory.New() + + d, err := dispatch.New( + dispatch.WithStore(st), + dispatch.WithConcurrency(1), + dispatch.WithQueues([]string{"default"}), + ) + if err != nil { + t.Fatalf("dispatch.New() error = %v", err) + } + + eng, err := engine.Build(d, + engine.WithWorkerCapacity(resource.Set{resource.Memory: 8 << 30})) + if err != nil { + t.Fatalf("engine.Build() error = %v", err) + } + + ctx := context.Background() + + if rErr := st.RegisterWorker(ctx, &cluster.Worker{ + ID: id.NewWorkerID(), + State: tt.state, + LastSeen: tt.seenAt, + Capacity: resource.Set{resource.Memory: 128 << 30}, + }); rErr != nil { + t.Fatalf("RegisterWorker() error = %v", rErr) + } + + if got := eng.MaxWorkerCapacity(ctx)[resource.Memory]; got != 8<<30 { + t.Errorf("MaxWorkerCapacity memory = %d, want 8 GiB; this worker cannot run anything", got) + } + + // And the capacity it advertised must not admit a job either. + def := job.NewDefinition("res.unusable."+tt.name, + func(context.Context, resPayload) error { return nil }, + job.WithResources(resource.MemoryGB(64)), + ) + engine.Register(eng, def) + + if _, err = engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}); !errors.Is(err, resource.ErrUnschedulable) { + t.Errorf("Enqueue() = %v, want ErrUnschedulable", err) + } + }) + } +} + +// errCluster is a cluster registry whose ListWorkers always fails. +// +// It embeds *memory.Store so it satisfies every store interface +// engine.Build type-asserts, and shadows the single method under test. +type errCluster struct { + *memory.Store +} + +var errListWorkers = errors.New("cluster registry unreachable") + +func (errCluster) ListWorkers(context.Context) ([]*cluster.Worker, error) { + return nil, errListWorkers +} + +// TestEnqueueSurvivesClusterStoreFailure pins the documented contract: +// capacity is advisory, so a registry that cannot be read downgrades to +// skipping the unschedulable check rather than failing the enqueue. +func TestEnqueueSurvivesClusterStoreFailure(t *testing.T) { + d, err := dispatch.New( + dispatch.WithStore(errCluster{memory.New()}), + dispatch.WithConcurrency(1), + dispatch.WithQueues([]string{"default"}), + ) + if err != nil { + t.Fatalf("dispatch.New() error = %v", err) + } + + // Capacity far below the job's requirement: were the registry + // readable, this enqueue would be rejected. + eng, err := engine.Build(d, + engine.WithWorkerCapacity(resource.Set{resource.Memory: 8 << 30})) + if err != nil { + t.Fatalf("engine.Build() error = %v", err) + } + + ctx := context.Background() + + if got := eng.MaxWorkerCapacity(ctx)[resource.Memory]; got != 8<<30 { + t.Errorf("MaxWorkerCapacity memory = %d; a failed read must still yield the local seed", got) + } + + def := job.NewDefinition("res.clusterdown", + func(context.Context, resPayload) error { return nil }, + job.WithResources(resource.MemoryGB(4)), + ) + engine.Register(eng, def) + + j, err := engine.Enqueue(ctx, eng, def.Name, resPayload{N: 1}) + if err != nil { + t.Fatalf("Enqueue() error = %v; a cluster read failure must not fail an enqueue", err) + } + + if j.Resources[resource.Memory] != 4<<30 { + t.Errorf("memory = %d, want 4 GiB; resolution still runs", j.Resources[resource.Memory]) + } +} + +func TestInputSizesTieBreaksDeterministically(t *testing.T) { + // Two inputs of equal size must always yield the same primary hash, + // or the same job enqueued twice would advertise different locality. + bindings := map[string]artifact.Ref{ + "zulu": {ID: id.NewArtifactID(), Size: 100, ContentHash: "blake3:zz"}, + "alpha": {ID: id.NewArtifactID(), Size: 100, ContentHash: "blake3:aa"}, + } + + for range 20 { + _, total, primary := engine.InputSizesForTest(bindings) + if total != 200 { + t.Fatalf("total = %d, want 200", total) + } + + if primary != "blake3:aa" { + t.Fatalf("primary = %q, want the name-ordered winner", primary) + } + } +} diff --git a/errors.go b/errors.go index b9fb1b9..87355f0 100644 --- a/errors.go +++ b/errors.go @@ -25,6 +25,27 @@ var ( ErrInvalidState = errors.New("dispatch: invalid state transition") ErrMaxRetriesExceeded = errors.New("dispatch: max retries exceeded") + // ErrPermanent marks a failure that retrying cannot resolve. A job + // whose handler returns an error wrapping it skips its remaining + // attempts and goes straight to the dead letter queue. + // + // Handlers wrap it to decline a retry they know is pointless: + // + // if !validPayload(p) { + // return fmt.Errorf("malformed payload: %w", dispatch.ErrPermanent) + // } + // + // The artifact plane wraps it for a missing or forbidden input, which + // is where it earns its keep: a job whose input was deleted would + // otherwise spend every attempt, and the whole backoff schedule + // between them, rediscovering that the object is still gone. + // + // Only mark a failure permanent when retrying is certain to fail the + // same way. Anything unrecognized stays retryable on purpose: a job + // retried needlessly costs some compute, while one dead-lettered by + // mistake loses work that would have succeeded. + ErrPermanent = errors.New("dispatch: permanent failure") + // Cluster errors. ErrLeadershipLost = errors.New("dispatch: leadership lost") ErrNotLeader = errors.New("dispatch: not the leader") diff --git a/exec/deps_test.go b/exec/deps_test.go new file mode 100644 index 0000000..8b30a4d --- /dev/null +++ b/exec/deps_test.go @@ -0,0 +1,41 @@ +package exec_test + +import ( + "go/build" + "strings" + "testing" +) + +// TestExecIsALeafPackage guards the import constraint the whole design +// rests on. job imports exec for Options.Execution, so exec importing job +// would be a cycle; importing worker or engine would drag the store, and +// with it the credentials, into a package the sandbox links. +func TestExecIsALeafPackage(t *testing.T) { + const self = "github.com/xraph/dispatch/exec" + + allowed := map[string]bool{ + "github.com/xraph/dispatch": true, + "github.com/xraph/dispatch/id": true, + "github.com/xraph/dispatch/scope": true, + "github.com/xraph/dispatch/artifact": true, + // Request.ResourceLimits carries job.Job.ResourceLimits across the + // execution boundary so exec/subprocess can enforce it per job. + // resource is a leaf like id and artifact: it imports neither job + // nor exec, so this adds no cycle. + "github.com/xraph/dispatch/resource": true, + } + + pkg, err := build.Import(self, "", 0) + if err != nil { + t.Fatalf("import %s: %v", self, err) + } + + for _, imp := range pkg.Imports { + if !strings.HasPrefix(imp, "github.com/xraph/dispatch") { + continue // standard library and third-party are fine + } + if !allowed[imp] { + t.Errorf("exec imports %q, which breaks the leaf constraint", imp) + } + } +} diff --git a/exec/doc.go b/exec/doc.go new file mode 100644 index 0000000..5ca7974 --- /dev/null +++ b/exec/doc.go @@ -0,0 +1,17 @@ +// Package exec defines the execution boundary between the Dispatch worker +// and a job handler. +// +// Today a handler is an ordinary Go function called in-process, sharing the +// worker's memory, credentials, and network. Handlers that parse untrusted +// bytes with memory-unsafe native libraries need more than that, so exec +// generalises the call into an [Executor] with implementations forming an +// escalating ladder: in-process, subprocess, OCI container, and Kubernetes +// Job-per-task. +// +// exec is a leaf package. It imports only id and artifact within this +// module — never job, worker, engine, scope, or the root dispatch +// package — so that job.Options can carry an execution [Policy] without +// an import cycle. This mirrors how artifact is itself positioned for +// input declarations. See deps_test.go for the guard that keeps this +// list accurate. +package exec diff --git a/exec/example_test.go b/exec/example_test.go new file mode 100644 index 0000000..89ef5ce --- /dev/null +++ b/exec/example_test.go @@ -0,0 +1,40 @@ +package exec_test + +import ( + "context" + "fmt" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/inproc" + "github.com/xraph/dispatch/job" +) + +type modelInput struct { + Detail int `json:"detail"` +} + +// ExampleRegistry_Select shows how a definition's declared isolation +// chooses the executor that runs it. +func ExampleRegistry_Select() { + registry := job.NewRegistry() + + // A handler that parses untrusted geometry declares that it needs a + // separate address space at minimum. + job.NewDefinition("tessellate.model", + func(context.Context, modelInput) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(registry) + + executors := exec.NewRegistry(inproc.New(registry)) + + _, err := executors.Select(registry.Policy("tessellate.model")) + fmt.Println(err != nil) + + // A handler that declares nothing runs in-process, as it always has. + e, err := executors.Select(registry.Policy("send.email")) + fmt.Println(e.Name(), err) + + // Output: + // true + // inprocess +} diff --git a/exec/exectest/doc.go b/exec/exectest/doc.go new file mode 100644 index 0000000..8c97208 --- /dev/null +++ b/exec/exectest/doc.go @@ -0,0 +1,12 @@ +// Package exectest is the conformance suite every exec.Executor must pass. +// +// The rungs of the isolation ladder are meant to be interchangeable: the +// same handler, the same payload, and the same declared inputs must behave +// the same way whether the handler runs in-process or in a pod. One shared +// table-driven suite is how that stays true, and it is what lets a new rung +// land without redesigning the ones before it. +// +// Rungs differ in what they can enforce — in-process cannot kill a handler +// that ignores its deadline — so a rung declares its Capabilities and the +// suite asserts the enforcement cases only against rungs that claim them. +package exectest diff --git a/exec/exectest/handlers.go b/exec/exectest/handlers.go new file mode 100644 index 0000000..a2780d9 --- /dev/null +++ b/exec/exectest/handlers.go @@ -0,0 +1,117 @@ +package exectest + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/xraph/dispatch/job" +) + +// Job names the suite installs. Every executor under test must be able to +// run all of them. +// +// Artifact-carrying fixtures (writing outputs, reading staged inputs) +// arrive in Phase 2 alongside the suite cases that exercise them. Adding +// them now would ship handlers no case runs. +const ( + JobOK = "exectest.ok" + JobError = "exectest.error" + JobPanic = "exectest.panic" + JobSlow = "exectest.slow" + JobEcho = "exectest.echo" +) + +// ErrIntentional is what JobError returns, so tests can match it exactly. +var ErrIntentional = errors.New("intentional failure") + +// EchoPayload is the payload JobEcho round-trips. Want, when non-zero, is +// the byte length the handler asserts Value has, which is how the suite +// proves a large payload crossed the boundary without truncation. +type EchoPayload struct { + Value string `json:"value"` + Want int `json:"want"` +} + +// SlowPayload controls how long JobSlow sleeps. +type SlowPayload struct { + SleepMillis int `json:"sleep_millis"` + + // IgnoreCtx makes the handler deaf to cancellation entirely, standing + // in for a native library that has stopped honouring it. Only a rung + // that can kill the process will stop this one. + IgnoreCtx bool `json:"ignore_ctx"` + + // SwallowCancel only matters when IgnoreCtx is false. A cooperative + // handler normally returns ctx.Err() once its context is done; setting + // this makes it catch that and return nil instead, the shape of a + // handler whose own cleanup swallows context.Canceled. The shim then + // reports StatusOK in its Result frame despite the deadline having + // fired, which is exactly the case that must not fool the executor + // into reporting success — the executor's own timedOut bookkeeping, + // not the frame's contents, has to be what decides the Status. + SwallowCancel bool `json:"swallow_cancel"` +} + +// Handlers returns the fixture handler set. Registering these is all an +// executor needs to be run through the suite. +func Handlers() []job.Registrable { + return []job.Registrable{ + job.NewDefinition(JobOK, func(context.Context, struct{}) error { + return nil + }), + job.NewDefinition(JobError, func(context.Context, struct{}) error { + return ErrIntentional + }), + job.NewDefinition(JobPanic, func(context.Context, struct{}) error { + panic("intentional panic") + }), + job.NewDefinition(JobSlow, func(ctx context.Context, p SlowPayload) error { + d := time.Duration(p.SleepMillis) * time.Millisecond + if p.IgnoreCtx { + // Stands in for a native library that has stopped + // honouring cancellation. Only a rung that can kill + // will stop this. + time.Sleep(d) + return nil + } + select { + case <-time.After(d): + return nil + case <-ctx.Done(): + if p.SwallowCancel { + return nil + } + + return ctx.Err() + } + }), + // Echo proves the payload crossed the boundary intact. It + // validates by round-tripping rather than by recording to a + // package-level variable, which an out-of-process rung could + // never observe anyway. + job.NewDefinition(JobEcho, func(_ context.Context, p EchoPayload) error { + if p.Value == "" { + return errors.New("exectest: echo received an empty payload") + } + if p.Want != 0 && len(p.Value) != p.Want { + return fmt.Errorf("exectest: echo got %d bytes, want %d", len(p.Value), p.Want) + } + + return nil + }), + } +} + +// HandlerNames returns the fixture job names, which is what a fingerprint +// is derived from. +func HandlerNames() []string { + defs := Handlers() + names := make([]string, 0, len(defs)) + for _, d := range defs { + names = append(names, d.JobName()) + } + + return names +} diff --git a/exec/exectest/suite.go b/exec/exectest/suite.go new file mode 100644 index 0000000..c2e0bdf --- /dev/null +++ b/exec/exectest/suite.go @@ -0,0 +1,378 @@ +package exectest + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "testing" + "time" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/id" +) + +// Capabilities describes what a rung can actually do, so the suite asserts +// enforcement only against rungs that provide it. +// +// It describes variation, not an opt-out. RunSuite cross-checks it against +// the executor's own Level: anything claiming LevelProcess or above must +// enforce deadlines and isolate panics, and saying otherwise fails the suite +// rather than skipping the tests that prove it. +type Capabilities struct { + // Enforces means the rung can stop a handler that ignores its + // deadline. Only out-of-process rungs can. + Enforces bool + + // ReportsUsage means the rung measures CPU time and peak memory + // rather than only wall time. + ReportsUsage bool + + // IsolatesPanic means a panicking handler does not take the caller + // down, so the rung reports it as a failed attempt rather than + // relying on the worker's recover middleware. + IsolatesPanic bool +} + +// RunSuite runs the conformance suite against one executor implementation. +// +// newExecutor is called per subtest so each case gets a clean executor. +// The returned executor must already have the fixture Handlers registered. +func RunSuite(t *testing.T, name string, newExecutor func(*testing.T) exec.Executor, caps Capabilities) { + t.Helper() + + t.Run(name, func(t *testing.T) { + t.Run("CapabilitiesMatchLevel", func(t *testing.T) { + testCapabilitiesMatchLevel(t, name, newExecutor, caps) + }) + t.Run("Identity", func(t *testing.T) { testIdentity(t, newExecutor) }) + t.Run("Success", func(t *testing.T) { testSuccess(t, newExecutor) }) + t.Run("HandlerError", func(t *testing.T) { testHandlerError(t, newExecutor) }) + t.Run("UnknownHandler", func(t *testing.T) { testUnknownHandler(t, newExecutor) }) + t.Run("InvalidRequest", func(t *testing.T) { testInvalidRequest(t, newExecutor) }) + t.Run("PayloadRoundTrip", func(t *testing.T) { testPayloadRoundTrip(t, newExecutor) }) + t.Run("LargePayload", func(t *testing.T) { testLargePayload(t, newExecutor) }) + t.Run("Cancellation", func(t *testing.T) { testCancellation(t, newExecutor) }) + t.Run("WallTimeRecorded", func(t *testing.T) { testWallTime(t, newExecutor) }) + t.Run("Reclaim", func(t *testing.T) { testReclaim(t, newExecutor) }) + + if caps.Enforces { + t.Run("DeadlineEnforced", func(t *testing.T) { testDeadlineEnforced(t, newExecutor) }) + t.Run("DeadlineEnforcedCooperative", func(t *testing.T) { testDeadlineEnforcedCooperative(t, newExecutor) }) + t.Run("DeadlineEnforcedSwallowedCancellation", func(t *testing.T) { + testDeadlineEnforcedSwallowedCancellation(t, newExecutor) + }) + } + if caps.IsolatesPanic { + t.Run("PanicIsolated", func(t *testing.T) { testPanicIsolated(t, newExecutor) }) + } + if caps.ReportsUsage { + t.Run("UsageReported", func(t *testing.T) { testUsageReported(t, newExecutor) }) + } + }) +} + +func request(name string, payload any) *exec.Request { + // payload is always one of this file's fixture payload types, which + // are all marshalable, so the error is unreachable in practice. + raw, _ := json.Marshal(payload) //nolint:errcheck // fixture payload types always marshal + + return &exec.Request{ + JobID: id.NewJobID(), + Name: name, + Payload: raw, + Fingerprint: exec.Fingerprint(HandlerNames()), + Policy: exec.NewPolicy(), + } +} + +// CheckCapabilities reports whether the capabilities a rung claims are +// consistent with the isolation it advertises, returning nil when they are. +// +// Without this check Capabilities is an escape hatch: a rung that reports +// LevelProcess but sets Enforces false skips DeadlineEnforced — the only +// test proving it can stop a handler that ignores cancellation, which is the +// property the whole isolation ladder exists to provide — and still passes +// the suite clean. An executor may run handlers where it cannot kill them, +// or it may claim LevelProcess; it may not do both. +// +// RunSuite calls this. It is exported so a rung's own tests can assert the +// same consistency without running the full suite. +func CheckCapabilities(name string, level exec.Level, caps Capabilities) error { + if level < exec.LevelProcess { + return nil + } + + var errs []error + if !caps.Enforces { + errs = append(errs, fmt.Errorf( + "executor %q reports Level %s but Capabilities.Enforces = false: "+ + "a rung running handlers out of process must be able to kill one that ignores its deadline", + name, level)) + } + if !caps.IsolatesPanic { + errs = append(errs, fmt.Errorf( + "executor %q reports Level %s but Capabilities.IsolatesPanic = false: "+ + "a handler panicking in another address space cannot take the worker down", + name, level)) + } + + return errors.Join(errs...) +} + +func testCapabilitiesMatchLevel( + t *testing.T, + name string, + newExecutor func(*testing.T) exec.Executor, + caps Capabilities, +) { + if err := CheckCapabilities(name, newExecutor(t).Level(), caps); err != nil { + t.Errorf("CheckCapabilities() = %v, want nil", err) + } +} + +func testIdentity(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + e := newExecutor(t) + if e.Name() == "" { + t.Error("Name() is empty") + } + if err := e.Close(); err != nil { + t.Errorf("Close() = %v, want nil", err) + } +} + +func testSuccess(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + res, err := newExecutor(t).Run(context.Background(), request(JobOK, struct{}{})) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Status != exec.StatusOK { + t.Errorf("Status = %q, want %q (handler err: %q)", res.Status, exec.StatusOK, res.HandlerErr) + } + if res.Err() != nil { + t.Errorf("Err() = %v, want nil", res.Err()) + } +} + +func testHandlerError(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + res, err := newExecutor(t).Run(context.Background(), request(JobError, struct{}{})) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Status != exec.StatusHandlerError { + t.Fatalf("Status = %q, want %q", res.Status, exec.StatusHandlerError) + } + if res.HandlerErr != ErrIntentional.Error() { + t.Errorf("HandlerErr = %q, want %q", res.HandlerErr, ErrIntentional.Error()) + } + if !errors.Is(res.Err(), exec.ErrHandler) { + t.Errorf("Err() = %v, want it to wrap ErrHandler", res.Err()) + } +} + +func testUnknownHandler(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + res, err := newExecutor(t).Run(context.Background(), request("exectest.absent", struct{}{})) + if err != nil { + t.Fatalf("Run() error = %v, want a Result", err) + } + if res.Status != exec.StatusLaunchFailed { + t.Fatalf("Status = %q, want %q", res.Status, exec.StatusLaunchFailed) + } + if res.Status.CountsAgainstRetries() { + t.Error("an unknown handler must not consume the retry budget") + } +} + +func testInvalidRequest(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + _, err := newExecutor(t).Run(context.Background(), &exec.Request{}) + if !errors.Is(err, exec.ErrInvalidRequest) { + t.Fatalf("Run() error = %v, want ErrInvalidRequest", err) + } +} + +func testPayloadRoundTrip(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + const want = "hello boundary" + + res, err := newExecutor(t).Run(context.Background(), + request(JobEcho, EchoPayload{Value: want, Want: len(want)})) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Status != exec.StatusOK { + t.Fatalf("Status = %q, want %q (handler err: %q)", res.Status, exec.StatusOK, res.HandlerErr) + } +} + +func testLargePayload(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + // Large enough to exceed a pipe buffer, so any rung that frames the + // request over a descriptor is exercised rather than accidentally + // fitting in one write. + big := make([]byte, 1<<20) + for i := range big { + big[i] = byte('a' + i%26) + } + + res, err := newExecutor(t).Run(context.Background(), + request(JobEcho, EchoPayload{Value: string(big), Want: len(big)})) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + // The handler compares the received length against Want, so a rung + // that truncated the payload in transit fails here rather than + // silently passing. + if res.Status != exec.StatusOK { + t.Errorf("Status = %q, want %q (handler err: %q)", res.Status, exec.StatusOK, res.HandlerErr) + } +} + +func testCancellation(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(20 * time.Millisecond) + cancel() + }() + + start := time.Now() + res, err := newExecutor(t).Run(ctx, + request(JobSlow, SlowPayload{SleepMillis: 5000, IgnoreCtx: false})) + elapsed := time.Since(start) + + // Whatever shape the failure takes, cancellation must actually cut the + // attempt short. Asserting only "it failed" would pass for a rung that + // ignored the cancel and let the handler run its full five seconds. + if elapsed > 3*time.Second { + t.Errorf("Run() took %v, want cancellation to cut it short", elapsed) + } + if err != nil { + // An out-of-process rung may surface cancellation as a launch + // error rather than a Result. Both shapes are acceptable. + return + } + if res.Status == exec.StatusOK { + t.Error("Status = ok, want a failure after cancellation") + } +} + +func testWallTime(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + res, err := newExecutor(t).Run(context.Background(), + request(JobSlow, SlowPayload{SleepMillis: 20})) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Usage.WallTime <= 0 { + t.Errorf("Usage.WallTime = %v, want > 0", res.Usage.WallTime) + } +} + +func testReclaim(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + // Reclaim must be safe to call when there is nothing to reclaim, + // because the pool calls it unconditionally at startup. + if err := newExecutor(t).Reclaim(context.Background(), id.NewWorkerID()); err != nil { + t.Errorf("Reclaim() = %v, want nil", err) + } +} + +func testDeadlineEnforced(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + req := request(JobSlow, SlowPayload{SleepMillis: 30000, IgnoreCtx: true}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + req.Policy = exec.NewPolicy(exec.GracePeriod(200 * time.Millisecond)) + + start := time.Now() + res, err := newExecutor(t).Run(context.Background(), req) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Status != exec.StatusTimeout { + t.Errorf("Status = %q, want %q", res.Status, exec.StatusTimeout) + } + // The handler asked to sleep 30s and ignores cancellation. Anything + // close to that means the rung did not actually kill it. + if elapsed > 10*time.Second { + t.Errorf("Run() took %v, want the deadline to be enforced", elapsed) + } +} + +// testDeadlineEnforcedCooperative is testDeadlineEnforced's counterpart for +// the well-behaved path: a handler that actually honours ctx.Done() and +// returns ctx.Err() once the deadline fires. That is the shape every +// correctly-written handler has, and it exercises a different part of an +// Enforces rung than the uncooperative case does — for the subprocess rung, +// the SIGTERM half of the kill ladder reaches the shim before SIGKILL would, +// the shim's own context gets cancelled, the handler returns promptly, and +// the shim writes a Result frame and exits 0. That produces a clean, +// unsignalled frame arriving at (or after) the same moment the executor's +// own deadline bookkeeping independently expired — and the executor still +// has to report StatusTimeout, not trust the frame's StatusOK. A rung that +// gets this ordering backwards makes StatusTimeout effectively unreachable +// for every handler that cooperates, which defeats the whole point of +// enforcement for the common case rather than the pathological one +// testDeadlineEnforced covers. +func testDeadlineEnforcedCooperative(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + req := request(JobSlow, SlowPayload{SleepMillis: 30000, IgnoreCtx: false}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + // Generous, deliberately: the point of this case is that the handler + // wins the race and exits cleanly well before any SIGKILL would fire, + // so the grace period needs real headroom for a cold, possibly + // race-detector-instrumented child to start up, decode the request, + // notice cancellation, and write its Result frame. A tight grace here + // would not make the assertion stricter, only flaky. + req.Policy = exec.NewPolicy(exec.GracePeriod(3 * time.Second)) + + res, err := newExecutor(t).Run(context.Background(), req) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Status != exec.StatusTimeout { + t.Errorf("Status = %q, want %q — a handler that returns ctx.Err() after the deadline must still report a blown deadline, not ok", + res.Status, exec.StatusTimeout) + } +} + +// testDeadlineEnforcedSwallowedCancellation covers the shape that fails +// silently rather than loudly: a handler whose own cleanup catches +// context.Canceled and returns nil instead of propagating it. The shim +// still reports StatusOK in that case — from the handler's point of view it +// really did return successfully — so an executor that ever lets a clean +// frame override its own deadline bookkeeping would report a blown deadline +// as StatusOK, and nothing about that would look like a failure to whoever +// is reading Result.Status. +func testDeadlineEnforcedSwallowedCancellation(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + req := request(JobSlow, SlowPayload{SleepMillis: 30000, IgnoreCtx: false, SwallowCancel: true}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + req.Policy = exec.NewPolicy(exec.GracePeriod(3 * time.Second)) + + res, err := newExecutor(t).Run(context.Background(), req) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Status != exec.StatusTimeout { + t.Errorf("Status = %q, want %q — a handler swallowing context.Canceled must not turn a blown deadline into ok", + res.Status, exec.StatusTimeout) + } +} + +func testPanicIsolated(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + // Reaching this line at all is half the assertion: a rung claiming + // IsolatesPanic must not let the handler's panic unwind into the + // caller and fail the test binary. + res, err := newExecutor(t).Run(context.Background(), request(JobPanic, struct{}{})) + if err != nil { + return // a launch-shaped error is acceptable + } + if res.Status != exec.StatusKilled && res.Status != exec.StatusHandlerError { + t.Errorf("Status = %q, want killed or handler_error for a panicking handler", res.Status) + } +} + +func testUsageReported(t *testing.T, newExecutor func(*testing.T) exec.Executor) { + res, err := newExecutor(t).Run(context.Background(), + request(JobSlow, SlowPayload{SleepMillis: 50})) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Usage.PeakRSS <= 0 { + t.Errorf("Usage.PeakRSS = %d, want > 0", res.Usage.PeakRSS) + } +} diff --git a/exec/exectest/suite_test.go b/exec/exectest/suite_test.go new file mode 100644 index 0000000..fccd484 --- /dev/null +++ b/exec/exectest/suite_test.go @@ -0,0 +1,83 @@ +package exectest_test + +import ( + "strings" + "testing" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/exectest" + "github.com/xraph/dispatch/exec/inproc" + "github.com/xraph/dispatch/job" +) + +func TestInProcessConformance(t *testing.T) { + exectest.RunSuite(t, "inprocess", func(*testing.T) exec.Executor { + r := job.NewRegistry() + for _, d := range exectest.Handlers() { + d.Register(r) + } + + return inproc.New(r) + }, exectest.Capabilities{ + // In-process enforces nothing: it cannot kill a handler that + // ignores cancellation, it has no separate address space to + // measure, and a panic propagates to the caller, which is what + // the worker's recover middleware is for. + Enforces: false, + ReportsUsage: false, + IsolatesPanic: false, + }) +} + +func TestCheckCapabilities(t *testing.T) { + // Capabilities describes variation between rungs, not an opt-out: a + // rung claiming out-of-process isolation cannot also claim it is unable + // to kill a handler, and so skip the test that proves it. + tests := []struct { + name string + level exec.Level + caps exectest.Capabilities + wantErr bool + }{ + { + name: "in-process may claim nothing", + level: exec.LevelNone, + caps: exectest.Capabilities{}, + }, + { + name: "out-of-process claiming both is consistent", + level: exec.LevelProcess, + caps: exectest.Capabilities{Enforces: true, IsolatesPanic: true}, + }, + { + name: "out-of-process disclaiming enforcement is not", + level: exec.LevelProcess, + caps: exectest.Capabilities{Enforces: false, IsolatesPanic: true}, + wantErr: true, + }, + { + name: "out-of-process disclaiming panic isolation is not", + level: exec.LevelProcess, + caps: exectest.Capabilities{Enforces: true, IsolatesPanic: false}, + wantErr: true, + }, + { + name: "a sandboxed rung is held to the same bar", + level: exec.LevelSandboxed, + caps: exectest.Capabilities{}, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := exectest.CheckCapabilities("subject", tt.level, tt.caps) + if gotErr := err != nil; gotErr != tt.wantErr { + t.Fatalf("CheckCapabilities() = %v, want error: %v", err, tt.wantErr) + } + if tt.wantErr && !strings.Contains(err.Error(), "subject") { + t.Errorf("CheckCapabilities() = %q, want it to name the executor", err) + } + }) + } +} diff --git a/exec/executor.go b/exec/executor.go new file mode 100644 index 0000000..b164d12 --- /dev/null +++ b/exec/executor.go @@ -0,0 +1,40 @@ +package exec + +import ( + "context" + + "github.com/xraph/dispatch/id" +) + +// Executor runs one job attempt. Implementations form an escalating ladder +// of isolation, and every one of them must pass the shared conformance +// suite in exec/exectest. +type Executor interface { + // Name identifies the executor in configuration, logs, and metrics. + Name() string + + // Level reports the isolation this executor actually provides, which + // is what Registry.Select matches a Policy against. + Level() Level + + // Run executes one attempt. + // + // The returned error is reserved for failures to launch — the handler + // never ran. A handler that ran and failed is reported through + // Result.Status, so the caller can tell a business failure from a + // dead sandbox without inspecting error text. + Run(ctx context.Context, req *Request) (*Result, error) + + // Reclaim releases sandboxes this worker leaked across a restart. The + // pool calls it once at startup, for every registered executor, and a + // failure is logged rather than fatal. + // + // A later phase runs the same sweep on the leader's behalf for workers + // the cluster has declared dead; that caller does not exist yet. + Reclaim(ctx context.Context, workerID id.WorkerID) error + + // Close releases the executor's own resources. The engine calls it for + // every registered executor when it stops, after the pool has finished + // its in-flight attempts. + Close() error +} diff --git a/exec/fingerprint.go b/exec/fingerprint.go new file mode 100644 index 0000000..332c2c7 --- /dev/null +++ b/exec/fingerprint.go @@ -0,0 +1,57 @@ +package exec + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "runtime/debug" + "sort" +) + +// FingerprintOf derives a stable identifier for a handler set and the build +// that contains it. +// +// A sandbox verifies this before running anything. When the sandbox re-execs +// the worker's own binary the check always passes and costs one comparison; +// its purpose is the Policy.Image override, where a stale image would +// otherwise run an old handler and report success. Drift becomes an +// immediate, correctly-classified launch failure instead of a silent wrong +// answer. +func FingerprintOf(names []string, revision string) string { + sorted := make([]string, len(names)) + copy(sorted, names) + sort.Strings(sorted) + + h := sha256.New() + // Length-prefix every element. Joining on a separator would let a + // handler named "a\nb" hash identically to the pair {"a", "b"}. + fmt.Fprintf(h, "%d:%s\n", len(revision), revision) + for _, n := range sorted { + fmt.Fprintf(h, "%d:%s\n", len(n), n) + } + + return hex.EncodeToString(h.Sum(nil)) +} + +// Fingerprint derives the identifier for a handler set using this binary's +// VCS revision. When the revision is unavailable — a build without VCS +// stamping — it falls back to the empty revision, so the fingerprint still +// covers the handler names. +func Fingerprint(names []string) string { + return FingerprintOf(names, buildRevision()) +} + +// buildRevision returns the VCS revision this binary was built from. +func buildRevision() string { + info, ok := debug.ReadBuildInfo() + if !ok { + return "" + } + for _, s := range info.Settings { + if s.Key == "vcs.revision" { + return s.Value + } + } + + return "" +} diff --git a/exec/fingerprint_test.go b/exec/fingerprint_test.go new file mode 100644 index 0000000..59ea145 --- /dev/null +++ b/exec/fingerprint_test.go @@ -0,0 +1,51 @@ +package exec_test + +import ( + "testing" + + "github.com/xraph/dispatch/exec" +) + +func TestFingerprintOf_StableAcrossOrder(t *testing.T) { + a := exec.FingerprintOf([]string{"b.job", "a.job", "c.job"}, "abc123") + b := exec.FingerprintOf([]string{"a.job", "b.job", "c.job"}, "abc123") + + if a != b { + t.Errorf("fingerprint depends on order: %q != %q", a, b) + } +} + +func TestFingerprintOf_ChangesWithNames(t *testing.T) { + a := exec.FingerprintOf([]string{"a.job"}, "abc123") + b := exec.FingerprintOf([]string{"a.job", "b.job"}, "abc123") + + if a == b { + t.Error("fingerprint did not change when a handler was added") + } +} + +func TestFingerprintOf_ChangesWithRevision(t *testing.T) { + a := exec.FingerprintOf([]string{"a.job"}, "abc123") + b := exec.FingerprintOf([]string{"a.job"}, "def456") + + if a == b { + t.Error("fingerprint did not change with the build revision") + } +} + +func TestFingerprintOf_DoesNotCollideOnSeparatorAmbiguity(t *testing.T) { + // {"a", "b"} and {"a\nb"} must not hash the same, or a handler named + // with an embedded separator could impersonate a two-handler set. + a := exec.FingerprintOf([]string{"a", "b"}, "r") + b := exec.FingerprintOf([]string{"a\nb"}, "r") + + if a == b { + t.Error("separator ambiguity produced a collision") + } +} + +func TestFingerprintOf_Empty(t *testing.T) { + if got := exec.FingerprintOf(nil, "r"); got == "" { + t.Error("FingerprintOf(nil) = empty, want a hash") + } +} diff --git a/exec/inproc/doc.go b/exec/inproc/doc.go new file mode 100644 index 0000000..52ad3fa --- /dev/null +++ b/exec/inproc/doc.go @@ -0,0 +1,9 @@ +// Package inproc runs job handlers in the worker process. +// +// This is Dispatch's original behaviour and remains the default. It +// provides no isolation: the handler shares the worker's memory, +// credentials, file descriptors, and network. That is the right trade for +// handlers that do not touch untrusted bytes, where launching a process +// per job would be pure overhead, and the wrong one for anything parsing +// a customer upload with a memory-unsafe library. +package inproc diff --git a/exec/inproc/inproc.go b/exec/inproc/inproc.go new file mode 100644 index 0000000..331d989 --- /dev/null +++ b/exec/inproc/inproc.go @@ -0,0 +1,82 @@ +package inproc + +import ( + "context" + "errors" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// Name is the identifier this executor registers under. +const Name = "inprocess" + +// Executor runs handlers in the worker process. +type Executor struct { + registry *job.Registry +} + +var _ exec.Executor = (*Executor)(nil) + +// New creates an in-process executor backed by a handler registry. +func New(r *job.Registry) *Executor { + return &Executor{registry: r} +} + +// Name identifies the executor. +func (e *Executor) Name() string { return Name } + +// Level reports that this executor provides no isolation. +func (e *Executor) Level() exec.Level { return exec.LevelNone } + +// Run looks the handler up by name and calls it. +func (e *Executor) Run(ctx context.Context, req *exec.Request) (*exec.Result, error) { + if err := req.Validate(); err != nil { + return nil, err + } + + handler, ok := e.registry.Get(req.Name) + if !ok { + // The handler never ran, so this is a launch failure rather than + // a job failure, and must not consume the retry budget. + return &exec.Result{ + Status: exec.StatusLaunchFailed, + HandlerErr: "no handler registered for job " + req.Name, + }, nil + } + + start := time.Now() + err := handler(ctx, req.Payload) + elapsed := time.Since(start) + + res := &exec.Result{ + Status: exec.StatusOK, + Usage: exec.Usage{WallTime: elapsed}, + } + if err != nil { + res.Status = exec.StatusHandlerError + res.HandlerErr = err.Error() + // The handler ran here, so its error chain is still intact: keep it + // whole rather than flattening it to a string, or errors.Is against + // dispatch.ErrPermanent — and against every sentinel an extension + // owns — would stop matching the moment a job is routed through an + // executor. + res.Cause = err + // Also set the flag an out-of-process rung would have to send in + // the chain's place, so the worker reads permanence the same way + // whichever rung produced the Result. + res.Permanent = errors.Is(err, dispatch.ErrPermanent) + } + + return res, nil +} + +// Reclaim is a no-op. An in-process handler cannot outlive the worker +// that called it, so there is never anything to reclaim. +func (e *Executor) Reclaim(context.Context, id.WorkerID) error { return nil } + +// Close is a no-op. The executor owns no resources of its own. +func (e *Executor) Close() error { return nil } diff --git a/exec/inproc/inproc_test.go b/exec/inproc/inproc_test.go new file mode 100644 index 0000000..c7a37ab --- /dev/null +++ b/exec/inproc/inproc_test.go @@ -0,0 +1,235 @@ +package inproc_test + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/inproc" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +type payload struct { + Value int `json:"value"` +} + +func TestExecutor_Identity(t *testing.T) { + e := inproc.New(job.NewRegistry()) + + if got := e.Name(); got != "inprocess" { + t.Errorf("Name() = %q, want %q", got, "inprocess") + } + if got := e.Level(); got != exec.LevelNone { + t.Errorf("Level() = %v, want %v", got, exec.LevelNone) + } +} + +func TestExecutor_Run(t *testing.T) { + sentinel := errors.New("boom") + + tests := []struct { + name string + handler func(context.Context, payload) error + wantStatus exec.Status + wantErrMsg string + }{ + { + name: "success", + handler: func(context.Context, payload) error { return nil }, + wantStatus: exec.StatusOK, + }, + { + name: "handler error", + handler: func(context.Context, payload) error { return sentinel }, + wantStatus: exec.StatusHandlerError, + wantErrMsg: "boom", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r := job.NewRegistry() + job.NewDefinition("test.job", tt.handler).Register(r) + e := inproc.New(r) + + res, err := e.Run(context.Background(), &exec.Request{ + JobID: id.NewJobID(), + Name: "test.job", + Payload: []byte(`{"value":7}`), + }) + if err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } + if res.Status != tt.wantStatus { + t.Errorf("Status = %q, want %q", res.Status, tt.wantStatus) + } + if res.HandlerErr != tt.wantErrMsg { + t.Errorf("HandlerErr = %q, want %q", res.HandlerErr, tt.wantErrMsg) + } + }) + } +} + +func TestExecutor_RunKeepsTheHandlerErrorWhole(t *testing.T) { + // In-process there is no boundary to lose the chain at, so the Result + // carries the handler's error itself. Without this, errors.Is against + // dispatch.ErrPermanent — and against any sentinel an extension owns — + // stops matching the moment a job is routed through an executor. + sentinel := errors.New("upstream gone") + + r := job.NewRegistry() + job.NewDefinition("test.job", func(context.Context, payload) error { + return fmt.Errorf("fetch: %w", sentinel) + }).Register(r) + + res, err := inproc.New(r).Run(context.Background(), &exec.Request{ + JobID: id.NewJobID(), + Name: "test.job", + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !errors.Is(res.Cause, sentinel) { + t.Errorf("errors.Is(Cause, sentinel) = false, want true (Cause = %v)", res.Cause) + } + if res.Permanent { + t.Error("Permanent = true for an ordinary handler error, want false") + } + if !errors.Is(res.Err(), sentinel) { + t.Errorf("errors.Is(Err(), sentinel) = false, want true (Err() = %v)", res.Err()) + } +} + +func TestExecutor_RunFlagsPermanentFailures(t *testing.T) { + // The flag is what an out-of-process rung will have to send instead of + // an error chain, so the in-process rung computes it too and the worker + // reads permanence one way for every rung. + r := job.NewRegistry() + job.NewDefinition("test.job", func(context.Context, payload) error { + return fmt.Errorf("malformed payload: %w", dispatch.ErrPermanent) + }).Register(r) + + res, err := inproc.New(r).Run(context.Background(), &exec.Request{ + JobID: id.NewJobID(), + Name: "test.job", + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !res.Permanent { + t.Error("Permanent = false for an error wrapping dispatch.ErrPermanent, want true") + } + if !errors.Is(res.Err(), dispatch.ErrPermanent) { + t.Errorf("errors.Is(Err(), ErrPermanent) = false, want true (Err() = %v)", res.Err()) + } + if got, want := res.Err().Error(), "malformed payload: dispatch: permanent failure"; got != want { + t.Errorf("Err() = %q, want %q — the handler's own text, unframed", got, want) + } +} + +func TestExecutor_RunPassesPayload(t *testing.T) { + var got payload + r := job.NewRegistry() + job.NewDefinition("test.job", func(_ context.Context, p payload) error { + got = p + return nil + }).Register(r) + + _, err := inproc.New(r).Run(context.Background(), &exec.Request{ + JobID: id.NewJobID(), + Name: "test.job", + Payload: []byte(`{"value":42}`), + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got.Value != 42 { + t.Errorf("payload.Value = %d, want 42", got.Value) + } +} + +func TestExecutor_RunUnknownHandlerIsALaunchFailure(t *testing.T) { + // The handler never ran, so this must not consume the retry budget. + res, err := inproc.New(job.NewRegistry()).Run(context.Background(), &exec.Request{ + JobID: id.NewJobID(), + Name: "absent", + }) + if err != nil { + t.Fatalf("Run() error = %v, want a Result", err) + } + if res.Status != exec.StatusLaunchFailed { + t.Fatalf("Status = %q, want %q", res.Status, exec.StatusLaunchFailed) + } + if res.Status.CountsAgainstRetries() { + t.Error("an unknown handler must not consume the retry budget") + } +} + +func TestExecutor_RunInvalidRequest(t *testing.T) { + _, err := inproc.New(job.NewRegistry()).Run(context.Background(), &exec.Request{}) + if !errors.Is(err, exec.ErrInvalidRequest) { + t.Fatalf("Run() error = %v, want %v", err, exec.ErrInvalidRequest) + } +} + +func TestExecutor_RunCancelledContext(t *testing.T) { + r := job.NewRegistry() + job.NewDefinition("test.job", func(ctx context.Context, _ payload) error { + <-ctx.Done() + return ctx.Err() + }).Register(r) + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(10 * time.Millisecond) + cancel() + }() + + res, err := inproc.New(r).Run(ctx, &exec.Request{ + JobID: id.NewJobID(), + Name: "test.job", + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + // In-process cancellation is cooperative: the handler chose to + // return, so this is a handler error, not an enforced timeout. + if res.Status != exec.StatusHandlerError { + t.Errorf("Status = %q, want %q", res.Status, exec.StatusHandlerError) + } +} + +func TestExecutor_RunRecordsWallTime(t *testing.T) { + r := job.NewRegistry() + job.NewDefinition("test.job", func(context.Context, payload) error { + time.Sleep(5 * time.Millisecond) + return nil + }).Register(r) + + res, err := inproc.New(r).Run(context.Background(), &exec.Request{ + JobID: id.NewJobID(), + Name: "test.job", + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if res.Usage.WallTime <= 0 { + t.Errorf("Usage.WallTime = %v, want > 0", res.Usage.WallTime) + } +} + +func TestExecutor_ReclaimAndClose(t *testing.T) { + e := inproc.New(job.NewRegistry()) + + if err := e.Reclaim(context.Background(), id.NewWorkerID()); err != nil { + t.Errorf("Reclaim() = %v, want nil", err) + } + if err := e.Close(); err != nil { + t.Errorf("Close() = %v, want nil", err) + } +} diff --git a/exec/policy.go b/exec/policy.go new file mode 100644 index 0000000..c9166db --- /dev/null +++ b/exec/policy.go @@ -0,0 +1,117 @@ +package exec + +import ( + "fmt" + "time" +) + +// DefaultGracePeriod is how long a sandbox is given to exit after being +// asked politely, before it is killed outright. +const DefaultGracePeriod = 30 * time.Second + +// Level is the minimum isolation a job definition requires. The levels are +// ordered, so a deployment offering a stronger level satisfies a definition +// asking for a weaker one. +type Level int + +const ( + // LevelNone runs the handler in the worker process. This is the + // default and it provides no isolation of any kind. + LevelNone Level = iota + + // LevelProcess runs the handler in a separate address space, so an + // exploited parser cannot read the worker's credentials. + LevelProcess + + // LevelSandboxed adds mount, network, PID, and user namespaces, a + // seccomp filter, and dropped capabilities. + LevelSandboxed + + // LevelVM adds an independent kernel — gVisor or Kata — so a Linux + // privilege escalation is not by itself an escape. + LevelVM +) + +// String renders the level for configuration, logs, and errors. +func (l Level) String() string { + switch l { + case LevelNone: + return "none" + case LevelProcess: + return "process" + case LevelSandboxed: + return "sandboxed" + case LevelVM: + return "vm" + default: + return fmt.Sprintf("Level(%d)", int(l)) + } +} + +// Policy is a job definition's execution declaration. It states the minimum +// isolation the handler requires, not the executor it runs on: which rung +// satisfies the requirement is a deployment decision. +type Policy struct { + // Level is the minimum isolation required. + Level Level + + // GracePeriod is how long the sandbox has to exit after SIGTERM + // before it is killed. + GracePeriod time.Duration + + // AllowDowngrade permits running at a weaker level than Level when + // the deployment cannot provide it. Without it, a deployment that + // cannot satisfy the policy fails at registration rather than + // silently running the handler unisolated. + AllowDowngrade bool + + // Image overrides the container image for out-of-process rungs. + // Empty means the worker's own image, which is the correct default + // because the sandbox re-execs the same binary. + Image string +} + +// PolicyOption configures a Policy. +type PolicyOption func(*Policy) + +// NewPolicy builds a Policy from options, starting from the defaults: +// no isolation and a 30-second grace period. +func NewPolicy(opts ...PolicyOption) Policy { + p := Policy{ + Level: LevelNone, + GracePeriod: DefaultGracePeriod, + } + for _, opt := range opts { + opt(&p) + } + + return p +} + +// Isolate sets the minimum isolation level the handler requires. +func Isolate(l Level) PolicyOption { + return func(p *Policy) { p.Level = l } +} + +// GracePeriod sets how long the sandbox has to exit cleanly after being +// signalled. Non-positive durations are ignored, because a zero grace +// period reduces the kill ladder to an immediate SIGKILL and loses any +// chance of a clean shutdown. +func GracePeriod(d time.Duration) PolicyOption { + return func(p *Policy) { + if d > 0 { + p.GracePeriod = d + } + } +} + +// AllowDowngrade permits running below the declared level when the +// deployment cannot satisfy it. +func AllowDowngrade() PolicyOption { + return func(p *Policy) { p.AllowDowngrade = true } +} + +// Image overrides the container image used by out-of-process rungs. +func Image(ref string) PolicyOption { + return func(p *Policy) { p.Image = ref } +} diff --git a/exec/policy_test.go b/exec/policy_test.go new file mode 100644 index 0000000..875aae0 --- /dev/null +++ b/exec/policy_test.go @@ -0,0 +1,80 @@ +package exec_test + +import ( + "testing" + "time" + + "github.com/xraph/dispatch/exec" +) + +func TestNewPolicy_Defaults(t *testing.T) { + p := exec.NewPolicy() + + if p.Level != exec.LevelNone { + t.Errorf("Level = %v, want %v", p.Level, exec.LevelNone) + } + if p.GracePeriod != 30*time.Second { + t.Errorf("GracePeriod = %v, want %v", p.GracePeriod, 30*time.Second) + } + if p.AllowDowngrade { + t.Error("AllowDowngrade = true, want false") + } + if p.Image != "" { + t.Errorf("Image = %q, want empty", p.Image) + } +} + +func TestNewPolicy_Options(t *testing.T) { + p := exec.NewPolicy( + exec.Isolate(exec.LevelSandboxed), + exec.GracePeriod(90*time.Second), + exec.AllowDowngrade(), + exec.Image("twinos/worker:v3"), + ) + + if p.Level != exec.LevelSandboxed { + t.Errorf("Level = %v, want %v", p.Level, exec.LevelSandboxed) + } + if p.GracePeriod != 90*time.Second { + t.Errorf("GracePeriod = %v, want %v", p.GracePeriod, 90*time.Second) + } + if !p.AllowDowngrade { + t.Error("AllowDowngrade = false, want true") + } + if p.Image != "twinos/worker:v3" { + t.Errorf("Image = %q, want %q", p.Image, "twinos/worker:v3") + } +} + +func TestNewPolicy_NonPositiveGracePeriodKeepsDefault(t *testing.T) { + // A zero or negative grace period would make the kill ladder in later + // rungs degenerate into an immediate SIGKILL, losing every chance of a + // clean shutdown. Reject it at construction rather than at kill time. + for _, d := range []time.Duration{0, -1 * time.Second} { + p := exec.NewPolicy(exec.GracePeriod(d)) + if p.GracePeriod != 30*time.Second { + t.Errorf("GracePeriod(%v) = %v, want default %v", d, p.GracePeriod, 30*time.Second) + } + } +} + +func TestLevel_String(t *testing.T) { + tests := []struct { + level exec.Level + want string + }{ + {exec.LevelNone, "none"}, + {exec.LevelProcess, "process"}, + {exec.LevelSandboxed, "sandboxed"}, + {exec.LevelVM, "vm"}, + {exec.Level(99), "Level(99)"}, + } + + for _, tt := range tests { + t.Run(tt.want, func(t *testing.T) { + if got := tt.level.String(); got != tt.want { + t.Errorf("String() = %q, want %q", got, tt.want) + } + }) + } +} diff --git a/exec/registry.go b/exec/registry.go new file mode 100644 index 0000000..d551bc4 --- /dev/null +++ b/exec/registry.go @@ -0,0 +1,139 @@ +package exec + +import ( + "errors" + "fmt" + "sort" + "sync" +) + +// ErrNoExecutor marks a policy no configured executor can satisfy. +var ErrNoExecutor = errors.New("no executor satisfies the policy") + +// Registry holds the executors a deployment has configured and matches +// job policies against them. +// +// It is safe for concurrent use, though in practice it is built once at +// startup and only read afterwards. +type Registry struct { + mu sync.RWMutex + def Executor + byName map[string]Executor +} + +// NewRegistry creates a registry with a default executor, which is the one +// used by any job that declares no isolation requirement. +func NewRegistry(def Executor) *Registry { + r := &Registry{ + def: def, + byName: make(map[string]Executor), + } + if def != nil { + r.byName[def.Name()] = def + } + + return r +} + +// Add registers an executor, replacing any existing one with the same name. +func (r *Registry) Add(e Executor) { + if e == nil { + return + } + + r.mu.Lock() + defer r.mu.Unlock() + r.byName[e.Name()] = e +} + +// Default returns the executor used when a job declares no requirement. +func (r *Registry) Default() Executor { + r.mu.RLock() + defer r.mu.RUnlock() + + return r.def +} + +// Executors returns every registered executor, ordered by name so callers +// and tests see a stable list. +func (r *Registry) Executors() []Executor { + r.mu.RLock() + defer r.mu.RUnlock() + + names := make([]string, 0, len(r.byName)) + for n := range r.byName { + names = append(names, n) + } + sort.Strings(names) + + out := make([]Executor, 0, len(names)) + for _, n := range names { + out = append(out, r.byName[n]) + } + + return out +} + +// Select returns the executor that should run a job with this policy. +// +// It picks the weakest executor that still satisfies the declared level, +// so a job needing a separate process is not handed a Kubernetes pod +// merely because one is configured. When nothing satisfies the policy the +// call fails rather than quietly running the handler with less isolation +// than it asked for — unless the policy opted into a downgrade. +func (r *Registry) Select(p Policy) (Executor, error) { + r.mu.RLock() + defer r.mu.RUnlock() + + if p.Level == LevelNone { + if r.def == nil { + return nil, fmt.Errorf("%w: no default executor configured", ErrNoExecutor) + } + + return r.def, nil + } + + var best Executor + for _, e := range r.byName { + if e.Level() < p.Level { + continue + } + if best == nil || e.Level() < best.Level() || + (e.Level() == best.Level() && e.Name() < best.Name()) { + best = e + } + } + if best != nil { + return best, nil + } + + if p.AllowDowngrade && r.def != nil { + return r.def, nil + } + + return nil, fmt.Errorf( + "%w: policy requires level %s, configured executors are %s", + ErrNoExecutor, p.Level, r.describeLocked(), + ) +} + +// describeLocked renders the configured executors for an error message. +// The caller must hold at least a read lock. +func (r *Registry) describeLocked() string { + if len(r.byName) == 0 { + return "(none)" + } + + names := make([]string, 0, len(r.byName)) + for n, e := range r.byName { + names = append(names, fmt.Sprintf("%s(%s)", n, e.Level())) + } + sort.Strings(names) + + out := names[0] + for _, n := range names[1:] { + out += ", " + n + } + + return out +} diff --git a/exec/registry_test.go b/exec/registry_test.go new file mode 100644 index 0000000..6aec24c --- /dev/null +++ b/exec/registry_test.go @@ -0,0 +1,128 @@ +package exec_test + +import ( + "context" + "errors" + "testing" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/id" +) + +// fakeExecutor is a minimal Executor for registry tests. +type fakeExecutor struct { + name string + level exec.Level +} + +func (f fakeExecutor) Name() string { return f.name } +func (f fakeExecutor) Level() exec.Level { return f.level } + +func (f fakeExecutor) Run(context.Context, *exec.Request) (*exec.Result, error) { + return &exec.Result{Status: exec.StatusOK}, nil +} + +func (f fakeExecutor) Reclaim(context.Context, id.WorkerID) error { return nil } +func (f fakeExecutor) Close() error { return nil } + +func TestRegistry_SelectPicksWeakestSufficient(t *testing.T) { + r := exec.NewRegistry(fakeExecutor{name: "inprocess", level: exec.LevelNone}) + r.Add(fakeExecutor{name: "subprocess", level: exec.LevelProcess}) + r.Add(fakeExecutor{name: "k8s", level: exec.LevelVM}) + + // A job needing process isolation must not be handed the Kubernetes + // rung when a cheaper sufficient one exists. + got, err := r.Select(exec.NewPolicy(exec.Isolate(exec.LevelProcess))) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + if got.Name() != "subprocess" { + t.Errorf("Select() = %q, want %q", got.Name(), "subprocess") + } +} + +func TestRegistry_SelectEscalatesWhenExactRungAbsent(t *testing.T) { + r := exec.NewRegistry(fakeExecutor{name: "inprocess", level: exec.LevelNone}) + r.Add(fakeExecutor{name: "k8s", level: exec.LevelVM}) + + got, err := r.Select(exec.NewPolicy(exec.Isolate(exec.LevelSandboxed))) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + if got.Name() != "k8s" { + t.Errorf("Select() = %q, want %q", got.Name(), "k8s") + } +} + +func TestRegistry_SelectRefusesSilentDowngrade(t *testing.T) { + r := exec.NewRegistry(fakeExecutor{name: "inprocess", level: exec.LevelNone}) + + _, err := r.Select(exec.NewPolicy(exec.Isolate(exec.LevelSandboxed))) + if !errors.Is(err, exec.ErrNoExecutor) { + t.Fatalf("Select() error = %v, want %v", err, exec.ErrNoExecutor) + } +} + +func TestRegistry_SelectAllowsExplicitDowngrade(t *testing.T) { + r := exec.NewRegistry(fakeExecutor{name: "inprocess", level: exec.LevelNone}) + + got, err := r.Select(exec.NewPolicy( + exec.Isolate(exec.LevelSandboxed), + exec.AllowDowngrade(), + )) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + if got.Name() != "inprocess" { + t.Errorf("Select() = %q, want %q", got.Name(), "inprocess") + } +} + +func TestRegistry_SelectDefaultForLevelNone(t *testing.T) { + r := exec.NewRegistry(fakeExecutor{name: "inprocess", level: exec.LevelNone}) + r.Add(fakeExecutor{name: "subprocess", level: exec.LevelProcess}) + + got, err := r.Select(exec.NewPolicy()) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + if got.Name() != "inprocess" { + t.Errorf("Select() = %q, want the default %q", got.Name(), "inprocess") + } +} + +func TestRegistry_AddReplacesSameName(t *testing.T) { + r := exec.NewRegistry(fakeExecutor{name: "inprocess", level: exec.LevelNone}) + r.Add(fakeExecutor{name: "subprocess", level: exec.LevelProcess}) + r.Add(fakeExecutor{name: "subprocess", level: exec.LevelSandboxed}) + + if n := len(r.Executors()); n != 2 { + t.Fatalf("len(Executors()) = %d, want 2", n) + } + got, err := r.Select(exec.NewPolicy(exec.Isolate(exec.LevelSandboxed))) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + if got.Name() != "subprocess" { + t.Errorf("Select() = %q, want %q", got.Name(), "subprocess") + } +} + +func TestRegistry_SelectTieBreaksByName(t *testing.T) { + // Two executors at the same level: the choice must be deterministic, + // because Go randomises map iteration and a job must not land on a + // different isolation rung from one run to the next. + r := exec.NewRegistry(fakeExecutor{name: "inprocess", level: exec.LevelNone}) + r.Add(fakeExecutor{name: "zeta", level: exec.LevelProcess}) + r.Add(fakeExecutor{name: "alpha", level: exec.LevelProcess}) + + for range 20 { + got, err := r.Select(exec.NewPolicy(exec.Isolate(exec.LevelProcess))) + if err != nil { + t.Fatalf("Select() error = %v", err) + } + if got.Name() != "alpha" { + t.Fatalf("Select() = %q, want %q", got.Name(), "alpha") + } + } +} diff --git a/exec/request.go b/exec/request.go new file mode 100644 index 0000000..3e065b5 --- /dev/null +++ b/exec/request.go @@ -0,0 +1,113 @@ +package exec + +import ( + "errors" + "fmt" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" +) + +// ErrInvalidRequest marks a Request that cannot be executed as given. +var ErrInvalidRequest = errors.New("invalid execution request") + +// InputSlot maps a declared input name to its location within InputDir. +// The path is relative, so the same Request describes the inputs whether +// the sandbox mounts them at /dispatch/in or reads them where they lie. +type InputSlot struct { + Name string + Path string +} + +// PriorOutput is an artifact an earlier attempt of this job committed. +// +// A sandbox keeps its artifact rows in memory and cannot query the store, +// so without these Accessor.Existing would always answer "no" and a +// retried handler would redo work it had already finished. The output +// would still be correct, which is exactly why this is worth carrying +// explicitly: nothing would fail, it would just quietly cost twice. +type PriorOutput struct { + Name string + Ref artifact.Ref +} + +// Request is one execution attempt, fully described. Everything the +// handler needs crosses the boundary in this value. The environment is +// not inherited wholesale from the worker: an out-of-process rung +// constructs the child's environment from Env plus a small fixed +// allowlist (PATH, HOME, TMPDIR) copied from the worker's own — see Env +// below and exec/subprocess.Executor.buildEnv. +type Request struct { + JobID id.JobID + Name string + Payload []byte + Attempt int + + // Deadline is when the attempt must be killed. Zero means no deadline. + Deadline time.Time + + // Fingerprint identifies the handler set the caller expects. + Fingerprint string + + // InputDir holds staged inputs and is read-only to the handler. + InputDir string + // OutputDir is where the handler writes artifacts. + OutputDir string + + Inputs []InputSlot + PriorOutputs []PriorOutput + + Policy Policy + + // ResourceLimits is the job's resolved enforcement ceiling + // (job.Job.ResourceLimits, see job.WithResourceLimits), carried + // across the execution boundary so a rung that can enforce something + // has the per-job numbers to enforce it with. A key absent or zero + // here means the job declared no ceiling for that dimension; it is + // not this type's business to say what a rung does about that — + // exec/subprocess.Executor falls back to its own deployment-wide + // default in that case, but exec itself has no opinion. + ResourceLimits resource.Set + + // ScopeAppID and ScopeOrgID label the attempt for logs and metrics. + // They are identifiers, never credentials. + ScopeAppID string + ScopeOrgID string + + // Env is passed to out-of-process rungs. It is not the worker's + // os.Environ() handed through: exec/subprocess.Executor.buildEnv + // builds the child's environment from this map plus its own + // configured base, never starting from the worker's full environment. + // It does still copy a fixed allowlist of PATH, HOME, and TMPDIR from + // the worker's own environment ahead of this map — HOME in + // particular is what locates ~/.aws and similar credential paths, so + // the dedicated uid this rung requires, not environment exclusion + // alone, is what actually keeps the child from reading them. + Env map[string]string +} + +// Validate reports whether the request is well formed. +func (r *Request) Validate() error { + if r.Name == "" { + return fmt.Errorf("%w: empty job name", ErrInvalidRequest) + } + if r.Attempt < 0 { + return fmt.Errorf("%w: negative attempt %d", ErrInvalidRequest, r.Attempt) + } + + return nil +} + +// InputPath returns the relative path of a declared input, or an empty +// string when the request carries no such input. +func (r *Request) InputPath(name string) string { + for _, in := range r.Inputs { + if in.Name == name { + return in.Path + } + } + + return "" +} diff --git a/exec/request_test.go b/exec/request_test.go new file mode 100644 index 0000000..bd93f69 --- /dev/null +++ b/exec/request_test.go @@ -0,0 +1,77 @@ +package exec_test + +import ( + "errors" + "testing" + "time" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/id" +) + +func validRequest() *exec.Request { + return &exec.Request{ + JobID: id.NewJobID(), + Name: "tessellate.model", + Payload: []byte(`{"detail":3}`), + Attempt: 0, + Deadline: time.Now().Add(time.Hour), + } +} + +func TestRequest_Validate(t *testing.T) { + tests := []struct { + name string + mutate func(*exec.Request) + wantErr error + }{ + { + name: "valid", + mutate: func(*exec.Request) {}, + }, + { + name: "missing name", + mutate: func(r *exec.Request) { r.Name = "" }, + wantErr: exec.ErrInvalidRequest, + }, + { + name: "negative attempt", + mutate: func(r *exec.Request) { r.Attempt = -1 }, + wantErr: exec.ErrInvalidRequest, + }, + { + name: "zero deadline is allowed", + mutate: func(r *exec.Request) { r.Deadline = time.Time{} }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := validRequest() + tt.mutate(req) + + err := req.Validate() + if tt.wantErr == nil { + if err != nil { + t.Fatalf("Validate() = %v, want nil", err) + } + return + } + if !errors.Is(err, tt.wantErr) { + t.Fatalf("Validate() = %v, want %v", err, tt.wantErr) + } + }) + } +} + +func TestRequest_InputPathLookup(t *testing.T) { + req := validRequest() + req.Inputs = []exec.InputSlot{{Name: "model", Path: "model/scene.ifc"}} + + if got := req.InputPath("model"); got != "model/scene.ifc" { + t.Errorf("InputPath(model) = %q, want %q", got, "model/scene.ifc") + } + if got := req.InputPath("absent"); got != "" { + t.Errorf("InputPath(absent) = %q, want empty", got) + } +} diff --git a/exec/result.go b/exec/result.go new file mode 100644 index 0000000..0253905 --- /dev/null +++ b/exec/result.go @@ -0,0 +1,185 @@ +package exec + +import ( + "errors" + "fmt" + "time" +) + +// Status sentinels, so callers can classify a failure with errors.Is +// rather than by comparing strings. +var ( + // ErrHandler marks an error the handler itself returned. + ErrHandler = errors.New("handler error") + // ErrTimeout marks an attempt killed for exceeding its deadline. + ErrTimeout = errors.New("execution timeout") + // ErrOOMKilled marks an attempt killed for exceeding a memory limit. + ErrOOMKilled = errors.New("out of memory") + // ErrKilled marks an attempt whose process died on a signal. + ErrKilled = errors.New("killed by signal") + // ErrLaunchFailed marks a sandbox that never started. + ErrLaunchFailed = errors.New("launch failed") +) + +// Usage records what an attempt consumed. Every rung above in-process +// accounts these anyway, so collecting them costs nothing and gives the +// resource model its measurements. +type Usage struct { + WallTime time.Duration + CPUTime time.Duration + PeakRSS int64 + DiskWritten int64 +} + +// OutputFile describes one artifact the handler produced, as claimed by +// the sandbox. The worker never reads this to decide what to commit: it +// is driven entirely by collectOutputEntries walking req.OutputDir on +// disk (worker/runner.go's commitOutputs), so a claim listed here that +// does not correspond to a real file on disk is simply never consulted, +// not checked and rejected. This field exists for whatever future +// consumer wants the sandbox's own account of what it wrote — logging, a +// mismatch warning — not as an input to the commit decision. +type OutputFile struct { + Name string + Size int64 + Hash string + ContentType string +} + +// Result reports how one execution attempt ended. +type Result struct { + // Status classifies the outcome. + Status Status + + // HandlerErr is the handler's error string, or a diagnostic for a + // launch failure. Empty on success. + HandlerErr string + + // ExitCode is the sandbox process's exit status, where one applies. + ExitCode int + + // Signal is the signal number that killed the process, or zero. + // Stored as an int rather than a syscall.Signal so this leaf package + // stays free of syscall. + Signal int + + // Usage records what the attempt consumed. + Usage Usage + + // Outputs lists the artifacts the sandbox claims to have written. + Outputs []OutputFile + + // Permanent means retrying cannot change the outcome, so the job skips + // its remaining attempts. This is the wire-visible half of permanence: + // a rung that ran the handler in another process cannot send back a Go + // error chain, but it can send back this flag. + Permanent bool + + // Cause is the handler's own error, kept whole. It exists so an + // in-process attempt loses nothing: errors.Is and errors.As against a + // caller's sentinels keep working through the exec layer. + // + // It is deliberately not serialised — an error chain does not survive a + // process boundary — which is why Permanent exists alongside it. A rung + // that marshals a Result leaves this nil and sets Permanent instead. + Cause error `json:"-"` +} + +// Err converts a Result into the error the worker propagates. It returns +// nil for StatusOK and an *Error otherwise. +func (r *Result) Err() error { + if r == nil || r.Status == StatusOK { + return nil + } + + return &Error{ + Status: r.Status, + Msg: r.HandlerErr, + ExitCode: r.ExitCode, + Signal: r.Signal, + Permanent: r.Permanent, + Cause: r.Cause, + } +} + +// Error is a failed execution attempt. It carries the Status so retry +// policy can branch on how the attempt failed rather than parsing text. +type Error struct { + Status Status + Msg string + ExitCode int + Signal int + + // Permanent means the job must not be retried, however the attempt was + // carried. The worker treats it exactly as it treats an error wrapping + // dispatch.ErrPermanent. + Permanent bool + + // Cause is the handler's own error when the attempt ran in this + // process. It is what errors.Is and errors.As reach through Unwrap. + Cause error +} + +// Error implements the error interface. +// +// When a cause survived — an in-process attempt — the handler's own text is +// returned verbatim. The job's LastError, its DLQ entry, and the worker's +// logs then read exactly as they did before execution isolation existed, +// and the exec framing is not stacked on top of an error the caller already +// understands. Only a failure with no cause to speak for it, which is every +// out-of-process shape, is rendered with the status. +func (e *Error) Error() string { + if e.Cause != nil { + return e.Cause.Error() + } + + switch { + case e.Msg != "": + return fmt.Sprintf("dispatch/exec: %s: %s", e.Status, e.Msg) + case e.Signal != 0: + return fmt.Sprintf("dispatch/exec: %s: signal %d", e.Status, e.Signal) + case e.ExitCode != 0: + return fmt.Sprintf("dispatch/exec: %s: exit %d", e.Status, e.ExitCode) + default: + return fmt.Sprintf("dispatch/exec: %s", e.Status) + } +} + +// Unwrap returns the status sentinel together with the handler's own error +// when one survived, so errors.Is reaches ErrHandler and errors.Is/errors.As +// reach the caller's own sentinels and error types through the same value. +// +// The multi-error form (Go 1.20) is what lets both hold at once. Returning +// only the sentinel would silently destroy user error identity for every +// in-process attempt, and dispatch.ErrPermanent with it. +func (e *Error) Unwrap() []error { + var errs []error + if sentinel := e.sentinel(); sentinel != nil { + errs = append(errs, sentinel) + } + if e.Cause != nil { + errs = append(errs, e.Cause) + } + + return errs +} + +// sentinel maps the status onto its package-level error value. +func (e *Error) sentinel() error { + switch e.Status { + case StatusHandlerError: + return ErrHandler + case StatusTimeout: + return ErrTimeout + case StatusOOMKilled: + return ErrOOMKilled + case StatusKilled: + return ErrKilled + case StatusLaunchFailed: + return ErrLaunchFailed + case StatusOK: + return nil + default: + return nil + } +} diff --git a/exec/result_test.go b/exec/result_test.go new file mode 100644 index 0000000..9344d5f --- /dev/null +++ b/exec/result_test.go @@ -0,0 +1,206 @@ +package exec_test + +import ( + "errors" + "fmt" + "strings" + "testing" + "time" + + "github.com/xraph/dispatch/exec" +) + +func TestResult_Err(t *testing.T) { + tests := []struct { + name string + result exec.Result + wantNil bool + wantIs error + wantText string + }{ + { + name: "ok returns nil", + result: exec.Result{Status: exec.StatusOK}, + wantNil: true, + }, + { + name: "handler error carries the handler message", + result: exec.Result{Status: exec.StatusHandlerError, HandlerErr: "bad IFC header"}, + wantIs: exec.ErrHandler, + wantText: "bad IFC header", + }, + { + name: "exit code without a handler message", + result: exec.Result{Status: exec.StatusHandlerError, ExitCode: 137}, + wantIs: exec.ErrHandler, + wantText: "exit 137", + }, + { + name: "timeout", + result: exec.Result{Status: exec.StatusTimeout}, + wantIs: exec.ErrTimeout, + wantText: "timeout", + }, + { + name: "oom killed", + result: exec.Result{Status: exec.StatusOOMKilled}, + wantIs: exec.ErrOOMKilled, + wantText: "oom_killed", + }, + { + name: "killed by signal", + result: exec.Result{Status: exec.StatusKilled, Signal: 11}, + wantIs: exec.ErrKilled, + wantText: "signal 11", + }, + { + name: "launch failed", + result: exec.Result{Status: exec.StatusLaunchFailed, HandlerErr: "image pull backoff"}, + wantIs: exec.ErrLaunchFailed, + wantText: "image pull backoff", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.result.Err() + + if tt.wantNil { + if err != nil { + t.Fatalf("Err() = %v, want nil", err) + } + return + } + if err == nil { + t.Fatal("Err() = nil, want error") + } + if !errors.Is(err, tt.wantIs) { + t.Errorf("errors.Is(%v, %v) = false, want true", err, tt.wantIs) + } + if !strings.Contains(err.Error(), tt.wantText) { + t.Errorf("Err() = %q, want it to contain %q", err.Error(), tt.wantText) + } + }) + } +} + +func TestResult_ErrPreservesTheCause(t *testing.T) { + // The point of Cause: an in-process attempt must lose nothing. A rung + // that flattened the handler's error to a string would break every + // errors.Is and errors.As an extension or the worker performs on it. + sentinel := errors.New("upstream unavailable") + cause := fmt.Errorf("fetch config: %w", sentinel) + + err := (&exec.Result{ + Status: exec.StatusHandlerError, + HandlerErr: cause.Error(), + Cause: cause, + }).Err() + + if !errors.Is(err, sentinel) { + t.Errorf("errors.Is(%v, sentinel) = false, want true", err) + } + if !errors.Is(err, exec.ErrHandler) { + t.Errorf("errors.Is(%v, ErrHandler) = false, want true — the status sentinel must still match", err) + } + // The handler's own text, with no exec framing: this is what lands in + // job.LastError, the DLQ entry, and the logs. + if got, want := err.Error(), cause.Error(); got != want { + t.Errorf("Error() = %q, want %q", got, want) + } + + // errors.As has to reach a caller's own error type through the same + // value, since that is how extensions inspect a failure. + var target *typedError + typed := &typedError{field: "name"} + err = (&exec.Result{Status: exec.StatusHandlerError, Cause: fmt.Errorf("wrapped: %w", typed)}).Err() + if !errors.As(err, &target) { + t.Fatalf("errors.As(%v, **typedError) = false, want true", err) + } + if target.field != "name" { + t.Errorf("target.field = %q, want %q", target.field, "name") + } +} + +// typedError is a caller-defined error type, standing in for the ones an +// extension matches with errors.As. +type typedError struct{ field string } + +func (e *typedError) Error() string { return "invalid field " + e.field } + +func TestResult_ErrCarriesPermanence(t *testing.T) { + // Permanent is the wire-visible signal: a rung that ran the handler in + // another process cannot send back an error chain, so this flag is how + // it declines the retry schedule. + err := (&exec.Result{ + Status: exec.StatusHandlerError, + HandlerErr: "malformed input", + Permanent: true, + }).Err() + + var execErr *exec.Error + if !errors.As(err, &execErr) { + t.Fatalf("errors.As(%v, **exec.Error) = false, want true", err) + } + if !execErr.Permanent { + t.Error("Error.Permanent = false, want true") + } + + notPermanent := (&exec.Result{Status: exec.StatusHandlerError, HandlerErr: "transient"}).Err() + if !errors.As(notPermanent, &execErr) { + t.Fatalf("errors.As(%v, **exec.Error) = false, want true", notPermanent) + } + if execErr.Permanent { + t.Error("Error.Permanent = true for a plain handler error, want false") + } +} + +func TestStatus_CountsAgainstRetries(t *testing.T) { + // A launch failure is infrastructure, not a property of the work. + // Letting it consume the retry budget means one bad node sends real + // customer work to the DLQ. + tests := []struct { + status exec.Status + want bool + }{ + {exec.StatusOK, false}, + {exec.StatusHandlerError, true}, + {exec.StatusTimeout, true}, + {exec.StatusOOMKilled, true}, + {exec.StatusKilled, true}, + {exec.StatusLaunchFailed, false}, + } + + for _, tt := range tests { + t.Run(string(tt.status), func(t *testing.T) { + if got := tt.status.CountsAgainstRetries(); got != tt.want { + t.Errorf("CountsAgainstRetries() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestStatus_IsFailure(t *testing.T) { + if exec.StatusOK.IsFailure() { + t.Error("StatusOK.IsFailure() = true, want false") + } + for _, s := range []exec.Status{ + exec.StatusHandlerError, exec.StatusTimeout, + exec.StatusOOMKilled, exec.StatusKilled, exec.StatusLaunchFailed, + } { + if !s.IsFailure() { + t.Errorf("%s.IsFailure() = false, want true", s) + } + } +} + +func TestUsage_ZeroValueIsUsable(t *testing.T) { + var u exec.Usage + if u.WallTime != 0 || u.CPUTime != 0 || u.PeakRSS != 0 || u.DiskWritten != 0 { + t.Errorf("zero Usage = %+v, want all zero", u) + } + u.WallTime = time.Second + if u.WallTime != time.Second { + t.Errorf("WallTime = %v, want %v", u.WallTime, time.Second) + } +} diff --git a/exec/shim/accessor.go b/exec/shim/accessor.go new file mode 100644 index 0000000..3c4fd2d --- /dev/null +++ b/exec/shim/accessor.go @@ -0,0 +1,163 @@ +package shim + +import ( + "context" + "fmt" + "io" + "os" + "path/filepath" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/exec" +) + +// accessor is the artifact.Accessor handed to a handler running inside the +// shim. It closes over the service, the owner, and the attempt — mirroring +// artifact/staging's own accessor — but resolves inputs against the +// Request's InputSlots instead of a staged-inputs map, since the shim +// receives its request as a wire value rather than building one from a +// staging plan. +type accessor struct { + svc *artifact.Service + req *exec.Request + owner artifact.OwnerRef + attempt int +} + +var _ artifact.Accessor = (*accessor)(nil) + +// newAccessorService builds the artifact service a sandboxed handler runs +// against: a real artifact.Service over a local directory and a minimal +// in-memory store local to this package. +// +// The handler therefore exercises the genuine Create/Commit/IfAbsent code +// path and cannot tell which side of the boundary it is on, while holding +// no backend credential and reaching no database. The in-memory rows are +// not a record of truth; they exist so Commit can return a Ref and +// Existing can answer within the attempt. The worker outside verifies what +// actually landed in the directory. +// +// memStore, not store/memory, is deliberate: store/memory is the full +// composite store and transitively imports workflow -> scope -> +// xraph/forge, which would link go-redis, k8s.io/client-go, a config +// loader, and an HTTP/QUIC/Prometheus stack into the sandbox binary. The +// phase's central claim -- this process holds no credential -- has to be +// provable by reading the import graph, not by tracing which +// package-level variables happen not to be constructed. +func newAccessorService(req *exec.Request) *artifact.Service { + return artifact.NewService( + newMemStore(), + NewLocalFS(req.OutputDir), + artifact.WithDefaultBucket("shim"), + ) +} + +// newAccessor builds the Accessor for req, scoped to owner and attempt. +func newAccessor(svc *artifact.Service, req *exec.Request, owner artifact.OwnerRef, attempt int) *accessor { + return &accessor{svc: svc, req: req, owner: owner, attempt: attempt} +} + +// Path returns the local file path of a declared input, or an empty +// string when the request carries no such input. +func (a *accessor) Path(name string) string { + for _, in := range a.req.Inputs { + if in.Name == name { + return filepath.Join(a.req.InputDir, in.Path) + } + } + + return "" +} + +// Open opens a declared input from local disk. +func (a *accessor) Open(_ context.Context, name string) (io.ReadCloser, error) { + path := a.Path(name) + if path == "" { + return nil, fmt.Errorf("dispatch/exec/shim: open %q: %w", name, artifact.ErrUnbound) + } + + // #nosec G304 -- path comes from the accessor's own binding table via + // a.Path, and an unbound name is rejected above. + f, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("dispatch/exec/shim: open %q: %w", name, err) + } + + return f, nil +} + +// Ref always reports no bound ref. +// +// exec.InputSlot carries only a Name and a Path — inputs are not staged +// through the artifact plane for out-of-process rungs yet, so there is no +// Ref to hand back. This is a known Phase 2 limitation, not a bug, and +// Path is no better off: nothing populates req.Inputs or req.InputDir for +// a declared input today (see execution-isolation.mdx, "Inputs are not +// yet staged out-of-process" — Phase 3 work), so a handler that calls Ref +// for a declared input gets false here for the same reason Path returns +// "" and Open returns artifact.ErrUnbound. +func (a *accessor) Ref(string) (artifact.Ref, bool) { + return artifact.Ref{}, false +} + +// Create begins writing an output owned by the running job. +func (a *accessor) Create( + ctx context.Context, + name string, + opts ...artifact.CreateOption, +) (*artifact.CommitWriter, error) { + return a.svc.Create(ctx, a.owner, a.attempt, name, opts...) +} + +// Existing reports an artifact a previous attempt committed under this +// name, letting a retried handler skip work already done. +func (a *accessor) Existing(ctx context.Context, name string) (artifact.Ref, bool) { + ref, err := a.svc.FindExisting(ctx, a.owner, name) + if err != nil { + return artifact.Ref{}, false + } + + return ref, true +} + +// seedPriorOutputs inserts req.PriorOutputs into svc's in-memory store as +// links for owner, so Existing and IfAbsent answer correctly for work an +// earlier attempt finished. +// +// Each prior output is seeded at Attempt 0. FindLinkByName returns the +// link with the highest attempt for a given name, and there is exactly one +// seeded link per name, so 0 only has to be lower than the attempt +// currently running — never equal to it — so a fresh Create in this +// attempt does not collide with the seed. +func seedPriorOutputs(ctx context.Context, svc *artifact.Service, owner artifact.OwnerRef, prior []exec.PriorOutput) error { + store := svc.Store() + + for _, po := range prior { + a := &artifact.Artifact{ + ID: po.Ref.ID, + Backend: po.Ref.Backend, + Bucket: po.Ref.Bucket, + Key: po.Ref.Key, + Size: po.Ref.Size, + ContentHash: po.Ref.ContentHash, + Lifecycle: artifact.Ephemeral, + CreatedAt: time.Now().UTC(), + } + link := &artifact.Link{ + ArtifactID: po.Ref.ID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleOutput, + Name: po.Name, + Attempt: 0, + CreatedAt: time.Now().UTC(), + } + + if err := store.CreateArtifact(ctx, a, link); err != nil { + return fmt.Errorf("dispatch/exec/shim: seed prior output %q: %w", po.Name, err) + } + } + + return nil +} diff --git a/exec/shim/doc.go b/exec/shim/doc.go new file mode 100644 index 0000000..f61b726 --- /dev/null +++ b/exec/shim/doc.go @@ -0,0 +1,9 @@ +// Package shim is the child side of out-of-process execution. +// +// A sandboxed process re-execs the worker's own binary, which calls Main. +// Main builds a bare job.Registry and a credential-free artifact.Service +// over a local directory, reads a request, runs the handler, and writes a +// result. It never constructs an engine, a store, or a DI container, so +// the process that parses an untrusted file holds no database credential +// and no object-store client. +package shim diff --git a/exec/shim/internal_test.go b/exec/shim/internal_test.go new file mode 100644 index 0000000..c9dcc50 --- /dev/null +++ b/exec/shim/internal_test.go @@ -0,0 +1,244 @@ +//go:build unix + +package shim + +// This file is package shim (internal), not shim_test, deliberately +// breaking the external-tests-only convention the rest of this package +// follows. mainExitCode and fdFromEnv are unexported, and the only way to +// exercise mainExitCode's exit-code derivation and its signal-handling +// goroutine without forking a real subprocess — which os.Exit inside Main +// would otherwise force — is to call it directly. File descriptors are +// process-scoped, so os.Pipe plus t.Setenv reaches it in-process. +// +// The build tag is here because callMainExitCode below calls +// syscall.Dup, which does not exist in Go's syscall package on Windows — +// this file predates the rest of this package's build-tag split (it's +// from Task 3) and was missed when that split happened. TestFDFromEnv +// itself needs nothing platform-specific, but it lives in the same file +// as callMainExitCode's other callers, so it is unix-only along with +// them rather than split out on its own; this package has no real +// non-Unix target (see procattr_other.go / limits_other.go in +// exec/subprocess, which refuse the whole rung outside Unix), so the +// coverage this loses there is nothing this rung claims to provide +// anyway. + +import ( + "context" + "errors" + "os" + "runtime" + "strconv" + "syscall" + "testing" + "time" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/wire" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// internalReq builds a minimal request for name, with its own OutputDir so +// runs don't interfere with each other. +func internalReq(t *testing.T, name string) *exec.Request { + t.Helper() + + return &exec.Request{ + JobID: id.NewJobID(), + Name: name, + Payload: []byte("{}"), + OutputDir: t.TempDir(), + } +} + +// callMainExitCode wires up a pair of in-process pipes standing in for fd +// 3 and fd 4, points EnvRequestFD/EnvResultFD at duplicates of them, +// encodes req into the request pipe, and returns whatever +// mainExitCode(defs) returns. +// +// The duplication matters. mainExitCode wraps whatever fd number it is +// given in its own *os.File via os.NewFile, and that wrapper carries a GC +// finalizer that closes the underlying fd once the wrapper becomes +// unreachable — confirmed empirically: wrapping a pipe fd, dropping the +// wrapper, and forcing a GC leaves the original os.Pipe halves reading +// "bad file descriptor". A real re-exec'd child does not hit this, +// because fd 3 and fd 4 there are the child's own fd-table entries +// (inherited via dup2 at exec time), independent of whatever fd number +// the parent's pipe end happened to have — closing them never touches the +// parent's descriptors. Handing mainExitCode this test's *own* reqR/resW +// fd numbers directly would violate that: a leftover wrapper from an +// earlier call in this same loop can get GC'd and finalize-close a fd +// number the OS has since reused for a *later* call's pipe. Dup'ing +// before the call gives mainExitCode an fd it exclusively owns, matching +// the real setup and making the finalizer harmless — it closes only the +// duplicate, never the original this function still holds. +// +// t.Setenv scopes the environment override to this test and restores it +// afterward. This function's own reqR/reqW/resR/resW are closed before it +// returns, so a test calling this in a loop does not exhaust descriptors; +// the dup'd fds handed to mainExitCode are left for its finalizer, which +// is exactly the ownership split described above. +func callMainExitCode(t *testing.T, defs []job.Registrable, req *exec.Request) int { + t.Helper() + + reqR, reqW, err := os.Pipe() + if err != nil { + t.Fatalf("Pipe() = %v", err) + } + defer reqR.Close() + defer reqW.Close() + + resR, resW, err := os.Pipe() + if err != nil { + t.Fatalf("Pipe() = %v", err) + } + defer resR.Close() + defer resW.Close() + + if eerr := wire.Encode(reqW, &wire.Frame{Kind: wire.KindRequest, Request: req}); eerr != nil { + t.Fatalf("Encode() = %v", eerr) + } + + reqDup, err := syscall.Dup(int(reqR.Fd())) + if err != nil { + t.Fatalf("Dup() = %v", err) + } + + resDup, err := syscall.Dup(int(resW.Fd())) + if err != nil { + t.Fatalf("Dup() = %v", err) + } + + t.Setenv(EnvRequestFD, strconv.Itoa(reqDup)) + t.Setenv(EnvResultFD, strconv.Itoa(resDup)) + + return mainExitCode(defs) +} + +// TestMainExitCode pins the exit-code contract mainExitCode must satisfy: +// a handler outcome, success or failure, is exit 0 because the parent +// reads it from the Result frame; a launch failure — the shim never +// finding a matching handler — is exit non-zero because Run reports that +// case by writing a Result and returning a nil error, and mainExitCode +// has to look past that nil to the Result's Status to catch it. Before the +// fix this whole table would have reported 0. +func TestMainExitCode(t *testing.T) { + defs := []job.Registrable{ + job.NewDefinition("internal.ok", func(context.Context, struct{}) error { return nil }), + job.NewDefinition("internal.err", func(context.Context, struct{}) error { + return errors.New("handler said no") + }), + } + + tests := []struct { + name string + job string + want int + }{ + {"handler success is exit 0", "internal.ok", 0}, + {"handler error is exit 0 -- StatusHandlerError is a business outcome", "internal.err", 0}, + {"unknown handler is a launch failure and must be exit non-zero", "internal.absent", 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := callMainExitCode(t, defs, internalReq(t, tt.job)); got != tt.want { + t.Errorf("mainExitCode() = %d, want %d", got, tt.want) + } + }) + } +} + +// TestMainExitCode_FingerprintMismatchIsNonZero covers the other launch +// failure path — a fingerprint the registry does not match — since it +// goes through a different branch of run than an unknown handler does. +func TestMainExitCode_FingerprintMismatchIsNonZero(t *testing.T) { + defs := []job.Registrable{ + job.NewDefinition("internal.ok", func(context.Context, struct{}) error { return nil }), + } + + r := internalReq(t, "internal.ok") + r.Fingerprint = "not-the-right-fingerprint" + + if got := callMainExitCode(t, defs, r); got != 1 { + t.Errorf("mainExitCode() = %d, want 1", got) + } +} + +// TestMainExitCode_NoGoroutineLeak proves the fix for the SIGTERM-handler +// goroutine leak: mainExitCode must not leave its signal-watching +// goroutine parked on <-sigCh forever when no signal ever arrives, which +// is the common case (Main's own os.Exit used to mask this by killing the +// process before a leak could accumulate). Called 20 times with no signal +// sent, the goroutine count must return to its starting point once each +// call's done channel unblocks the watcher. +func TestMainExitCode_NoGoroutineLeak(t *testing.T) { + defs := []job.Registrable{ + job.NewDefinition("internal.ok", func(context.Context, struct{}) error { return nil }), + } + + // One warm-up call first, outside the measured baseline: the Go + // runtime lazily starts its own permanent signal-forwarding goroutine + // the first time any code calls signal.Notify, and that goroutine + // lives for the rest of the process. Counting goroutines before this + // warm-up would mistake that one-time, one-goroutine runtime cost for + // a leak on the very first iteration, when the thing under test is + // growth across repeated calls, not that fixed cost. + if got := callMainExitCode(t, defs, internalReq(t, "internal.ok")); got != 0 { + t.Fatalf("mainExitCode() warm-up call = %d, want 0", got) + } + + before := runtime.NumGoroutine() + + const calls = 20 + for i := 0; i < calls; i++ { + if got := callMainExitCode(t, defs, internalReq(t, "internal.ok")); got != 0 { + t.Fatalf("mainExitCode() call %d = %d, want 0", i, got) + } + } + + // The done channel closes as mainExitCode returns, but the goroutine + // parked in select still needs a scheduler quantum to wake up and + // exit. Poll with a deadline rather than a fixed sleep so the test is + // both fast on the common path and not flaky under load. + deadline := time.Now().Add(time.Second) + for runtime.NumGoroutine() > before && time.Now().Before(deadline) { + runtime.Gosched() + time.Sleep(time.Millisecond) + } + + if got := runtime.NumGoroutine(); got > before { + t.Errorf("NumGoroutine() = %d after %d calls with no signal, want <= %d (leak in the SIGTERM handler goroutine)", + got, calls, before) + } +} + +// TestFDFromEnv exercises fdFromEnv directly: an unset variable and an +// unparsable one both fall back to the default, and a valid one wins. +func TestFDFromEnv(t *testing.T) { + t.Setenv(EnvRequestFD, "42") + if got := fdFromEnv(EnvRequestFD, defaultRequestFD); got != 42 { + t.Errorf("fdFromEnv() = %d, want 42", got) + } + + t.Setenv(EnvRequestFD, "not-a-number") + if got := fdFromEnv(EnvRequestFD, defaultRequestFD); got != defaultRequestFD { + t.Errorf("fdFromEnv() = %d, want default %d", got, defaultRequestFD) + } + + // A negative value parses fine but is not a descriptor, and callers convert + // to uintptr, where it would wrap to an enormous bogus fd rather than fail. + for _, v := range []string{"-1", "-99"} { + t.Setenv(EnvRequestFD, v) + if got := fdFromEnv(EnvRequestFD, defaultRequestFD); got != defaultRequestFD { + t.Errorf("fdFromEnv(%q) = %d, want default %d", v, got, defaultRequestFD) + } + } + + if err := os.Unsetenv("DISPATCH_EXEC_SHIM_TEST_UNSET"); err != nil { + t.Fatalf("Unsetenv() = %v", err) + } + if got := fdFromEnv("DISPATCH_EXEC_SHIM_TEST_UNSET", defaultResultFD); got != defaultResultFD { + t.Errorf("fdFromEnv() = %d, want default %d", got, defaultResultFD) + } +} diff --git a/exec/shim/localfs.go b/exec/shim/localfs.go new file mode 100644 index 0000000..7e7c944 --- /dev/null +++ b/exec/shim/localfs.go @@ -0,0 +1,288 @@ +package shim + +import ( + "context" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + + "github.com/zeebo/blake3" + + "github.com/xraph/dispatch/artifact" +) + +// dirMode is the permission bits for directories LocalFS creates. It +// grants the owner and group read/execute, and nothing to everyone else, +// so a sandboxed handler does not leave world-readable output behind. +const dirMode = 0o750 + +// fileMode is the permission bits for files LocalFS writes. It grants the +// owner and group read/write, and nothing to everyone else. +const fileMode = 0o640 + +// hashPrefix labels a content digest with its algorithm, matching the +// format artifact/cache uses for Ref.ContentHash so a hash computed here +// is directly comparable to one computed there. +const hashPrefix = "blake3:" + +// tempPattern names the scratch file Create writes into before it is +// renamed into place. The leading dot hides it from casual directory +// listings of the output tree. +const tempPattern = ".tmp-*" + +// LocalFS is a directory-backed artifact.Backend. It stores every object +// as a plain file under root, so the process holding it needs no object +// store credential — only a directory it can read and write. +// +// LocalFS is built for the exec/shim boundary: a sandboxed child process +// writes its outputs here, and the parent collects them from the +// filesystem once the child exits, without the child ever holding a +// network client for the real object store. +type LocalFS struct { + root string +} + +// Compile-time check that LocalFS satisfies the contract it exists for. +var _ artifact.Backend = (*LocalFS)(nil) + +// NewLocalFS returns a Backend that stores objects as files under root. +// The caller is responsible for root existing and being writable; LocalFS +// creates subdirectories under it as needed but never creates root itself. +func NewLocalFS(root string) *LocalFS { + return &LocalFS{root: root} +} + +// Name identifies this backend. +func (fs *LocalFS) Name() string { return "localfs" } + +// resolve maps a bucket-relative key onto a path under root, rejecting +// any key that would place the result outside root. +// +// A key is attacker-influenced in the general case: this backend runs +// inside the process that is parsing a possibly-malicious file, and a key +// that escaped root would let that process write or read anywhere the +// sandbox UID can reach. resolve is therefore the single choke point +// every method routes through, and it fails closed rather than silently +// clamping a traversal into some other path inside root: a caller that +// asked for "../escape" gets an error, not a different object. +// +// A key must always denote something strictly inside root, never root +// itself: Create's caller expects a file, not the directory it lives in, +// and Delete's caller expects to remove one object, not the whole scratch +// tree. filepath.Clean collapses every self-cancelling key ("", ".", +// "a/..", "a/b/../..") to exactly ".", so that single value is the one +// case to reject here. +func (fs *LocalFS) resolve(key string) (string, error) { + if filepath.IsAbs(key) { + return "", fmt.Errorf("shim: key %q must not be absolute", key) + } + + cleaned := filepath.Clean(key) + if cleaned == "." { + return "", fmt.Errorf("shim: key %q resolves to the output root", key) + } + + full := filepath.Join(fs.root, cleaned) + + rel, err := filepath.Rel(fs.root, full) + if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("shim: key %q escapes root", key) + } + + return full, nil +} + +// Open returns a reader over the object's bytes. +func (fs *LocalFS) Open(_ context.Context, ref artifact.Ref) (io.ReadCloser, error) { + path, err := fs.resolve(ref.Key) + if err != nil { + return nil, fmt.Errorf("shim: open %s/%s: %w", ref.Bucket, ref.Key, err) + } + + // #nosec G304 -- path is confined to fs.root by resolve's containment check. + f, err := os.Open(path) + if err != nil { + if os.IsNotExist(err) { + return nil, fmt.Errorf("shim: open %s/%s: %w", ref.Bucket, ref.Key, artifact.ErrNotFound) + } + + return nil, fmt.Errorf("shim: open %s/%s: %w", ref.Bucket, ref.Key, err) + } + + return f, nil +} + +// Create begins writing a new object. The bytes are not visible until +// Commit. +func (fs *LocalFS) Create(_ context.Context, bucket, key string) (artifact.Writer, error) { + path, err := fs.resolve(key) + if err != nil { + return nil, fmt.Errorf("shim: create %s/%s: %w", bucket, key, err) + } + + dir := filepath.Dir(path) + if mkerr := os.MkdirAll(dir, dirMode); mkerr != nil { + return nil, fmt.Errorf("shim: create %s/%s: %w", bucket, key, mkerr) + } + + // Create the temp file in the same directory as the final path so the + // rename in Commit is an atomic, same-filesystem operation. + tmp, err := os.CreateTemp(dir, tempPattern) + if err != nil { + return nil, fmt.Errorf("shim: create %s/%s: %w", bucket, key, err) + } + + if cherr := tmp.Chmod(fileMode); cherr != nil { + _ = tmp.Close() + // tmp.Name() is a sibling of path inside the directory resolve + // already confined to fs.root; nothing here reads attacker input. + // Safe to remove: the path came from resolve's containment check, not + // from a raw key. + _ = os.Remove(tmp.Name()) + + return nil, fmt.Errorf("shim: create %s/%s: %w", bucket, key, cherr) + } + + return &localWriter{ + file: tmp, + hash: blake3.New(), + bucket: bucket, + key: key, + final: path, + }, nil +} + +// Stat reports the object's size without reading it. +func (fs *LocalFS) Stat(_ context.Context, ref artifact.Ref) (artifact.ObjectInfo, error) { + path, err := fs.resolve(ref.Key) + if err != nil { + return artifact.ObjectInfo{}, fmt.Errorf("shim: stat %s/%s: %w", ref.Bucket, ref.Key, err) + } + + info, err := os.Stat(path) + if err != nil { + if os.IsNotExist(err) { + return artifact.ObjectInfo{}, fmt.Errorf("shim: stat %s/%s: %w", ref.Bucket, ref.Key, artifact.ErrNotFound) + } + + return artifact.ObjectInfo{}, fmt.Errorf("shim: stat %s/%s: %w", ref.Bucket, ref.Key, err) + } + + return artifact.ObjectInfo{Size: info.Size()}, nil +} + +// Delete removes the object. Deleting a missing object is not an error. +func (fs *LocalFS) Delete(_ context.Context, ref artifact.Ref) error { + path, err := fs.resolve(ref.Key) + if err != nil { + return fmt.Errorf("shim: delete %s/%s: %w", ref.Bucket, ref.Key, err) + } + + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("shim: delete %s/%s: %w", ref.Bucket, ref.Key, err) + } + + return nil +} + +// errWriterFinished means Commit or Abort was called on a localWriter +// that had already been finalised by one or the other. +var errWriterFinished = errors.New("shim: writer already finished") + +// localWriter accumulates bytes in a temp file and renames it into place +// on Commit. It hashes the bytes as they stream through so Commit never +// has to re-read the file to report a content hash. +type localWriter struct { + file *os.File + hash *blake3.Hasher + bucket string + key string + final string + + written int64 + done bool +} + +// Write appends bytes to the pending object. +func (w *localWriter) Write(p []byte) (int, error) { + n, err := w.file.Write(p) + w.written += int64(n) + + if n > 0 { + // hash.Hash.Write never returns an error, per the io.Writer + // contract documented on the standard library's hash.Hash. + if _, herr := w.hash.Write(p[:n]); herr != nil { + return n, herr + } + } + + return n, err +} + +// Commit finalises the object: it flushes and closes the temp file, +// renames it into place, and reports the logical size and content hash. +func (w *localWriter) Commit(_ context.Context) (artifact.ObjectInfo, error) { + if w.done { + return artifact.ObjectInfo{}, fmt.Errorf("shim: commit %s/%s: %w", w.bucket, w.key, errWriterFinished) + } + + w.done = true + tmpName := w.file.Name() + + // tmpName is the temp file this writer created inside the directory + // resolve already confined to fs.root; it is not attacker input. + if err := w.file.Sync(); err != nil { + _ = w.file.Close() + _ = os.Remove(tmpName) + + return artifact.ObjectInfo{}, fmt.Errorf("shim: commit %s/%s: %w", w.bucket, w.key, err) + } + + if err := w.file.Close(); err != nil { + _ = os.Remove(tmpName) + + return artifact.ObjectInfo{}, fmt.Errorf("shim: commit %s/%s: %w", w.bucket, w.key, err) + } + + // tmpName and w.final both passed through resolve's containment check + // (w.final at Create time; tmpName is a sibling CreateTemp made inside + // that same, already-contained directory). + // Both paths are confined to fs.root by resolve. + if err := os.Rename(tmpName, w.final); err != nil { + _ = os.Remove(tmpName) + + return artifact.ObjectInfo{}, fmt.Errorf("shim: commit %s/%s: %w", w.bucket, w.key, err) + } + + sum := hex.EncodeToString(w.hash.Sum(nil)) + + return artifact.ObjectInfo{ + Size: w.written, + ETag: hashPrefix + sum, + }, nil +} + +// Abort discards the partial object. It is a no-op after Commit. +func (w *localWriter) Abort() error { + if w.done { + return nil + } + + w.done = true + name := w.file.Name() + + _ = w.file.Close() + + // name is this writer's own temp file, created inside the directory + // resolve already confined to fs.root at Create time. + // name is our own temp file under the resolved, contained directory. + if err := os.Remove(name); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("shim: abort %s/%s: %w", w.bucket, w.key, err) + } + + return nil +} diff --git a/exec/shim/localfs_test.go b/exec/shim/localfs_test.go new file mode 100644 index 0000000..b8db771 --- /dev/null +++ b/exec/shim/localfs_test.go @@ -0,0 +1,200 @@ +package shim_test + +import ( + "context" + "errors" + "io" + "os" + "path/filepath" + "testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/exec/shim" +) + +func TestLocalFS_CreateThenOpen(t *testing.T) { + root := t.TempDir() + be := shim.NewLocalFS(root) + ctx := context.Background() + + w, err := be.Create(ctx, "b", "ephemeral/job/j1/0/mesh.glb") + if err != nil { + t.Fatalf("Create() = %v", err) + } + if _, werr := io.WriteString(w, "hello"); werr != nil { + t.Fatalf("Write() = %v", werr) + } + if _, cerr := w.Commit(ctx); cerr != nil { + t.Fatalf("Commit() = %v", cerr) + } + + // The bytes must be a real file under root, at the key's path, so the + // parent can collect outputs without cooperating with the child. + onDisk := filepath.Join(root, "ephemeral/job/j1/0/mesh.glb") + got, err := os.ReadFile(onDisk) + if err != nil { + t.Fatalf("expected a file at %s: %v", onDisk, err) + } + if string(got) != "hello" { + t.Errorf("file = %q, want %q", got, "hello") + } + + rc, err := be.Open(ctx, artifact.Ref{Backend: "localfs", Bucket: "b", Key: "ephemeral/job/j1/0/mesh.glb"}) + if err != nil { + t.Fatalf("Open() = %v", err) + } + defer rc.Close() + back, _ := io.ReadAll(rc) + if string(back) != "hello" { + t.Errorf("Open() = %q, want %q", back, "hello") + } +} + +func TestLocalFS_OpenMissingIsNotFound(t *testing.T) { + be := shim.NewLocalFS(t.TempDir()) + _, err := be.Open(context.Background(), artifact.Ref{Bucket: "b", Key: "nope"}) + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("Open(missing) = %v, want ErrNotFound", err) + } +} + +func TestLocalFS_AbortLeavesNoFile(t *testing.T) { + root := t.TempDir() + be := shim.NewLocalFS(root) + + w, err := be.Create(context.Background(), "b", "k") + if err != nil { + t.Fatalf("Create() = %v", err) + } + if _, err := io.WriteString(w, "partial"); err != nil { + t.Fatalf("Write() = %v", err) + } + if err := w.Abort(); err != nil { + t.Fatalf("Abort() = %v", err) + } + + if _, err := os.Stat(filepath.Join(root, "k")); !os.IsNotExist(err) { + t.Error("Abort() left a file behind; an aborted write must be invisible") + } +} + +func TestLocalFS_RejectsEscapingKey(t *testing.T) { + // A key is attacker-influenced in the general case. It must never + // resolve outside root. A key that lexically collapses to root itself + // ("", ".", "a/..", "a/b/../..") is equally dangerous even though it + // never leaves root: Create would write into root's parent (Dir of + // root), and Delete would remove the whole scratch directory instead + // of one object. + be := shim.NewLocalFS(t.TempDir()) + ctx := context.Background() + keys := []string{ + "../escape", "a/../../escape", "/absolute", + "", ".", "a/..", "a/b/../..", + } + for _, key := range keys { + if _, err := be.Create(ctx, "b", key); err == nil { + t.Errorf("Create(%q) = nil, want a rejection", key) + } + if err := be.Delete(ctx, artifact.Ref{Bucket: "b", Key: key}); err == nil { + t.Errorf("Delete(%q) = nil, want a rejection", key) + } + } +} + +func TestLocalFS_Name(t *testing.T) { + be := shim.NewLocalFS(t.TempDir()) + if got := be.Name(); got != "localfs" { + t.Errorf("Name() = %q, want %q", got, "localfs") + } +} + +func TestLocalFS_DeleteMissingIsNotError(t *testing.T) { + be := shim.NewLocalFS(t.TempDir()) + err := be.Delete(context.Background(), artifact.Ref{Bucket: "b", Key: "nope"}) + if err != nil { + t.Fatalf("Delete(missing) = %v, want nil", err) + } +} + +func TestLocalFS_StatReportsSize(t *testing.T) { + root := t.TempDir() + be := shim.NewLocalFS(root) + ctx := context.Background() + + w, err := be.Create(ctx, "b", "k") + if err != nil { + t.Fatalf("Create() = %v", err) + } + if _, werr := io.WriteString(w, "hello world"); werr != nil { + t.Fatalf("Write() = %v", werr) + } + info, err := w.Commit(ctx) + if err != nil { + t.Fatalf("Commit() = %v", err) + } + if info.Size != int64(len("hello world")) { + t.Errorf("Commit() size = %d, want %d", info.Size, len("hello world")) + } + if info.ETag == "" { + t.Errorf("Commit() ETag is empty, want a content hash") + } + + stat, err := be.Stat(ctx, artifact.Ref{Bucket: "b", Key: "k"}) + if err != nil { + t.Fatalf("Stat() = %v", err) + } + if stat.Size != int64(len("hello world")) { + t.Errorf("Stat() size = %d, want %d", stat.Size, len("hello world")) + } +} + +func TestLocalFS_CommitAbortIsNoOp(t *testing.T) { + root := t.TempDir() + be := shim.NewLocalFS(root) + ctx := context.Background() + + w, err := be.Create(ctx, "b", "k") + if err != nil { + t.Fatalf("Create() = %v", err) + } + if _, err := io.WriteString(w, "hello"); err != nil { + t.Fatalf("Write() = %v", err) + } + if _, err := w.Commit(ctx); err != nil { + t.Fatalf("Commit() = %v", err) + } + if err := w.Abort(); err != nil { + t.Fatalf("Abort() after Commit() = %v, want nil", err) + } + + // The committed file must still be present; Abort after Commit must + // not remove it. + if _, err := os.Stat(filepath.Join(root, "k")); err != nil { + t.Errorf("file gone after post-commit Abort(): %v", err) + } +} + +func TestLocalFS_FileModeIsNotWorldReadable(t *testing.T) { + root := t.TempDir() + be := shim.NewLocalFS(root) + ctx := context.Background() + + w, err := be.Create(ctx, "b", "k") + if err != nil { + t.Fatalf("Create() = %v", err) + } + if _, werr := io.WriteString(w, "hello"); werr != nil { + t.Fatalf("Write() = %v", werr) + } + if _, cerr := w.Commit(ctx); cerr != nil { + t.Fatalf("Commit() = %v", cerr) + } + + info, err := os.Stat(filepath.Join(root, "k")) + if err != nil { + t.Fatalf("Stat() = %v", err) + } + if info.Mode().Perm()&0o007 != 0 { + t.Errorf("file mode = %v, must not be world-accessible", info.Mode().Perm()) + } +} diff --git a/exec/shim/main.go b/exec/shim/main.go new file mode 100644 index 0000000..f6290b9 --- /dev/null +++ b/exec/shim/main.go @@ -0,0 +1,400 @@ +package shim + +import ( + "context" + "errors" + "fmt" + "io" + "io/fs" + "os" + "os/signal" + "path/filepath" + "strconv" + "strings" + "syscall" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/wire" + "github.com/xraph/dispatch/job" +) + +const ( + // EnvRequestFD names the environment variable that overrides which + // file descriptor Main reads the exec.Request from. Unset, Main reads + // fd 3 — the parent's convention for the first descriptor past + // stdin/stdout/stderr. + EnvRequestFD = "DISPATCH_EXEC_REQUEST_FD" + + // EnvResultFD names the environment variable that overrides which + // file descriptor Main writes the exec.Result to. Unset, Main writes + // fd 4. + EnvResultFD = "DISPATCH_EXEC_RESULT_FD" + + // ArgName is the argv[1] marker a parent sets when it re-execs its own + // binary into the shim, distinguishing that invocation from an + // ordinary run of the worker. argv[0] is still the binary path itself, + // same as any other invocation — subprocess.Executor.Run builds the + // child's argv as [binary, ArgName, ...WithArgs], so this is the first + // argument after the binary, not the zeroth. + ArgName = "dispatch-exec" + + // defaultRequestFD is the descriptor Main reads from absent an + // EnvRequestFD override. + defaultRequestFD = 3 + + // defaultResultFD is the descriptor Main writes to absent an + // EnvResultFD override. + defaultResultFD = 4 + + // EnvRlimitAS, EnvRlimitNoFile, EnvRlimitNProc, EnvRlimitCore, and + // EnvRlimitFSize name the environment variables the parent uses to + // pass POSIX resource limits into the child. Go cannot set a child's + // rlimits through os/exec's SysProcAttr, so the parent (see buildEnv + // in exec/subprocess) passes the desired values here, and Main + // applies them itself via syscall.Setrlimit — see applyRlimits — + // before running anything. Each is unset when the parent has no + // value to send, except EnvRlimitCore, which the parent always sends + // as "0". + EnvRlimitAS = "DISPATCH_RLIMIT_AS" + EnvRlimitNoFile = "DISPATCH_RLIMIT_NOFILE" + EnvRlimitNProc = "DISPATCH_RLIMIT_NPROC" + EnvRlimitCore = "DISPATCH_RLIMIT_CORE" + EnvRlimitFSize = "DISPATCH_RLIMIT_FSIZE" + + // EnvRlimitStrict names the environment variable the parent sets + // (see subprocess.WithStrictRlimits) to ask Main to fail the launch + // outright when an rlimit it was actually asked to apply fails for a + // reason other than "this platform does not support that limit at + // all" — see applyRlimits for the distinction. Set to any non-empty + // value to enable; unset (the default) keeps every rlimit failure + // non-fatal. + EnvRlimitStrict = "DISPATCH_RLIMIT_STRICT" +) + +// Main is the sandboxed child's entrypoint. It builds a bare job.Registry +// from defs and a credential-free artifact.Service over a local directory, +// reads one exec.Request from its request descriptor, runs the matching +// handler, and writes one exec.Result to its result descriptor. +// +// Main never returns: it calls os.Exit. A handler returning an error is +// exit 0, since that is a business outcome the Result frame already +// carries as StatusHandlerError. A nonzero exit means the shim itself +// failed the attempt before or instead of running the handler — a request +// that failed to decode, or one whose fingerprint or handler name did not +// resolve — which is what lets the parent distinguish those cases from a +// clean run without having to inspect the frame first: the exit code +// corroborates what the Result already says. +func Main(defs ...job.Registrable) { + os.Exit(mainExitCode(defs)) +} + +// mainExitCode does the real work of Main and returns the process exit +// code, rather than calling os.Exit itself. os.Exit does not run deferred +// calls, so a single call to it right at the top of Main — after every +// defer here has already unwound — is what lets the signal handler and +// the cancel func clean up on every path. +// +// It calls run rather than Run because Run's public contract discards the +// Result on success, and the exit code has to be derived from that Result: +// a fingerprint mismatch or an unknown handler makes Run return a nil +// error (the frame is the report, by design — see Run's doc comment), so +// err == nil alone cannot tell mainExitCode apart from a clean success. +// Only StatusHandlerError keeps exit 0; every other non-OK status, +// including one Run reported without an error, is a nonzero exit. +func mainExitCode(defs []job.Registrable) int { + // fdFromEnv guarantees a non-negative descriptor, so the uintptr + // conversion cannot wrap. + in := os.NewFile(uintptr(fdFromEnv(EnvRequestFD, defaultRequestFD)), "dispatch-exec-request") + out := os.NewFile(uintptr(fdFromEnv(EnvResultFD, defaultResultFD)), "dispatch-exec-result") + + // Applied before anything else touches the request: RLIMIT_CORE in + // particular exists to stop a segfaulting handler from dumping the + // input that crashed it (and the process's own memory) to disk, which + // only holds if the limit is in place before the handler ever runs. + // + // failures excludes anything applyRlimits judged "this platform does + // not support that limit at all" — Darwin's blanket refusal of + // RLIMIT_AS being the standing example — since that is a structural, + // permanent fact about the kernel this process is running under, not + // a misconfiguration. What is left is failures a correctly-configured + // limit should not produce on a platform that claims to support it — + // EPERM because the requested value exceeds the hard limit is the + // common shape — which is what EnvRlimitStrict (see + // subprocess.WithStrictRlimits) opts into treating as a launch + // failure rather than a warning. + if failures := applyRlimits(); len(failures) > 0 && os.Getenv(EnvRlimitStrict) != "" { + res := &exec.Result{ + Status: exec.StatusLaunchFailed, + HandlerErr: fmt.Sprintf("dispatch/exec/shim: rlimits requested by WithStrictRlimits failed to apply: %s", joinRlimitFailures(failures)), + } + _ = writeResult(out, res) //nolint:errcheck // best-effort: the process is exiting non-zero regardless of whether the frame made it across + + return 1 + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGTERM) + defer signal.Stop(sigCh) + + // done unblocks the goroutine below once mainExitCode is about to + // return. signal.Stop alone only unregisters future deliveries; it + // neither closes sigCh nor wakes a goroutine already parked on + // <-sigCh, so without this the goroutine leaks on every call that + // never receives a SIGTERM — which, outside of tests, is every call, + // since Main's own os.Exit would otherwise mask the leak by killing + // the process before it could accumulate. + done := make(chan struct{}) + defer close(done) + + go func() { + select { + case <-sigCh: + cancel() + case <-done: + } + }() + + res, err := run(ctx, in, out, defs) + if err != nil { + return 1 + } + + // StatusHandlerError stays exit 0: it is a business outcome, and the + // parent reads it from the Result frame, not the exit code. + // StatusLaunchFailed is the shim's own failure and must not look like + // a clean exit. + if res.Status == exec.StatusLaunchFailed { + return 1 + } + + return 0 +} + +// rlimitFailure describes one configured rlimit that applyRlimits could +// not apply for a reason other than "this platform is known not to +// support this limit at all" — see applyRlimits (rlimit_unix.go) for how +// that distinction is made. Declared here, without a build tag, rather +// than alongside applyRlimits itself, because mainExitCode needs the type +// on every platform: rlimit_other.go's non-Unix stub returns the same +// []rlimitFailure (always empty) so mainExitCode does not need its own +// build-tagged branch just to call applyRlimits. +type rlimitFailure struct { + label string + err error +} + +// joinRlimitFailures renders failures as a single diagnostic string for +// the launch-failure Result mainExitCode writes when EnvRlimitStrict is +// set. Every failure that reaches here already went to stderr individually +// by applyRlimits; this is the summary that ends up in the Result the +// caller actually sees, not a replacement for those lines. +func joinRlimitFailures(failures []rlimitFailure) string { + parts := make([]string, len(failures)) + for i, f := range failures { + parts[i] = fmt.Sprintf("%s: %v", f.label, f.err) + } + + return strings.Join(parts, "; ") +} + +// fdFromEnv reads a file descriptor number from the named environment +// variable, falling back to def when the variable is unset or unparsable. +func fdFromEnv(name string, def int) int { + v, ok := os.LookupEnv(name) + if !ok { + return def + } + + n, err := strconv.Atoi(v) + if err != nil { + return def + } + + // A descriptor is never negative, and callers convert to uintptr, where a + // negative would wrap to an enormous bogus fd instead of failing. + if n < 0 { + return def + } + + return n +} + +// Run is the testable core of the shim: it reads one exec.Request from in, +// runs the matching handler out of defs, and writes one exec.Result to +// out. Splitting it from Main is what lets tests drive it with in-memory +// buffers instead of real file descriptors. +// +// Run returns a non-nil error only when it could not produce a Result +// frame at all — the request failed to decode, or was malformed enough +// that there is no attempt to report on. Every failure that happens once a +// well-formed request is in hand — an unknown handler, a fingerprint +// mismatch, the handler itself erroring — is reported by writing a Result +// frame and returning nil, because the frame is the report, not the +// error: the parent reads Status, not the shim's exit code, to learn what +// happened. +// +// Run discards the Result it produces, keeping its signature exactly what +// callers (and the brief's test) expect. mainExitCode needs that Result to +// derive an exit code, which Run's error alone cannot give it — see run. +func Run(ctx context.Context, in io.Reader, out io.Writer, defs []job.Registrable) error { + _, err := run(ctx, in, out, defs) + + return err +} + +// run is Run's implementation, plus the Result it wrote. mainExitCode is +// the reason this exists as a separate, unexported function: Run's error +// return is nil for a launch failure by design (the frame is the report), +// so the only way mainExitCode can tell a launch failure apart from a +// clean success is to inspect the Result's Status directly, which Run's +// public contract does not expose. +func run(ctx context.Context, in io.Reader, out io.Writer, defs []job.Registrable) (*exec.Result, error) { + frame, err := wire.Decode(in) + if err != nil { + return nil, fmt.Errorf("dispatch/exec/shim: read request: %w", err) + } + + req := frame.Request + if req == nil { + return nil, errors.New("dispatch/exec/shim: frame carries no request") + } + + if verr := req.Validate(); verr != nil { + return nil, fmt.Errorf("dispatch/exec/shim: %w", verr) + } + + registry := job.NewRegistry() + for _, d := range defs { + d.Register(registry) + } + + if req.Fingerprint != "" { + if got := exec.Fingerprint(registry.Names()); got != req.Fingerprint { + res := &exec.Result{ + Status: exec.StatusLaunchFailed, + HandlerErr: fmt.Sprintf( + "fingerprint mismatch: request wants %s, this binary's handler set is %s", + req.Fingerprint, got, + ), + } + + return res, writeResult(out, res) + } + } + + handler, ok := registry.Get(req.Name) + if !ok { + res := &exec.Result{ + Status: exec.StatusLaunchFailed, + HandlerErr: fmt.Sprintf("no handler registered for job %q", req.Name), + } + + return res, writeResult(out, res) + } + + svc := newAccessorService(req) + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: req.JobID.String()} + + if len(req.PriorOutputs) > 0 { + if serr := seedPriorOutputs(ctx, svc, owner, req.PriorOutputs); serr != nil { + return nil, fmt.Errorf("dispatch/exec/shim: %w", serr) + } + } + + ctx = artifact.WithAccessor(ctx, newAccessor(svc, req, owner, req.Attempt)) + + if !req.Deadline.IsZero() { + var cancel context.CancelFunc + + ctx, cancel = context.WithDeadline(ctx, req.Deadline) + defer cancel() + } + + start := time.Now() + handlerErr := handler(ctx, req.Payload) + wallTime := time.Since(start) + + res := &exec.Result{Usage: exec.Usage{WallTime: wallTime}} + if handlerErr != nil { + res.Status = exec.StatusHandlerError + res.HandlerErr = handlerErr.Error() + res.Permanent = errors.Is(handlerErr, dispatch.ErrPermanent) + } else { + res.Status = exec.StatusOK + } + + outputs, err := collectOutputs(req.OutputDir) + if err != nil { + return nil, fmt.Errorf("dispatch/exec/shim: collect outputs: %w", err) + } + + res.Outputs = outputs + + return res, writeResult(out, res) +} + +// writeResult encodes and writes the single result frame Run ever +// produces. +func writeResult(out io.Writer, res *exec.Result) error { + if err := wire.Encode(out, &wire.Frame{Kind: wire.KindResult, Result: res}); err != nil { + return fmt.Errorf("dispatch/exec/shim: encode result: %w", err) + } + + return nil +} + +// collectOutputs walks dir and reports every regular file found as an +// artifact the handler produced. +// +// It reads the filesystem rather than any manifest the handler could have +// populated itself, so a handler cannot claim an output it did not +// actually write. Dot-prefixed entries are skipped: LocalFS's Create +// writes into a hidden temp file before renaming it into place on Commit, +// so a leftover one (a write that was never committed or aborted) is not +// mistaken for a finished artifact. A missing directory is treated as no +// outputs rather than an error, since a handler that wrote nothing never +// causes LocalFS to create it. +func collectOutputs(dir string) ([]exec.OutputFile, error) { + var outputs []exec.OutputFile + + // dir is req.OutputDir, a path the parent chose and mounted for this + // attempt, never a value read out of the untrusted payload the + // handler parses. + // dir is the request's own OutputDir, not attacker-controlled. + err := filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + + if d.IsDir() || strings.HasPrefix(d.Name(), ".") { + return nil + } + + info, err := d.Info() + if err != nil { + return fmt.Errorf("stat %s: %w", path, err) + } + + outputs = append(outputs, exec.OutputFile{ + Name: d.Name(), + Size: info.Size(), + }) + + return nil + }) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + return nil, nil + } + + return nil, err + } + + return outputs, nil +} diff --git a/exec/shim/main_test.go b/exec/shim/main_test.go new file mode 100644 index 0000000..e5390d5 --- /dev/null +++ b/exec/shim/main_test.go @@ -0,0 +1,250 @@ +package shim_test + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "testing" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/shim" + "github.com/xraph/dispatch/exec/wire" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +type shimPayload struct { + Mode string `json:"mode"` +} + +func shimHandlers(t *testing.T) []job.Registrable { + t.Helper() + + return []job.Registrable{ + job.NewDefinition("shim.ok", func(context.Context, shimPayload) error { return nil }), + job.NewDefinition("shim.err", func(context.Context, shimPayload) error { + return errors.New("handler said no") + }), + job.NewDefinition("shim.permanent", func(context.Context, shimPayload) error { + return dispatch.ErrPermanent + }), + } +} + +func runShim(t *testing.T, req *exec.Request, defs []job.Registrable) *exec.Result { + t.Helper() + + var in, out bytes.Buffer + if err := wire.Encode(&in, &wire.Frame{Kind: wire.KindRequest, Request: req}); err != nil { + t.Fatalf("Encode() = %v", err) + } + if err := shim.Run(context.Background(), &in, &out, defs); err != nil { + t.Fatalf("Run() = %v", err) + } + f, err := wire.Decode(&out) + if err != nil { + t.Fatalf("Decode() = %v", err) + } + + return f.Result +} + +func req(t *testing.T, name string) *exec.Request { + t.Helper() + raw, _ := json.Marshal(shimPayload{Mode: "x"}) + + return &exec.Request{ + JobID: id.NewJobID(), + Name: name, + Payload: raw, + OutputDir: t.TempDir(), + } +} + +func TestRun(t *testing.T) { + defs := shimHandlers(t) + + tests := []struct { + name string + job string + wantStatus exec.Status + wantPerm bool + }{ + {"success", "shim.ok", exec.StatusOK, false}, + {"handler error", "shim.err", exec.StatusHandlerError, false}, + {"permanent crosses as a flag", "shim.permanent", exec.StatusHandlerError, true}, + {"unknown handler is a launch failure", "shim.absent", exec.StatusLaunchFailed, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := runShim(t, req(t, tt.job), defs) + if got.Status != tt.wantStatus { + t.Errorf("Status = %q, want %q (err %q)", got.Status, tt.wantStatus, got.HandlerErr) + } + if got.Permanent != tt.wantPerm { + t.Errorf("Permanent = %v, want %v", got.Permanent, tt.wantPerm) + } + if got.Cause != nil { + t.Error("Cause must be nil across the boundary") + } + }) + } +} + +func TestRunFingerprintMismatch(t *testing.T) { + r := req(t, "shim.ok") + r.Fingerprint = "not-the-right-fingerprint" + + got := runShim(t, r, shimHandlers(t)) + if got.Status != exec.StatusLaunchFailed { + t.Fatalf("Status = %q, want %q", got.Status, exec.StatusLaunchFailed) + } +} + +func TestRunMatchingFingerprintPasses(t *testing.T) { + defs := shimHandlers(t) + names := make([]string, 0, len(defs)) + for _, d := range defs { + names = append(names, d.JobName()) + } + + r := req(t, "shim.ok") + r.Fingerprint = exec.Fingerprint(names) + + if got := runShim(t, r, defs); got.Status != exec.StatusOK { + t.Fatalf("Status = %q, want %q", got.Status, exec.StatusOK) + } +} + +func TestRunRecordsOutputs(t *testing.T) { + dir := t.TempDir() + defs := []job.Registrable{ + job.NewDefinition("shim.writes", func(ctx context.Context, _ shimPayload) error { + w, err := artifact.From(ctx).Create(ctx, "mesh.glb") + if err != nil { + return err + } + if _, werr := w.Write([]byte("meshbytes")); werr != nil { + return werr + } + _, err = w.Commit(ctx) + + return err + }), + } + + r := req(t, "shim.writes") + r.OutputDir = dir + + got := runShim(t, r, defs) + if got.Status != exec.StatusOK { + t.Fatalf("Status = %q, want ok (err %q)", got.Status, got.HandlerErr) + } + if len(got.Outputs) != 1 || got.Outputs[0].Name != "mesh.glb" { + t.Fatalf("Outputs = %+v, want one named mesh.glb", got.Outputs) + } + if got.Outputs[0].Size != int64(len("meshbytes")) { + t.Errorf("Size = %d, want %d", got.Outputs[0].Size, len("meshbytes")) + } +} + +// TestRunSeedsPriorOutputsForExisting proves that a PriorOutput on the +// request makes Accessor.Existing answer true for that name — the whole +// point of seeding — while a name that was never seeded still answers +// false, so the seed is not mistaken for a wildcard "everything exists". +func TestRunSeedsPriorOutputsForExisting(t *testing.T) { + defs := []job.Registrable{ + job.NewDefinition("shim.checks_existing", func(ctx context.Context, _ shimPayload) error { + if _, ok := artifact.From(ctx).Existing(ctx, "mesh.glb"); !ok { + return errors.New("expected mesh.glb to resolve as an existing prior output") + } + if _, ok := artifact.From(ctx).Existing(ctx, "unseeded.glb"); ok { + return errors.New("unseeded.glb must not resolve; nothing seeded it") + } + + return nil + }), + } + + r := req(t, "shim.checks_existing") + r.PriorOutputs = []exec.PriorOutput{ + { + Name: "mesh.glb", + Ref: artifact.Ref{ + ID: id.NewArtifactID(), + Backend: "localfs", + Bucket: "shim", + Key: "ephemeral/job/prior-attempt/0/mesh.glb", + Size: 9, + }, + }, + } + + got := runShim(t, r, defs) + if got.Status != exec.StatusOK { + t.Fatalf("Status = %q, want ok (err %q)", got.Status, got.HandlerErr) + } +} + +// TestRunResolvesDeclaredInput proves the accessor's Path and Open work +// against a declared, staged input, and that Ref reports the documented +// Phase 2 limitation (no ref travels with an out-of-process InputSlot) +// rather than panicking or lying. +func TestRunResolvesDeclaredInput(t *testing.T) { + inputDir := t.TempDir() + if err := os.WriteFile(filepath.Join(inputDir, "in.txt"), []byte("input bytes"), 0o600); err != nil { + t.Fatalf("WriteFile() = %v", err) + } + + defs := []job.Registrable{ + job.NewDefinition("shim.reads_input", func(ctx context.Context, _ shimPayload) error { + acc := artifact.From(ctx) + + if got := acc.Path("in.txt"); got != filepath.Join(inputDir, "in.txt") { + return errors.New("Path() did not resolve to the staged file") + } + if got := acc.Path("undeclared.txt"); got != "" { + return errors.New("Path() of an undeclared input must be empty") + } + + rc, err := acc.Open(ctx, "in.txt") + if err != nil { + return err + } + defer rc.Close() + + buf := make([]byte, len("input bytes")) + if _, err := rc.Read(buf); err != nil { + return err + } + if string(buf) != "input bytes" { + return errors.New("Open() did not read the staged file's bytes") + } + + if _, ok := acc.Ref("in.txt"); ok { + return errors.New("Ref() must report false: Phase 2 does not stage a ref for InputSlot") + } + + if _, err := acc.Open(ctx, "undeclared.txt"); !errors.Is(err, artifact.ErrUnbound) { + return errors.New("Open() of an undeclared input must wrap artifact.ErrUnbound") + } + + return nil + }), + } + + r := req(t, "shim.reads_input") + r.InputDir = inputDir + r.Inputs = []exec.InputSlot{{Name: "in.txt", Path: "in.txt"}} + + got := runShim(t, r, defs) + if got.Status != exec.StatusOK { + t.Fatalf("Status = %q, want ok (err %q)", got.Status, got.HandlerErr) + } +} diff --git a/exec/shim/rlimit_as_openbsd.go b/exec/shim/rlimit_as_openbsd.go new file mode 100644 index 0000000..b9b026f --- /dev/null +++ b/exec/shim/rlimit_as_openbsd.go @@ -0,0 +1,16 @@ +//go:build openbsd + +package shim + +// rlimitAS reports RLIMIT_AS as unusable on OpenBSD. Go's syscall +// package exports RLIMIT_CORE, RLIMIT_CPU, RLIMIT_DATA, RLIMIT_FSIZE, +// RLIMIT_NOFILE, and RLIMIT_STACK for this platform (confirmed against +// its generated zerrors_openbsd_*.go tables) but not RLIMIT_AS — the +// symbol simply is not there, so referencing syscall.RLIMIT_AS the way +// rlimit_as_unix.go does for every other unix platform would fail to +// compile here. applyRlimits treats ok=false as "skip this one, warn if +// the operator asked for it" rather than guessing a raw resource number +// this package has not verified for OpenBSD. +func rlimitAS() (resource int, ok bool) { + return 0, false +} diff --git a/exec/shim/rlimit_as_unix.go b/exec/shim/rlimit_as_unix.go new file mode 100644 index 0000000..102e088 --- /dev/null +++ b/exec/shim/rlimit_as_unix.go @@ -0,0 +1,16 @@ +//go:build unix && !openbsd + +package shim + +import "syscall" + +// rlimitAS reports the RLIMIT_AS resource number and whether it is +// usable on this platform. Everywhere this file builds, that is simply +// syscall.RLIMIT_AS: Go resolves the named constant to the correct +// number for the platform and architecture being compiled for, which is +// verified true for every unix platform except one — see +// rlimit_as_openbsd.go, the build-tag complement of this file, for why +// OpenBSD needs its own implementation instead of this one. +func rlimitAS() (resource int, ok bool) { + return syscall.RLIMIT_AS, true +} diff --git a/exec/shim/rlimit_other.go b/exec/shim/rlimit_other.go new file mode 100644 index 0000000..6798b01 --- /dev/null +++ b/exec/shim/rlimit_other.go @@ -0,0 +1,14 @@ +//go:build !unix + +package shim + +// applyRlimits is a no-op outside Unix: RLIMIT_* and syscall.Setrlimit +// are POSIX concepts this rung does not emulate anywhere else. It is +// never reached in practice — exec/subprocess's checkLaunch refuses to +// launch the child at all on a non-Unix platform (limits_other.go, in +// that package) — but this stub still needs to exist so this package +// compiles there: main.go's call to applyRlimits is unconditional, since +// mainExitCode itself has no reason to know which platform it is running +// on. Returning no failures means EnvRlimitStrict never trips here +// either, which is moot in practice for the same reason. +func applyRlimits() []rlimitFailure { return nil } diff --git a/exec/shim/rlimit_unix.go b/exec/shim/rlimit_unix.go new file mode 100644 index 0000000..7a60db5 --- /dev/null +++ b/exec/shim/rlimit_unix.go @@ -0,0 +1,223 @@ +//go:build unix + +package shim + +import ( + "errors" + "fmt" + "os" + "runtime" + "strconv" + "syscall" +) + +// rlimitSpec pairs the environment variable the parent sets with the raw +// setrlimit(2) resource number and a label for diagnostics. resourceOK is +// false when this platform's exact resource number for the limit has not +// been verified — see rlimitAS and rlimitNProc — in which case +// applyRlimits skips the syscall.Setrlimit call entirely rather than risk +// handing it the wrong resource number, and treats that skip as a +// failure like any other: this is Dispatch not having checked the +// constant for this platform, a library gap, not the platform itself +// refusing the limit — see isKnownUnsupported for that distinct case, +// which only applies to an actual syscall.Setrlimit failure. +type rlimitSpec struct { + env string + resource int + resourceOK bool + label string +} + +// rlimitSpecs lists every limit applyRlimits knows how to apply. +// +// RLIMIT_CORE, RLIMIT_NOFILE, and RLIMIT_FSIZE are taken from the +// syscall package's own named constants, which is safe everywhere this +// file builds (verified against Go's generated zerrors tables for every +// platform under the "unix" build constraint: aix, darwin, dragonfly, +// freebsd, illumos/solaris, linux — including its mips/mips64 variants, +// which differ from every other linux arch — netbsd, and openbsd): the +// symbol is present on all of them, and Go resolves it to the correct +// platform- and arch-specific number automatically. RLIMIT_AS and +// RLIMIT_NPROC do not have that property — see rlimitAS and rlimitNProc +// for why each needs its own platform-aware resolution instead of a bare +// syscall.RLIMIT_* reference. +func rlimitSpecs() []rlimitSpec { + as, asOK := rlimitAS() + nproc, nprocOK := rlimitNProc() + + return []rlimitSpec{ + // Applied first: this is the one limit the parent always sends, + // and it is the one the worker's own security promise depends + // on most directly. + {EnvRlimitCore, syscall.RLIMIT_CORE, true, "RLIMIT_CORE"}, + {EnvRlimitAS, as, asOK, "RLIMIT_AS"}, + {EnvRlimitNoFile, syscall.RLIMIT_NOFILE, true, "RLIMIT_NOFILE"}, + {EnvRlimitFSize, syscall.RLIMIT_FSIZE, true, "RLIMIT_FSIZE"}, + {EnvRlimitNProc, nproc, nprocOK, "RLIMIT_NPROC"}, + } +} + +// rlimitNProc reports the raw RLIMIT_NPROC resource number for the +// current platform and GOARCH, and whether it is verified. +// +// Go's syscall package does not export RLIMIT_NPROC on any platform — it +// was trimmed from the generated zerrors tables along with +// RLIMIT_MEMLOCK and RLIMIT_RSS — so every value below is a raw resource +// number taken from each platform's own layout, not a +// named constant. Getting this wrong is worse than getting RLIMIT_AS +// wrong: an incorrect resource number does not fail to compile or even +// fail at runtime, it just silently caps a *different* resource. This +// task's original version hardcoded 6 for every non-Darwin platform, +// which is only correct for Linux's "asm-generic" architectures — on +// Linux/mips and Linux/mips64, position 6 is RLIMIT_AS, not +// RLIMIT_NPROC, so a configured NProc limit would have silently become +// an AS limit there instead, killing the child on its first allocation +// while the logs claimed NProc was applied fine. +// +// Verified per (GOOS, GOARCH) against Go's own generated +// syscall/zerrors_*.go tables plus the reported values for platforms Go +// does not export the surrounding constants for at all: +// - linux, non-mips (amd64, arm64, 386, arm, riscv64, ppc64, ppc64le, +// s390x, loong64): 6, matching — the same +// ordering behind Go's own AS=9/NOFILE=7 on every one of these arches. +// - linux, mips family (mips, mipsle, mips64, mips64le): 8, matching +// that family's own distinct resource.h layout (also the reason its +// AS=6 and NOFILE=5 diverge from every other linux arch). +// - darwin: 7. +// - freebsd: 7. +// +// Every other platform under the "unix" build constraint — netbsd, +// openbsd, dragonfly, solaris/illumos, aix, android, ios — returns +// ok=false rather than a guessed number: this rung has not verified +// RLIMIT_NPROC's position for any of them, and a wrong guess here is a +// silent security regression, not a build failure or a loud one, so +// applyRlimits skips it with a warning instead of risking that. +func rlimitNProc() (resource int, ok bool) { + switch runtime.GOOS { + case "linux": + switch runtime.GOARCH { + case "mips", "mipsle", "mips64", "mips64le": + return 8, true + default: + return 6, true + } + case "darwin": + return 7, true + case "freebsd": + return 7, true + default: + return 0, false + } +} + +// applyRlimits reads every limit the parent set in the environment (see +// EnvRlimitAS and friends) and applies it via syscall.Setrlimit before the +// handler ever runs. It returns every failure that was not judged "the +// platform itself refuses this limit" — see isKnownUnsupported — for +// mainExitCode to act on when EnvRlimitStrict (subprocess.WithStrictRlimits) +// is set; regardless of that, every failure is always logged to stderr +// here, whether it counts toward strict mode or not. +// +// A limit whose env var is unset is left alone — that is how buildEnv +// says "no opinion" for anything but RLIMIT_CORE, which it always sends. +// A limit whose resource number is not verified for this platform +// (resourceOK false — see rlimitAS and rlimitNProc) is skipped without +// ever calling Setrlimit, since there is no safe resource number to +// pass — but this counts as a failure like any other, not as +// known-unsupported: it is Dispatch not having checked the constant for +// this platform, a library gap that an update to rlimitAS/rlimitNProc +// could close, which is a different thing from the platform's own kernel +// refusing the limit (isKnownUnsupported, below) — an operator who asked +// for WithStrictRlimits gets told either way, since "we haven't verified +// this" is exactly the kind of silent gap strict mode exists to surface. +// A malformed value the parent sent (which should never happen — buildEnv +// only ever writes strconv.FormatInt output — but is not trusted blindly +// here regardless) is the same: a failure, unconditionally. +// +// By default — without WithStrictRlimits — every failure here, of any +// kind, is logged and otherwise ignored rather than treated as a launch +// failure, on the reasoning that an individual rlimit is one layer among +// several this rung provides — the uid boundary and the process-group +// kill hold regardless — and because Darwin's kernel rejects +// setrlimit(RLIMIT_AS, ...) outright (EINVAL) no matter what value is +// requested, so treating every failure as fatal by default would make +// this rung unusable there in practice. WithStrictRlimits opts a specific +// Executor out of that default when an operator has decided a limit not +// applying — for any reason, verified-but-refused or simply unverified — +// is worse than the attempt not running at all. +// +// Every message here goes to stderr, not through the logger the request +// carries: WithLogger's own default is a no-op logger, so anything routed +// through it would be invisible unless the operator opted in, which is +// backwards for a message that exists specifically to warn the operator. +func applyRlimits() []rlimitFailure { + var failures []rlimitFailure + + for _, s := range rlimitSpecs() { + v, ok := os.LookupEnv(s.env) + if !ok { + continue + } + + if f, hasFailure := applyOne(s, v); hasFailure { + failures = append(failures, f) + } + } + + return failures +} + +// applyOne applies a single spec whose env var is set to v, and reports +// the failure it produced, if any. Split out from applyRlimits so the +// resourceOK-false path — reachable in practice only on a platform this +// package has not verified any spec's resource number for, none of which +// this rung's CI or dev machine are — can be exercised directly by a +// test on any platform, by constructing a synthetic rlimitSpec, rather +// than requiring an actual unverified platform to prove it produces a +// failure instead of a silent skip. +func applyOne(s rlimitSpec, v string) (rlimitFailure, bool) { + if !s.resourceOK { + err := fmt.Errorf("this platform's resource number for %s is not verified by dispatch, refusing to guess", s.label) + fmt.Fprintf(os.Stderr, "dispatch/exec/shim: %s=%s requested but %v, skipping\n", s.label, v, err) + + return rlimitFailure{s.label, err}, true + } + + n, err := strconv.ParseInt(v, 10, 64) + if err != nil || n < 0 { + err = fmt.Errorf("value %q is invalid", v) + fmt.Fprintf(os.Stderr, "dispatch/exec/shim: %s %v, skipping\n", s.label, err) + + return rlimitFailure{s.label, err}, true + } + + if err := syscall.Setrlimit(s.resource, newRlimit(n)); err != nil { + fmt.Fprintf(os.Stderr, "dispatch/exec/shim: setrlimit %s=%d failed, continuing without it: %v\n", s.label, n, err) + + if !isKnownUnsupported(s.label, err) { + return rlimitFailure{s.label, err}, true + } + } + + return rlimitFailure{}, false +} + +// isKnownUnsupported reports whether an actual syscall.Setrlimit failure +// — one that ran against a resource number this package has verified, +// via a resourceOK spec — is a structural, permanent fact about the +// current kernel rather than a misconfiguration. The standing example is +// Darwin, whose kernel rejects setrlimit(RLIMIT_AS, ...) with EINVAL +// unconditionally, for any value: the kernel is saying no, not Dispatch +// declining to guess. Anything else — most concretely EPERM because the +// requested value exceeds the process's own hard limit — is treated as +// unexpected, since on a platform that does support the limit, that +// shape of failure means the configured value itself is the problem. +// +// This function only ever sees a failure from a real Setrlimit call. +// resourceOK-false skips (rlimitSpecs, applyRlimits) never reach here at +// all — those are a library gap, not a platform refusal, and applyRlimits +// already counts them as failures unconditionally without asking this +// function anything. +func isKnownUnsupported(label string, err error) bool { + return runtime.GOOS == "darwin" && label == "RLIMIT_AS" && errors.Is(err, syscall.EINVAL) +} diff --git a/exec/shim/rlimit_unix_test.go b/exec/shim/rlimit_unix_test.go new file mode 100644 index 0000000..1c09e06 --- /dev/null +++ b/exec/shim/rlimit_unix_test.go @@ -0,0 +1,180 @@ +//go:build unix + +package shim + +// package shim (internal), not shim_test — same rationale as +// internal_test.go: mainExitCode's strict-vs-warning routing and +// isKnownUnsupported/joinRlimitFailures are unexported. This file is +// unix-tagged for the same practical reason internal_test.go is (see its +// own doc comment): the rlimit path this file exercises — env vars, +// mainExitCode's early-return branch, applyOne, isKnownUnsupported — is +// built from syscall.Setrlimit and the RLIMIT_* constants, which do not +// exist in Go's syscall package on Windows. +// +// What this file does NOT do: call applyRlimits, or applyOne directly, +// with a value that would actually succeed against a real resource like +// RLIMIT_AS or RLIMIT_NOFILE. callMainExitCode (internal_test.go) runs +// mainExitCode in this test binary's own process, not a forked child — a +// rlimit that actually took effect here would permanently lower it for +// every subsequent test in this same test binary run, since rlimits can +// only be lowered without privilege, never raised back. Most cases below +// use a value applyRlimits rejects before ever calling Setrlimit (a +// negative number); TestApplyOneUnverifiedResourceIsAFailure is the +// exception — it uses a positive, plausible-looking value ("12345") and +// is still safe, because its synthetic rlimitSpec sets resourceOK false, +// which makes applyOne return a failure without ever reaching Setrlimit +// regardless of what the value is. The cases that need a real Setrlimit +// outcome — proving WithStrictRlimits doesn't fire for a platform's own +// structural refusal, or does fire for a real failure — are +// exec/subprocess's TestStrictRlimitsFailsLaunchOnUnexpectedFailure and +// TestStrictRlimitsToleratesKnownUnsupported (limits_unix_test.go), which +// fork a real child and so cannot pollute this process. + +import ( + "context" + "errors" + "fmt" + "runtime" + "syscall" + "testing" + + "github.com/xraph/dispatch/job" +) + +func TestIsKnownUnsupported(t *testing.T) { + tests := []struct { + name string + label string + err error + want bool + }{ + { + name: "darwin RLIMIT_AS EINVAL is known-unsupported", + label: "RLIMIT_AS", + err: syscall.EINVAL, + want: isDarwin(), + }, + { + name: "RLIMIT_AS EPERM is not known-unsupported, even on Darwin", + label: "RLIMIT_AS", + err: syscall.EPERM, + want: false, + }, + { + name: "a different label with EINVAL is not known-unsupported", + label: "RLIMIT_NOFILE", + err: syscall.EINVAL, + want: false, + }, + { + name: "a wrapped EINVAL still matches, via errors.Is", + label: "RLIMIT_AS", + err: fmt.Errorf("setrlimit: %w", syscall.EINVAL), + want: isDarwin(), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isKnownUnsupported(tt.label, tt.err); got != tt.want { + t.Errorf("isKnownUnsupported(%q, %v) = %v, want %v", tt.label, tt.err, got, tt.want) + } + }) + } +} + +func TestJoinRlimitFailures(t *testing.T) { + tests := []struct { + name string + failures []rlimitFailure + want string + }{ + {name: "empty", failures: nil, want: ""}, + { + name: "one", + failures: []rlimitFailure{{"RLIMIT_NOFILE", errors.New("value \"-1\" is invalid")}}, + want: `RLIMIT_NOFILE: value "-1" is invalid`, + }, + { + name: "two, joined with semicolons", + failures: []rlimitFailure{ + {"RLIMIT_NOFILE", errors.New("boom")}, + {"RLIMIT_FSIZE", errors.New("also boom")}, + }, + want: "RLIMIT_NOFILE: boom; RLIMIT_FSIZE: also boom", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := joinRlimitFailures(tt.failures); got != tt.want { + t.Errorf("joinRlimitFailures() = %q, want %q", got, tt.want) + } + }) + } +} + +// TestMainExitCodeStrictRlimitFailsLaunch proves the EnvRlimitStrict +// wiring end to end through mainExitCode: an rlimit value that fails +// before Setrlimit is ever reached (invalid, safe to run in this +// process — see the file doc comment) becomes exit 1 with strict mode +// set, and stays exit 0 (a warning, per applyRlimits) without it. +func TestMainExitCodeStrictRlimitFailsLaunch(t *testing.T) { + defs := []job.Registrable{ + job.NewDefinition("internal.ok", func(context.Context, struct{}) error { return nil }), + } + + t.Run("default: invalid rlimit is a warning, launch proceeds", func(t *testing.T) { + t.Setenv(EnvRlimitNoFile, "-1") + + if got := callMainExitCode(t, defs, internalReq(t, "internal.ok")); got != 0 { + t.Errorf("mainExitCode() = %d, want 0", got) + } + }) + + t.Run("strict: invalid rlimit fails the launch", func(t *testing.T) { + t.Setenv(EnvRlimitNoFile, "-1") + t.Setenv(EnvRlimitStrict, "1") + + if got := callMainExitCode(t, defs, internalReq(t, "internal.ok")); got != 1 { + t.Errorf("mainExitCode() = %d, want 1", got) + } + }) +} + +// TestApplyOneUnverifiedResourceIsAFailure is the regression test for the +// gap the review round 2 found: resourceOK false used to log and +// `continue` without ever appending to the failures slice, which meant +// WithStrictRlimits gave no guarantee at all on a platform this package +// has not verified a resource number for — the one case an operator +// opting into strictness needs it most. A synthetic rlimitSpec is used +// rather than a real unverified platform (this rung's CI and dev machine +// are both verified for every field WithRlimits exposes today), which is +// exactly what applyOne being split out of applyRlimits is for. +func TestApplyOneUnverifiedResourceIsAFailure(t *testing.T) { + spec := rlimitSpec{env: "DISPATCH_TEST_UNVERIFIED", resource: 0, resourceOK: false, label: "RLIMIT_TEST"} + + f, hasFailure := applyOne(spec, "12345") + if !hasFailure { + t.Fatal("applyOne() reported no failure for an unverified resource number, want a failure so WithStrictRlimits actually refuses the launch") + } + if f.label != "RLIMIT_TEST" { + t.Errorf("failure label = %q, want %q", f.label, "RLIMIT_TEST") + } + if f.err == nil { + t.Error("failure err = nil, want a non-nil reason") + } +} + +// A resourceOK-true / real-Setrlimit contrast case (proving +// isKnownUnsupported actually reaches applyOne's result rather than only +// being unit-tested in isolation) deliberately does not live here: doing +// that against RLIMIT_AS in-process, in this test binary, would risk +// exactly the pollution this file's doc comment describes avoiding — on +// Linux, a generous-enough value would not fail at all, and would then +// apply for real, for the rest of this test binary's life. +// TestStrictRlimitsToleratesKnownUnsupported (exec/subprocess, +// limits_unix_test.go) covers that case instead, through a real forked +// child, which cannot pollute anything once it exits. + +func isDarwin() bool { return runtime.GOOS == "darwin" } diff --git a/exec/shim/rlimit_value_bsd64.go b/exec/shim/rlimit_value_bsd64.go new file mode 100644 index 0000000..d025bea --- /dev/null +++ b/exec/shim/rlimit_value_bsd64.go @@ -0,0 +1,14 @@ +//go:build freebsd || dragonfly + +package shim + +import "syscall" + +// newRlimit builds a syscall.Rlimit with both Cur and Max set to n. +// FreeBSD and Dragonfly are the two platforms under the "unix" build +// constraint whose generated syscall.Rlimit uses int64 fields rather +// than uint64 — everywhere else newRlimit needs the uint64 conversion +// this file's build-tag complement (rlimit_value_unix.go) provides. +func newRlimit(n int64) *syscall.Rlimit { + return &syscall.Rlimit{Cur: n, Max: n} +} diff --git a/exec/shim/rlimit_value_unix.go b/exec/shim/rlimit_value_unix.go new file mode 100644 index 0000000..6cf5432 --- /dev/null +++ b/exec/shim/rlimit_value_unix.go @@ -0,0 +1,15 @@ +//go:build unix && !freebsd && !dragonfly + +package shim + +import "syscall" + +// newRlimit builds a syscall.Rlimit with both Cur and Max set to n. +// Everywhere this file builds, Rlimit's fields are uint64 (verified +// against Go's generated ztypes_*.go for aix, darwin, linux, netbsd, +// openbsd, and solaris/illumos) — see rlimit_value_bsd64.go for the two +// platforms where that is not true. +func newRlimit(n int64) *syscall.Rlimit { + //nolint:gosec // G115: n is validated non-negative by applyRlimits before this is called. + return &syscall.Rlimit{Cur: uint64(n), Max: uint64(n)} +} diff --git a/exec/shim/store.go b/exec/shim/store.go new file mode 100644 index 0000000..9610baf --- /dev/null +++ b/exec/shim/store.go @@ -0,0 +1,216 @@ +package shim + +import ( + "context" + "sync" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// memStore is a minimal, in-process artifact.Store for the sandboxed +// shim. It exists so exec/shim never has to import store/memory, which is +// the full composite store and transitively pulls the worker's +// database, object-store, and config-loader clients into the sandbox +// binary. +// +// The shim's use of a store is single-process, single-attempt, and +// short-lived: it backs one artifact.Service for the lifetime of one +// handler invocation, only ever seeded by seedPriorOutputs and read or +// written by that handler through artifact.Accessor. It is not a record +// of truth — the worker outside the sandbox verifies what actually landed +// on disk — so memStore implements real map semantics only for the +// methods that path reaches (CreateArtifact, GetArtifact, +// FindArtifactByKey, LinkArtifact, FindLinkByName) and honest stubs for +// the rest. A sandboxed handler never lists, sweeps, or purges: lifecycle +// management is the worker's job, not the sandbox's. +type memStore struct { + mu sync.Mutex + artifacts map[string]*artifact.Artifact + links []*artifact.Link +} + +var _ artifact.Store = (*memStore)(nil) + +// newMemStore builds an empty memStore. +func newMemStore() *memStore { + return &memStore{ + artifacts: make(map[string]*artifact.Artifact), + } +} + +// CreateArtifact inserts an artifact and, when link is non-nil, its first +// link, under the store's lock. Returns artifact.ErrExists if an artifact +// already exists at the same backend, bucket, and key. +func (s *memStore) CreateArtifact(_ context.Context, a *artifact.Artifact, link *artifact.Link) error { + s.mu.Lock() + defer s.mu.Unlock() + + for _, existing := range s.artifacts { + if existing.DeletedAt != nil { + continue + } + + if existing.Backend == a.Backend && existing.Bucket == a.Bucket && existing.Key == a.Key { + return artifact.ErrExists + } + } + + s.artifacts[a.ID.String()] = a.Clone() + + if link != nil { + s.appendLinkLocked(link) + } + + return nil +} + +// appendLinkLocked adds link unless an identical one is already present. +// Callers must hold s.mu. +func (s *memStore) appendLinkLocked(link *artifact.Link) { + for _, existing := range s.links { + if existing.ArtifactID == link.ArtifactID && + existing.OwnerKind == link.OwnerKind && + existing.OwnerID == link.OwnerID && + existing.Name == link.Name && + existing.Attempt == link.Attempt { + return + } + } + + s.links = append(s.links, link.Clone()) +} + +// GetArtifact retrieves an artifact by ID. Returns artifact.ErrNotFound if +// it does not exist or has been soft-deleted. +func (s *memStore) GetArtifact(_ context.Context, artifactID id.ArtifactID) (*artifact.Artifact, error) { + s.mu.Lock() + defer s.mu.Unlock() + + a, ok := s.artifacts[artifactID.String()] + if !ok || a.DeletedAt != nil { + return nil, artifact.ErrNotFound + } + + return a.Clone(), nil +} + +// FindArtifactByKey retrieves an artifact by its storage coordinates. +// Returns artifact.ErrNotFound if none exists. +func (s *memStore) FindArtifactByKey(_ context.Context, backend, bucket, key string) (*artifact.Artifact, error) { + s.mu.Lock() + defer s.mu.Unlock() + + for _, a := range s.artifacts { + if a.DeletedAt != nil { + continue + } + + if a.Backend == backend && a.Bucket == bucket && a.Key == key { + return a.Clone(), nil + } + } + + return nil, artifact.ErrNotFound +} + +// UpdateArtifact is an honest stub: no path the shim exercises calls it, +// since updating size, hash, content type, or expiry after the fact is a +// worker-side concern once the object has landed on the real backend. It +// returns nil rather than an error so a caller that reaches it anyway is +// not broken by a store that only the sandbox uses. +func (s *memStore) UpdateArtifact(_ context.Context, _ *artifact.Artifact) error { + return nil +} + +// ListArtifacts is an honest stub: the sandbox never lists, since it has +// no notion of "every artifact" beyond what this one attempt created or +// was seeded with. It returns an empty result rather than an error. +func (s *memStore) ListArtifacts(_ context.Context, _ artifact.ListOpts) ([]*artifact.Artifact, error) { + return nil, nil +} + +// LinkArtifact records that an owner references an artifact. Linking the +// same artifact, owner, name, and attempt twice is a no-op. +func (s *memStore) LinkArtifact(_ context.Context, link *artifact.Link) error { + s.mu.Lock() + defer s.mu.Unlock() + + s.appendLinkLocked(link) + + return nil +} + +// ListLinks is an honest stub: nothing in the shim's path enumerates an +// owner's links wholesale — FindLinkByName answers the one question +// IfAbsent needs. It returns an empty result rather than an error. +func (s *memStore) ListLinks(_ context.Context, _ artifact.OwnerRef) ([]*artifact.Link, error) { + return nil, nil +} + +// FindLinkByName returns the link for an owner and name with the highest +// attempt number. Returns artifact.ErrNotFound if no attempt has produced +// it. +func (s *memStore) FindLinkByName(_ context.Context, owner artifact.OwnerRef, name string) (*artifact.Link, error) { + s.mu.Lock() + defer s.mu.Unlock() + + var best *artifact.Link + + for _, l := range s.links { + if l.OwnerKind != owner.Kind || l.OwnerID != owner.ID || l.Name != name { + continue + } + + if best == nil || l.Attempt > best.Attempt { + best = l + } + } + + if best == nil { + return nil, artifact.ErrNotFound + } + + return best.Clone(), nil +} + +// ListArtifactsByOwner is an honest stub: the shim's accessor resolves +// inputs from the request's InputSlots on local disk, not by asking the +// store what is linked to the owner. It returns an empty result rather +// than an error. +func (s *memStore) ListArtifactsByOwner( + _ context.Context, + _ artifact.OwnerRef, + _ artifact.Role, +) ([]*artifact.Artifact, error) { + return nil, nil +} + +// SweepEphemeral is an honest stub: the sandbox never sweeps. The worker +// outside the sandbox owns ephemeral-artifact lifecycle once the real +// object has landed on its backend. It returns an empty result rather +// than an error. +func (s *memStore) SweepEphemeral(_ context.Context, _ artifact.SweepOpts) ([]*artifact.Artifact, error) { + return nil, nil +} + +// SweepOrphans is an honest stub for the same reason as SweepEphemeral: +// orphan reclamation is the worker's business, not a single short-lived +// sandbox attempt's. It returns an empty result rather than an error. +func (s *memStore) SweepOrphans(_ context.Context, _ time.Time, _ int) ([]*artifact.Artifact, error) { + return nil, nil +} + +// ListPurgeable is an honest stub: purging soft-deleted rows is lifecycle +// management the worker performs, never the sandbox. It returns an empty +// result rather than an error. +func (s *memStore) ListPurgeable(_ context.Context, _ time.Duration, _ int) ([]*artifact.Artifact, error) { + return nil, nil +} + +// PurgeArtifact is an honest stub for the same reason as ListPurgeable. It +// returns nil rather than an error. +func (s *memStore) PurgeArtifact(_ context.Context, _ id.ArtifactID) error { + return nil +} diff --git a/exec/shim/store_test.go b/exec/shim/store_test.go new file mode 100644 index 0000000..d6a668f --- /dev/null +++ b/exec/shim/store_test.go @@ -0,0 +1,369 @@ +package shim + +// This file is package shim (internal), not shim_test, deliberately +// breaking the external-tests-only convention the rest of this package +// follows -- the same exception internal_test.go documents. memStore and +// newMemStore are unexported: the point of these tests is the sandbox's +// own persistence semantics, not just what artifact.Service exposes +// through them, so there is no way to reach the type under test from an +// external package. + +import ( + "context" + "errors" + osexec "os/exec" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +func newTestArtifact(key string) *artifact.Artifact { + return &artifact.Artifact{ + ID: id.NewArtifactID(), + Backend: "memory", + Bucket: "b", + Key: key, + Lifecycle: artifact.Ephemeral, + CreatedAt: time.Now().UTC(), + } +} + +func newTestOwner() artifact.OwnerRef { + return artifact.OwnerRef{Kind: artifact.OwnerJob, ID: "job-1"} +} + +func newTestLink(artifactID id.ArtifactID, owner artifact.OwnerRef, name string, attempt int) *artifact.Link { + return &artifact.Link{ + ArtifactID: artifactID, + OwnerKind: owner.Kind, + OwnerID: owner.ID, + Role: artifact.RoleOutput, + Name: name, + Attempt: attempt, + CreatedAt: time.Now().UTC(), + } +} + +func TestMemStore_CreateArtifact(t *testing.T) { + tests := []struct { + name string + seed bool + wantErr error + }{ + {name: "insert succeeds"}, + {name: "duplicate backend/bucket/key returns ErrExists", seed: true, wantErr: artifact.ErrExists}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := newMemStore() + ctx := context.Background() + owner := newTestOwner() + a := newTestArtifact("mesh.glb") + + if tt.seed { + seed := newTestArtifact("mesh.glb") + if err := s.CreateArtifact(ctx, seed, nil); err != nil { + t.Fatalf("seed CreateArtifact() error = %v, want nil", err) + } + } + + link := newTestLink(a.ID, owner, "mesh", 0) + + err := s.CreateArtifact(ctx, a, link) + if !errors.Is(err, tt.wantErr) { + t.Fatalf("CreateArtifact() error = %v, want %v", err, tt.wantErr) + } + + if tt.wantErr != nil { + return + } + + got, gerr := s.GetArtifact(ctx, a.ID) + if gerr != nil { + t.Fatalf("GetArtifact() error = %v, want nil", gerr) + } + + if got.Key != a.Key { + t.Errorf("GetArtifact().Key = %q, want %q", got.Key, a.Key) + } + }) + } +} + +func TestMemStore_GetArtifact(t *testing.T) { + tests := []struct { + name string + lookup func(seeded id.ArtifactID) id.ArtifactID + wantErr error + }{ + { + name: "known id returns the artifact", + lookup: func(seeded id.ArtifactID) id.ArtifactID { return seeded }, + }, + { + name: "unknown id returns ErrNotFound", + lookup: func(id.ArtifactID) id.ArtifactID { return id.NewArtifactID() }, + wantErr: artifact.ErrNotFound, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := newMemStore() + ctx := context.Background() + a := newTestArtifact("mesh.glb") + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact() error = %v, want nil", err) + } + + _, err := s.GetArtifact(ctx, tt.lookup(a.ID)) + if !errors.Is(err, tt.wantErr) { + t.Fatalf("GetArtifact() error = %v, want %v", err, tt.wantErr) + } + }) + } +} + +func TestMemStore_FindArtifactByKey(t *testing.T) { + tests := []struct { + name string + lookupKey string + wantErr error + }{ + {name: "finds what CreateArtifact inserted", lookupKey: "mesh.glb"}, + {name: "unknown key returns ErrNotFound", lookupKey: "missing.glb", wantErr: artifact.ErrNotFound}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := newMemStore() + ctx := context.Background() + a := newTestArtifact("mesh.glb") + + if err := s.CreateArtifact(ctx, a, nil); err != nil { + t.Fatalf("CreateArtifact() error = %v, want nil", err) + } + + got, err := s.FindArtifactByKey(ctx, a.Backend, a.Bucket, tt.lookupKey) + if !errors.Is(err, tt.wantErr) { + t.Fatalf("FindArtifactByKey() error = %v, want %v", err, tt.wantErr) + } + + if tt.wantErr != nil { + return + } + + if got.ID != a.ID { + t.Errorf("FindArtifactByKey().ID = %v, want %v", got.ID, a.ID) + } + }) + } +} + +func TestMemStore_FindLinkByName(t *testing.T) { + t.Run("returns the highest attempt when two links share a name", func(t *testing.T) { + s := newMemStore() + ctx := context.Background() + owner := newTestOwner() + + older := newTestArtifact("v0.glb") + newer := newTestArtifact("v1.glb") + + if err := s.CreateArtifact(ctx, older, newTestLink(older.ID, owner, "mesh", 0)); err != nil { + t.Fatalf("CreateArtifact(older) error = %v, want nil", err) + } + + if err := s.CreateArtifact(ctx, newer, newTestLink(newer.ID, owner, "mesh", 1)); err != nil { + t.Fatalf("CreateArtifact(newer) error = %v, want nil", err) + } + + got, err := s.FindLinkByName(ctx, owner, "mesh") + if err != nil { + t.Fatalf("FindLinkByName() error = %v, want nil", err) + } + + if got.ArtifactID != newer.ID { + t.Errorf("FindLinkByName().ArtifactID = %v, want %v", got.ArtifactID, newer.ID) + } + + if got.Attempt != 1 { + t.Errorf("FindLinkByName().Attempt = %d, want 1", got.Attempt) + } + }) + + t.Run("no attempt has produced the name returns ErrNotFound", func(t *testing.T) { + s := newMemStore() + + _, err := s.FindLinkByName(context.Background(), newTestOwner(), "mesh") + if !errors.Is(err, artifact.ErrNotFound) { + t.Fatalf("FindLinkByName() error = %v, want %v", err, artifact.ErrNotFound) + } + }) +} + +// TestMemStore_ConcurrentCreateArtifact proves memStore's mutex actually +// guards the map: several goroutines each creating a distinct artifact +// must not corrupt the store or trip the race detector, and every +// artifact created must be readable afterward. +func TestMemStore_ConcurrentCreateArtifact(t *testing.T) { + s := newMemStore() + ctx := context.Background() + owner := newTestOwner() + + const n = 32 + + artifacts := make([]*artifact.Artifact, n) + for i := range artifacts { + artifacts[i] = newTestArtifact("out-" + strconv.Itoa(i) + ".glb") + } + + var wg sync.WaitGroup + + errs := make([]error, n) + + for i := range artifacts { + wg.Add(1) + + go func(i int) { + defer wg.Done() + + link := newTestLink(artifacts[i].ID, owner, artifacts[i].Key, 0) + errs[i] = s.CreateArtifact(ctx, artifacts[i], link) + }(i) + } + + wg.Wait() + + for i, err := range errs { + if err != nil { + t.Fatalf("CreateArtifact(%d) error = %v, want nil", i, err) + } + } + + for i, a := range artifacts { + if _, err := s.GetArtifact(ctx, a.ID); err != nil { + t.Errorf("GetArtifact(%d) error = %v, want nil", i, err) + } + } +} + +// shimAllowedModules lists the exact third-party (and xraph/go-utils) +// module roots exec/shim's dependency closure may reach, besides the Go +// standard library and dispatch's own packages (shimAllowedDispatchPackages, +// below). A dependency matches an entry here if it equals the entry or +// sits under it (e.g. "go.uber.org/zap/zapcore" under "go.uber.org/zap"): +// that is ordinary internal structure of an already-allowed module, not a +// new one, and stays permissive across a version bump that adds or renames +// an internal subpackage of something already here. +// +// This list is small and meant to stay that way — see +// TestShimLinksNoInfrastructure's own doc comment for why it exists as an +// allowlist rather than a denylist. +var shimAllowedModules = []string{ + "github.com/vmihailenco/msgpack", // wire.Encode/Decode's frame encoding + "github.com/vmihailenco/tagparser", // msgpack's own struct-tag parser + "github.com/zeebo/blake3", // artifact content hashing + "github.com/klauspost/cpuid", // blake3's SIMD dispatch + // cpuid v2.4.0 reads CPU features through syscalls (AT_HWCAP on + // linux, sysctl on darwin) where older versions used CPUID + // instructions alone, so this arrived with a version bump rather + // than with any change here. Raw syscall bindings are the furthest + // thing from an infrastructure client, which is what this list + // exists to keep out, and the prefix rule above does not cover it + // because x/sys is its own module rather than a subpackage of + // something already allowed. + "golang.org/x/sys", // cpuid's CPU feature detection + "go.jetify.com/typeid", // ID generation + "github.com/gofrs/uuid", // typeid's underlying uuid generator + "go.uber.org/zap", // logging + "go.uber.org/multierr", // zap's error-joining helper + "github.com/xraph/go-utils/log", // the log.Logger interface dispatch itself uses +} + +// shimAllowedDispatchPackages lists the exact dispatch packages exec/shim's +// dependency closure may reach. Deliberately NOT a prefix match on +// "github.com/xraph/dispatch": that would let github.com/xraph/dispatch/store +// pass silently just because it shares the module root, which is exactly +// the shape of gap this test exists to close (see the doc comment below). +var shimAllowedDispatchPackages = map[string]bool{ + "github.com/xraph/dispatch": true, + "github.com/xraph/dispatch/artifact": true, + "github.com/xraph/dispatch/exec": true, + "github.com/xraph/dispatch/exec/shim": true, // the package under test itself + "github.com/xraph/dispatch/exec/wire": true, + "github.com/xraph/dispatch/id": true, + "github.com/xraph/dispatch/job": true, + "github.com/xraph/dispatch/resource": true, +} + +// shimDepIsStdlib reports whether dep is a standard-library import, +// including its own internal/... and vendor/... packages (crypto/tls's +// vendored copy of x/crypto, for instance). A module path always has a +// dot in its first path segment (a domain, e.g. "github.com" or +// "go.uber.org"); the standard library never does. +func shimDepIsStdlib(dep string) bool { + first, _, _ := strings.Cut(dep, "/") + + return !strings.Contains(first, ".") +} + +// TestShimLinksNoInfrastructure fails if the sandbox binary gains an +// import that could reach a credential, a socket, or a config file. The +// phase's central claim is that this process holds none of those, and +// that claim should be checkable by inspection rather than by tracing +// which package-level variables happen not to be constructed. +// +// This is an allowlist, not a denylist. A denylist of forbidden substrings +// ("go-redis", "client-go", ...) only catches infrastructure clients this +// test's author already thought to name — a review round added +// `import _ "github.com/jackc/pgx/v5"` to this package and the previous +// denylist version of this test passed clean: ten pgx packages entered the +// closure, including pgconn (which holds connection credentials), and none +// of them matched "go-redis", "client-go", "confy", or "xraph/forge". An +// allowlist inverts the failure mode: anything not already known to be one +// of the handful of packages this sandbox actually needs — msgpack, +// blake3, typeid, zap, go-utils/log, and dispatch's own small, explicitly +// enumerated internal set — fails closed instead of open. +func TestShimLinksNoInfrastructure(t *testing.T) { + out, err := osexec.CommandContext(context.Background(), "go", "list", "-deps", "github.com/xraph/dispatch/exec/shim").Output() + if err != nil { + t.Skipf("go list unavailable: %v", err) + } + + for _, dep := range strings.Split(strings.TrimSpace(string(out)), "\n") { + if dep == "" || shimDepIsStdlib(dep) { + continue + } + + if strings.HasPrefix(dep, "github.com/xraph/dispatch") { + if !shimAllowedDispatchPackages[dep] { + t.Errorf("exec/shim links dispatch package %q, which is outside the sandbox's allowed "+ + "closure (shimAllowedDispatchPackages) — a sandbox process must not reach dispatch/store "+ + "or any other package that can hold a credential", dep) + } + + continue + } + + allowed := false + for _, prefix := range shimAllowedModules { + if dep == prefix || strings.HasPrefix(dep, prefix+"/") { + allowed = true + break + } + } + if !allowed { + t.Errorf("exec/shim links %q, which is outside the sandbox's allowed closure "+ + "(shimAllowedModules) — a sandbox process must hold no infrastructure client "+ + "(a database driver, a cache client, a config loader); if this is a deliberate new "+ + "dependency, add its module root here only after confirming it holds none", dep) + } + } +} diff --git a/exec/status.go b/exec/status.go new file mode 100644 index 0000000..aafc6c8 --- /dev/null +++ b/exec/status.go @@ -0,0 +1,56 @@ +package exec + +// Status classifies how an execution attempt ended. +// +// A bare error cannot express this. In-process, a handler returning an +// error and a handler dying are the same value; out-of-process they are +// different events needing different handling, and only some of them are +// the handler's fault. +type Status string + +const ( + // StatusOK means the handler ran and returned nil. + StatusOK Status = "ok" + + // StatusHandlerError means the handler ran and returned an error. + // This is a business failure and follows the normal retry path. + StatusHandlerError Status = "handler_error" + + // StatusTimeout means the deadline expired and the sandbox was + // killed. Unlike a cancelled context, this is enforced. + StatusTimeout Status = "timeout" + + // StatusOOMKilled means a memory limit was hit. The handler did not + // choose this and may succeed with a larger allocation. + StatusOOMKilled Status = "oom_killed" + + // StatusKilled means the process died on a signal — a SIGSEGV from a + // memory-unsafe parser, or a seccomp trap. It is security-relevant. + StatusKilled Status = "killed" + + // StatusLaunchFailed means the sandbox never started: an image pull + // failure, an exhausted quota, a missing runtime. The handler never + // ran, so this is infrastructure rather than work. + StatusLaunchFailed Status = "launch_failed" +) + +// IsFailure reports whether the status represents anything other than +// success. +func (s Status) IsFailure() bool { return s != StatusOK } + +// CountsAgainstRetries reports whether an attempt ending in this status +// should consume the job's retry budget. +// +// Launch failures do not. An ImagePullBackOff or a FailedScheduling says +// nothing about the work, and burning three retries on one bad node would +// send healthy jobs to the DLQ. +func (s Status) CountsAgainstRetries() bool { + switch s { + case StatusHandlerError, StatusTimeout, StatusOOMKilled, StatusKilled: + return true + case StatusOK, StatusLaunchFailed: + return false + default: + return true + } +} diff --git a/exec/subprocess/conformance_test.go b/exec/subprocess/conformance_test.go new file mode 100644 index 0000000..1ee9cb1 --- /dev/null +++ b/exec/subprocess/conformance_test.go @@ -0,0 +1,34 @@ +package subprocess_test + +import ( + "os" + "testing" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/exectest" + "github.com/xraph/dispatch/exec/subprocess" +) + +// TestSubprocessConformance runs this rung through the shared conformance +// suite with enforcement claimed, not just inherited: unlike the in-process +// rung, this one can actually kill a handler that ignores cancellation and +// cannot let a handler's panic reach the caller, so it claims Enforces and +// IsolatesPanic and the suite's DeadlineEnforced and PanicIsolated cases +// (plus the cooperative-deadline cases gated on Enforces) run for real +// instead of being skipped. ReportsUsage stays false: cgroups and PeakRSS +// are Phase 3, and a capability flag that lies is worse than one that is +// absent, because later rungs (OCI, Kubernetes) copy this file as their +// starting point. +func TestSubprocessConformance(t *testing.T) { + exectest.RunSuite(t, "subprocess", func(*testing.T) exec.Executor { + return subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithAllowSameUser(), // CI cannot drop privileges + ) + }, exectest.Capabilities{ + Enforces: true, + IsolatesPanic: true, + ReportsUsage: false, // cgroups and PeakRSS are Phase 3 + }) +} diff --git a/exec/subprocess/doc.go b/exec/subprocess/doc.go new file mode 100644 index 0000000..1fe8afb --- /dev/null +++ b/exec/subprocess/doc.go @@ -0,0 +1,54 @@ +// Package subprocess runs job handlers in a re-exec'd child process. +// +// The worker launches its own binary again with argv[1] set to +// shim.ArgName, so the child ends up running shim.Main with the same +// handler registry the worker itself has, by construction — there is no +// second binary to build, ship, or keep in sync. The request crosses to +// the child on fd 3 and the result crosses back on fd 4, using the +// wire package's length-prefixed frames; stdout and stderr are left free +// for the handler's own output and for whatever a native library writes, +// and are streamed to the configured logger instead. +// +// This is Dispatch's exec.LevelProcess rung: a crash, a panic, or a +// memory-unsafe parser going off the rails takes down the child, not the +// worker, and the child does not receive the worker's environment +// wholesale — buildEnv constructs it from the request's own Env plus a +// small fixed allowlist (PATH, HOME, TMPDIR) copied from the worker's, +// never a plain pass-through of os.Environ() — so it cannot read a +// credential that lived only in a worker environment variable outside +// that allowlist. It is not a sandbox in the mount/network/seccomp sense +// — that is exec.LevelSandboxed, a stronger rung built on the same wire +// protocol. +// +// # The uid/gid boundary +// +// WithUser configures a dedicated, low-privilege uid/gid for the child. +// It is required, not advisory: without one, the child would run as the +// worker's own uid and could read anything the worker can — the Dispatch +// config file on disk among it, which is where database credentials +// typically live, not just the worker's environment — so Run refuses to +// start at all unless WithAllowSameUser opts into that explicitly. See +// WithUser and WithAllowSameUser. +// +// A genuine uid drop — WithUser naming anything other than the worker's +// own uid — clears supplementary groups along with it: dropping to a +// different uid or gid already requires the same privilege setgroups(2) +// itself needs (CAP_SETUID/CAP_SETGID on Linux, root on Darwin), so +// os/exec's own "clear supplementary groups whenever Credential is set +// and NoSetGroups is false" behaviour applies in full (see sysProcAttr, +// procattr_unix.go). A worker running as root with, say, "docker" in its +// supplementary groups (a common shape for a systemd unit that also +// manages containers) does not hand that membership to a genuinely +// dropped child — the child keeps only the low-privilege uid/gid's own +// groups. +// +// The one case where supplementary groups are NOT cleared is +// WithAllowSameUser's same-uid, same-gid path, where Credential is +// otherwise a no-op: NoSetGroups is set there specifically because +// setgroups needs privilege this process does not have when it is not +// already root (every dev machine and CI run), and there is nothing to +// clear anyway, since the child is running as the worker's own account. +// That path already defeats most of this rung's purpose for the uid/gid +// boundary itself — see WithUser and WithAllowSameUser — so its +// supplementary-group behaviour is the smaller of the two problems. +package subprocess diff --git a/exec/subprocess/executor.go b/exec/subprocess/executor.go new file mode 100644 index 0000000..acb3366 --- /dev/null +++ b/exec/subprocess/executor.go @@ -0,0 +1,828 @@ +package subprocess + +import ( + "context" + "fmt" + "os" + osexec "os/exec" + "sort" + "strconv" + "sync" + "syscall" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/shim" + "github.com/xraph/dispatch/exec/wire" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" +) + +// Name is the identifier this executor registers under. +const Name = "subprocess" + +// requestFD and resultFD are the descriptors the child sees fd 3 and fd 4 +// as, once os/exec appends ExtraFiles after stdin/stdout/stderr. They are +// a single source of truth for both the ExtraFiles ordering below and the +// environment variables that tell the shim where to look, so the two can +// never drift apart. +const ( + requestFD = 3 + resultFD = 4 +) + +// options holds every Option's effect. The configured user and Rlimits are +// now enforced: checkLaunch (limits_unix.go / limits_other.go) refuses to +// start when no uid is configured, or when the configured uid matches the +// worker's own, in either case without AllowSameUser, sysProcAttr +// (procattr_unix.go) sets Credential from uid/gid, and +// buildEnv below passes rlimits to the child, which shim.Main applies via +// syscall.Setrlimit — except RLIMIT_AS, which buildEnv lets a job's own +// resource.Memory ceiling (Request.ResourceLimits) override per attempt, +// falling back to rlimits.AddressSpace only when the job declares +// nothing. The kill ladder's SIGTERM-then-grace-period-then-SIGKILL +// sequence runs in terminate (kill_unix.go), called from killProcess +// below. +type options struct { + binary string + args []string + env map[string]string + uid int + gid int + hasUser bool + allowSameUser bool + logger log.Logger + rlimits Rlimits + hasRlimits bool + strictRlimits bool + scratchDir string +} + +// Option configures an Executor. +type Option func(*options) + +// WithBinary sets the path to the binary the executor re-execs for every +// attempt. In production this is the worker's own executable, found via +// os.Executable; tests pass os.Args[0] so the re-exec'd child is the test +// binary itself, running the shim instead of go test. +func WithBinary(path string) Option { + return func(o *options) { o.binary = path } +} + +// WithArgs sets extra arguments appended after shim.ArgName when launching +// the child. Most deployments need none of these — the marker argument is +// enough for the binary to know to run the shim. +func WithArgs(args ...string) Option { + return func(o *options) { + o.args = append([]string(nil), args...) + } +} + +// WithEnv supplies the base environment for the child. It is merged with +// Request.Env, which wins on any key both sides set, and with a small +// fixed allowlist (PATH, HOME, TMPDIR) copied from the worker's own +// environment. The child never inherits os.Environ() wholesale — that is +// the entire point of this rung, since the worker's environment is where +// its own credentials tend to live. +func WithEnv(env map[string]string) Option { + return func(o *options) { + m := make(map[string]string, len(env)) + for k, v := range env { + m[k] = v + } + o.env = m + } +} + +// WithUser configures the uid and gid the child runs as, dropped via +// Credential on sysProcAttr before exec. Run refuses to start when no uid +// is configured at all, and also when one is configured that matches the +// worker's own, unless WithAllowSameUser is also given — see its doc +// comment for why. +// +// A uid or gid that genuinely differs from the worker's own also clears +// the child's supplementary groups, along with the primary uid/gid — see +// the package doc comment's "The uid/gid boundary" section for why that +// falls out of the same privilege check rather than needing separate +// handling, and for the one path (WithAllowSameUser) where it does not +// happen. +func WithUser(uid, gid int) Option { + return func(o *options) { + o.uid = uid + o.gid = gid + o.hasUser = true + } +} + +// WithAllowSameUser is the single opt-out for running this rung +// unisolated on the uid boundary. Without it, Run refuses to start in +// either of the two shapes that leave the child running as the worker's +// own uid: WithUser never called at all, or WithUser naming the worker's +// own uid explicitly. Both leave the child able to read every credential +// the isolation exists to hide — ~/.aws, /var/run/secrets, the Dispatch +// config itself — so both share this one switch rather than each getting +// its own. +func WithAllowSameUser() Option { + return func(o *options) { o.allowSameUser = true } +} + +// WithLogger sets where the child's stdout and stderr are streamed, each +// line tagged with the job's id and name. The default is a no-op logger, +// so output is silently discarded rather than reaching os.Stdout, which +// would interleave a handler's own output with the worker's. +func WithLogger(l log.Logger) Option { + return func(o *options) { o.logger = l } +} + +// WithRlimits configures POSIX resource limits for the child. buildEnv +// passes non-zero fields to the child as environment variables, and +// shim.Main applies them via syscall.Setrlimit before running the +// handler — see the Rlimits doc comment for why that happens child-side. +func WithRlimits(r Rlimits) Option { + return func(o *options) { + o.rlimits = r + o.hasRlimits = true + } +} + +// WithStrictRlimits makes a configured rlimit that did not actually take +// effect a launch failure instead of a warning — with one exception. +// shim.Main (applyRlimits) still treats the current kernel's own +// structural refusal to support a given limit at all as a warning +// regardless of this option: Darwin rejecting setrlimit(RLIMIT_AS, ...) +// unconditionally is the standing example, and there is nothing a +// caller-supplied value could have done differently about that, so +// making it fatal here would only make WithRlimits{AddressSpace: ...} +// unusable on Darwin rather than catch a real misconfiguration. +// +// Everything else this rung can fail on, it does catch, including two +// shapes worth naming explicitly: a value that exceeds the process's own +// hard limit (EPERM on a platform that does support the resource), and — +// distinct from the kernel refusing the limit — Dispatch itself not +// having verified the raw resource number for the current platform at +// all (RLIMIT_NPROC on most non-Linux/Darwin/FreeBSD Unixes, currently; +// see rlimitNProc in exec/shim/rlimit_unix.go). That second case is a +// library gap, not a platform fact, and without this option it would +// otherwise be indistinguishable, from the outside, from the limit +// simply having applied. +// +// Without this, a configured rlimit that fails to apply is logged to the +// child's stderr and otherwise ignored — see the Rlimits doc comment — +// which is invisible by default unless WithLogger is also configured, +// since WithLogger's own default discards output silently. An operator +// who wants a guarantee that a configured limit actually took effect, +// not just an attempt at one, should use this rather than relying on +// stderr being watched. +func WithStrictRlimits() Option { + return func(o *options) { o.strictRlimits = true } +} + +// WithScratchDir sets the directory under which each attempt gets a fresh +// working directory for the child's process (Cmd.Dir). It defaults to +// os.TempDir(). This is distinct from Request.OutputDir: that is where the +// handler writes artifacts through the accessor, while this is just a safe +// place for the process to start in, so a handler that writes a relative +// path outside the artifact API lands somewhere disposable instead of the +// worker's own working directory. +func WithScratchDir(path string) Option { + return func(o *options) { o.scratchDir = path } +} + +// Rlimits configures the POSIX resource limits applied to the child +// process. Go cannot set a child's rlimits through SysProcAttr, so +// buildEnv passes these child-side as environment variables (see +// shim.EnvRlimitAS and friends), and shim.Main applies them via +// syscall.Setrlimit before running the handler. Zero means "leave the +// limit at whatever the worker itself runs with." +type Rlimits struct { + // AddressSpace caps RLIMIT_AS in bytes. Note Darwin's kernel rejects + // setrlimit(RLIMIT_AS, ...) outright (EINVAL) regardless of value — + // shim.Main logs and continues rather than failing the attempt when + // that happens, since the rest of the isolation (uid boundary, + // process-group kill, the other limits) still holds. + AddressSpace int64 + // NoFile caps RLIMIT_NOFILE, the open file descriptor count. + NoFile int64 + // NProc caps RLIMIT_NPROC, the number of processes the child's uid + // may run — a second line of defence against a forking exploit even + // with the process group killed on timeout. + NProc int64 + // Core caps RLIMIT_CORE. buildEnv forces the child's actual limit to + // zero unconditionally, regardless of what is set here, so a + // segfaulting parser cannot dump the input that crashed it, and the + // worker's memory alongside it, to disk. This field is accepted for + // API symmetry with the other limits but has no effect. + Core int64 + // FSize caps RLIMIT_FSIZE in bytes, bounding how much a runaway + // handler can write before the kernel kills it outright. + FSize int64 +} + +// Executor runs handlers in a child process, re-exec'ing the configured +// binary as the shim for every attempt. It satisfies exec.Executor at +// exec.LevelProcess: a crash or a memory-unsafe parser going off the rails +// takes the child down, not the worker, and the child receives a +// constructed environment rather than the worker's own. +type Executor struct { + opts options +} + +var _ exec.Executor = (*Executor)(nil) + +// New creates a subprocess executor from options. Absent WithLogger, child +// output is discarded rather than reaching the worker's own stdout/stderr. +// Absent WithScratchDir, os.TempDir() is used. +func New(opts ...Option) *Executor { + o := options{ + logger: log.NewNoopLogger(), + scratchDir: os.TempDir(), + } + for _, opt := range opts { + opt(&o) + } + + return &Executor{opts: o} +} + +// Name identifies the executor. +func (e *Executor) Name() string { return Name } + +// Level reports that this executor isolates the handler into its own +// address space, but does not sandbox it — no mount, network, or PID +// namespace, no seccomp filter. +func (e *Executor) Level() exec.Level { return exec.LevelProcess } + +// Run launches the child, feeds it req over fd 3, and waits for either a +// result frame on fd 4 or the process ending on its own. A returned error +// means the child never started; every other outcome, including the child +// crashing or missing its deadline, is reported through Result.Status. +func (e *Executor) Run(ctx context.Context, req *exec.Request) (*exec.Result, error) { + if err := req.Validate(); err != nil { + return nil, fmt.Errorf("dispatch/exec/subprocess: invalid request: %w", err) + } + + // checkLaunch refuses before any pipe or process exists: on Unix, no + // uid configured at all, or a configured uid matching the worker's + // own, neither without WithAllowSameUser; on every other platform, + // unconditionally, since this rung has no isolation to offer there. + // See limits_unix.go / limits_other.go. + if err := checkLaunch(e.opts); err != nil { + return &exec.Result{ + Status: exec.StatusLaunchFailed, + HandlerErr: err.Error(), + }, nil + } + + // The request pipe: reqR becomes the child's fd 3, reqW is ours to + // write the frame on. The result pipe: resW becomes the child's fd 4, + // resR is ours to read the frame from. + reqR, reqW, err := os.Pipe() + if err != nil { + return nil, fmt.Errorf("dispatch/exec/subprocess: create request pipe: %w", err) + } + resR, resW, err := os.Pipe() + if err != nil { + _ = reqR.Close() + _ = reqW.Close() + + return nil, fmt.Errorf("dispatch/exec/subprocess: create result pipe: %w", err) + } + // A second, independent pair per stream, so draining stdout/stderr is + // entirely decoupled from Cmd.Wait()'s own bookkeeping. Cmd.StdoutPipe + // ties draining to Wait in a way that is unsafe to run concurrently + // with a Wait call happening in another goroutine (its own docs say + // so); assigning a plain *os.File to Cmd.Stdout/Cmd.Stderr instead + // means os/exec just dups the fd into the child and otherwise leaves + // it alone, so we can read our own end on our own schedule. + outR, outW, err := os.Pipe() + if err != nil { + _ = reqR.Close() + _ = reqW.Close() + _ = resR.Close() + _ = resW.Close() + + return nil, fmt.Errorf("dispatch/exec/subprocess: create stdout pipe: %w", err) + } + errR, errW, err := os.Pipe() + if err != nil { + _ = reqR.Close() + _ = reqW.Close() + _ = resR.Close() + _ = resW.Close() + _ = outR.Close() + _ = outW.Close() + + return nil, fmt.Errorf("dispatch/exec/subprocess: create stderr pipe: %w", err) + } + + scratch, err := os.MkdirTemp(e.opts.scratchDir, "dispatch-exec-") + if err != nil { + _ = reqR.Close() + _ = reqW.Close() + _ = resR.Close() + _ = resW.Close() + _ = outR.Close() + _ = outW.Close() + _ = errR.Close() + _ = errW.Close() + + return nil, fmt.Errorf("dispatch/exec/subprocess: create scratch dir: %w", err) + } + defer os.RemoveAll(scratch) // best-effort cleanup; a leftover empty scratch dir is not worth failing the attempt over + + args := append([]string{shim.ArgName}, e.opts.args...) + // CommandContext rather than Command to satisfy noctx; the context + // passed here is intentionally context.Background(), not the caller's + // ctx, because cancellation is handled explicitly below by the + // waitLoop select and killProcess. Wiring the caller's ctx in here too + // would give os/exec its own independent kill-on-cancel path (with its + // own WaitDelay semantics) racing the one this function already owns. + // #nosec G204 -- binary and args come from operator configuration + // (WithBinary/WithArgs), never from the untrusted job payload. + cmd := osexec.CommandContext(context.Background(), e.opts.binary, args...) + cmd.Env = e.buildEnv(req) + cmd.Dir = scratch + cmd.ExtraFiles = []*os.File{reqR, resW} // index 0 -> fd 3, index 1 -> fd 4, matching requestFD/resultFD above + cmd.Stdout = outW + cmd.Stderr = errW + cmd.SysProcAttr = sysProcAttr(e.opts) // Setpgid, so killProcess below can reach the whole group, not just this one process; Credential when a user is configured + + if err := cmd.Start(); err != nil { + _ = reqR.Close() + _ = reqW.Close() + _ = resR.Close() + _ = resW.Close() + _ = outR.Close() + _ = outW.Close() + _ = errR.Close() + _ = errW.Close() + + return &exec.Result{ + Status: exec.StatusLaunchFailed, + HandlerErr: err.Error(), + }, nil + } + + // The child inherited its own copies of these four fds across + // fork/exec. Ours are now redundant, and keeping them open is actively + // harmful: as long as our copy of resW stays open, resR can never see + // EOF, even after the child exits and closes its own copy — Run would + // block forever reading a result frame from a process that is already + // gone. The same reasoning applies to outW/errW for the stdio pipes. + // reqR only matters for symmetry; nothing reads from our copy anyway. + _ = reqR.Close() + _ = resW.Close() + _ = outW.Close() + _ = errW.Close() + + // Each of these closes exactly once no matter which of several racing + // paths gets there first: the dedicated writer goroutine below always + // closes reqW itself once it is done with it, and the drain-grace + // timeout further down can also force any of the four closed to + // unblock a reader or writer stuck on a descendant the child left + // behind. sync.OnceFunc is what keeps that from ever double-closing a + // file — recycled fd numbers make a double-close silently break an + // unrelated descriptor rather than just returning a harmless error. + closeReqW := sync.OnceFunc(func() { _ = reqW.Close() }) + closeOutR := sync.OnceFunc(func() { _ = outR.Close() }) + closeErrR := sync.OnceFunc(func() { _ = errR.Close() }) + closeResR := sync.OnceFunc(func() { _ = resR.Close() }) + defer closeReqW() + defer closeOutR() + defer closeErrR() + defer closeResR() + + var stdioWG sync.WaitGroup + stdioWG.Add(2) + go func() { + defer stdioWG.Done() + streamOutput(outR, e.opts.logger, req, "stdout") + }() + go func() { + defer stdioWG.Done() + streamOutput(errR, e.opts.logger, req, "stderr") + }() + + type frameRead struct { + frame *wire.Frame + err error + } + frameCh := make(chan frameRead, 1) + go func() { + f, ferr := wire.Decode(resR) + frameCh <- frameRead{f, ferr} + }() + + // writeRequest runs on its own goroutine instead of blocking Run's own + // goroutine here: a synchronous write large enough to fill the pipe + // buffer (64KB on Linux, often 16KB on macOS) would block until the + // child drains it, and if the child is alive but has not yet reached + // wire.Decode — busy, stopped, traced, or simply not a real shim — + // that block would sit outside the waitLoop select below, unreachable + // by both the deadline and ctx. Once the child is killed (or exits on + // its own), the last reader of reqR goes away and the pending Write + // unblocks with EPIPE on its own, without this needing to close reqW + // itself for that to happen. + encodeCh := make(chan error, 1) + go func() { + err := writeRequest(reqW, req) + closeReqW() + encodeCh <- err + }() + + waitCh := make(chan error, 1) + go func() { + waitCh <- cmd.Wait() + }() + + var deadlineCh <-chan time.Time + if !req.Deadline.IsZero() { + d := time.Until(req.Deadline) + if d < 0 { + d = 0 + } + timer := time.NewTimer(d) + defer timer.Stop() + deadlineCh = timer.C + } + + // grace is Policy.GracePeriod, falling back to exec.DefaultGracePeriod + // when the request carries a zero value — a Request built without + // exec.NewPolicy (which applies that same default itself) leaves + // Policy as its zero value, and a zero grace period would collapse + // the ladder in kill_unix.go's terminate back into an immediate + // SIGKILL, silently losing the whole point of Task 6 for any caller + // that did not opt in explicitly. + grace := req.Policy.GracePeriod + if grace <= 0 { + grace = exec.DefaultGracePeriod + } + + var ( + timedOut bool + callerDone bool + ctxDoneCh = ctx.Done() + ) + +waitLoop: + for { + select { + case <-waitCh: + break waitLoop + case <-deadlineCh: + deadlineCh = nil // this case must not fire again once handled + // select "chooses uniformly at random among those that can + // proceed" (Go spec), so if the process finished on its own in + // the same instant the timer fired, this case can still win + // even though waitCh is already deliverable. A cooperative + // handler makes that a live outcome, not a theoretical one: + // the shim traps SIGTERM and cancels its own handler context, + // and killProcess below sends SIGTERM as the first rung of + // its kill ladder, so "the tracked process exits right as the + // deadline fires" only gets more common, not less. Checking + // waitCh non-blockingly resolves the tie deterministically in + // favour of what actually happened to the process, instead of + // leaving classify to guess from a frame and a signal after + // the fact. + select { + case <-waitCh: + break waitLoop + default: + } + timedOut = true + killProcess(cmd, grace) + case <-ctxDoneCh: + ctxDoneCh = nil // ditto, so we do not spin once ctx is done + select { + case <-waitCh: + break waitLoop + default: + } + callerDone = true + killProcess(cmd, grace) + } + } + + // The tracked process is reaped, but that guarantees nothing about + // outR/errR/resR/reqW seeing EOF or completing: anything the handler + // spawned — a shelled-out ffmpeg, a stray background job — inherits + // stdout, stderr, and both wire descriptors, and keeps its own copy of + // each open for as long as it runs. killProcess only reaches the + // tracked process (Setpgid extends that to its whole group, but only + // when killProcess actually runs — nothing kills a grandchild left + // behind by a process that exited cleanly on its own, before any + // deadline or cancellation ever fired). drainGrace bounds how long the + // four operations below wait for their own copies to close on their + // own before this forces them closed instead: long enough for a + // handler's own trailing writes to flush, short enough that a leaked + // descendant cannot wedge Run past this point indefinitely. + const drainGrace = 3 * time.Second + + stdioDone := make(chan struct{}) + go func() { + stdioWG.Wait() + close(stdioDone) + }() + + drainTimer := time.NewTimer(drainGrace) + defer drainTimer.Stop() + + var ( + fr frameRead + encodeErr error + stdioPending = true + framePending = true + encodePending = true + timerCh = drainTimer.C + ) + + for stdioPending || framePending || encodePending { + select { + case <-stdioDone: + stdioPending = false + case fr = <-frameCh: + framePending = false + case encodeErr = <-encodeCh: + encodePending = false + case <-timerCh: + // Force every reader/writer still blocked to return, so the + // goroutines above can finish and this loop can exit rather + // than waiting on a descendant process that may never close + // these fds on its own. timerCh is a one-shot channel — it + // cannot fire twice — so this only ever forces the closes + // once, then lets the now-unblocked goroutines deliver + // through the cases above on the next iterations. + closeOutR() + closeErrR() + closeResR() + closeReqW() + timerCh = nil + } + } + + return e.classify(req, fr.frame, fr.err, encodeErr, cmd.ProcessState, timedOut, callerDone), nil +} + +// killProcess best-effort runs the kill ladder (terminate, kill_unix.go) +// against the started process's whole group: SIGTERM to the group, up to +// grace for the whole group — not just the tracked leader — to empty out +// on its own, then SIGKILL to the group if any of it is still there once +// grace elapses. Setpgid (sysProcAttr, procattr_unix.go) is what makes +// "the group" reach anything the tracked process forked, not just the one +// process this package started directly — without it, even the SIGKILL +// half would leave a native library's forked helpers running. +// +// This blocks the waitLoop select for up to grace, which is deliberate: +// the alternative is racing terminate against the very channels that +// triggered it, and there is nothing useful for waitLoop to do with a +// second deadline or cancellation signal while a kill is already in +// flight — see terminate's own doc comment for why grace runs from here, +// not from whatever triggered this call. +func killProcess(cmd *osexec.Cmd, grace time.Duration) { + if cmd.Process == nil { + return + } + + terminate(cmd, grace) +} + +// writeRequest encodes and writes the single request frame Run ever +// sends. It exists mainly so Run's own body does not have to build the +// wire.Frame inline. +func writeRequest(w *os.File, req *exec.Request) error { + if err := wire.Encode(w, &wire.Frame{Kind: wire.KindRequest, Request: req}); err != nil { + return fmt.Errorf("dispatch/exec/subprocess: write request: %w", err) + } + + return nil +} + +// buildEnv constructs the child's environment. It never starts from +// os.Environ(): only a fixed allowlist (PATH, HOME, TMPDIR) is copied from +// the worker's own environment, then the executor's configured base +// (WithEnv), then the request's own Env, which wins any conflict since it +// is the most specific source. The fd and rlimit variables are set last +// and cannot be overridden by any of the above, because their values are +// fixed by how ExtraFiles and options.rlimits were built, not something +// any caller should influence. +func (e *Executor) buildEnv(req *exec.Request) []string { + merged := make(map[string]string, len(e.opts.env)+len(req.Env)+10) + + for _, k := range [...]string{"PATH", "HOME", "TMPDIR"} { + if v, ok := os.LookupEnv(k); ok { + merged[k] = v + } + } + for k, v := range e.opts.env { + merged[k] = v + } + for k, v := range req.Env { + merged[k] = v + } + + merged[shim.EnvRequestFD] = strconv.Itoa(requestFD) + merged[shim.EnvResultFD] = strconv.Itoa(resultFD) + + // RLIMIT_CORE is forced to zero unconditionally, regardless of + // whether WithRlimits was ever called: a segfaulting parser would + // otherwise dump a core containing the malicious input and the + // process's own memory to disk. The rest are only sent when + // configured, and only per field — zero means "leave the limit at + // whatever the worker itself runs with" (see the Rlimits doc + // comment). shim.Main applies all of these child-side via + // syscall.Setrlimit, since Go cannot set a child's rlimits through + // SysProcAttr. + merged[shim.EnvRlimitCore] = "0" + + // RLIMIT_AS is the one dimension job.WithResourceLimits can actually + // reach today: resource.Memory maps onto it unambiguously. The job's + // own ceiling wins when it declares one — a job that asked for a + // tight limit must get it even when the deployment's WithRlimits sets + // a looser one — and the deployment-wide AddressSpace is the fallback + // for the (overwhelming majority of) jobs that declare nothing, so + // nobody who never touches the resource model loses the protection + // WithRlimits already gave them. + // + // resource.CPU has no comparably clean rlimit — RLIMIT_CPU caps total + // CPU *time*, not the instantaneous share a millicore budget + // describes — so it is deliberately left unmapped rather than given + // invented semantics. Every other Rlimits field (NoFile, NProc, + // FSize) has no per-job resource.Set counterpart at all, so those + // stay deployment-wide only, exactly as before. + addressSpace := int64(0) + if e.opts.hasRlimits { + addressSpace = e.opts.rlimits.AddressSpace + } + if v := req.ResourceLimits[resource.Memory]; v > 0 { + addressSpace = v + } + if addressSpace != 0 { + merged[shim.EnvRlimitAS] = strconv.FormatInt(addressSpace, 10) + } + + if e.opts.hasRlimits { + if e.opts.rlimits.NoFile != 0 { + merged[shim.EnvRlimitNoFile] = strconv.FormatInt(e.opts.rlimits.NoFile, 10) + } + if e.opts.rlimits.NProc != 0 { + merged[shim.EnvRlimitNProc] = strconv.FormatInt(e.opts.rlimits.NProc, 10) + } + if e.opts.rlimits.FSize != 0 { + merged[shim.EnvRlimitFSize] = strconv.FormatInt(e.opts.rlimits.FSize, 10) + } + } + if e.opts.strictRlimits { + merged[shim.EnvRlimitStrict] = "1" + } + + out := make([]string, 0, len(merged)) + for k, v := range merged { + out = append(out, k+"="+v) + } + sort.Strings(out) // deterministic, so the same request produces the same argv/env for debugging + + return out +} + +// classify turns what the child reported and what happened to the process +// into a Result. +// +// The rule: the parent trusts the child's frame for what the handler did, +// and its own wait status for what happened to the process. A deadline +// expiry is reported as StatusTimeout regardless of what the frame says — +// timedOut is only ever set by the waitLoop above after it has already +// checked, non-blockingly, that waitCh was not already deliverable, so the +// uniform-random select tie that used to reach classify directly is +// resolved before timedOut is ever set. That check is on the channel, not +// the process: if the child happened to exit microseconds earlier and the +// Wait() goroutine simply had not yet delivered to waitCh, timedOut still +// ends up true and killProcess no-ops on an already-gone process. The +// window this leaves is reap-and-deliver latency, not the width of a +// select's random pick — narrowed, not eliminated. Short of that, a +// decoded frame is authoritative for the status it reports, unless the +// process still died on a signal — a frame claiming StatusOK from a +// process that was, in fact, killed is not to be believed, so the process +// status wins that disagreement. Absent a usable frame entirely — the shim +// never got the chance to report, or its report was cut off mid-write — +// the process's own exit code or signal is all there is to classify by. +func (e *Executor) classify( + req *exec.Request, + frame *wire.Frame, + frameErr error, + encodeErr error, + ps *os.ProcessState, + timedOut, callerCanceled bool, +) *exec.Result { + exitCode, signal, signaled := processOutcome(ps) + + if timedOut { + return &exec.Result{ + Status: exec.StatusTimeout, + HandlerErr: fmt.Sprintf("dispatch/exec/subprocess: deadline %s exceeded", req.Deadline.Format(time.RFC3339)), + ExitCode: exitCode, + Signal: signal, + } + } + + if frameErr == nil && frame != nil && frame.Result != nil { + res := *frame.Result + if signaled { + return &exec.Result{ + Status: exec.StatusKilled, + HandlerErr: fmt.Sprintf( + "dispatch/exec/subprocess: process killed by signal %d after reporting %s", + signal, res.Status, + ), + ExitCode: exitCode, + Signal: signal, + Usage: res.Usage, + // Permanent carries through even though the process was + // signalled: a permanent failure the handler already + // flagged should not silently turn retryable just because + // the kill signal arrived a moment after the report did — + // Result.Err converts this into exec.Error.Permanent, + // which worker/runner.go's retry check reads directly. + // + // Outputs is deliberately NOT carried through, unlike an + // earlier version of this branch claimed ("should not + // become invisible"): it does become invisible regardless + // of what this field holds, because the only commit call + // site (worker/runner.go) gates on Status == StatusOK, and + // prepareOutputDir's own deferred cleanup removes the + // scratch directory those outputs point into before + // anything downstream could read them. That is a + // deliberate choice, not an oversight this comment is + // papering over: a handler killed moments after writing an + // artifact may have left it mid-write, and committing a + // truncated file under the job's real output name is worse + // than committing nothing. Reviving this would need two + // changes made together, not one — populating Outputs here + // again AND widening worker/runner.go's commit gate to + // include StatusKilled — so a future change to either side + // alone does not silently start committing partial output. + Permanent: res.Permanent, + } + } + res.ExitCode = exitCode + res.Signal = 0 + + return &res + } + + reason := "no result frame" + switch { + case callerCanceled: + reason = "context canceled" + case frameErr != nil: + reason = frameErr.Error() + case encodeErr != nil: + reason = encodeErr.Error() + } + + return &exec.Result{ + Status: exec.StatusKilled, + HandlerErr: fmt.Sprintf( + "dispatch/exec/subprocess: process ended without a usable result (%s); exit=%d signal=%d", + reason, exitCode, signal, + ), + ExitCode: exitCode, + Signal: signal, + } +} + +// processOutcome reads the exit code and, on the platforms this rung +// supports, the signal that ended the process. signaled reports whether +// the process died on a signal rather than exiting on its own; when it is +// true, exitCode is meaningless and signal is what actually happened. +func processOutcome(ps *os.ProcessState) (exitCode, signal int, signaled bool) { + if ps == nil { + return -1, 0, false + } + + ws, ok := ps.Sys().(syscall.WaitStatus) + if !ok { + return ps.ExitCode(), 0, false + } + if ws.Signaled() { + return -1, int(ws.Signal()), true + } + + return ws.ExitStatus(), 0, false +} + +// Reclaim is a no-op for this rung. That is no longer because nothing can +// leak: the child used to share the worker's own process group, which +// would let a signal aimed at the worker's whole group reach it too, but +// now that it is its own group leader (see sysProcAttr in +// procattr_unix.go, Setpgid), a worker that dies mid-attempt can leave an +// orphaned child running with nothing left to signal it. Actually sweeping +// those up is deferred — it needs a worker identity stable across +// restarts, which this rung does not have today; see the SDD ledger's +// Phase 3 note on the same gap for the stronger rungs. +func (e *Executor) Reclaim(context.Context, id.WorkerID) error { return nil } + +// Close releases the executor's own resources. Subprocess holds none — +// every pipe and process it creates is scoped to a single Run. +func (e *Executor) Close() error { return nil } diff --git a/exec/subprocess/executor_test.go b/exec/subprocess/executor_test.go new file mode 100644 index 0000000..9577c20 --- /dev/null +++ b/exec/subprocess/executor_test.go @@ -0,0 +1,259 @@ +package subprocess_test + +import ( + "context" + "encoding/json" + "os" + "strings" + "testing" + "time" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/exectest" + "github.com/xraph/dispatch/exec/subprocess" + "github.com/xraph/dispatch/id" +) + +func newExecutor(t *testing.T) *subprocess.Executor { + t.Helper() + + return subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; these tests are not about the uid boundary + ) +} + +func request(t *testing.T, name string, payload any) *exec.Request { + t.Helper() + raw, err := json.Marshal(payload) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + return &exec.Request{ + JobID: id.NewJobID(), + Name: name, + Payload: raw, + OutputDir: t.TempDir(), + Policy: exec.NewPolicy(exec.GracePeriod(time.Second)), + } +} + +func TestIdentity(t *testing.T) { + e := newExecutor(t) + if e.Name() != "subprocess" { + t.Errorf("Name() = %q, want %q", e.Name(), "subprocess") + } + if e.Level() != exec.LevelProcess { + t.Errorf("Level() = %v, want %v", e.Level(), exec.LevelProcess) + } +} + +func TestRunSuccess(t *testing.T) { + res, err := newExecutor(t).Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusOK { + t.Fatalf("Status = %q, want ok (err %q)", res.Status, res.HandlerErr) + } +} + +func TestRunHandlerErrorIsExitZero(t *testing.T) { + res, err := newExecutor(t).Run(context.Background(), request(t, exectest.JobError, struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusHandlerError { + t.Fatalf("Status = %q, want handler_error", res.Status) + } + if res.ExitCode != 0 { + t.Errorf("ExitCode = %d, want 0 — a handler saying no is not the shim failing", res.ExitCode) + } +} + +func TestRunPanicIsKilled(t *testing.T) { + // The child dies; the parent must report it as killed rather than + // letting the panic reach the worker, which is the whole point. + res, err := newExecutor(t).Run(context.Background(), request(t, exectest.JobPanic, struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusKilled && res.Status != exec.StatusHandlerError { + t.Errorf("Status = %q, want killed or handler_error", res.Status) + } + if res.Status == exec.StatusOK { + t.Error("a panicking handler must not report success") + } +} + +func TestRunDeadlineKillsAHandlerThatIgnoresCancellation(t *testing.T) { + req := request(t, exectest.JobSlow, exectest.SlowPayload{SleepMillis: 30000, IgnoreCtx: true}) + req.Deadline = time.Now().Add(500 * time.Millisecond) + + start := time.Now() + res, err := newExecutor(t).Run(context.Background(), req) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusTimeout { + t.Errorf("Status = %q, want timeout", res.Status) + } + // This is the assertion the whole phase exists for: the handler asked + // to sleep 30s and ignores cancellation, so anything near that means + // the deadline was still advisory. + if elapsed > 10*time.Second { + t.Errorf("Run() took %v; the deadline was not enforced", elapsed) + } +} + +func TestRunUnknownHandlerIsLaunchFailure(t *testing.T) { + res, err := newExecutor(t).Run(context.Background(), request(t, "subprocess.absent", struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusLaunchFailed { + t.Fatalf("Status = %q, want launch_failed", res.Status) + } + if res.Status.CountsAgainstRetries() { + t.Error("an unknown handler must not consume the retry budget") + } +} + +func TestRunMissingBinaryIsLaunchFailure(t *testing.T) { + e := subprocess.New(subprocess.WithBinary("/nonexistent/dispatch-worker"), subprocess.WithAllowSameUser()) + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + + // Either shape is acceptable, but it must be classified as a launch + // failure and must not consume a retry. + switch { + case err != nil: + // A raw error from Run is treated as a launch failure by the Runner. + case res.Status != exec.StatusLaunchFailed: + t.Fatalf("Status = %q, want launch_failed", res.Status) + } +} + +// runBounded runs e.Run in the background and fails the test fast if it +// does not return within bound, rather than letting a regression hang +// until go test's own multi-minute panic timeout does the job for it. +// Returning within a bound despite something misbehaving downstream is +// exactly the property the C1/C2 regression tests exist to check, so they +// assert on it directly instead of only on the returned Status. +func runBounded( + ctx context.Context, t *testing.T, e *subprocess.Executor, req *exec.Request, bound time.Duration, +) (*exec.Result, time.Duration) { + t.Helper() + + type outcome struct { + res *exec.Result + err error + } + done := make(chan outcome, 1) + start := time.Now() + go func() { + res, err := e.Run(ctx, req) + done <- outcome{res, err} + }() + + select { + case o := <-done: + if o.err != nil { + t.Fatalf("Run() = %v", o.err) + } + + return o.res, time.Since(start) + case <-time.After(bound): + t.Fatalf("Run() did not return within %v", bound) + + return nil, 0 // unreachable; Fatalf stops the goroutine + } +} + +// TestRunGrandchildCannotWedgeTheDrain reproduces C1: a handler's own +// subprocess (a shelled-out ffmpeg, a stray background job) can exit +// cleanly and leave a grandchild running behind it. There is no deadline +// and no cancellation here, so killProcess never runs — the tracked +// process was never killed, it just exited on its own — which is exactly +// the case where nothing upstream would otherwise stop that grandchild +// from holding stdout/stderr open indefinitely. Only the post-wait drain +// grace stands between that and Run hanging for as long as the +// grandchild runs. +func TestRunGrandchildCannotWedgeTheDrain(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{envLeakChild: "1"}), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; this test is not about the uid boundary + ) + + res, elapsed := runBounded(context.Background(), t, e, request(t, exectest.JobOK, struct{}{}), 8*time.Second) + + if res.Status != exec.StatusOK { + t.Errorf("Status = %q, want ok (err %q)", res.Status, res.HandlerErr) + } + // fixtureSleep is 30s; anything well under that proves Run did not + // wait on the grandchild to finish on its own. + if elapsed > 6*time.Second { + t.Errorf("Run() took %v; a leaked grandchild wedged the post-wait drain", elapsed) + } +} + +// TestRunRequestWriteIsInterruptedByDeadline reproduces C2: a request +// large enough to fill the pipe buffer (64KB on Linux, often 16KB on +// macOS), written to a child that never reads fd 3 at all. Without the +// fix, that write blocks in Run's own goroutine before the deadline timer +// even exists, so nothing can interrupt it until the child eventually +// exits on its own — here, only after fixtureSleep, far past the +// deadline this test sets. +func TestRunRequestWriteIsInterruptedByDeadline(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{envSleepOnly: "1"}), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; this test is not about the uid boundary + ) + + req := request(t, exectest.JobOK, struct{ Value string }{Value: strings.Repeat("x", 1<<20)}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + + res, elapsed := runBounded(context.Background(), t, e, req, 8*time.Second) + + if res.Status != exec.StatusTimeout && res.Status != exec.StatusKilled { + t.Errorf("Status = %q, want timeout or killed", res.Status) + } + if elapsed > 6*time.Second { + t.Errorf("Run() took %v; the deadline did not reach a request write blocked outside the select loop", elapsed) + } +} + +// TestRunContextCancellationKillsChild is not one of the brief's listed +// cases, but it proves a constraint the brief states in prose: cancelling +// the caller's context must actually stop the child, not just make Run +// return while the process keeps running. IgnoreCtx makes the handler deaf +// to its own context, so the only way this finishes quickly is the parent +// killing the OS process from the outside. +func TestRunContextCancellationKillsChild(t *testing.T) { + req := request(t, exectest.JobSlow, exectest.SlowPayload{SleepMillis: 30000, IgnoreCtx: true}) + + ctx, cancel := context.WithCancel(context.Background()) + start := time.Now() + go func() { + time.Sleep(200 * time.Millisecond) + cancel() + }() + + res, err := newExecutor(t).Run(ctx, req) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status == exec.StatusOK { + t.Error("a cancelled attempt must not report success") + } + if elapsed > 10*time.Second { + t.Errorf("Run() took %v; context cancellation did not stop the child", elapsed) + } +} diff --git a/exec/subprocess/internal_test.go b/exec/subprocess/internal_test.go new file mode 100644 index 0000000..2db2aa9 --- /dev/null +++ b/exec/subprocess/internal_test.go @@ -0,0 +1,204 @@ +package subprocess + +// This file is package subprocess (internal), not subprocess_test, +// deliberately breaking the external-tests-only convention the rest of +// this package follows — precedent already set by +// exec/shim/internal_test.go for the same reason: classify is unexported, +// and the regression this covers cannot be forced from outside the +// package. The race is in Go's own select scheduling (which of two +// simultaneously-ready channels it picks), not in anything a black-box +// caller of Run can control or observe reliably; a black-box repro of it +// against the pre-fix code took a 20-iteration loop to reproduce (first +// failing on the ninth run), so a single-shot black-box version would +// pass most of the time and be worse than no test at all. Calling +// classify directly with synthetic inputs pins its contract down instead. + +import ( + "slices" + "strings" + "testing" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/shim" + "github.com/xraph/dispatch/exec/wire" + "github.com/xraph/dispatch/resource" +) + +func TestClassifyTimedOutOverridesADecodedFrame(t *testing.T) { + e := &Executor{} + frame := &wire.Frame{ + Kind: wire.KindResult, + Result: &exec.Result{Status: exec.StatusOK}, + } + + tests := []struct { + name string + timedOut bool + want exec.Status + }{ + { + // The regression itself: round 1 had classify trust an + // unsignalled decoded frame over timedOut, which made + // StatusTimeout unreachable for any cooperative handler — + // the shim already traps SIGTERM and writes a Result frame + // on its way out, and Task 6 adds the parent's SIGTERM half + // of the kill ladder, so "the tracked process finishes + // right as the deadline fires" is a live shape, not a + // theoretical one. waitLoop now only ever sets timedOut + // after confirming, non-blockingly, that waitCh had not + // already delivered — see its comment in executor.go — so + // by the time classify sees timedOut, it must win + // regardless of what the frame says. + name: "timed out overrides a clean frame", + timedOut: true, + want: exec.StatusTimeout, + }, + { + // Contrast case: without timedOut, the same frame is + // trusted, proving the table above is exercising the + // branch it claims to, not returning a constant. + name: "no timeout trusts the frame", + timedOut: false, + want: exec.StatusOK, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // ps == nil is already tolerated by processOutcome (exit + // code -1, signal 0, signalled false), so no synthesised + // *os.ProcessState is needed to reach either branch here. + res := e.classify(&exec.Request{}, frame, nil, nil, nil, tt.timedOut, false) + if res.Status != tt.want { + t.Errorf("Status = %q, want %q", res.Status, tt.want) + } + }) + } +} + +// TestBuildEnvCarriesRlimits pins buildEnv's construction of the +// DISPATCH_RLIMIT_* variables directly, without spawning a real child — +// the child-side enforcement itself is covered end to end by +// TestRlimitsAreAppliedChildSide (limits_unix_test.go), but that test can +// only observe the *effect* of a limit landing, not which ones buildEnv +// actually decided to send. This is the regression net for the two rules +// that effect can't distinguish: RLIMIT_CORE always going out as "0" +// regardless of what's configured, and every other field being omitted +// entirely rather than sent as "0" when left at its zero value. +func TestBuildEnvCarriesRlimits(t *testing.T) { + tests := []struct { + name string + hasRlimits bool + rlimits Rlimits + requestLimits resource.Set + wantHas []string + wantAbsent []string + }{ + { + name: "no WithRlimits call still forces core to zero", + hasRlimits: false, + wantHas: []string{shim.EnvRlimitCore + "=0"}, + wantAbsent: []string{shim.EnvRlimitAS, shim.EnvRlimitNoFile, shim.EnvRlimitNProc, shim.EnvRlimitFSize}, + }, + { + name: "configured fields are sent, zero fields are omitted", + hasRlimits: true, + rlimits: Rlimits{AddressSpace: 1 << 30, NoFile: 64, Core: 999}, + wantHas: []string{ + shim.EnvRlimitCore + "=0", // Core forced to zero even though 999 was configured + shim.EnvRlimitAS + "=1073741824", + shim.EnvRlimitNoFile + "=64", + }, + wantAbsent: []string{shim.EnvRlimitNProc, shim.EnvRlimitFSize}, // left at zero, so omitted + }, + { + // job.WithResourceLimits(resource.MemoryBytes(...)) resolved at + // enqueue and carried across on Request.ResourceLimits. Per-job + // must win over the deployment-wide WithRlimits default. + name: "a job's own memory limit wins over the deployment-wide default", + hasRlimits: true, + rlimits: Rlimits{AddressSpace: 1 << 30}, + requestLimits: resource.Set{resource.Memory: 2 << 30}, + wantHas: []string{shim.EnvRlimitAS + "=2147483648"}, + }, + { + // The overwhelming majority of jobs declare nothing: the + // deployment-wide default must still apply exactly as before. + name: "deployment-wide default still applies when the job declares nothing", + hasRlimits: true, + rlimits: Rlimits{AddressSpace: 1 << 30}, + requestLimits: nil, + wantHas: []string{shim.EnvRlimitAS + "=1073741824"}, + }, + { + // A job's own limit must apply even with no deployment-wide + // WithRlimits call at all — the per-job ceiling is not merely + // an override of a configured default, it is enforced on its + // own. + name: "a job's own memory limit applies with no deployment-wide default configured", + hasRlimits: false, + requestLimits: resource.Set{resource.Memory: 512 << 20}, + wantHas: []string{shim.EnvRlimitAS + "=536870912"}, + }, + { + // resource.CPU has no clean rlimit mapping (see buildEnv's + // comment) and must never leak into RLIMIT_AS. + name: "a CPU limit alone never produces RLIMIT_AS", + hasRlimits: false, + requestLimits: resource.Set{resource.CPU: 2000}, + wantAbsent: []string{shim.EnvRlimitAS}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + e := &Executor{opts: options{rlimits: tt.rlimits, hasRlimits: tt.hasRlimits}} + env := e.buildEnv(&exec.Request{ResourceLimits: tt.requestLimits}) + + for _, want := range tt.wantHas { + if !slices.Contains(env, want) { + t.Errorf("buildEnv() = %v, want to contain %q", env, want) + } + } + for _, prefix := range tt.wantAbsent { + for _, kv := range env { + if strings.HasPrefix(kv, prefix+"=") { + t.Errorf("buildEnv() contains %q, want %s absent", kv, prefix) + } + } + } + }) + } +} + +// TestBuildEnvDoesNotInheritTheParentEnvironment proves buildEnv's fixed +// PATH/HOME/TMPDIR allowlist (see its own doc comment, and spec §6's "the +// child's environment is constructed, not inherited") actually excludes +// everything else in the worker's own environment, not just everything +// this test happens to think of. A sentinel set in this test process's own +// environment stands in for a credential the worker might carry — a DSN, +// an API key — and must not reach the child's constructed environment. +// +// Nothing else in this package's suite would catch a regression here: the +// rlimit tests above never configure WithEnv and don't need to (rlimit +// vars are unconditional), TestRunSuccess and its neighbors do configure +// WithEnv but only ever assert on the attempt's Status, never on what +// buildEnv actually sent, and a child that received the entire parent +// os.Environ() would still run those fixtures to completion successfully. +// A review round demonstrated exactly that: replacing the allowlist with a +// full os.Environ() copy passed the whole phase test suite clean. +func TestBuildEnvDoesNotInheritTheParentEnvironment(t *testing.T) { + const sentinel = "DISPATCH_TEST_PARENT_ONLY_SENTINEL" + t.Setenv(sentinel, "leaked-credential-value") + + e := &Executor{} + env := e.buildEnv(&exec.Request{}) + + for _, kv := range env { + if strings.HasPrefix(kv, sentinel+"=") { + t.Fatalf("buildEnv() = %v, contains %q — the child's environment must be "+ + "constructed from an allowlist, not inherited from the worker's own os.Environ()", + env, kv) + } + } +} diff --git a/exec/subprocess/kill_other.go b/exec/subprocess/kill_other.go new file mode 100644 index 0000000..1c5c142 --- /dev/null +++ b/exec/subprocess/kill_other.go @@ -0,0 +1,23 @@ +//go:build !unix + +package subprocess + +import ( + osexec "os/exec" + "time" +) + +// terminate falls back to an immediate, non-graceful kill outside Unix, +// since there is neither a process group nor a SIGTERM to build a grace +// period ladder out of on this build. grace is accepted only for +// signature symmetry with the Unix build's terminate (kill_unix.go) and +// is otherwise unused. This path is never reached in practice — +// checkLaunch (limits_other.go) refuses to start this rung at all outside +// Unix — so it exists only to keep the package compiling here. +func terminate(cmd *osexec.Cmd, _ time.Duration) { + if cmd.Process == nil { + return + } + + _ = killGroup(cmd, 0) //nolint:errcheck // best-effort; unreachable in practice, see doc comment above +} diff --git a/exec/subprocess/kill_unix.go b/exec/subprocess/kill_unix.go new file mode 100644 index 0000000..be2431e --- /dev/null +++ b/exec/subprocess/kill_unix.go @@ -0,0 +1,158 @@ +//go:build unix + +package subprocess + +import ( + "errors" + osexec "os/exec" + "syscall" + "time" +) + +// pollInterval is how often waitGroupEmpty re-probes the process group +// during the grace period. Short enough that the common case — the group +// empties out promptly once it has been asked to — is detected quickly +// rather than riding out the whole grace period; long enough not to +// matter as CPU overhead for the case where something in the group +// ignores SIGTERM and this polls for the full duration instead. +const pollInterval = 10 * time.Millisecond + +// terminate runs the kill ladder: SIGTERM to the child's whole process +// group, then up to grace for the group to empty out on its own, +// escalating to SIGKILL — again to the whole group — only if it has not +// by the time grace elapses. +// +// Signalling the negative pid (killGroup, procattr_unix.go) is what +// reaches a native library's forked helper as well as the tracked process +// itself, and that matters differently for each half of the ladder. A +// helper that does not trap SIGTERM simply ignores the first signal and +// rides out the grace period, the same as it would have under the old +// direct-SIGKILL behaviour minus the wait; it dies on the SIGKILL that +// follows regardless, because SIGKILL cannot be trapped or ignored by +// anything. Sending SIGTERM to the group rather than just the tracked +// process is what gives a *cooperative* helper — one that does trap +// SIGTERM, unlike the shim's own handler process — the same chance to +// shut down cleanly that the tracked process gets; addressing only the +// leader would leave such a helper to be killed outright a grace period +// later instead, for no reason beyond which process happened to fork it. +// +// grace is measured from here, when the SIGTERM is actually sent, not +// from whatever triggered this call — a deadline or a cancelled caller +// context. That makes Policy.GracePeriod additive on top of the deadline +// rather than carved out of it: an operator who configures a six-hour +// deadline and a 30-second grace period gets six hours of run time +// followed by up to 30 more seconds for a cooperative shutdown, not a +// grace period that eats into the six hours and races whatever caused the +// deadline to fire. Carving it out instead would shrink the handler's +// effective budget by an amount that has nothing to do with the +// handler's own behaviour, purely as an artifact of how the ladder +// happens to be implemented — the opposite of what "give it 30 seconds to +// shut down after six hours" is asking for. +// +// The escalation decision below is made about the *group*, not the +// leader — this is the fix for a real bug an earlier version of this +// function had. Setpgid without an explicit Pgid makes the tracked +// process its own group leader, so its pid also names the group; pgid is +// captured once here, up front, because — unlike the leader's own pid, +// which killGroup's probe treats as unsafe to reuse the instant the +// leader is reaped — a pgid stays valid, and safe to keep addressing, +// for as long as *any* member of the group remains alive. The production +// shape this rung exists for is a handler that forks a native helper +// (an OpenCASCADE-style worker, say) and then itself exits cleanly on +// SIGTERM: the leader is gone almost immediately, well inside grace, +// while the helper it left behind is not. A version of this function +// that asked only "has the leader exited" would read that as "done, +// nothing left to escalate to SIGKILL" — and it did, silently leaving +// the helper running for as long as it liked. Asking whether the *group* +// is empty instead keeps waiting exactly as long as anything is still in +// it, leader or not, which is what makes the final SIGKILL below +// actually fire for this case rather than being skipped as though there +// were nothing left to reach. +func terminate(cmd *osexec.Cmd, grace time.Duration) { + if cmd.Process == nil { + return + } + + // Setpgid without an explicit Pgid (sysProcAttr, procattr_unix.go) + // makes the tracked process its own group leader, so its pid is also + // the group's pgid at the moment this call begins — captured once, + // before either signal, so the rest of this function keeps addressing + // the same group even once the leader itself has been reaped. + pgid := cmd.Process.Pid + + // Errors from both signal sends are deliberately discarded, for the + // same reason killGroup's own doc comment gives for falling through + // on a non-ErrProcessDone probe error: this function's caller + // (killProcess) has no way to make Run itself fail differently based + // on whether a signal landed, and a kill that was attempted and + // failed is still strictly better-off than one skipped entirely — + // classify has the process's actual wait status to fall back on + // either way. + _ = killGroup(cmd, syscall.SIGTERM) //nolint:errcheck // best-effort; see comment above + + if waitGroupEmpty(pgid, grace) { + // The common case: every process in the group — the tracked + // leader and anything cooperative it forked — honoured SIGTERM + // and exited before grace ran out, so there is nothing left to + // escalate to SIGKILL. + return + } + + // Signalling raw here, rather than through killGroup, is deliberate + // and is the other half of this function's fix: killGroup's own probe + // is keyed to cmd.Process specifically and skips the send outright + // once that process is reaped, which is exactly wrong for this call — + // the tracked leader having already exited is the expected shape of + // the bug this escalation exists to catch, not a reason to skip it. + // waitGroupEmpty having just reported the group as non-empty stands in + // for that probe instead: the kernel does not hand a pgid back out + // for reuse while any process is still using it as its group, so a + // non-ESRCH result from that same kind of check moments ago means + // pgid was, at that moment, still this attempt's own group and not a + // number the kernel had already recycled. That does not close the + // window entirely — a last member could exit in the interval between + // that check and this send, freeing the pgid for reuse before the + // signal lands — the same kind of gap killGroup's own doc comment + // already accepts for the single-process case, narrowed here to the + // width of one syscall rather than removed. + _ = syscall.Kill(-pgid, syscall.SIGKILL) //nolint:errcheck // best-effort; see comment above +} + +// waitGroupEmpty reports whether every process in the group named by +// pgid has exited by the time grace elapses. +// +// It probes the group directly — syscall.Kill(-pgid, 0) — rather than +// asking only whether the tracked leader is still alive. That distinction +// is the point: syscall.Kill with a negative pid succeeds as long as the +// caller has permission to signal at least one member of the group, and +// only fails with ESRCH once none are left, so this notices a helper the +// leader forked and left running exactly the same way it would notice an +// uncooperative leader — checking the leader alone would report the +// group "empty" the moment a *cooperative* leader exits, even while +// something it forked is still very much alive. +// +// This cannot use cmd.Wait() to find out when the leader specifically has +// gone: Run's own dedicated goroutine already owns the one call to Wait() +// this *osexec.Cmd will ever get, and os/exec panics if Wait is invoked a +// second time. Polling the group directly sidesteps that entirely, since +// it never touches cmd.Process at all. +func waitGroupEmpty(pgid int, grace time.Duration) bool { + deadline := time.Now().Add(grace) + + for { + if err := syscall.Kill(-pgid, 0); errors.Is(err, syscall.ESRCH) { + return true + } + + remaining := time.Until(deadline) + if remaining <= 0 { + return false + } + + if remaining < pollInterval { + time.Sleep(remaining) + } else { + time.Sleep(pollInterval) + } + } +} diff --git a/exec/subprocess/kill_unix_test.go b/exec/subprocess/kill_unix_test.go new file mode 100644 index 0000000..edda3ea --- /dev/null +++ b/exec/subprocess/kill_unix_test.go @@ -0,0 +1,306 @@ +//go:build unix + +package subprocess_test + +import ( + "context" + "errors" + "os" + "path/filepath" + "strconv" + "strings" + "syscall" + "testing" + "time" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/exectest" + "github.com/xraph/dispatch/exec/subprocess" +) + +// TestKillLadderReachesAHandlerIgnoringSIGTERM proves the ladder still +// bounds Run's return even when the handler ignores the SIGTERM half of +// it entirely: IgnoreCtx makes JobSlow deaf to context cancellation (and +// so, transitively, to the shim's own SIGTERM trap, which only cancels +// that context), so the only thing that can end this attempt within any +// reasonable bound is the SIGKILL that follows the grace period. +func TestKillLadderReachesAHandlerIgnoringSIGTERM(t *testing.T) { + req := request(t, exectest.JobSlow, exectest.SlowPayload{SleepMillis: 60000, IgnoreCtx: true}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + req.Policy = exec.NewPolicy(exec.GracePeriod(300 * time.Millisecond)) + + start := time.Now() + res, err := newExecutor(t).Run(context.Background(), req) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusTimeout { + t.Errorf("Status = %q, want timeout", res.Status) + } + if elapsed > 5*time.Second { + t.Errorf("Run() took %v; SIGKILL did not follow the grace period", elapsed) + } +} + +// TestKillLadderKillsTheWholeProcessGroup is the assertion the brief asks +// for by name: that the process group is actually gone once Run returns, +// not merely the one process this package tracks directly. The +// envGroupKill fixture (main_test.go) ignores SIGTERM outright and forks +// a grandchild that does nothing but sleep, so the only way both ever die +// is the ladder's SIGKILL half reaching the whole group. +// +// A missing Setpgid would not fail silently on just the grandchild: with +// no Setpgid the child stays a member of the worker's own process group +// rather than becoming its own group leader, so terminate's pgid — the +// child's pid — names no group at all. killGroup's SIGTERM send and the +// escalation's raw syscall.Kill(-pgid, SIGKILL) both then return ESRCH, +// waitGroupEmpty's very first probe reports the group already "empty", +// and terminate returns without ever escalating. Neither the leader nor +// the grandchild is signalled by this package at all in that case — Run +// still returns once the deadline's own bookkeeping decides the attempt, +// not because anything here reaped either process. +func TestKillLadderKillsTheWholeProcessGroup(t *testing.T) { + req := request(t, exectest.JobOK, struct{}{}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + req.Policy = exec.NewPolicy(exec.GracePeriod(300 * time.Millisecond)) + + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{envGroupKill: "1"}), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; this test is not about the uid boundary + ) + + start := time.Now() + _, err := e.Run(context.Background(), req) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Run() = %v", err) + } + if elapsed > 5*time.Second { + t.Errorf("Run() took %v; SIGKILL did not follow the grace period", elapsed) + } + + pidPath := filepath.Join(req.OutputDir, "grandchild.pid") + raw, rerr := os.ReadFile(pidPath) + if rerr != nil { + t.Fatalf("read grandchild pid file: %v", rerr) + } + + pid, perr := strconv.Atoi(strings.TrimSpace(string(raw))) + if perr != nil { + t.Fatalf("parse grandchild pid %q: %v", raw, perr) + } + + // A signal-0 probe reports ESRCH once the kernel has no process left + // at that pid to deliver to. This is the assertion that catches a + // missing Setpgid or a kill sent to the wrong target: without the + // process-group signal actually reaching the grandchild, this pid + // would still be alive here, sleeping out longSleep on its own — + // which is deliberately much longer than this test's own bound, so + // that surviving would show up as "still alive," not as "happened to + // exit on its own around the same time," see envLongSleep's doc + // comment (main_test.go). + if kerr := syscall.Kill(pid, 0); kerr != syscall.ESRCH { + t.Errorf("syscall.Kill(%d, 0) = %v, want ESRCH — grandchild pid %d is still alive", pid, kerr, pid) + } +} + +// TestKillLadderReapsAHelperAfterACooperativeLeaderExits is the C1 +// regression test: it pins the bug where terminate decided whether to +// escalate to SIGKILL by asking only whether the tracked *leader* had +// exited. The envLeaderExitsHelperSurvives fixture's leader exits almost +// immediately on SIGTERM (its default disposition, since this fixture +// installs no handler for it) while the helper it forks first ignores +// SIGTERM outright — so a leader-only liveness check reports "done" +// within a poll interval of the leader dying, milliseconds after SIGTERM +// is sent, and never reaches the SIGKILL that the surviving helper +// actually needs. Measured against the pre-fix implementation: terminate +// returned in ~12ms and the helper was still alive afterwards +// (syscall.Kill(pid, 0) == nil). The fix (waitGroupEmpty, kill_unix.go) +// asks whether the whole *group* is empty instead, which keeps waiting +// through the helper's presence and does send the group SIGKILL once +// grace elapses. +func TestKillLadderReapsAHelperAfterACooperativeLeaderExits(t *testing.T) { + req := request(t, exectest.JobOK, struct{}{}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + req.Policy = exec.NewPolicy(exec.GracePeriod(time.Second)) + + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{envLeaderExitsHelperSurvives: "1"}), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; this test is not about the uid boundary + ) + + start := time.Now() + _, err := e.Run(context.Background(), req) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Run() = %v", err) + } + // The leader itself dies within a poll interval or two of SIGTERM + // landing, well under 300ms after the deadline fires. Naively, a + // regression back to leader-only liveness ought to make Run() return + // that quickly too — deadline (300ms) plus a handful of milliseconds — + // since it would treat the leader's own exit as "nothing left to wait + // for" and never reach the group SIGKILL. In practice this lower bound + // does NOT reliably catch that regression, and crediting it as the + // thing that does (an earlier version of this comment did) is false: + // Run() also waits, separately, on drainGrace (executor.go) for + // stdout/stderr to close, and this fixture's surviving helper keeps + // those descriptors open regardless of which way terminate's own + // liveness check goes. Mutating waitGroupEmpty back to leader-only — + // reintroducing the exact bug this test exists to catch — measured + // elapsed at 3.31s here, held entirely by that unrelated 3s drainGrace + // floor, comfortably clearing 1100ms despite the regression. The + // assertion that actually catches it is the ESRCH check below: a + // leader-only liveness check never sends the helper SIGKILL, so it is + // still alive when Run() returns, which only that check observes. This + // bound is kept anyway as a sanity check that grace is actually + // awaited rather than skipped outright (see the upper bound just + // below it) — it is just not this bug's regression net. + if elapsed < 1100*time.Millisecond { + t.Errorf("Run() took %v; returned before the group SIGKILL had a chance to run", elapsed) + } + if elapsed > 8*time.Second { + t.Errorf("Run() took %v; grace was not bounded", elapsed) + } + + pidPath := filepath.Join(req.OutputDir, "helper.pid") + raw, rerr := os.ReadFile(pidPath) + if rerr != nil { + t.Fatalf("read helper pid file: %v", rerr) + } + + pid, perr := strconv.Atoi(strings.TrimSpace(string(raw))) + if perr != nil { + t.Fatalf("parse helper pid %q: %v", raw, perr) + } + + // The helper is a grandchild from this test process's point of view — + // forked by the leader, not by us — so once the leader is gone it is + // reparented and reaped by whatever subreaper the OS hands it to, not + // necessarily promptly relative to Run() returning. syscall.Kill(pid, + // 0) reports success, not ESRCH, for a zombie that has already died + // to SIGKILL but not yet been reaped: the process table entry is + // still there. A single check right after Run() returns raced that + // reaping under parallel-package load and flaked once in CI despite + // the helper genuinely having been killed. Polling for up to a + // couple of seconds waits out the reap without weakening what this + // asserts — the helper still has to actually be gone, just not + // instantaneously. + deadline := time.Now().Add(2 * time.Second) + var kerr error + for { + kerr = syscall.Kill(pid, 0) + if errors.Is(kerr, syscall.ESRCH) || time.Now().After(deadline) { + break + } + time.Sleep(10 * time.Millisecond) + } + if !errors.Is(kerr, syscall.ESRCH) { + t.Errorf("syscall.Kill(%d, 0) = %v, want ESRCH — a leader that exits on SIGTERM must not let its own uncooperative helper survive", pid, kerr) + } +} + +// TestKillLadderSendsSIGTERMBeforeGraceElapses is the C3 regression test: +// TestKillLadderReachesAHandlerIgnoringSIGTERM and +// TestKillLadderKillsTheWholeProcessGroup both use fixtures that ignore +// SIGTERM, so both pass identically whether or not the ladder's SIGTERM +// rung runs at all — a build that skipped straight to SIGKILL after grace +// would pass them too. The envSigtermMarker fixture closes that gap: it +// carries no timer of its own and can only end by receiving and handling +// an actual SIGTERM, so the marker file existing, and Run() returning +// well inside the grace period rather than only after it, is evidence +// that could not be produced any other way. +func TestKillLadderSendsSIGTERMBeforeGraceElapses(t *testing.T) { + req := request(t, exectest.JobOK, struct{}{}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + req.Policy = exec.NewPolicy(exec.GracePeriod(5 * time.Second)) + + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{envSigtermMarker: "1"}), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; this test is not about the uid boundary + ) + + start := time.Now() + res, err := e.Run(context.Background(), req) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("Run() = %v", err) + } + // The generous 5s grace period is deliberately not what bounds this: + // if SIGTERM is actually sent and handled, this returns well under + // it, not anywhere near it. A build that never sent SIGTERM and + // instead waited for SIGKILL after the full grace period would take + // close to 5.3s here — well past this bound — since the fixture, + // having received no SIGTERM, would still be blocked on <-sigCh right + // up until SIGKILL ends it. The margin below grace is kept wide + // (rather than a tight multiple of the 300ms deadline) so a cold, + // possibly race-detector-instrumented child's own startup latency + // cannot make this test flaky. + if elapsed > 4*time.Second { + t.Errorf("Run() took %v; wanted well under the 5s grace period, suggesting SIGTERM was never sent", elapsed) + } + // A process ended by SIGKILL reports Signaled() == true; this fixture + // only ever exits cleanly (os.Exit(0) after handling SIGTERM), so a + // signalled result here would itself mean the marker-writing path was + // never reached and SIGKILL ended it directly instead. + if res.Signal != 0 { + t.Errorf("Signal = %d, want 0 — the fixture exits cleanly after handling SIGTERM, it does not die by SIGKILL", res.Signal) + } + + markerPath := filepath.Join(req.OutputDir, "sigterm-received") + if _, serr := os.Stat(markerPath); serr != nil { + t.Errorf("sigterm-received marker missing (%v) — the fixture only writes it after actually receiving SIGTERM", serr) + } +} + +// TestKillLadderClassifiesACooperativeTimeoutCorrectly is the C2 +// regression test, package-local to exec/subprocess. Every OTHER timeout +// test in this file uses a handler that ignores cancellation (IgnoreCtx: +// true); this one does not, so within this file it is the only test that +// drives classify's timedOut-overrides-the-frame rule (executor.go) +// through a real process instead of the synthetic inputs +// internal_test.go uses. It is not the only such test in the module, +// though: exectest's conformance suite (exec/exectest/suite.go) runs the +// same shape twice more against this rung — testDeadlineEnforcedCooperative +// and testDeadlineEnforcedSwallowedCancellation — as part of +// TestSubprocessConformance. With IgnoreCtx: false, JobSlow honours +// ctx.Done() and returns promptly, the shim writes a Result frame and +// exits 0 — frameOK && !signaled — while the parent's own deadline still +// independently fires and sets timedOut, which classify must let win +// regardless of what the frame says. +func TestKillLadderClassifiesACooperativeTimeoutCorrectly(t *testing.T) { + req := request(t, exectest.JobSlow, exectest.SlowPayload{SleepMillis: 60000, IgnoreCtx: false}) + req.Deadline = time.Now().Add(300 * time.Millisecond) + // A generous grace period, not a tight one: this test's whole point is + // that the handler wins the cooperative race and exits cleanly before + // any SIGKILL, so it needs real headroom for a cold, possibly + // race-detector-instrumented child to start up, decode the request, + // notice cancellation, and write its Result frame — a tight grace + // here does not make the test stricter, it just makes it flaky by + // occasionally forcing the SIGKILL half of the ladder to win instead, + // which asserts nothing about the timedOut-overrides-the-frame rule + // this test exists to pin. + req.Policy = exec.NewPolicy(exec.GracePeriod(5 * time.Second)) + + res, err := newExecutor(t).Run(context.Background(), req) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusTimeout { + t.Errorf("Status = %q, want timeout — timedOut must override a clean frame even when the handler cooperates", res.Status) + } + if res.Signal != 0 { + t.Errorf("Signal = %d, want 0 — a cooperative handler exits via the shim's own Result frame, not a signal", res.Signal) + } + if res.ExitCode != 0 { + t.Errorf("ExitCode = %d, want 0 — the shim exits 0 on a handler error, which ctx.Err() is", res.ExitCode) + } +} diff --git a/exec/subprocess/limits_other.go b/exec/subprocess/limits_other.go new file mode 100644 index 0000000..f21ae60 --- /dev/null +++ b/exec/subprocess/limits_other.go @@ -0,0 +1,35 @@ +//go:build !unix + +package subprocess + +import "errors" + +// checkLaunch refuses unconditionally outside Unix. Process groups, +// Credential-based uid dropping, and rlimits are all POSIX concepts this +// rung does not emulate anywhere else — running the child unconfined would +// make Run look like it succeeded while providing none of the isolation +// this package exists for, which is worse than failing loudly. So it fails +// loudly instead, regardless of what o carries. +func checkLaunch(options) error { + return errors.New("dispatch/exec/subprocess: the subprocess rung requires a Unix platform") +} + +// SameUserRefused always reports false outside Unix. Available, below, +// already refuses the whole rung on this platform, so configuration code +// that calls Available first — as it must — never reaches this question +// in practice; it exists here only so callers that are not platform- +// specific themselves (extension.resolveExecutionOptions) compile on +// every platform this package supports. See checkLaunch and its Unix +// counterpart in limits_unix.go, the one that matters. +func SameUserRefused(int, bool) bool { return false } + +// Available reports that this platform cannot run the subprocess rung at +// all, for the identical reason checkLaunch refuses above. Configuration +// code should call this before ever constructing an Executor: checkLaunch +// catches the same condition too, but only once Run is actually called +// for a job attempt, which turns a startup misconfiguration into a +// per-job launch failure discovered in production instead of a single +// clear error at boot. +func Available() error { + return errors.New("dispatch/exec/subprocess: the subprocess rung requires a Unix platform") +} diff --git a/exec/subprocess/limits_unix.go b/exec/subprocess/limits_unix.go new file mode 100644 index 0000000..f3a6a34 --- /dev/null +++ b/exec/subprocess/limits_unix.go @@ -0,0 +1,73 @@ +//go:build unix + +package subprocess + +import ( + "fmt" + "os" +) + +// checkLaunch runs at the top of Run, before any pipe or process is +// created, and refuses to launch when the configured options would gut +// this rung's isolation. +// +// On Unix, running the child as the worker's own uid — whether because no +// uid was configured at all, or because one was configured that happens to +// match the worker's own — leaves it able to read ~/.aws, /var/run/secrets, +// and the Dispatch config, which removes most of the value of this rung. +// Both shapes are refused unless the caller opts in explicitly via +// WithAllowSameUser: there is one switch for "I know this is unisolated," +// not a separate one for each way of ending up unisolated. Rlimits have no +// equivalent launch-time check — they cannot fail until they are actually +// applied, which happens child-side in shim.Main (see EnvRlimitAS and +// friends in exec/shim), since Go cannot set a child's rlimits through +// SysProcAttr. +func checkLaunch(o options) error { + if !o.hasUser { + if o.allowSameUser { + return nil + } + + return fmt.Errorf( + "dispatch/exec/subprocess: no uid configured; the child would run as the worker's own uid, " + + "which defeats this rung's isolation — pass WithUser to configure a dedicated uid, or " + + "WithAllowSameUser to accept running unisolated", + ) + } + + if SameUserRefused(o.uid, o.allowSameUser) { + return fmt.Errorf( + "dispatch/exec/subprocess: configured uid %d matches the worker's own uid; "+ + "running the child as the worker defeats this rung's isolation — pass WithAllowSameUser to allow it", + o.uid, + ) + } + + return nil +} + +// SameUserRefused reports whether checkLaunch would refuse to launch a +// child configured with this uid and allowSameUser setting because the uid +// matches the worker's own — one of the two conditions checkLaunch checks, +// factored out so configuration code can ask the same question before ever +// constructing an Executor, without keeping a second, independent copy of +// checkLaunch's own logic that could drift from it silently. It does not +// check o.hasUser, and so does not answer checkLaunch's other condition — +// no uid configured at all — which callers must check separately; see +// resolveExecutionOptions (extension/execution.go) for the shape that +// checking both looks like. +func SameUserRefused(uid int, allowSameUser bool) bool { + return !allowSameUser && uid == os.Getuid() +} + +// Available reports whether this platform can run the subprocess rung at +// all — always true on Unix. It only answers the platform question; +// checkLaunch (and SameUserRefused, above) answers different ones — does +// THIS configuration name a uid at all, and if so, is it the worker's own — +// which only checkLaunch itself catches on the actual launch path, once Run +// is called for a job attempt. Configuration code should call Available +// before ever constructing an Executor, so a deployment on an unsupported +// platform fails once, loudly, at startup — not job by job, on every +// attempt's launch failure, once it is already running. Check hasUser and +// call SameUserRefused alongside it for the same reason. +func Available() error { return nil } diff --git a/exec/subprocess/limits_unix_test.go b/exec/subprocess/limits_unix_test.go new file mode 100644 index 0000000..9270544 --- /dev/null +++ b/exec/subprocess/limits_unix_test.go @@ -0,0 +1,235 @@ +//go:build unix + +package subprocess_test + +import ( + "context" + "os" + "strings" + "testing" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/exectest" + "github.com/xraph/dispatch/exec/subprocess" +) + +// TestSameUserIsRefusedByDefault proves checkLaunch (limits_unix.go) +// refuses to start the child when the configured uid matches the +// worker's own: without this, the uid boundary this rung exists to +// provide is silently absent, since the child ends up able to read +// everything the worker itself can. +func TestSameUserIsRefusedByDefault(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithUser(os.Getuid(), os.Getgid()), + ) + + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + switch { + case err != nil: // acceptable: refused at launch + case res.Status != exec.StatusLaunchFailed: + t.Fatalf("Status = %q, want launch_failed — running as the worker's own UID guts this rung", res.Status) + } +} + +// TestSameUserAllowedExplicitly proves WithAllowSameUser lifts the refusal +// above and the run proceeds normally. Credential.NoSetGroups matters +// here specifically: without it, os/exec's setgroups(2) call requires +// privilege this test process (running as an ordinary, non-root user, as +// every dev machine and CI run here does) does not have, and this case +// would fail to launch even though the uid/gid themselves are unchanged. +func TestSameUserAllowedExplicitly(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithUser(os.Getuid(), os.Getgid()), + subprocess.WithAllowSameUser(), + ) + + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusOK { + t.Fatalf("Status = %q, want ok (err %q)", res.Status, res.HandlerErr) + } +} + +// TestNoUserIsRefusedByDefault proves checkLaunch refuses to start when no +// uid is configured at all — not just when a configured uid happens to +// match the worker's own. Without this, `execution.subprocess.enabled: +// true` with no `user` set would run the child as the worker's own uid +// silently: it defeats the uid boundary in exactly the same way a +// configured same-uid does, so it is refused the same way, not merely +// warned about. +func TestNoUserIsRefusedByDefault(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + ) + + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + switch { + case err != nil: // acceptable: refused at launch + case res.Status != exec.StatusLaunchFailed: + t.Fatalf("Status = %q, want launch_failed — no uid configured guts this rung", res.Status) + } +} + +// TestNoUserAllowedExplicitly proves WithAllowSameUser lifts the no-uid +// refusal above too, not just the same-uid one: it is the single opt-out +// for both shapes of running this rung unisolated. +func TestNoUserAllowedExplicitly(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithAllowSameUser(), + ) + + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusOK { + t.Fatalf("Status = %q, want ok (err %q)", res.Status, res.HandlerErr) + } +} + +// TestSameUserRefusedMatchesCheckLaunch pins subprocess.SameUserRefused +// to checkLaunch's own actual launch-time decision. checkLaunch calls +// SameUserRefused itself (see limits_unix.go), so in principle this test +// cannot fail unless a future edit gives the two functions independent +// logic again — which is exactly the drift extension.resolveExecutionOptions +// (extension/execution.go) depends on not happening: it calls +// SameUserRefused to fail fast at startup, and that check is only correct +// for as long as it asks the identical question checkLaunch asks at +// Run() time. +func TestSameUserRefusedMatchesCheckLaunch(t *testing.T) { + tests := []struct { + name string + allowSameUser bool + }{ + {"refused without AllowSameUser", false}, + {"allowed with AllowSameUser", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + predicted := subprocess.SameUserRefused(os.Getuid(), tt.allowSameUser) + + opts := []subprocess.Option{ + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithUser(os.Getuid(), os.Getgid()), + } + if tt.allowSameUser { + opts = append(opts, subprocess.WithAllowSameUser()) + } + + e := subprocess.New(opts...) + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + + refused := err != nil || res.Status == exec.StatusLaunchFailed + if refused != predicted { + t.Fatalf("SameUserRefused(%d, %v) = %v, but the actual launch's refusal = %v — "+ + "the two have drifted apart", os.Getuid(), tt.allowSameUser, predicted, refused) + } + }) + } +} + +// TestRlimitsAreAppliedChildSide proves a configured Rlimits value +// actually reaches the child: RLIMIT_NOFILE is used rather than +// RLIMIT_AS or RLIMIT_CORE because it is the one limit in this set that +// (a) Darwin actually allows setrlimit to lower, unlike RLIMIT_AS (see +// the AddressSpace doc comment on Rlimits), and (b) this test process can +// still observe from the parent side: a child capped well below the +// number of file descriptors JobOK itself needs (fd 3 and 4 for the wire +// protocol, plus whatever the Go runtime and os/exec's own stdio setup +// already hold open) fails outright, which is a clear, portable signal +// that the limit landed rather than being silently ignored. +func TestRlimitsAreAppliedChildSide(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithRlimits(subprocess.Rlimits{NoFile: 3}), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; this test is not about the uid boundary + ) + + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + // Lowering RLIMIT_NOFILE does not close or invalidate descriptors + // already open at the time it takes effect — fd 0/1/2 (inherited) + // and fd 3/4 (the wire protocol, opened by the parent before exec) + // all survive a soft limit of 3 that comes later, in applyRlimits. + // What a limit of 3 does is make the *next* open() the shim's own + // runtime needs (a network poller fd, a temp file, anything) fail, + // which is early and unconditional enough in a live Go program that + // the shim cannot reach StatusOK afterwards. That is the actual + // mechanism this asserts on: not descriptors 3/4 becoming unusable, + // but nothing further being allocatable. + if res.Status == exec.StatusOK { + t.Error("Status = ok; a RLIMIT_NOFILE of 3 should have made the child unable to run at all, so the limit was not applied") + } +} + +// TestStrictRlimitsFailsLaunchOnUnexpectedFailure proves WithStrictRlimits +// turns an rlimit failure into StatusLaunchFailed end to end, through a +// real forked child rather than shim's own in-process unit tests (see +// exec/shim/rlimit_unix_test.go for why those stick to values that never +// reach a real Setrlimit call). NoFile: -1 is deliberately a value +// applyRlimits rejects before ever calling Setrlimit — see that same +// file's TestApplyOneUnverifiedResourceIsAFailure and its neighbors for +// why a value that actually depends on kernel-specific hard limits (a +// real NoFile ceiling, say) would be less portable across this rung's +// supported platforms than a value rejected up front. +func TestStrictRlimitsFailsLaunchOnUnexpectedFailure(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithRlimits(subprocess.Rlimits{NoFile: -1}), + subprocess.WithStrictRlimits(), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; this test is not about the uid boundary + ) + + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusLaunchFailed { + t.Errorf("Status = %q, want launch_failed", res.Status) + } + if !strings.Contains(res.HandlerErr, "RLIMIT_NOFILE") { + t.Errorf("HandlerErr = %q, want it to name RLIMIT_NOFILE", res.HandlerErr) + } +} + +// TestStrictRlimitsToleratesKnownUnsupported proves WithStrictRlimits does +// not turn a platform's own structural refusal of a limit into a launch +// failure. The AddressSpace value here is deliberately generous (2GiB) so +// that on a platform where RLIMIT_AS is actually settable (Linux, in +// particular this rung's CI) it simply succeeds rather than crashing this +// small handler; on Darwin, syscall.Setrlimit(RLIMIT_AS, ...) fails with +// EINVAL unconditionally regardless of value, which isKnownUnsupported +// classifies as structural rather than a misconfiguration WithStrictRlimits +// should catch. Either way the expected outcome is the same: StatusOK. +func TestStrictRlimitsToleratesKnownUnsupported(t *testing.T) { + e := subprocess.New( + subprocess.WithBinary(os.Args[0]), + subprocess.WithEnv(map[string]string{"DISPATCH_EXEC_SHIM_TEST": "1"}), + subprocess.WithRlimits(subprocess.Rlimits{AddressSpace: 2 << 30}), + subprocess.WithStrictRlimits(), + subprocess.WithAllowSameUser(), // CI cannot drop privileges; this test is not about the uid boundary + ) + + res, err := e.Run(context.Background(), request(t, exectest.JobOK, struct{}{})) + if err != nil { + t.Fatalf("Run() = %v", err) + } + if res.Status != exec.StatusOK { + t.Fatalf("Status = %q, want ok (err %q) — a platform's own refusal to support RLIMIT_AS must not be treated as a WithStrictRlimits failure", res.Status, res.HandlerErr) + } +} diff --git a/exec/subprocess/main_test.go b/exec/subprocess/main_test.go new file mode 100644 index 0000000..06b7096 --- /dev/null +++ b/exec/subprocess/main_test.go @@ -0,0 +1,343 @@ +package subprocess_test + +import ( + "context" + "os" + osexec "os/exec" + "os/signal" + "path/filepath" + "strconv" + "syscall" + "testing" + "time" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/exectest" + "github.com/xraph/dispatch/exec/shim" + "github.com/xraph/dispatch/exec/wire" +) + +// Env vars gating the fixture branches below. Each stands in for a shim +// this binary is not, so the C1/C2 regression tests can drive a real +// process tree without needing a real handler that misbehaves this way. +const ( + // envLeakChild selects a fixture that behaves like a shim whose + // handler shelled out to something and left it running: it reads the + // request, spawns a grandchild that inherits stdout, stderr, and both + // wire descriptors and then just sleeps, writes a normal successful + // result, and exits promptly. Reproduces C1: the tracked process + // exits cleanly and quickly, but something it spawned is still + // holding the pipes open. + envLeakChild = "DISPATCH_EXEC_LEAK_CHILD_TEST" + + // envSleepOnly selects a fixture that does nothing but sleep and + // exit. It never touches fd 3 or fd 4 at all, so used as the direct + // child it reproduces C2: a child that is alive but has not reached + // wire.Decode. Used as envLeakChild's grandchild, it is what keeps + // C1's pipes open past the tracked process's own exit. + envSleepOnly = "DISPATCH_EXEC_SLEEP_ONLY_TEST" + + // envGroupKill selects a fixture for the kill ladder's own test + // (kill_unix_test.go): it ignores SIGTERM itself FIRST — before + // forking anything, the safer order, closing the window where a + // SIGTERM landing between fork and signal.Ignore would kill it + // outright — then forks a grandchild that ignores SIGTERM and sleeps + // far longer than this fixture's own SIGTERM-ignoring sleep + // (envLongSleep, not envSleepOnly — see its own doc comment for why), + // writes that grandchild's pid to a file in the request's OutputDir, + // and sleeps. Only the ladder's SIGKILL half — sent to the whole + // process group, not just this fixture — can end either process, + // which is what makes this the fixture that catches a missing + // Setpgid: without it, this process never becomes its own group + // leader, terminate's pgid (its pid) names no real group, both the + // SIGTERM and the escalating SIGKILL come back ESRCH, and NEITHER + // process — not this one, not the grandchild — is ever signalled by + // the ladder at all. + envGroupKill = "DISPATCH_EXEC_GROUP_KILL_TEST" + + // envLongSleep selects a fixture that ignores SIGTERM and sleeps for + // longSleep, far longer than fixtureSleep. It exists specifically as + // envGroupKill's grandchild: that fixture's own leader also ignores + // SIGTERM and sleeps fixtureSleep (30s), and under a hypothetical + // Setpgid regression neither process is ever actually signalled by + // the kill ladder at all, so both would simply run out their own + // timers and exit on their own. If the grandchild's timer were also + // fixtureSleep, it would exit at roughly the same wall-clock moment + // the leader's own timer does — which is also roughly when Run() + // finally returns in that broken scenario — making the test's ESRCH + // check on the grandchild's pid pass for the wrong reason (it expired + // on its own, not because any signal reached it) instead of catching + // the regression. Giving it a much longer timer means that if the + // kill ladder never actually reaches it, it is still provably alive + // when the test checks, and the ESRCH assertion is load-bearing on + // its own rather than riding on coincidental timing. + envLongSleep = "DISPATCH_EXEC_LONG_SLEEP_TEST" + + // envIgnoreSigtermLongSleep selects a fixture like envLongSleep, but + // used as a *helper* left behind by a cooperative leader + // (envLeaderExitsHelperSurvives) rather than as envGroupKill's + // grandchild. Functionally identical to envLongSleep; kept as a + // separate fixture rather than reused so each test's intent reads + // clearly from which env var its process tree is built out of. + envIgnoreSigtermLongSleep = "DISPATCH_EXEC_IGNORE_SIGTERM_LONG_SLEEP_TEST" + + // envLeaderExitsHelperSurvives selects the fixture for the C1 + // regression test (kill_unix_test.go): it forks a helper that ignores + // SIGTERM and sleeps far longer than any bound the test asserts on + // (envIgnoreSigtermLongSleep), writes that helper's pid to a file in + // the request's OutputDir, and then — unlike envGroupKill's own + // fixture — does *not* ignore SIGTERM itself, so the signal's default + // disposition ends this leader almost immediately once the ladder's + // first rung arrives. This is the shape terminate's escalation used + // to get wrong: a leader that exits promptly on SIGTERM (standing in + // for the production shim's own cooperative shutdown) while something + // it forked keeps running. A version of the ladder that decided + // whether to escalate to SIGKILL by asking only "has the leader + // exited" would answer yes almost immediately and never send it, + // leaving the helper alive indefinitely. + envLeaderExitsHelperSurvives = "DISPATCH_EXEC_LEADER_EXITS_HELPER_SURVIVES_TEST" + + // envSigtermMarker selects a fixture for the C3 regression test + // (kill_unix_test.go): unlike every other fixture in this file, it + // carries no timeout or sleep of its own at all — it reads the + // request only to learn OutputDir, then blocks indefinitely until it + // receives an actual SIGTERM, at which point it writes a marker file + // and exits promptly. That is what makes it load-bearing evidence + // that the ladder's SIGTERM rung actually ran: this process has no + // other way to end quickly, so the marker existing (and Run() + // returning well under the grace period) can only mean SIGTERM was + // sent and handled — not, for instance, that a deadline embedded in + // the request itself caused a cooperative exit through some unrelated + // path, and not that the ladder skipped straight to SIGKILL, which + // this fixture cannot trap or write a marker in response to. + envSigtermMarker = "DISPATCH_EXEC_SIGTERM_MARKER_TEST" + + // fixtureSleep is deliberately much longer than any bound the C1/C2 + // tests assert on, so a regression is caught by the test's own + // timeout rather than by this sleep ever completing. + fixtureSleep = 30 * time.Second + + // longSleep is deliberately much longer than fixtureSleep — see + // envLongSleep's own doc comment for why that gap matters. + longSleep = 5 * time.Minute +) + +// TestMain lets this test binary act as its own sandbox child. The +// subprocess rung re-execs the running binary, so under test the child is +// this binary again; when the marker env var is set it runs the shim and +// exits instead of running tests. +func TestMain(m *testing.M) { + switch { + case os.Getenv("DISPATCH_EXEC_SHIM_TEST") != "": + shim.Main(exectest.Handlers()...) + return // unreachable; Main exits + case os.Getenv(envLeakChild) != "": + runLeakChild() + return // unreachable; runLeakChild exits + case os.Getenv(envSleepOnly) != "": + // Stands in for a grandchild left running by a handler's own + // subprocess: it holds whatever fds it inherited open and does + // nothing else. + time.Sleep(fixtureSleep) + os.Exit(0) + return + case os.Getenv(envGroupKill) != "": + runGroupKillFixture() + return // unreachable; runGroupKillFixture exits + case os.Getenv(envLongSleep) != "": + signal.Ignore(syscall.SIGTERM) + time.Sleep(longSleep) + os.Exit(0) + return + case os.Getenv(envIgnoreSigtermLongSleep) != "": + signal.Ignore(syscall.SIGTERM) + time.Sleep(longSleep) + os.Exit(0) + return + case os.Getenv(envLeaderExitsHelperSurvives) != "": + runLeaderExitsHelperSurvivesFixture() + return // unreachable; runLeaderExitsHelperSurvivesFixture exits + case os.Getenv(envSigtermMarker) != "": + runSigtermMarkerFixture() + return // unreachable; runSigtermMarkerFixture exits + } + + os.Exit(m.Run()) +} + +// runLeakChild is the envLeakChild fixture body. See its doc comment +// above for what it reproduces. +func runLeakChild() { + in := os.NewFile(uintptr(fdFromEnv(shim.EnvRequestFD, 3)), "dispatch-exec-request") + out := os.NewFile(uintptr(fdFromEnv(shim.EnvResultFD, 4)), "dispatch-exec-result") + + // Drain the request so the parent's write does not block on this + // fixture; the content itself does not matter here. + _, _ = wire.Decode(in) + + // CommandContext with context.Background() rather than Command, purely + // to satisfy noctx; this fixture never cancels the grandchild via ctx + // — it is deliberately left running, see the Start comment below. + grandchild := osexec.CommandContext(context.Background(), os.Args[0]) + // Deliberately NOT append(os.Environ(), ...): this process's own + // environment still carries envLeakChild (it is inherited, not + // consumed), and os.Environ() would carry it straight into the + // grandchild too — which TestMain's switch checks first, so the + // grandchild would decide it is another leak-child fixture and spawn + // its own grandchild, and so on without ever bottoming out. An + // explicit, minimal env keeps this fixture to exactly the two + // generations it means to create. + grandchild.Env = []string{envSleepOnly + "=1"} + grandchild.Stdout = os.Stdout + grandchild.Stderr = os.Stderr + grandchild.ExtraFiles = []*os.File{in, out} + // Deliberately not waited on: it must outlive this process to + // reproduce C1, exactly as a detached background job a handler + // started would. + _ = grandchild.Start() + + res := &exec.Result{Status: exec.StatusOK} + _ = wire.Encode(out, &wire.Frame{Kind: wire.KindResult, Result: res}) + + os.Exit(0) +} + +// runGroupKillFixture is the envGroupKill fixture body. See its doc +// comment above for what it reproduces. +func runGroupKillFixture() { + in := os.NewFile(uintptr(fdFromEnv(shim.EnvRequestFD, 3)), "dispatch-exec-request") + + frame, err := wire.Decode(in) + if err != nil || frame.Request == nil { + os.Exit(1) + return + } + + // Unlike the shim, which traps SIGTERM to cancel its handler's + // context and shut down cleanly, this fixture ignores it outright — + // standing in for a handler process that does not cooperate with the + // first rung of the ladder at all, so only the SIGKILL half can end + // it, and only if that SIGKILL actually reaches this process's whole + // group rather than just its leader. + signal.Ignore(syscall.SIGTERM) + + // CommandContext with context.Background() rather than Command, purely + // to satisfy noctx; this fixture never cancels the grandchild via ctx + // — it must outlive this process's own SIGTERM handling, exactly as a + // native library's forked helper would. + grandchild := osexec.CommandContext(context.Background(), os.Args[0]) + // Deliberately NOT append(os.Environ(), ...) — see runLeakChild's own + // comment on the same line for why: this process's environment still + // carries envGroupKill, and inheriting it wholesale would make the + // grandchild decide it is another instance of this same fixture. + // + // envLongSleep, not envSleepOnly: see envLongSleep's own doc comment + // for why this grandchild needs a sleep clearly longer than this + // fixture's own fixtureSleep. + grandchild.Env = []string{envLongSleep + "=1"} + if err := grandchild.Start(); err != nil { + os.Exit(1) + return + } + + // The parent test process reads this file after Run returns to learn + // which pid to probe for liveness — it has no other way to learn a + // pid this deep in a process tree it does not control directly. + pidPath := filepath.Join(frame.Request.OutputDir, "grandchild.pid") + if err := os.WriteFile(pidPath, []byte(strconv.Itoa(grandchild.Process.Pid)), 0o600); err != nil { + os.Exit(1) + return + } + + // No result frame is written: this fixture is killed before it gets + // the chance to, which is the point — the parent's classify call has + // nothing to trust here but the process's own wait status. + time.Sleep(fixtureSleep) + os.Exit(0) +} + +// runLeaderExitsHelperSurvivesFixture is the envLeaderExitsHelperSurvives +// fixture body. See its doc comment above for what it reproduces. +func runLeaderExitsHelperSurvivesFixture() { + in := os.NewFile(uintptr(fdFromEnv(shim.EnvRequestFD, 3)), "dispatch-exec-request") + + frame, err := wire.Decode(in) + if err != nil || frame.Request == nil { + os.Exit(1) + return + } + + // CommandContext with context.Background() rather than Command, purely + // to satisfy noctx; this fixture never cancels the helper via ctx — it + // must outlive this leader, exactly as a native library's forked + // helper would. + helper := osexec.CommandContext(context.Background(), os.Args[0]) + // Deliberately NOT append(os.Environ(), ...) — see runLeakChild's own + // comment on the same line for why. + helper.Env = []string{envIgnoreSigtermLongSleep + "=1"} + if err := helper.Start(); err != nil { + os.Exit(1) + return + } + + pidPath := filepath.Join(frame.Request.OutputDir, "helper.pid") + if err := os.WriteFile(pidPath, []byte(strconv.Itoa(helper.Process.Pid)), 0o600); err != nil { + os.Exit(1) + return + } + + // Deliberately no signal.Ignore call here, unlike runGroupKillFixture: + // this leader lets SIGTERM's default disposition (terminate) end it + // almost immediately, standing in for the production shim's own + // cooperative shutdown. The exact mechanism does not matter to this + // fixture, only that the leader exits promptly, well before grace + // elapses, while the helper it just forked (which does ignore + // SIGTERM) does not. + time.Sleep(fixtureSleep) + os.Exit(0) +} + +// runSigtermMarkerFixture is the envSigtermMarker fixture body. See its +// doc comment above for what it reproduces. +func runSigtermMarkerFixture() { + in := os.NewFile(uintptr(fdFromEnv(shim.EnvRequestFD, 3)), "dispatch-exec-request") + + frame, err := wire.Decode(in) + if err != nil || frame.Request == nil { + os.Exit(1) + return + } + + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGTERM) + + // Blocks until the parent's kill ladder sends SIGTERM. Nothing else + // in this fixture can end it — no deadline, no fixed sleep — so + // reaching the line below is itself proof that a real SIGTERM arrived + // and was handled, not a side effect of some unrelated timer. + <-sigCh + + markerPath := filepath.Join(frame.Request.OutputDir, "sigterm-received") + _ = os.WriteFile(markerPath, []byte("1"), 0o600) + + os.Exit(0) +} + +// fdFromEnv mirrors shim's own unexported helper of the same name: it +// reads a file descriptor number from the named environment variable, +// falling back to def when unset or unparsable. Duplicated here rather +// than imported because shim does not export it, and this fixture is not +// part of the shim package. +func fdFromEnv(name string, def int) int { + v, ok := os.LookupEnv(name) + if !ok { + return def + } + + n, err := strconv.Atoi(v) + if err != nil { + return def + } + + return n +} diff --git a/exec/subprocess/procattr_other.go b/exec/subprocess/procattr_other.go new file mode 100644 index 0000000..c280c99 --- /dev/null +++ b/exec/subprocess/procattr_other.go @@ -0,0 +1,27 @@ +//go:build !unix + +package subprocess + +import ( + osexec "os/exec" + "syscall" +) + +// sysProcAttr is a no-op outside Unix: process groups and Credential-based +// uid dropping are both POSIX concepts this rung does not emulate anywhere +// else, and this package does not otherwise claim to support a non-Unix +// platform. checkLaunch (limits_other.go) is what actually refuses Run on +// this platform, so this stub is never reached in practice — it exists +// only so the package still compiles here. +func sysProcAttr(options) *syscall.SysProcAttr { return nil } + +// killGroup falls back to killing the process directly outside Unix, +// since there is no process group to address as a whole. sig exists only +// for signature symmetry with the Unix build's killGroup: os.Process.Kill +// is the only process-ending call this build can make through os/exec +// regardless of which signal the ladder asked for, and this path is never +// reached in practice anyway — checkLaunch (limits_other.go) refuses to +// start this rung at all outside Unix. +func killGroup(cmd *osexec.Cmd, _ syscall.Signal) error { + return cmd.Process.Kill() +} diff --git a/exec/subprocess/procattr_unix.go b/exec/subprocess/procattr_unix.go new file mode 100644 index 0000000..6d14d93 --- /dev/null +++ b/exec/subprocess/procattr_unix.go @@ -0,0 +1,152 @@ +//go:build unix + +package subprocess + +import ( + "errors" + "os" + osexec "os/exec" + "syscall" +) + +// sysProcAttr puts the child in its own process group. The kill ladder +// (Task 6) needs this to reach a native library's forked helpers, and +// killGroup below needs it right now: without it, killing the tracked +// process leaves anything it spawned running, which is the classic silent +// failure of this design — the deadline appears to have worked while the +// real work continues in the background. +// +// When a user is configured, this also sets Credential so the child drops +// to that uid/gid before exec. checkLaunch (limits_unix.go) is what +// refuses to reach here at all when the configured uid matches the +// worker's own without WithAllowSameUser; sysProcAttr itself does not +// re-derive that policy, it just builds the attribute struct. +// +// Credential.NoSetGroups is true only when both uid and gid already equal +// the caller's own — the WithAllowSameUser dev/CI path, where the +// Credential is otherwise a no-op. That is the only case where it is +// needed: os/exec calls setgroups(2) to clear supplementary groups +// whenever NoSetGroups is false and Credential.Groups is nil (see +// exec_linux.go and exec_libc2.go's shared "if !cred.NoSetGroups" +// guard), and setgroups requires privilege (CAP_SETGID on Linux, root on +// Darwin) independent of whether Uid/Gid actually change — so without +// this exception, even WithAllowSameUser's same-uid case would fail to +// launch whenever the worker itself is not root, which is every dev +// machine and every CI run here. +// +// Whenever uid or gid genuinely differ from the caller's own, this is +// false, so setgroups does run and the child's supplementary groups are +// actually cleared rather than silently inherited. That launch already +// requires the same privilege setgroups does (only root/CAP_SETUID can +// change to a different uid at all), so this does not introduce a new +// privilege requirement — it only skips the clear in the one case where +// the process doing the dropping has no such privilege to begin with, +// and dropping is a no-op anyway. Getting this backwards is a real +// containment gap, not a cosmetic one: a worker running as root with +// supplementary group "docker" (common for a systemd unit that also +// manages containers) would otherwise hand every "sandboxed" child that +// same group membership — and therefore group-write access to +// /var/run/docker.sock, i.e. root on the host — regardless of the uid it +// was dropped to. +func sysProcAttr(o options) *syscall.SysProcAttr { + attr := &syscall.SysProcAttr{Setpgid: true} + if o.hasUser { + attr.Credential = &syscall.Credential{ + Uid: uint32(o.uid), //nolint:gosec // G115: operator-configured via WithUser, never attacker input. + Gid: uint32(o.gid), //nolint:gosec // G115: operator-configured via WithUser, never attacker input. + NoSetGroups: o.uid == os.Getuid() && o.gid == os.Getgid(), + } + } + + return attr +} + +// killGroup sends sig to the child's whole process group rather than just +// the child itself. Setpgid without an explicit Pgid makes the child its +// own group leader, so its pid doubles as its pgid, and signalling the +// negative pid — syscall.Kill(-pid, sig) — is how POSIX addresses a group +// rather than one process. Whatever signal is sent reaches every +// descendant the tracked process forked, not only the process this +// package started directly. +// +// It takes the signal as a parameter, rather than being a SIGKILL-only +// function, so that the probe below — the part that is actually +// delicate — has one general-purpose home instead of being copied for +// each signal that might need it. terminate (kill_unix.go) calls this +// for the ladder's SIGTERM leg, sent right when a decision to terminate +// has just been made, which is exactly the situation this probe is +// suited to: the tracked process is overwhelmingly likely to still be +// the one this package started, and the only race worth guarding against +// is the few-Go-statements gap the doc comment below describes. +// +// terminate's SIGKILL leg deliberately does *not* come through here — +// see its own doc comment for why: by the time grace has elapsed, the +// tracked leader having already exited is the expected shape of the +// exact bug that escalation exists to catch, not a rare race, so gating +// that signal on cmd.Process specifically being still alive would silence +// it in precisely the case it is needed. This function's probe is +// correct for a signal sent at the *start* of a kill; it is the wrong +// tool for one decided after a wait, which is what changed between "this +// function used to be terminate's only signal path" and now. +// +// The probe before the kill exists because a raw syscall.Kill(-pid, ...) +// has no idea whether pid still names the process this package started. +// os.Process.wait marks the process done (doRelease(statusDone)) and takes +// a write lock on its own sigMu *before* it calls wait4 — the syscall that +// actually lets the kernel hand the pid to someone else — specifically so +// that Process.Signal, which read-locks the same sigMu, cannot land on a +// reused pid once that has happened. Routing through cmd.Process.Signal +// first reuses that same fencing instead of re-deriving it, which +// syscall.Kill alone has none of. It does not close the window entirely — +// there is still a gap between this check succeeding and the raw +// syscall.Kill call below — but it narrows it from "however long the +// waitLoop select takes to notice the process exited" down to a couple of +// Go statements, and the waitLoop fix above (checking waitCh before +// setting timedOut) removes the specific interleaving that used to make +// that gap wide enough to matter in practice. +// +// Two trades this makes, both deliberate: +// +// A non-ErrProcessDone probe error falls through to attempt the group +// signal anyway rather than returning it. waitLoop has no way to make +// progress other than waitCh eventually firing, so a killGroup that gives +// up here would not fail the attempt — it would hang Run indefinitely +// instead, waiting on a kill that was never sent to a process that is +// never going to exit on its own. That is strictly worse than attempting +// a signal that might itself fail: signalling this package's own child +// used to have no path that produces a non-ErrProcessDone error here, but +// sysProcAttr above now sets Credential with a dedicated uid when one is +// configured, and that uid boundary makes EPERM a real possibility. A +// failed signal is recoverable — classify still has the process's actual +// wait status to report from, whatever it turns out to be; a signal that +// was never attempted is not. +// +// An ErrProcessDone probe result returns immediately, without ever +// reaching the group signal below. For the SIGTERM call this function is +// actually used for today, that is a narrow, accepted gap: a leader +// reaped in the couple of Go statements between waitLoop's own check and +// this probe landing leaves any surviving grandchildren un-signalled by +// *this* call, where an unconditional syscall.Kill(-pid, ...) (pid == +// pgid here) would still have reached them, since a pgid stays valid as +// long as any member of the group is still alive, leader or not. +// Signalling anyway in that case was considered and rejected: it would +// mean signalling a pgid derived from a pid the kernel may already have +// handed to an unrelated process group, which is the exact hazard this +// probe exists to avoid — reaching a stray grandchild on the SIGTERM leg +// is not worth reintroducing that, because it is not the last word: +// terminate's SIGKILL escalation (kill_unix.go) does not route through +// this function at all, specifically so that a leader reaped by the time +// grace elapses — the expected shape of a cooperative exit, not a rare +// race — cannot make the *final* signal a no-op the same way. This +// narrows an already-partial guarantee rather than removing a complete +// one: a grandchild left behind by a tracked process that exited cleanly +// on its own, before killProcess was ever called at all, was already +// unreachable by this function (see the +// drainGrace comment in Run). +func killGroup(cmd *osexec.Cmd, sig syscall.Signal) error { + if err := cmd.Process.Signal(syscall.Signal(0)); err != nil && errors.Is(err, os.ErrProcessDone) { + return nil + } + + return syscall.Kill(-cmd.Process.Pid, sig) +} diff --git a/exec/subprocess/procattr_unix_internal_test.go b/exec/subprocess/procattr_unix_internal_test.go new file mode 100644 index 0000000..63ef2df --- /dev/null +++ b/exec/subprocess/procattr_unix_internal_test.go @@ -0,0 +1,116 @@ +//go:build unix + +package subprocess + +// This file is package subprocess (internal), not subprocess_test — same +// rationale as internal_test.go: sysProcAttr is unexported, and what it +// builds (Setpgid, Credential.Uid/Gid, and the NoSetGroups policy) cannot +// be observed from outside the package without actually spawning a +// process and inspecting its privileges, which for a genuinely different +// uid needs root and is exactly what the brief says not to write. Calling +// sysProcAttr directly with synthetic options pins its contract instead, +// the same way internal_test.go pins classify's. +// +// This exists because TestSameUserIsRefusedByDefault and +// TestSameUserAllowedExplicitly (limits_unix_test.go) exercise +// checkLaunch, a policy function with no dependency on sysProcAttr at +// all — replacing sysProcAttr's entire body with a bare `return +// &syscall.SysProcAttr{Setpgid: true}` (dropping Credential unconditionally) +// leaves both of those tests, and the rest of this package's suite, +// green. That is the regression this file exists to catch: a change +// that silently stops dropping privileges at all, the one thing this +// task exists to make happen. + +import ( + "os" + "testing" +) + +func TestSysProcAttrSetsSetpgid(t *testing.T) { + attr := sysProcAttr(options{}) + if attr == nil { + t.Fatal("sysProcAttr(options{}) = nil") + } + if !attr.Setpgid { + t.Error("Setpgid = false, want true — Task 4's whole-group kill depends on this") + } +} + +func TestSysProcAttrNoUserConfiguredSetsNoCredential(t *testing.T) { + attr := sysProcAttr(options{}) + if attr.Credential != nil { + t.Errorf("Credential = %+v, want nil when no user is configured", attr.Credential) + } +} + +// TestSysProcAttrCredential covers both branches of the NoSetGroups +// policy directly, without spawning a process: the same-uid case CI can +// actually exercise end to end (see TestSameUserAllowedExplicitly), and +// the differing-uid case, which needs root to launch for real and so is +// asserted here as a value, per the brief's own guidance not to write a +// root-only test. +func TestSysProcAttrCredential(t *testing.T) { + self := os.Getuid() + selfGid := os.Getgid() + + tests := []struct { + name string + uid, gid int + wantNoSetGroups bool + }{ + { + name: "uid and gid both match the caller's own", + uid: self, + gid: selfGid, + wantNoSetGroups: true, // the WithAllowSameUser path: Credential is a no-op, setgroups needs privilege this process does not have + }, + { + name: "uid differs", + uid: self + 1, + gid: selfGid, + wantNoSetGroups: false, // a genuine drop: already needs CAP_SETUID/root, so setgroups actually runs and clears supplementary groups + }, + { + name: "gid differs", + uid: self, + gid: selfGid + 1, + wantNoSetGroups: false, + }, + { + name: "both differ", + uid: self + 1, + gid: selfGid + 1, + wantNoSetGroups: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + attr := sysProcAttr(options{hasUser: true, uid: tt.uid, gid: tt.gid}) + + if attr.Credential == nil { + t.Fatal("Credential = nil, want non-nil when hasUser is true") + } + if attr.Credential.Uid != uint32(tt.uid) { + t.Errorf("Credential.Uid = %d, want %d", attr.Credential.Uid, tt.uid) + } + if attr.Credential.Gid != uint32(tt.gid) { + t.Errorf("Credential.Gid = %d, want %d", attr.Credential.Gid, tt.gid) + } + if attr.Credential.NoSetGroups != tt.wantNoSetGroups { + t.Errorf("Credential.NoSetGroups = %v, want %v", attr.Credential.NoSetGroups, tt.wantNoSetGroups) + } + }) + } +} + +// TestSysProcAttrAlwaysSetsSetpgidWithUser guards against a regression +// where adding Credential handling accidentally drops Setpgid — the two +// are independent fields on the same struct, and Task 4's whole-group-kill +// guarantee must survive regardless of whether a user is configured. +func TestSysProcAttrAlwaysSetsSetpgidWithUser(t *testing.T) { + attr := sysProcAttr(options{hasUser: true, uid: os.Getuid(), gid: os.Getgid()}) + if !attr.Setpgid { + t.Error("Setpgid = false, want true even when a user is configured") + } +} diff --git a/exec/subprocess/stdio.go b/exec/subprocess/stdio.go new file mode 100644 index 0000000..e1a1cd4 --- /dev/null +++ b/exec/subprocess/stdio.go @@ -0,0 +1,61 @@ +package subprocess + +import ( + "bufio" + "io" + "strings" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/exec" +) + +// shimDiagnosticPrefix marks a line on the child's stderr as coming from +// the shim itself rather than from the handler or whatever it shelled +// out to — applyRlimits (exec/shim/rlimit_unix.go) is the current +// producer. streamOutput logs these at Warn rather than Info so they do +// not blend into ordinary handler chatter, which is otherwise all Info: +// a shim diagnostic is Dispatch telling the operator something about the +// isolation itself, not output the handler chose to produce. +const shimDiagnosticPrefix = "dispatch/exec/shim: " + +// streamOutput copies r line by line into logger, tagging each line with +// the job's id and name and which stream it came from. It runs until r +// returns EOF or another read error, which happens once the child's copy +// of the underlying pipe is closed — at process exit, or once the parent +// kills it and the process is reaped. +// +// Reading is line-oriented, through a bufio.Reader rather than a +// bufio.Scanner, so a handler or a native library writing one very long +// line of unstructured output cannot exceed Scanner's default token limit +// and silently drop the rest of the stream; ReadString has no such cap. +// +// This is still gated on logger: WithLogger's own default is a no-op +// logger, so a shim diagnostic is invisible here regardless of level +// unless the caller configured one — see WithStrictRlimits for the +// rlimit case specifically, which gets a guarantee that does not depend +// on a logger being configured at all. +func streamOutput(r io.Reader, logger log.Logger, req *exec.Request, stream string) { + reader := bufio.NewReader(r) + + for { + line, err := reader.ReadString('\n') + if line != "" { + trimmed := strings.TrimSuffix(line, "\n") + fields := []log.Field{ + log.String("job_id", req.JobID.String()), + log.String("job_name", req.Name), + log.String("stream", stream), + } + + if strings.HasPrefix(trimmed, shimDiagnosticPrefix) { + logger.Warn(trimmed, fields...) + } else { + logger.Info(trimmed, fields...) + } + } + if err != nil { + return + } + } +} diff --git a/exec/wire/codec.go b/exec/wire/codec.go new file mode 100644 index 0000000..c700c3f --- /dev/null +++ b/exec/wire/codec.go @@ -0,0 +1,100 @@ +package wire + +import ( + "encoding/binary" + "errors" + "fmt" + "io" + + "github.com/vmihailenco/msgpack/v5" +) + +// ErrShortFrame marks a frame whose declared length exceeded the bytes +// actually available. The writer died mid-write. +var ErrShortFrame = errors.New("short frame") + +// Encode writes one length-prefixed frame. +func Encode(w io.Writer, f *Frame) error { + body, err := msgpack.Marshal(sanitize(f)) + if err != nil { + return fmt.Errorf("dispatch/exec/wire: marshal frame: %w", err) + } + if len(body) > MaxFrameBytes { + return fmt.Errorf("dispatch/exec/wire: frame of %d bytes exceeds the %d limit", + len(body), MaxFrameBytes) + } + + var hdr [4]byte + binary.BigEndian.PutUint32(hdr[:], uint32(len(body))) //nolint:gosec // guarded by the MaxFrameBytes check above + if _, err := w.Write(hdr[:]); err != nil { + return fmt.Errorf("dispatch/exec/wire: write header: %w", err) + } + if _, err := w.Write(body); err != nil { + return fmt.Errorf("dispatch/exec/wire: write body: %w", err) + } + + return nil +} + +// sanitize returns a Frame safe to marshal. msgpack/v5 has built-in +// support for the error interface: it encodes a non-nil error as its +// Error() string and, on decode, reconstructs a new error from that +// string. Left alone, that would silently turn Result.Cause into a +// lookalike error on the far side of the wire instead of nil, which is +// exactly what exec.Result's own doc comment says a marshaling rung must +// not do: "A rung that marshals a Result leaves this nil and sets +// Permanent instead." Stripping it here, on a copy, keeps that contract +// without touching exec's types or mutating the caller's Frame. +func sanitize(f *Frame) *Frame { + if f.Result == nil || f.Result.Cause == nil { + return f + } + + r := *f.Result + r.Cause = nil + out := *f + out.Result = &r + + return &out +} + +// Decode reads one length-prefixed frame. +// +// It returns io.EOF when the stream is empty, which means the writer +// produced nothing at all, and ErrShortFrame when a header was read but +// the body was incomplete. +func Decode(r io.Reader) (*Frame, error) { + var hdr [4]byte + if _, err := io.ReadFull(r, hdr[:]); err != nil { + if errors.Is(err, io.EOF) { + return nil, io.EOF + } + if errors.Is(err, io.ErrUnexpectedEOF) { + return nil, fmt.Errorf("dispatch/exec/wire: header: %w", ErrShortFrame) + } + + return nil, fmt.Errorf("dispatch/exec/wire: read header: %w", err) + } + + n := binary.BigEndian.Uint32(hdr[:]) + if n > MaxFrameBytes { + return nil, fmt.Errorf("dispatch/exec/wire: declared frame of %d bytes exceeds the %d limit", + n, MaxFrameBytes) + } + + body := make([]byte, n) + if _, err := io.ReadFull(r, body); err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return nil, fmt.Errorf("dispatch/exec/wire: body: %w", ErrShortFrame) + } + + return nil, fmt.Errorf("dispatch/exec/wire: read body: %w", err) + } + + var f Frame + if err := msgpack.Unmarshal(body, &f); err != nil { + return nil, fmt.Errorf("dispatch/exec/wire: unmarshal frame: %w", err) + } + + return &f, nil +} diff --git a/exec/wire/codec_test.go b/exec/wire/codec_test.go new file mode 100644 index 0000000..ac7f188 --- /dev/null +++ b/exec/wire/codec_test.go @@ -0,0 +1,136 @@ +package wire_test + +import ( + "bytes" + "errors" + "io" + "testing" + "time" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/wire" + "github.com/xraph/dispatch/id" +) + +func TestRoundTripRequest(t *testing.T) { + want := &exec.Request{ + JobID: id.NewJobID(), + Name: "tessellate.model", + Payload: []byte(`{"detail":3}`), + Attempt: 2, + Deadline: time.Now().Add(time.Hour).UTC().Truncate(time.Second), + Fingerprint: "abc123", + InputDir: "/dispatch/in", + OutputDir: "/dispatch/out", + Inputs: []exec.InputSlot{{Name: "model", Path: "model/scene.ifc"}}, + Env: map[string]string{"TMPDIR": "/tmp"}, + } + + var buf bytes.Buffer + if err := wire.Encode(&buf, &wire.Frame{Kind: wire.KindRequest, Request: want}); err != nil { + t.Fatalf("Encode() = %v", err) + } + + got, err := wire.Decode(&buf) + if err != nil { + t.Fatalf("Decode() = %v", err) + } + if got.Kind != wire.KindRequest { + t.Fatalf("Kind = %v, want %v", got.Kind, wire.KindRequest) + } + if got.Request.Name != want.Name || got.Request.Attempt != want.Attempt { + t.Errorf("Request = %+v, want name %q attempt %d", got.Request, want.Name, want.Attempt) + } + if string(got.Request.Payload) != string(want.Payload) { + t.Errorf("Payload = %q, want %q", got.Request.Payload, want.Payload) + } + if got.Request.JobID != want.JobID { + t.Errorf("JobID = %v, want %v", got.Request.JobID, want.JobID) + } + if len(got.Request.Inputs) != 1 || got.Request.Inputs[0].Name != "model" { + t.Errorf("Inputs = %+v, want one slot named model", got.Request.Inputs) + } +} + +func TestRoundTripResult(t *testing.T) { + want := &exec.Result{ + Status: exec.StatusHandlerError, + HandlerErr: "bad IFC header", + Permanent: true, + Usage: exec.Usage{WallTime: 1500 * time.Millisecond}, + Outputs: []exec.OutputFile{{Name: "mesh.glb", Size: 42, Hash: "blake3:9f"}}, + } + + var buf bytes.Buffer + if err := wire.Encode(&buf, &wire.Frame{Kind: wire.KindResult, Result: want}); err != nil { + t.Fatalf("Encode() = %v", err) + } + got, err := wire.Decode(&buf) + if err != nil { + t.Fatalf("Decode() = %v", err) + } + if got.Result.Status != want.Status || got.Result.HandlerErr != want.HandlerErr { + t.Errorf("Result = %+v, want %+v", got.Result, want) + } + if !got.Result.Permanent { + t.Error("Permanent = false, want true — the wire permanence signal must survive") + } + if got.Result.Usage.WallTime != want.Usage.WallTime { + t.Errorf("WallTime = %v, want %v", got.Result.Usage.WallTime, want.Usage.WallTime) + } + if len(got.Result.Outputs) != 1 || got.Result.Outputs[0].Name != "mesh.glb" { + t.Errorf("Outputs = %+v, want one named mesh.glb", got.Result.Outputs) + } +} + +func TestResultCauseIsNotSerialised(t *testing.T) { + // Cause is a live Go error and cannot cross a process boundary. It must + // not break encoding, and it must come back nil. + in := &exec.Result{Status: exec.StatusHandlerError, Cause: errors.New("boom")} + + var buf bytes.Buffer + if err := wire.Encode(&buf, &wire.Frame{Kind: wire.KindResult, Result: in}); err != nil { + t.Fatalf("Encode() = %v", err) + } + got, err := wire.Decode(&buf) + if err != nil { + t.Fatalf("Decode() = %v", err) + } + if got.Result.Cause != nil { + t.Errorf("Cause = %v, want nil after a round trip", got.Result.Cause) + } +} + +func TestDecodeEmptyStreamIsEOF(t *testing.T) { + // The child wrote nothing at all — it crashed before producing a result. + // This must be distinguishable from a truncated frame. + _, err := wire.Decode(bytes.NewReader(nil)) + if !errors.Is(err, io.EOF) { + t.Fatalf("Decode(empty) = %v, want io.EOF", err) + } +} + +func TestDecodeTruncatedFrame(t *testing.T) { + var buf bytes.Buffer + if err := wire.Encode(&buf, &wire.Frame{ + Kind: wire.KindResult, + Result: &exec.Result{Status: exec.StatusOK}, + }); err != nil { + t.Fatalf("Encode() = %v", err) + } + + full := buf.Bytes() + _, err := wire.Decode(bytes.NewReader(full[:len(full)-2])) + if !errors.Is(err, wire.ErrShortFrame) { + t.Fatalf("Decode(truncated) = %v, want ErrShortFrame", err) + } +} + +func TestDecodeRejectsAbsurdLength(t *testing.T) { + // A corrupt or hostile header must not make the parent allocate + // gigabytes. Header is a 4-byte big-endian length. + hdr := []byte{0xFF, 0xFF, 0xFF, 0xFF} + if _, err := wire.Decode(bytes.NewReader(hdr)); err == nil { + t.Fatal("Decode(absurd length) = nil, want an error") + } +} diff --git a/exec/wire/doc.go b/exec/wire/doc.go new file mode 100644 index 0000000..c42f808 --- /dev/null +++ b/exec/wire/doc.go @@ -0,0 +1,9 @@ +// Package wire carries execution requests and results across a process +// boundary. +// +// The envelope is a 4-byte big-endian length followed by a msgpack body. +// Framing rather than a bare value stream is what lets a parent tell a +// child that produced nothing (EOF before a header — it crashed) from one +// that produced a partial write (header then short body — corruption). +// Those are different failures and get different statuses. +package wire diff --git a/exec/wire/frame.go b/exec/wire/frame.go new file mode 100644 index 0000000..81e9276 --- /dev/null +++ b/exec/wire/frame.go @@ -0,0 +1,26 @@ +package wire + +import "github.com/xraph/dispatch/exec" + +// MaxFrameBytes caps a decoded frame. A payload larger than this is +// refused rather than allocated, so a corrupt or hostile length header +// cannot exhaust the reader's memory. +const MaxFrameBytes = 64 << 20 + +// Kind identifies which side of the exchange a frame carries. +type Kind uint8 + +const ( + // KindRequest is a parent-to-child execution request. + KindRequest Kind = 1 + // KindResult is a child-to-parent execution result. + KindResult Kind = 2 +) + +// Frame is one message. Exactly one of Request or Result is set, +// matching Kind. +type Frame struct { + Kind Kind `msgpack:"kind"` + Request *exec.Request `msgpack:"request,omitempty"` + Result *exec.Result `msgpack:"result,omitempty"` +} diff --git a/extension/artifact.go b/extension/artifact.go new file mode 100644 index 0000000..eb30277 --- /dev/null +++ b/extension/artifact.go @@ -0,0 +1,125 @@ +package extension + +import ( + "fmt" + + "github.com/xraph/forge" + trovelib "github.com/xraph/trove" + "github.com/xraph/vessel" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/cache" + troveadapter "github.com/xraph/dispatch/artifact/trove" + "github.com/xraph/dispatch/resource" +) + +// resolveArtifactBackend finds the object store backing the artifact +// plane, in three tiers: an explicitly supplied backend, a named Trove +// store from configuration, then the default Trove instance in the +// container. +// +// This mirrors how resolveGroveDB discovers a database, and shares its +// most important property: when nothing is found it returns nil rather +// than an error, so an application that never asked for artifacts runs +// exactly as it did before. +// +// Trove's own extension registers *trove.Trove both unnamed and once per +// named store, so multi-store setups are covered by the named lookup +// without Dispatch importing trove/extension at all. +func (e *Extension) resolveArtifactBackend(fapp forge.App) (artifact.Backend, error) { + if e.artifactBackend != nil { + return e.artifactBackend, nil + } + + if !e.config.Artifacts.Enabled { + return nil, nil //nolint:nilnil // no backend is a valid, disabled state + } + + opts := []troveadapter.Option{} + + if name := e.config.Artifacts.TroveStore; name != "" { + t, err := vessel.InjectNamed[*trovelib.Trove](fapp.Container(), name) + if err != nil { + return nil, fmt.Errorf("dispatch: trove store %q not found in container: %w", name, err) + } + + return troveadapter.New(t, append(opts, troveadapter.WithName(name))...), nil + } + + t, err := vessel.Inject[*trovelib.Trove](fapp.Container()) + if err != nil { + // Artifacts were requested but no Trove is mounted. That is a + // configuration mistake worth reporting rather than silently + // disabling a feature the operator asked for. + return nil, fmt.Errorf( + "dispatch: artifacts are enabled but no *trove.Trove is registered in the container; "+ + "mount the trove extension or supply a backend with WithArtifactBackend: %w", err) + } + + e.Logger().Info("dispatch: auto-discovered trove from container") + + return troveadapter.New(t, opts...), nil +} + +// buildArtifactPlane constructs the artifact service and staging cache, +// returning nils when no backend is configured. +// +// resources is the shared admission ledger, or nil when the resource +// model is off. Passed in rather than read off the extension so the one +// coupling that matters is visible at the call site: this cache and the +// worker pool must be handed the SAME manager. +func (e *Extension) buildArtifactPlane( + fapp forge.App, resources resource.Manager, +) (*artifact.Service, *cache.Cache, error) { + backend, err := e.resolveArtifactBackend(fapp) + if err != nil { + return nil, nil, err + } + + if backend == nil { + return nil, nil, nil + } + + cfg := e.config.Artifacts + + svc := artifact.NewService(e.artifactStore, backend, + artifact.WithDefaultBucket(cfg.Bucket), + artifact.WithEphemeralPrefix(cfg.EphemeralPrefix), + artifact.WithRetention(cfg.Retention), + ) + + cacheOpts := []cache.Option{} + if e.logger != nil { + cacheOpts = append(cacheOpts, cache.WithLogger(e.logger)) + } + + // WithBudget configures the cache's PRIVATE manager and is ignored + // once WithManager supplies one. With the resource model on, the + // configured budget has already been routed into the shared ledger's + // disk capacity by stagingBudget, so setting it here too would be a + // second ceiling that does nothing. + if resources != nil { + cacheOpts = append(cacheOpts, cache.WithManager(resources)) + } else if cfg.Cache.Budget > 0 { + cacheOpts = append(cacheOpts, cache.WithBudget(cfg.Cache.Budget)) + } + + c, err := cache.New(cfg.Cache.Dir, backend, cacheOpts...) + if err != nil { + return nil, nil, fmt.Errorf("dispatch: create staging cache: %w", err) + } + + if perr := vessel.Provide(fapp.Container(), func() (*artifact.Service, error) { + return svc, nil + }); perr != nil { + return nil, nil, fmt.Errorf("dispatch: register artifact service in container: %w", perr) + } + + if perr := vessel.Provide(fapp.Container(), func() (*cache.Cache, error) { + return c, nil + }); perr != nil { + return nil, nil, fmt.Errorf("dispatch: register staging cache in container: %w", perr) + } + + return svc, c, nil +} diff --git a/extension/artifact_test.go b/extension/artifact_test.go new file mode 100644 index 0000000..ebdeddd --- /dev/null +++ b/extension/artifact_test.go @@ -0,0 +1,59 @@ +package extension_test + +import ( + "testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/extension" + "github.com/xraph/dispatch/store/memory" +) + +// TestArtifactsDisabledByDefault pins the backward-compatibility +// guarantee: an application that never asked for artifacts gets none, and +// nothing about its behaviour changes. +func TestArtifactsDisabledByDefault(t *testing.T) { + ext := extension.New() + + if ext.Artifacts() != nil { + t.Fatal("artifact service exists before Register — the plane must be opt-in") + } + + if ext.ArtifactCache() != nil { + t.Fatal("staging cache exists before Register") + } +} + +// TestWithArtifactBackendEnablesPlane checks that supplying a backend +// explicitly is enough — no container, no Trove, no Forge. +func TestWithArtifactBackendEnablesPlane(t *testing.T) { + backend := artifacttest.NewBackend() + + ext := extension.New( + extension.WithArtifactBackend(backend), + extension.WithArtifactStore(memory.New()), + extension.WithArtifactCacheDir(t.TempDir()), + extension.WithArtifactCacheBudget(1<<20), + ) + + if ext == nil { + t.Fatal("New returned nil") + } + + // The plane is constructed during Register, which needs a Forge app. + // What is verifiable here is that the options applied without panic + // and left the extension in a usable state. + if ext.Artifacts() != nil { + t.Fatal("service should not exist until Register runs") + } +} + +// TestArtifactBackendSatisfiesInterface is a compile-time guard in test +// form: the test double and the real adapter must stay interchangeable. +func TestArtifactBackendSatisfiesInterface(t *testing.T) { + var b artifact.Backend = artifacttest.NewBackend() + + if b.Name() != "memory" { + t.Fatalf("Name() = %q, want %q", b.Name(), "memory") + } +} diff --git a/extension/config.go b/extension/config.go index 599693c..4cb17b2 100644 --- a/extension/config.go +++ b/extension/config.go @@ -1,6 +1,11 @@ package extension -import "github.com/xraph/dispatch" +import ( + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/resource" +) // Config holds configuration for the Dispatch Forge extension. // Fields can be set programmatically via Option functions or loaded from @@ -16,7 +21,21 @@ type Config struct { // DisableMigrate disables auto-migration on start. DisableMigrate bool `default:"false" json:"disable_migrate" mapstructure:"disable_migrate" yaml:"disable_migrate"` - // Dispatch holds the core dispatcher configuration. + // Dispatch is parsed from YAML (extensions.dispatch.dispatch.* or + // dispatch.dispatch.*) but not currently applied anywhere: nothing in + // this package reads e.config.Dispatch, and dispatch.Config's own + // fields carry no yaml/mapstructure/json struct tags of their own, so + // even a populated value here binds by Go field name at best. Setting + // concurrency, queues, poll interval, or any other dispatch.Config + // field under this key parses without error and has no effect. + // + // The dispatcher itself IS configurable through this extension — + // just not through this field. WithConcurrency, WithQueues, + // WithPollInterval, WithMaxPollInterval, WithHeartbeatInterval, + // WithStaleJobThreshold, WithWorkerStoreCallTimeout, and the + // WithCron* options (options.go) each translate one dispatch.Config + // field into the matching dispatch.With* functional option; use + // those from Go, not this field from YAML. Dispatch dispatch.Config `json:"dispatch" mapstructure:"dispatch" yaml:"dispatch"` // GroveDatabase is the name of a grove.DB registered in the DI container. @@ -31,6 +50,17 @@ type Config struct { // (unnamed) kv.Store is used. GroveKV string `json:"grove_kv" mapstructure:"grove_kv" yaml:"grove_kv"` + // Artifacts configures the artifact plane. + Artifacts ArtifactConfig `json:"artifacts" mapstructure:"artifacts" yaml:"artifacts"` + + // Resources configures the worker's resource model. + Resources ResourceConfig `json:"resources" mapstructure:"resources" yaml:"resources"` + + // Execution configures which isolation rungs beyond the in-process + // default are available to job definitions that declare a stronger + // minimum via job.WithExecution. + Execution ExecutionConfig `json:"execution" mapstructure:"execution" yaml:"execution"` + // EnableDWP enables the Dispatch Wire Protocol for real-time // client communication (WebSocket, SSE, HTTP RPC). EnableDWP bool `default:"false" json:"enable_dwp" mapstructure:"enable_dwp" yaml:"enable_dwp"` @@ -50,3 +80,216 @@ func DefaultConfig() Config { BasePath: "/dispatch", } } + +// ArtifactConfig configures the artifact plane — Dispatch's tracked +// object storage for job inputs and outputs. +// +// The plane is entirely opt-in. With no backend resolved, Dispatch +// behaves exactly as it did before artifacts existed. +type ArtifactConfig struct { + // Enabled turns the artifact plane on. When false, no backend is + // resolved even if a Trove instance is present in the container. + Enabled bool `default:"false" json:"enabled" mapstructure:"enabled" yaml:"enabled"` + + // TroveStore is the name of a *trove.Trove registered in the DI + // container. Empty resolves the default (unnamed) instance, which is + // what a single-store Trove extension provides. + TroveStore string `json:"trove_store" mapstructure:"trove_store" yaml:"trove_store"` + + // Bucket is where Dispatch writes the ephemeral artifacts it owns. + Bucket string `default:"dispatch-artifacts" json:"bucket" mapstructure:"bucket" yaml:"bucket"` + + // EphemeralPrefix is the key prefix for Dispatch-owned objects. + EphemeralPrefix string `default:"ephemeral" json:"ephemeral_prefix" mapstructure:"ephemeral_prefix" yaml:"ephemeral_prefix"` + + // Retention is how long an ephemeral artifact survives after every + // owner reaches a terminal state. + Retention time.Duration `default:"168h" json:"retention" mapstructure:"retention" yaml:"retention"` + + // PurgeGrace is how long a soft-deleted artifact's bytes survive + // before the purge pass removes them. It is the window in which a + // mistaken sweep can still be caught. + PurgeGrace time.Duration `default:"24h" json:"purge_grace" mapstructure:"purge_grace" yaml:"purge_grace"` + + // Cache configures the worker-local staging cache. + Cache ArtifactCacheConfig `json:"cache" mapstructure:"cache" yaml:"cache"` +} + +// ArtifactCacheConfig configures the worker-local staging cache. +type ArtifactCacheConfig struct { + // Dir is where staged artifacts are held on local disk. + Dir string `default:"/var/lib/dispatch/cache" json:"dir" mapstructure:"dir" yaml:"dir"` + + // Budget caps the bytes the cache may hold. A job needing more + // staging space than is free waits rather than filling the volume. + // + // With the resource model on this becomes the shared ledger's disk + // capacity rather than a private ceiling inside the cache, so it is + // the same number a job's disk requirement is admitted against. + Budget int64 `json:"budget" mapstructure:"budget" yaml:"budget"` +} + +// ExecutionConfig configures which execution rungs beyond the always- +// present in-process default a deployment makes available. +// +// A job definition declares the isolation it needs with +// job.WithExecution(exec.Isolate(...)); this block decides which rungs +// EXIST to satisfy that declaration — it never changes what any +// definition asks for, and it never lets a declaration that cannot be +// satisfied run anyway. engine.RegisterChecked already refuses a policy +// nothing configured here can satisfy (exec.ErrNoExecutor), unless the +// definition itself opted into exec.AllowDowngrade — that per-definition +// choice is deliberately not something this block can override, because a +// config-wide override would be exactly the silent downgrade +// RegisterChecked exists to prevent. +// +// The whole block is additive and opt-in: a deployment that configures +// none of it registers no extra executor, and every job keeps running +// in-process exactly as it does today. +type ExecutionConfig struct { + // Subprocess configures the out-of-process rung (exec.LevelProcess) — + // the handler runs in a re-exec'd child process instead of the + // worker's own address space. A zero value (Enabled: false, the + // default) registers nothing. + Subprocess SubprocessConfig `json:"subprocess" mapstructure:"subprocess" yaml:"subprocess"` +} + +// SubprocessConfig configures exec/subprocess.Executor. +// +// This rung refuses to launch outside Unix (exec/subprocess's checkLaunch +// and Available); the extension checks Available itself at startup, so +// enabling this on an unsupported platform fails registration once, +// loudly, instead of failing every job's first launch attempt at +// runtime. +type SubprocessConfig struct { + // Enabled registers the subprocess executor with the engine. Without + // it, nothing else in this struct has any effect. + Enabled bool `default:"false" json:"enabled" mapstructure:"enabled" yaml:"enabled"` + + // Binary overrides the path to the binary the executor re-execs for + // every attempt. Empty resolves os.Executable() — the worker's own + // binary — which is correct for every deployment that has not split + // the sandboxed handlers into a separate build. + Binary string `json:"binary" mapstructure:"binary" yaml:"binary"` + + // User and Group are the uid/gid the child process runs as + // (subprocess.WithUser). Both must be set together, or neither: a + // uid with no configured gid is rejected at startup rather than + // silently running the child under the worker's own primary group. + // + // Zero means "not configured" rather than uid/gid 0 — this config + // surface has no way to request running the child as root, which is + // deliberate: it is never the isolation this rung exists to provide. + User int `json:"user" mapstructure:"user" yaml:"user"` + Group int `json:"group" mapstructure:"group" yaml:"group"` + + // AllowSameUser is the single opt-out for running this rung unisolated + // on the uid boundary (subprocess.WithAllowSameUser): it permits User + // to name the worker's own uid, and it permits leaving User unset + // entirely. Without it, either shape refuses at startup — + // resolveExecutionOptions rejects it during Register, before this + // worker ever starts processing jobs — rather than passing cleanly + // and only then failing every attempt's launch, forever, once the + // deployment is already running. That is a deliberate security + // default (see WithAllowSameUser and checkLaunch, which enforces the + // same rule again at Run() for callers that build subprocess.Executor + // directly instead of through this config) that this config surface + // passes through rather than working around: nothing here defaults it + // to true, so a configuration mistake cannot silently defeat it. + AllowSameUser bool `default:"false" json:"allow_same_user" mapstructure:"allow_same_user" yaml:"allow_same_user"` + + // ScratchDir is the root directory both the child process's working + // directory (subprocess.WithScratchDir) and, when the artifact plane + // is also enabled, the Runner's scratch OutputDir + // (engine.WithScratchRoot) are created under. Empty means + // os.TempDir() for both, their own independent defaults. + // + // Configuring this with the artifact plane OFF is not an error: the + // child still gets a scratch working directory, but nothing commits + // its outputs and PriorOutputs stays empty, exactly as + // worker.Runner.WithArtifacts documents for a nil service — the + // extension logs a warning at startup so that is a deliberate choice, + // not a silent one. + ScratchDir string `json:"scratch_dir" mapstructure:"scratch_dir" yaml:"scratch_dir"` + + // Rlimits configures POSIX resource limits applied to the child + // (subprocess.WithRlimits). Fields are in bytes (AddressSpace, FSize, + // Core) or counts (NoFile, NProc); zero leaves that limit at whatever + // the worker itself runs with — except Core, which buildEnv forces + // to zero unconditionally regardless of this value, so a configured + // Core is always ignored (see RlimitsConfig.Core below). There is no + // unit-suffixed string parsing here (no "16GB") — this repo takes no + // new dependency to provide one, and every other byte-valued config + // field (ArtifactCacheConfig.Budget, resource.Set) is already a plain + // integer for the same reason. + Rlimits RlimitsConfig `json:"rlimits" mapstructure:"rlimits" yaml:"rlimits"` + + // StrictRlimits makes a configured rlimit that did not actually take + // effect a launch failure instead of a silently ignored warning (see + // subprocess.WithStrictRlimits). + StrictRlimits bool `default:"false" json:"strict_rlimits" mapstructure:"strict_rlimits" yaml:"strict_rlimits"` +} + +// RlimitsConfig configures the child process's POSIX resource limits. See +// subprocess.Rlimits for what each field does; the field names and units +// here mirror it exactly. +type RlimitsConfig struct { + // AddressSpace caps RLIMIT_AS in bytes. + AddressSpace int64 `json:"address_space" mapstructure:"address_space" yaml:"address_space"` + // NoFile caps RLIMIT_NOFILE, the open file descriptor count. + NoFile int64 `json:"nofile" mapstructure:"nofile" yaml:"nofile"` + // NProc caps RLIMIT_NPROC, the number of processes the child's uid + // may run. + NProc int64 `json:"nproc" mapstructure:"nproc" yaml:"nproc"` + // Core caps RLIMIT_CORE. Accepted for API symmetry; subprocess forces + // the child's actual core limit to zero unconditionally regardless of + // this value — see subprocess.Rlimits.Core. + Core int64 `json:"core" mapstructure:"core" yaml:"core"` + // FSize caps RLIMIT_FSIZE in bytes. + FSize int64 `json:"fsize" mapstructure:"fsize" yaml:"fsize"` +} + +// ResourceConfig configures how this worker's capacity is derived and +// whether jobs are admitted against it at all. +// +// Off by default, and off means the RUNTIME does nothing new: no manager +// is constructed, the pool offers no dequeue budget, every store backend +// skips its fit predicate, and the staging cache keeps the private disk +// budget it has always had. A deployment that does not set this behaves +// at runtime as it did before the resource model existed. +// +// Two things happen regardless, and neither is under this flag. The +// resource columns are added by migration on every upgrade, because a +// job row has to be readable by every worker in a mixed-version fleet; +// on Postgres that includes an index build, which is why it is done +// CONCURRENTLY. And every enqueue writes those columns, at their zero +// values when nothing declares a requirement. +type ResourceConfig struct { + // Enabled turns on capacity detection and resource-aware admission. + Enabled bool `default:"false" json:"enabled" mapstructure:"enabled" yaml:"enabled"` + + // CPUOvercommit multiplies the detected core count. CPU is + // compressible — exceeding it makes jobs slow rather than dead — so + // values above 1.0 are a legitimate throughput trade. Zero means 1.0. + // + // There is deliberately no memory equivalent. Overcommitting memory + // is how a box enters the OOM cascade this model exists to prevent. + CPUOvercommit float64 `default:"1.0" json:"cpu_overcommit" mapstructure:"cpu_overcommit" yaml:"cpu_overcommit"` + + // MemoryFraction is the share of the detected memory limit to + // advertise, leaving the rest for the Go runtime, the page cache, and + // everything else on the box. Zero means 0.8. + MemoryFraction float64 `default:"0.8" json:"memory_fraction" mapstructure:"memory_fraction" yaml:"memory_fraction"` + + // Explicit overrides detection per key, and is the ONLY way to + // declare a custom resource: there is no detection for "fpga". + // Quantities are canonical units — cpu in millicores, memory and disk + // in bytes, gpu in milli-devices, custom keys in whatever integer the + // declaring job means by them. + Explicit resource.Set `json:"explicit" mapstructure:"explicit" yaml:"explicit"` + + // CustomKeys narrows the custom resource keys this worker advertises + // at dequeue. Empty advertises every custom key in the detected + // capacity, which is the honest default. + CustomKeys []string `json:"custom_keys" mapstructure:"custom_keys" yaml:"custom_keys"` +} diff --git a/extension/config_internal_test.go b/extension/config_internal_test.go new file mode 100644 index 0000000..c177b29 --- /dev/null +++ b/extension/config_internal_test.go @@ -0,0 +1,244 @@ +package extension + +import ( + "reflect" + "testing" + + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/store/memory" +) + +// TestResourceConfigYAMLShape pins the keys an operator writes. They are +// the published interface — a typo in a struct tag silently produces a +// worker running on detected defaults while the config file says +// otherwise, which is the failure mode this whole task is about. +func TestResourceConfigYAMLShape(t *testing.T) { + want := map[string]string{ + "Enabled": "enabled", + "CPUOvercommit": "cpu_overcommit", + "MemoryFraction": "memory_fraction", + "Explicit": "explicit", + "CustomKeys": "custom_keys", + } + + rt := reflect.TypeOf(ResourceConfig{}) + + for i := range rt.NumField() { + f := rt.Field(i) + + key, known := want[f.Name] + if !known { + t.Errorf("field %s has no expected config key; update this test", f.Name) + + continue + } + + for _, tag := range []string{"yaml", "mapstructure", "json"} { + if got := f.Tag.Get(tag); got != key { + t.Errorf("%s: %s tag = %q, want %q", f.Name, tag, got, key) + } + } + } + + if got := reflect.TypeOf(Config{}).Field(fieldIndex(t, Config{}, "Resources")).Tag.Get("yaml"); got != "resources" { + t.Errorf("Config.Resources yaml tag = %q, want %q", got, "resources") + } +} + +func fieldIndex(t *testing.T, v any, name string) int { + t.Helper() + + f, ok := reflect.TypeOf(v).FieldByName(name) + if !ok { + t.Fatalf("no field %q", name) + } + + return f.Index[0] +} + +// TestMergeResourceConfig covers the precedence rules: YAML wins where it +// spoke, programmatic options fill the gaps, and enabling is an OR so a +// binary built with WithResources cannot be silently switched off by a +// config file that simply does not mention it. +func TestMergeResourceConfig(t *testing.T) { + t.Run("programmatic enable survives silent yaml", func(t *testing.T) { + got := mergeResourceConfig(ResourceConfig{}, ResourceConfig{Enabled: true}) + if !got.Enabled { + t.Error("WithResources was dropped by a config file that said nothing") + } + }) + + t.Run("yaml wins on scalars", func(t *testing.T) { + got := mergeResourceConfig( + ResourceConfig{CPUOvercommit: 2, MemoryFraction: 0.5}, + ResourceConfig{CPUOvercommit: 4, MemoryFraction: 0.9}, + ) + + if got.CPUOvercommit != 2 || got.MemoryFraction != 0.5 { + t.Errorf("programmatic values overrode yaml: %+v", got) + } + }) + + t.Run("explicit capacity merges per key", func(t *testing.T) { + got := mergeResourceConfig( + ResourceConfig{Explicit: resource.Set{resource.Memory: 1 << 30}}, + ResourceConfig{Explicit: resource.Set{"fpga": 2, resource.Memory: 8 << 30}}, + ) + + // The binary knows how many FPGAs it was built to talk to; the + // operator knows how much memory to hand this pod. Neither erases + // the other, and on a key they both set, the file wins. + if got.Explicit["fpga"] != 2 { + t.Errorf("programmatic custom key lost: %v", got.Explicit) + } + + if got.Explicit[resource.Memory] != 1<<30 { + t.Errorf("yaml did not win on memory: %v", got.Explicit) + } + }) +} + +// TestWithWorkerCustomKeysMergesAndCopies keeps the option consistent +// with WithExplicitCapacity beside it and engine.WithWorkerCustomKeys +// below it: keys accumulate, duplicates collapse, and the caller's slice +// is not retained. +func TestWithWorkerCustomKeysMergesAndCopies(t *testing.T) { + caller := []string{"fpga", "tpu"} + + e := New( + WithWorkerCustomKeys(caller...), + WithWorkerCustomKeys("fpga", "npu"), + ) + + want := []string{"fpga", "tpu", "npu"} + + got := e.config.Resources.CustomKeys + if len(got) != len(want) { + t.Fatalf("CustomKeys = %v, want %v", got, want) + } + + for i := range want { + if got[i] != want[i] { + t.Fatalf("CustomKeys = %v, want %v", got, want) + } + } + + // The caller keeps no handle on extension state. + caller[0] = "mutated" + + if e.config.Resources.CustomKeys[0] != "fpga" { + t.Errorf("the caller's slice is aliased: %v", e.config.Resources.CustomKeys) + } +} + +// TestStagingBudgetRouting pins where the cache budget ends up. +// +// cache.WithBudget is ignored once a manager is supplied, so if the +// configured number does not arrive here it does not arrive anywhere, and +// an operator who wrote a budget gets whatever Detect chose instead. +func TestStagingBudgetRouting(t *testing.T) { + cases := []struct { + name string + artifact ArtifactConfig + noStore bool + want int64 + }{ + { + name: "artifacts off omits disk entirely", + want: 0, + }, + { + // init only builds the plane when the dispatcher's store + // implements artifact.Store. Configured-but-not-built has to + // omit disk too, or the ledger advertises 20 GiB with no cache + // behind it and no reclaimer registered for it. + name: "artifacts configured but no artifact store", + artifact: ArtifactConfig{Enabled: true}, + noStore: true, + want: 0, + }, + { + name: "configured budget", + artifact: ArtifactConfig{Enabled: true, Cache: ArtifactCacheConfig{Budget: 200 << 30}}, + want: 200 << 30, + }, + { + name: "unset budget falls back to the cache default", + artifact: ArtifactConfig{Enabled: true}, + want: cache.DefaultBudget, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + e := &Extension{config: Config{Artifacts: tc.artifact}} + if !tc.noStore { + e.artifactStore = memory.New() + } + + if got := e.stagingBudget(); got != tc.want { + t.Errorf("stagingBudget() = %d, want %d", got, tc.want) + } + }) + } +} + +// TestDiskOverrideWarning pins when the operator gets told that two +// config keys are setting the same number. +// +// resources.explicit.disk is not just what the worker advertises: the +// cache reads its eviction ceiling off the ledger's disk capacity, so an +// explicit value is what the cache WRITES against. Above the volume that +// is ENOSPC, and before this it could only be reached by the cache budget +// key that names itself. +func TestDiskOverrideWarning(t *testing.T) { + cases := []struct { + name string + explicit resource.Set + staging int64 + want bool + }{ + { + name: "both set and disagreeing", + explicit: resource.Set{resource.Disk: 500 << 30}, + staging: 200 << 30, + want: true, + }, + { + name: "both set and agreeing", + explicit: resource.Set{resource.Disk: 200 << 30}, + staging: 200 << 30, + }, + { + name: "only the cache budget", + staging: 200 << 30, + }, + { + // No staging cache to disagree with: explicit disk is then the + // only source there is, which is not an override of anything. + name: "only explicit, no cache", + explicit: resource.Set{resource.Disk: 500 << 30}, + }, + { + name: "an unrelated explicit key", + explicit: resource.Set{resource.Memory: 8 << 30}, + staging: 200 << 30, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, conflict := diskOverride(tc.explicit, tc.staging) + + if conflict != tc.want { + t.Fatalf("diskOverride() conflict = %v, want %v", conflict, tc.want) + } + + if conflict && got != tc.explicit[resource.Disk] { + t.Errorf("reported explicit disk = %d, want %d", + got, tc.explicit[resource.Disk]) + } + }) + } +} diff --git a/extension/execution.go b/extension/execution.go new file mode 100644 index 0000000..b0f3fe0 --- /dev/null +++ b/extension/execution.go @@ -0,0 +1,232 @@ +package extension + +import ( + "errors" + "fmt" + "os" + + "github.com/xraph/dispatch/engine" + "github.com/xraph/dispatch/exec/subprocess" +) + +// mergeExecutionConfig folds programmatic execution settings into what +// YAML supplied, following the same precedence rule mergeResourceConfig +// documents: YAML wins where it said something, programmatic options +// fill the gaps, and every enable-shaped flag is an OR rather than an +// override. Without this, a binary that called WithConfig to turn on +// the subprocess rung — the one thing standing between a malicious +// upload and the worker's own credentials — had that silently discarded +// the moment ANY YAML "dispatch" key existed, because loadConfiguration +// replaces e.config wholesale with mergeConfigurations' result and this +// block was the one Config field that function never touched. The +// deployment would then run every job in-process, believing it was +// sandboxed, with no error or warning anywhere. +func mergeExecutionConfig(yamlCfg, programmatic ExecutionConfig) ExecutionConfig { + yamlCfg.Subprocess = mergeSubprocessConfig(yamlCfg.Subprocess, programmatic.Subprocess) + + return yamlCfg +} + +// mergeSubprocessConfig applies mergeExecutionConfig's precedence rule +// field by field. +// +// User and Group are merged as a pair, not independently: buildSubprocessOptions +// already refuses a config that sets one without the other, so filling +// them from different sources here could silently manufacture exactly +// that invalid combination. YAML naming either one counts as YAML having +// spoken on the pair; only when YAML sets neither does the programmatic +// pair fill the gap. +func mergeSubprocessConfig(yamlCfg, programmatic SubprocessConfig) SubprocessConfig { + if programmatic.Enabled { + yamlCfg.Enabled = true + } + + if yamlCfg.Binary == "" && programmatic.Binary != "" { + yamlCfg.Binary = programmatic.Binary + } + + if yamlCfg.User == 0 && yamlCfg.Group == 0 && (programmatic.User != 0 || programmatic.Group != 0) { + yamlCfg.User = programmatic.User + yamlCfg.Group = programmatic.Group + } + + if programmatic.AllowSameUser { + yamlCfg.AllowSameUser = true + } + + if yamlCfg.ScratchDir == "" && programmatic.ScratchDir != "" { + yamlCfg.ScratchDir = programmatic.ScratchDir + } + + if yamlCfg.Rlimits == (RlimitsConfig{}) && programmatic.Rlimits != (RlimitsConfig{}) { + yamlCfg.Rlimits = programmatic.Rlimits + } + + if programmatic.StrictRlimits { + yamlCfg.StrictRlimits = true + } + + return yamlCfg +} + +// resolveExecutionOptions turns the execution config block into engine +// options that register additional isolation rungs beyond the always- +// present in-process default. +// +// This mirrors resolveArtifactBackend's shape: nothing configured returns +// no options and no error, so a deployment that never asks for out-of- +// process isolation runs exactly as it did before this existed. Unlike +// resolveArtifactBackend, a request this deployment or this platform +// cannot actually satisfy is refused HERE, at startup — not deferred to +// engine.RegisterChecked finding no executor for a policy, and not +// deferred further still to a job's first launch failure. Failing early +// and once is the whole point: the alternative is a job that requeues +// itself every poll interval forever, discovering the same +// misconfiguration on every attempt. +func (e *Extension) resolveExecutionOptions() ([]engine.Option, error) { + cfg := e.config.Execution.Subprocess + + if !cfg.Enabled { + if cfg.ScratchDir != "" && e.Logger() != nil { + e.Logger().Warn("dispatch: execution.subprocess.scratch_dir is set but " + + "execution.subprocess.enabled is false; the value has no effect") + } + + return nil, nil + } + + // Available only answers a platform question — can this OS run the + // rung at all — not a per-configuration one. It does NOT catch a + // configured uid equal to the worker's own; that is checkLaunch's own + // refusal, and on its own it would not run until the first job + // actually launches. Caught once, at startup, instead of on every + // job's first launch attempt once the deployment is already running. + if err := subprocess.Available(); err != nil { + return nil, fmt.Errorf("dispatch: execution.subprocess is enabled but %w", err) + } + + // The uid question, asked here through the identical helpers + // checkLaunch itself uses (subprocess.SameUserRefused, and hasUser via + // cfg.User == 0), so this check and checkLaunch's cannot drift apart + // silently. Without it, a config naming the worker's own uid — or + // naming none at all — passes startup clean and then fails to launch + // every single job forever — requeued, reaped off its expired lease, + // and re-leased, until it exhausts the launch-attempt cap and DLQs, + // which happens per job, in production, long after this function + // returned nil. + // + // Both shapes — no uid configured, and a configured uid equal to the + // worker's own — are refused unless allow_same_user opts in: one + // switch for "I know this is unisolated," not a warning for one shape + // and a hard refusal for the other. See checkLaunch + // (exec/subprocess/limits_unix.go) for why they are treated alike. + if cfg.User == 0 { + if !cfg.AllowSameUser { + return nil, errors.New( + "dispatch: execution.subprocess is enabled with no user configured; the sandboxed " + + "child would run as the worker's own uid, with the worker's own read access to " + + "its credentials and filesystem — set execution.subprocess.user, or " + + "execution.subprocess.allow_same_user to accept running unisolated") + } + } else if subprocess.SameUserRefused(cfg.User, cfg.AllowSameUser) { + return nil, fmt.Errorf( + "dispatch: execution.subprocess.user %d matches the worker's own uid; "+ + "running the child as the worker defeats this rung's isolation — set "+ + "execution.subprocess.allow_same_user to allow it", cfg.User) + } + + opts, err := e.buildSubprocessOptions(cfg) + if err != nil { + return nil, err + } + + engOpts := []engine.Option{engine.WithExecutor(subprocess.New(opts...))} + + // ScratchDir only does anything once the artifact plane is also + // configured — see WithScratchRoot and the ScratchDir field's own + // doc comment. Setting it unconditionally here is still correct: the + // warning above already told the operator when it will be a no-op, + // and engine.Build itself only reads eng.scratchRoot when the + // artifact plane resolved to a non-nil service. + if cfg.ScratchDir != "" { + engOpts = append(engOpts, engine.WithScratchRoot(cfg.ScratchDir)) + + if e.artifacts == nil && e.Logger() != nil { + e.Logger().Warn("dispatch: execution.subprocess.scratch_dir is configured but " + + "the artifact plane is disabled; the out-of-process rung still gets a " + + "scratch working directory, but nothing it writes is committed and " + + "PriorOutputs stays empty — see artifacts.enabled") + } + } + + return engOpts, nil +} + +// buildSubprocessOptions translates a SubprocessConfig into the +// subprocess.Option values its Executor is constructed from. +func (e *Extension) buildSubprocessOptions(cfg SubprocessConfig) ([]subprocess.Option, error) { + binary := cfg.Binary + if binary == "" { + resolved, err := os.Executable() + if err != nil { + return nil, fmt.Errorf("dispatch: resolve the worker's own binary for the subprocess rung: %w", err) + } + binary = resolved + } + + opts := []subprocess.Option{ + subprocess.WithBinary(binary), + } + + // Only when a logger is actually available: subprocess.New defaults + // to a safe no-op logger on its own, and passing through a nil + // interface here — e.Logger() before Register runs, or an app whose + // own Logger() returns nil — would silently replace that default + // with a nil Logger that panics the moment the child's stdout or + // stderr is streamed through it. + if logger := e.Logger(); logger != nil { + opts = append(opts, subprocess.WithLogger(logger)) + } + + // WithUser takes uid and gid together; a config declaring one without + // the other is ambiguous about what it wants and is refused rather + // than guessed at — silently running the child under the worker's + // own primary group while only the uid was pinned would quietly + // weaken the boundary the operator thought they configured. + if (cfg.User == 0) != (cfg.Group == 0) { + return nil, errors.New( + "dispatch: execution.subprocess.user and execution.subprocess.group must both be set, or neither") + } + if cfg.User != 0 { + opts = append(opts, subprocess.WithUser(cfg.User, cfg.Group)) + } + + // Deliberately unconditional on cfg.AllowSameUser alone: nothing in + // this function ever adds WithAllowSameUser on the operator's behalf, + // so a config that leaves it false keeps subprocess.checkLaunch's own + // refusal — the worker's own uid is not an acceptable default here — + // fully in force. + if cfg.AllowSameUser { + opts = append(opts, subprocess.WithAllowSameUser()) + } + + if cfg.ScratchDir != "" { + opts = append(opts, subprocess.WithScratchDir(cfg.ScratchDir)) + } + + if cfg.Rlimits != (RlimitsConfig{}) { + opts = append(opts, subprocess.WithRlimits(subprocess.Rlimits{ + AddressSpace: cfg.Rlimits.AddressSpace, + NoFile: cfg.Rlimits.NoFile, + NProc: cfg.Rlimits.NProc, + Core: cfg.Rlimits.Core, + FSize: cfg.Rlimits.FSize, + })) + } + + if cfg.StrictRlimits { + opts = append(opts, subprocess.WithStrictRlimits()) + } + + return opts, nil +} diff --git a/extension/execution_internal_test.go b/extension/execution_internal_test.go new file mode 100644 index 0000000..8e3d237 --- /dev/null +++ b/extension/execution_internal_test.go @@ -0,0 +1,367 @@ +package extension + +import ( + "context" + "os" + "reflect" + "testing" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/subprocess" + "github.com/xraph/dispatch/id" +) + +// TestResolveExecutionOptionsDisabledByDefault pins the backward- +// compatibility guarantee: a Config that never mentions execution +// registers no extra executor at all. +func TestResolveExecutionOptionsDisabledByDefault(t *testing.T) { + e := New() + + opts, err := e.resolveExecutionOptions() + if err != nil { + t.Fatalf("resolveExecutionOptions() = %v, want nil", err) + } + if len(opts) != 0 { + t.Fatalf("resolveExecutionOptions() = %d options, want 0", len(opts)) + } +} + +// TestResolveExecutionOptionsEnablesSubprocess checks that enabling the +// block produces exactly one engine.Option (WithExecutor; no scratch dir +// configured here to add a second). The actual registration behaviour — +// that this is what lets exec.Isolate(exec.LevelProcess) satisfy — is +// covered end to end in engine/execution_subprocess_test.go, which +// exercises engine.WithExecutor directly and needs no Forge app. +func TestResolveExecutionOptionsEnablesSubprocess(t *testing.T) { + e := New() + e.config.Execution.Subprocess.Enabled = true + e.config.Execution.Subprocess.AllowSameUser = true // no user configured; see TestResolveExecutionOptionsRefusesMissingUserAtStartup + + opts, err := e.resolveExecutionOptions() + if err != nil { + t.Fatalf("resolveExecutionOptions() = %v, want nil", err) + } + if len(opts) != 1 { + t.Fatalf("resolveExecutionOptions() = %d options, want 1", len(opts)) + } +} + +// TestResolveExecutionOptionsScratchDirAddsSecondOption checks that a +// configured scratch_dir produces the extra engine.WithScratchRoot +// option alongside WithExecutor. +func TestResolveExecutionOptionsScratchDirAddsSecondOption(t *testing.T) { + e := New() + e.config.Execution.Subprocess.Enabled = true + e.config.Execution.Subprocess.ScratchDir = t.TempDir() + e.config.Execution.Subprocess.AllowSameUser = true // no user configured; see TestResolveExecutionOptionsRefusesMissingUserAtStartup + + opts, err := e.resolveExecutionOptions() + if err != nil { + t.Fatalf("resolveExecutionOptions() = %v, want nil", err) + } + if len(opts) != 2 { + t.Fatalf("resolveExecutionOptions() = %d options, want 2 (WithExecutor + WithScratchRoot)", len(opts)) + } +} + +// TestResolveExecutionOptionsRefusesSameUserAtStartup proves the +// resolveExecutionOptions fix: a config naming the worker's own uid, with +// no allow_same_user, must fail HERE, at startup — not pass cleanly and +// then fail launch on every single job, forever, discovered only in +// production. subprocess.Available alone cannot catch this (it is a +// platform check, not a per-config one); this exercises the separate +// same-uid check added alongside it. +func TestResolveExecutionOptionsRefusesSameUserAtStartup(t *testing.T) { + e := New() + e.config.Execution.Subprocess.Enabled = true + e.config.Execution.Subprocess.User = os.Getuid() + e.config.Execution.Subprocess.Group = os.Getgid() + + _, err := e.resolveExecutionOptions() + if err == nil { + t.Fatal("resolveExecutionOptions() = nil error, want one for a uid matching the worker's own") + } +} + +// TestResolveExecutionOptionsAllowsSameUserExplicitly proves the same-uid +// startup check does not fire when allow_same_user opted in — the +// deliberate escape hatch stays available, exactly as it does at Run() +// time in checkLaunch. +func TestResolveExecutionOptionsAllowsSameUserExplicitly(t *testing.T) { + e := New() + e.config.Execution.Subprocess.Enabled = true + e.config.Execution.Subprocess.User = os.Getuid() + e.config.Execution.Subprocess.Group = os.Getgid() + e.config.Execution.Subprocess.AllowSameUser = true + + if _, err := e.resolveExecutionOptions(); err != nil { + t.Fatalf("resolveExecutionOptions() = %v, want nil", err) + } +} + +// TestResolveExecutionOptionsRefusesMissingUserAtStartup proves enabling +// the rung with no uid configured — the minimal `execution: {subprocess: +// {enabled: true}}` — refuses to start HERE, at startup, rather than +// warning and letting every job silently run unisolated. This used to be a +// WARN, on the theory that leaving `user` unset was a legitimate if weak +// choice; it is treated the same as a configured uid equal to the +// worker's own instead, because both leave the child with full read +// access to the worker's credentials and filesystem — see checkLaunch +// (exec/subprocess/limits_unix.go) and doc.go's "uid/gid boundary" +// section. +func TestResolveExecutionOptionsRefusesMissingUserAtStartup(t *testing.T) { + e := New() + e.config.Execution.Subprocess.Enabled = true + + _, err := e.resolveExecutionOptions() + if err == nil { + t.Fatal("resolveExecutionOptions() = nil error, want one for no uid configured") + } +} + +// TestResolveExecutionOptionsAllowsMissingUserExplicitly is the mirror of +// TestResolveExecutionOptionsAllowsSameUserExplicitly for the other shape +// checkLaunch refuses: allow_same_user lifts the no-uid refusal too, since +// it is the single escape hatch for both. +func TestResolveExecutionOptionsAllowsMissingUserExplicitly(t *testing.T) { + e := New() + e.config.Execution.Subprocess.Enabled = true + e.config.Execution.Subprocess.AllowSameUser = true + + if _, err := e.resolveExecutionOptions(); err != nil { + t.Fatalf("resolveExecutionOptions() = %v, want nil", err) + } +} + +// TestResolveExecutionOptionsSucceedsWhenUserConfigured is the negative +// case for the above: a properly configured uid must not be refused. +func TestResolveExecutionOptionsSucceedsWhenUserConfigured(t *testing.T) { + e := New() + e.config.Execution.Subprocess.Enabled = true + e.config.Execution.Subprocess.User = 65532 + e.config.Execution.Subprocess.Group = 65532 + + if _, err := e.resolveExecutionOptions(); err != nil { + t.Fatalf("resolveExecutionOptions() = %v, want nil", err) + } +} + +// TestBuildSubprocessOptionsRejectsLopsidedUserGroup pins the guard that +// keeps a config from silently running the child under the worker's own +// primary group when only a uid was configured. +func TestBuildSubprocessOptionsRejectsLopsidedUserGroup(t *testing.T) { + e := New() + + if _, err := e.buildSubprocessOptions(SubprocessConfig{User: 65532}); err == nil { + t.Fatal("buildSubprocessOptions() = nil error, want one for a uid with no configured gid") + } + if _, err := e.buildSubprocessOptions(SubprocessConfig{Group: 65532}); err == nil { + t.Fatal("buildSubprocessOptions() = nil error, want one for a gid with no configured uid") + } +} + +// TestBuildSubprocessOptionsNeverDefaultsAllowSameUser proves config +// cannot accidentally defeat WithUser's own same-uid refusal: nothing in +// buildSubprocessOptions adds WithAllowSameUser unless the config +// explicitly asked for it, so an executor built from a config that names +// the worker's own uid must still refuse to launch. +// +// checkLaunch (exec/subprocess/limits_unix.go) runs before any pipe or +// process is created, so this reaches the refusal without actually +// spawning anything — the launch failure comes back through +// Result.Status, not a returned error, exactly like every other launch +// failure this rung reports. +func TestBuildSubprocessOptionsNeverDefaultsAllowSameUser(t *testing.T) { + e := New() + + uid := os.Getuid() + opts, err := e.buildSubprocessOptions(SubprocessConfig{User: uid, Group: os.Getgid()}) + if err != nil { + t.Fatalf("buildSubprocessOptions() = %v, want nil", err) + } + + ex := subprocess.New(opts...) + + res, runErr := ex.Run(context.Background(), &exec.Request{ + JobID: id.NewJobID(), + Name: "test.subprocess.same-uid", + Payload: []byte("{}"), + OutputDir: t.TempDir(), + Policy: exec.NewPolicy(), + }) + if runErr != nil { + t.Fatalf("Run() = %v, want nil (a launch refusal reports through Result, not error)", runErr) + } + if res.Status != exec.StatusLaunchFailed { + t.Fatalf("Result.Status = %q, want %q — the same-uid refusal must survive config translation", + res.Status, exec.StatusLaunchFailed) + } +} + +// TestConfigExecutionYAMLShape pins the YAML/mapstructure/json keys an +// operator writes for execution.subprocess — see +// TestResourceConfigYAMLShape in config_internal_test.go for why this +// matters: a struct-tag typo silently accepts a config key that does +// nothing. +func TestConfigExecutionYAMLShape(t *testing.T) { + want := map[string]string{ + "Enabled": "enabled", + "Binary": "binary", + "User": "user", + "Group": "group", + "AllowSameUser": "allow_same_user", + "ScratchDir": "scratch_dir", + "Rlimits": "rlimits", + "StrictRlimits": "strict_rlimits", + } + + rt := reflect.TypeOf(SubprocessConfig{}) + for i := range rt.NumField() { + f := rt.Field(i) + + key, known := want[f.Name] + if !known { + t.Errorf("field %s has no expected config key; update this test", f.Name) + continue + } + + for _, tag := range []string{"yaml", "mapstructure", "json"} { + if got := f.Tag.Get(tag); got != key { + t.Errorf("%s: %s tag = %q, want %q", f.Name, tag, got, key) + } + } + } + + if got := reflect.TypeOf(Config{}).Field(fieldIndex(t, Config{}, "Execution")).Tag.Get("yaml"); got != "execution" { + t.Errorf("Config.Execution yaml tag = %q, want %q", got, "execution") + } +} + +// TestMergeConfigurationsKeepsProgrammaticExecution is the regression +// test for the bug mergeConfigurations had: it merged Artifacts and +// Resources but never touched Execution at all, so a binary that called +// WithConfig(Config{Execution: ...}) to turn on the subprocess rung had +// that silently dropped the instant loadConfiguration found ANY YAML +// "dispatch" key — even one that says nothing whatsoever about +// execution. The deployment would then run every job in-process, +// believing it was sandboxed, with no error or warning anywhere. +// +// Reproduces the exact shape from the report: YAML sets an unrelated +// field (BasePath) and says nothing about execution at all. +func TestMergeConfigurationsKeepsProgrammaticExecution(t *testing.T) { + e := New() + + yamlConfig := Config{BasePath: "/custom"} + programmaticConfig := Config{ + Execution: ExecutionConfig{ + Subprocess: SubprocessConfig{ + Enabled: true, + User: 65532, + Group: 65532, + }, + }, + } + + got := e.mergeConfigurations(yamlConfig, programmaticConfig) + + if !got.Execution.Subprocess.Enabled { + t.Error("WithConfig's execution.subprocess.enabled was dropped by a config file that said nothing about execution") + } + if got.Execution.Subprocess.User != 65532 || got.Execution.Subprocess.Group != 65532 { + t.Errorf("Execution.Subprocess.User/Group = %d/%d, want 65532/65532", + got.Execution.Subprocess.User, got.Execution.Subprocess.Group) + } +} + +// TestMergeExecutionConfig covers mergeExecutionConfig's precedence +// rules directly, mirroring TestMergeResourceConfig in +// config_internal_test.go. +func TestMergeExecutionConfig(t *testing.T) { + t.Run("programmatic enable survives silent yaml", func(t *testing.T) { + got := mergeExecutionConfig( + ExecutionConfig{}, + ExecutionConfig{Subprocess: SubprocessConfig{Enabled: true}}, + ) + if !got.Subprocess.Enabled { + t.Error("enabled was dropped by a config file that said nothing") + } + }) + + t.Run("yaml enabled survives silent programmatic config", func(t *testing.T) { + got := mergeExecutionConfig( + ExecutionConfig{Subprocess: SubprocessConfig{Enabled: true}}, + ExecutionConfig{}, + ) + if !got.Subprocess.Enabled { + t.Error("yaml's enabled was lost") + } + }) + + t.Run("yaml wins on scalars", func(t *testing.T) { + got := mergeExecutionConfig( + ExecutionConfig{Subprocess: SubprocessConfig{Binary: "/yaml/bin", ScratchDir: "/yaml/scratch"}}, + ExecutionConfig{Subprocess: SubprocessConfig{Binary: "/programmatic/bin", ScratchDir: "/programmatic/scratch"}}, + ) + if got.Subprocess.Binary != "/yaml/bin" { + t.Errorf("Binary = %q, want yaml value", got.Subprocess.Binary) + } + if got.Subprocess.ScratchDir != "/yaml/scratch" { + t.Errorf("ScratchDir = %q, want yaml value", got.Subprocess.ScratchDir) + } + }) + + t.Run("programmatic fills gaps left by yaml", func(t *testing.T) { + got := mergeExecutionConfig( + ExecutionConfig{}, + ExecutionConfig{Subprocess: SubprocessConfig{Binary: "/programmatic/bin", ScratchDir: "/programmatic/scratch"}}, + ) + if got.Subprocess.Binary != "/programmatic/bin" { + t.Errorf("Binary = %q, want programmatic value", got.Subprocess.Binary) + } + if got.Subprocess.ScratchDir != "/programmatic/scratch" { + t.Errorf("ScratchDir = %q, want programmatic value", got.Subprocess.ScratchDir) + } + }) + + t.Run("user and group merge as a pair, never independently", func(t *testing.T) { + got := mergeExecutionConfig( + ExecutionConfig{Subprocess: SubprocessConfig{User: 100, Group: 100}}, + ExecutionConfig{Subprocess: SubprocessConfig{User: 200, Group: 200}}, + ) + if got.Subprocess.User != 100 || got.Subprocess.Group != 100 { + t.Errorf("User/Group = %d/%d, want yaml's 100/100", got.Subprocess.User, got.Subprocess.Group) + } + + got = mergeExecutionConfig( + ExecutionConfig{}, + ExecutionConfig{Subprocess: SubprocessConfig{User: 200, Group: 200}}, + ) + if got.Subprocess.User != 200 || got.Subprocess.Group != 200 { + t.Errorf("User/Group = %d/%d, want programmatic's 200/200 filling an empty yaml pair", + got.Subprocess.User, got.Subprocess.Group) + } + }) + + t.Run("allow_same_user and strict_rlimits are an OR", func(t *testing.T) { + got := mergeExecutionConfig( + ExecutionConfig{}, + ExecutionConfig{Subprocess: SubprocessConfig{AllowSameUser: true, StrictRlimits: true}}, + ) + if !got.Subprocess.AllowSameUser { + t.Error("programmatic AllowSameUser was dropped") + } + if !got.Subprocess.StrictRlimits { + t.Error("programmatic StrictRlimits was dropped") + } + }) + + t.Run("rlimits: yaml wins wholesale when it set any field", func(t *testing.T) { + got := mergeExecutionConfig( + ExecutionConfig{Subprocess: SubprocessConfig{Rlimits: RlimitsConfig{NoFile: 10}}}, + ExecutionConfig{Subprocess: SubprocessConfig{Rlimits: RlimitsConfig{NoFile: 20, NProc: 5}}}, + ) + if got.Subprocess.Rlimits != (RlimitsConfig{NoFile: 10}) { + t.Errorf("Rlimits = %+v, want yaml's {NoFile: 10}", got.Subprocess.Rlimits) + } + }) +} diff --git a/extension/extension.go b/extension/extension.go index 8b28bda..0c3b0ea 100644 --- a/extension/extension.go +++ b/extension/extension.go @@ -26,12 +26,16 @@ import ( "github.com/xraph/dispatch" "github.com/xraph/dispatch/api" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/artifact/sweeper" "github.com/xraph/dispatch/backoff" dispatchdash "github.com/xraph/dispatch/dashboard" "github.com/xraph/dispatch/dwp" "github.com/xraph/dispatch/engine" "github.com/xraph/dispatch/ext" mw "github.com/xraph/dispatch/middleware" + "github.com/xraph/dispatch/resource" mongostore "github.com/xraph/dispatch/store/mongo" pgstore "github.com/xraph/dispatch/store/postgres" redisstore "github.com/xraph/dispatch/store/redis" @@ -71,6 +75,19 @@ type Extension struct { useGrove bool useGroveKV bool enableDWP bool + + // artifactBackend is an explicitly supplied backend, taking priority + // over anything discovered in the container. + artifactBackend artifact.Backend + // artifactStore is the store the artifact service persists through. + artifactStore artifact.Store + artifacts *artifact.Service + artifactCache *cache.Cache + sweeper *sweeper.Sweeper + + // resources is the single admission ledger shared by the staging + // cache and the worker pool. Nil when the resource model is off. + resources resource.Manager } // New creates a Dispatch Forge extension with the given options. @@ -91,6 +108,17 @@ func (e *Extension) Engine() *engine.Engine { return e.eng } // API returns the API handler. func (e *Extension) API() *api.API { return e.apiHandler } +// Artifacts returns the artifact service, or nil when the plane is off. +func (e *Extension) Artifacts() *artifact.Service { return e.artifacts } + +// ArtifactCache returns the staging cache, or nil when the plane is off. +func (e *Extension) ArtifactCache() *cache.Cache { return e.artifactCache } + +// Resources returns the shared admission ledger, or nil when the +// resource model is off. It is the same instance the staging cache and +// the worker pool were built with — which is the point of exposing it. +func (e *Extension) Resources() resource.Manager { return e.resources } + // DWPServer returns the DWP server, or nil if DWP is not enabled. func (e *Extension) DWPServer() *dwp.Server { return e.dwpServer } @@ -187,6 +215,59 @@ func (e *Extension) init(fapp forge.App) error { engOpts = append(engOpts, engine.WithStreamBroker()) } + // Resolve the artifact store before the ledger, not after. The + // ledger's disk capacity is the staging budget, and there is no + // staging cache unless this resolves — a store that does not + // implement artifact.Store turns `artifacts.enabled: true` into a + // plane that is configured and never built. + if e.artifactStore == nil { + if as, ok := d.Store().(artifact.Store); ok { + e.artifactStore = as + } + } + + // The admission ledger is built next, because both of the things that + // follow have to be given the SAME instance: the staging cache holds + // a lease per cached entry and registers itself as the ledger's disk + // reclaimer, and the worker pool offers disk at dequeue as free PLUS + // what that reclaimer could evict. Two managers and the second half + // of that budget is permanently zero. + e.resources = e.buildResourceManager() + + // The artifact plane is built before the engine, because the staging + // middleware has to be in the chain the engine constructs. + if e.artifactStore != nil { + svc, artCache, aerr := e.buildArtifactPlane(fapp, e.resources) + if aerr != nil { + return aerr + } + + if svc != nil { + e.artifacts = svc + e.artifactCache = artCache + engOpts = append(engOpts, engine.WithArtifacts(svc, artCache)) + } + } + + // Execution rungs beyond the in-process default are resolved after + // the artifact plane, not before: resolveExecutionOptions checks + // e.artifacts to decide whether a configured scratch directory will + // actually be read (see its own comment), and that has to reflect + // the artifact plane's FINAL state, not a guess made ahead of it. + execOpts, execErr := e.resolveExecutionOptions() + if execErr != nil { + return execErr + } + engOpts = append(engOpts, execOpts...) + + if e.resources != nil { + engOpts = append(engOpts, engine.WithResourceManager(e.resources)) + + if keys := e.config.Resources.CustomKeys; len(keys) > 0 { + engOpts = append(engOpts, engine.WithWorkerCustomKeys(keys)) + } + } + e.eng, err = engine.Build(d, engOpts...) if err != nil { return fmt.Errorf("dispatch: build engine: %w", err) @@ -275,16 +356,77 @@ func (e *Extension) Start(ctx context.Context) error { return err } + e.startSweeper(ctx) + e.MarkStarted() return nil } +// startSweeper begins reclaiming Dispatch-owned storage. +// +// It runs on the elected leader only, so a fleet does not race to delete +// the same objects, and it is skipped entirely when the artifact plane is +// off. +func (e *Extension) startSweeper(ctx context.Context) { + if e.artifacts == nil || !e.artifacts.Enabled() { + return + } + + logger := e.logger + if logger == nil { + logger = e.App().Logger() + } + + cfg := e.config.Artifacts + + e.sweeper = sweeper.New(e.artifactStore, e.artifacts.Backend(), + sweeper.WithRetention(cfg.Retention), + sweeper.WithPurgeGrace(cfg.PurgeGrace), + sweeper.WithLogger(logger), + sweeper.WithLeaderCheck(e.isClusterLeader), + ) + + if serr := e.sweeper.Start(ctx); serr != nil { + logger.Warn("dispatch: could not start the artifact sweeper", + log.String("error", serr.Error())) + } +} + +// isClusterLeader reports whether this instance holds cluster leadership. +// A single-instance deployment has no cluster store and is always the +// leader by default. +func (e *Extension) isClusterLeader() bool { + cls := e.eng.ClusterStore() + if cls == nil { + return true + } + + leader, err := cls.GetLeader(context.Background()) + if err != nil || leader == nil { + return false + } + + self := e.eng.WorkerID() + if self.IsNil() { + return false + } + + return leader.ID.String() == self.String() +} + // Stop gracefully shuts down the dispatch engine. func (e *Extension) Stop(ctx context.Context) error { if e.eng == nil { e.MarkStopped() return nil } + if e.sweeper != nil { + if serr := e.sweeper.Stop(ctx); serr != nil { + e.Logger().Warn("dispatch: artifact sweeper did not stop cleanly", + forge.F("error", serr.Error())) + } + } + err := e.eng.Stop(ctx) e.MarkStopped() return err @@ -401,6 +543,40 @@ func (e *Extension) mergeWithDefaults(cfg Config) Config { if cfg.BasePath == "" { cfg.BasePath = defaults.BasePath } + + if cfg.Artifacts.Bucket == "" { + cfg.Artifacts.Bucket = "dispatch-artifacts" + } + + if cfg.Artifacts.EphemeralPrefix == "" { + cfg.Artifacts.EphemeralPrefix = artifact.DefaultEphemeralPrefix + } + + if cfg.Artifacts.Retention == 0 { + cfg.Artifacts.Retention = 168 * time.Hour + } + + if cfg.Artifacts.PurgeGrace == 0 { + cfg.Artifacts.PurgeGrace = 24 * time.Hour + } + + if cfg.Artifacts.Cache.Dir == "" { + cfg.Artifacts.Cache.Dir = "/var/lib/dispatch/cache" + } + + // Only filled in when the model is on. A zero CPUOvercommit on a + // disabled config must stay zero, so a later `enabled: true` in YAML + // cannot be silently reinterpreted as "someone chose these numbers". + if cfg.Resources.Enabled { + if cfg.Resources.CPUOvercommit <= 0 { + cfg.Resources.CPUOvercommit = resource.DefaultCPUOvercommit + } + + if cfg.Resources.MemoryFraction <= 0 { + cfg.Resources.MemoryFraction = resource.DefaultMemoryFraction + } + } + return cfg } @@ -419,6 +595,25 @@ func (e *Extension) mergeConfigurations(yamlConfig, programmaticConfig Config) C yamlConfig.EnableDWP = true } + if programmaticConfig.Artifacts.Enabled { + yamlConfig.Artifacts.Enabled = true + } + + if yamlConfig.Artifacts.TroveStore == "" && programmaticConfig.Artifacts.TroveStore != "" { + yamlConfig.Artifacts.TroveStore = programmaticConfig.Artifacts.TroveStore + } + + if yamlConfig.Artifacts.Cache.Dir == "" && programmaticConfig.Artifacts.Cache.Dir != "" { + yamlConfig.Artifacts.Cache.Dir = programmaticConfig.Artifacts.Cache.Dir + } + + if yamlConfig.Artifacts.Cache.Budget == 0 && programmaticConfig.Artifacts.Cache.Budget != 0 { + yamlConfig.Artifacts.Cache.Budget = programmaticConfig.Artifacts.Cache.Budget + } + + yamlConfig.Resources = mergeResourceConfig(yamlConfig.Resources, programmaticConfig.Resources) + yamlConfig.Execution = mergeExecutionConfig(yamlConfig.Execution, programmaticConfig.Execution) + // String fields: YAML takes precedence. if yamlConfig.BasePath == "" && programmaticConfig.BasePath != "" { yamlConfig.BasePath = programmaticConfig.BasePath diff --git a/extension/options.go b/extension/options.go index dce8a36..02d8cb4 100644 --- a/extension/options.go +++ b/extension/options.go @@ -1,15 +1,18 @@ package extension import ( + "slices" "time" log "github.com/xraph/go-utils/log" "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" "github.com/xraph/dispatch/backoff" "github.com/xraph/dispatch/dwp" "github.com/xraph/dispatch/ext" mw "github.com/xraph/dispatch/middleware" + "github.com/xraph/dispatch/resource" ) // ExtOption configures the Dispatch Forge extension. @@ -211,3 +214,113 @@ func WithDWP(opts ...dwp.Option) ExtOption { e.dwpOpts = append(e.dwpOpts, opts...) } } + +// WithArtifactBackend supplies the object store backing the artifact +// plane explicitly, bypassing container discovery. +// +// Use it outside Forge, or when the backend is not Trove. +func WithArtifactBackend(b artifact.Backend) ExtOption { + return func(e *Extension) { + e.artifactBackend = b + e.config.Artifacts.Enabled = true + } +} + +// WithArtifactStore sets the store the artifact plane persists through. +// It defaults to the dispatcher's store when that store implements +// artifact.Store, which every bundled backend does. +func WithArtifactStore(s artifact.Store) ExtOption { + return func(e *Extension) { e.artifactStore = s } +} + +// WithArtifacts enables the artifact plane, resolving a Trove instance +// from the DI container. Pass a store name for multi-store Trove setups, +// or the empty string for the default instance. +func WithArtifacts(troveStore string) ExtOption { + return func(e *Extension) { + e.config.Artifacts.Enabled = true + e.config.Artifacts.TroveStore = troveStore + } +} + +// WithArtifactCacheDir sets where staged artifacts are held on disk. +func WithArtifactCacheDir(dir string) ExtOption { + return func(e *Extension) { e.config.Artifacts.Cache.Dir = dir } +} + +// WithArtifactCacheBudget caps the bytes the staging cache may hold. +// +// With the resource model enabled this becomes the shared ledger's disk +// capacity, so it is the figure a job's declared disk requirement is +// admitted against rather than a private ceiling inside the cache. +func WithArtifactCacheBudget(bytes int64) ExtOption { + return func(e *Extension) { e.config.Artifacts.Cache.Budget = bytes } +} + +// WithResources turns on capacity detection and resource-aware +// admission. +// +// The extension then builds one resource.Manager over the detected +// capacity and hands the same instance to the staging cache and the +// worker pool: the cache holds a lease per cached entry and reclaims disk +// on demand, the pool admits every claimed job against what is actually +// free. Without this the pool dequeues unbounded and nothing changes. +// +// Detection is cgroup-first, so a container with a two-core quota +// advertises two cores rather than the host's sixty-four. +func WithResources() ExtOption { + return func(e *Extension) { e.config.Resources.Enabled = true } +} + +// WithCPUOvercommit multiplies the detected core count (default 1.0). +// +// CPU is compressible: exceeding it makes jobs slow, not dead. There is +// no memory equivalent, deliberately — overcommitting memory is how a box +// enters the OOM cascade this model prevents. +func WithCPUOvercommit(factor float64) ExtOption { + return func(e *Extension) { e.config.Resources.CPUOvercommit = factor } +} + +// WithMemoryFraction sets the share of detected memory to advertise +// (default 0.8), leaving the remainder for the Go runtime, the page +// cache, and everything else sharing the box. +func WithMemoryFraction(fraction float64) ExtOption { + return func(e *Extension) { e.config.Resources.MemoryFraction = fraction } +} + +// WithExplicitCapacity overrides detection per key and is the only way to +// declare a custom resource — nothing detects "fpga". +// +// Quantities are canonical units: cpu in millicores, memory and disk in +// bytes, gpu in milli-devices. Sets merge per key across calls. +func WithExplicitCapacity(sets ...resource.Set) ExtOption { + return func(e *Extension) { + for _, s := range sets { + if e.config.Resources.Explicit == nil { + e.config.Resources.Explicit = make(resource.Set, len(s)) + } + + for k, v := range s { + e.config.Resources.Explicit[k] = v + } + } + } +} + +// WithWorkerCustomKeys narrows the custom resource keys this worker +// advertises at dequeue. Empty advertises every custom key it has +// capacity for, which is usually what you want; this exists so a worker +// draining a device can stop attracting work for it. +// +// Keys accumulate across calls and duplicates collapse, matching +// WithExplicitCapacity above. The slice is copied, so the caller keeps no +// handle on extension state. +func WithWorkerCustomKeys(keys ...string) ExtOption { + return func(e *Extension) { + for _, k := range keys { + if !slices.Contains(e.config.Resources.CustomKeys, k) { + e.config.Resources.CustomKeys = append(e.config.Resources.CustomKeys, k) + } + } + } +} diff --git a/extension/resource.go b/extension/resource.go new file mode 100644 index 0000000..38dbffb --- /dev/null +++ b/extension/resource.go @@ -0,0 +1,166 @@ +package extension + +import ( + "slices" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/resource" +) + +// mergeResourceConfig folds programmatic resource settings into what YAML +// supplied, following the same rule as the rest of the extension: YAML +// wins where it said something, programmatic options fill the gaps, and +// the enable flag is an OR rather than an override. +// +// Explicit capacity merges per key rather than wholesale, so an operator +// can pin `memory` in YAML while the binary declares the `fpga` count it +// was built to talk to, and neither erases the other. +func mergeResourceConfig(yamlCfg, programmatic ResourceConfig) ResourceConfig { + if programmatic.Enabled { + yamlCfg.Enabled = true + } + + if yamlCfg.CPUOvercommit <= 0 && programmatic.CPUOvercommit > 0 { + yamlCfg.CPUOvercommit = programmatic.CPUOvercommit + } + + if yamlCfg.MemoryFraction <= 0 && programmatic.MemoryFraction > 0 { + yamlCfg.MemoryFraction = programmatic.MemoryFraction + } + + if len(programmatic.Explicit) > 0 { + merged := programmatic.Explicit.Clone() + for k, v := range yamlCfg.Explicit { + merged[k] = v + } + + yamlCfg.Explicit = merged + } + + if len(yamlCfg.CustomKeys) == 0 && len(programmatic.CustomKeys) > 0 { + yamlCfg.CustomKeys = slices.Clone(programmatic.CustomKeys) + } + + return yamlCfg +} + +// buildResourceManager derives this worker's capacity and returns the one +// admission ledger the staging cache and the worker pool share. +// +// One instance, built here, before anything that needs it. The cache +// constructs a private manager for itself when none is supplied, so the +// mistake this function exists to make impossible is handing the cache +// one manager and the pool another: the pool's Reclaimable() would then +// be permanently zero, the disk the cache is sitting on would never be +// offered back to the dequeue budget, and the worker would go quiet +// without logging anything wrong. +// +// A nil return is the disabled path and is not an error. Every consumer +// treats nil as "no resource model", which is the behaviour that predates +// it. +func (e *Extension) buildResourceManager() resource.Manager { + cfg := e.config.Resources + if !cfg.Enabled { + return nil + } + + staging := e.stagingBudget() + + e.warnOnDiskOverride(staging) + + capacity := resource.Detect(resource.CapacityConfig{ + CPUOvercommit: cfg.CPUOvercommit, + MemoryFraction: cfg.MemoryFraction, + DiskBytes: staging, + Explicit: cfg.Explicit.Clone(), + }) + + e.Logger().Info("dispatch: resource model enabled", + log.Any("capacity", capacity)) + + return resource.NewManager(capacity) +} + +// warnOnDiskOverride reports two config keys setting the same number with +// the less obvious one winning. +// +// resources.explicit.disk is not merely what this worker advertises. The +// cache reads its budget straight off the shared ledger's disk capacity +// (Cache.Budget), so an explicit value becomes the ceiling the cache +// evicts against and therefore the ceiling it WRITES against. Pin it +// above the volume and the worker fills the disk rather than evicting, +// and the first symptom is ENOSPC in a job that has nothing to do with +// caching. +// +// The precedence is deliberate — explicit means explicit, and an operator +// who wants a ledger disk figure that differs from the cache allowance +// has legitimate reasons — so this warns rather than fails. What it +// refuses to do is let the disagreement stay silent. +func (e *Extension) warnOnDiskOverride(staging int64) { + explicit, conflict := diskOverride(e.config.Resources.Explicit, staging) + if !conflict { + return + } + + e.Logger().Warn("dispatch: resources.explicit.disk overrides the staging cache budget", + log.Int64("explicit_disk", explicit), + log.Int64("cache_budget", staging), + log.String("effect", "the staging cache evicts against explicit_disk, so a value "+ + "above the volume's free space fills the disk instead of reclaiming")) +} + +// diskOverride reports whether an explicit disk capacity disagrees with +// the staging budget it is about to replace, and what it is. +// +// Nothing to say when there is no staging cache to disagree with, when +// the operator set only one of the two, or when they set both to the same +// number. +func diskOverride(explicit resource.Set, staging int64) (int64, bool) { + v, set := explicit[resource.Disk] + if !set || staging <= 0 || v == staging { + return 0, false + } + + return v, true +} + +// stagingBudget is the disk capacity the shared ledger advertises. +// +// It has to come from here rather than from the cache, because +// cache.WithBudget is ignored the moment a manager is supplied — a shared +// ledger's disk capacity IS the allowance, and a second ceiling +// underneath it would only be somewhere for the two to disagree. Routing +// the configured budget in as the manager's disk capacity is what keeps +// an operator who wrote a number in the config from silently getting +// whatever Detect chose instead. +// +// Zero when there is no staging cache at all, which omits the disk key +// entirely rather than advertising capacity nothing can reclaim. An +// explicit `disk` in the resources config still overrides this, since +// Detect applies Explicit last. +// +// The artifactStore check is the one that is easy to miss: init only +// builds the plane when the dispatcher's store implements artifact.Store, +// so `artifacts.enabled: true` on a store that does not is a configured +// plane that never exists. Testing the config alone would advertise +// 20 GiB of disk with no cache behind it and no reclaimer registered for +// it — the exact invariant TestNoArtifactPlaneOmitsDisk pins. This +// requires the store to be resolved before the ledger is built, which is +// why init does that first. +func (e *Extension) stagingBudget() int64 { + if e.artifactStore == nil { + return 0 + } + + if !e.config.Artifacts.Enabled && e.artifactBackend == nil { + return 0 + } + + if b := e.config.Artifacts.Cache.Budget; b > 0 { + return b + } + + return cache.DefaultBudget +} diff --git a/extension/resource_test.go b/extension/resource_test.go new file mode 100644 index 0000000..06ec29f --- /dev/null +++ b/extension/resource_test.go @@ -0,0 +1,378 @@ +package extension_test + +import ( + "context" + "sync" + "testing" + "time" + + forgetesting "github.com/xraph/forge/testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/extension" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/store/memory" +) + +const mib = int64(1) << 20 + +// registerWithResources builds an extension with the artifact plane and +// the resource model both on, and registers it against a test app. +func registerWithResources(t *testing.T, opts ...extension.ExtOption) *extension.Extension { + t.Helper() + + base := make([]extension.ExtOption, 0, 7+len(opts)) + base = append(base, + extension.WithStore(memory.New()), + extension.WithArtifactBackend(artifacttest.NewBackend()), + extension.WithArtifactStore(memory.New()), + extension.WithArtifactCacheDir(t.TempDir()), + extension.WithArtifactCacheBudget(64*mib), + extension.WithResources(), + extension.WithDisableRoutes(), + ) + + ext := extension.New(append(base, opts...)...) + + if err := ext.Register(forgetesting.NewTestApp("test-app", "0.1.0")); err != nil { + t.Fatalf("Register: %v", err) + } + + return ext +} + +// TestCacheAndEngineShareOneManager is the end-to-end form of the +// coupling worker.TestCacheAndPoolShareOneLedger pins in isolation: +// through the real extension wiring, the staging cache and the worker +// pool must end up holding the same resource.Manager. +// +// It asserts behaviourally rather than by pointer, because the symptom of +// getting this wrong is behavioural. Bytes staged through the cache have +// to show up as reclaimable in the ledger the engine handed the pool. If +// the cache ever falls back to the private manager it builds when none is +// supplied, that number stays zero and the worker silently stops claiming +// disk-hungry work. +func TestCacheAndEngineShareOneManager(t *testing.T) { + ext := registerWithResources(t) + + mgr := ext.Resources() + if mgr == nil { + t.Fatal("resource manager is nil after Register with WithResources") + } + + if got := ext.Engine().Resources(); got != mgr { + t.Fatal("the engine was given a different manager than the extension built") + } + + c := ext.ArtifactCache() + if c == nil { + t.Fatal("staging cache is nil") + } + + // The cache's own view of its budget is the shared ledger's disk + // capacity, which is where the configured cache budget was routed. + if got, want := c.Budget(), 64*mib; got != want { + t.Fatalf("cache budget = %d, want %d — the configured budget did not "+ + "reach the shared ledger's disk capacity", got, want) + } + + if got, want := mgr.Capacity()[resource.Disk], 64*mib; got != want { + t.Fatalf("ledger disk capacity = %d, want %d", got, want) + } + + // Stage real bytes and read the ledger the pool reads. + backend, ok := ext.Artifacts().Backend().(*artifacttest.Backend) + if !ok { + t.Fatalf("backend is %T, want *artifacttest.Backend", ext.Artifacts().Backend()) + } + + const staged = 4 * mib + + backend.Put("stage", "model.bin", make([]byte, staged)) + + _, _, release, err := c.Stage(context.Background(), artifact.Ref{ + Backend: backend.Name(), + Bucket: "stage", + Key: "model.bin", + Size: staged, + }) + if err != nil { + t.Fatalf("stage: %v", err) + } + + if got, want := mgr.Free()[resource.Disk], 64*mib-staged; got != want { + t.Fatalf("while pinned: ledger free disk = %d, want %d — "+ + "the cache took its lease against a different manager", got, want) + } + + release() + + if got := mgr.Reclaimable()[resource.Disk]; got != staged { + t.Fatalf("after release: ledger reclaimable disk = %d, want %d — "+ + "the cache is not registered as this ledger's disk reclaimer", got, staged) + } +} + +// TestExplicitCapacityDeclaresCustomResources checks the only path a +// custom resource can arrive by. Nothing detects an FPGA. +func TestExplicitCapacityDeclaresCustomResources(t *testing.T) { + ext := registerWithResources(t, + extension.WithExplicitCapacity(resource.Set{"fpga": 2}, resource.GPUs(4)), + ) + + capacity := ext.Resources().Capacity() + + if got := capacity["fpga"]; got != 2 { + t.Errorf("fpga capacity = %d, want 2", got) + } + + if got := capacity[resource.GPU]; got != 4*resource.MilliScale { + t.Errorf("gpu capacity = %d, want %d", got, 4*resource.MilliScale) + } + + // Explicit overrides detection, so an operator pinning memory gets + // exactly that rather than a fraction of it. + ext2 := registerWithResources(t, + extension.WithExplicitCapacity(resource.MemoryBytes(7*mib)), + ) + + if got := ext2.Resources().Capacity()[resource.Memory]; got != 7*mib { + t.Errorf("explicit memory capacity = %d, want %d", got, 7*mib) + } +} + +// TestResourcesDisabledIsTodaysBehaviour pins the backward-compatibility +// guarantee. Absent config means no ledger anywhere: the pool dequeues +// unbounded, every backend skips its fit predicate, and the staging cache +// keeps the private disk budget it has always had. +func TestResourcesDisabledIsTodaysBehaviour(t *testing.T) { + ext := extension.New( + extension.WithStore(memory.New()), + extension.WithArtifactBackend(artifacttest.NewBackend()), + extension.WithArtifactStore(memory.New()), + extension.WithArtifactCacheDir(t.TempDir()), + extension.WithArtifactCacheBudget(32*mib), + extension.WithDisableRoutes(), + ) + + if err := ext.Register(forgetesting.NewTestApp("test-app", "0.1.0")); err != nil { + t.Fatalf("Register: %v", err) + } + + if ext.Resources() != nil { + t.Fatal("a manager was built without WithResources") + } + + if ext.Engine().Resources() != nil { + t.Fatal("the engine has a manager without WithResources") + } + + // The cache still honours WithBudget, because with no shared ledger + // its private manager is the only ceiling there is. + if got, want := ext.ArtifactCache().Budget(), 32*mib; got != want { + t.Fatalf("cache budget = %d, want %d", got, want) + } +} + +// dequeueSpy records the DequeueOpts the extension's own worker pool +// sends to the store. It is the only vantage point from which the +// degradation guarantee can actually be observed: the pool is not +// exported, so what a disabled resource model does has to be read off the +// wire it writes to. +type dequeueSpy struct { + *memory.Store + + mu sync.Mutex + opts []job.DequeueOpts +} + +func (s *dequeueSpy) DequeueJobs(ctx context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + s.mu.Lock() + s.opts = append(s.opts, opts) + s.mu.Unlock() + + return s.Store.DequeueJobs(ctx, opts) +} + +// seen returns a copy of everything recorded so far. +func (s *dequeueSpy) seen() []job.DequeueOpts { + s.mu.Lock() + defer s.mu.Unlock() + + return append([]job.DequeueOpts(nil), s.opts...) +} + +// runAndCaptureDequeues starts the extension, waits for the pool to poll +// at least once, and returns what it asked the store for. +func runAndCaptureDequeues(t *testing.T, opts ...extension.ExtOption) []job.DequeueOpts { + t.Helper() + + spy := &dequeueSpy{Store: memory.New()} + + base := make([]extension.ExtOption, 0, 3+len(opts)) + base = append(base, + extension.WithStore(spy), + extension.WithDisableRoutes(), + extension.WithPollInterval(20*time.Millisecond), + ) + + ext := extension.New(append(base, opts...)...) + + if err := ext.Register(forgetesting.NewTestApp("test-app", "0.1.0")); err != nil { + t.Fatalf("Register: %v", err) + } + + ctx := t.Context() + + if err := ext.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + + deadline := time.Now().Add(5 * time.Second) + for len(spy.seen()) == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + + stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := ext.Stop(stopCtx); err != nil { + t.Fatalf("Stop: %v", err) + } + + seen := spy.seen() + if len(seen) == 0 { + t.Fatal("the pool never polled the store") + } + + return seen +} + +// TestDisabledPoolDequeuesUnbounded is the degradation guarantee stated +// where it is actually observable: with no resource model configured, +// every dequeue this worker issues carries no budget and no custom keys, +// so DequeueOpts.IsUnbounded() holds and every store backend skips its +// fit predicate entirely. +// +// This is the single behaviour every deployment that predates the +// resource model depends on, so it is asserted through the extension's +// own pool rather than inferred from the manager being nil. +func TestDisabledPoolDequeuesUnbounded(t *testing.T) { + for i, opts := range runAndCaptureDequeues(t) { + if !opts.IsUnbounded() { + t.Fatalf("dequeue %d: IsUnbounded() = false, want true "+ + "(Budget=%v CustomKeys=%v)", i, opts.Budget, opts.CustomKeys) + } + } +} + +// TestEnabledPoolDequeuesBounded is its mirror: turning the model on has +// to change what goes over that same wire, or the pool is holding a +// ledger it never consults. +func TestEnabledPoolDequeuesBounded(t *testing.T) { + seen := runAndCaptureDequeues(t, + extension.WithResources(), + extension.WithExplicitCapacity(resource.Set{resource.Memory: 4 * mib}), + ) + + for i, opts := range seen { + if opts.IsUnbounded() { + t.Fatalf("dequeue %d: IsUnbounded() = true, want false with a ledger", i) + } + + if opts.Budget[resource.Memory] != 4*mib { + t.Fatalf("dequeue %d: budget memory = %d, want %d", + i, opts.Budget[resource.Memory], 4*mib) + } + } +} + +// TestCacheBudgetDefaultsWhenUnset covers the gap that would otherwise +// hand the ledger a zero disk capacity: artifacts on, budget unstated. +// The cache's own default has to be what the ledger advertises, or every +// disk-declaring job becomes unschedulable on a worker that can in fact +// stage 20 GiB. +func TestCacheBudgetDefaultsWhenUnset(t *testing.T) { + ext := extension.New( + extension.WithStore(memory.New()), + extension.WithArtifactBackend(artifacttest.NewBackend()), + extension.WithArtifactStore(memory.New()), + extension.WithArtifactCacheDir(t.TempDir()), + extension.WithResources(), + extension.WithDisableRoutes(), + ) + + if err := ext.Register(forgetesting.NewTestApp("test-app", "0.1.0")); err != nil { + t.Fatalf("Register: %v", err) + } + + if got := ext.Resources().Capacity()[resource.Disk]; got <= 0 { + t.Fatalf("ledger disk capacity = %d, want the cache default", got) + } + + if got, want := ext.ArtifactCache().Budget(), + ext.Resources().Capacity()[resource.Disk]; got != want { + t.Fatalf("cache budget = %d, ledger disk = %d — these must be one number", got, want) + } +} + +// TestNoArtifactPlaneOmitsDisk checks that a worker with no staging cache +// advertises no disk at all, rather than capacity nothing can reclaim. +func TestNoArtifactPlaneOmitsDisk(t *testing.T) { + ext := extension.New( + extension.WithStore(memory.New()), + extension.WithResources(), + extension.WithDisableRoutes(), + ) + + if err := ext.Register(forgetesting.NewTestApp("test-app", "0.1.0")); err != nil { + t.Fatalf("Register: %v", err) + } + + capacity := ext.Resources().Capacity() + + if _, present := capacity[resource.Disk]; present { + t.Fatalf("disk advertised with no staging cache: %v", capacity) + } + + // Detection still produced the two keys that always exist. + if capacity[resource.CPU] <= 0 || capacity[resource.Memory] <= 0 { + t.Fatalf("detection produced no cpu/memory: %v", capacity) + } +} + +// TestEnablingResourcesDoesNotEnableTheUnschedulableCheck pins the +// other half of the wiring, and it is the inverse of what this test +// asserted before: `resources.enabled: true` in a config file is a +// statement about THIS pod, and must not be read as a statement about +// the largest worker in the fleet. +// +// The unschedulable check compares a job's requirement against the +// biggest worker anywhere. Deriving that ceiling from the local ledger +// answered it with the wrong number and could not recover: cluster +// workers only round-trip Capacity on store/memory, so on postgres, +// sqlite, mongo and redis the fleet view is permanently empty and the +// ceiling is permanently this process. A light API pod would then reject +// at enqueue the tessellation job the heavy tier runs perfectly well — +// the opening scenario of the design, inverted. +// +// So a Forge deployment that turns resources on gets local admission and +// a resource-aware dequeue, and nothing that can fail an enqueue. The +// check is engine.WithWorkerCapacity, and it is the operator's to +// declare. +func TestEnablingResourcesDoesNotEnableTheUnschedulableCheck(t *testing.T) { + ext := registerWithResources(t, + extension.WithExplicitCapacity(resource.Set{"fpga": 3}), + ) + + if ledger := ext.Resources().Capacity(); ledger["fpga"] != 3 { + t.Fatalf("ledger capacity = %v, want the explicit fpga declaration", ledger) + } + + if published := ext.Engine().MaxWorkerCapacity(context.Background()); len(published) != 0 { + t.Errorf("MaxWorkerCapacity = %v, want empty: enabling resources must not turn a "+ + "fleet-wide enqueue-time check on with one pod's numbers", published) + } +} diff --git a/go.mod b/go.mod index 5764efa..e2f17a2 100644 --- a/go.mod +++ b/go.mod @@ -7,28 +7,27 @@ require ( github.com/jackc/pgx/v5 v5.10.0 github.com/redis/go-redis/v9 v9.21.0 github.com/robfig/cron/v3 v3.0.1 - github.com/testcontainers/testcontainers-go v0.42.0 - github.com/testcontainers/testcontainers-go/modules/mongodb v0.42.0 - github.com/testcontainers/testcontainers-go/modules/postgres v0.42.0 - github.com/testcontainers/testcontainers-go/modules/redis v0.42.0 - github.com/uptrace/bun v1.2.16 - github.com/uptrace/bun/dialect/pgdialect v1.2.16 - github.com/uptrace/bun/driver/pgdriver v1.2.16 + github.com/testcontainers/testcontainers-go v0.44.0 + github.com/testcontainers/testcontainers-go/modules/mongodb v0.44.0 + github.com/testcontainers/testcontainers-go/modules/postgres v0.44.0 + github.com/testcontainers/testcontainers-go/modules/redis v0.44.0 github.com/vmihailenco/msgpack/v5 v5.4.1 - github.com/xraph/forge v1.9.2 + github.com/xraph/forge v1.9.8 github.com/xraph/forgeui v1.4.1 - github.com/xraph/grove v1.6.0 - github.com/xraph/grove/drivers/mongodriver v1.6.0 - github.com/xraph/grove/drivers/pgdriver v1.6.0 - github.com/xraph/grove/drivers/sqlitedriver v1.6.0 - github.com/xraph/grove/kv v1.6.0 - github.com/xraph/grove/kv/drivers/redisdriver v1.6.0 - github.com/xraph/relay v1.6.0 + github.com/xraph/grove v1.6.1 + github.com/xraph/grove/drivers/mongodriver v1.6.1 + github.com/xraph/grove/drivers/pgdriver v1.6.1 + github.com/xraph/grove/drivers/sqlitedriver v1.6.1 + github.com/xraph/grove/kv v1.6.1 + github.com/xraph/grove/kv/drivers/redisdriver v1.6.1 + github.com/xraph/relay v1.6.2 + github.com/xraph/trove v1.6.4 github.com/xraph/vessel v1.0.4 + github.com/zeebo/blake3 v0.2.4 go.jetify.com/typeid/v2 v2.0.0-alpha.3 - go.mongodb.org/mongo-driver/v2 v2.5.0 + go.mongodb.org/mongo-driver/v2 v2.8.0 go.opentelemetry.io/otel v1.44.0 - go.opentelemetry.io/otel/sdk v1.43.0 + go.opentelemetry.io/otel/sdk v1.44.0 go.opentelemetry.io/otel/trace v1.44.0 golang.org/x/sync v0.22.0 golang.org/x/time v0.15.0 @@ -50,18 +49,18 @@ require ( github.com/cpuguy83/dockercfg v0.3.2 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/distribution/reference v0.6.0 // indirect - github.com/docker/go-connections v0.6.0 // indirect + github.com/docker/go-connections v0.7.0 // indirect github.com/docker/go-units v0.5.0 // indirect - github.com/ebitengine/purego v0.10.0 // indirect + github.com/ebitengine/purego v0.10.1 // indirect github.com/emicklei/go-restful/v3 v3.13.0 // indirect github.com/fatih/color v1.19.0 // indirect - github.com/felixge/httpsnoop v1.0.4 // indirect + github.com/felixge/httpsnoop v1.1.0 // indirect github.com/fsnotify/fsnotify v1.10.1 // indirect github.com/fxamacker/cbor/v2 v2.9.2 // indirect github.com/gabriel-vasile/mimetype v1.4.15 // indirect github.com/go-logr/logr v1.4.4 // indirect github.com/go-logr/stdr v1.2.2 // indirect - github.com/go-ole/go-ole v1.2.6 // indirect + github.com/go-ole/go-ole v1.3.0 // indirect github.com/go-openapi/jsonpointer v1.0.0 // indirect github.com/go-openapi/jsonreference v1.0.0 // indirect github.com/go-openapi/swag v0.28.0 // indirect @@ -89,7 +88,7 @@ require ( github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.19.1 // indirect github.com/leodido/go-urn v1.5.0 // indirect - github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect + github.com/lufia/plan9stats v0.0.0-20260330125221-c963978e514e // indirect github.com/magiconair/properties v1.8.10 // indirect github.com/mattn/go-colorable v0.1.15 // indirect github.com/mattn/go-isatty v0.0.24 // indirect @@ -97,7 +96,7 @@ require ( github.com/moby/docker-image-spec v1.3.1 // indirect github.com/moby/go-archive v0.2.0 // indirect github.com/moby/patternmatcher v0.6.1 // indirect - github.com/moby/sys/sequential v0.6.0 // indirect + github.com/moby/sys/sequential v0.7.0 // indirect github.com/moby/sys/user v0.4.0 // indirect github.com/moby/sys/userns v0.1.0 // indirect github.com/moby/term v0.5.2 // indirect @@ -111,19 +110,19 @@ require ( github.com/quic-go/qpack v0.6.0 // indirect github.com/quic-go/quic-go v0.61.0 // indirect github.com/quic-go/webtransport-go v0.12.0 // indirect - github.com/shirou/gopsutil/v4 v4.26.3 // indirect + github.com/shirou/gopsutil/v4 v4.26.6 // indirect github.com/sirupsen/logrus v1.9.4 // indirect github.com/spf13/pflag v1.0.10 // indirect github.com/stretchr/testify v1.11.1 // indirect - github.com/tklauser/go-sysconf v0.3.16 // indirect - github.com/tklauser/numcpus v0.11.0 // indirect + github.com/tklauser/go-sysconf v0.4.0 // indirect + github.com/tklauser/numcpus v0.12.0 // indirect github.com/uptrace/bunrouter v1.0.23 // indirect github.com/x448/float16 v0.8.4 // indirect github.com/xraph/confy v1.0.2 // indirect github.com/xraph/go-utils v1.1.6 github.com/yusufpapurcu/wmi v1.2.4 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect - go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0 // indirect go.uber.org/multierr v1.11.0 // indirect go.uber.org/zap v1.28.0 // indirect go.yaml.in/yaml/v2 v2.4.4 // indirect @@ -167,22 +166,20 @@ require ( github.com/go-openapi/swag/stringutils v0.28.0 // indirect github.com/go-openapi/swag/typeutils v0.28.0 // indirect github.com/go-openapi/swag/yamlutils v0.28.0 // indirect - github.com/gofrs/uuid/v5 v5.3.2 // indirect + github.com/gofrs/uuid/v5 v5.5.1 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 // indirect github.com/hashicorp/go-metrics v0.6.1 // indirect - github.com/jinzhu/inflection v1.0.0 // indirect + github.com/klauspost/cpuid/v2 v2.4.0 // indirect github.com/mdelapenya/tlscert v0.2.0 // indirect - github.com/moby/moby/api v1.54.1 // indirect - github.com/moby/moby/client v0.4.0 // indirect + github.com/moby/moby/api v1.55.0 // indirect + github.com/moby/moby/client v0.5.0 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect github.com/prometheus/client_golang v1.24.1 // indirect github.com/prometheus/client_model v0.6.2 // indirect github.com/prometheus/common v0.70.1 // indirect github.com/prometheus/procfs v0.21.1 // indirect - github.com/puzpuzpuz/xsync/v3 v3.5.1 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect - github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc // indirect github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect github.com/xdg-go/pbkdf2 v1.0.0 // indirect github.com/xdg-go/scram v1.2.0 // indirect @@ -197,7 +194,6 @@ require ( google.golang.org/genproto/googleapis/api v0.0.0-20260414002931-afd174a4e478 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect google.golang.org/grpc v1.82.1 // indirect - mellium.im/sasl v0.3.2 // indirect modernc.org/libc v1.68.0 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect diff --git a/go.sum b/go.sum index cdf9762..76bdef6 100644 --- a/go.sum +++ b/go.sum @@ -56,23 +56,23 @@ github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5Qvfr github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI= github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= -github.com/docker/go-connections v0.6.0 h1:LlMG9azAe1TqfR7sO+NJttz1gy6KO7VJBh+pMmjSD94= -github.com/docker/go-connections v0.6.0/go.mod h1:AahvXYshr6JgfUJGdDCs2b5EZG/vmaMAntpSFH5BFKE= +github.com/docker/go-connections v0.7.0 h1:6SsRfJddP22WMrCkj19x9WKjEDTB+ahsdiGYf0mN39c= +github.com/docker/go-connections v0.7.0/go.mod h1:no1qkHdjq7kLMGUXYAduOhYPSJxxvgWBh7ogVvptn3Q= github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= github.com/dunglas/httpsfv v1.1.0 h1:Jw76nAyKWKZKFrpMMcL76y35tOpYHqQPzHQiwDvpe54= github.com/dunglas/httpsfv v1.1.0/go.mod h1:zID2mqw9mFsnt7YC3vYQ9/cjq30q41W+1AnDwH8TiMg= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= -github.com/ebitengine/purego v0.10.0 h1:QIw4xfpWT6GWTzaW5XEKy3HXoqrJGx1ijYHzTF0/ISU= -github.com/ebitengine/purego v0.10.0/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ= +github.com/ebitengine/purego v0.10.1 h1:dewVBCBT2GaMu1SrNTYxQhgQBethzfhiwvZiLGP/qyY= +github.com/ebitengine/purego v0.10.1/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ= github.com/emicklei/go-restful/v3 v3.13.0 h1:C4Bl2xDndpU6nJ4bc1jXd+uTmYPVUwkD6bFY/oTyCes= github.com/emicklei/go-restful/v3 v3.13.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk= github.com/fatih/color v1.19.0 h1:Zp3PiM21/9Ld6FzSKyL5c/BULoe/ONr9KlbYVOfG8+w= github.com/fatih/color v1.19.0/go.mod h1:zNk67I0ZUT1bEGsSGyCZYZNrHuTkJJB+r6Q9VuMi0LE= -github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= -github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= +github.com/felixge/httpsnoop v1.1.0 h1:3YtUj32ZZkqZtt3sZZsClsymw/QDuVfpNhoA31zeORc= +github.com/felixge/httpsnoop v1.1.0/go.mod h1:Zqxgdd+1Rkcz8euOqdr7lqgCRJztwr5hp9vDSi5UZCE= github.com/fsnotify/fsnotify v1.10.1 h1:b0/UzAf9yR5rhf3RPm9gf3ehBPpf0oZKIjtpKrx59Ho= github.com/fsnotify/fsnotify v1.10.1/go.mod h1:TLheqan6HD6GBK6PrDWyDPBaEV8LspOxvPSjC+bVfgo= github.com/fxamacker/cbor/v2 v2.9.2 h1:X4Ksno9+x3cz0TZv69ec1hxP/+tymuR8PXQJyDwfh78= @@ -88,8 +88,9 @@ github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8= github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= -github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= +github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE= +github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78= github.com/go-openapi/jsonpointer v1.0.0 h1:kR9tHqY0CtZaOPVFm622dPVNhrvYpwr4uCxgL3h1H8s= github.com/go-openapi/jsonpointer v1.0.0/go.mod h1:Z3rw7dWu1p9IgitXCFamSlA5lmDiklEB6vkaxcNZW5Y= github.com/go-openapi/jsonreference v1.0.0 h1:jlmTr6torcd1YgDQvSfNmRtKzYDO4FGBkrAdlAVWnpY= @@ -141,8 +142,8 @@ github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og= github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw= github.com/gobwas/ws v1.4.0 h1:CTaoG1tojrh4ucGPcoJFiAQUAsEWekEWvLy7GsVNqGs= github.com/gobwas/ws v1.4.0/go.mod h1:G3gNqMNtPppf5XUz7O4shetPpcZ1VJ7zt18dlUeakrc= -github.com/gofrs/uuid/v5 v5.3.2 h1:2jfO8j3XgSwlz/wHqemAEugfnTlikAYHhnqQ8Xh4fE0= -github.com/gofrs/uuid/v5 v5.3.2/go.mod h1:CDOjlDMVAtN56jqyRUZh58JT31Tiw7/oQyEXZV+9bD8= +github.com/gofrs/uuid/v5 v5.5.1 h1:z1Ce19/JwNidXpy3tOQc3241lnJLKdKyq/xlNvlD4Ng= +github.com/gofrs/uuid/v5 v5.5.1/go.mod h1:bbAA98EoIlxyRHIVg6ektCSsZ5n8mSbwgEhvhMYlZgg= github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= @@ -155,7 +156,6 @@ github.com/google/gnostic-models v0.7.1 h1:SisTfuFKJSKM5CPZkffwi6coztzzeYUhc3v4y github.com/google/gnostic-models v0.7.1/go.mod h1:whL5G0m6dmc5cPxKc5bdKdEN3UjI7OUGxBlw57miDrQ= github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= -github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= @@ -213,8 +213,6 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= -github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCVDaaPEHmU= github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= @@ -222,8 +220,8 @@ github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHm github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w= github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk= github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= -github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= -github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV913Dw= +github.com/klauspost/cpuid/v2 v2.4.0/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= @@ -239,8 +237,8 @@ github.com/leodido/go-urn v1.5.0 h1:pLqT2kq1zpHW/1D18QMjMpdtX7cekxqtJJjg5ANyWw0= github.com/leodido/go-urn v1.5.0/go.mod h1:9BORnCDhdPBJNDEX+w1bJisa8yOKYi116VeO96s4ifE= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= -github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= -github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= +github.com/lufia/plan9stats v0.0.0-20260330125221-c963978e514e h1:Q6MvJtQK/iRcRtzAscm/zF23XxJlbECiGPyRicsX+Ak= +github.com/lufia/plan9stats v0.0.0-20260330125221-c963978e514e/go.mod h1:autxFIvghDt3jPTLoqZ9OZ7s9qTGNAWmYCjVFWPX/zg= github.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8SYxI99mE= github.com/magiconair/properties v1.8.10/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc= @@ -262,14 +260,14 @@ github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3N github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo= github.com/moby/go-archive v0.2.0 h1:zg5QDUM2mi0JIM9fdQZWC7U8+2ZfixfTYoHL7rWUcP8= github.com/moby/go-archive v0.2.0/go.mod h1:mNeivT14o8xU+5q1YnNrkQVpK+dnNe/K6fHqnTg4qPU= -github.com/moby/moby/api v1.54.1 h1:TqVzuJkOLsgLDDwNLmYqACUuTehOHRGKiPhvH8V3Nn4= -github.com/moby/moby/api v1.54.1/go.mod h1:+RQ6wluLwtYaTd1WnPLykIDPekkuyD/ROWQClE83pzs= -github.com/moby/moby/client v0.4.0 h1:S+2XegzHQrrvTCvF6s5HFzcrywWQmuVnhOXe2kiWjIw= -github.com/moby/moby/client v0.4.0/go.mod h1:QWPbvWchQbxBNdaLSpoKpCdf5E+WxFAgNHogCWDoa7g= +github.com/moby/moby/api v1.55.0 h1:2/sexvQyqIWS8pRSCFddBfpW2qE7vR7FCL+vN8pxwMc= +github.com/moby/moby/api v1.55.0/go.mod h1:+RQ6wluLwtYaTd1WnPLykIDPekkuyD/ROWQClE83pzs= +github.com/moby/moby/client v0.5.0 h1:5XhyPk2fuOWf6RlSFa3MkIIgDZkF25xToXW8Q/BH7cc= +github.com/moby/moby/client v0.5.0/go.mod h1:rcVpF8ncl9vo5gaIBdol6CnbEtSj1uxMvEV/UrykF/s= github.com/moby/patternmatcher v0.6.1 h1:qlhtafmr6kgMIJjKJMDmMWq7WLkKIo23hsrpR3x084U= github.com/moby/patternmatcher v0.6.1/go.mod h1:hDPoyOpDY7OrrMDLaYoY3hf52gNCR/YOUYxkhApJIxc= -github.com/moby/sys/sequential v0.6.0 h1:qrx7XFUd/5DxtqcoH1h438hF5TmOvzC/lspjy7zgvCU= -github.com/moby/sys/sequential v0.6.0/go.mod h1:uyv8EUTrca5PnDsdMGXhZe6CCe8U/UiTWd+lL+7b/Ko= +github.com/moby/sys/sequential v0.7.0 h1:ASQNGNROJSuOO6LL6bPHbKvuZu6NU8P4ldPWk31zj/8= +github.com/moby/sys/sequential v0.7.0/go.mod h1:NfSTAp6V3fw4tmkD62PEcOKeZKquXT8VKCkf7aVR79o= github.com/moby/sys/user v0.4.0 h1:jhcMKit7SA80hivmFJcbB1vqmw//wU61Zdui2eQXuMs= github.com/moby/sys/user v0.4.0/go.mod h1:bG+tYYYJgaMtRKgEmuueC0hJEAZWwtIbZTB+85uoHjs= github.com/moby/sys/userns v0.1.0 h1:tVLXkFOxVu9A64/yh59slHVv9ahO9UIev4JZusOLG/g= @@ -323,8 +321,6 @@ github.com/prometheus/procfs v0.0.2/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsT github.com/prometheus/procfs v0.0.8/go.mod h1:7Qr8sr6344vo1JqZ6HhLceV9o3AJ1Ff+GxbHq6oeK9A= github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI= github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY= -github.com/puzpuzpuz/xsync/v3 v3.5.1 h1:GJYJZwO6IdxN/IKbneznS6yPkVC+c3zyY/j19c++5Fg= -github.com/puzpuzpuz/xsync/v3 v3.5.1/go.mod h1:VjzYrABPabuM4KyBh1Ftq6u8nhwY5tBPKP9jpmh0nnA= github.com/quic-go/go-ossfuzz-seeds v0.1.0 h1:APacT+iIaNF6fd8AGEiN3bT/Jtkd2jz4v4TzM7MFjy0= github.com/quic-go/go-ossfuzz-seeds v0.1.0/go.mod h1:3IOHRbJIc+L6YKMwfDtJAM9Vj9k0YY4muhuyUYk5tbk= github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= @@ -347,8 +343,8 @@ github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEV github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU= github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529 h1:nn5Wsu0esKSJiIVhscUtVbo7ada43DJhG55ua/hjS5I= github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc= -github.com/shirou/gopsutil/v4 v4.26.3 h1:2ESdQt90yU3oXF/CdOlRCJxrP+Am1aBYubTMTfxJ1qc= -github.com/shirou/gopsutil/v4 v4.26.3/go.mod h1:LZ6ewCSkBqUpvSOf+LsTGnRinC6iaNUNMGBtDkJBaLQ= +github.com/shirou/gopsutil/v4 v4.26.6 h1:Mzr/npDtQC/xpeEuQKHZt8Zo9CmPvhTj8nkR8w5TLDs= +github.com/shirou/gopsutil/v4 v4.26.6/go.mod h1:LZ6ewCSkBqUpvSOf+LsTGnRinC6iaNUNMGBtDkJBaLQ= github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo= github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE= github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w= @@ -366,27 +362,19 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/ github.com/stretchr/testify v1.7.2/go.mod h1:R6va5+xMeoiuVRoj+gSkQ7d3FALtqAAGI1FQKckRals= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= -github.com/testcontainers/testcontainers-go v0.42.0 h1:He3IhTzTZOygSXLJPMX7n44XtK+qhjat1nI9cneBbUY= -github.com/testcontainers/testcontainers-go v0.42.0/go.mod h1:vZjdY1YmUA1qEForxOIOazfsrdyORJAbhi0bp8plN30= -github.com/testcontainers/testcontainers-go/modules/mongodb v0.42.0 h1:jX10Aprgf1L+Ov+KxcheZ/1JXdiJ/3wdevfWFSkxm6s= -github.com/testcontainers/testcontainers-go/modules/mongodb v0.42.0/go.mod h1:Ph+xH0hAC6djPFTjPgLa3VmSfE4h82kzVIKxTj3n2o4= -github.com/testcontainers/testcontainers-go/modules/postgres v0.42.0 h1:GCbb1ndrF7OTDiIvxXyItaDab4qkzTFJ48LKFdM7EIo= -github.com/testcontainers/testcontainers-go/modules/postgres v0.42.0/go.mod h1:IRPBaI8jXdrNfD0e4Zm7Fbcgaz5shKxOQv4axiL09xs= -github.com/testcontainers/testcontainers-go/modules/redis v0.42.0 h1:id/6LH8ZeDrtAUVSuNvZUAJ1kVpb82y1pr9yweAWsRg= -github.com/testcontainers/testcontainers-go/modules/redis v0.42.0/go.mod h1:uF0jI8FITagQpBNOgweGBmPf6rP4K0SeL1XFPbsZSSY= -github.com/tklauser/go-sysconf v0.3.16 h1:frioLaCQSsF5Cy1jgRBrzr6t502KIIwQ0MArYICU0nA= -github.com/tklauser/go-sysconf v0.3.16/go.mod h1:/qNL9xxDhc7tx3HSRsLWNnuzbVfh3e7gh/BmM179nYI= -github.com/tklauser/numcpus v0.11.0 h1:nSTwhKH5e1dMNsCdVBukSZrURJRoHbSEQjdEbY+9RXw= -github.com/tklauser/numcpus v0.11.0/go.mod h1:z+LwcLq54uWZTX0u/bGobaV34u6V7KNlTZejzM6/3MQ= -github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc h1:9lRDQMhESg+zvGYmW5DyG0UqvY96Bu5QYsTLvCHdrgo= -github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc/go.mod h1:bciPuU6GHm1iF1pBvUfxfsH0Wmnc2VbpgvbI9ZWuIRs= +github.com/testcontainers/testcontainers-go v0.44.0 h1:/Fwh6HY1mIikhnm9e7HwoxGycx0lzRAE0f5VQpjFxzI= +github.com/testcontainers/testcontainers-go v0.44.0/go.mod h1:IcnwQrYTO86xHXu5bvMaBH7ATlbS3Qn1M1QWW3c66rE= +github.com/testcontainers/testcontainers-go/modules/mongodb v0.44.0 h1:VSPDFiumAtt0CkZEVbmAkEmYVRvsJpKJy9oF3exRKYg= +github.com/testcontainers/testcontainers-go/modules/mongodb v0.44.0/go.mod h1:kHfzrY1cYP/zr9H4TdqAxbP836A1C2fyUojlHidhFGI= +github.com/testcontainers/testcontainers-go/modules/postgres v0.44.0 h1:8fdv/9y3JMxjQ+ULAcOG8RtgeNu5t9XF9LolSXDuTwM= +github.com/testcontainers/testcontainers-go/modules/postgres v0.44.0/go.mod h1:CFr2LncGYokw+OKjXcr8ARCKG1SaC2UEnGxFBovE86g= +github.com/testcontainers/testcontainers-go/modules/redis v0.44.0 h1:43EH7N6yB5B2tY/9uhPit487tMLm5iQiyKQaXWXNbnk= +github.com/testcontainers/testcontainers-go/modules/redis v0.44.0/go.mod h1:k4nnCSzm3z8yRMBKBn3rhsllbFjjhVn/2JjWNxxArg8= +github.com/tklauser/go-sysconf v0.4.0 h1:7H0uAN+7RkwWRaxhYXDLqa5V3LPrJeV8wmD9dRUgPQU= +github.com/tklauser/go-sysconf v0.4.0/go.mod h1:8mTNWyog7H+MpKijp4VmKJAd2bbYQ2zuUwkYRbUArPI= +github.com/tklauser/numcpus v0.12.0 h1:NR85qdvHA9pFse3x3weVZ0r0ST8R6l5RHbZrlRaqob4= +github.com/tklauser/numcpus v0.12.0/go.mod h1:ABHeXzJnr/qqwguhClkZKT1/8VABcYrsyUiUGobwWJg= github.com/tv42/httpunix v0.0.0-20150427012821-b75d8614f926/go.mod h1:9ESjWnEqriFuLhtthL60Sar/7RFoluCcXsuvEwTV5KM= -github.com/uptrace/bun v1.2.16 h1:QlObi6ZIK5Ao7kAALnh91HWYNZUBbVwye52fmlQM9kc= -github.com/uptrace/bun v1.2.16/go.mod h1:jMoNg2n56ckaawi/O/J92BHaECmrz6IRjuMWqlMaMTM= -github.com/uptrace/bun/dialect/pgdialect v1.2.16 h1:KFNZ0LxAyczKNfK/IJWMyaleO6eI9/Z5tUv3DE1NVL4= -github.com/uptrace/bun/dialect/pgdialect v1.2.16/go.mod h1:IJdMeV4sLfh0LDUZl7TIxLI0LipF1vwTK3hBC7p5qLo= -github.com/uptrace/bun/driver/pgdriver v1.2.16 h1:b1kpXKUxtTSGYow5Vlsb+dKV3z0R7aSAJNfMfKp61ZU= -github.com/uptrace/bun/driver/pgdriver v1.2.16/go.mod h1:H6lUZ9CBfp1X5Vq62YGSV7q96/v94ja9AYFjKvdoTk0= github.com/uptrace/bunrouter v1.0.23 h1:Bi7NKw3uCQkcA/GUCtDNPq5LE5UdR9pe+UyWbjHB/wU= github.com/uptrace/bunrouter v1.0.23/go.mod h1:O3jAcl+5qgnF+ejhgkmbceEk0E/mqaK+ADOocdNpY8M= github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8= @@ -403,26 +391,28 @@ github.com/xdg-go/stringprep v1.0.4 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6 github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM= github.com/xraph/confy v1.0.2 h1:90jmVLdw9J0uqJOnfDxqatOVJHeZ/IM1R89zd3df1FA= github.com/xraph/confy v1.0.2/go.mod h1:/jKqCF8cMpCatuNO2uFQ/7VClDBP2+V74fM5KFM42o8= -github.com/xraph/forge v1.9.2 h1:Z/da9yvAdVyV5yL+4CRDE8L4WZkirvKfDbqdmCy50XM= -github.com/xraph/forge v1.9.2/go.mod h1:7YxmuSpTIqoFbNNZXbLRsH0gmRhOv/PbeRpp0wNpZoE= +github.com/xraph/forge v1.9.8 h1:GkWQm/bj58FGMTF07qn3v8AMw0QvXD0XsfCyY74tmWE= +github.com/xraph/forge v1.9.8/go.mod h1:/AoHxHQtooM3LMsmlsYA+zM9g2DgsYFPKV+lV/eJQ2c= github.com/xraph/forgeui v1.4.1 h1:LHK1t/sZ+9zL+MNUZralO9/rc0f5UCa19dpbWTuRMNg= github.com/xraph/forgeui v1.4.1/go.mod h1:rH/+wb1tt2pXSHotWAvoP+Lt846xlIjuwPDSpS5K5mw= github.com/xraph/go-utils v1.1.6 h1:eFN4dnTnC4+dzbD+6/4Bhlx7fGawi3flsLEVx9YFakQ= github.com/xraph/go-utils v1.1.6/go.mod h1:Mckdi+nR0bI4bUESKSYajJq4tNSPsvZiuLRYJ0+qDQw= -github.com/xraph/grove v1.6.0 h1:r6AZmOHlQtu1Ech589DTtQGmMcapMCX0qXuWnKymPzU= -github.com/xraph/grove v1.6.0/go.mod h1:bgjHNhnmyfEyzbdpcppRt+Zf24nNcbGKlo450Mi4giI= -github.com/xraph/grove/drivers/mongodriver v1.6.0 h1:/VbgV5ZdVTfSpb3Vv1UY4w4npURGvhnXgEJz3xj1g3I= -github.com/xraph/grove/drivers/mongodriver v1.6.0/go.mod h1:xojoSuw3qSm3NIe3xBmasocRQ7oyju3XbjvdlCXVXIU= -github.com/xraph/grove/drivers/pgdriver v1.6.0 h1:MQct9OUbGnr5UJ29LfjiBnquWzMj9jiAt7n9NgkZABE= -github.com/xraph/grove/drivers/pgdriver v1.6.0/go.mod h1:BBCrIuCBrRajE/dB5IKC0HshU/bUHPNS/etWtmX19qo= -github.com/xraph/grove/drivers/sqlitedriver v1.6.0 h1:KZx6A0mO6q5+v3Z4T2m+9dMmQLBMzQCuxdqfTNAGVLw= -github.com/xraph/grove/drivers/sqlitedriver v1.6.0/go.mod h1:xzHewWROOPVn0Luu8/sWEAhSgmDnw3xmNrRw0xcefnM= -github.com/xraph/grove/kv v1.6.0 h1:2C6dZ7G1AxmzV5przqM8S3RToNmjUR5ZcRyRR1QkKow= -github.com/xraph/grove/kv v1.6.0/go.mod h1:rdwh/Ja2WOhZ3YwrIp8vBeDNATAOj+2vrK+XVuRHrJU= -github.com/xraph/grove/kv/drivers/redisdriver v1.6.0 h1:ums0CN5/pnWpED17CtGHUx2UFg3WCboBJN0Zdw940Tc= -github.com/xraph/grove/kv/drivers/redisdriver v1.6.0/go.mod h1:15TFWsrEvCTHKiqgn40x06VzQTFEaiWyjpVoA2m3FFU= -github.com/xraph/relay v1.6.0 h1:EeXiBCjPOLRPVhuvIwP8pxQlEONBmwPTiCxt24cHXhI= -github.com/xraph/relay v1.6.0/go.mod h1:1Mhdv06q+jV3ahcTrKdhZCrIiFPMtdn2O2TGlSVWL4E= +github.com/xraph/grove v1.6.1 h1:mlQM7j4yAbOoPIOE7ut9lQbDq2zhOGKyeB00T74XjuM= +github.com/xraph/grove v1.6.1/go.mod h1:bgjHNhnmyfEyzbdpcppRt+Zf24nNcbGKlo450Mi4giI= +github.com/xraph/grove/drivers/mongodriver v1.6.1 h1:bRCsHqXtFeY8cX4mRNYbxyTfhoxrkfNmULAayOiyciE= +github.com/xraph/grove/drivers/mongodriver v1.6.1/go.mod h1:8Y05f6EATbbeMNoSiX2UrdxOVMsT2P740I0oFvltP0Y= +github.com/xraph/grove/drivers/pgdriver v1.6.1 h1:aqkxGL4EdIpDzaB882CGEcGDwVmmHLf4zhVylfxEgh8= +github.com/xraph/grove/drivers/pgdriver v1.6.1/go.mod h1:LXpNNoIOVci3uFXfwYp0Ovzm42hq5OukTKO0JPwGixE= +github.com/xraph/grove/drivers/sqlitedriver v1.6.1 h1:C6nRXmr6t6urDZdnWMD+HO7K0nOqQYcNZ6M+iG1k5hA= +github.com/xraph/grove/drivers/sqlitedriver v1.6.1/go.mod h1:xzHewWROOPVn0Luu8/sWEAhSgmDnw3xmNrRw0xcefnM= +github.com/xraph/grove/kv v1.6.1 h1:TnavtgRqa4eqV2AZgjHVCxJaR7XZLqKJwE8wJVu/hNo= +github.com/xraph/grove/kv v1.6.1/go.mod h1:rdwh/Ja2WOhZ3YwrIp8vBeDNATAOj+2vrK+XVuRHrJU= +github.com/xraph/grove/kv/drivers/redisdriver v1.6.1 h1:Mlv8c8Xlu2xTKE65lWmZuyC6V0WXb/rq9zOnlEcdGMQ= +github.com/xraph/grove/kv/drivers/redisdriver v1.6.1/go.mod h1:aZxHux8IHL1+LMR+aZC6yU4iHhb6JXg0vSPyfdq5vz4= +github.com/xraph/relay v1.6.2 h1:V4xYS5fTxo0ihvi+XLUc9AhxfiJ0auvys97ioFd+DhU= +github.com/xraph/relay v1.6.2/go.mod h1:dUMwMP2i8Sc8zJlQlrg4xm3fps4l4A5LSpzbP8LdLok= +github.com/xraph/trove v1.6.4 h1:OeIeAO+58BK8YLqIUIMEMAyZpSNJEKOOc1b1+LMXX24= +github.com/xraph/trove v1.6.4/go.mod h1:bLs3RZEd6OlBY4xjmLaaacNKR7M5RdM75GQvEIAPVCo= github.com/xraph/vessel v1.0.4 h1:YVG80hm7bAKTwogwLXZMlp46tMHHJThUcvCUOWao9Yk= github.com/xraph/vessel v1.0.4/go.mod h1:5hgrMbuczxu2kRIM3iVJ3wPSb6HOvbMhV4nkRFaPNqs= github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 h1:ilQV1hzziu+LLM3zUTJ0trRztfwgjqKnBWNtSRkbmwM= @@ -430,16 +420,22 @@ github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78/go.mod h1:aL8wCCfTfS github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0= github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +github.com/zeebo/assert v1.1.0 h1:hU1L1vLTHsnO8x8c9KAR5GmM5QscxHg5RNU5z5qbUWY= +github.com/zeebo/assert v1.1.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= +github.com/zeebo/blake3 v0.2.4 h1:KYQPkhpRtcqh0ssGYcKLG1JYvddkEA8QwCM/yBqhaZI= +github.com/zeebo/blake3 v0.2.4/go.mod h1:7eeQ6d2iXWRGF6npfaxl2CU+xy2Fjo2gxeyZGCRUjcE= +github.com/zeebo/pcg v1.0.1 h1:lyqfGeWiv4ahac6ttHs+I5hwtH/+1mrhlCtVNQM2kHo= +github.com/zeebo/pcg v1.0.1/go.mod h1:09F0S9iiKrwn9rlI5yjLkmrug154/YRW6KnnXVDM/l4= github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= go.jetify.com/typeid/v2 v2.0.0-alpha.3 h1:T6RPx6bNl10lp0JN2Xz/XcgLZWSlVmL58Xqy9cgTCcc= go.jetify.com/typeid/v2 v2.0.0-alpha.3/go.mod h1:zfD1ZDHDJNgXZANsO9jDOD81XRRQ0zAOnDBEHmIV/Gw= -go.mongodb.org/mongo-driver/v2 v2.5.0 h1:yXUhImUjjAInNcpTcAlPHiT7bIXhshCTL3jVBkF3xaE= -go.mongodb.org/mongo-driver/v2 v2.5.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0= +go.mongodb.org/mongo-driver/v2 v2.8.0 h1:CxWDGQYY8QQwNjAl/aq2sfWakdnWZynnqJ9F4DhHbP8= +go.mongodb.org/mongo-driver/v2 v2.8.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRNDSWJOTobXh5HyQKjq6wUC5tNybqjIqDpAY4CU= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0/go.mod h1:69uWxva0WgAA/4bu2Yy70SLDBwZXuQ6PbBpbsa5iZrQ= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0 h1:8tvICD4vSTOOsNrsI4Ljf6C+6UKvpTEH5XY3JMoyPoo= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0/go.mod h1:z9+yiacE0IHRqM4qFfkbt/JYlmYXgss8GY/jXoNuPJI= go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= go.opentelemetry.io/otel/exporters/jaeger v1.17.0 h1:D7UpUy2Xc2wsi1Ras6V40q806WM07rqoCWzXu7Sqy+4= @@ -450,10 +446,10 @@ go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.43.0 h1:3iZJK go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.43.0/go.mod h1:/G+nUPfhq2e+qiXMGxMwumDrP5jtzU+mWN7/sjT2rak= go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= -go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= -go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg= -go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= -go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= +go.opentelemetry.io/otel/sdk v1.44.0 h1:nHYwb9lK+fJPU/dnT6s7W7Z8itMWyqrnVfbheVYrZ58= +go.opentelemetry.io/otel/sdk v1.44.0/go.mod h1:Osuydd3Se74nqjAKxid74N5eC+jfEqfTegHRnq58oK0= +go.opentelemetry.io/otel/sdk/metric v1.44.0 h1:3LlKgI+VjbVsjNRFZJZAJ30WjXC5VkNRks6si09iEfI= +go.opentelemetry.io/otel/sdk/metric v1.44.0/go.mod h1:5B5pMARnXxKhltooO4xUuCBorl65a4EpnTalObqOigA= go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= go.opentelemetry.io/proto/otlp v1.10.0 h1:IQRWgT5srOCYfiWnpqUYz9CVmbO8bFmKcwYxpuCSL2g= @@ -515,6 +511,7 @@ golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= @@ -579,8 +576,6 @@ k8s.io/kube-openapi v0.0.0-20260721132016-d427ff9ee9ad h1:oXImqH8mQNk7PmvzKhmN3d k8s.io/kube-openapi v0.0.0-20260721132016-d427ff9ee9ad/go.mod h1:0/mqHCVhlumdJ3BhCfnjSZQE037nAhNodh1/hK0T8/I= k8s.io/utils v0.0.0-20260707023825-cf1189d6abe3 h1:jVkFFVfXdXP74B/zbO3hM3hpSFD0xvhQ5U686DPurkE= k8s.io/utils v0.0.0-20260707023825-cf1189d6abe3/go.mod h1:M2s5JB1lIYP3jzZdorPLHXIPJzt9vv2muW5a6L9DtNM= -mellium.im/sasl v0.3.2 h1:PT6Xp7ccn9XaXAnJ03FcEjmAn7kK1x7aoXV6F+Vmrl0= -mellium.im/sasl v0.3.2/go.mod h1:NKXDi1zkr+BlMHLQjY3ofYuU4KSPFxknb8mfEu6SveY= modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis= modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0= modernc.org/ccgo/v4 v4.30.2 h1:4yPaaq9dXYXZ2V8s1UgrC3KIj580l2N4ClrLwnbv2so= diff --git a/id/id.go b/id/id.go index ba2a605..804ec32 100644 --- a/id/id.go +++ b/id/id.go @@ -32,6 +32,7 @@ const ( PrefixDLQ Prefix = "dlq" PrefixEvent Prefix = "evt" PrefixWorker Prefix = "wkr" + PrefixArtifact Prefix = "art" ) // ID is the primary identifier type for all Dispatch entities. @@ -136,6 +137,9 @@ type EventID = ID // WorkerID is a type-safe identifier for workers (prefix: "wkr"). type WorkerID = ID +// ArtifactID is a type-safe identifier for artifacts (prefix: "art"). +type ArtifactID = ID + // AnyID is a type alias that accepts any valid prefix. type AnyID = ID @@ -167,6 +171,9 @@ func NewEventID() ID { return New(PrefixEvent) } // NewWorkerID generates a new unique worker ID. func NewWorkerID() ID { return New(PrefixWorker) } +// NewArtifactID generates a new unique artifact ID. +func NewArtifactID() ID { return New(PrefixArtifact) } + // ────────────────────────────────────────────────── // Convenience parsers // ────────────────────────────────────────────────── @@ -195,6 +202,9 @@ func ParseEventID(s string) (ID, error) { return ParseWithPrefix(s, PrefixEvent) // ParseWorkerID parses a string and validates the "wkr" prefix. func ParseWorkerID(s string) (ID, error) { return ParseWithPrefix(s, PrefixWorker) } +// ParseArtifactID parses a string and validates the "art" prefix. +func ParseArtifactID(s string) (ID, error) { return ParseWithPrefix(s, PrefixArtifact) } + // ParseAny parses a string into an ID without type checking the prefix. func ParseAny(s string) (ID, error) { return Parse(s) } diff --git a/id/id_test.go b/id/id_test.go index 141c0d7..c2dfbff 100644 --- a/id/id_test.go +++ b/id/id_test.go @@ -276,3 +276,25 @@ func TestBSONUnmarshalInvalidType(t *testing.T) { t.Error("expected error for invalid BSON type, got nil") } } + +func TestArtifactID(t *testing.T) { + got := id.NewArtifactID() + if got.Prefix() != id.PrefixArtifact { + t.Fatalf("prefix = %q, want %q", got.Prefix(), id.PrefixArtifact) + } + if got.IsNil() { + t.Fatal("NewArtifactID returned nil ID") + } + + parsed, err := id.ParseArtifactID(got.String()) + if err != nil { + t.Fatalf("ParseArtifactID(%q) error = %v", got.String(), err) + } + if parsed.String() != got.String() { + t.Fatalf("round trip = %q, want %q", parsed.String(), got.String()) + } + + if _, err := id.ParseArtifactID("job_01h2xcejqtf2nbrexx3vqjhp41"); err == nil { + t.Fatal("ParseArtifactID accepted a job ID, want error") + } +} diff --git a/job/dequeue_opts_test.go b/job/dequeue_opts_test.go new file mode 100644 index 0000000..282ae19 --- /dev/null +++ b/job/dequeue_opts_test.go @@ -0,0 +1,429 @@ +package job_test + +import ( + "errors" + "testing" + "time" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +func TestDequeueOptsIsUnbounded(t *testing.T) { + reserved := id.NewJobID() + + tests := []struct { + name string + opts job.DequeueOpts + want bool + }{ + {"zero value", job.DequeueOpts{}, true}, + {"queues and limit only", job.DequeueOpts{Queues: []string{"default"}, Limit: 8}, true}, + {"nil budget map", job.DequeueOpts{Budget: nil}, true}, + {"empty budget map", job.DequeueOpts{Budget: resource.Set{}}, true}, + // An explicit zero is a worker with nothing free, not an absent + // constraint. Treating it as unbounded would hand that worker the + // whole queue. + {"explicit zero budget key", job.DequeueOpts{Budget: resource.Set{resource.Memory: 0}}, false}, + {"budget", job.DequeueOpts{Budget: resource.Set{resource.CPU: 1000}}, false}, + {"custom keys", job.DequeueOpts{CustomKeys: []string{"fpga"}}, false}, + // Locality answers "should I order?", never "should I filter?". + // Counting it here would make PreferHashes-only opts bounded, and + // the empty-CustomKeys rule would then reject custom-resource jobs + // on its behalf. + {"prefer hashes alone", job.DequeueOpts{PreferHashes: []string{"blake3:a"}}, true}, + { + "prefer hashes with a budget", + job.DequeueOpts{ + Budget: resource.Set{resource.CPU: 1000}, + PreferHashes: []string{"blake3:a"}, + }, + false, + }, + {"reserved for", job.DequeueOpts{ReservedFor: &reserved}, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.opts.IsUnbounded(); got != tt.want { + t.Errorf("IsUnbounded() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestDequeueOptsAllows(t *testing.T) { + const gib = int64(1) << 30 + + other := id.NewJobID() + + newJob := func(res resource.Set) *job.Job { + return &job.Job{ID: id.NewJobID(), Resources: res} + } + + tests := []struct { + name string + opts job.DequeueOpts + j *job.Job + want bool + }{ + { + "zero opts allow an oversized job", + job.DequeueOpts{}, + newJob(resource.Set{resource.Memory: 512 * gib, "fpga": 4}), + true, + }, + { + "exact fit is allowed", + job.DequeueOpts{Budget: resource.Set{resource.Memory: 4 * gib}}, + newJob(resource.Set{resource.Memory: 4 * gib}), + true, + }, + { + "one byte over is not", + job.DequeueOpts{Budget: resource.Set{resource.Memory: 4 * gib}}, + newJob(resource.Set{resource.Memory: 4*gib + 1}), + false, + }, + { + "absent budget key is unconstrained", + job.DequeueOpts{Budget: resource.Set{resource.Memory: 4 * gib}}, + newJob(resource.Set{resource.GPU: 8000}), + true, + }, + { + "explicit zero budget key filters", + job.DequeueOpts{Budget: resource.Set{resource.GPU: 0}}, + newJob(resource.Set{resource.GPU: 1}), + false, + }, + { + "zero requirement fits a zero budget", + job.DequeueOpts{Budget: resource.Set{resource.GPU: 0}}, + newJob(nil), + true, + }, + { + "custom key not offered", + job.DequeueOpts{CustomKeys: []string{"tpu"}}, + newJob(resource.Set{"fpga": 1}), + false, + }, + { + "custom key offered", + job.DequeueOpts{CustomKeys: []string{"tpu"}}, + newJob(resource.Set{"tpu": 4}), + true, + }, + { + "custom quantity is not compared at dequeue", + job.DequeueOpts{CustomKeys: []string{"tpu"}, Budget: resource.Set{"tpu": 1}}, + newJob(resource.Set{"tpu": 64}), + true, + }, + { + "prefix does not match a longer offered key", + job.DequeueOpts{CustomKeys: []string{"fpga-large"}}, + newJob(resource.Set{"fpga": 1}), + false, + }, + { + "prefix does not match a shorter offered key", + job.DequeueOpts{CustomKeys: []string{"fpga"}}, + newJob(resource.Set{"fpga-large": 1}), + false, + }, + { + "subset of an interleaved offer", + job.DequeueOpts{CustomKeys: []string{"fpga", "nvme", "tpu"}}, + newJob(resource.Set{"fpga": 1, "tpu": 1}), + true, + }, + { + "partial offer is not enough", + job.DequeueOpts{CustomKeys: []string{"fpga"}}, + newJob(resource.Set{"fpga": 1, "tpu": 1}), + false, + }, + { + // Bounded opts with an empty offer: the caller is + // resource-aware and has no custom resources. + "empty offer on bounded opts rejects a custom requirement", + job.DequeueOpts{Budget: resource.Set{resource.Memory: 4 * gib}}, + newJob(resource.Set{"fpga": 1}), + false, + }, + { + // Unbounded opts with an empty offer: the caller does not use + // the resource model, and must keep claiming what it always did. + "empty offer on unbounded opts claims a custom requirement", + job.DequeueOpts{}, + newJob(resource.Set{"fpga": 1}), + true, + }, + { + "a zero-quantity custom key is not a requirement", + job.DequeueOpts{CustomKeys: []string{"tpu"}}, + newJob(resource.Set{"fpga": 0}), + true, + }, + { + // PreferHashes must never filter, not even transitively through + // the empty-CustomKeys rule. + "prefer hashes alone claims a custom requirement", + job.DequeueOpts{PreferHashes: []string{"blake3:a"}}, + newJob(resource.Set{"fpga": 1, resource.Memory: 1 << 40}), + true, + }, + { + "reserved for another job", + job.DequeueOpts{ReservedFor: &other}, + newJob(nil), + false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.opts.Allows(tt.j); got != tt.want { + t.Errorf("Allows() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestDequeueOptsAllowsReservedJob(t *testing.T) { + j := &job.Job{ID: id.NewJobID(), Resources: resource.Set{resource.Memory: 8}} + opts := job.DequeueOpts{ReservedFor: &j.ID} + + if !opts.Allows(j) { + t.Error("Allows(reserved job) = false, want true") + } + + // A reservation must not be able to bypass the budget, or the one path + // a scheduler uses to place a large job is the path with no ceiling. + opts.Budget = resource.Set{resource.Memory: 4} + if opts.Allows(j) { + t.Error("Allows(reserved job over budget) = true, want false") + } +} + +func TestDequeueOptsLess(t *testing.T) { + base := time.Now().UTC() + + mk := func(priority int, offset time.Duration, hash string) *job.Job { + return &job.Job{ + ID: id.NewJobID(), + Priority: priority, + RunAt: base.Add(offset), + PrimaryInputHash: hash, + } + } + + opts := job.DequeueOpts{PreferHashes: []string{"blake3:local"}} + + highRemote := mk(9, time.Minute, "blake3:remote") + lowLocal := mk(1, 0, "blake3:local") + lowRemoteEarly := mk(1, time.Second, "") + lowRemoteLate := mk(1, time.Minute, "") + + // Every rule is asserted in BOTH directions, and that is not + // belt-and-braces: `func Less(a, b *Job) bool { return true }` + // satisfies every one-directional assertion here, and a comparator + // that is true both ways is not an ordering at all — sort.SliceStable + // would return an arbitrary permutation and the backends that sort in + // Go would hand back arbitrary work. + // + // The equal case is asserted for the same reason: Less must be a + // strict weak ordering, so two jobs that tie on all three terms must + // report false in both directions. + for _, tc := range []struct { + name string + hi *job.Job + lo *job.Job + }{ + // Priority outranks locality: a stream of locally cached work + // must not be able to starve a high-priority job. + {"priority outranks locality", highRemote, lowLocal}, + // Within a priority band, locality wins even against an earlier + // RunAt. + {"locality outranks RunAt within a band", lowLocal, lowRemoteEarly}, + // Beyond that, RunAt ascending. + {"RunAt breaks the remaining tie", lowRemoteEarly, lowRemoteLate}, + } { + t.Run(tc.name, func(t *testing.T) { + if !opts.Less(tc.hi, tc.lo) { + t.Errorf("Less(%s) = false, want true", tc.name) + } + + if opts.Less(tc.lo, tc.hi) { + t.Errorf("Less is true in BOTH directions for %s, which is not an ordering", + tc.name) + } + }) + } + + // A genuine tie: same priority, same locality answer, same RunAt. + tieA := mk(1, time.Second, "") + tieB := mk(1, time.Second, "") + + if opts.Less(tieA, tieB) || opts.Less(tieB, tieA) { + t.Error("Less reports an order between two jobs that tie on every term; " + + "a comparator that never returns false for equal elements is not a strict " + + "weak ordering and sort will produce an arbitrary permutation") + } + + // Two jobs the caller has BOTH staged tie on locality, so RunAt + // decides — locality is a boolean, never a ranking among preferred + // jobs. + bothLocalEarly := mk(1, 0, "blake3:local") + bothLocalLate := mk(1, time.Minute, "blake3:local") + + if !opts.Less(bothLocalEarly, bothLocalLate) || opts.Less(bothLocalLate, bothLocalEarly) { + t.Error("two equally preferred jobs must be separated by RunAt, not by hash value") + } + + if opts.Prefers(mk(1, 0, "")) { + t.Error("Prefers(job with no hash) = true, want false") + } +} + +// TestDequeueOptsPreferredHashes pins the normalisation every backend +// binds instead of PreferHashes. +// +// The empty entry is the one that matters. primary_input_hash is a plain +// string column, so a job that was never hashed stores ” rather than +// NULL, and an empty entry bound verbatim makes `” = ANY('{""}')` true +// for every one of them. Under a tight Limit that is not a reordering +// but a filter: the jobs the caller actually staged stop being claimed. +func TestDequeueOptsPreferredHashes(t *testing.T) { + opts := job.DequeueOpts{ + PreferHashes: []string{"blake3:b", "", "blake3:a", "blake3:b", ""}, + } + + got := opts.PreferredHashes() + want := []string{"blake3:b", "blake3:a"} + + if len(got) != len(want) { + t.Fatalf("PreferredHashes() = %v, want %v", got, want) + } + + for i := range want { + if got[i] != want[i] { + t.Fatalf("PreferredHashes() = %v, want %v (caller order, deduplicated, no empties)", + got, want) + } + } + + // A list of nothing but empties carries no locality signal at all, + // and must not leave a backend emitting a term that matches every + // unhashed job. + empties := job.DequeueOpts{PreferHashes: []string{"", ""}} + if got = empties.PreferredHashes(); len(got) != 0 { + t.Errorf("PreferredHashes() = %v for a list of empty strings, want none", got) + } + + var zero job.DequeueOpts + if zero.PreferredHashes() != nil { + t.Error("PreferredHashes() on zero opts must be nil") + } +} + +func TestDequeueOptsOfferedCustomKeys(t *testing.T) { + opts := job.DequeueOpts{CustomKeys: []string{"tpu", "fpga", "tpu", ""}} + + got := opts.OfferedCustomKeys() + want := []string{"fpga", "tpu"} + + if len(got) != len(want) { + t.Fatalf("OfferedCustomKeys() = %v, want %v", got, want) + } + + for i := range want { + if got[i] != want[i] { + t.Fatalf("OfferedCustomKeys() = %v, want %v", got, want) + } + } + + if none := (job.DequeueOpts{}).OfferedCustomKeys(); none != nil { + t.Errorf("OfferedCustomKeys() on zero opts = %v, want nil", none) + } +} + +func TestDequeueOptsGrants(t *testing.T) { + worker := id.NewWorkerID() + until := time.Date(2026, 8, 12, 12, 0, 0, 0, time.UTC) + + tests := []struct { + name string + opts job.DequeueOpts + want bool + }{ + // The zero value is the pool as it is configured today, and it + // must never grant: that is the backward-compatibility guarantee + // every backend reads off this method. + {"zero value", job.DequeueOpts{}, false}, + {"queues and limit only", job.DequeueOpts{Queues: []string{"default"}, Limit: 8}, false}, + // A worker id alone is not a request for a lease. Grants keys on + // LeaseUntil only, so a caller that identifies itself without + // asking for a lease still gets its lease columns left alone. + {"worker without expiry", job.DequeueOpts{WorkerID: worker}, false}, + {"expiry alone", job.DequeueOpts{LeaseUntil: until}, true}, + {"worker and expiry", job.DequeueOpts{WorkerID: worker, LeaseUntil: until}, true}, + // A past expiry is a real grant, not a no-op: it is how a caller + // hands over a job it wants reclaimed at the next sweep, and the + // conformance suite's fencing case relies on it. + { + "expiry in the past", + job.DequeueOpts{WorkerID: worker, LeaseUntil: until.Add(-time.Hour)}, + true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.opts.Grants(); got != tt.want { + t.Errorf("Grants() = %v, want %v", got, tt.want) + } + }) + } +} + +// TestDequeueOptsValidate pins the one input every backend must refuse. +// +// A lease granted to the zero worker can never be renewed, because +// RenewLease matches on worker ID — so the job would be claimed, expire, +// be reclaimed, and go round again forever. That is a queue that never +// drains rather than an error, which is why it is rejected here instead +// of tolerated. +func TestDequeueOptsValidate(t *testing.T) { + worker := id.NewWorkerID() + until := time.Date(2026, 8, 12, 12, 0, 0, 0, time.UTC) + + tests := []struct { + name string + opts job.DequeueOpts + want error + }{ + {"zero value", job.DequeueOpts{}, nil}, + // No grant means WorkerID is ignored, both ways round. + {"no grant, no worker", job.DequeueOpts{Queues: []string{"default"}, Limit: 1}, nil}, + {"no grant, worker set", job.DequeueOpts{WorkerID: worker}, nil}, + {"grant with worker", job.DequeueOpts{WorkerID: worker, LeaseUntil: until}, nil}, + {"grant without worker", job.DequeueOpts{LeaseUntil: until}, job.ErrLeaseWithoutWorker}, + { + "grant with explicitly zero worker", + job.DequeueOpts{WorkerID: id.WorkerID{}, LeaseUntil: until}, + job.ErrLeaseWithoutWorker, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.opts.Validate(); !errors.Is(got, tt.want) { + t.Errorf("Validate() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/job/errors.go b/job/errors.go new file mode 100644 index 0000000..2442fab --- /dev/null +++ b/job/errors.go @@ -0,0 +1,26 @@ +package job + +import "errors" + +// Lease sentinels. +// +// These live in job rather than the root dispatch package — where every +// other sentinel lives — because the root package already exports +// ErrLeadershipLost for cluster leadership. A sibling ErrLeaseLost for job +// leases one line away would be a standing invitation to grab the wrong +// one. Qualified as job.ErrLeaseLost, the call site is unambiguous. +var ( + // ErrLeaseLost means the worker no longer holds the lease it tried to + // act on: the job was reclaimed, reassigned, or deleted while the + // worker believed it was still running it. A worker receiving this must + // stop working on the job immediately — someone else owns it now. + ErrLeaseLost = errors.New("dispatch/job: lease lost") + + // ErrLeaseWithoutWorker means a dequeue asked for a lease + // (DequeueOpts.LeaseUntil) without naming the worker that would hold + // it. It is a programming error, not a degenerate case: RenewLease + // matches on worker ID, so a lease held by the zero worker can never + // be renewed and the job would be claimed and reclaimed on every + // cycle forever. Backends refuse the claim rather than granting it. + ErrLeaseWithoutWorker = errors.New("dispatch/job: DequeueOpts.LeaseUntil set without WorkerID") +) diff --git a/job/job.go b/job/job.go index 4996b69..7ab647a 100644 --- a/job/job.go +++ b/job/job.go @@ -5,6 +5,7 @@ import ( "github.com/xraph/dispatch" "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" ) // State represents the lifecycle state of a job. @@ -46,4 +47,91 @@ type Job struct { CompletedAt *time.Time `json:"completed_at,omitempty"` HeartbeatAt *time.Time `json:"heartbeat_at,omitempty"` Timeout time.Duration `json:"timeout,omitempty"` + + // ArtifactBindings carries the encoded map of declared input names to + // artifact refs. It travels with the job because Payload is opaque to + // the engine: bindings placed inside it would be invisible to the + // scheduler and to the staging middleware. + ArtifactBindings []byte `json:"artifact_bindings,omitempty"` + + // Resources is the resolved requirement, computed once at enqueue. + // Scheduling reads this rather than calling user code. + Resources resource.Set `json:"resources,omitempty"` + + // ResourceLimits is the resolved enforcement ceiling. + ResourceLimits resource.Set `json:"resource_limits,omitempty"` + + // ResourceClass is forwarded to the isolation backend uninterpreted. + ResourceClass string `json:"resource_class,omitempty"` + + // InputBytes is the total size of the declared artifact inputs. It + // is the estimator's primary feature and the measurement bucket key. + InputBytes int64 `json:"input_bytes,omitempty"` + + // PrimaryInputHash is the content hash of the largest declared + // input, used as the locality-scheduling signal. Often empty: the + // artifact plane fills content_hash at first staging, not at + // registration, so locality helps from an artifact's second use on. + PrimaryInputHash string `json:"primary_input_hash,omitempty"` + + // LeaseEpoch is the fencing token for the current lease. It increments + // on every grant and every reclamation. RenewLease, the grant inside + // DequeueJobs, ReclaimExpiredLeases, and UpdateLeasedJob all check + // it, so a worker holding a stale epoch fails to renew — and the pool + // cancels the job within one heartbeat interval — or has its terminal + // write refused outright. UpdateJob does not check it: that is a + // whole-row write with no epoch predicate, so a caller that wants the + // fence must use UpdateLeasedJob instead. + // + // Without the renewal check, a worker resuming from a long GC pause + // would carry on renewing a lease on a job another worker now owns. + LeaseEpoch int `json:"lease_epoch"` + + // LeaseExpiresAt is when the current lease lapses if not renewed. + // Nil means no lease is held. + LeaseExpiresAt *time.Time `json:"lease_expires_at,omitempty"` + + // LeaseTTL is how long each renewal extends the lease for this job, + // copied at enqueue from the definition's WithLeaseTTL, or from the + // enqueue site's when that supplies one. Zero means the pool's + // default. + // + // This is what makes per-definition thresholds work: a 30-second job + // and a six-hour job carry different values on their own rows, so one + // reclaim query serves both. + LeaseTTL time.Duration `json:"lease_ttl,omitempty"` + + // EvictCount is how many times this job has lost a worker to + // infrastructure — a reclaimed lease, or later a graceful drain. It is + // deliberately separate from RetryCount: a preempted job has not + // failed, and charging preemptions to the retry budget would send a + // healthy job to the DLQ having never once errored. + EvictCount int `json:"evict_count"` +} + +// ClearOwnership drops every field recording who was running the job and +// under what lease, so the row can safely go back to a runnable state. +// +// Any path returning a job to pending from outside the lease machinery +// has to call this, and forgetting to is not a cosmetic bug. UpdateJob +// writes the whole row on every backend, lease columns included, so a +// stale LeaseExpiresAt survives the transition. The row is then pending +// with an expiry already in the past, which is harmless right up until +// the job is claimed by a caller that does not grant a lease +// (DequeueOpts.Grants() is false whenever LeaseUntil is zero). That claim +// writes state and worker but never touches the expiry, so the job lands +// in running carrying a lapsed lease and the very next sweep reclaims it. +// It is claimed and reclaimed forever, never running to completion, with +// EvictCount climbing on every pass. +// +// LeaseEpoch is deliberately left alone. It is a fencing token and must +// never move backwards, and there is nothing to fence here: a job being +// returned to pending by this path is not running, so no holder exists to +// invalidate. ReclaimExpiredLeases increments it because it is taking the +// job away from a live holder, which is a different situation. +func (j *Job) ClearOwnership() { + j.WorkerID = id.WorkerID{} + j.StartedAt = nil + j.HeartbeatAt = nil + j.LeaseExpiresAt = nil } diff --git a/job/lease.go b/job/lease.go new file mode 100644 index 0000000..623c28c --- /dev/null +++ b/job/lease.go @@ -0,0 +1,74 @@ +package job + +import ( + "time" +) + +// DefaultLeaseTTL is how long a lease survives without renewal when +// neither the definition nor the pool specifies otherwise. It matches the +// historical Config.StaleJobThreshold so that adopting leases does not +// change reclamation timing for an existing deployment. +const DefaultLeaseTTL = 30 * time.Second + +// UnleasedReclaimGrace is how long a running job carrying no lease at all +// must have been silent before reclamation will adopt it. +// +// It exists because a null expiry has two very different causes and the +// reclaim predicate cannot tell them apart from the expiry alone. One is a +// job left running by a build that predates leases, which nothing will +// ever look at again: Lease.IsExpired reports false for a zero expiry, the +// pool stopped calling ReapStaleJobs for a store implementing LeaseStore, +// and dequeue claims only pending and retrying rows. The other is a live, +// perfectly healthy job, because DequeueOpts.Grants() is false whenever +// LeaseUntil is zero, so any caller claiming through Store without lease +// options holds a running job with no lease by design. +// +// Silence is what separates them, which is why every backend gates the +// exception on heartbeat_at, falling back to started_at for a worker that +// died before its first beat, rather than on the null expiry alone. A row +// with neither timestamp is never adopted: there is nothing to measure age +// against, and guessing would mean guessing against a running job. +// +// The value is arbitrary and an operator cannot tune it, which is worth +// saying plainly rather than burying. ReclaimExpiredLeases(ctx, limit) +// carries no threshold, and widening that signature would change all five +// backends. Fifteen minutes is chosen to be conservative rather than +// precise: before leases these same rows were reaped at +// Config.StaleJobThreshold, 30 seconds by default, so any value well above +// that is strictly less aggressive than what already shipped. Overshooting +// costs only how long an abandoned job waits to come back. +const UnleasedReclaimGrace = 15 * time.Minute + +// Lease is the grant a worker holds over a running job. +// +// It carries only the expiry, because that is the single question the +// backends ask of it: IsExpired is the one authority on whether a lease +// has lapsed, and the zero-means-never-leased rule below is the reason +// that question cannot simply be asked of a time.Time inline. The holder +// and the fencing token live on the job row itself, as Job.WorkerID and +// Job.LeaseEpoch, which is where every writer already reads them. +type Lease struct { + // ExpiresAt is when the lease lapses if not renewed. A zero value + // means no lease has been granted. + ExpiresAt time.Time +} + +// IsExpired reports whether the lease has lapsed as of now. +// +// A zero ExpiresAt reports false: no lease was ever granted, which is +// "not held" rather than "expired". Reporting true would let the reclaim +// loop steal jobs that were never leased. +// +// Reclamation does eventually take such a job, but never through this +// function. It applies a separate and much coarser rule, gated on the row +// having gone silent for UnleasedReclaimGrace, precisely so that the +// question this function answers stays "is a lease lapsed" rather than +// blurring into "is a job abandoned". The two are not the same, and the +// backends depend on them staying separate. +func (l Lease) IsExpired(now time.Time) bool { + if l.ExpiresAt.IsZero() { + return false + } + + return !now.Before(l.ExpiresAt) +} diff --git a/job/lease_test.go b/job/lease_test.go new file mode 100644 index 0000000..66b7a11 --- /dev/null +++ b/job/lease_test.go @@ -0,0 +1,119 @@ +package job_test + +import ( + "testing" + "time" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +func TestWithLeaseTTL(t *testing.T) { + tests := []struct { + name string + give time.Duration + want time.Duration + }{ + {name: "positive is applied", give: 6 * time.Hour, want: 6 * time.Hour}, + // A zero or negative TTL would mean the lease expires the instant it + // is granted, so every job would be reclaimed before its first + // heartbeat and nothing would ever complete. Ignore it rather than + // persist it. + {name: "zero keeps the default", give: 0, want: 0}, + {name: "negative keeps the default", give: -time.Second, want: 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + opts := job.DefaultOptions() + job.WithLeaseTTL(tt.give)(&opts) + + if opts.LeaseTTL != tt.want { + t.Errorf("LeaseTTL = %v, want %v", opts.LeaseTTL, tt.want) + } + }) + } +} + +func TestDefaultOptions_LeaseTTLIsUnsetByDefault(t *testing.T) { + // Zero means "use the pool's default", which preserves the existing + // StaleJobThreshold semantics for every definition that says nothing. + if got := job.DefaultOptions().LeaseTTL; got != 0 { + t.Errorf("default LeaseTTL = %v, want 0", got) + } +} + +func TestLease_IsExpired(t *testing.T) { + base := time.Date(2026, 8, 12, 12, 0, 0, 0, time.UTC) + + tests := []struct { + name string + lease job.Lease + now time.Time + want bool + }{ + { + name: "not yet expired", + lease: job.Lease{ExpiresAt: base.Add(time.Minute)}, + now: base, + want: false, + }, + { + name: "expired", + lease: job.Lease{ExpiresAt: base.Add(-time.Second)}, + now: base, + want: true, + }, + { + name: "exactly at expiry is expired", + lease: job.Lease{ExpiresAt: base}, + now: base, + want: true, + }, + { + // A zero ExpiresAt means no lease was ever granted. Treating it + // as expired would let the reclaim loop steal jobs that were + // never leased, so it must read as "not held" rather than + // "expired". + name: "zero expiry is never expired", + lease: job.Lease{}, + now: base, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.lease.IsExpired(tt.now); got != tt.want { + t.Errorf("IsExpired(%v) = %v, want %v", tt.now, got, tt.want) + } + }) + } +} + +func TestJob_CarriesLeaseFields(t *testing.T) { + expires := time.Date(2026, 8, 12, 12, 0, 0, 0, time.UTC) + j := &job.Job{ + ID: id.NewJobID(), + LeaseEpoch: 3, + LeaseExpiresAt: &expires, + LeaseTTL: 6 * time.Hour, + EvictCount: 2, + } + + if j.ID.IsNil() { + t.Errorf("ID = nil, want a generated ID") + } + if j.LeaseEpoch != 3 { + t.Errorf("LeaseEpoch = %d, want 3", j.LeaseEpoch) + } + if j.LeaseExpiresAt == nil || !j.LeaseExpiresAt.Equal(expires) { + t.Errorf("LeaseExpiresAt = %v, want %v", j.LeaseExpiresAt, expires) + } + if j.LeaseTTL != 6*time.Hour { + t.Errorf("LeaseTTL = %v, want %v", j.LeaseTTL, 6*time.Hour) + } + if j.EvictCount != 2 { + t.Errorf("EvictCount = %d, want 2", j.EvictCount) + } +} diff --git a/job/options.go b/job/options.go index 01897a5..08615a4 100644 --- a/job/options.go +++ b/job/options.go @@ -1,6 +1,12 @@ package job -import "time" +import ( + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/resource" +) // Options configures per-job behavior such as retries, queue, and priority. type Options struct { @@ -16,8 +22,49 @@ type Options struct { // Timeout is the maximum duration a job may run before being cancelled. Timeout time.Duration + // LeaseTTL is how long this job's lease survives without renewal. + // Zero means the worker pool's default. Set it above the expected gap + // between heartbeats, not above the expected runtime — a lease is a + // liveness window, not a time limit. + LeaseTTL time.Duration + // RunAt schedules the job for future execution. Zero means immediate. RunAt time.Time + + // Inputs declares the artifacts this job consumes. Declaring them, + // rather than burying refs in the opaque payload, is what lets the + // engine size the job before scheduling it, validate bindings at + // enqueue, and stage the bytes before the handler runs. + Inputs []artifact.InputSpec + + // Bindings maps declared input names to the artifacts supplied at + // enqueue. The engine validates them against Inputs before the job is + // persisted. + Bindings map[string]artifact.Ref + + // Resources declares what this job needs. It is the floor: the + // engine may raise it via ResourceFunc or a configured estimator, + // and a per-enqueue WithResources call overrides both. + Resources resource.Set + + // ResourceLimits is the enforcement ceiling. When unset, memory and + // the other incompressible keys default to their request and CPU is + // left unbounded. + ResourceLimits resource.Set + + // ResourceFunc computes the requirement from the enqueue-time + // request, for jobs whose footprint scales with their input. It runs + // once, in the enqueuing process, and never on the scheduling path. + ResourceFunc resource.ResourceFunc + + // ResourceClass is an opaque scheduling class the isolation backend + // interprets. Core never reads it. + ResourceClass string + + // Execution declares the minimum isolation this job's handler + // requires. The zero value runs in-process, which is what every + // existing definition gets. + Execution exec.Policy } // DefaultOptions returns Options with sensible defaults. @@ -27,6 +74,7 @@ func DefaultOptions() Options { Queue: "default", Priority: 0, Timeout: 5 * time.Minute, + Execution: exec.NewPolicy(), } } @@ -67,3 +115,98 @@ func WithRunAt(t time.Time) Option { o.RunAt = t } } + +// WithArtifactInputs declares the artifact inputs a job consumes. +// +// The engine validates every binding against these declarations at +// enqueue and stages the declared inputs before the handler runs. +func WithArtifactInputs(specs ...artifact.InputSpec) Option { + return func(o *Options) { + o.Inputs = append(o.Inputs, specs...) + } +} + +// WithResources declares the resources a job needs. Multiple sets are +// merged per key, so the common form reads as a list: +// +// job.WithResources(resource.CPUs(4), resource.MemoryGB(16)) +// +// Passed at enqueue instead of on the definition, it overrides every +// other source, including a configured estimator. +func WithResources(sets ...resource.Set) Option { + return func(o *Options) { + for _, s := range sets { + if o.Resources == nil { + o.Resources = make(resource.Set, len(s)) + } + + for k, v := range s { + o.Resources[k] = v + } + } + } +} + +// WithResourceLimits sets the enforcement ceiling explicitly. +func WithResourceLimits(sets ...resource.Set) Option { + return func(o *Options) { + for _, s := range sets { + if o.ResourceLimits == nil { + o.ResourceLimits = make(resource.Set, len(s)) + } + + for k, v := range s { + o.ResourceLimits[k] = v + } + } + } +} + +// WithResourceFunc computes the requirement from the job's input. +// +// This is what lets one definition serve a 40 MB model and a 4 GB one: +// the artifact plane knows the input size at enqueue, so the function +// sees it before the job is ever scheduled. +func WithResourceFunc(fn resource.ResourceFunc) Option { + return func(o *Options) { o.ResourceFunc = fn } +} + +// WithResourceClass sets an opaque scheduling class for the isolation +// backend. Core stores and forwards it without interpretation. +func WithResourceClass(class string) Option { + return func(o *Options) { o.ResourceClass = class } +} + +// WithLeaseTTL sets how long this job's lease survives without renewal. +// +// It works on a definition and at an enqueue, and the enqueue wins: a +// definition declares the TTL its work needs, and a caller who knows this +// particular job is different overrides it, the same way WithResources +// overrides a declared requirement. +// +// A lease TTL is a liveness window, not a time limit: it should be a small +// multiple of the heartbeat interval regardless of how long the work takes. +// Non-positive durations are ignored, because a zero TTL would expire the +// lease the instant it was granted and the job would be reclaimed before +// its first heartbeat. +func WithLeaseTTL(d time.Duration) Option { + return func(o *Options) { + if d > 0 { + o.LeaseTTL = d + } + } +} + +// WithExecution declares the isolation this job's handler requires. +// +// It mirrors WithArtifactInputs: the exec package builds the value and +// job adapts it, which is what keeps exec a leaf that never imports job. +func WithExecution(opts ...exec.PolicyOption) Option { + return func(o *Options) { + p := o.Execution + for _, opt := range opts { + opt(&p) + } + o.Execution = p + } +} diff --git a/job/registrable.go b/job/registrable.go new file mode 100644 index 0000000..497ce5e --- /dev/null +++ b/job/registrable.go @@ -0,0 +1,32 @@ +package job + +import "github.com/xraph/dispatch/exec" + +// Registrable is a job definition that can register itself into a Registry +// without the caller knowing its payload type. +// +// Go forbids generic methods, but a method on a generic type is legal, so +// Definition[T] can satisfy this non-generic interface. That is what lets +// definitions with different payload types live in one slice — and a slice +// is what an out-of-process entrypoint can be handed, since it cannot be +// given the engine that would otherwise do the registering. +type Registrable interface { + // Register adds this definition's handler to the registry. + Register(r *Registry) + + // JobName returns the name the definition registers under. + JobName() string + + // Policy returns the execution declaration, so a caller can check + // that the deployment can satisfy it before registering anything. + Policy() exec.Policy +} + +// Register adds the definition's handler to the registry. +func (d *Definition[T]) Register(r *Registry) { RegisterDefinition(r, d) } + +// JobName returns the name this definition registers under. +func (d *Definition[T]) JobName() string { return d.Name } + +// Policy returns this definition's execution declaration. +func (d *Definition[T]) Policy() exec.Policy { return d.Opts.Execution } diff --git a/job/registrable_test.go b/job/registrable_test.go new file mode 100644 index 0000000..95d654b --- /dev/null +++ b/job/registrable_test.go @@ -0,0 +1,92 @@ +package job_test + +import ( + "context" + "testing" + "time" + + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/job" +) + +type meshPayload struct { + Detail int `json:"detail"` +} + +func TestDefinition_ImplementsRegistrable(t *testing.T) { + // The whole out-of-process design depends on this compiling: a + // heterogeneous slice of definitions with different payload types. + defs := []job.Registrable{ + job.NewDefinition("send-email", func(_ context.Context, _ emailPayload) error { return nil }), + job.NewDefinition("tessellate", func(_ context.Context, _ meshPayload) error { return nil }), + } + + r := job.NewRegistry() + for _, d := range defs { + d.Register(r) + } + + for _, want := range []string{"send-email", "tessellate"} { + if _, ok := r.Get(want); !ok { + t.Errorf("handler %q not registered", want) + } + } +} + +func TestDefinition_JobName(t *testing.T) { + d := job.NewDefinition("tessellate", func(_ context.Context, _ meshPayload) error { return nil }) + + if got := d.JobName(); got != "tessellate" { + t.Errorf("JobName() = %q, want %q", got, "tessellate") + } +} + +func TestWithExecution(t *testing.T) { + d := job.NewDefinition("tessellate", + func(_ context.Context, _ meshPayload) error { return nil }, + job.WithExecution( + exec.Isolate(exec.LevelSandboxed), + exec.GracePeriod(90*time.Second), + ), + ) + + if d.Opts.Execution.Level != exec.LevelSandboxed { + t.Errorf("Level = %v, want %v", d.Opts.Execution.Level, exec.LevelSandboxed) + } + if d.Opts.Execution.GracePeriod != 90*time.Second { + t.Errorf("GracePeriod = %v, want %v", d.Opts.Execution.GracePeriod, 90*time.Second) + } +} + +func TestDefaultOptions_HasUsableExecutionPolicy(t *testing.T) { + // A definition that says nothing about execution must still carry a + // usable grace period, or later rungs would kill instantly. + d := job.NewDefinition("plain", func(_ context.Context, _ meshPayload) error { return nil }) + + if d.Opts.Execution.Level != exec.LevelNone { + t.Errorf("Level = %v, want %v", d.Opts.Execution.Level, exec.LevelNone) + } + if d.Opts.Execution.GracePeriod != exec.DefaultGracePeriod { + t.Errorf("GracePeriod = %v, want %v", d.Opts.Execution.GracePeriod, exec.DefaultGracePeriod) + } +} + +func TestRegistry_Policy(t *testing.T) { + r := job.NewRegistry() + d := job.NewDefinition("tessellate", + func(_ context.Context, _ meshPayload) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelVM)), + ) + d.Register(r) + + if got := r.Policy("tessellate").Level; got != exec.LevelVM { + t.Errorf("Policy(tessellate).Level = %v, want %v", got, exec.LevelVM) + } + // An unregistered name yields the zero policy with usable defaults. + if got := r.Policy("absent").Level; got != exec.LevelNone { + t.Errorf("Policy(absent).Level = %v, want %v", got, exec.LevelNone) + } + if got := r.Policy("absent").GracePeriod; got != exec.DefaultGracePeriod { + t.Errorf("Policy(absent).GracePeriod = %v, want %v", got, exec.DefaultGracePeriod) + } +} diff --git a/job/registry.go b/job/registry.go index 728cf8f..2052a06 100644 --- a/job/registry.go +++ b/job/registry.go @@ -4,9 +4,39 @@ import ( "context" "encoding/json" "fmt" + "slices" "sync" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/resource" ) +// ResourceDecl is what a definition declares about its resource needs, +// captured at registration. +// +// The declaration has to be reachable by job name because the typed +// definition is gone by the time a job is enqueued through EnqueueRaw, +// exactly as with Inputs. +type ResourceDecl struct { + // Requests is the declared requirement. It is a floor: the engine may + // raise it, and a per-enqueue override replaces it per key. + Requests resource.Set + // Limits is the declared enforcement ceiling, if any. + Limits resource.Set + // Func computes the requirement from the enqueue-time request. + Func resource.ResourceFunc + // Class is the opaque scheduling class for the isolation backend. + Class string +} + +// IsZero reports whether the definition declares nothing about +// resources, which is how every job behaves before this feature is used. +func (d ResourceDecl) IsZero() bool { + return d.Requests.IsZero() && d.Limits.IsZero() && d.Func == nil && d.Class == "" +} + // HandlerFunc is a type-erased job handler that accepts raw JSON payload. // The typed Definition[T] is converted to a HandlerFunc at registration // time by closing over JSON unmarshal + the typed handler. @@ -17,12 +47,36 @@ type HandlerFunc func(ctx context.Context, payload []byte) error type Registry struct { mu sync.RWMutex handlers map[string]HandlerFunc + + // inputs holds each job's artifact declarations. The staging + // middleware needs them keyed by job name, because by the time a job + // is executing the typed definition is long gone. + inputs map[string][]artifact.InputSpec + + // resources holds each job's resource declaration, for the same + // reason: enqueue works from a job name and a payload. + resources map[string]ResourceDecl + + // leaseTTLs holds each job's declared lease TTL, for the same reason + // again: EnqueueRaw has a name and a payload, not the typed + // definition, so a per-definition TTL is unreachable unless it is + // keyed by name here. Absent means the definition declared none. + leaseTTLs map[string]time.Duration + + // policies holds each job's execution declaration. The worker needs + // it keyed by name for the same reason inputs are: at execution time + // the typed definition is long gone. + policies map[string]exec.Policy } // NewRegistry creates an empty job registry. func NewRegistry() *Registry { return &Registry{ - handlers: make(map[string]HandlerFunc), + handlers: make(map[string]HandlerFunc), + inputs: make(map[string][]artifact.InputSpec), + resources: make(map[string]ResourceDecl), + leaseTTLs: make(map[string]time.Duration), + policies: make(map[string]exec.Policy), } } @@ -46,6 +100,101 @@ func RegisterDefinition[T any](r *Registry, def *Definition[T]) { r.mu.Lock() defer r.mu.Unlock() r.handlers[def.Name] = handler + + if len(def.Opts.Inputs) > 0 { + specs := make([]artifact.InputSpec, len(def.Opts.Inputs)) + copy(specs, def.Opts.Inputs) + r.inputs[def.Name] = specs + } + + // The sets are cloned, not aliased: a definition's Options stay + // reachable by the caller, and a resolved requirement must not change + // under a job that was already enqueued. + decl := ResourceDecl{ + Requests: def.Opts.Resources.Clone(), + Limits: def.Opts.ResourceLimits.Clone(), + Func: def.Opts.ResourceFunc, + Class: def.Opts.ResourceClass, + } + + if !decl.IsZero() { + r.resources[def.Name] = decl + } + + // Stored under the same non-zero guard as the resource declaration, + // and for the same reason: zero already means "the pool's default" + // everywhere downstream, so recording it would only make an absent + // declaration indistinguishable from a deliberate one without + // changing what any caller does with it. + if def.Opts.LeaseTTL > 0 { + r.leaseTTLs[def.Name] = def.Opts.LeaseTTL + } + + // Unlike inputs and resources, the policy is stored unconditionally: + // DefaultOptions gives every definition a non-zero grace period, so a + // zero-guard here would never skip anything and would only obscure + // intent. + r.policies[def.Name] = def.Opts.Execution +} + +// Resources returns the resource declaration for a job, or the zero +// ResourceDecl when it declares none. +// +// The sets are cloned on the way out as well as in: returning the stored +// maps by reference would let one caller's mutation rewrite the +// requirement every future job of that name resolves from. +func (r *Registry) Resources(name string) ResourceDecl { + r.mu.RLock() + defer r.mu.RUnlock() + + decl := r.resources[name] + decl.Requests = decl.Requests.Clone() + decl.Limits = decl.Limits.Clone() + + return decl +} + +// LeaseTTL returns the lease TTL a definition declared, or zero when it +// declared none. +// +// Zero is not a sentinel this has to distinguish: it is what Job.LeaseTTL +// already means everywhere downstream — Pool.leaseTTLFor falls through to +// the pool default, then the stale-job threshold, then +// job.DefaultLeaseTTL — so an unregistered name and a definition that +// declared nothing correctly resolve the same way. +func (r *Registry) LeaseTTL(name string) time.Duration { + r.mu.RLock() + defer r.mu.RUnlock() + + return r.leaseTTLs[name] +} + +// Policy returns the execution declaration for a job. An unregistered name +// yields a default policy rather than a zero one, so callers always get a +// usable grace period. +func (r *Registry) Policy(name string) exec.Policy { + r.mu.RLock() + defer r.mu.RUnlock() + + if p, ok := r.policies[name]; ok { + return p + } + + return exec.NewPolicy() +} + +// Inputs returns the artifact declarations for a job, or nil when it +// declares none. +// +// The slice is copied, matching the copy RegisterDefinition makes on the +// way in: a caller that mutated the stored declaration would change what +// every future job of that name validates against. InputSpec is all +// value fields, so a shallow copy is a complete one. +func (r *Registry) Inputs(name string) []artifact.InputSpec { + r.mu.RLock() + defer r.mu.RUnlock() + + return slices.Clone(r.inputs[name]) } // Get returns the handler for the given job name. diff --git a/job/resource_options_test.go b/job/resource_options_test.go new file mode 100644 index 0000000..564d2ef --- /dev/null +++ b/job/resource_options_test.go @@ -0,0 +1,78 @@ +package job_test + +import ( + "context" + "testing" + + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +func TestWithResourcesMergesVariadicSets(t *testing.T) { + opts := job.DefaultOptions() + + for _, o := range []job.Option{ + job.WithResources(resource.CPUs(4), resource.MemoryGB(16)), + } { + o(&opts) + } + + if opts.Resources[resource.CPU] != 4000 { + t.Errorf("cpu = %d, want 4000", opts.Resources[resource.CPU]) + } + if opts.Resources[resource.Memory] != 16<<30 { + t.Errorf("memory = %d, want 16 GiB", opts.Resources[resource.Memory]) + } +} + +func TestWithResourcesAccumulatesAcrossCalls(t *testing.T) { + opts := job.DefaultOptions() + + job.WithResources(resource.CPUs(2))(&opts) + job.WithResources(resource.MemoryGB(8))(&opts) + + if opts.Resources[resource.CPU] != 2000 || opts.Resources[resource.Memory] != 8<<30 { + t.Errorf("got %v, want both keys retained", opts.Resources) + } +} + +func TestWithResourcesLaterCallWinsPerKey(t *testing.T) { + opts := job.DefaultOptions() + + job.WithResources(resource.MemoryGB(8))(&opts) + job.WithResources(resource.MemoryGB(32))(&opts) + + if opts.Resources[resource.Memory] != 32<<30 { + t.Errorf("memory = %d, want the later 32 GiB", + opts.Resources[resource.Memory]) + } +} + +func TestWithResourceFuncIsStored(t *testing.T) { + opts := job.DefaultOptions() + + job.WithResourceFunc(func(_ context.Context, r resource.Request) (resource.Set, error) { + return resource.MemoryBytes(r.InputBytes * 3), nil + })(&opts) + + if opts.ResourceFunc == nil { + t.Fatal("ResourceFunc was not stored") + } + + got, err := opts.ResourceFunc(context.Background(), resource.Request{InputBytes: 100}) + if err != nil { + t.Fatalf("ResourceFunc() error = %v", err) + } + if got[resource.Memory] != 300 { + t.Errorf("got %v, want 300 bytes", got) + } +} + +func TestDefaultOptionsDeclareNoResources(t *testing.T) { + opts := job.DefaultOptions() + + if !opts.Resources.IsZero() { + t.Errorf("Resources = %v; the default must declare nothing so "+ + "existing jobs are unaffected", opts.Resources) + } +} diff --git a/job/store.go b/job/store.go index e4eba55..c9b595f 100644 --- a/job/store.go +++ b/job/store.go @@ -2,9 +2,11 @@ package job import ( "context" + "sort" "time" "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" ) // ListOpts controls pagination and filtering for job list queries. @@ -25,15 +27,404 @@ type CountOpts struct { State State } +// budgetedKeys are the canonical dimensions a store compares numerically +// at dequeue. They are exactly the keys every backend persists as its own +// indexed scalar column, which is what lets the fit test be an indexable +// range predicate rather than a document comparison. +// +// Custom keys are deliberately absent: they are matched by containment +// (see DequeueOpts.CustomKeys), not by quantity. +var budgetedKeys = [...]string{resource.CPU, resource.Memory, resource.Disk, resource.GPU} + +// DequeueOpts narrows a dequeue to the jobs the caller can actually run. +// +// The fit predicate lives in the query rather than in the worker because +// DequeueJobs claims a job and marks it running atomically: by the time a +// worker can read a job's requirements it already owns it. Filtering after +// the claim would mean requeueing, and a 32 GB job would then bounce +// between small workers, burning a write on every bounce and delaying +// exactly the job that is hardest to place. Every constraint here must +// therefore be evaluated as part of the claim, never applied to the rows +// the claim returned. +type DequeueOpts struct { + // Queues restricts the claim to these queue names. + // + // An EMPTY list is not a portable request and callers should not + // send one. The backends genuinely disagree, and the disagreement + // predates this option: + // + // memory claims from every queue + // postgres claims nothing (queue = ANY(NULL) matches no row) + // sqlite claims nothing, by an explicit early return before the + // query is built + // redis claims nothing — the index is one sorted set per queue + // name and there has never been a cross-queue index + // mongo claims nothing, by an explicit early return + // + // Mongo's guard is the one that had to be added: {queue: {$in: nil}} + // marshals to BSON null and the server rejects the whole query with + // "$in needs an array", so an empty list was a hard error rather + // than any behaviour at all. It returns empty to match the three + // backends a real deployment would be running, rather than inventing + // an all-queues scan that no other persistent backend performs. + // + // The conformance suite does not exercise an empty list and no + // caller in this repository sends one. Unifying the five would be a + // behaviour change to store/memory and is deliberately not made + // here; documenting the split is what stops a caller assuming it. + Queues []string + + // Limit is the maximum number of jobs to claim. It counts eligible + // jobs only: a job excluded by Budget or CustomKeys must not consume + // a slot, or one oversized job at the head of the queue would starve + // a worker that had capacity for everything behind it. + // + // Limit <= 0 claims nothing and returns (nil, nil); this is unified + // across all five backends, checked before Validate. + Limit int + + // Budget is the free capacity the caller is offering, in canonical + // units (cpu millicores, memory and disk bytes, gpu milli-devices). + // A job is eligible on a dimension when its requirement is <= the + // budget for that dimension. + // + // An ABSENT key is unconstrained, not zero. This inverts + // resource.Set.Fits, which treats absent capacity as zero, and the + // inversion is deliberate: a worker that declares only memory must + // still claim GPU-requiring jobs, because otherwise adding a + // dimension to one worker's config would silently strand work on + // every worker that had not been updated yet. + // + // A key present with the value zero is a real constraint — a worker + // with no free memory — and excludes any job requiring more than + // zero of it. That is why IsUnbounded tests key presence rather than + // resource.Set.IsZero. + // + // Custom keys in Budget are ignored by the predicate. Quantity + // matching on a custom dimension would need a document comparison or + // a join table in five backends to serve a rare case; offer custom + // keys through CustomKeys instead. + Budget resource.Set + + // CustomKeys are the custom resource keys the caller offers, + // typically free.CustomKeys(). A job is eligible only if every custom + // key it requires appears here. + // + // An empty list is read against IsUnbounded rather than on its own, + // because "no custom keys" means two different things to two + // different callers. If the whole of o is unbounded, the caller does + // not use the resource model at all and claims everything, custom + // keys included — that is the backward-compatibility guarantee. If o + // is bounded in any other way, an empty list means this worker + // genuinely has no custom resources, and a job requiring an fpga must + // not be handed to it. + // + // Reading an empty list as unconstrained in the bounded case would + // let a resource-aware worker claim work it cannot possibly run; + // reading it as "offers nothing" in the unbounded case would strand + // every custom-key job in the fleet the day this option shipped. + // + // Only key containment is tested at dequeue; the quantity is enforced + // locally after the claim, by the admission path that already owns + // the accounting. Matching a quantity here would need a document + // comparison or a join table in five backends to serve a rare case. + // + // Backends match against the delimited string + // resource.EncodeCustomKeys produced at enqueue, whose leading and + // trailing separators are what stop ",fpga," matching a worker that + // only offers ",fpga-large,". The test is subset, not substring: a + // job needing {fpga, tpu} is eligible for a caller offering + // {fpga, nvme, tpu}, even though the offered list interleaves a key + // the job does not want. + CustomKeys []string + + // PreferHashes are PrimaryInputHash values the caller already has + // staged locally. A job whose PrimaryInputHash appears here sorts + // ahead of jobs at the same priority, saving a re-download. + // + // Backends must bind PreferredHashes rather than this field: an + // empty string here is not a locality signal and must not become one. + // + // This is advisory and must NEVER filter, and must never outrank + // priority: locality that could reorder across priority bands would + // let a steady stream of locally cached work starve the high-priority + // job the pool exists to run first. The full ordering is priority + // descending, then preferred before unpreferred, then RunAt + // ascending. + // + // It is deliberately NOT a term of IsUnbounded. If it were, opts + // carrying only PreferHashes would count as bounded, and the empty + // CustomKeys rule would then reject every custom-resource job — this + // field would filter transitively, contradicting the paragraph above. + PreferHashes []string + + // ReservedFor restricts the claim to a single job. When set, no other + // job may be returned, and that job is still subject to every other + // constraint here — a targeted claim that could bypass Budget would + // reintroduce exactly the overcommit this predicate prevents. + ReservedFor *id.JobID + + // WorkerID is the worker taking the lease. It is required when + // LeaseUntil is set and ignored otherwise. + WorkerID id.WorkerID + + // LeaseUntil, when non-zero, makes the claim grant a lease: the + // claimed rows get WorkerID, this expiry, and an incremented + // lease_epoch, in the same statement that claims them. + // + // A zero value grants no lease and leaves every lease column + // untouched, which is exactly how DequeueJobs behaved before leases + // existed. That is the backward-compatibility guarantee: a caller + // that does not opt in cannot be affected by this. + // + // The grant must be part of the claim, not a second write, and the + // reason is the opposite of the obvious one. Reclamation cannot + // rescue a half-granted job on the timing that leases exist to + // provide: Lease.IsExpired reports false for a zero ExpiresAt + // precisely so the reclaim loop never steals a job that was never + // leased. A crash between a claim and a separate grant would leave a + // row running with no expiry, and the fine-grained lease machinery + // cannot see it at all. + // + // Such a row is recovered eventually, because reclamation adopts an + // unleased running job once it has been silent for + // UnleasedReclaimGrace, but that clause is a coarse compatibility + // fallback measured in minutes, not the per-job TTL a lease buys. A + // half-granted job would wait it out, which is precisely the coarse + // stale-job behaviour leases exist to replace. + // + // One write means a claimed job always carries a lease something can + // act on promptly. + LeaseUntil time.Time +} + +// Grants reports whether o asks the claim to grant a lease. Backends test +// intent through this rather than open-coding the zero check five times. +func (o DequeueOpts) Grants() bool { + return !o.LeaseUntil.IsZero() +} + +// Validate reports whether o is a coherent request, and is the one place +// every backend checks before it writes anything. +// +// The only incoherent combination is a grant with no holder. A lease +// granted to the zero worker can never be renewed, because RenewLease +// matches on worker ID — so the job would be claimed, left to expire, and +// reclaimed on every cycle forever, which presents as a queue that never +// drains rather than as an error. Refusing the claim turns a silent +// livelock into a caller bug reported at the call that caused it. +func (o DequeueOpts) Validate() error { + if o.Grants() && o.WorkerID.IsNil() { + return ErrLeaseWithoutWorker + } + + return nil +} + +// IsUnbounded reports whether o restricts WHICH jobs may be claimed. +// +// A dequeue asks two independent questions, and this answers only the +// first: +// +// should I filter? — Budget, CustomKeys, ReservedFor. IsUnbounded. +// should I order? — PreferHashes. len(o.PreferHashes) > 0. +// +// Do not reuse this for the second. A backend skips the fit predicate +// when IsUnbounded is true, and separately adds the locality term to its +// ORDER BY whenever PreferHashes is non-empty. A caller that sets only +// PreferHashes therefore gets locality ordering over an unfiltered +// candidate set, which is precisely what "advisory, never a filter" +// means. Folding PreferHashes in here would make it bounded, and the +// empty-CustomKeys rule would then reject every custom-resource job on +// its behalf. +// +// It tests Budget for key presence rather than calling +// resource.Set.IsZero: a Budget of {"memory": 0} is an exhausted worker, +// which must claim nothing that needs memory. Treating it as unbounded +// would hand that worker a job it cannot run. +func (o DequeueOpts) IsUnbounded() bool { + return len(o.Budget) == 0 && + len(o.CustomKeys) == 0 && + o.ReservedFor == nil +} + +// Allows reports whether j satisfies every constraint in o except Queues, +// Limit, and ordering. It is the executable definition of the fit +// predicate: backends that select candidates in Go should call it instead +// of reimplementing the rules, and backends that express the predicate in +// their query language must return the same answer for every job. +// +// It must be applied BEFORE the claim. Claiming a job and then rejecting +// it here is not an implementation of this contract — it is the +// claim-then-requeue behaviour the whole option exists to avoid. +func (o DequeueOpts) Allows(j *Job) bool { + if j == nil { + return false + } + + // A caller that constrains nothing claims everything, including jobs + // declaring custom resources. Backends should reach the same result + // by skipping the predicate entirely on IsUnbounded. + if o.IsUnbounded() { + return true + } + + if o.ReservedFor != nil && j.ID != *o.ReservedFor { + return false + } + + for _, k := range budgetedKeys { + budget, declared := o.Budget[k] + if !declared { + continue + } + + if j.Resources[k] > budget { + return false + } + } + + // o is bounded by this point, so an empty offer means the caller has + // no custom resources — not that it declined to say. See CustomKeys. + required := j.Resources.CustomKeys() + if len(required) == 0 { + return true + } + + offered := make(map[string]struct{}, len(o.CustomKeys)) + for _, k := range o.CustomKeys { + offered[k] = struct{}{} + } + + for _, k := range required { + if _, ok := offered[k]; !ok { + return false + } + } + + return true +} + +// Prefers reports whether j's PrimaryInputHash is one the caller already +// has staged. It is the sort key backends apply after priority. +func (o DequeueOpts) Prefers(j *Job) bool { + if j == nil || j.PrimaryInputHash == "" { + return false + } + + for _, h := range o.PreferHashes { + if h == j.PrimaryInputHash { + return true + } + } + + return false +} + +// Less orders two eligible jobs the way every backend must return them: +// priority descending, then preferred-by-locality before not, then RunAt +// ascending. Ties beyond that are unspecified. +func (o DequeueOpts) Less(a, b *Job) bool { + if a.Priority != b.Priority { + return a.Priority > b.Priority + } + + if pa, pb := o.Prefers(a), o.Prefers(b); pa != pb { + return pa + } + + return a.RunAt.Before(b.RunAt) +} + +// PreferredHashes returns the locality hashes worth matching on: +// deduplicated, in caller order, with the empty string dropped. It is +// what every backend must bind, never PreferHashes itself. +// +// The empty string is the case that matters. Prefers reports false for a +// job with no PrimaryInputHash, so an empty entry offers no information +// — but primary_input_hash is a plain string column, and a job that was +// never hashed stores ” rather than NULL on the SQL backends. Bound +// verbatim, ” = ANY('{""}') is TRUE, so a single empty entry would make +// EVERY unhashed job "locally staged" and hand it the head of its +// priority band. Under a tight Limit that is not a reordering, it is a +// filter: the jobs the caller actually has staged stop being claimed at +// all. Mongo already stripped empties; postgres and sqlite bound them. +func (o DequeueOpts) PreferredHashes() []string { + if len(o.PreferHashes) == 0 { + return nil + } + + seen := make(map[string]struct{}, len(o.PreferHashes)) + out := make([]string, 0, len(o.PreferHashes)) + + for _, h := range o.PreferHashes { + if _, dup := seen[h]; dup || h == "" { + continue + } + + seen[h] = struct{}{} + + out = append(out, h) + } + + return out +} + +// OfferedCustomKeys returns CustomKeys sorted, for backends that build a +// delimited parameter and need a stable, deduplicated order. +func (o DequeueOpts) OfferedCustomKeys() []string { + if len(o.CustomKeys) == 0 { + return nil + } + + seen := make(map[string]struct{}, len(o.CustomKeys)) + out := make([]string, 0, len(o.CustomKeys)) + + for _, k := range o.CustomKeys { + if _, dup := seen[k]; dup || k == "" { + continue + } + + seen[k] = struct{}{} + + out = append(out, k) + } + + sort.Strings(out) + + return out +} + // Store defines the persistence contract for jobs. type Store interface { // EnqueueJob persists a new job in pending state. EnqueueJob(ctx context.Context, j *Job) error - // DequeueJobs atomically claims up to limit pending jobs from the given - // queues, sets them to running, and returns them. Jobs are ordered by - // priority (descending) then RunAt (ascending). - DequeueJobs(ctx context.Context, queues []string, limit int) ([]*Job, error) + // DequeueJobs atomically claims up to opts.Limit ready jobs from + // opts.Queues that fit opts, sets them to running, and returns them + // ordered by priority descending, then locality-preferred first, then + // RunAt ascending. + // + // The fit test is part of the claim, not a filter over claimed rows. + // A job that does not fit stays pending and untouched, available to + // the next worker that does have room for it. + // + // When opts.Grants() the same statement also grants a lease: the + // claimed rows get opts.WorkerID, opts.LeaseUntil, and an incremented + // lease_epoch, and the returned jobs carry the epoch they were + // granted. The grant travels in the claiming write itself, never as a + // follow-up: a row left running with no expiry is outside the lease + // machinery entirely, so a crash between two writes would leave the + // job waiting out the far coarser UnleasedReclaimGrace instead of the + // TTL it was given. See DequeueOpts.LeaseUntil. + // + // Opts that do not grant leave every lease column untouched. A grant + // with no WorkerID is refused with ErrLeaseWithoutWorker and claims + // nothing. + // + // Every backend must pass storetest.RunDequeueSuite, which is the + // contract this signature only sketches. + DequeueJobs(ctx context.Context, opts DequeueOpts) ([]*Job, error) // GetJob retrieves a job by ID. GetJob(ctx context.Context, jobID id.JobID) (*Job, error) @@ -58,3 +449,115 @@ type Store interface { // CountJobs returns the number of jobs matching the given options. CountJobs(ctx context.Context, opts CountOpts) (int64, error) } + +// LeaseStore is the opt-in lease capability. +// +// It is deliberately not part of Store. A backend that implements Store +// alone keeps compiling and keeps behaving exactly as it does today, +// reaped on the pool's single global threshold. A backend that also +// implements LeaseStore gets per-definition lease TTLs, epoch fencing, +// and atomic reclamation. This mirrors the capability idiom the artifact +// backend already uses for RangeReader and Presigner. +// +// Where an expiry is passed at all — RenewLease here, and +// DequeueOpts.LeaseUntil for the grant — it is an absolute timestamp +// rather than a TTL. If the store computed now+ttl it would need +// per-dialect interval arithmetic over a nanosecond integer, and SQLite, +// Mongo, and Redis have no interval type at all. Passing a timestamp +// means every backend only writes a value, and lease policy lives in one +// place. ReclaimExpiredLeases takes no expiry: it acts on the ones +// already written. +// +// The GRANT is deliberately not here. It travels on DequeueOpts instead +// (WorkerID and LeaseUntil), so a leased claim is an ordinary claim that +// also writes the lease columns and therefore carries the budget, the +// custom-key containment, the locality preference and the reservation +// like any other. A second dequeue entry point taking (queues, limit) +// existed here until the two paths had to be kept in sync by hand; it +// carried none of those, so turning leases and resources on together — +// the natural upgrade — claimed jobs that could not fit and requeued +// them on every poll. This interface is what a backend adds ON TOP of +// Store to make that grant renewable and reclaimable. +type LeaseStore interface { + // RenewLease extends the lease to leaseUntil, but only if the job is + // still running, still assigned to workerID, and still at epoch. + // + // It returns ErrLeaseLost when that condition does not hold. That + // return is the entire fencing mechanism: a worker that was reclaimed + // while paused learns it no longer owns the job within one heartbeat + // interval, instead of continuing to write for hours. + RenewLease( + ctx context.Context, + jobID id.JobID, + workerID id.WorkerID, + epoch int, + leaseUntil time.Time, + ) error + + // ReclaimExpiredLeases returns to pending every running job whose + // lease has expired, clearing the worker assignment, incrementing + // lease_epoch to fence the previous holder, and incrementing + // evict_count. RetryCount is never touched — a lost lease is + // infrastructure, not a handler failure. + // + // It also adopts a running job carrying no lease at all, but only + // once that job has been silent for UnleasedReclaimGrace. See that + // constant for why the exception exists and why it is gated on + // silence rather than on the missing expiry alone. + // + // The claim and the read are one atomic statement, so two pools + // reclaiming concurrently cannot both take the same job. + // + // A non-positive limit claims nothing and returns (nil, nil), checked + // before any query runs. This matches DequeueOpts.Limit, so the two + // methods on this interface agree. + // + // ORDER IS NOT PART OF THE CONTRACT, and the backends genuinely + // differ. Postgres, SQLite and Mongo take the longest-expired jobs + // first, because each already reads through an index on the expiry + // and ordering it costs nothing. Memory iterates a map and Redis + // walks an unordered set member list, so both return an arbitrary + // subset. + // + // This only becomes visible when limit is smaller than the number of + // jobs eligible at that instant, where the ordered backends drain + // oldest-first and the other two drain arbitrarily. Nothing starves + // either way, since each pass moves what it takes out of the eligible + // set and the reaper runs on a timer. + // + // Redis is the reason this is documented rather than unified. It can + // stop scanning as soon as it has claimed `limit` jobs; ordering would + // force every call to read the entire job-id set first, turning a + // bounded scan into a full one on the largest deployments, which is a + // steep price for an ordering no caller has asked for. + ReclaimExpiredLeases(ctx context.Context, limit int) ([]*Job, error) + + // UpdateLeasedJob persists j only while the caller still holds the + // lease. + // + // The write applies iff the row is still running, still assigned to + // workerID, and still at epoch. Otherwise it returns ErrLeaseLost and + // leaves the row untouched — the lease has moved on, and a worker + // that no longer owns a job has no coherent claim to make about it. + // + // It writes the same columns as UpdateJob EXCEPT the lease-owned + // ones: lease_epoch, lease_expires_at, worker_id, and heartbeat_at + // are never written here. Those have exactly three writers — the + // grant in DequeueJobs, RenewLease, and ReclaimExpiredLeases — and + // this method deliberately is not a fourth. + // + // j is the caller's claim-time (or otherwise stale) snapshot, so + // j.LeaseExpiresAt is whatever the expiry was when this worker last + // read the row — every renewal since has pushed the real value + // forward. A whole-row write that passed the epoch predicate would + // still roll the expiry backwards, shortening the current holder's + // lease by however long this caller ran. Excluding the lease-owned + // columns is what makes the epoch predicate a genuine fence rather + // than a check that only looks like one. + // + // A missing row returns dispatch.ErrJobNotFound rather than + // ErrLeaseLost: zero rows affected by the fence predicate means "the + // lease moved on", but a row that no longer exists was never a + // question of who holds it. + UpdateLeasedJob(ctx context.Context, j *Job, workerID id.WorkerID, epoch int) error +} diff --git a/options.go b/options.go index 2eb445d..5b9ae8e 100644 --- a/options.go +++ b/options.go @@ -2,6 +2,7 @@ package dispatch import ( "context" + "sync" "time" log "github.com/xraph/go-utils/log" @@ -48,6 +49,15 @@ type Dispatcher struct { // started tracks whether Start has been called. started bool + + // stopOnce makes Stop idempotent. Engine.Stop calls it, and a service + // shutting down from both a signal handler and a deferred cleanup + // reaches it twice; without this, extensions saw two shutdown events + // and the store was closed twice. A sync.Once rather than a flag + // because those two callers are usually different goroutines, which + // an unsynchronised bool would not separate. This mirrors + // cron.Scheduler, which guards its own Stop the same way. + stopOnce sync.Once } // New creates a new Dispatcher with the given options. @@ -93,18 +103,26 @@ func (d *Dispatcher) Start(ctx context.Context) error { // Stop gracefully shuts down the dispatcher. func (d *Dispatcher) Stop(ctx context.Context) error { - if d.pool != nil && d.started { - if err := d.pool.Stop(ctx); err != nil { - d.logger.Error("pool stop error", log.String("error", err.Error())) + // A second call returns nil rather than repeating the first call's + // error, matching Pool.Stop, which also reports nothing once it has + // already stopped. + var err error + + d.stopOnce.Do(func() { + if d.pool != nil && d.started { + if poolErr := d.pool.Stop(ctx); poolErr != nil { + d.logger.Error("pool stop error", log.String("error", poolErr.Error())) + } } - } - if d.extensions != nil { - d.extensions.EmitShutdown(ctx) - } - if d.store != nil { - return d.store.Close() - } - return nil + if d.extensions != nil { + d.extensions.EmitShutdown(ctx) + } + if d.store != nil { + err = d.store.Close() + } + }) + + return err } // WithConcurrency sets the maximum number of concurrent job processors. @@ -171,10 +189,16 @@ func WithHeartbeatInterval(d time.Duration) Option { } } -// WithStaleJobThreshold sets how long without a heartbeat before a -// job is considered stale and reaped. Default 30s. The reaper runs -// at this interval too, so larger values reduce the reap query rate. -// Set to 0 to disable stale-job reaping entirely. +// WithStaleJobThreshold sets how often the reaper runs, and how long +// without a heartbeat a job may go before it is reaped. Default 30s. Set +// to 0 to disable reaping entirely. +// +// The second half of that no longer applies to a store implementing +// job.LeaseStore, which is all five built-in backends. There the window +// before a job is taken back is the lease TTL, per job, from +// WithLeaseTTL or the pool default, and this value only decides how +// frequently the reaper looks and whether it runs at all. It still +// governs the reap window outright for any other store. func WithStaleJobThreshold(d time.Duration) Option { return func(disp *Dispatcher) error { disp.config.StaleJobThreshold = d diff --git a/resource/capacity.go b/resource/capacity.go new file mode 100644 index 0000000..b976160 --- /dev/null +++ b/resource/capacity.go @@ -0,0 +1,101 @@ +package resource + +import ( + "math" + "runtime" +) + +// Default capacity tuning. +const ( + // DefaultCPUOvercommit is 1.0: no overcommit unless asked for. + DefaultCPUOvercommit = 1.0 + // DefaultMemoryFraction leaves 20% for the Go runtime, the OS page + // cache, and everything else sharing the box. + DefaultMemoryFraction = 0.8 + // fallbackMemoryBytes is used when the host total cannot be read. + // Deliberately small: under-advertising costs throughput, while + // over-advertising costs the OOM cascade this package prevents. + fallbackMemoryBytes = 2 << 30 +) + +// CapacityConfig controls how a worker's capacity is derived. +// +// There is no MemoryOvercommit. Overcommitting memory is how a box gets +// into the OOM cascade this package exists to prevent, and a knob whose +// only outcome is an incident should not exist. +type CapacityConfig struct { + // CPUOvercommit multiplies detected cores. CPU is compressible: + // exceeding it makes jobs slow, not dead, so overcommit is safe + // where memory overcommit is not. + CPUOvercommit float64 + // MemoryFraction is the share of the detected limit to advertise. + MemoryFraction float64 + // DiskBytes is the staging cache budget. Zero omits the key. + DiskBytes int64 + // Explicit overrides detection per key, and is the only way to + // declare a custom resource. + Explicit Set +} + +// DefaultCapacityConfig returns the conservative defaults. +func DefaultCapacityConfig() CapacityConfig { + return CapacityConfig{ + CPUOvercommit: DefaultCPUOvercommit, + MemoryFraction: DefaultMemoryFraction, + } +} + +// Detect derives a worker's capacity, with Explicit overriding any +// autodetected key. +func Detect(cfg CapacityConfig) Set { + if cfg.CPUOvercommit <= 0 { + cfg.CPUOvercommit = DefaultCPUOvercommit + } + + if cfg.MemoryFraction <= 0 { + cfg.MemoryFraction = DefaultMemoryFraction + } + + out := Set{ + CPU: int64(math.Floor(float64(detectCPUMillis()) * cfg.CPUOvercommit)), + Memory: int64(math.Floor(float64(detectMemoryBytes()) * cfg.MemoryFraction)), + } + + if cfg.DiskBytes > 0 { + out[Disk] = cfg.DiskBytes + } + + for k, v := range cfg.Explicit { + out[k] = v + } + + return out +} + +// detectCPUMillis prefers the cgroup quota over the host core count. +// +// In a container with a two-core quota, runtime.NumCPU reports the +// host's 64, and every capacity derived from it is wrong by a factor of +// 32 — which on a resource-aware scheduler means admitting 32× the work +// the box can actually run. +func detectCPUMillis() int64 { + if quota, ok := cgroupCPUMillis(); ok && quota > 0 { + return quota + } + + return int64(runtime.NumCPU()) * MilliScale +} + +// detectMemoryBytes prefers the cgroup limit over host total, for the +// same reason. +func detectMemoryBytes() int64 { + if limit, ok := cgroupMemoryBytes(); ok && limit > 0 { + return limit + } + + if total, ok := hostMemoryBytes(); ok && total > 0 { + return total + } + + return fallbackMemoryBytes +} diff --git a/resource/capacity_linux.go b/resource/capacity_linux.go new file mode 100644 index 0000000..ef9ce94 --- /dev/null +++ b/resource/capacity_linux.go @@ -0,0 +1,89 @@ +//go:build linux + +package resource + +import ( + "bufio" + "math" + "os" + "strconv" + "strings" +) + +const ( + cgroupCPUMaxPath = "/sys/fs/cgroup/cpu.max" + cgroupMemoryMaxPath = "/sys/fs/cgroup/memory.max" + procMemInfoPath = "/proc/meminfo" +) + +// cgroupCPUMillis reads cgroup v2 cpu.max, formatted " " +// where quota may be the literal "max" for unlimited. +func cgroupCPUMillis() (int64, bool) { + data, err := os.ReadFile(cgroupCPUMaxPath) + if err != nil { + return 0, false + } + + fields := strings.Fields(string(data)) + if len(fields) != 2 || fields[0] == "max" { + return 0, false + } + + quota, err := strconv.ParseInt(fields[0], 10, 64) + if err != nil { + return 0, false + } + + period, err := strconv.ParseInt(fields[1], 10, 64) + if err != nil || period <= 0 { + return 0, false + } + + return int64(math.Floor(float64(quota) / float64(period) * MilliScale)), true +} + +// cgroupMemoryBytes reads cgroup v2 memory.max, "max" when unlimited. +func cgroupMemoryBytes() (int64, bool) { + data, err := os.ReadFile(cgroupMemoryMaxPath) + if err != nil { + return 0, false + } + + text := strings.TrimSpace(string(data)) + if text == "max" { + return 0, false + } + + limit, err := strconv.ParseInt(text, 10, 64) + if err != nil { + return 0, false + } + + return limit, true +} + +// hostMemoryBytes reads MemTotal from /proc/meminfo, reported in kB. +func hostMemoryBytes() (int64, bool) { + f, err := os.Open(procMemInfoPath) + if err != nil { + return 0, false + } + defer f.Close() + + scanner := bufio.NewScanner(f) + for scanner.Scan() { + fields := strings.Fields(scanner.Text()) + if len(fields) < 2 || fields[0] != "MemTotal:" { + continue + } + + kb, err := strconv.ParseInt(fields[1], 10, 64) + if err != nil { + return 0, false + } + + return kb * 1024, true + } + + return 0, false +} diff --git a/resource/capacity_other.go b/resource/capacity_other.go new file mode 100644 index 0000000..abc7086 --- /dev/null +++ b/resource/capacity_other.go @@ -0,0 +1,14 @@ +//go:build !linux + +package resource + +// cgroupCPUMillis has no cgroup to read outside Linux. +func cgroupCPUMillis() (int64, bool) { return 0, false } + +// cgroupMemoryBytes has no cgroup to read outside Linux. +func cgroupMemoryBytes() (int64, bool) { return 0, false } + +// hostMemoryBytes is not implemented outside Linux, so Detect falls +// back to fallbackMemoryBytes. A non-Linux worker running heavy jobs +// should configure Explicit memory rather than rely on detection. +func hostMemoryBytes() (int64, bool) { return 0, false } diff --git a/resource/capacity_test.go b/resource/capacity_test.go new file mode 100644 index 0000000..f63e76d --- /dev/null +++ b/resource/capacity_test.go @@ -0,0 +1,75 @@ +package resource_test + +import ( + "testing" + + "github.com/xraph/dispatch/resource" +) + +func TestDetectExplicitWins(t *testing.T) { + got := resource.Detect(resource.CapacityConfig{ + CPUOvercommit: 1.0, + MemoryFraction: 0.8, + Explicit: resource.Set{resource.Memory: 42, "fpga": 2}, + }) + + if got[resource.Memory] != 42 { + t.Errorf("explicit memory = %d, want 42", got[resource.Memory]) + } + if got["fpga"] != 2 { + t.Errorf("custom key must be carried through, got %v", got) + } + if got[resource.CPU] == 0 { + t.Error("cpu should still be autodetected when not explicit") + } +} + +func TestDetectAppliesOvercommitAndFraction(t *testing.T) { + base := resource.Detect(resource.CapacityConfig{ + CPUOvercommit: 1.0, MemoryFraction: 1.0, + }) + doubled := resource.Detect(resource.CapacityConfig{ + CPUOvercommit: 2.0, MemoryFraction: 0.5, + }) + + if doubled[resource.CPU] != base[resource.CPU]*2 { + t.Errorf("cpu overcommit not applied: %d vs %d", + doubled[resource.CPU], base[resource.CPU]) + } + if doubled[resource.Memory] > base[resource.Memory]/2+1 { + t.Errorf("memory fraction not applied: %d vs %d", + doubled[resource.Memory], base[resource.Memory]) + } +} + +func TestDefaultCapacityConfigIsConservative(t *testing.T) { + cfg := resource.DefaultCapacityConfig() + + if cfg.CPUOvercommit != 1.0 { + t.Errorf("CPUOvercommit = %v, want 1.0", cfg.CPUOvercommit) + } + if cfg.MemoryFraction >= 1.0 { + t.Errorf("MemoryFraction = %v; must leave runtime and OS headroom", + cfg.MemoryFraction) + } +} + +func TestDetectDiskFromConfig(t *testing.T) { + got := resource.Detect(resource.CapacityConfig{ + CPUOvercommit: 1.0, MemoryFraction: 0.8, DiskBytes: 200 << 30, + }) + if got[resource.Disk] != 200<<30 { + t.Errorf("disk = %d, want %d", got[resource.Disk], int64(200)<<30) + } +} + +func TestDetectNeverReturnsZeroCPUOrMemory(t *testing.T) { + got := resource.Detect(resource.CapacityConfig{}) + + if got[resource.CPU] <= 0 { + t.Errorf("cpu = %d; a zero-value config must still autodetect", got[resource.CPU]) + } + if got[resource.Memory] <= 0 { + t.Errorf("memory = %d; a zero-value config must still autodetect", got[resource.Memory]) + } +} diff --git a/resource/codec.go b/resource/codec.go new file mode 100644 index 0000000..6918af9 --- /dev/null +++ b/resource/codec.go @@ -0,0 +1,87 @@ +package resource + +import ( + "encoding/json" + "strings" +) + +// CustomKeySep delimits the custom-resource key list produced by +// EncodeCustomKeys. The list is stored as a delimited string rather than +// an array so every store backend can express the containment test in +// its own idiom (a LIKE/GLOB pattern, a substring match, ...) without a +// schema translation. +const CustomKeySep = "," + +// EncodeSet marshals a Set for a byte-oriented JSON column (e.g. +// Postgres JSONB). A zero Set — nil or every quantity zero — encodes to +// nil rather than "{}", so an undeclared job's row stays indistinguishable +// from one written before these columns existed: a genuine SQL NULL, not +// an empty JSON object. +func EncodeSet(s Set) ([]byte, error) { + if s.IsZero() { + return nil, nil + } + + return json.Marshal(s) +} + +// DecodeSet unmarshals a column produced by EncodeSet, treating both a +// NULL column (nil slice) and an empty one (zero-length slice) as +// "unset" rather than an error. +func DecodeSet(b []byte) (Set, error) { + if len(b) == 0 { + return nil, nil + } + + var s Set + if err := json.Unmarshal(b, &s); err != nil { + return nil, err + } + + return s, nil +} + +// EncodeSetString is EncodeSet for backends whose JSON column is a +// string-typed TEXT rather than a byte column — SQLite has no JSONB +// type, so its column is *string, where nil is what stores as SQL NULL. +// The encoding rules are identical to EncodeSet: a zero Set produces a +// nil pointer, never a pointer to "" or "{}". +func EncodeSetString(s Set) (*string, error) { + b, err := EncodeSet(s) + if err != nil { + return nil, err + } + + if b == nil { + return nil, nil + } + + js := string(b) + + return &js, nil +} + +// DecodeSetString is DecodeSet for the *string TEXT representation, +// treating both a NULL column (nil pointer) and an empty string as +// "unset". +func DecodeSetString(s *string) (Set, error) { + if s == nil { + return DecodeSet(nil) + } + + return DecodeSet([]byte(*s)) +} + +// EncodeCustomKeys renders the non-canonical keys of s carrying a +// nonzero quantity as a delimited string with a leading and trailing +// separator, so a containment test can match on ",fpga," and never +// partially match ",fpga-large,". A Set with no custom keys — including +// a nil or zero Set — encodes to "". +func EncodeCustomKeys(s Set) string { + keys := s.CustomKeys() + if len(keys) == 0 { + return "" + } + + return CustomKeySep + strings.Join(keys, CustomKeySep) + CustomKeySep +} diff --git a/resource/codec_test.go b/resource/codec_test.go new file mode 100644 index 0000000..650f8c1 --- /dev/null +++ b/resource/codec_test.go @@ -0,0 +1,245 @@ +package resource + +import "testing" + +// TestEncodeCustomKeys pins the leading/trailing separator encoding that +// the dequeue predicate depends on. A bare join ("fpga" instead of +// ",fpga,") would let a containment match on "fpga" partially match +// "fpga-large" — that's the case the first subtest exists to catch. This +// is the single shared copy: store/postgres and store/sqlite used to +// each pin their own duplicate of this behavior; now both call +// EncodeCustomKeys directly, so one suite covers every backend. +func TestEncodeCustomKeys(t *testing.T) { + tests := []struct { + name string + set Set + want string + }{ + { + name: "prefix collision case: fpga must not partially match fpga-large", + set: Set{"fpga": 2, "fpga-large": 1}, + want: ",fpga,fpga-large,", + }, + { + name: "canonical keys only encodes to empty string", + set: Set{ + CPU: 4000, + Memory: 16 << 30, + Disk: 1 << 30, + GPU: 1000, + }, + want: "", + }, + { + name: "zero-quantity custom key is excluded", + set: Set{"fpga": 0, "gpu-slot": 3}, + want: ",gpu-slot,", + }, + { + name: "nil set encodes to empty string", + set: nil, + want: "", + }, + { + name: "empty (zero) set encodes to empty string", + set: Set{}, + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := EncodeCustomKeys(tt.set) + if got != tt.want { + t.Errorf("EncodeCustomKeys(%v) = %q, want %q", tt.set, got, tt.want) + } + }) + } +} + +// TestEncodeDecodeSetRoundTrip exercises EncodeSet/DecodeSet directly, +// without going through any store. The zero-Set-encodes-to-nil behavior +// is load-bearing: it is what makes an undeclared job indistinguishable +// from a row written before this migration (NULL, not "{}"). +func TestEncodeDecodeSetRoundTrip(t *testing.T) { + t.Run("zero set encodes to nil, not {}", func(t *testing.T) { + b, err := EncodeSet(Set{}) + if err != nil { + t.Fatalf("EncodeSet() error = %v", err) + } + if b != nil { + t.Errorf("EncodeSet(zero Set) = %q, want nil", b) + } + }) + + t.Run("nil set encodes to nil", func(t *testing.T) { + b, err := EncodeSet(nil) + if err != nil { + t.Fatalf("EncodeSet() error = %v", err) + } + if b != nil { + t.Errorf("EncodeSet(nil) = %q, want nil", b) + } + }) + + t.Run("nonzero set survives encode then decode", func(t *testing.T) { + want := Set{ + CPU: 4000, + Memory: 16 << 30, + "fpga": 2, + } + + b, err := EncodeSet(want) + if err != nil { + t.Fatalf("EncodeSet() error = %v", err) + } + if b == nil { + t.Fatal("EncodeSet(nonzero Set) = nil, want encoded bytes") + } + + got, err := DecodeSet(b) + if err != nil { + t.Fatalf("DecodeSet() error = %v", err) + } + + if len(got) != len(want) { + t.Fatalf("DecodeSet() = %v, want %v", got, want) + } + for k, v := range want { + if got[k] != v { + t.Errorf("DecodeSet()[%q] = %d, want %d", k, got[k], v) + } + } + }) + + t.Run("DecodeSet(nil) returns nil, not an error", func(t *testing.T) { + got, err := DecodeSet(nil) + if err != nil { + t.Fatalf("DecodeSet(nil) error = %v", err) + } + if got != nil { + t.Errorf("DecodeSet(nil) = %v, want nil", got) + } + }) + + t.Run("DecodeSet(empty slice) returns nil, not an error", func(t *testing.T) { + got, err := DecodeSet([]byte{}) + if err != nil { + t.Fatalf("DecodeSet(empty) error = %v", err) + } + if got != nil { + t.Errorf("DecodeSet(empty) = %v, want nil", got) + } + }) + + t.Run("DecodeSet rejects malformed JSON", func(t *testing.T) { + if _, err := DecodeSet([]byte("not json")); err == nil { + t.Error("DecodeSet(malformed) error = nil, want non-nil") + } + }) +} + +// TestEncodeDecodeSetStringRoundTrip covers the *string TEXT-column +// representation (SQLite has no JSONB type) with the same rules as +// EncodeSet/DecodeSet, including that a zero Set still produces a nil +// pointer (SQL NULL), not a pointer to "" or "{}". +func TestEncodeDecodeSetStringRoundTrip(t *testing.T) { + t.Run("zero set encodes to nil, not a pointer to {}", func(t *testing.T) { + s, err := EncodeSetString(Set{}) + if err != nil { + t.Fatalf("EncodeSetString() error = %v", err) + } + if s != nil { + t.Errorf("EncodeSetString(zero Set) = %q, want nil", *s) + } + }) + + t.Run("nil set encodes to nil", func(t *testing.T) { + s, err := EncodeSetString(nil) + if err != nil { + t.Fatalf("EncodeSetString() error = %v", err) + } + if s != nil { + t.Errorf("EncodeSetString(nil) = %q, want nil", *s) + } + }) + + t.Run("nonzero set survives encode then decode", func(t *testing.T) { + want := Set{ + CPU: 4000, + Memory: 16 << 30, + "fpga": 2, + } + + s, err := EncodeSetString(want) + if err != nil { + t.Fatalf("EncodeSetString() error = %v", err) + } + if s == nil { + t.Fatal("EncodeSetString(nonzero Set) = nil, want an encoded string") + } + + got, err := DecodeSetString(s) + if err != nil { + t.Fatalf("DecodeSetString() error = %v", err) + } + + if len(got) != len(want) { + t.Fatalf("DecodeSetString() = %v, want %v", got, want) + } + for k, v := range want { + if got[k] != v { + t.Errorf("DecodeSetString()[%q] = %d, want %d", k, got[k], v) + } + } + }) + + t.Run("DecodeSetString(nil) returns nil, not an error", func(t *testing.T) { + got, err := DecodeSetString(nil) + if err != nil { + t.Fatalf("DecodeSetString(nil) error = %v", err) + } + if got != nil { + t.Errorf("DecodeSetString(nil) = %v, want nil", got) + } + }) + + t.Run("DecodeSetString(empty string) returns nil, not an error", func(t *testing.T) { + empty := "" + got, err := DecodeSetString(&empty) + if err != nil { + t.Fatalf("DecodeSetString(empty) error = %v", err) + } + if got != nil { + t.Errorf("DecodeSetString(empty) = %v, want nil", got) + } + }) + + t.Run("DecodeSetString rejects malformed JSON", func(t *testing.T) { + bad := "not json" + if _, err := DecodeSetString(&bad); err == nil { + t.Error("DecodeSetString(malformed) error = nil, want non-nil") + } + }) + + t.Run("EncodeSet and EncodeSetString agree byte-for-byte", func(t *testing.T) { + set := Set{CPU: 2000, "fpga": 1} + + b, err := EncodeSet(set) + if err != nil { + t.Fatalf("EncodeSet() error = %v", err) + } + + s, err := EncodeSetString(set) + if err != nil { + t.Fatalf("EncodeSetString() error = %v", err) + } + if s == nil { + t.Fatal("EncodeSetString() = nil, want an encoded string") + } + + if string(b) != *s { + t.Errorf("EncodeSet() = %q, EncodeSetString() = %q, want identical JSON", b, *s) + } + }) +} diff --git a/resource/doc.go b/resource/doc.go new file mode 100644 index 0000000..dc6d122 --- /dev/null +++ b/resource/doc.go @@ -0,0 +1,24 @@ +// Package resource defines Dispatch's resource model: what a job needs, +// what a worker has, and the admission control between them. +// +// # Leaf package +// +// resource may import only the root dispatch package, dispatch/id, and +// the standard library. job imports resource for Options.Resources, so +// any edge back to job — or to artifact, worker, engine, or store — is +// an import cycle. This is why the estimator's input is plain data +// ([]InputSize) rather than []artifact.Ref: the caller translates. +// +// # Units +// +// Every quantity is an int64 in a canonical unit, so accounting that +// adds and subtracts the same values thousands of times cannot drift: +// +// cpu millicores 1 core = 1000 +// memory bytes +// disk bytes +// gpu milli-devices 1 device = 1000 +// +// Any other key is a custom resource with user-defined semantics, in +// the style of Ray's resource dict. Custom quantities are integers. +package resource diff --git a/resource/errors.go b/resource/errors.go new file mode 100644 index 0000000..681ba13 --- /dev/null +++ b/resource/errors.go @@ -0,0 +1,13 @@ +package resource + +import "errors" + +// ErrCapacityExceeded means an acquisition could not be satisfied: either +// the request is larger than total capacity and never can be, or every +// holder still held its share when the caller's context ended. +var ErrCapacityExceeded = errors.New("dispatch/resource: capacity exceeded") + +// ErrUnschedulable means a requirement exceeds the largest known worker +// capacity, so no worker could ever run it. Returned at enqueue so the +// job fails on a developer's machine rather than pending forever. +var ErrUnschedulable = errors.New("dispatch/resource: no worker can satisfy the requirement") diff --git a/resource/manager.go b/resource/manager.go new file mode 100644 index 0000000..81b73ae --- /dev/null +++ b/resource/manager.go @@ -0,0 +1,435 @@ +package resource + +import ( + "context" + "fmt" + "sort" + "sync" + "time" +) + +// Reclaimer frees capacity for one key on the manager's behalf. +// +// It exists because the resource dimensions are not symmetric: cached +// bytes on disk can be evicted to make room, while memory held by a +// running job cannot. Registering a reclaimer for a key turns blocking +// into "reclaim, then block only if that was not enough". A key with no +// reclaimer — memory, CPU — can only wait for a release. +// +// A reclaimer holds a Manager lease for what it caches, and returns +// units by releasing that lease — it never credits the ledger directly. +// The artifact cache holds one lease per cached entry, so evicting an +// entry releases that entry's lease and the bytes come back through the +// same path a finished job's would. +type Reclaimer interface { + // Reclaim frees up to need units of key, returning how many it + // actually freed. Returning zero means nothing more is reclaimable. + // + // The units are returned to the manager by releasing the Lease that + // held them. The return value is only a "something changed, re-check" + // signal — the manager does not add it to the ledger, because those + // units are still recorded against a live lease until that lease is + // released, and counting them in both places would invent capacity. + // + // Reclaim is called with the manager's lock dropped, so releasing a + // lease from inside it is safe and is the expected implementation. + Reclaim(ctx context.Context, key string, need int64) (int64, error) + // Available reports how much could be freed without blocking. + Available(key string) int64 +} + +// Lease is a held allocation. Release is idempotent. +type Lease interface { + Held() Set + Owner() string + Release() +} + +// LeaseInfo is a point-in-time view of one lease, for the capacity API. +type LeaseInfo struct { + Owner string `json:"owner"` + Held Set `json:"held"` + AcquiredAt time.Time `json:"acquired_at"` +} + +// Manager admits work against a fixed capacity. +// +// It generalizes the artifact cache's single-key disk budget to N keys: +// same mutex and condition variable, same context-bounded wait, same +// reclaim-then-wait loop — but keyed, so memory and CPU are accounted +// by the mechanism that already worked for disk rather than by a +// second one built alongside it. +type Manager interface { + // Acquire blocks until want fits, reclaiming where a Reclaimer is + // registered. It is bounded by ctx, so a blocked job can never + // outlive its deadline. A want larger than total capacity fails + // immediately: no release or reclamation could ever satisfy it, so + // waiting would only defer a certain error. + Acquire(ctx context.Context, owner string, want Set) (Lease, error) + // TryAcquire is the non-blocking form. It never reclaims, because a + // caller that cannot wait also cannot afford eviction I/O. + TryAcquire(owner string, want Set) (Lease, bool) + + // Free is what is immediately available. + Free() Set + // Reclaimable is what registered reclaimers could free on top of Free. + Reclaimable() Set + // Capacity is the configured total. + Capacity() Set + // Leases is a snapshot of what is currently held. + Leases() []LeaseInfo + + // RegisterReclaimer installs the reclaim policy for one key. + RegisterReclaimer(key string, r Reclaimer) + // Wake tells blocked acquirers that a reclaimer's holdings changed + // — that something now reclaimable was not a moment ago. + // + // The manager broadcasts when a lease is released and on nothing + // else, and it cannot see inside a reclaimer. An artifact cache + // whose last stager lets go of an entry has just made those bytes + // evictable without releasing anything, so the acquirer that asked + // for them, was told "everything is pinned", and went to sleep will + // sleep to its deadline unless it is told to look again. + Wake() +} + +// ManagerOption configures a manager. +type ManagerOption func(*manager) + +// WithClock replaces the time source, for deterministic tests. +func WithClock(now func() time.Time) ManagerOption { + return func(m *manager) { + if now != nil { + m.now = now + } + } +} + +type manager struct { + mu sync.Mutex + cond *sync.Cond + now func() time.Time + + capacity Set + used Set + reclaimers map[string]Reclaimer + // reclaimGen counts Wake calls. A broadcast only reaches a waiter + // already inside cond.Wait, and the window this has to cross — + // reclaimLocked running with the lock dropped — is precisely when + // the acquirer is not yet waiting. So the signal is left as state + // the acquirer re-checks under the lock rather than as an event it + // has to be present for. + reclaimGen uint64 + + nextID int64 + leases map[int64]*lease +} + +// NewManager builds a manager over a fixed capacity. +func NewManager(capacity Set, opts ...ManagerOption) Manager { + m := &manager{ + now: time.Now, + capacity: capacity.Clone(), + used: make(Set, len(capacity)), + reclaimers: make(map[string]Reclaimer), + leases: make(map[int64]*lease), + } + m.cond = sync.NewCond(&m.mu) + + for _, opt := range opts { + opt(m) + } + + if m.capacity == nil { + m.capacity = make(Set) + } + + return m +} + +func (m *manager) RegisterReclaimer(key string, r Reclaimer) { + m.mu.Lock() + defer m.mu.Unlock() + + if r == nil { + delete(m.reclaimers, key) + + return + } + + m.reclaimers[key] = r +} + +// Wake records that a reclaimer's holdings changed and wakes the +// waiters. +// +// The counter is the part that matters. reclaimLocked drops this lock +// across Reclaim, because eviction does I/O and calls back in to +// release a lease, and a caller that completes a whole wake cycle +// inside that window would otherwise broadcast to nobody: the acquirer +// is between "you have nothing to give me" and cond.Wait, and it goes +// to sleep having missed the one thing that would have helped it. +// Bumping a generation the acquirer re-reads under the lock turns that +// lost edge into state it cannot miss. +func (m *manager) Wake() { + m.mu.Lock() + defer m.mu.Unlock() + + m.reclaimGen++ + m.cond.Broadcast() +} + +func (m *manager) Capacity() Set { + m.mu.Lock() + defer m.mu.Unlock() + + return m.capacity.Clone() +} + +func (m *manager) Free() Set { + m.mu.Lock() + defer m.mu.Unlock() + + return m.freeLocked() +} + +func (m *manager) freeLocked() Set { + return m.capacity.Sub(m.used) +} + +func (m *manager) Reclaimable() Set { + m.mu.Lock() + reclaimers := make(map[string]Reclaimer, len(m.reclaimers)) + for k, r := range m.reclaimers { + reclaimers[k] = r + } + m.mu.Unlock() + + // The reclaimer map is snapshotted under our lock and then queried + // with it dropped. Available takes the reclaimer's own lock, and a + // reclaimer takes ours whenever it releases a lease to return units — + // taking both locks in both orders is a deadlock. + out := make(Set, len(reclaimers)) + for k, r := range reclaimers { + out[k] = r.Available(k) + } + + return out +} + +func (m *manager) Leases() []LeaseInfo { + m.mu.Lock() + defer m.mu.Unlock() + + out := make([]LeaseInfo, 0, len(m.leases)) + for _, l := range m.leases { + out = append(out, LeaseInfo{ + Owner: l.owner, + Held: l.held.Clone(), + AcquiredAt: l.acquiredAt, + }) + } + + sort.Slice(out, func(i, j int) bool { + return out[i].AcquiredAt.Before(out[j].AcquiredAt) + }) + + return out +} + +func (m *manager) TryAcquire(owner string, want Set) (Lease, bool) { + if want.IsZero() { + return m.grant(owner, want), true + } + + m.mu.Lock() + defer m.mu.Unlock() + + if !want.Fits(m.freeLocked()) { + return nil, false + } + + return m.grantLocked(owner, want), true +} + +func (m *manager) Acquire(ctx context.Context, owner string, want Set) (Lease, error) { + if want.IsZero() { + return m.grant(owner, want), nil + } + + m.mu.Lock() + defer m.mu.Unlock() + + // A request larger than total capacity can never be satisfied. + if over := want.Exceeds(m.capacity); len(over) > 0 { + return nil, fmt.Errorf("%w: %v exceeds worker capacity", ErrCapacityExceeded, over) + } + + stop := m.watchContext(ctx) + defer stop() + + for !want.Fits(m.freeLocked()) { + if err := ctx.Err(); err != nil { + return nil, fmt.Errorf("%w: waiting for %v: %w", + ErrCapacityExceeded, want.Exceeds(m.freeLocked()), err) + } + + // Read before reclaimLocked, compared after: that call is the + // only place this loop lets go of the lock, so it is the only + // window a Wake can land in unheard. + gen := m.reclaimGen + + if m.reclaimLocked(ctx, want) { + continue + } + + if m.reclaimGen != gen { + // A reclaimer gained something while we were asking a + // different question. Ask again rather than sleep on an + // answer that is already stale. + continue + } + + if ctx.Err() != nil { + // The deadline may have passed while reclaimLocked had this + // lock dropped, and watchContext broadcasts exactly once — + // to nobody, since we were not waiting yet. Going into Wait + // now would be waiting forever for a wake-up that has + // already happened. Round the loop instead and let the + // check at the top return the error. + continue + } + + // Nothing reclaimable on any short key. Only a release can help. + // Nothing above releases the lock, so the ctx watcher's + // broadcast cannot slip past between that check and this wait. + m.cond.Wait() + } + + return m.grantLocked(owner, want), nil +} + +// reclaimLocked asks the registered reclaimer for each short key to free +// the shortfall. It reports whether anything was freed, which is only a +// "something changed, re-check the ledger" signal for the caller's loop. +// +// The freed amount is deliberately not added back here. The manager's +// invariant is used == Σ live lease.held, and release is its only +// mutator: a reclaimer returns units by releasing the lease that held +// them, so by the time Reclaim returns the ledger already reflects the +// change. Crediting the return value on top would count those units +// twice — the lease still records them until it is released, and +// release subtracts held in full — which is how a worker would slowly +// invent capacity it does not have. +// +// The manager's lock is released across the reclaimer call. That is +// what makes the design work rather than an optimization: Reclaim calls +// back into the manager to return the bytes, so holding the lock +// through it would deadlock on the first eviction. It also keeps +// eviction I/O off the admission path for every other caller. +func (m *manager) reclaimLocked(ctx context.Context, want Set) bool { + free := m.freeLocked() + + var freedAny bool + + for _, key := range want.Exceeds(free) { + r, ok := m.reclaimers[key] + if !ok { + continue + } + + need := want[key] - free[key] + + m.mu.Unlock() + freed, err := r.Reclaim(ctx, key, need) + m.mu.Lock() + + if err == nil && freed > 0 { + freedAny = true + } + } + + return freedAny +} + +// watchContext broadcasts when ctx ends so a waiter is interruptible. +func (m *manager) watchContext(ctx context.Context) func() { + if ctx.Done() == nil { + return func() {} + } + + done := make(chan struct{}) + + go func() { + select { + case <-ctx.Done(): + m.mu.Lock() + m.cond.Broadcast() + m.mu.Unlock() + case <-done: + } + }() + + return func() { close(done) } +} + +func (m *manager) grant(owner string, want Set) Lease { + m.mu.Lock() + defer m.mu.Unlock() + + return m.grantLocked(owner, want) +} + +func (m *manager) grantLocked(owner string, want Set) *lease { + m.nextID++ + + l := &lease{ + id: m.nextID, + mgr: m, + owner: owner, + held: want.Clone(), + acquiredAt: m.now(), + } + + m.used = m.used.Add(want) + m.leases[l.id] = l + + return l +} + +// release returns a lease's resources. It is the only path that reduces +// used, reclamation included, which is what keeps the invariant +// used == Σ live lease.held true by construction. +// +// Idempotent: releasing twice must not credit the ledger twice, or a +// worker slowly invents capacity. The sync.Once on the lease and the +// liveness check here are belt and braces — either alone would hold for +// the paths that exist today, so both are kept deliberately. +func (m *manager) release(l *lease) { + m.mu.Lock() + defer m.mu.Unlock() + + if _, live := m.leases[l.id]; !live { + return + } + + delete(m.leases, l.id) + m.used = m.used.Sub(l.held) + m.cond.Broadcast() +} + +type lease struct { + id int64 + mgr *manager + owner string + held Set + acquiredAt time.Time + once sync.Once +} + +func (l *lease) Held() Set { return l.held.Clone() } +func (l *lease) Owner() string { return l.owner } + +func (l *lease) Release() { + l.once.Do(func() { l.mgr.release(l) }) +} diff --git a/resource/manager_test.go b/resource/manager_test.go new file mode 100644 index 0000000..f6377c8 --- /dev/null +++ b/resource/manager_test.go @@ -0,0 +1,443 @@ +package resource_test + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/xraph/dispatch/resource" +) + +func TestManagerTryAcquire(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Memory: 100, resource.CPU: 4000}) + + lease, ok := m.TryAcquire("job-1", resource.Set{resource.Memory: 60}) + if !ok { + t.Fatal("first acquire should succeed") + } + if m.Free()[resource.Memory] != 40 { + t.Errorf("free memory = %d, want 40", m.Free()[resource.Memory]) + } + + if _, ok := m.TryAcquire("job-2", resource.Set{resource.Memory: 60}); ok { + t.Error("second acquire should not fit") + } + + lease.Release() + if m.Free()[resource.Memory] != 100 { + t.Errorf("free memory after release = %d, want 100", m.Free()[resource.Memory]) + } +} + +func TestManagerDoubleReleaseIsSafe(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Memory: 100}) + lease, _ := m.TryAcquire("job-1", resource.Set{resource.Memory: 60}) + + lease.Release() + lease.Release() + + if m.Free()[resource.Memory] != 100 { + t.Errorf("double release corrupted the ledger: free = %d, want 100", + m.Free()[resource.Memory]) + } +} + +func TestManagerAcquireBlocksThenSucceeds(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Memory: 100}) + held, _ := m.TryAcquire("job-1", resource.Set{resource.Memory: 80}) + + done := make(chan error, 1) + go func() { + _, err := m.Acquire(context.Background(), "job-2", resource.Set{resource.Memory: 80}) + done <- err + }() + + select { + case err := <-done: + t.Fatalf("acquire returned %v while capacity was held", err) + case <-time.After(50 * time.Millisecond): + } + + held.Release() + + select { + case err := <-done: + if err != nil { + t.Fatalf("acquire after release: %v", err) + } + case <-time.After(time.Second): + t.Fatal("acquire did not wake after release") + } +} + +func TestManagerAcquireBoundedByContext(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Memory: 100}) + defer func() { _, _ = m.TryAcquire("x", nil) }() + + if _, ok := m.TryAcquire("job-1", resource.Set{resource.Memory: 100}); !ok { + t.Fatal("setup acquire failed") + } + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + _, err := m.Acquire(ctx, "job-2", resource.Set{resource.Memory: 50}) + if !errors.Is(err, resource.ErrCapacityExceeded) { + t.Fatalf("got %v, want ErrCapacityExceeded", err) + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("error should wrap the context cause, got %v", err) + } +} + +func TestManagerAcquireOverCapacityFailsImmediately(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Memory: 100}) + + start := time.Now() + _, err := m.Acquire(context.Background(), "job-1", resource.Set{resource.Memory: 200}) + if !errors.Is(err, resource.ErrCapacityExceeded) { + t.Fatalf("got %v, want ErrCapacityExceeded", err) + } + if elapsed := time.Since(start); elapsed > 100*time.Millisecond { + t.Errorf("a request larger than capacity must not block, waited %v", elapsed) + } +} + +// countingReclaimer models what the artifact cache becomes: a component +// that holds a Manager lease per cached entry and returns units by +// releasing those leases, never by crediting the ledger itself. +// +// Reclaim takes the reclaimer's lock and then the manager's, via +// Release. That is only safe because the manager drops its own lock +// across Reclaim and snapshots its reclaimer map before calling +// Available, so the manager never holds its lock while reaching for +// this one. The fake is built this way on purpose: it is the test of +// that lock ordering as much as of the accounting. +type countingReclaimer struct { + mu sync.Mutex + entries []resource.Lease + calls int +} + +// newCountingReclaimer acquires count leases of each units against m, +// standing in for count cached entries. +func newCountingReclaimer(t *testing.T, m resource.Manager, key string, count int, each int64) *countingReclaimer { + t.Helper() + + r := &countingReclaimer{} + + for i := 0; i < count; i++ { + l, ok := m.TryAcquire("cache-entry", resource.Set{key: each}) + if !ok { + t.Fatalf("reclaimer setup: entry %d of %d did not fit", i, count) + } + + r.entries = append(r.entries, l) + } + + return r +} + +func (r *countingReclaimer) Reclaim(_ context.Context, key string, need int64) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + + r.calls++ + + var freed int64 + + // Evict whole entries until the shortfall is covered. Releasing the + // lease is what actually returns the units to the manager; the + // returned count only tells it to re-check. + for freed < need && len(r.entries) > 0 { + last := len(r.entries) - 1 + entry := r.entries[last] + r.entries = r.entries[:last] + + freed += entry.Held()[key] + entry.Release() + } + + return freed, nil +} + +func (r *countingReclaimer) Available(key string) int64 { + r.mu.Lock() + defer r.mu.Unlock() + + var avail int64 + for _, entry := range r.entries { + avail += entry.Held()[key] + } + + return avail +} + +func TestManagerReclaimerFreesDisk(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Disk: 100}) + + // Eight cached entries of 10, then a job holding the rest, so the + // ledger is full and only reclamation can satisfy the next request. + rec := newCountingReclaimer(t, m, resource.Disk, 8, 10) + m.RegisterReclaimer(resource.Disk, rec) + + if _, ok := m.TryAcquire("other", resource.Set{resource.Disk: 20}); !ok { + t.Fatal("setup acquire failed") + } + + if got := m.Reclaimable()[resource.Disk]; got != 80 { + t.Errorf("Reclaimable() disk = %d, want 80", got) + } + if got := m.Free()[resource.Disk]; got != 0 { + t.Fatalf("setup left %d disk free, want 0", got) + } + + lease, err := m.Acquire(context.Background(), "job-1", resource.Set{resource.Disk: 50}) + if err != nil { + t.Fatalf("Acquire() error = %v", err) + } + if rec.calls == 0 { + t.Error("reclaimer was never called") + } + if lease.Held()[resource.Disk] != 50 { + t.Errorf("held = %v, want 50 disk", lease.Held()) + } + + // The invariant reclamation must not break: every unit that is not + // free is recorded against exactly one live lease. It holds because a + // reclaimer returns units by releasing the lease that held them, so + // release stays the ledger's only mutator. + var accounted int64 + for _, info := range m.Leases() { + accounted += info.Held[resource.Disk] + } + + if free := m.Free()[resource.Disk]; accounted+free != 100 { + t.Errorf("ledger invariant broken: %d held + %d free != 100 capacity", accounted, free) + } +} + +func TestManagerMemoryHasNoReclaimer(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Memory: 100, resource.Disk: 100}) + rec := newCountingReclaimer(t, m, resource.Disk, 10, 10) + m.RegisterReclaimer(resource.Disk, rec) + + if _, ok := m.TryAcquire("holder", resource.Set{resource.Memory: 100}); !ok { + t.Fatal("setup acquire failed") + } + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + if _, err := m.Acquire(ctx, "job-1", resource.Set{resource.Memory: 50}); err == nil { + t.Fatal("memory acquisition should not have been satisfied") + } + if rec.calls != 0 { + t.Errorf("the disk reclaimer was called for a memory request (%d times)", rec.calls) + } +} + +func TestManagerNeverExceedsCapacity(t *testing.T) { + const ( + capacity = 1000 + goroutines = 64 + iterations = 50 + ) + + m := resource.NewManager(resource.Set{resource.Memory: capacity}) + + var ( + mu sync.Mutex + held int64 + peak int64 + ) + + var wg sync.WaitGroup + for g := range goroutines { + wg.Add(1) + + go func(g int) { + defer wg.Done() + + want := int64((g%8)+1) * 50 + + for range iterations { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + lease, err := m.Acquire(ctx, "g", resource.Set{resource.Memory: want}) + cancel() + + if err != nil { + continue + } + + mu.Lock() + held += want + if held > peak { + peak = held + } + current := held + mu.Unlock() + + if current > capacity { + t.Errorf("held %d exceeds capacity %d", current, capacity) + } + + mu.Lock() + held -= want + mu.Unlock() + + lease.Release() + } + }(g) + } + + wg.Wait() + + if peak == 0 { + t.Fatal("no acquisition ever succeeded; the test proved nothing") + } + if m.Free()[resource.Memory] != capacity { + t.Errorf("ledger leaked: free = %d, want %d", m.Free()[resource.Memory], capacity) + } +} + +// pinnedReclaimer models the artifact cache at its most awkward: it +// holds one lease it will not give up while something is using the +// entry, and the thing that stops using it does so at the worst +// possible moment — inside Reclaim, while the manager has its lock +// dropped and the acquirer is not yet waiting. +type pinnedReclaimer struct { + mu sync.Mutex + mgr resource.Manager + entry resource.Lease + pinned bool + // unpinDuringReclaim releases the pin from inside the first Reclaim + // call, which is exactly the window a broadcast is lost in. + unpinDuringReclaim bool +} + +func (r *pinnedReclaimer) Reclaim(_ context.Context, key string, _ int64) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + + if r.pinned { + if r.unpinDuringReclaim { + r.unpinDuringReclaim = false + r.pinned = false + // A stager just let go. Nothing was released, so the only + // signal the manager can get is this one. + r.mgr.Wake() + } + + return 0, nil + } + + if r.entry == nil { + return 0, nil + } + + freed := r.entry.Held()[key] + r.entry.Release() + r.entry = nil + + return freed, nil +} + +func (r *pinnedReclaimer) Available(key string) int64 { + r.mu.Lock() + defer r.mu.Unlock() + + if r.pinned || r.entry == nil { + return 0 + } + + return r.entry.Held()[key] +} + +// TestWakeSurvivesTheReclaimWindow is the lost-wakeup guard. +// +// Wake broadcasts, but a broadcast only reaches a waiter that is +// already waiting, and the window it has to cross — reclaimLocked +// running with the lock dropped — is precisely when the acquirer is +// not. The generation counter is what the acquirer re-reads under the +// lock to notice it missed the edge; without it this test hangs to its +// deadline. +func TestWakeSurvivesTheReclaimWindow(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Disk: 10}) + + entry, ok := m.TryAcquire("cache-entry", resource.Set{resource.Disk: 10}) + if !ok { + t.Fatal("setup acquire failed") + } + + rec := &pinnedReclaimer{mgr: m, entry: entry, pinned: true, unpinDuringReclaim: true} + m.RegisterReclaimer(resource.Disk, rec) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + lease, err := m.Acquire(ctx, "job", resource.Set{resource.Disk: 10}) + if err != nil { + t.Fatalf("Acquire: %v — the unpin landed in the reclaim window and was missed", err) + } + + defer lease.Release() + + if free := m.Free()[resource.Disk]; free != 0 { + t.Fatalf("Free()[disk] = %d, want 0", free) + } +} + +// slowReclaimer takes longer to answer than the caller has to live. +type slowReclaimer struct{ delay time.Duration } + +func (r slowReclaimer) Reclaim(_ context.Context, _ string, _ int64) (int64, error) { + time.Sleep(r.delay) + + return 0, nil +} + +func (slowReclaimer) Available(string) int64 { return 0 } + +// TestAcquireDoesNotHangWhenContextEndsDuringReclaim covers the other +// edge the same window creates. watchContext broadcasts once when the +// deadline passes; if that happens while reclaimLocked has the lock +// dropped, the acquirer arrives at Wait after the only wake-up it was +// ever going to get. It has to re-check the context under the lock it +// is about to hand over, or it waits for ever on an expired deadline. +func TestAcquireDoesNotHangWhenContextEndsDuringReclaim(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Disk: 10}) + m.RegisterReclaimer(resource.Disk, slowReclaimer{delay: 300 * time.Millisecond}) + + held, ok := m.TryAcquire("other", resource.Set{resource.Disk: 10}) + if !ok { + t.Fatal("setup acquire failed") + } + + defer held.Release() + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + done := make(chan error, 1) + + go func() { + lease, err := m.Acquire(ctx, "job", resource.Set{resource.Disk: 10}) + if lease != nil { + lease.Release() + } + + done <- err + }() + + select { + case err := <-done: + if !errors.Is(err, resource.ErrCapacityExceeded) { + t.Fatalf("Acquire = %v, want ErrCapacityExceeded", err) + } + case <-time.After(5 * time.Second): + t.Fatal("Acquire hung past its deadline: the context ended inside the reclaim window") + } +} diff --git a/resource/resourcetest/doc.go b/resource/resourcetest/doc.go new file mode 100644 index 0000000..0175c6e --- /dev/null +++ b/resource/resourcetest/doc.go @@ -0,0 +1,7 @@ +// Package resourcetest provides fakes for testing resource-aware code: +// a deterministic clock, a lease-based reclaimer, and a scripted +// estimator. +// +// It mirrors artifact/artifacttest, so a test that needs both staging +// and admission fakes reaches for the same shapes in both packages. +package resourcetest diff --git a/resource/resourcetest/fakes.go b/resource/resourcetest/fakes.go new file mode 100644 index 0000000..26a93df --- /dev/null +++ b/resource/resourcetest/fakes.go @@ -0,0 +1,225 @@ +package resourcetest + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/xraph/dispatch/resource" +) + +// fakeReclaimerOwner is the lease owner name every FakeReclaimer +// acquires under, so a test inspecting resource.Manager.Leases() can +// tell the fake's holdings apart from the job leases it is competing +// against. +const fakeReclaimerOwner = "resourcetest.FakeReclaimer" + +// Clock is a manually advanced time source for resource.WithClock. +type Clock struct { + mu sync.Mutex + now time.Time +} + +// NewClock starts a clock at t. +func NewClock(t time.Time) *Clock { return &Clock{now: t} } + +// Now returns the current fake time. +func (c *Clock) Now() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + + return c.now +} + +// Advance moves the clock forward. +func (c *Clock) Advance(d time.Duration) { + c.mu.Lock() + defer c.mu.Unlock() + + c.now = c.now.Add(d) +} + +// FakeReclaimer stands in for a component like the artifact cache: it +// holds real resource.Manager leases for one key and frees capacity by +// releasing them, never by crediting a private counter. +// +// This matters because resource.Manager's invariant is +// used == Σ (held of every live lease); a Reclaimer that does not hold +// a lease has nothing to release, so under the real contract it can +// never actually return units to the ledger. A counter-based fake would +// make Manager.Acquire spin on a predicate that never becomes true once +// the counter is spent, and a test built on it would hang or pass +// without exercising reclamation at all. FakeReclaimer is built the way +// countingReclaimer in resource/manager_test.go is: it acquires the +// pool it represents as real leases up front and releases whole leases +// to satisfy Reclaim, so a test asserting against it is exercising the +// same admission and release path a live reclaimer would. +type FakeReclaimer struct { + mu sync.Mutex + mgr resource.Manager + key string + entries []resource.Lease + calls int + err error +} + +// NewFakeReclaimer builds a reclaimer for key by acquiring count leases +// of leaseSize units each from m, standing in for count evictable +// entries (an artifact cache would hold one lease per cached object). +// The pool Reclaim can free totals count*leaseSize units. +// +// Splitting the pool into count leases rather than one big one is what +// lets Reclaim satisfy a partial request the way a real evictor does: +// releasing whole entries until the shortfall is covered, rather than +// only being able to free everything or nothing. +// +// It returns an error rather than taking a *testing.T, because a +// construction failure here is a caller setup mistake (asking for more +// than the manager's capacity), not a test assertion — the caller +// decides whether that is fatal. +func NewFakeReclaimer(m resource.Manager, key string, leaseSize int64, count int) (*FakeReclaimer, error) { + r := &FakeReclaimer{mgr: m, key: key} + + for i := range count { + lease, ok := m.TryAcquire(fakeReclaimerOwner, resource.Set{key: leaseSize}) + if !ok { + // Every lease acquired in iterations before this one is real, + // live capacity taken from m — the only handle to it is + // r.entries, which is about to be discarded. Release them all + // before returning, or the caller's manager permanently loses + // that capacity: indistinguishable from a job that leaked. + for _, held := range r.entries { + held.Release() + } + + return nil, fmt.Errorf("resourcetest: acquire lease %d/%d of %d %s: capacity exhausted", + i+1, count, leaseSize, key) + } + + r.entries = append(r.entries, lease) + } + + return r, nil +} + +// SetError makes every subsequent Reclaim fail with err instead of +// releasing leases. +func (r *FakeReclaimer) SetError(err error) { + r.mu.Lock() + defer r.mu.Unlock() + + r.err = err +} + +// Refill acquires one more lease of n units from the manager and adds +// it to the pool Reclaim can release. It fails if the manager has no +// room to grant it: refill goes through the same admission check any +// other acquisition does, rather than pretending capacity that is not +// there. +func (r *FakeReclaimer) Refill(n int64) error { + r.mu.Lock() + key := r.key + mgr := r.mgr + r.mu.Unlock() + + lease, ok := mgr.TryAcquire(fakeReclaimerOwner, resource.Set{key: n}) + if !ok { + return fmt.Errorf("resourcetest: refill %d %s: capacity exhausted", n, key) + } + + r.mu.Lock() + r.entries = append(r.entries, lease) + r.mu.Unlock() + + return nil +} + +// Calls reports how many times Reclaim was called. +func (r *FakeReclaimer) Calls() int { + r.mu.Lock() + defer r.mu.Unlock() + + return r.calls +} + +// Reclaim releases whole held leases until at least need units of key +// have been freed, returning the total actually released. Releasing +// the lease is what returns the units to the manager; the return value +// is only the "something changed, re-check" signal Manager.Acquire +// uses, exactly as resource.Reclaimer documents. +// +// A key this reclaimer was not built for frees nothing: it owns no +// leases in that key, so there is nothing honest to release. +func (r *FakeReclaimer) Reclaim(_ context.Context, key string, need int64) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + + r.calls++ + + if r.err != nil { + return 0, r.err + } + + if key != r.key { + return 0, nil + } + + var freed int64 + + for freed < need && len(r.entries) > 0 { + last := len(r.entries) - 1 + entry := r.entries[last] + r.entries = r.entries[:last] + + freed += entry.Held()[key] + entry.Release() + } + + return freed, nil +} + +// Available reports how much this reclaimer could still free without +// blocking: the sum of what its live leases hold. A key it was not +// built for is never available through it. +func (r *FakeReclaimer) Available(key string) int64 { + r.mu.Lock() + defer r.mu.Unlock() + + if key != r.key { + return 0 + } + + var avail int64 + for _, entry := range r.entries { + avail += entry.Held()[key] + } + + return avail +} + +// FakeEstimator returns a scripted Set and records what it was asked. +type FakeEstimator struct { + Out resource.Set + Err error + Calls int + Last resource.Request +} + +// Estimate returns the scripted result. +func (e *FakeEstimator) Estimate(_ context.Context, r resource.Request) (resource.Set, error) { + e.Calls++ + e.Last = r + + if e.Err != nil { + return nil, e.Err + } + + return e.Out, nil +} + +// Compile-time proof the fakes satisfy the interfaces they stand in for. +var ( + _ resource.Reclaimer = (*FakeReclaimer)(nil) + _ resource.Estimator = (*FakeEstimator)(nil) +) diff --git a/resource/resourcetest/fakes_test.go b/resource/resourcetest/fakes_test.go new file mode 100644 index 0000000..4f76890 --- /dev/null +++ b/resource/resourcetest/fakes_test.go @@ -0,0 +1,184 @@ +package resourcetest_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/resource/resourcetest" +) + +func TestClockAdvances(t *testing.T) { + start := time.Date(2026, 8, 12, 0, 0, 0, 0, time.UTC) + c := resourcetest.NewClock(start) + + if !c.Now().Equal(start) { + t.Fatalf("Now() = %v, want %v", c.Now(), start) + } + + c.Advance(90 * time.Second) + + if want := start.Add(90 * time.Second); !c.Now().Equal(want) { + t.Fatalf("Now() = %v, want %v", c.Now(), want) + } +} + +// TestFakeReclaimerReturnsUnitsToManager is the point of the helper: a +// FakeReclaimer does not free capacity by decrementing a private +// counter, it frees capacity by releasing real leases it holds against +// a real Manager. This proves that end to end — the manager starts +// full, only reclamation can make room, and the ledger balances +// afterward — rather than only asserting on the fake's own bookkeeping. +func TestFakeReclaimerReturnsUnitsToManager(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Disk: 100}) + + // Ten entries of ten units each, standing in for ten cached objects, + // so a partial reclaim can be satisfied by releasing some but not all + // of them. + r, err := resourcetest.NewFakeReclaimer(m, resource.Disk, 10, 10) + if err != nil { + t.Fatalf("NewFakeReclaimer() error = %v", err) + } + m.RegisterReclaimer(resource.Disk, r) + + // The fake's own leases already account for the entire capacity, so + // nothing is free and only reclamation can satisfy the next request. + if got := m.Free()[resource.Disk]; got != 0 { + t.Fatalf("setup: free = %d, want 0", got) + } + if got := m.Reclaimable()[resource.Disk]; got != 100 { + t.Fatalf("Reclaimable() = %d, want 100", got) + } + + lease, err := m.Acquire(context.Background(), "job-1", resource.Set{resource.Disk: 60}) + if err != nil { + t.Fatalf("Acquire() error = %v", err) + } + if got := lease.Held()[resource.Disk]; got != 60 { + t.Errorf("held = %d, want 60", got) + } + + if r.Calls() == 0 { + t.Error("Calls() = 0, reclaimer was never invoked") + } + + // Reclaim released six of the fake's ten-unit leases to cover the + // sixty-unit shortfall, so it should have forty units left to offer. + if got := r.Available(resource.Disk); got != 40 { + t.Errorf("Available() = %d, want 40", got) + } + + // A key the fake was not built for is never available through it and + // never reclaimed from it. + if got := r.Available(resource.Memory); got != 0 { + t.Errorf("Available(memory) = %d, want 0", got) + } + callsBefore := r.Calls() + if freed, err := r.Reclaim(context.Background(), resource.Memory, 10); err != nil || freed != 0 { + t.Errorf("Reclaim(memory) = (%d, %v), want (0, nil)", freed, err) + } + if r.Calls() != callsBefore+1 { + t.Errorf("Calls() = %d, want %d", r.Calls(), callsBefore+1) + } + + // The manager's invariant must hold: every unit not free is recorded + // against exactly one live lease, held by the job or by the fake. + var accounted int64 + for _, info := range m.Leases() { + accounted += info.Held[resource.Disk] + } + + if free := m.Free()[resource.Disk]; accounted+free != 100 { + t.Errorf("ledger invariant broken: %d held + %d free != 100 capacity", accounted, free) + } +} + +// TestFakeReclaimerConstructionFailureLeaksNothing is the regression +// test for the leaked-lease bug: NewFakeReclaimer acquires leases one +// at a time, so a failure partway through must release everything it +// already took, or the manager permanently loses that capacity with no +// live job to account for it — indistinguishable from a leak. +func TestFakeReclaimerConstructionFailureLeaksNothing(t *testing.T) { + // 65 units of capacity is enough for six 10-unit leases (60) but not + // a seventh, so the tenth (of ten requested) never even gets tried: + // the failure happens on lease 7. + m := resource.NewManager(resource.Set{resource.Disk: 65}) + + if _, err := resourcetest.NewFakeReclaimer(m, resource.Disk, 10, 10); err == nil { + t.Fatal("NewFakeReclaimer() error = nil, want an error for capacity exhausted") + } + + if got := m.Free()[resource.Disk]; got != 65 { + t.Errorf("Free() after failed construction = %d, want 65 (nothing should leak)", got) + } + if leases := m.Leases(); len(leases) != 0 { + t.Errorf("Leases() after failed construction = %v, want none held", leases) + } +} + +func TestFakeReclaimerSetError(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Disk: 10}) + + r, err := resourcetest.NewFakeReclaimer(m, resource.Disk, 10, 1) + if err != nil { + t.Fatalf("NewFakeReclaimer() error = %v", err) + } + + boom := errors.New("boom") + r.SetError(boom) + + freed, err := r.Reclaim(context.Background(), resource.Disk, 5) + if !errors.Is(err, boom) { + t.Fatalf("Reclaim() error = %v, want %v", err, boom) + } + if freed != 0 { + t.Errorf("freed = %d, want 0", freed) + } + // The lease is still held: an error must not silently release it. + if got := r.Available(resource.Disk); got != 10 { + t.Errorf("Available() = %d, want 10 (lease must survive an error)", got) + } +} + +func TestFakeReclaimerRefill(t *testing.T) { + m := resource.NewManager(resource.Set{resource.Disk: 20}) + + r, err := resourcetest.NewFakeReclaimer(m, resource.Disk, 10, 1) + if err != nil { + t.Fatalf("NewFakeReclaimer() error = %v", err) + } + + if got := r.Available(resource.Disk); got != 10 { + t.Fatalf("Available() = %d, want 10", got) + } + + if err := r.Refill(10); err != nil { + t.Fatalf("Refill() error = %v", err) + } + if got := r.Available(resource.Disk); got != 20 { + t.Errorf("Available() after Refill = %d, want 20", got) + } + + // Refill goes through the manager's real admission check, so asking + // for more than is left must fail rather than invent capacity. + if err := r.Refill(1); err == nil { + t.Error("Refill() over capacity should have failed") + } +} + +func TestFakeEstimatorRecordsRequest(t *testing.T) { + e := &resourcetest.FakeEstimator{Out: resource.MemoryGB(4)} + + got, err := e.Estimate(context.Background(), resource.Request{JobName: "tessellate"}) + if err != nil { + t.Fatalf("Estimate() error = %v", err) + } + if got[resource.Memory] != 4<<30 { + t.Errorf("got %v, want 4 GiB", got) + } + if e.Calls != 1 || e.Last.JobName != "tessellate" { + t.Errorf("calls = %d, last = %+v", e.Calls, e.Last) + } +} diff --git a/resource/set.go b/resource/set.go new file mode 100644 index 0000000..2a47521 --- /dev/null +++ b/resource/set.go @@ -0,0 +1,214 @@ +package resource + +import ( + "math" + "sort" +) + +// Canonical resource keys. Any key outside this set is a custom +// resource: an integer quantity with user-defined semantics. +const ( + // CPU is measured in millicores. One core is 1000. + CPU = "cpu" + // Memory is measured in bytes. + Memory = "memory" + // Disk is measured in bytes. For a worker this is the staging cache + // budget; for a job, the bytes its inputs and outputs need locally. + Disk = "disk" + // GPU is measured in milli-devices. One device is 1000, so a + // fractional declaration is expressible. Kubernetes accepts only + // whole devices, so track C rounds up at translation. + GPU = "gpu" +) + +// MilliScale is the multiplier for the milli-denominated keys. +const MilliScale = 1000 + +// Set is a resource vector. An absent key is zero, so a Set never needs +// to enumerate keys it does not constrain. +// +// Set is a map, so it is not safe for concurrent mutation. Every method +// here returns a new Set rather than mutating the receiver, which makes +// a Set stored on a job or a lease safe to read from many goroutines. +type Set map[string]int64 + +// CPUs builds a Set from a core count. CPUs(2.5) is 2500 millicores. +func CPUs(n float64) Set { return Set{CPU: milli(n)} } + +// MemoryBytes builds a Set from a byte count. +func MemoryBytes(n int64) Set { return Set{Memory: n} } + +// MemoryGB builds a Set from a gibibyte count. +func MemoryGB(n int64) Set { return Set{Memory: n << 30} } + +// DiskBytes builds a Set from a byte count. +func DiskBytes(n int64) Set { return Set{Disk: n} } + +// DiskGB builds a Set from a gibibyte count. +func DiskGB(n int64) Set { return Set{Disk: n << 30} } + +// GPUs builds a Set from a device count. GPUs(0.5) is half a device. +func GPUs(n float64) Set { return Set{GPU: milli(n)} } + +// Custom builds a Set for a single custom resource key. +func Custom(key string, n int64) Set { return Set{key: n} } + +// milli converts a fractional count to its milli-denominated integer, +// rounding up so a request is never understated. +func milli(n float64) int64 { + return int64(math.Ceil(n * MilliScale)) +} + +// Clone returns an independent copy. +func (s Set) Clone() Set { + if s == nil { + return nil + } + + out := make(Set, len(s)) + for k, v := range s { + out[k] = v + } + + return out +} + +// Add returns the per-key sum. Neither operand is modified. +func (s Set) Add(o Set) Set { + out := s.Clone() + if out == nil { + out = make(Set, len(o)) + } + + for k, v := range o { + out[k] += v + } + + return out +} + +// Sub returns the per-key difference, clamped at zero. A resource +// vector is never negative: "owing" capacity is not a state the +// accounting can represent, and clamping keeps a double release from +// corrupting the ledger. +func (s Set) Sub(o Set) Set { + out := s.Clone() + if out == nil { + out = make(Set, len(o)) + } + + for k, v := range o { + if out[k] -= v; out[k] < 0 { + out[k] = 0 + } + } + + return out +} + +// Max returns the per-key maximum, used to merge a floor into a +// resolved requirement. +func (s Set) Max(o Set) Set { + out := s.Clone() + if out == nil { + out = make(Set, len(o)) + } + + for k, v := range o { + if v > out[k] { + out[k] = v + } + } + + return out +} + +// Scale multiplies every quantity by f, rounding up. Used for safety +// factors and OOM retry escalation, where rounding down would produce +// the same failure again. +func (s Set) Scale(f float64) Set { + out := make(Set, len(s)) + for k, v := range s { + out[k] = int64(math.Ceil(float64(v) * f)) + } + + return out +} + +// Fits reports whether every quantity in s is within capacity. An +// absent capacity key is zero, so demanding a resource the capacity +// does not list never fits. +func (s Set) Fits(capacity Set) bool { + for k, v := range s { + if v > capacity[k] { + return false + } + } + + return true +} + +// Exceeds returns the keys on which s does not fit capacity, sorted. +// It is what turns a failed admission into an error naming the +// dimension that did not fit rather than a bare "does not fit". +func (s Set) Exceeds(capacity Set) []string { + var over []string + + for k, v := range s { + if v > capacity[k] { + over = append(over, k) + } + } + + sort.Strings(over) + + return over +} + +// Keys returns every key, sorted. +func (s Set) Keys() []string { + keys := make([]string, 0, len(s)) + for k := range s { + keys = append(keys, k) + } + + sort.Strings(keys) + + return keys +} + +// CustomKeys returns the non-canonical keys carrying a nonzero +// quantity, sorted. This is the set persisted on the job row and +// matched by containment at dequeue. +func (s Set) CustomKeys() []string { + var keys []string + + for k, v := range s { + if v == 0 { + continue + } + + switch k { + case CPU, Memory, Disk, GPU: + default: + keys = append(keys, k) + } + } + + sort.Strings(keys) + + return keys +} + +// IsZero reports whether every quantity is zero. A zero Set means "no +// declared requirement", which is how every job behaves before this +// feature is configured. +func (s Set) IsZero() bool { + for _, v := range s { + if v != 0 { + return false + } + } + + return true +} diff --git a/resource/set_test.go b/resource/set_test.go new file mode 100644 index 0000000..d18d82c --- /dev/null +++ b/resource/set_test.go @@ -0,0 +1,138 @@ +package resource_test + +import ( + "testing" + + "github.com/xraph/dispatch/resource" +) + +func TestSetArithmetic(t *testing.T) { + tests := []struct { + name string + op func() resource.Set + want resource.Set + }{ + { + name: "add merges disjoint keys", + op: func() resource.Set { return resource.CPUs(2).Add(resource.MemoryGB(4)) }, + want: resource.Set{resource.CPU: 2000, resource.Memory: 4 << 30}, + }, + { + name: "add sums shared keys", + op: func() resource.Set { return resource.CPUs(2).Add(resource.CPUs(1.5)) }, + want: resource.Set{resource.CPU: 3500}, + }, + { + name: "sub clamps at zero", + op: func() resource.Set { return resource.CPUs(1).Sub(resource.CPUs(4)) }, + want: resource.Set{resource.CPU: 0}, + }, + { + name: "max takes the larger per key", + op: func() resource.Set { + return resource.Set{resource.CPU: 1000, resource.Memory: 100}. + Max(resource.Set{resource.CPU: 500, resource.Memory: 900}) + }, + want: resource.Set{resource.CPU: 1000, resource.Memory: 900}, + }, + { + name: "scale rounds up so a request is never understated", + op: func() resource.Set { return resource.Set{resource.Memory: 3}.Scale(1.5) }, + want: resource.Set{resource.Memory: 5}, + }, + { + name: "custom keys participate in arithmetic", + op: func() resource.Set { return resource.Custom("fpga", 1).Add(resource.Custom("fpga", 2)) }, + want: resource.Set{"fpga": 3}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := tt.op() + if len(got) != len(tt.want) { + t.Fatalf("got %v, want %v", got, tt.want) + } + for k, v := range tt.want { + if got[k] != v { + t.Errorf("key %q: got %d, want %d", k, got[k], v) + } + } + }) + } +} + +func TestSetFits(t *testing.T) { + tests := []struct { + name string + want resource.Set + capacity resource.Set + fits bool + }{ + { + name: "empty set fits anything", + want: resource.Set{}, + capacity: resource.Set{resource.CPU: 1}, + fits: true, + }, + { + name: "exact fit", + want: resource.Set{resource.Memory: 100}, + capacity: resource.Set{resource.Memory: 100}, + fits: true, + }, + { + name: "over on one key fails", + want: resource.Set{resource.CPU: 1, resource.Memory: 101}, + capacity: resource.Set{resource.CPU: 8, resource.Memory: 100}, + fits: false, + }, + { + name: "absent capacity key is zero, so any demand fails", + want: resource.Set{"fpga": 1}, + capacity: resource.Set{resource.CPU: 8}, + fits: false, + }, + { + name: "zero demand on an absent key still fits", + want: resource.Set{"fpga": 0}, + capacity: resource.Set{resource.CPU: 8}, + fits: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.want.Fits(tt.capacity); got != tt.fits { + t.Errorf("Fits() = %v, want %v", got, tt.fits) + } + }) + } +} + +func TestSetCustomKeysSorted(t *testing.T) { + s := resource.Set{ + resource.CPU: 1000, "zebra": 1, resource.Memory: 2, "alpha": 3, + } + got := s.CustomKeys() + want := []string{"alpha", "zebra"} + if len(got) != len(want) { + t.Fatalf("got %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("got %v, want %v", got, want) + } + } +} + +func TestSetImmutability(t *testing.T) { + a := resource.CPUs(1) + b := a.Add(resource.CPUs(1)) + if a[resource.CPU] != 1000 { + t.Errorf("Add mutated the receiver: %v", a) + } + if b[resource.CPU] != 2000 { + t.Errorf("Add returned %v", b) + } +} diff --git a/resource/spec.go b/resource/spec.go new file mode 100644 index 0000000..ba23dd7 --- /dev/null +++ b/resource/spec.go @@ -0,0 +1,241 @@ +package resource + +import ( + "context" + "fmt" + "strings" +) + +// Spec is the resolved, immutable resource contract for one job. It is +// produced once at enqueue and written to the job row, so scheduling +// reads columns and never calls user code. +// +// Spec is also the contract track C consumes to build a pod: Requests +// map to resource requests, Limits to limits, and Class to whatever +// scheduling class the isolation backend uses. +type Spec struct { + // Requests is what admission accounts for and what a pod requests. + Requests Set `json:"requests,omitempty"` + // Limits is the enforcement ceiling. Memory defaults to Requests + // (guaranteed); CPU is left unset (burstable), because overrunning + // CPU makes a job slow while overrunning memory makes it dead. + Limits Set `json:"limits,omitempty"` + // Class is an opaque scheduling class for the isolation backend. + Class string `json:"class,omitempty"` +} + +// IsZero reports whether the spec constrains nothing. +func (s Spec) IsZero() bool { + return s.Requests.IsZero() && s.Limits.IsZero() && s.Class == "" +} + +// InputSize describes one declared input at enqueue time. +// +// It is plain data rather than an artifact.Ref because resource is a +// leaf package (see doc.go). The engine translates bindings into these, +// which also makes an Estimator testable with a struct literal. +type InputSize struct { + // Name is the declared input slot name. + Name string `json:"name"` + // Bytes is the input's size. + Bytes int64 `json:"bytes"` + // Hash may be empty: the artifact plane fills content_hash + // opportunistically at first staging, not at registration. + Hash string `json:"hash,omitempty"` +} + +// Request is everything an estimator may consider. +type Request struct { + JobName string + Queue string + Payload []byte + Inputs []InputSize + InputBytes int64 + Declared Set + Attempt int + ScopeOrgID string +} + +// ResourceFunc computes a requirement from the enqueue-time request. It +// runs once, in the enqueuing process, never on the scheduling path. +// +//nolint:revive // ResourceFunc is the name job definitions register under (see Tasks 7-8); Func would collide with ResolveInput.Func. +type ResourceFunc func(ctx context.Context, r Request) (Set, error) + +// Estimator infers a requirement from historical measurement. +// +// There is no built-in implementation. The interface exists so that a +// rollup over observed usage, or a learned predictor, can be slotted in +// with nothing else moving; until one is installed, requirements come +// from the declaration and the enqueue-time override alone. +type Estimator interface { + Estimate(ctx context.Context, r Request) (Set, error) +} + +// The sources ResolveInput.OnError names. +const ( + // SourceFunc is a job definition's or enqueue's ResourceFunc. + SourceFunc = "func" + // SourceEstimator is the installed Estimator. + SourceEstimator = "estimator" +) + +// ResolveInput carries every source a requirement can come from, +// lowest precedence first. +type ResolveInput struct { + GlobalDefault Set + QueueDefault Set + Declared Set + Func ResourceFunc + Estimator Estimator + Override Set + + DeclaredLimits Set + OverrideLimits Set + + Class string + Request Request + + // MaxCapacity is the largest single-worker capacity known to the + // engine. Empty disables the unschedulable check, which is correct + // for a single-process engine with no registered workers. + MaxCapacity Set + + // OnError is called when Func or Estimator returns an error, with + // the source that failed (SourceFunc or SourceEstimator). + // + // Neither may fail an enqueue — that contract is unchanged, and the + // lower-precedence value stands. But a fallback with no signal is + // how a misconfigured estimator under-sizes every job in the fleet + // forever: the jobs enqueue, they run, they OOM, and nothing + // anywhere says the estimator was the reason. Callers should log it. + // + // Optional. Nil discards, which is the right default for a test that + // only cares about precedence. Called synchronously inside Resolve on + // the enqueue path: it must not block or panic, or it fails or stalls + // the enqueue that triggered it. + OnError func(source string, err error) +} + +// report hands err to OnError when one was supplied. +func (in ResolveInput) report(source string, err error) { + if in.OnError != nil { + in.OnError(source, err) + } +} + +// Resolve collapses every source into one Spec. +// +// Precedence is a per-key overlay, lowest first: +// +// global default → queue default → declaration → func → estimator → override +// +// Per-key rather than whole-set replacement, so an estimator that +// predicts only memory leaves a declared CPU value intact. +// +// Neither the func nor the estimator may fail enqueue: a failure there +// is handed to OnError and the lower-precedence value stands. The one +// error Resolve does return is ErrUnschedulable, because a job no worker +// can run must fail loudly and immediately. +func Resolve(ctx context.Context, in ResolveInput) (Spec, error) { + req := make(Set) + + overlay := func(o Set) { + for k, v := range o { + req[k] = v + } + } + + overlay(in.GlobalDefault) + overlay(in.QueueDefault) + overlay(in.Declared) + + // Both dynamic sources see the declaration, so either can choose to + // defer to it by returning it unchanged. + r := in.Request + r.Declared = req.Clone() + + // A failure in either dynamic source falls back to the declaration + // and is REPORTED. Swallowing it was the whole hazard: both run once + // per enqueue in the enqueuing process, so a broken one is broken for + // every job of that name from then on, and the symptom — jobs sized + // from the static declaration and OOMing under real input — points + // nowhere near the estimator. + if in.Func != nil { + out, err := in.Func(ctx, r) + if err != nil { + in.report(SourceFunc, err) + } else { + overlay(out) + } + } + + if in.Estimator != nil { + out, err := in.Estimator.Estimate(ctx, r) + if err != nil { + in.report(SourceEstimator, err) + } else { + overlay(out) + } + } + + overlay(in.Override) + + spec := Spec{ + Requests: req, + Limits: defaultLimits(req, in.DeclaredLimits, in.OverrideLimits), + Class: in.Class, + } + + if err := checkSchedulable(spec.Requests, in.MaxCapacity); err != nil { + return Spec{}, err + } + + return spec, nil +} + +// defaultLimits gives every incompressible key a limit equal to its +// request and leaves CPU unset. Explicit limits override both. +func defaultLimits(requests, declared, override Set) Set { + limits := make(Set, len(requests)) + + for k, v := range requests { + if k == CPU || v == 0 { + continue + } + + limits[k] = v + } + + for k, v := range declared { + limits[k] = v + } + + for k, v := range override { + limits[k] = v + } + + return limits +} + +// checkSchedulable rejects a requirement no worker could ever satisfy. +// An empty maxCapacity means capacity is unknown, which disables the +// check rather than rejecting everything. +func checkSchedulable(requests, maxCapacity Set) error { + if len(maxCapacity) == 0 { + return nil + } + + over := requests.Exceeds(maxCapacity) + if len(over) == 0 { + return nil + } + + parts := make([]string, 0, len(over)) + for _, k := range over { + parts = append(parts, fmt.Sprintf("%s: need %d, largest worker has %d", + k, requests[k], maxCapacity[k])) + } + + return fmt.Errorf("%w (%s)", ErrUnschedulable, strings.Join(parts, "; ")) +} diff --git a/resource/spec_test.go b/resource/spec_test.go new file mode 100644 index 0000000..7590b68 --- /dev/null +++ b/resource/spec_test.go @@ -0,0 +1,204 @@ +package resource_test + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/xraph/dispatch/resource" +) + +// stubEstimator returns a fixed Set, or an error when set. +type stubEstimator struct { + out resource.Set + err error +} + +func (s stubEstimator) Estimate(context.Context, resource.Request) (resource.Set, error) { + return s.out, s.err +} + +func TestResolvePrecedence(t *testing.T) { + tests := []struct { + name string + in resource.ResolveInput + want resource.Set + }{ + { + name: "nothing declared resolves to zero", + in: resource.ResolveInput{}, + want: resource.Set{}, + }, + { + name: "queue default overlays global default per key", + in: resource.ResolveInput{ + GlobalDefault: resource.Set{resource.CPU: 1000, resource.Memory: 1 << 30}, + QueueDefault: resource.Set{resource.Memory: 4 << 30}, + }, + want: resource.Set{resource.CPU: 1000, resource.Memory: 4 << 30}, + }, + { + name: "declaration beats queue default", + in: resource.ResolveInput{ + QueueDefault: resource.Set{resource.Memory: 4 << 30}, + Declared: resource.Set{resource.Memory: 16 << 30}, + }, + want: resource.Set{resource.Memory: 16 << 30}, + }, + { + name: "estimator beats declaration but only on keys it returns", + in: resource.ResolveInput{ + Declared: resource.Set{resource.CPU: 4000, resource.Memory: 16 << 30}, + Estimator: stubEstimator{out: resource.Set{resource.Memory: 6 << 30}}, + }, + want: resource.Set{resource.CPU: 4000, resource.Memory: 6 << 30}, + }, + { + name: "enqueue override beats the estimator", + in: resource.ResolveInput{ + Declared: resource.Set{resource.Memory: 16 << 30}, + Estimator: stubEstimator{out: resource.Set{resource.Memory: 6 << 30}}, + Override: resource.Set{resource.Memory: 48 << 30}, + }, + want: resource.Set{resource.Memory: 48 << 30}, + }, + { + name: "resource func beats declaration and is fed InputBytes", + in: resource.ResolveInput{ + Declared: resource.Set{resource.CPU: 4000}, + Request: resource.Request{InputBytes: 2 << 30}, + Func: func(_ context.Context, r resource.Request) (resource.Set, error) { + return resource.MemoryBytes(r.InputBytes * 3), nil + }, + }, + want: resource.Set{resource.CPU: 4000, resource.Memory: 6 << 30}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := resource.Resolve(context.Background(), tt.in) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if len(got.Requests) != len(tt.want) { + t.Fatalf("got %v, want %v", got.Requests, tt.want) + } + for k, v := range tt.want { + if got.Requests[k] != v { + t.Errorf("key %q: got %d, want %d", k, got.Requests[k], v) + } + } + }) + } +} + +func TestResolveEstimatorErrorFallsBack(t *testing.T) { + got, err := resource.Resolve(context.Background(), resource.ResolveInput{ + Declared: resource.Set{resource.Memory: 8 << 30}, + Estimator: stubEstimator{err: errors.New("rollup unavailable")}, + }) + if err != nil { + t.Fatalf("an estimator error must never fail enqueue: %v", err) + } + if got.Requests[resource.Memory] != 8<<30 { + t.Errorf("got %v, want the declaration preserved", got.Requests) + } +} + +// TestResolveReportsSourceErrors pins the other half of "never fails an +// enqueue": it must not fail SILENTLY either. +// +// Both dynamic sources run once per enqueue, in the enqueuing process, +// so one that is misconfigured is misconfigured for every job of that +// name from then on. Falling back to the declaration with no signal +// means the fleet under-sizes those jobs indefinitely and the symptom — +// OOM kills under real input — points at the handler, not at the +// estimator that stopped answering. +func TestResolveReportsSourceErrors(t *testing.T) { + funcErr := errors.New("input service unreachable") + estErr := errors.New("rollup unavailable") + + reported := map[string]error{} + + got, err := resource.Resolve(context.Background(), resource.ResolveInput{ + Declared: resource.Set{resource.Memory: 8 << 30}, + Func: func(context.Context, resource.Request) (resource.Set, error) { + return nil, funcErr + }, + Estimator: stubEstimator{err: estErr}, + OnError: func(source string, err error) { + reported[source] = err + }, + }) + if err != nil { + t.Fatalf("a reported source error must still never fail enqueue: %v", err) + } + + if got.Requests[resource.Memory] != 8<<30 { + t.Errorf("got %v, want the declaration preserved", got.Requests) + } + + if !errors.Is(reported[resource.SourceFunc], funcErr) { + t.Errorf("func error reported as %v, want %v", reported[resource.SourceFunc], funcErr) + } + + if !errors.Is(reported[resource.SourceEstimator], estErr) { + t.Errorf("estimator error reported as %v, want %v", + reported[resource.SourceEstimator], estErr) + } + + // A source that succeeds reports nothing. + reported = map[string]error{} + + if _, err = resource.Resolve(context.Background(), resource.ResolveInput{ + Estimator: stubEstimator{out: resource.Set{resource.Memory: 1 << 30}}, + OnError: func(source string, err error) { reported[source] = err }, + }); err != nil { + t.Fatalf("Resolve: %v", err) + } + + if len(reported) != 0 { + t.Errorf("a successful estimator reported %v, want nothing", reported) + } +} + +func TestResolveRejectsUnschedulable(t *testing.T) { + _, err := resource.Resolve(context.Background(), resource.ResolveInput{ + Declared: resource.Set{resource.Memory: 64 << 30}, + MaxCapacity: resource.Set{resource.Memory: 32 << 30}, + }) + if !errors.Is(err, resource.ErrUnschedulable) { + t.Fatalf("got %v, want ErrUnschedulable", err) + } + if !strings.Contains(err.Error(), resource.Memory) { + t.Errorf("error must name the dimension that does not fit: %v", err) + } +} + +func TestResolveSkipsCapacityCheckWhenUnknown(t *testing.T) { + // A single-process engine with no registered workers has no known + // capacity. Rejecting on an empty MaxCapacity would reject everything. + _, err := resource.Resolve(context.Background(), resource.ResolveInput{ + Declared: resource.Set{resource.Memory: 64 << 30}, + }) + if err != nil { + t.Fatalf("empty MaxCapacity must disable the check, got %v", err) + } +} + +func TestResolveLimitsDefaultToRequestsExceptCPU(t *testing.T) { + got, err := resource.Resolve(context.Background(), resource.ResolveInput{ + Declared: resource.Set{resource.CPU: 4000, resource.Memory: 8 << 30}, + }) + if err != nil { + t.Fatalf("Resolve() error = %v", err) + } + if got.Limits[resource.Memory] != 8<<30 { + t.Errorf("memory limit should default to the request, got %v", got.Limits) + } + if _, ok := got.Limits[resource.CPU]; ok { + t.Errorf("CPU limit should be unset (burstable), got %v", got.Limits) + } +} diff --git a/store/memory/artifact.go b/store/memory/artifact.go new file mode 100644 index 0000000..a917b19 --- /dev/null +++ b/store/memory/artifact.go @@ -0,0 +1,495 @@ +package memory + +import ( + "context" + "sort" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// CreateArtifact inserts an artifact and, when link is non-nil, its first +// link atomically under the store lock. +func (s *Store) CreateArtifact(_ context.Context, a *artifact.Artifact, link *artifact.Link) error { + s.mu.Lock() + defer s.mu.Unlock() + + for _, existing := range s.artifacts { + if existing.DeletedAt != nil { + continue + } + + if existing.Backend == a.Backend && existing.Bucket == a.Bucket && existing.Key == a.Key { + return artifact.ErrExists + } + } + + s.artifacts[a.ID.String()] = a.Clone() + + if link != nil { + s.appendLinkLocked(link) + } + + return nil +} + +// appendLinkLocked adds a link if an identical one is not already present. +// Callers must hold s.mu. +func (s *Store) appendLinkLocked(link *artifact.Link) { + for _, existing := range s.artifactLinks { + if existing.ArtifactID == link.ArtifactID && + existing.OwnerKind == link.OwnerKind && + existing.OwnerID == link.OwnerID && + existing.Name == link.Name && + existing.Attempt == link.Attempt { + return + } + } + + s.artifactLinks = append(s.artifactLinks, link.Clone()) +} + +// GetArtifact retrieves an artifact by ID, excluding soft-deleted ones. +func (s *Store) GetArtifact(_ context.Context, artifactID id.ArtifactID) (*artifact.Artifact, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + a, ok := s.artifacts[artifactID.String()] + if !ok || a.DeletedAt != nil { + return nil, artifact.ErrNotFound + } + + return a.Clone(), nil +} + +// FindArtifactByKey retrieves an artifact by its storage coordinates. +func (s *Store) FindArtifactByKey(_ context.Context, backend, bucket, key string) (*artifact.Artifact, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + for _, a := range s.artifacts { + if a.DeletedAt != nil { + continue + } + + if a.Backend == backend && a.Bucket == bucket && a.Key == key { + return a.Clone(), nil + } + } + + return nil, artifact.ErrNotFound +} + +// UpdateArtifact persists changes to size, hash, content type, and expiry. +// Lifecycle is deliberately not updatable. +func (s *Store) UpdateArtifact(_ context.Context, a *artifact.Artifact) error { + s.mu.Lock() + defer s.mu.Unlock() + + existing, ok := s.artifacts[a.ID.String()] + if !ok { + return artifact.ErrNotFound + } + + updated := a.Clone() + updated.Lifecycle = existing.Lifecycle + updated.CreatedAt = existing.CreatedAt + updated.DeletedAt = existing.DeletedAt + s.artifacts[a.ID.String()] = updated + + return nil +} + +// ListArtifacts returns artifacts matching the given options, newest first. +func (s *Store) ListArtifacts(_ context.Context, opts artifact.ListOpts) ([]*artifact.Artifact, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + out := make([]*artifact.Artifact, 0, len(s.artifacts)) + + for _, a := range s.artifacts { + if a.DeletedAt != nil && !opts.IncludeDeleted { + continue + } + + if opts.Lifecycle != "" && a.Lifecycle != opts.Lifecycle { + continue + } + + if opts.ScopeAppID != "" && a.ScopeAppID != opts.ScopeAppID { + continue + } + + if opts.ScopeOrgID != "" && a.ScopeOrgID != opts.ScopeOrgID { + continue + } + + out = append(out, a.Clone()) + } + + sort.Slice(out, func(i, j int) bool { + if out[i].CreatedAt.Equal(out[j].CreatedAt) { + return out[i].ID.String() < out[j].ID.String() + } + + return out[i].CreatedAt.After(out[j].CreatedAt) + }) + + return paginate(out, opts.Offset, opts.Limit), nil +} + +func paginate(in []*artifact.Artifact, offset, limit int) []*artifact.Artifact { + if offset > 0 { + if offset >= len(in) { + return nil + } + + in = in[offset:] + } + + if limit > 0 && limit < len(in) { + in = in[:limit] + } + + return in +} + +// LinkArtifact records that an owner references an artifact. Linking the +// same tuple twice is a no-op. +func (s *Store) LinkArtifact(_ context.Context, link *artifact.Link) error { + s.mu.Lock() + defer s.mu.Unlock() + + s.appendLinkLocked(link) + + return nil +} + +// ListLinks returns every link belonging to the given owner. +func (s *Store) ListLinks(_ context.Context, owner artifact.OwnerRef) ([]*artifact.Link, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + return s.linksForOwnerLocked(owner), nil +} + +// linksForOwnerLocked collects an owner's links. Callers must hold s.mu. +func (s *Store) linksForOwnerLocked(owner artifact.OwnerRef) []*artifact.Link { + var out []*artifact.Link + + for _, l := range s.artifactLinks { + if l.OwnerKind == owner.Kind && l.OwnerID == owner.ID { + out = append(out, l.Clone()) + } + } + + return out +} + +// FindLinkByName returns the link for an owner and name with the highest +// attempt number, breaking ties by CreatedAt descending — see the +// artifact.Store doc comment for why ties happen and what the tie-break +// does and does not fix. +func (s *Store) FindLinkByName(_ context.Context, owner artifact.OwnerRef, name string) (*artifact.Link, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + var best *artifact.Link + + for _, l := range s.artifactLinks { + if l.OwnerKind != owner.Kind || l.OwnerID != owner.ID || l.Name != name { + continue + } + + if best == nil || l.Attempt > best.Attempt || + (l.Attempt == best.Attempt && l.CreatedAt.After(best.CreatedAt)) { + best = l + } + } + + if best == nil { + return nil, artifact.ErrNotFound + } + + return best.Clone(), nil +} + +// ListArtifactsByOwner returns the artifacts linked to an owner, +// optionally filtered by role. +func (s *Store) ListArtifactsByOwner( + _ context.Context, + owner artifact.OwnerRef, + role artifact.Role, +) ([]*artifact.Artifact, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + var out []*artifact.Artifact + + seen := make(map[string]bool) + + for _, l := range s.artifactLinks { + if l.OwnerKind != owner.Kind || l.OwnerID != owner.ID { + continue + } + + if role != "" && l.Role != role { + continue + } + + key := l.ArtifactID.String() + if seen[key] { + continue + } + + a, ok := s.artifacts[key] + if !ok || a.DeletedAt != nil { + continue + } + + seen[key] = true + + out = append(out, a.Clone()) + } + + return out, nil +} + +// SweepEphemeral marks eligible ephemeral artifacts as deleted. +// +// An artifact is eligible when every owner that links it is terminal and +// the retention window has elapsed since the last of them finished. The +// lifecycle guard is the first statement of the loop body and is written +// as a literal, mirroring the SQL backends. +func (s *Store) SweepEphemeral(_ context.Context, opts artifact.SweepOpts) ([]*artifact.Artifact, error) { + s.mu.Lock() + defer s.mu.Unlock() + + now := time.Now().UTC() + + var out []*artifact.Artifact + + for _, a := range s.artifacts { + if a.Lifecycle != artifact.Ephemeral { + continue + } + + if a.DeletedAt != nil { + continue + } + + if opts.Limit > 0 && len(out) >= opts.Limit { + break + } + + links := s.linksForArtifactLocked(a.ID) + if len(links) == 0 { + // Orphans are SweepOrphans' business, not ours. + continue + } + + terminalAt, ok := s.ownersTerminalAtLocked(links) + if !ok { + continue + } + + if a.ExpiresAt != nil { + if a.ExpiresAt.After(now) { + continue + } + } else if terminalAt.Add(opts.Retention).After(now) { + continue + } + + if !opts.DryRun { + deleted := now + a.DeletedAt = &deleted + } + + out = append(out, a.Clone()) + } + + return out, nil +} + +// linksForArtifactLocked collects links pointing at an artifact. +// Callers must hold s.mu. +func (s *Store) linksForArtifactLocked(artifactID id.ArtifactID) []*artifact.Link { + var out []*artifact.Link + + for _, l := range s.artifactLinks { + if l.ArtifactID == artifactID { + out = append(out, l) + } + } + + return out +} + +// ownersTerminalAtLocked reports the latest terminal time across every +// owner in links, and whether all of them are in fact terminal. An owner +// that no longer exists counts as terminal at the link's creation time — +// its job or run row was purged, so it cannot still be running. +// Callers must hold s.mu. +func (s *Store) ownersTerminalAtLocked(links []*artifact.Link) (time.Time, bool) { + var latest time.Time + + for _, l := range links { + at, ok := s.ownerTerminalAtLocked(l) + if !ok { + return time.Time{}, false + } + + if at.After(latest) { + latest = at + } + } + + return latest, true +} + +func (s *Store) ownerTerminalAtLocked(l *artifact.Link) (time.Time, bool) { + switch l.OwnerKind { + case artifact.OwnerJob: + j, ok := s.jobs[l.OwnerID] + if !ok { + return l.CreatedAt, true + } + + if !isTerminalJobState(string(j.State)) { + return time.Time{}, false + } + + if j.CompletedAt != nil { + return *j.CompletedAt, true + } + + return j.UpdatedAt, true + + case artifact.OwnerRun, artifact.OwnerStep: + r, ok := s.runs[l.OwnerID] + if !ok { + return l.CreatedAt, true + } + + if !isTerminalRunState(string(r.State)) { + return time.Time{}, false + } + + if r.CompletedAt != nil { + return *r.CompletedAt, true + } + + return r.UpdatedAt, true + + default: + return l.CreatedAt, true + } +} + +func isTerminalJobState(state string) bool { + switch state { + case "completed", "failed", "cancelled": + return true + default: + return false + } +} + +func isTerminalRunState(state string) bool { + switch state { + case "completed", "failed", "cancelled": + return true + default: + return false + } +} + +// SweepOrphans marks ephemeral artifacts with no links at all that were +// created before the cutoff. +func (s *Store) SweepOrphans(_ context.Context, cutoff time.Time, limit int) ([]*artifact.Artifact, error) { + s.mu.Lock() + defer s.mu.Unlock() + + now := time.Now().UTC() + + var out []*artifact.Artifact + + for _, a := range s.artifacts { + if a.Lifecycle != artifact.Ephemeral { + continue + } + + if a.DeletedAt != nil { + continue + } + + if limit > 0 && len(out) >= limit { + break + } + + if !a.CreatedAt.Before(cutoff) { + continue + } + + if len(s.linksForArtifactLocked(a.ID)) > 0 { + continue + } + + deleted := now + a.DeletedAt = &deleted + + out = append(out, a.Clone()) + } + + return out, nil +} + +// ListPurgeable returns soft-deleted artifacts older than grace. +func (s *Store) ListPurgeable(_ context.Context, grace time.Duration, limit int) ([]*artifact.Artifact, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + now := time.Now().UTC() + + var out []*artifact.Artifact + + for _, a := range s.artifacts { + if a.DeletedAt == nil { + continue + } + + if limit > 0 && len(out) >= limit { + break + } + + if a.DeletedAt.Add(grace).After(now) { + continue + } + + out = append(out, a.Clone()) + } + + return out, nil +} + +// PurgeArtifact hard-deletes an artifact and its links. +func (s *Store) PurgeArtifact(_ context.Context, artifactID id.ArtifactID) error { + s.mu.Lock() + defer s.mu.Unlock() + + delete(s.artifacts, artifactID.String()) + + kept := s.artifactLinks[:0] + + for _, l := range s.artifactLinks { + if l.ArtifactID != artifactID { + kept = append(kept, l) + } + } + + s.artifactLinks = kept + + return nil +} diff --git a/store/memory/artifact_test.go b/store/memory/artifact_test.go new file mode 100644 index 0000000..337b48e --- /dev/null +++ b/store/memory/artifact_test.go @@ -0,0 +1,13 @@ +package memory_test + +import ( + "testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/store/memory" +) + +func TestArtifactStoreConformance(t *testing.T) { + artifacttest.RunStoreSuite(t, func(*testing.T) artifact.Store { return memory.New() }) +} diff --git a/store/memory/dequeue_test.go b/store/memory/dequeue_test.go new file mode 100644 index 0000000..1250f1c --- /dev/null +++ b/store/memory/dequeue_test.go @@ -0,0 +1,20 @@ +package memory_test + +import ( + "testing" + + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" + "github.com/xraph/dispatch/store/storetest" +) + +// TestDequeueConformance runs the resource-aware dequeue suite against the +// memory store — the reference implementation the SQL and document backends +// (Tasks 15-18) are checked against. +func TestDequeueConformance(t *testing.T) { + storetest.RunDequeueSuite(t, func(t *testing.T) job.Store { + t.Helper() + + return memory.New() + }) +} diff --git a/store/memory/lease.go b/store/memory/lease.go new file mode 100644 index 0000000..2928aff --- /dev/null +++ b/store/memory/lease.go @@ -0,0 +1,161 @@ +package memory + +import ( + "context" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// Compile-time check that the memory store provides the lease capability. +// +// The grant itself is not here: it travels on job.DequeueOpts and is +// applied by DequeueJobs, under the same write lock that performs the +// claim. This file holds only what a lease needs afterwards. +var _ job.LeaseStore = (*Store)(nil) + +// RenewLease extends the lease only if the caller still holds it. +func (m *Store) RenewLease( + _ context.Context, + jobID id.JobID, + workerID id.WorkerID, + epoch int, + leaseUntil time.Time, +) error { + m.mu.Lock() + defer m.mu.Unlock() + + j, ok := m.jobs[jobID.String()] + if !ok { + return job.ErrLeaseLost + } + if j.State != job.StateRunning || j.WorkerID != workerID || j.LeaseEpoch != epoch { + return job.ErrLeaseLost + } + + now := time.Now().UTC() + until := leaseUntil + beat := now + + j.LeaseExpiresAt = &until + j.HeartbeatAt = &beat + j.UpdatedAt = now + + return nil +} + +// reclaimable reports whether a running job should be taken back. +// +// The first clause is the actual rule, and job.Lease.IsExpired remains its +// only authority: a lease was granted and has lapsed. +// +// The second is the same narrow exception the four persistent backends +// apply to a running job carrying no lease at all — see +// job.UnleasedReclaimGrace for why it is gated on silence rather than on +// the null expiry alone. Only half of that rationale reaches this backend: +// memory loses every row on restart, so it has no pre-lease build to +// inherit rows from, but a claim through DequeueJobs without lease options +// still produces an unleased running job, and abandoning one would +// otherwise strand it for the life of the process. +func reclaimable(j *job.Job, now time.Time) bool { + if j.LeaseExpiresAt != nil { + return job.Lease{ExpiresAt: *j.LeaseExpiresAt}.IsExpired(now) + } + + // Heartbeat first, falling back to the claim time for a caller that + // stopped before its first beat — the same order ReapStaleJobs used. + silent := j.HeartbeatAt + if silent == nil { + silent = j.StartedAt + } + if silent == nil { + return false + } + + return silent.Before(now.Add(-job.UnleasedReclaimGrace)) +} + +// ReclaimExpiredLeases returns expired-lease jobs to pending, fencing +// their previous holders. +// +// A non-positive limit claims nothing, matching DequeueOpts.Limit's +// behavior rather than reading zero or negative as unlimited. +func (m *Store) ReclaimExpiredLeases(_ context.Context, limit int) ([]*job.Job, error) { + if limit <= 0 { + return nil, nil + } + + m.mu.Lock() + defer m.mu.Unlock() + + now := time.Now().UTC() + + reclaimed := make([]*job.Job, 0, len(m.jobs)) + for _, j := range m.jobs { + if len(reclaimed) >= limit { + break + } + if j.State != job.StateRunning { + continue + } + if !reclaimable(j, now) { + continue + } + + j.State = job.StatePending + j.RunAt = now + j.WorkerID = id.WorkerID{} + j.StartedAt = nil + j.HeartbeatAt = nil + j.LeaseExpiresAt = nil + j.LeaseEpoch++ + j.EvictCount++ + j.UpdatedAt = now + + reclaimed = append(reclaimed, cloneJob(j)) + } + + return reclaimed, nil +} + +// UpdateLeasedJob persists j only while the caller still holds the +// lease — still running, still assigned to workerID, still at epoch. +// Otherwise it returns job.ErrLeaseLost and leaves the stored job +// untouched. +// +// Only the business fields move. lease_epoch, lease_expires_at, +// worker_id, and heartbeat_at are copied from the CURRENTLY STORED job, +// never from j: j is the caller's stale snapshot, and every renewal +// since it was taken has pushed the real expiry forward. Overwriting +// that with j's copy would roll the winner's lease backwards even +// though the epoch check passed — see job.LeaseStore.UpdateLeasedJob. +func (m *Store) UpdateLeasedJob(_ context.Context, j *job.Job, workerID id.WorkerID, epoch int) error { + m.mu.Lock() + defer m.mu.Unlock() + + cur, ok := m.jobs[j.ID.String()] + if !ok { + return dispatch.ErrJobNotFound + } + if cur.State != job.StateRunning || cur.WorkerID != workerID || cur.LeaseEpoch != epoch { + return job.ErrLeaseLost + } + + leaseEpoch := cur.LeaseEpoch + leaseExpiresAt := cur.LeaseExpiresAt + leaseWorkerID := cur.WorkerID + heartbeatAt := cur.HeartbeatAt + + cp := cloneJob(j) + cp.LeaseEpoch = leaseEpoch + cp.LeaseExpiresAt = leaseExpiresAt + cp.WorkerID = leaseWorkerID + cp.HeartbeatAt = heartbeatAt + cp.UpdatedAt = time.Now().UTC() + + m.jobs[j.ID.String()] = cp + + return nil +} diff --git a/store/memory/lease_test.go b/store/memory/lease_test.go new file mode 100644 index 0000000..918723f --- /dev/null +++ b/store/memory/lease_test.go @@ -0,0 +1,349 @@ +package memory_test + +import ( + "context" + "testing" + "time" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/store/memory" + "github.com/xraph/dispatch/store/storetest" +) + +// findJob returns the job with the given ID from jobs, or nil. +// +// ReclaimExpiredLeases is not queue-scoped (see storetest.Contains), so a +// test asserting on a specific job must find it by ID rather than assume +// it is alone in the result. +func findJob(jobs []*job.Job, jobID id.JobID) *job.Job { + for _, j := range jobs { + if j.ID == jobID { + return j + } + } + + return nil +} + +func TestLeaseConformance(t *testing.T) { + storetest.RunLeaseSuite(t, func(t *testing.T) storetest.LeaseStore { + t.Helper() + + return memory.New() + }) +} + +// TestLeaseStoreDoesNotAliasResourceMap covers the same class of bug as +// TestMemoryStoreDoesNotAliasResourceMap (resource_test.go), but for the +// lease-granting paths: the leased claim and ReclaimExpiredLeases both +// used to hand back a job built from a shallow struct copy, aliasing +// Resources, ResourceLimits, Payload, and ArtifactBindings against the +// stored job. A worker mutating its leased job's Resources would silently +// rewrite the stored requirement. +// +// The shared lease conformance suite in store/storetest/lease.go does not +// check this — it runs against backends where a shallow struct copy isn't +// even the mechanism, so this case lives here instead. +func TestLeaseStoreDoesNotAliasResourceMap(t *testing.T) { + t.Run("DequeueJobsWithGrant", func(t *testing.T) { + st := memory.New() + ctx := context.Background() + worker := id.NewWorkerID() + const queue = "lease-alias-dequeue" + + j := storetest.PendingJob("alias-dequeue", queue, 0) + j.Resources = resource.Set{resource.CPU: 1000, resource.Memory: 8 << 30} + if err := st.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := st.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: time.Now().UTC().Add(time.Minute), + }) + if err != nil { + t.Fatalf("DequeueJobs() error = %v", err) + } + if len(got) != 1 { + t.Fatalf("DequeueJobs() returned %d jobs, want 1", len(got)) + } + + // A caller mutating what the claim handed back must not rewrite the + // stored job. + got[0].Resources[resource.Memory] = 1 + + stored, err := st.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if stored.Resources[resource.Memory] != 8<<30 { + t.Fatalf("the leased claim aliased the stored map: memory = %d", + stored.Resources[resource.Memory]) + } + }) + + t.Run("ReclaimExpiredLeases", func(t *testing.T) { + st := memory.New() + ctx := context.Background() + const queue = "lease-alias-reclaim" + + j := storetest.RunningJob("alias-reclaim", queue, 0) + j.Resources = resource.Set{resource.CPU: 1000, resource.Memory: 8 << 30} + if err := st.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + reclaimed, err := st.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("ReclaimExpiredLeases() error = %v", err) + } + if !storetest.Contains(reclaimed, j.ID) { + t.Fatalf("reclaimed set does not contain %s", j.ID) + } + + target := findJob(reclaimed, j.ID) + if target == nil { + t.Fatalf("could not find reclaimed job %s in result", j.ID) + } + + // A caller mutating what ReclaimExpiredLeases handed back must not + // rewrite the stored job. + target.Resources[resource.Memory] = 1 + + stored, err := st.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if stored.Resources[resource.Memory] != 8<<30 { + t.Fatalf("ReclaimExpiredLeases aliased the stored map: memory = %d", + stored.Resources[resource.Memory]) + } + }) +} + +// TestReclaimAdoptsRunningJobsWithoutLease covers a running job carrying +// no lease at all, matching the rule the four persistent backends apply. +// +// Memory has no upgrade to survive, since it loses every row on restart, +// so the pre-lease-build half of the problem cannot reach it. The other +// half can: DequeueOpts.Grants() is false whenever LeaseUntil is zero, so +// a caller claiming through job.Store without lease options holds a +// running job with no lease, and if that caller stops without completing +// the job, nothing reclaims it for the life of the process. Reclamation +// skips a zero expiry and dequeue claims only pending and retrying rows. +// +// The negative cases carry the safety argument: a null expiry does not by +// itself mean the job was abandoned, so silence is what separates a dead +// claim from a live one. +func TestReclaimAdoptsRunningJobsWithoutLease(t *testing.T) { + s := memory.New() + ctx := context.Background() + now := time.Now().UTC() + + // unleased builds a running job with no lease fields at all, the shape + // a claim through DequeueJobs without lease options leaves behind. + unleased := func(name string, started, beat *time.Time) *job.Job { + return &job.Job{ + ID: id.NewJobID(), + Name: name, + Queue: "default", + Payload: []byte(`{}`), + State: job.StateRunning, + MaxRetries: 3, + RunAt: now, + StartedAt: started, + HeartbeatAt: beat, + } + } + + at := func(d time.Duration) *time.Time { + t := now.Add(-d) + + return &t + } + + silent := job.UnleasedReclaimGrace + time.Minute + + cases := []struct { + j *job.Job + want bool + why string + }{ + { + j: unleased("silent-heartbeat", at(2*silent), at(silent)), + want: true, + why: "abandoned by a caller that stopped reporting", + }, + { + // Old claim, current heartbeat: pins that heartbeat wins over + // started_at rather than both needing to be fresh. + j: unleased("fresh-heartbeat", at(2*silent), at(0)), + want: false, + why: "still reporting, so it belongs to a live caller", + }, + { + j: unleased("no-heartbeat-old-start", at(silent), nil), + want: true, + why: "claimed long ago and never heartbeated", + }, + { + j: unleased("no-heartbeat-fresh-start", at(0), nil), + want: false, + why: "just claimed; its first heartbeat is not due yet", + }, + { + j: unleased("no-times", nil, nil), + want: false, + why: "no timestamp to establish age from", + }, + } + + for _, c := range cases { + if err := s.EnqueueJob(ctx, c.j); err != nil { + t.Fatalf("enqueue %s: %v", c.j.Name, err) + } + } + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + + for _, c := range cases { + got := findJob(reclaimed, c.j.ID) != nil + if got == c.want { + continue + } + if c.want { + t.Errorf("%s was not reclaimed but should have been: %s", c.j.Name, c.why) + } else { + t.Errorf("%s was reclaimed but must not be: %s", c.j.Name, c.why) + } + } +} + +// TestClearOwnershipStopsTheRequeueLivelock reproduces what a stale lease +// column does to a job returned to pending from outside the lease +// machinery, which is the shape api.retryJob wrote before it called +// job.Job.ClearOwnership. +// +// UpdateJob writes the whole row on every backend, lease columns included, +// so a failed job put back to pending keeps the LeaseExpiresAt of its +// failed run. Nothing notices while it sits pending. It goes wrong at the +// next claim that grants no lease, which is a supported call +// (DequeueOpts.Grants() is false whenever LeaseUntil is zero): that claim +// writes state and worker but never touches the expiry, so the job enters +// running already holding a lapsed lease and the next sweep takes it +// straight back. Claimed, reclaimed, claimed again, forever. +// +// The memory store is used because the bug is in the shared job row rather +// than in any one backend's SQL, and this keeps the reproduction in-process. +func TestClearOwnershipStopsTheRequeueLivelock(t *testing.T) { + ctx := context.Background() + + // requeued builds the row a retry path produces, with or without the + // ownership reset, and returns it after one no-lease claim. + requeued := func(t *testing.T, resetOwnership bool) (*memory.Store, *job.Job) { + t.Helper() + + s := memory.New() + lapsed := time.Now().UTC().Add(-time.Hour) + started := lapsed.Add(-time.Minute) + j := &job.Job{ + ID: id.NewJobID(), + Name: "retried", + Queue: "default", + Payload: []byte(`{}`), + State: job.StateFailed, + MaxRetries: 3, + WorkerID: id.NewWorkerID(), + StartedAt: &started, + HeartbeatAt: &lapsed, + LeaseExpiresAt: &lapsed, + LeaseEpoch: 4, + } + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + // What api.retryJob does to a failed job. + j.State = job.StatePending + j.RetryCount = 0 + j.LastError = "" + j.RunAt = time.Now().UTC() + j.CompletedAt = nil + if resetOwnership { + j.ClearOwnership() + } else { + j.StartedAt = nil // the old code cleared only this + } + if err := s.UpdateJob(ctx, j); err != nil { + t.Fatalf("update: %v", err) + } + + // A claim that grants no lease: legal, and it never writes the + // lease columns. + claimed, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + }) + if err != nil { + t.Fatalf("dequeue: %v", err) + } + if len(claimed) != 1 { + t.Fatalf("claimed %d jobs, want 1", len(claimed)) + } + + return s, j + } + + t.Run("without the reset the claim is immediately reclaimed", func(t *testing.T) { + s, j := requeued(t, false) + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 10) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + if findJob(reclaimed, j.ID) == nil { + t.Skip("the stale expiry no longer reaches running; this reproduction is obsolete") + } + }) + + t.Run("with the reset the claim survives", func(t *testing.T) { + s, j := requeued(t, true) + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 10) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + if findJob(reclaimed, j.ID) != nil { + t.Fatal("a freshly claimed job was reclaimed: the retry path left a lapsed " + + "lease on the row, so it can never run to completion") + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.State != job.StateRunning { + t.Errorf("State = %s, want running", got.State) + } + // The fencing token must not have been rolled back by the reset. + if got.LeaseEpoch < 4 { + t.Errorf("LeaseEpoch = %d, want >= 4: a fencing token must never move backwards", + got.LeaseEpoch) + } + }) +} + +func TestDLQConformance(t *testing.T) { + storetest.RunDLQSuite(t, func(t *testing.T) storetest.DLQStore { + t.Helper() + + return memory.New() + }) +} diff --git a/store/memory/resource_test.go b/store/memory/resource_test.go new file mode 100644 index 0000000..257c62b --- /dev/null +++ b/store/memory/resource_test.go @@ -0,0 +1,83 @@ +package memory_test + +import ( + "context" + "testing" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/store/memory" +) + +func newResourceJob(t *testing.T) *job.Job { + t.Helper() + + return &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "tessellate.model", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + Resources: resource.Set{resource.CPU: 4000, resource.Memory: 16 << 30}, + InputBytes: 4 << 30, + } +} + +func TestMemoryStoreRoundTripsResources(t *testing.T) { + st := memory.New() + ctx := context.Background() + + j := newResourceJob(t) + if err := st.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := st.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + + if got.Resources[resource.Memory] != 16<<30 { + t.Errorf("memory = %d, want 16 GiB", got.Resources[resource.Memory]) + } + if got.InputBytes != 4<<30 { + t.Errorf("InputBytes = %d, want 4 GiB", got.InputBytes) + } +} + +func TestMemoryStoreDoesNotAliasResourceMap(t *testing.T) { + st := memory.New() + ctx := context.Background() + + j := newResourceJob(t) + if err := st.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + // A caller mutating its own copy must not rewrite the stored job. + j.Resources[resource.Memory] = 1 + + got, err := st.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if got.Resources[resource.Memory] != 16<<30 { + t.Fatalf("the store aliased the caller's map: memory = %d", + got.Resources[resource.Memory]) + } + + // And a caller mutating what it read must not rewrite it either. + got.Resources[resource.Memory] = 2 + + again, err := st.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if again.Resources[resource.Memory] != 16<<30 { + t.Fatalf("the store returned an aliased map: memory = %d", + again.Resources[resource.Memory]) + } +} diff --git a/store/memory/store.go b/store/memory/store.go index ad9392a..e726218 100644 --- a/store/memory/store.go +++ b/store/memory/store.go @@ -2,11 +2,13 @@ package memory import ( "context" + "fmt" "sort" "sync" "time" "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" "github.com/xraph/dispatch/cluster" "github.com/xraph/dispatch/cron" "github.com/xraph/dispatch/dlq" @@ -25,6 +27,7 @@ var ( _ dlq.Store = (*Store)(nil) _ event.Store = (*Store)(nil) _ cluster.Store = (*Store)(nil) + _ artifact.Store = (*Store)(nil) ) // Store is a fully in-memory implementation of store.Store. @@ -40,6 +43,9 @@ type Store struct { events map[string]*event.Event workers map[string]*cluster.Worker + artifacts map[string]*artifact.Artifact + artifactLinks []*artifact.Link + // leader tracks the current cluster leader worker ID string. leader string leaderUntil time.Time @@ -55,6 +61,7 @@ func New() *Store { dlqs: make(map[string]*dlq.Entry), events: make(map[string]*event.Event), workers: make(map[string]*cluster.Worker), + artifacts: make(map[string]*artifact.Artifact), } } @@ -75,6 +82,39 @@ func (m *Store) Close() error { return nil } // Job Store // ────────────────────────────────────────────────── +// cloneJob deep-copies the fields that are reference types. +// +// Every copy in this store used to be a shallow struct copy, which is +// correct for scalars and wrong for maps and slices: the copy would +// alias the caller's underlying data, so a handler mutating its own job +// (or a caller mutating what it read back) would silently rewrite the +// stored job. +// +// StartedAt, CompletedAt, HeartbeatAt, and LeaseExpiresAt are *time.Time +// and are deliberately left pointer-shared rather than deep-copied: +// time.Time has no exported method that mutates the value in place, and +// nothing in this codebase writes through a *time.Time (it's always +// reassigned via `j.Field = &newTime`, never `*j.Field = newTime`), so +// two Job structs sharing the same pointee cannot observe each other's +// changes. +func cloneJob(j *job.Job) *job.Job { + out := *j + out.Resources = j.Resources.Clone() + out.ResourceLimits = j.ResourceLimits.Clone() + + if j.Payload != nil { + out.Payload = make([]byte, len(j.Payload)) + copy(out.Payload, j.Payload) + } + + if j.ArtifactBindings != nil { + out.ArtifactBindings = make([]byte, len(j.ArtifactBindings)) + copy(out.ArtifactBindings, j.ArtifactBindings) + } + + return &out +} + // EnqueueJob persists a new job in pending state. func (m *Store) EnqueueJob(_ context.Context, j *job.Job) error { m.mu.Lock() @@ -84,19 +124,43 @@ func (m *Store) EnqueueJob(_ context.Context, j *job.Job) error { if _, exists := m.jobs[key]; exists { return dispatch.ErrJobAlreadyExists } - cp := *j - m.jobs[key] = &cp + m.jobs[key] = cloneJob(j) return nil } -// DequeueJobs atomically claims up to limit pending jobs from the given -// queues, sets them to running, and returns them. -func (m *Store) DequeueJobs(_ context.Context, queues []string, limit int) ([]*job.Job, error) { +// DequeueJobs atomically claims up to opts.Limit ready jobs from +// opts.Queues that fit opts, sets them to running, and returns them +// ordered by priority descending, then locality-preferred first, then +// RunAt ascending. +// +// A non-positive Limit claims nothing: a worker computing zero free slots +// must claim zero jobs, matching the SQL backends' `LIMIT 0` behavior +// rather than reading zero as "unlimited". The fit predicate itself is +// job.DequeueOpts.Allows / Less, not reimplemented here, so this store +// stays the reference the SQL backends are checked against. +// +// When opts.Grants() the claim also grants a lease, under the one write +// lock that already performs the claim. When it does not, the claim writes +// a running job with no lease at all, which is a supported shape rather +// than a broken one. Such a job is not invisible to ReclaimExpiredLeases: +// job.Lease.IsExpired still reports false for its zero expiry, but reclaim +// adopts it once it has gone silent for job.UnleasedReclaimGrace, so +// abandoning one no longer strands it for the life of the process. See +// job.DequeueOpts.LeaseUntil. +func (m *Store) DequeueJobs(_ context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + if opts.Limit <= 0 { + return nil, nil + } + + if err := opts.Validate(); err != nil { + return nil, fmt.Errorf("dispatch/memory: dequeue jobs: %w", err) + } + m.mu.Lock() defer m.mu.Unlock() - queueSet := make(map[string]struct{}, len(queues)) - for _, q := range queues { + queueSet := make(map[string]struct{}, len(opts.Queues)) + for _, q := range opts.Queues { queueSet[q] = struct{}{} } @@ -116,19 +180,26 @@ func (m *Store) DequeueJobs(_ context.Context, queues []string, limit int) ([]*j continue } } + // IsUnbounded skips the fit predicate entirely: a caller not + // using the resource model claims everything, including jobs + // declaring custom resources. This must be evaluated as part of + // the claim below, never applied after — a job that does not fit + // stays pending and untouched. + if !opts.IsUnbounded() && !opts.Allows(j) { + continue + } candidates = append(candidates, j) } - // Sort: priority DESC, RunAt ASC. + // Order, then truncate: priority DESC, then locality-preferred before + // not (a tiebreak strictly within a priority band, never above it), + // then RunAt ASC. sort.Slice(candidates, func(i, k int) bool { - if candidates[i].Priority != candidates[k].Priority { - return candidates[i].Priority > candidates[k].Priority - } - return candidates[i].RunAt.Before(candidates[k].RunAt) + return opts.Less(candidates[i], candidates[k]) }) - if limit > 0 && len(candidates) > limit { - candidates = candidates[:limit] + if len(candidates) > opts.Limit { + candidates = candidates[:opts.Limit] } result := make([]*job.Job, len(candidates)) @@ -136,9 +207,21 @@ func (m *Store) DequeueJobs(_ context.Context, queues []string, limit int) ([]*j j.State = job.StateRunning n := now j.StartedAt = &n - // Return a copy so callers can mutate without racing with the store. - cp := *j - result[i] = &cp + j.UpdatedAt = now + + if opts.Grants() { + until := opts.LeaseUntil + j.WorkerID = opts.WorkerID + j.LeaseEpoch++ + j.LeaseExpiresAt = &until + } + + // Return a copy so callers can mutate without racing with the + // store. cloneJob deep-copies Resources, ResourceLimits, Payload + // and ArtifactBindings; a shallow struct copy here would let a + // worker mutating its claimed job rewrite the stored requirement + // (TestLeaseStoreDoesNotAliasResourceMap). + result[i] = cloneJob(j) } return result, nil @@ -153,8 +236,7 @@ func (m *Store) GetJob(_ context.Context, jobID id.JobID) (*job.Job, error) { if !ok { return nil, dispatch.ErrJobNotFound } - cp := *j - return &cp, nil + return cloneJob(j), nil } // UpdateJob persists changes to an existing job. @@ -166,9 +248,9 @@ func (m *Store) UpdateJob(_ context.Context, j *job.Job) error { if _, ok := m.jobs[key]; !ok { return dispatch.ErrJobNotFound } - cp := *j + cp := cloneJob(j) cp.UpdatedAt = time.Now().UTC() - m.jobs[key] = &cp + m.jobs[key] = cp return nil } @@ -198,8 +280,7 @@ func (m *Store) ListJobsByState(_ context.Context, state job.State, opts job.Lis if opts.Queue != "" && j.Queue != opts.Queue { continue } - cp := *j - result = append(result, &cp) + result = append(result, cloneJob(j)) } // Sort by CreatedAt for deterministic output. @@ -251,8 +332,7 @@ func (m *Store) ReapStaleJobs(_ context.Context, threshold time.Duration) ([]*jo expired := (j.HeartbeatAt != nil && j.HeartbeatAt.Before(cutoff)) || (j.HeartbeatAt == nil && j.StartedAt != nil && j.StartedAt.Before(cutoff)) if expired { - cp := *j - stale = append(stale, &cp) + stale = append(stale, cloneJob(j)) } } return stale, nil diff --git a/store/memory/store_test.go b/store/memory/store_test.go index b65402d..f78828e 100644 --- a/store/memory/store_test.go +++ b/store/memory/store_test.go @@ -152,7 +152,7 @@ func TestJobDequeue(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - jobs, err := s.DequeueJobs(ctx, tt.queues, tt.limit) + jobs, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: tt.queues, Limit: tt.limit}) if err != nil { t.Fatalf("DequeueJobs: %v", err) } @@ -188,7 +188,7 @@ func TestJobDequeueLimitAndRunAt(t *testing.T) { } } - jobs, err := s.DequeueJobs(ctx, []string{"default"}, 10) + jobs, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 10}) if err != nil { t.Fatalf("DequeueJobs: %v", err) } diff --git a/store/mongo/artifact.go b/store/mongo/artifact.go new file mode 100644 index 0000000..43b0b24 --- /dev/null +++ b/store/mongo/artifact.go @@ -0,0 +1,564 @@ +package mongo + +import ( + "context" + "fmt" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo/options" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// defaultSweepLimit bounds an unbounded sweep so a single pass can never +// touch an unbounded number of documents. +const defaultSweepLimit = 1000 + +// ephemeralOnly is the lifecycle guard shared by every sweep path. +// +// It is a function returning a fresh literal rather than a package +// variable so no caller can mutate the shared map and widen a sweep to +// reach durable artifacts. +func ephemeralOnly() bson.M { + return bson.M{ + "lifecycle": string(artifact.Ephemeral), + "deleted_at": nil, + } +} + +// CreateArtifact inserts an artifact and, when link is non-nil, its first +// link. +// +// Mongo multi-document atomicity needs a replica set. When the deployment +// supports transactions the pair is written in a session; otherwise the +// artifact is written first and the link second, and a crash between them +// leaves an orphan that the orphan sweep collects. +func (s *Store) CreateArtifact(ctx context.Context, a *artifact.Artifact, link *artifact.Link) error { + _, err := s.mdb.Collection(colArtifacts).InsertOne(ctx, toArtifactModel(a)) + if err != nil { + if isDuplicateKey(err) { + return artifact.ErrExists + } + + return fmt.Errorf("dispatch/mongo: create artifact: %w", err) + } + + if link == nil { + return nil + } + + if lerr := s.LinkArtifact(ctx, link); lerr != nil { + // The artifact exists but is unlinked. Rather than leave a + // permanent orphan, drop it so the caller can retry cleanly. + if _, derr := s.mdb.Collection(colArtifacts). + DeleteOne(ctx, bson.M{"_id": a.ID.String()}); derr != nil { + s.logger.Warn("dispatch/mongo: could not roll back unlinked artifact", + log.String("artifact_id", a.ID.String()), + log.String("error", derr.Error()), + ) + } + + return lerr + } + + return nil +} + +// GetArtifact retrieves a live artifact by ID. +func (s *Store) GetArtifact(ctx context.Context, artifactID id.ArtifactID) (*artifact.Artifact, error) { + var m artifactModel + + err := s.mdb.Collection(colArtifacts). + FindOne(ctx, bson.M{"_id": artifactID.String(), "deleted_at": nil}). + Decode(&m) + if err != nil { + if isNoDocuments(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf("dispatch/mongo: get artifact: %w", err) + } + + return fromArtifactModel(&m) +} + +// FindArtifactByKey retrieves a live artifact by its storage coordinates. +func (s *Store) FindArtifactByKey(ctx context.Context, backend, bucket, key string) (*artifact.Artifact, error) { + var m artifactModel + + filter := bson.M{"backend": backend, "bucket": bucket, "key": key, "deleted_at": nil} + + err := s.mdb.Collection(colArtifacts).FindOne(ctx, filter).Decode(&m) + if err != nil { + if isNoDocuments(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf("dispatch/mongo: find artifact by key: %w", err) + } + + return fromArtifactModel(&m) +} + +// UpdateArtifact persists size, hash, content type, and expiry. Lifecycle, +// created_at, and deleted_at are deliberately not updatable here. +func (s *Store) UpdateArtifact(ctx context.Context, a *artifact.Artifact) error { + update := bson.M{"$set": bson.M{ + "size": a.Size, + "content_hash": a.ContentHash, + "content_type": a.ContentType, + "expires_at": a.ExpiresAt, + }} + + res, err := s.mdb.Collection(colArtifacts). + UpdateOne(ctx, bson.M{"_id": a.ID.String()}, update) + if err != nil { + return fmt.Errorf("dispatch/mongo: update artifact: %w", err) + } + + if res.MatchedCount == 0 { + return artifact.ErrNotFound + } + + return nil +} + +// ListArtifacts returns artifacts matching the given options, newest first. +func (s *Store) ListArtifacts(ctx context.Context, opts artifact.ListOpts) ([]*artifact.Artifact, error) { + filter := bson.M{} + + if !opts.IncludeDeleted { + filter["deleted_at"] = nil + } + + if opts.Lifecycle != "" { + filter["lifecycle"] = string(opts.Lifecycle) + } + + if opts.ScopeAppID != "" { + filter["scope_app_id"] = opts.ScopeAppID + } + + if opts.ScopeOrgID != "" { + filter["scope_org_id"] = opts.ScopeOrgID + } + + find := options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}, {Key: "_id", Value: 1}}) + + if opts.Limit > 0 { + find = find.SetLimit(int64(opts.Limit)) + } + + if opts.Offset > 0 { + find = find.SetSkip(int64(opts.Offset)) + } + + return s.findArtifacts(ctx, filter, find) +} + +// findArtifacts runs a find and decodes the whole cursor. +func (s *Store) findArtifacts( + ctx context.Context, + filter bson.M, + opts ...options.Lister[options.FindOptions], +) ([]*artifact.Artifact, error) { + cur, err := s.mdb.Collection(colArtifacts).Find(ctx, filter, opts...) + if err != nil { + return nil, fmt.Errorf("dispatch/mongo: find artifacts: %w", err) + } + + var models []artifactModel + if err := cur.All(ctx, &models); err != nil { + return nil, fmt.Errorf("dispatch/mongo: decode artifacts: %w", err) + } + + return fromArtifactModels(models) +} + +// LinkArtifact records that an owner references an artifact, idempotently. +func (s *Store) LinkArtifact(ctx context.Context, link *artifact.Link) error { + m := toLinkModel(link) + + filter := bson.M{ + "artifact_id": m.ArtifactID, + "owner_kind": m.OwnerKind, + "owner_id": m.OwnerID, + "name": m.Name, + "attempt": m.Attempt, + } + + _, err := s.mdb.Collection(colArtifactLinks). + UpdateOne(ctx, filter, bson.M{"$setOnInsert": m}, options.UpdateOne().SetUpsert(true)) + if err != nil { + if isDuplicateKey(err) { + // A concurrent upsert won the race; the link exists either way. + return nil + } + + return fmt.Errorf("dispatch/mongo: link artifact: %w", err) + } + + return nil +} + +// ListLinks returns every link belonging to the given owner. +func (s *Store) ListLinks(ctx context.Context, owner artifact.OwnerRef) ([]*artifact.Link, error) { + filter := bson.M{"owner_kind": string(owner.Kind), "owner_id": owner.ID} + sort := options.Find().SetSort(bson.D{{Key: "name", Value: 1}, {Key: "attempt", Value: 1}}) + + return s.findLinks(ctx, filter, sort) +} + +func (s *Store) findLinks( + ctx context.Context, + filter bson.M, + opts ...options.Lister[options.FindOptions], +) ([]*artifact.Link, error) { + cur, err := s.mdb.Collection(colArtifactLinks).Find(ctx, filter, opts...) + if err != nil { + return nil, fmt.Errorf("dispatch/mongo: find links: %w", err) + } + + var models []artifactLinkModel + if err := cur.All(ctx, &models); err != nil { + return nil, fmt.Errorf("dispatch/mongo: decode links: %w", err) + } + + out := make([]*artifact.Link, 0, len(models)) + + for i := range models { + l, cerr := fromLinkModel(&models[i]) + if cerr != nil { + return nil, cerr + } + + out = append(out, l) + } + + return out, nil +} + +// FindLinkByName returns the highest-attempt link for an owner and name, +// breaking ties by created_at descending — see the artifact.Store doc +// comment for why ties happen and what the tie-break does and does not +// fix. +func (s *Store) FindLinkByName( + ctx context.Context, + owner artifact.OwnerRef, + name string, +) (*artifact.Link, error) { + var m artifactLinkModel + + filter := bson.M{"owner_kind": string(owner.Kind), "owner_id": owner.ID, "name": name} + opt := options.FindOne().SetSort(bson.D{{Key: "attempt", Value: -1}, {Key: "created_at", Value: -1}}) + + if err := s.mdb.Collection(colArtifactLinks).FindOne(ctx, filter, opt).Decode(&m); err != nil { + if isNoDocuments(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf("dispatch/mongo: find link by name: %w", err) + } + + return fromLinkModel(&m) +} + +// ListArtifactsByOwner returns live artifacts linked to an owner. +func (s *Store) ListArtifactsByOwner( + ctx context.Context, + owner artifact.OwnerRef, + role artifact.Role, +) ([]*artifact.Artifact, error) { + filter := bson.M{"owner_kind": string(owner.Kind), "owner_id": owner.ID} + if role != "" { + filter["role"] = string(role) + } + + links, err := s.findLinks(ctx, filter) + if err != nil { + return nil, err + } + + if len(links) == 0 { + return nil, nil + } + + ids := make([]string, 0, len(links)) + seen := make(map[string]bool, len(links)) + + for _, l := range links { + key := l.ArtifactID.String() + if seen[key] { + continue + } + + seen[key] = true + + ids = append(ids, key) + } + + return s.findArtifacts(ctx, bson.M{"_id": bson.M{"$in": ids}, "deleted_at": nil}) +} + +// SweepEphemeral marks eligible ephemeral artifacts as deleted. +// +// Mongo cannot join links to jobs and runs in a single expressive +// statement the way SQL can, so eligibility is computed in two steps: +// candidates are narrowed by the ephemeral guard, then each candidate's +// owners are resolved and checked. The lifecycle guard is applied both +// when selecting candidates and again on the write. +func (s *Store) SweepEphemeral( + ctx context.Context, + opts artifact.SweepOpts, +) ([]*artifact.Artifact, error) { + limit := opts.Limit + if limit <= 0 { + limit = defaultSweepLimit + } + + candidates, err := s.findArtifacts(ctx, ephemeralOnly(), + options.Find().SetSort(bson.D{{Key: "created_at", Value: 1}})) + if err != nil { + return nil, err + } + + nowAt := now() + + var eligible []*artifact.Artifact + + for _, a := range candidates { + if len(eligible) >= limit { + break + } + + links, lerr := s.findLinks(ctx, bson.M{"artifact_id": a.ID.String()}) + if lerr != nil { + return nil, lerr + } + + if len(links) == 0 { + // Orphans are SweepOrphans' business. + continue + } + + terminalAt, ok, terr := s.ownersTerminalAt(ctx, links) + if terr != nil { + return nil, terr + } + + if !ok { + continue + } + + if a.ExpiresAt != nil { + if a.ExpiresAt.After(nowAt) { + continue + } + } else if terminalAt.Add(opts.Retention).After(nowAt) { + continue + } + + eligible = append(eligible, a) + } + + if len(eligible) == 0 || opts.DryRun { + return eligible, nil + } + + return s.markDeleted(ctx, eligible, nowAt) +} + +// ownersTerminalAt reports the latest terminal time across an artifact's +// owners, and whether all of them are terminal. An owner document that no +// longer exists counts as terminal at the link's creation time: its job or +// run was purged, so it cannot still be running. +func (s *Store) ownersTerminalAt(ctx context.Context, links []*artifact.Link) (time.Time, bool, error) { + var latest time.Time + + for _, l := range links { + at, ok, err := s.ownerTerminalAt(ctx, l) + if err != nil { + return time.Time{}, false, err + } + + if !ok { + return time.Time{}, false, nil + } + + if at.After(latest) { + latest = at + } + } + + return latest, true, nil +} + +func (s *Store) ownerTerminalAt(ctx context.Context, l *artifact.Link) (time.Time, bool, error) { + col, ok := ownerCollection(l.OwnerKind) + if !ok { + return l.CreatedAt, true, nil + } + + var doc struct { + State string `bson:"state"` + CompletedAt *time.Time `bson:"completed_at"` + UpdatedAt time.Time `bson:"updated_at"` + } + + err := s.mdb.Collection(col).FindOne(ctx, bson.M{"_id": l.OwnerID}).Decode(&doc) + if err != nil { + if isNoDocuments(err) { + return l.CreatedAt, true, nil + } + + return time.Time{}, false, fmt.Errorf("dispatch/mongo: resolve artifact owner: %w", err) + } + + if !isTerminalOwnerState(doc.State) { + return time.Time{}, false, nil + } + + if doc.CompletedAt != nil { + return *doc.CompletedAt, true, nil + } + + return doc.UpdatedAt, true, nil +} + +func ownerCollection(kind artifact.OwnerKind) (string, bool) { + switch kind { + case artifact.OwnerJob: + return colJobs, true + case artifact.OwnerRun, artifact.OwnerStep: + return colWorkflowRuns, true + default: + return "", false + } +} + +func isTerminalOwnerState(state string) bool { + switch state { + case "completed", "failed", "cancelled": + return true + default: + return false + } +} + +// SweepOrphans marks link-less ephemeral artifacts created before cutoff. +func (s *Store) SweepOrphans( + ctx context.Context, + cutoff time.Time, + limit int, +) ([]*artifact.Artifact, error) { + if limit <= 0 { + limit = defaultSweepLimit + } + + filter := ephemeralOnly() + filter["created_at"] = bson.M{"$lt": cutoff} + + candidates, err := s.findArtifacts(ctx, filter, + options.Find().SetSort(bson.D{{Key: "created_at", Value: 1}})) + if err != nil { + return nil, err + } + + var orphans []*artifact.Artifact + + for _, a := range candidates { + if len(orphans) >= limit { + break + } + + n, cerr := s.mdb.Collection(colArtifactLinks). + CountDocuments(ctx, bson.M{"artifact_id": a.ID.String()}) + if cerr != nil { + return nil, fmt.Errorf("dispatch/mongo: count artifact links: %w", cerr) + } + + if n > 0 { + continue + } + + orphans = append(orphans, a) + } + + if len(orphans) == 0 { + return nil, nil + } + + return s.markDeleted(ctx, orphans, now()) +} + +// markDeleted soft-deletes the given artifacts. The ephemeral guard is +// repeated on the write so the update itself, not only the selection +// above it, refuses to touch a durable artifact. +func (s *Store) markDeleted( + ctx context.Context, + artifacts []*artifact.Artifact, + at time.Time, +) ([]*artifact.Artifact, error) { + ids := make([]string, 0, len(artifacts)) + for _, a := range artifacts { + ids = append(ids, a.ID.String()) + } + + filter := ephemeralOnly() + filter["_id"] = bson.M{"$in": ids} + + _, err := s.mdb.Collection(colArtifacts). + UpdateMany(ctx, filter, bson.M{"$set": bson.M{"deleted_at": at}}) + if err != nil { + return nil, fmt.Errorf("dispatch/mongo: mark artifacts deleted: %w", err) + } + + out := make([]*artifact.Artifact, 0, len(artifacts)) + + for _, a := range artifacts { + clone := a.Clone() + deleted := at + clone.DeletedAt = &deleted + out = append(out, clone) + } + + return out, nil +} + +// ListPurgeable returns soft-deleted artifacts older than grace. +func (s *Store) ListPurgeable( + ctx context.Context, + grace time.Duration, + limit int, +) ([]*artifact.Artifact, error) { + if limit <= 0 { + limit = defaultSweepLimit + } + + filter := bson.M{"deleted_at": bson.M{"$ne": nil, "$lte": now().Add(-grace)}} + + find := options.Find(). + SetSort(bson.D{{Key: "deleted_at", Value: 1}}). + SetLimit(int64(limit)) + + return s.findArtifacts(ctx, filter, find) +} + +// PurgeArtifact hard-deletes an artifact and its links. +func (s *Store) PurgeArtifact(ctx context.Context, artifactID id.ArtifactID) error { + if _, err := s.mdb.Collection(colArtifactLinks). + DeleteMany(ctx, bson.M{"artifact_id": artifactID.String()}); err != nil { + return fmt.Errorf("dispatch/mongo: purge artifact links: %w", err) + } + + if _, err := s.mdb.Collection(colArtifacts). + DeleteOne(ctx, bson.M{"_id": artifactID.String()}); err != nil { + return fmt.Errorf("dispatch/mongo: purge artifact: %w", err) + } + + return nil +} diff --git a/store/mongo/artifact_models.go b/store/mongo/artifact_models.go new file mode 100644 index 0000000..b6cdd3e --- /dev/null +++ b/store/mongo/artifact_models.go @@ -0,0 +1,129 @@ +package mongo + +import ( + "time" + + "github.com/xraph/grove" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// ── Artifact model ──────────────────────────────────────────────── + +type artifactModel struct { + grove.BaseModel `grove:"table:dispatch_artifacts"` + + ID string `grove:"id,pk" bson:"_id"` + Backend string `bson:"backend"` + Bucket string `bson:"bucket"` + Key string `bson:"key"` + Size int64 `bson:"size"` + ContentHash string `bson:"content_hash,omitempty"` + ContentType string `bson:"content_type,omitempty"` + Lifecycle string `bson:"lifecycle"` + ScopeAppID string `bson:"scope_app_id,omitempty"` + ScopeOrgID string `bson:"scope_org_id,omitempty"` + ExpiresAt *time.Time `bson:"expires_at,omitempty"` + CreatedAt time.Time `bson:"created_at"` + DeletedAt *time.Time `bson:"deleted_at,omitempty"` +} + +func toArtifactModel(a *artifact.Artifact) *artifactModel { + return &artifactModel{ + ID: a.ID.String(), + Backend: a.Backend, + Bucket: a.Bucket, + Key: a.Key, + Size: a.Size, + ContentHash: a.ContentHash, + ContentType: a.ContentType, + Lifecycle: string(a.Lifecycle), + ScopeAppID: a.ScopeAppID, + ScopeOrgID: a.ScopeOrgID, + ExpiresAt: a.ExpiresAt, + CreatedAt: a.CreatedAt, + DeletedAt: a.DeletedAt, + } +} + +func fromArtifactModel(m *artifactModel) (*artifact.Artifact, error) { + aid, err := id.ParseArtifactID(m.ID) + if err != nil { + return nil, err + } + + return &artifact.Artifact{ + ID: aid, + Backend: m.Backend, + Bucket: m.Bucket, + Key: m.Key, + Size: m.Size, + ContentHash: m.ContentHash, + ContentType: m.ContentType, + Lifecycle: artifact.Lifecycle(m.Lifecycle), + ScopeAppID: m.ScopeAppID, + ScopeOrgID: m.ScopeOrgID, + ExpiresAt: m.ExpiresAt, + CreatedAt: m.CreatedAt, + DeletedAt: m.DeletedAt, + }, nil +} + +func fromArtifactModels(models []artifactModel) ([]*artifact.Artifact, error) { + out := make([]*artifact.Artifact, 0, len(models)) + + for i := range models { + a, err := fromArtifactModel(&models[i]) + if err != nil { + return nil, err + } + + out = append(out, a) + } + + return out, nil +} + +// ── Artifact link model ─────────────────────────────────────────── + +type artifactLinkModel struct { + grove.BaseModel `grove:"table:dispatch_artifact_links"` + + ArtifactID string `grove:"artifact_id,pk" bson:"artifact_id"` + OwnerKind string `bson:"owner_kind"` + OwnerID string `bson:"owner_id"` + Name string `bson:"name"` + Attempt int `bson:"attempt"` + Role string `bson:"role"` + CreatedAt time.Time `bson:"created_at"` +} + +func toLinkModel(l *artifact.Link) *artifactLinkModel { + return &artifactLinkModel{ + ArtifactID: l.ArtifactID.String(), + OwnerKind: string(l.OwnerKind), + OwnerID: l.OwnerID, + Name: l.Name, + Attempt: l.Attempt, + Role: string(l.Role), + CreatedAt: l.CreatedAt, + } +} + +func fromLinkModel(m *artifactLinkModel) (*artifact.Link, error) { + aid, err := id.ParseArtifactID(m.ArtifactID) + if err != nil { + return nil, err + } + + return &artifact.Link{ + ArtifactID: aid, + OwnerKind: artifact.OwnerKind(m.OwnerKind), + OwnerID: m.OwnerID, + Role: artifact.Role(m.Role), + Name: m.Name, + Attempt: m.Attempt, + CreatedAt: m.CreatedAt, + }, nil +} diff --git a/store/mongo/artifact_test.go b/store/mongo/artifact_test.go new file mode 100644 index 0000000..5e025ff --- /dev/null +++ b/store/mongo/artifact_test.go @@ -0,0 +1,54 @@ +package mongo_test + +import ( + "context" + "testing" + + "go.mongodb.org/mongo-driver/v2/bson" + mongod "go.mongodb.org/mongo-driver/v2/mongo" + "go.mongodb.org/mongo-driver/v2/mongo/options" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" +) + +// TestArtifactStoreConformance runs the shared artifact.Store suite +// against MongoDB. +// +// One container serves every subtest; the artifact collections are +// emptied between them because the suite asserts absolute counts. +// Documents are deleted rather than the collections dropped so the +// indexes created by Migrate — in particular the partial unique index on +// live keys — stay in force for every subtest. +func TestArtifactStoreConformance(t *testing.T) { + ctx := context.Background() + uri := startMongo(t) + store := openStore(t, uri) + + if err := store.Migrate(ctx); err != nil { + t.Fatalf("migrate: %v", err) + } + + client, err := mongod.Connect(options.Client().ApplyURI(uri)) + if err != nil { + t.Fatalf("connect raw mongo client: %v", err) + } + + t.Cleanup(func() { + if derr := client.Disconnect(ctx); derr != nil { + t.Errorf("disconnect raw mongo client: %v", derr) + } + }) + + db := client.Database(testDBName) + + artifacttest.RunStoreSuite(t, func(t *testing.T) artifact.Store { + for _, col := range []string{"dispatch_artifact_links", "dispatch_artifacts"} { + if _, derr := db.Collection(col).DeleteMany(ctx, bson.M{}); derr != nil { + t.Fatalf("clear %s: %v", col, derr) + } + } + + return store + }) +} diff --git a/store/mongo/dequeue.go b/store/mongo/dequeue.go new file mode 100644 index 0000000..373e6a2 --- /dev/null +++ b/store/mongo/dequeue.go @@ -0,0 +1,179 @@ +package mongo + +import ( + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// preferredField is the name the locality flag is computed into before +// the candidate sort. It is prefixed so it can never collide with a +// stored field, and it never leaves the aggregation — the $project at the +// end of the pipeline keeps only _id. +const preferredField = "__dispatch_preferred" + +// dequeueBudgetFields maps each canonical dimension the fit predicate +// compares to the scalar BSON field holding it. +// +// These are exactly the dimensions job.DequeueOpts.Allows loops over, and +// exactly the fields jobModel declares for the purpose. The comparison +// deliberately does NOT read the full-fidelity resource_requests +// subdocument: BSON comparison of a subdocument member would have to +// contend with the member being absent on every job that does not declare +// that dimension, and with the numeric type the driver happened to write. +// The scalars are plain int64s that both write paths always emit. +var dequeueBudgetFields = []struct { + key string + field string +}{ + {resource.CPU, "req_cpu_milli"}, + {resource.Memory, "req_memory_bytes"}, + {resource.Disk, "req_disk_bytes"}, + {resource.GPU, "req_gpu_milli"}, +} + +// dequeueFilter compiles opts into the query that decides WHICH jobs may +// be claimed. It is the BSON expression of job.DequeueOpts.Allows and +// must answer identically for every job. +// +// The same filter is used for the candidate read and, with _id pinned, +// for each claiming FindOneAndUpdate. +func dequeueFilter(opts job.DequeueOpts, t time.Time) bson.M { + filter := bson.M{ + "state": bson.M{"$in": []string{string(job.StatePending), string(job.StateRetrying)}}, + "queue": bson.M{"$in": opts.Queues}, + "run_at": bson.M{"$lte": t}, + } + + // Unbounded opts emit the original query verbatim: a caller that does + // not use the resource model claims everything, including jobs + // declaring custom resources it could not possibly satisfy. Anything + // else strands work the day this option ships. PreferHashes is + // deliberately not consulted here — it orders, it never filters. + if opts.IsUnbounded() { + return filter + } + + if opts.ReservedFor != nil { + filter["_id"] = opts.ReservedFor.String() + } + + conjuncts := make([]bson.M, 0, len(dequeueBudgetFields)+1) + + // An absent budget key is unconstrained, not zero, so only declared + // dimensions produce a comparison. A key present with the value zero + // is a real constraint and still emits one — that is an exhausted + // worker, which must claim nothing that needs the dimension. + // + // The test is requirement <= budget: a job needing exactly the free + // capacity is claimable, or the last slot on every worker is + // permanently unusable. + for _, dim := range dequeueBudgetFields { + budget, declared := opts.Budget[dim.key] + if !declared { + continue + } + + // The null branch is the Mongo-specific half. Range operators are + // type-bracketed: {$lte: 0} matches numbers only, so it rejects a + // document where the field is null AND one where it is absent — + // and absent is exactly the shape of any job written before these + // scalar fields existed, which declares no requirement at all and + // must stay claimable. A plain equality against nil covers both + // of those shapes in one clause, which is why there is no + // $exists test here. + // + // None of the req_* fields are indexed, so this $or costs nothing + // the planner would otherwise have had: the index serves + // queue/state/priority/run_at and every resource clause is a + // residual filter either way. + conjuncts = append(conjuncts, bson.M{"$or": []bson.M{ + {dim.field: bson.M{"$lte": budget}}, + {dim.field: nil}, + }}) + } + + conjuncts = append(conjuncts, customKeyFilter(opts)) + + filter["$and"] = conjuncts + + return filter +} + +// customKeyFilter renders custom-resource containment as a genuine SUBSET +// test. +// +// req_custom_keys holds resource.EncodeCustomKeys' output — the sorted +// required keys wrapped in leading and trailing separators, e.g. +// ",fpga,tpu," — or "" when the job needs none. The obvious formulation, +// a substring match of the stored list inside the offered one, passes +// every single-key case including the prefix collision and then silently +// strands a job needing {fpga,tpu} from a caller offering {fpga,nvme,tpu}, +// because the interleaved key breaks the contiguous run. The job it +// strands is the specialised one that is hardest to place anywhere else. +// +// Mongo can state the real thing instead: split the stored list back into +// keys and ask $setIsSubset. That is exact by construction — no ordering +// assumption, no prefix hazard, no arithmetic on separators. The SQL +// backends need nested REPLACE only because they have no set operator. +// +// $ifNull guards the split: $split raises on a non-string input, so a +// document written before req_custom_keys existed would fail the whole +// query rather than being read as "requires nothing". $filter then drops +// the empty elements the wrapping separators produce. +// +// An empty offer reaches here only for opts that are bounded some other +// way — IsUnbounded already returned above — so it correctly means "this +// worker has no custom resources": $setIsSubset then admits jobs +// requiring none and rejects every other. +func customKeyFilter(opts job.DequeueOpts) bson.M { + required := bson.M{"$filter": bson.M{ + "input": bson.M{"$split": bson.A{ + bson.M{"$ifNull": bson.A{"$req_custom_keys", ""}}, + resource.CustomKeySep, + }}, + "cond": bson.M{"$ne": bson.A{"$$this", ""}}, + }} + + offered := opts.OfferedCustomKeys() + if offered == nil { + offered = []string{} + } + + // $literal, because a bare array in an aggregation expression has its + // elements evaluated — a resource key beginning with "$" would + // otherwise be read as a field path. + return bson.M{"$expr": bson.M{"$setIsSubset": bson.A{required, bson.M{"$literal": offered}}}} +} + +// preferredExpr computes 1 for a job the caller already has staged and 0 +// for every other, so a descending sort on it puts preferred first. +// +// It is a computed 0/1 rather than a sort on primary_input_hash itself, +// which is the trap Mongo sets here: findAndModify and $sort both take +// field paths, so "preferred first" reads as though it could be spelled +// {primary_input_hash: -1}. That sorts by hash VALUE, which has nothing +// to do with whether the caller staged it — any job whose hash happens to +// collate above the staged one outranks it — and it also decides the +// no-signal case by BSON collation order rather than by the contract, +// since null sorts below strings ascending and therefore above them +// descending. +// +// $ifNull is belt-and-braces on top of that: $in already evaluates to +// false for a missing or null field path, so the guard is not what makes +// the null case lose — the computed flag is. It is kept because it makes +// "unknown means not preferred" explicit at the point of decision rather +// than a property of $in a later reader has to re-derive. +func preferredExpr(hashes []string) bson.M { + return bson.M{"$cond": bson.A{ + bson.M{"$in": bson.A{ + bson.M{"$ifNull": bson.A{"$primary_input_hash", ""}}, + bson.M{"$literal": hashes}, + }}, + 1, + 0, + }} +} diff --git a/store/mongo/dequeue_conformance_test.go b/store/mongo/dequeue_conformance_test.go new file mode 100644 index 0000000..757969b --- /dev/null +++ b/store/mongo/dequeue_conformance_test.go @@ -0,0 +1,388 @@ +package mongo_test + +import ( + "context" + "testing" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/store/storetest" +) + +// TestDequeueConformance runs the resource-aware dequeue suite against the +// Mongo store. +// +// One container is stood up and shared by every subtest, which the suite +// documents as safe: each case enqueues onto its own queue and asserts +// only on the jobs it created. A container per subtest would dominate the +// runtime of the whole package. +func TestDequeueConformance(t *testing.T) { + uri := startMongo(t) + shared := openStore(t, uri) + + storetest.RunDequeueSuite(t, func(t *testing.T) job.Store { + t.Helper() + + return shared + }) +} + +// TestDequeueOrdersNullPrimaryInputHashAsUnpreferred covers the one +// document shape the shared suite cannot produce: jobModel. +// PrimaryInputHash is a plain string, so no Go path can write a null or +// leave the key out. +// +// It matters because Mongo would get this backwards by default. A missing +// field reads as BSON null, and null sorts BEFORE strings ascending — +// under the descending locality term an uncoalesced sort on the raw hash +// would rank exactly the documents with NO locality signal above the ones +// the caller has already staged, inverting the optimization. The +// implementation therefore sorts on a computed 0/1 whose $ifNull maps +// both shapes to "not preferred". +// +// Both shapes are written here with the raw driver: an explicit null and +// an absent key. Both carry an earlier RunAt than the staged job, so they +// win any tie the ordering fails to break — if the guard is dropped they +// come back first. +// +// The fourth fixture, "remote", is what stops the case being satisfied by +// accident. Its hash is a real string that sorts ABOVE the staged one +// descending, so a backend that sorted on primary_input_hash itself — the +// obvious way to write "preferred first" with a field-path sort, and the +// one Mongo will happily accept — hands back a job the caller has no +// local copy of, ahead of the one it staged. Without it, a raw-field sort +// produces the correct answer for the wrong reason: BSON null orders +// below every string, so the single staged hash would lead regardless. +func TestDequeueOrdersNullPrimaryInputHashAsUnpreferred(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + rawDB := rawDatabase(t, uri) + ctx := context.Background() + + const ( + queue = "null-hash-order" + local = "blake3:staged-here" + ) + + base := time.Now().UTC().Add(-time.Hour).Truncate(time.Millisecond) + + nullHash := newMongoHashFixture("null-hash", queue, base) + missingHash := newMongoHashFixture("missing-hash", queue, base.Add(time.Minute)) + cached := newMongoHashFixture("cached", queue, base.Add(2*time.Minute)) + cached.PrimaryInputHash = local + remote := newMongoHashFixture("remote", queue, base.Add(3*time.Minute)) + remote.PrimaryInputHash = "zzz:never-staged" + + for _, j := range []*job.Job{nullHash, missingHash, cached, remote} { + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue %s: %v", j.Name, err) + } + } + + col := rawDB.Collection("dispatch_jobs") + + if _, err := col.UpdateByID(ctx, nullHash.ID.String(), + bson.M{"$set": bson.M{"primary_input_hash": nil}}); err != nil { + t.Fatalf("null out primary_input_hash: %v", err) + } + + if _, err := col.UpdateByID(ctx, missingHash.ID.String(), + bson.M{"$unset": bson.M{"primary_input_hash": ""}}); err != nil { + t.Fatalf("unset primary_input_hash: %v", err) + } + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + PreferHashes: []string{local}, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 4 { + t.Fatalf("claimed %d jobs, want 4 — locality must never filter", len(got)) + } + + want := []string{"cached", "null-hash", "missing-hash", "remote"} + for i, name := range want { + if got[i].Name != name { + t.Fatalf("claimed %v, want %v: a null or absent primary_input_hash must sort "+ + "as NOT preferred, never ahead of the job the caller has staged", + mongoJobNames(got), want) + } + } +} + +// TestUndeclaredJobMatchesNullEqualityOnBothWritePaths is the empirical +// proof behind the filter's null branch, and behind the warning carried +// on jobModel.ResourceRequests. +// +// Mongo stores two different shapes for the same logical "declares +// nothing" state, because the two write paths differ: EnqueueJob goes +// through grove's structToMapInsert, which reflects over grove tags and +// never consults the bson tag, so `omitempty` has no effect and the key +// is written PRESENT-and-null; UpdateJob hands the struct to the raw +// driver's ReplaceOne, which honours `omitempty` and drops the key +// ENTIRELY. +// +// The test drives both paths and then asserts, against the real +// documents, that: +// +// - a single plain equality against nil matches BOTH shapes, so one +// clause covers both write paths and no $or is needed; +// - $exists:false matches only the ReplaceOne shape, which is why the +// filter does not use it — a predicate built on $exists:false would +// silently drop every job still on its original inserted document; +// - a type-bracketed range operator matches NEITHER, which is the +// reason the scalar comparisons in dequeueFilter carry a null branch +// at all rather than relying on {$lte: n} alone; +// - the four scalar req_* fields and req_custom_keys are present as +// real values on BOTH paths, which is what makes them, rather than +// the resource_requests subdocument, the safe basis for the fit +// predicate. +func TestUndeclaredJobMatchesNullEqualityOnBothWritePaths(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + rawDB := rawDatabase(t, uri) + ctx := context.Background() + + const queue = "null-shape-proof" + + inserted := newMongoHashFixture("inserted", queue, time.Now().UTC().Add(-time.Hour)) + replaced := newMongoHashFixture("replaced", queue, time.Now().UTC().Add(-time.Hour)) + + for _, j := range []*job.Job{inserted, replaced} { + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue %s: %v", j.Name, err) + } + } + + // Round-trip the second one through the ReplaceOne write path. + stored, err := s.GetJob(ctx, replaced.ID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + + if err = s.UpdateJob(ctx, stored); err != nil { + t.Fatalf("UpdateJob: %v", err) + } + + col := rawDB.Collection("dispatch_jobs") + ids := bson.M{"$in": []string{inserted.ID.String(), replaced.ID.String()}} + + // The two shapes really are different, or the rest proves nothing. + var insertedDoc, replacedDoc bson.M + + if err = col.FindOne(ctx, bson.M{"_id": inserted.ID.String()}).Decode(&insertedDoc); err != nil { + t.Fatalf("read inserted doc: %v", err) + } + + if err = col.FindOne(ctx, bson.M{"_id": replaced.ID.String()}).Decode(&replacedDoc); err != nil { + t.Fatalf("read replaced doc: %v", err) + } + + if v, ok := insertedDoc["resource_requests"]; !ok || v != nil { + t.Fatalf("inserted resource_requests = %#v (present=%t), want present-and-null", v, ok) + } + + if v, ok := replacedDoc["resource_requests"]; ok { + t.Fatalf("replaced resource_requests = %#v, want key absent", v) + } + + // One equality clause, both shapes. + count, err := col.CountDocuments(ctx, bson.M{"_id": ids, "resource_requests": nil}) + if err != nil { + t.Fatalf("count null-equality: %v", err) + } + + if count != 2 { + t.Errorf("{resource_requests: nil} matched %d/2 documents; a single equality against "+ + "null must cover both the present-and-null and the absent shape", count) + } + + // $exists:false sees only the ReplaceOne shape — the trap. + count, err = col.CountDocuments(ctx, + bson.M{"_id": ids, "resource_requests": bson.M{"$exists": false}}) + if err != nil { + t.Fatalf("count $exists:false: %v", err) + } + + if count != 1 { + t.Errorf("{resource_requests: {$exists: false}} matched %d/2 documents, want 1; "+ + "if this ever matches both, the write paths converged", count) + } + + // Range operators are type-bracketed: null is outside the numeric + // bracket, so {$lte: n} matches neither a null nor an absent field. + // This is exactly why dequeueFilter pairs each scalar comparison with + // a null branch. + count, err = col.CountDocuments(ctx, + bson.M{"_id": ids, "resource_requests": bson.M{"$lte": bson.M{}}}) + if err != nil { + t.Fatalf("count type-bracketed range: %v", err) + } + + if count != 0 { + t.Errorf("a range operator matched %d/2 null-or-absent documents, want 0", count) + } + + // And the scalars the fit predicate actually compares are real values + // on both write paths. + for name, doc := range map[string]bson.M{"inserted": insertedDoc, "replaced": replacedDoc} { + for _, field := range []string{ + "req_cpu_milli", "req_memory_bytes", "req_disk_bytes", "req_gpu_milli", + } { + v, ok := doc[field] + if !ok { + t.Errorf("%s document is missing %s; the fit predicate compares it numerically", name, field) + + continue + } + + if got := toInt64(v); got != 0 { + t.Errorf("%s document %s = %v, want 0", name, field, v) + } + } + + if v, ok := doc["req_custom_keys"]; !ok || v != "" { + t.Errorf("%s document req_custom_keys = %#v (present=%t), want present and empty", name, v, ok) + } + } +} + +// TestDequeueClaimsDocumentsWrittenBeforeResourceFieldsExisted is what +// makes the filter's null branches load-bearing. +// +// A job written by a build that predates the req_* fields carries none of +// them, and a job carrying none of them declares no requirement at all — +// it must stay claimable by every worker, exactly like the freshly +// enqueued undeclared job. Mongo does not agree by default: range +// operators are type-bracketed, so {$lte: 0} rejects an absent field, and +// $split raises outright on a null input rather than reading it as the +// empty list. Either would strand every pre-existing job in the +// collection the moment a resource-aware worker started polling — a queue +// that simply stops draining, with nothing in the system reporting why. +// +// The document is degraded with the raw driver because no Go path can +// produce it: toJobModel always populates all five fields. +func TestDequeueClaimsDocumentsWrittenBeforeResourceFieldsExisted(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + rawDB := rawDatabase(t, uri) + ctx := context.Background() + + const queue = "pre-resource-fields" + + legacy := newMongoHashFixture("legacy", queue, time.Now().UTC().Add(-time.Hour)) + if err := s.EnqueueJob(ctx, legacy); err != nil { + t.Fatalf("enqueue: %v", err) + } + + if _, err := rawDB.Collection("dispatch_jobs").UpdateByID(ctx, legacy.ID.String(), + bson.M{"$unset": bson.M{ + "req_cpu_milli": "", + "req_memory_bytes": "", + "req_disk_bytes": "", + "req_gpu_milli": "", + "req_custom_keys": "", + "resource_requests": "", + }}); err != nil { + t.Fatalf("strip resource fields: %v", err) + } + + // An exhausted, resource-aware worker: every dimension present and + // zero, no custom resources offered. + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{ + resource.CPU: 0, + resource.Memory: 0, + resource.Disk: 0, + resource.GPU: 0, + }, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 1 || got[0].Name != "legacy" { + t.Fatalf("claimed %v, want [legacy]: a document with no req_* fields declares "+ + "no requirement and must fit any budget", mongoJobNames(got)) + } +} + +// TestDequeueClaimsCustomKeySupersetJobsUnderRealDocuments is a +// belt-and-braces check on the $setIsSubset containment test against +// documents that also carry canonical dimensions, which is the shape a +// substring formulation gets wrong in production but not in the suite's +// custom-key-only fixtures. +func TestDequeueClaimsCustomKeySupersetJobsUnderRealDocuments(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + ctx := context.Background() + + const queue = "custom-subset-mixed" + + base := time.Now().UTC().Add(-time.Hour) + + multi := newMongoHashFixture("needs-fpga-and-tpu", queue, base) + multi.Resources = resource.Set{ + resource.CPU: 2 * resource.MilliScale, + resource.Memory: storetest.GiB, + "fpga": 1, + "tpu": 1, + } + + prefix := newMongoHashFixture("needs-fpga-large", queue, base.Add(time.Minute)) + prefix.Resources = resource.Set{resource.Memory: storetest.GiB, "fpga-large": 1} + + for _, j := range []*job.Job{multi, prefix} { + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue %s: %v", j.Name, err) + } + } + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.CPU: 4 * resource.MilliScale, resource.Memory: 4 * storetest.GiB}, + CustomKeys: []string{"fpga", "nvme", "tpu"}, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 1 || got[0].Name != "needs-fpga-and-tpu" { + t.Fatalf("claimed %v, want [needs-fpga-and-tpu]: containment is a subset test, "+ + "and \"fpga\" must not match a worker offering only \"fpga-large\"", mongoJobNames(got)) + } +} + +func newMongoHashFixture(name, queue string, runAt time.Time) *job.Job { + return &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: name, + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: runAt, + } +} + +func mongoJobNames(jobs []*job.Job) []string { + out := make([]string, 0, len(jobs)) + for _, j := range jobs { + out = append(out, j.Name) + } + + return out +} diff --git a/store/mongo/dequeue_test.go b/store/mongo/dequeue_test.go index 2aa64d0..22848ce 100644 --- a/store/mongo/dequeue_test.go +++ b/store/mongo/dequeue_test.go @@ -113,7 +113,7 @@ func TestDequeueJobsIdleIssuesNoWriteCommands(t *testing.T) { } for range 5 { - jobs, err := s.DequeueJobs(ctx, []string{"default"}, 4) + jobs, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 4}) if err != nil { t.Fatalf("dequeue: %v", err) } @@ -149,7 +149,7 @@ func TestDequeueJobsClaimsPendingJob(t *testing.T) { t.Fatalf("enqueue: %v", err) } - jobs, err := s.DequeueJobs(ctx, []string{"default"}, 4) + jobs, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 4}) if err != nil { t.Fatalf("dequeue: %v", err) } @@ -163,3 +163,51 @@ func TestDequeueJobsClaimsPendingJob(t *testing.T) { t.Fatalf("expected state running, got %s", jobs[0].State) } } + +// TestDequeueWithNoQueuesClaimsNothing pins the guard that replaced a +// server error. +// +// The driver marshals a nil []string to BSON null, so an empty +// DequeueOpts.Queues used to produce {queue: {$in: null}} and Mongo +// rejected the whole query with "$in needs an array". That is not +// "claims everything" or "claims nothing" — it is a dequeue that ERRORS +// on every poll, which backs the worker pool off exponentially on a +// configuration that merely named no queues, and logs a message about +// BSON to an operator who changed a queue list. +// +// Postgres, SQLite and Redis all return nothing for the same input. This +// matches them. See job.DequeueOpts.Queues for the full per-backend +// split, which is deliberately documented rather than unified: doing +// otherwise would change store/memory's behaviour. +func TestDequeueWithNoQueuesClaimsNothing(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + ctx := context.Background() + + now := time.Now().UTC() + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "no-queue-filter", + Queue: "default", + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: now.Add(-time.Second), + } + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + for _, queues := range [][]string{nil, {}} { + got, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: queues, Limit: 4}) + if err != nil { + t.Fatalf("DequeueJobs with %#v queues returned an error rather than an empty "+ + "result: %v", queues, err) + } + + if len(got) != 0 { + t.Errorf("DequeueJobs with %#v queues claimed %d jobs, want none", queues, len(got)) + } + } +} diff --git a/store/mongo/job.go b/store/mongo/job.go index 41eead0..be062ce 100644 --- a/store/mongo/job.go +++ b/store/mongo/job.go @@ -8,6 +8,7 @@ import ( "time" "go.mongodb.org/mongo-driver/v2/bson" + mongod "go.mongodb.org/mongo-driver/v2/mongo" "go.mongodb.org/mongo-driver/v2/mongo/options" "github.com/xraph/dispatch" @@ -28,73 +29,203 @@ func (s *Store) EnqueueJob(ctx context.Context, j *job.Job) error { return nil } -// DequeueJobs atomically claims up to limit pending jobs from the given -// queues. Each claim is a FindOneAndUpdate (atomic per-doc), but for limit > 1 -// the claims are issued in parallel so wall-clock cost stays close to a single -// round-trip even on a slow connection. -func (s *Store) DequeueJobs(ctx context.Context, queues []string, limit int) ([]*job.Job, error) { - if limit <= 0 { +// maxDequeueRounds bounds the read-then-claim retry below. +// +// A round claims nothing only when every candidate it read was taken by a +// competing claimer in between. Returning an empty batch then would be a +// lie a drain loop believes — storetest's concurrency case has each +// claimer stop on the first empty result — so a contended round is +// retried rather than reported. The bound keeps a pathological loser +// terminating instead of spinning. +const maxDequeueRounds = 8 + +// DequeueJobs atomically claims up to opts.Limit ready jobs from +// opts.Queues that fit opts, sets them to running, and returns them +// ordered by priority descending, then locality-preferred first, then +// RunAt ascending. +// +// Mongo has no statement that can order, limit, and claim many documents +// in one shot, so the claim is composed of two parts: +// +// - one ordered candidate read, which applies the fit predicate, the +// full contract ordering, and the limit — order THEN truncate, +// server-side, over the whole eligible set; +// - one FindOneAndUpdate per candidate, keyed by _id and carrying the +// SAME fit predicate plus the state and run_at guards. findAndModify +// is atomic per document, so two workers racing for one job produce +// exactly one winner; the loser's call matches nothing and it simply +// gets no job, never a second claim of the same one. +// +// The read is also the write gate the previous probe provided: when it +// finds no candidate, not a single write command is sent. A job that does +// not fit is never written to — the predicate is a conjunct of the +// claiming update itself, not a filter over claimed documents. +// +// When opts.Grants() the lease fields are part of that same per-document +// update, never a follow-up write. Per-document atomicity is what makes +// the claim exclusive, and it is also what keeps the grant recoverable: +// ReclaimExpiredLeases matches lease_expires_at {$ne: nil}, so a document +// left running with no expiry is not one at risk of reclamation — it is +// one reclamation can never see. See job.DequeueOpts.LeaseUntil. +func (s *Store) DequeueJobs(ctx context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + // A worker computing zero free slots must claim zero jobs, never the + // whole queue. Matches the SQL backends' LIMIT 0. + if opts.Limit <= 0 { return nil, nil } - // Probe with a cheap indexed read before claiming. findAndModify is a - // write command even when it matches nothing, so without this gate - // idle pollers generate constant write traffic (collection write - // locks, profiler noise, billed write ops). The FindOneAndUpdate - // claims below remain the atomic gatekeepers; losing the race after a - // positive probe just yields an empty batch. - probeCol := s.mdb.Collection(colJobs) - probeFilter := bson.M{ - "state": bson.M{"$in": []string{string(job.StatePending), string(job.StateRetrying)}}, - "queue": bson.M{"$in": queues}, - "run_at": bson.M{"$lte": now()}, - } - probeOpts := options.FindOne().SetProjection(bson.M{"_id": 1}) - probeErr := withRetry(ctx, defaultRetry, func(ctx context.Context) error { - return probeCol.FindOne(ctx, probeFilter, probeOpts).Err() - }) - if probeErr != nil { - if isNoDocuments(probeErr) { + // Ahead of the empty-Queues guard, so an incoherent grant is reported + // rather than swallowed by a return that happens to be silent here. + // A caller that names no queues still deserves to hear that its lease + // has no holder, and the five backends must agree on which inputs are + // errors — see job.DequeueOpts.Validate. + if err := opts.Validate(); err != nil { + return nil, fmt.Errorf("dispatch/mongo: dequeue jobs: %w", err) + } + + // An empty queue list is a guard, not a query. The driver marshals a + // nil []string to BSON null, so {queue: {$in: null}} reaches the + // server and is rejected outright — "$in needs an array" — which + // surfaces as a dequeue ERROR every poll, backing the pool off on a + // configuration that merely names no queues. Postgres, SQLite and + // Redis all claim nothing for the same input; this matches them + // rather than inventing an all-queues scan no other persistent + // backend performs. See job.DequeueOpts.Queues. + if len(opts.Queues) == 0 { + return nil, nil + } + + for range maxDequeueRounds { + t := now() + + ids, err := s.dequeueCandidates(ctx, opts, t) + if err != nil { + return nil, err + } + + if len(ids) == 0 { return nil, nil } - return nil, fmt.Errorf("dispatch/mongo: dequeue probe: %w", probeErr) - } - if limit == 1 { - j, err := s.dequeueOne(ctx, queues, now()) - if err != nil || j == nil { + jobs, err := s.claimCandidates(ctx, opts, ids, t) + if err != nil { return nil, err } - return []*job.Job{j}, nil + + if len(jobs) > 0 { + return jobs, nil + } } - t := now() - results := make([]*job.Job, limit) - errsCh := make(chan error, limit) - var wg sync.WaitGroup + return nil, nil +} + +// dequeueCandidates returns the _ids of the top opts.Limit eligible jobs, +// already in contract order. +// +// Ordering happens BEFORE truncation and on the server, over every +// eligible document — not over an arbitrary slice of them. That is the +// whole reason the candidates are read rather than letting N independent +// FindOneAndUpdates each pick their own document: with a locality term +// they could not, because Mongo's findAndModify sort takes field paths +// only and locality is a computed predicate. +func (s *Store) dequeueCandidates( + ctx context.Context, + opts job.DequeueOpts, + t time.Time, +) ([]string, error) { + pipeline := mongod.Pipeline{ + bson.D{{Key: "$match", Value: dequeueFilter(opts, t)}}, + } + + sortDoc := bson.D{{Key: "priority", Value: -1}} + + // Locality is applied whenever PreferHashes is non-empty, including + // on otherwise-unbounded opts: IsUnbounded governs FILTERING only. + // And it ranks strictly BELOW priority — above it, a steady stream of + // locally staged low-priority work would starve the high-priority job + // the pool exists to run first. + if hashes := opts.PreferredHashes(); len(hashes) > 0 { + pipeline = append(pipeline, bson.D{{Key: "$addFields", Value: bson.M{ + preferredField: preferredExpr(hashes), + }}}) + + sortDoc = append(sortDoc, bson.E{Key: preferredField, Value: -1}) + } + + sortDoc = append(sortDoc, bson.E{Key: "run_at", Value: 1}) + + pipeline = append(pipeline, + // $sort immediately followed by $limit is coalesced into a + // bounded top-k sort, so the locality term — which no index can + // serve — still costs memory proportional to the limit, not to + // the size of the pending queue. + bson.D{{Key: "$sort", Value: sortDoc}}, + bson.D{{Key: "$limit", Value: int64(opts.Limit)}}, + bson.D{{Key: "$project", Value: bson.M{"_id": 1}}}, + ) + + var rows []struct { + ID string `bson:"_id"` + } + + err := withRetry(ctx, defaultRetry, func(ctx context.Context) error { + cursor, aggErr := s.mdb.Collection(colJobs).Aggregate(ctx, pipeline) + if aggErr != nil { + return aggErr + } + defer cursor.Close(ctx) + + rows = rows[:0] + + return cursor.All(ctx, &rows) + }) + if err != nil { + return nil, fmt.Errorf("dispatch/mongo: dequeue candidates: %w", err) + } + + ids := make([]string, 0, len(rows)) + for _, r := range rows { + ids = append(ids, r.ID) + } + + return ids, nil +} - // Once one worker hits ErrNoDocuments the queue is empty; cancel the rest - // to avoid pointless round-trips against an empty queue. - cctx, cancel := context.WithCancel(ctx) - defer cancel() +// claimCandidates claims each candidate in parallel and returns the ones +// it won, still in candidate order. +// +// Order is preserved by writing each result into its candidate's slot and +// compacting afterwards, never by appending in completion order — the +// claims race each other, so completion order is arbitrary. +func (s *Store) claimCandidates( + ctx context.Context, + opts job.DequeueOpts, + ids []string, + t time.Time, +) ([]*job.Job, error) { + results := make([]*job.Job, len(ids)) + errsCh := make(chan error, len(ids)) - for i := 0; i < limit; i++ { - i := i + var wg sync.WaitGroup + + for i, jobID := range ids { wg.Add(1) + go func() { defer wg.Done() - j, err := s.dequeueOne(cctx, queues, t) + + j, err := s.claimOne(ctx, opts, jobID, t) if err != nil { errsCh <- err + return } - if j == nil { - cancel() - return - } + results[i] = j }() } + wg.Wait() close(errsCh) @@ -104,53 +235,77 @@ func (s *Store) DequeueJobs(ctx context.Context, queues []string, limit int) ([] } } - jobs := make([]*job.Job, 0, limit) + jobs := make([]*job.Job, 0, len(ids)) + for _, j := range results { if j != nil { jobs = append(jobs, j) } } + return jobs, nil } -// dequeueOne claims a single job atomically. Returns (nil, nil) when no -// claimable job exists. Wrapped in withRetry so transient network blips -// don't bubble up as dequeue errors. -func (s *Store) dequeueOne(ctx context.Context, queues []string, t time.Time) (*job.Job, error) { - col := s.mdb.Collection(colJobs) - filter := bson.M{ - "state": bson.M{"$in": []string{string(job.StatePending), string(job.StateRetrying)}}, - "queue": bson.M{"$in": queues}, - "run_at": bson.M{"$lte": t}, - } - update := bson.M{ - "$set": bson.M{ - "state": string(job.StateRunning), - "started_at": t, - "updated_at": t, - }, - } - opts := options.FindOneAndUpdate(). - SetReturnDocument(options.After). - SetSort(bson.D{ - {Key: "priority", Value: -1}, - {Key: "run_at", Value: 1}, - }) +// claimOne atomically claims one candidate. It returns (nil, nil) when +// the document no longer matches — claimed by someone else, or no longer +// eligible — which is a lost race, not an error. +// +// The filter is the FULL dequeue filter with _id pinned, not just the +// _id: the fit predicate must be evaluated as part of the claim, so a job +// that does not fit is never written to even if it somehow reached the +// candidate list. +// +// The lease grant, when opts asks for one, is part of THIS update +// document. findAndModify applies the whole document atomically, so the +// winner of the race is leased in the instant it is claimed and there is +// no moment at which the job is running without a lease. +func (s *Store) claimOne( + ctx context.Context, + opts job.DequeueOpts, + jobID string, + t time.Time, +) (*job.Job, error) { + filter := dequeueFilter(opts, t) + filter["_id"] = jobID + + set := bson.M{ + "state": string(job.StateRunning), + "started_at": t, + "updated_at": t, + } + update := bson.M{"$set": set} + + if opts.Grants() { + set["worker_id"] = opts.WorkerID.String() + set["lease_expires_at"] = opts.LeaseUntil.UTC() + // $inc rather than a computed value: the epoch is the fence and + // must advance from whatever the document currently holds, which + // only the document knows. + update["$inc"] = bson.M{"lease_epoch": 1} + } + + updateOpts := options.FindOneAndUpdate().SetReturnDocument(options.After) var m jobModel + err := withRetry(ctx, defaultRetry, func(ctx context.Context) error { - return col.FindOneAndUpdate(ctx, filter, update, opts).Decode(&m) + return s.mdb.Collection(colJobs). + FindOneAndUpdate(ctx, filter, update, updateOpts). + Decode(&m) }) if err != nil { if isNoDocuments(err) { return nil, nil } + return nil, fmt.Errorf("dispatch/mongo: dequeue jobs: %w", err) } + j, convErr := fromJobModel(&m) if convErr != nil { return nil, fmt.Errorf("dispatch/mongo: dequeue convert: %w", convErr) } + return j, nil } diff --git a/store/mongo/lease.go b/store/mongo/lease.go new file mode 100644 index 0000000..4490ff0 --- /dev/null +++ b/store/mongo/lease.go @@ -0,0 +1,269 @@ +package mongo + +import ( + "context" + "fmt" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo/options" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// The compile-time check that this store provides the lease capability +// lives in store.go alongside the other interface assertions. +// +// The grant itself is not in this file: it travels on job.DequeueOpts and +// is applied by claimOne, inside the same FindOneAndUpdate that claims the +// document, so a leased claim carries the fit predicate and the ordering +// like any other. This file holds only what a lease needs afterwards. + +// RenewLease extends the lease only if the caller still holds it. +func (s *Store) RenewLease( + ctx context.Context, + jobID id.JobID, + workerID id.WorkerID, + epoch int, + leaseUntil time.Time, +) error { + t := now() + col := s.mdb.Collection(colJobs) + + filter := bson.M{ + "_id": jobID.String(), + "state": string(job.StateRunning), + "worker_id": workerID.String(), + "lease_epoch": epoch, + } + update := bson.M{"$set": bson.M{ + "lease_expires_at": leaseUntil.UTC(), + "heartbeat_at": t, + "updated_at": t, + }} + + var matched int64 + err := withRetry(ctx, defaultRetry, func(ctx context.Context) error { + r, updErr := col.UpdateOne(ctx, filter, update) + if updErr != nil { + return updErr + } + matched = r.MatchedCount + + return nil + }) + if err != nil { + return fmt.Errorf("dispatch/mongo: renew lease: %w", err) + } + if matched == 0 { + return job.ErrLeaseLost + } + + return nil +} + +// ReclaimExpiredLeases returns expired-lease jobs to pending, fencing +// their previous holders. +// +// Each job is claimed by its own conditional FindOneAndUpdate keyed on the +// epoch it was seen at, so two pools reclaiming concurrently cannot both +// take the same job — the loser's filter no longer matches. +func (s *Store) ReclaimExpiredLeases(ctx context.Context, limit int) ([]*job.Job, error) { + if limit <= 0 { + return nil, nil + } + + t := now() + col := s.mdb.Collection(colJobs) + + // The first branch is the actual rule: a lease was granted and has + // lapsed. The other two are the narrow exception for a running job + // carrying no lease at all, gated on silence rather than on the null + // expiry alone — see job.UnleasedReclaimGrace for why a null expiry + // does not by itself mean the job was abandoned, and why a row with + // neither timestamp is deliberately left alone. + // + // Testing lease_expires_at against null rather than $exists is + // load-bearing: this collection holds both shapes for the same absent + // value, because grove's insert path writes an explicit null while the + // driver's own encoder honors omitempty and drops the key (see the + // comment on jobModel.ResourceRequests). Plain null equality is the one + // test that matches both. + silent := t.Add(-job.UnleasedReclaimGrace) + filter := bson.M{ + "state": string(job.StateRunning), + "$or": bson.A{ + bson.M{"lease_expires_at": bson.M{"$ne": nil, "$lte": t}}, + bson.M{ + "lease_expires_at": nil, + "heartbeat_at": bson.M{"$ne": nil, "$lte": silent}, + }, + bson.M{ + "lease_expires_at": nil, + "heartbeat_at": nil, + "started_at": bson.M{"$ne": nil, "$lte": silent}, + }, + }, + } + update := bson.M{ + "$set": bson.M{ + "state": string(job.StatePending), + "run_at": t, + "worker_id": "", + "started_at": nil, + "heartbeat_at": nil, + "lease_expires_at": nil, + "updated_at": t, + }, + "$inc": bson.M{ + "lease_epoch": 1, + "evict_count": 1, + }, + } + // Mongo sorts null and missing before every date, so the unleased rows + // matched by the exception above are taken first. That is the right + // order (they have been stranded longest) and it cannot starve the + // leased ones, because each claim moves the row to pending and it stops + // matching the filter. + opts := options.FindOneAndUpdate(). + SetReturnDocument(options.After). + SetSort(bson.D{{Key: "lease_expires_at", Value: 1}}) + + jobs := make([]*job.Job, 0, limit) + for len(jobs) < limit { + var m jobModel + err := withRetry(ctx, defaultRetry, func(ctx context.Context) error { + return col.FindOneAndUpdate(ctx, filter, update, opts).Decode(&m) + }) + if err != nil { + if isNoDocuments(err) { + break + } + + return nil, fmt.Errorf("dispatch/mongo: reclaim expired leases: %w", err) + } + + j, convErr := fromJobModel(&m) + if convErr != nil { + return nil, fmt.Errorf("dispatch/mongo: reclaim convert: %w", convErr) + } + jobs = append(jobs, j) + } + + return jobs, nil +} + +// UpdateLeasedJob persists j only while the caller still holds the +// lease. +// +// Unlike UpdateJob's ReplaceOne, this is an UpdateOne with the fence +// predicate IN THE FILTER and a $set of named business fields — no +// read-modify-write needed, and the document is never replaced wholesale. +// worker_id, lease_epoch, lease_expires_at, and heartbeat_at are never +// assigned in $set: j is the caller's stale snapshot, so writing its copy +// of lease_expires_at back would roll the current holder's expiry +// backwards even though the filter's epoch check passed. +// +// resource_requests and resource_limits mirror toJobModel's zero-Set +// handling: a zero Set is $unset rather than $set to an empty document, +// matching the "absent key, not an empty subdocument" contract UpdateJob +// keeps via bson "omitempty" on ReplaceOne. +func (s *Store) UpdateLeasedJob(ctx context.Context, j *job.Job, workerID id.WorkerID, epoch int) error { + m := toJobModel(j) + m.UpdatedAt = now() + + col := s.mdb.Collection(colJobs) + + filter := bson.M{ + "_id": m.ID, + "state": string(job.StateRunning), + "worker_id": workerID.String(), + "lease_epoch": epoch, + } + + set := bson.M{ + "name": m.Name, + "queue": m.Queue, + "payload": m.Payload, + "state": m.State, + "priority": m.Priority, + "max_retries": m.MaxRetries, + "retry_count": m.RetryCount, + "last_error": m.LastError, + "scope_app_id": m.ScopeAppID, + "scope_org_id": m.ScopeOrgID, + "run_at": m.RunAt, + "started_at": m.StartedAt, + "completed_at": m.CompletedAt, + "timeout": m.Timeout, + "created_at": m.CreatedAt, + "updated_at": m.UpdatedAt, + "lease_ttl": m.LeaseTTL, + "evict_count": m.EvictCount, + "req_cpu_milli": m.ReqCPUMilli, + "req_memory_bytes": m.ReqMemoryBytes, + "req_disk_bytes": m.ReqDiskBytes, + "req_gpu_milli": m.ReqGPUMilli, + "req_custom_keys": m.ReqCustomKeys, + "resource_class": m.ResourceClass, + "input_bytes": m.InputBytes, + "primary_input_hash": m.PrimaryInputHash, + } + + unset := bson.M{} + if m.ResourceRequests.IsZero() { + unset["resource_requests"] = "" + } else { + set["resource_requests"] = m.ResourceRequests + } + if m.ResourceLimits.IsZero() { + unset["resource_limits"] = "" + } else { + set["resource_limits"] = m.ResourceLimits + } + + update := bson.M{"$set": set} + if len(unset) > 0 { + update["$unset"] = unset + } + + // No withRetry here, deliberately, unlike RenewLease above (and + // UpdateJob never uses it either). RenewLease's filter is safe to + // retry because its own $set never touches the fields the filter + // tests — state, worker_id, lease_epoch — so re-matching after an + // already-applied write reapplies the same lease_expires_at and + // heartbeat_at values, which is idempotent by construction. This + // filter is not: it requires state:"running" while $set moves state + // to whatever terminal (or retrying) value the caller asked for, so a + // retry after the first attempt actually landed would find zero + // matching documents, fall through to the existence check below, and + // return ErrLeaseLost for a write that already succeeded — worst case + // on sendToDLQ, where the row is genuinely marked failed but the + // runner, believing the lease was lost, returns before ever reaching + // dlqService.Push. Matching UpdateJob's no-retry choice here is what + // keeps this method's error exactly as trustworthy as the write it + // reports on. + r, err := col.UpdateOne(ctx, filter, update) + if err != nil { + return fmt.Errorf("dispatch/mongo: update leased job: %w", err) + } + if r.MatchedCount > 0 { + return nil + } + + // Zero matches means either the fence predicate failed (the lease + // moved on) or the document is gone. Only the latter is + // ErrJobNotFound; the former is ErrLeaseLost, the entire point of + // this method. + count, countErr := col.CountDocuments(ctx, bson.M{"_id": m.ID}) + if countErr != nil { + return fmt.Errorf("dispatch/mongo: update leased job existence check: %w", countErr) + } + if count == 0 { + return dispatch.ErrJobNotFound + } + + return job.ErrLeaseLost +} diff --git a/store/mongo/lease_test.go b/store/mongo/lease_test.go new file mode 100644 index 0000000..0cb4586 --- /dev/null +++ b/store/mongo/lease_test.go @@ -0,0 +1,150 @@ +package mongo_test + +import ( + "context" + "testing" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/storetest" +) + +// TestReclaimAdoptsRunningJobsWithoutLease covers a running job carrying +// no lease at all. Two things produce one. A job already running when a +// fleet upgraded to a lease-aware build has no lease_expires_at, and +// job.Lease.IsExpired reports false for a zero expiry, so reclamation +// skips it forever while dequeue — which claims only pending and retrying +// rows — never looks at it again. A caller claiming through job.Store +// without lease options produces the same row shape at any time, because +// DequeueOpts.Grants() is false when LeaseUntil is zero. +// +// The negative cases carry the safety argument and matter more than the +// positive ones: the second kind of row is perfectly healthy, and evicting +// live work would be worse than the bug being fixed. Silence is the only +// thing separating the two. +// +// Both null shapes are exercised because this collection genuinely holds +// both, for the reason documented at jobModel.ResourceRequests: EnqueueJob +// goes through grove's structToMapInsert and writes an explicit BSON null, +// while UpdateJob hands the struct to the driver's own encoder, which +// honors "omitempty" and drops the key. A filter matching only one of them +// would strand half the collection. +func TestReclaimAdoptsRunningJobsWithoutLease(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + ctx := context.Background() + col := rawDatabase(t, uri).Collection("dispatch_jobs") + + withHeartbeat := func(name string, startedAgo, beatAgo time.Duration) *job.Job { + j := runningJob(name, startedAgo) + beat := time.Now().UTC().Add(-beatAgo) + j.HeartbeatAt = &beat + + return j + } + + // The two cases below need their lease fields set up before they can be + // listed alongside the rest. + live := runningJob("live-lease", time.Minute) + until := time.Now().UTC().Add(10 * time.Minute) + live.LeaseExpiresAt = &until + live.LeaseEpoch = 1 + ageless := runningJob("no-times", 0) + + type leaseCase struct { + j *job.Job + want bool + why string + } + + cases := []leaseCase{ + { + j: withHeartbeat("stale-heartbeat", 30*time.Minute, 20*time.Minute), + want: true, + why: "abandoned by a worker that stopped reporting", + }, + { + // Old claim, current heartbeat: pins that heartbeat_at wins + // over started_at rather than both needing to be fresh. + j: withHeartbeat("fresh-heartbeat", 30*time.Minute, 0), + want: false, + why: "still reporting, so it belongs to a healthy worker", + }, + { + j: runningJob("no-heartbeat-old-start", 20*time.Minute), + want: true, + why: "claimed long ago and never heartbeated: died before its first beat", + }, + { + j: runningJob("no-heartbeat-fresh-start", 0), + want: false, + why: "just claimed; its first heartbeat is not due yet", + }, + {j: live, want: false, why: "holds a lease that has not lapsed"}, + {j: ageless, want: false, why: "no timestamp to establish age from"}, + } + + for _, c := range cases { + if err := s.EnqueueJob(ctx, c.j); err != nil { + t.Fatalf("enqueue %s: %v", c.j.Name, err) + } + } + + // The ABSENT-key shape, on a row that must still be adopted. Enqueue + // wrote an explicit null for every row above; this is the only way to + // produce the other shape, and a $exists-based filter would miss it. + if _, err := col.UpdateOne(ctx, + bson.M{"_id": cases[2].j.ID.String()}, + bson.M{"$unset": bson.M{"lease_expires_at": ""}}, + ); err != nil { + t.Fatalf("unset lease_expires_at: %v", err) + } + if _, err := col.UpdateOne(ctx, + bson.M{"_id": ageless.ID.String()}, + bson.M{"$unset": bson.M{"started_at": "", "heartbeat_at": ""}}, + ); err != nil { + t.Fatalf("unset timestamps: %v", err) + } + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + + for _, c := range cases { + got := storetest.Contains(reclaimed, c.j.ID) + if got == c.want { + continue + } + if c.want { + t.Errorf("%s was not reclaimed but should have been: %s", c.j.Name, c.why) + } else { + t.Errorf("%s was reclaimed but must not be: %s", c.j.Name, c.why) + } + } +} + +func TestLeaseConformance(t *testing.T) { + // One container for the whole suite — startMongo spins a testcontainer + // and doing that eleven times would dominate the runtime. The suite is + // written to tolerate a shared store. + uri := startMongo(t) + + storetest.RunLeaseSuite(t, func(t *testing.T) storetest.LeaseStore { + t.Helper() + + return openStore(t, uri) + }) +} + +func TestDLQConformance(t *testing.T) { + uri := startMongo(t) + + storetest.RunDLQSuite(t, func(t *testing.T) storetest.DLQStore { + t.Helper() + + return openStore(t, uri) + }) +} diff --git a/store/mongo/migrations.go b/store/mongo/migrations.go index 836e9a4..f2a2cd6 100644 --- a/store/mongo/migrations.go +++ b/store/mongo/migrations.go @@ -205,5 +205,121 @@ func init() { return mexec.DropCollection(ctx, (*workerModel)(nil)) }, }, + &migrate.Migration{ + Name: "create_dispatch_artifacts", + Version: "20260811000001", + Up: func(ctx context.Context, exec migrate.Executor) error { + mexec, ok := exec.(*mongomigrate.Executor) + if !ok { + return fmt.Errorf("expected mongomigrate executor, got %T", exec) + } + + if err := mexec.CreateCollection(ctx, (*artifactModel)(nil)); err != nil { + return err + } + + // The key uniqueness index is partial: only live rows + // collide, so a purged key becomes reusable. + err := mexec.CreateIndexes(ctx, colArtifacts, []mongo.IndexModel{ + { + Keys: bson.D{ + {Key: "backend", Value: 1}, + {Key: "bucket", Value: 1}, + {Key: "key", Value: 1}, + }, + Options: options.Index(). + SetUnique(true). + SetPartialFilterExpression(bson.M{"deleted_at": bson.M{"$eq": nil}}), + }, + {Keys: bson.D{{Key: "lifecycle", Value: 1}, {Key: "created_at", Value: 1}}}, + {Keys: bson.D{{Key: "deleted_at", Value: 1}}}, + {Keys: bson.D{{Key: "content_hash", Value: 1}}}, + {Keys: bson.D{{Key: "scope_app_id", Value: 1}, {Key: "scope_org_id", Value: 1}}}, + }) + if err != nil { + return err + } + + if err := mexec.CreateCollection(ctx, (*artifactLinkModel)(nil)); err != nil { + return err + } + + return mexec.CreateIndexes(ctx, colArtifactLinks, []mongo.IndexModel{ + { + Keys: bson.D{ + {Key: "artifact_id", Value: 1}, + {Key: "owner_kind", Value: 1}, + {Key: "owner_id", Value: 1}, + {Key: "name", Value: 1}, + {Key: "attempt", Value: 1}, + }, + Options: options.Index().SetUnique(true), + }, + {Keys: bson.D{{Key: "owner_kind", Value: 1}, {Key: "owner_id", Value: 1}}}, + {Keys: bson.D{{Key: "artifact_id", Value: 1}}}, + }) + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + mexec, ok := exec.(*mongomigrate.Executor) + if !ok { + return fmt.Errorf("expected mongomigrate executor, got %T", exec) + } + + if err := mexec.DropCollection(ctx, (*artifactLinkModel)(nil)); err != nil { + return err + } + + return mexec.DropCollection(ctx, (*artifactModel)(nil)) + }, + }, + &migrate.Migration{ + Name: "add_job_lease_index", + Version: "20260812000001", + Up: func(ctx context.Context, exec migrate.Executor) error { + mexec, ok := exec.(*mongomigrate.Executor) + if !ok { + return fmt.Errorf("expected mongomigrate executor, got %T", exec) + } + + // Mongo is schemaless, so the lease fields need no + // migration — only the index the reclaim scan reads. + return mexec.CreateIndexes(ctx, colJobs, []mongo.IndexModel{ + {Keys: bson.D{{Key: "state", Value: 1}, {Key: "lease_expires_at", Value: 1}}}, + }) + }, + Down: func(_ context.Context, _ migrate.Executor) error { + // Dropping an index is not worth failing a rollback over. + return nil + }, + }, + &migrate.Migration{ + Name: "add_job_resource_index", + Version: "20260812140000", + Up: func(ctx context.Context, exec migrate.Executor) error { + mexec, ok := exec.(*mongomigrate.Executor) + if !ok { + return fmt.Errorf("expected mongomigrate executor, got %T", exec) + } + + // Mongo is schemaless, so the new resource fields need no + // migration -- only the compound index a future dequeue + // predicate reads: the same (queue, priority DESC, run_at + // ASC) ordering as the SQL backends' covering index, plus + // req_memory_bytes so the predicate's numeric memory + // comparison stays index-assisted too. + return mexec.CreateIndexes(ctx, colJobs, []mongo.IndexModel{ + {Keys: bson.D{ + {Key: "queue", Value: 1}, + {Key: "priority", Value: -1}, + {Key: "run_at", Value: 1}, + {Key: "req_memory_bytes", Value: 1}, + }}, + }) + }, + Down: func(_ context.Context, _ migrate.Executor) error { + // Dropping an index is not worth failing a rollback over. + return nil + }, + }, ) } diff --git a/store/mongo/models.go b/store/mongo/models.go index 0a73e7c..8d8353d 100644 --- a/store/mongo/models.go +++ b/store/mongo/models.go @@ -13,6 +13,7 @@ import ( "github.com/xraph/dispatch/event" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" "github.com/xraph/dispatch/workflow" ) @@ -40,10 +41,63 @@ type jobModel struct { Timeout int64 `grove:"timeout,notnull" bson:"timeout"` CreatedAt time.Time `grove:"created_at,notnull" bson:"created_at"` UpdatedAt time.Time `grove:"updated_at,notnull" bson:"updated_at"` + + LeaseEpoch int `grove:"lease_epoch,notnull" bson:"lease_epoch"` + LeaseExpiresAt *time.Time `grove:"lease_expires_at" bson:"lease_expires_at,omitempty"` + LeaseTTL int64 `grove:"lease_ttl,notnull" bson:"lease_ttl"` + EvictCount int `grove:"evict_count,notnull" bson:"evict_count"` + + // The four canonical dimensions get their own fields because a + // later dequeue predicate compares them numerically and must + // behave identically across all five backends; JSON/BSON comparison + // semantics are not portable. They are derived from Resources by + // toJobModel -- the caller never sets them directly. + ReqCPUMilli int64 `grove:"req_cpu_milli,notnull,default:0" bson:"req_cpu_milli"` + ReqMemoryBytes int64 `grove:"req_memory_bytes,notnull,default:0" bson:"req_memory_bytes"` + ReqDiskBytes int64 `grove:"req_disk_bytes,notnull,default:0" bson:"req_disk_bytes"` + ReqGPUMilli int64 `grove:"req_gpu_milli,notnull,default:0" bson:"req_gpu_milli"` + ReqCustomKeys string `grove:"req_custom_keys,notnull,default:''" bson:"req_custom_keys"` + + // ResourceRequests and ResourceLimits are the full-fidelity copy of + // Resources / ResourceLimits, including custom keys the scalar + // fields above do not carry; fromJobModel reads Resources back from + // here, not from the scalars. Unlike the SQL backends this needs no + // resource.EncodeSet/DecodeSet codec wrapper: the BSON driver + // marshals resource.Set (a map[string]int64) as a native subdocument, + // so toJobModel leaves the field at its Go zero value (nil map) for + // a zero Set rather than assigning it -- never an empty subdocument. + // + // The "omitempty" bson tag below does NOT behave the same on both + // write paths, and a future reader relying on only one of them will + // be wrong for the other: + // - EnqueueJob (NewInsert) goes through grove's structToMapInsert, + // which builds the document by reflecting over the grove tags + // and unconditionally sets doc[column] = value -- it never looks + // at the bson tag, so "omitempty" has no effect here. A zero Set + // is written as an explicit BSON null; the key IS present. + // - UpdateJob (ReplaceOne) hands the struct straight to the raw + // driver, whose native bson encoder DOES honor "omitempty": a + // zero Set drops the key entirely; the key is ABSENT. + // Both are "never {}" and both decode back to a nil Set, so reads + // are unaffected. A query written directly against Mongo, though, + // must not test for only one shape: a plain equality test against + // null covers both, $exists:false covers only the ReplaceOne one, + // and a type-bracketed range operator covers neither. All three are + // asserted against real documents by + // TestUndeclaredJobMatchesNullEqualityOnBothWritePaths. + // + // The dequeue fit predicate sidesteps the asymmetry entirely by + // comparing the scalar fields above, which both write paths always + // emit -- see dequeueFilter in dequeue.go. + ResourceRequests resource.Set `grove:"resource_requests" bson:"resource_requests,omitempty"` + ResourceLimits resource.Set `grove:"resource_limits" bson:"resource_limits,omitempty"` + ResourceClass string `grove:"resource_class,notnull,default:''" bson:"resource_class"` + InputBytes int64 `grove:"input_bytes,notnull,default:0" bson:"input_bytes"` + PrimaryInputHash string `grove:"primary_input_hash" bson:"primary_input_hash"` } func toJobModel(j *job.Job) *jobModel { - return &jobModel{ + m := &jobModel{ ID: j.ID.String(), Name: j.Name, Queue: j.Queue, @@ -63,7 +117,34 @@ func toJobModel(j *job.Job) *jobModel { Timeout: j.Timeout.Nanoseconds(), CreatedAt: j.CreatedAt, UpdatedAt: j.UpdatedAt, + + LeaseEpoch: j.LeaseEpoch, + LeaseExpiresAt: j.LeaseExpiresAt, + LeaseTTL: j.LeaseTTL.Nanoseconds(), + EvictCount: j.EvictCount, + + ReqCPUMilli: j.Resources[resource.CPU], + ReqMemoryBytes: j.Resources[resource.Memory], + ReqDiskBytes: j.Resources[resource.Disk], + ReqGPUMilli: j.Resources[resource.GPU], + ReqCustomKeys: resource.EncodeCustomKeys(j.Resources), + ResourceClass: j.ResourceClass, + InputBytes: j.InputBytes, + PrimaryInputHash: j.PrimaryInputHash, + } + + // A zero Set (nil, or every quantity zero) is left as the field's + // Go zero value -- nil -- rather than assigned j.Resources/ + // j.ResourceLimits verbatim, so "omitempty" drops the key instead + // of writing an empty subdocument. + if !j.Resources.IsZero() { + m.ResourceRequests = j.Resources + } + if !j.ResourceLimits.IsZero() { + m.ResourceLimits = j.ResourceLimits } + + return m } func fromJobModel(m *jobModel) (*job.Job, error) { @@ -93,6 +174,17 @@ func fromJobModel(m *jobModel) (*job.Job, error) { CompletedAt: m.CompletedAt, HeartbeatAt: m.HeartbeatAt, Timeout: time.Duration(m.Timeout), + + LeaseEpoch: m.LeaseEpoch, + LeaseExpiresAt: m.LeaseExpiresAt, + LeaseTTL: time.Duration(m.LeaseTTL), + EvictCount: m.EvictCount, + + Resources: m.ResourceRequests, + ResourceLimits: m.ResourceLimits, + ResourceClass: m.ResourceClass, + InputBytes: m.InputBytes, + PrimaryInputHash: m.PrimaryInputHash, } if m.WorkerID != "" { @@ -290,6 +382,20 @@ type dlqEntryModel struct { FailedAt time.Time `grove:"failed_at,notnull" bson:"failed_at"` ReplayedAt *time.Time `grove:"replayed_at" bson:"replayed_at,omitempty"` CreatedAt time.Time `grove:"created_at,notnull" bson:"created_at"` + + // Carried so Replay can rebuild a job that behaves like the failed + // one; see the dlq.Entry doc. Mongo is schemaless, so these need no + // migration, and resource.Set marshals as a native BSON subdocument + // exactly as it does on jobModel. + Priority int `grove:"priority,notnull,default:0" bson:"priority"` + Timeout int64 `grove:"timeout,notnull,default:0" bson:"timeout"` + LeaseTTL int64 `grove:"lease_ttl,notnull,default:0" bson:"lease_ttl"` + ArtifactBindings []byte `grove:"artifact_bindings" bson:"artifact_bindings,omitempty"` + Resources resource.Set `grove:"resources" bson:"resources,omitempty"` + ResourceLimits resource.Set `grove:"resource_limits" bson:"resource_limits,omitempty"` + ResourceClass string `grove:"resource_class" bson:"resource_class"` + InputBytes int64 `grove:"input_bytes,notnull,default:0" bson:"input_bytes"` + PrimaryInputHash string `grove:"primary_input_hash" bson:"primary_input_hash"` } func toDLQModel(e *dlq.Entry) *dlqEntryModel { @@ -307,6 +413,16 @@ func toDLQModel(e *dlq.Entry) *dlqEntryModel { FailedAt: e.FailedAt, ReplayedAt: e.ReplayedAt, CreatedAt: e.CreatedAt, + + Priority: e.Priority, + Timeout: int64(e.Timeout), + LeaseTTL: int64(e.LeaseTTL), + ArtifactBindings: e.ArtifactBindings, + Resources: e.Resources, + ResourceLimits: e.ResourceLimits, + ResourceClass: e.ResourceClass, + InputBytes: e.InputBytes, + PrimaryInputHash: e.PrimaryInputHash, } } @@ -335,6 +451,16 @@ func fromDLQModel(m *dlqEntryModel) (*dlq.Entry, error) { FailedAt: m.FailedAt, ReplayedAt: m.ReplayedAt, CreatedAt: m.CreatedAt, + + Priority: m.Priority, + Timeout: time.Duration(m.Timeout), + LeaseTTL: time.Duration(m.LeaseTTL), + ArtifactBindings: m.ArtifactBindings, + Resources: m.Resources, + ResourceLimits: m.ResourceLimits, + ResourceClass: m.ResourceClass, + InputBytes: m.InputBytes, + PrimaryInputHash: m.PrimaryInputHash, }, nil } diff --git a/store/mongo/resource_test.go b/store/mongo/resource_test.go new file mode 100644 index 0000000..37985fe --- /dev/null +++ b/store/mongo/resource_test.go @@ -0,0 +1,234 @@ +package mongo_test + +import ( + "context" + "testing" + + "go.mongodb.org/mongo-driver/v2/bson" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// TestMongoRoundTripsResources mirrors store/postgres's and store/sqlite's +// resource round-trip test: a job with cpu/memory/a custom key/limits/ +// class/input signal must come back identical, and the custom key +// specifically must survive since it only lives in the full-fidelity +// resource_requests subdocument, not the scalar fields. +func TestMongoRoundTripsResources(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "tessellate.model", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + Resources: resource.Set{ + resource.CPU: 4000, resource.Memory: 16 << 30, "fpga": 2, + }, + ResourceLimits: resource.Set{resource.Memory: 16 << 30}, + ResourceClass: "heavy", + InputBytes: 4 << 30, + PrimaryInputHash: "blake3:9f2a", + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + + if got.Resources[resource.CPU] != 4000 { + t.Errorf("cpu = %d, want 4000", got.Resources[resource.CPU]) + } + if got.Resources["fpga"] != 2 { + t.Errorf("custom key lost in round-trip: %v", got.Resources) + } + if got.ResourceLimits[resource.Memory] != 16<<30 { + t.Errorf("limits = %v", got.ResourceLimits) + } + if got.ResourceClass != "heavy" { + t.Errorf("class = %q, want heavy", got.ResourceClass) + } + if got.InputBytes != 4<<30 || got.PrimaryInputHash != "blake3:9f2a" { + t.Errorf("input signal lost: bytes=%d hash=%q", + got.InputBytes, got.PrimaryInputHash) + } +} + +// TestMongoJobWithNoResourcesRoundTrips pins the backward-compatibility +// contract: a job enqueued without any resource declaration must come back +// with a zero Set, not a Set containing zero-valued canonical keys. +func TestMongoJobWithNoResourcesRoundTrips(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "notify.user", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if !got.Resources.IsZero() { + t.Errorf("Resources = %v, want zero", got.Resources) + } +} + +// TestMongoUndeclaredJobHasNoResourceSubdocument closes the gap a +// Resources.IsZero() assertion cannot: decoding an absent +// resource_requests key and decoding an empty {} subdocument both yield +// IsZero() == true, so that assertion alone does not prove the stored +// value is genuinely null rather than an empty subdocument. This test +// reads the raw BSON document straight from the driver, bypassing +// jobModel entirely, and asserts resource_requests/resource_limits are +// BSON null -- never an empty subdocument -- which is Mongo's exact +// analogue of the SQL backends' NULL column. +// +// It asserts null rather than "key absent", which is a deliberate, +// empirically-verified choice: EnqueueJob writes through grove's +// mongodriver.NewInsert, whose structToMapInsert builds the document by +// reflecting over every grove-tagged field and always assigning +// doc[column] = value.Interface() -- it has no concept of the field's +// bson struct tag, so "omitempty" on ResourceRequests/ResourceLimits +// has no effect on this path and a nil resource.Set is written as an +// explicit BSON null, not omitted. (The bson tag DOES take effect on +// UpdateJob, which calls the raw driver's ReplaceOne(ctx, filter, m) +// directly -- there a nil Set genuinely drops the key. The two write +// paths are asymmetric in this one respect; both leave the field +// non-existent-or-null, never "{}", so IsZero() on read is unaffected +// either way.) +func TestMongoUndeclaredJobHasNoResourceSubdocument(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + rawDB := rawDatabase(t, uri) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "notify.user", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + var raw bson.M + if err := rawDB.Collection("dispatch_jobs"). + FindOne(ctx, bson.M{"_id": j.ID.String()}). + Decode(&raw); err != nil { + t.Fatalf("raw FindOne() error = %v", err) + } + + reqVal, reqOK := raw["resource_requests"] + if !reqOK { + t.Error("resource_requests key missing from raw document, want present-and-null") + } else if reqVal != nil { + t.Errorf("resource_requests = %#v (%T), want BSON null, not an empty subdocument", reqVal, reqVal) + } + + limitsVal, limitsOK := raw["resource_limits"] + if !limitsOK { + t.Error("resource_limits key missing from raw document, want present-and-null") + } else if limitsVal != nil { + t.Errorf("resource_limits = %#v (%T), want BSON null, not an empty subdocument", limitsVal, limitsVal) + } + + // The scalar/class/hash fields are always-present columns in the + // SQL backends (NOT NULL DEFAULT), so their Mongo analogues should + // still be written -- only the full-fidelity subdocuments are + // null-when-zero. + if v, ok := raw["req_cpu_milli"]; !ok || toInt64(v) != 0 { + t.Errorf("req_cpu_milli = %v, want present and 0", raw["req_cpu_milli"]) + } + if v, ok := raw["req_custom_keys"]; !ok || v != "" { + t.Errorf("req_custom_keys = %v, want present and empty", raw["req_custom_keys"]) + } +} + +// TestMongoUpdateJobDropsResourceKeyEntirely documents the asymmetry +// called out above: UpdateJob's write path (ReplaceOne with the raw +// driver, not grove's insert helper) DOES honor the bson "omitempty" +// tag, so replacing a job down to a zero Set drops the key from the +// document entirely rather than leaving it null. Both representations +// -- absent key, or present-and-null -- decode back to a nil Set, so +// this is a proof of the write path's actual behavior, not a +// requirement the contract imposes. +func TestMongoUpdateJobDropsResourceKeyEntirely(t *testing.T) { + uri := startMongo(t) + s := openStore(t, uri) + rawDB := rawDatabase(t, uri) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "notify.user", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + Resources: resource.Set{resource.CPU: 1000}, + } + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + j.Resources = nil + if err := s.UpdateJob(ctx, j); err != nil { + t.Fatalf("UpdateJob() error = %v", err) + } + + var raw bson.M + if err := rawDB.Collection("dispatch_jobs"). + FindOne(ctx, bson.M{"_id": j.ID.String()}). + Decode(&raw); err != nil { + t.Fatalf("raw FindOne() error = %v", err) + } + + if v, ok := raw["resource_requests"]; ok { + t.Errorf("resource_requests = %#v, want key entirely absent after ReplaceOne", v) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if !got.Resources.IsZero() { + t.Errorf("Resources = %v, want zero after clearing", got.Resources) + } +} + +func toInt64(v any) int64 { + switch n := v.(type) { + case int32: + return int64(n) + case int64: + return n + default: + return -1 + } +} diff --git a/store/mongo/store.go b/store/mongo/store.go index 2851f16..7a7c67b 100644 --- a/store/mongo/store.go +++ b/store/mongo/store.go @@ -16,6 +16,7 @@ import ( "github.com/xraph/grove" "github.com/xraph/grove/drivers/mongodriver" + "github.com/xraph/dispatch/artifact" "github.com/xraph/dispatch/cluster" "github.com/xraph/dispatch/cron" "github.com/xraph/dispatch/dlq" @@ -26,13 +27,15 @@ import ( // Collection name constants. const ( - colJobs = "dispatch_jobs" - colWorkflowRuns = "dispatch_workflow_runs" - colCheckpoints = "dispatch_checkpoints" - colCronEntries = "dispatch_cron_entries" - colDLQ = "dispatch_dlq" - colEvents = "dispatch_events" - colWorkers = "dispatch_workers" + colJobs = "dispatch_jobs" + colWorkflowRuns = "dispatch_workflow_runs" + colCheckpoints = "dispatch_checkpoints" + colCronEntries = "dispatch_cron_entries" + colDLQ = "dispatch_dlq" + colEvents = "dispatch_events" + colWorkers = "dispatch_workers" + colArtifacts = "dispatch_artifacts" + colArtifactLinks = "dispatch_artifact_links" ) // Ensure Store implements all subsystem interfaces at compile time. @@ -43,6 +46,8 @@ var ( _ dlq.Store = (*Store)(nil) _ event.Store = (*Store)(nil) _ cluster.Store = (*Store)(nil) + _ artifact.Store = (*Store)(nil) + _ job.LeaseStore = (*Store)(nil) ) // Store is a grove ORM implementation of store.Store using MongoDB driver. @@ -86,6 +91,16 @@ func (s *Store) DB() *grove.DB { // // CreateMany is itself idempotent — mongo silently no-ops indexes that already // exist with matching specs — so this is safe to call on every boot. +// +// Note this is the whole of the mongo backend's migration path: the grove +// migration group in migrations.go is not run from anywhere, so anything +// that must happen on upgrade belongs here rather than there. Jobs left +// running by a pre-lease build are deliberately NOT handled here. A +// one-shot backfill cannot see a job an old pod claims after the migration +// has already run, and it would evict jobs those pods are still running, +// because it seeds an expiry in the past that their heartbeats do not know +// to push. ReclaimExpiredLeases adopts them instead, continuously and +// gated on silence — see job.UnleasedReclaimGrace. func (s *Store) Migrate(ctx context.Context) error { indexes := migrationIndexes() @@ -137,6 +152,44 @@ func isDuplicateKey(err error) bool { // migrationIndexes returns the index definitions for all dispatch collections. func migrationIndexes() map[string][]mongod.IndexModel { return map[string][]mongod.IndexModel{ + colArtifacts: { + // Partial unique index on the storage coordinates: only live + // rows collide, so a purged key becomes reusable. + { + Keys: bson.D{ + {Key: "backend", Value: 1}, + {Key: "bucket", Value: 1}, + {Key: "key", Value: 1}, + }, + Options: options.Index(). + SetName("dispatch_artifacts_unique_live_key"). + SetUnique(true). + SetPartialFilterExpression(bson.M{"deleted_at": bson.M{"$eq": nil}}), + }, + {Keys: bson.D{{Key: "lifecycle", Value: 1}, {Key: "created_at", Value: 1}}}, + {Keys: bson.D{{Key: "deleted_at", Value: 1}}}, + {Keys: bson.D{{Key: "content_hash", Value: 1}}}, + {Keys: bson.D{ + {Key: "scope_app_id", Value: 1}, + {Key: "scope_org_id", Value: 1}, + }}, + }, + colArtifactLinks: { + { + Keys: bson.D{ + {Key: "artifact_id", Value: 1}, + {Key: "owner_kind", Value: 1}, + {Key: "owner_id", Value: 1}, + {Key: "name", Value: 1}, + {Key: "attempt", Value: 1}, + }, + Options: options.Index(). + SetName("dispatch_artifact_links_unique"). + SetUnique(true), + }, + {Keys: bson.D{{Key: "owner_kind", Value: 1}, {Key: "owner_id", Value: 1}}}, + {Keys: bson.D{{Key: "artifact_id", Value: 1}}}, + }, colJobs: { // Dequeue index: queue + state + priority + run_at. {Keys: bson.D{ @@ -157,6 +210,8 @@ func migrationIndexes() map[string][]mongod.IndexModel { {Key: "state", Value: 1}, {Key: "heartbeat_at", Value: 1}, }}, + // Lease index for the expired-lease reclaim scan. + {Keys: bson.D{{Key: "state", Value: 1}, {Key: "lease_expires_at", Value: 1}}}, }, colWorkflowRuns: { {Keys: bson.D{{Key: "state", Value: 1}}}, diff --git a/store/postgres/artifact.go b/store/postgres/artifact.go new file mode 100644 index 0000000..b28ecac --- /dev/null +++ b/store/postgres/artifact.go @@ -0,0 +1,463 @@ +package postgres + +import ( + "context" + "database/sql" + "errors" + "fmt" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// CreateArtifact inserts an artifact and, when link is non-nil, its first +// link in a single transaction. +func (s *Store) CreateArtifact(ctx context.Context, a *artifact.Artifact, link *artifact.Link) error { + if link == nil { + _, err := s.pgdb.NewInsert(toArtifactModel(a)).Exec(ctx) + if err != nil { + if isDuplicateKey(err) { + return artifact.ErrExists + } + + return fmt.Errorf(errPrefix+"create artifact: %w", err) + } + + return nil + } + + tx, err := s.pgdb.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf(errPrefix+"create artifact: begin: %w", err) + } + + defer func() { + if rerr := tx.Rollback(); rerr != nil && !errors.Is(rerr, sql.ErrTxDone) { + s.logger.Warn(errPrefix+"artifact tx rollback", log.String("error", rerr.Error())) + } + }() + + if _, err := tx.Exec(ctx, insertArtifactSQL, artifactInsertArgs(a)...); err != nil { + if isDuplicateKey(err) { + return artifact.ErrExists + } + + return fmt.Errorf(errPrefix+"create artifact: %w", err) + } + + if _, err := tx.Exec(ctx, insertLinkSQL, linkInsertArgs(link)...); err != nil { + return fmt.Errorf(errPrefix+"create artifact link: %w", err) + } + + if err := tx.Commit(); err != nil { + return fmt.Errorf(errPrefix+"create artifact: commit: %w", err) + } + + return nil +} + +const insertArtifactSQL = ` + INSERT INTO dispatch_artifacts + (id, backend, bucket, key, size, content_hash, content_type, + lifecycle, scope_app_id, scope_org_id, expires_at, created_at, deleted_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)` + +func artifactInsertArgs(a *artifact.Artifact) []any { + return []any{ + a.ID.String(), a.Backend, a.Bucket, a.Key, a.Size, + nullString(a.ContentHash), nullString(a.ContentType), + string(a.Lifecycle), nullString(a.ScopeAppID), nullString(a.ScopeOrgID), + a.ExpiresAt, a.CreatedAt, a.DeletedAt, + } +} + +const insertLinkSQL = ` + INSERT INTO dispatch_artifact_links + (artifact_id, owner_kind, owner_id, name, attempt, role, created_at) + VALUES ($1, $2, $3, $4, $5, $6, $7) + ON CONFLICT (artifact_id, owner_kind, owner_id, name, attempt) DO NOTHING` + +func linkInsertArgs(l *artifact.Link) []any { + return []any{ + l.ArtifactID.String(), string(l.OwnerKind), l.OwnerID, + l.Name, l.Attempt, string(l.Role), l.CreatedAt, + } +} + +// nullString maps the empty string to SQL NULL so nullable text columns +// stay NULL rather than storing an empty value. +func nullString(s string) any { + if s == "" { + return nil + } + + return s +} + +// GetArtifact retrieves a live artifact by ID. +func (s *Store) GetArtifact(ctx context.Context, artifactID id.ArtifactID) (*artifact.Artifact, error) { + var m artifactModel + + err := s.pgdb.NewSelect(&m). + Where("id = ?", artifactID.String()). + Where("deleted_at IS NULL"). + Scan(ctx) + if err != nil { + if isNoRows(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf(errPrefix+"get artifact: %w", err) + } + + return fromArtifactModel(&m) +} + +// FindArtifactByKey retrieves a live artifact by its storage coordinates. +func (s *Store) FindArtifactByKey(ctx context.Context, backend, bucket, key string) (*artifact.Artifact, error) { + var m artifactModel + + err := s.pgdb.NewSelect(&m). + Where("backend = ?", backend). + Where("bucket = ?", bucket). + Where("key = ?", key). + Where("deleted_at IS NULL"). + Scan(ctx) + if err != nil { + if isNoRows(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf(errPrefix+"find artifact by key: %w", err) + } + + return fromArtifactModel(&m) +} + +// UpdateArtifact persists size, hash, content type, and expiry. Lifecycle, +// created_at, and deleted_at are deliberately not updatable here. +func (s *Store) UpdateArtifact(ctx context.Context, a *artifact.Artifact) error { + res, err := s.pgdb.NewRaw(` + UPDATE dispatch_artifacts + SET size = $1, content_hash = $2, content_type = $3, expires_at = $4 + WHERE id = $5`, + a.Size, nullString(a.ContentHash), nullString(a.ContentType), a.ExpiresAt, a.ID.String(), + ).Exec(ctx) + if err != nil { + return fmt.Errorf(errPrefix+"update artifact: %w", err) + } + + n, err := res.RowsAffected() + if err == nil && n == 0 { + return artifact.ErrNotFound + } + + return nil +} + +// ListArtifacts returns artifacts matching the given options, newest first. +func (s *Store) ListArtifacts(ctx context.Context, opts artifact.ListOpts) ([]*artifact.Artifact, error) { + var models []artifactModel + + q := s.pgdb.NewSelect(&models) + + if !opts.IncludeDeleted { + q = q.Where("deleted_at IS NULL") + } + + if opts.Lifecycle != "" { + q = q.Where("lifecycle = ?", string(opts.Lifecycle)) + } + + if opts.ScopeAppID != "" { + q = q.Where("scope_app_id = ?", opts.ScopeAppID) + } + + if opts.ScopeOrgID != "" { + q = q.Where("scope_org_id = ?", opts.ScopeOrgID) + } + + q = q.OrderExpr("created_at DESC, id ASC") + + if opts.Limit > 0 { + q = q.Limit(opts.Limit) + } + + if opts.Offset > 0 { + q = q.Offset(opts.Offset) + } + + if err := q.Scan(ctx); err != nil { + return nil, fmt.Errorf(errPrefix+"list artifacts: %w", err) + } + + return fromArtifactModels(models) +} + +// LinkArtifact records that an owner references an artifact, idempotently. +func (s *Store) LinkArtifact(ctx context.Context, link *artifact.Link) error { + _, err := s.pgdb.Exec(ctx, insertLinkSQL, linkInsertArgs(link)...) + if err != nil { + return fmt.Errorf(errPrefix+"link artifact: %w", err) + } + + return nil +} + +// ListLinks returns every link belonging to the given owner. +func (s *Store) ListLinks(ctx context.Context, owner artifact.OwnerRef) ([]*artifact.Link, error) { + var models []artifactLinkModel + + err := s.pgdb.NewSelect(&models). + Where("owner_kind = ?", string(owner.Kind)). + Where("owner_id = ?", owner.ID). + OrderExpr("name ASC, attempt ASC"). + Scan(ctx) + if err != nil { + return nil, fmt.Errorf(errPrefix+"list links: %w", err) + } + + out := make([]*artifact.Link, 0, len(models)) + + for i := range models { + l, cerr := fromLinkModel(&models[i]) + if cerr != nil { + return nil, cerr + } + + out = append(out, l) + } + + return out, nil +} + +// FindLinkByName returns the highest-attempt link for an owner and name, +// breaking ties by created_at descending — see the artifact.Store doc +// comment for why ties happen and what the tie-break does and does not +// fix. +func (s *Store) FindLinkByName( + ctx context.Context, + owner artifact.OwnerRef, + name string, +) (*artifact.Link, error) { + var m artifactLinkModel + + err := s.pgdb.NewSelect(&m). + Where("owner_kind = ?", string(owner.Kind)). + Where("owner_id = ?", owner.ID). + Where("name = ?", name). + OrderExpr("attempt DESC, created_at DESC"). + Limit(1). + Scan(ctx) + if err != nil { + if isNoRows(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf(errPrefix+"find link by name: %w", err) + } + + return fromLinkModel(&m) +} + +// ListArtifactsByOwner returns live artifacts linked to an owner. +func (s *Store) ListArtifactsByOwner( + ctx context.Context, + owner artifact.OwnerRef, + role artifact.Role, +) ([]*artifact.Artifact, error) { + var models []artifactModel + + query := ` + SELECT DISTINCT a.* FROM dispatch_artifacts a + JOIN dispatch_artifact_links l ON l.artifact_id = a.id + WHERE l.owner_kind = $1 AND l.owner_id = $2 AND a.deleted_at IS NULL` + + args := []any{string(owner.Kind), owner.ID} + + if role != "" { + query += ` AND l.role = $3` + + args = append(args, string(role)) + } + + if err := s.pgdb.NewRaw(query, args...).Scan(ctx, &models); err != nil { + return nil, fmt.Errorf(errPrefix+"list artifacts by owner: %w", err) + } + + return fromArtifactModels(models) +} + +// terminalJobStates and terminalRunStates are the states after which an +// owner can no longer touch its artifacts. +const ( + terminalJobStatesSQL = `('completed', 'failed', 'cancelled')` + terminalRunStatesSQL = `('completed', 'failed', 'cancelled')` +) + +// eligibleEphemeralSQL selects ephemeral artifacts whose every linked +// owner is terminal and whose retention window has elapsed. +// +// The lifecycle predicate is a literal. It is never bound from a +// parameter, so no caller can widen this statement to reach a durable +// artifact. +// +// An owner row that no longer exists counts as terminal at the link's +// creation time: its job or run was purged, so it cannot still be running. +const eligibleEphemeralSQL = ` + SELECT a.id + FROM dispatch_artifacts a + JOIN dispatch_artifact_links l ON l.artifact_id = a.id + LEFT JOIN dispatch_jobs j + ON l.owner_kind = 'job' AND j.id = l.owner_id + LEFT JOIN dispatch_workflow_runs r + ON l.owner_kind IN ('run', 'step') AND r.id = l.owner_id + WHERE a.lifecycle = 'ephemeral' + AND a.deleted_at IS NULL + GROUP BY a.id, a.expires_at + HAVING bool_and( + CASE + WHEN l.owner_kind = 'job' + THEN j.id IS NULL OR j.state IN ` + terminalJobStatesSQL + ` + WHEN l.owner_kind IN ('run', 'step') + THEN r.id IS NULL OR r.state IN ` + terminalRunStatesSQL + ` + ELSE TRUE + END + ) + AND ( + CASE + WHEN a.expires_at IS NOT NULL THEN a.expires_at <= NOW() + ELSE MAX( + COALESCE(j.completed_at, j.updated_at, r.completed_at, r.updated_at, l.created_at) + ) + make_interval(secs => $1::double precision) <= NOW() + END + )` + +// SweepEphemeral marks eligible ephemeral artifacts as deleted. +func (s *Store) SweepEphemeral( + ctx context.Context, + opts artifact.SweepOpts, +) ([]*artifact.Artifact, error) { + limit := opts.Limit + if limit <= 0 { + limit = defaultSweepLimit + } + + selectSQL := eligibleEphemeralSQL + ` + LIMIT $2` + + if opts.DryRun { + var models []artifactModel + + query := ` + SELECT * FROM dispatch_artifacts + WHERE id IN (` + selectSQL + `) + ORDER BY created_at ASC` + + if err := s.pgdb.NewRaw(query, opts.Retention.Seconds(), limit).Scan(ctx, &models); err != nil { + return nil, fmt.Errorf(errPrefix+"sweep ephemeral (dry run): %w", err) + } + + return fromArtifactModels(models) + } + + var models []artifactModel + + query := ` + UPDATE dispatch_artifacts + SET deleted_at = NOW() + WHERE lifecycle = 'ephemeral' + AND deleted_at IS NULL + AND id IN (` + selectSQL + `) + RETURNING *` + + if err := s.pgdb.NewRaw(query, opts.Retention.Seconds(), limit).Scan(ctx, &models); err != nil { + return nil, fmt.Errorf(errPrefix+"sweep ephemeral: %w", err) + } + + return fromArtifactModels(models) +} + +// defaultSweepLimit bounds an unbounded sweep so a single pass can never +// lock an unbounded number of rows. +const defaultSweepLimit = 1000 + +// SweepOrphans marks link-less ephemeral artifacts created before cutoff. +func (s *Store) SweepOrphans( + ctx context.Context, + cutoff time.Time, + limit int, +) ([]*artifact.Artifact, error) { + if limit <= 0 { + limit = defaultSweepLimit + } + + var models []artifactModel + + // lifecycle = 'ephemeral' is a literal here for the same reason as in + // eligibleEphemeralSQL: durable artifacts must be unreachable. + query := ` + UPDATE dispatch_artifacts + SET deleted_at = NOW() + WHERE lifecycle = 'ephemeral' + AND deleted_at IS NULL + AND id IN ( + SELECT a.id FROM dispatch_artifacts a + WHERE a.lifecycle = 'ephemeral' + AND a.deleted_at IS NULL + AND a.created_at < $1 + AND NOT EXISTS ( + SELECT 1 FROM dispatch_artifact_links l WHERE l.artifact_id = a.id + ) + ORDER BY a.created_at ASC + LIMIT $2 + ) + RETURNING *` + + if err := s.pgdb.NewRaw(query, cutoff, limit).Scan(ctx, &models); err != nil { + return nil, fmt.Errorf(errPrefix+"sweep orphans: %w", err) + } + + return fromArtifactModels(models) +} + +// ListPurgeable returns soft-deleted artifacts older than grace. +func (s *Store) ListPurgeable( + ctx context.Context, + grace time.Duration, + limit int, +) ([]*artifact.Artifact, error) { + if limit <= 0 { + limit = defaultSweepLimit + } + + var models []artifactModel + + query := ` + SELECT * FROM dispatch_artifacts + WHERE deleted_at IS NOT NULL + AND deleted_at + make_interval(secs => $1::double precision) <= NOW() + ORDER BY deleted_at ASC + LIMIT $2` + + if err := s.pgdb.NewRaw(query, grace.Seconds(), limit).Scan(ctx, &models); err != nil { + return nil, fmt.Errorf(errPrefix+"list purgeable: %w", err) + } + + return fromArtifactModels(models) +} + +// PurgeArtifact hard-deletes an artifact. Links cascade. +func (s *Store) PurgeArtifact(ctx context.Context, artifactID id.ArtifactID) error { + _, err := s.pgdb.NewRaw( + `DELETE FROM dispatch_artifacts WHERE id = $1`, artifactID.String(), + ).Exec(ctx) + if err != nil { + return fmt.Errorf(errPrefix+"purge artifact: %w", err) + } + + return nil +} diff --git a/store/postgres/artifact_models.go b/store/postgres/artifact_models.go new file mode 100644 index 0000000..0bdf41f --- /dev/null +++ b/store/postgres/artifact_models.go @@ -0,0 +1,117 @@ +package postgres + +import ( + "time" + + "github.com/xraph/grove" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// ── Artifact model ──────────────────────────────────────────────── + +type artifactModel struct { + grove.BaseModel `grove:"table:dispatch_artifacts"` + + ID string `grove:"id,pk"` + Backend string `grove:"backend,notnull"` + Bucket string `grove:"bucket,notnull"` + Key string `grove:"key,notnull"` + Size int64 `grove:"size,notnull,default:0"` + ContentHash string `grove:"content_hash"` + ContentType string `grove:"content_type"` + Lifecycle string `grove:"lifecycle,notnull"` + ScopeAppID string `grove:"scope_app_id"` + ScopeOrgID string `grove:"scope_org_id"` + ExpiresAt *time.Time `grove:"expires_at"` + CreatedAt time.Time `grove:"created_at,notnull,default:current_timestamp"` + DeletedAt *time.Time `grove:"deleted_at"` +} + +func toArtifactModel(a *artifact.Artifact) *artifactModel { + return &artifactModel{ + ID: a.ID.String(), + Backend: a.Backend, + Bucket: a.Bucket, + Key: a.Key, + Size: a.Size, + ContentHash: a.ContentHash, + ContentType: a.ContentType, + Lifecycle: string(a.Lifecycle), + ScopeAppID: a.ScopeAppID, + ScopeOrgID: a.ScopeOrgID, + ExpiresAt: a.ExpiresAt, + CreatedAt: a.CreatedAt, + DeletedAt: a.DeletedAt, + } +} + +func fromArtifactModel(m *artifactModel) (*artifact.Artifact, error) { + aid, err := id.ParseArtifactID(m.ID) + if err != nil { + return nil, err + } + + return &artifact.Artifact{ + ID: aid, + Backend: m.Backend, + Bucket: m.Bucket, + Key: m.Key, + Size: m.Size, + ContentHash: m.ContentHash, + ContentType: m.ContentType, + Lifecycle: artifact.Lifecycle(m.Lifecycle), + ScopeAppID: m.ScopeAppID, + ScopeOrgID: m.ScopeOrgID, + ExpiresAt: m.ExpiresAt, + CreatedAt: m.CreatedAt, + DeletedAt: m.DeletedAt, + }, nil +} + +func fromArtifactModels(models []artifactModel) ([]*artifact.Artifact, error) { + out := make([]*artifact.Artifact, 0, len(models)) + + for i := range models { + a, err := fromArtifactModel(&models[i]) + if err != nil { + return nil, err + } + + out = append(out, a) + } + + return out, nil +} + +// ── Artifact link model ─────────────────────────────────────────── + +type artifactLinkModel struct { + grove.BaseModel `grove:"table:dispatch_artifact_links"` + + ArtifactID string `grove:"artifact_id,pk"` + OwnerKind string `grove:"owner_kind,pk"` + OwnerID string `grove:"owner_id,pk"` + Name string `grove:"name,pk"` + Attempt int `grove:"attempt,pk"` + Role string `grove:"role,notnull"` + CreatedAt time.Time `grove:"created_at,notnull,default:current_timestamp"` +} + +func fromLinkModel(m *artifactLinkModel) (*artifact.Link, error) { + aid, err := id.ParseArtifactID(m.ArtifactID) + if err != nil { + return nil, err + } + + return &artifact.Link{ + ArtifactID: aid, + OwnerKind: artifact.OwnerKind(m.OwnerKind), + OwnerID: m.OwnerID, + Role: artifact.Role(m.Role), + Name: m.Name, + Attempt: m.Attempt, + CreatedAt: m.CreatedAt, + }, nil +} diff --git a/store/postgres/artifact_test.go b/store/postgres/artifact_test.go new file mode 100644 index 0000000..771170c --- /dev/null +++ b/store/postgres/artifact_test.go @@ -0,0 +1,22 @@ +//go:build integration + +package postgres_test + +import ( + "testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" +) + +// TestArtifactStoreConformance runs the shared artifact.Store suite +// against Postgres. +// +// Each subtest gets its own container. That is slow, but the suite +// asserts absolute row counts, so it needs a genuinely empty store, and +// this path only runs under the integration build tag. +func TestArtifactStoreConformance(t *testing.T) { + artifacttest.RunStoreSuite(t, func(t *testing.T) artifact.Store { + return setupTestStore(t) + }) +} diff --git a/store/postgres/cluster.go b/store/postgres/cluster.go index 78904b1..ac0d983 100644 --- a/store/postgres/cluster.go +++ b/store/postgres/cluster.go @@ -24,7 +24,7 @@ func (s *Store) RegisterWorker(ctx context.Context, w *cluster.Worker) error { Set("metadata = EXCLUDED.metadata"). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: register worker: %w", err) + return fmt.Errorf(errPrefix+"register worker: %w", err) } return nil } @@ -35,7 +35,7 @@ func (s *Store) DeregisterWorker(ctx context.Context, workerID id.WorkerID) erro Where("id = ?", workerID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: deregister worker: %w", err) + return fmt.Errorf(errPrefix+"deregister worker: %w", err) } if rows, _ := res.RowsAffected(); rows == 0 { //nolint:errcheck // driver always returns nil return dispatch.ErrWorkerNotFound @@ -50,7 +50,7 @@ func (s *Store) HeartbeatWorker(ctx context.Context, workerID id.WorkerID) error Where("id = ?", workerID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: heartbeat worker: %w", err) + return fmt.Errorf(errPrefix+"heartbeat worker: %w", err) } if rows, _ := res.RowsAffected(); rows == 0 { //nolint:errcheck // driver always returns nil return dispatch.ErrWorkerNotFound @@ -65,14 +65,14 @@ func (s *Store) ListWorkers(ctx context.Context) ([]*cluster.Worker, error) { OrderExpr("created_at ASC"). Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: list workers: %w", err) + return nil, fmt.Errorf(errPrefix+"list workers: %w", err) } workers := make([]*cluster.Worker, 0, len(models)) for i := range models { w, convErr := fromWorkerModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: list workers convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"list workers convert: %w", convErr) } workers = append(workers, w) } @@ -86,11 +86,11 @@ func (s *Store) DeleteStaleWorkers(ctx context.Context, threshold time.Duration) Where("last_seen < NOW() - ?::interval", threshold.String()). Exec(ctx) if err != nil { - return 0, fmt.Errorf("dispatch/bun: delete stale workers: %w", err) + return 0, fmt.Errorf(errPrefix+"delete stale workers: %w", err) } n, err := res.RowsAffected() if err != nil { - return 0, fmt.Errorf("dispatch/bun: delete stale workers rows affected: %w", err) + return 0, fmt.Errorf(errPrefix+"delete stale workers rows affected: %w", err) } return n, nil } @@ -103,14 +103,14 @@ func (s *Store) ReapDeadWorkers(ctx context.Context, threshold time.Duration) ([ Where("last_seen < NOW() - ?::interval", threshold.String()). Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: reap dead workers: %w", err) + return nil, fmt.Errorf(errPrefix+"reap dead workers: %w", err) } workers := make([]*cluster.Worker, 0, len(models)) for i := range models { w, convErr := fromWorkerModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: reap dead workers convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"reap dead workers convert: %w", convErr) } workers = append(workers, w) } @@ -130,19 +130,23 @@ func (s *Store) AcquireLeadership(ctx context.Context, workerID id.WorkerID, ttl Where("is_leader = TRUE AND leader_until < NOW()"). Exec(ctx) if err != nil { - return false, fmt.Errorf("dispatch/bun: clear expired leader: %w", err) + return false, fmt.Errorf(errPrefix+"clear expired leader: %w", err) } // Step 2: Check if there's already an active leader that isn't us. + // + // The whole row is selected rather than just the id column: the driver + // requires the number of returned columns to match the number of + // destination fields, so projecting a single column into the full + // model fails at scan time. var leader workerModel err = s.pgdb.NewSelect(&leader). - Column("id"). Where("is_leader = TRUE AND leader_until >= NOW()"). Limit(1). Scan(ctx) if err != nil { if !isNoRows(err) { - return false, fmt.Errorf("dispatch/bun: check leader: %w", err) + return false, fmt.Errorf(errPrefix+"check leader: %w", err) } // No active leader — proceed to claim. } else if leader.ID != wID { @@ -156,7 +160,7 @@ func (s *Store) AcquireLeadership(ctx context.Context, workerID id.WorkerID, ttl Where("id = ?", wID). Exec(ctx) if claimErr != nil { - return false, fmt.Errorf("dispatch/bun: claim leadership: %w", claimErr) + return false, fmt.Errorf(errPrefix+"claim leadership: %w", claimErr) } if rows, _ := res.RowsAffected(); rows == 0 { //nolint:errcheck // driver always returns nil return false, nil @@ -174,7 +178,7 @@ func (s *Store) RenewLeadership(ctx context.Context, workerID id.WorkerID, ttl t Where("id = ? AND is_leader = TRUE", workerID.String()). Exec(ctx) if err != nil { - return false, fmt.Errorf("dispatch/bun: renew leadership: %w", err) + return false, fmt.Errorf(errPrefix+"renew leadership: %w", err) } if rows, _ := res.RowsAffected(); rows == 0 { //nolint:errcheck // driver always returns nil return false, nil @@ -193,7 +197,7 @@ func (s *Store) GetLeader(ctx context.Context) (*cluster.Worker, error) { if isNoRows(err) { return nil, nil } - return nil, fmt.Errorf("dispatch/bun: get leader: %w", err) + return nil, fmt.Errorf(errPrefix+"get leader: %w", err) } return fromWorkerModel(m) } diff --git a/store/postgres/cron.go b/store/postgres/cron.go index cc76fca..bf5cda6 100644 --- a/store/postgres/cron.go +++ b/store/postgres/cron.go @@ -19,7 +19,7 @@ func (s *Store) RegisterCron(ctx context.Context, entry *cron.Entry) error { if isDuplicateKey(err) { return dispatch.ErrDuplicateCron } - return fmt.Errorf("dispatch/bun: register cron: %w", err) + return fmt.Errorf(errPrefix+"register cron: %w", err) } return nil } @@ -35,7 +35,7 @@ func (s *Store) GetCron(ctx context.Context, entryID id.CronID) (*cron.Entry, er if isNoRows(err) { return nil, dispatch.ErrCronNotFound } - return nil, fmt.Errorf("dispatch/bun: get cron: %w", err) + return nil, fmt.Errorf(errPrefix+"get cron: %w", err) } return fromCronModel(m) } @@ -47,14 +47,14 @@ func (s *Store) ListCrons(ctx context.Context) ([]*cron.Entry, error) { OrderExpr("created_at ASC"). Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: list crons: %w", err) + return nil, fmt.Errorf(errPrefix+"list crons: %w", err) } entries := make([]*cron.Entry, 0, len(models)) for i := range models { e, convErr := fromCronModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: list crons convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"list crons convert: %w", convErr) } entries = append(entries, e) } @@ -77,7 +77,7 @@ func (s *Store) AcquireCronLock(ctx context.Context, entryID id.CronID, workerID Where("(locked_by IS NULL OR locked_until < ? OR locked_by = ?)", now, wID). Exec(ctx) if err != nil { - return false, fmt.Errorf("dispatch/bun: acquire cron lock: %w", err) + return false, fmt.Errorf(errPrefix+"acquire cron lock: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil @@ -87,7 +87,7 @@ func (s *Store) AcquireCronLock(ctx context.Context, entryID id.CronID, workerID Where("id = ?", entryID.String()). Count(ctx) if existErr != nil { - return false, fmt.Errorf("dispatch/bun: check cron exists: %w", existErr) + return false, fmt.Errorf(errPrefix+"check cron exists: %w", existErr) } if count == 0 { return false, dispatch.ErrCronNotFound @@ -109,7 +109,7 @@ func (s *Store) ReleaseCronLock(ctx context.Context, entryID id.CronID, workerID Where("locked_by = ?", workerID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: release cron lock: %w", err) + return fmt.Errorf(errPrefix+"release cron lock: %w", err) } return nil } @@ -122,7 +122,7 @@ func (s *Store) UpdateCronLastRun(ctx context.Context, entryID id.CronID, at tim Where("id = ?", entryID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: update cron last run: %w", err) + return fmt.Errorf(errPrefix+"update cron last run: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { @@ -137,7 +137,7 @@ func (s *Store) UpdateCronEntry(ctx context.Context, entry *cron.Entry) error { m.UpdatedAt = time.Now().UTC() res, err := s.pgdb.NewUpdate(m).WherePK().Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: update cron entry: %w", err) + return fmt.Errorf(errPrefix+"update cron entry: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { @@ -152,7 +152,7 @@ func (s *Store) DeleteCron(ctx context.Context, entryID id.CronID) error { Where("id = ?", entryID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: delete cron: %w", err) + return fmt.Errorf(errPrefix+"delete cron: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { diff --git a/store/postgres/dequeue_sql_test.go b/store/postgres/dequeue_sql_test.go new file mode 100644 index 0000000..544a398 --- /dev/null +++ b/store/postgres/dequeue_sql_test.go @@ -0,0 +1,362 @@ +package postgres + +import ( + "strconv" + "strings" + "testing" + "time" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// The conformance suite that actually proves this backend correct lives +// behind //go:build integration and needs Docker. These tests deliberately +// do not, so `go test ./store/postgres/...` on a machine with no container +// runtime still checks the shape of the statement rather than reporting a +// green package that ran nothing. +// +// They assert the compiled SQL, not behaviour — the database is the only +// thing that can answer for behaviour, and dequeue_test.go asks it. + +// TestBuildDequeueQueryUnboundedEmitsOriginalStatement is the +// backward-compatibility guarantee in its narrowest form: opts that +// constrain nothing must compile to exactly the statement that shipped +// before DequeueOpts existed, with no fit predicate at all. A worker not +// using the resource model claims everything, jobs declaring custom +// resources included. +func TestBuildDequeueQueryUnboundedEmitsOriginalStatement(t *testing.T) { + opts := job.DequeueOpts{Queues: []string{"default"}, Limit: 10} + + if !opts.IsUnbounded() { + t.Fatalf("DequeueOpts%+v.IsUnbounded() = false, want true", opts) + } + + query, args := buildDequeueQuery(opts) + + for _, banned := range []string{"req_", "REPLACE", "primary_input_hash"} { + if strings.Contains(query, banned) { + t.Errorf("unbounded dequeue emitted %q:\n%s", banned, query) + } + } + + if !strings.Contains(query, "ORDER BY priority DESC, run_at ASC\n") { + t.Errorf("unbounded dequeue lost the original ordering:\n%s", query) + } + + // Queues and limit, nothing else. + if len(args) != 2 { + t.Fatalf("unbounded dequeue bound %d args, want 2: %v", len(args), args) + } + + if args[1] != 10 { + t.Errorf("limit bound as %v, want 10", args[1]) + } +} + +// TestBuildDequeueQueryAppliesLocalityToUnboundedOpts pins the split +// IsUnbounded exists to make: it governs FILTERING only. Opts carrying +// nothing but PreferHashes are unbounded, so no fit predicate is emitted — +// and the locality term is applied anyway. A backend that derived "should +// I order?" from IsUnbounded would silently drop the signal here. +func TestBuildDequeueQueryAppliesLocalityToUnboundedOpts(t *testing.T) { + opts := job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 4, + PreferHashes: []string{"blake3:staged"}, + } + + if !opts.IsUnbounded() { + t.Fatal("opts carrying only PreferHashes report IsUnbounded() = false") + } + + query, _ := buildDequeueQuery(opts) + + if strings.Contains(query, "req_") { + t.Errorf("PreferHashes emitted a fit predicate — locality must never filter:\n%s", query) + } + + if !strings.Contains(query, "COALESCE(primary_input_hash = ANY(") { + t.Errorf("locality term missing from unbounded opts:\n%s", query) + } +} + +// TestBuildDequeueQueryOrdersLocalityBelowPriority pins the one ordering +// mistake that would not show up as a wrong answer, only as starvation: +// priority must come first. Locality above it would let a steady stream of +// low-priority jobs with staged inputs beat a high-priority job with cold +// ones. +// +// It also pins that the ordering is applied to the inner candidate SELECT +// as well as the outer one, so the LIMIT truncates an ordered set. +func TestBuildDequeueQueryOrdersLocalityBelowPriority(t *testing.T) { + query, _ := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 2, + PreferHashes: []string{"blake3:staged"}, + }) + + const want = "priority DESC, COALESCE(primary_input_hash = ANY($2), FALSE) DESC, run_at ASC" + + if n := strings.Count(query, want); n != 2 { + t.Fatalf("ordering %q appears %d times, want 2 (inner candidate SELECT and outer SELECT):\n%s", + want, n, query) + } + + // The LIMIT must sit after the inner ORDER BY, or it truncates an + // unordered set and the ordering above is decoration. + orderAt := strings.Index(query, want) + limitAt := strings.Index(query, "LIMIT ") + + if orderAt > limitAt { + t.Errorf("inner LIMIT precedes the inner ORDER BY:\n%s", query) + } +} + +// TestBuildDequeueQueryBindsEveryValue is the injection check. Every value +// the caller controls — queue names, budgets, custom keys, the reserved +// id, hashes, the limit — must reach Postgres as a bind parameter. The +// only thing concatenated into the statement is a column name from +// budgetColumns, all of which are compile-time constants. +func TestBuildDequeueQueryBindsEveryValue(t *testing.T) { + reserved := id.NewJobID() + + query, args := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"q'; DROP TABLE dispatch_jobs; --"}, + Limit: 3, + Budget: resource.Set{resource.Memory: 4 << 30}, + CustomKeys: []string{"fpga'); --"}, + PreferHashes: []string{"blake3:x"}, + ReservedFor: &reserved, + }) + + for _, hostile := range []string{"DROP TABLE", "fpga", reserved.String(), "blake3:x"} { + if strings.Contains(query, hostile) { + t.Errorf("value %q was interpolated into the statement:\n%s", hostile, query) + } + } + + // queues, reserved id, memory budget, separator, one custom key, + // prefer hashes, limit. + if len(args) != 7 { + t.Fatalf("bound %d args, want 7: %v", len(args), args) + } +} + +// TestBuildDequeueQueryBudgetPredicate pins the three rules a budget +// comparison has to get right at once: only declared keys are compared, a +// declared zero is still a real constraint, and the comparison is <= so an +// exact fit is claimable. +func TestBuildDequeueQueryBudgetPredicate(t *testing.T) { + query, _ := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + Budget: resource.Set{resource.Memory: 4 << 30, resource.GPU: 0}, + }) + + for _, want := range []string{"req_memory_bytes <= $", "req_gpu_milli <= $"} { + if !strings.Contains(query, want) { + t.Errorf("missing %q:\n%s", want, query) + } + } + + // CPU and disk were never declared, so they are unconstrained — not + // compared against zero. + for _, banned := range []string{"req_cpu_milli", "req_disk_bytes"} { + if strings.Contains(query, banned) { + t.Errorf("undeclared dimension %q was constrained:\n%s", banned, query) + } + } + + if strings.Contains(query, "req_memory_bytes < $") && + !strings.Contains(query, "req_memory_bytes <= $") { + t.Error("budget compared with < rather than <=; an exact fit must be claimable") + } +} + +// TestBuildDequeueQueryCustomKeysAreASubsetTest pins containment as +// nested REPLACE rather than LIKE. One REPLACE per offered key, each +// stripping ",key," and putting the separator back, with the surviving +// string required to be empty or a lone separator. +// +// The LIKE formulation this replaces passes every single-key case in the +// conformance suite and then silently strands multi-key jobs, so the shape +// is worth pinning here as well as behaviourally. +func TestBuildDequeueQueryCustomKeysAreASubsetTest(t *testing.T) { + query, args := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + CustomKeys: []string{"tpu", "fpga", "nvme"}, + }) + + if strings.Contains(query, "LIKE") { + t.Errorf("custom-key containment used LIKE — that is a substring test:\n%s", query) + } + + if n := strings.Count(query, "REPLACE("); n != 3 { + t.Errorf("emitted %d REPLACE calls for 3 offered keys:\n%s", n, query) + } + + if !strings.Contains(query, "IN ('', $") { + t.Errorf("subset test does not end in the empty-or-separator check:\n%s", query) + } + + // Keys are bound wrapped in separators, which is what stops ",fpga," + // matching a job that needs ",fpga-large,". + for _, want := range []string{",fpga,", ",nvme,", ",tpu,"} { + if !hasArg(args, want) { + t.Errorf("offered key %q not bound wrapped in separators: %v", want, args) + } + } +} + +// TestBuildDequeueQueryBoundedWithNoCustomKeysExcludesCustomJobs is the +// half of the empty-offer rule a backend gets wrong. Bounded opts with an +// empty offer are a resource-aware worker with no custom resources, so a +// job requiring an fpga must not be claimable — the predicate still has to +// be emitted, with no REPLACE wrapping it. +func TestBuildDequeueQueryBoundedWithNoCustomKeysExcludesCustomJobs(t *testing.T) { + query, _ := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + Budget: resource.Set{resource.Memory: 4 << 30}, + }) + + if strings.Contains(query, "REPLACE(") { + t.Errorf("an empty offer emitted a REPLACE:\n%s", query) + } + + if !strings.Contains(query, "req_custom_keys IN ('', $") { + t.Errorf("bounded opts with no offered keys must still exclude custom-key jobs:\n%s", query) + } +} + +func hasArg(args []any, want string) bool { + for _, a := range args { + if s, ok := a.(string); ok && s == want { + return true + } + } + + return false +} + +// argFor returns the value the statement reads at the assignment starting +// with prefix, by resolving the $N it names against the args slice. +// +// Postgres numbers its placeholders, so the SQLite bind-order hazard does +// not exist here — but the equivalent one does: buildDequeueQuery's bind +// closure derives each number from len(args) at the moment it is called, +// so a helper that writes its text and appends its values in different +// orders would emit a number naming somebody else's value. Reading the +// number back out of the finished statement is what catches that. +func argFor(t *testing.T, query string, args []any, prefix string) any { + t.Helper() + + i := strings.Index(query, prefix) + if i < 0 { + t.Fatalf("statement has no %q assignment:\n%s", prefix, query) + } + + rest := query[i+len(prefix):] + if rest == "" || rest[0] != '$' { + t.Fatalf("%q is not read from a bind parameter:\n%s", prefix, query) + } + + end := 1 + for end < len(rest) && rest[end] >= '0' && rest[end] <= '9' { + end++ + } + + n, err := strconv.Atoi(rest[1:end]) + if err != nil { + t.Fatalf("%q names an unparseable placeholder %q", prefix, rest[:end]) + } + + if n < 1 || n > len(args) { + t.Fatalf("%q names $%d but only %d args were bound: %v", prefix, n, len(args), args) + } + + return args[n-1] +} + +// TestBuildDequeueQueryGrantsLeaseInTheClaim pins the two halves of the +// grant: it is absent unless the caller asks for it, and when present it +// is part of the claiming UPDATE's SET clause rather than a second +// statement — which is what makes a claimed job always carry a lease +// (see job.DequeueOpts.LeaseUntil). +// +// The epoch is pinned as `lease_epoch + 1` specifically. Binding a +// computed successor instead would need a prior read, and the read is +// what the single statement exists to avoid. +func TestBuildDequeueQueryGrantsLeaseInTheClaim(t *testing.T) { + worker := id.NewWorkerID() + + // Deliberately NOT UTC. lease_expires_at is a timestamptz and the + // driver would convert either way, but every other timestamp this + // package writes is normalized before it is bound, and a caller that + // hands over a wall-clock time in its own zone is the ordinary case. + until := time.Date(2026, 8, 12, 12, 1, 30, 0, time.FixedZone("UTC-5", -5*60*60)) + + base := job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 3, + Budget: resource.Set{resource.Memory: 4 << 30}, + } + + grantOpts := base + grantOpts.WorkerID = worker + grantOpts.LeaseUntil = until + + t.Run("no grant leaves the lease columns alone", func(t *testing.T) { + query, _ := buildDequeueQuery(base) + + for _, banned := range []string{"worker_id", "lease_epoch", "lease_expires_at"} { + if strings.Contains(query, banned) { + t.Errorf("opts granting no lease still wrote %q:\n%s", banned, query) + } + } + }) + + t.Run("grant rides in the SET clause", func(t *testing.T) { + query, args := buildDequeueQuery(grantOpts) + + setEnd := strings.Index(query, "WHERE id IN (") + if setEnd < 0 { + t.Fatalf("statement lost its claim shape:\n%s", query) + } + + // Every lease assignment must fall inside the UPDATE's SET clause, + // which is the part of the text before the candidate subquery. + for _, want := range []string{"worker_id = $", "lease_epoch = lease_epoch + 1", "lease_expires_at = $"} { + at := strings.Index(query, want) + if at < 0 { + t.Errorf("granting opts did not emit %q:\n%s", want, query) + + continue + } + + if at > setEnd { + t.Errorf("%q is outside the claiming SET clause:\n%s", want, query) + } + } + + if got := argFor(t, query, args, "worker_id = "); got != worker.String() { + t.Errorf("worker_id reads %v, want %s", got, worker) + } + + got, ok := argFor(t, query, args, "lease_expires_at = ").(time.Time) + if !ok { + t.Fatalf("lease_expires_at was not bound as a time.Time: %v", args) + } + + if !got.Equal(until) { + t.Errorf("lease_expires_at reads %v, want %v", got, until) + } + + if got.Location() != time.UTC { + t.Errorf("lease_expires_at bound in %v, want UTC", got.Location()) + } + }) +} diff --git a/store/postgres/dequeue_test.go b/store/postgres/dequeue_test.go new file mode 100644 index 0000000..2acd1e8 --- /dev/null +++ b/store/postgres/dequeue_test.go @@ -0,0 +1,278 @@ +//go:build integration + +package postgres_test + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/xraph/grove/drivers/pgdriver" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/store/storetest" +) + +// TestDequeueConformance runs the resource-aware dequeue suite against the +// Postgres store — the reference SQL implementation the SQLite backend +// follows. +// +// The container is stood up once and shared by every subtest. The suite +// documents that this is safe: each case enqueues onto its own queue and +// asserts only on the jobs it created, and standing up a fresh container +// per subtest would dominate the runtime of the whole file. +func TestDequeueConformance(t *testing.T) { + shared := setupTestStore(t) + + storetest.RunDequeueSuite(t, func(t *testing.T) job.Store { + t.Helper() + + return shared + }) +} + +// TestDequeueOrdersNullPrimaryInputHashAsUnpreferred covers the one row +// shape the shared suite cannot produce. +// +// primary_input_hash is nullable, and rows written before the resource +// migration carry a genuine SQL NULL rather than the empty string the +// current insert path writes. `NULL = ANY(...)` evaluates to NULL, not +// false, and Postgres sorts NULLs FIRST under DESC — so an uncoalesced +// locality term would rank exactly the rows with no locality signal ABOVE +// the ones the caller has already staged, inverting the optimization. +// +// The NULL is written with raw SQL because no Go path can produce one: +// jobModel.PrimaryInputHash is a plain string. +func TestDequeueOrdersNullPrimaryInputHashAsUnpreferred(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + const ( + queue = "null-hash-order" + local = "blake3:staged-here" + ) + + base := time.Now().UTC().Add(-time.Hour).Truncate(time.Millisecond) + + // nullHash is the earlier of the two by RunAt, so it wins any tie the + // ordering fails to break: if the locality term is not NULL-safe it + // comes back first. + nullHash := newHashFixture("null-hash", queue, base) + cached := newHashFixture("cached", queue, base.Add(time.Minute)) + cached.PrimaryInputHash = local + + for _, j := range []*job.Job{nullHash, cached} { + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue %s: %v", j.Name, err) + } + } + + if _, err := pgdriver.Unwrap(s.DB()).NewRaw( + `UPDATE dispatch_jobs SET primary_input_hash = NULL WHERE id = $1`, + nullHash.ID.String(), + ).Exec(ctx); err != nil { + t.Fatalf("null out primary_input_hash: %v", err) + } + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + PreferHashes: []string{local}, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 2 { + t.Fatalf("claimed %d jobs, want 2", len(got)) + } + + if got[0].Name != "cached" || got[1].Name != "null-hash" { + t.Fatalf("claimed [%s %s], want [cached null-hash]: a NULL hash must sort as "+ + "not preferred, not ahead of the job the caller has staged", + got[0].Name, got[1].Name) + } +} + +// TestDequeueBoundedQueryPlanUsesDequeueIndex proves the fit predicate did +// not cost the candidate scan its index. +// +// idx_dispatch_jobs_dequeue_res is what keeps that scan from reading every +// pending row on a busy queue, and a predicate written so the planner +// *could not* use it — a function wrapped around a compared column, a +// value cast the wrong way — would surface only as a latency regression +// under load, long after this change shipped. The plan is taken for the +// inner candidate SELECT, which is the half the index serves. +// +// The check asks whether the predicate CAN use the index, not which plan +// the planner happens to cost lowest: at any table size a test can +// populate in a second, a sequential scan over dispatch_jobs is genuinely +// cheaper, so a bare EXPLAIN would only be measuring the fixture. So +// sequential scans are disabled and the two indexes that would otherwise +// win on cost are dropped inside a transaction that is always rolled +// back — DDL is transactional in Postgres, so the schema is untouched. +// All of it runs on one dedicated connection, or the pool could hand the +// EXPLAIN a session that never saw the setup. +func TestDequeueBoundedQueryPlanUsesDequeueIndex(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + const ( + queue = "plan-check" + rows = 500 + ) + + base := time.Now().UTC().Add(-time.Hour) + + for i := range rows { + j := newHashFixture("plan", queue, base) + j.Resources = resource.Set{resource.Memory: storetest.GiB} + j.Priority = i % 7 + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + } + + conn, err := pgdriver.Unwrap(s.DB()).AcquireConn(ctx) + if err != nil { + t.Fatalf("acquire dedicated conn: %v", err) + } + + defer conn.Release() + + for _, stmt := range []string{ + `ANALYZE dispatch_jobs`, + `SET enable_seqscan = off`, + `BEGIN`, + // idx_dispatch_jobs_state is cheaper on a table this small, so it + // would satisfy a name check without proving anything. + // + // idx_dispatch_jobs_dequeue is NOT dropped here because the + // resource migration drops it for good: it is the same key list + // and the same partial predicate as + // idx_dispatch_jobs_dequeue_res, differing only by the INCLUDE + // payload, so shipping both meant two B-tree writes per enqueue + // and two per claim on the busiest table in the schema for a plan + // the planner would never pick. Its absence is asserted below. + `DROP INDEX idx_dispatch_jobs_state`, + } { + if _, execErr := conn.Exec(ctx, stmt); execErr != nil { + t.Fatalf("%s: %v", stmt, execErr) + } + } + + defer func() { + if _, rbErr := conn.Exec(ctx, `ROLLBACK`); rbErr != nil { + t.Errorf("rollback index drops: %v", rbErr) + } + }() + + var plan []byte + + // The fully bounded predicate: all four dimensions plus the nested + // REPLACE containment test. + err = conn.QueryRow(ctx, ` + EXPLAIN (FORMAT JSON) + SELECT id FROM dispatch_jobs + WHERE state IN ('pending', 'retrying') + AND queue = ANY($1) + AND run_at <= NOW() + AND req_cpu_milli <= $2 + AND req_memory_bytes <= $3 + AND req_disk_bytes <= $4 + AND req_gpu_milli <= $5 + AND REPLACE(req_custom_keys, $6, $7) IN ('', $7) + ORDER BY priority DESC, run_at ASC + LIMIT 4`, + []string{queue}, + 8*resource.MilliScale, 4*storetest.GiB, 100*storetest.GiB, 4*resource.MilliScale, + ",fpga,", ",", + ).Scan(&plan) + if err != nil { + t.Fatalf("explain: %v", err) + } + + if !strings.Contains(string(plan), "idx_dispatch_jobs_dequeue_res") { + t.Errorf("bounded dequeue plan does not use idx_dispatch_jobs_dequeue_res:\n%s", plan) + } + + t.Logf("bounded dequeue plan:\n%s", plan) +} + +// TestResourceMigrationDropsTheRedundantDequeueIndex pins that the +// migration retires the index its own covering index supersedes. +// +// idx_dispatch_jobs_dequeue (migration 001) and +// idx_dispatch_jobs_dequeue_res (migration 009) have the IDENTICAL key +// list — (queue, priority DESC, run_at ASC) — and the IDENTICAL partial +// predicate. The only difference is the INCLUDE payload, which makes the +// second a strict superset. Every deployment that kept both paid a +// second index insert on every enqueue and a second delete on every +// claim, forever, on the hottest table in the schema, to serve a plan +// the planner has no reason to choose. +func TestResourceMigrationDropsTheRedundantDequeueIndex(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + conn, err := pgdriver.Unwrap(s.DB()).AcquireConn(ctx) + if err != nil { + t.Fatalf("acquire dedicated conn: %v", err) + } + + defer conn.Release() + + for _, tc := range []struct { + index string + want bool + }{ + {"idx_dispatch_jobs_dequeue_res", true}, + {"idx_dispatch_jobs_dequeue", false}, + } { + var present bool + + if err = conn.QueryRow(ctx, + `SELECT EXISTS (SELECT 1 FROM pg_class WHERE relname = $1 AND relkind = 'i')`, + tc.index).Scan(&present); err != nil { + t.Fatalf("look up %s: %v", tc.index, err) + } + + if present != tc.want { + t.Errorf("index %s present = %v, want %v", tc.index, present, tc.want) + } + } + + // The covering index must also be VALID: CREATE INDEX CONCURRENTLY + // leaves an unusable-but-present index behind when it fails, which + // the planner ignores and IF NOT EXISTS would then skip forever. + var valid bool + + if err = conn.QueryRow(ctx, ` + SELECT i.indisvalid + FROM pg_class c JOIN pg_index i ON i.indexrelid = c.oid + WHERE c.relname = 'idx_dispatch_jobs_dequeue_res'`).Scan(&valid); err != nil { + t.Fatalf("read indisvalid: %v", err) + } + + if !valid { + t.Error("idx_dispatch_jobs_dequeue_res is INVALID: a failed CONCURRENTLY build was " + + "left in place, so the dequeue scan has no usable index and nothing reports it") + } +} + +func newHashFixture(name, queue string, runAt time.Time) *job.Job { + return &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: name, + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: runAt, + } +} diff --git a/store/postgres/dlq.go b/store/postgres/dlq.go index c6cf091..858cb8d 100644 --- a/store/postgres/dlq.go +++ b/store/postgres/dlq.go @@ -12,10 +12,12 @@ import ( // PushDLQ adds a failed job entry to the dead letter queue. func (s *Store) PushDLQ(ctx context.Context, entry *dlq.Entry) error { - m := toDLQModel(entry) - _, err := s.pgdb.NewInsert(m).Exec(ctx) + m, err := toDLQModel(entry) if err != nil { - return fmt.Errorf("dispatch/bun: push dlq: %w", err) + return err + } + if _, err = s.pgdb.NewInsert(m).Exec(ctx); err != nil { + return fmt.Errorf(errPrefix+"push dlq: %w", err) } return nil } @@ -40,14 +42,14 @@ func (s *Store) ListDLQ(ctx context.Context, opts dlq.ListOpts) ([]*dlq.Entry, e err := q.Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: list dlq: %w", err) + return nil, fmt.Errorf(errPrefix+"list dlq: %w", err) } entries := make([]*dlq.Entry, 0, len(models)) for i := range models { e, convErr := fromDLQModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: list dlq convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"list dlq convert: %w", convErr) } entries = append(entries, e) } @@ -65,7 +67,7 @@ func (s *Store) GetDLQ(ctx context.Context, entryID id.DLQID) (*dlq.Entry, error if isNoRows(err) { return nil, dispatch.ErrDLQNotFound } - return nil, fmt.Errorf("dispatch/bun: get dlq: %w", err) + return nil, fmt.Errorf(errPrefix+"get dlq: %w", err) } return fromDLQModel(m) } @@ -77,7 +79,7 @@ func (s *Store) ReplayDLQ(ctx context.Context, entryID id.DLQID) error { Where("id = ?", entryID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: replay dlq: %w", err) + return fmt.Errorf(errPrefix+"replay dlq: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { @@ -93,7 +95,7 @@ func (s *Store) PurgeDLQ(ctx context.Context, before time.Time) (int64, error) { Where("failed_at < ?", before). Exec(ctx) if err != nil { - return 0, fmt.Errorf("dispatch/bun: purge dlq: %w", err) + return 0, fmt.Errorf(errPrefix+"purge dlq: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil return rows, nil @@ -104,7 +106,7 @@ func (s *Store) CountDLQ(ctx context.Context) (int64, error) { count, err := s.pgdb.NewSelect((*dlqEntryModel)(nil)). Count(ctx) if err != nil { - return 0, fmt.Errorf("dispatch/bun: count dlq: %w", err) + return 0, fmt.Errorf(errPrefix+"count dlq: %w", err) } return count, nil } diff --git a/store/postgres/event.go b/store/postgres/event.go index cebb7b3..0afba07 100644 --- a/store/postgres/event.go +++ b/store/postgres/event.go @@ -15,7 +15,7 @@ func (s *Store) PublishEvent(ctx context.Context, evt *event.Event) error { m := toEventModel(evt) _, err := s.pgdb.NewInsert(m).Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: publish event: %w", err) + return fmt.Errorf(errPrefix+"publish event: %w", err) } return nil } @@ -52,12 +52,12 @@ func (s *Store) SubscribeEvent(ctx context.Context, name string, timeout time.Du sleepCtx(ctx, 50*time.Millisecond) continue } - return nil, fmt.Errorf("dispatch/bun: subscribe event: %w", err) + return nil, fmt.Errorf(errPrefix+"subscribe event: %w", err) } evt, convErr := fromEventModel(m) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: subscribe event convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"subscribe event convert: %w", convErr) } return evt, nil } @@ -70,7 +70,7 @@ func (s *Store) AckEvent(ctx context.Context, eventID id.EventID) error { Where("id = ?", eventID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: ack event: %w", err) + return fmt.Errorf(errPrefix+"ack event: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { diff --git a/store/postgres/helpers.go b/store/postgres/helpers.go index bd834f4..014513d 100644 --- a/store/postgres/helpers.go +++ b/store/postgres/helpers.go @@ -7,6 +7,13 @@ import ( "github.com/jackc/pgx/v5/pgconn" ) +// errPrefix identifies this storage layer in wrapped errors and log messages. +// Concatenate it into format strings rather than passing it as an argument so +// the result stays a compile-time constant and go vet keeps checking the verbs: +// +// fmt.Errorf(errPrefix+"get job: %w", err) +const errPrefix = "dispatch/postgres: " + // isNoRows returns true when err indicates no rows were found. func isNoRows(err error) bool { return errors.Is(err, sql.ErrNoRows) diff --git a/store/postgres/job.go b/store/postgres/job.go index 0430252..a90b395 100644 --- a/store/postgres/job.go +++ b/store/postgres/job.go @@ -3,63 +3,276 @@ package postgres import ( "context" "fmt" + "strconv" + "strings" "time" "github.com/xraph/dispatch" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" ) // EnqueueJob persists a new job in pending state. func (s *Store) EnqueueJob(ctx context.Context, j *job.Job) error { - m := toJobModel(j) - _, err := s.pgdb.NewInsert(m).Exec(ctx) + m, err := toJobModel(j) + if err != nil { + return err + } + + _, err = s.pgdb.NewInsert(m).Exec(ctx) if err != nil { if isDuplicateKey(err) { return dispatch.ErrJobAlreadyExists } - return fmt.Errorf("dispatch/bun: enqueue job: %w", err) + return fmt.Errorf(errPrefix+"enqueue job: %w", err) } s.notifyWake(ctx) return nil } -// DequeueJobs atomically claims up to limit pending jobs from the given -// queues, sets them to running, and returns them. Uses SELECT FOR UPDATE -// SKIP LOCKED for concurrent-safe dequeue via raw SQL. -func (s *Store) DequeueJobs(ctx context.Context, queues []string, limit int) ([]*job.Job, error) { +// DequeueJobs atomically claims up to opts.Limit ready jobs from +// opts.Queues that fit opts, sets them to running, and returns them +// ordered by priority descending, then locality-preferred first, then +// RunAt ascending. +// +// The fit predicate is compiled into the same statement that performs +// the claim, so a job that does not fit is never written to: it stays +// pending and untouched for the next worker that does have room. The +// UPDATE ... WHERE id IN (SELECT ... FOR UPDATE SKIP LOCKED) shape that +// makes the claim atomic is unchanged; the predicate is simply another +// conjunct of the inner SELECT's WHERE. +// +// When opts.Grants() the lease columns are additional assignments in that +// same UPDATE's SET clause, never a follow-up statement. ReclaimExpiredLeases +// requires lease_expires_at IS NOT NULL, so a row left running with a null +// expiry is not a row at risk of reclamation — it is one reclamation can +// never see. See job.DequeueOpts.LeaseUntil. +func (s *Store) DequeueJobs(ctx context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + // A worker computing zero free slots must claim zero jobs, never the + // whole queue. Postgres would already return nothing for LIMIT 0, but + // a negative LIMIT is an error rather than an empty result, and + // neither is worth a round trip. + if opts.Limit <= 0 { + return nil, nil + } + + if err := opts.Validate(); err != nil { + return nil, fmt.Errorf(errPrefix+"dequeue jobs: %w", err) + } + + query, args := buildDequeueQuery(opts) + var models []jobModel - err := s.pgdb.NewRaw(` + + err := s.pgdb.NewRaw(query, args...).Scan(ctx, &models) + if err != nil { + return nil, fmt.Errorf(errPrefix+"dequeue jobs: %w", err) + } + + jobs := make([]*job.Job, 0, len(models)) + for i := range models { + j, convErr := fromJobModel(&models[i]) + if convErr != nil { + return nil, fmt.Errorf(errPrefix+"dequeue convert: %w", convErr) + } + jobs = append(jobs, j) + } + return jobs, nil +} + +// budgetColumns maps each canonical dimension the dequeue predicate +// compares to the scalar column that holds it. These are exactly the +// dimensions job.DequeueOpts.Allows loops over, and exactly the columns +// idx_dispatch_jobs_dequeue_res INCLUDEs, so each comparison is a scalar +// range test the index can answer from its own tuples rather than a JSON +// probe into resource_requests. +// +// The column names are compile-time constants and are the only +// identifiers ever concatenated into the statement below; every value +// travels as a bind parameter. +var budgetColumns = []struct { + key string + column string +}{ + {resource.CPU, "req_cpu_milli"}, + {resource.Memory, "req_memory_bytes"}, + {resource.Disk, "req_disk_bytes"}, + {resource.GPU, "req_gpu_milli"}, +} + +// dequeueSQL is the claim statement with four things filled in: the lease +// grant, the fit predicate, the ordering, and the limit placeholder. The +// ordering is substituted twice because the inner SELECT decides WHICH +// rows the LIMIT keeps and the outer SELECT decides the order they come +// back in — ordering only the outer one would hand a small-limit worker +// an arbitrary slice of the eligible set in tidy order. +// +// The grant is a suffix of the SET clause rather than a statement of its +// own, which is what makes "claimed" and "leased" the same event. +const dequeueSQL = ` WITH dequeued AS ( UPDATE dispatch_jobs - SET state = 'running', started_at = NOW(), updated_at = NOW() + SET state = 'running', started_at = NOW(), updated_at = NOW()%s WHERE id IN ( SELECT id FROM dispatch_jobs WHERE state IN ('pending', 'retrying') AND queue = ANY($1) - AND run_at <= NOW() - ORDER BY priority DESC, run_at ASC + AND run_at <= NOW()%s + ORDER BY %s FOR UPDATE SKIP LOCKED - LIMIT $2 + LIMIT %s ) RETURNING * ) - SELECT * FROM dequeued ORDER BY priority DESC, run_at ASC`, - queues, limit, - ).Scan(ctx, &models) - if err != nil { - return nil, fmt.Errorf("dispatch/bun: dequeue jobs: %w", err) + SELECT * FROM dequeued ORDER BY %s` + +// buildDequeueQuery compiles opts into the claim statement and its bind +// parameters. It is the SQL expression of job.DequeueOpts.Allows and +// Less, and must answer identically for every job. +func buildDequeueQuery(opts job.DequeueOpts) (query string, args []any) { + args = []any{opts.Queues} + + // bind appends v and returns the placeholder that reads it. Values + // never reach the statement text. + bind := func(v any) string { + args = append(args, v) + + return "$" + strconv.Itoa(len(args)) } - jobs := make([]*job.Job, 0, len(models)) - for i := range models { - j, convErr := fromJobModel(&models[i]) - if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: dequeue convert: %w", convErr) + grant := buildLeaseGrant(opts, bind) + fit := buildFitPredicate(opts, bind) + order := buildDequeueOrder(opts, bind) + limit := bind(opts.Limit) + + return fmt.Sprintf(dequeueSQL, grant, fit, order, limit, order), args +} + +// buildLeaseGrant renders the lease assignments appended to the claim's +// SET clause, or "" when opts grants no lease. +// +// Empty is the whole backward-compatibility guarantee: a caller that does +// not opt in emits the statement it emitted before leases existed and +// leaves worker_id, lease_epoch and lease_expires_at exactly as they were. +// +// lease_epoch = lease_epoch + 1 rather than a bound value, because the +// epoch is the fence: it must advance from whatever the row currently +// holds, which only the row knows. Reading it and writing back a computed +// successor would be the read-modify-write this statement exists to avoid. +func buildLeaseGrant(opts job.DequeueOpts, bind func(any) string) string { + if !opts.Grants() { + return "" + } + + return ",\n\t\t\t worker_id = " + bind(opts.WorkerID.String()) + + ",\n\t\t\t lease_epoch = lease_epoch + 1" + + ",\n\t\t\t lease_expires_at = " + bind(opts.LeaseUntil.UTC()) +} + +// buildFitPredicate renders the conjuncts that decide WHICH jobs may be +// claimed, or "" when opts constrains nothing. +func buildFitPredicate(opts job.DequeueOpts, bind func(any) string) string { + // Unbounded opts emit the original query verbatim: a caller that does + // not use the resource model claims everything, including jobs + // declaring custom resources it could not possibly satisfy. Anything + // else strands work the day this option ships. + if opts.IsUnbounded() { + return "" + } + + var b strings.Builder + + if opts.ReservedFor != nil { + b.WriteString("\n\t\t\t\t AND id = " + bind(opts.ReservedFor.String())) + } + + // An absent budget key is unconstrained, not zero, so only declared + // dimensions produce a comparison. A key present with the value zero + // is a real constraint and still emits one — that is an exhausted + // worker, which must claim nothing that needs the dimension. + // + // The test is requirement <= budget: a job needing exactly the free + // capacity is claimable, or the last slot on every worker is + // permanently unusable. + for _, dim := range budgetColumns { + budget, declared := opts.Budget[dim.key] + if !declared { + continue } - jobs = append(jobs, j) + + b.WriteString("\n\t\t\t\t AND " + dim.column + " <= " + bind(budget)) } - return jobs, nil + + b.WriteString("\n\t\t\t\t AND " + buildCustomKeyPredicate(opts, bind)) + + return b.String() +} + +// buildCustomKeyPredicate renders custom-resource containment as a +// genuine SUBSET test. +// +// req_custom_keys holds resource.EncodeCustomKeys' output — the sorted +// required keys wrapped in leading and trailing separators, e.g. +// ",fpga,tpu,". The obvious formulation, LIKE '%' || req_custom_keys || +// '%' against the offered list, is a SUBSTRING test: it passes every +// single-key case including the prefix collision, then silently strands a +// job needing {fpga,tpu} from a caller offering {fpga,nvme,tpu}, because +// the interleaved key breaks the contiguous run. The job it strands is +// the specialised one that is hardest to place anywhere else. +// +// Instead each offered key is stripped from the stored list by a nested +// REPLACE of ",key," with ",", which restores the separator the removal +// consumed and so composes in any order. What remains is "" or a lone +// separator exactly when every required key was offered. That is the +// portable formulation SQLite can copy verbatim; string_to_array(...) <@ +// ARRAY[...] would be the Postgres-native alternative, but it has to +// filter the empty elements the wrapping separators produce, and keeping +// the two SQL backends identical is worth more than the array operator. +func buildCustomKeyPredicate(opts job.DequeueOpts, bind func(any) string) string { + sep := bind(resource.CustomKeySep) + offered := opts.OfferedCustomKeys() + + // Bounded opts with an empty offer are a resource-aware worker that + // genuinely has no custom resources, so only jobs requiring none are + // eligible. This is the case IsUnbounded above has already excluded. + expr := "req_custom_keys" + for _, k := range offered { + expr = "REPLACE(" + expr + ", " + bind(resource.CustomKeySep+k+resource.CustomKeySep) + ", " + sep + ")" + } + + return expr + " IN ('', " + sep + ")" +} + +// buildDequeueOrder renders the ordering every backend must return: +// priority descending, then locality-preferred before not, then RunAt +// ascending. +// +// Locality ranks strictly BELOW priority. Above it, a steady stream of +// low-priority jobs whose inputs are already staged would beat a +// high-priority job with cold inputs — an optimization overriding +// user-expressed intent, and the exact starvation the predicate exists to +// prevent. A preferred job jumps its own priority band and no further. +// +// The term is applied whenever PreferHashes is non-empty, including on +// otherwise-unbounded opts: IsUnbounded governs filtering only. +func buildDequeueOrder(opts job.DequeueOpts, bind func(any) string) string { + // PreferredHashes, not PreferHashes: an empty entry would bind as + // '' and match every unhashed job, since primary_input_hash is a + // plain string and an unhashed job stores '' rather than NULL. See + // job.DequeueOpts.PreferredHashes. + hashes := opts.PreferredHashes() + if len(hashes) == 0 { + return "priority DESC, run_at ASC" + } + + // primary_input_hash is nullable — rows written before the resource + // migration have no value — and NULL = ANY(...) is NULL, not false. + // Postgres sorts NULLs FIRST under DESC, so an uncoalesced term would + // rank exactly the rows with no locality signal ABOVE the ones the + // caller has staged. COALESCE makes "unknown" mean "not preferred". + return "priority DESC, COALESCE(primary_input_hash = ANY(" + + bind(hashes) + "), FALSE) DESC, run_at ASC" } // GetJob retrieves a job by ID. @@ -73,18 +286,22 @@ func (s *Store) GetJob(ctx context.Context, jobID id.JobID) (*job.Job, error) { if isNoRows(err) { return nil, dispatch.ErrJobNotFound } - return nil, fmt.Errorf("dispatch/bun: get job: %w", err) + return nil, fmt.Errorf(errPrefix+"get job: %w", err) } return fromJobModel(m) } // UpdateJob persists changes to an existing job. func (s *Store) UpdateJob(ctx context.Context, j *job.Job) error { - m := toJobModel(j) + m, err := toJobModel(j) + if err != nil { + return err + } + m.UpdatedAt = time.Now().UTC() res, err := s.pgdb.NewUpdate(m).WherePK().Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: update job: %w", err) + return fmt.Errorf(errPrefix+"update job: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { @@ -99,7 +316,7 @@ func (s *Store) DeleteJob(ctx context.Context, jobID id.JobID) error { Where("id = ?", jobID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: delete job: %w", err) + return fmt.Errorf(errPrefix+"delete job: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { @@ -129,14 +346,14 @@ func (s *Store) ListJobsByState(ctx context.Context, state job.State, opts job.L err := q.Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: list jobs by state: %w", err) + return nil, fmt.Errorf(errPrefix+"list jobs by state: %w", err) } jobs := make([]*job.Job, 0, len(models)) for i := range models { j, convErr := fromJobModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: list jobs convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"list jobs convert: %w", convErr) } jobs = append(jobs, j) } @@ -151,7 +368,7 @@ func (s *Store) HeartbeatJob(ctx context.Context, jobID id.JobID, _ id.WorkerID) Where("id = ?", jobID.String()). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: heartbeat job: %w", err) + return fmt.Errorf(errPrefix+"heartbeat job: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { @@ -171,14 +388,14 @@ func (s *Store) ReapStaleJobs(ctx context.Context, threshold time.Duration) ([]* Where("COALESCE(heartbeat_at, started_at) < NOW() - ?::interval", threshold.String()). Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: reap stale jobs: %w", err) + return nil, fmt.Errorf(errPrefix+"reap stale jobs: %w", err) } jobs := make([]*job.Job, 0, len(models)) for i := range models { j, convErr := fromJobModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: reap stale convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"reap stale convert: %w", convErr) } jobs = append(jobs, j) } @@ -198,7 +415,7 @@ func (s *Store) CountJobs(ctx context.Context, opts job.CountOpts) (int64, error count, err := q.Count(ctx) if err != nil { - return 0, fmt.Errorf("dispatch/bun: count jobs: %w", err) + return 0, fmt.Errorf(errPrefix+"count jobs: %w", err) } return count, nil } diff --git a/store/postgres/lease.go b/store/postgres/lease.go new file mode 100644 index 0000000..43c9d77 --- /dev/null +++ b/store/postgres/lease.go @@ -0,0 +1,196 @@ +package postgres + +import ( + "context" + "fmt" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// The compile-time check that this store provides the lease capability +// lives in store.go alongside the other interface assertions, matching the +// mongo and sqlite backends. +// +// The grant itself is not here either: it travels on job.DequeueOpts and is +// compiled into DequeueJobs' claim statement by buildLeaseGrant, so a +// leased claim carries the fit predicate and the ordering like any other. +// This file holds only what a lease needs afterwards. + +// RenewLease extends the lease only if the caller still holds it. +func (s *Store) RenewLease( + ctx context.Context, + jobID id.JobID, + workerID id.WorkerID, + epoch int, + leaseUntil time.Time, +) error { + res, err := s.pgdb.NewRaw(` + UPDATE dispatch_jobs + SET lease_expires_at = $1, + heartbeat_at = NOW(), + updated_at = NOW() + WHERE id = $2 + AND state = 'running' + AND worker_id = $3 + AND lease_epoch = $4`, + leaseUntil.UTC(), jobID.String(), workerID.String(), epoch, + ).Exec(ctx) + if err != nil { + return fmt.Errorf(errPrefix+"renew lease: %w", err) + } + + rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil + if rows == 0 { + // Deleted, reclaimed, or reassigned — in every case this worker no + // longer owns the job and must stop. + return job.ErrLeaseLost + } + + return nil +} + +// ReclaimExpiredLeases returns expired-lease jobs to pending, fencing +// their previous holders. +// +// The claim and the read are one statement. The old select-then-update +// reaper let two pools both see the same stale job and both reset it. +func (s *Store) ReclaimExpiredLeases(ctx context.Context, limit int) ([]*job.Job, error) { + // A non-positive limit claims nothing. Postgres would already return + // nothing for LIMIT 0, but a negative LIMIT is a runtime error rather + // than an empty result, and neither is worth a round trip. + if limit <= 0 { + return nil, nil + } + + // The first branch is the actual rule: a lease was granted and has + // lapsed. The second is the narrow exception for a running job carrying + // no lease at all, gated on silence rather than on the null expiry + // alone — see job.UnleasedReclaimGrace for why a null expiry does not + // by itself mean the job was abandoned, and why COALESCE returning NULL + // (neither timestamp set) must not be adopted. + // + // The cutoff is bound rather than computed as NOW() - INTERVAL so that + // all four backends read the same constant from one place. + silent := time.Now().UTC().Add(-job.UnleasedReclaimGrace) + + var models []jobModel + err := s.pgdb.NewRaw(` + WITH expired AS ( + SELECT id FROM dispatch_jobs + WHERE state = 'running' + AND ( (lease_expires_at IS NOT NULL AND lease_expires_at <= NOW()) + OR (lease_expires_at IS NULL + AND COALESCE(heartbeat_at, started_at) IS NOT NULL + AND COALESCE(heartbeat_at, started_at) <= $2) ) + ORDER BY lease_expires_at ASC NULLS FIRST + FOR UPDATE SKIP LOCKED + LIMIT $1 + ) + UPDATE dispatch_jobs + SET state = 'pending', + run_at = NOW(), + worker_id = NULL, + started_at = NULL, + heartbeat_at = NULL, + lease_expires_at = NULL, + lease_epoch = lease_epoch + 1, + evict_count = evict_count + 1, + updated_at = NOW() + WHERE id IN (SELECT id FROM expired) + RETURNING *`, + limit, silent, + ).Scan(ctx, &models) + if err != nil { + return nil, fmt.Errorf(errPrefix+"reclaim expired leases: %w", err) + } + + jobs := make([]*job.Job, 0, len(models)) + for i := range models { + j, convErr := fromJobModel(&models[i]) + if convErr != nil { + return nil, fmt.Errorf(errPrefix+"reclaim convert: %w", convErr) + } + jobs = append(jobs, j) + } + + return jobs, nil +} + +// updateLeasedJobSQL writes every business column UpdateJob writes, +// fenced on the row being still running, still assigned to the caller's +// workerID, and still at the caller's epoch. +// +// lease_epoch, lease_expires_at, worker_id, and heartbeat_at are +// deliberately absent from the SET list — not merely bound to their old +// values, but never mentioned — because j is the caller's stale +// snapshot. Even behind a passing epoch predicate, writing +// j.LeaseExpiresAt back would roll the real expiry backwards by however +// long the caller ran since it last read the row. Those four columns +// have exactly three writers (the grant in DequeueJobs, RenewLease, and +// ReclaimExpiredLeases); this statement is deliberately not a fourth. +const updateLeasedJobSQL = ` + UPDATE dispatch_jobs + SET name = $1, queue = $2, payload = $3, state = $4, priority = $5, + max_retries = $6, retry_count = $7, last_error = $8, + scope_app_id = $9, scope_org_id = $10, run_at = $11, + started_at = $12, completed_at = $13, timeout = $14, + lease_ttl = $15, evict_count = $16, created_at = $17, + updated_at = $18, + req_cpu_milli = $19, req_memory_bytes = $20, req_disk_bytes = $21, + req_gpu_milli = $22, req_custom_keys = $23, + resource_requests = $24, resource_limits = $25, + resource_class = $26, input_bytes = $27, primary_input_hash = $28 + WHERE id = $29 + AND state = 'running' + AND worker_id = $30 + AND lease_epoch = $31` + +// UpdateLeasedJob persists j only while the caller still holds the +// lease. +func (s *Store) UpdateLeasedJob(ctx context.Context, j *job.Job, workerID id.WorkerID, epoch int) error { + m, err := toJobModel(j) + if err != nil { + return err + } + + now := time.Now().UTC() + + res, err := s.pgdb.NewRaw(updateLeasedJobSQL, + m.Name, m.Queue, m.Payload, m.State, m.Priority, + m.MaxRetries, m.RetryCount, m.LastError, + m.ScopeAppID, m.ScopeOrgID, m.RunAt, + m.StartedAt, m.CompletedAt, m.Timeout, + m.LeaseTTL, m.EvictCount, m.CreatedAt, + now, + m.ReqCPUMilli, m.ReqMemoryBytes, m.ReqDiskBytes, + m.ReqGPUMilli, m.ReqCustomKeys, + m.ResourceRequests, m.ResourceLimits, + m.ResourceClass, m.InputBytes, m.PrimaryInputHash, + m.ID, workerID.String(), epoch, + ).Exec(ctx) + if err != nil { + return fmt.Errorf(errPrefix+"update leased job: %w", err) + } + + rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil + if rows > 0 { + return nil + } + + // Zero rows means either the fence predicate failed (the lease moved + // on) or the row is gone. Only the latter is ErrJobNotFound; the + // former is ErrLeaseLost, the entire point of this method. + exists := new(jobModel) + existErr := s.pgdb.NewSelect(exists).Where("id = ?", m.ID).Limit(1).Scan(ctx) + if existErr != nil { + if isNoRows(existErr) { + return dispatch.ErrJobNotFound + } + return fmt.Errorf(errPrefix+"update leased job existence check: %w", existErr) + } + + return job.ErrLeaseLost +} diff --git a/store/postgres/lease_test.go b/store/postgres/lease_test.go new file mode 100644 index 0000000..e857d38 --- /dev/null +++ b/store/postgres/lease_test.go @@ -0,0 +1,29 @@ +package postgres_test + +import ( + "testing" + + "github.com/xraph/dispatch/store/storetest" +) + +func TestLeaseConformance(t *testing.T) { + dsn := startWakePostgres(t) + + storetest.RunLeaseSuite(t, func(t *testing.T) storetest.LeaseStore { + t.Helper() + + return openWakeStore(t, dsn) + }) +} + +func TestDLQConformance(t *testing.T) { + // One container for the suite, like TestLeaseConformance above; the + // cases work only on entries they created, so a shared store is fine. + dsn := startWakePostgres(t) + + storetest.RunDLQSuite(t, func(t *testing.T) storetest.DLQStore { + t.Helper() + + return openWakeStore(t, dsn) + }) +} diff --git a/store/postgres/migrations.go b/store/postgres/migrations.go index 97e5c32..6165ee2 100644 --- a/store/postgres/migrations.go +++ b/store/postgres/migrations.go @@ -325,5 +325,442 @@ func init() { return err }, }, + + // 007: Create artifacts and artifact links tables. + &migrate.Migration{ + Name: "create_artifacts_tables", + Version: "20260811120000", + Up: func(ctx context.Context, exec migrate.Executor) error { + if _, err := exec.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS dispatch_artifacts ( + id TEXT PRIMARY KEY, + backend TEXT NOT NULL, + bucket TEXT NOT NULL, + key TEXT NOT NULL, + size BIGINT NOT NULL DEFAULT 0, + content_hash TEXT, + content_type TEXT, + lifecycle TEXT NOT NULL, + scope_app_id TEXT, + scope_org_id TEXT, + expires_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + deleted_at TIMESTAMPTZ + )`); err != nil { + return err + } + + // Partial unique index rather than a table constraint: a + // purged key must be reusable, so only live rows collide. + if _, err := exec.Exec(ctx, ` + CREATE UNIQUE INDEX IF NOT EXISTS uq_dispatch_artifacts_key + ON dispatch_artifacts (backend, bucket, key) + WHERE deleted_at IS NULL`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifacts_sweep + ON dispatch_artifacts (lifecycle, created_at) + WHERE deleted_at IS NULL`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifacts_purge + ON dispatch_artifacts (deleted_at) + WHERE deleted_at IS NOT NULL`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifacts_hash + ON dispatch_artifacts (content_hash) + WHERE content_hash IS NOT NULL`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS dispatch_artifact_links ( + artifact_id TEXT NOT NULL REFERENCES dispatch_artifacts(id) ON DELETE CASCADE, + owner_kind TEXT NOT NULL, + owner_id TEXT NOT NULL, + name TEXT NOT NULL, + attempt INTEGER NOT NULL DEFAULT 0, + role TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (artifact_id, owner_kind, owner_id, name, attempt) + )`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifact_links_owner + ON dispatch_artifact_links (owner_kind, owner_id)`); err != nil { + return err + } + + _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifact_links_artifact + ON dispatch_artifact_links (artifact_id)`) + + return err + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + if _, err := exec.Exec(ctx, `DROP TABLE IF EXISTS dispatch_artifact_links`); err != nil { + return err + } + + _, err := exec.Exec(ctx, `DROP TABLE IF EXISTS dispatch_artifacts`) + + return err + }, + }, + + // 008: Lease columns. Execution becomes a lease: lease_ttl on the + // row is what lets one reclaim query serve a 30-second job and a + // six-hour one, and lease_epoch fences a worker that was reclaimed + // while it was merely paused. + // + // This migration runs against a LIVE fleet — extension.Start calls + // Migrate by default, so the first upgraded pod executes it while + // every old pod is still enqueueing, claiming and completing on + // dispatch_jobs. See migration 009 below, which hit the identical + // hazards; the same three mechanisms are used here, and 008 runs + // first so it has to be at least as safe. + &migrate.Migration{ + Name: "add_job_lease_columns", + Version: "20260812120000", + Up: func(ctx context.Context, exec migrate.Executor) error { + // Batched into one ALTER and bounded by lock_timeout: a + // pending ACCESS EXCLUSIVE request blocks the whole lock + // queue behind it, so an ALTER that queues behind one + // in-flight SELECT ... FOR UPDATE SKIP LOCKED stalls every + // enqueue and completion in the fleet. All four defaults + // are constants, so on PostgreSQL 11+ this is a catalog + // update with no table rewrite; acquiring the lock is the + // only part that can wait, and that is what the timeout + // bounds. + if err := withLockTimeout(ctx, exec, ` + ALTER TABLE dispatch_jobs + ADD COLUMN IF NOT EXISTS lease_epoch INTEGER NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS lease_expires_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS lease_ttl BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS evict_count INTEGER NOT NULL DEFAULT 0`); err != nil { + return err + } + + // There is deliberately no backfill of lease_expires_at for + // rows already running here. An earlier version of this + // migration seeded them from COALESCE(heartbeat_at, + // started_at, NOW()) so the first sweep would collect them, + // and that was too weak in one direction and too strong in + // the other. Too weak: it only sees rows that exist the + // moment it runs, so during a rolling upgrade every job an + // old pod claims afterwards is stranded exactly as before, + // and so is every job claimed later through job.Store + // without lease options, which is a supported call that + // never grants a lease. Too strong: seeding an expiry in the + // past evicts jobs old pods are still actively running, + // because an old binary's heartbeats do not push an expiry + // it does not know about. + // + // ReclaimExpiredLeases adopts those rows instead, on every + // sweep rather than once, and gated on silence so a worker + // that is still heartbeating is never touched. See + // job.UnleasedReclaimGrace. + + // CONCURRENTLY: a plain CREATE INDEX holds a SHARE lock + // for the whole build, blocking every INSERT, UPDATE and + // DELETE on dispatch_jobs until it finishes. Grove does + // not wrap Up in a transaction — migrate.Orchestrator + // calls m.Up directly and the pg executor runs autocommit + // on a pinned connection — so CONCURRENTLY, which cannot + // run inside one, is available here. Migration 009 relies + // on the same property. + // + // Its cost is that a failed build leaves an INVALID index + // that the planner ignores and IF NOT EXISTS would then + // skip forever, so an invalid leftover is dropped first. + // That is what makes a retry converge instead of + // reporting success over an unusable index. + // + // Built after the backfill so the write lands in the index + // as it is created rather than as a maintenance cost on + // top of it. + if err := dropIfInvalid(ctx, exec, "idx_dispatch_jobs_lease"); err != nil { + return err + } + + _, err := exec.Exec(ctx, ` + CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_dispatch_jobs_lease + ON dispatch_jobs (lease_expires_at) + WHERE state = 'running'`) + + return err + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + _, err := exec.Exec(ctx, + `DROP INDEX CONCURRENTLY IF EXISTS idx_dispatch_jobs_lease`) + if err != nil { + return err + } + + return withLockTimeout(ctx, exec, ` + ALTER TABLE dispatch_jobs + DROP COLUMN IF EXISTS lease_epoch, + DROP COLUMN IF EXISTS lease_expires_at, + DROP COLUMN IF EXISTS lease_ttl, + DROP COLUMN IF EXISTS evict_count`) + }, + }, + + // 009: Resource model columns for resource-aware scheduling. + // + // Every column defaults to zero or empty, so rows written before + // this migration remain dequeueable by every worker during a + // rolling deploy. + // + // This migration runs against a LIVE fleet. extension.Start calls + // Migrate by default, so the first upgraded pod executes it while + // every old pod is still enqueueing, claiming and completing on + // dispatch_jobs — the hottest table in the schema. Both halves are + // written for that: the DDL takes its exclusive lock under a + // timeout rather than queueing behind a long-running claim, and + // the index is built without blocking writes at all. + &migrate.Migration{ + Name: "job_resource_columns", + Version: "20260812130000", + Up: func(ctx context.Context, exec migrate.Executor) error { + // One ALTER, not ten. Each statement takes its own + // ACCESS EXCLUSIVE lock on dispatch_jobs, so ten of them + // are ten separate chances to queue behind an in-flight + // SELECT ... FOR UPDATE SKIP LOCKED — and every enqueue + // and completion in the fleet queues behind THAT, because + // a waiting ACCESS EXCLUSIVE request blocks the lock + // queue ahead of it. Migration 008 already batches its + // four this way. + // + // All ten defaults are constants, so on PostgreSQL 11+ + // this is a catalog update and does not rewrite the + // table; the lock is held for microseconds once acquired. + // Acquiring it is the part that can wait, which is what + // the lock_timeout below bounds. + if err := withLockTimeout(ctx, exec, ` + ALTER TABLE dispatch_jobs + ADD COLUMN IF NOT EXISTS req_cpu_milli BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS req_memory_bytes BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS req_disk_bytes BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS req_gpu_milli BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS req_custom_keys TEXT NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS resource_requests JSONB, + ADD COLUMN IF NOT EXISTS resource_limits JSONB, + ADD COLUMN IF NOT EXISTS resource_class TEXT NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS input_bytes BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS primary_input_hash TEXT`); err != nil { + return err + } + + // Covering index: the dequeue predicate reads all four + // scalars for every candidate row, so including them + // keeps the scan index-only. + // + // CONCURRENTLY because a plain CREATE INDEX takes a SHARE + // lock for the whole build, which blocks every INSERT, + // UPDATE and DELETE on dispatch_jobs — the entire fleet's + // enqueues, claims and completions, for as long as the + // build takes on a production-sized queue. Grove does not + // wrap Up in a transaction (migrate.Orchestrator.Migrate + // calls m.Up directly, and the pg executor runs + // autocommit on a pinned connection), so CONCURRENTLY, + // which cannot run inside one, is available here. + // + // The cost of CONCURRENTLY is that a failed build leaves + // an INVALID index behind, which the planner ignores and + // IF NOT EXISTS would then silently skip forever. So an + // invalid leftover is dropped first, which makes a retry + // after a failed migration converge instead of wedging. + if err := dropIfInvalid(ctx, exec, "idx_dispatch_jobs_dequeue_res"); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_dispatch_jobs_dequeue_res + ON dispatch_jobs (queue, priority DESC, run_at ASC) + INCLUDE (req_cpu_milli, req_memory_bytes, + req_disk_bytes, req_gpu_milli) + WHERE state IN ('pending', 'retrying')`); err != nil { + return err + } + + // idx_dispatch_jobs_dequeue (migration 001) has the + // IDENTICAL key list and the IDENTICAL partial predicate; + // the index just created differs only by an INCLUDE + // payload, which makes it a strict superset. Keeping both + // costs a second B-tree insert on every enqueue and a + // second delete on every claim, forever, on the hottest + // table in the schema — for a plan the planner would + // never choose. + // + // Dropped AFTER the replacement is valid, so there is no + // instant at which the dequeue statement has no index. + _, err := exec.Exec(ctx, + `DROP INDEX CONCURRENTLY IF EXISTS idx_dispatch_jobs_dequeue`) + + return err + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + // Restore 001's index before dropping its superset, same + // ordering rule in reverse: never leave the dequeue + // statement unindexed. + if err := dropIfInvalid(ctx, exec, "idx_dispatch_jobs_dequeue"); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_dispatch_jobs_dequeue + ON dispatch_jobs (queue, priority DESC, run_at ASC) + WHERE state IN ('pending', 'retrying')`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, + `DROP INDEX CONCURRENTLY IF EXISTS idx_dispatch_jobs_dequeue_res`); err != nil { + return err + } + + return withLockTimeout(ctx, exec, ` + ALTER TABLE dispatch_jobs + DROP COLUMN IF EXISTS req_cpu_milli, + DROP COLUMN IF EXISTS req_memory_bytes, + DROP COLUMN IF EXISTS req_disk_bytes, + DROP COLUMN IF EXISTS req_gpu_milli, + DROP COLUMN IF EXISTS req_custom_keys, + DROP COLUMN IF EXISTS resource_requests, + DROP COLUMN IF EXISTS resource_limits, + DROP COLUMN IF EXISTS resource_class, + DROP COLUMN IF EXISTS input_bytes, + DROP COLUMN IF EXISTS primary_input_hash`) + }, + }, + // Replay rebuilds a job from the DLQ row and enqueues it directly, + // so anything the row does not carry the replayed job silently takes + // a default for. Losing lease_ttl is the worst of them: a six-hour + // job replayed on the pool default lease lapses mid-run and is + // reclaimed and restarted forever without finishing. See dlq.Entry. + &migrate.Migration{ + Name: "dlq_job_execution_columns", + Version: "20260817120000", + Up: func(ctx context.Context, exec migrate.Executor) error { + // One batched ALTER under a lock timeout, for the reason + // spelled out on job_resource_columns: each statement takes + // its own ACCESS EXCLUSIVE lock, and a waiting request for + // one blocks every other writer queued behind it. Every + // default here is a constant, so this is a catalog update + // rather than a table rewrite. + // + // dispatch_dlq is written once per dead job and read by the + // dashboard, so it is far less contended than dispatch_jobs. + // The timeout is kept anyway: cheap when uncontended, and the + // failure it prevents is a deploy that never finishes. + return withLockTimeout(ctx, exec, ` + ALTER TABLE dispatch_dlq + ADD COLUMN IF NOT EXISTS priority INT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS timeout BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS lease_ttl BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS artifact_bindings BYTEA, + ADD COLUMN IF NOT EXISTS resources JSONB, + ADD COLUMN IF NOT EXISTS resource_limits JSONB, + ADD COLUMN IF NOT EXISTS resource_class TEXT NOT NULL DEFAULT '', + ADD COLUMN IF NOT EXISTS input_bytes BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS primary_input_hash TEXT`) + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + return withLockTimeout(ctx, exec, ` + ALTER TABLE dispatch_dlq + DROP COLUMN IF EXISTS priority, + DROP COLUMN IF EXISTS timeout, + DROP COLUMN IF EXISTS lease_ttl, + DROP COLUMN IF EXISTS artifact_bindings, + DROP COLUMN IF EXISTS resources, + DROP COLUMN IF EXISTS resource_limits, + DROP COLUMN IF EXISTS resource_class, + DROP COLUMN IF EXISTS input_bytes, + DROP COLUMN IF EXISTS primary_input_hash`) + }, + }, ) } + +// ddlLockTimeout bounds how long a DDL statement waits for its ACCESS +// EXCLUSIVE lock on a table the fleet is actively writing. +// +// Without it the ALTER waits indefinitely behind whatever claim happens +// to hold a row lock, AND — because a pending ACCESS EXCLUSIVE request +// blocks everything queued behind it — every enqueue and completion in +// the fleet waits behind the ALTER. A migration that cannot get the lock +// promptly must fail and be retried, not convert one slow query into a +// fleet-wide stall. +// +// Three seconds is long enough to win an uncontended queue comfortably +// and short enough that a failed attempt is a blip rather than an +// outage. Migrate is idempotent here, so a retry is free. +const ddlLockTimeout = "3s" + +// withLockTimeout runs one DDL statement under ddlLockTimeout. +// +// SET rather than SET LOCAL because grove runs Up outside a transaction, +// where SET LOCAL is a no-op with a warning. The pg executor pins one +// connection for the whole migration run (see pgmigrate.Executor), so +// the setting would otherwise leak into every migration after this one — +// hence the reset, which runs even when the statement failed. +func withLockTimeout(ctx context.Context, exec migrate.Executor, stmt string) error { + if _, err := exec.Exec(ctx, `SET lock_timeout = '`+ddlLockTimeout+`'`); err != nil { + return err + } + + _, execErr := exec.Exec(ctx, stmt) + + if _, err := exec.Exec(ctx, `SET lock_timeout = DEFAULT`); err != nil && execErr == nil { + return err + } + + return execErr +} + +// dropIfInvalid removes an index left INVALID by a CREATE INDEX +// CONCURRENTLY that failed partway. +// +// Such an index exists in the catalog but is ignored by the planner, so +// CREATE INDEX CONCURRENTLY IF NOT EXISTS sees it, does nothing, and the +// table permanently has no usable index while every subsequent migration +// run reports success. Dropping it first is what makes a retry after a +// failed migration converge. +func dropIfInvalid(ctx context.Context, exec migrate.Executor, name string) error { + rows, err := exec.Query(ctx, ` + SELECT 1 + FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + WHERE c.relname = $1 AND NOT i.indisvalid`, name) + if err != nil { + return err + } + + invalid := rows.Next() + + if closeErr := rows.Close(); closeErr != nil { + return closeErr + } + + if rowsErr := rows.Err(); rowsErr != nil { + return rowsErr + } + + if !invalid { + return nil + } + + _, err = exec.Exec(ctx, `DROP INDEX CONCURRENTLY IF EXISTS `+name) + + return err +} diff --git a/store/postgres/migrations_test.go b/store/postgres/migrations_test.go new file mode 100644 index 0000000..bfff5dc --- /dev/null +++ b/store/postgres/migrations_test.go @@ -0,0 +1,376 @@ +//go:build integration + +package postgres_test + +import ( + "context" + "testing" + "time" + + log "github.com/xraph/go-utils/log" + "github.com/xraph/grove/driver" + "github.com/xraph/grove/drivers/pgdriver" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/postgres" + "github.com/xraph/dispatch/store/storetest" +) + +// leaseMigrationVersion is the version string of the lease migration, +// restated here so a test can delete its bookkeeping row and force the +// re-run every restarting pod performs after a failure. +const leaseMigrationVersion = "20260812120000" + +// remigrate re-runs the migration group against the same database, which +// is what a pod does on every start. +func remigrate(t *testing.T, s *postgres.Store) { + t.Helper() + + if err := postgres.New(s.DB(), postgres.WithLogger(log.NewNoopLogger())). + Migrate(context.Background()); err != nil { + t.Fatalf("re-running the migration group must succeed, got: %v", err) + } +} + +// forgetLeaseMigration deletes the lease migration's grove_migrations row, +// reproducing a crash between Up and RecordApplied — a real window, +// because migrate.Orchestrator calls them as two separate steps. +func forgetLeaseMigration(t *testing.T, conn driver.DedicatedConn) { + t.Helper() + + if _, err := conn.Exec(context.Background(), + `DELETE FROM grove_migrations WHERE version = $1`, leaseMigrationVersion); err != nil { + t.Fatalf("delete migration row: %v", err) + } +} + +// indexIsValid reports whether the named index exists and is usable. +// +// The two are different states and the difference is the whole hazard of +// CREATE INDEX CONCURRENTLY: a build that fails leaves the index present +// in the catalog but INVALID, which the planner ignores and IF NOT EXISTS +// would then skip forever. +func indexIsValid(t *testing.T, conn driver.DedicatedConn, name string) (exists, valid bool) { + t.Helper() + + rows, err := conn.Query(context.Background(), ` + SELECT i.indisvalid + FROM pg_class c JOIN pg_index i ON i.indexrelid = c.oid + WHERE c.relname = $1`, name) + if err != nil { + t.Fatalf("read indisvalid for %s: %v", name, err) + } + + defer rows.Close() + + if !rows.Next() { + return false, false + } + + if err = rows.Scan(&valid); err != nil { + t.Fatalf("scan indisvalid: %v", err) + } + + if err = rows.Err(); err != nil { + t.Fatalf("iterate indisvalid: %v", err) + } + + return true, valid +} + +// TestLeaseMigrationBuildsAValidIndex pins that the lease index survives +// a normal run as a USABLE index. +// +// It is not a tautology. The migration builds it CONCURRENTLY, because a +// plain CREATE INDEX holds a SHARE lock for the whole build and blocks +// every INSERT, UPDATE and DELETE on dispatch_jobs while the rest of the +// fleet is still enqueueing and completing on it. The price is that a +// failed CONCURRENTLY build leaves an INVALID index that nothing reports: +// migrations keep succeeding, the planner keeps ignoring it, and the +// reclaim sweep degrades to a sequential scan of the whole table. +func TestLeaseMigrationBuildsAValidIndex(t *testing.T) { + s := setupTestStore(t) + + conn, err := pgdriver.Unwrap(s.DB()).AcquireConn(context.Background()) + if err != nil { + t.Fatalf("acquire dedicated conn: %v", err) + } + + defer conn.Release() + + exists, valid := indexIsValid(t, conn, "idx_dispatch_jobs_lease") + if !exists { + t.Fatal("idx_dispatch_jobs_lease is missing after a clean migration") + } + + if !valid { + t.Error("idx_dispatch_jobs_lease is INVALID: a failed CONCURRENTLY build was left " + + "in place, so the reclaim sweep has no usable index and nothing reports it") + } +} + +// TestLeaseMigrationConvergesFromAnInvalidIndex is the reason the +// migration drops an invalid leftover before building. +// +// CREATE INDEX CONCURRENTLY IF NOT EXISTS sees an INVALID index, decides +// there is nothing to do, and returns success. Without the pre-drop the +// table would then be permanently without a usable lease index while +// every subsequent migration run reported that everything was fine. +func TestLeaseMigrationConvergesFromAnInvalidIndex(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + conn, err := pgdriver.Unwrap(s.DB()).AcquireConn(ctx) + if err != nil { + t.Fatalf("acquire dedicated conn: %v", err) + } + + defer conn.Release() + + // The catalog state a failed CONCURRENTLY build leaves behind. There + // is no way to reach it through DDL, so it is written directly. + if _, err = conn.Exec(ctx, ` + UPDATE pg_index SET indisvalid = false + WHERE indexrelid = 'idx_dispatch_jobs_lease'::regclass`); err != nil { + t.Fatalf("mark the index invalid: %v", err) + } + + if _, valid := indexIsValid(t, conn, "idx_dispatch_jobs_lease"); valid { + t.Fatal("fixture is wrong: the index should be INVALID") + } + + forgetLeaseMigration(t, conn) + remigrate(t, s) + + exists, valid := indexIsValid(t, conn, "idx_dispatch_jobs_lease") + if !exists { + t.Fatal("idx_dispatch_jobs_lease is missing after the retry") + } + + if !valid { + t.Error("the retry left the index INVALID: CREATE INDEX CONCURRENTLY IF NOT EXISTS " + + "skipped the unusable leftover instead of replacing it, and reported success") + } +} + +// TestLeaseMigrationSurvivesAPartialApplication is the postgres half of +// the re-runnability proof. +// +// Grove executes Up outside any transaction, so a process killed partway +// through leaves some of the change applied and no row in +// grove_migrations. Every statement in Up is written to be re-runnable +// from any point for exactly that reason; this asserts it rather than +// trusting the IF NOT EXISTS clauses by inspection. +func TestLeaseMigrationSurvivesAPartialApplication(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + conn, err := pgdriver.Unwrap(s.DB()).AcquireConn(ctx) + if err != nil { + t.Fatalf("acquire dedicated conn: %v", err) + } + + defer conn.Release() + + // What a crash after the ALTER's first two columns would have left. + for _, stmt := range []string{ + `ALTER TABLE dispatch_jobs DROP COLUMN lease_ttl`, + `ALTER TABLE dispatch_jobs DROP COLUMN evict_count`, + `DROP INDEX IF EXISTS idx_dispatch_jobs_lease`, + } { + if _, err = conn.Exec(ctx, stmt); err != nil { + t.Fatalf("%s: %v", stmt, err) + } + } + + forgetLeaseMigration(t, conn) + remigrate(t, s) + + for _, col := range []string{ + "lease_epoch", "lease_expires_at", "lease_ttl", "evict_count", + } { + var present bool + + if err = conn.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'dispatch_jobs' AND column_name = $1)`, + col).Scan(&present); err != nil { + t.Fatalf("look up %s: %v", col, err) + } + + if !present { + t.Errorf("column %s missing after the retry", col) + } + } + + // And the schema is usable, not merely present. + j := storetestPendingJob("after-lease-retry", "lease-retry", 6*time.Hour) + if err = s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob after the retry: %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob after the retry: %v", err) + } + + if got.LeaseTTL != 6*time.Hour { + t.Errorf("LeaseTTL = %v, want the declaration intact", got.LeaseTTL) + } +} + +// TestReclaimAdoptsRunningJobsWithoutLease covers a running job carrying +// no lease at all, which is the row shape that used to be stranded +// permanently. lease_expires_at is NULL, ReclaimExpiredLeases required a +// non-NULL expiry to consider a row at all (job.Lease.IsExpired +// deliberately reads a zero expiry as "never leased", not "expired"), the +// pool's reaper no longer calls ReapStaleJobs once the backend implements +// job.LeaseStore, and dequeue claims only pending and retrying rows. So +// nothing looked at these again and they held their slots forever. +// +// Migration 008 used to seed an expiry for them. It no longer does, for +// the reasons given at that migration: a one-shot backfill misses every +// row created after it runs, including anything an old pod claims later in +// a rolling upgrade and anything claimed through job.Store without lease +// options, which never grants a lease at all. Reclamation adopts them +// instead, on every sweep and gated on silence. +// +// The negative cases are the ones that matter. A NULL expiry does not by +// itself mean the job was abandoned, so evicting on that alone would take +// live work away from a healthy caller. +func TestReclaimAdoptsRunningJobsWithoutLease(t *testing.T) { + silent := time.Now().UTC().Add(-job.UnleasedReclaimGrace - time.Minute) + older := silent.Add(-time.Hour) + fresh := time.Now().UTC() + + tests := []struct { + name string + heartbeat time.Time + started time.Time + want bool + why string + }{ + { + name: "silent heartbeat is adopted", + heartbeat: silent, + started: older, + want: true, + why: "abandoned by a worker that stopped reporting", + }, + { + name: "fresh heartbeat is left alone", + heartbeat: fresh, + started: older, + want: false, + why: "still reporting, so it belongs to a healthy worker", + }, + { + name: "silent started_at is adopted when it never heartbeated", + started: silent, + want: true, + why: "died before its first beat", + }, + { + name: "fresh started_at is left alone", + started: fresh, + want: false, + why: "just claimed; its first heartbeat is not due yet", + }, + { + name: "neither timestamp is left alone", + want: false, + why: "no timestamp to establish age from", + }, + } + + // One container for every case, not one each. Each case asserts on the + // job it created rather than on the size of the result, so a shared + // database is safe, and five containers back to back is both slow and + // a needless chance for one of them to fail to come up. + dsn := startWakePostgres(t) + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := openWakeStore(t, dsn) + ctx := context.Background() + + conn, err := pgdriver.Unwrap(s.DB()).AcquireConn(ctx) + if err != nil { + t.Fatalf("acquire dedicated conn: %v", err) + } + + defer conn.Release() + + j := storetestPendingJob("mid-flight", "adopt", 0) + if err = s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob: %v", err) + } + + // The row an upgrading fleet finds, or that a caller without + // lease options writes: running, with no lease at all. + if _, err = conn.Exec(ctx, ` + UPDATE dispatch_jobs + SET state = 'running', heartbeat_at = $1, started_at = $2, + lease_expires_at = NULL + WHERE id = $3`, + nullableTime(tt.heartbeat), nullableTime(tt.started), j.ID.String()); err != nil { + t.Fatalf("rewind the row: %v", err) + } + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 10) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + + // Membership, not length: the store is shared across cases, + // and reclamation is not queue-scoped. + got := storetest.Contains(reclaimed, j.ID) + if got != tt.want { + if tt.want { + t.Fatalf("job was not reclaimed but should have been: %s", tt.why) + } + + t.Fatalf("job was reclaimed but must not be: %s", tt.why) + } + if tt.want { + for _, r := range reclaimed { + if r.ID == j.ID && r.State != job.StatePending { + t.Errorf("reclaimed job state = %v, want pending", r.State) + } + } + } + }) + } +} + +// nullableTime renders the zero time as a NULL bind rather than as year +// one, so a test case can express "this column was never written". +func nullableTime(t time.Time) any { + if t.IsZero() { + return nil + } + + return t +} + +// storetestPendingJob builds a pending job directly rather than through +// storetest.PendingJob, so these tests do not depend on the conformance +// suite's fixture shape. +func storetestPendingJob(name, queue string, ttl time.Duration) *job.Job { + now := time.Now().UTC() + + return &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: name, + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: now.Add(-time.Second), + LeaseTTL: ttl, + } +} diff --git a/store/postgres/models.go b/store/postgres/models.go index d422cb1..bd4ba9c 100644 --- a/store/postgres/models.go +++ b/store/postgres/models.go @@ -13,6 +13,7 @@ import ( "github.com/xraph/dispatch/event" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" "github.com/xraph/dispatch/workflow" ) @@ -21,55 +22,115 @@ import ( type jobModel struct { grove.BaseModel `grove:"table:dispatch_jobs"` - ID string `grove:"id,pk"` - Name string `grove:"name,notnull"` - Queue string `grove:"queue,notnull,default:'default'"` - Payload []byte `grove:"payload,notnull,type:bytea"` - State string `grove:"state,notnull,default:'pending'"` - Priority int `grove:"priority,notnull,default:0"` - MaxRetries int `grove:"max_retries,notnull,default:3"` - RetryCount int `grove:"retry_count,notnull,default:0"` - LastError string `grove:"last_error"` - ScopeAppID string `grove:"scope_app_id"` - ScopeOrgID string `grove:"scope_org_id"` - WorkerID string `grove:"worker_id"` - RunAt time.Time `grove:"run_at,notnull,default:current_timestamp"` - StartedAt *time.Time `grove:"started_at"` - CompletedAt *time.Time `grove:"completed_at"` - HeartbeatAt *time.Time `grove:"heartbeat_at"` - Timeout int64 `grove:"timeout,notnull,default:0"` - CreatedAt time.Time `grove:"created_at,notnull,default:current_timestamp"` - UpdatedAt time.Time `grove:"updated_at,notnull,default:current_timestamp"` + ID string `grove:"id,pk"` + Name string `grove:"name,notnull"` + Queue string `grove:"queue,notnull,default:'default'"` + Payload []byte `grove:"payload,notnull,type:bytea"` + State string `grove:"state,notnull,default:'pending'"` + Priority int `grove:"priority,notnull,default:0"` + MaxRetries int `grove:"max_retries,notnull,default:3"` + RetryCount int `grove:"retry_count,notnull,default:0"` + LastError string `grove:"last_error"` + ScopeAppID string `grove:"scope_app_id"` + ScopeOrgID string `grove:"scope_org_id"` + WorkerID string `grove:"worker_id"` + RunAt time.Time `grove:"run_at,notnull,default:current_timestamp"` + StartedAt *time.Time `grove:"started_at"` + CompletedAt *time.Time `grove:"completed_at"` + HeartbeatAt *time.Time `grove:"heartbeat_at"` + Timeout int64 `grove:"timeout,notnull,default:0"` + LeaseEpoch int `grove:"lease_epoch,notnull,default:0"` + LeaseExpiresAt *time.Time `grove:"lease_expires_at"` + LeaseTTL int64 `grove:"lease_ttl,notnull,default:0"` + EvictCount int `grove:"evict_count,notnull,default:0"` + CreatedAt time.Time `grove:"created_at,notnull,default:current_timestamp"` + UpdatedAt time.Time `grove:"updated_at,notnull,default:current_timestamp"` + + // The four canonical dimensions get real scalar columns because the + // dequeue predicate compares them and must behave identically across + // five backends; JSON comparison semantics are not portable. They are + // derived from Resources by toJobModel — the caller never sets them + // directly. + ReqCPUMilli int64 `grove:"req_cpu_milli,notnull,default:0"` + ReqMemoryBytes int64 `grove:"req_memory_bytes,notnull,default:0"` + ReqDiskBytes int64 `grove:"req_disk_bytes,notnull,default:0"` + ReqGPUMilli int64 `grove:"req_gpu_milli,notnull,default:0"` + ReqCustomKeys string `grove:"req_custom_keys,notnull,default:''"` + + // ResourceRequests and ResourceLimits are the full-fidelity JSON copy + // of Resources / ResourceLimits, including custom keys the scalar + // columns above do not carry. fromJobModel reads Resources back from + // here, not from the scalars. + ResourceRequests []byte `grove:"resource_requests,type:jsonb"` + ResourceLimits []byte `grove:"resource_limits,type:jsonb"` + ResourceClass string `grove:"resource_class,notnull,default:''"` + InputBytes int64 `grove:"input_bytes,notnull,default:0"` + PrimaryInputHash string `grove:"primary_input_hash"` } -func toJobModel(j *job.Job) *jobModel { - return &jobModel{ - ID: j.ID.String(), - Name: j.Name, - Queue: j.Queue, - Payload: j.Payload, - State: string(j.State), - Priority: j.Priority, - MaxRetries: j.MaxRetries, - RetryCount: j.RetryCount, - LastError: j.LastError, - ScopeAppID: j.ScopeAppID, - ScopeOrgID: j.ScopeOrgID, - WorkerID: j.WorkerID.String(), - RunAt: j.RunAt, - StartedAt: j.StartedAt, - CompletedAt: j.CompletedAt, - HeartbeatAt: j.HeartbeatAt, - Timeout: j.Timeout.Nanoseconds(), - CreatedAt: j.CreatedAt, - UpdatedAt: j.UpdatedAt, +func toJobModel(j *job.Job) (*jobModel, error) { + reqJSON, err := resource.EncodeSet(j.Resources) + if err != nil { + return nil, fmt.Errorf(errPrefix+"marshal job resources: %w", err) + } + + limitsJSON, err := resource.EncodeSet(j.ResourceLimits) + if err != nil { + return nil, fmt.Errorf(errPrefix+"marshal job resource limits: %w", err) } + + return &jobModel{ + ID: j.ID.String(), + Name: j.Name, + Queue: j.Queue, + Payload: j.Payload, + State: string(j.State), + Priority: j.Priority, + MaxRetries: j.MaxRetries, + RetryCount: j.RetryCount, + LastError: j.LastError, + ScopeAppID: j.ScopeAppID, + ScopeOrgID: j.ScopeOrgID, + WorkerID: j.WorkerID.String(), + RunAt: j.RunAt, + StartedAt: j.StartedAt, + CompletedAt: j.CompletedAt, + HeartbeatAt: j.HeartbeatAt, + Timeout: j.Timeout.Nanoseconds(), + LeaseEpoch: j.LeaseEpoch, + LeaseExpiresAt: j.LeaseExpiresAt, + LeaseTTL: j.LeaseTTL.Nanoseconds(), + EvictCount: j.EvictCount, + CreatedAt: j.CreatedAt, + UpdatedAt: j.UpdatedAt, + + ReqCPUMilli: j.Resources[resource.CPU], + ReqMemoryBytes: j.Resources[resource.Memory], + ReqDiskBytes: j.Resources[resource.Disk], + ReqGPUMilli: j.Resources[resource.GPU], + ReqCustomKeys: resource.EncodeCustomKeys(j.Resources), + ResourceRequests: reqJSON, + ResourceLimits: limitsJSON, + ResourceClass: j.ResourceClass, + InputBytes: j.InputBytes, + PrimaryInputHash: j.PrimaryInputHash, + }, nil } func fromJobModel(m *jobModel) (*job.Job, error) { parsedID, err := id.ParseJobID(m.ID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse job id %q: %w", m.ID, err) + return nil, fmt.Errorf(errPrefix+"parse job id %q: %w", m.ID, err) + } + + resources, err := resource.DecodeSet(m.ResourceRequests) + if err != nil { + return nil, fmt.Errorf(errPrefix+"unmarshal job resources: %w", err) + } + + limits, err := resource.DecodeSet(m.ResourceLimits) + if err != nil { + return nil, fmt.Errorf(errPrefix+"unmarshal job resource limits: %w", err) } j := &job.Job{ @@ -77,22 +138,32 @@ func fromJobModel(m *jobModel) (*job.Job, error) { CreatedAt: m.CreatedAt, UpdatedAt: m.UpdatedAt, }, - ID: parsedID, - Name: m.Name, - Queue: m.Queue, - Payload: m.Payload, - State: job.State(m.State), - Priority: m.Priority, - MaxRetries: m.MaxRetries, - RetryCount: m.RetryCount, - LastError: m.LastError, - ScopeAppID: m.ScopeAppID, - ScopeOrgID: m.ScopeOrgID, - RunAt: m.RunAt, - StartedAt: m.StartedAt, - CompletedAt: m.CompletedAt, - HeartbeatAt: m.HeartbeatAt, - Timeout: time.Duration(m.Timeout), + ID: parsedID, + Name: m.Name, + Queue: m.Queue, + Payload: m.Payload, + State: job.State(m.State), + Priority: m.Priority, + MaxRetries: m.MaxRetries, + RetryCount: m.RetryCount, + LastError: m.LastError, + ScopeAppID: m.ScopeAppID, + ScopeOrgID: m.ScopeOrgID, + RunAt: m.RunAt, + StartedAt: m.StartedAt, + CompletedAt: m.CompletedAt, + HeartbeatAt: m.HeartbeatAt, + Timeout: time.Duration(m.Timeout), + LeaseEpoch: m.LeaseEpoch, + LeaseExpiresAt: m.LeaseExpiresAt, + LeaseTTL: time.Duration(m.LeaseTTL), + EvictCount: m.EvictCount, + + Resources: resources, + ResourceLimits: limits, + ResourceClass: m.ResourceClass, + InputBytes: m.InputBytes, + PrimaryInputHash: m.PrimaryInputHash, } if m.WorkerID != "" { @@ -144,7 +215,7 @@ func toRunModel(r *workflow.Run) *workflowRunModel { func fromRunModel(m *workflowRunModel) (*workflow.Run, error) { parsedID, err := id.ParseRunID(m.ID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse run id %q: %w", m.ID, err) + return nil, fmt.Errorf(errPrefix+"parse run id %q: %w", m.ID, err) } return &workflow.Run{ @@ -180,12 +251,12 @@ type checkpointModel struct { func fromCheckpointModel(m *checkpointModel) (*workflow.Checkpoint, error) { parsedID, err := id.ParseCheckpointID(m.ID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse checkpoint id %q: %w", m.ID, err) + return nil, fmt.Errorf(errPrefix+"parse checkpoint id %q: %w", m.ID, err) } parsedRunID, err := id.ParseRunID(m.RunID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse run id %q: %w", m.RunID, err) + return nil, fmt.Errorf(errPrefix+"parse run id %q: %w", m.RunID, err) } return &workflow.Checkpoint{ @@ -245,7 +316,7 @@ func toCronModel(e *cron.Entry) *cronEntryModel { func fromCronModel(m *cronEntryModel) (*cron.Entry, error) { parsedID, err := id.ParseCronID(m.ID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse cron id %q: %w", m.ID, err) + return nil, fmt.Errorf(errPrefix+"parse cron id %q: %w", m.ID, err) } e := &cron.Entry{ @@ -290,9 +361,33 @@ type dlqEntryModel struct { FailedAt time.Time `grove:"failed_at,notnull,default:current_timestamp"` ReplayedAt *time.Time `grove:"replayed_at"` CreatedAt time.Time `grove:"created_at,notnull,default:current_timestamp"` + + // Carried so Replay can rebuild a job that behaves like the failed + // one; see the dlq.Entry doc. The two resource sets are stored with + // the same jsonb codec jobModel uses rather than as scalar columns: + // nothing queries a DLQ row by resource requirement, so the scalar + // split that dequeue's fit predicate needs buys nothing here. + Priority int `grove:"priority,notnull,default:0"` + Timeout int64 `grove:"timeout,notnull,default:0"` + LeaseTTL int64 `grove:"lease_ttl,notnull,default:0"` + ArtifactBindings []byte `grove:"artifact_bindings,type:bytea"` + Resources []byte `grove:"resources,type:jsonb"` + ResourceLimits []byte `grove:"resource_limits,type:jsonb"` + ResourceClass string `grove:"resource_class,notnull,default:''"` + InputBytes int64 `grove:"input_bytes,notnull,default:0"` + PrimaryInputHash string `grove:"primary_input_hash"` } -func toDLQModel(e *dlq.Entry) *dlqEntryModel { +func toDLQModel(e *dlq.Entry) (*dlqEntryModel, error) { + resources, err := resource.EncodeSet(e.Resources) + if err != nil { + return nil, fmt.Errorf(errPrefix+"encode dlq resources: %w", err) + } + limits, err := resource.EncodeSet(e.ResourceLimits) + if err != nil { + return nil, fmt.Errorf(errPrefix+"encode dlq resource limits: %w", err) + } + return &dlqEntryModel{ ID: e.ID.String(), JobID: e.JobID.String(), @@ -307,18 +402,37 @@ func toDLQModel(e *dlq.Entry) *dlqEntryModel { FailedAt: e.FailedAt, ReplayedAt: e.ReplayedAt, CreatedAt: e.CreatedAt, - } + + Priority: e.Priority, + Timeout: int64(e.Timeout), + LeaseTTL: int64(e.LeaseTTL), + ArtifactBindings: e.ArtifactBindings, + Resources: resources, + ResourceLimits: limits, + ResourceClass: e.ResourceClass, + InputBytes: e.InputBytes, + PrimaryInputHash: e.PrimaryInputHash, + }, nil } func fromDLQModel(m *dlqEntryModel) (*dlq.Entry, error) { parsedID, err := id.ParseDLQID(m.ID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse dlq id %q: %w", m.ID, err) + return nil, fmt.Errorf(errPrefix+"parse dlq id %q: %w", m.ID, err) } parsedJobID, err := id.ParseJobID(m.JobID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse job id %q: %w", m.JobID, err) + return nil, fmt.Errorf(errPrefix+"parse job id %q: %w", m.JobID, err) + } + + resources, err := resource.DecodeSet(m.Resources) + if err != nil { + return nil, fmt.Errorf(errPrefix+"decode dlq resources: %w", err) + } + limits, err := resource.DecodeSet(m.ResourceLimits) + if err != nil { + return nil, fmt.Errorf(errPrefix+"decode dlq resource limits: %w", err) } return &dlq.Entry{ @@ -335,6 +449,16 @@ func fromDLQModel(m *dlqEntryModel) (*dlq.Entry, error) { FailedAt: m.FailedAt, ReplayedAt: m.ReplayedAt, CreatedAt: m.CreatedAt, + + Priority: m.Priority, + Timeout: time.Duration(m.Timeout), + LeaseTTL: time.Duration(m.LeaseTTL), + ArtifactBindings: m.ArtifactBindings, + Resources: resources, + ResourceLimits: limits, + ResourceClass: m.ResourceClass, + InputBytes: m.InputBytes, + PrimaryInputHash: m.PrimaryInputHash, }, nil } @@ -367,7 +491,7 @@ func toEventModel(evt *event.Event) *eventModel { func fromEventModel(m *eventModel) (*event.Event, error) { parsedID, err := id.ParseEventID(m.ID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse event id %q: %w", m.ID, err) + return nil, fmt.Errorf(errPrefix+"parse event id %q: %w", m.ID, err) } return &event.Event{ @@ -416,7 +540,7 @@ func toWorkerModel(w *cluster.Worker) *workerModel { func fromWorkerModel(m *workerModel) (*cluster.Worker, error) { parsedID, err := id.ParseWorkerID(m.ID) if err != nil { - return nil, fmt.Errorf("dispatch/bun: parse worker id %q: %w", m.ID, err) + return nil, fmt.Errorf(errPrefix+"parse worker id %q: %w", m.ID, err) } return &cluster.Worker{ diff --git a/store/postgres/resource_test.go b/store/postgres/resource_test.go new file mode 100644 index 0000000..7afc7fe --- /dev/null +++ b/store/postgres/resource_test.go @@ -0,0 +1,138 @@ +//go:build integration + +package postgres_test + +import ( + "context" + "testing" + + "github.com/xraph/grove/drivers/pgdriver" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +func TestPostgresRoundTripsResources(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "tessellate.model", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + Resources: resource.Set{ + resource.CPU: 4000, resource.Memory: 16 << 30, "fpga": 2, + }, + ResourceLimits: resource.Set{resource.Memory: 16 << 30}, + ResourceClass: "heavy", + InputBytes: 4 << 30, + PrimaryInputHash: "blake3:9f2a", + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + + if got.Resources[resource.CPU] != 4000 { + t.Errorf("cpu = %d, want 4000", got.Resources[resource.CPU]) + } + if got.Resources["fpga"] != 2 { + t.Errorf("custom key lost in round-trip: %v", got.Resources) + } + if got.ResourceLimits[resource.Memory] != 16<<30 { + t.Errorf("limits = %v", got.ResourceLimits) + } + if got.ResourceClass != "heavy" { + t.Errorf("class = %q, want heavy", got.ResourceClass) + } + if got.InputBytes != 4<<30 || got.PrimaryInputHash != "blake3:9f2a" { + t.Errorf("input signal lost: bytes=%d hash=%q", + got.InputBytes, got.PrimaryInputHash) + } +} + +// TestPostgresJobWithNoResourcesRoundTrips pins the backward-compatibility +// contract, and it asserts the STORED value rather than only the decoded +// one. +// +// Resources.IsZero() is true for a nil Set, for Set{}, and for +// Set{"memory": 0} alike, so a decoded-side assertion cannot tell an +// undeclared job from one that stored an empty JSON document. The +// difference is not cosmetic: NULL is what a worker running the +// pre-resource code reads as "no requirement", and '{}' is what a +// half-migrated fleet would have to interpret. Redis and Mongo already +// assert on the raw stored value; this is the SQL analogue, and it is +// the column NULL itself that a rolling deploy depends on. +func TestPostgresJobWithNoResourcesRoundTrips(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "notify.user", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if !got.Resources.IsZero() { + t.Errorf("Resources = %v, want zero", got.Resources) + } + + var ( + requestsNull bool + limitsNull bool + cpu int64 + customKeys string + class string + inputBytes int64 + ) + + if err = pgdriver.Unwrap(s.DB()).QueryRow(ctx, ` + SELECT resource_requests IS NULL, + resource_limits IS NULL, + req_cpu_milli, + req_custom_keys, + resource_class, + input_bytes + FROM dispatch_jobs WHERE id = $1`, j.ID.String()). + Scan(&requestsNull, &limitsNull, &cpu, &customKeys, &class, &inputBytes); err != nil { + t.Fatalf("raw column read: %v", err) + } + + if !requestsNull { + t.Error("resource_requests is not NULL for an undeclared job; " + + "an empty JSONB document decodes to the same zero Set but is a different row") + } + + if !limitsNull { + t.Error("resource_limits is not NULL for an undeclared job") + } + + // The scalar columns are NOT NULL DEFAULT, so they must be present + // and zero — that is what lets the dequeue comparisons be bare + // rather than COALESCEd. + if cpu != 0 || customKeys != "" || class != "" || inputBytes != 0 { + t.Errorf("scalar columns = cpu %d, keys %q, class %q, input %d; want all zero/empty", + cpu, customKeys, class, inputBytes) + } +} diff --git a/store/postgres/store.go b/store/postgres/store.go index 2de5ecb..15e74fc 100644 --- a/store/postgres/store.go +++ b/store/postgres/store.go @@ -11,6 +11,7 @@ import ( _ "github.com/xraph/grove/drivers/pgdriver/pgmigrate" // register pg migration executor "github.com/xraph/grove/migrate" + "github.com/xraph/dispatch/artifact" "github.com/xraph/dispatch/cluster" "github.com/xraph/dispatch/cron" "github.com/xraph/dispatch/dlq" @@ -22,11 +23,13 @@ import ( // Ensure Store implements all subsystem interfaces at compile time. var ( _ job.Store = (*Store)(nil) + _ job.LeaseStore = (*Store)(nil) _ workflow.Store = (*Store)(nil) _ cron.Store = (*Store)(nil) _ dlq.Store = (*Store)(nil) _ event.Store = (*Store)(nil) _ cluster.Store = (*Store)(nil) + _ artifact.Store = (*Store)(nil) ) // Store is a grove ORM implementation of store.Store using PostgreSQL dialect. @@ -70,11 +73,11 @@ func (s *Store) DB() *grove.DB { func (s *Store) Migrate(ctx context.Context) error { executor, err := migrate.NewExecutorFor(s.pgdb) if err != nil { - return fmt.Errorf("dispatch/bun: create migration executor: %w", err) + return fmt.Errorf(errPrefix+"create migration executor: %w", err) } orch := migrate.NewOrchestrator(executor, Migrations) if _, err := orch.Migrate(ctx); err != nil { - return fmt.Errorf("dispatch/bun: migration failed: %w", err) + return fmt.Errorf(errPrefix+"migration failed: %w", err) } return nil } diff --git a/store/postgres/store_test.go b/store/postgres/store_test.go index 30524fc..6567ef7 100644 --- a/store/postgres/store_test.go +++ b/store/postgres/store_test.go @@ -4,7 +4,6 @@ package postgres_test import ( "context" - "database/sql" "errors" "fmt" "testing" @@ -13,9 +12,10 @@ import ( "github.com/testcontainers/testcontainers-go" pgmodule "github.com/testcontainers/testcontainers-go/modules/postgres" "github.com/testcontainers/testcontainers-go/wait" - "github.com/uptrace/bun" - "github.com/uptrace/bun/dialect/pgdialect" - "github.com/uptrace/bun/driver/pgdriver" + + "github.com/xraph/grove" + "github.com/xraph/grove/drivers/pgdriver" + _ "github.com/xraph/grove/drivers/pgdriver/pgmigrate" // registers the pg migrate executor "github.com/xraph/dispatch" "github.com/xraph/dispatch/cluster" @@ -30,7 +30,7 @@ import ( log "github.com/xraph/go-utils/log" ) -// setupTestStore creates a Postgres container and returns a connected Bun Store. +// setupTestStore creates a Postgres container and returns a connected store. func setupTestStore(t *testing.T) *postgres.Store { t.Helper() @@ -61,9 +61,15 @@ func setupTestStore(t *testing.T) *postgres.Store { t.Fatalf("get connection string: %v", err) } - // Create Bun DB from pgdriver. - sqldb := sql.OpenDB(pgdriver.NewConnector(pgdriver.WithDSN(connStr))) - db := bun.NewDB(sqldb, pgdialect.New()) + drv := pgdriver.New() + if openErr := drv.Open(ctx, connStr); openErr != nil { + t.Fatalf("open pgdriver: %v", openErr) + } + + db, err := grove.Open(drv) + if err != nil { + t.Fatalf("grove open: %v", err) + } t.Cleanup(func() { _ = db.Close() @@ -144,6 +150,14 @@ func TestJobStore_DequeueSkipLocked(t *testing.T) { ctx := context.Background() // Enqueue 3 jobs with different priorities. + // + // run_at is backdated rather than set to now. Dequeue eligibility is + // `run_at <= NOW()`, where run_at comes from this process's clock and + // NOW() from the database's. When the two disagree by even a fraction + // of a millisecond -- routine when Postgres runs in a container -- a + // job stamped "now" is briefly in the future and is skipped, which + // made this test flake between 3, 2, and 0 eligible jobs. Backdating + // keeps the test about priority ordering and SKIP LOCKED. for i := 0; i < 3; i++ { j := &job.Job{ Entity: dispatch.NewEntity(), @@ -154,7 +168,7 @@ func TestJobStore_DequeueSkipLocked(t *testing.T) { State: job.StatePending, Priority: i, // 0, 1, 2 MaxRetries: 3, - RunAt: time.Now().UTC(), + RunAt: time.Now().UTC().Add(-time.Minute), } if err := s.EnqueueJob(ctx, j); err != nil { t.Fatalf("enqueue job-%d: %v", i, err) @@ -162,7 +176,7 @@ func TestJobStore_DequeueSkipLocked(t *testing.T) { } // Dequeue 2 — should get highest priority first. - dequeued, err := s.DequeueJobs(ctx, []string{"default"}, 2) + dequeued, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 2}) if err != nil { t.Fatalf("dequeue: %v", err) } @@ -177,7 +191,7 @@ func TestJobStore_DequeueSkipLocked(t *testing.T) { } // Dequeue remaining — should get 1 job. - remaining, err := s.DequeueJobs(ctx, []string{"default"}, 10) + remaining, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 10}) if err != nil { t.Fatalf("dequeue remaining: %v", err) } diff --git a/store/postgres/wake.go b/store/postgres/wake.go index 3986613..a11aa07 100644 --- a/store/postgres/wake.go +++ b/store/postgres/wake.go @@ -43,7 +43,7 @@ func (s *Store) StartWakeListener(ctx context.Context, wake func()) (func(), err l, err := s.pgdb.Listen(ctx, wakeChannel, handler) if err != nil { cancel() - return nil, fmt.Errorf("dispatch/postgres: start wake listener: %w", err) + return nil, fmt.Errorf(errPrefix+"start wake listener: %w", err) } done := make(chan struct{}) diff --git a/store/postgres/workflow.go b/store/postgres/workflow.go index 34b1ac8..fcccb86 100644 --- a/store/postgres/workflow.go +++ b/store/postgres/workflow.go @@ -18,7 +18,7 @@ func (s *Store) CreateRun(ctx context.Context, run *workflow.Run) error { if isDuplicateKey(err) { return dispatch.ErrJobAlreadyExists } - return fmt.Errorf("dispatch/bun: create run: %w", err) + return fmt.Errorf(errPrefix+"create run: %w", err) } return nil } @@ -34,7 +34,7 @@ func (s *Store) GetRun(ctx context.Context, runID id.RunID) (*workflow.Run, erro if isNoRows(err) { return nil, dispatch.ErrRunNotFound } - return nil, fmt.Errorf("dispatch/bun: get run: %w", err) + return nil, fmt.Errorf(errPrefix+"get run: %w", err) } return fromRunModel(m) } @@ -45,7 +45,7 @@ func (s *Store) UpdateRun(ctx context.Context, run *workflow.Run) error { m.UpdatedAt = time.Now().UTC() res, err := s.pgdb.NewUpdate(m).WherePK().Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: update run: %w", err) + return fmt.Errorf(errPrefix+"update run: %w", err) } rows, _ := res.RowsAffected() //nolint:errcheck // driver always returns nil if rows == 0 { @@ -74,14 +74,14 @@ func (s *Store) ListRuns(ctx context.Context, opts workflow.ListOpts) ([]*workfl err := q.Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: list runs: %w", err) + return nil, fmt.Errorf(errPrefix+"list runs: %w", err) } runs := make([]*workflow.Run, 0, len(models)) for i := range models { r, convErr := fromRunModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: list runs convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"list runs convert: %w", convErr) } runs = append(runs, r) } @@ -104,7 +104,7 @@ func (s *Store) SaveCheckpoint(ctx context.Context, runID id.RunID, stepName str Set("created_at = EXCLUDED.created_at"). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: save checkpoint: %w", err) + return fmt.Errorf(errPrefix+"save checkpoint: %w", err) } return nil } @@ -122,7 +122,7 @@ func (s *Store) GetCheckpoint(ctx context.Context, runID id.RunID, stepName stri if isNoRows(err) { return nil, nil // no checkpoint is not an error } - return nil, fmt.Errorf("dispatch/bun: get checkpoint: %w", err) + return nil, fmt.Errorf(errPrefix+"get checkpoint: %w", err) } return m.Data, nil } @@ -135,14 +135,14 @@ func (s *Store) ListCheckpoints(ctx context.Context, runID id.RunID) ([]*workflo OrderExpr("created_at ASC"). Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: list checkpoints: %w", err) + return nil, fmt.Errorf(errPrefix+"list checkpoints: %w", err) } checkpoints := make([]*workflow.Checkpoint, 0, len(models)) for i := range models { cp, convErr := fromCheckpointModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: list checkpoints convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"list checkpoints convert: %w", convErr) } checkpoints = append(checkpoints, cp) } @@ -157,14 +157,14 @@ func (s *Store) ListChildRuns(ctx context.Context, parentRunID id.RunID) ([]*wor OrderExpr("created_at ASC"). Scan(ctx) if err != nil { - return nil, fmt.Errorf("dispatch/bun: list child runs: %w", err) + return nil, fmt.Errorf(errPrefix+"list child runs: %w", err) } runs := make([]*workflow.Run, 0, len(models)) for i := range models { r, convErr := fromRunModel(&models[i]) if convErr != nil { - return nil, fmt.Errorf("dispatch/bun: list child runs convert: %w", convErr) + return nil, fmt.Errorf(errPrefix+"list child runs convert: %w", convErr) } runs = append(runs, r) } @@ -179,7 +179,7 @@ func (s *Store) DeleteCheckpointsAfter(ctx context.Context, runID id.RunID, afte Where("created_at > (SELECT created_at FROM dispatch_checkpoints WHERE run_id = ? AND step_name = ?)", runID.String(), afterStep). Exec(ctx) if err != nil { - return fmt.Errorf("dispatch/bun: delete checkpoints after: %w", err) + return fmt.Errorf(errPrefix+"delete checkpoints after: %w", err) } return nil } diff --git a/store/redis/artifact.go b/store/redis/artifact.go new file mode 100644 index 0000000..6441917 --- /dev/null +++ b/store/redis/artifact.go @@ -0,0 +1,842 @@ +package redis + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sort" + "strconv" + "time" + + goredis "github.com/redis/go-redis/v9" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// defaultSweepLimit bounds an unbounded sweep so a single pass can never +// touch an unbounded number of keys. +const defaultSweepLimit = 1000 + +// artifactEntity is the JSON shape stored under an artifact key. +type artifactEntity struct { + ID string `json:"id"` + Backend string `json:"backend"` + Bucket string `json:"bucket"` + Key string `json:"key"` + Size int64 `json:"size"` + ContentHash string `json:"content_hash,omitempty"` + ContentType string `json:"content_type,omitempty"` + Lifecycle string `json:"lifecycle"` + ScopeAppID string `json:"scope_app_id,omitempty"` + ScopeOrgID string `json:"scope_org_id,omitempty"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` + CreatedAt time.Time `json:"created_at"` + DeletedAt *time.Time `json:"deleted_at,omitempty"` +} + +func toArtifactEntity(a *artifact.Artifact) *artifactEntity { + return &artifactEntity{ + ID: a.ID.String(), + Backend: a.Backend, + Bucket: a.Bucket, + Key: a.Key, + Size: a.Size, + ContentHash: a.ContentHash, + ContentType: a.ContentType, + Lifecycle: string(a.Lifecycle), + ScopeAppID: a.ScopeAppID, + ScopeOrgID: a.ScopeOrgID, + ExpiresAt: a.ExpiresAt, + CreatedAt: a.CreatedAt, + DeletedAt: a.DeletedAt, + } +} + +func fromArtifactEntity(e *artifactEntity) (*artifact.Artifact, error) { + aid, err := id.ParseArtifactID(e.ID) + if err != nil { + return nil, err + } + + return &artifact.Artifact{ + ID: aid, + Backend: e.Backend, + Bucket: e.Bucket, + Key: e.Key, + Size: e.Size, + ContentHash: e.ContentHash, + ContentType: e.ContentType, + Lifecycle: artifact.Lifecycle(e.Lifecycle), + ScopeAppID: e.ScopeAppID, + ScopeOrgID: e.ScopeOrgID, + ExpiresAt: e.ExpiresAt, + CreatedAt: e.CreatedAt, + DeletedAt: e.DeletedAt, + }, nil +} + +// linkEntity is the JSON shape stored per link. +type linkEntity struct { + ArtifactID string `json:"artifact_id"` + OwnerKind string `json:"owner_kind"` + OwnerID string `json:"owner_id"` + Name string `json:"name"` + Attempt int `json:"attempt"` + Role string `json:"role"` + CreatedAt time.Time `json:"created_at"` +} + +func toLinkEntity(l *artifact.Link) *linkEntity { + return &linkEntity{ + ArtifactID: l.ArtifactID.String(), + OwnerKind: string(l.OwnerKind), + OwnerID: l.OwnerID, + Name: l.Name, + Attempt: l.Attempt, + Role: string(l.Role), + CreatedAt: l.CreatedAt, + } +} + +func fromLinkEntity(e *linkEntity) (*artifact.Link, error) { + aid, err := id.ParseArtifactID(e.ArtifactID) + if err != nil { + return nil, err + } + + return &artifact.Link{ + ArtifactID: aid, + OwnerKind: artifact.OwnerKind(e.OwnerKind), + OwnerID: e.OwnerID, + Role: artifact.Role(e.Role), + Name: e.Name, + Attempt: e.Attempt, + CreatedAt: e.CreatedAt, + }, nil +} + +// linkField is the hash field identifying one link within an owner's or +// artifact's link hash. +func linkField(name string, attempt int) string { + return name + "\x00" + strconv.Itoa(attempt) +} + +// CreateArtifact inserts an artifact and, when link is non-nil, its first +// link. +// +// The live-key guard is a SETNX on a dedicated key, which is what makes +// two concurrent creates at the same coordinates resolve to one winner +// and one ErrExists. +func (s *Store) CreateArtifact(ctx context.Context, a *artifact.Artifact, link *artifact.Link) error { + guard := artifactKeyGuard(a.Backend, a.Bucket, a.Key) + + ok, err := s.rdb.SetNX(ctx, guard, a.ID.String(), 0).Result() + if err != nil { + return fmt.Errorf("dispatch/redis: create artifact guard: %w", err) + } + + if !ok { + return artifact.ErrExists + } + + if err := s.setEntity(ctx, artifactKey(a.ID.String()), toArtifactEntity(a)); err != nil { + // Release the guard so the coordinates are not permanently burned. + s.rdb.Del(ctx, guard) + + return fmt.Errorf("dispatch/redis: create artifact: %w", err) + } + + if err := s.rdb.SAdd(ctx, artifactIDsKey, a.ID.String()).Err(); err != nil { + return fmt.Errorf("dispatch/redis: index artifact: %w", err) + } + + // Only ephemeral artifacts enter the sweep index. This is the Redis + // form of the SQL lifecycle literal: the sweeps read this index and + // no durable artifact is ever a member. + if a.Lifecycle == artifact.Ephemeral { + score := float64(a.CreatedAt.UnixNano()) + if err := s.rdb.ZAdd(ctx, artifactEphemeralKey, goredis.Z{Score: score, Member: a.ID.String()}).Err(); err != nil { + return fmt.Errorf("dispatch/redis: index ephemeral artifact: %w", err) + } + } + + if link == nil { + return nil + } + + return s.LinkArtifact(ctx, link) +} + +// GetArtifact retrieves a live artifact by ID. +func (s *Store) GetArtifact(ctx context.Context, artifactID id.ArtifactID) (*artifact.Artifact, error) { + a, err := s.loadArtifact(ctx, artifactID.String()) + if err != nil { + return nil, err + } + + if a.IsDeleted() { + return nil, artifact.ErrNotFound + } + + return a, nil +} + +// loadArtifact reads an artifact regardless of soft-deletion. +func (s *Store) loadArtifact(ctx context.Context, artifactID string) (*artifact.Artifact, error) { + var e artifactEntity + + if err := s.getEntity(ctx, artifactKey(artifactID), &e); err != nil { + if isNotFound(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf("dispatch/redis: get artifact: %w", err) + } + + return fromArtifactEntity(&e) +} + +// FindArtifactByKey retrieves a live artifact by its storage coordinates. +func (s *Store) FindArtifactByKey(ctx context.Context, backend, bucket, key string) (*artifact.Artifact, error) { + got, err := s.rdb.Get(ctx, artifactKeyGuard(backend, bucket, key)).Result() + if err != nil { + return nil, artifact.ErrNotFound + } + + return s.GetArtifact(ctx, id.MustParse(got)) +} + +// UpdateArtifact persists size, hash, content type, and expiry. Lifecycle, +// created_at, and deleted_at are deliberately not updatable here. +func (s *Store) UpdateArtifact(ctx context.Context, a *artifact.Artifact) error { + existing, err := s.loadArtifact(ctx, a.ID.String()) + if err != nil { + return err + } + + existing.Size = a.Size + existing.ContentHash = a.ContentHash + existing.ContentType = a.ContentType + existing.ExpiresAt = a.ExpiresAt + + if err := s.setEntity(ctx, artifactKey(a.ID.String()), toArtifactEntity(existing)); err != nil { + return fmt.Errorf("dispatch/redis: update artifact: %w", err) + } + + return nil +} + +// ListArtifacts returns artifacts matching the given options, newest first. +func (s *Store) ListArtifacts(ctx context.Context, opts artifact.ListOpts) ([]*artifact.Artifact, error) { + ids, err := s.rdb.SMembers(ctx, artifactIDsKey).Result() + if err != nil { + return nil, fmt.Errorf("dispatch/redis: list artifact ids: %w", err) + } + + out := make([]*artifact.Artifact, 0, len(ids)) + + for _, got := range ids { + a, lerr := s.loadArtifact(ctx, got) + if lerr != nil { + if errors.Is(lerr, artifact.ErrNotFound) { + continue + } + + return nil, lerr + } + + if a.IsDeleted() && !opts.IncludeDeleted { + continue + } + + if opts.Lifecycle != "" && a.Lifecycle != opts.Lifecycle { + continue + } + + if opts.ScopeAppID != "" && a.ScopeAppID != opts.ScopeAppID { + continue + } + + if opts.ScopeOrgID != "" && a.ScopeOrgID != opts.ScopeOrgID { + continue + } + + out = append(out, a) + } + + sort.Slice(out, func(i, j int) bool { + if out[i].CreatedAt.Equal(out[j].CreatedAt) { + return out[i].ID.String() < out[j].ID.String() + } + + return out[i].CreatedAt.After(out[j].CreatedAt) + }) + + if opts.Offset > 0 { + if opts.Offset >= len(out) { + return nil, nil + } + + out = out[opts.Offset:] + } + + if opts.Limit > 0 && opts.Limit < len(out) { + out = out[:opts.Limit] + } + + return out, nil +} + +// LinkArtifact records that an owner references an artifact, idempotently. +// HSet on the same field is naturally idempotent. +func (s *Store) LinkArtifact(ctx context.Context, link *artifact.Link) error { + raw, err := json.Marshal(toLinkEntity(link)) + if err != nil { + return fmt.Errorf("dispatch/redis: marshal link: %w", err) + } + + owner := artifact.OwnerRef{Kind: link.OwnerKind, ID: link.OwnerID} + field := linkField(link.Name, link.Attempt) + + if err := s.rdb.HSet(ctx, ownerLinksKey(string(owner.Kind), owner.ID), field, raw).Err(); err != nil { + return fmt.Errorf("dispatch/redis: link artifact: %w", err) + } + + member := string(link.OwnerKind) + "\x00" + link.OwnerID + "\x00" + field + if err := s.rdb.SAdd(ctx, artifactLinksKey(link.ArtifactID.String()), member).Err(); err != nil { + return fmt.Errorf("dispatch/redis: index artifact link: %w", err) + } + + return nil +} + +// ListLinks returns every link belonging to the given owner. +func (s *Store) ListLinks(ctx context.Context, owner artifact.OwnerRef) ([]*artifact.Link, error) { + vals, err := s.rdb.HGetAll(ctx, ownerLinksKey(string(owner.Kind), owner.ID)).Result() + if err != nil { + return nil, fmt.Errorf("dispatch/redis: list links: %w", err) + } + + out := make([]*artifact.Link, 0, len(vals)) + + for _, raw := range vals { + var e linkEntity + if uerr := json.Unmarshal([]byte(raw), &e); uerr != nil { + return nil, fmt.Errorf("dispatch/redis: unmarshal link: %w", uerr) + } + + l, cerr := fromLinkEntity(&e) + if cerr != nil { + return nil, cerr + } + + out = append(out, l) + } + + sort.Slice(out, func(i, j int) bool { + if out[i].Name == out[j].Name { + return out[i].Attempt < out[j].Attempt + } + + return out[i].Name < out[j].Name + }) + + return out, nil +} + +// FindLinkByName returns the highest-attempt link for an owner and name, +// breaking ties by CreatedAt descending — see the artifact.Store doc +// comment for why ties happen and what the tie-break does and does not +// fix. +// +// In practice a tie cannot reach this loop on this backend: LinkArtifact +// stores each link in a hash keyed by linkField(name, attempt), so a +// second write to the same (owner, name, attempt) overwrites the first +// rather than coexisting as a second row the way the SQL and document +// backends allow. The CreatedAt compare is kept anyway, both so this +// backend agrees with the other four if that storage shape ever changes, +// and because ListLinks' own sort does not otherwise guarantee which of +// two equal-attempt entries — however they arose — comes first. +func (s *Store) FindLinkByName( + ctx context.Context, + owner artifact.OwnerRef, + name string, +) (*artifact.Link, error) { + links, err := s.ListLinks(ctx, owner) + if err != nil { + return nil, err + } + + var best *artifact.Link + + for _, l := range links { + if l.Name != name { + continue + } + + if best == nil || l.Attempt > best.Attempt || + (l.Attempt == best.Attempt && l.CreatedAt.After(best.CreatedAt)) { + best = l + } + } + + if best == nil { + return nil, artifact.ErrNotFound + } + + return best, nil +} + +// ListArtifactsByOwner returns live artifacts linked to an owner. +func (s *Store) ListArtifactsByOwner( + ctx context.Context, + owner artifact.OwnerRef, + role artifact.Role, +) ([]*artifact.Artifact, error) { + links, err := s.ListLinks(ctx, owner) + if err != nil { + return nil, err + } + + var out []*artifact.Artifact + + seen := make(map[string]bool, len(links)) + + for _, l := range links { + if role != "" && l.Role != role { + continue + } + + key := l.ArtifactID.String() + if seen[key] { + continue + } + + seen[key] = true + + a, lerr := s.loadArtifact(ctx, key) + if lerr != nil { + if errors.Is(lerr, artifact.ErrNotFound) { + continue + } + + return nil, lerr + } + + if a.IsDeleted() { + continue + } + + out = append(out, a) + } + + return out, nil +} + +// SweepEphemeral marks eligible ephemeral artifacts as deleted. +// +// Candidates come from the ephemeral sorted set, which by construction +// contains no durable artifact. The lifecycle is re-checked after loading +// each candidate so the guard does not rest on index hygiene alone. +func (s *Store) SweepEphemeral( + ctx context.Context, + opts artifact.SweepOpts, +) ([]*artifact.Artifact, error) { + limit := opts.Limit + if limit <= 0 { + limit = defaultSweepLimit + } + + ids, err := s.rdb.ZRange(ctx, artifactEphemeralKey, 0, -1).Result() + if err != nil { + return nil, fmt.Errorf("dispatch/redis: sweep ephemeral: %w", err) + } + + nowAt := time.Now().UTC() + + var eligible []*artifact.Artifact + + for _, got := range ids { + if len(eligible) >= limit { + break + } + + a, lerr := s.loadArtifact(ctx, got) + if lerr != nil { + if errors.Is(lerr, artifact.ErrNotFound) { + continue + } + + return nil, lerr + } + + if a.Lifecycle != artifact.Ephemeral || a.IsDeleted() { + continue + } + + links, llerr := s.linksForArtifact(ctx, a.ID.String()) + if llerr != nil { + return nil, llerr + } + + if len(links) == 0 { + // Orphans are SweepOrphans' business. + continue + } + + terminalAt, ok, terr := s.ownersTerminalAt(ctx, links) + if terr != nil { + return nil, terr + } + + if !ok { + continue + } + + if a.ExpiresAt != nil { + if a.ExpiresAt.After(nowAt) { + continue + } + } else if terminalAt.Add(opts.Retention).After(nowAt) { + continue + } + + eligible = append(eligible, a) + } + + if len(eligible) == 0 || opts.DryRun { + return eligible, nil + } + + return s.markDeleted(ctx, eligible, nowAt) +} + +// linksForArtifact resolves every link pointing at an artifact. +func (s *Store) linksForArtifact(ctx context.Context, artifactID string) ([]*artifact.Link, error) { + members, err := s.rdb.SMembers(ctx, artifactLinksKey(artifactID)).Result() + if err != nil { + return nil, fmt.Errorf("dispatch/redis: list artifact links: %w", err) + } + + out := make([]*artifact.Link, 0, len(members)) + + for _, m := range members { + kind, ownerID, field, ok := splitLinkMember(m) + if !ok { + continue + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerKind(kind), ID: ownerID} + + raw, herr := s.rdb.HGet(ctx, ownerLinksKey(string(owner.Kind), owner.ID), field).Result() + if herr != nil { + continue + } + + var e linkEntity + if uerr := json.Unmarshal([]byte(raw), &e); uerr != nil { + return nil, fmt.Errorf("dispatch/redis: unmarshal link: %w", uerr) + } + + l, cerr := fromLinkEntity(&e) + if cerr != nil { + return nil, cerr + } + + out = append(out, l) + } + + return out, nil +} + +// splitLinkMember parses "kind\x00ownerID\x00name\x00attempt". +func splitLinkMember(m string) (kind, ownerID, field string, ok bool) { + first := indexByteFrom(m, 0) + if first < 0 { + return "", "", "", false + } + + second := indexByteFrom(m, first+1) + if second < 0 { + return "", "", "", false + } + + return m[:first], m[first+1 : second], m[second+1:], true +} + +func indexByteFrom(s string, from int) int { + for i := from; i < len(s); i++ { + if s[i] == 0 { + return i + } + } + + return -1 +} + +// ownersTerminalAt reports the latest terminal time across an artifact's +// owners, and whether all of them are terminal. An owner that no longer +// exists counts as terminal at the link's creation time. +func (s *Store) ownersTerminalAt(ctx context.Context, links []*artifact.Link) (time.Time, bool, error) { + var latest time.Time + + for _, l := range links { + at, ok, err := s.ownerTerminalAt(ctx, l) + if err != nil { + return time.Time{}, false, err + } + + if !ok { + return time.Time{}, false, nil + } + + if at.After(latest) { + latest = at + } + } + + return latest, true, nil +} + +func (s *Store) ownerTerminalAt(ctx context.Context, l *artifact.Link) (time.Time, bool, error) { + switch l.OwnerKind { + case artifact.OwnerJob: + var e jobEntity + if err := s.getEntity(ctx, jobKey(l.OwnerID), &e); err != nil { + if isNotFound(err) { + return l.CreatedAt, true, nil + } + + return time.Time{}, false, fmt.Errorf("dispatch/redis: resolve artifact owner job: %w", err) + } + + if !isTerminalOwnerState(e.State) { + return time.Time{}, false, nil + } + + if e.CompletedAt != nil { + return *e.CompletedAt, true, nil + } + + return e.UpdatedAt, true, nil + + case artifact.OwnerRun, artifact.OwnerStep: + var e runEntity + if err := s.getEntity(ctx, runKey(l.OwnerID), &e); err != nil { + if isNotFound(err) { + return l.CreatedAt, true, nil + } + + return time.Time{}, false, fmt.Errorf("dispatch/redis: resolve artifact owner run: %w", err) + } + + if !isTerminalOwnerState(e.State) { + return time.Time{}, false, nil + } + + if e.CompletedAt != nil { + return *e.CompletedAt, true, nil + } + + return e.UpdatedAt, true, nil + + default: + return l.CreatedAt, true, nil + } +} + +func isTerminalOwnerState(state string) bool { + switch state { + case "completed", "failed", "cancelled": + return true + default: + return false + } +} + +// SweepOrphans marks link-less ephemeral artifacts created before cutoff. +func (s *Store) SweepOrphans( + ctx context.Context, + cutoff time.Time, + limit int, +) ([]*artifact.Artifact, error) { + if limit <= 0 { + limit = defaultSweepLimit + } + + // The ephemeral index is scored by creation time, so the cutoff is a + // range query rather than a scan. + ids, err := s.rdb.ZRangeByScore(ctx, artifactEphemeralKey, &goredis.ZRangeBy{ + Min: "-inf", + Max: strconv.FormatInt(cutoff.UnixNano(), 10), + }).Result() + if err != nil { + return nil, fmt.Errorf("dispatch/redis: sweep orphans: %w", err) + } + + var orphans []*artifact.Artifact + + for _, got := range ids { + if len(orphans) >= limit { + break + } + + a, lerr := s.loadArtifact(ctx, got) + if lerr != nil { + if errors.Is(lerr, artifact.ErrNotFound) { + continue + } + + return nil, lerr + } + + if a.Lifecycle != artifact.Ephemeral || a.IsDeleted() { + continue + } + + if !a.CreatedAt.Before(cutoff) { + continue + } + + n, cerr := s.rdb.SCard(ctx, artifactLinksKey(a.ID.String())).Result() + if cerr != nil { + return nil, fmt.Errorf("dispatch/redis: count artifact links: %w", cerr) + } + + if n > 0 { + continue + } + + orphans = append(orphans, a) + } + + if len(orphans) == 0 { + return nil, nil + } + + return s.markDeleted(ctx, orphans, time.Now().UTC()) +} + +// markDeleted soft-deletes the given artifacts, re-checking the lifecycle +// on each so the write itself refuses to touch a durable artifact. +func (s *Store) markDeleted( + ctx context.Context, + artifacts []*artifact.Artifact, + at time.Time, +) ([]*artifact.Artifact, error) { + out := make([]*artifact.Artifact, 0, len(artifacts)) + + for _, a := range artifacts { + if a.Lifecycle != artifact.Ephemeral { + continue + } + + clone := a.Clone() + deleted := at + clone.DeletedAt = &deleted + + if err := s.setEntity(ctx, artifactKey(clone.ID.String()), toArtifactEntity(clone)); err != nil { + return nil, fmt.Errorf("dispatch/redis: mark artifact deleted: %w", err) + } + + // Release the live-key guard so the coordinates become reusable, + // and index the deletion time for the purge pass. + s.rdb.Del(ctx, artifactKeyGuard(clone.Backend, clone.Bucket, clone.Key)) + + if err := s.rdb.ZAdd(ctx, artifactDeletedKey, + goredis.Z{Score: float64(at.UnixNano()), Member: clone.ID.String()}).Err(); err != nil { + return nil, fmt.Errorf("dispatch/redis: index deleted artifact: %w", err) + } + + if err := s.rdb.ZRem(ctx, artifactEphemeralKey, clone.ID.String()).Err(); err != nil { + return nil, fmt.Errorf("dispatch/redis: deindex ephemeral artifact: %w", err) + } + + out = append(out, clone) + } + + return out, nil +} + +// ListPurgeable returns soft-deleted artifacts older than grace. +func (s *Store) ListPurgeable( + ctx context.Context, + grace time.Duration, + limit int, +) ([]*artifact.Artifact, error) { + if limit <= 0 { + limit = defaultSweepLimit + } + + cutoff := time.Now().UTC().Add(-grace) + + ids, err := s.rdb.ZRangeByScore(ctx, artifactDeletedKey, &goredis.ZRangeBy{ + Min: "-inf", + Max: strconv.FormatInt(cutoff.UnixNano(), 10), + }).Result() + if err != nil { + return nil, fmt.Errorf("dispatch/redis: list purgeable: %w", err) + } + + out := make([]*artifact.Artifact, 0, len(ids)) + + for _, got := range ids { + if len(out) >= limit { + break + } + + a, lerr := s.loadArtifact(ctx, got) + if lerr != nil { + if errors.Is(lerr, artifact.ErrNotFound) { + continue + } + + return nil, lerr + } + + if !a.IsDeleted() { + continue + } + + out = append(out, a) + } + + return out, nil +} + +// PurgeArtifact hard-deletes an artifact and its link index entries. +func (s *Store) PurgeArtifact(ctx context.Context, artifactID id.ArtifactID) error { + key := artifactID.String() + + links, err := s.linksForArtifact(ctx, key) + if err != nil { + return err + } + + for _, l := range links { + owner := artifact.OwnerRef{Kind: l.OwnerKind, ID: l.OwnerID} + if herr := s.rdb.HDel(ctx, ownerLinksKey(string(owner.Kind), owner.ID), linkField(l.Name, l.Attempt)).Err(); herr != nil { + return fmt.Errorf("dispatch/redis: purge artifact link: %w", herr) + } + } + + a, err := s.loadArtifact(ctx, key) + if err == nil { + s.rdb.Del(ctx, artifactKeyGuard(a.Backend, a.Bucket, a.Key)) + } + + pipe := s.rdb.TxPipeline() + pipe.Del(ctx, artifactKey(key)) + pipe.Del(ctx, artifactLinksKey(key)) + pipe.SRem(ctx, artifactIDsKey, key) + pipe.ZRem(ctx, artifactEphemeralKey, key) + pipe.ZRem(ctx, artifactDeletedKey, key) + + if _, err := pipe.Exec(ctx); err != nil { + return fmt.Errorf("dispatch/redis: purge artifact: %w", err) + } + + return nil +} diff --git a/store/redis/artifact_test.go b/store/redis/artifact_test.go new file mode 100644 index 0000000..9b6f7ae --- /dev/null +++ b/store/redis/artifact_test.go @@ -0,0 +1,18 @@ +//go:build integration + +package redis_test + +import ( + "testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" +) + +// TestArtifactStoreConformance runs the shared artifact.Store suite +// against Redis. +func TestArtifactStoreConformance(t *testing.T) { + artifacttest.RunStoreSuite(t, func(t *testing.T) artifact.Store { + return setupTestStore(t) + }) +} diff --git a/store/redis/dequeue.go b/store/redis/dequeue.go new file mode 100644 index 0000000..59277fa --- /dev/null +++ b/store/redis/dequeue.go @@ -0,0 +1,488 @@ +package redis + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sort" + "time" + + goredis "github.com/redis/go-redis/v9" + + "github.com/xraph/dispatch/job" +) + +// Where the fit predicate lives, and why it is not in Lua. +// +// Redis has no query language, so the predicate had two possible homes: +// Go, or a Lua script running on the server. Lua would move less data +// over the wire; it would also be a SECOND implementation of +// job.DequeueOpts.Allows/Less, written in a language with no tests of its +// own, in a repository whose whole resource-model track exists to stop the +// five backends drifting apart — the shared resource/codec.go and the +// shared storetest conformance suite are both there for exactly that +// reason. A Lua re-statement of "absent budget key is unconstrained", +// "requirement <= budget", "custom keys are a subset test, not a +// substring one" and "locality ranks below priority" is the single +// highest-drift-risk thing this track could contain. +// +// So the predicate stays in Go and is job.DequeueOpts' own code, called +// directly — the same choice store/memory makes, and for the same reason: +// one expression of the contract, zero drift. What Redis pays for that is +// reading candidate entities it may then discard. +// +// The CLAIM is a different question from the predicate, and it does NOT +// move to Go. See claimCandidates. + +// maxDequeueRounds bounds the retry loop that runs when every candidate +// this call selected was claimed by a competing worker first. Without a +// bound, a busy queue could keep a caller scanning indefinitely; with it, +// a fully contended call returns empty and the pool simply polls again. +// Mirrors store/mongo, which composes its claim the same way. +const maxDequeueRounds = 3 + +// dequeueScanBatch is how many job entities one pipelined read fetches. +// A scan that reads many members issues them as pipelined GETs in batches +// rather than as one GET per round trip. A pipeline (not MGET) is +// deliberate: go-redis splits a pipeline across cluster nodes by slot, +// whereas a multi-key MGET spanning slots is a CROSSSLOT error. +const dequeueScanBatch = 256 + +// The two scan modes, and why the cheap one is not just an optimization. +// +// A caller that constrains nothing — no Budget, no CustomKeys, no +// ReservedFor, no PreferHashes — is the pool as it is configured when +// nobody has turned the resource model on, which is the overwhelming +// majority of deployments and the one the whole track promised not to +// regress. Such a caller wants the first Limit ready members of the +// index in score order, which is exactly what the pre-track ZPopMin +// gave it, at a cost proportional to Limit rather than to the backlog. +// +// The full scan exists only because the fit predicate and the locality +// term are properties the score cannot express: a job that does not fit +// may sit anywhere in the index, so "the first Limit members" is not an +// answer to a bounded caller's question. Those callers opted in, and +// they pay a cost proportional to the depth of the queues they named. +// +// Charging an unconfigured worker that cost is what made a deep queue a +// self-inflicted outage: Redis is single-threaded, so one worker reading +// a 100k-member index every poll interval blocks every enqueue, +// heartbeat and lease renewal behind it, and it does so precisely when +// the queue is deepest. +const ( + // unboundedScanFloor is the smallest window the bounded scan reads, + // so a Limit of 1 still tolerates a little junk at the head of the + // index without a second round trip. + unboundedScanFloor = 16 + + // unboundedScanCeiling caps the total members one bounded scan may + // read from one queue. Reached when more than this many + // higher-priority members sit at the head of a single queue's index + // and are not yet eligible — scheduled ahead of RunAt, or in retry + // backoff — with ready work behind them. Measured: with 3000 such + // members ahead of 10 ready jobs, three consecutive unbounded polls + // each returned zero jobs (2056 commands, 13 ms each), and no + // progress was made until the not-yet-eligible members became due. + // The window is bounded and cannot livelock, but it is starvation, + // not slowdown. This window did not exist in the previous full-scan + // implementation; it is a deliberate trade against the full-scan cost + // documented above, not an oversight. + unboundedScanCeiling = 2048 +) + +// dequeueCandidate is one job that passed the fit predicate, together +// with the queue index member that has to be won to claim it. +type dequeueCandidate struct { + id string + queue string + job *job.Job +} + +// DequeueJobs atomically claims up to opts.Limit ready jobs from +// opts.Queues that fit opts, sets them to running, and returns them +// ordered by priority descending, then locality-preferred first, then +// RunAt ascending. +// +// The call is a candidate scan followed by an exclusive claim: +// +// - dequeueCandidates reads the queue index, decodes each entity, +// applies job.DequeueOpts.Allows, sorts the SURVIVORS with +// job.DequeueOpts.Less and only then truncates to Limit. Order +// precedes truncation, over the whole eligible set — a scan that took +// Limit members first and sorted within them would hand a worker with +// a small limit arbitrary low-priority work forever, which is what +// storetest's LimitTruncatesAfterOrdering pins. How much of the index +// is read depends on whether the caller filters anything: see the +// unboundedScan* constants. +// - claimCandidates wins each survivor by removing it from the queue +// index, which is what makes the claim exclusive, then writes it back +// as running — carrying the lease grant when opts.Grants(), in that +// same write. +// +// A job that does not fit is never removed from the index and never +// written to: it stays pending and untouched for the next worker that +// does have room. That is the whole point of the predicate — a job +// claimed and then requeued would bounce between small workers, delaying +// exactly the job that is hardest to place. +// +// A non-positive Limit claims nothing, matching the SQL backends' LIMIT 0. +// A worker computing zero free slots must claim zero jobs, never the +// whole queue. +// +// Empty opts.Queues claims nothing here, which is this backend's existing +// "all queues" behaviour: the queue index is one sorted set per queue +// name and there has never been a cross-queue index to scan. The +// conformance suite never exercises empty Queues, and no caller in this +// repository sends it. +func (s *Store) DequeueJobs(ctx context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + if opts.Limit <= 0 { + return nil, nil + } + + if err := opts.Validate(); err != nil { + return nil, fmt.Errorf("dispatch/redis: dequeue jobs: %w", err) + } + + for range maxDequeueRounds { + candidates, err := s.dequeueCandidates(ctx, opts) + if err != nil { + return nil, err + } + + if len(candidates) == 0 { + return nil, nil + } + + claimed, err := s.claimCandidates(ctx, opts, candidates) + if err != nil { + return nil, err + } + + if len(claimed) > 0 { + return claimed, nil + } + } + + return nil, nil +} + +// dequeueCandidates returns the top opts.Limit eligible jobs across +// opts.Queues, already in contract order. +// +// The score is a legacy ordering hint — a float packing negated priority +// and RunAt into one number — and the returned ORDER never trusts it: +// the contract order includes a locality term the score cannot express, +// and the float's RunAt component loses resolution as priority grows. +// Ordering is decided by job.DequeueOpts.Less over decoded jobs. +// +// What the score does decide is WHICH members are read, and that is +// where the two modes differ. See the unboundedScan* constants: a caller +// that filters and orders nothing reads a bounded window from the head +// of the index; every other caller reads the whole index, because a job +// it can accept may sit anywhere in it. +func (s *Store) dequeueCandidates(ctx context.Context, opts job.DequeueOpts) ([]dequeueCandidate, error) { + t := now() + candidates := make([]dequeueCandidate, 0, opts.Limit) + + for _, q := range opts.Queues { + found, err := s.scanQueue(ctx, opts, q, t) + if err != nil { + return nil, err + } + + candidates = append(candidates, found...) + } + + // Order THEN truncate, across every queue named, exactly as + // store/memory does: priority descending, then locality-preferred + // before not — a tiebreak strictly within a priority band, never + // above it — then RunAt ascending. + // + // Do not "optimize" this by truncating first. A measured caveat for + // whoever tries: swapping these two statements alone does NOT fail + // storetest's LimitTruncatesAfterOrdering, because ZRange happens to + // return members in score order and the score happens to encode + // priority. It DOES fail + // LocalityDecidesWhichRowsSurviveATightLimit, which was added for + // exactly this reason: locality cannot be baked into a score that + // was written before the caller's staged set was known, so no scan + // order can make truncate-before-sort accidentally right there. + sort.SliceStable(candidates, func(i, k int) bool { + return opts.Less(candidates[i].job, candidates[k].job) + }) + + if len(candidates) > opts.Limit { + candidates = candidates[:opts.Limit] + } + + return candidates, nil +} + +// scanQueue returns the eligible members of one queue's index, in the +// order the index yielded them. +// +// A bounded caller reads the whole index. An unbounded one reads a +// window that starts near Limit and doubles, stopping as soon as it has +// Limit eligible jobs — so the common case is one ZRANGE and one +// pipeline of Limit-ish GETs, and the widening only pays for junk the +// state/RunAt gate actually rejected. +func (s *Store) scanQueue( + ctx context.Context, + opts job.DequeueOpts, + q string, + t time.Time, +) ([]dequeueCandidate, error) { + key := queueKey(q) + full := !opts.IsUnbounded() || len(opts.PreferredHashes()) > 0 + + var ( + out []dequeueCandidate + scanned int64 + window = int64(max(opts.Limit, unboundedScanFloor)) + ) + + for { + stop := int64(-1) // the whole index + if !full { + stop = scanned + window - 1 + } + + ids, err := s.rdb.ZRange(ctx, key, scanned, stop).Result() + if err != nil && !isRedisNil(err) { + return nil, fmt.Errorf("dispatch/redis: dequeue scan %q: %w", q, err) + } + + if len(ids) == 0 { + return out, nil + } + + found, err := s.eligibleIn(ctx, opts, q, t, ids) + if err != nil { + return nil, err + } + + out = append(out, found...) + scanned += int64(len(ids)) + + // The full scan asked for everything and got it in one call. + if full { + return out, nil + } + + if len(out) >= opts.Limit || + int64(len(ids)) < window || // index exhausted + scanned >= unboundedScanCeiling { + return out, nil + } + + window = min(window*2, unboundedScanCeiling-scanned) + if window <= 0 { + return out, nil + } + } +} + +// eligibleIn decodes the named index members and returns those a worker +// may claim, positionally in ids order. +func (s *Store) eligibleIn( + ctx context.Context, + opts job.DequeueOpts, + q string, + t time.Time, + ids []string, +) ([]dequeueCandidate, error) { + out := make([]dequeueCandidate, 0, len(ids)) + + for start := 0; start < len(ids); start += dequeueScanBatch { + end := min(start+dequeueScanBatch, len(ids)) + batch := ids[start:end] + + entities, readErr := s.readJobEntities(ctx, batch) + if readErr != nil { + return nil, readErr + } + + for i, e := range entities { + if e == nil { + continue // indexed but gone, or unreadable + } + + // The index is not a state filter: EnqueueJob adds every job + // it writes, including one handed to it already running, and + // ReclaimExpiredLeases re-adds jobs it returned to pending. + // Only a ready job may be claimed. + if st := job.State(e.State); st != job.StatePending && st != job.StateRetrying { + continue + } + + if !e.RunAt.IsZero() && e.RunAt.After(t) { + continue + } + + j, convErr := fromJobEntity(e) + if convErr != nil { + continue + } + + // IsUnbounded skips the fit predicate entirely: a caller not + // using the resource model claims everything, including jobs + // declaring custom resources it could not possibly satisfy. + // It governs FILTERING only — PreferHashes still orders + // below, even here. + if !opts.IsUnbounded() && !opts.Allows(j) { + continue + } + + out = append(out, dequeueCandidate{id: batch[i], queue: q, job: j}) + } + } + + return out, nil +} + +// readJobEntities fetches one batch of job entities by ID, returning a +// slice positionally aligned with ids and holding nil where the entity +// was missing or could not be decoded. +// +// The read is a pipeline of GETs rather than one MGET so it stays correct +// against a clustered client, where the job keys of one queue are spread +// over many slots. +func (s *Store) readJobEntities(ctx context.Context, ids []string) ([]*jobEntity, error) { + pipe := s.rdb.Pipeline() + + cmds := make([]*goredis.StringCmd, len(ids)) + for i, jID := range ids { + cmds[i] = pipe.Get(ctx, jobKey(jID)) + } + + // A missing key makes Exec report goredis.Nil for the batch as a + // whole; the per-command results below distinguish the misses, and a + // job that vanished between the index read and this one is simply not + // a candidate. + if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, goredis.Nil) { + return nil, fmt.Errorf("dispatch/redis: dequeue read entities: %w", err) + } + + out := make([]*jobEntity, len(ids)) + + for i, cmd := range cmds { + raw, err := cmd.Bytes() + if err != nil { + continue + } + + var e jobEntity + if json.Unmarshal(raw, &e) != nil { + continue + } + + out[i] = &e + } + + return out, nil +} + +// claimCandidates takes exclusive ownership of each candidate and returns +// the ones this caller won, still in candidate order. +// +// THE CLAIM IS WHERE ATOMICITY LIVES, and it is unchanged in kind from +// what this store did before the predicate existed. Removal from the +// queue's sorted set is the claim: ZREM is a single Redis command, so the +// server executes it indivisibly, and of any number of workers racing for +// one member exactly one gets a reply of 1. Everyone else gets 0 and +// moves on empty-handed — never a second claim of the same job. That is +// the same guarantee the ZPopMin this store used before gave: also one +// atomic command, also removal-is-the-claim. +// +// ZREM rather than ZPopMin because ZPopMin chooses its own members by +// score, which would mean popping jobs this caller has already decided do +// not fit and then having to put them back — the claim-then-requeue the +// predicate exists to prevent. ZREM lets Go nominate exactly the jobs that +// passed. +// +// Only after winning the removal is the entity read and rewritten as +// running. Between those two steps the job is reachable by no other +// dequeue path, so the read-modify-write needs no compare-and-set of its +// own; this is the identical window the previous ZPopMin implementation +// had, and a crash inside it leaves the job exactly as a crash after +// ZPopMin did. +// +// THE LEASE GRANT BELONGS TO THAT SAME WRITE, and stays a plain +// read-modify-write for exactly the reason above rather than becoming a +// second SET or a Lua compare-and-set. Winning the ZREM already removed +// the job from every path any other worker could reach it by, so there is +// nothing left to race against and no epoch to compare. What must not +// happen is granting in a SECOND write after this one: ReclaimExpiredLeases +// skips any entity whose job.Lease.IsExpired is false, and that is false +// for a zero expiry, so an entity written as running with no lease is one +// reclamation can never see rather than one it might take. A crash between +// the two writes would strand the job until the coarse global stale-job +// threshold noticed it — the mechanism leases exist to replace. See +// job.DequeueOpts.LeaseUntil. +func (s *Store) claimCandidates( + ctx context.Context, + opts job.DequeueOpts, + candidates []dequeueCandidate, +) ([]*job.Job, error) { + pipe := s.rdb.Pipeline() + + rems := make([]*goredis.IntCmd, len(candidates)) + for i, c := range candidates { + rems[i] = pipe.ZRem(ctx, queueKey(c.queue), c.id) + } + + if _, err := pipe.Exec(ctx); err != nil { + return nil, fmt.Errorf("dispatch/redis: dequeue claim: %w", err) + } + + t := now() + claimed := make([]*job.Job, 0, len(candidates)) + + for i, c := range candidates { + won, err := rems[i].Result() + if err != nil || won != 1 { + continue // another worker removed it first + } + + key := jobKey(c.id) + + var e jobEntity + if getErr := s.getEntity(ctx, key, &e); getErr != nil { + continue // won the index entry but the entity is gone + } + + // Re-read rather than reusing the scanned copy, so the blob + // written back is built on the freshest state. The state check is + // belt-and-braces: winning the ZREM already excludes every other + // dequeue path, so this only fires if some other writer moved the + // job out of pending while it sat in the index. + if st := job.State(e.State); st != job.StatePending && st != job.StateRetrying { + continue + } + + e.State = string(job.StateRunning) + e.StartedAt = &t + e.UpdatedAt = t + + if opts.Grants() { + until := opts.LeaseUntil.UTC() + e.WorkerID = opts.WorkerID.String() + e.LeaseEpoch++ + e.LeaseExpiresAt = &until + } + + if setErr := s.setEntity(ctx, key, &e); setErr != nil { + return nil, fmt.Errorf("dispatch/redis: dequeue update: %w", setErr) + } + + j, convErr := fromJobEntity(&e) + if convErr != nil { + return nil, convErr + } + + claimed = append(claimed, j) + } + + return claimed, nil +} diff --git a/store/redis/dequeue_test.go b/store/redis/dequeue_test.go new file mode 100644 index 0000000..6ce3015 --- /dev/null +++ b/store/redis/dequeue_test.go @@ -0,0 +1,454 @@ +//go:build integration + +package redis_test + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "testing" + "time" + + goredis "github.com/redis/go-redis/v9" + + "github.com/xraph/grove/kv/drivers/redisdriver" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + redisstore "github.com/xraph/dispatch/store/redis" + "github.com/xraph/dispatch/store/storetest" +) + +// rawJobKey mirrors the unexported jobKey helper. These tests live in the +// external test package and have to reach the stored blob directly to +// produce shapes no Go write path can produce, so the key convention is +// restated here rather than exported from the package under test. +func rawJobKey(jobID id.JobID) string { return "dispatch:job:" + jobID.String() } + +// TestDequeueConformance runs the resource-aware dequeue suite against +// the Redis store. +// +// One container is stood up and shared by every subtest, which the suite +// documents as safe: each case enqueues onto its own queue and asserts +// only on the jobs it created. A container per subtest would dominate the +// runtime of the whole package. +func TestDequeueConformance(t *testing.T) { + shared := setupTestStore(t) + + storetest.RunDequeueSuite(t, func(t *testing.T) job.Store { + t.Helper() + + return shared + }) +} + +// newRawFitJob builds a pending job ready to run now, for the two cases +// below that patch the stored blob afterwards. +func newRawFitJob(name, queue string, runAt time.Time) *job.Job { + return &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: name, + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: runAt, + } +} + +// patchJobBlob rewrites the stored JSON for one job by applying mutate to +// its decoded object form. Going through a map rather than jobEntity is +// the point: it can delete a key entirely, which no Go struct write path +// in this package can do. +func patchJobBlob(t *testing.T, s *redisstore.Store, jobID id.JobID, mutate func(map[string]any)) { + t.Helper() + + ctx := context.Background() + client := redisdriver.UnwrapClient(s.KV()) + key := rawJobKey(jobID) + + raw, err := client.Get(ctx, key).Bytes() + if err != nil { + t.Fatalf("read raw blob for %s: %v", jobID, err) + } + + var blob map[string]any + if err = json.Unmarshal(raw, &blob); err != nil { + t.Fatalf("decode raw blob for %s: %v", jobID, err) + } + + mutate(blob) + + patched, err := json.Marshal(blob) + if err != nil { + t.Fatalf("encode patched blob for %s: %v", jobID, err) + } + + if err = client.Set(ctx, key, patched, 0).Err(); err != nil { + t.Fatalf("write patched blob for %s: %v", jobID, err) + } +} + +func rawJobNames(jobs []*job.Job) []string { + out := make([]string, 0, len(jobs)) + for _, j := range jobs { + out = append(out, j.Name) + } + + return out +} + +// TestDequeueClaimsJobsWrittenBeforeTheResourceFields is the first of the +// two gaps the shared conformance suite structurally cannot cover: it +// only ever creates jobs through EnqueueJob, so every blob it produces +// carries the full current field set. A deployment that has been running +// since before the resource model shipped has blobs that do not. +// +// Such a blob has no req_cpu_milli, no req_memory_bytes, no +// req_disk_bytes, no req_gpu_milli, no req_custom_keys, and no +// resource_requests key at all. It declares NO requirement, so a +// resource-aware worker must still claim it. The two ways to get that +// wrong are to read an absent field as a parse failure and drop the job, +// or to read an absent custom-key list as "unknown, therefore reject" +// under bounded opts — either one silently strands every pre-upgrade job +// in the queue, with nothing reporting it. +// +// Redis reaches the right answer through the decode path rather than +// through a filter clause, which is why this case is worth pinning +// explicitly: encoding/json leaves an absent field at its zero value, and +// resource.DecodeSet maps both a null and an absent resource_requests to +// a nil Set, so the reconstructed job requires nothing and +// job.DequeueOpts.Allows admits it. Mongo needed explicit {field: nil} +// branches for the same property because its range operators are +// type-bracketed; Redis needs none, and this test is what keeps that true. +// +// The opts are deliberately BOUNDED — a budget plus an empty CustomKeys +// offer — because unbounded opts skip the predicate entirely and would +// claim the job no matter how badly the decode handled the missing keys. +func TestDequeueClaimsJobsWrittenBeforeTheResourceFields(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + const queue = "redis-legacy-blob" + + base := time.Now().UTC().Add(-time.Hour).Truncate(time.Millisecond) + + legacy := newRawFitJob("legacy-no-resource-fields", queue, base) + nulled := newRawFitJob("legacy-null-resource-requests", queue, base.Add(time.Minute)) + + for _, j := range []*job.Job{legacy, nulled} { + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue %s: %v", j.Name, err) + } + } + + // The pre-upgrade shape: none of the resource keys exist. + patchJobBlob(t, s, legacy.ID, func(blob map[string]any) { + for _, field := range []string{ + "req_cpu_milli", "req_memory_bytes", "req_disk_bytes", "req_gpu_milli", + "req_custom_keys", "resource_requests", "resource_limits", + } { + delete(blob, field) + } + }) + + // The other shape a hand-written or older blob can carry: the key is + // present and explicitly null. resource.DecodeSet must read it the + // same way it reads an absent key. + patchJobBlob(t, s, nulled.ID, func(blob map[string]any) { + blob["resource_requests"] = nil + blob["req_custom_keys"] = nil + }) + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 4 * storetest.GiB}, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 2 { + t.Fatalf("claimed %v, want both legacy jobs: a job written before the req_* fields "+ + "existed declares no requirement and must stay claimable", rawJobNames(got)) + } + + for _, j := range got { + if len(j.Resources) != 0 { + t.Errorf("job %s came back requiring %v, want nothing", j.Name, j.Resources) + } + } +} + +// TestDequeueOrdersNullPrimaryInputHashAsUnpreferred is the second gap. +// jobEntity.PrimaryInputHash is a plain string, so no Go write path in +// this package can produce a blob whose primary_input_hash is null or +// absent — the shared suite cannot construct the shape at all. +// +// A job with no hash has no locality signal and must therefore sort as +// NOT preferred: behind the job the caller has already staged, and +// otherwise by the ordinary rules. The failure mode this guards is +// implementing "preferred first" as a sort on the hash VALUE, which has +// nothing to do with whether the caller staged it. +// +// The "remote" fixture is what stops the case passing for the wrong +// reason, and it is not optional. Mongo found that transplanting +// Postgres's version of this test produced a FALSE POSITIVE: with only +// one hash value in play, the empty/null hash collates BELOW it, so a +// broken value sort still put the staged job first. "remote" carries +// "zzz:never-staged", which collates ABOVE "blake3:staged-here", so a +// value sort hands back the job the caller has NO local copy of, first. +// Both fixtures with no signal are also enqueued EARLIER than the staged +// one, so they win any tie an unimplemented locality term fails to break. +// +// Mutation-verified: replacing the contract comparator with a descending +// sort on PrimaryInputHash within a priority band yields +// [remote cached null-hash missing-hash] and fails here; treating any +// non-empty hash as preferred yields [cached remote null-hash +// missing-hash] and fails here too. +func TestDequeueOrdersNullPrimaryInputHashAsUnpreferred(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + const ( + queue = "redis-null-hash-order" + local = "blake3:staged-here" + ) + + base := time.Now().UTC().Add(-time.Hour).Truncate(time.Millisecond) + + nullHash := newRawFitJob("null-hash", queue, base) + missingHash := newRawFitJob("missing-hash", queue, base.Add(time.Minute)) + + cached := newRawFitJob("cached", queue, base.Add(2*time.Minute)) + cached.PrimaryInputHash = local + + remote := newRawFitJob("remote", queue, base.Add(3*time.Minute)) + remote.PrimaryInputHash = "zzz:never-staged" + + for _, j := range []*job.Job{nullHash, missingHash, cached, remote} { + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue %s: %v", j.Name, err) + } + } + + patchJobBlob(t, s, nullHash.ID, func(blob map[string]any) { + blob["primary_input_hash"] = nil + }) + + patchJobBlob(t, s, missingHash.ID, func(blob map[string]any) { + delete(blob, "primary_input_hash") + }) + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + PreferHashes: []string{local}, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 4 { + t.Fatalf("claimed %d jobs (%v), want 4 — locality must never filter", + len(got), rawJobNames(got)) + } + + want := []string{"cached", "null-hash", "missing-hash", "remote"} + for i, name := range want { + if got[i].Name != name { + t.Fatalf("claimed %v, want %v: a null or absent primary_input_hash must sort as "+ + "NOT preferred, never ahead of the job the caller has staged", + rawJobNames(got), want) + } + } +} + +// ────────────────────────────────────────────────── +// Scan cost +// ────────────────────────────────────────────────── + +// countingHook tallies every command this client sends, pipelined ones +// individually. A pipeline is one round trip but N commands of Redis +// CPU, and Redis is single-threaded, so it is the command count — not +// the round-trip count — that decides whether a deep queue starves +// enqueues, heartbeats and lease renewals. +type countingHook struct { + mu sync.Mutex + n int +} + +func (h *countingHook) add(n int) { + h.mu.Lock() + defer h.mu.Unlock() + + h.n += n +} + +func (h *countingHook) count() int { + h.mu.Lock() + defer h.mu.Unlock() + + return h.n +} + +func (h *countingHook) reset() { + h.mu.Lock() + defer h.mu.Unlock() + + h.n = 0 +} + +func (h *countingHook) DialHook(next goredis.DialHook) goredis.DialHook { return next } + +func (h *countingHook) ProcessHook(next goredis.ProcessHook) goredis.ProcessHook { + return func(ctx context.Context, cmd goredis.Cmder) error { + h.add(1) + + return next(ctx, cmd) + } +} + +func (h *countingHook) ProcessPipelineHook(next goredis.ProcessPipelineHook) goredis.ProcessPipelineHook { + return func(ctx context.Context, cmds []goredis.Cmder) error { + h.add(len(cmds)) + + return next(ctx, cmds) + } +} + +// TestUnboundedDequeueDoesNotScanTheWholeQueue pins the promise the rest +// of this file cannot: with no resource configuration, a poll must cost +// what it cost before the resource model existed. +// +// The opts here are exactly what worker.Pool sends when nobody has +// configured anything — Queues and a Limit, no Budget, no CustomKeys, no +// PreferHashes — so job.DequeueOpts.IsUnbounded is true and the fit +// predicate is skipped. A caller that filters nothing has no reason to +// read a job it is not going to return. +// +// The bound is on COMMANDS, not on wall time, because the failure this +// guards is Redis CPU rather than latency at any one worker. The +// implementation this replaced issued one ZRANGE over the whole index +// plus one pipelined GET per pending member: at the backlog below that +// is one ZRANGE and ~2000 GETs against a Limit of 4, growing linearly +// with the depth of the queue, on every poll of every worker. +// +// Mutation-verified: restoring the `ZRange(ctx, key, 0, -1)` full scan +// fails here with a count above 2000. +func TestUnboundedDequeueDoesNotScanTheWholeQueue(t *testing.T) { + const ( + queue = "redis-scan-cost" + backlog = 2000 + limit = 4 + + // Generous on purpose: the point is the difference between a + // constant and a term linear in backlog, not a tight count that + // breaks the next time the claim grows a command. + maxCommands = 150 + ) + + s := setupTestStore(t) + ctx := context.Background() + + base := time.Now().UTC().Add(-time.Hour).Truncate(time.Millisecond) + + for i := range backlog { + j := newRawFitJob(fmt.Sprintf("backlog-%04d", i), queue, + base.Add(time.Duration(i)*time.Millisecond)) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue %d: %v", i, err) + } + } + + hook := &countingHook{} + redisdriver.UnwrapClient(s.KV()).AddHook(hook) + + hook.reset() + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: limit, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != limit { + t.Fatalf("claimed %d jobs, want %d", len(got), limit) + } + + // The bounded window still has to answer the ordering contract over + // what it read: these were enqueued at ascending RunAt and equal + // priority, so the head of the index is the head of the answer. + for i, j := range got { + if want := fmt.Sprintf("backlog-%04d", i); j.Name != want { + t.Fatalf("claimed %v, want the first %d by RunAt", rawJobNames(got), limit) + } + } + + if n := hook.count(); n > maxCommands { + t.Fatalf("an unbounded dequeue of %d jobs from a %d-deep queue cost %d Redis commands "+ + "(want <= %d): the scan is proportional to the backlog, so the deeper the queue "+ + "the more Redis CPU every poll of every worker burns", + limit, backlog, n, maxCommands) + } + + t.Logf("unbounded dequeue: %d Redis commands at a backlog of %d", hook.count(), backlog) +} + +// TestBoundedDequeueStillSeesPastTheWindow is the other half of the +// contract above: the bounded window is for callers that filter nothing. +// A caller with a budget may have to look past the head of the index, +// because the one job it can run may sit anywhere in it — so the full +// scan must survive for exactly those callers. +func TestBoundedDequeueStillSeesPastTheWindow(t *testing.T) { + const ( + queue = "redis-scan-past-window" + ahead = 400 + ) + + s := setupTestStore(t) + ctx := context.Background() + + base := time.Now().UTC().Add(-time.Hour).Truncate(time.Millisecond) + + // A wall of jobs no small worker can take... + for i := range ahead { + j := newRawFitJob(fmt.Sprintf("huge-%03d", i), queue, + base.Add(time.Duration(i)*time.Millisecond)) + j.Resources = resource.Set{resource.Memory: 64 * storetest.GiB} + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue huge %d: %v", i, err) + } + } + + // ...and one it can, far enough back that no bounded window reaches it. + small := newRawFitJob("small", queue, base.Add(time.Duration(ahead)*time.Millisecond)) + small.Resources = resource.Set{resource.Memory: storetest.GiB} + + if err := s.EnqueueJob(ctx, small); err != nil { + t.Fatalf("enqueue small: %v", err) + } + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 4, + Budget: resource.Set{resource.Memory: 2 * storetest.GiB}, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 1 || got[0].Name != "small" { + t.Fatalf("claimed %v, want [small]: a bounded caller must scan past the jobs it "+ + "cannot run, or a wall of oversized work strands every worker behind it", + rawJobNames(got)) + } +} diff --git a/store/redis/dlq.go b/store/redis/dlq.go index 32b1a48..165f573 100644 --- a/store/redis/dlq.go +++ b/store/redis/dlq.go @@ -8,6 +8,7 @@ import ( "github.com/xraph/dispatch" "github.com/xraph/dispatch/dlq" "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" ) // ── JSON model for KV storage ── @@ -26,6 +27,19 @@ type dlqEntity struct { FailedAt time.Time `json:"failed_at"` ReplayedAt *time.Time `json:"replayed_at,omitempty"` CreatedAt time.Time `json:"created_at"` + + // Carried so Replay can rebuild a job that behaves like the failed + // one; see the dlq.Entry doc. resource.Set is a map[string]int64 and + // marshals natively, like every other field here. + Priority int `json:"priority,omitempty"` + Timeout time.Duration `json:"timeout,omitempty"` + LeaseTTL time.Duration `json:"lease_ttl,omitempty"` + ArtifactBindings []byte `json:"artifact_bindings,omitempty"` + Resources resource.Set `json:"resources,omitempty"` + ResourceLimits resource.Set `json:"resource_limits,omitempty"` + ResourceClass string `json:"resource_class,omitempty"` + InputBytes int64 `json:"input_bytes,omitempty"` + PrimaryInputHash string `json:"primary_input_hash,omitempty"` } func toDLQEntity(e *dlq.Entry) *dlqEntity { @@ -43,6 +57,16 @@ func toDLQEntity(e *dlq.Entry) *dlqEntity { FailedAt: e.FailedAt, ReplayedAt: e.ReplayedAt, CreatedAt: e.CreatedAt, + + Priority: e.Priority, + Timeout: e.Timeout, + LeaseTTL: e.LeaseTTL, + ArtifactBindings: e.ArtifactBindings, + Resources: e.Resources, + ResourceLimits: e.ResourceLimits, + ResourceClass: e.ResourceClass, + InputBytes: e.InputBytes, + PrimaryInputHash: e.PrimaryInputHash, } } @@ -68,6 +92,16 @@ func fromDLQEntity(e *dlqEntity) (*dlq.Entry, error) { FailedAt: e.FailedAt, ReplayedAt: e.ReplayedAt, CreatedAt: e.CreatedAt, + + Priority: e.Priority, + Timeout: e.Timeout, + LeaseTTL: e.LeaseTTL, + ArtifactBindings: e.ArtifactBindings, + Resources: e.Resources, + ResourceLimits: e.ResourceLimits, + ResourceClass: e.ResourceClass, + InputBytes: e.InputBytes, + PrimaryInputHash: e.PrimaryInputHash, }, nil } diff --git a/store/redis/job.go b/store/redis/job.go index e8009ca..adaab4f 100644 --- a/store/redis/job.go +++ b/store/redis/job.go @@ -12,6 +12,7 @@ import ( "github.com/xraph/dispatch" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" ) // ── JSON model for KV storage ── @@ -36,9 +37,48 @@ type jobEntity struct { Timeout int64 `json:"timeout"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` + + LeaseEpoch int `json:"lease_epoch"` + LeaseExpiresAt *time.Time `json:"lease_expires_at,omitempty"` + LeaseTTL int64 `json:"lease_ttl"` + EvictCount int `json:"evict_count"` + + // The four canonical dimensions get their own fields because a + // later dequeue predicate compares them numerically and must + // behave identically across all five backends. They are derived + // from Resources by toJobEntity -- never independently settable. + ReqCPUMilli int64 `json:"req_cpu_milli"` + ReqMemoryBytes int64 `json:"req_memory_bytes"` + ReqDiskBytes int64 `json:"req_disk_bytes"` + ReqGPUMilli int64 `json:"req_gpu_milli"` + ReqCustomKeys string `json:"req_custom_keys"` + + // ResourceRequests and ResourceLimits are the full-fidelity encoded + // copy of Resources / ResourceLimits produced by resource.EncodeSet, + // including custom keys the scalar fields above do not carry. + // fromJobEntity reconstructs Resources from here, not from the + // scalars. "omitempty" on this []byte -- nil for a zero Set, per + // EncodeSet's contract -- is what keeps an undeclared job's JSON + // blob free of the key entirely, mirroring the SQL backends' NULL + // column: an absent value, not "{}" or "". + ResourceRequests []byte `json:"resource_requests,omitempty"` + ResourceLimits []byte `json:"resource_limits,omitempty"` + ResourceClass string `json:"resource_class"` + InputBytes int64 `json:"input_bytes"` + PrimaryInputHash string `json:"primary_input_hash"` } -func toJobEntity(j *job.Job) *jobEntity { +func toJobEntity(j *job.Job) (*jobEntity, error) { + reqJSON, err := resource.EncodeSet(j.Resources) + if err != nil { + return nil, fmt.Errorf("dispatch/redis: marshal job resources: %w", err) + } + + limitsJSON, err := resource.EncodeSet(j.ResourceLimits) + if err != nil { + return nil, fmt.Errorf("dispatch/redis: marshal job resource limits: %w", err) + } + return &jobEntity{ ID: j.ID.String(), Name: j.Name, @@ -59,7 +99,23 @@ func toJobEntity(j *job.Job) *jobEntity { Timeout: j.Timeout.Nanoseconds(), CreatedAt: j.CreatedAt, UpdatedAt: j.UpdatedAt, - } + + LeaseEpoch: j.LeaseEpoch, + LeaseExpiresAt: j.LeaseExpiresAt, + LeaseTTL: j.LeaseTTL.Nanoseconds(), + EvictCount: j.EvictCount, + + ReqCPUMilli: j.Resources[resource.CPU], + ReqMemoryBytes: j.Resources[resource.Memory], + ReqDiskBytes: j.Resources[resource.Disk], + ReqGPUMilli: j.Resources[resource.GPU], + ReqCustomKeys: resource.EncodeCustomKeys(j.Resources), + ResourceRequests: reqJSON, + ResourceLimits: limitsJSON, + ResourceClass: j.ResourceClass, + InputBytes: j.InputBytes, + PrimaryInputHash: j.PrimaryInputHash, + }, nil } func fromJobEntity(e *jobEntity) (*job.Job, error) { @@ -68,6 +124,16 @@ func fromJobEntity(e *jobEntity) (*job.Job, error) { return nil, fmt.Errorf("dispatch/redis: parse job id: %w", err) } + resources, err := resource.DecodeSet(e.ResourceRequests) + if err != nil { + return nil, fmt.Errorf("dispatch/redis: unmarshal job resources: %w", err) + } + + limits, err := resource.DecodeSet(e.ResourceLimits) + if err != nil { + return nil, fmt.Errorf("dispatch/redis: unmarshal job resource limits: %w", err) + } + j := &job.Job{ Entity: dispatch.Entity{ CreatedAt: e.CreatedAt, @@ -89,6 +155,17 @@ func fromJobEntity(e *jobEntity) (*job.Job, error) { CompletedAt: e.CompletedAt, HeartbeatAt: e.HeartbeatAt, Timeout: time.Duration(e.Timeout), + + LeaseEpoch: e.LeaseEpoch, + LeaseExpiresAt: e.LeaseExpiresAt, + LeaseTTL: time.Duration(e.LeaseTTL), + EvictCount: e.EvictCount, + + Resources: resources, + ResourceLimits: limits, + ResourceClass: e.ResourceClass, + InputBytes: e.InputBytes, + PrimaryInputHash: e.PrimaryInputHash, } if e.WorkerID != "" { @@ -115,7 +192,10 @@ func (s *Store) EnqueueJob(ctx context.Context, j *job.Job) error { return dispatch.ErrJobAlreadyExists } - e := toJobEntity(j) + e, err := toJobEntity(j) + if err != nil { + return err + } if setErr := s.setEntity(ctx, key, e); setErr != nil { return fmt.Errorf("dispatch/redis: enqueue set entity: %w", setErr) } @@ -135,53 +215,8 @@ func (s *Store) EnqueueJob(ctx context.Context, j *job.Job) error { return nil } -// DequeueJobs atomically pops up to limit jobs from the given queues. -func (s *Store) DequeueJobs(ctx context.Context, queues []string, limit int) ([]*job.Job, error) { - t := now() - var jobs []*job.Job - - for _, q := range queues { - if len(jobs) >= limit { - break - } - remaining := limit - len(jobs) - qk := queueKey(q) - - // Pop from sorted set (lowest score = highest priority + earliest RunAt). - members, err := s.rdb.ZPopMin(ctx, qk, int64(remaining)).Result() - if err != nil { - return nil, fmt.Errorf("dispatch/redis: dequeue zpopmin: %w", err) - } - - for _, z := range members { - jID, ok := z.Member.(string) - if !ok { - continue - } - - key := jobKey(jID) - var e jobEntity - if getErr := s.getEntity(ctx, key, &e); getErr != nil { - continue // skip missing - } - - // Update state to running. - e.State = string(job.StateRunning) - e.StartedAt = &t - e.UpdatedAt = t - if setErr := s.setEntity(ctx, key, &e); setErr != nil { - return nil, fmt.Errorf("dispatch/redis: dequeue update: %w", setErr) - } - - j, convErr := fromJobEntity(&e) - if convErr != nil { - return nil, convErr - } - jobs = append(jobs, j) - } - } - return jobs, nil -} +// DequeueJobs lives in dequeue.go, where the fit predicate and the claim +// are documented together. // GetJob retrieves a job by ID. func (s *Store) GetJob(ctx context.Context, jobID id.JobID) (*job.Job, error) { @@ -195,7 +230,34 @@ func (s *Store) GetJob(ctx context.Context, jobID id.JobID) (*job.Job, error) { return fromJobEntity(&e) } -// UpdateJob persists changes to an existing job. +// UpdateJob persists changes to an existing job and keeps the queue index +// in step with the state it just wrote. +// +// The index is what makes a job visible to DequeueJobs at all: EnqueueJob +// adds a member, the claim removes it, and every state transition after +// that arrives HERE. A retry from Runner.scheduleRetry, a job handed back +// by Pool.requeueRateLimited or Pool.requeueUndispatched, a stale job reset +// by Pool.reapStaleJobs — all of them just set a runnable state and call +// this. Writing that state without restoring the member leaves a job that +// GetJob reports as pending and no dequeue can ever see again: it looks +// healthy in the dashboard and runs nowhere. The other four backends have +// no equivalent hazard because they have no index — they re-derive +// candidacy from the row on every query. +// +// The two index writes sit on OPPOSITE sides of the entity write, which is +// deliberate. The stored entity is authoritative — dequeue.go re-checks +// state and RunAt against it — so a member that should not be there is +// inert, while a member that is missing is a job that never runs again. +// Ordering each write so the index errs towards the harmless side means a +// lost second write cannot strand anything: +// +// becoming runnable — ZADD first, so a failed entity write leaves a +// spare member the state gate ignores. +// becoming final — ZREM last, so a failed entity write leaves the job +// both runnable and still indexed. +// +// ZADD on a member already present only updates its score, so this is also +// safe for a job that was never claimed. func (s *Store) UpdateJob(ctx context.Context, j *job.Job) error { jID := j.ID.String() key := jobKey(jID) @@ -208,9 +270,37 @@ func (s *Store) UpdateJob(ctx context.Context, j *job.Job) error { return dispatch.ErrJobNotFound } - e := toJobEntity(j) + e, err := toJobEntity(j) + if err != nil { + return err + } e.UpdatedAt = now() - return s.setEntity(ctx, key, e) + + // Nothing in this repository moves a job between queues after enqueue, + // so j.Queue is the queue it was indexed under. If that ever changes, + // the old queue keeps a member pointing at this job, and this function + // has to read the stored entity to learn which queue to clear. + qk := queueKey(j.Queue) + runnable := j.State == job.StatePending || j.State == job.StateRetrying + + if runnable { + z := goredis.Z{Score: jobScore(j.Priority, j.RunAt), Member: jID} + if zErr := s.rdb.ZAdd(ctx, qk, z).Err(); zErr != nil { + return fmt.Errorf("dispatch/redis: update job index add: %w", zErr) + } + } + + if setErr := s.setEntity(ctx, key, e); setErr != nil { + return fmt.Errorf("dispatch/redis: update job set entity: %w", setErr) + } + + if !runnable { + if zErr := s.rdb.ZRem(ctx, qk, jID).Err(); zErr != nil { + return fmt.Errorf("dispatch/redis: update job index remove: %w", zErr) + } + } + + return nil } // DeleteJob removes a job by ID. diff --git a/store/redis/keys.go b/store/redis/keys.go index 961a99f..bcb37b3 100644 --- a/store/redis/keys.go +++ b/store/redis/keys.go @@ -73,3 +73,38 @@ const workerIDsKey = keyPrefix + "worker_ids" // leaderKey stores the current leader worker ID. const leaderKey = keyPrefix + "leader" + +// ── Artifact keys ── + +// artifactKey returns the key for an artifact entity: dispatch:artifact:{id} +func artifactKey(id string) string { return keyPrefix + "artifact:" + id } + +// artifactIDsKey is the Set tracking all artifact IDs for enumeration. +const artifactIDsKey = keyPrefix + "artifact_ids" + +// artifactKeyGuard maps live storage coordinates to an artifact ID. It is +// claimed with SETNX so concurrent creates at the same coordinates resolve +// to one winner, and released on soft-delete so a purged key is reusable. +func artifactKeyGuard(backend, bucket, key string) string { + return fmt.Sprintf("%sartifact_key:%s:%s:%s", keyPrefix, backend, bucket, key) +} + +// artifactEphemeralKey is the Sorted Set of ephemeral artifact IDs scored +// by creation time. Durable artifacts are never members, which is this +// backend's form of the SQL "lifecycle = 'ephemeral'" literal. +const artifactEphemeralKey = keyPrefix + "artifact_ephemeral" + +// artifactDeletedKey is the Sorted Set of soft-deleted artifact IDs scored +// by deletion time, driving the purge pass. +const artifactDeletedKey = keyPrefix + "artifact_deleted" + +// artifactLinksKey is the Set of link members pointing at an artifact. +func artifactLinksKey(artifactID string) string { + return keyPrefix + "artifact_links:" + artifactID +} + +// ownerLinksKey is the Hash of an owner's artifact links, keyed by +// "name\x00attempt". +func ownerLinksKey(kind, ownerID string) string { + return fmt.Sprintf("%sartifact_owner_links:%s:%s", keyPrefix, kind, ownerID) +} diff --git a/store/redis/lease.go b/store/redis/lease.go new file mode 100644 index 0000000..762237e --- /dev/null +++ b/store/redis/lease.go @@ -0,0 +1,483 @@ +package redis + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + goredis "github.com/redis/go-redis/v9" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// Both scripts below decode the stored job blob with cjson, but only to +// CHECK three scalar fields (state, worker_id, lease_epoch) — never to +// mutate and re-encode it. That split exists because of a bug found in +// review: an earlier version of this file had each script decode the +// whole blob, mutate a few fields in Lua, and cjson.encode the result +// back. That looked safe on paper — jobEntity has no slice/map fields, so +// the classic cjson empty-array/empty-object ambiguity never applied, and +// every *time.Time field is a string, not a number — and it was still +// wrong. cjson represents every JSON number as a Lua double, and Redis's +// cjson renders a double that large in scientific notation on encode +// (`9999999999999999` came back as `1e+16`). encoding/json then refuses +// to parse that back into an int64 at all on the next read — not a +// rounding error, a hard unmarshal failure — for Timeout or LeaseTTL past +// 2^53ns (~104 days), on every renewal or reclaim that touched the row, +// whether or not it cared about those fields. See +// TestLeaseLargeDurationRoundTrip for the reproduction, and note that in +// ReclaimExpiredLeases specifically, that unmarshal error aborted the +// whole scan, abandoning every other expired job the call had already +// found before it ever reached the queue re-add. +// +// The fix: Go now owns all serialization. RenewLease and +// ReclaimExpiredLeases each read the current entity, compute the fully +// updated entity in Go, and json.Marshal it themselves — the same path +// every other write in this store already uses. The script's job shrinks +// to being the compare-and-set: decode just enough to check +// state/worker_id/lease_epoch against what the caller expects, and if +// they match, SET the pre-built blob Go handed it. cjson.encode is never +// called by either script now, so no field can be reshaped by it — the +// whole corruption class is closed, not just the two fields that +// happened to trip it first. +// +// That fix has its own tradeoff, and it is deliberate, not overlooked. +// Between Go's read and the script's SET there is a real window — a full +// round trip — during which some other writer could change a field on +// this same job that the lease check doesn't cover. UpdateJob is the +// concrete example: it doesn't touch lease_epoch, so the epoch check +// inside these scripts would still pass, and the blind SET would +// overwrite whatever UpdateJob just wrote with Go's now-stale copy of +// that field. The epoch check makes the *lease* compare-and-set atomic; +// it does not make every write to the row serialize with every other +// write. +// +// This paragraph used to claim that window was "currently unreachable, +// not closed," because nothing called UpdateJob concurrently with +// RenewLease or ReclaimExpiredLeases on the same job. That premise was +// wrong, and worker/runner.go's terminal writes are the disproof: a +// worker whose lease had already been reclaimed — a live holder was +// mid-attempt, renewing on schedule — still called the unfenced UpdateJob +// from handleSuccess, scheduleRetry, or sendToDLQ, using a claim-time job +// snapshot. It won the race, rolled lease_epoch backwards, and marked the +// job completed while the reclaiming worker was still executing it. The +// "lease-aware pool loop" this comment said would be later work is +// exactly what RenewLease-on-heartbeat already was; the missing piece was +// never the pool loop, it was a fenced write for the runner to call. +// +// UpdateLeasedJob, below, is that write. It reuses renewLeaseScript +// itself — the fence it needs (still running, still this worker, still +// this epoch) is identical to RenewLease's — so the runner's terminal +// writes now go through the same compare-and-set FAMILY as renewal and +// reclaim: whichever of the three lands second on a given job sees a +// state, worker_id, or lease_epoch that has already moved and backs off +// cleanly, instead of blind-overwriting a write it never saw. Concretely, +// that means a genuine RECLAIM can no longer be lost to a zombie's +// terminal write racing it — the split-brain this whole track exists to +// close. +// +// It does NOT close the general read/SET race described two paragraphs +// up, and does not attempt to: UpdateLeasedJob is itself a Go-side +// read-modify-write, with the identical round-trip window between its own +// GET and its own script call. A RenewLease landing in that window passes +// UpdateLeasedJob's compare-and-set (state/worker_id/lease_epoch are +// exactly what it expects) and has the expiry it just pushed forward +// silently overwritten by UpdateLeasedJob's older copy — the same failure +// mode UpdateJob always had against these scripts, just between two +// different methods that both now use them. What's closed is one +// specific, previously-reachable instance: a zombie's own terminal write +// no longer wins against the RECLAIM that fenced it. UpdateJob itself +// remains unfenced by design (see job.LeaseStore.UpdateLeasedJob — the +// reaper's legacy path, rate-limit and shutdown requeues must stay +// unfenced, or reclaiming a dead worker's job would deadlock on that +// worker's own epoch), so a caller that reaches UpdateJob directly for a +// job whose lease has moved on — bypassing the pool and UpdateLeasedJob +// entirely — can still race every script in this file. A full-blob +// compare-and-swap (checking the entire previous blob byte-for-byte, not +// just three fields, before the SET) would narrow both windows further, +// but is still rejected for the reason it always was: it would make +// renewal fail on any unrelated concurrent write, including perfectly +// legitimate ones, and a spurious ErrLeaseLost is exactly what makes a +// pool cancel a healthy running job. + +// renewLeaseScript extends a lease only when the caller still holds it. +// +// KEYS[1] job key. ARGV[1] worker id, ARGV[2] expected epoch, ARGV[3] the +// complete updated entity, pre-serialized by Go (see the file comment +// above for why Lua never re-serializes it itself). +// Returns 1 on renewal, 0 when the lease is no longer held. +var renewLeaseScript = goredis.NewScript(` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return 0 +end +local j = cjson.decode(raw) +if j.state ~= 'running' then + return 0 +end +if j.worker_id ~= ARGV[1] then + return 0 +end +if tostring(j.lease_epoch) ~= ARGV[2] then + return 0 +end +redis.call('SET', KEYS[1], ARGV[3]) +return 1 +`) + +// reclaimScript resets one job to pending only if it is still running at +// the expected epoch. +// +// Reclamation does not need to re-derive "should this be taken back" +// inside Lua: that decision was already made correctly in Go, using real +// time.Time comparison (see reclaimable), before this script was ever +// called. This script re-verifies only equality — still running, still at +// the epoch Go observed — which is enough to make the claim exclusive: if +// another caller (or a fresh grant) already moved the job, the epoch or +// state check fails and this caller loses, cleanly. +// +// KEYS[1] job key. ARGV[1] expected epoch, ARGV[2] the complete +// pending-state entity, pre-serialized by Go with lease_epoch already +// incremented and evict_count already incremented (see the file comment +// above for why Lua never mutates or re-serializes it itself). +// Returns 1 when this caller took the job, 0 when someone else did (or +// the job moved out of running between Go's read and this script). +var reclaimScript = goredis.NewScript(` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return 0 +end +local j = cjson.decode(raw) +if j.state ~= 'running' then + return 0 +end +if tostring(j.lease_epoch) ~= ARGV[1] then + return 0 +end +redis.call('SET', KEYS[1], ARGV[2]) +return 1 +`) + +// The grant is not in this file: it travels on job.DequeueOpts and is +// applied by claimCandidates, in the same read-modify-write that writes +// the claimed job as running, so a leased claim carries the fit predicate +// and the ordering like any other. See store/redis/dequeue.go for why +// that write needs no compare-and-set of its own while the two below do. + +// RenewLease extends the lease only if the caller still holds it. +// +// Go reads the current entity, mutates only the lease/heartbeat fields, +// and serializes the whole thing with encoding/json — the same path +// every other write in this store uses. The script's only job is the +// compare-and-set: verify state/worker_id/lease_epoch still match what +// this read saw, and if so, SET the blob Go built. See the file comment +// above for the ABA tradeoff this introduces and why it's accepted. +func (s *Store) RenewLease( + ctx context.Context, + jobID id.JobID, + workerID id.WorkerID, + epoch int, + leaseUntil time.Time, +) error { + key := jobKey(jobID.String()) + + var e jobEntity + if getErr := s.getEntity(ctx, key, &e); getErr != nil { + if isNotFound(getErr) { + return job.ErrLeaseLost + } + return fmt.Errorf("dispatch/redis: renew lease get: %w", getErr) + } + + t := now() + until := leaseUntil.UTC() + e.LeaseExpiresAt = &until + e.HeartbeatAt = &t + e.UpdatedAt = t + + blob, marshalErr := json.Marshal(&e) + if marshalErr != nil { + return fmt.Errorf("dispatch/redis: renew lease marshal: %w", marshalErr) + } + + res, err := renewLeaseScript.Run(ctx, s.rdb, + []string{key}, + workerID.String(), + epoch, + blob, + ).Int64() + if err != nil && !errors.Is(err, goredis.Nil) { + return fmt.Errorf("dispatch/redis: renew lease: %w", err) + } + if res != 1 { + return job.ErrLeaseLost + } + + return nil +} + +// reclaimable reports whether a running job should be taken back. +// +// The first clause is the actual rule, and job.Lease.IsExpired remains its +// only authority: a lease was granted and has lapsed. +// +// The second is a deliberate, narrow exception to the invariant documented +// at job/lease.go, which is that a zero expiry means "never leased" rather +// than "expired" precisely so that reclamation cannot steal a job nobody +// ever leased. See job.UnleasedReclaimGrace for why the exception is +// needed and why it is gated on silence rather than on the null expiry +// alone; the four persistent backends all apply the same rule, three of +// them in SQL or a query filter and this one here. +func reclaimable(e *jobEntity, t time.Time) bool { + if e.LeaseExpiresAt != nil { + return job.Lease{ExpiresAt: *e.LeaseExpiresAt}.IsExpired(t) + } + + // Heartbeat first, falling back to the claim time for a worker that + // died before its first beat — the same order ReapStaleJobs used. + silent := e.HeartbeatAt + if silent == nil { + silent = e.StartedAt + } + if silent == nil { + return false + } + + return silent.Before(t.Add(-job.UnleasedReclaimGrace)) +} + +// ReclaimExpiredLeases returns expired-lease jobs to pending, fencing +// their previous holders. +// +// Reclamation walks the job-id set rather than a sorted index, matching +// ReapStaleJobs — there is no secondary index of running-with-expired- +// lease jobs in this backend. Each candidate is filtered here in Go by +// reclaimable, the pending-state entity is computed in Go, +// and claimExpired does the compare-and-set: keyed on the epoch this call +// observed, so two pools scanning concurrently cannot both take the same +// job — whichever script call runs second sees an epoch (or state) that +// no longer matches and backs off. +func (s *Store) ReclaimExpiredLeases(ctx context.Context, limit int) ([]*job.Job, error) { + // A non-positive limit claims nothing, matching the SQL backends' + // LIMIT 0 and DequeueOpts.Limit's behavior. + if limit <= 0 { + return nil, nil + } + + t := now() + + ids, err := s.rdb.SMembers(ctx, jobIDsKey).Result() + if err != nil { + return nil, fmt.Errorf("dispatch/redis: reclaim smembers: %w", err) + } + + reclaimed := make([]*job.Job, 0, limit) + for _, jID := range ids { + if len(reclaimed) >= limit { + break + } + + var e jobEntity + if getErr := s.getEntity(ctx, jobKey(jID), &e); getErr != nil { + continue // gone by the time we looked + } + if job.State(e.State) != job.StateRunning { + continue + } + if !reclaimable(&e, t) { + continue + } + + expectedEpoch := e.LeaseEpoch + + // The pending-state entity, computed entirely in Go. Lua only + // checks state/lease_epoch against expectedEpoch and blind-SETs + // this blob — see the file comment above for why. + after := e + after.State = string(job.StatePending) + after.RunAt = t + after.UpdatedAt = t + after.WorkerID = "" + after.StartedAt = nil + after.HeartbeatAt = nil + after.LeaseExpiresAt = nil + after.LeaseEpoch = expectedEpoch + 1 + after.EvictCount++ + + blob, marshalErr := json.Marshal(&after) + if marshalErr != nil { + return nil, fmt.Errorf("dispatch/redis: reclaim marshal: %w", marshalErr) + } + + claimed, claimErr := s.claimExpired(ctx, jID, expectedEpoch, blob) + if claimErr != nil { + return nil, claimErr + } + if !claimed { + continue // another pool got there first + } + + // The claim reset the job's entity to pending, but ZPopMin already + // removed it from queueKey at grant time — it never went back on + // its own. Without this it becomes invisible to every future + // dequeue: reset to pending but unreachable, forever. ZADD on a + // member already present just updates its score, so this is safe + // even for a job that was enqueued straight into running (as the + // conformance suite's RunningJob helper does) and was therefore + // never popped in the first place. + zErr := s.rdb.ZAdd(ctx, queueKey(after.Queue), + goredis.Z{Score: jobScore(after.Priority, after.RunAt), Member: jID}).Err() + if zErr != nil { + return nil, fmt.Errorf("dispatch/redis: reclaim requeue: %w", zErr) + } + + j, convErr := fromJobEntity(&after) + if convErr != nil { + continue + } + reclaimed = append(reclaimed, j) + } + + return reclaimed, nil +} + +// claimExpired atomically SETs one expired job's pre-built pending-state +// blob, but only if it is still running at the expected epoch, reporting +// whether this caller was the one that took it. +// +// There is deliberately no re-read after the claim. Go already knows +// exactly what the row now says, because Go built the blob it just wrote. +// An earlier version of this function re-read the entity after a +// successful claim — which meant decoding whatever cjson.encode had just +// produced, and that was precisely the step that turned a large Timeout +// or LeaseTTL into a scientific-notation string encoding/json couldn't +// parse. Because that failure happened inside ReclaimExpiredLeases' loop, +// it aborted the whole scan and abandoned every other expired job already +// found. Go no longer needs to ask Redis what the row says; it already +// knows, because it wrote it. +func (s *Store) claimExpired(ctx context.Context, jID string, epoch int, blob []byte) (bool, error) { + res, err := reclaimScript.Run(ctx, s.rdb, + []string{jobKey(jID)}, + epoch, + blob, + ).Int64() + if err != nil && !errors.Is(err, goredis.Nil) { + return false, fmt.Errorf("dispatch/redis: reclaim claim: %w", err) + } + + return res == 1, nil +} + +// UpdateLeasedJob persists j only while the caller still holds the +// lease. +// +// This is the read-modify-write the column-exclusion shape forces: a +// whole-entity SET can't express "every field except these four," so Go +// reads the current entity, overlays only j's business fields onto it — +// never a pre-serialized blob built from j alone — and hands the result +// to renewLeaseScript for the same compare-and-set RenewLease uses. The +// fence the two need is identical (still running, still this worker, +// still this epoch), so reusing the script means the runner's terminal +// writes now go through the exact same compare-and-set family as +// renewal and reclaim. See the file comment above for exactly what that +// does and does not close. +// +// lease_epoch, lease_expires_at, worker_id, and heartbeat_at are copied +// from the entity Go just read, never from j: j is the caller's stale +// snapshot, and every renewal since it was taken has pushed the real +// expiry forward. Writing j's copy of any of those back would roll the +// current holder's lease backwards even though the script's epoch check +// passes — see job.LeaseStore.UpdateLeasedJob. +// +// It also mirrors UpdateJob's queue-index discipline (store/redis/job.go), +// which this method must not skip just because its own write is +// conditional. DequeueJobs claims a job with a ZPopMin off queueKey and +// never puts the member back on its own — see dequeue.go — so ANY write +// that lands a job in a runnable state (pending or retrying, the outcome +// of scheduleRetry) must restore the queue member, or the job becomes +// permanently unreachable: visible through GetJob, claimed by no future +// dequeue. handleSuccess and sendToDLQ happen to write terminal states, +// which never need the index touched at all, but that is a property of +// today's callers, not a license for this method to assume it. +func (s *Store) UpdateLeasedJob(ctx context.Context, j *job.Job, workerID id.WorkerID, epoch int) error { + key := jobKey(j.ID.String()) + + var cur jobEntity + if getErr := s.getEntity(ctx, key, &cur); getErr != nil { + if isNotFound(getErr) { + return dispatch.ErrJobNotFound + } + return fmt.Errorf("dispatch/redis: update leased job get: %w", getErr) + } + + next, err := toJobEntity(j) + if err != nil { + return err + } + + next.LeaseEpoch = cur.LeaseEpoch + next.LeaseExpiresAt = cur.LeaseExpiresAt + next.WorkerID = cur.WorkerID + next.HeartbeatAt = cur.HeartbeatAt + next.UpdatedAt = now() + + blob, marshalErr := json.Marshal(next) + if marshalErr != nil { + return fmt.Errorf("dispatch/redis: update leased job marshal: %w", marshalErr) + } + + // ZADD happens BEFORE the script; ZREM happens AFTER it — the same + // asymmetric ordering UpdateJob uses, and for the same reason. The + // stored entity is authoritative (dequeue re-checks state against + // it), so a spare member sitting ahead of a write that has not + // landed yet is inert; removing a member before the write is + // confirmed would risk stranding a job that turns out to still be + // legitimately runnable. Unlike UpdateJob's unconditional write, this + // one can fail its compare-and-set — but that failure mode is exactly + // as harmless for the ZADD side: a job whose fenced write was refused + // is still 'running', which was never a queue member to begin with + // (the claim popped it), so a spare pending/retrying-scored member + // pointing at a running entity is inert until dequeue's own state + // check discards it. + jID := j.ID.String() + qk := queueKey(next.Queue) + runnable := job.State(next.State) == job.StatePending || job.State(next.State) == job.StateRetrying + + if runnable { + z := goredis.Z{Score: jobScore(next.Priority, next.RunAt), Member: jID} + if zErr := s.rdb.ZAdd(ctx, qk, z).Err(); zErr != nil { + return fmt.Errorf("dispatch/redis: update leased job index add: %w", zErr) + } + } + + res, err := renewLeaseScript.Run(ctx, s.rdb, + []string{key}, + workerID.String(), + epoch, + blob, + ).Int64() + if err != nil && !errors.Is(err, goredis.Nil) { + return fmt.Errorf("dispatch/redis: update leased job: %w", err) + } + if res != 1 { + // The read above found the row, so a failed compare-and-set here + // means the lease moved on between that read and the script + // running — not that the row is gone. dispatch.ErrJobNotFound is + // reserved for the case caught above, where the row was already + // missing. Nothing was written, so any ZADD above is left as the + // harmless spare member described there; there is nothing to undo. + return job.ErrLeaseLost + } + + if !runnable { + if zErr := s.rdb.ZRem(ctx, qk, jID).Err(); zErr != nil { + return fmt.Errorf("dispatch/redis: update leased job index remove: %w", zErr) + } + } + + return nil +} diff --git a/store/redis/lease_test.go b/store/redis/lease_test.go new file mode 100644 index 0000000..415a050 --- /dev/null +++ b/store/redis/lease_test.go @@ -0,0 +1,247 @@ +package redis_test + +import ( + "context" + "testing" + "time" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/storetest" +) + +func TestLeaseConformance(t *testing.T) { + // One container, shared keyspace — do not use openReapRedis here, which + // calls startRedis on every invocation and would spin twelve containers. + connStr := startRedis(t) + + storetest.RunLeaseSuite(t, func(t *testing.T) storetest.LeaseStore { + t.Helper() + + return openRedisStore(t, connStr) + }) +} + +// TestLeaseLargeDurationRoundTrip is a regression test for a corruption +// bug in the original decode-mutate-cjson.encode-SET implementation of +// RenewLease and ReclaimExpiredLeases: cjson represents every JSON number +// as a Lua double, and doubles stop representing int64 nanosecond +// durations exactly past 2^53 (~104 days). Redis's cjson goes further and +// silently reformats a value that large as scientific notation (e.g. +// `1e+16`) on encode, which encoding/json then refuses to parse back into +// an int64 at all — not a rounding error, a hard unmarshal failure on the +// very next read of the row. +// +// None of the 20 conformance cases exercise a duration anywhere near that +// size, which is why the suite never caught it. This test uses a +// Timeout/LeaseTTL of 200 days (comfortably past the 2^53ns boundary) and +// asserts both fields come back byte-for-byte exact after a renewal and +// after a reclaim — plus that the reclaimed job is actually back on the +// queue, since the bug's failure mode aborted ReclaimExpiredLeases before +// it reached the requeue step. +func TestLeaseLargeDurationRoundTrip(t *testing.T) { + s := openReapRedis(t) + ctx := context.Background() + + // 200 days in nanoseconds is ~1.728e16, comfortably past 2^53 + // (~9.007e15, ~104.25 days) where cjson's double-precision numbers + // stop representing int64 nanosecond counts exactly. + const bigDuration = 200 * 24 * time.Hour + const queue = "lease-large-duration" + + j := storetest.PendingJob("large-duration", queue, bigDuration) + j.Timeout = bigDuration + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + worker := id.NewWorkerID() + now := time.Now().UTC() + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + if got[0].Timeout != bigDuration { + t.Fatalf("Timeout after dequeue = %v, want %v", got[0].Timeout, bigDuration) + } + if got[0].LeaseTTL != bigDuration { + t.Fatalf("LeaseTTL after dequeue = %v, want %v", got[0].LeaseTTL, bigDuration) + } + + // 1. Renew the lease, then read the job back and check both large + // durations survived byte-for-byte. + if renewErr := s.RenewLease(ctx, got[0].ID, worker, got[0].LeaseEpoch, now.Add(time.Hour)); renewErr != nil { + t.Fatalf("RenewLease: %v", renewErr) + } + + afterRenew, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get after renew: %v", err) + } + if afterRenew.Timeout != bigDuration { + t.Errorf("Timeout after renew = %v, want %v (exact)", afterRenew.Timeout, bigDuration) + } + if afterRenew.LeaseTTL != bigDuration { + t.Errorf("LeaseTTL after renew = %v, want %v (exact)", afterRenew.LeaseTTL, bigDuration) + } + + // Force the lease into the past so it is eligible for reclamation. + // RenewLease only checks state/worker/epoch, not that leaseUntil is + // in the future, so this is a legitimate way to simulate expiry + // without reaching into the store's internals. + expired := now.Add(-time.Second) + if renewErr := s.RenewLease(ctx, j.ID, worker, afterRenew.LeaseEpoch, expired); renewErr != nil { + t.Fatalf("RenewLease into the past: %v", renewErr) + } + + // 2. Reclaim the expired lease, then read the job back and check both + // large durations again, and confirm it is actually back on the queue. + reclaimed, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + if !storetest.Contains(reclaimed, j.ID) { + t.Fatalf("reclaimed set does not contain %s", j.ID) + } + + afterReclaim, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get after reclaim: %v", err) + } + if afterReclaim.Timeout != bigDuration { + t.Errorf("Timeout after reclaim = %v, want %v (exact)", afterReclaim.Timeout, bigDuration) + } + if afterReclaim.LeaseTTL != bigDuration { + t.Errorf("LeaseTTL after reclaim = %v, want %v (exact)", afterReclaim.LeaseTTL, bigDuration) + } + + requeued, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil { + t.Fatalf("DequeueJobs after reclaim: %v", err) + } + if !storetest.Contains(requeued, j.ID) { + t.Fatalf("reclaimed job %s was not requeued", j.ID) + } +} + +// TestReclaimAdoptsPreUpgradeRunningJobs covers the fleet upgrade: a job +// that was already running when the lease feature shipped has no lease +// expiry at all, so job.Lease.IsExpired reports false for it (a zero +// expiry means "never leased", not "expired"), the pool no longer calls +// ReapStaleJobs for a lease-capable store, and dequeue claims only pending +// and retrying rows. Nothing would ever look at such a job again. +// +// Redis gets no backfill because it has no migration mechanism to hang one +// on — Migrate is a no-op — so reclamation itself carries a narrow +// compatibility clause instead. The cases below are the boundary of that +// clause, and the negative ones matter more than the positive ones: a +// null expiry does NOT by itself mean the job is abandoned. Any caller +// using job.Store directly without lease options claims a perfectly +// healthy running job with no lease at all (DequeueOpts.Grants() is false +// when LeaseUntil is zero), and stealing that job would be a far worse bug +// than the one being fixed. +func TestReclaimAdoptsPreUpgradeRunningJobs(t *testing.T) { + s := openReapRedis(t) + ctx := context.Background() + now := time.Now().UTC() + + // withHeartbeat returns a pre-upgrade running job — no lease fields — + // whose last heartbeat is beatAgo old. + withHeartbeat := func(name string, startedAgo, beatAgo time.Duration) *job.Job { + j := runningJob(name, startedAgo) + beat := now.Add(-beatAgo) + j.HeartbeatAt = &beat + + return j + } + + // A job with neither timestamp set, so there is nothing to measure age + // against. Reclaiming on a null expiry alone would take it; the + // staleness gate is what stops it. + withoutTimes := func(name string) *job.Job { + j := runningJob(name, 0) + j.StartedAt = nil + + return j + } + + cases := []struct { + j *job.Job + want bool + why string + }{ + { + j: withHeartbeat("stale-heartbeat", 30*time.Minute, 20*time.Minute), + want: true, + why: "abandoned by a worker that stopped reporting; this is the bug being fixed", + }, + { + // The one that protects a live no-lease caller. Note the start + // time is old: only the heartbeat says this worker is alive, so + // this also pins that heartbeat_at takes precedence over + // started_at rather than both having to be fresh. + j: withHeartbeat("fresh-heartbeat", 30*time.Minute, 0), + want: false, + why: "still reporting, so it belongs to a healthy worker", + }, + { + j: runningJob("no-heartbeat-old-start", 20*time.Minute), + want: true, + why: "claimed long ago and never heartbeated: died before its first beat", + }, + { + j: runningJob("no-heartbeat-fresh-start", 0), + want: false, + why: "just claimed; its first heartbeat is not due yet", + }, + { + j: withoutTimes("no-times"), + want: false, + why: "no timestamp to establish age from", + }, + } + + for _, c := range cases { + if err := s.EnqueueJob(ctx, c.j); err != nil { + t.Fatalf("enqueue %s: %v", c.j.Name, err) + } + } + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + + for _, c := range cases { + got := storetest.Contains(reclaimed, c.j.ID) + if got == c.want { + continue + } + if c.want { + t.Errorf("%s was not reclaimed but should have been: %s", c.j.Name, c.why) + } else { + t.Errorf("%s was reclaimed but must not be: %s", c.j.Name, c.why) + } + } +} + +func TestDLQConformance(t *testing.T) { + connStr := startRedis(t) + + storetest.RunDLQSuite(t, func(t *testing.T) storetest.DLQStore { + t.Helper() + + return openRedisStore(t, connStr) + }) +} diff --git a/store/redis/requeue_test.go b/store/redis/requeue_test.go new file mode 100644 index 0000000..20d715e --- /dev/null +++ b/store/redis/requeue_test.go @@ -0,0 +1,225 @@ +package redis_test + +import ( + "context" + "testing" + "time" + + goredis "github.com/redis/go-redis/v9" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + redisstore "github.com/xraph/dispatch/store/redis" +) + +// This file pins the invariant that separates this backend from the other +// four: the queue sorted set is a CANDIDATE INDEX, and a job absent from it +// is never even considered, whatever its state says. Postgres, SQLite, +// Mongo and memory have no such structure — they re-derive candidacy per +// query from the row itself, so returning a job to pending is enough there +// and means nothing here. +// +// dequeue.go re-checks state and RunAt against the decoded entity, so the +// index does not decide ELIGIBILITY. It decides VISIBILITY, and it is +// maintained by only three writers: EnqueueJob and ReclaimExpiredLeases add, +// the claim removes. UpdateJob is the fourth writer every post-claim state +// transition flows through — Runner.scheduleRetry, Pool.requeueRateLimited, +// Pool.requeueUndispatched and Pool.reapStaleJobs — and it is the one this +// file exists for. +// +// The suites below assert on DequeueJobs rather than on the sorted set, +// because "the job runs again" is the property those four callers depend +// on. TestUpdateJob_MaintainsQueueIndex is the deliberate exception, and +// says why. + +func requeueJob(name, queue string, runAt time.Time) *job.Job { + return &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: name, + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: runAt, + } +} + +// dueNow is a RunAt far enough in the past to be unambiguously due, since +// dequeue gates on run_at <= now. +func dueNow() time.Time { return time.Now().UTC().Add(-time.Second) } + +// dequeueOnce claims from queue and reports whether jobID came back. +func dequeueOnce(t *testing.T, s *redisstore.Store, queue string, jobID id.JobID) bool { + t.Helper() + + got, err := s.DequeueJobs(context.Background(), job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + }) + if err != nil { + t.Fatalf("dequeue from %s: %v", queue, err) + } + + for _, j := range got { + if j.ID == jobID { + return true + } + } + + return false +} + +// TestUpdateJob_ReturnsRunnableJobToQueue covers every path that hands a +// claimed job back for another attempt. All four express themselves +// identically — set a runnable state on a job that has already been +// claimed, then UpdateJob — so all four stand or fall on this one rule. +// +// The retrying case is the load-bearing one: without it, no retry ever runs +// a second time on this backend. +func TestUpdateJob_ReturnsRunnableJobToQueue(t *testing.T) { + s := openRedisStore(t, startRedis(t)) + ctx := context.Background() + + tests := []struct { + name string + state job.State + queue string + runnable bool + }{ + {name: "pending", state: job.StatePending, queue: "requeue-pending", runnable: true}, + {name: "retrying", state: job.StateRetrying, queue: "requeue-retrying", runnable: true}, + {name: "completed", state: job.StateCompleted, queue: "requeue-completed", runnable: false}, + {name: "failed", state: job.StateFailed, queue: "requeue-failed", runnable: false}, + {name: "cancelled", state: job.StateCancelled, queue: "requeue-cancelled", runnable: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + j := requeueJob(tt.name, tt.queue, dueNow()) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + // The claim is what removes the job from the index, so it is a + // precondition of the bug, not incidental setup. + if !dequeueOnce(t, s, tt.queue, j.ID) { + t.Fatalf("precondition: freshly enqueued job was not claimed") + } + + j.State = tt.state + j.RunAt = dueNow() + if err := s.UpdateJob(ctx, j); err != nil { + t.Fatalf("update to %s: %v", tt.state, err) + } + + if got := dequeueOnce(t, s, tt.queue, j.ID); got != tt.runnable { + t.Fatalf("after UpdateJob to %s: dequeued=%v, want %v", tt.state, got, tt.runnable) + } + }) + } +} + +// TestUpdateJob_RequeuedRetryWaitsForBackoff guards the half of the fix +// that is easy to lose. Restoring a retry's visibility must not also make +// it due: scheduleRetry sets RunAt to now+backoff, and nothing in worker/ +// re-checks RunAt after the claim, so the store is the only thing standing +// between a failing job and a hot retry loop. +// +// The two dequeues differ only in the job's RunAt, which is what makes this +// a test of the time gate rather than of the index. +func TestUpdateJob_RequeuedRetryWaitsForBackoff(t *testing.T) { + const queue = "requeue-backoff" + + s := openRedisStore(t, startRedis(t)) + ctx := context.Background() + + j := requeueJob("backoff", queue, dueNow()) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + if !dequeueOnce(t, s, queue, j.ID) { + t.Fatalf("precondition: freshly enqueued job was not claimed") + } + + // Exactly what scheduleRetry writes: retrying, due after a backoff. + j.State = job.StateRetrying + j.RetryCount = 1 + j.RunAt = time.Now().UTC().Add(time.Hour) + if err := s.UpdateJob(ctx, j); err != nil { + t.Fatalf("update to retrying: %v", err) + } + + if dequeueOnce(t, s, queue, j.ID) { + t.Fatal("retry was claimed before its backoff elapsed") + } + + // Standing in for the backoff elapsing, so the test needs no sleep. + j.RunAt = dueNow() + if err := s.UpdateJob(ctx, j); err != nil { + t.Fatalf("update run_at: %v", err) + } + + if !dequeueOnce(t, s, queue, j.ID) { + t.Fatal("retry was not claimed after its backoff elapsed") + } +} + +// TestUpdateJob_MaintainsQueueIndex is the one white-box test here, because +// the property is itself white-box: a job that reaches a terminal state +// without ever being claimed keeps its index member, and nothing else ever +// removes it. dequeue.go's state gate means that member is harmless to +// CORRECTNESS, which is exactly why no behavioural assertion can see it — +// and why it would otherwise accumulate, one dead member per cancelled job, +// forever, inside the bounded window scanQueue reads from the head of the +// index. +// +// Asserting presence before asserting absence is deliberate: a test that +// only checked absence would pass just as happily against a mistyped key +// that never existed. +func TestUpdateJob_MaintainsQueueIndex(t *testing.T) { + const ( + queue = "index-hygiene" + queueKey = "dispatch:queue:" + queue + ) + + connStr := startRedis(t) + s := openRedisStore(t, connStr) + ctx := context.Background() + + opt, err := goredis.ParseURL(connStr) + if err != nil { + t.Fatalf("parse redis url: %v", err) + } + rdb := goredis.NewClient(opt) + t.Cleanup(func() { _ = rdb.Close() }) + + indexed := func() bool { + t.Helper() + + members, mErr := rdb.ZRange(ctx, queueKey, 0, -1).Result() + if mErr != nil { + t.Fatalf("read queue index: %v", mErr) + } + + return len(members) > 0 + } + + j := requeueJob("cancelled-before-claim", queue, dueNow()) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + if !indexed() { + t.Fatalf("precondition: enqueued job is not in %s", queueKey) + } + + j.State = job.StateCancelled + if err := s.UpdateJob(ctx, j); err != nil { + t.Fatalf("update to cancelled: %v", err) + } + + if indexed() { + t.Fatal("cancelled job still holds a member in the queue index") + } +} diff --git a/store/redis/resource_test.go b/store/redis/resource_test.go new file mode 100644 index 0000000..bf1e817 --- /dev/null +++ b/store/redis/resource_test.go @@ -0,0 +1,152 @@ +//go:build integration + +package redis_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// TestRedisRoundTripsResources mirrors store/postgres's and store/sqlite's +// resource round-trip test: a job with cpu/memory/a custom key/limits/ +// class/input signal must come back identical, and the custom key +// specifically must survive since it only lives in the full-fidelity +// resource_requests JSON, not the scalar fields. +func TestRedisRoundTripsResources(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "tessellate.model", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + Resources: resource.Set{ + resource.CPU: 4000, resource.Memory: 16 << 30, "fpga": 2, + }, + ResourceLimits: resource.Set{resource.Memory: 16 << 30}, + ResourceClass: "heavy", + InputBytes: 4 << 30, + PrimaryInputHash: "blake3:9f2a", + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + + if got.Resources[resource.CPU] != 4000 { + t.Errorf("cpu = %d, want 4000", got.Resources[resource.CPU]) + } + if got.Resources["fpga"] != 2 { + t.Errorf("custom key lost in round-trip: %v", got.Resources) + } + if got.ResourceLimits[resource.Memory] != 16<<30 { + t.Errorf("limits = %v", got.ResourceLimits) + } + if got.ResourceClass != "heavy" { + t.Errorf("class = %q, want heavy", got.ResourceClass) + } + if got.InputBytes != 4<<30 || got.PrimaryInputHash != "blake3:9f2a" { + t.Errorf("input signal lost: bytes=%d hash=%q", + got.InputBytes, got.PrimaryInputHash) + } +} + +// TestRedisJobWithNoResourcesRoundTrips pins the backward-compatibility +// contract: a job enqueued without any resource declaration must come back +// with a zero Set, not a Set containing zero-valued canonical keys. +func TestRedisJobWithNoResourcesRoundTrips(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "notify.user", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if !got.Resources.IsZero() { + t.Errorf("Resources = %v, want zero", got.Resources) + } +} + +// TestRedisUndeclaredJobHasNoResourceKeyInRawJSON closes the gap a +// Resources.IsZero() assertion cannot: decoding an absent JSON key and +// decoding an empty "{}" both yield IsZero() == true, so that assertion +// alone does not prove the stored value is genuinely absent rather than +// an empty document. This test reads the raw JSON blob straight from the +// KV store, bypassing jobEntity entirely, and asserts the +// resource_requests/resource_limits keys are not present at all -- the +// Redis analogue of the SQL backends' NULL-column proof. +func TestRedisUndeclaredJobHasNoResourceKeyInRawJSON(t *testing.T) { + s := setupTestStore(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "notify.user", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + // "dispatch:job:" mirrors the unexported jobKey format in + // keys.go (keyPrefix + "job:" + id) -- there is no other way to + // read the raw entity from outside the package. + raw, err := s.KV().GetRaw(ctx, "dispatch:job:"+j.ID.String()) + if err != nil { + t.Fatalf("GetRaw() error = %v", err) + } + + var fields map[string]json.RawMessage + if err := json.Unmarshal(raw, &fields); err != nil { + t.Fatalf("unmarshal raw entity: %v", err) + } + + if _, ok := fields["resource_requests"]; ok { + t.Errorf("resource_requests present in raw JSON: %s", fields["resource_requests"]) + } + if _, ok := fields["resource_limits"]; ok { + t.Errorf("resource_limits present in raw JSON: %s", fields["resource_limits"]) + } + + // The scalar/class/hash fields are always-present in the SQL + // backends (NOT NULL DEFAULT), so their Redis analogues should + // still be written -- only the full-fidelity JSON is + // absent-when-zero. + if v, ok := fields["req_cpu_milli"]; !ok || string(v) != "0" { + t.Errorf("req_cpu_milli = %s, want present and 0", v) + } + if v, ok := fields["req_custom_keys"]; !ok || string(v) != `""` { + t.Errorf("req_custom_keys = %s, want present and empty", v) + } +} diff --git a/store/redis/store.go b/store/redis/store.go index 1d16489..5e21227 100644 --- a/store/redis/store.go +++ b/store/redis/store.go @@ -14,6 +14,7 @@ import ( "github.com/xraph/grove/kv" "github.com/xraph/grove/kv/drivers/redisdriver" + "github.com/xraph/dispatch/artifact" "github.com/xraph/dispatch/cluster" "github.com/xraph/dispatch/cron" "github.com/xraph/dispatch/dlq" @@ -25,11 +26,13 @@ import ( // Compile-time interface checks. var ( _ job.Store = (*Store)(nil) + _ job.LeaseStore = (*Store)(nil) _ workflow.Store = (*Store)(nil) _ cron.Store = (*Store)(nil) _ dlq.Store = (*Store)(nil) _ event.Store = (*Store)(nil) _ cluster.Store = (*Store)(nil) + _ artifact.Store = (*Store)(nil) ) // Option configures the Store. diff --git a/store/redis/store_test.go b/store/redis/store_test.go index b06779e..274809c 100644 --- a/store/redis/store_test.go +++ b/store/redis/store_test.go @@ -156,7 +156,7 @@ func TestJobStore_DequeueJobs(t *testing.T) { } // Dequeue 2 -- should get highest priority first. - dequeued, err := s.DequeueJobs(ctx, []string{"default"}, 2) + dequeued, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 2}) if err != nil { t.Fatalf("dequeue: %v", err) } @@ -171,7 +171,7 @@ func TestJobStore_DequeueJobs(t *testing.T) { } // Dequeue remaining -- should get 1 job. - remaining, err := s.DequeueJobs(ctx, []string{"default"}, 10) + remaining, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 10}) if err != nil { t.Fatalf("dequeue remaining: %v", err) } diff --git a/store/sqlite/artifact.go b/store/sqlite/artifact.go new file mode 100644 index 0000000..c1dea90 --- /dev/null +++ b/store/sqlite/artifact.go @@ -0,0 +1,524 @@ +package sqlite + +import ( + "context" + "database/sql" + "errors" + "fmt" + "strings" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// defaultSweepLimit bounds an unbounded sweep so a single pass can never +// touch an unbounded number of rows. +const defaultSweepLimit = 1000 + +const insertArtifactSQL = ` + INSERT INTO dispatch_artifacts + (id, backend, bucket, key, size, content_hash, content_type, + lifecycle, scope_app_id, scope_org_id, expires_at, created_at, deleted_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)` + +func artifactInsertArgs(a *artifact.Artifact) []any { + return []any{ + a.ID.String(), a.Backend, a.Bucket, a.Key, a.Size, + nullString(a.ContentHash), nullString(a.ContentType), + string(a.Lifecycle), nullString(a.ScopeAppID), nullString(a.ScopeOrgID), + a.ExpiresAt, a.CreatedAt, a.DeletedAt, + } +} + +const insertLinkSQL = ` + INSERT INTO dispatch_artifact_links + (artifact_id, owner_kind, owner_id, name, attempt, role, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (artifact_id, owner_kind, owner_id, name, attempt) DO NOTHING` + +func linkInsertArgs(l *artifact.Link) []any { + return []any{ + l.ArtifactID.String(), string(l.OwnerKind), l.OwnerID, + l.Name, l.Attempt, string(l.Role), l.CreatedAt, + } +} + +// nullString maps the empty string to SQL NULL so nullable text columns +// stay NULL rather than storing an empty value. +func nullString(s string) any { + if s == "" { + return nil + } + + return s +} + +// CreateArtifact inserts an artifact and, when link is non-nil, its first +// link in a single transaction. +func (s *Store) CreateArtifact(ctx context.Context, a *artifact.Artifact, link *artifact.Link) error { + if link == nil { + if _, err := s.sdb.Exec(ctx, insertArtifactSQL, artifactInsertArgs(a)...); err != nil { + if isDuplicateKey(err) { + return artifact.ErrExists + } + + return fmt.Errorf("dispatch/sqlite: create artifact: %w", err) + } + + return nil + } + + tx, err := s.sdb.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("dispatch/sqlite: create artifact: begin: %w", err) + } + + defer func() { + if rerr := tx.Rollback(); rerr != nil && !errors.Is(rerr, sql.ErrTxDone) { + s.logger.Warn("dispatch/sqlite: artifact tx rollback", log.String("error", rerr.Error())) + } + }() + + if _, err := tx.Exec(ctx, insertArtifactSQL, artifactInsertArgs(a)...); err != nil { + if isDuplicateKey(err) { + return artifact.ErrExists + } + + return fmt.Errorf("dispatch/sqlite: create artifact: %w", err) + } + + if _, err := tx.Exec(ctx, insertLinkSQL, linkInsertArgs(link)...); err != nil { + return fmt.Errorf("dispatch/sqlite: create artifact link: %w", err) + } + + if err := tx.Commit(); err != nil { + return fmt.Errorf("dispatch/sqlite: create artifact: commit: %w", err) + } + + return nil +} + +// GetArtifact retrieves a live artifact by ID. +func (s *Store) GetArtifact(ctx context.Context, artifactID id.ArtifactID) (*artifact.Artifact, error) { + m := new(artifactModel) + + err := s.sdb.NewSelect(m). + Where("id = ?", artifactID.String()). + Where("deleted_at IS NULL"). + Limit(1). + Scan(ctx) + if err != nil { + if isNoRows(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf("dispatch/sqlite: get artifact: %w", err) + } + + return fromArtifactModel(m) +} + +// FindArtifactByKey retrieves a live artifact by its storage coordinates. +func (s *Store) FindArtifactByKey(ctx context.Context, backend, bucket, key string) (*artifact.Artifact, error) { + m := new(artifactModel) + + err := s.sdb.NewSelect(m). + Where("backend = ?", backend). + Where("bucket = ?", bucket). + Where("key = ?", key). + Where("deleted_at IS NULL"). + Limit(1). + Scan(ctx) + if err != nil { + if isNoRows(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf("dispatch/sqlite: find artifact by key: %w", err) + } + + return fromArtifactModel(m) +} + +// UpdateArtifact persists size, hash, content type, and expiry. Lifecycle, +// created_at, and deleted_at are deliberately not updatable here. +func (s *Store) UpdateArtifact(ctx context.Context, a *artifact.Artifact) error { + res, err := s.sdb.Exec(ctx, ` + UPDATE dispatch_artifacts + SET size = ?, content_hash = ?, content_type = ?, expires_at = ? + WHERE id = ?`, + a.Size, nullString(a.ContentHash), nullString(a.ContentType), a.ExpiresAt, a.ID.String(), + ) + if err != nil { + return fmt.Errorf("dispatch/sqlite: update artifact: %w", err) + } + + if n, rerr := res.RowsAffected(); rerr == nil && n == 0 { + return artifact.ErrNotFound + } + + return nil +} + +// ListArtifacts returns artifacts matching the given options, newest first. +func (s *Store) ListArtifacts(ctx context.Context, opts artifact.ListOpts) ([]*artifact.Artifact, error) { + var models []artifactModel + + q := s.sdb.NewSelect(&models) + + if !opts.IncludeDeleted { + q = q.Where("deleted_at IS NULL") + } + + if opts.Lifecycle != "" { + q = q.Where("lifecycle = ?", string(opts.Lifecycle)) + } + + if opts.ScopeAppID != "" { + q = q.Where("scope_app_id = ?", opts.ScopeAppID) + } + + if opts.ScopeOrgID != "" { + q = q.Where("scope_org_id = ?", opts.ScopeOrgID) + } + + q = q.OrderExpr("created_at DESC, id ASC") + + if opts.Limit > 0 { + q = q.Limit(opts.Limit) + } + + if opts.Offset > 0 { + q = q.Offset(opts.Offset) + } + + if err := q.Scan(ctx); err != nil { + return nil, fmt.Errorf("dispatch/sqlite: list artifacts: %w", err) + } + + return fromArtifactModels(models) +} + +// LinkArtifact records that an owner references an artifact, idempotently. +func (s *Store) LinkArtifact(ctx context.Context, link *artifact.Link) error { + if _, err := s.sdb.Exec(ctx, insertLinkSQL, linkInsertArgs(link)...); err != nil { + return fmt.Errorf("dispatch/sqlite: link artifact: %w", err) + } + + return nil +} + +// ListLinks returns every link belonging to the given owner. +func (s *Store) ListLinks(ctx context.Context, owner artifact.OwnerRef) ([]*artifact.Link, error) { + var models []artifactLinkModel + + err := s.sdb.NewSelect(&models). + Where("owner_kind = ?", string(owner.Kind)). + Where("owner_id = ?", owner.ID). + OrderExpr("name ASC, attempt ASC"). + Scan(ctx) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: list links: %w", err) + } + + return fromLinkModels(models) +} + +// FindLinkByName returns the highest-attempt link for an owner and name, +// breaking ties by created_at descending — see the artifact.Store doc +// comment for why ties happen and what the tie-break does and does not +// fix. +func (s *Store) FindLinkByName( + ctx context.Context, + owner artifact.OwnerRef, + name string, +) (*artifact.Link, error) { + m := new(artifactLinkModel) + + err := s.sdb.NewSelect(m). + Where("owner_kind = ?", string(owner.Kind)). + Where("owner_id = ?", owner.ID). + Where("name = ?", name). + OrderExpr("attempt DESC, created_at DESC"). + Limit(1). + Scan(ctx) + if err != nil { + if isNoRows(err) { + return nil, artifact.ErrNotFound + } + + return nil, fmt.Errorf("dispatch/sqlite: find link by name: %w", err) + } + + return fromLinkModel(m) +} + +// ListArtifactsByOwner returns live artifacts linked to an owner. +func (s *Store) ListArtifactsByOwner( + ctx context.Context, + owner artifact.OwnerRef, + role artifact.Role, +) ([]*artifact.Artifact, error) { + var models []artifactModel + + query := ` + SELECT DISTINCT a.* FROM dispatch_artifacts a + JOIN dispatch_artifact_links l ON l.artifact_id = a.id + WHERE l.owner_kind = ? AND l.owner_id = ? AND a.deleted_at IS NULL` + + args := []any{string(owner.Kind), owner.ID} + + if role != "" { + query += ` AND l.role = ?` + + args = append(args, string(role)) + } + + if err := s.sdb.NewRaw(query, args...).Scan(ctx, &models); err != nil { + return nil, fmt.Errorf("dispatch/sqlite: list artifacts by owner: %w", err) + } + + return fromArtifactModels(models) +} + +// eligibleEphemeralSQL selects ephemeral artifacts whose every linked +// owner is terminal and whose retention window has elapsed. +// +// The lifecycle predicate is a literal. It is never bound from a +// parameter, so no caller can widen this statement to reach a durable +// artifact. +// +// SQLite has no bool_and, so the all-terminal test is expressed as +// MIN(CASE ... END) = 1. Timestamps are ISO8601 text and compare +// lexicographically, so the retention cutoff is computed in Go and bound +// as a formatted time rather than built with SQL interval arithmetic. +// +// An owner row that no longer exists counts as terminal at the link's +// creation time: its job or run was purged, so it cannot still be running. +const eligibleEphemeralSQL = ` + SELECT a.id + FROM dispatch_artifacts a + JOIN dispatch_artifact_links l ON l.artifact_id = a.id + LEFT JOIN dispatch_jobs j + ON l.owner_kind = 'job' AND j.id = l.owner_id + LEFT JOIN dispatch_workflow_runs r + ON l.owner_kind IN ('run', 'step') AND r.id = l.owner_id + WHERE a.lifecycle = 'ephemeral' + AND a.deleted_at IS NULL + GROUP BY a.id, a.expires_at + HAVING MIN( + CASE + WHEN l.owner_kind = 'job' + THEN CASE WHEN j.id IS NULL + OR j.state IN ('completed', 'failed', 'cancelled') THEN 1 ELSE 0 END + WHEN l.owner_kind IN ('run', 'step') + THEN CASE WHEN r.id IS NULL + OR r.state IN ('completed', 'failed', 'cancelled') THEN 1 ELSE 0 END + ELSE 1 + END + ) = 1 + AND ( + CASE + WHEN a.expires_at IS NOT NULL THEN a.expires_at <= ? + ELSE MAX( + COALESCE(j.completed_at, j.updated_at, r.completed_at, r.updated_at, l.created_at) + ) <= ? + END + ) + LIMIT ?` + +// SweepEphemeral marks eligible ephemeral artifacts as deleted. +func (s *Store) SweepEphemeral( + ctx context.Context, + opts artifact.SweepOpts, +) ([]*artifact.Artifact, error) { + limit := opts.Limit + if limit <= 0 { + limit = defaultSweepLimit + } + + now := time.Now().UTC() + cutoff := now.Add(-opts.Retention) + + ids, err := s.selectIDs(ctx, eligibleEphemeralSQL, now, cutoff, limit) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: sweep ephemeral: %w", err) + } + + if len(ids) == 0 { + return nil, nil + } + + if opts.DryRun { + return s.artifactsByIDs(ctx, ids) + } + + if err := s.markDeleted(ctx, ids, now); err != nil { + return nil, fmt.Errorf("dispatch/sqlite: sweep ephemeral: %w", err) + } + + return s.artifactsByIDs(ctx, ids) +} + +// SweepOrphans marks link-less ephemeral artifacts created before cutoff. +func (s *Store) SweepOrphans( + ctx context.Context, + cutoff time.Time, + limit int, +) ([]*artifact.Artifact, error) { + if limit <= 0 { + limit = defaultSweepLimit + } + + // lifecycle = 'ephemeral' is a literal for the same reason as above. + const query = ` + SELECT a.id FROM dispatch_artifacts a + WHERE a.lifecycle = 'ephemeral' + AND a.deleted_at IS NULL + AND a.created_at < ? + AND NOT EXISTS ( + SELECT 1 FROM dispatch_artifact_links l WHERE l.artifact_id = a.id + ) + ORDER BY a.created_at ASC + LIMIT ?` + + ids, err := s.selectIDs(ctx, query, cutoff, limit) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: sweep orphans: %w", err) + } + + if len(ids) == 0 { + return nil, nil + } + + if err := s.markDeleted(ctx, ids, time.Now().UTC()); err != nil { + return nil, fmt.Errorf("dispatch/sqlite: sweep orphans: %w", err) + } + + return s.artifactsByIDs(ctx, ids) +} + +// ListPurgeable returns soft-deleted artifacts older than grace. +func (s *Store) ListPurgeable( + ctx context.Context, + grace time.Duration, + limit int, +) ([]*artifact.Artifact, error) { + if limit <= 0 { + limit = defaultSweepLimit + } + + var models []artifactModel + + cutoff := time.Now().UTC().Add(-grace) + + err := s.sdb.NewSelect(&models). + Where("deleted_at IS NOT NULL"). + Where("deleted_at <= ?", cutoff). + OrderExpr("deleted_at ASC"). + Limit(limit). + Scan(ctx) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: list purgeable: %w", err) + } + + return fromArtifactModels(models) +} + +// PurgeArtifact hard-deletes an artifact. Links cascade. +func (s *Store) PurgeArtifact(ctx context.Context, artifactID id.ArtifactID) error { + if _, err := s.sdb.Exec(ctx, + `DELETE FROM dispatch_artifact_links WHERE artifact_id = ?`, artifactID.String(), + ); err != nil { + return fmt.Errorf("dispatch/sqlite: purge artifact links: %w", err) + } + + if _, err := s.sdb.Exec(ctx, + `DELETE FROM dispatch_artifacts WHERE id = ?`, artifactID.String(), + ); err != nil { + return fmt.Errorf("dispatch/sqlite: purge artifact: %w", err) + } + + return nil +} + +// selectIDs runs a query whose first column is an artifact ID. +func (s *Store) selectIDs(ctx context.Context, query string, args ...any) ([]string, error) { + rows, err := s.sdb.Query(ctx, query, args...) + if err != nil { + return nil, err + } + + defer func() { + if cerr := rows.Close(); cerr != nil { + s.logger.Warn("dispatch/sqlite: close artifact id rows", log.String("error", cerr.Error())) + } + }() + + var ids []string + + for rows.Next() { + var got string + if serr := rows.Scan(&got); serr != nil { + return nil, serr + } + + ids = append(ids, got) + } + + return ids, rows.Err() +} + +// markDeleted soft-deletes the given artifacts. The lifecycle literal is +// repeated here so the write itself, not only the selection above it, +// refuses to touch a durable artifact. +func (s *Store) markDeleted(ctx context.Context, ids []string, at time.Time) error { + query := ` + UPDATE dispatch_artifacts + SET deleted_at = ? + WHERE lifecycle = 'ephemeral' + AND deleted_at IS NULL + AND id IN (` + placeholders(len(ids)) + `)` + + args := make([]any, 0, len(ids)+1) + args = append(args, at) + + for _, got := range ids { + args = append(args, got) + } + + _, err := s.sdb.Exec(ctx, query, args...) + + return err +} + +// artifactsByIDs loads artifacts by ID, including soft-deleted ones, so a +// sweep can return what it just marked. +func (s *Store) artifactsByIDs(ctx context.Context, ids []string) ([]*artifact.Artifact, error) { + var models []artifactModel + + query := `SELECT * FROM dispatch_artifacts WHERE id IN (` + placeholders(len(ids)) + `)` + + args := make([]any, 0, len(ids)) + for _, got := range ids { + args = append(args, got) + } + + if err := s.sdb.NewRaw(query, args...).Scan(ctx, &models); err != nil { + return nil, fmt.Errorf("dispatch/sqlite: load artifacts by id: %w", err) + } + + return fromArtifactModels(models) +} + +// placeholders builds "?, ?, ?" for an IN clause of n values. +func placeholders(n int) string { + if n == 0 { + return "NULL" + } + + return strings.TrimSuffix(strings.Repeat("?, ", n), ", ") +} diff --git a/store/sqlite/artifact_models.go b/store/sqlite/artifact_models.go new file mode 100644 index 0000000..20a942d --- /dev/null +++ b/store/sqlite/artifact_models.go @@ -0,0 +1,114 @@ +package sqlite + +import ( + "time" + + "github.com/xraph/grove" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/id" +) + +// ── Artifact model ──────────────────────────────────────────────── + +type artifactModel struct { + grove.BaseModel `grove:"table:dispatch_artifacts"` + + ID string `grove:"id,pk"` + Backend string `grove:"backend,notnull"` + Bucket string `grove:"bucket,notnull"` + Key string `grove:"key,notnull"` + Size int64 `grove:"size,notnull,default:0"` + ContentHash string `grove:"content_hash"` + ContentType string `grove:"content_type"` + Lifecycle string `grove:"lifecycle,notnull"` + ScopeAppID string `grove:"scope_app_id"` + ScopeOrgID string `grove:"scope_org_id"` + ExpiresAt *time.Time `grove:"expires_at"` + CreatedAt time.Time `grove:"created_at,notnull,default:current_timestamp"` + DeletedAt *time.Time `grove:"deleted_at"` +} + +func fromArtifactModel(m *artifactModel) (*artifact.Artifact, error) { + aid, err := id.ParseArtifactID(m.ID) + if err != nil { + return nil, err + } + + return &artifact.Artifact{ + ID: aid, + Backend: m.Backend, + Bucket: m.Bucket, + Key: m.Key, + Size: m.Size, + ContentHash: m.ContentHash, + ContentType: m.ContentType, + Lifecycle: artifact.Lifecycle(m.Lifecycle), + ScopeAppID: m.ScopeAppID, + ScopeOrgID: m.ScopeOrgID, + ExpiresAt: m.ExpiresAt, + CreatedAt: m.CreatedAt, + DeletedAt: m.DeletedAt, + }, nil +} + +func fromArtifactModels(models []artifactModel) ([]*artifact.Artifact, error) { + out := make([]*artifact.Artifact, 0, len(models)) + + for i := range models { + a, err := fromArtifactModel(&models[i]) + if err != nil { + return nil, err + } + + out = append(out, a) + } + + return out, nil +} + +// ── Artifact link model ─────────────────────────────────────────── + +type artifactLinkModel struct { + grove.BaseModel `grove:"table:dispatch_artifact_links"` + + ArtifactID string `grove:"artifact_id,pk"` + OwnerKind string `grove:"owner_kind,pk"` + OwnerID string `grove:"owner_id,pk"` + Name string `grove:"name,pk"` + Attempt int `grove:"attempt,pk"` + Role string `grove:"role,notnull"` + CreatedAt time.Time `grove:"created_at,notnull,default:current_timestamp"` +} + +func fromLinkModel(m *artifactLinkModel) (*artifact.Link, error) { + aid, err := id.ParseArtifactID(m.ArtifactID) + if err != nil { + return nil, err + } + + return &artifact.Link{ + ArtifactID: aid, + OwnerKind: artifact.OwnerKind(m.OwnerKind), + OwnerID: m.OwnerID, + Role: artifact.Role(m.Role), + Name: m.Name, + Attempt: m.Attempt, + CreatedAt: m.CreatedAt, + }, nil +} + +func fromLinkModels(models []artifactLinkModel) ([]*artifact.Link, error) { + out := make([]*artifact.Link, 0, len(models)) + + for i := range models { + l, err := fromLinkModel(&models[i]) + if err != nil { + return nil, err + } + + out = append(out, l) + } + + return out, nil +} diff --git a/store/sqlite/artifact_test.go b/store/sqlite/artifact_test.go new file mode 100644 index 0000000..899455b --- /dev/null +++ b/store/sqlite/artifact_test.go @@ -0,0 +1,17 @@ +package sqlite_test + +import ( + "testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" +) + +// TestArtifactStoreConformance runs the shared artifact.Store suite +// against SQLite. Each subtest gets its own in-memory database because +// the suite asserts absolute row counts. +func TestArtifactStoreConformance(t *testing.T) { + artifacttest.RunStoreSuite(t, func(t *testing.T) artifact.Store { + return openSqliteStore(t) + }) +} diff --git a/store/sqlite/dequeue_sql_test.go b/store/sqlite/dequeue_sql_test.go new file mode 100644 index 0000000..83d5cc2 --- /dev/null +++ b/store/sqlite/dequeue_sql_test.go @@ -0,0 +1,402 @@ +package sqlite + +import ( + "fmt" + "strings" + "testing" + "time" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// These tests assert the compiled statement, not behaviour — +// dequeue_test.go's conformance run asks the database for behaviour. They +// exist because two properties of this statement are invisible to the +// shared suite on SQLite: +// +// - `?` binds POSITIONALLY, so the args slice must be built in the same +// order the placeholders appear in the text. Get that wrong and a +// budget is compared against a queue name. Nothing in the suite reads +// the statement, and a mis-binding shows up as a wrong answer only for +// the option combinations the suite happens to exercise. +// - The candidate SELECT's ORDER BY is what makes the LIMIT truncate an +// ORDERED set. Re-measured against the current 21-case suite: +// removing it entirely fails exactly ONE case, +// LocalityDecidesWhichRowsSurviveATightLimit. The other twenty pass, +// because SQLite answers the scan from idx_dispatch_jobs_dequeue_res, +// whose leading key order happens to match priority DESC, run_at ASC, +// so the right rows come back for the wrong reason. +// TestBuildDequeueQueryOrdersLocalityBelowPriority is the pin that +// does not depend on which index the planner chose. + +// render substitutes each bind parameter into the statement in order, so +// a test can read the finished SQL the way SQLite reads it. It is a test +// helper only: the production path never interpolates a value. +func render(t *testing.T, query string, args []any) string { + t.Helper() + + if n := strings.Count(query, "?"); n != len(args) { + t.Fatalf("statement has %d placeholders but %d args were bound:\n%s\n%v", + n, len(args), query, args) + } + + var b strings.Builder + + i := 0 + + for _, r := range query { + if r != '?' { + b.WriteRune(r) + + continue + } + + switch v := args[i].(type) { + case string: + b.WriteString("'" + strings.ReplaceAll(v, "'", "''") + "'") + case time.Time: + b.WriteString("'" + v.Format(time.RFC3339Nano) + "'") + default: + fmt.Fprintf(&b, "%v", v) + } + + i++ + } + + return b.String() +} + +func fixedNow() time.Time { + return time.Date(2026, 8, 12, 12, 0, 0, 0, time.UTC) +} + +// TestBuildDequeueQueryUnboundedEmitsOriginalStatement is the +// backward-compatibility guarantee in its narrowest form: opts that +// constrain nothing must compile to the statement that shipped before +// DequeueOpts existed, with no fit predicate at all. A worker not using +// the resource model claims everything, jobs declaring custom resources +// included. +func TestBuildDequeueQueryUnboundedEmitsOriginalStatement(t *testing.T) { + opts := job.DequeueOpts{Queues: []string{"default"}, Limit: 10} + + if !opts.IsUnbounded() { + t.Fatalf("DequeueOpts%+v.IsUnbounded() = false, want true", opts) + } + + query, args := buildDequeueQuery(opts, fixedNow()) + + for _, banned := range []string{"req_", "REPLACE", "primary_input_hash"} { + if strings.Contains(query, banned) { + t.Errorf("unbounded dequeue emitted %q:\n%s", banned, query) + } + } + + if !strings.Contains(query, "ORDER BY priority DESC, run_at ASC\n") { + t.Errorf("unbounded dequeue lost the original ordering:\n%s", query) + } + + // started_at, updated_at, one queue, run_at, limit. + if len(args) != 5 { + t.Fatalf("unbounded dequeue bound %d args, want 5: %v", len(args), args) + } + + if args[4] != 10 { + t.Errorf("limit bound as %v, want 10", args[4]) + } +} + +// TestBuildDequeueQueryAppliesLocalityToUnboundedOpts pins the split +// IsUnbounded exists to make: it governs FILTERING only. Opts carrying +// nothing but PreferHashes are unbounded, so no fit predicate is emitted — +// and the locality term is applied anyway. A backend that derived "should +// I order?" from IsUnbounded would silently drop the signal here. +func TestBuildDequeueQueryAppliesLocalityToUnboundedOpts(t *testing.T) { + opts := job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 4, + PreferHashes: []string{"blake3:staged"}, + } + + if !opts.IsUnbounded() { + t.Fatal("opts carrying only PreferHashes report IsUnbounded() = false") + } + + query, _ := buildDequeueQuery(opts, fixedNow()) + + if strings.Contains(query, "req_") { + t.Errorf("PreferHashes emitted a fit predicate — locality must never filter:\n%s", query) + } + + if !strings.Contains(query, "COALESCE(primary_input_hash IN (") { + t.Errorf("locality term missing from unbounded opts:\n%s", query) + } +} + +// TestBuildDequeueQueryOrdersLocalityBelowPriority pins two things at +// once. +// +// First, priority comes before locality. The reverse would not show up as +// a wrong answer, only as starvation: a steady stream of low-priority +// jobs with staged inputs beating a high-priority job with cold ones. +// +// Second, the ORDER BY sits inside the candidate SELECT and before its +// LIMIT, so the LIMIT truncates an ordered set. That one is the reason +// this file exists — see the note at the top: with the ORDER BY deleted, +// every conformance case still passes on SQLite. +func TestBuildDequeueQueryOrdersLocalityBelowPriority(t *testing.T) { + query, args := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 2, + PreferHashes: []string{"blake3:staged"}, + }, fixedNow()) + + const want = "priority DESC, COALESCE(primary_input_hash IN (?), 0) DESC, run_at ASC" + + if n := strings.Count(query, want); n != 1 { + t.Fatalf("ordering %q appears %d times, want 1:\n%s", want, n, query) + } + + orderAt := strings.Index(query, "ORDER BY ") + limitAt := strings.Index(query, "LIMIT ") + + if orderAt < 0 || limitAt < 0 || orderAt > limitAt { + t.Fatalf("the candidate LIMIT must follow its ORDER BY, or it truncates "+ + "an unordered set and the ordering is decoration:\n%s", query) + } + + // And the locality term reads the hash the caller staged, not something + // bound out of position. + if got := render(t, query, args); !strings.Contains(got, + "COALESCE(primary_input_hash IN ('blake3:staged'), 0) DESC") { + t.Errorf("locality term bound the wrong value:\n%s", got) + } +} + +// TestBuildDequeueQueryBindsInTextualOrder is the SQLite-specific +// hazard. Postgres numbers its placeholders, so buildDequeueQuery could +// bind in any order there and still be correct. Here `?` is positional: +// the nth value bound is read by the nth `?` in the text, so the builders +// must bind exactly as they write. +// +// Rendering the statement is the only way to see that, and a swap is +// silent otherwise — a budget compared against a queue name is a +// perfectly valid SQLite expression. +// +// The granting case is the one that most needs this. buildLeaseGrant +// binds BETWEEN updated_at and the queue list, which is the middle of the +// sequence rather than either end, so getting it wrong shifts every +// later value by two: the queue list would read the worker id and the +// expiry, and the statement would still run. +func TestBuildDequeueQueryBindsInTextualOrder(t *testing.T) { + reserved := id.NewJobID() + worker := id.NewWorkerID() + now := fixedNow() + until := now.Add(90 * time.Second) + stamp := "'" + now.Format(time.RFC3339Nano) + "'" + leaseStamp := "'" + until.Format(time.RFC3339Nano) + "'" + + base := job.DequeueOpts{ + Queues: []string{"alpha", "beta"}, + Limit: 3, + Budget: resource.Set{resource.Memory: 4 << 30, resource.GPU: 0}, + CustomKeys: []string{"tpu", "fpga"}, + PreferHashes: []string{"blake3:staged"}, + ReservedFor: &reserved, + } + + // Every case asserts the whole sequence, not just its own addition: a + // mis-bound grant corrupts the values AFTER it, so the queue list and + // the limit are the assertions that actually catch it. + common := []string{ + "AND queue IN ('alpha','beta')", + "AND run_at <= " + stamp, + "AND id = '" + reserved.String() + "'", + "AND req_memory_bytes <= 4294967296", + "AND req_gpu_milli <= 0", + "REPLACE(REPLACE(req_custom_keys, ',fpga,', ','), ',tpu,', ',') IN ('', ',')", + "COALESCE(primary_input_hash IN ('blake3:staged'), 0) DESC", + "LIMIT 3", + } + + grantOpts := base + grantOpts.WorkerID = worker + grantOpts.LeaseUntil = until + + tests := []struct { + name string + opts job.DequeueOpts + want []string + banned []string + }{ + { + name: "no grant", + opts: base, + want: append([]string{ + "SET state = 'running', started_at = " + stamp + ", updated_at = " + stamp + "\n", + }, common...), + // A caller that did not ask for a lease must not have one + // written, and the statement is where that is decided. + banned: []string{"worker_id", "lease_epoch", "lease_expires_at"}, + }, + { + name: "grant", + opts: grantOpts, + want: append([]string{ + "SET state = 'running', started_at = " + stamp + ", updated_at = " + stamp + ",", + "worker_id = '" + worker.String() + "'", + "lease_epoch = lease_epoch + 1", + "lease_expires_at = " + leaseStamp, + }, common...), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + query, args := buildDequeueQuery(tt.opts, now) + got := render(t, query, args) + + for _, want := range tt.want { + if !strings.Contains(got, want) { + t.Errorf("rendered statement is missing %q — a value was bound out of "+ + "position:\n%s", want, got) + } + } + + for _, banned := range tt.banned { + if strings.Contains(got, banned) { + t.Errorf("rendered statement contains %q for opts that grant no "+ + "lease:\n%s", banned, got) + } + } + }) + } +} + +// TestBuildDequeueQueryBindsEveryValue is the injection check. Every +// value the caller controls — queue names, budgets, custom keys, the +// reserved id, hashes, the limit — must reach SQLite as a bind parameter. +// The only things concatenated into the statement are column names from +// budgetColumns, all compile-time constants. +func TestBuildDequeueQueryBindsEveryValue(t *testing.T) { + reserved := id.NewJobID() + + query, args := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"q'; DROP TABLE dispatch_jobs; --"}, + Limit: 3, + Budget: resource.Set{resource.Memory: 4 << 30}, + CustomKeys: []string{"fpga'); --"}, + PreferHashes: []string{"blake3:x"}, + ReservedFor: &reserved, + }, fixedNow()) + + for _, hostile := range []string{"DROP TABLE", "fpga", reserved.String(), "blake3:x"} { + if strings.Contains(query, hostile) { + t.Errorf("value %q was interpolated into the statement:\n%s", hostile, query) + } + } + + // started_at, updated_at, queue, run_at, reserved id, memory budget, + // one custom key + its separator, the closing separator, hash, limit. + if len(args) != 11 { + t.Fatalf("bound %d args, want 11: %v", len(args), args) + } +} + +// TestBuildDequeueQueryBudgetPredicate pins the three rules a budget +// comparison has to get right at once: only declared keys are compared, a +// declared zero is still a real constraint, and the comparison is <= so an +// exact fit is claimable. +func TestBuildDequeueQueryBudgetPredicate(t *testing.T) { + query, _ := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + Budget: resource.Set{resource.Memory: 4 << 30, resource.GPU: 0}, + }, fixedNow()) + + for _, want := range []string{"req_memory_bytes <= ?", "req_gpu_milli <= ?"} { + if !strings.Contains(query, want) { + t.Errorf("missing %q:\n%s", want, query) + } + } + + // CPU and disk were never declared, so they are unconstrained — not + // compared against zero. + for _, banned := range []string{"req_cpu_milli", "req_disk_bytes"} { + if strings.Contains(query, banned) { + t.Errorf("undeclared dimension %q was constrained:\n%s", banned, query) + } + } +} + +// TestBuildDequeueQueryCustomKeysAreASubsetTest pins containment as +// nested REPLACE rather than LIKE or GLOB. One REPLACE per offered key, +// each stripping ",key," and putting the separator back, with the +// surviving string required to be empty or a lone separator. +// +// The substring formulation this replaces passes every single-key case in +// the conformance suite — including the prefix collision — and then +// silently strands multi-key jobs, so the shape is worth pinning here as +// well as behaviourally. +func TestBuildDequeueQueryCustomKeysAreASubsetTest(t *testing.T) { + query, args := buildDequeueQuery(job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + CustomKeys: []string{"tpu", "fpga", "nvme"}, + }, fixedNow()) + + for _, banned := range []string{"LIKE", "GLOB", "INSTR"} { + if strings.Contains(query, banned) { + t.Errorf("custom-key containment used %s — that is a substring test:\n%s", + banned, query) + } + } + + if n := strings.Count(query, "REPLACE("); n != 3 { + t.Errorf("emitted %d REPLACE calls for 3 offered keys:\n%s", n, query) + } + + // Keys are bound wrapped in separators, which is what stops ",fpga," + // matching a job that needs ",fpga-large,". + rendered := render(t, query, args) + if !strings.Contains(rendered, + "REPLACE(REPLACE(REPLACE(req_custom_keys, ',fpga,', ','), ',nvme,', ','), ',tpu,', ',') IN ('', ',')") { + t.Errorf("subset test is not the nested-REPLACE recipe:\n%s", rendered) + } +} + +// TestBuildDequeueQueryBoundedWithNoCustomKeysExcludesCustomJobs is the +// half of the empty-offer rule a backend gets wrong. Bounded opts with an +// empty offer are a resource-aware worker with no custom resources, so a +// job requiring an fpga must not be claimable — the predicate still has to +// be emitted, with no REPLACE wrapping it. +// +// It doubles as the empty-IN-list guard: SQLite's `IN ()` is a syntax +// error, so an offer of no keys must not expand into one. +func TestBuildDequeueQueryBoundedWithNoCustomKeysExcludesCustomJobs(t *testing.T) { + opts := job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + Budget: resource.Set{resource.Memory: 4 << 30}, + } + + if opts.IsUnbounded() { + t.Fatal("opts carrying a memory budget report IsUnbounded() = true") + } + + query, args := buildDequeueQuery(opts, fixedNow()) + + if strings.Contains(query, "REPLACE(") { + t.Errorf("an empty custom-key offer emitted a REPLACE:\n%s", query) + } + + if got := render(t, query, args); !strings.Contains(got, "req_custom_keys IN ('', ',')") { + t.Errorf("bounded opts with no offered keys must still exclude custom "+ + "requirements:\n%s", got) + } + + if strings.Contains(query, "IN ()") { + t.Errorf("emitted an empty IN list, which SQLite rejects as a syntax error:\n%s", query) + } +} diff --git a/store/sqlite/dequeue_test.go b/store/sqlite/dequeue_test.go new file mode 100644 index 0000000..1bdbc57 --- /dev/null +++ b/store/sqlite/dequeue_test.go @@ -0,0 +1,436 @@ +package sqlite_test + +import ( + "context" + "testing" + "time" + + "github.com/xraph/grove/drivers/sqlitedriver" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + sqlitestore "github.com/xraph/dispatch/store/sqlite" + "github.com/xraph/dispatch/store/storetest" +) + +// TestDequeueConformance runs the shared resource-aware dequeue contract +// against SQLite. openSqliteStore (store/sqlite/reap_test.go:19) opens a +// migrated store on a per-test temp directory, so every subtest gets its +// own database file for free. +func TestDequeueConformance(t *testing.T) { + storetest.RunDequeueSuite(t, func(t *testing.T) job.Store { + t.Helper() + + return openSqliteStore(t) + }) +} + +// ────────────────────────────────────────────────── +// The two SQLite dialect traps, stated as tests +// ────────────────────────────────────────────────── + +// TestNegativeLimitWouldBeUnlimitedInSQLite is the reason DequeueJobs +// returns early on Limit <= 0 rather than letting the statement handle it. +// +// The suite's NonPositiveLimitClaimsNothing pins the behaviour; this pins +// the hazard behind it, which is dialect-specific and easy to "simplify" +// away. On Postgres a negative LIMIT is an error. On SQLite it means +// UNLIMITED, so deleting the early return would hand a worker that just +// computed zero free slots the entire queue. +func TestNegativeLimitWouldBeUnlimitedInSQLite(t *testing.T) { + s := openSqliteStore(t) + ctx := context.Background() + + const queue = "negative-limit" + + for i := range 3 { + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "job", + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + RunAt: time.Now().UTC().Add(-time.Hour).Add(time.Duration(i) * time.Minute), + } + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob: %v", err) + } + } + + var reachable int + + err := sqlitedriver.Unwrap(s.DB()).NewRaw(` + SELECT COUNT(*) FROM ( + SELECT id FROM dispatch_jobs WHERE queue = ? LIMIT -1 + )`, queue, + ).Scan(ctx, &reachable) + if err != nil { + t.Fatalf("probe LIMIT -1: %v", err) + } + + if reachable != 3 { + t.Fatalf("LIMIT -1 returned %d of 3 rows; if SQLite no longer reads a "+ + "negative limit as unlimited, the early return in DequeueJobs can be "+ + "re-justified — until then it is load-bearing", reachable) + } + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{queue}, Limit: -1}) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 0 { + t.Fatalf("a negative limit claimed %d jobs, want 0", len(got)) + } +} + +// TestDequeueWithNoQueuesClaimsNothing covers the other dialect trap: +// expanding an empty queue list would emit `queue IN ()`, which SQLite +// rejects as a SYNTAX ERROR where Postgres's `queue = ANY('{}')` is +// merely false. The caller must get the same empty result Postgres gives, +// not an error. +func TestDequeueWithNoQueuesClaimsNothing(t *testing.T) { + s := openSqliteStore(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "unreachable", + Queue: "no-queues", + Payload: []byte(`{}`), + State: job.StatePending, + RunAt: time.Now().UTC().Add(-time.Hour), + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob: %v", err) + } + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{Limit: 10}) + if err != nil { + t.Fatalf("DequeueJobs with no queues: %v", err) + } + + if len(got) != 0 { + t.Fatalf("claimed %v with an empty queue list, want nothing", names(got)) + } + + // And the job was left alone for a caller that does name its queue. + got, err = s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"no-queues"}, Limit: 10}) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 1 { + t.Fatalf("claimed %v, want the one job", names(got)) + } +} + +// ────────────────────────────────────────────────── +// Gap 1: rows written before the resource columns existed +// ────────────────────────────────────────────────── + +// TestDequeueClaimsRowsWrittenBeforeTheResourceColumns covers the shape +// the shared suite cannot construct: a row inserted before migration +// 20260812130000 added req_cpu_milli and friends. +// +// It matters because a bare `req_memory_bytes <= ?` against a NULL +// evaluates to NULL, which is not true, so such a row would be silently +// dropped from every bounded dequeue — a job that was fine yesterday +// becoming unclaimable after an upgrade, with nothing reporting it. +// +// SQLite's ALTER TABLE ADD COLUMN ... NOT NULL DEFAULT 0 is what stops +// that: existing rows read back the default rather than NULL. This test +// inserts a row naming only the pre-migration columns — the only way to +// reproduce a legacy row, since no Go write path in this package can omit +// a column — asserts the defaults landed, and then claims it under +// BOUNDED opts. +// +// Bounded is load-bearing. Unbounded opts emit no fit predicate at all, +// so they would claim the row however badly the columns read, and the +// test would prove nothing. +func TestDequeueClaimsRowsWrittenBeforeTheResourceColumns(t *testing.T) { + s := openSqliteStore(t) + ctx := context.Background() + + const queue = "legacy-resource-columns" + + legacyID := id.NewJobID() + runAt := time.Now().UTC().Add(-time.Hour) + raw := sqlitedriver.Unwrap(s.DB()) + + _, err := raw.NewRaw(` + INSERT INTO dispatch_jobs + (id, name, queue, payload, state, priority, max_retries, + retry_count, run_at, created_at, updated_at) + VALUES (?, ?, ?, ?, 'pending', 0, 3, 0, ?, ?, ?)`, + legacyID.String(), "written-before-the-columns", queue, + []byte(`{}`), runAt, runAt, runAt, + ).Exec(ctx) + if err != nil { + t.Fatalf("insert legacy row: %v", err) + } + + // Precondition: the row really does carry the added columns' defaults, + // and really does carry SQL NULL in the nullable JSON/hash columns — + // otherwise the claim below would be proving something easier. + var ( + cpu, mem, disk, gpu int64 + custom string + requestsNull int + hashNull int + ) + + // -1 is a sentinel a NULL column would produce: scanning a real NULL + // into an int64 is an error, and erroring here would hide the failure + // that actually matters, which is the claim below coming back empty. + err = raw.NewRaw(` + SELECT COALESCE(req_cpu_milli, -1), COALESCE(req_memory_bytes, -1), + COALESCE(req_disk_bytes, -1), COALESCE(req_gpu_milli, -1), + COALESCE(req_custom_keys, ''), + resource_requests IS NULL, primary_input_hash IS NULL + FROM dispatch_jobs WHERE id = ?`, legacyID.String(), + ).Scan(ctx, &cpu, &mem, &disk, &gpu, &custom, &requestsNull, &hashNull) + if err != nil { + t.Fatalf("read back legacy row: %v", err) + } + + if cpu != 0 || mem != 0 || disk != 0 || gpu != 0 || custom != "" { + t.Errorf("legacy row read back req_* = %d/%d/%d/%d/%q, want 0/0/0/0/\"\" "+ + "(-1 or means the column lost its NOT NULL DEFAULT and the "+ + "bare comparison in buildFitPredicate now evaluates to NULL)", + cpu, mem, disk, gpu, custom) + } + + if requestsNull != 1 || hashNull != 1 { + t.Fatalf("legacy row is not the pre-migration shape: "+ + "resource_requests IS NULL = %d, primary_input_hash IS NULL = %d, want 1/1", + requestsNull, hashNull) + } + + opts := job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 4 * storetest.GiB}, + } + + if opts.IsUnbounded() { + t.Fatal("opts carrying a memory budget report IsUnbounded() = true; " + + "an unbounded dequeue skips the predicate and would prove nothing here") + } + + got, err := s.DequeueJobs(ctx, opts) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + if len(got) != 1 || got[0].ID != legacyID { + t.Fatalf("claimed %v, want the one legacy row: a job written before the "+ + "req_* columns existed declares no requirement and must stay claimable", + names(got)) + } +} + +// TestResourceColumnsRejectNull pins the schema invariant the dequeue +// predicate leans on. The budget comparisons are deliberately NOT wrapped +// in COALESCE — that would cost the dequeue index, since SQLite cannot +// answer an expression from a plain column index — and they are only safe +// bare because NULL cannot reach those columns. +// +// If a future migration relaxes NOT NULL on any of them, this test fails +// and points at buildFitPredicate, rather than the change landing quietly +// and stranding rows. +func TestResourceColumnsRejectNull(t *testing.T) { + s := openSqliteStore(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "not-null-probe", + Queue: "not-null-probe", + Payload: []byte(`{}`), + State: job.StatePending, + RunAt: time.Now().UTC().Add(-time.Hour), + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob: %v", err) + } + + raw := sqlitedriver.Unwrap(s.DB()) + + for _, column := range []string{ + "req_cpu_milli", "req_memory_bytes", "req_disk_bytes", + "req_gpu_milli", "req_custom_keys", + } { + // column comes from the compile-time list above, never from input. + _, err := raw.NewRaw( + `UPDATE dispatch_jobs SET `+column+` = NULL WHERE id = ?`, j.ID.String(), + ).Exec(ctx) + if err == nil { + t.Errorf("%s accepted NULL; the dequeue predicate compares it bare, "+ + "so a NULL there silently drops the row from every bounded claim", column) + } + } +} + +// ────────────────────────────────────────────────── +// Gap 2: NULL primary_input_hash in the locality ordering +// ────────────────────────────────────────────────── + +const stagedHash = "blake3:staged-here" + +// nullHashFixtures enqueues four same-priority jobs on queue and then +// rewrites one job's primary_input_hash to a genuine SQL NULL, which the +// Go write path cannot produce (jobModel.PrimaryInputHash is a plain +// string, so an unset hash stores ”). +// +// The fixture set is chosen so a broken locality term cannot pass by +// accident, which is the trap two earlier backends fell into: +// +// - "remote-high" carries a hash that collates ABOVE the staged one. With +// only one hash value present, a NULL and an empty string both collate +// BELOW every real hash, so a sort on the hash VALUE would put the +// staged job first for entirely the wrong reason. "zzz:never-staged" +// is what makes that mutation visible. +// - The staged job is the LAST by RunAt, so it must jump the whole band; +// every unpreferred fixture beats it on any tie the locality term +// fails to break. +// - The NULL-hash job is the EARLIEST by RunAt, so it must come first +// among the unpreferred. Dropping the COALESCE would sort it below +// them all — SQLite puts NULL last under DESC — even though it is no +// less preferred than an empty or unmatched hash. +func nullHashFixtures(t *testing.T, s *sqlitestore.Store, queue string) { + t.Helper() + + ctx := context.Background() + base := time.Now().UTC().Add(-time.Hour).Truncate(time.Millisecond) + + fixtures := []struct { + name string + hash string + offset time.Duration + }{ + {"null-hash", "", 0}, + {"remote-high", "zzz:never-staged", time.Minute}, + {"empty-hash", "", 2 * time.Minute}, + {"cached", stagedHash, 3 * time.Minute}, + } + + var nullHashID id.JobID + + for _, f := range fixtures { + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: f.name, + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: base.Add(f.offset), + PrimaryInputHash: f.hash, + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob(%s): %v", f.name, err) + } + + if f.name == "null-hash" { + nullHashID = j.ID + } + } + + _, err := sqlitedriver.Unwrap(s.DB()).NewRaw( + `UPDATE dispatch_jobs SET primary_input_hash = NULL WHERE id = ?`, + nullHashID.String(), + ).Exec(ctx) + if err != nil { + t.Fatalf("null out primary_input_hash: %v", err) + } +} + +// TestDequeueOrdersNullPrimaryInputHashAsUnpreferred covers the returned +// ORDER: a row with no locality signal at all is neither preferred nor +// penalised, it is simply unpreferred, and RunAt then separates it from +// the other unpreferred rows. +func TestDequeueOrdersNullPrimaryInputHashAsUnpreferred(t *testing.T) { + s := openSqliteStore(t) + ctx := context.Background() + + const queue = "null-hash-order" + + nullHashFixtures(t, s, queue) + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + PreferHashes: []string{stagedHash}, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + wantSequence(t, got, "cached", "null-hash", "remote-high", "empty-hash") +} + +// TestDequeueSelectsPreferredOverNullHashUnderLimit covers the other half, +// and it is the half only SQL can answer: WHICH jobs a tight limit keeps. +// +// The returned slice is ordered in Go by job.DequeueOpts.Less, so the test +// above would still pass with no locality term in the statement at all. +// Here only two of four eligible jobs may be claimed, so the choice is +// made entirely by the statement's ORDER BY, and the COALESCE that decides +// where the NULL-hash row sits is what picks the second one. +func TestDequeueSelectsPreferredOverNullHashUnderLimit(t *testing.T) { + s := openSqliteStore(t) + ctx := context.Background() + + const queue = "null-hash-limit" + + nullHashFixtures(t, s, queue) + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 2, + PreferHashes: []string{stagedHash}, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + wantSequence(t, got, "cached", "null-hash") +} + +// ────────────────────────────────────────────────── +// Helpers +// ────────────────────────────────────────────────── + +func names(jobs []*job.Job) []string { + out := make([]string, 0, len(jobs)) + for _, j := range jobs { + out = append(out, j.Name) + } + + return out +} + +func wantSequence(t *testing.T, got []*job.Job, want ...string) { + t.Helper() + + gotNames := names(got) + if len(gotNames) != len(want) { + t.Fatalf("claimed %v, want %v", gotNames, want) + } + + for i := range want { + if gotNames[i] != want[i] { + t.Fatalf("claimed %v, want %v (differs at index %d)", gotNames, want, i) + } + } +} diff --git a/store/sqlite/dlq.go b/store/sqlite/dlq.go index a2d1c1b..7c23ba5 100644 --- a/store/sqlite/dlq.go +++ b/store/sqlite/dlq.go @@ -12,9 +12,11 @@ import ( // PushDLQ adds a failed job entry to the dead letter queue. func (s *Store) PushDLQ(ctx context.Context, entry *dlq.Entry) error { - m := toDLQModel(entry) - _, err := s.sdb.NewInsert(m).Exec(ctx) + m, err := toDLQModel(entry) if err != nil { + return err + } + if _, err = s.sdb.NewInsert(m).Exec(ctx); err != nil { return fmt.Errorf("dispatch/sqlite: push dlq: %w", err) } return nil diff --git a/store/sqlite/doc.go b/store/sqlite/doc.go index 9a1610c..550f9c6 100644 --- a/store/sqlite/doc.go +++ b/store/sqlite/doc.go @@ -13,4 +13,38 @@ // db, _ := grove.Open(ctx, "sqlite", dsn) // store := sqlite.New(db) // store.Migrate(ctx) +// +// # Write concurrency +// +// SQLite allows one writer at a time for the whole database, and this +// store does more of its work through writes than a reader would expect: +// claiming a job, renewing a lease and reclaiming an expired one are all +// writes, and a busy pool performs them continuously. +// +// Two settings normally smooth that over, and neither is reachable from +// here. Grove's sqlitedriver enables WAL but sets no busy_timeout, so a +// writer that loses the race fails immediately with SQLITE_BUSY rather +// than waiting for the lock, and the driver does not expose the underlying +// *sql.DB, so this package cannot call SetMaxOpenConns to keep more than +// one connection from trying at once. The store compensates in Go by +// retrying SQLITE_BUSY with a jittered backoff (see withBusyRetry), which +// is enough for ordinary contention. +// +// It is a mitigation, not a substitute. If a deployment is write-heavy +// enough to see SQLITE_BUSY surface as an error after the retries are +// exhausted, the fixes are, in order of preference: +// +// - Open the database with busy_timeout set in the DSN, for example +// "file:dispatch.db?_pragma=busy_timeout(5000)", so SQLite itself +// blocks on the lock instead of failing fast. The exact parameter +// name depends on the driver build. +// - Constrain the pool to a single connection if the driver in use +// allows configuring it, which serialises writers before they reach +// SQLite rather than after. +// - Move to postgres. SQLite's single-writer model is a property of the +// engine, and a queue with several busy pools is the workload it +// suits least. +// +// A single process with one worker pool, which is what embedded and CLI +// deployments usually are, will not meaningfully encounter this. package sqlite diff --git a/store/sqlite/job.go b/store/sqlite/job.go index dd3b33c..fec24e2 100644 --- a/store/sqlite/job.go +++ b/store/sqlite/job.go @@ -3,18 +3,24 @@ package sqlite import ( "context" "fmt" + "sort" "strings" "time" "github.com/xraph/dispatch" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" ) // EnqueueJob persists a new job in pending state. func (s *Store) EnqueueJob(ctx context.Context, j *job.Job) error { - m := toJobModel(j) - _, err := s.sdb.NewInsert(m).Exec(ctx) + m, err := toJobModel(j) + if err != nil { + return err + } + + _, err = s.sdb.NewInsert(m).Exec(ctx) if err != nil { if isDuplicateKey(err) { return dispatch.ErrJobAlreadyExists @@ -24,39 +30,70 @@ func (s *Store) EnqueueJob(ctx context.Context, j *job.Job) error { return nil } -// DequeueJobs atomically claims up to limit pending jobs from the given -// queues. SQLite doesn't support FOR UPDATE SKIP LOCKED, so we use -// BEGIN IMMEDIATE with a subquery + UPDATE pattern. -func (s *Store) DequeueJobs(ctx context.Context, queues []string, limit int) ([]*job.Job, error) { - now := time.Now().UTC() +// DequeueJobs atomically claims up to opts.Limit ready jobs from +// opts.Queues that fit opts, sets them to running, and returns them +// ordered by priority descending, then locality-preferred first, then +// RunAt ascending. +// +// SQLite doesn't support FOR UPDATE SKIP LOCKED, so the claim is a single +// UPDATE ... WHERE id IN (SELECT ... ORDER BY ... LIMIT ?) RETURNING *: +// one statement, which SQLite runs inside an implicit immediate +// transaction with the database write lock held, so two concurrent +// claimers cannot select the same candidate. That mechanism is unchanged +// here — the fit predicate is simply another conjunct of the inner +// SELECT's WHERE, so a job that does not fit is never written to. It stays +// pending and untouched for the next worker that does have room. +// +// When opts.Grants() the lease columns are additional assignments in that +// same UPDATE's SET clause, never a follow-up statement. ReclaimExpiredLeases +// requires lease_expires_at IS NOT NULL, so a row left running with a null +// expiry is not a row at risk of reclamation — it is one reclamation can +// never see. See job.DequeueOpts.LeaseUntil. +func (s *Store) DequeueJobs(ctx context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + // A worker computing zero free slots must claim zero jobs, never the + // whole queue. This early return is load-bearing on SQLite rather than + // a saved round trip: `LIMIT -1` means UNLIMITED here, the exact + // opposite of Postgres, where it is an error. Without this, an + // exhausted worker asking for -1 would claim the entire queue. + if opts.Limit <= 0 { + return nil, nil + } - // Build queue placeholders for raw SQL. - placeholders := make([]string, len(queues)) - args := make([]any, 0, len(queues)+3) - args = append(args, now, now) // started_at, updated_at - for i, q := range queues { - placeholders[i] = "?" - args = append(args, q) + // Ahead of the empty-Queues guard, so an incoherent grant is reported + // rather than swallowed by a return that happens to be silent here. + // A caller that names no queues still deserves to hear that its lease + // has no holder, and the five backends must agree on which inputs are + // errors — see job.DequeueOpts.Validate. + if err := opts.Validate(); err != nil { + return nil, fmt.Errorf("dispatch/sqlite: dequeue jobs: %w", err) } - args = append(args, now, limit) // run_at <=, LIMIT - query := fmt.Sprintf(` - UPDATE dispatch_jobs - SET state = 'running', started_at = ?, updated_at = ? - WHERE id IN ( - SELECT id FROM dispatch_jobs - WHERE state IN ('pending', 'retrying') - AND queue IN (%s) - AND run_at <= ? - ORDER BY priority DESC, run_at ASC - LIMIT ? - ) - RETURNING *`, - strings.Join(placeholders, ","), - ) + // `queue IN ()` is a SQLite syntax error, where Postgres's + // `queue = ANY('{}')` is merely false. Claiming nothing is what + // store/postgres does for the same input, and there is no existing + // "all queues" behaviour to preserve: this backend has never had a + // query that could run without a queue list. + if len(opts.Queues) == 0 { + return nil, nil + } + query, args := buildDequeueQuery(opts, time.Now().UTC()) + + // SQLite serializes writers with a single database-wide write lock, and + // grove's sqlitedriver sets no busy_timeout, so a claimer that loses the + // race for that lock fails immediately with SQLITE_BUSY rather than + // blocking. Retrying is what turns "another worker is claiming right + // now" back into "claim shortly", which is what the atomicity guarantee + // looks like from the caller's side; the claim itself is atomic either + // way, since only one statement can hold the write lock. Same helper the + // lease writes use (store/sqlite/lease.go:38). var models []jobModel - err := s.sdb.NewRaw(query, args...).Scan(ctx, &models) + + err := withBusyRetry(ctx, func() error { + models = nil + + return s.sdb.NewRaw(query, args...).Scan(ctx, &models) + }) if err != nil { return nil, fmt.Errorf("dispatch/sqlite: dequeue jobs: %w", err) } @@ -69,9 +106,273 @@ func (s *Store) DequeueJobs(ctx context.Context, queues []string, limit int) ([] } jobs = append(jobs, j) } + + // SQLite defines no order for the rows an UPDATE ... RETURNING emits, + // and unlike Postgres it has no data-modifying CTE to wrap the claim in + // and order the output of. So the statement's ORDER BY decides WHICH + // jobs the LIMIT keeps — the part that must happen inside the claim — + // and the returned slice is ordered here, by the contract's own + // comparator rather than a fifth restatement of it. + sort.SliceStable(jobs, func(a, b int) bool { return opts.Less(jobs[a], jobs[b]) }) + return jobs, nil } +// budgetColumns maps each canonical dimension the dequeue predicate +// compares to the scalar column that holds it. These are exactly the +// dimensions job.DequeueOpts.Allows loops over, and exactly the columns +// idx_dispatch_jobs_dequeue_res carries in its key list — SQLite has no +// INCLUDE clause, so the migration folds them into the key instead, and +// each comparison is still a scalar range test rather than a probe into +// the resource_requests JSON. +// +// Every column is NOT NULL DEFAULT 0, which is what lets the comparisons +// below be bare rather than COALESCEd: a row written before migration +// 20260812130000 reads back the default 0, never NULL, so +// `req_memory_bytes <= ?` cannot silently evaluate to NULL and drop a +// legacy job. TestDequeueClaimsRowsWrittenBeforeTheResourceColumns and +// TestResourceColumnsRejectNull pin both halves of that. +// +// The column names are compile-time constants and are the only +// identifiers ever concatenated into the statement below; every value +// travels as a bind parameter. +var budgetColumns = []struct { + key string + column string +}{ + {resource.CPU, "req_cpu_milli"}, + {resource.Memory, "req_memory_bytes"}, + {resource.Disk, "req_disk_bytes"}, + {resource.GPU, "req_gpu_milli"}, +} + +// dequeueSQL is the claim statement with everything the caller decides +// filled in: the started_at and updated_at placeholders, the lease grant, +// the queue list, the run_at placeholder, the fit predicate, the ordering, +// and the limit placeholder. +// +// The grant is a suffix of the SET clause rather than a statement of its +// own, which is what makes "claimed" and "leased" the same event. +// +// Unlike the Postgres statement this mirrors, the ordering appears once, +// not twice: there is no outer SELECT to order because SQLite has no +// data-modifying CTE. The one occurrence is the load-bearing one — it +// decides which rows the LIMIT keeps. +// +// Do not delete it on the grounds that the returned slice is sorted in +// Go anyway. Re-measured against the current 21-case +// storetest.RunDequeueSuite, with this ORDER BY removed entirely: +// exactly ONE case fails, LocalityDecidesWhichRowsSurviveATightLimit. +// Every other case still passes, LimitTruncatesAfterOrdering included, +// because SQLite answers the candidate scan from +// idx_dispatch_jobs_dequeue_res and that index's leading key order is +// priority DESC, run_at ASC — so twenty of twenty-one right answers +// arrive for the wrong reason and would stop arriving the moment the +// planner picked another index. +// +// The single case that does catch it is the one that cannot be baked +// into any index: locality depends on the caller's staged set, which was +// not known when the rows were written. That case was added for exactly +// this hole. TestBuildDequeueQueryOrdersLocalityBelowPriority and +// TestDequeueSelectsPreferredOverNullHashUnderLimit pin it here too. +const dequeueSQL = ` + UPDATE dispatch_jobs + SET state = 'running', started_at = %s, updated_at = %s%s + WHERE id IN ( + SELECT id FROM dispatch_jobs + WHERE state IN ('pending', 'retrying') + AND queue IN (%s) + AND run_at <= %s%s + ORDER BY %s + LIMIT %s + ) + RETURNING *` + +// buildDequeueQuery compiles opts into the claim statement and its bind +// parameters. It is the SQL expression of job.DequeueOpts.Allows and +// Less, and must answer identically for every job. +// +// SQLite binds `?` parameters positionally, so args must be appended in +// the order the placeholders appear in the finished statement. Every +// helper below therefore binds as it writes, and they are called in +// textual order: SET, queues, run_at, fit predicate, ORDER BY, LIMIT. +func buildDequeueQuery(opts job.DequeueOpts, now time.Time) (query string, args []any) { + // +10 covers the fixed binds: started_at, updated_at, the grant's + // worker id and expiry, run_at, the custom-key separator, and the + // limit, with headroom. + args = make([]any, 0, len(opts.Queues)+len(opts.CustomKeys)*2+len(opts.PreferHashes)+10) + + // bind appends v and returns the placeholder that reads it. Values + // never reach the statement text. + bind := func(v any) string { + args = append(args, v) + + return "?" + } + + startedAt, updatedAt := bind(now), bind(now) + grant := buildLeaseGrant(opts, bind) + + queues := make([]string, len(opts.Queues)) + for i, q := range opts.Queues { + queues[i] = bind(q) + } + + runAt := bind(now) + fit := buildFitPredicate(opts, bind) + order := buildDequeueOrder(opts, bind) + limit := bind(opts.Limit) + + return fmt.Sprintf(dequeueSQL, + startedAt, updatedAt, grant, strings.Join(queues, ","), runAt, fit, order, limit, + ), args +} + +// buildLeaseGrant renders the lease assignments appended to the claim's +// SET clause, or "" when opts grants no lease. +// +// Empty is the whole backward-compatibility guarantee: a caller that does +// not opt in emits the statement it emitted before leases existed and +// leaves worker_id, lease_epoch and lease_expires_at exactly as they were. +// +// lease_epoch = lease_epoch + 1 rather than a bound value, because the +// epoch is the fence: it must advance from whatever the row currently +// holds, which only the row knows. Reading it and writing back a computed +// successor would be the read-modify-write this statement exists to avoid. +// +// It binds between updated_at and the queue list because `?` is +// positional here — see buildDequeueQuery. +func buildLeaseGrant(opts job.DequeueOpts, bind func(any) string) string { + if !opts.Grants() { + return "" + } + + return ",\n\t\t worker_id = " + bind(opts.WorkerID.String()) + + ",\n\t\t lease_epoch = lease_epoch + 1" + + ",\n\t\t lease_expires_at = " + bind(opts.LeaseUntil.UTC()) +} + +// buildFitPredicate renders the conjuncts that decide WHICH jobs may be +// claimed, or "" when opts constrains nothing. +func buildFitPredicate(opts job.DequeueOpts, bind func(any) string) string { + // Unbounded opts emit the original query verbatim: a caller that does + // not use the resource model claims everything, including jobs + // declaring custom resources it could not possibly satisfy. Anything + // else strands work the day this option ships. + if opts.IsUnbounded() { + return "" + } + + var b strings.Builder + + if opts.ReservedFor != nil { + b.WriteString("\n\t\t\t AND id = " + bind(opts.ReservedFor.String())) + } + + // An absent budget key is unconstrained, not zero, so only declared + // dimensions produce a comparison. A key present with the value zero + // is a real constraint and still emits one — that is an exhausted + // worker, which must claim nothing that needs the dimension. + // + // The test is requirement <= budget: a job needing exactly the free + // capacity is claimable, or the last slot on every worker is + // permanently unusable. + for _, dim := range budgetColumns { + budget, declared := opts.Budget[dim.key] + if !declared { + continue + } + + b.WriteString("\n\t\t\t AND " + dim.column + " <= " + bind(budget)) + } + + b.WriteString("\n\t\t\t AND " + buildCustomKeyPredicate(opts, bind)) + + return b.String() +} + +// buildCustomKeyPredicate renders custom-resource containment as a +// genuine SUBSET test, character for character the recipe store/postgres +// uses — it was written in nested REPLACE rather than an array operator +// precisely so this backend could copy it. +// +// req_custom_keys holds resource.EncodeCustomKeys' output — the sorted +// required keys wrapped in leading and trailing separators, e.g. +// ",fpga,tpu,". The obvious formulation, a LIKE/GLOB containment test +// against the offered list, is a SUBSTRING test: it passes every +// single-key case including the prefix collision, then silently strands a +// job needing {fpga,tpu} from a caller offering {fpga,nvme,tpu}, because +// the interleaved key breaks the contiguous run. The job it strands is +// the specialised one that is hardest to place anywhere else. +// +// Instead each offered key is stripped from the stored list by a nested +// REPLACE of ",key," with ",", which restores the separator the removal +// consumed and so composes in any order. What remains is "" or a lone +// separator exactly when every required key was offered. +// +// The nesting order is what keeps the bindings positional: the innermost +// REPLACE is written leftmost, so binding the key and then the separator +// once per wrapper, in loop order, matches the order SQLite reads the +// placeholders in. The separator cannot be bound once and reused the way +// Postgres reuses $7 — `?` has no number to refer back to. +func buildCustomKeyPredicate(opts job.DequeueOpts, bind func(any) string) string { + // Bounded opts with an empty offer are a resource-aware worker that + // genuinely has no custom resources, so only jobs requiring none are + // eligible. This is the case IsUnbounded above has already excluded. + expr := "req_custom_keys" + for _, k := range opts.OfferedCustomKeys() { + expr = "REPLACE(" + expr + ", " + + bind(resource.CustomKeySep+k+resource.CustomKeySep) + ", " + + bind(resource.CustomKeySep) + ")" + } + + return expr + " IN ('', " + bind(resource.CustomKeySep) + ")" +} + +// buildDequeueOrder renders the ordering every backend must return: +// priority descending, then locality-preferred before not, then RunAt +// ascending. +// +// Locality ranks strictly BELOW priority. Above it, a steady stream of +// low-priority jobs whose inputs are already staged would beat a +// high-priority job with cold inputs — an optimization overriding +// user-expressed intent, and the exact starvation the predicate exists to +// prevent. A preferred job jumps its own priority band and no further. +// +// The term is applied whenever PreferHashes is non-empty, including on +// otherwise-unbounded opts: IsUnbounded governs filtering only. +func buildDequeueOrder(opts job.DequeueOpts, bind func(any) string) string { + // PreferredHashes, not PreferHashes: an empty entry would bind as '' + // and match every unhashed job, since primary_input_hash is a plain + // string and an unhashed job stores '' rather than NULL. See + // job.DequeueOpts.PreferredHashes. + preferred := opts.PreferredHashes() + if len(preferred) == 0 { + return "priority DESC, run_at ASC" + } + + hashes := make([]string, len(preferred)) + for i, h := range preferred { + hashes[i] = bind(h) + } + + // primary_input_hash is nullable — rows written before the resource + // migration have no value — and `NULL IN (...)` is NULL, not 0. + // + // The COALESCE is needed for a different reason than the identical + // call in store/postgres. Postgres sorts NULLs FIRST under DESC, so + // there an uncoalesced term would rank the rows with no locality + // signal ABOVE the staged ones. SQLite sorts NULLs LAST under DESC, + // which looks harmless — but it would sort a NULL-hash job below every + // merely-unpreferred one, and those two are equally unpreferred under + // job.DequeueOpts.Prefers, so RunAt is what must separate them. + // COALESCE makes "unknown" mean exactly "not preferred", nothing + // stronger. Note the result is the integer 0/1 rather than a boolean, + // which DESC orders correctly but nothing may compare to TRUE. + return "priority DESC, COALESCE(primary_input_hash IN (" + + strings.Join(hashes, ",") + "), 0) DESC, run_at ASC" +} + // GetJob retrieves a job by ID. func (s *Store) GetJob(ctx context.Context, jobID id.JobID) (*job.Job, error) { m := new(jobModel) @@ -90,7 +391,11 @@ func (s *Store) GetJob(ctx context.Context, jobID id.JobID) (*job.Job, error) { // UpdateJob persists changes to an existing job. func (s *Store) UpdateJob(ctx context.Context, j *job.Job) error { - m := toJobModel(j) + m, err := toJobModel(j) + if err != nil { + return err + } + m.UpdatedAt = time.Now().UTC() res, err := s.sdb.NewUpdate(m).WherePK().Exec(ctx) if err != nil { diff --git a/store/sqlite/lease.go b/store/sqlite/lease.go new file mode 100644 index 0000000..965c00a --- /dev/null +++ b/store/sqlite/lease.go @@ -0,0 +1,292 @@ +package sqlite + +import ( + "context" + "fmt" + "math/rand/v2" + "strings" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// maxLeaseBusyRetries bounds how many times a lease write retries after +// SQLITE_BUSY before giving up. +// +// Grove's sqlitedriver configures WAL mode but no busy_timeout, and does +// not expose the underlying *sql.DB for this package to configure one +// itself. Two connections attempting a write at the same instant therefore +// have the loser fail immediately with SQLITE_BUSY instead of blocking. +// SQLite still serializes writes at the engine level — retrying closes the +// gap between "not grantable this instant" and "will succeed shortly" that +// ReclaimExpiredLeases's atomicity guarantee depends on under concurrency. +const maxLeaseBusyRetries = 100 + +// leaseBusyRetryDelay is the mean pause between retries. It is small +// because a write against this schema completes in well under a +// millisecond; the retry exists to ride out a burst of contention, not to +// wait out something the caller should instead be timed out for. +const leaseBusyRetryDelay = time.Millisecond + +// busyRetryDelay returns the next pause, jittered across half to one and a +// half times leaseBusyRetryDelay. +// +// A fixed delay is what makes contention here self-sustaining rather than +// self-clearing. SQLite takes one write lock, so of N writers that collide +// exactly one wins and the other N-1 all sleep the identical interval and +// wake together to collide again. The loser set stays in lockstep for as +// long as it takes one of them to win each round, which is the worst +// possible shape for a backoff. Spreading the wake-ups decorrelates them +// after the first collision. +// +// The jitter is centred on the old constant rather than added to it, so +// the expected time to exhaust maxLeaseBusyRetries is unchanged and this +// cannot quietly turn a fast failure into a slow one. Non-crypto rand is +// the right tool, as it is for backoff.Jitter. +func busyRetryDelay() time.Duration { + half := leaseBusyRetryDelay / 2 + + // #nosec G404 -- retry jitter only needs to spread contention, not resist + // prediction, so math/rand is the right tool here. + return half + time.Duration(rand.Float64()*float64(leaseBusyRetryDelay)) +} + +// isSQLiteBusy reports whether err is the driver's SQLITE_BUSY, meaning +// another connection currently holds SQLite's single write lock. Matched +// on the error message the same way isDuplicateKey matches its error. +func isSQLiteBusy(err error) bool { + return err != nil && strings.Contains(err.Error(), "SQLITE_BUSY") +} + +// withBusyRetry runs fn, retrying while it returns SQLITE_BUSY, up to +// maxLeaseBusyRetries times or until ctx is done. +func withBusyRetry(ctx context.Context, fn func() error) error { + var err error + for range maxLeaseBusyRetries { + err = fn() + if err == nil || !isSQLiteBusy(err) { + return err + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(busyRetryDelay()): + } + } + + return err +} + +// The compile-time check that this store provides the lease capability +// lives in store.go alongside the other interface assertions. +// +// The grant itself is not in this file: it travels on job.DequeueOpts and +// is compiled into DequeueJobs' claim statement by buildLeaseGrant, so a +// leased claim carries the fit predicate and the ordering like any other. +// This file holds only what a lease needs afterwards, plus the SQLITE_BUSY +// retry those writes and the claim share. + +// RenewLease extends the lease only if the caller still holds it. +func (s *Store) RenewLease( + ctx context.Context, + jobID id.JobID, + workerID id.WorkerID, + epoch int, + leaseUntil time.Time, +) error { + now := time.Now().UTC() + + var rows int64 + err := withBusyRetry(ctx, func() error { + res, execErr := s.sdb.NewUpdate((*jobModel)(nil)). + Set("lease_expires_at = ?", leaseUntil.UTC()). + Set("heartbeat_at = ?", now). + Set("updated_at = ?", now). + Where("id = ?", jobID.String()). + Where("state = 'running'"). + Where("worker_id = ?", workerID.String()). + Where("lease_epoch = ?", epoch). + Exec(ctx) + if execErr != nil { + return execErr + } + rows, _ = res.RowsAffected() //nolint:errcheck // driver always returns nil + return nil + }) + if err != nil { + return fmt.Errorf("dispatch/sqlite: renew lease: %w", err) + } + + if rows == 0 { + // Deleted, reclaimed, or reassigned — in every case this worker no + // longer owns the job and must stop. + return job.ErrLeaseLost + } + + return nil +} + +// ReclaimExpiredLeases returns expired-lease jobs to pending, fencing +// their previous holders. +func (s *Store) ReclaimExpiredLeases(ctx context.Context, limit int) ([]*job.Job, error) { + // A non-positive limit claims nothing. This early return is + // load-bearing on SQLite rather than a saved round trip: `LIMIT -1` + // means UNLIMITED here, the exact opposite of Postgres, where it is a + // runtime error. Without this, a negative limit would reclaim the + // entire table. + if limit <= 0 { + return nil, nil + } + + now := time.Now().UTC() + + // The reclaim predicate below adds a branch for a running job carrying + // no lease at all, gated on silence rather than on the null expiry + // alone — see job.UnleasedReclaimGrace for why a null expiry does not + // by itself mean the job was abandoned, and why COALESCE returning NULL + // (neither timestamp set) must not be adopted. + // + // silent is a bound time.Time and must stay one. SQLite has no + // timestamp type, so every comparison here is a string comparison + // against whatever grove's sqlitedriver wrote, and the driver renders a + // time.Time with Go's default layout rather than ISO-8601. Formatting + // this value instead would sort it above every driver-written timestamp + // ('T' > ' '), which is the bug the migration 008 backfill shipped with + // before it was removed. + // + // Here it is worse than it was there, and in a way that inverts. The + // backfill compared in the direction that made a formatted value match + // NOTHING, so the failure was jobs staying stranded: bad, but the same + // outcome as having no backfill. This predicate compares the other way, + // so a formatted value is greater than every stored timestamp and + // matches EVERYTHING, reclaiming healthy running jobs out from under + // live workers. The same mistake fails open here rather than closed. + silent := now.Add(-job.UnleasedReclaimGrace) + + var models []jobModel + err := withBusyRetry(ctx, func() error { + models = nil + return s.sdb.NewRaw(` + UPDATE dispatch_jobs + SET state = 'pending', + run_at = ?, + worker_id = NULL, + started_at = NULL, + heartbeat_at = NULL, + lease_expires_at = NULL, + lease_epoch = lease_epoch + 1, + evict_count = evict_count + 1, + updated_at = ? + WHERE id IN ( + SELECT id FROM dispatch_jobs + WHERE state = 'running' + AND ( (lease_expires_at IS NOT NULL AND lease_expires_at <= ?) + OR (lease_expires_at IS NULL + AND COALESCE(heartbeat_at, started_at) IS NOT NULL + AND COALESCE(heartbeat_at, started_at) <= ?) ) + ORDER BY lease_expires_at ASC + LIMIT ? + ) + RETURNING *`, + now, now, now, silent, limit, + ).Scan(ctx, &models) + }) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: reclaim expired leases: %w", err) + } + + jobs := make([]*job.Job, 0, len(models)) + for i := range models { + j, convErr := fromJobModel(&models[i]) + if convErr != nil { + return nil, fmt.Errorf("dispatch/sqlite: reclaim convert: %w", convErr) + } + jobs = append(jobs, j) + } + + return jobs, nil +} + +// updateLeasedJobSQL writes every business column UpdateJob writes, +// fenced on the row being still running, still assigned to the caller's +// workerID, and still at the caller's epoch. Same shape as Postgres' +// equivalent statement, character for character in intent. +// +// lease_epoch, lease_expires_at, worker_id, and heartbeat_at are +// deliberately absent from the SET list. j is the caller's stale +// snapshot, and writing j.LeaseExpiresAt back would roll the real expiry +// backwards even behind a passing epoch predicate. Those four columns +// have exactly three writers (the grant in DequeueJobs, RenewLease, and +// ReclaimExpiredLeases); this statement is deliberately not a fourth. +const updateLeasedJobSQL = ` + UPDATE dispatch_jobs + SET name = ?, queue = ?, payload = ?, state = ?, priority = ?, + max_retries = ?, retry_count = ?, last_error = ?, + scope_app_id = ?, scope_org_id = ?, run_at = ?, + started_at = ?, completed_at = ?, timeout = ?, + lease_ttl = ?, evict_count = ?, created_at = ?, + updated_at = ?, + req_cpu_milli = ?, req_memory_bytes = ?, req_disk_bytes = ?, + req_gpu_milli = ?, req_custom_keys = ?, + resource_requests = ?, resource_limits = ?, + resource_class = ?, input_bytes = ?, primary_input_hash = ? + WHERE id = ? + AND state = 'running' + AND worker_id = ? + AND lease_epoch = ?` + +// UpdateLeasedJob persists j only while the caller still holds the +// lease. +func (s *Store) UpdateLeasedJob(ctx context.Context, j *job.Job, workerID id.WorkerID, epoch int) error { + m, err := toJobModel(j) + if err != nil { + return err + } + + now := time.Now().UTC() + + var rows int64 + execErr := withBusyRetry(ctx, func() error { + res, err := s.sdb.NewRaw(updateLeasedJobSQL, + m.Name, m.Queue, m.Payload, m.State, m.Priority, + m.MaxRetries, m.RetryCount, m.LastError, + m.ScopeAppID, m.ScopeOrgID, m.RunAt, + m.StartedAt, m.CompletedAt, m.Timeout, + m.LeaseTTL, m.EvictCount, m.CreatedAt, + now, + m.ReqCPUMilli, m.ReqMemoryBytes, m.ReqDiskBytes, + m.ReqGPUMilli, m.ReqCustomKeys, + m.ResourceRequests, m.ResourceLimits, + m.ResourceClass, m.InputBytes, m.PrimaryInputHash, + m.ID, workerID.String(), epoch, + ).Exec(ctx) + if err != nil { + return err + } + rows, _ = res.RowsAffected() //nolint:errcheck // driver always returns nil + return nil + }) + if execErr != nil { + return fmt.Errorf("dispatch/sqlite: update leased job: %w", execErr) + } + + if rows > 0 { + return nil + } + + // Zero rows means either the fence predicate failed (the lease moved + // on) or the row is gone. Only the latter is ErrJobNotFound; the + // former is ErrLeaseLost, the entire point of this method. + exists := new(jobModel) + existErr := s.sdb.NewSelect(exists).Where("id = ?", m.ID).Limit(1).Scan(ctx) + if existErr != nil { + if isNoRows(existErr) { + return dispatch.ErrJobNotFound + } + return fmt.Errorf("dispatch/sqlite: update leased job existence check: %w", existErr) + } + + return job.ErrLeaseLost +} diff --git a/store/sqlite/lease_jitter_test.go b/store/sqlite/lease_jitter_test.go new file mode 100644 index 0000000..996b981 --- /dev/null +++ b/store/sqlite/lease_jitter_test.go @@ -0,0 +1,63 @@ +package sqlite + +import ( + "testing" + "time" +) + +// TestBusyRetryDelayIsJitteredAroundTheBase covers the two properties the +// retry backoff depends on, both of which a plain constant would break. +// +// SQLite takes one write lock for the whole database, so of N writers that +// collide exactly one wins and the rest retry. With a fixed delay those +// losers sleep the identical interval and wake together to collide again, +// staying in lockstep for as long as it takes them to drain one at a time. +// Spreading the wake-ups is the entire point, so a delay that is always +// the same value is the bug this guards against. +// +// The bound matters just as much in the other direction: the jitter is +// centred on leaseBusyRetryDelay rather than added to it, so that +// maxLeaseBusyRetries attempts still take about as long as they did +// before. A jitter that only ever extended the delay would quietly double +// how long a caller waits before a busy database is reported as an error. +func TestBusyRetryDelayIsJitteredAroundTheBase(t *testing.T) { + const ( + samples = 200 + low = leaseBusyRetryDelay / 2 + high = leaseBusyRetryDelay + leaseBusyRetryDelay/2 + ) + + seen := make(map[time.Duration]struct{}, samples) + var total time.Duration + + for range samples { + d := busyRetryDelay() + + if d < low || d >= high { + t.Fatalf("delay %v outside [%v, %v)", d, low, high) + } + + seen[d] = struct{}{} + total += d + } + + // A constant would produce exactly one distinct value. The threshold is + // deliberately far below `samples` so this cannot flake on collisions. + if len(seen) < samples/4 { + t.Errorf("only %d distinct delays in %d samples: contending writers "+ + "would retry in lockstep", len(seen), samples) + } + + // The mean should sit near the base. Tolerance is wide because this is + // a real random source and the test must not flake; it is here to catch + // a jitter that shifted the centre, not to measure the distribution. + mean := total / samples + drift := mean - leaseBusyRetryDelay + if drift < 0 { + drift = -drift + } + if drift > leaseBusyRetryDelay/4 { + t.Errorf("mean delay %v drifted from base %v: the retry budget is no "+ + "longer what maxLeaseBusyRetries was tuned for", mean, leaseBusyRetryDelay) + } +} diff --git a/store/sqlite/lease_test.go b/store/sqlite/lease_test.go new file mode 100644 index 0000000..4b87ed9 --- /dev/null +++ b/store/sqlite/lease_test.go @@ -0,0 +1,26 @@ +package sqlite_test + +import ( + "testing" + + "github.com/xraph/dispatch/store/storetest" +) + +func TestLeaseConformance(t *testing.T) { + // openSqliteStore already opens a migrated store on a per-test temp + // directory (store/sqlite/reap_test.go:19), so every subtest gets its + // own database for free. + storetest.RunLeaseSuite(t, func(t *testing.T) storetest.LeaseStore { + t.Helper() + + return openSqliteStore(t) + }) +} + +func TestDLQConformance(t *testing.T) { + storetest.RunDLQSuite(t, func(t *testing.T) storetest.DLQStore { + t.Helper() + + return openSqliteStore(t) + }) +} diff --git a/store/sqlite/migrations.go b/store/sqlite/migrations.go index d359952..6459061 100644 --- a/store/sqlite/migrations.go +++ b/store/sqlite/migrations.go @@ -289,5 +289,375 @@ func init() { return err }, }, + + // 006: Create artifacts and artifact links tables. + &migrate.Migration{ + Name: "create_artifacts_tables", + Version: "20260811120000", + Up: func(ctx context.Context, exec migrate.Executor) error { + if _, err := exec.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS dispatch_artifacts ( + id TEXT PRIMARY KEY, + backend TEXT NOT NULL, + bucket TEXT NOT NULL, + key TEXT NOT NULL, + size INTEGER NOT NULL DEFAULT 0, + content_hash TEXT, + content_type TEXT, + lifecycle TEXT NOT NULL, + scope_app_id TEXT, + scope_org_id TEXT, + expires_at TEXT, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + deleted_at TEXT + )`); err != nil { + return err + } + + // Partial unique index rather than a table constraint: a + // purged key must be reusable, so only live rows collide. + if _, err := exec.Exec(ctx, ` + CREATE UNIQUE INDEX IF NOT EXISTS uq_dispatch_artifacts_key + ON dispatch_artifacts (backend, bucket, key) + WHERE deleted_at IS NULL`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifacts_sweep + ON dispatch_artifacts (lifecycle, created_at) + WHERE deleted_at IS NULL`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifacts_purge + ON dispatch_artifacts (deleted_at) + WHERE deleted_at IS NOT NULL`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE TABLE IF NOT EXISTS dispatch_artifact_links ( + artifact_id TEXT NOT NULL REFERENCES dispatch_artifacts(id) ON DELETE CASCADE, + owner_kind TEXT NOT NULL, + owner_id TEXT NOT NULL, + name TEXT NOT NULL, + attempt INTEGER NOT NULL DEFAULT 0, + role TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), + PRIMARY KEY (artifact_id, owner_kind, owner_id, name, attempt) + )`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifact_links_owner + ON dispatch_artifact_links (owner_kind, owner_id)`); err != nil { + return err + } + + _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_artifact_links_artifact + ON dispatch_artifact_links (artifact_id)`) + + return err + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + if _, err := exec.Exec(ctx, `DROP TABLE IF EXISTS dispatch_artifact_links`); err != nil { + return err + } + + _, err := exec.Exec(ctx, `DROP TABLE IF EXISTS dispatch_artifacts`) + + return err + }, + }, + + // Lease columns. See the postgres migration of the same name for + // why the lease lives on the row. + // + // Every ADD COLUMN is guarded, for the reason spelled out on the + // resource migration below: SQLite has no ADD COLUMN IF NOT + // EXISTS and grove runs Up outside any transaction, so a failure + // at the third of the four would leave two columns added and no + // row in grove_migrations, and every retry from every pod would + // then die on "duplicate column name" forever, with no recovery + // short of hand-written DDL against a production database. The + // guard is the whole difference between a failed migration and a + // wedged one. + &migrate.Migration{ + Name: "add_job_lease_columns", + Version: "20260812120000", + Up: func(ctx context.Context, exec migrate.Executor) error { + for _, c := range []struct{ name, ddl string }{ + {"lease_epoch", `INTEGER NOT NULL DEFAULT 0`}, + // TEXT, not a numeric timestamp: SQLite has no + // timestamp type, every other time column here is + // text, and the reclaim predicate compares + // lease_expires_at against the driver's own rendering + // of a time.Time. See ReclaimExpiredLeases for why that + // rendering, not ISO-8601, is what has to be matched. + {"lease_expires_at", `TEXT`}, + {"lease_ttl", `INTEGER NOT NULL DEFAULT 0`}, + {"evict_count", `INTEGER NOT NULL DEFAULT 0`}, + } { + if err := addColumnIfMissing(ctx, exec, + "dispatch_jobs", c.name, c.ddl); err != nil { + return err + } + } + + // There is deliberately no backfill of lease_expires_at for + // rows already running here, matching the postgres migration + // of the same name. A one-shot backfill only sees rows that + // exist the moment it runs, so it misses every job an old pod + // claims later in a rolling upgrade, and every job claimed + // through job.Store without lease options, which never grants + // a lease at all. It also seeds an expiry in the past, which + // evicts jobs old pods are still actively running. + // ReclaimExpiredLeases adopts those rows on every sweep + // instead, gated on silence. See job.UnleasedReclaimGrace. + + _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_jobs_lease + ON dispatch_jobs (lease_expires_at) WHERE state = 'running'`) + + return err + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + if _, err := exec.Exec(ctx, `DROP INDEX IF EXISTS idx_dispatch_jobs_lease`); err != nil { + return err + } + + // Guarded for the same reason Up is: a Down that fails + // halfway must be re-runnable. + for _, col := range []string{ + "lease_epoch", "lease_expires_at", "lease_ttl", "evict_count", + } { + if err := dropColumnIfPresent(ctx, exec, "dispatch_jobs", col); err != nil { + return err + } + } + + return nil + }, + }, + + // Resource model columns for resource-aware scheduling. See the + // postgres migration of the same name for why the four canonical + // dimensions get real scalar columns alongside the JSON copy. + // + // Every column defaults to zero or empty, so rows written before + // this migration remain dequeueable by every worker during a + // rolling deploy. + &migrate.Migration{ + Name: "job_resource_columns", + Version: "20260812130000", + // Every ADD COLUMN is guarded, and the guard is not + // belt-and-braces — it is the only thing standing between a + // partial failure and a permanently wedged database. + // + // SQLite has no ADD COLUMN IF NOT EXISTS and grove runs Up + // bare, outside any transaction. So a failure at the sixth of + // ten statements — a disk-full, a SIGKILL mid-deploy, a + // cancelled context — leaves five columns added and no row in + // grove_migrations. The retry then dies on "duplicate column + // name: req_cpu_milli" and dies the same way forever: every + // pod that starts reports the same error, and there is no + // recovery short of an operator hand-writing DDL against a + // production database. Postgres avoided this with IF NOT + // EXISTS throughout; this is SQLite's equivalent. + // + // A transaction would have been the other answer. It is not + // taken because the sqlite migrate executor routes through + // the pooled *sql.DB rather than a pinned connection, so a + // BEGIN issued via Exec is not guaranteed to be the same + // session as the statements that follow it. Guarding each + // statement needs no assumption about connection affinity and + // makes the whole Up re-runnable from any point, which is + // strictly stronger than atomic-or-nothing. + Up: func(ctx context.Context, exec migrate.Executor) error { + columns := []struct{ name, ddl string }{ + {"req_cpu_milli", `INTEGER NOT NULL DEFAULT 0`}, + {"req_memory_bytes", `INTEGER NOT NULL DEFAULT 0`}, + {"req_disk_bytes", `INTEGER NOT NULL DEFAULT 0`}, + {"req_gpu_milli", `INTEGER NOT NULL DEFAULT 0`}, + {"req_custom_keys", `TEXT NOT NULL DEFAULT ''`}, + // No JSONB in SQLite: plain TEXT columns hold the + // full-fidelity JSON copy that fromJobModel reads + // Resources/ResourceLimits back from. + {"resource_requests", `TEXT`}, + {"resource_limits", `TEXT`}, + {"resource_class", `TEXT NOT NULL DEFAULT ''`}, + {"input_bytes", `INTEGER NOT NULL DEFAULT 0`}, + {"primary_input_hash", `TEXT`}, + } + + for _, c := range columns { + if err := addColumnIfMissing(ctx, exec, + "dispatch_jobs", c.name, c.ddl); err != nil { + return err + } + } + + // SQLite has no INCLUDE clause for a covering index, so + // the scalar columns the dequeue predicate reads go + // directly in the key list instead. + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_jobs_dequeue_res + ON dispatch_jobs (queue, priority DESC, run_at ASC, + req_cpu_milli, req_memory_bytes, + req_disk_bytes, req_gpu_milli) + WHERE state IN ('pending', 'retrying')`); err != nil { + return err + } + + // idx_dispatch_jobs_dequeue (the initial migration) is a + // strict key PREFIX of the index just created, over the + // same partial predicate, so every query it could serve + // the new one serves too. Keeping both costs a second + // B-tree insert on every enqueue and a second delete on + // every claim, forever, for a plan SQLite would never + // choose. Postgres drops its equivalent for the same + // reason. + // + // Dropped after the replacement exists, so there is no + // instant at which the dequeue statement has no index. + _, err := exec.Exec(ctx, `DROP INDEX IF EXISTS idx_dispatch_jobs_dequeue`) + + return err + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + // Restore the prefix index before dropping its superset, + // same ordering rule in reverse. + if _, err := exec.Exec(ctx, ` + CREATE INDEX IF NOT EXISTS idx_dispatch_jobs_dequeue + ON dispatch_jobs (queue, priority DESC, run_at ASC) + WHERE state IN ('pending', 'retrying')`); err != nil { + return err + } + + if _, err := exec.Exec(ctx, + `DROP INDEX IF EXISTS idx_dispatch_jobs_dequeue_res`); err != nil { + return err + } + + // Guarded for the same reason Up is: a Down that fails + // halfway must be re-runnable. + for _, col := range []string{ + "req_cpu_milli", "req_memory_bytes", "req_disk_bytes", + "req_gpu_milli", "req_custom_keys", "resource_requests", + "resource_limits", "resource_class", "input_bytes", + "primary_input_hash", + } { + if err := dropColumnIfPresent(ctx, exec, "dispatch_jobs", col); err != nil { + return err + } + } + + return nil + }, + }, + // Replay rebuilds a job from the DLQ row and enqueues it directly, + // so anything the row does not carry the replayed job silently takes + // a default for. Losing lease_ttl is the worst of them: a six-hour + // job replayed on the pool default lease lapses mid-run and is + // reclaimed and restarted forever without finishing. See dlq.Entry. + &migrate.Migration{ + Name: "dlq_job_execution_columns", + Version: "20260817120000", + Up: func(ctx context.Context, exec migrate.Executor) error { + // Every ADD COLUMN is guarded, for the reason spelled out on + // the lease and resource migrations above: SQLite has no ADD + // COLUMN IF NOT EXISTS and grove runs Up outside any + // transaction, so a failure partway through would leave some + // columns added and no row in grove_migrations, and every + // retry from every pod would then die on "duplicate column + // name" forever. + for _, c := range []struct{ name, ddl string }{ + {"priority", `INTEGER NOT NULL DEFAULT 0`}, + {"timeout", `INTEGER NOT NULL DEFAULT 0`}, + {"lease_ttl", `INTEGER NOT NULL DEFAULT 0`}, + {"artifact_bindings", `BLOB`}, + // Nullable TEXT, matching dispatch_jobs: resource. + // EncodeSetString writes NULL for a zero Set rather than + // an empty object, so a job that declared nothing reads + // back as a nil Set rather than an empty one. + {"resources", `TEXT`}, + {"resource_limits", `TEXT`}, + {"resource_class", `TEXT NOT NULL DEFAULT ''`}, + {"input_bytes", `INTEGER NOT NULL DEFAULT 0`}, + {"primary_input_hash", `TEXT`}, + } { + if err := addColumnIfMissing(ctx, exec, + "dispatch_dlq", c.name, c.ddl); err != nil { + return err + } + } + + return nil + }, + Down: func(ctx context.Context, exec migrate.Executor) error { + // Guarded for the same reason Up is: a Down that fails + // halfway must be re-runnable. + for _, col := range []string{ + "priority", "timeout", "lease_ttl", "artifact_bindings", + "resources", "resource_limits", "resource_class", + "input_bytes", "primary_input_hash", + } { + if err := dropColumnIfPresent(ctx, exec, "dispatch_dlq", col); err != nil { + return err + } + } + + return nil + }, + }, ) } + +// columnExists reports whether table already has the named column. +// +// pragma_table_info is the table-valued form of PRAGMA table_info, which +// means it can be queried with bind parameters like any other relation +// rather than string-formatted into a PRAGMA statement. +func columnExists(ctx context.Context, exec migrate.Executor, table, column string) (bool, error) { + rows, err := exec.Query(ctx, + `SELECT 1 FROM pragma_table_info(?) WHERE name = ?`, table, column) + if err != nil { + return false, err + } + + found := rows.Next() + + if closeErr := rows.Close(); closeErr != nil { + return false, closeErr + } + + return found, rows.Err() +} + +// addColumnIfMissing is SQLite's stand-in for ADD COLUMN IF NOT EXISTS. +func addColumnIfMissing(ctx context.Context, exec migrate.Executor, table, column, ddl string) error { + present, err := columnExists(ctx, exec, table, column) + if err != nil || present { + return err + } + + _, err = exec.Exec(ctx, `ALTER TABLE `+table+` ADD COLUMN `+column+` `+ddl) + + return err +} + +// dropColumnIfPresent is SQLite's stand-in for DROP COLUMN IF EXISTS. +func dropColumnIfPresent(ctx context.Context, exec migrate.Executor, table, column string) error { + present, err := columnExists(ctx, exec, table, column) + if err != nil || !present { + return err + } + + _, err = exec.Exec(ctx, `ALTER TABLE `+table+` DROP COLUMN `+column) + + return err +} diff --git a/store/sqlite/migrations_test.go b/store/sqlite/migrations_test.go new file mode 100644 index 0000000..66a557b --- /dev/null +++ b/store/sqlite/migrations_test.go @@ -0,0 +1,457 @@ +package sqlite_test + +import ( + "context" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/xraph/grove" + "github.com/xraph/grove/driver" + "github.com/xraph/grove/drivers/sqlitedriver" + _ "github.com/xraph/grove/drivers/sqlitedriver/sqlitemigrate" // registers the sqlite migrate executor + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + sqlitestore "github.com/xraph/dispatch/store/sqlite" +) + +// resourceMigrationVersion is the version string of the migration under +// test, restated here so the test can delete its bookkeeping row. +const resourceMigrationVersion = "20260812130000" + +// leaseMigrationVersion is the same for the lease migration, which runs +// immediately before it. +const leaseMigrationVersion = "20260812120000" + +// leaseColumns is every column the lease migration adds, in the order it +// adds them. +var leaseColumns = []string{ + "lease_epoch", "lease_expires_at", "lease_ttl", "evict_count", +} + +// openMigratedWithDriver is openSqliteStore with the raw driver handed +// back too, so a test can reach past the store and mutate the schema +// into shapes no code path produces. +func openMigratedWithDriver(t *testing.T) (*sqlitestore.Store, driver.Driver, *grove.DB) { + t.Helper() + + ctx := context.Background() + + drv := sqlitedriver.New() + if err := drv.Open(ctx, filepath.Join(t.TempDir(), "dispatch.db")); err != nil { + t.Fatalf("open sqlitedriver: %v", err) + } + + db, err := grove.Open(drv) + if err != nil { + t.Fatalf("grove open: %v", err) + } + + t.Cleanup(func() { _ = db.Close() }) + + s := sqlitestore.New(db) + if err := s.Migrate(ctx); err != nil { + t.Fatalf("migrate: %v", err) + } + + return s, drv, db +} + +func mustExec(t *testing.T, drv driver.Driver, stmt string, args ...any) { + t.Helper() + + if _, err := drv.Exec(context.Background(), stmt, args...); err != nil { + t.Fatalf("exec %q: %v", stmt, err) + } +} + +// hasColumn reports whether dispatch_jobs currently has the named column. +// +// The table is fixed rather than a parameter: every migration in this file +// changes the one table, and a parameter that only ever receives one value +// reads as generality the tests do not have. +func hasColumn(t *testing.T, drv driver.Driver, column string) bool { + t.Helper() + + rows, err := drv.Query(context.Background(), + `SELECT 1 FROM pragma_table_info('dispatch_jobs') WHERE name = ?`, column) + if err != nil { + t.Fatalf("pragma_table_info: %v", err) + } + + found := rows.Next() + + if err = rows.Close(); err != nil { + t.Fatalf("close pragma rows: %v", err) + } + + return found +} + +// TestResourceMigrationSurvivesAPartialApplication is the proof that the +// resource migration can no longer wedge a SQLite deployment +// permanently. +// +// The failure it reproduces is not hypothetical and has no workaround. +// SQLite has no ADD COLUMN IF NOT EXISTS, and grove executes Up outside +// any transaction (migrate.Orchestrator.Migrate calls m.Up and only then +// RecordApplied). So a process killed — or a disk filled, or a context +// cancelled — partway through the ten ADD COLUMNs leaves some columns +// added and NO row in grove_migrations. The next start re-runs Up from +// the top, hits ALTER TABLE ADD COLUMN on a column that already exists, +// and fails with "duplicate column name". It then fails identically on +// every subsequent start, forever, on every pod, with no recovery short +// of an operator hand-writing DDL against a live database. +// +// The state below is exactly what a crash after the fourth statement +// leaves behind: the first four columns present, the remaining six +// absent, the index absent, and the migration unrecorded. A retry must +// complete it and the store must work afterwards. +// +// Mutation-verified: reverting Up to unguarded ALTER TABLE ADD COLUMN +// fails here with "duplicate column name: req_cpu_milli". +func TestResourceMigrationSurvivesAPartialApplication(t *testing.T) { + s, drv, db := openMigratedWithDriver(t) + ctx := context.Background() + + // Rewind to the halfway state a crash would have left. + for _, col := range []string{ + "req_custom_keys", "resource_requests", "resource_limits", + "resource_class", "input_bytes", "primary_input_hash", + } { + mustExec(t, drv, `ALTER TABLE dispatch_jobs DROP COLUMN `+col) + } + + mustExec(t, drv, `DROP INDEX IF EXISTS idx_dispatch_jobs_dequeue_res`) + mustExec(t, drv, `DELETE FROM grove_migrations WHERE version = '`+resourceMigrationVersion+`'`) + + for _, col := range []string{"req_cpu_milli", "req_gpu_milli"} { + if !hasColumn(t, drv, col) { + t.Fatalf("fixture is wrong: %s should still be present", col) + } + } + + if hasColumn(t, drv, "primary_input_hash") { + t.Fatal("fixture is wrong: primary_input_hash should have been dropped") + } + + // The retry every restarting pod performs. + retried := sqlitestore.New(db) + if err := retried.Migrate(ctx); err != nil { + t.Fatalf("re-running a half-applied migration must succeed, got: %v\n"+ + "a SQLite deployment that failed partway through this migration would be "+ + "unrecoverable without hand-written DDL", err) + } + + for _, col := range []string{ + "req_cpu_milli", "req_memory_bytes", "req_disk_bytes", "req_gpu_milli", + "req_custom_keys", "resource_requests", "resource_limits", + "resource_class", "input_bytes", "primary_input_hash", + } { + if !hasColumn(t, drv, col) { + t.Errorf("column %s missing after the retry", col) + } + } + + // And the schema is usable, not merely present. + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "after-retry", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + Resources: resource.Set{resource.Memory: 8 << 30, "fpga": 1}, + PrimaryInputHash: "blake3:abc", + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob after the retry: %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob after the retry: %v", err) + } + + if got.Resources[resource.Memory] != 8<<30 || got.Resources["fpga"] != 1 { + t.Errorf("resources = %v, want the declaration intact", got.Resources) + } +} + +// TestResourceMigrationIsFullyIdempotent covers the other partial shape: +// a crash AFTER the last statement but BEFORE RecordApplied, which is a +// real window because the two are separate calls. Re-running Up with +// every column and the index already in place must be a no-op. +func TestResourceMigrationIsFullyIdempotent(t *testing.T) { + _, drv, db := openMigratedWithDriver(t) + ctx := context.Background() + + mustExec(t, drv, `DELETE FROM grove_migrations WHERE version = '`+resourceMigrationVersion+`'`) + + if err := sqlitestore.New(db).Migrate(ctx); err != nil { + t.Fatalf("re-running a fully applied migration must be a no-op, got: %v", err) + } +} + +// TestResourceMigrationDropsTheRedundantDequeueIndex pins the other half +// of the change: idx_dispatch_jobs_dequeue is a strict key prefix of +// idx_dispatch_jobs_dequeue_res over the identical partial predicate, so +// shipping both means every enqueue pays two B-tree inserts and every +// claim two deletes, forever, for a plan SQLite would never choose. +func TestResourceMigrationDropsTheRedundantDequeueIndex(t *testing.T) { + _, drv, _ := openMigratedWithDriver(t) + + rows, err := drv.Query(context.Background(), + `SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'dispatch_jobs'`) + if err != nil { + t.Fatalf("read sqlite_master: %v", err) + } + + var names []string + + for rows.Next() { + var name string + if err = rows.Scan(&name); err != nil { + t.Fatalf("scan index name: %v", err) + } + + names = append(names, name) + } + + if err = rows.Close(); err != nil { + t.Fatalf("close sqlite_master rows: %v", err) + } + + var sawSuperset bool + + for _, n := range names { + if n == "idx_dispatch_jobs_dequeue" { + t.Errorf("idx_dispatch_jobs_dequeue survived the resource migration; it is a key "+ + "prefix of idx_dispatch_jobs_dequeue_res, so both are maintained on every "+ + "write for no plan. Indexes present: %s", strings.Join(names, ", ")) + } + + if n == "idx_dispatch_jobs_dequeue_res" { + sawSuperset = true + } + } + + if !sawSuperset { + t.Fatalf("idx_dispatch_jobs_dequeue_res is missing; indexes present: %s", + strings.Join(names, ", ")) + } +} + +// TestLeaseMigrationSurvivesAPartialApplication is the lease migration's +// version of the proof above, and it needs its own because 008 runs +// BEFORE 009: a deployment wedged here never reaches the resource +// migration at all. +// +// Same mechanism, same consequence. SQLite has no ADD COLUMN IF NOT +// EXISTS and grove executes Up outside any transaction (Orchestrator +// calls m.Up and only then RecordApplied), so a process killed partway +// through the four ADD COLUMNs leaves some columns added and no row in +// grove_migrations. Unguarded, the next start re-runs Up from the top and +// dies on "duplicate column name" — identically, forever, on every pod. +// +// The state below is what a crash after the second statement leaves: the +// first two lease columns present, the last two absent, the index absent, +// and the migration unrecorded. +// +// Mutation-verified: reverting Up to bare ALTER TABLE ADD COLUMN fails +// here with "duplicate column name: lease_epoch". +func TestLeaseMigrationSurvivesAPartialApplication(t *testing.T) { + s, drv, db := openMigratedWithDriver(t) + ctx := context.Background() + + for _, col := range []string{"lease_ttl", "evict_count"} { + mustExec(t, drv, `ALTER TABLE dispatch_jobs DROP COLUMN `+col) + } + + mustExec(t, drv, `DROP INDEX IF EXISTS idx_dispatch_jobs_lease`) + mustExec(t, drv, `DELETE FROM grove_migrations WHERE version = '`+leaseMigrationVersion+`'`) + + if !hasColumn(t, drv, "lease_epoch") { + t.Fatal("fixture is wrong: lease_epoch should still be present") + } + + if hasColumn(t, drv, "evict_count") { + t.Fatal("fixture is wrong: evict_count should have been dropped") + } + + // The retry every restarting pod performs. + if err := sqlitestore.New(db).Migrate(ctx); err != nil { + t.Fatalf("re-running a half-applied lease migration must succeed, got: %v\n"+ + "a SQLite deployment that failed partway through this migration would be "+ + "unrecoverable without hand-written DDL", err) + } + + for _, col := range leaseColumns { + if !hasColumn(t, drv, col) { + t.Errorf("column %s missing after the retry", col) + } + } + + // And the schema is usable, not merely present. + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "after-lease-retry", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + LeaseTTL: 6 * time.Hour, + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob after the retry: %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob after the retry: %v", err) + } + + if got.LeaseTTL != 6*time.Hour { + t.Errorf("LeaseTTL = %v, want the declaration intact", got.LeaseTTL) + } +} + +// TestLeaseMigrationIsFullyIdempotent covers the other partial shape: a +// crash AFTER the last statement but BEFORE RecordApplied, which is a +// real window because the two are separate calls. +func TestLeaseMigrationIsFullyIdempotent(t *testing.T) { + _, drv, db := openMigratedWithDriver(t) + + mustExec(t, drv, `DELETE FROM grove_migrations WHERE version = '`+leaseMigrationVersion+`'`) + + if err := sqlitestore.New(db).Migrate(context.Background()); err != nil { + t.Fatalf("re-running a fully applied lease migration must be a no-op, got: %v", err) + } +} + +// TestReclaimAdoptsRunningJobsWithoutLease covers a running job carrying +// no lease at all, which is the row shape that used to be stranded +// permanently. lease_expires_at is NULL, ReclaimExpiredLeases required a +// non-NULL expiry to consider a row at all (job.Lease.IsExpired +// deliberately reads a zero expiry as "never leased", not "expired"), the +// pool's reaper no longer calls ReapStaleJobs once the backend implements +// job.LeaseStore, and dequeue claims only pending and retrying rows. +// +// Migration 008 used to seed an expiry for these rows and no longer does; +// see that migration for why a one-shot backfill was the wrong mechanism. +// +// Asserting that reclaim COLLECTS the row, rather than that some column +// changed, is what makes this test worth having on SQLite specifically. +// There is no timestamp type here, so the predicate is a string comparison +// against the driver's own rendering of a time.Time, and the backfill this +// replaced originally used strftime, which sorts above every +// driver-written timestamp ('T' > ' ') and silently matched nothing. Only +// an assertion on the collected row catches that class of mistake. +func TestReclaimAdoptsRunningJobsWithoutLease(t *testing.T) { + silent := time.Now().UTC().Add(-job.UnleasedReclaimGrace - time.Minute) + older := silent.Add(-time.Hour) + fresh := time.Now().UTC() + + tests := []struct { + name string + heartbeat time.Time + started time.Time + want bool + why string + }{ + { + name: "silent heartbeat is adopted", + heartbeat: silent, + started: older, + want: true, + why: "abandoned by a worker that stopped reporting", + }, + { + name: "fresh heartbeat is left alone", + heartbeat: fresh, + started: older, + want: false, + why: "still reporting, so it belongs to a healthy worker", + }, + { + name: "silent started_at is adopted when it never heartbeated", + started: silent, + want: true, + why: "died before its first beat", + }, + { + name: "fresh started_at is left alone", + started: fresh, + want: false, + why: "just claimed; its first heartbeat is not due yet", + }, + { + name: "neither timestamp is left alone", + want: false, + why: "no timestamp to establish age from", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s, drv, _ := openMigratedWithDriver(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "mid-flight", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + } + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob: %v", err) + } + + // The row an upgrading fleet finds, or that a caller without + // lease options writes: running, with no lease at all. The + // timestamps are bound as time.Time so they are rendered by the + // same driver path the store itself writes through. + mustExec(t, drv, ` + UPDATE dispatch_jobs + SET state = 'running', worker_id = 'w-1', + heartbeat_at = ?, started_at = ?, lease_expires_at = NULL + WHERE id = ?`, + nullableTime(tt.heartbeat), nullableTime(tt.started), j.ID.String()) + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 10) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + + got := len(reclaimed) == 1 && reclaimed[0].ID == j.ID + if got != tt.want { + if tt.want { + t.Fatalf("job was not reclaimed but should have been: %s", tt.why) + } + + t.Fatalf("job was reclaimed but must not be: %s", tt.why) + } + if tt.want && reclaimed[0].State != job.StatePending { + t.Errorf("reclaimed job state = %v, want pending", reclaimed[0].State) + } + }) + } +} + +// nullableTime renders the zero time as a NULL bind rather than as year +// one, so a test case can express "this column was never written". +func nullableTime(t time.Time) any { + if t.IsZero() { + return nil + } + + return t +} diff --git a/store/sqlite/models.go b/store/sqlite/models.go index 2e37da4..b6fda68 100644 --- a/store/sqlite/models.go +++ b/store/sqlite/models.go @@ -14,6 +14,7 @@ import ( "github.com/xraph/dispatch/event" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" "github.com/xraph/dispatch/workflow" ) @@ -41,9 +42,49 @@ type jobModel struct { Timeout int64 `grove:"timeout,notnull,default:0"` CreatedAt time.Time `grove:"created_at,notnull"` UpdatedAt time.Time `grove:"updated_at,notnull"` + + LeaseEpoch int `grove:"lease_epoch,notnull,default:0"` + LeaseExpiresAt *time.Time `grove:"lease_expires_at"` + LeaseTTL int64 `grove:"lease_ttl,notnull,default:0"` + EvictCount int `grove:"evict_count,notnull,default:0"` + + // The four canonical dimensions get real scalar columns because the + // dequeue predicate compares them and must behave identically across + // five backends; JSON comparison semantics are not portable. They are + // derived from Resources by toJobModel -- the caller never sets them + // directly. + ReqCPUMilli int64 `grove:"req_cpu_milli,notnull,default:0"` + ReqMemoryBytes int64 `grove:"req_memory_bytes,notnull,default:0"` + ReqDiskBytes int64 `grove:"req_disk_bytes,notnull,default:0"` + ReqGPUMilli int64 `grove:"req_gpu_milli,notnull,default:0"` + ReqCustomKeys string `grove:"req_custom_keys,notnull,default:''"` + + // ResourceRequests and ResourceLimits are the full-fidelity JSON copy + // of Resources / ResourceLimits, including custom keys the scalar + // columns above do not carry. fromJobModel reads Resources back from + // here, not from the scalars. SQLite has no JSONB type, so these are + // plain TEXT columns; *string rather than []byte or string so a NULL + // column (undeclared job) round-trips as nil instead of an empty + // string, mirroring resource.EncodeSetString's NULL-for-zero-Set + // contract. + ResourceRequests *string `grove:"resource_requests"` + ResourceLimits *string `grove:"resource_limits"` + ResourceClass string `grove:"resource_class,notnull,default:''"` + InputBytes int64 `grove:"input_bytes,notnull,default:0"` + PrimaryInputHash string `grove:"primary_input_hash"` } -func toJobModel(j *job.Job) *jobModel { +func toJobModel(j *job.Job) (*jobModel, error) { + reqJSON, err := resource.EncodeSetString(j.Resources) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: marshal job resources: %w", err) + } + + limitsJSON, err := resource.EncodeSetString(j.ResourceLimits) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: marshal job resource limits: %w", err) + } + return &jobModel{ ID: j.ID.String(), Name: j.Name, @@ -64,7 +105,23 @@ func toJobModel(j *job.Job) *jobModel { Timeout: j.Timeout.Nanoseconds(), CreatedAt: j.CreatedAt, UpdatedAt: j.UpdatedAt, - } + + LeaseEpoch: j.LeaseEpoch, + LeaseExpiresAt: j.LeaseExpiresAt, + LeaseTTL: j.LeaseTTL.Nanoseconds(), + EvictCount: j.EvictCount, + + ReqCPUMilli: j.Resources[resource.CPU], + ReqMemoryBytes: j.Resources[resource.Memory], + ReqDiskBytes: j.Resources[resource.Disk], + ReqGPUMilli: j.Resources[resource.GPU], + ReqCustomKeys: resource.EncodeCustomKeys(j.Resources), + ResourceRequests: reqJSON, + ResourceLimits: limitsJSON, + ResourceClass: j.ResourceClass, + InputBytes: j.InputBytes, + PrimaryInputHash: j.PrimaryInputHash, + }, nil } func fromJobModel(m *jobModel) (*job.Job, error) { @@ -73,6 +130,16 @@ func fromJobModel(m *jobModel) (*job.Job, error) { return nil, fmt.Errorf("dispatch/sqlite: parse job id %q: %w", m.ID, err) } + resources, err := resource.DecodeSetString(m.ResourceRequests) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: unmarshal job resources: %w", err) + } + + limits, err := resource.DecodeSetString(m.ResourceLimits) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: unmarshal job resource limits: %w", err) + } + j := &job.Job{ Entity: dispatch.Entity{ CreatedAt: m.CreatedAt, @@ -94,6 +161,17 @@ func fromJobModel(m *jobModel) (*job.Job, error) { CompletedAt: m.CompletedAt, HeartbeatAt: m.HeartbeatAt, Timeout: time.Duration(m.Timeout), + + LeaseEpoch: m.LeaseEpoch, + LeaseExpiresAt: m.LeaseExpiresAt, + LeaseTTL: time.Duration(m.LeaseTTL), + EvictCount: m.EvictCount, + + Resources: resources, + ResourceLimits: limits, + ResourceClass: m.ResourceClass, + InputBytes: m.InputBytes, + PrimaryInputHash: m.PrimaryInputHash, } if m.WorkerID != "" { @@ -291,9 +369,34 @@ type dlqEntryModel struct { FailedAt time.Time `grove:"failed_at,notnull"` ReplayedAt *time.Time `grove:"replayed_at"` CreatedAt time.Time `grove:"created_at,notnull"` + + // Carried so Replay can rebuild a job that behaves like the failed + // one; see the dlq.Entry doc. The resource sets use the same nullable + // TEXT encoding jobModel uses (resource.EncodeSetString writes NULL + // for a zero Set) rather than scalar columns: nothing queries a DLQ + // row by resource requirement, so the scalar split that dequeue's fit + // predicate needs buys nothing here. + Priority int `grove:"priority,notnull,default:0"` + Timeout int64 `grove:"timeout,notnull,default:0"` + LeaseTTL int64 `grove:"lease_ttl,notnull,default:0"` + ArtifactBindings []byte `grove:"artifact_bindings"` + Resources *string `grove:"resources"` + ResourceLimits *string `grove:"resource_limits"` + ResourceClass string `grove:"resource_class,notnull,default:''"` + InputBytes int64 `grove:"input_bytes,notnull,default:0"` + PrimaryInputHash string `grove:"primary_input_hash"` } -func toDLQModel(e *dlq.Entry) *dlqEntryModel { +func toDLQModel(e *dlq.Entry) (*dlqEntryModel, error) { + resources, err := resource.EncodeSetString(e.Resources) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: encode dlq resources: %w", err) + } + limits, err := resource.EncodeSetString(e.ResourceLimits) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: encode dlq resource limits: %w", err) + } + return &dlqEntryModel{ ID: e.ID.String(), JobID: e.JobID.String(), @@ -308,7 +411,17 @@ func toDLQModel(e *dlq.Entry) *dlqEntryModel { FailedAt: e.FailedAt, ReplayedAt: e.ReplayedAt, CreatedAt: e.CreatedAt, - } + + Priority: e.Priority, + Timeout: int64(e.Timeout), + LeaseTTL: int64(e.LeaseTTL), + ArtifactBindings: e.ArtifactBindings, + Resources: resources, + ResourceLimits: limits, + ResourceClass: e.ResourceClass, + InputBytes: e.InputBytes, + PrimaryInputHash: e.PrimaryInputHash, + }, nil } func fromDLQModel(m *dlqEntryModel) (*dlq.Entry, error) { @@ -322,6 +435,15 @@ func fromDLQModel(m *dlqEntryModel) (*dlq.Entry, error) { return nil, fmt.Errorf("dispatch/sqlite: parse job id %q: %w", m.JobID, err) } + resources, err := resource.DecodeSetString(m.Resources) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: decode dlq resources: %w", err) + } + limits, err := resource.DecodeSetString(m.ResourceLimits) + if err != nil { + return nil, fmt.Errorf("dispatch/sqlite: decode dlq resource limits: %w", err) + } + return &dlq.Entry{ ID: parsedID, JobID: parsedJobID, @@ -336,6 +458,16 @@ func fromDLQModel(m *dlqEntryModel) (*dlq.Entry, error) { FailedAt: m.FailedAt, ReplayedAt: m.ReplayedAt, CreatedAt: m.CreatedAt, + + Priority: m.Priority, + Timeout: time.Duration(m.Timeout), + LeaseTTL: time.Duration(m.LeaseTTL), + ArtifactBindings: m.ArtifactBindings, + Resources: resources, + ResourceLimits: limits, + ResourceClass: m.ResourceClass, + InputBytes: m.InputBytes, + PrimaryInputHash: m.PrimaryInputHash, }, nil } diff --git a/store/sqlite/resource_test.go b/store/sqlite/resource_test.go new file mode 100644 index 0000000..e91f58e --- /dev/null +++ b/store/sqlite/resource_test.go @@ -0,0 +1,142 @@ +package sqlite_test + +import ( + "context" + "testing" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// TestSqliteRoundTripsResources mirrors store/postgres's resource round-trip +// test: a job with cpu/memory/a custom key/limits/class/input signal must +// come back identical, and the custom key specifically must survive since +// it only lives in the JSON column, not the scalar columns. +func TestSqliteRoundTripsResources(t *testing.T) { + s := openSqliteStore(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "tessellate.model", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + Resources: resource.Set{ + resource.CPU: 4000, resource.Memory: 16 << 30, "fpga": 2, + }, + ResourceLimits: resource.Set{resource.Memory: 16 << 30}, + ResourceClass: "heavy", + InputBytes: 4 << 30, + PrimaryInputHash: "blake3:9f2a", + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + + if got.Resources[resource.CPU] != 4000 { + t.Errorf("cpu = %d, want 4000", got.Resources[resource.CPU]) + } + if got.Resources["fpga"] != 2 { + t.Errorf("custom key lost in round-trip: %v", got.Resources) + } + if got.ResourceLimits[resource.Memory] != 16<<30 { + t.Errorf("limits = %v", got.ResourceLimits) + } + if got.ResourceClass != "heavy" { + t.Errorf("class = %q, want heavy", got.ResourceClass) + } + if got.InputBytes != 4<<30 || got.PrimaryInputHash != "blake3:9f2a" { + t.Errorf("input signal lost: bytes=%d hash=%q", + got.InputBytes, got.PrimaryInputHash) + } +} + +// TestSqliteJobWithNoResourcesRoundTrips pins the backward-compatibility +// contract: a job enqueued without any resource declaration must come back +// with a zero Set, not a Set containing zero-valued canonical keys -- those +// are indistinguishable to a caller checking IsZero(), but the row-level +// NULL-vs-"{}" distinction is what a rolling deploy depends on. +// +// So this asserts the STORED value, not just the decoded one. IsZero() is +// true for nil, for Set{} and for Set{"memory": 0} alike, which means the +// decoded-side check the comment above describes cannot actually see the +// distinction it names. Reading the column back raw can: NULL is what a +// worker still running the pre-resource code reads as "no requirement". +// Redis and Mongo already assert on their raw stored value; this is +// SQLite's. +func TestSqliteJobWithNoResourcesRoundTrips(t *testing.T) { + s, drv, _ := openMigratedWithDriver(t) + ctx := context.Background() + + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: "notify.user", + Queue: "default", + State: job.StatePending, + Payload: []byte(`{}`), + } + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob() error = %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob() error = %v", err) + } + if !got.Resources.IsZero() { + t.Errorf("Resources = %v, want zero", got.Resources) + } + + var ( + requestsNull int + limitsNull int + cpu int64 + customKeys string + class string + inputBytes int64 + ) + + row := drv.QueryRow(ctx, ` + SELECT resource_requests IS NULL, + resource_limits IS NULL, + req_cpu_milli, + req_custom_keys, + resource_class, + input_bytes + FROM dispatch_jobs WHERE id = ?`, j.ID.String()) + + if err = row.Scan(&requestsNull, &limitsNull, &cpu, + &customKeys, &class, &inputBytes); err != nil { + t.Fatalf("raw column read: %v", err) + } + + if requestsNull != 1 { + t.Error(`resource_requests is not NULL for an undeclared job; a stored "{}" decodes ` + + `to the same zero Set but is a different row, and NULL is what a worker running ` + + `the pre-resource code reads as "no requirement"`) + } + + if limitsNull != 1 { + t.Error("resource_limits is not NULL for an undeclared job") + } + + // The scalar columns are NOT NULL DEFAULT, so they must be present + // and zero — that is what lets the dequeue comparisons be bare + // rather than COALESCEd. + if cpu != 0 || customKeys != "" || class != "" || inputBytes != 0 { + t.Errorf("scalar columns = cpu %d, keys %q, class %q, input %d; want all zero/empty", + cpu, customKeys, class, inputBytes) + } +} diff --git a/store/sqlite/store.go b/store/sqlite/store.go index d581dd3..97b45aa 100644 --- a/store/sqlite/store.go +++ b/store/sqlite/store.go @@ -14,6 +14,7 @@ import ( _ "github.com/xraph/grove/drivers/sqlitedriver/sqlitemigrate" // register sqlite migration executor "github.com/xraph/grove/migrate" + "github.com/xraph/dispatch/artifact" "github.com/xraph/dispatch/cluster" "github.com/xraph/dispatch/cron" "github.com/xraph/dispatch/dlq" @@ -30,6 +31,8 @@ var ( _ dlq.Store = (*Store)(nil) _ event.Store = (*Store)(nil) _ cluster.Store = (*Store)(nil) + _ artifact.Store = (*Store)(nil) + _ job.LeaseStore = (*Store)(nil) ) // Store is a grove ORM implementation of store.Store using SQLite dialect. diff --git a/store/store.go b/store/store.go index 03a3221..2a4bcec 100644 --- a/store/store.go +++ b/store/store.go @@ -7,6 +7,7 @@ package store import ( "context" + "github.com/xraph/dispatch/artifact" "github.com/xraph/dispatch/cluster" "github.com/xraph/dispatch/cron" "github.com/xraph/dispatch/dlq" @@ -37,6 +38,7 @@ type Store interface { dlq.Store event.Store cluster.Store + artifact.Store // Migrate runs all schema migrations. Migrate(ctx context.Context) error diff --git a/store/storetest/dequeue.go b/store/storetest/dequeue.go new file mode 100644 index 0000000..e8a89e6 --- /dev/null +++ b/store/storetest/dequeue.go @@ -0,0 +1,1047 @@ +package storetest + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// GiB is one gibibyte in bytes, the unit the memory and disk cases use. +const GiB = int64(1) << 30 + +// RunDequeueSuite runs the resource-aware dequeue conformance suite +// against a backend. +// +// newStore is called once per subtest and receives that subtest's +// *testing.T, so a backend that stands up a container or opens a +// database can register teardown on the T that owns it. Closing over the +// parent T instead would hold every subtest's resources open until the +// whole suite finished, and would turn a setup t.Fatalf into a FailNow on +// a parent test. +// +// newStore may return the same underlying store on every call. Each case +// enqueues onto its own queue and asserts only on the jobs it created, so +// cases cannot interfere — which matters because starting a fresh +// Postgres, Mongo, or Redis container per subtest would dominate the +// runtime of the whole suite. +func RunDequeueSuite(t *testing.T, newStore func(t *testing.T) job.Store) { + t.Helper() + + cases := []struct { + name string + fn func(*testing.T, job.Store) + }{ + {"ZeroBudgetSelectsEverything", testZeroBudgetSelectsEverything}, + {"MemoryBudgetFilters", testMemoryBudgetFilters}, + {"CPUBudgetFilters", testCPUBudgetFilters}, + {"DiskBudgetFilters", testDiskBudgetFilters}, + {"GPUBudgetFilters", testGPUBudgetFilters}, + {"AbsentBudgetKeyIsUnconstrained", testAbsentBudgetKeyIsUnconstrained}, + { + "BoundedBudgetWithNoCustomKeysRejectsCustomRequirement", + testBoundedBudgetWithNoCustomKeysRejectsCustomRequirement, + }, + {"ExplicitZeroBudgetKeyStillFilters", testExplicitZeroBudgetKeyStillFilters}, + {"ZeroRequirementAlwaysFits", testZeroRequirementAlwaysFits}, + {"ExactFitIsClaimable", testExactFitIsClaimable}, + {"CustomKeyContainmentFilters", testCustomKeyContainmentFilters}, + {"CustomKeyPrefixDoesNotFalselyMatch", testCustomKeyPrefixDoesNotFalselyMatch}, + {"CustomKeySubsetOfOfferedKeysIsClaimable", testCustomKeySubsetIsClaimable}, + {"PriorityOrderingPreservedWithinBudget", testPriorityOrderingPreservedWithinBudget}, + {"LimitTruncatesAfterOrdering", testLimitTruncatesAfterOrdering}, + {"PreferHashesAloneOrdersWithoutFiltering", testPreferHashesAloneOrdersWithoutFiltering}, + { + "PreferHashesSortWithinPriorityBandAndNeverFilter", + testPreferHashesSortWithinPriorityBand, + }, + { + "LocalityDecidesWhichRowsSurviveATightLimit", + testLocalityDecidesWhichRowsSurviveATightLimit, + }, + {"ReservedForRestrictsToOneJob", testReservedForRestrictsToOneJob}, + {"NonPositiveLimitClaimsNothing", testNonPositiveLimitClaimsNothing}, + {"ClaimIsAtomicUnderConcurrency", testClaimIsAtomicUnderConcurrency}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + tc.fn(t, newStore(t)) + }) + } +} + +// ────────────────────────────────────────────────── +// Fixtures +// ────────────────────────────────────────────────── + +// runAtBase is the anchor every fixture's RunAt is expressed against. It +// is an hour in the past so every job is ready immediately, and every +// case that asserts ordering derives its RunAt from this anchor with an +// explicit offset. Ordering is therefore a property of the data, never of +// how long the backend took to answer the previous call. +func runAtBase() time.Time { + return time.Now().UTC().Add(-time.Hour).Truncate(time.Millisecond) +} + +// option mutates a fixture before it is enqueued. +type fitOption func(*job.Job) + +// withPriority sets the job's scheduling priority. +func withPriority(p int) fitOption { + return func(j *job.Job) { j.Priority = p } +} + +// withRunAtOffset moves the job's RunAt forward from the shared anchor. +// Offsets must stay under an hour so the job remains ready to run. +func withRunAtOffset(d time.Duration) fitOption { + return func(j *job.Job) { j.RunAt = runAtBase().Add(d) } +} + +// withHash sets the locality signal PreferHashes matches against. +func withHash(h string) fitOption { + return func(j *job.Job) { j.PrimaryInputHash = h } +} + +// newJob builds a pending job that is ready to run now, on the given +// queue, requiring res. name is echoed in every failure message, so it +// should describe the job's role in the case. +func newFitJob(name, queue string, res resource.Set, opts ...fitOption) *job.Job { + j := &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: name, + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: runAtBase(), + Resources: res, + } + + for _, opt := range opts { + opt(j) + } + + return j +} + +// ────────────────────────────────────────────────── +// Assertions +// ────────────────────────────────────────────────── + +func mustEnqueue(t *testing.T, s job.Store, jobs ...*job.Job) { + t.Helper() + + for _, j := range jobs { + if err := s.EnqueueJob(context.Background(), j); err != nil { + t.Fatalf("EnqueueJob(%s): %v", j.Name, err) + } + } +} + +func mustDequeue(t *testing.T, s job.Store, opts job.DequeueOpts) []*job.Job { + t.Helper() + + got, err := s.DequeueJobs(context.Background(), opts) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + return got +} + +func jobNames(jobs []*job.Job) []string { + out := make([]string, 0, len(jobs)) + for _, j := range jobs { + out = append(out, j.Name) + } + + return out +} + +// wantExactly asserts the claimed set is exactly want, ignoring order. +func wantExactly(t *testing.T, got []*job.Job, want ...string) { + t.Helper() + + seen := make(map[string]int, len(got)) + for _, j := range got { + seen[j.Name]++ + } + + for _, w := range want { + switch n := seen[w]; { + case n == 0: + t.Errorf("job %q was not claimed; claimed set = %v, want %v", w, jobNames(got), want) + case n > 1: + t.Errorf("job %q claimed %d times; claimed set = %v", w, n, jobNames(got)) + } + + delete(seen, w) + } + + for extra := range seen { + t.Errorf("job %q was claimed but does not fit; claimed set = %v, want %v", + extra, jobNames(got), want) + } +} + +// wantOrder asserts the claimed jobs are exactly want, in that order. +func wantOrder(t *testing.T, got []*job.Job, want ...string) { + t.Helper() + + gotNames := jobNames(got) + if len(gotNames) != len(want) { + t.Fatalf("claimed %v, want %v", gotNames, want) + } + + for i := range want { + if gotNames[i] != want[i] { + t.Fatalf("claimed %v, want %v (differs at index %d)", gotNames, want, i) + } + } +} + +// wantStillClaimable proves a job the predicate rejected was left alone +// rather than claimed and put back. A backend that claims first and +// filters afterwards can pass a set assertion while still having burned a +// write on the job it rejected; here the job must come back on the very +// next unconstrained call. +func wantStillClaimable(t *testing.T, s job.Store, queue string, want ...string) { + t.Helper() + + got := mustDequeue(t, s, job.DequeueOpts{Queues: []string{queue}, Limit: 100}) + wantExactly(t, got, want...) +} + +// ────────────────────────────────────────────────── +// Cases +// ────────────────────────────────────────────────── + +// testZeroBudgetSelectsEverything is the backward-compatibility +// guarantee and the most important case in the suite. A caller that +// declares no budget must see exactly what the two-argument call used to +// return, including jobs whose requirements no worker could ever satisfy. +// Anything else silently strands work the moment this option ships. +func testZeroBudgetSelectsEverything(t *testing.T, s job.Store) { + const queue = "fit-zero-budget" + + undeclared := newFitJob("undeclared", queue, nil, withRunAtOffset(0)) + small := newFitJob("small", queue, resource.Set{resource.Memory: GiB}, withRunAtOffset(time.Minute)) + huge := newFitJob("huge", queue, resource.Set{ + resource.CPU: 64 * resource.MilliScale, + resource.Memory: 512 * GiB, + resource.Disk: 4096 * GiB, + resource.GPU: 8 * resource.MilliScale, + "fpga": 2, + }, withRunAtOffset(2*time.Minute)) + + mustEnqueue(t, s, undeclared, small, huge) + + opts := job.DequeueOpts{Queues: []string{queue}, Limit: 10} + if !opts.IsUnbounded() { + t.Fatalf("DequeueOpts%+v.IsUnbounded() = false, want true", opts) + } + + wantExactly(t, mustDequeue(t, s, opts), "undeclared", "small", "huge") +} + +func testMemoryBudgetFilters(t *testing.T, s job.Store) { + runDimensionCase(t, s, "fit-memory", resource.Memory, 4*GiB, 2*GiB, 8*GiB) +} + +func testCPUBudgetFilters(t *testing.T, s job.Store) { + runDimensionCase(t, s, "fit-cpu", resource.CPU, + 4*resource.MilliScale, 2*resource.MilliScale, 8*resource.MilliScale) +} + +func testDiskBudgetFilters(t *testing.T, s job.Store) { + runDimensionCase(t, s, "fit-disk", resource.Disk, 100*GiB, 10*GiB, 400*GiB) +} + +func testGPUBudgetFilters(t *testing.T, s job.Store) { + runDimensionCase(t, s, "fit-gpu", resource.GPU, + 2*resource.MilliScale, resource.MilliScale, 4*resource.MilliScale) +} + +// runDimensionCase proves one dimension filters on its own. Each +// dimension gets its own case because every backend stores the four as +// four separate indexed columns, and a copy-paste slip that compares +// req_memory_bytes twice and req_disk_bytes never would otherwise show up +// only in production. +func runDimensionCase(t *testing.T, s job.Store, queue, key string, budget, fitting, exceeding int64) { + t.Helper() + + fits := newFitJob("fits", queue, resource.Set{key: fitting}, withRunAtOffset(0)) + exceeds := newFitJob("exceeds", queue, resource.Set{key: exceeding}, withRunAtOffset(time.Minute)) + + mustEnqueue(t, s, fits, exceeds) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{key: budget}, + }) + + wantExactly(t, got, "fits") + wantStillClaimable(t, s, queue, "exceeds") +} + +// testAbsentBudgetKeyIsUnconstrained pins the inversion of +// resource.Set.Fits that the store-side predicate deliberately makes: an +// absent budget key means "not constrained", not "zero available". +// +// A worker that declares only memory must still claim GPU-requiring +// jobs. Otherwise adding a dimension to one worker's config would +// silently strand work on every worker whose config had not been updated +// yet, and the failure would look like a queue that stopped draining for +// no reason. +func testAbsentBudgetKeyIsUnconstrained(t *testing.T, s job.Store) { + const queue = "fit-absent-key" + + gpuHeavy := newFitJob("gpu-heavy", queue, resource.Set{ + resource.Memory: GiB, + resource.GPU: 8 * resource.MilliScale, + }, withRunAtOffset(0)) + + // The declared dimension must keep filtering. An implementation that + // read "absent key is unconstrained" as "any missing key disables the + // predicate" would claim this one too. + tooBig := newFitJob("too-big", queue, resource.Set{ + resource.Memory: 64 * GiB, + resource.GPU: resource.MilliScale, + }, withRunAtOffset(time.Minute)) + + mustEnqueue(t, s, gpuHeavy, tooBig) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 4 * GiB}, + }) + + wantExactly(t, got, "gpu-heavy") + wantStillClaimable(t, s, queue, "too-big") +} + +// testBoundedBudgetWithNoCustomKeysRejectsCustomRequirement is the half +// of the empty-offer rule a backend will get wrong. +// +// An empty CustomKeys list means two different things depending on the +// rest of the opts, and the difference is decided by IsUnbounded: +// +// - Unbounded opts — no budget, no keys, no hashes, no reservation — +// are a caller that does not use the resource model. It claims +// everything, custom resources included. ZeroBudgetSelectsEverything +// covers that half, and it is the backward-compatibility guarantee. +// - Opts bounded in any other way, here by a memory budget, are a +// resource-aware caller. An empty offer then means this worker has no +// custom resources at all, so a job requiring an fpga must NOT be +// claimed. +// +// A backend that reads an empty offer as "unconstrained" in both cases +// hands a resource-aware worker specialised work it cannot possibly run. +// A backend that reads it as "offers nothing" in both cases strands every +// custom-key job in the fleet the day this option ships. The gate is +// IsUnbounded, nothing else. +func testBoundedBudgetWithNoCustomKeysRejectsCustomRequirement(t *testing.T, s job.Store) { + const queue = "fit-bounded-no-custom" + + needsFPGA := newFitJob("needs-fpga", queue, resource.Set{ + resource.Memory: GiB, + "fpga": 1, + }, withRunAtOffset(0)) + plain := newFitJob("plain", queue, resource.Set{resource.Memory: GiB}, withRunAtOffset(time.Minute)) + + mustEnqueue(t, s, needsFPGA, plain) + + opts := job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 4 * GiB}, + } + + if opts.IsUnbounded() { + t.Fatal("opts carrying a memory budget report IsUnbounded() = true") + } + + wantExactly(t, mustDequeue(t, s, opts), "plain") + + // And the unbounded caller still gets it, so the rejection above was + // the offer being empty and bounded, not the job being unclaimable. + wantStillClaimable(t, s, queue, "needs-fpga") +} + +// testExplicitZeroBudgetKeyStillFilters is the other half of the absent +// key rule. A key present with the value zero is a real constraint — a +// worker with no free memory — and must exclude anything needing more +// than zero of it. A backend that decides "unbounded" by asking whether +// the budget is all zeros (resource.Set.IsZero) instead of whether the +// key is present would hand an exhausted worker the whole queue. +func testExplicitZeroBudgetKeyStillFilters(t *testing.T, s job.Store) { + const queue = "fit-explicit-zero" + + needsMemory := newFitJob("needs-memory", queue, resource.Set{resource.Memory: 1}, withRunAtOffset(0)) + needsNothing := newFitJob("needs-nothing", queue, nil, withRunAtOffset(time.Minute)) + + mustEnqueue(t, s, needsMemory, needsNothing) + + opts := job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 0}, + } + + if opts.IsUnbounded() { + t.Fatalf("DequeueOpts with an explicit zero memory budget reports IsUnbounded() = true") + } + + wantExactly(t, mustDequeue(t, s, opts), "needs-nothing") + wantStillClaimable(t, s, queue, "needs-memory") +} + +// testZeroRequirementAlwaysFits covers the job every deployment has most +// of: one that declares nothing. It must be claimable under any budget, +// including a budget of zero on every dimension. +// +// The two fixtures differ only in their write path, and that difference +// is not cosmetic. Mongo's insert path writes resource_requests as an +// explicit BSON null while its update path omits the key entirely, so a +// predicate that tests only one of the two shapes passes here on a +// freshly enqueued job and drops every job that has ever been updated — +// which, in production, is every job that was ever retried. +func testZeroRequirementAlwaysFits(t *testing.T, s job.Store) { + ctx := context.Background() + + const queue = "fit-zero-requirement" + + fresh := newFitJob("never-updated", queue, nil, withRunAtOffset(0)) + updated := newFitJob("updated-after-enqueue", queue, nil, withRunAtOffset(time.Minute)) + + mustEnqueue(t, s, fresh, updated) + + stored, err := s.GetJob(ctx, updated.ID) + if err != nil { + t.Fatalf("GetJob(%s): %v", updated.Name, err) + } + + if err = s.UpdateJob(ctx, stored); err != nil { + t.Fatalf("UpdateJob(%s): %v", updated.Name, err) + } + + // An exhausted worker: every dimension present, every one at zero. + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{ + resource.CPU: 0, + resource.Memory: 0, + resource.Disk: 0, + resource.GPU: 0, + }, + }) + + wantExactly(t, got, "never-updated", "updated-after-enqueue") +} + +// testExactFitIsClaimable pins the comparison as requirement <= budget, +// not <. A job needing exactly the free capacity must be claimable, or +// the last slot on every worker is silently unusable — a rounding error +// that costs a fixed fraction of the fleet forever. +func testExactFitIsClaimable(t *testing.T, s job.Store) { + const queue = "fit-exact" + + exact := newFitJob("exact", queue, resource.Set{ + resource.CPU: 2 * resource.MilliScale, + resource.Memory: 4 * GiB, + }, withRunAtOffset(0)) + + // One byte over the same budget. If this is claimed the comparison is + // the wrong way round; if "exact" is dropped the comparison is <. + overByOne := newFitJob("over-by-one", queue, resource.Set{ + resource.CPU: 2 * resource.MilliScale, + resource.Memory: 4*GiB + 1, + }, withRunAtOffset(time.Minute)) + + mustEnqueue(t, s, exact, overByOne) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{ + resource.CPU: 2 * resource.MilliScale, + resource.Memory: 4 * GiB, + }, + }) + + wantExactly(t, got, "exact") + wantStillClaimable(t, s, queue, "over-by-one") +} + +// testCustomKeyContainmentFilters proves a job needing a custom key the +// caller does not offer is not claimed, and that offering a key is what +// makes it claimable. Only key containment is tested at dequeue; the +// quantity is settled locally after the claim. +func testCustomKeyContainmentFilters(t *testing.T, s job.Store) { + const queue = "fit-custom-containment" + + plain := newFitJob("plain", queue, resource.Set{resource.Memory: GiB}, withRunAtOffset(0)) + needsTPU := newFitJob("needs-tpu", queue, resource.Set{ + resource.Memory: GiB, + "tpu": 1, + }, withRunAtOffset(time.Minute)) + needsFPGA := newFitJob("needs-fpga", queue, resource.Set{ + resource.Memory: GiB, + "fpga": 1, + }, withRunAtOffset(2*time.Minute)) + + mustEnqueue(t, s, plain, needsTPU, needsFPGA) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 8 * GiB}, + CustomKeys: []string{"tpu"}, + }) + + // The quantity is deliberately not compared: needs-tpu asks for one + // TPU and the caller offered the key without a count. + wantExactly(t, got, "plain", "needs-tpu") + wantStillClaimable(t, s, queue, "needs-fpga") +} + +// testCustomKeyPrefixDoesNotFalselyMatch is why +// resource.EncodeCustomKeys wraps its output in leading and trailing +// separators. A caller offering "fpga-large" must not claim a job needing +// "fpga", and a caller offering "fpga" must not claim a job needing +// "fpga-large". +// +// A backend that implements containment as a bare LIKE '%fpga%' — or as +// strings.Contains on the unwrapped list — passes every other custom-key +// case in this suite and fails this one. Both directions are checked in a +// single call, because a backend can get one right by accident. +func testCustomKeyPrefixDoesNotFalselyMatch(t *testing.T, s job.Store) { + const queue = "fit-custom-prefix" + + needsFPGA := newFitJob("needs-fpga", queue, resource.Set{"fpga": 1}, withRunAtOffset(0)) + needsFPGALarge := newFitJob("needs-fpga-large", queue, + resource.Set{"fpga-large": 1}, withRunAtOffset(time.Minute)) + + mustEnqueue(t, s, needsFPGA, needsFPGALarge) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + CustomKeys: []string{"fpga-large"}, + }) + + wantExactly(t, got, "needs-fpga-large") + wantStillClaimable(t, s, queue, "needs-fpga") +} + +// testCustomKeySubsetIsClaimable pins containment as a genuine subset +// test rather than a substring one. +// +// If this case is failing, LIKE is the thing that broke it, and reaching +// for another LIKE will not fix it. Both the job's required keys and the +// caller's offered keys are stored sorted, so +// +// :offered LIKE '%' || req_custom_keys || '%' +// +// passes every single-key case in this suite — including the prefix +// collision — and then silently drops a job needing {fpga, tpu} from a +// caller offering {fpga, nvme, tpu}: the job encodes to ",fpga,tpu,", +// which is not a substring of ",fpga,nvme,tpu,", because the interleaved +// key breaks the contiguous run. The job stranded is precisely the +// specialised one that is hardest to place anywhere else, and nothing in +// the system reports it. +// +// A portable exact formulation for SQL backends: strip each offered key +// from the stored list with nested REPLACE calls — one per offered key, +// built in Go since the offered set is a parameter — always replacing +// ",key," with ",", and require that what remains is "" or ",". With an +// offer of {fpga, nvme, tpu}: +// +// req_custom_keys = '' +// OR REPLACE(REPLACE(REPLACE(req_custom_keys, ',fpga,', ','), +// ',nvme,', ','), +// ',tpu,', ',') IN ('', ',') +// +// Postgres may prefer string_to_array(...) <@ ARRAY[...]; Mongo can use +// $expr with $setIsSubset over $split; Redis and memory should just call +// job.DequeueOpts.Allows. +func testCustomKeySubsetIsClaimable(t *testing.T, s job.Store) { + const ( + superset = "fit-custom-superset" + partial = "fit-custom-partial" + ) + + both := newFitJob("needs-fpga-and-tpu", superset, resource.Set{ + "fpga": 1, + "tpu": 1, + }, withRunAtOffset(0)) + + mustEnqueue(t, s, both) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{superset}, + Limit: 10, + CustomKeys: []string{"fpga", "nvme", "tpu"}, + }) + + wantExactly(t, got, "needs-fpga-and-tpu") + + // The other half: offering some of what a job needs is not enough. + half := newFitJob("needs-both-offered-one", partial, resource.Set{ + "fpga": 1, + "tpu": 1, + }, withRunAtOffset(0)) + + mustEnqueue(t, s, half) + + none := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{partial}, + Limit: 10, + CustomKeys: []string{"fpga"}, + }) + + wantExactly(t, none) + wantStillClaimable(t, s, partial, "needs-both-offered-one") +} + +// testPriorityOrderingPreservedWithinBudget proves the fit predicate did +// not cost the queue its ordering, and — through the oversized fixture — +// that the predicate runs inside the claim rather than over the rows the +// claim returned. +// +// The oversized job carries the highest priority and does not fit. With a +// limit of four, a backend that claims the top four rows and then drops +// the ones that do not fit returns three jobs and leaves the low-priority +// one behind, so the length assertion alone catches claim-then-filter. +func testPriorityOrderingPreservedWithinBudget(t *testing.T, s job.Store) { + const queue = "fit-priority-order" + + oversized := newFitJob("oversized", queue, resource.Set{resource.Memory: 64 * GiB}, + withPriority(100), withRunAtOffset(0)) + high := newFitJob("high", queue, resource.Set{resource.Memory: GiB}, + withPriority(9), withRunAtOffset(time.Minute)) + midEarly := newFitJob("mid-early", queue, resource.Set{resource.Memory: GiB}, + withPriority(5), withRunAtOffset(2*time.Minute)) + midLate := newFitJob("mid-late", queue, resource.Set{resource.Memory: GiB}, + withPriority(5), withRunAtOffset(3*time.Minute)) + low := newFitJob("low", queue, resource.Set{resource.Memory: GiB}, + withPriority(1), withRunAtOffset(4*time.Minute)) + + mustEnqueue(t, s, oversized, high, midEarly, midLate, low) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 4, + Budget: resource.Set{resource.Memory: 4 * GiB}, + }) + + wantOrder(t, got, "high", "mid-early", "mid-late", "low") + wantStillClaimable(t, s, queue, "oversized") +} + +// testLimitTruncatesAfterOrdering pins WHICH jobs the limit keeps, not +// just how many. +// +// Every other ordering case has as many eligible jobs as the limit +// allows, so all of them come back and any order bug shows up as a +// permutation. A backend that truncates BEFORE sorting — take an +// arbitrary N of the eligible set, then sort within it — returns the +// right count in the right relative order and passes all of them. Here +// six jobs are eligible and only two may be claimed, so the top two must +// be the two highest priorities and nothing else. +// +// Left unpinned, a worker with a small limit could be handed low-priority +// work indefinitely while high-priority jobs wait — the same starvation +// the locality ordering rule exists to prevent, arriving through a +// different door. The shapes most at risk are a Redis candidate scan that +// takes N before ordering, and a Mongo dequeueOne issuing N parallel +// FindOneAndUpdates whose individual sorts do not compose into a global +// one. +func testLimitTruncatesAfterOrdering(t *testing.T, s job.Store) { + const ( + queue = "fit-limit-order" + eligible = 6 + ) + + // Enqueued lowest priority first, so a backend that ignores ordering + // and takes the first two it finds is likely to return the wrong pair. + for i := range eligible { + mustEnqueue(t, s, newFitJob( + fmt.Sprintf("prio-%d", i), queue, + resource.Set{resource.Memory: GiB}, + withPriority(i), + withRunAtOffset(time.Duration(i)*time.Minute), + )) + } + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 2, + Budget: resource.Set{resource.Memory: 4 * GiB}, + }) + + wantOrder(t, got, "prio-5", "prio-4") + + // The four that lost are untouched, and the next call takes the next + // two by the same rule. + next := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 2, + Budget: resource.Set{resource.Memory: 4 * GiB}, + }) + + wantOrder(t, next, "prio-3", "prio-2") + wantStillClaimable(t, s, queue, "prio-1", "prio-0") +} + +// testPreferHashesAloneOrdersWithoutFiltering pins the split between the +// two questions a dequeue asks: "should I filter?" and "should I order?". +// +// PreferHashes only ever answers the second, so it is not a term of +// IsUnbounded. Opts carrying nothing but PreferHashes are unbounded: the +// backend skips the fit predicate entirely and still applies the locality +// term to its ORDER BY. If a backend derives "should I order?" from +// IsUnbounded, or folds PreferHashes into it, then these opts become +// bounded, the empty-CustomKeys rule fires, and the fpga job below +// vanishes — PreferHashes filtering transitively, which is exactly what +// "advisory, never a filter" forbids. +func testPreferHashesAloneOrdersWithoutFiltering(t *testing.T, s job.Store) { + const ( + queue = "fit-prefer-only" + local = "blake3:staged-here" + ) + + // Requires a custom resource the caller never offers, and is enqueued + // last, so it can only come back if nothing filtered it. + exotic := newFitJob("exotic", queue, resource.Set{ + resource.Memory: 512 * GiB, + "fpga": 4, + }, withRunAtOffset(2*time.Minute)) + plain := newFitJob("plain", queue, nil, withRunAtOffset(0)) + cached := newFitJob("cached", queue, nil, withRunAtOffset(time.Minute), withHash(local)) + + mustEnqueue(t, s, plain, cached, exotic) + + opts := job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + PreferHashes: []string{local}, + } + + if !opts.IsUnbounded() { + t.Fatal("opts carrying only PreferHashes report IsUnbounded() = false; " + + "locality is an ordering signal, not a constraint") + } + + // Ordering still applies — cached jumps ahead of the earlier plain — + // and nothing is filtered, including the job needing an unoffered fpga + // and more memory than any worker has. + wantOrder(t, mustDequeue(t, s, opts), "cached", "plain", "exotic") +} + +// testPreferHashesSortWithinPriorityBand covers the locality signal, and +// exists mostly to stop five ORDER BY clauses being written the obvious +// wrong way. +// +// "Preferred jobs sort first" does NOT mean +// +// ORDER BY (primary_input_hash = ANY($h)) DESC, priority DESC, run_at ASC +// +// It means +// +// ORDER BY priority DESC, (primary_input_hash = ANY($h)) DESC, run_at ASC +// +// Locality is an optimization; priority is user-expressed intent, and an +// optimization does not override intent. With locality above priority, a +// steady stream of low-priority jobs whose inputs are already cached +// beats a high-priority job whose inputs are cold — locality causing +// exactly the starvation this predicate exists to prevent. So a +// preferred job jumps its own priority band and no further. +// +// The other half: PreferHashes must never filter. All four jobs come +// back, including the two the caller has no local copy of, or every job +// with cold inputs would be stranded. +func testPreferHashesSortWithinPriorityBand(t *testing.T, s job.Store) { + const ( + queue = "fit-prefer-hashes" + local = "blake3:cached-locally" + ) + + urgent := newFitJob("urgent-remote", queue, resource.Set{resource.Memory: GiB}, + withPriority(5), withRunAtOffset(0), withHash("blake3:elsewhere")) + early := newFitJob("early-remote", queue, resource.Set{resource.Memory: GiB}, + withPriority(1), withRunAtOffset(time.Minute), withHash("blake3:also-elsewhere")) + mid := newFitJob("mid-unhashed", queue, resource.Set{resource.Memory: GiB}, + withPriority(1), withRunAtOffset(2*time.Minute)) + late := newFitJob("late-local", queue, resource.Set{resource.Memory: GiB}, + withPriority(1), withRunAtOffset(3*time.Minute), withHash(local)) + + mustEnqueue(t, s, urgent, early, mid, late) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 4 * GiB}, + PreferHashes: []string{local}, + }) + + // urgent outranks late-local on priority even though late-local is the + // one already staged; late-local then jumps its own priority band; the + // remaining two keep RunAt order. Nothing is filtered out. + wantOrder(t, got, "urgent-remote", "late-local", "early-remote", "mid-unhashed") +} + +// testLocalityDecidesWhichRowsSurviveATightLimit closes the gap neither +// LimitTruncatesAfterOrdering nor either PreferHashes case above can +// reach: WHICH rows a tight Limit keeps when only locality — not +// priority, not a static index, not a cached score — can tell them +// apart. +// +// LimitTruncatesAfterOrdering pins order-then-truncate using priority +// alone, and on two backends that is decorative: SQLite's dequeue index +// is keyed priority DESC, run_at ASC, so its natural scan order already +// matches the contract with no ORDER BY at all, and Redis's ZRange +// returns score order, which encodes priority the same way. Neither +// PreferHashes case above can catch a truncate-before-sort bug either, +// because both use a Limit of 10 against fewer eligible jobs than that — +// every eligible job comes back, so which ones a tight limit keeps is +// never exercised. +// +// PreferHashes closes that hole precisely because it cannot be baked +// into a static index or a precomputed score: it is supplied by the +// caller at call time, fresh on every call. A backend that truncates to +// Limit before applying the full ordering has no way to get this case +// right by accident, on any backend. +// +// All five jobs share one priority band, so priority cannot decide the +// winners; only Prefers can. The two preferred jobs carry the LATEST +// RunAt of the five, so a backend that orders by priority then RunAt +// alone — which is what "the index already matches" and "the score +// already matches" both reduce to — keeps the two EARLIEST non-preferred +// jobs instead. Limit is exactly the preferred count, so the winning set +// is decided entirely by locality, not merely reordered within it. +// +// The losers are deliberately not one hash shape: +// +// - noHash carries no hash at all, the common case, proving locality +// does not depend on every row having something to compare. +// - sortsAbovePreferred carries a hash that is lexicographically +// GREATER than the preferred hash. A backend that orders by the raw +// hash value instead of a preferred/not-preferred membership test — +// the wrong-way-round mistake job.DequeueOpts.Prefers exists to +// prevent — ranks it above a preferred job and is caught here too. +// - coldRemote carries an ordinary unmatched hash, so the case is not +// merely "empty vs. non-empty". +// +// A fixture set carrying only one non-empty hash value would risk a +// false pass: an empty/absent hash collates below strings on some +// engines, so a broken sort that happens to agree with the correct +// answer only because "no hash" always sorts last would pass here for +// the wrong reason. noHash and sortsAbovePreferred together rule that +// out — sortsAbovePreferred forces a raw-value sort to prefer a loser +// over a winner regardless of where empty hashes collate. +func testLocalityDecidesWhichRowsSurviveATightLimit(t *testing.T, s job.Store) { + const ( + queue = "fit-locality-limit" + preferred = "blake3:locally-cached" + ) + + // Earliest RunAt of the five: what a priority+RunAt-only sort would + // keep under Limit 2, and must NOT win here. + noHash := newFitJob("no-hash", queue, nil, withPriority(5), withRunAtOffset(0)) + sortsAbovePreferred := newFitJob("sorts-above-preferred", queue, nil, + withPriority(5), withRunAtOffset(time.Minute), withHash("zzz:never-staged")) + coldRemote := newFitJob("cold-remote", queue, nil, + withPriority(5), withRunAtOffset(2*time.Minute), withHash("blake3:elsewhere")) + + // Latest RunAt of the five: must win anyway, purely on locality. + preferredEarly := newFitJob("preferred-early", queue, nil, + withPriority(5), withRunAtOffset(3*time.Minute), withHash(preferred)) + preferredLate := newFitJob("preferred-late", queue, nil, + withPriority(5), withRunAtOffset(4*time.Minute), withHash(preferred)) + + mustEnqueue(t, s, noHash, sortsAbovePreferred, coldRemote, preferredEarly, preferredLate) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 2, + PreferHashes: []string{preferred}, + }) + + wantOrder(t, got, "preferred-early", "preferred-late") + wantStillClaimable(t, s, queue, "no-hash", "sorts-above-preferred", "cold-remote") +} + +// testReservedForRestrictsToOneJob proves a targeted claim returns that +// job and nothing else — and that it is still subject to the budget. A +// reservation that could bypass the fit test would reintroduce exactly +// the overcommit the predicate exists to prevent, through the one code +// path a scheduler is most likely to use for a large job. +func testReservedForRestrictsToOneJob(t *testing.T, s job.Store) { + const queue = "fit-reserved" + + first := newFitJob("first", queue, resource.Set{resource.Memory: GiB}, withRunAtOffset(0)) + target := newFitJob("target", queue, resource.Set{resource.Memory: GiB}, withRunAtOffset(time.Minute)) + third := newFitJob("third", queue, resource.Set{resource.Memory: GiB}, withRunAtOffset(2*time.Minute)) + + mustEnqueue(t, s, first, target, third) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 4 * GiB}, + ReservedFor: &target.ID, + }) + + wantExactly(t, got, "target") + + // Reserved, but the caller has no room for it. + none := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 0}, + ReservedFor: &third.ID, + }) + + wantExactly(t, none) + wantStillClaimable(t, s, queue, "first", "third") +} + +// testNonPositiveLimitClaimsNothing pins the controller ruling on Limit +// <= 0: it claims NOTHING, not "unlimited". +// +// A worker computes its Limit from free capacity, so a Limit of zero means +// it has zero free slots right now. Reading zero as unlimited would hand +// that exhausted worker the entire queue instead of the empty result its +// own accounting asked for — a worker with no room claiming everything is +// strictly worse than a worker that briefly polls too conservatively. +// Postgres and SQLite already emit `LIMIT 0` and claim nothing; every +// other backend must agree, including the negative case, which a caller +// should never send but a store must still handle safely rather than +// looping or claiming without bound. +func testNonPositiveLimitClaimsNothing(t *testing.T, s job.Store) { + const queue = "fit-non-positive-limit" + + only := newFitJob("only", queue, nil, withRunAtOffset(0)) + + mustEnqueue(t, s, only) + + zero := mustDequeue(t, s, job.DequeueOpts{Queues: []string{queue}, Limit: 0}) + wantExactly(t, zero) + + negative := mustDequeue(t, s, job.DequeueOpts{Queues: []string{queue}, Limit: -1}) + wantExactly(t, negative) + + // The job was never claimed by either call above. + wantStillClaimable(t, s, queue, "only") +} + +// testClaimIsAtomicUnderConcurrency proves the predicate did not cost the +// claim its atomicity — the one property the whole store contract rests +// on, since a job handed to two workers is run twice. +// +// The assertion is an invariant, not a timing guess: whichever way the +// goroutines interleave, every job must end up claimed exactly once. A +// correct backend can never violate that, so a correct backend can never +// flake here. A backend that selects candidates and then updates them in +// a separate statement violates it whenever two scans overlap, which is +// what makes the case worth running. +func testClaimIsAtomicUnderConcurrency(t *testing.T, s job.Store) { + const ( + queue = "fit-concurrent" + jobCount = 20 + claimers = 4 + batch = 2 + ) + + mine := make(map[id.JobID]string, jobCount) + + for i := range jobCount { + j := newFitJob(fmt.Sprintf("concurrent-%d", i), queue, + resource.Set{resource.Memory: GiB}, withRunAtOffset(time.Duration(i)*time.Second)) + + mustEnqueue(t, s, j) + + mine[j.ID] = j.Name + } + + var ( + mu sync.Mutex + claims = make(map[id.JobID]int, jobCount) + wg sync.WaitGroup + ) + + errCh := make(chan error, claimers) + + for range claimers { + wg.Add(1) + + go func() { + defer wg.Done() + + // Each claimer drains the queue in small batches rather than + // asking for everything once. A single large claim per goroutine + // lets four claimers serialize — the first takes all 20 and the + // rest find an empty queue, so the interleaving that exposes a + // non-atomic claim never happens. Looping keeps every claimer + // contending until the queue is actually empty, which is what + // makes this a reliable detector rather than a probabilistic one. + for round := 0; round <= jobCount; round++ { + got, err := s.DequeueJobs(context.Background(), job.DequeueOpts{ + Queues: []string{queue}, + Limit: batch, + Budget: resource.Set{resource.Memory: 4 * GiB}, + }) + if err != nil { + errCh <- err + + return + } + + if len(got) == 0 { + return + } + + mu.Lock() + + for _, j := range got { + claims[j.ID]++ + } + + mu.Unlock() + } + + // Only reachable if the backend keeps handing out jobs after the + // queue should be empty, which the per-job assertions below will + // also catch. Bounding the loop keeps that a failure, not a hang. + errCh <- fmt.Errorf("claimer still draining %q after %d rounds", queue, jobCount) + }() + } + + wg.Wait() + close(errCh) + + for err := range errCh { + t.Fatalf("concurrent DequeueJobs: %v", err) + } + + for jobID, name := range mine { + switch n := claims[jobID]; { + case n == 0: + t.Errorf("job %s was never claimed", name) + case n > 1: + t.Errorf("job %s claimed %d times, want exactly 1 — the claim is not atomic", name, n) + } + } +} diff --git a/store/storetest/dequeue_test.go b/store/storetest/dequeue_test.go new file mode 100644 index 0000000..a925419 --- /dev/null +++ b/store/storetest/dequeue_test.go @@ -0,0 +1,205 @@ +package storetest_test + +import ( + "context" + "sort" + "sync" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/storetest" +) + +// TestSuiteAgainstReference runs the conformance suite against the +// reference store below. +// +// The suite ships before any backend implements the widened signature, so +// without this it would be an untested specification. The reference is +// deliberately the most literal implementation of the contract there is — +// it filters with job.DequeueOpts.Allows and orders with +// job.DequeueOpts.Less — which makes this test a check that the suite's +// cases are mutually consistent and that the exported predicate helpers +// actually satisfy them. +func TestSuiteAgainstReference(t *testing.T) { + storetest.RunDequeueSuite(t, func(_ *testing.T) job.Store { + return newReferenceStore() + }) +} + +// referenceStore is a minimal in-process job.Store. It exists only to +// exercise the suite; the real backends live under store/. +type referenceStore struct { + mu sync.Mutex + jobs map[id.JobID]*job.Job +} + +func newReferenceStore() *referenceStore { + return &referenceStore{jobs: make(map[id.JobID]*job.Job)} +} + +func cloneRefJob(j *job.Job) *job.Job { + out := *j + out.Resources = j.Resources.Clone() + out.ResourceLimits = j.ResourceLimits.Clone() + + return &out +} + +func (r *referenceStore) EnqueueJob(_ context.Context, j *job.Job) error { + r.mu.Lock() + defer r.mu.Unlock() + + if _, exists := r.jobs[j.ID]; exists { + return dispatch.ErrJobAlreadyExists + } + + r.jobs[j.ID] = cloneRefJob(j) + + return nil +} + +// DequeueJobs selects, orders, limits, and only then claims — the order +// the contract requires. +// +// A non-positive Limit claims nothing: it counts eligible jobs only, and +// zero of them may ever be claimed, matching the SQL backends' `LIMIT 0` +// rather than reading zero as unlimited. +func (r *referenceStore) DequeueJobs(_ context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + r.mu.Lock() + defer r.mu.Unlock() + + if opts.Limit <= 0 { + return nil, nil + } + + queues := make(map[string]struct{}, len(opts.Queues)) + for _, q := range opts.Queues { + queues[q] = struct{}{} + } + + now := time.Now().UTC() + + candidates := make([]*job.Job, 0, len(r.jobs)) + + for _, j := range r.jobs { + if j.State != job.StatePending && j.State != job.StateRetrying { + continue + } + + if !j.RunAt.IsZero() && j.RunAt.After(now) { + continue + } + + if len(queues) > 0 { + if _, ok := queues[j.Queue]; !ok { + continue + } + } + + if !opts.Allows(j) { + continue + } + + candidates = append(candidates, j) + } + + sort.SliceStable(candidates, func(a, b int) bool { + return opts.Less(candidates[a], candidates[b]) + }) + + if len(candidates) > opts.Limit { + candidates = candidates[:opts.Limit] + } + + claimed := make([]*job.Job, 0, len(candidates)) + + for _, j := range candidates { + started := now + j.State = job.StateRunning + j.StartedAt = &started + + claimed = append(claimed, cloneRefJob(j)) + } + + return claimed, nil +} + +func (r *referenceStore) GetJob(_ context.Context, jobID id.JobID) (*job.Job, error) { + r.mu.Lock() + defer r.mu.Unlock() + + j, ok := r.jobs[jobID] + if !ok { + return nil, dispatch.ErrJobNotFound + } + + return cloneRefJob(j), nil +} + +func (r *referenceStore) UpdateJob(_ context.Context, j *job.Job) error { + r.mu.Lock() + defer r.mu.Unlock() + + if _, ok := r.jobs[j.ID]; !ok { + return dispatch.ErrJobNotFound + } + + r.jobs[j.ID] = cloneRefJob(j) + + return nil +} + +func (r *referenceStore) DeleteJob(_ context.Context, jobID id.JobID) error { + r.mu.Lock() + defer r.mu.Unlock() + + delete(r.jobs, jobID) + + return nil +} + +func (r *referenceStore) ListJobsByState( + _ context.Context, state job.State, _ job.ListOpts, +) ([]*job.Job, error) { + r.mu.Lock() + defer r.mu.Unlock() + + out := make([]*job.Job, 0, len(r.jobs)) + + for _, j := range r.jobs { + if j.State == state { + out = append(out, cloneRefJob(j)) + } + } + + return out, nil +} + +func (r *referenceStore) HeartbeatJob(_ context.Context, jobID id.JobID, _ id.WorkerID) error { + r.mu.Lock() + defer r.mu.Unlock() + + j, ok := r.jobs[jobID] + if !ok { + return dispatch.ErrJobNotFound + } + + beat := time.Now().UTC() + j.HeartbeatAt = &beat + + return nil +} + +func (r *referenceStore) ReapStaleJobs(_ context.Context, _ time.Duration) ([]*job.Job, error) { + return nil, nil +} + +func (r *referenceStore) CountJobs(_ context.Context, _ job.CountOpts) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + + return int64(len(r.jobs)), nil +} diff --git a/store/storetest/dlq.go b/store/storetest/dlq.go new file mode 100644 index 0000000..5c4d92d --- /dev/null +++ b/store/storetest/dlq.go @@ -0,0 +1,202 @@ +package storetest + +import ( + "bytes" + "context" + "testing" + "time" + + "github.com/xraph/dispatch/dlq" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/resource" +) + +// DLQStore is the capability this suite exercises. +type DLQStore interface { + dlq.Store +} + +// RunDLQSuite asserts the DLQ contract every backend has to satisfy. +// +// newStore may return a shared store, so every case works on entries it +// created itself and never asserts on the total count. +func RunDLQSuite(t *testing.T, newStore func(t *testing.T) DLQStore) { + t.Helper() + + t.Run("PushDLQPreservesExecutionFields", func(t *testing.T) { + testDLQPreservesExecutionFields(t, newStore(t)) + }) + + t.Run("PushDLQPreservesAbsentResourceSets", func(t *testing.T) { + testDLQPreservesAbsentResourceSets(t, newStore(t)) + }) +} + +// testDLQPreservesExecutionFields is the round trip that stops a backend +// from silently dropping a column. +// +// dlq.Replay rebuilds a job from the stored entry and enqueues it +// directly, without going back through the engine, so nothing re-derives +// any of these values for it. A backend that fails to persist one does not +// report an error: the entry reads back with a zero in that field and the +// replayed job quietly runs with a default instead. For LeaseTTL that +// default is short enough to make a long job unrunnable, since its lease +// lapses mid-run and reclamation restarts it forever. +// +// Every value below is deliberately non-zero and distinct, so a mapper +// that drops a field, or crosses two of them, cannot produce a passing +// result by accident. +func testDLQPreservesExecutionFields(t *testing.T, s DLQStore) { + t.Helper() + + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Millisecond) + + want := &dlq.Entry{ + ID: id.NewDLQID(), + JobID: id.NewJobID(), + JobName: "long-render", + Queue: "dlq-fidelity", + Payload: []byte(`{"frame":42}`), + Error: "handler exploded", + RetryCount: 3, + MaxRetries: 3, + ScopeAppID: "app_1", + ScopeOrgID: "org_1", + FailedAt: now, + CreatedAt: now, + + Priority: 7, + Timeout: 90 * time.Minute, + // Six hours: the case the per-job TTL exists for, and far enough + // from any pool default that a dropped value is unmistakable. + LeaseTTL: 6 * time.Hour, + ArtifactBindings: []byte(`{"input":"artifact_abc"}`), + Resources: resource.Set{"cpu_milli": 2000, "memory_bytes": 1 << 30}, + ResourceLimits: resource.Set{"cpu_milli": 4000, "memory_bytes": 2 << 30}, + ResourceClass: "gpu-large", + InputBytes: 4096, + PrimaryInputHash: "sha256:deadbeef", + } + + if err := s.PushDLQ(ctx, want); err != nil { + t.Fatalf("PushDLQ: %v", err) + } + + got, err := s.GetDLQ(ctx, want.ID) + if err != nil { + t.Fatalf("GetDLQ: %v", err) + } + + if got.Priority != want.Priority { + t.Errorf("Priority = %d, want %d", got.Priority, want.Priority) + } + if got.Timeout != want.Timeout { + t.Errorf("Timeout = %v, want %v", got.Timeout, want.Timeout) + } + if got.LeaseTTL != want.LeaseTTL { + t.Errorf("LeaseTTL = %v, want %v: a replayed job would fall back to the "+ + "pool default and be reclaimed mid-run forever", got.LeaseTTL, want.LeaseTTL) + } + if !bytes.Equal(got.ArtifactBindings, want.ArtifactBindings) { + t.Errorf("ArtifactBindings = %q, want %q", got.ArtifactBindings, want.ArtifactBindings) + } + if !resourceSetEqual(got.Resources, want.Resources) { + t.Errorf("Resources = %v, want %v", got.Resources, want.Resources) + } + if !resourceSetEqual(got.ResourceLimits, want.ResourceLimits) { + t.Errorf("ResourceLimits = %v, want %v", got.ResourceLimits, want.ResourceLimits) + } + if got.ResourceClass != want.ResourceClass { + t.Errorf("ResourceClass = %q, want %q", got.ResourceClass, want.ResourceClass) + } + if got.InputBytes != want.InputBytes { + t.Errorf("InputBytes = %d, want %d", got.InputBytes, want.InputBytes) + } + if got.PrimaryInputHash != want.PrimaryInputHash { + t.Errorf("PrimaryInputHash = %q, want %q", got.PrimaryInputHash, want.PrimaryInputHash) + } + + // ListDLQ decodes through the same mapper but a different query, and + // on at least one backend that is a genuinely separate code path. + listed, err := s.ListDLQ(ctx, dlq.ListOpts{Queue: want.Queue, Limit: 50}) + if err != nil { + t.Fatalf("ListDLQ: %v", err) + } + + var found *dlq.Entry + for _, e := range listed { + if e.ID == want.ID { + found = e + + break + } + } + if found == nil { + t.Fatalf("ListDLQ did not return the pushed entry") + } + if found.LeaseTTL != want.LeaseTTL { + t.Errorf("ListDLQ LeaseTTL = %v, want %v", found.LeaseTTL, want.LeaseTTL) + } +} + +// testDLQPreservesAbsentResourceSets pins that a job which declared no +// resources reads back with none, rather than with an empty set. +// +// The distinction is not cosmetic on the SQL backends: resource. +// EncodeSetString writes NULL for a zero Set specifically so the two stay +// distinguishable, and a mapper that turns absent into empty would make +// every replayed job look like it had declared an explicit empty +// requirement. +func testDLQPreservesAbsentResourceSets(t *testing.T, s DLQStore) { + t.Helper() + + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Millisecond) + + bare := &dlq.Entry{ + ID: id.NewDLQID(), + JobID: id.NewJobID(), + JobName: "plain", + Queue: "dlq-fidelity-bare", + Payload: []byte(`{}`), + Error: "boom", + MaxRetries: 3, + FailedAt: now, + CreatedAt: now, + } + + if err := s.PushDLQ(ctx, bare); err != nil { + t.Fatalf("PushDLQ: %v", err) + } + + got, err := s.GetDLQ(ctx, bare.ID) + if err != nil { + t.Fatalf("GetDLQ: %v", err) + } + + if len(got.Resources) != 0 { + t.Errorf("Resources = %v, want none", got.Resources) + } + if len(got.ResourceLimits) != 0 { + t.Errorf("ResourceLimits = %v, want none", got.ResourceLimits) + } + if got.LeaseTTL != 0 { + t.Errorf("LeaseTTL = %v, want 0", got.LeaseTTL) + } +} + +// resourceSetEqual compares two sets by content, treating nil and empty as +// the same thing. +func resourceSetEqual(a, b resource.Set) bool { + if len(a) != len(b) { + return false + } + for k, v := range a { + if b[k] != v { + return false + } + } + + return true +} diff --git a/store/storetest/lease.go b/store/storetest/lease.go new file mode 100644 index 0000000..a2e923c --- /dev/null +++ b/store/storetest/lease.go @@ -0,0 +1,1119 @@ +package storetest + +import ( + "context" + "errors" + "fmt" + "reflect" + "sync" + "testing" + "time" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// RunLeaseSuite runs the lease conformance suite against a backend. +// +// newStore is called once per subtest. It may return the same underlying +// store every time: each case enqueues onto its own queue and asserts on +// the jobs it created, so cases do not interfere. That matters because +// starting a fresh Postgres or Redis container per subtest would dominate +// the runtime of the whole suite. +// +// The lease is GRANTED by job.Store.DequeueJobs, through +// job.DequeueOpts.WorkerID and LeaseUntil — there is no separate leased +// dequeue. That is why every case below claims with DequeueJobs: a +// backend cannot pass this suite with a grant path that skips the fit +// predicate, because there is only one path. +func RunLeaseSuite(t *testing.T, newStore func(t *testing.T) LeaseStore) { + t.Helper() + + t.Run("DequeueGrantsLeaseAndBumpsEpoch", func(t *testing.T) { + testDequeueGrantsLeaseAndBumpsEpoch(t, newStore(t)) + }) + t.Run("DequeueGrantsNoLeaseWhenLeaseUntilZero", func(t *testing.T) { + testDequeueGrantsNoLeaseWhenLeaseUntilZero(t, newStore(t)) + }) + t.Run("DequeueRejectsLeaseWithoutWorker", func(t *testing.T) { + testDequeueRejectsLeaseWithoutWorker(t, newStore(t)) + }) + t.Run("DequeueComposesBudgetAndLeaseGrant", func(t *testing.T) { + testDequeueComposesBudgetAndLeaseGrant(t, newStore(t)) + }) + t.Run("RenewLeaseExtends", func(t *testing.T) { + testRenewLeaseExtends(t, newStore(t)) + }) + t.Run("RenewLeaseRejectsStaleEpoch", func(t *testing.T) { + testRenewLeaseRejectsStaleEpoch(t, newStore(t)) + }) + t.Run("RenewLeaseRejectsWrongWorker", func(t *testing.T) { + testRenewLeaseRejectsWrongWorker(t, newStore(t)) + }) + t.Run("RenewLeaseRejectsMissingJob", func(t *testing.T) { + testRenewLeaseRejectsMissingJob(t, newStore(t)) + }) + t.Run("ReclaimExpiredLeases", func(t *testing.T) { + testReclaimExpiredLeases(t, newStore(t)) + }) + t.Run("ReclaimSkipsLiveLease", func(t *testing.T) { + testReclaimSkipsLiveLease(t, newStore(t)) + }) + t.Run("ReclaimFencesPreviousHolder", func(t *testing.T) { + testReclaimFencesPreviousHolder(t, newStore(t)) + }) + t.Run("ReclaimPreservesRetryCount", func(t *testing.T) { + testReclaimPreservesRetryCount(t, newStore(t)) + }) + t.Run("ReclaimIsExclusive", func(t *testing.T) { + testReclaimIsExclusive(t, newStore(t)) + }) + t.Run("ReclaimIsExclusiveUnderConcurrency", func(t *testing.T) { + testReclaimIsExclusiveUnderConcurrency(t, newStore(t)) + }) + t.Run("LeaseTTLRoundTrips", func(t *testing.T) { + testLeaseTTLRoundTrips(t, newStore(t)) + }) + t.Run("UpdateLeasedJobAppliesAtHeldEpoch", func(t *testing.T) { + testUpdateLeasedJobAppliesAtHeldEpoch(t, newStore(t)) + }) + t.Run("UpdateLeasedJobRejectsStaleEpoch", func(t *testing.T) { + testUpdateLeasedJobRejectsStaleEpoch(t, newStore(t)) + }) + t.Run("UpdateLeasedJobRejectsWrongWorker", func(t *testing.T) { + testUpdateLeasedJobRejectsWrongWorker(t, newStore(t)) + }) + t.Run("UpdateLeasedJobRejectsWhenNoLongerRunning", func(t *testing.T) { + testUpdateLeasedJobRejectsWhenNoLongerRunning(t, newStore(t)) + }) + t.Run("UpdateLeasedJobPreservesLeaseColumnsAgainstAStaleSnapshot", func(t *testing.T) { + testUpdateLeasedJobPreservesLeaseColumnsAgainstAStaleSnapshot(t, newStore(t)) + }) + t.Run("UpdateLeasedJobRunnableWriteIsDequeueable", func(t *testing.T) { + testUpdateLeasedJobRunnableWriteIsDequeueable(t, newStore(t)) + }) + t.Run("ReclaimNonPositiveLimitReclaimsNothing", func(t *testing.T) { + testReclaimNonPositiveLimitReclaimsNothing(t, newStore(t)) + }) +} + +func testDequeueGrantsLeaseAndBumpsEpoch(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + until := time.Now().UTC().Add(time.Minute) + const queue = "lease-grant" + + j := PendingJob("grant", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + WorkerID: worker, + LeaseUntil: until, + }) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + if len(got) != 1 { + t.Fatalf("DequeueJobs returned %d jobs, want 1", len(got)) + } + + d := got[0] + if d.State != job.StateRunning { + t.Errorf("State = %s, want %s", d.State, job.StateRunning) + } + if d.LeaseEpoch != 1 { + t.Errorf("LeaseEpoch = %d, want 1", d.LeaseEpoch) + } + if d.WorkerID != worker { + t.Errorf("WorkerID = %s, want %s", d.WorkerID, worker) + } + if d.LeaseExpiresAt == nil { + t.Fatal("LeaseExpiresAt = nil, want the granted expiry") + } + if diff := d.LeaseExpiresAt.Sub(until); diff > time.Second || diff < -time.Second { + t.Errorf("LeaseExpiresAt = %v, want within 1s of %v", d.LeaseExpiresAt, until) + } + if d.StartedAt == nil { + t.Error("StartedAt = nil, want it set at dequeue") + } + + // The grant must have been PERSISTED by the claim, not merely decorated + // onto the returned copy. A backend that granted as a follow-up write + // would still pass every assertion above; this is the one that fails if + // the row itself is running with no lease — the state no reclaimer can + // see, because every backend ignores a null expiry, and therefore the + // state nothing recovers from. See job.DequeueOpts.LeaseUntil. + stored, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if stored.LeaseEpoch != d.LeaseEpoch { + t.Errorf("stored LeaseEpoch = %d, want %d", stored.LeaseEpoch, d.LeaseEpoch) + } + if stored.WorkerID != worker { + t.Errorf("stored WorkerID = %s, want %s", stored.WorkerID, worker) + } + if stored.LeaseExpiresAt == nil { + t.Error("stored LeaseExpiresAt = nil, want the granted expiry") + } +} + +// testDequeueGrantsNoLeaseWhenLeaseUntilZero pins the backward- +// compatibility guarantee: an opt-out caller must see exactly the +// behaviour DequeueJobs had before leases existed. +// +// Without this case a backend that granted unconditionally — writing the +// zero worker and bumping the epoch on every claim — passes every other +// case in this file, and every caller that never asked for a lease would +// have its jobs reclaimed out from under it on the next sweep. +func testDequeueGrantsNoLeaseWhenLeaseUntilZero(t *testing.T, s LeaseStore) { + ctx := context.Background() + const queue = "lease-not-granted" + + j := PendingJob("not-granted", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + // Read the baseline back from the store rather than trusting the + // in-memory job: a backend that defaults lease_epoch differently at + // insert would otherwise look like it bumped. + before, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get before claim: %v", err) + } + + // No WorkerID and no LeaseUntil — the opts a pool sends today. + got, err := s.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{queue}, Limit: 1}) + if err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + if len(got) != 1 { + t.Fatalf("DequeueJobs returned %d jobs, want 1", len(got)) + } + + d := got[0] + if d.State != job.StateRunning { + t.Errorf("State = %s, want %s", d.State, job.StateRunning) + } + if d.LeaseEpoch != before.LeaseEpoch { + t.Errorf("LeaseEpoch = %d, want it unchanged at %d", d.LeaseEpoch, before.LeaseEpoch) + } + if d.LeaseExpiresAt != nil { + t.Errorf("LeaseExpiresAt = %v, want nil — no lease was asked for", d.LeaseExpiresAt) + } + if !d.WorkerID.IsNil() { + t.Errorf("WorkerID = %s, want it unset — no lease was asked for", d.WorkerID) + } + + stored, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get after claim: %v", err) + } + if stored.LeaseEpoch != before.LeaseEpoch { + t.Errorf("stored LeaseEpoch = %d, want it unchanged at %d", + stored.LeaseEpoch, before.LeaseEpoch) + } + if stored.LeaseExpiresAt != nil { + t.Errorf("stored LeaseExpiresAt = %v, want nil", stored.LeaseExpiresAt) + } + if !stored.WorkerID.IsNil() { + t.Errorf("stored WorkerID = %s, want it unset", stored.WorkerID) + } +} + +// testDequeueRejectsLeaseWithoutWorker covers the one incoherent request: +// a grant with no holder. +// +// RenewLease matches on worker ID, so a lease held by the zero worker can +// never be renewed — the job would be claimed, expire, be reclaimed, and +// go round again forever. That presents as a queue that never drains +// rather than as an error, so the store refuses the claim instead. +func testDequeueRejectsLeaseWithoutWorker(t *testing.T, s LeaseStore) { + ctx := context.Background() + const queue = "lease-no-worker" + + j := PendingJob("no-worker", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + // WorkerID deliberately unset. + LeaseUntil: time.Now().UTC().Add(time.Minute), + }) + if !errors.Is(err, job.ErrLeaseWithoutWorker) { + t.Fatalf("DequeueJobs with LeaseUntil and no WorkerID = %v, want %v", + err, job.ErrLeaseWithoutWorker) + } + if len(got) != 0 { + t.Errorf("DequeueJobs returned %d jobs, want 0 — a refused claim must claim nothing", + len(got)) + } + + // The refusal must come before any write: the job is still there for a + // correctly configured worker to take. + after, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if after.State != job.StatePending { + t.Errorf("State = %s, want %s — the refused claim wrote to the job", + after.State, job.StatePending) + } +} + +// testDequeueComposesBudgetAndLeaseGrant proves the two halves of +// DequeueOpts a production caller always sends together actually work +// together. worker/pool.go's fetchLoop sends Budget, CustomKeys, +// WorkerID, and LeaseUntil in ONE DequeueOpts on every call — there is no +// path that grants a lease without also carrying whatever budget the +// worker has — but RunDequeueSuite (job/store.go's resource-aware fit +// predicate) and the rest of this suite (the lease grant) had exercised +// those two contracts only apart. A backend could satisfy both suites +// while still getting the composed shape wrong, for example by granting +// the lease before the fit predicate ran, or by having the two features +// implemented against different code paths that silently disagree once +// both parameters are non-zero. +func testDequeueComposesBudgetAndLeaseGrant(t *testing.T, s LeaseStore) { + const queue = "lease-and-budget" + + fits := newFitJob("fits", queue, resource.Set{resource.Memory: GiB}, withRunAtOffset(0)) + exceeds := newFitJob("exceeds", queue, + resource.Set{resource.Memory: 8 * GiB}, withRunAtOffset(time.Minute)) + + mustEnqueue(t, s, fits, exceeds) + + worker := id.NewWorkerID() + until := time.Now().UTC().Add(time.Minute) + + got := mustDequeue(t, s, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 10, + Budget: resource.Set{resource.Memory: 4 * GiB}, + WorkerID: worker, + LeaseUntil: until, + }) + + // The fit predicate ran: only "fits" was claimed. + wantExactly(t, got, "fits") + + d := got[0] + if d.State != job.StateRunning { + t.Errorf("State = %s, want %s", d.State, job.StateRunning) + } + if d.WorkerID != worker { + t.Errorf("WorkerID = %s, want %s", d.WorkerID, worker) + } + if d.LeaseEpoch != 1 { + t.Errorf("LeaseEpoch = %d, want 1", d.LeaseEpoch) + } + if d.LeaseExpiresAt == nil { + t.Fatal("LeaseExpiresAt = nil, want the granted expiry") + } + if diff := d.LeaseExpiresAt.Sub(until); diff > time.Second || diff < -time.Second { + t.Errorf("LeaseExpiresAt = %v, want within 1s of %v", d.LeaseExpiresAt, until) + } + + // The lease grant did not bypass the fit predicate for the oversized + // job: it stays fully pending, not merely unleased. + wantStillClaimable(t, s, queue, "exceeds") +} + +func testRenewLeaseExtends(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + now := time.Now().UTC() + const queue = "lease-renew" + + j := PendingJob("renew", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(30 * time.Second), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + + extended := now.Add(10 * time.Minute) + if renewErr := s.RenewLease(ctx, got[0].ID, worker, got[0].LeaseEpoch, extended); renewErr != nil { + t.Fatalf("RenewLease: %v", renewErr) + } + + after, err := s.GetJob(ctx, got[0].ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if after.LeaseExpiresAt == nil { + t.Fatal("LeaseExpiresAt = nil after renewal") + } + if diff := after.LeaseExpiresAt.Sub(extended); diff > time.Second || diff < -time.Second { + t.Errorf("LeaseExpiresAt = %v, want within 1s of %v", after.LeaseExpiresAt, extended) + } + // Renewal must not bump the epoch — only grant and reclaim do. If it + // did, the holder's own next renewal would fence itself. + if after.LeaseEpoch != got[0].LeaseEpoch { + t.Errorf("LeaseEpoch = %d after renewal, want it unchanged at %d", + after.LeaseEpoch, got[0].LeaseEpoch) + } +} + +func testRenewLeaseRejectsStaleEpoch(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + now := time.Now().UTC() + + const queue = "lease-stale-epoch" + + j := PendingJob("stale-epoch", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + + err = s.RenewLease(ctx, got[0].ID, worker, got[0].LeaseEpoch-1, now.Add(time.Hour)) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("RenewLease with stale epoch = %v, want %v", err, job.ErrLeaseLost) + } +} + +func testRenewLeaseRejectsWrongWorker(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + other := id.NewWorkerID() + now := time.Now().UTC() + + const queue = "lease-wrong-worker" + + j := PendingJob("wrong-worker", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + + err = s.RenewLease(ctx, got[0].ID, other, got[0].LeaseEpoch, now.Add(time.Hour)) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("RenewLease from another worker = %v, want %v", err, job.ErrLeaseLost) + } +} + +func testRenewLeaseRejectsMissingJob(t *testing.T, s LeaseStore) { + ctx := context.Background() + + err := s.RenewLease(ctx, id.NewJobID(), id.NewWorkerID(), 1, time.Now().UTC().Add(time.Hour)) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("RenewLease on a missing job = %v, want %v", err, job.ErrLeaseLost) + } +} + +func testReclaimExpiredLeases(t *testing.T, s LeaseStore) { + ctx := context.Background() + + j := RunningJob("expired", "lease-reclaim", 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + got, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + if !Contains(got, j.ID) { + t.Fatalf("reclaimed set does not contain %s", j.ID) + } + + after, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if after.State != job.StatePending { + t.Errorf("State = %s, want %s", after.State, job.StatePending) + } + if after.EvictCount != j.EvictCount+1 { + t.Errorf("EvictCount = %d, want %d", after.EvictCount, j.EvictCount+1) + } + if !after.WorkerID.IsNil() { + t.Errorf("WorkerID = %s, want it cleared", after.WorkerID) + } + if after.StartedAt != nil { + t.Errorf("StartedAt = %v, want nil", after.StartedAt) + } + if after.HeartbeatAt != nil { + t.Errorf("HeartbeatAt = %v, want nil", after.HeartbeatAt) + } + if after.LeaseExpiresAt != nil { + t.Errorf("LeaseExpiresAt = %v, want nil", after.LeaseExpiresAt) + } +} + +func testReclaimSkipsLiveLease(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + const queue = "lease-live" + + j := PendingJob("live", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + if _, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: time.Now().UTC().Add(time.Hour), + }); err != nil { + t.Fatalf("DequeueJobs: %v", err) + } + + got, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + if Contains(got, j.ID) { + t.Fatal("a live lease was reclaimed") + } +} + +func testReclaimFencesPreviousHolder(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + now := time.Now().UTC() + + // This is the split-brain case the whole phase exists to close. A + // worker holds a lease, the lease expires while the worker is paused, + // the reaper reclaims it — and the worker then wakes and tries to + // renew. It must be refused. + const queue = "lease-fenced" + + j := PendingJob("fenced", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(-time.Second), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + heldEpoch := got[0].LeaseEpoch + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + if !Contains(reclaimed, j.ID) { + t.Fatalf("reclaimed set does not contain %s", j.ID) + } + + afterReclaim, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if afterReclaim.LeaseEpoch <= heldEpoch { + t.Errorf("LeaseEpoch = %d after reclaim, want > %d", afterReclaim.LeaseEpoch, heldEpoch) + } + + // The zombie wakes up. + err = s.RenewLease(ctx, j.ID, worker, heldEpoch, now.Add(time.Hour)) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("zombie RenewLease = %v, want %v", err, job.ErrLeaseLost) + } +} + +func testReclaimPreservesRetryCount(t *testing.T, s LeaseStore) { + ctx := context.Background() + + j := RunningJob("retry-untouched", "lease-retry", 0) + j.RetryCount = 2 + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + if _, err := s.ReclaimExpiredLeases(ctx, 100); err != nil { + t.Fatalf("ReclaimExpiredLeases: %v", err) + } + + after, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + // Losing a lease is infrastructure, not a handler failure. Charging it + // to the retry budget would DLQ a job that never once errored. + if after.RetryCount != 2 { + t.Errorf("RetryCount = %d, want it unchanged at 2", after.RetryCount) + } +} + +func testReclaimIsExclusive(t *testing.T, s LeaseStore) { + ctx := context.Background() + + j := RunningJob("exclusive", "lease-exclusive", 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + // Two pools reclaiming concurrently must not both take the job, or two + // workers would run it. Sequential calls prove the same invariant: the + // second call cannot see this job, because the first cleared its lease. + first, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("first ReclaimExpiredLeases: %v", err) + } + second, err := s.ReclaimExpiredLeases(ctx, 100) + if err != nil { + t.Fatalf("second ReclaimExpiredLeases: %v", err) + } + + if !Contains(first, j.ID) { + t.Errorf("first reclaim did not take %s", j.ID) + } + if Contains(second, j.ID) { + t.Errorf("second reclaim took %s again — reclamation is not exclusive", j.ID) + } +} + +func testReclaimIsExclusiveUnderConcurrency(t *testing.T, s LeaseStore) { + ctx := context.Background() + const ( + queue = "lease-concurrent" + jobCount = 20 + reclaimers = 4 + + // Deliberately not jobCount. Reclamation is not queue-scoped (see + // Contains), and newStore may hand this case a store other cases + // have already left expired jobs in. A limit of jobCount would + // let four reclaimers fill their quotas with those instead, leave + // some of `mine` unclaimed, and fail the n == 0 branch below for + // a reason that has nothing to do with exclusivity. The limit + // only has to be too large to ever be the binding constraint; + // what this case measures is double-claiming, not throughput. + reclaimLimit = 10_000 + ) + + mine := make(map[id.JobID]bool, jobCount) + for i := range jobCount { + j := RunningJob(fmt.Sprintf("concurrent-%d", i), queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + mine[j.ID] = true + } + + var ( + mu sync.Mutex + claims = make(map[id.JobID]int) + wg sync.WaitGroup + ) + errCh := make(chan error, reclaimers) + + for range reclaimers { + wg.Add(1) + go func() { + defer wg.Done() + + got, err := s.ReclaimExpiredLeases(ctx, reclaimLimit) + if err != nil { + errCh <- err + + return + } + + mu.Lock() + defer mu.Unlock() + for _, j := range got { + claims[j.ID]++ + } + }() + } + + wg.Wait() + close(errCh) + for err := range errCh { + t.Fatalf("concurrent ReclaimExpiredLeases: %v", err) + } + + // The invariant, not a timing guess: a job handed to two reclaimers + // would be run by two workers. A correct backend never violates this, + // so a correct backend never flakes here. A select-then-update backend + // violates it whenever two scans overlap. + for jobID := range mine { + switch n := claims[jobID]; { + case n == 0: + t.Errorf("job %s was never claimed", jobID) + case n > 1: + t.Errorf("job %s claimed %d times, want exactly 1 — reclamation is not atomic", jobID, n) + } + } +} + +func testLeaseTTLRoundTrips(t *testing.T, s LeaseStore) { + ctx := context.Background() + + j := PendingJob("ttl", "lease-ttl", 6*time.Hour) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + // The pool reads LeaseTTL off the dequeued row to compute each + // renewal's expiry, so a backend that drops it silently reverts every + // job to the default. + if got.LeaseTTL != 6*time.Hour { + t.Errorf("LeaseTTL = %v, want %v", got.LeaseTTL, 6*time.Hour) + } +} + +// timeEqual compares two *time.Time fields the way GetJob round-trips +// them: both nil, or both non-nil and equal instants. Plain == on the +// pointers would compare addresses, and reflect.DeepEqual on a bare +// time.Time can trip over a monotonic reading that a store round-trip +// already strips — Equal is the contract these fields actually promise. +func timeEqual(a, b *time.Time) bool { + if a == nil || b == nil { + return a == b + } + + return a.Equal(*b) +} + +// testUpdateLeasedJobAppliesAtHeldEpoch is the positive case: a worker +// that still holds the epoch it was granted at claim time gets its +// terminal write applied. +func testUpdateLeasedJobAppliesAtHeldEpoch(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + now := time.Now().UTC() + const queue = "lease-update-applies" + + j := PendingJob("update-applies", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + + claimed := got[0] + claimed.State = job.StateCompleted + completedAt := now + claimed.CompletedAt = &completedAt + claimed.LastError = "" // a business field, round-tripped like any other + + if updErr := s.UpdateLeasedJob(ctx, claimed, worker, claimed.LeaseEpoch); updErr != nil { + t.Fatalf("UpdateLeasedJob: %v", updErr) + } + + after, err := s.GetJob(ctx, claimed.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if after.State != job.StateCompleted { + t.Errorf("State = %s, want %s", after.State, job.StateCompleted) + } + if after.CompletedAt == nil { + t.Fatal("CompletedAt = nil, want it set") + } + // The write must have gone through at the epoch the caller named — + // UpdateLeasedJob never bumps it, only the grant and reclaim do. + if after.LeaseEpoch != claimed.LeaseEpoch { + t.Errorf("LeaseEpoch = %d after the write, want it unchanged at %d", + after.LeaseEpoch, claimed.LeaseEpoch) + } +} + +// testUpdateLeasedJobRejectsStaleEpoch covers a worker presenting an +// epoch older than the one the row currently holds. The row must be +// byte-identical afterwards — a refused fenced write is not merely +// "did not apply the intended change," it is "touched nothing at all." +func testUpdateLeasedJobRejectsStaleEpoch(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + now := time.Now().UTC() + const queue = "lease-update-stale-epoch" + + j := PendingJob("update-stale-epoch", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + claimed := got[0] + + before, err := s.GetJob(ctx, claimed.ID) + if err != nil { + t.Fatalf("get before: %v", err) + } + + attempt := *claimed + attempt.State = job.StateCompleted + + err = s.UpdateLeasedJob(ctx, &attempt, worker, claimed.LeaseEpoch-1) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("UpdateLeasedJob with stale epoch = %v, want %v", err, job.ErrLeaseLost) + } + + after, err := s.GetJob(ctx, claimed.ID) + if err != nil { + t.Fatalf("get after: %v", err) + } + if !reflect.DeepEqual(before, after) { + t.Errorf("row changed after a refused fenced write:\nbefore = %+v\nafter = %+v", before, after) + } +} + +// testUpdateLeasedJobRejectsWrongWorker covers a worker presenting the +// held epoch but the wrong worker ID — the same epoch reused by a +// process that never legitimately held it, or a claim-time snapshot +// misattributed to the wrong caller. +func testUpdateLeasedJobRejectsWrongWorker(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + other := id.NewWorkerID() + now := time.Now().UTC() + const queue = "lease-update-wrong-worker" + + j := PendingJob("update-wrong-worker", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + claimed := got[0] + + before, err := s.GetJob(ctx, claimed.ID) + if err != nil { + t.Fatalf("get before: %v", err) + } + + attempt := *claimed + attempt.State = job.StateCompleted + + err = s.UpdateLeasedJob(ctx, &attempt, other, claimed.LeaseEpoch) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("UpdateLeasedJob from another worker = %v, want %v", err, job.ErrLeaseLost) + } + + after, err := s.GetJob(ctx, claimed.ID) + if err != nil { + t.Fatalf("get after: %v", err) + } + if !reflect.DeepEqual(before, after) { + t.Errorf("row changed after a refused fenced write:\nbefore = %+v\nafter = %+v", before, after) + } +} + +// testUpdateLeasedJobRejectsWhenNoLongerRunning covers a job that has +// already left the running state — completion does not bump lease_epoch +// or reassign worker_id, so a naive predicate that checked only those two +// would let a second write land on an already-terminal row. +func testUpdateLeasedJobRejectsWhenNoLongerRunning(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + now := time.Now().UTC() + const queue = "lease-update-not-running" + + j := PendingJob("update-not-running", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + claimed := got[0] + epoch := claimed.LeaseEpoch + + claimed.State = job.StateCompleted + if updErr := s.UpdateLeasedJob(ctx, claimed, worker, epoch); updErr != nil { + t.Fatalf("first UpdateLeasedJob: %v", updErr) + } + + // The job is completed now: still assigned to worker, still at + // epoch — completion touches neither — but no longer running. A + // second fenced write must be refused on state alone. + retry := *claimed + retry.LastError = "a second write attempt" + + err = s.UpdateLeasedJob(ctx, &retry, worker, epoch) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("UpdateLeasedJob on a non-running job = %v, want %v", err, job.ErrLeaseLost) + } +} + +// testUpdateLeasedJobPreservesLeaseColumnsAgainstAStaleSnapshot is the +// case a naive whole-row-plus-predicate implementation fails. The lease +// is granted, renewed several times so lease_expires_at moves well +// ahead of the claim-time value, and then the fenced write is issued +// with the ORIGINAL stale snapshot — exactly what worker/runner.go does, +// since j is captured once at claim time and never refreshed. The +// passing epoch predicate must not be mistaken for permission to write +// j's copy of the lease columns back: lease_expires_at, lease_epoch, +// worker_id, and heartbeat_at must come out exactly as they were right +// before this call, not rolled back to what the stale snapshot said, +// while the business column this call actually changed must have moved. +func testUpdateLeasedJobPreservesLeaseColumnsAgainstAStaleSnapshot(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + now := time.Now().UTC() + const queue = "lease-update-stale-expiry" + + j := PendingJob("update-stale-expiry", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + // Grant properly through DequeueJobs rather than hand-constructing + // the divergence — hand-constructing would test the fixture, not + // this method. + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(30 * time.Second), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + staleSnapshot := got[0] + epoch := staleSnapshot.LeaseEpoch + + // Renew several times so the store's lease_expires_at (and + // heartbeat_at, which every backend's RenewLease also advances) + // moves well past what staleSnapshot remembers. + extended := now + for range 3 { + extended = extended.Add(time.Hour) + if renewErr := s.RenewLease(ctx, staleSnapshot.ID, worker, epoch, extended); renewErr != nil { + t.Fatalf("RenewLease: %v", renewErr) + } + } + + beforeWrite, err := s.GetJob(ctx, staleSnapshot.ID) + if err != nil { + t.Fatalf("get before write: %v", err) + } + + // The worker finishes and hands the fenced write its ORIGINAL, + // now-stale claim-time snapshot. + stale := *staleSnapshot + stale.State = job.StateCompleted + completedAt := now + stale.CompletedAt = &completedAt + + if updErr := s.UpdateLeasedJob(ctx, &stale, worker, epoch); updErr != nil { + t.Fatalf("UpdateLeasedJob: %v", updErr) + } + + after, err := s.GetJob(ctx, staleSnapshot.ID) + if err != nil { + t.Fatalf("get after: %v", err) + } + + // The business column this call actually changed moved. + if after.State != job.StateCompleted { + t.Errorf("State = %s, want %s", after.State, job.StateCompleted) + } + + // Every lease-owned column is exactly what it was right before the + // fenced write — not what the stale claim-time snapshot said. + if !timeEqual(after.LeaseExpiresAt, beforeWrite.LeaseExpiresAt) { + t.Errorf("LeaseExpiresAt = %v, want it unchanged at %v (not rolled back to the claim-time %v)", + after.LeaseExpiresAt, beforeWrite.LeaseExpiresAt, staleSnapshot.LeaseExpiresAt) + } + if after.LeaseEpoch != beforeWrite.LeaseEpoch { + t.Errorf("LeaseEpoch = %d, want it unchanged at %d", after.LeaseEpoch, beforeWrite.LeaseEpoch) + } + if after.WorkerID != beforeWrite.WorkerID { + t.Errorf("WorkerID = %s, want it unchanged at %s", after.WorkerID, beforeWrite.WorkerID) + } + if !timeEqual(after.HeartbeatAt, beforeWrite.HeartbeatAt) { + t.Errorf("HeartbeatAt = %v, want it unchanged at %v", after.HeartbeatAt, beforeWrite.HeartbeatAt) + } +} + +// testUpdateLeasedJobRunnableWriteIsDequeueable covers the case every +// prior UpdateLeasedJob case missed: all of them write a TERMINAL state, +// and a backend whose fenced write persists the row's own storage without +// also restoring whatever index the original claim removed the job from +// would pass every one of them while still stranding a retried job +// forever. scheduleRetry is the call site that writes a runnable state +// (StateRetrying) through this method, so this claims a job, fences a +// write back to StateRetrying with a due RunAt, and asserts a second +// DequeueJobs on the SAME queue returns it. Dequeue is queue-scoped, so +// counting its result is safe per the suite's own rule even though this +// runs against a store shared with every other case. +func testUpdateLeasedJobRunnableWriteIsDequeueable(t *testing.T, s LeaseStore) { + ctx := context.Background() + worker := id.NewWorkerID() + now := time.Now().UTC() + const queue = "lease-update-runnable" + + j := PendingJob("update-runnable", queue, 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + got, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: worker, + LeaseUntil: now.Add(time.Minute), + }) + if err != nil || len(got) != 1 { + t.Fatalf("DequeueJobs: %v (n=%d)", err, len(got)) + } + claimed := got[0] + + claimed.State = job.StateRetrying + claimed.RunAt = now.Add(-time.Second) // already due + claimed.LastError = "transient, retrying" + + if updErr := s.UpdateLeasedJob(ctx, claimed, worker, claimed.LeaseEpoch); updErr != nil { + t.Fatalf("UpdateLeasedJob: %v", updErr) + } + + // A row report alone is not enough — GetJob would show StateRetrying + // whether or not the backend's own claim index (a Redis ZSET; nothing + // analogous on the SQL/Mongo backends, which re-derive candidacy from + // the row) still points at it. Only a second dequeue proves the job + // is actually reachable again. + again, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{queue}, + Limit: 1, + WorkerID: id.NewWorkerID(), + LeaseUntil: now.Add(time.Minute), + }) + if err != nil { + t.Fatalf("second DequeueJobs: %v", err) + } + if len(again) != 1 { + t.Fatalf("second DequeueJobs returned %d jobs, want 1 — the retried job is unreachable", len(again)) + } + if again[0].ID != claimed.ID { + t.Errorf("second DequeueJobs returned job %s, want %s", again[0].ID, claimed.ID) + } +} + +// testReclaimNonPositiveLimitReclaimsNothing pins the non-positive-limit +// contract: a limit <= 0 claims nothing and returns (nil, nil), and, +// critically, leaves the expired job still reclaimable, so a later call +// with a positive limit still returns it. That second half is what +// separates "declined" from "silently consumed", and only it can tell the +// unified guard apart from a backend that swallowed the job. +// +// The guard is a real behaviour change on three of the five backends, and +// each of them broke differently, which is why this belongs in the shared +// suite rather than in one backend's tests: +// +// postgres a negative limit reached `LIMIT $1` and Postgres rejected +// the whole statement with "LIMIT must not be negative" +// (SQLSTATE 2201W), so reclamation ERRORED rather than +// returning empty +// sqlite SQLite defines a negative LIMIT as "no limit", so the same +// call reclaimed the ENTIRE table +// mongo already returned (nil, nil) before any query, as a side +// effect of a fix for a negative-capacity make() panic +// memory read zero and negative as unlimited +// redis read zero and negative as unlimited +// +// Two of those are opposites of each other: the same input that errored on +// Postgres reclaimed everything on SQLite. A backend implementing this +// interface outside the repository gets the contract checked for free by +// running this suite, which is the point of moving it here. +func testReclaimNonPositiveLimitReclaimsNothing(t *testing.T, s LeaseStore) { + ctx := context.Background() + + for _, limit := range []int{0, -1} { + // A queue per limit, because the suite may be handed a shared + // store and reclamation is not queue-scoped. + j := RunningJob("expired", fmt.Sprintf("reclaim-nonpositive-%d", limit), 0) + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("limit=%d: enqueue: %v", limit, err) + } + + got, err := s.ReclaimExpiredLeases(ctx, limit) + if err != nil { + t.Fatalf("limit=%d: ReclaimExpiredLeases: %v", limit, err) + } + if len(got) != 0 { + t.Fatalf("limit=%d: reclaimed %d jobs, want 0", limit, len(got)) + } + + after, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("limit=%d: get: %v", limit, err) + } + if after.State != job.StateRunning { + t.Fatalf("limit=%d: State = %s, want still running (nothing reclaimed)", + limit, after.State) + } + + // The job must still be reclaimable: a non-positive limit must not + // have silently consumed it. + reclaimed, err := s.ReclaimExpiredLeases(ctx, 10) + if err != nil { + t.Fatalf("limit=%d: follow-up ReclaimExpiredLeases: %v", limit, err) + } + if !Contains(reclaimed, j.ID) { + t.Fatalf("limit=%d: job not reclaimed by a follow-up call with a positive limit", + limit) + } + } +} diff --git a/store/storetest/storetest.go b/store/storetest/storetest.go new file mode 100644 index 0000000..487d7df --- /dev/null +++ b/store/storetest/storetest.go @@ -0,0 +1,114 @@ +// Package storetest provides conformance suites that every Dispatch store +// backend must pass. The suites are shared so five implementations cannot +// quietly disagree about semantics that only one of them has tests for. +// +// RunLeaseSuite covers the opt-in lease capability. RunDequeueSuite +// covers the resource-aware dequeue contract, where disagreement is not +// cosmetic: the same job would become eligible on different workers +// depending only on which store the operator chose, and the dimension +// that silently drifts is the one deciding whether a 32 GB job lands on a +// 4 GB machine. Its two load-bearing cases are +// ZeroBudgetSelectsEverything, the guarantee that an unconstrained caller +// still sees exactly what it saw before the option existed, and +// ClaimIsAtomicUnderConcurrency, which proves the fit predicate did not +// cost the claim its atomicity. +// +// Known limitations of RunDequeueSuite, so a backend author knows what is +// unpinned rather than guaranteed: +// +// - ZeroBudgetSelectsEverything asserts set membership, not order. The +// ordering contract is pinned by PriorityOrderingPreservedWithinBudget, +// LimitTruncatesAfterOrdering, and the two PreferHashes cases; a +// backend that ordered correctly only when a budget was present would +// not be caught. +// - Every case names its queues explicitly, so empty DequeueOpts.Queues +// — "all queues" — is never exercised. It cannot be, while the suite +// supports backends that share one store across subtests: an +// all-queues claim would take other cases' jobs. +// - Requirements are built with resource.Set literals or nil, so a +// non-nil empty Set is never round-tripped through a backend here. +// +// Run RunDequeueSuite with `go test -race`. ClaimIsAtomicUnderConcurrency +// is the case that proves the fit predicate did not cost the claim its +// atomicity, and without -race it is only probabilistic: a realistic +// mutant that drops atomicity (select-then-update as two statements +// instead of one) was caught in just 6 of 20 runs under `go test` alone. +// -race is the configuration where the interleaving that double-claims a +// job reliably triggers the detector instead of getting lucky. +// +// The package depends only on job, resource, id, and the root package. It +// must never import a store backend: the backends import this, not the +// reverse. +package storetest + +import ( + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// LeaseStore is what the lease suite requires of a backend: the base job +// store plus the opt-in lease capability. +type LeaseStore interface { + job.Store + job.LeaseStore +} + +// PendingJob builds a job ready to be dequeued now, on the given queue and +// with the given lease TTL. A zero ttl leaves LeaseTTL unset, meaning the +// pool default. +// +// The queue is a parameter because the suite may run against one shared +// store — spinning a fresh Postgres or Redis container per subtest would +// cost more than the coverage is worth — so each case dequeues from its own +// queue to stay isolated from its neighbours. +func PendingJob(name, queue string, ttl time.Duration) *job.Job { + now := time.Now().UTC() + + return &job.Job{ + Entity: dispatch.NewEntity(), + ID: id.NewJobID(), + Name: name, + Queue: queue, + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: now.Add(-time.Second), + LeaseTTL: ttl, + } +} + +// RunningJob builds a job already in the running state with an expired +// lease held by an unknown worker, for testing reclamation directly. +func RunningJob(name, queue string, ttl time.Duration) *job.Job { + now := time.Now().UTC() + started := now.Add(-time.Minute) + expired := now.Add(-time.Second) + + j := PendingJob(name, queue, ttl) + j.State = job.StateRunning + j.StartedAt = &started + j.LeaseExpiresAt = &expired + j.LeaseEpoch = 1 + j.WorkerID = id.NewWorkerID() + + return j +} + +// Contains reports whether jobs includes the given ID. +// +// Reclamation is not queue-scoped, so cases that exercise it must assert on +// the job they created rather than on the length of the returned slice — +// otherwise a shared store makes every such case depend on what its +// neighbours left behind. +func Contains(jobs []*job.Job, jobID id.JobID) bool { + for _, j := range jobs { + if j.ID == jobID { + return true + } + } + + return false +} diff --git a/worker/admission.go b/worker/admission.go new file mode 100644 index 0000000..6565fa9 --- /dev/null +++ b/worker/admission.go @@ -0,0 +1,290 @@ +package worker + +import ( + "context" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" +) + +// admitted is a claimed job together with the local resource lease held +// on its behalf, as handed from the fetcher to a worker. +// +// The lease travels with the job rather than being looked up later, +// because the two are acquired in different goroutines: the fetcher +// admits the job, a worker runs it, and nothing between them may lose +// track of the capacity that was reserved. A nil lease means no manager +// is configured — the degraded path where the pool behaves exactly as it +// did before the resource model existed. +type admitted struct { + job *job.Job + lease resource.Lease +} + +// inflight is the pool's record of one job it currently owns. +// +// It holds the job's cancel func and its resource lease together so the +// single release path can return everything the job took. Keeping the +// lease here rather than on job.Job is deliberate: job.Job is the +// persisted row, shared with the store and serialized to it, and a live +// lease is process-local state that must never be written down. +type inflight struct { + cancel context.CancelCauseFunc + lease resource.Lease + + // leaseEpoch is the lease epoch this worker was granted when it + // claimed the job. Every renewal presents it, and a renewal that no + // longer matches the row means another worker owns the job now. + // + // Unlike the resource lease above — which is live process state that + // must never be written down — this is the opposite: a process-local + // copy of a persisted value, held so renewal need not re-read the row. + leaseEpoch int + + // leaseTTL is how far each renewal pushes the expiry out, resolved + // once at claim time from the job's own LeaseTTL and the pool's + // defaults. Held here so the renewal loop does not issue a GetJob per + // job per interval purely to recover a duration it already knew. + leaseTTL time.Duration +} + +// dequeueBudget is the capacity ceiling this worker offers the store, +// or nil when no resource manager is configured. +// +// Every key is the manager's free capacity EXCEPT disk, which is free +// plus what a registered reclaimer could evict. The asymmetry is the +// whole point of the reclaimer interface, and it is wrong in both +// directions: +// +// - Disk that is cached but unleased is available to a new job, since +// staging evicts to make room. Offering only Free() there would let a +// warm cache stop the worker claiming anything, which is exactly +// backwards — a full cache is a healthy cache. +// - Memory (and CPU, and GPU) held by a running job cannot be handed +// to a second job by any amount of eviction. Offering Free() + +// Reclaimable() there would admit work the box cannot run, which is +// the OOM cascade this whole model exists to prevent. A reclaimer +// registered for memory is therefore ignored HERE on purpose; it +// still serves Manager.Acquire, which can afford to wait. +// +// Custom keys are passed through untouched. The store's fit predicate +// ignores quantities on custom dimensions (see job.DequeueOpts.Budget) +// and matches them by key through CustomKeys instead, so the quantity a +// custom key carries here is informational — it is enforced locally by +// admit, after the claim. +func (p *Pool) dequeueBudget() resource.Set { + if p.resources == nil { + return nil + } + + budget := p.resources.Free() + if budget == nil { + budget = make(resource.Set) + } + + if extra := p.resources.Reclaimable()[resource.Disk]; extra > 0 { + budget[resource.Disk] += extra + } + + return budget +} + +// offeredCustomKeys is the set of custom resource keys this worker +// advertises at dequeue. +// +// An explicit WithWorkerCustomKeys wins, so a worker can offer a +// capability it does not meter. Otherwise the keys are derived from the +// manager's capacity, which is the honest default: a worker configured +// with 2 fpga has 2 fpga to offer, and one configured with none must not +// claim work that needs one. +func (p *Pool) offeredCustomKeys() []string { + if len(p.customKeys) > 0 { + return p.customKeys + } + + if p.resources == nil { + return nil + } + + return p.resources.Capacity().CustomKeys() +} + +// admit reserves local capacity for a job that has already been claimed. +// +// It uses Acquire under a short deadline, NOT TryAcquire, and the +// distinction is the difference between working and deadlocking. +// TryAcquire never reclaims — its own doc says a caller that cannot wait +// cannot afford eviction I/O either — but dequeueBudget offers disk as +// free PLUS reclaimable. Pairing the two would tell the store that 100 +// GiB is available, take delivery of the job it sends back, and then +// refuse it against free alone. Every poll. Forever. The budget's promise +// has to be redeemable by the thing that redeems it, so admission has to +// be able to evict. +// +// Acquire reclaims first and only waits if reclamation was not enough, so +// a deadline turns "block until someone finishes" into "evict if you can, +// then give up". The instinct behind TryAcquire — never block the fetcher, +// which is sitting on claimed, running jobs whose heartbeats are ticking — +// is served by the deadline instead of by refusing to reclaim. +// +// The deadline is pollInterval, derived rather than invented: a refusal +// only costs one requeue and the next poll retries, so the fetcher should +// never stall longer than the cadence it would have waited anyway. +// +// ctx is the budget for the WHOLE batch, not for this job — see +// admissionBudget. Each successive job gets whatever is left of it. +// +// A failure is reachable in normal operation even though the store already +// applied a fit predicate. Dequeue matches custom resources by key only, +// never by quantity, so a worker offering "fpga" can legitimately claim a +// job wanting four of them. It is also reachable on the canonical keys, +// because the budget was computed before the claim and another job may +// have been admitted since. The returned error names the dimensions that +// did not fit, which is what the requeue path logs. +func (p *Pool) admit(ctx context.Context, j *job.Job) (resource.Lease, error) { + if p.resources == nil { + return nil, nil + } + + return p.resources.Acquire(ctx, j.ID.String(), j.Resources) +} + +// admissionBudget bounds how long the fetcher may spend reclaiming for +// one batch of claimed jobs. +// +// One budget for the batch, not one per job. Per job, the worst case is +// batch size × deadline, and both terms are independently tunable: a +// concurrency of 20 with a 5s poll interval is a 100s stall against a +// 30s stale-job threshold — the reaper would start reclaiming jobs this +// fetcher is still holding, in running state and not yet heartbeating. +// Sharing one deadline makes the worst case the deadline itself, whatever +// the batch size. +// +// Spending the budget does not poison the whole batch, but it does cost +// the jobs behind more than just the wait, and the difference matters. +// +// Anything that FITS is granted outright, expired context or not: +// resource.Manager.Acquire returns before it ever looks at ctx when the +// request fits current free capacity. So a batch of jobs this worker has +// room for is admitted in full however long the first one took. +// +// What an exhausted budget also costs is RECLAMATION. Acquire's loop +// checks ctx.Err() at the top and returns before calling reclaimLocked +// below it, so once the budget is spent a job that would have fitted +// after evicting some cached artifact bytes is no longer given the +// chance: it is refused and requeued as a misfit even though the disk it +// needed was reclaimable. That is a correct outcome — the job stays +// pending, nothing is lost, and the next poll gets a fresh budget — but +// it is a throughput cost paid by every job after the one that burned +// the budget, not "the fetcher merely stopped waiting". +// +// It mirrors callCtx: no manager or no jobs means no budget to spend, and +// the caller still gets a cancel func so it can defer uniformly. +func (p *Pool) admissionBudget(batch int) (context.Context, context.CancelFunc) { + if p.resources == nil || batch == 0 { + return p.cancelCtx, func() {} + } + + return context.WithTimeout(p.cancelCtx, p.admitTimeout()) +} + +// admitTimeout is the admission budget's duration. A non-positive poll +// interval would expire the context before Acquire's first iteration, +// degrading it back into the TryAcquire behaviour that cannot redeem the +// disk budget, so it floors at something small rather than at zero. +func (p *Pool) admitTimeout() time.Duration { + if p.pollInterval > 0 { + return p.pollInterval + } + + return time.Millisecond +} + +// requeueLocalMisfit returns a job this worker claimed but could not +// admit to pending, so another worker — or this one, later — can run it. +// +// It reuses the rate-limited requeue path verbatim: same state, same +// short delay. A job that no worker in the fleet can ever fit will bounce +// on that delay rather than run; detecting that condition is the job of +// unschedulable sweeping, which is a later phase and deliberately not +// approximated here. What does pace it is the fetch loop, which treats a +// batch that dispatched nothing as an empty poll and backs off. +// +// cause is logged rather than recomputed. resource.Manager already names +// the dimensions that did not fit in its error, and re-deriving them here +// would take the manager's mutex and call every reclaimer's Available on +// the misfit path — to produce a worse answer, since the natural thing to +// compare against is the ceiling the store was offered rather than the +// free capacity the acquisition actually failed on. +func (p *Pool) requeueLocalMisfit(j *job.Job, cause error) { + if p.cancelCtx.Err() != nil { + // Admission was interrupted by shutdown, not by a shortfall. The + // rate-limited path would write through the pool's own cancelled + // context and silently fail, stranding a running job with no + // worker until the reaper; the undispatched path uses a fresh one. + p.requeueUndispatched(j) + + return + } + + p.logger.Debug("job does not fit local capacity, returning to pending", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + log.Any("required", j.Resources), + log.String("error", cause.Error()), + ) + + p.requeueRateLimited(j) +} + +// releaseQueueSlot returns the queue/tenant token acquired for j, if the +// pool has a queue manager. Safe to call for a job that never ran. +func (p *Pool) releaseQueueSlot(j *job.Job) { + if p.queueManager != nil { + p.queueManager.Release(j.Queue, j.ScopeOrgID) + } +} + +// abandon gives back everything an undispatched job holds during +// shutdown: its row goes back to pending, and its queue token and +// resource lease are released. Without this a stopping pool would leave +// capacity spoken for by a job it never ran. +func (p *Pool) abandon(a admitted) { + p.requeueUndispatched(a.job) + p.releaseQueueSlot(a.job) + + if a.lease != nil { + a.lease.Release() + } +} + +// finishJob returns everything one attempt held: the tracking entry and +// its cancel func, the queue/tenant token, the resource lease, and the +// worker slot. +// +// It runs from a single defer in runJob so a panicking handler cannot +// leak capacity. A pool without middleware.Recover installed will still +// crash on that panic — that is the caller's choice — but it will not +// first strand a lease that nothing else can release, leaving the worker +// permanently short of the memory that job was holding. +// +// The lease is taken from the admitted value rather than read back out +// of the in-flight record, so it is released even if the panic happened +// before the record was ever written. Lease.Release is idempotent, so the +// two paths cannot double-credit the ledger. +func (p *Pool) finishJob(a admitted) { + if rec := p.untrackJob(a.job.ID.String()); rec != nil { + rec.cancel(nil) + } + + p.releaseQueueSlot(a.job) + + if a.lease != nil { + a.lease.Release() + } + + p.slots <- struct{}{} +} diff --git a/worker/admission_test.go b/worker/admission_test.go new file mode 100644 index 0000000..1ca8211 --- /dev/null +++ b/worker/admission_test.go @@ -0,0 +1,1044 @@ +package worker + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/dlq" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/middleware" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/resource/resourcetest" + "github.com/xraph/dispatch/store/memory" +) + +const gib = int64(1) << 30 + +// TestDequeueBudgetUsesFreePlusReclaimableForDiskOnly pins the single +// asymmetry the whole admission model rests on. +// +// Disk that a cache holds but no running job leases is available to a new +// job, because staging evicts to make room; counting only Free() there +// would let a warm cache stop the worker claiming anything. Memory held +// by a running job is available to nobody, whatever is registered against +// it; counting Reclaimable() there would admit work the box cannot run. +// +// Both halves are asserted against ONE manager with a reclaimer on BOTH +// keys, so the test fails if the rule is ever applied by key-agnostic +// arithmetic in either direction. +func TestDequeueBudgetUsesFreePlusReclaimableForDiskOnly(t *testing.T) { + mgr := resource.NewManager(resource.Set{ + resource.Memory: 8 * gib, + resource.Disk: 100 * gib, + }) + + // A staging cache holding 40 GiB across four evictable entries. + diskCache, err := resourcetest.NewFakeReclaimer(mgr, resource.Disk, 10*gib, 4) + if err != nil { + t.Fatalf("disk reclaimer: %v", err) + } + + mgr.RegisterReclaimer(resource.Disk, diskCache) + + // A reclaimer on memory, which must make no difference whatsoever. + // Registering one is legal — Manager.Acquire will use it, because it + // can afford to wait for the eviction — but the dequeue budget must + // not, because a non-blocking claim cannot. + memCache, err := resourcetest.NewFakeReclaimer(mgr, resource.Memory, gib, 2) + if err != nil { + t.Fatalf("memory reclaimer: %v", err) + } + + mgr.RegisterReclaimer(resource.Memory, memCache) + + // One job already running on this worker. + running, ok := mgr.TryAcquire("job-running", resource.Set{ + resource.Memory: 4 * gib, + resource.Disk: 20 * gib, + }) + if !ok { + t.Fatal("TryAcquire for the running job did not fit") + } + + defer running.Release() + + free := mgr.Free() + if got, want := free[resource.Disk], 40*gib; got != want { + t.Fatalf("setup: free disk = %d, want %d", got, want) + } + + if got, want := free[resource.Memory], 2*gib; got != want { + t.Fatalf("setup: free memory = %d, want %d", got, want) + } + + p := &Pool{resources: mgr} + budget := p.dequeueBudget() + + // 40 GiB free + 40 GiB the cache can evict. + if got, want := budget[resource.Disk], 80*gib; got != want { + t.Errorf("budget disk = %d, want %d (free + reclaimable)", got, want) + } + + // 2 GiB free. The 2 GiB the memory reclaimer holds is NOT offered. + if got, want := budget[resource.Memory], 2*gib; got != want { + t.Errorf("budget memory = %d, want %d (free only, never reclaimable)", got, want) + } +} + +// TestDequeueBudgetWithoutReclaimerIsFree is the control for the disk +// rule: with nothing registered to evict, disk is plain free capacity. +func TestDequeueBudgetWithoutReclaimerIsFree(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100 * gib}) + + held, ok := mgr.TryAcquire("job-running", resource.Set{resource.Disk: 30 * gib}) + if !ok { + t.Fatal("TryAcquire did not fit") + } + + defer held.Release() + + p := &Pool{resources: mgr} + + if got, want := p.dequeueBudget()[resource.Disk], 70*gib; got != want { + t.Errorf("budget disk = %d, want %d", got, want) + } +} + +// TestOfferedCustomKeysDerivesFromCapacity pins the default: a worker +// advertises exactly the custom keys its manager meters, and an explicit +// option overrides that. +func TestOfferedCustomKeysDerivesFromCapacity(t *testing.T) { + mgr := resource.NewManager(resource.Set{ + resource.Memory: gib, + "fpga": 2, + "nvme": 1, + }) + + p := &Pool{resources: mgr} + + got := p.offeredCustomKeys() + if len(got) != 2 || got[0] != "fpga" || got[1] != "nvme" { + t.Errorf("offeredCustomKeys() = %v, want [fpga nvme]", got) + } + + p.customKeys = []string{"fpga"} + + if got := p.offeredCustomKeys(); len(got) != 1 || got[0] != "fpga" { + t.Errorf("offeredCustomKeys() with override = %v, want [fpga]", got) + } +} + +// TestLeaseReleasedAfterExecution proves capacity comes back on every +// exit path an attempt has: success, handler error, and panic. +func TestLeaseReleasedAfterExecution(t *testing.T) { + t.Run("success", func(t *testing.T) { + assertLeaseReturned(t, nil) + }) + + t.Run("handler error", func(t *testing.T) { + assertLeaseReturned(t, errBoom) + }) + + // The panic case drives runJob directly rather than through a started + // pool. A handler that panics past a pool with no Recover middleware + // takes the process with it — there is no surviving pool to observe + // afterwards — but the defer that returns the lease is the same one + // either way, and this recovers the panic at the boundary so the + // ledger can be read on the other side of it. + t.Run("panic", func(t *testing.T) { + h := newLeaseHarness(t, false) + + job.RegisterDefinition(h.registry, job.NewDefinition("panicker", + func(_ context.Context, _ struct{}) error { + panic("handler exploded") + })) + + j := newResourceJob("panicker", resource.Set{resource.Memory: gib}) + + lease, err := h.admitOne(j) + if err != nil { + t.Fatalf("admit refused a job that fits: %v", err) + } + + // Pre-condition, not the assertion: the lease is genuinely held + // going in, so the drained check afterwards means something. + if got := h.manager.Free()[resource.Memory]; got != 3*gib { + t.Fatalf("free memory while admitted = %d, want %d", got, 3*gib) + } + + var panicked bool + + func() { + defer func() { + if r := recover(); r != nil { + panicked = true + } + }() + + h.pool.runJob(admitted{job: j, lease: lease}) + }() + + if !panicked { + t.Fatal("expected the handler panic to propagate out of runJob") + } + + h.assertDrained() + }) +} + +// assertLeaseReturned runs one job end to end through a started pool and +// asserts the manager is back at full capacity afterwards. +func assertLeaseReturned(t *testing.T, handlerErr error) { + t.Helper() + + h := newLeaseHarness(t, true) + + var ran atomic.Bool + + job.RegisterDefinition(h.registry, job.NewDefinition("leased", + func(_ context.Context, _ struct{}) error { + ran.Store(true) + + return handlerErr + })) + + j := newResourceJob("leased", resource.Set{resource.Memory: gib}) + if err := h.store.EnqueueJob(context.Background(), j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + h.start() + waitFor(t, "job to run", ran.Load) + h.stop() + + if got := h.manager.Free()[resource.Memory]; got != 4*gib { + t.Errorf("free memory after execution = %d, want %d", got, 4*gib) + } + + h.assertDrained() +} + +// TestAdmitReclaimsDiskBeforeRefusing closes the loop between the two +// halves of the disk rule. +// +// dequeueBudget offers disk as free PLUS what the cache can evict, so +// admission has to be able to evict, or the worker promises the store +// capacity it will then refuse to honour — claiming the job, bouncing it, +// and doing it again on the next poll forever. Nothing registers a +// reclaimer on the shared manager yet, which is the only reason that is +// not already happening in production. +func TestAdmitReclaimsDiskBeforeRefusing(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100 * gib}) + + // A staging cache holding 60 GiB across six evictable entries, so 40 + // GiB is free and 60 GiB is reclaimable. + cache, err := resourcetest.NewFakeReclaimer(mgr, resource.Disk, 10*gib, 6) + if err != nil { + t.Fatalf("disk reclaimer: %v", err) + } + + mgr.RegisterReclaimer(resource.Disk, cache) + + h := newHarness(t, mgr, true) + + if got := h.pool.dequeueBudget()[resource.Disk]; got != 100*gib { + t.Fatalf("setup: budget disk = %d, want %d", got, 100*gib) + } + + if free := mgr.Free()[resource.Disk]; free != 40*gib { + t.Fatalf("setup: free disk = %d, want %d", free, 40*gib) + } + + // A job sized to the budget the store was given. TryAcquire would + // refuse this against the 40 GiB that is free right now. + j := newResourceJob("staging-hog", resource.Set{resource.Disk: 100 * gib}) + + lease, err := h.admitOne(j) + if err != nil { + t.Fatalf("admit refused a job the dequeue budget promised: %v", err) + } + + if cache.Calls() == 0 { + t.Error("admit took the lease without reclaiming; the budget was redeemed by luck, not eviction") + } + + if got := mgr.Free()[resource.Disk]; got != 0 { + t.Errorf("free disk while the job holds everything = %d, want 0", got) + } + + lease.Release() + + if got := mgr.Free()[resource.Disk]; got != 100*gib { + t.Errorf("free disk after release = %d, want %d", got, 100*gib) + } +} + +// TestAdmitRefusesLargerThanCapacityImmediately covers the cheap +// refusal: a want bigger than the whole worker short-circuits in Acquire +// before it waits at all, so the permanently-impossible case never costs +// the admission deadline. +// +// Note what this does NOT prove: because it short-circuits, it never +// reaches the deadline path. TestAdmitRefusalIsBounded is the guard for +// that. +func TestAdmitRefusesLargerThanCapacityImmediately(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 10 * gib}) + h := newHarness(t, mgr, true, WithPollInterval(2*time.Second)) + + j := newResourceJob("too-big", resource.Set{resource.Disk: 40 * gib}) + + start := time.Now() + + if _, err := h.admitOne(j); err == nil { + t.Fatal("admit accepted a job larger than total capacity") + } + + // The deadline is 2s here precisely so that waiting would be obvious. + if elapsed := time.Since(start); elapsed > 500*time.Millisecond { + t.Errorf("admit blocked for %v on an impossible job; want an immediate refusal", elapsed) + } +} + +// TestAdmitRefusalIsBounded exercises the deadline itself: a want that +// fits the worker's capacity but not its free capacity, with nothing +// reclaimable to cover the gap, must give up at the deadline instead of +// waiting for a running job to finish. +// +// The reclaimer is wedged rather than absent, so the refusal goes the +// long way round — reclaim is attempted, frees nothing, and the wait is +// what ends it. +func TestAdmitRefusalIsBounded(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100 * gib}) + + // 80 GiB held by something that is not a reclaimer: a running job. + // Nothing can take this back within the deadline. + running, ok := mgr.TryAcquire("job-running", resource.Set{resource.Disk: 80 * gib}) + if !ok { + t.Fatal("setup: running job did not fit") + } + + defer running.Release() + + cache, err := resourcetest.NewFakeReclaimer(mgr, resource.Disk, gib, 5) + if err != nil { + t.Fatalf("disk reclaimer: %v", err) + } + + cache.SetError(errBoom) // the evictor is wedged + mgr.RegisterReclaimer(resource.Disk, cache) + + const deadline = 120 * time.Millisecond + + h := newHarness(t, mgr, true, WithPollInterval(deadline)) + + // 15 GiB free, 5 GiB stuck behind a broken evictor, 30 GiB wanted. + j := newResourceJob("wont-fit-yet", resource.Set{resource.Disk: 30 * gib}) + + start := time.Now() + _, admitErr := h.admitOne(j) + elapsed := time.Since(start) + + if admitErr == nil { + t.Fatal("admit granted a lease no eviction could cover") + } + + if cache.Calls() == 0 { + t.Error("admit never attempted reclamation") + } + + if elapsed < deadline { + t.Errorf("admit gave up after %v, before the %v deadline; it is not waiting at all", elapsed, deadline) + } + + if elapsed > 4*deadline { + t.Errorf("admit blocked for %v against a %v deadline", elapsed, deadline) + } + + // A refusal must leave nothing behind. + if free := mgr.Free()[resource.Disk]; free != 15*gib { + t.Errorf("free disk after a refusal = %d, want %d", free, 15*gib) + } +} + +// TestAdmissionBudgetIsPerBatchNotPerJob pins the fetcher's worst-case +// stall. +// +// Per job, a batch costs batch size × deadline, and the fetcher spends +// that stall holding jobs the store has already marked running and that +// nobody is heartbeating yet. Concurrency 20 with a 5s poll interval +// would be a 100s stall against a 30s stale-job threshold — the reaper +// reclaiming jobs the fetcher still has in hand. One budget for the batch +// makes the worst case the deadline, whatever the batch size. +func TestAdmissionBudgetIsPerBatchNotPerJob(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Disk: 100 * gib}) + + running, ok := mgr.TryAcquire("job-running", resource.Set{resource.Disk: 80 * gib}) + if !ok { + t.Fatal("setup: running job did not fit") + } + + defer running.Release() + + cache, err := resourcetest.NewFakeReclaimer(mgr, resource.Disk, gib, 5) + if err != nil { + t.Fatalf("disk reclaimer: %v", err) + } + + cache.SetError(errBoom) + mgr.RegisterReclaimer(resource.Disk, cache) + + const deadline = 120 * time.Millisecond + + h := newHarness(t, mgr, true, WithPollInterval(deadline)) + + // Four jobs that each have to wait out the budget. + batch := make([]*job.Job, 0, 4) + for range 4 { + batch = append(batch, newResourceJob("wont-fit-yet", resource.Set{resource.Disk: 30 * gib})) + } + + ctx, cancel := h.pool.admissionBudget(len(batch)) + defer cancel() + + start := time.Now() + + for i, j := range batch { + if _, admitErr := h.pool.admit(ctx, j); admitErr == nil { + t.Fatalf("job %d was admitted against a full worker", i) + } + } + + elapsed := time.Since(start) + + t.Logf("4 unadmittable jobs under one %v budget took %v", deadline, elapsed) + + // Per job this is 4 × 120ms = 480ms. Shared, the first job spends the + // budget and the other three fail on the expired context immediately. + if elapsed > 2*deadline { + t.Errorf("batch admission took %v against a %v budget; the deadline is being spent per job", + elapsed, deadline) + } + + // The spent budget must not poison the batch: Acquire only consults + // the context when a request does not fit, so a job with room is still + // admitted after the budget is gone. + small := newResourceJob("fits-anyway", resource.Set{resource.Disk: gib}) + + lease, admitErr := h.pool.admit(ctx, small) + if admitErr != nil { + t.Fatalf("a job that fits was refused on an expired batch budget: %v", admitErr) + } + + lease.Release() +} + +// TestPoolPacesUnfittableBacklog pins the pacing of a batch that +// dispatched nothing. +// +// The "there may be more ready work, poll again immediately" fast path +// was only safe because every returned job went through the blocking +// hand-off, which paced the loop at the rate work completes. A job +// refused locally never touches that channel and returns instantly, so +// reading len(jobs) as productive spins the claim/requeue cycle as fast +// as the store can serve it — each turn costing a dequeue and an +// UpdateJob against a backlog nothing can run. +func TestPoolPacesUnfittableBacklog(t *testing.T) { + mgr := resource.NewManager(resource.Set{"fpga": 1}) + h := newHarness(t, mgr, true, + WithPoolConcurrency(2), + WithMaxPollInterval(50*time.Millisecond), + ) + + job.RegisterDefinition(h.registry, job.NewDefinition("never-fits", + func(_ context.Context, _ struct{}) error { return nil })) + + // Every one of these passes the store's filter — dequeue matches + // custom keys, not quantities — and fails admission. + for range 60 { + j := newResourceJob("never-fits", resource.Set{"fpga": 4}) + if err := h.store.EnqueueJob(context.Background(), j); err != nil { + t.Fatalf("enqueue: %v", err) + } + } + + h.start() + time.Sleep(500 * time.Millisecond) + + calls := h.store.calls() + + h.stop() + + t.Logf("DequeueJobs calls in 500ms against a 60-job unfittable backlog = %d", calls) + + // One poll per 10ms pollInterval is ~50 in 500ms, and the cadence caps + // it there by construction. Reverting the fix on this exact test + // measures 713 — the loop polling as fast as the store can answer. + if calls > 60 { + t.Errorf("DequeueJobs calls in 500ms = %d, want <= 60; a batch that dispatched nothing is not being paced", calls) + } +} + +// TestPoolRunsRunnableJobBehindUnfittableBacklog is the other half of the +// pacing contract, and the reason a refused batch is paced at the poll +// interval rather than handed to the idle backoff. +// +// The store returns the queue head, so an unrunnable backlog is re-claimed +// a Limit at a time and only clears by being claimed. Backing off +// exponentially between those claims strands a job the worker CAN run +// behind jobs it cannot, for the whole ramp — measured at 8.5s with a 2s +// cap, against ~94ms at the poll cadence. That is a worse failure than the +// spin it would be fixing. +func TestPoolRunsRunnableJobBehindUnfittableBacklog(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Memory: 8 * gib, "fpga": 1}) + h := newHarness(t, mgr, true, + WithPoolConcurrency(2), + WithMaxPollInterval(2*time.Second), + ) + + var ran atomic.Bool + + job.RegisterDefinition(h.registry, job.NewDefinition("never-fits", + func(_ context.Context, _ struct{}) error { return nil })) + job.RegisterDefinition(h.registry, job.NewDefinition("fits", + func(_ context.Context, _ struct{}) error { + ran.Store(true) + + return nil + })) + + // A backlog of jobs this worker can never run, all sorting ahead of + // what comes next by RunAt. + for range 20 { + j := newResourceJob("never-fits", resource.Set{"fpga": 4}) + if err := h.store.EnqueueJob(context.Background(), j); err != nil { + t.Fatalf("enqueue: %v", err) + } + } + + h.start() + time.Sleep(200 * time.Millisecond) // let the loop settle into the backlog + + good := newResourceJob("fits", resource.Set{resource.Memory: gib}) + if err := h.store.EnqueueJob(context.Background(), good); err != nil { + t.Fatalf("enqueue: %v", err) + } + + start := time.Now() + h.pool.Wake() + + deadline := time.Now().Add(3 * time.Second) + for !ran.Load() && time.Now().Before(deadline) { + time.Sleep(2 * time.Millisecond) + } + + elapsed := time.Since(start) + + h.stop() + + t.Logf("runnable job ran %v after Wake, behind a 20-job unfittable backlog", elapsed) + + if !ran.Load() { + t.Fatal("a runnable job never ran; it is stuck behind the unfittable backlog") + } + + // Draining 20 misfits two at a time at a 10ms cadence is ~100ms. One + // second is far past that and far under the 2s cap a single backed-off + // wait would have cost. + if elapsed > time.Second { + t.Errorf("runnable job took %v to run; the refused backlog is being paced by the idle backoff", elapsed) + } +} + +// TestPoolRequeuesJobThatDoesNotFitLocally covers the gap the store +// cannot close: dequeue matches custom resources by KEY, never by +// quantity, so a worker offering "fpga" legitimately claims a job wanting +// four of them. That job must go back to pending, not run — and the +// queue/tenant token taken for it must come back too. +func TestPoolRequeuesJobThatDoesNotFitLocally(t *testing.T) { + mgr := resource.NewManager(resource.Set{"fpga": 1}) + qm := newCountingQueueManager() + h := newHarness(t, mgr, true, WithQueueManager(qm)) + + var ran atomic.Bool + + job.RegisterDefinition(h.registry, job.NewDefinition("needs-four-fpga", + func(_ context.Context, _ struct{}) error { + ran.Store(true) + + return nil + })) + + j := newResourceJob("needs-four-fpga", resource.Set{"fpga": 4}) + beforeRunAt := j.RunAt + + if err := h.store.EnqueueJob(context.Background(), j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + h.start() + + var got *job.Job + + // requeueRateLimited (which requeueLocalMisfit reuses verbatim) now + // clears the assignment fields on its way back to pending, so + // StartedAt no longer survives as proof the store handed the job + // over. RunAt does: requeueRateLimited unconditionally pushes it to + // time.Now().Add(pollInterval), which nothing else in this test ever + // touches, so RunAt moving past its enqueue-time value is proof of + // exactly one thing — a real claim-then-requeue cycle ran. + waitFor(t, "job to be claimed and returned to pending", func() bool { + fetched, err := h.store.GetJob(context.Background(), j.ID) + if err != nil { + return false + } + + got = fetched + + return fetched.State == job.StatePending && fetched.RunAt.After(beforeRunAt) + }) + + h.stop() + + // RunAt pushed forward proves the store DID hand the job over — the + // key filter let it through, exactly as documented — and pending + // proves the worker refused it locally on quantity. + if !got.RunAt.After(beforeRunAt) || got.State != job.StatePending { + t.Fatalf("job state = %q, RunAt = %v (before %v); want pending with RunAt pushed forward after a claim", + got.State, got.RunAt, beforeRunAt) + } + + // And the fields that record a worker assignment must NOT survive — + // this is the property the shared clearJobAssignment helper exists + // for: a pending job must never look claimed by a worker that no + // longer holds it. + if got.StartedAt != nil { + t.Errorf("StartedAt = %v, want nil — a job returned to pending must not still look claimed", got.StartedAt) + } + if !got.WorkerID.IsNil() { + t.Errorf("WorkerID = %s, want cleared", got.WorkerID) + } + + if ran.Load() { + t.Error("handler ran for a job that does not fit local capacity") + } + + if free := mgr.Free()["fpga"]; free != 1 { + t.Errorf("free fpga = %d, want 1 (a refused claim must take no lease)", free) + } + + acquired, released := qm.counts() + if acquired == 0 { + t.Fatal("queue manager was never consulted; the misfit path is not being exercised") + } + + if acquired != released { + t.Errorf("queue tokens acquired = %d, released = %d; a refused job must give its token back", + acquired, released) + } + + h.assertDrained() +} + +// TestAbandonReturnsEverything covers the shutdown path, which no +// realistic race can be made to hit on demand: a job claimed by the +// fetcher and cancelled before the hand-off must give back its row, its +// queue token, and its lease. +func TestAbandonReturnsEverything(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Memory: 4 * gib}) + qm := newCountingQueueManager() + h := newHarness(t, mgr, true, WithQueueManager(qm)) + + ctx := context.Background() + + j := newResourceJob("interrupted", resource.Set{resource.Memory: gib}) + if err := h.store.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + // Walk the fetcher's steps by hand, up to the hand-off it never wins. + claimed, err := h.store.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 1}) + if err != nil || len(claimed) != 1 { + t.Fatalf("dequeue: %v (%d jobs)", err, len(claimed)) + } + + if !h.pool.queueManager.Acquire(claimed[0].Queue, claimed[0].ScopeOrgID) { + t.Fatal("queue manager refused") + } + + lease, err := h.admitOne(claimed[0]) + if err != nil { + t.Fatalf("admit: %v", err) + } + + h.pool.abandon(admitted{job: claimed[0], lease: lease}) + + got, err := h.store.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get job: %v", err) + } + + if got.State != job.StatePending { + t.Errorf("job state = %q, want %q", got.State, job.StatePending) + } + + if got.StartedAt != nil { + t.Errorf("StartedAt = %v, want nil (the attempt never started)", got.StartedAt) + } + + if acquired, released := qm.counts(); acquired != released { + t.Errorf("queue tokens acquired = %d, released = %d", acquired, released) + } + + if free := mgr.Free()[resource.Memory]; free != 4*gib { + t.Errorf("free memory = %d, want %d", free, 4*gib) + } + + h.assertDrained() +} + +// TestRequeueLocalMisfitDuringShutdownUsesFreshContext pins the branch +// that keeps a stopping pool from stranding claimed jobs. +// +// Admission is bounded by the pool's own context, so shutdown makes it +// fail for every job the fetcher is holding. Requeueing those through the +// rate-limited path would write through that same cancelled context, the +// UpdateJob would fail, and the job would sit in running with no worker +// until the reaper noticed. +func TestRequeueLocalMisfitDuringShutdownUsesFreshContext(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Memory: 4 * gib}) + h := newHarness(t, mgr, true) + + ctx := context.Background() + + j := newResourceJob("interrupted", resource.Set{resource.Memory: gib}) + if err := h.store.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + claimed, err := h.store.DequeueJobs(ctx, job.DequeueOpts{Queues: []string{"default"}, Limit: 1}) + if err != nil || len(claimed) != 1 { + t.Fatalf("dequeue: %v (%d jobs)", err, len(claimed)) + } + + // Another job took the whole worker between the budget and the claim, + // so admission would have to wait — the only way a cancelled context + // can be observed, since a job that fits is granted outright. + blocker, ok := mgr.TryAcquire("other-job", resource.Set{resource.Memory: 4 * gib}) + if !ok { + t.Fatal("setup: blocker did not fit") + } + + defer blocker.Release() + + // The pool is stopping: its context is dead and every store call made + // through it would fail. + h.pool.cancelFunc() + + _, admitErr := h.admitOne(claimed[0]) + if admitErr == nil { + t.Fatal("admit succeeded against a cancelled pool context with no free capacity") + } + + h.pool.requeueLocalMisfit(claimed[0], admitErr) + + got, err := h.store.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get job: %v", err) + } + + if got.State != job.StatePending { + t.Errorf("job state = %q, want %q; the requeue wrote through the cancelled context", + got.State, job.StatePending) + } +} + +// TestPoolWithoutManagerPassesZeroBudget is the degradation guarantee: a +// pool with no resource manager sends opts every backend treats as +// unbounded, so it claims exactly what it claimed before this model +// existed. +func TestPoolWithoutManagerPassesZeroBudget(t *testing.T) { + h := newHarness(t, nil, true) + + h.start() + waitFor(t, "a dequeue", func() bool { return h.store.calls() > 0 }) + h.stop() + + opts := h.store.lastOpts() + + if !opts.IsUnbounded() { + t.Errorf("DequeueOpts.IsUnbounded() = false, want true (Budget=%v CustomKeys=%v ReservedFor=%v)", + opts.Budget, opts.CustomKeys, opts.ReservedFor) + } + + if opts.Budget != nil { + t.Errorf("Budget = %v, want nil", opts.Budget) + } + + if opts.CustomKeys != nil { + t.Errorf("CustomKeys = %v, want nil", opts.CustomKeys) + } + + if opts.Limit == 0 { + t.Error("Limit = 0, want the free slot count (the rest of the opts must still be wired)") + } +} + +// TestPoolWithManagerPassesBoundedBudget is the other side of the +// degradation test: with a manager the opts must actually constrain, or +// the wiring above would be untested by it. +func TestPoolWithManagerPassesBoundedBudget(t *testing.T) { + mgr := resource.NewManager(resource.Set{resource.Memory: 4 * gib, "fpga": 2}) + h := newHarness(t, mgr, true) + + h.start() + waitFor(t, "a dequeue", func() bool { return h.store.calls() > 0 }) + h.stop() + + opts := h.store.lastOpts() + + if opts.IsUnbounded() { + t.Fatal("DequeueOpts.IsUnbounded() = true, want false with a resource manager") + } + + if got, want := opts.Budget[resource.Memory], 4*gib; got != want { + t.Errorf("Budget[memory] = %d, want %d", got, want) + } + + if len(opts.CustomKeys) != 1 || opts.CustomKeys[0] != "fpga" { + t.Errorf("CustomKeys = %v, want [fpga]", opts.CustomKeys) + } +} + +// ────────────────────────────────────────────────── +// Harness +// ────────────────────────────────────────────────── + +var errBoom = errors.New("boom") + +// recordingOptsStore is a memory store that remembers the DequeueOpts it +// was called with, which is how the wiring tests read what the fetcher +// built. +type recordingOptsStore struct { + *memory.Store + + mu sync.Mutex + last job.DequeueOpts + count int +} + +func (r *recordingOptsStore) DequeueJobs(ctx context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + r.mu.Lock() + r.last = opts + r.count++ + r.mu.Unlock() + + return r.Store.DequeueJobs(ctx, opts) +} + +func (r *recordingOptsStore) lastOpts() job.DequeueOpts { + r.mu.Lock() + defer r.mu.Unlock() + + return r.last +} + +func (r *recordingOptsStore) calls() int { + r.mu.Lock() + defer r.mu.Unlock() + + return r.count +} + +// countingQueueManager is a QueueManager that admits everything and +// counts both sides of the token. +// +// It exists because WithQueueManager had no test call site anywhere in +// the repo, so nothing could observe a token that was never released — +// which is how two new release call sites landed in a blind spot. +type countingQueueManager struct { + acquires atomic.Int64 + releases atomic.Int64 + refuse atomic.Bool +} + +func newCountingQueueManager() *countingQueueManager { return &countingQueueManager{} } + +func (q *countingQueueManager) Acquire(_, _ string) bool { + if q.refuse.Load() { + return false + } + + q.acquires.Add(1) + + return true +} + +func (q *countingQueueManager) Release(_, _ string) { q.releases.Add(1) } + +func (q *countingQueueManager) counts() (acquired, released int64) { + return q.acquires.Load(), q.releases.Load() +} + +type leaseHarness struct { + t *testing.T + pool *Pool + store *recordingOptsStore + registry *job.Registry + manager resource.Manager + started bool +} + +// newLeaseHarness builds a pool over a 4 GiB memory manager. +func newLeaseHarness(t *testing.T, withRecover bool) *leaseHarness { + t.Helper() + + return newHarness(t, resource.NewManager(resource.Set{resource.Memory: 4 * gib}), withRecover) +} + +// newHarness builds a pool with concurrency 1 over mgr, which may be nil. +// withRecover installs middleware.Recover; the panic test omits it on +// purpose, and then never starts the pool. Options in extra are applied +// last, so they override the defaults here. +func newHarness(t *testing.T, mgr resource.Manager, withRecover bool, extra ...PoolOption) *leaseHarness { + t.Helper() + + logger := log.NewNoopLogger() + s := &recordingOptsStore{Store: memory.New()} + reg := job.NewRegistry() + extensions := ext.NewRegistry(logger) + + var mws []middleware.Middleware + if withRecover { + mws = append(mws, middleware.Recover(logger)) + } + + runner := NewExecutor(reg, extensions, s, dlq.NewService(s, s), + backoff.NewConstant(10*time.Millisecond), logger, mws...) + + opts := []PoolOption{ + WithPoolConcurrency(1), + WithPollInterval(10 * time.Millisecond), + WithMaxPollInterval(10 * time.Millisecond), + WithPoolQueues([]string{"default"}), + } + if mgr != nil { + opts = append(opts, WithResourceManager(mgr)) + } + + opts = append(opts, extra...) + + h := &leaseHarness{ + t: t, + pool: NewPool(s, runner, extensions, logger, opts...), + store: s, + registry: reg, + manager: mgr, + } + + // The panic test calls runJob without Start, so give it the two + // fields Start would have set. Start overwrites both, so a harness + // that is later started is unaffected. + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + h.pool.cancelCtx, h.pool.cancelFunc = ctx, cancel + h.pool.slots = make(chan struct{}, 1) + + t.Cleanup(h.stop) + + return h +} + +// admitOne admits a single job under a one-job admission budget, the way +// the fetcher would for a batch of one. +func (h *leaseHarness) admitOne(j *job.Job) (resource.Lease, error) { + ctx, cancel := h.pool.admissionBudget(1) + defer cancel() + + return h.pool.admit(ctx, j) +} + +func (h *leaseHarness) start() { + h.t.Helper() + + if err := h.pool.Start(context.Background()); err != nil { + h.t.Fatalf("start: %v", err) + } + + h.started = true +} + +func (h *leaseHarness) stop() { + if !h.started { + return + } + + h.started = false + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := h.pool.Stop(ctx); err != nil { + h.t.Errorf("stop: %v", err) + } +} + +// assertDrained checks the manager holds no leases at all, which catches +// a release that returned the right quantity under the wrong lease. +func (h *leaseHarness) assertDrained() { + h.t.Helper() + + if h.manager == nil { + return + } + + if held := h.manager.Leases(); len(held) != 0 { + h.t.Errorf("manager still holds %d lease(s): %+v", len(held), held) + } +} + +func newResourceJob(name string, req resource.Set) *job.Job { + now := time.Now().UTC() + + j := &job.Job{ + ID: id.NewJobID(), + Name: name, + Queue: "default", + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 0, + RunAt: now, + Resources: req, + } + j.CreatedAt = now + j.UpdatedAt = now + + return j +} + +func waitFor(t *testing.T, what string, cond func() bool) { + t.Helper() + + deadline := time.After(5 * time.Second) + + for { + if cond() { + return + } + + select { + case <-deadline: + t.Fatalf("timed out waiting for %s", what) + case <-time.After(5 * time.Millisecond): + } + } +} diff --git a/worker/executor.go b/worker/executor.go deleted file mode 100644 index 2d15a66..0000000 --- a/worker/executor.go +++ /dev/null @@ -1,174 +0,0 @@ -// Package worker provides the job execution engine — an Executor that -// invokes registered handlers through middleware, and a Pool that -// manages concurrent worker goroutines polling for jobs. -package worker - -import ( - "context" - "fmt" - "time" - - log "github.com/xraph/go-utils/log" - - "github.com/xraph/dispatch/backoff" - "github.com/xraph/dispatch/dlq" - "github.com/xraph/dispatch/ext" - "github.com/xraph/dispatch/job" - "github.com/xraph/dispatch/middleware" -) - -// Executor runs a single job through middleware and the registered handler, -// then handles retry logic, DLQ push, state updates, and lifecycle events. -type Executor struct { - registry *job.Registry - extensions *ext.Registry - store job.Store - dlqService *dlq.Service - backoff backoff.Strategy - mw middleware.Middleware - logger log.Logger -} - -// NewExecutor creates an Executor with the given dependencies. -func NewExecutor( - registry *job.Registry, - extensions *ext.Registry, - store job.Store, - dlqService *dlq.Service, - bo backoff.Strategy, - logger log.Logger, - mws ...middleware.Middleware, -) *Executor { - return &Executor{ - registry: registry, - extensions: extensions, - store: store, - dlqService: dlqService, - backoff: bo, - mw: middleware.Chain(mws...), - logger: logger, - } -} - -// Execute runs a job through the middleware chain and handler. -// On success: marks completed, emits JobCompleted. -// On failure with retries remaining: marks retrying with backoff, emits JobRetrying. -// On failure with retries exhausted: marks failed, pushes to DLQ, emits JobFailed + JobDLQ. -func (e *Executor) Execute(ctx context.Context, j *job.Job) error { - handler, ok := e.registry.Get(j.Name) - if !ok { - return fmt.Errorf("no handler registered for job %q", j.Name) - } - - start := time.Now() - - // The terminal handler that calls the registered job handler. - terminal := func(ctx context.Context) error { - return handler(ctx, j.Payload) - } - - // Run through middleware chain. - err := e.mw(ctx, j, terminal) - elapsed := time.Since(start) - - now := time.Now().UTC() - j.UpdatedAt = now - - if err != nil { - return e.handleFailure(ctx, j, err, now) - } - - return e.handleSuccess(ctx, j, now, elapsed) -} - -// handleSuccess marks the job as completed and emits the lifecycle event. -func (e *Executor) handleSuccess(ctx context.Context, j *job.Job, now time.Time, elapsed time.Duration) error { - j.State = job.StateCompleted - j.CompletedAt = &now - - if updateErr := e.store.UpdateJob(ctx, j); updateErr != nil { - e.logger.Error("failed to update job after success", - log.String("job_id", j.ID.String()), - log.String("job_name", j.Name), - log.String("error", updateErr.Error()), - ) - return updateErr - } - - e.extensions.EmitJobCompleted(ctx, j, elapsed) - return nil -} - -// handleFailure increments the retry counter and either retries or sends to DLQ. -func (e *Executor) handleFailure(ctx context.Context, j *job.Job, handlerErr error, now time.Time) error { - j.RetryCount++ - j.LastError = handlerErr.Error() - - if j.RetryCount <= j.MaxRetries { - return e.scheduleRetry(ctx, j, now) - } - - return e.sendToDLQ(ctx, j, handlerErr) -} - -// scheduleRetry sets the job to StateRetrying with a backoff delay. -func (e *Executor) scheduleRetry(ctx context.Context, j *job.Job, now time.Time) error { - delay := e.backoff.Delay(j.RetryCount) - nextRunAt := now.Add(delay) - j.RunAt = nextRunAt - j.State = job.StateRetrying - - if updateErr := e.store.UpdateJob(ctx, j); updateErr != nil { - e.logger.Error("failed to update job for retry", - log.String("job_id", j.ID.String()), - log.String("error", updateErr.Error()), - ) - return updateErr - } - - e.extensions.EmitJobRetrying(ctx, j, j.RetryCount, nextRunAt) - - e.logger.Info("job scheduled for retry", - log.String("job_id", j.ID.String()), - log.String("job_name", j.Name), - log.Int("attempt", j.RetryCount), - log.Int("max_retries", j.MaxRetries), - log.Duration("delay", delay), - ) - - return fmt.Errorf("job %s retry %d/%d: %w", j.Name, j.RetryCount, j.MaxRetries, fmt.Errorf("%s", j.LastError)) -} - -// sendToDLQ marks the job as failed, pushes it to the DLQ, and emits events. -func (e *Executor) sendToDLQ(ctx context.Context, j *job.Job, handlerErr error) error { - j.State = job.StateFailed - - if updateErr := e.store.UpdateJob(ctx, j); updateErr != nil { - e.logger.Error("failed to update job as failed", - log.String("job_id", j.ID.String()), - log.String("error", updateErr.Error()), - ) - return updateErr - } - - if e.dlqService != nil { - if dlqErr := e.dlqService.Push(ctx, j, handlerErr); dlqErr != nil { - e.logger.Error("failed to push job to DLQ", - log.String("job_id", j.ID.String()), - log.String("error", dlqErr.Error()), - ) - } - } - - e.extensions.EmitJobFailed(ctx, j, handlerErr) - e.extensions.EmitJobDLQ(ctx, j, handlerErr) - - e.logger.Warn("job moved to DLQ after exhausting retries", - log.String("job_id", j.ID.String()), - log.String("job_name", j.Name), - log.Int("retry_count", j.RetryCount), - log.String("error", handlerErr.Error()), - ) - - return handlerErr -} diff --git a/worker/executor_compat.go b/worker/executor_compat.go new file mode 100644 index 0000000..262a347 --- /dev/null +++ b/worker/executor_compat.go @@ -0,0 +1,37 @@ +package worker + +import ( + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/dlq" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/middleware" +) + +// Executor is the former name of Runner. +// +// The type was renamed because it orchestrates an attempt — middleware, +// retry, DLQ, state, events — and was never the thing that invokes the +// handler. That is now exec.Executor. This alias keeps existing code +// compiling. +// +// Deprecated: use Runner. +type Executor = Runner + +// NewExecutor creates a Runner with no executor registry, so handlers are +// called directly in-process exactly as before. +// +// Deprecated: use NewRunner, which takes an *exec.Registry. +func NewExecutor( + registry *job.Registry, + extensions *ext.Registry, + store job.Store, + dlqService *dlq.Service, + bo backoff.Strategy, + logger log.Logger, + mws ...middleware.Middleware, +) *Runner { + return NewRunner(registry, extensions, store, dlqService, bo, nil, logger, mws...) +} diff --git a/worker/executor_test.go b/worker/executor_test.go new file mode 100644 index 0000000..e717e67 --- /dev/null +++ b/worker/executor_test.go @@ -0,0 +1,170 @@ +package worker_test + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/dlq" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/middleware" + "github.com/xraph/dispatch/store/memory" + "github.com/xraph/dispatch/worker" + + log "github.com/xraph/go-utils/log" +) + +// newFailingExecutor builds an executor whose only registered job fails with +// handlerErr, and returns it with the store so the test can read back the +// job's state and the DLQ. +func newFailingExecutor(t *testing.T, handlerErr error) (*worker.Executor, *memory.Store) { + t.Helper() + + logger := log.NewNoopLogger() + s := memory.New() + reg := job.NewRegistry() + extensions := ext.NewRegistry(logger) + + job.RegisterDefinition(reg, job.NewDefinition("failing", + func(context.Context, struct{}) error { return handlerErr })) + + executor := worker.NewExecutor( + reg, extensions, s, dlq.NewService(s, s), + backoff.NewConstant(time.Hour), logger, + middleware.Recover(logger), + ) + + return executor, s +} + +// enqueueFailing stores a job with retries remaining, so a retry is what +// would happen if nothing intervened. +func enqueueFailing(t *testing.T, s *memory.Store) *job.Job { + t.Helper() + + now := time.Now().UTC() + + j := &job.Job{ + Entity: dispatch.Entity{CreatedAt: now, UpdatedAt: now}, + ID: id.NewJobID(), + Name: "failing", + Queue: "default", + State: job.StateRunning, + MaxRetries: 5, + RunAt: now, + } + + if err := s.EnqueueJob(context.Background(), j); err != nil { + t.Fatalf("EnqueueJob: %v", err) + } + + return j +} + +// TestExecutePermanentFailureSkipsRetries is the behaviour the whole +// classification chain exists for. With five retries left, a permanent +// failure must still land in the DLQ on the first attempt: the retry +// schedule is there to outlast a transient fault, and a condition that +// cannot change is not one. +func TestExecutePermanentFailureSkipsRetries(t *testing.T) { + ctx := context.Background() + + handlerErr := fmt.Errorf("stage input %q: %w", "model", dispatch.ErrPermanent) + executor, s := newFailingExecutor(t, handlerErr) + j := enqueueFailing(t, s) + + if err := executor.Execute(ctx, j); !errors.Is(err, dispatch.ErrPermanent) { + t.Fatalf("Execute = %v, want the permanent error back", err) + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + + if got.State != job.StateFailed { + t.Fatalf("state = %v, want %v (permanent failures must not be scheduled for retry)", + got.State, job.StateFailed) + } + + if got.RetryCount != 1 { + t.Fatalf("RetryCount = %d, want 1 — the remaining %d attempts must be skipped", + got.RetryCount, j.MaxRetries-1) + } + + entries, err := s.ListDLQ(ctx, dlq.ListOpts{Limit: 10}) + if err != nil { + t.Fatalf("ListDLQ: %v", err) + } + + if len(entries) != 1 { + t.Fatalf("DLQ has %d entries, want 1 — a permanent failure goes straight there", len(entries)) + } +} + +// TestExecuteTransientFailureStillRetries is the other half of the contract. +// An unclassified error keeps its retries, because dead-lettering something +// that would have succeeded costs more than retrying something that will not. +func TestExecuteTransientFailureStillRetries(t *testing.T) { + ctx := context.Background() + + executor, s := newFailingExecutor(t, errors.New("dial tcp: connection refused")) + j := enqueueFailing(t, s) + + if err := executor.Execute(ctx, j); err == nil { + t.Fatal("Execute returned nil for a failing handler") + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + + if got.State != job.StateRetrying { + t.Fatalf("state = %v, want %v — an unclassified failure must keep its retries", + got.State, job.StateRetrying) + } + + entries, err := s.ListDLQ(ctx, dlq.ListOpts{Limit: 10}) + if err != nil { + t.Fatalf("ListDLQ: %v", err) + } + + if len(entries) != 0 { + t.Fatalf("DLQ has %d entries, want 0 — retries were still available", len(entries)) + } +} + +// TestExecutePermanentFailureUnwrapsThroughLayers covers the shape the error +// actually has in production: the sentinel sits under several layers of +// context added between the storage backend and the executor. +func TestExecutePermanentFailureUnwrapsThroughLayers(t *testing.T) { + ctx := context.Background() + + deep := fmt.Errorf("stage input %q: %w", "model", + fmt.Errorf("stage art/gone.bin: %w", + fmt.Errorf("dispatch/artifact: not found: %w", dispatch.ErrPermanent))) + + executor, s := newFailingExecutor(t, deep) + j := enqueueFailing(t, s) + + if err := executor.Execute(ctx, j); err == nil { + t.Fatal("Execute returned nil for a failing handler") + } + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + + if got.State != job.StateFailed { + t.Fatalf("state = %v, want %v — the sentinel must be found through the wrapping", + got.State, job.StateFailed) + } +} diff --git a/worker/export_test.go b/worker/export_test.go new file mode 100644 index 0000000..9c78779 --- /dev/null +++ b/worker/export_test.go @@ -0,0 +1,72 @@ +package worker + +import ( + "context" + "time" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// LeaseTTLFor exposes the unexported TTL resolution to worker_test. +func (p *Pool) LeaseTTLFor(j *job.Job) time.Duration { return p.leaseTTLFor(j) } + +// ReapInterval exposes the resolved scan cadence to worker_test. +func (p *Pool) ReapInterval() time.Duration { return p.resolvedReapInterval() } + +// ReclaimOnce runs a single reclamation pass against ctx. +// +// reapStaleJobs reaches the store through callCtx, which derives from +// cancelCtx — normally set by Start. This wires it directly so a test can +// drive one pass without running the pool's goroutines. +func (p *Pool) ReclaimOnce(ctx context.Context) { + p.cancelCtx, p.cancelFunc = context.WithCancel(ctx) + defer p.cancelFunc() + p.reapStaleJobs() +} + +// TrackJob creates a cancellable context for jobID and records it via the +// pool's own trackJob, exactly as runJob does for a real attempt. It +// returns the context so a test can observe whether — and via +// context.Cause, why — a later heartbeat/renewal pass cancelled it, +// without running the pool's goroutines or the executor. +func (p *Pool) TrackJob(jobID string, leaseEpoch int, leaseTTL time.Duration) context.Context { + ctx, cancel := context.WithCancelCause(context.Background()) + p.trackJob(jobID, cancel, nil, leaseEpoch, leaseTTL) + + return ctx +} + +// HeartbeatOnce runs a single heartbeat/renewal pass against ctx. +// +// sendHeartbeats reaches the store through callCtx, which derives from +// cancelCtx — normally set by Start. This wires it directly so a test can +// drive one pass without running the pool's goroutines, mirroring +// ReclaimOnce above. +func (p *Pool) HeartbeatOnce(ctx context.Context) { + p.cancelCtx, p.cancelFunc = context.WithCancel(ctx) + defer p.cancelFunc() + p.sendHeartbeats() +} + +// WithLeaseFenceForTest attaches a lease fence to ctx exactly as runJob +// does before calling Execute, so a runner_test case can drive +// handleSuccess / scheduleRetry / sendToDLQ through the fenced +// UpdateLeasedJob path without spinning up a full Pool. +func WithLeaseFenceForTest(ctx context.Context, store job.LeaseStore, workerID id.WorkerID, epoch int) context.Context { + return withLeaseFence(ctx, leaseFence{store: store, workerID: workerID, epoch: epoch}) +} + +// RequeueRateLimitedForTest runs requeueRateLimited against ctx, wiring +// cancelCtx exactly as ReclaimOnce does, so a test can drive it without +// running the pool's goroutines. +func (p *Pool) RequeueRateLimitedForTest(ctx context.Context, j *job.Job) { + p.cancelCtx, p.cancelFunc = context.WithCancel(ctx) + defer p.cancelFunc() + p.requeueRateLimited(j) +} + +// RequeueUndispatchedForTest exposes requeueUndispatched to worker_test. +func (p *Pool) RequeueUndispatchedForTest(j *job.Job) { + p.requeueUndispatched(j) +} diff --git a/worker/fifo_unix_test.go b/worker/fifo_unix_test.go new file mode 100644 index 0000000..80511b7 --- /dev/null +++ b/worker/fifo_unix_test.go @@ -0,0 +1,77 @@ +//go:build unix + +package worker_test + +import ( + "context" + "path/filepath" + "syscall" + "testing" + "time" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/inproc" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// TestRunner_OpeningAFIFOWouldHangSoItIsNeverOpened is C2's actual +// reproduction: a FIFO with no writer on the other end blocks a plain +// os.Open forever. Since the worker walks OutputDir strictly after the +// sandbox's own Run has already returned — well past whatever deadline +// governed the attempt itself — nothing about the walk is itself +// context-aware, so one mkfifo left in OutputDir would otherwise wedge +// a worker goroutine permanently, and enough of them would exhaust a +// pool. This drives Execute in a goroutine with a hard wall-clock +// timeout and fails loudly if it does not return well within it, rather +// than actually hanging the test suite the way the bug would hang a +// worker. +func TestRunner_OpeningAFIFOWouldHangSoItIsNeverOpened(t *testing.T) { + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"real.txt": "kept"}, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + rec.beforeReturn = func() { + fifoPath := filepath.Join(rec.got.OutputDir, "pipe") + if err := syscall.Mkfifo(fifoPath, 0o600); err != nil { + t.Fatalf("mkfifo: %v", err) + } + // Deliberately no writer is ever opened on the other end: that + // absence is exactly what makes a blocking os.Open on this path + // hang forever, which is the bug being proven fixed. + } + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", MaxRetries: 3} + + done := make(chan error, 1) + go func() { + done <- runner.Execute(context.Background(), j) + }() + + select { + case err := <-done: + if err != nil { + t.Fatalf("Execute() = %v, want nil — a FIFO must be skipped, not fail the attempt", err) + } + case <-time.After(3 * time.Second): + t.Fatal("HUNG: Execute() did not return within 3s — a FIFO in OutputDir was opened and blocked forever") + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + if len(links) != 1 || links[0].Name != "real.txt" { + t.Errorf("committed links = %+v, want exactly [real.txt]", links) + } +} diff --git a/worker/lease_fence.go b/worker/lease_fence.go new file mode 100644 index 0000000..7050a29 --- /dev/null +++ b/worker/lease_fence.go @@ -0,0 +1,58 @@ +package worker + +import ( + "context" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" +) + +// leaseFenceKey is the unexported context key leaseFence values are +// stored under. An unexported type keyed to this package is what stops +// a caller outside worker from colliding with — or forging — the fence. +type leaseFenceKey struct{} + +// leaseFence carries what a fenced terminal write needs: the store's +// lease capability, this worker's own ID, and the epoch it was granted +// at claim time. +// +// The epoch travels here rather than being re-read from the job at +// write time on purpose. inflight.leaseEpoch is a process-local copy +// taken once, at the moment DequeueJobs granted it; the job value a +// handler returns is that same claim-time snapshot passed by reference +// through the whole attempt, so reading j.LeaseEpoch here would just be +// a slower way to reach the identical number. What it must NOT become is +// a value re-read from the store: that would defeat the fence by asking +// the row what epoch to check itself against. +type leaseFence struct { + store job.LeaseStore + workerID id.WorkerID + epoch int +} + +// withLeaseFence attaches f to ctx for the duration of one job attempt. +// +// The pool calls this once, in runJob, before handing ctx to Execute — +// never inside Runner itself, which has no notion of "the pool's +// tracked epoch" and must not grow one. A Runner used directly, without +// a Pool (see NewExecutor, and every runner_test.go case that calls +// Execute against a bare context), simply never sees a fence and keeps +// today's unfenced UpdateJob — see leaseFenceFromContext. +func withLeaseFence(ctx context.Context, f leaseFence) context.Context { + return context.WithValue(ctx, leaseFenceKey{}, f) +} + +// leaseFenceFromContext returns the fence attached by withLeaseFence, if +// ctx carries one. +// +// ok is false whenever no fence was attached — a store that does not +// implement job.LeaseStore, or a Runner driven directly without a Pool +// — and every caller here treats that identically to "use the unfenced +// path," which is the backward-compatibility guarantee: a store that +// implements only job.Store must keep behaving exactly as it did before +// this method existed. +func leaseFenceFromContext(ctx context.Context) (leaseFence, bool) { + f, ok := ctx.Value(leaseFenceKey{}).(leaseFence) + + return f, ok +} diff --git a/worker/lease_fence_test.go b/worker/lease_fence_test.go new file mode 100644 index 0000000..6c44378 --- /dev/null +++ b/worker/lease_fence_test.go @@ -0,0 +1,355 @@ +package worker_test + +import ( + "context" + "errors" + "testing" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/dlq" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/middleware" + "github.com/xraph/dispatch/store/memory" + "github.com/xraph/dispatch/worker" +) + +// fakeLeaseJobStore extends fakeJobStore (runner_test.go) with a +// scriptable UpdateLeasedJob and no-op RenewLease / ReclaimExpiredLeases +// stubs, so it satisfies job.LeaseStore. This is what lets a runner-level +// test drive the fenced path through a specific outcome — applied, or +// job.ErrLeaseLost — without a real store's timing, mirroring +// fakeLeaseStore in lease_test.go for the pool's own renewal path. +type fakeLeaseJobStore struct { + *fakeJobStore + + leasedErr error + leasedUpdates int + lastJob *job.Job + lastWorkerID id.WorkerID + lastEpoch int +} + +func (f *fakeLeaseJobStore) UpdateLeasedJob(_ context.Context, j *job.Job, workerID id.WorkerID, epoch int) error { + f.leasedUpdates++ + f.lastJob = j + f.lastWorkerID = workerID + f.lastEpoch = epoch + + return f.leasedErr +} + +func (f *fakeLeaseJobStore) RenewLease(context.Context, id.JobID, id.WorkerID, int, time.Time) error { + return nil +} + +func (f *fakeLeaseJobStore) ReclaimExpiredLeases(context.Context, int) ([]*job.Job, error) { + return nil, nil +} + +var _ job.LeaseStore = (*fakeLeaseJobStore)(nil) + +func TestRunner_HandleSuccess_RoutesThroughFencedWriteWhenAttached(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition("ok.job", func(context.Context, struct{}) error { return nil }).Register(reg) + + store := &fakeLeaseJobStore{fakeJobStore: newFakeJobStore()} + runner := worker.NewRunner( + reg, ext.NewRegistry(log.NewNoopLogger()), store, nil, + backoff.NewExponential(time.Second, time.Hour), nil, log.NewNoopLogger(), + ) + + workerID := id.NewWorkerID() + j := &job.Job{ID: id.NewJobID(), Name: "ok.job", MaxRetries: 3} + ctx := worker.WithLeaseFenceForTest(context.Background(), store, workerID, 7) + + if err := runner.Execute(ctx, j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + if store.leasedUpdates != 1 { + t.Errorf("UpdateLeasedJob calls = %d, want 1", store.leasedUpdates) + } + if store.updates != 0 { + t.Errorf("UpdateJob calls = %d, want 0 — a fenced attempt must not also use the unfenced path", + store.updates) + } + if store.lastWorkerID != workerID || store.lastEpoch != 7 { + t.Errorf("fenced write used (worker=%s epoch=%d), want (worker=%s epoch=%d)", + store.lastWorkerID, store.lastEpoch, workerID, 7) + } + if store.lastJob.State != job.StateCompleted { + t.Errorf("fenced write's State = %s, want %s", store.lastJob.State, job.StateCompleted) + } +} + +func TestRunner_Execute_FallsBackToUnfencedWriteWithNoFenceAttached(t *testing.T) { + // No context fence at all — a Runner used without a Pool (NewExecutor, + // or any direct caller) must keep calling the plain UpdateJob it + // always has. This is the backward-compatibility guarantee: adopting + // job.LeaseStore on a backend must not become a hard requirement of + // using Runner. + reg := job.NewRegistry() + job.NewDefinition("ok.job", func(context.Context, struct{}) error { return nil }).Register(reg) + + store := &fakeLeaseJobStore{fakeJobStore: newFakeJobStore()} + runner := worker.NewRunner( + reg, ext.NewRegistry(log.NewNoopLogger()), store, nil, + backoff.NewExponential(time.Second, time.Hour), nil, log.NewNoopLogger(), + ) + + j := &job.Job{ID: id.NewJobID(), Name: "ok.job", MaxRetries: 3} + + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + if store.updates != 1 { + t.Errorf("UpdateJob calls = %d, want 1", store.updates) + } + if store.leasedUpdates != 0 { + t.Errorf("UpdateLeasedJob calls = %d, want 0 — no fence was attached", store.leasedUpdates) + } +} + +// TestRunner_TerminalWrites_AbandonOnLeaseLost exercises all four +// fenced call sites — success, retry, DLQ, and the launch-failure +// requeue — against a store scripted to return job.ErrLeaseLost, and +// checks the one behaviour the whole fix exists for: the runner does +// not retry, DLQ, or touch the row again, and the loss is observable +// through the extension registry with no new plumbing. +func TestRunner_TerminalWrites_AbandonOnLeaseLost(t *testing.T) { + tests := []struct { + name string + jobName string + handler func(context.Context, struct{}) error + maxRetries int + retryCount int + }{ + { + name: "handleSuccess", + jobName: "ok.job", + handler: func(context.Context, struct{}) error { return nil }, + }, + { + name: "scheduleRetry", + jobName: "retry.job", + handler: func(context.Context, struct{}) error { return errors.New("boom") }, + maxRetries: 3, + retryCount: 0, + }, + { + name: "sendToDLQ", + jobName: "dlq.job", + handler: func(context.Context, struct{}) error { return errors.New("boom") }, + maxRetries: 1, + retryCount: 1, // already at the ceiling, so handleFailure routes straight to sendToDLQ + }, + { + // A launch failure never reaches the middleware chain's + // ordinary error path — it is *exec.Error carrying a status + // that CountsAgainstRetries() reports false for, which + // handleFailure routes to requeueAfterLaunchFailure instead of + // scheduleRetry. Before this fix that call went through the + // plain, unfenced store.UpdateJob no matter what was attached + // to ctx; this case only proves anything once it does not. + name: "requeueAfterLaunchFailure", + jobName: "launch.job", + handler: func(context.Context, struct{}) error { + return &exec.Error{Status: exec.StatusLaunchFailed, Msg: "boom"} + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition(tt.jobName, tt.handler).Register(reg) + + store := &fakeLeaseJobStore{fakeJobStore: newFakeJobStore(), leasedErr: job.ErrLeaseLost} + extensions := ext.NewRegistry(log.NewNoopLogger()) + tracker := &trackingExt{} + extensions.Register(tracker) + + runner := worker.NewRunner( + reg, extensions, store, nil, + backoff.NewConstant(time.Millisecond), nil, log.NewNoopLogger(), + ) + + workerID := id.NewWorkerID() + j := &job.Job{ + ID: id.NewJobID(), + Name: tt.jobName, + MaxRetries: tt.maxRetries, + RetryCount: tt.retryCount, + } + ctx := worker.WithLeaseFenceForTest(context.Background(), store, workerID, 5) + + err := runner.Execute(ctx, j) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("Execute() error = %v, want job.ErrLeaseLost", err) + } + + if store.leasedUpdates != 1 { + t.Errorf("UpdateLeasedJob calls = %d, want exactly 1 — no retry of the write itself", + store.leasedUpdates) + } + if store.updates != 0 { + t.Errorf("UpdateJob calls = %d, want 0 — a lost lease must never fall back to the unfenced write", + store.updates) + } + if !tracker.failed.Load() { + t.Error("OnJobFailed did not fire — audit_hook/relay_hook would not observe the lost lease") + } + }) + } +} + +// TestPool_LeaseLostDuringExecution_DoesNotClobberTheWinner reproduces +// the bug this whole track exists to close, end to end, through the real +// Pool + Runner + a real store: a worker's lease is reclaimed and handed +// to a second claimant while the first worker's handler is still running, +// and the first worker's handler then returns success. Before this fix +// that unfenced terminal write would still land — rolling lease_epoch +// backwards, marking the job completed, and fencing the legitimate new +// holder off its own job. After it, the write must be refused and the +// second claimant's row must survive completely untouched. +func TestPool_LeaseLostDuringExecution_DoesNotClobberTheWinner(t *testing.T) { + logger := log.NewNoopLogger() + s := memory.New() + reg := job.NewRegistry() + extensions := ext.NewRegistry(logger) + tracker := &trackingExt{} + extensions.Register(tracker) + + dlqSvc := dlq.NewService(s, s) + bo := backoff.NewConstant(10 * time.Millisecond) + + entered := make(chan struct{}) + release := make(chan struct{}) + job.RegisterDefinition(reg, job.NewDefinition("long-job", func(_ context.Context, _ struct{}) error { + close(entered) + <-release + + return nil + })) + + executor := worker.NewExecutor(reg, extensions, s, dlqSvc, bo, logger, middleware.Recover(logger)) + + pool := worker.NewPool(s, executor, extensions, logger, + worker.WithPoolConcurrency(1), + worker.WithPollInterval(10*time.Millisecond), + worker.WithPoolQueues([]string{"default"}), + // A short TTL and no heartbeat loop: the pool grants a lease at + // claim time but never renews it, so it's genuinely expired by + // the time this test reclaims it below — not a timing fiction. + worker.WithDefaultLeaseTTL(30*time.Millisecond), + ) + + j := &job.Job{ + ID: id.NewJobID(), + Name: "long-job", + Queue: "default", + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: time.Now().UTC(), + } + j.CreatedAt = time.Now().UTC() + j.UpdatedAt = j.CreatedAt + + if err := s.EnqueueJob(context.Background(), j); err != nil { + t.Fatalf("enqueue: %v", err) + } + + if err := pool.Start(context.Background()); err != nil { + t.Fatalf("start: %v", err) + } + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + _ = pool.Stop(ctx) + }() + + select { + case <-entered: + case <-time.After(5 * time.Second): + t.Fatal("handler never started") + } + + // The lease TTL (30ms) has certainly elapsed by now: entering the + // handler required a full poll-and-claim round trip already. Reclaim + // it out from under the running worker and hand it to a second + // claimant, exactly as an operator's reaper would after a pause. + time.Sleep(60 * time.Millisecond) + + reclaimed, err := s.ReclaimExpiredLeases(context.Background(), 10) + if err != nil { + t.Fatalf("reclaim: %v", err) + } + if !storetestContains(reclaimed, j.ID) { + t.Fatalf("reclaimed set does not contain %s", j.ID) + } + + winner := id.NewWorkerID() + claimed, err := s.DequeueJobs(context.Background(), job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + WorkerID: winner, + LeaseUntil: time.Now().UTC().Add(time.Hour), + }) + if err != nil || len(claimed) != 1 { + t.Fatalf("second claim: %v (n=%d)", err, len(claimed)) + } + winnerEpoch := claimed[0].LeaseEpoch + + // Let the original (now-zombie) attempt finish successfully. + close(release) + + deadline := time.After(3 * time.Second) + for !tracker.failed.Load() { + select { + case <-deadline: + t.Fatal("timed out waiting for the zombie's fenced write to be refused") + default: + time.Sleep(10 * time.Millisecond) + } + } + + after, err := s.GetJob(context.Background(), j.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if after.State != job.StateRunning { + t.Errorf("State = %s, want %s — the zombie's write must not have applied", + after.State, job.StateRunning) + } + if after.WorkerID != winner { + t.Errorf("WorkerID = %s, want %s — the winner's claim must survive untouched", + after.WorkerID, winner) + } + if after.LeaseEpoch != winnerEpoch { + t.Errorf("LeaseEpoch = %d, want %d — must not have been rolled back", after.LeaseEpoch, winnerEpoch) + } + if after.CompletedAt != nil { + t.Errorf("CompletedAt = %v, want nil — the job was never legitimately completed", after.CompletedAt) + } +} + +// storetestContains mirrors storetest.Contains without importing the +// storetest package's testing.T-shaped API into this one. +func storetestContains(jobs []*job.Job, jobID id.JobID) bool { + for _, j := range jobs { + if j.ID == jobID { + return true + } + } + + return false +} diff --git a/worker/lease_test.go b/worker/lease_test.go new file mode 100644 index 0000000..d159416 --- /dev/null +++ b/worker/lease_test.go @@ -0,0 +1,229 @@ +package worker_test + +import ( + "context" + "errors" + "testing" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" + "github.com/xraph/dispatch/worker" +) + +// TestPool_LeaseTTLFor exercises the TTL precedence chain: a job's own +// LeaseTTL, then the pool's WithDefaultLeaseTTL, then the legacy +// WithStaleJobThreshold, then job.DefaultLeaseTTL when nothing at all is +// configured. +func TestPool_LeaseTTLFor(t *testing.T) { + const ( + jobTTL = 5 * time.Second + poolDefault = 45 * time.Second + staleThresh = 90 * time.Second + ) + + tests := []struct { + name string + opts []worker.PoolOption + job *job.Job + want time.Duration + }{ + { + name: "job's own TTL wins", + opts: []worker.PoolOption{ + worker.WithDefaultLeaseTTL(poolDefault), + worker.WithStaleJobThreshold(staleThresh), + }, + job: &job.Job{LeaseTTL: jobTTL}, + want: jobTTL, + }, + { + name: "pool default when the job declares none", + opts: []worker.PoolOption{ + worker.WithDefaultLeaseTTL(poolDefault), + worker.WithStaleJobThreshold(staleThresh), + }, + job: &job.Job{}, + want: poolDefault, + }, + { + name: "stale-job threshold when the pool default is unset", + opts: []worker.PoolOption{ + worker.WithStaleJobThreshold(staleThresh), + }, + job: &job.Job{}, + want: staleThresh, + }, + { + name: "job.DefaultLeaseTTL when nothing is configured", + opts: nil, + job: &job.Job{}, + want: job.DefaultLeaseTTL, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + pool := worker.NewPool(memory.New(), nil, nil, log.NewNoopLogger(), tt.opts...) + + got := pool.LeaseTTLFor(tt.job) + if got != tt.want { + t.Errorf("LeaseTTLFor() = %v, want %v", got, tt.want) + } + }) + } +} + +// TestLeaseFencing is the store-level proof the pool's renewal path +// depends on: a lease granted at claim time, expired and reclaimed by +// another worker, must fence the original holder's next renewal with +// job.ErrLeaseLost. +func TestLeaseFencing(t *testing.T) { + ctx := context.Background() + s := memory.New() + + workerA := id.NewWorkerID() + + j := &job.Job{ + ID: id.NewJobID(), + Name: "fenced", + Queue: "default", + State: job.StatePending, + MaxRetries: 3, + RunAt: time.Now().UTC(), + } + j.CreatedAt = time.Now().UTC() + j.UpdatedAt = j.CreatedAt + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue error: %v", err) + } + + // Claim it with a lease that is already expired, so the very next + // ReclaimExpiredLeases sweep picks it up. + claimed, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + WorkerID: workerA, + LeaseUntil: time.Now().UTC().Add(-time.Minute), + }) + if err != nil { + t.Fatalf("dequeue error: %v", err) + } + if len(claimed) != 1 { + t.Fatalf("claimed = %d jobs, want 1", len(claimed)) + } + grantedEpoch := claimed[0].LeaseEpoch + + reclaimed, err := s.ReclaimExpiredLeases(ctx, 10) + if err != nil { + t.Fatalf("reclaim error: %v", err) + } + if len(reclaimed) != 1 { + t.Fatalf("reclaimed = %d jobs, want 1", len(reclaimed)) + } + + // Worker A, unaware it was reclaimed, tries to renew with the epoch + // it was originally granted. That must be fenced. + err = s.RenewLease(ctx, j.ID, workerA, grantedEpoch, time.Now().UTC().Add(time.Minute)) + if !errors.Is(err, job.ErrLeaseLost) { + t.Fatalf("RenewLease() error = %v, want job.ErrLeaseLost", err) + } +} + +// fakeLeaseStore embeds *memory.Store for every method except RenewLease, +// which it overrides to return a caller-configured error. This is what +// lets a test drive the pool's heartbeat/renewal path through a specific +// outcome — lease lost, or a transient store error — without needing a +// lease that has genuinely expired or a second worker to steal it. +type fakeLeaseStore struct { + *memory.Store + + renewErr error +} + +// RenewLease shadows the promoted *memory.Store method and always +// returns the configured error, ignoring every argument. +func (f *fakeLeaseStore) RenewLease( + _ context.Context, + _ id.JobID, + _ id.WorkerID, + _ int, + _ time.Time, +) error { + return f.renewErr +} + +// TestPool_SendHeartbeats_CancelsOnLeaseLost is the pool-level half of +// the fencing proof: TestLeaseFencing above shows the STORE returns +// job.ErrLeaseLost when the caller no longer holds the lease; this shows +// the POOL reacts to that return by cancelling the job's context with +// job.ErrLeaseLost as the cause — the one behaviour this whole task +// exists to add. Asserting only ctx.Err() != nil would not distinguish +// this from a shutdown cancellation, so the assertion goes through +// context.Cause. +func TestPool_SendHeartbeats_CancelsOnLeaseLost(t *testing.T) { + s := &fakeLeaseStore{Store: memory.New(), renewErr: job.ErrLeaseLost} + pool := worker.NewPool(s, nil, nil, log.NewNoopLogger()) + + jobID := id.NewJobID().String() + ctx := pool.TrackJob(jobID, 3, 30*time.Second) + + pool.HeartbeatOnce(context.Background()) + + if ctx.Err() == nil { + t.Fatalf("job context was not cancelled after RenewLease returned job.ErrLeaseLost") + } + + cause := context.Cause(ctx) + if !errors.Is(cause, job.ErrLeaseLost) { + t.Fatalf("context.Cause(ctx) = %v, want job.ErrLeaseLost", cause) + } +} + +// TestPool_SendHeartbeats_KeepsJobAliveOnTransientError is the negative +// case: a renewal failure that is NOT job.ErrLeaseLost — a transient +// store blip — must leave the job's context alive. A pool that cancelled +// on every renewal error would kill healthy jobs whenever the store +// hiccups, which is exactly what the WARN-not-cancel branch in +// sendHeartbeats exists to avoid. +func TestPool_SendHeartbeats_KeepsJobAliveOnTransientError(t *testing.T) { + s := &fakeLeaseStore{Store: memory.New(), renewErr: context.DeadlineExceeded} + pool := worker.NewPool(s, nil, nil, log.NewNoopLogger()) + + jobID := id.NewJobID().String() + ctx := pool.TrackJob(jobID, 3, 30*time.Second) + + pool.HeartbeatOnce(context.Background()) + + if err := ctx.Err(); err != nil { + t.Fatalf("job context was cancelled on a transient renewal error: cause = %v", context.Cause(ctx)) + } +} + +// storeOnly wraps a job.Store behind the bare interface, so it satisfies +// job.Store without also satisfying job.LeaseStore even though the +// concrete memory.Store underneath implements both. Embedding the +// interface value rather than the concrete type is what hides the extra +// method set: storeOnly's own method set is exactly job.Store's. +type storeOnly struct { + job.Store +} + +// TestPool_CapabilityLessStore proves a custom backend that implements +// only job.Store — not job.LeaseStore — keeps working: construction logs +// rather than panics, and TTL resolution still falls all the way back to +// job.DefaultLeaseTTL. +func TestPool_CapabilityLessStore(t *testing.T) { + s := storeOnly{Store: memory.New()} + + pool := worker.NewPool(s, nil, nil, log.NewNoopLogger()) + + got := pool.LeaseTTLFor(&job.Job{}) + if got != job.DefaultLeaseTTL { + t.Errorf("LeaseTTLFor() = %v, want job.DefaultLeaseTTL (%v)", got, job.DefaultLeaseTTL) + } +} diff --git a/worker/outputs_other.go b/worker/outputs_other.go new file mode 100644 index 0000000..6b7074f --- /dev/null +++ b/worker/outputs_other.go @@ -0,0 +1,42 @@ +//go:build !unix + +package worker + +import "os" + +// openRegularNoFollow opens path for reading. +// +// This platform has no portable, dependency-free equivalent of Unix's +// O_NOFOLLOW/O_NONBLOCK open flags in the standard library, so it +// cannot close the narrow TOCTOU window the unix build additionally +// closes (see outputs_unix.go's doc comment). That window requires a +// still-running process to swap a regular file for a symlink or FIFO +// between collectOutputEntries listing it and this function opening it +// — collectOutputEntries' own Lstat-based type filter, run moments +// earlier in the same synchronous walk, is what stops the case this +// package actually exists to prevent: a symlink or FIFO present in +// OutputDir all along is never opened at all, on any platform, because +// it never reaches this function to begin with. +func openRegularNoFollow(path string) (*os.File, error) { + return os.Open(path) +} + +// processAlive reports whether pid names a process that is currently +// running. On this platform, unlike Unix, os.FindProcess itself opens +// a handle to the process and fails if none exists at pid, so the +// FindProcess call is the actual check here — there is no separate +// Signal(0) probe to make it accurate. +func processAlive(pid int) bool { + if pid <= 0 { + return false + } + + proc, err := os.FindProcess(pid) + if err != nil { + return false + } + + _ = proc.Release() //nolint:errcheck // best-effort handle cleanup + + return true +} diff --git a/worker/outputs_unix.go b/worker/outputs_unix.go new file mode 100644 index 0000000..769271d --- /dev/null +++ b/worker/outputs_unix.go @@ -0,0 +1,65 @@ +//go:build unix + +package worker + +import ( + "os" + "syscall" +) + +// openRegularNoFollow opens path for reading, refusing to follow a +// symlink at the final path component and refusing to block if it has +// instead become a FIFO with no writer. +// +// collectOutputEntries already filters out a symlink or FIFO dirent by +// its Lstat-reported type before any path derived from it ever reaches +// here, so this is defense in depth against the narrow window between +// that listing and this open: something that was a regular file when +// listed but has since been replaced (a still-running process the +// sandbox left behind, racing this walk) would otherwise be followed — +// or blocked on — anyway. +// +// - O_NOFOLLOW makes the open fail with ELOOP if the final path +// component is now a symlink, rather than silently opening whatever +// it resolves to — which, unlike the file it replaced, could be +// anything this worker process can read: its own config, cloud +// credentials, a mounted service-account token. +// - O_NONBLOCK stops the open call itself from blocking if the final +// path component is now a FIFO with no writer on the other end. +// Without it, open() on such a FIFO blocks before this function +// even returns — well before the caller's later +// f.Stat().Mode().IsRegular() check ever gets a chance to run and +// reject it, so O_NOFOLLOW alone would not have closed this half of +// the same race. It is deliberately never cleared afterward: POSIX +// guarantees O_NONBLOCK has no effect on a regular file's +// read/write behaviour, which is the only kind of file this ever +// returns successfully — anything else is caught and rejected by +// the caller's own fstat check. +func openRegularNoFollow(path string) (*os.File, error) { + // #nosec G304 -- this function is the symlink-race hardening: O_NOFOLLOW + // and O_NONBLOCK are the mitigation, and the caller fstats for a regular + // file before using the handle. + return os.OpenFile(path, os.O_RDONLY|syscall.O_NOFOLLOW|syscall.O_NONBLOCK, 0) +} + +// processAlive reports whether pid names a process that is currently +// running, using the standard POSIX existence probe: sending signal 0 +// sends nothing but still fails with ESRCH if the process is gone. +// +// os.FindProcess itself cannot answer this on Unix — per its own +// documentation it "always succeeds and returns a Process for the +// given pid, regardless of whether the process exists" on this +// platform family — so the real check is the Signal(0) call, not the +// FindProcess call that precedes it. +func processAlive(pid int) bool { + if pid <= 0 { + return false + } + + proc, err := os.FindProcess(pid) + if err != nil { + return false + } + + return proc.Signal(syscall.Signal(0)) == nil +} diff --git a/worker/pool.go b/worker/pool.go index b77fa30..b7fdf39 100644 --- a/worker/pool.go +++ b/worker/pool.go @@ -3,6 +3,7 @@ package worker import ( "context" "errors" + "slices" "strings" "sync" "time" @@ -12,6 +13,7 @@ import ( "github.com/xraph/dispatch/ext" "github.com/xraph/dispatch/id" "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" ) // isTransientStoreErr reports whether err is a transient store failure — @@ -51,14 +53,31 @@ type QueueManager interface { // defaultStoreCallTimeout caps how long a single store roundtrip // (DequeueJobs, HeartbeatJob, ReapStaleJobs, UpdateJob) is allowed // to run before the worker abandons it. Without this, a slow Mongo / -// Postgres session selection would let dequeue calls stack on the -// shared driver pool until every connection is checked out — which -// is exactly the cascade that produces "context deadline exceeded" -// floods at boot. 5 seconds is generous enough for a healthy -// roundtrip and tight enough that 10 workers polling every second -// can't pile up more than 50 in-flight calls at once. +// Postgres session selection would let calls stack on the shared driver +// pool until every connection is checked out — which is exactly the +// cascade that produces "context deadline exceeded" floods at boot. +// DequeueJobs comes from the single fetchLoop poller (see its own doc +// comment below), not one call per worker goroutine, but +// HeartbeatJob/UpdateJob can still arrive from up to `concurrency` +// worker goroutines at once, each tending its own active job. 5 seconds +// is generous enough for a healthy roundtrip and tight enough that a +// pool of 10 concurrent workers can't pile up more than a few dozen +// in-flight calls at once. const defaultStoreCallTimeout = 5 * time.Second +// DefaultReapInterval is how often the pool scans for expired leases. +// +// It is deliberately independent of any lease TTL. The reaper used to +// tick at the stale-job threshold, so a five-minute threshold meant a +// dead job could sit for ten minutes before anyone looked; and with +// per-definition TTLs there is no single threshold left to tick at. +const DefaultReapInterval = 15 * time.Second + +// DefaultReclaimBatch caps how many expired leases one pass reclaims, so +// a backlog after an outage drains over several ticks instead of one +// statement that locks a large slice of the table. +const DefaultReclaimBatch = 100 + // Pool manages a set of concurrent worker goroutines fed by a single // fetcher that polls the store for jobs and executes them through the // Executor. @@ -83,6 +102,21 @@ type Pool struct { heartbeatInterval time.Duration staleJobThreshold time.Duration + // reapInterval is the reaper's scan cadence. Zero uses + // DefaultReapInterval; this is deliberately independent of + // staleJobThreshold and any lease TTL. See WithReapInterval. + reapInterval time.Duration + + // leaseStore is the store's optional lease capability. Nil means the + // backend implements only job.Store, and the pool keeps its previous + // behaviour: unleased claims, bare heartbeats, threshold reaping. + leaseStore job.LeaseStore + + // defaultLeaseTTL is how far a renewal pushes the expiry for jobs + // that declare no LeaseTTL of their own. Zero falls back to + // staleJobThreshold, then to job.DefaultLeaseTTL. + defaultLeaseTTL time.Duration + // storeCallTimeout caps how long a single store call (dequeue, // heartbeat, reap, update) may hold a driver-pool connection // before being abandoned. Zero means use defaultStoreCallTimeout; @@ -93,16 +127,26 @@ type Pool struct { // Queue manager (optional). queueManager QueueManager + // resources admits jobs against this worker's real capacity. Nil + // disables the resource model entirely: the pool offers no budget at + // dequeue and takes no lease, which is exactly how it behaved before + // the model existed. + resources resource.Manager + + // customKeys are the custom resource keys this worker advertises, + // overriding the keys derived from the manager's capacity. + customKeys []string + stopCh chan struct{} wakeCh chan struct{} // nudges the fetcher out of its idle backoff - jobCh chan *job.Job // hand-off from the fetcher to the workers + jobCh chan admitted // hand-off from the fetcher to the workers slots chan struct{} // free-worker tokens; capacity == concurrency cancelCtx context.Context // Cancelled on Stop to interrupt in-flight store operations cancelFunc context.CancelFunc // Cancels cancelCtx wg sync.WaitGroup mu sync.Mutex running bool - activeJobs map[string]context.CancelFunc + activeJobs map[string]*inflight activeMu sync.Mutex } @@ -119,7 +163,8 @@ func WithPoolQueues(queues []string) PoolOption { return func(p *Pool) { p.queues = queues } } -// WithPollInterval sets how often workers poll for new jobs. +// WithPollInterval sets how often the pool's single fetcher (fetchLoop) +// polls the store for new jobs. func WithPollInterval(d time.Duration) PoolOption { return func(p *Pool) { p.pollInterval = d } } @@ -137,18 +182,91 @@ func WithHeartbeatInterval(d time.Duration) PoolOption { } // WithStaleJobThreshold sets the threshold after which running jobs -// without a heartbeat are considered stale and reaped. A zero value -// disables stale job reaping. +// without a heartbeat are considered stale and reaped, on the legacy +// SELECT-then-UPDATE path a backend implementing only job.Store falls +// back to (reapStaleJobsLegacy). It has no effect on any first-party +// backend: every one of them (memory, mongo, postgres, redis, sqlite) +// implements job.LeaseStore, so reapStaleJobs routes to +// reclaimExpiredLeases instead, which reclaims purely by lease expiry +// (job.LeaseTTL / WithDefaultLeaseTTL) and never reads this value. A +// zero value disables reaping outright, on either path — see +// reapStaleJobs and reclaimExpiredLeases. func WithStaleJobThreshold(d time.Duration) PoolOption { return func(p *Pool) { p.staleJobThreshold = d } } +// WithReapInterval sets how often the reaper scans for expired leases / +// stale jobs. +// +// This is the scan cadence, not the lease duration — it does not control +// how long a lease survives without renewal. For that, see +// WithDefaultLeaseTTL and job.WithLeaseTTL. A zero value leaves +// DefaultReapInterval in place; it has no effect when WithStaleJobThreshold +// is zero, since that still disables reaping entirely. +func WithReapInterval(d time.Duration) PoolOption { + return func(p *Pool) { p.reapInterval = d } +} + +// resolvedReapInterval returns the configured reap interval, or +// DefaultReapInterval when unset. +func (p *Pool) resolvedReapInterval() time.Duration { + if p.reapInterval > 0 { + return p.reapInterval + } + + return DefaultReapInterval +} + // WithQueueManager sets the queue manager for rate limiting and // concurrency control. func WithQueueManager(m QueueManager) PoolOption { return func(p *Pool) { p.queueManager = m } } +// WithResourceManager makes the pool resource-aware. +// +// With a manager installed the fetcher offers its free capacity as the +// dequeue budget, so the store never hands this worker a job it cannot +// run, and every claimed job holds a lease for as long as it executes. +// Without one the pool passes an unbounded DequeueOpts and takes no +// leases — every backend skips its fit predicate and behaviour is +// identical to a pool that predates the resource model. +// +// Prefer engine.WithResourceManager to calling this directly. A manager +// makes admit able to stall the fetcher for up to one pollInterval per +// batch while it holds claimed, running-state, not-yet-heartbeating jobs, +// which puts pollInterval and staleJobThreshold into a relationship the +// pool cannot police: it is handed both as already-decided values and +// does not own the policy. engine.Build validates them together +// (checkReaperMargin) before it constructs a pool. Construct one here +// instead and that check does not run — a staleJobThreshold inside the +// stall lets the reaper reclaim a job this fetcher still holds, and the +// job runs twice. +func WithResourceManager(m resource.Manager) PoolOption { + return func(p *Pool) { p.resources = m } +} + +// WithWorkerCustomKeys sets the custom resource keys this worker offers +// at dequeue, overriding the keys derived from the resource manager's +// capacity. +// +// The derived default is usually what you want; this exists to narrow +// it, so a worker draining a device can stop attracting work for it +// without being reconfigured. Keep the list a subset of the manager's +// custom capacity: dequeue matches custom resources by key and never by +// quantity, so a key advertised here that the manager has no capacity +// for will pass the store's filter and then be refused locally, and the +// job will bounce back to pending on every attempt. +// +// Setting this without a resource manager makes the dequeue bounded on +// its own, which is a deliberate opt-in: the worker then claims only +// jobs whose custom keys it offers, with no quantity accounting at all. +// +// The slice is copied, so the caller keeps no handle on pool state. +func WithWorkerCustomKeys(keys []string) PoolOption { + return func(p *Pool) { p.customKeys = slices.Clone(keys) } +} + // WithStoreCallTimeout caps a single store roundtrip. Pass a positive // duration to override defaultStoreCallTimeout, zero to leave the // default in place, or a negative value to disable the timeout @@ -157,6 +275,40 @@ func WithStoreCallTimeout(d time.Duration) PoolOption { return func(p *Pool) { p.storeCallTimeout = d } } +// WithDefaultLeaseTTL sets how far each renewal pushes a job's lease +// expiry when the job declares no LeaseTTL of its own. +// +// This is a liveness window, not a time limit: it should be a small +// multiple of the heartbeat interval regardless of how long the work +// takes. When unset the pool uses the configured stale-job threshold, so +// an existing deployment sees the same reclamation timing it had before +// leases existed. +func WithDefaultLeaseTTL(d time.Duration) PoolOption { + return func(p *Pool) { p.defaultLeaseTTL = d } +} + +// leaseTTLFor resolves how far a renewal should push this job's lease: +// the job's own declaration, else the pool default, else the legacy +// stale-job threshold, else job.DefaultLeaseTTL. The threshold sits in +// the chain so a deployment that configured only StaleJobThreshold — +// which is every deployment predating leases — keeps its current timing. +// +// j may be nil: the initial grant at claim time has no job to consult +// yet, so it resolves straight to the pool-level fallbacks. +func (p *Pool) leaseTTLFor(j *job.Job) time.Duration { + if j != nil && j.LeaseTTL > 0 { + return j.LeaseTTL + } + if p.defaultLeaseTTL > 0 { + return p.defaultLeaseTTL + } + if p.staleJobThreshold > 0 { + return p.staleJobThreshold + } + + return job.DefaultLeaseTTL +} + // NewPool creates a worker pool. func NewPool( store job.Store, @@ -177,7 +329,7 @@ func NewPool( logger: logger, stopCh: make(chan struct{}), wakeCh: make(chan struct{}, 1), - activeJobs: make(map[string]context.CancelFunc), + activeJobs: make(map[string]*inflight), } for _, opt := range opts { opt(p) @@ -185,6 +337,12 @@ func NewPool( if p.maxPollInterval < p.pollInterval { p.maxPollInterval = p.pollInterval } + if ls, ok := store.(job.LeaseStore); ok { + p.leaseStore = ls + } else { + logger.Warn("store does not implement job.LeaseStore; " + + "per-definition lease TTLs and epoch fencing are disabled") + } return p } @@ -203,11 +361,11 @@ func (p *Pool) Wake() { func (p *Pool) WorkerID() id.WorkerID { return p.workerID } // Start launches the worker goroutines. It returns immediately. -func (p *Pool) Start(_ context.Context) error { +func (p *Pool) Start(ctx context.Context) error { p.mu.Lock() - defer p.mu.Unlock() if p.running { + p.mu.Unlock() return nil } p.running = true @@ -224,7 +382,7 @@ func (p *Pool) Start(_ context.Context) error { // one DequeueJobs call per cycle instead of `concurrency` concurrent // calls, which kept idle pools writing to the store every second and // could exhaust the shared driver pool on its own. - p.jobCh = make(chan *job.Job) + p.jobCh = make(chan admitted) p.slots = make(chan struct{}, p.concurrency) for range p.concurrency { p.slots <- struct{}{} @@ -249,9 +407,75 @@ func (p *Pool) Start(_ context.Context) error { go p.reaperLoop() } + p.mu.Unlock() + + // Sweep sandboxes this worker left behind across a restart before it + // takes new work. In-process reclaim is a no-op; an out-of-process rung + // would otherwise keep orphaned children or pods alive indefinitely. + // Best effort by design: a rung that cannot sweep must not stop the + // pool from running the jobs it can still execute. + // + // Run in the background, tracked by p.wg like every other pool + // goroutine, so Start returns immediately as documented even when a + // rung's Reclaim does real (and potentially slow) process or + // filesystem I/O. Outside p.mu: nothing else needs the lock held for + // this, and holding it here would block Stop and every other pool + // method on a sweep that has no bound of its own. + if p.executor != nil { + p.wg.Add(1) + go p.runReclaimSweep(ctx) + } + return nil } +// runReclaimSweep runs the executor reclaim sweep to completion or +// cancellation, whichever comes first, and reports it to p.wg. +// +// The sweep's context is cancelled by either the ctx Start was called with +// or Pool.Stop (via p.cancelCtx) — whichever fires first — so a blocked +// rung can always be interrupted: by the caller giving up on Start's ctx, +// or by the pool being stopped before the sweep finishes. +func (p *Pool) runReclaimSweep(ctx context.Context) { + defer p.wg.Done() + + sweepCtx, cancel := mergeDone(ctx, p.cancelCtx) + defer cancel() + + // Reclamation is best-effort cleanup, not a startup precondition: a + // rung that cannot sweep must not stop the pool from running the jobs + // it can still execute, so an error here is logged, never fatal. + if err := p.executor.Reclaim(sweepCtx, p.workerID); err != nil { + p.logger.Warn("executor reclaim failed", + log.String("worker_id", p.workerID.String()), + log.String("error", err.Error()), + ) + } +} + +// mergeDone returns a context cancelled as soon as either a or b is done, +// and a cancel func the caller must call once it no longer needs the +// merged context — otherwise the watcher goroutine backing it leaks until +// whichever of a or b is cancelled last. +func mergeDone(a, b context.Context) (context.Context, context.CancelFunc) { + merged, cancel := context.WithCancel(context.Background()) + + stopWatch := make(chan struct{}) + go func() { + select { + case <-a.Done(): + case <-b.Done(): + case <-stopWatch: + } + cancel() + }() + + return merged, func() { + cancel() + close(stopWatch) + } +} + // Stop signals all workers to stop and waits for them to finish. // If the context has a deadline, active jobs are cancelled when time runs out. func (p *Pool) Stop(ctx context.Context) error { @@ -337,7 +561,26 @@ func (p *Pool) fetchLoop() { } dqCtx, dqCancel := p.callCtx() - jobs, err := p.store.DequeueJobs(dqCtx, p.queues, held) + // The budget is recomputed every cycle rather than cached: jobs + // finish and reclaimers evict between polls, and a stale ceiling + // would either strand work or admit work this worker cannot run. + // With no resource manager configured both fields are empty, the + // opts are IsUnbounded, and every backend skips its fit predicate. + dqOpts := job.DequeueOpts{ + Queues: p.queues, + Limit: held, + Budget: p.dequeueBudget(), + CustomKeys: p.offeredCustomKeys(), + } + // The grant uses the POOL default TTL, not any per-job value: the + // store cannot apply per-job TTL arithmetic at claim time, and the + // grant only has to survive until the first renewal one heartbeat + // later — which then applies the job's real TTL. + if p.leaseStore != nil { + dqOpts.WorkerID = p.workerID + dqOpts.LeaseUntil = time.Now().UTC().Add(p.leaseTTLFor(nil)) + } + jobs, err := p.store.DequeueJobs(dqCtx, dqOpts) dqCancel() if err != nil { p.releaseSlots(held) @@ -358,6 +601,17 @@ func (p *Pool) fetchLoop() { continue } + // dispatched counts jobs actually handed to a worker, which is what + // paces the loop below — len(jobs) is not, because a refused job + // never reaches the blocking hand-off that used to do the pacing. + dispatched := 0 + + // One reclaim budget for the whole batch. Per job it would be + // batch size × the deadline, and the fetcher spends that stall + // holding claimed jobs that are already in running state and not + // yet heartbeating. + admitCtx, admitCancel := p.admissionBudget(len(jobs)) + for _, j := range jobs { // Check queue/tenant rate limit and concurrency. if p.queueManager != nil && !p.queueManager.Acquire(j.Queue, j.ScopeOrgID) { @@ -365,28 +619,75 @@ func (p *Pool) fetchLoop() { continue } + // Reserve local capacity between the claim and the hand-off, + // so no job reaches a worker without the resources it declared + // already accounted for. + lease, admitErr := p.admit(admitCtx, j) + if admitErr != nil { + p.releaseQueueSlot(j) + p.requeueLocalMisfit(j, admitErr) + + continue + } + + a := admitted{job: j, lease: lease} + select { - case p.jobCh <- j: + case p.jobCh <- a: held-- // The worker now owns this slot. + dispatched++ case <-p.stopCh: - p.requeueUndispatched(j) + p.abandon(a) p.releaseSlots(held) + admitCancel() + return case <-p.cancelCtx.Done(): - p.requeueUndispatched(j) + p.abandon(a) p.releaseSlots(held) + admitCancel() + return } } + + admitCancel() p.releaseSlots(held) - if len(jobs) > 0 { - // There may be more ready work; poll again immediately. + // Three cadences, because a poll has three outcomes and only one + // of them paces itself. + switch { + case dispatched > 0: + // Something went through the blocking hand-off, which throttles + // this loop to the rate work completes. There may be more ready + // work behind it, so poll again immediately. interval = p.pollInterval + continue + + case len(jobs) > 0: + // Rows came back and none could be dispatched: rate limited, or + // too big for what is free. Those return instantly, so the fast + // path above would spin the claim/requeue cycle as fast as the + // store can serve it — 713 dequeues in 500ms on the regression + // test, an UpdateJob on each, against a backlog nothing can run. + // + // The idle backoff below is the wrong pace too, and worse: the + // store returns the queue head, so a batch of unrunnable jobs is + // re-claimed a Limit at a time, and backing off exponentially + // leaves a runnable job sitting behind them for the whole ramp + // (8.5s in that same test, against 94ms at this cadence). The + // refused rows only clear by being claimed and pushed forward, + // so the loop has to keep claiming — just not faster than it was + // configured to poll. + interval = p.pollInterval + + default: + // Nothing ready at all. Back off; only a Wake or real work + // resets the cadence. + interval = min(interval*2, p.maxPollInterval) } - interval = min(interval*2, p.maxPollInterval) woken, ok := p.wait(interval) if !ok { return @@ -424,11 +725,31 @@ func (p *Pool) releaseSlots(n int) { } } +// clearJobAssignment clears the fields that record which worker holds a +// job's lease, so a job put back to pending does not keep looking +// assigned to a worker that no longer holds it — stale WorkerID and +// LeaseExpiresAt values an operator reads as "this pending job belongs to +// a worker" when it does not. +// +// ReclaimExpiredLeases is the canonical clearer and additionally bumps +// LeaseEpoch and EvictCount, which this helper deliberately leaves alone: +// those are lease-eviction bookkeeping for another worker reclaiming a +// job out from under its holder. Every caller here is the job's own +// current holder putting it back — a rate limit, a shutdown, a launch +// failure — not a reclamation, so there is no previous holder to fence. +func clearJobAssignment(j *job.Job) { + j.WorkerID = id.WorkerID{} + j.StartedAt = nil + j.HeartbeatAt = nil + j.LeaseExpiresAt = nil +} + // requeueRateLimited returns a job the queue manager refused to pending // with a small delay. func (p *Pool) requeueRateLimited(j *job.Job) { j.State = job.StatePending j.RunAt = time.Now().Add(p.pollInterval) + clearJobAssignment(j) updCtx, updCancel := p.callCtx() updateErr := p.store.UpdateJob(updCtx, j) updCancel() @@ -448,7 +769,7 @@ func (p *Pool) requeueUndispatched(j *job.Job) { defer cancel() j.State = job.StatePending j.RunAt = time.Now().UTC() - j.StartedAt = nil + clearJobAssignment(j) if err := p.store.UpdateJob(ctx, j); err != nil { p.logger.Warn("failed to return undispatched job to pending", log.String("job_id", j.ID.String()), @@ -457,8 +778,9 @@ func (p *Pool) requeueUndispatched(j *job.Job) { } } -// workerLoop executes jobs handed off by the fetcher, returning its slot -// token after each job. +// workerLoop executes jobs handed off by the fetcher. The slot token is +// returned by runJob's own defer, not here, so it comes back on every +// exit path the attempt can take. func (p *Pool) workerLoop() { defer p.wg.Done() @@ -468,18 +790,42 @@ func (p *Pool) workerLoop() { return case <-p.cancelCtx.Done(): return - case j := <-p.jobCh: - p.runJob(j) - p.slots <- struct{}{} + case a := <-p.jobCh: + p.runJob(a) } } } -func (p *Pool) runJob(j *job.Job) { +func (p *Pool) runJob(a admitted) { + // Everything this attempt holds — the worker slot, the queue/tenant + // token, and the resource lease — comes back through this one defer, + // so a handler that panics past a pool with no Recover middleware + // cannot leave capacity spoken for by a job that is no longer running. + defer p.finishJob(a) + + j := a.job + p.extensions.EmitJobStarted(p.cancelCtx, j) - ctx, cancel := context.WithCancel(p.cancelCtx) - p.trackJob(j.ID.String(), cancel) + ctx, cancel := context.WithCancelCause(p.cancelCtx) + p.trackJob(j.ID.String(), cancel, a.lease, j.LeaseEpoch, p.leaseTTLFor(j)) + + // A fenced terminal write needs the store's lease capability, this + // pool's worker ID, and the epoch j.LeaseEpoch carries from the + // claim — the same three values trackJob just recorded in inflight. + // Attaching them to ctx here, once, is what lets Runner.handleSuccess + // / scheduleRetry / sendToDLQ route through UpdateLeasedJob instead + // of the unfenced UpdateJob without Runner ever reaching back into + // Pool state. Nil leaseStore — a backend that implements only + // job.Store — attaches nothing, so Execute keeps calling UpdateJob + // exactly as before leases existed. + if p.leaseStore != nil { + ctx = withLeaseFence(ctx, leaseFence{ + store: p.leaseStore, + workerID: p.workerID, + epoch: j.LeaseEpoch, + }) + } execErr := p.executor.Execute(ctx, j) if execErr != nil { @@ -489,14 +835,6 @@ func (p *Pool) runJob(j *job.Job) { log.String("error", execErr.Error()), ) } - - p.untrackJob(j.ID.String()) - cancel() - - // Release the queue/tenant slot. - if p.queueManager != nil { - p.queueManager.Release(j.Queue, j.ScopeOrgID) - } } // heartbeatLoop periodically sends heartbeats for all active jobs. @@ -516,37 +854,95 @@ func (p *Pool) heartbeatLoop() { } } +// activeSnapshot is one job's liveness bookkeeping as of the moment +// sendHeartbeats took its snapshot of activeJobs: the epoch and TTL are +// copied out under the lock so the renewal loop below can run without +// holding it. +type activeSnapshot struct { + jobID string + leaseEpoch int + leaseTTL time.Duration +} + +// sendHeartbeats keeps every active job's ownership fresh in the store. +// +// When the store implements job.LeaseStore, this renews the job's lease +// with the epoch it was granted at claim time instead of sending a bare +// liveness heartbeat. A renewal that comes back job.ErrLeaseLost means +// another worker now holds the job — this one was paused past its lease +// and reclaimed — so the job is cancelled here, within one heartbeat +// interval, rather than left to run to completion racing its replacement. +// +// Without a lease-capable store this keeps calling HeartbeatJob exactly +// as before leases existed. func (p *Pool) sendHeartbeats() { p.activeMu.Lock() - jobIDs := make([]string, 0, len(p.activeJobs)) - for jobID := range p.activeJobs { - jobIDs = append(jobIDs, jobID) + snapshots := make([]activeSnapshot, 0, len(p.activeJobs)) + for jobID, rec := range p.activeJobs { + snapshots = append(snapshots, activeSnapshot{ + jobID: jobID, + leaseEpoch: rec.leaseEpoch, + leaseTTL: rec.leaseTTL, + }) } p.activeMu.Unlock() - for _, jobIDStr := range jobIDs { - parsedID, parseErr := id.ParseJobID(jobIDStr) + for _, snap := range snapshots { + parsedID, parseErr := id.ParseJobID(snap.jobID) if parseErr != nil { - p.logger.Warn("heartbeat: invalid job id", log.String("job_id", jobIDStr)) + p.logger.Warn("heartbeat: invalid job id", log.String("job_id", snap.jobID)) continue } + + if p.leaseStore == nil { + hbCtx, hbCancel := p.callCtx() + err := p.store.HeartbeatJob(hbCtx, parsedID, p.workerID) + hbCancel() + if err != nil { + p.logger.Warn("heartbeat failed", + log.String("job_id", snap.jobID), + log.String("error", err.Error()), + ) + } + continue + } + hbCtx, hbCancel := p.callCtx() - err := p.store.HeartbeatJob(hbCtx, parsedID, p.workerID) + renewErr := p.leaseStore.RenewLease( + hbCtx, parsedID, p.workerID, snap.leaseEpoch, time.Now().UTC().Add(snap.leaseTTL), + ) hbCancel() - if err != nil { - p.logger.Warn("heartbeat failed", - log.String("job_id", jobIDStr), - log.String("error", err.Error()), + if renewErr == nil { + continue + } + + if errors.Is(renewErr, job.ErrLeaseLost) { + p.logger.Warn("lease lost, cancelling job", + log.String("job_id", snap.jobID), ) + p.cancelJob(snap.jobID, job.ErrLeaseLost) + continue } + + // A transient store blip must not cancel a healthy job; only a + // definitive ErrLeaseLost above does that. + p.logger.Warn("lease renewal failed", + log.String("job_id", snap.jobID), + log.String("error", renewErr.Error()), + ) } } -// reaperLoop periodically reaps stale jobs whose heartbeat has expired. +// reaperLoop periodically returns jobs whose worker has gone away. +// +// What "gone away" means depends on the store. For one implementing +// job.LeaseStore this reclaims lapsed leases, and the heartbeat only +// matters because renewing the lease writes it; for any other store it +// falls back to the old heartbeat-age sweep. reapStaleJobs picks. func (p *Pool) reaperLoop() { defer p.wg.Done() - ticker := time.NewTicker(p.staleJobThreshold) + ticker := time.NewTicker(p.resolvedReapInterval()) defer ticker.Stop() for { @@ -559,7 +955,51 @@ func (p *Pool) reaperLoop() { } } +// reapStaleJobs reclaims jobs that have gone dark, through the store's +// atomic path when it has one, falling back to the legacy SELECT-then- +// UPDATE path for a backend that implements only job.Store. func (p *Pool) reapStaleJobs() { + if p.leaseStore != nil { + p.reclaimExpiredLeases() + + return + } + + p.reapStaleJobsLegacy() +} + +// reclaimExpiredLeases takes back jobs whose lease has lapsed. The store +// does the claim and the read in one statement, so unlike the legacy +// path two pools cannot both reset the same job. +func (p *Pool) reclaimExpiredLeases() { + reapCtx, reapCancel := p.callCtx() + reclaimed, err := p.leaseStore.ReclaimExpiredLeases(reapCtx, DefaultReclaimBatch) + reapCancel() + if err != nil { + if isTransientStoreErr(err) { + p.logger.Warn("reclaim expired leases transient error", log.String("error", err.Error())) + } else { + p.logger.Error("reclaim expired leases error", log.String("error", err.Error())) + } + + return + } + + for _, j := range reclaimed { + p.logger.Info("reclaimed expired lease", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + log.Int("evict_count", j.EvictCount), + log.Int("lease_epoch", j.LeaseEpoch), + ) + } + + if len(reclaimed) > 0 { + p.Wake() + } +} + +func (p *Pool) reapStaleJobsLegacy() { reapCtx, reapCancel := p.callCtx() stale, err := p.store.ReapStaleJobs(reapCtx, p.staleJobThreshold) reapCancel() @@ -579,9 +1019,7 @@ func (p *Pool) reapStaleJobs() { for _, j := range stale { j.State = job.StatePending j.RunAt = time.Now().UTC() - j.WorkerID = id.WorkerID{} // Clear the worker assignment. - j.HeartbeatAt = nil - j.StartedAt = nil + clearJobAssignment(j) updCtx, updCancel := p.callCtx() updateErr := p.store.UpdateJob(updCtx, j) @@ -608,23 +1046,65 @@ func (p *Pool) reapStaleJobs() { } } -func (p *Pool) trackJob(jobID string, cancel context.CancelFunc) { +// trackJob records one job's cancel func, resource lease, and job-lease +// bookkeeping so the heartbeat loop and shutdown path can act on it. +// +// leaseEpoch and leaseTTL are the job lease's epoch and renewal window, +// resolved once here at claim time rather than re-read from the store on +// every heartbeat. +func (p *Pool) trackJob( + jobID string, + cancel context.CancelCauseFunc, + lease resource.Lease, + leaseEpoch int, + leaseTTL time.Duration, +) { p.activeMu.Lock() - p.activeJobs[jobID] = cancel + p.activeJobs[jobID] = &inflight{ + cancel: cancel, + lease: lease, + leaseEpoch: leaseEpoch, + leaseTTL: leaseTTL, + } p.activeMu.Unlock() } -func (p *Pool) untrackJob(jobID string) { +// untrackJob removes and returns the in-flight record, or nil if the job +// was never tracked. The caller owns the record's cancel func from here. +func (p *Pool) untrackJob(jobID string) *inflight { p.activeMu.Lock() + defer p.activeMu.Unlock() + + rec, ok := p.activeJobs[jobID] + if !ok { + return nil + } + delete(p.activeJobs, jobID) - p.activeMu.Unlock() + + return rec } func (p *Pool) cancelActiveJobs() { p.activeMu.Lock() defer p.activeMu.Unlock() - for jobID, cancel := range p.activeJobs { + + for jobID, rec := range p.activeJobs { p.logger.Warn("cancelling active job", log.String("job_id", jobID)) - cancel() + rec.cancel(context.Canceled) } } + +// cancelJob cancels one in-flight job with a cause the executor and the +// handler can tell apart from a timeout or a shutdown. +func (p *Pool) cancelJob(jobID string, cause error) { + p.activeMu.Lock() + rec, ok := p.activeJobs[jobID] + p.activeMu.Unlock() + + if !ok { + return + } + + rec.cancel(cause) +} diff --git a/worker/pool_test.go b/worker/pool_test.go index f802da2..b8ae217 100644 --- a/worker/pool_test.go +++ b/worker/pool_test.go @@ -330,7 +330,8 @@ func newRecordingStore() *recordingStore { return rs } -func (r *recordingStore) DequeueJobs(ctx context.Context, queues []string, limit int) ([]*job.Job, error) { +func (r *recordingStore) DequeueJobs(ctx context.Context, opts job.DequeueOpts) ([]*job.Job, error) { + limit := opts.Limit r.dequeueCalls.Add(1) cur := r.inFlight.Add(1) defer r.inFlight.Add(-1) @@ -352,7 +353,7 @@ func (r *recordingStore) DequeueJobs(ctx context.Context, queues []string, limit break } } - return r.Store.DequeueJobs(ctx, queues, limit) + return r.Store.DequeueJobs(ctx, opts) } func setupRecordingPool(t *testing.T, rs *recordingStore, opts ...worker.PoolOption) (*worker.Pool, *job.Registry) { @@ -486,6 +487,12 @@ func TestPool_WakeResetsBackoff(t *testing.T) { // TestPool_ReaperWakesFetcher verifies that when the reaper resets a stale // job to pending, it wakes the fetcher so the job is retried immediately // instead of waiting out the inflated idle poll interval. +// +// WithReapInterval is set explicitly because the reaper's scan cadence is +// deliberately independent of WithStaleJobThreshold (see +// worker.DefaultReapInterval) — the two used to be the same value, but a +// test that relied on that coupling would time out against today's 15s +// default cadence. func TestPool_ReaperWakesFetcher(t *testing.T) { rs := newRecordingStore() pool, reg := setupRecordingPool(t, rs, @@ -493,6 +500,7 @@ func TestPool_ReaperWakesFetcher(t *testing.T) { worker.WithPollInterval(10*time.Millisecond), worker.WithMaxPollInterval(5*time.Second), worker.WithStaleJobThreshold(100*time.Millisecond), + worker.WithReapInterval(100*time.Millisecond), worker.WithPoolQueues([]string{"default"}), ) @@ -515,18 +523,24 @@ func TestPool_ReaperWakesFetcher(t *testing.T) { time.Sleep(1200 * time.Millisecond) // Inject a stale running job — a worker elsewhere died mid-execution. + // recordingStore embeds *memory.Store, which implements job.LeaseStore, + // so the pool reaps through reclaimExpiredLeases rather than the + // legacy heartbeat-threshold path; LeaseExpiresAt must be set (and in + // the past) for the injected job to be visible to that path — a zero + // LeaseExpiresAt reads as "never leased", not "expired". now := time.Now().UTC() old := now.Add(-time.Hour) j := &job.Job{ - ID: newTestJobID(), - Name: "reaped", - Queue: "default", - Payload: []byte(`{}`), - State: job.StateRunning, - MaxRetries: 3, - RunAt: old, - StartedAt: &old, - HeartbeatAt: &old, + ID: newTestJobID(), + Name: "reaped", + Queue: "default", + Payload: []byte(`{}`), + State: job.StateRunning, + MaxRetries: 3, + RunAt: old, + StartedAt: &old, + HeartbeatAt: &old, + LeaseExpiresAt: &old, } j.CreatedAt = old j.UpdatedAt = old diff --git a/worker/reclaim_test.go b/worker/reclaim_test.go new file mode 100644 index 0000000..f200e9d --- /dev/null +++ b/worker/reclaim_test.go @@ -0,0 +1,356 @@ +package worker_test + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/dlq" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/inproc" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" + "github.com/xraph/dispatch/worker" + + log "github.com/xraph/go-utils/log" +) + +// TestPool_ReapInterval is the regression guard for the bug this task +// fixes: the reaper used to tick at the stale-job threshold, so a 5-minute +// threshold meant a dead job could sit for up to 10 minutes before anyone +// looked. The scan cadence must be independent of that threshold. +func TestPool_ReapInterval(t *testing.T) { + tests := []struct { + name string + opts []worker.PoolOption + want time.Duration + }{ + { + name: "explicit WithReapInterval is honoured", + opts: []worker.PoolOption{worker.WithReapInterval(3 * time.Second)}, + want: 3 * time.Second, + }, + { + name: "stale-job threshold alone does not set the cadence", + opts: []worker.PoolOption{worker.WithStaleJobThreshold(5 * time.Minute)}, + want: worker.DefaultReapInterval, + }, + { + name: "nothing configured falls back to the default", + opts: nil, + want: worker.DefaultReapInterval, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + pool := worker.NewPool(memory.New(), nil, nil, log.NewNoopLogger(), tt.opts...) + + if got := pool.ReapInterval(); got != tt.want { + t.Errorf("ReapInterval() = %v, want %v", got, tt.want) + } + }) + } +} + +// TestPool_ReclaimOnce_ReturnsExpiredLeaseToPending proves the new atomic +// path resets a job whose lease has lapsed without touching its retry +// budget: losing a lease is infrastructure, not a handler failure, so +// charging it to RetryCount would DLQ a job that never once errored. +func TestPool_ReclaimOnce_ReturnsExpiredLeaseToPending(t *testing.T) { + ctx := context.Background() + s := memory.New() + + j := &job.Job{ + ID: id.NewJobID(), + Name: "reclaim-me", + Queue: "default", + State: job.StatePending, + MaxRetries: 3, + RetryCount: 1, + RunAt: time.Now().UTC(), + } + j.CreatedAt = time.Now().UTC() + j.UpdatedAt = j.CreatedAt + + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("enqueue error: %v", err) + } + + // Claim it with a lease already in the past, so the very next + // reclaim sweep picks it up as expired. + claimed, err := s.DequeueJobs(ctx, job.DequeueOpts{ + Queues: []string{"default"}, + Limit: 1, + WorkerID: id.NewWorkerID(), + LeaseUntil: time.Now().UTC().Add(-time.Minute), + }) + if err != nil { + t.Fatalf("dequeue error: %v", err) + } + if len(claimed) != 1 { + t.Fatalf("claimed = %d jobs, want 1", len(claimed)) + } + + pool := worker.NewPool(s, nil, nil, log.NewNoopLogger()) + pool.ReclaimOnce(ctx) + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("get job error: %v", err) + } + + if got.State != job.StatePending { + t.Errorf("State = %v, want %v", got.State, job.StatePending) + } + if got.EvictCount != 1 { + t.Errorf("EvictCount = %d, want 1", got.EvictCount) + } + if got.RetryCount != 1 { + t.Errorf("RetryCount = %d, want 1 (unchanged — a lost lease is not a handler failure)", got.RetryCount) + } +} + +// TestPool_ReclaimOnce_LegacyPathForCapabilityLessStore proves the legacy +// SELECT-then-UPDATE path still runs, without panicking, against a backend +// that implements only job.Store — not job.LeaseStore. +func TestPool_ReclaimOnce_LegacyPathForCapabilityLessStore(t *testing.T) { + s := storeOnly{Store: memory.New()} + + pool := worker.NewPool(s, nil, nil, log.NewNoopLogger(), + worker.WithStaleJobThreshold(time.Minute), + ) + + pool.ReclaimOnce(context.Background()) + t.Log("legacy reap path ran without panicking against a job.Store-only backend") +} + +// blockingReclaimExecutor is an exec.Executor whose Reclaim blocks until +// either the test releases it or its context is cancelled. It stands in for +// an out-of-process rung doing real (and potentially slow) process or +// filesystem I/O during the pool's startup sweep. +type blockingReclaimExecutor struct { + release chan struct{} + calledCh chan struct{} + + mu sync.Mutex + unblockedBy string // "release" or "ctx" +} + +func newBlockingReclaimExecutor() *blockingReclaimExecutor { + return &blockingReclaimExecutor{ + release: make(chan struct{}), + calledCh: make(chan struct{}), + } +} + +func (e *blockingReclaimExecutor) Name() string { return "blocking" } +func (e *blockingReclaimExecutor) Level() exec.Level { return exec.LevelProcess } + +func (e *blockingReclaimExecutor) Run(_ context.Context, _ *exec.Request) (*exec.Result, error) { + return &exec.Result{Status: exec.StatusOK}, nil +} + +// Reclaim blocks until release is closed or ctx is done, recording which +// one unblocked it so tests can tell a leaked block from an interrupted one. +func (e *blockingReclaimExecutor) Reclaim(ctx context.Context, _ id.WorkerID) error { + close(e.calledCh) + + select { + case <-e.release: + e.mu.Lock() + e.unblockedBy = "release" + e.mu.Unlock() + case <-ctx.Done(): + e.mu.Lock() + e.unblockedBy = "ctx" + e.mu.Unlock() + } + + return nil +} + +func (e *blockingReclaimExecutor) Close() error { return nil } + +func (e *blockingReclaimExecutor) waitCalled(t *testing.T, d time.Duration) { + t.Helper() + select { + case <-e.calledCh: + case <-time.After(d): + t.Fatal("Reclaim was never called") + } +} + +func (e *blockingReclaimExecutor) getUnblockedBy() string { + e.mu.Lock() + defer e.mu.Unlock() + return e.unblockedBy +} + +// erroringReclaimExecutor is an exec.Executor whose Reclaim always fails +// immediately, without blocking. +type erroringReclaimExecutor struct{} + +func (e *erroringReclaimExecutor) Name() string { return "erroring" } +func (e *erroringReclaimExecutor) Level() exec.Level { return exec.LevelProcess } + +func (e *erroringReclaimExecutor) Run(_ context.Context, _ *exec.Request) (*exec.Result, error) { + return &exec.Result{Status: exec.StatusOK}, nil +} + +func (e *erroringReclaimExecutor) Reclaim(context.Context, id.WorkerID) error { + return errors.New("sandbox cleanup failed") +} + +func (e *erroringReclaimExecutor) Close() error { return nil } + +// setupPoolWithExecutor builds a pool whose Runner is wired to an +// exec.Registry containing the in-process default plus extra, so a job +// with no declared isolation still runs (through in-process) while the +// pool's Reclaim sweep also visits extra. +func setupPoolWithExecutor(t *testing.T, extra exec.Executor) ( + *worker.Pool, *memory.Store, *job.Registry, +) { + t.Helper() + logger := log.NewNoopLogger() + s := memory.New() + reg := job.NewRegistry() + extensions := ext.NewRegistry(logger) + + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(extra) + + dlqSvc := dlq.NewService(s, s) + bo := backoff.NewConstant(10 * time.Millisecond) + runner := worker.NewRunner(reg, extensions, s, dlqSvc, bo, executors, logger) + + pool := worker.NewPool(s, runner, extensions, logger, + worker.WithPoolConcurrency(1), + worker.WithPollInterval(10*time.Millisecond), + worker.WithPoolQueues([]string{"default"}), + ) + + return pool, s, reg +} + +// TestPool_StartReturnsPromptlyWhenReclaimBlocks proves the defect: Start's +// doc comment says "it returns immediately", but before the fix the +// Reclaim sweep ran synchronously, inside p.mu, before Start returned. A +// rung whose Reclaim blocks (a subprocess rung doing real I/O, not the +// in-process no-op) would hang Start forever. +func TestPool_StartReturnsPromptlyWhenReclaimBlocks(t *testing.T) { + fake := newBlockingReclaimExecutor() + pool, _, _ := setupPoolWithExecutor(t, fake) + + done := make(chan error, 1) + go func() { + done <- pool.Start(context.Background()) + }() + + select { + case err := <-done: + if err != nil { + t.Fatalf("Start returned error: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Start did not return within 2s while Reclaim was blocked") + } + + // Let Stop interrupt the still-blocked sweep rather than leaking it + // past the end of the test. + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := pool.Stop(ctx); err != nil { + t.Fatalf("Stop error: %v", err) + } +} + +// TestPool_StopInterruptsBlockedReclaim verifies Stop can always interrupt +// a Reclaim sweep still in flight, and that it is released through context +// cancellation — not by the test's own release channel, which is never +// closed here — so the goroutine backing it does not leak past Stop. +func TestPool_StopInterruptsBlockedReclaim(t *testing.T) { + fake := newBlockingReclaimExecutor() + pool, _, _ := setupPoolWithExecutor(t, fake) + + if err := pool.Start(context.Background()); err != nil { + t.Fatalf("Start error: %v", err) + } + + // Make sure the sweep has actually started blocking before we stop. + fake.waitCalled(t, 2*time.Second) + + stopDone := make(chan error, 1) + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + stopDone <- pool.Stop(ctx) + }() + + select { + case err := <-stopDone: + if err != nil { + t.Fatalf("Stop returned error: %v", err) + } + case <-time.After(3 * time.Second): + t.Fatal("Stop did not complete while Reclaim was blocked") + } + + if got := fake.getUnblockedBy(); got != "ctx" { + t.Errorf("Reclaim unblocked by %q, want %q (context cancellation, not the release channel)", got, "ctx") + } +} + +// TestPool_ReclaimErrorDoesNotBlockProcessing verifies reclamation is +// best-effort: a Reclaim that returns an error must be logged and not stop +// the pool from starting and running jobs. +func TestPool_ReclaimErrorDoesNotBlockProcessing(t *testing.T) { + pool, s, reg := setupPoolWithExecutor(t, &erroringReclaimExecutor{}) + + var processed atomic.Bool + job.RegisterDefinition(reg, job.NewDefinition("reclaim-error-job", func(_ context.Context, _ struct{}) error { + processed.Store(true) + return nil + })) + + j := &job.Job{ + ID: id.NewJobID(), + Name: "reclaim-error-job", + Queue: "default", + Payload: []byte(`{}`), + State: job.StatePending, + MaxRetries: 3, + RunAt: time.Now().UTC(), + } + j.CreatedAt = time.Now().UTC() + j.UpdatedAt = j.CreatedAt + if err := s.EnqueueJob(context.Background(), j); err != nil { + t.Fatalf("enqueue error: %v", err) + } + + if err := pool.Start(context.Background()); err != nil { + t.Fatalf("Start error: %v", err) + } + + deadline := time.After(5 * time.Second) + for !processed.Load() { + select { + case <-deadline: + t.Fatal("timed out waiting for job to be processed despite the reclaim error") + default: + time.Sleep(10 * time.Millisecond) + } + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := pool.Stop(ctx); err != nil { + t.Fatalf("Stop error: %v", err) + } +} diff --git a/worker/requeue_assignment_test.go b/worker/requeue_assignment_test.go new file mode 100644 index 0000000..6913fa5 --- /dev/null +++ b/worker/requeue_assignment_test.go @@ -0,0 +1,146 @@ +package worker_test + +import ( + "context" + "testing" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" + "github.com/xraph/dispatch/worker" +) + +// assignedJob builds a job that looks like it is currently assigned to a +// worker: a running job with WorkerID, StartedAt, HeartbeatAt, and +// LeaseExpiresAt all set, mirroring what a real claim leaves on the row. +func assignedJob(name string) *job.Job { + started := time.Now().UTC().Add(-time.Minute) + leaseExpires := time.Now().UTC().Add(time.Hour) + + return &job.Job{ + ID: id.NewJobID(), + Name: name, + Queue: "default", + State: job.StateRunning, + MaxRetries: 3, + WorkerID: id.NewWorkerID(), + StartedAt: &started, + HeartbeatAt: &started, + LeaseExpiresAt: &leaseExpires, + } +} + +// wantAssignmentCleared fails the test unless every field that records a +// job's worker assignment has been cleared. This is the one property all +// four requeue paths tested below must share, even though each keeps its +// own distinct delay/retry behaviour. +func wantAssignmentCleared(t *testing.T, j *job.Job) { + t.Helper() + + if j.WorkerID != (id.WorkerID{}) { + t.Errorf("WorkerID = %s, want cleared", j.WorkerID) + } + if j.StartedAt != nil { + t.Errorf("StartedAt = %v, want nil", j.StartedAt) + } + if j.HeartbeatAt != nil { + t.Errorf("HeartbeatAt = %v, want nil", j.HeartbeatAt) + } + if j.LeaseExpiresAt != nil { + t.Errorf("LeaseExpiresAt = %v, want nil", j.LeaseExpiresAt) + } + if j.State != job.StatePending { + t.Errorf("State = %s, want %s", j.State, job.StatePending) + } +} + +// TestRequeueRateLimited_ClearsJobAssignment covers worker/pool.go's +// requeueRateLimited: before the fix it only touched State and RunAt, +// leaving a rate-limited job looking assigned to the worker that never +// even started it. +func TestRequeueRateLimited_ClearsJobAssignment(t *testing.T) { + store := newFakeJobStore() + pool := worker.NewPool(store, nil, nil, log.NewNoopLogger()) + + j := assignedJob("rate-limited") + pool.RequeueRateLimitedForTest(context.Background(), j) + + wantAssignmentCleared(t, j) +} + +// TestRequeueUndispatched_ClearsJobAssignment covers requeueUndispatched, +// the best-effort shutdown path for a claimed-but-not-started job. Before +// the fix it cleared StartedAt but left WorkerID, HeartbeatAt, and +// LeaseExpiresAt in place. +func TestRequeueUndispatched_ClearsJobAssignment(t *testing.T) { + store := newFakeJobStore() + pool := worker.NewPool(store, nil, nil, log.NewNoopLogger()) + + j := assignedJob("undispatched") + pool.RequeueUndispatchedForTest(j) + + wantAssignmentCleared(t, j) +} + +// TestReapStaleJobsLegacy_ClearsJobAssignment covers the legacy +// SELECT-then-UPDATE reap path a backend implementing only job.Store +// falls back to. Before the fix it cleared WorkerID/HeartbeatAt/StartedAt +// but not LeaseExpiresAt. +// +// storeOnly hides memory.Store's job.LeaseStore methods so the pool +// routes reapStaleJobs to the legacy path instead of +// reclaimExpiredLeases, exactly as it would against a real +// capability-less backend. +func TestReapStaleJobsLegacy_ClearsJobAssignment(t *testing.T) { + ctx := context.Background() + s := storeOnly{Store: memory.New()} + + j := assignedJob("stale-legacy") + if err := s.EnqueueJob(ctx, j); err != nil { + t.Fatalf("EnqueueJob: %v", err) + } + + pool := worker.NewPool(s, nil, nil, log.NewNoopLogger(), + worker.WithStaleJobThreshold(time.Second)) + pool.ReclaimOnce(ctx) + + got, err := s.GetJob(ctx, j.ID) + if err != nil { + t.Fatalf("GetJob: %v", err) + } + + wantAssignmentCleared(t, got) +} + +// TestRequeueAfterLaunchFailure_ClearsJobAssignment covers +// Runner.requeueAfterLaunchFailure. Before the fix it only touched State +// and RunAt, so a job whose sandbox failed to launch went back to +// pending still carrying the WorkerID and lease fields of the attempt +// that never ran. +func TestRequeueAfterLaunchFailure_ClearsJobAssignment(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition("launch.job", func(context.Context, struct{}) error { + return &exec.Error{Status: exec.StatusLaunchFailed, Msg: "boom"} + }).Register(reg) + + store := newFakeJobStore() + runner := worker.NewRunner( + reg, ext.NewRegistry(log.NewNoopLogger()), store, nil, + backoff.NewConstant(time.Millisecond), nil, log.NewNoopLogger(), + ) + + j := assignedJob("launch.job") + + err := runner.Execute(context.Background(), j) + if err == nil { + t.Fatal("Execute() = nil, want the launch-failure error") + } + + wantAssignmentCleared(t, j) +} diff --git a/worker/runner.go b/worker/runner.go new file mode 100644 index 0000000..e852e84 --- /dev/null +++ b/worker/runner.go @@ -0,0 +1,1223 @@ +// Package worker provides the job execution engine — a Runner that +// orchestrates a single job attempt through middleware and an exec.Executor, +// and a Pool that manages concurrent worker goroutines polling for jobs. +package worker + +import ( + "context" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/dlq" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/middleware" +) + +// maxLaunchAttempts bounds how many times one job may fail to launch on this +// worker before it is sent to the dead letter queue. +// +// A launch failure deliberately does not consume the retry budget: an image +// that will not pull says nothing about the work. But nothing else bounds it +// either — the job returns to pending at the first step of the backoff curve +// and is dequeued again about a second later, forever, costing a store write +// and a worker slot each time. Five attempts rides out the case this +// leniency exists for, a job reaching a worker that has not been deployed +// its handler yet, and stops a genuinely undeployable job from spinning. +const maxLaunchAttempts = 5 + +// launchAttemptTTL is how long a job's launch count survives without being +// touched. +// +// The count is normally deleted the moment the job succeeds, retries, or is +// dead-lettered here. A job requeued by this worker and then picked up by a +// worker that does have its handler never comes back for that deletion, so +// entries expire as well and the map cannot grow with every such job for the +// life of the process. +const launchAttemptTTL = 30 * time.Minute + +// scratchDirPrefix names every scratch directory prepareOutputDir +// creates. sweepStaleScratchDirs matches on it so Reclaim only ever +// removes directories this package itself created, never an unrelated +// entry that happens to share a temp root. +const scratchDirPrefix = "dispatch-out-" + +// staleScratchDirAge is how old a leftover scratch directory must be, +// on top of its owning PID no longer being alive, before +// sweepStaleScratchDirs removes it. It is a courtesy tie-break only — +// what actually protects a currently running sibling's directory is +// processAlive(pid) matching the PID sweepStaleScratchDirs parses out of +// the directory's own name (parseScratchDirPID), not age: a live process +// writing inside its directory does not advance the directory's own +// mtime, so age alone cannot tell a stale directory from one a live +// process is still using (see sweepStaleScratchDirs' own doc comment). +// This constant only narrows the sliver of a window between MkdirTemp +// creating the directory and the owning process actually beginning to +// use it, once ownership has already ruled the PID dead. +const staleScratchDirAge = time.Hour + +// errFenceLost marks a commit-outputs failure caused specifically by +// the lease fence, as opposed to an ordinary artifact-plane failure — +// see commitOutputs and terminalFor's classification of its error. +var errFenceLost = errors.New("dispatch/worker: lease fence lost") + +// errDuplicateOutputName marks two files under one OutputDir that would +// commit under the identical name — see collectOutputEntries. +var errDuplicateOutputName = errors.New("dispatch/worker: duplicate output name") + +// Runner executes a single job attempt: it selects an executor from the +// job's policy, runs the attempt through the middleware chain, then +// handles retry logic, DLQ push, state updates, and lifecycle events. +// +// Runner orchestrates the attempt. It does not itself invoke the handler — +// that is exec.Executor's job, which is what lets the same attempt run +// in-process or in a pod without this file changing. +type Runner struct { + registry *job.Registry + extensions *ext.Registry + store job.Store + dlqService *dlq.Service + backoff backoff.Strategy + executors *exec.Registry + mw middleware.Middleware + logger log.Logger + + // artifacts is the artifact plane an out-of-process rung's outputs + // are committed through, and PriorOutputs is resolved from. Nil + // leaves both off — see WithArtifacts. + artifacts *artifact.Service + + // scratchRoot is the directory an out-of-process attempt's scratch + // OutputDir is created under. Empty means os.TempDir(), resolved at + // request time rather than here so a change to the process's temp + // directory after construction still takes effect. + scratchRoot string + + // launchMu guards launches. One Runner is shared by every worker + // goroutine in the pool, so the counter is mutex-guarded rather than + // living on the job value. + launchMu sync.Mutex + launches map[string]launchAttempt +} + +// launchAttempt is one job's running launch-failure count on this worker. +type launchAttempt struct { + count int + seen time.Time +} + +// NewRunner creates a Runner with the given dependencies. +// +// A nil executors registry means handlers are called directly, which is +// the behaviour the deprecated NewExecutor preserves. +func NewRunner( + registry *job.Registry, + extensions *ext.Registry, + store job.Store, + dlqService *dlq.Service, + bo backoff.Strategy, + executors *exec.Registry, + logger log.Logger, + mws ...middleware.Middleware, +) *Runner { + return &Runner{ + registry: registry, + extensions: extensions, + store: store, + dlqService: dlqService, + backoff: bo, + executors: executors, + mw: middleware.Chain(mws...), + logger: logger, + launches: make(map[string]launchAttempt), + } +} + +// WithArtifacts configures the artifact plane an out-of-process rung +// commits its outputs through, and the root directory its scratch +// OutputDir is created under. It returns r so callers can chain it onto +// NewRunner. +// +// Never calling this, or passing a nil or disabled svc, leaves output +// committing off: an out-of-process attempt still gets a fresh, empty +// OutputDir that is removed once the attempt ends, but PriorOutputs +// stays empty and nothing the sandbox wrote is committed — exactly +// Runner's behaviour before this existed. An empty scratchRoot defaults +// to os.TempDir(). +func (r *Runner) WithArtifacts(svc *artifact.Service, scratchRoot string) *Runner { + r.artifacts = svc + r.scratchRoot = scratchRoot + + return r +} + +// Reclaim asks every configured executor to release sandboxes this worker +// leaked across a restart, and removes stale scratch output directories +// a previous process left behind. The pool calls it once at startup. +// +// Failures are joined rather than fatal: a rung that cannot sweep should not +// stop the worker from running the jobs it can still execute. +func (r *Runner) Reclaim(ctx context.Context, workerID id.WorkerID) error { + // Unconditional, and deliberately so: prepareOutputDir creates a + // scratch directory for ANY out-of-process attempt (executor.Level() + // > exec.LevelNone), whether or not this Runner has an artifact plane + // configured — see WithArtifacts' own doc comment. A Runner with no + // artifact plane still leaks exactly the same "dispatch-out-…" + // directories on a hard crash, so the sweep has to run regardless of + // r.artifacts to reclaim them; gating it the way commitOutputs is + // gated would leave those directories on disk forever. The sweep is + // safe to run unconditionally because it only ever touches entries + // matching scratchDirPrefix — nothing else sharing the scratch root + // is at risk. + r.sweepStaleScratchDirs() + + if r.executors == nil { + return nil + } + + var errs []error + for _, e := range r.executors.Executors() { + if err := e.Reclaim(ctx, workerID); err != nil { + errs = append(errs, fmt.Errorf("%s: %w", e.Name(), err)) + } + } + + return errors.Join(errs...) +} + +// Close releases every configured executor's own resources. +// +// An engine closes its own registry when it stops, since it outlives the +// pool; this is the equivalent for a caller that assembles a Runner itself. +// Call it after the pool has finished its in-flight attempts. +func (r *Runner) Close() error { + if r.executors == nil { + return nil + } + + var errs []error + for _, e := range r.executors.Executors() { + if err := e.Close(); err != nil { + errs = append(errs, fmt.Errorf("%s: %w", e.Name(), err)) + } + } + + return errors.Join(errs...) +} + +// Execute runs a job through the middleware chain and its executor, then +// routes the outcome through handleFailure/handleSuccess. The terminal +// write depends on which of several shapes the outcome takes: +// - Success: marks completed, emits JobCompleted. +// - A launch failure (the handler never ran) with launch attempts +// remaining: sets StatePending with a backoff delay and emits +// nothing — see requeueAfterLaunchFailure. It does not consume the +// job's own retry budget. +// - A launch failure that exhausts maxLaunchAttempts, or any failure +// marked permanent (dispatch.ErrPermanent, or exec.Error.Permanent +// for a rung that cannot carry a Go error chain): DLQs immediately — +// on the FIRST attempt if that is when permanence was discovered, +// not only "after exhausting retries." Marks failed, pushes to DLQ, +// emits JobFailed + JobDLQ. +// - An ordinary failure with retries remaining: marks retrying with +// backoff, emits JobRetrying. +// - An ordinary failure with retries exhausted: marks failed, pushes +// to DLQ, emits JobFailed + JobDLQ, same as the permanent case. +// - Any of the terminal writes above losing the race to +// job.ErrLeaseLost: abandonLostLease writes nothing to the store at +// all — the current lease holder's own write must stand untouched — +// and only emits JobFailed so extensions still observe the loss. +func (r *Runner) Execute(ctx context.Context, j *job.Job) error { + terminal, err := r.terminalFor(j) + if err != nil { + // terminalFor fails before the middleware chain — and therefore + // r.mw — ever runs, so returning err bare here, as this used to, + // bypasses handleFailure entirely: no LastError, no retry + // bookkeeping, no DLQ, nothing but a Debug log from Pool.runJob. + // The job sits at StateRunning until its lease expires, gets + // reaped, and is re-leased to try again — forever, for a policy + // that cannot become satisfiable no matter how many times it is + // retried. Routing it through handleFailure here is what turns + // that into a real terminal write; see terminalFor's own wrap of + // exec.Registry.Select's error with dispatch.ErrPermanent for why + // this then reaches the DLQ in one step rather than being retried. + now := time.Now().UTC() + j.UpdatedAt = now + + return r.handleFailure(ctx, j, err, now) + } + + start := time.Now() + execErr := r.mw(ctx, j, terminal) + elapsed := time.Since(start) + + now := time.Now().UTC() + j.UpdatedAt = now + + if execErr != nil { + return r.handleFailure(ctx, j, execErr, now) + } + + return r.handleSuccess(ctx, j, now, elapsed) +} + +// terminalFor builds the innermost handler for this job. +// +// Everything cross-cutting — recover, tracing, metrics, logging, scope, +// timeout, and artifact staging — wraps this closure, which is precisely +// why staging keeps running in the worker process and an out-of-process +// handler receives a directory rather than storage credentials. +func (r *Runner) terminalFor(j *job.Job) (middleware.Handler, error) { + if r.executors == nil { + handler, ok := r.registry.Get(j.Name) + if !ok { + return nil, fmt.Errorf("no handler registered for job %q", j.Name) + } + + return func(ctx context.Context) error { + return handler(ctx, j.Payload) + }, nil + } + + policy := r.registry.Policy(j.Name) + executor, err := r.executors.Select(policy) + if err != nil { + // A policy no configured executor satisfies cannot become + // satisfiable by retrying — the same reasoning already applied to + // exec.ErrInvalidRequest below, wrapped with dispatch.ErrPermanent + // so Execute's caller sends this straight to the DLQ instead of + // requeuing a job that will discover the identical + // misconfiguration on every future attempt. + return nil, fmt.Errorf("%w: dispatch/worker: select executor for job %q: %w", dispatch.ErrPermanent, j.Name, err) + } + + return func(ctx context.Context) error { + req := r.request(j, policy) + + // A rung above in-process gets a scratch directory to write its + // outputs into and, when the artifact plane is configured, every + // output an earlier attempt of this job already committed. Without + // the latter the sandbox's in-memory store has no notion of prior + // attempts and Existing/IfAbsent would answer "no" every time, + // silently redoing work a previous attempt already finished. + if executor.Level() > exec.LevelNone { + dir, cleanup, dirErr := r.prepareOutputDir(j) + if dirErr != nil { + return &exec.Error{Status: exec.StatusLaunchFailed, Msg: dirErr.Error()} + } + defer cleanup() + + req.OutputDir = dir + + if r.artifacts != nil && r.artifacts.Enabled() { + prior, priorErr := r.resolvePriorOutputs(ctx, j) + if priorErr != nil { + return &exec.Error{Status: exec.StatusLaunchFailed, Msg: priorErr.Error()} + } + + req.PriorOutputs = prior + } else { + r.logger.Debug("artifact plane disabled; running without prior outputs", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + ) + } + } + + res, runErr := executor.Run(ctx, req) + if runErr != nil { + // Run reserves its error return for launch failures: the handler + // never ran, so the retry budget must not pay for it. + // + // An invalid request is the exception. It is a caller programming + // error that will fail identically on every attempt, and since a + // launch failure never increments RetryCount it would requeue + // forever. Fail it permanently instead. + if errors.Is(runErr, exec.ErrInvalidRequest) { + return fmt.Errorf("%w: %w", dispatch.ErrPermanent, runErr) + } + + return &exec.Error{Status: exec.StatusLaunchFailed, Msg: runErr.Error()} + } + + // Commit what the sandbox actually left on disk before reporting + // the attempt as done. This runs ahead of the lease-fenced terminal + // write Execute makes afterward (see abandonLostLease), not gated + // on it. + // + // What actually protects against two workers each believing they + // hold the same job at the same RetryCount at once — a lease + // reclaim racing a zombie that has not yet noticed its lease + // expired — is CreateFenced: commitOutputs commits under this + // worker's lease epoch as the fence token when one is available, + // so two holders at the same nominal attempt resolve to different + // backend objects, never the same one. That protection does not + // depend on either holder knowing anything is wrong. + // + // commitOutputs ALSO checks context.Cause(ctx) for job.ErrLeaseLost + // specifically before committing anything, and rechecks it before + // every individual file. Be precise about what this catches and + // what it does not: the pool's heartbeat loop only sets that cause + // once a RenewLease call comes back with a *definitive* + // job.ErrLeaseLost (Pool.sendHeartbeats), which can lag the actual + // reclaim by up to one heartbeat interval, and — the case that + // matters most — never fires at all while this worker cannot reach + // the store, since a transient renewal error deliberately does not + // cancel a healthy job. A network partition, the canonical way a + // lease is actually lost, is exactly what this check cannot see. + // It is worth having anyway: it is nearly free, and it turns the + // common case (the heartbeat noticed before Run even returned) + // into an outright refusal to commit rather than relying solely on + // CreateFenced to make the collision harmless. Call it what it is + // — an opportunistic check for a signal the pool may already have + // — not a fence of its own. + // + // Only a genuinely failed attempt (res.Status != StatusOK) skips + // this outright. + if executor.Level() > exec.LevelNone && res.Status == exec.StatusOK { + if commitErr := r.commitOutputs(ctx, j, req); commitErr != nil { + if errors.Is(commitErr, errFenceLost) { + // This used to matter for safety: StatusLaunchFailed + // routed through requeueAfterLaunchFailure, which wrote + // via the plain, unfenced store.UpdateJob and could + // stomp whatever the real lease holder had already done + // to the row. requeueAfterLaunchFailure now goes + // through the same fenced updateJob as every other + // terminal write (and routes ErrLeaseLost to + // abandonLostLease identically), so either + // classification would be safe today. This still + // avoids StatusLaunchFailed anyway: the handler ran to + // completion here, so the ordinary retry path is the + // more honest description of what happened, not a + // workaround for the other one being unsafe. + return fmt.Errorf("dispatch/worker: job %s: %w", j.ID, commitErr) + } + + // An ordinary commit failure — a duplicate output name, + // a backend error — is an artifact-plane fault, not a + // verdict on the handler's own work: the handler already + // ran to completion and reported success. Routing it + // through StatusLaunchFailed keeps it off the job's real + // retry budget, since a non-idempotent handler should not + // pay for storage being unavailable, and bounds a + // deterministic failure (see collectOutputEntries' own + // duplicate-name check) at maxLaunchAttempts instead of + // burning the whole retry schedule on something retrying + // can never fix. + return &exec.Error{Status: exec.StatusLaunchFailed, Msg: commitErr.Error()} + } + } + + return res.Err() + }, nil +} + +// request builds the execution request for one attempt. +func (r *Runner) request(j *job.Job, policy exec.Policy) *exec.Request { + req := &exec.Request{ + JobID: j.ID, + Name: j.Name, + Payload: j.Payload, + Attempt: j.RetryCount, + // The fingerprint states which handler set this attempt was built + // for. An out-of-process rung compares it against the set it + // actually links, so a stale image running an old handler is a + // launch failure rather than a silent wrong answer. Computed per + // attempt rather than cached: registration is a startup activity + // in practice, but a stale fingerprint would disable exactly the + // check it exists to make. + Fingerprint: exec.Fingerprint(r.registry.Names()), + Policy: policy, + // The job's resolved enforcement ceiling (see job.WithResourceLimits + // and engine.resolveResources), so a rung that can enforce + // something per job — exec/subprocess.Executor maps + // resource.Memory to RLIMIT_AS — has the number to enforce. + // j.ResourceLimits is nil for the overwhelming majority of jobs + // today, which is indistinguishable from "declared nothing". + ResourceLimits: j.ResourceLimits, + ScopeAppID: j.ScopeAppID, + ScopeOrgID: j.ScopeOrgID, + } + if j.Timeout > 0 { + req.Deadline = time.Now().Add(j.Timeout) + } + + return req +} + +// scratchDirPattern returns the os.MkdirTemp pattern prepareOutputDir +// uses for one job's scratch directory. It embeds this process's PID +// between scratchDirPrefix and jobID so sweepStaleScratchDirs can later +// decide who might still be using a leftover directory by asking +// whether that PID is a live process, rather than guessing from age — +// see parseScratchDirPID and sweepStaleScratchDirs. +func scratchDirPattern(jobID string) string { + return fmt.Sprintf("%s%d-%s-", scratchDirPrefix, os.Getpid(), jobID) +} + +// parseScratchDirPID extracts the PID scratchDirPattern embedded in a +// scratch directory's name. ok is false for anything that does not +// match the expected shape — never one of ours, or corrupted — which +// sweepStaleScratchDirs treats identically to "not ours": left alone. +func parseScratchDirPID(name string) (pid int, ok bool) { + rest, hasPrefix := strings.CutPrefix(name, scratchDirPrefix) + if !hasPrefix { + return 0, false + } + + sep := strings.IndexByte(rest, '-') + if sep <= 0 { + return 0, false + } + + n, err := strconv.Atoi(rest[:sep]) + if err != nil || n <= 0 { + return 0, false + } + + return n, true +} + +// prepareOutputDir creates a fresh, empty scratch directory for one +// out-of-process attempt to write its outputs into, under r.scratchRoot +// — os.TempDir() when that is unset. +// +// The returned cleanup removes the directory and must be deferred by the +// caller regardless of how the attempt ends: a stray directory per +// out-of-process attempt would otherwise accumulate on disk for the life +// of the worker process. +func (r *Runner) prepareOutputDir(j *job.Job) (dir string, cleanup func(), err error) { + root := r.scratchRoot + if root == "" { + root = os.TempDir() + } + + dir, err = os.MkdirTemp(root, scratchDirPattern(j.ID.String())) + if err != nil { + return "", func() {}, fmt.Errorf("dispatch/worker: create output directory: %w", err) + } + + cleanup = func() { + if rmErr := os.RemoveAll(dir); rmErr != nil { + r.logger.Warn("failed to remove scratch output directory", + log.String("job_id", j.ID.String()), + log.String("dir", dir), + log.String("error", rmErr.Error()), + ) + } + } + + return dir, cleanup, nil +} + +// sweepStaleScratchDirs removes scratch directories prepareOutputDir +// left behind because the worker process that created them died before +// its own deferred cleanup ran. It is best-effort: a removal failure is +// logged, not returned, since one stuck directory must not stop Reclaim +// from doing the rest of what it does at startup. +// +// It runs regardless of whether this Runner has an artifact plane +// configured: prepareOutputDir creates a scratch directory for any +// out-of-process attempt independent of r.artifacts (only committing +// those outputs is gated on it — see WithArtifacts and commitOutputs), +// so a Runner without an artifact plane leaks the identical +// "dispatch-out-…" directories on a crash and needs the identical sweep. +// This is safe to run unconditionally, including for a Runner that never +// configured an out-of-process executor at all: only entries matching +// scratchDirPrefix with a well-formed embedded PID (parseScratchDirPID) +// are ever touched, so a scratch root this Runner never wrote into is +// left alone by construction, not by a config check. +// +// Ownership, not age, is what decides whether a directory is touched: +// a directory whose embedded PID (see parseScratchDirPID) belongs to a +// process that is still alive — this one or a sibling sharing the same +// scratch root — is left alone regardless of how old it looks, because +// a live process rewriting the files inside it does not advance the +// directory's own mtime. Age is only a courtesy tie-break once the +// owning PID is confirmed gone, against the sliver of a window between +// MkdirTemp creating the directory and the owning process actually +// beginning to use it. +func (r *Runner) sweepStaleScratchDirs() { + root := r.scratchRoot + if root == "" { + root = os.TempDir() + } + + entries, err := os.ReadDir(root) + if err != nil { + // Best-effort: an unreadable or (already-gone) scratch root is + // not something Reclaim should fail startup over. + return + } + + cutoff := time.Now().Add(-staleScratchDirAge) + + for _, entry := range entries { + if !entry.IsDir() { + continue + } + + pid, ok := parseScratchDirPID(entry.Name()) + if !ok { + continue // never one of ours + } + + if processAlive(pid) { + continue + } + + info, infoErr := entry.Info() + if infoErr != nil || info.ModTime().After(cutoff) { + continue + } + + stale := filepath.Join(root, entry.Name()) + if rmErr := os.RemoveAll(stale); rmErr != nil { + r.logger.Warn("failed to remove stale scratch directory", + log.String("dir", stale), + log.String("error", rmErr.Error()), + ) + } + } +} + +// resolvePriorOutputs returns one PriorOutput per name any earlier +// attempt of j already committed, keeping the highest-attempt link when +// more than one attempt produced the same name — the same tie-break +// FindLinkByName applies for a single-name lookup. +// +// This is the worker-side half of PriorOutputs (see exec.PriorOutput): +// an out-of-process rung's artifact store is in-memory and local to one +// attempt, with no notion of earlier ones, so without this a retried +// handler's Existing/IfAbsent check would answer "no" every time and +// quietly redo work a previous attempt had already finished. +func (r *Runner) resolvePriorOutputs(ctx context.Context, j *job.Job) ([]exec.PriorOutput, error) { + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: j.ID.String()} + + links, err := r.artifacts.Store().ListLinks(ctx, owner) + if err != nil { + return nil, fmt.Errorf("dispatch/worker: list prior links for job %s: %w", j.ID, err) + } + + best := make(map[string]*artifact.Link, len(links)) + for _, link := range links { + if link.Role != artifact.RoleOutput { + continue + } + + if cur, ok := best[link.Name]; !ok || link.Attempt > cur.Attempt { + best[link.Name] = link + } + } + + if len(best) == 0 { + return nil, nil + } + + // Sorted so the request a given job history produces is deterministic + // rather than following map iteration order. + names := make([]string, 0, len(best)) + for name := range best { + names = append(names, name) + } + sort.Strings(names) + + prior := make([]exec.PriorOutput, 0, len(names)) + for _, name := range names { + a, getErr := r.artifacts.Get(ctx, best[name].ArtifactID) + if getErr != nil { + return nil, fmt.Errorf("dispatch/worker: resolve prior output %q for job %s: %w", name, j.ID, getErr) + } + + prior = append(prior, exec.PriorOutput{Name: name, Ref: a.Ref()}) + } + + return prior, nil +} + +// outputEntry is one regular file collectOutputEntries found on disk, +// ready to commit under name. +type outputEntry struct { + name string + path string +} + +// commitOutputs commits the regular files the sandbox actually left in +// req.OutputDir through the artifact service, linking each to j as an +// output of this attempt. +// +// It is driven entirely by what collectOutputEntries finds really on +// disk, never by anything the sandbox itself reported: a claim crossed +// a process boundary a compromised handler fully controls, so nothing +// about it — a name, a size, a hash, a content type — is evidence that +// anything actually landed anywhere. A handler that claims an output it +// never wrote gets no artifact row for it, because nothing here ever +// reads such a claim to decide what to commit. +// +// If the artifact plane is disabled, this logs once at debug and does +// nothing: the sandbox's outputs are discarded along with the scratch +// directory, exactly as they were before out-of-process committing +// existed. +func (r *Runner) commitOutputs(ctx context.Context, j *job.Job, req *exec.Request) error { + if r.artifacts == nil || !r.artifacts.Enabled() { + r.logger.Debug("artifact plane disabled; not committing sandbox outputs", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + ) + + return nil + } + + // Checked before anything else, and specifically for job.ErrLeaseLost + // — not any cancellation. ctx is also cancelled for reasons that have + // nothing to do with who owns the job (a graceful shutdown mid-commit + // being the obvious one), and misreporting THAT as "lease fence lost" + // would be actively wrong: it is not what happened, and it changes + // how the resulting failure gets classified below. Only the specific, + // known signal the pool's heartbeat loop sets — job.ErrLeaseLost, the + // moment RenewLease comes back definitive (Pool.sendHeartbeats / + // cancelJob) — is treated as a fence loss here. See the longer + // comment above this closure for what this check does and does not + // cover; it is read-only regardless: nothing here renews or rewrites + // anything lease_fence.go or the pool itself owns. + if cause := context.Cause(ctx); errors.Is(cause, job.ErrLeaseLost) { + return fmt.Errorf("%w: %w", errFenceLost, cause) + } + + entries, err := collectOutputEntries(req.OutputDir) + if err != nil { + return err + } + + if len(entries) == 0 { + return nil + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: j.ID.String()} + + return r.commitOutputEntries(ctx, owner, j.RetryCount, fenceToken(ctx), entries) +} + +// collectOutputEntries walks dir for the regular files a sandbox left +// behind, returning one outputEntry per unique base name in a +// deterministic order. +// +// Non-regular entries — symlinks, FIFOs, sockets, devices — are skipped +// without ever being opened. WalkDir reports each entry's type from an +// Lstat taken at listing time, so a symlink is identified and skipped +// here, before any path derived from it is ever handed to a file-open +// call anywhere in this package. Opening what a symlink resolves to +// would let a compromised handler point one at anything this worker +// process can read — its own config, cloud credentials, a mounted +// service-account token — and have the bytes published as an ordinary +// job output; opening a FIFO with no writer on the other end blocks +// forever, wedging a worker goroutine, which is just as fatal on a +// smaller scale. Both are excluded by the same type check, which is +// what makes it sufficient on its own: a directory-entry check alone +// (d.IsDir()) catches neither, since both report false for it. +// +// Dot-prefixed files are a rung's own uncommitted temp files (see +// exec/shim's LocalFS) or otherwise hidden by convention. A dot-prefixed +// directory is skipped with fs.SkipDir specifically, not a bare nil: +// WalkDir descends into a directory regardless of what the callback +// returns for it unless told SkipDir, so returning nil for a hidden +// directory would still walk — and still commit — whatever non-hidden +// files happen to live inside it. +// +// Two files at different paths sharing one base name are reported as +// errDuplicateOutputName rather than letting the second silently +// resolve to the same committed name as the first: Create's own name +// parameter may not contain a path separator, so any nested directory +// structure under OutputDir is necessarily flattened to its leaf name +// by the time it reaches the artifact plane, and two leaves colliding +// is a structural problem this function must surface, not paper over +// by committing whichever one the walk happened to visit last. +// +// dir not existing at all is a hard error, never an empty result. This +// function's only caller creates dir itself before the sandbox ever +// runs (see prepareOutputDir), so the ONLY way it can be missing here +// is something having removed it after the fact — the sandbox +// deleting its own OutputDir, or, before sweepStaleScratchDirs was +// made ownership-based, another process's sweep colliding with a +// still-running attempt. A vanished output directory is never a +// legitimate "the handler produced nothing": that is an empty, +// EXISTING directory, which WalkDir already reports as zero entries +// with no error, handled below without reaching this branch at all. +func collectOutputEntries(dir string) ([]outputEntry, error) { + var entries []outputEntry + seenAt := make(map[string]string) + + walkErr := filepath.WalkDir(dir, func(path string, d fs.DirEntry, walkEntryErr error) error { + if walkEntryErr != nil { + return walkEntryErr + } + + if d.IsDir() { + if strings.HasPrefix(d.Name(), ".") { + return fs.SkipDir + } + + return nil + } + + if strings.HasPrefix(d.Name(), ".") { + return nil + } + + if !d.Type().IsRegular() { + return nil + } + + name := d.Name() + if prior, dup := seenAt[name]; dup { + return fmt.Errorf("dispatch/worker: %q and %q would both commit as output %q: %w", + prior, path, name, errDuplicateOutputName) + } + seenAt[name] = path + + entries = append(entries, outputEntry{name: name, path: path}) + + return nil + }) + if walkErr != nil { + if errors.Is(walkErr, fs.ErrNotExist) { + return nil, fmt.Errorf("dispatch/worker: output directory %q no longer exists: %w", dir, walkErr) + } + + return nil, fmt.Errorf("dispatch/worker: walk output directory: %w", walkErr) + } + + // Sorted so a partial commit's boundary — which entries land before a + // later one fails, and so which ones commitOutputFile's own + // FindCommitted check will recognise as already-done on a retry — is + // a deterministic function of name, not of the filesystem's own + // directory-listing order, which is unspecified and can vary between + // platforms or even between runs on the same one. + sort.Slice(entries, func(i, k int) bool { return entries[i].name < entries[k].name }) + + return entries, nil +} + +// fenceToken returns the lease epoch ctx carries, stringified, or "" if +// ctx carries no fence at all — a bare Runner driven without a Pool, or +// a store that does not implement job.LeaseStore. It is read-only: this +// neither renews nor otherwise touches anything leaseFenceFromContext's +// own package (lease_fence.go) owns. +func fenceToken(ctx context.Context) string { + fence, ok := leaseFenceFromContext(ctx) + if !ok { + return "" + } + + return strconv.Itoa(fence.epoch) +} + +// commitOutputEntries commits each entry through the artifact service +// under token — see artifact.Service.CreateFenced — checking the lease +// fence again before every individual commit. +// +// A failure partway through leaves whatever already landed in place: +// nothing here rolls a prior success back. That is deliberate, not an +// omission — an earlier version of this function did roll back, by +// deleting the backend bytes it had just written, which left the +// artifact row's link behind pointing at nothing. resolvePriorOutputs, +// FindExisting, and IfAbsent all read that link, not the bytes, so a +// retried handler consulting PriorOutputs was told "already done, skip +// regenerating it" for data that had just been deleted. Leaving a +// partial success in place instead, and giving commitOutputFile +// (FindCommitted) a way to recognise its own earlier work as a no-op +// rather than a collision, is what a retry actually converges against. A +// losing attempt may therefore end this call with some entries committed +// and others not; what matters is that every name it DOES leave +// committed stays valid and reusable, by this same attempt's own retry +// or by resolvePriorOutputs on a later one. +func (r *Runner) commitOutputEntries( + ctx context.Context, + owner artifact.OwnerRef, + attempt int, + token string, + entries []outputEntry, +) error { + for _, entry := range entries { + if cause := context.Cause(ctx); errors.Is(cause, job.ErrLeaseLost) { + return fmt.Errorf("%w: %w", errFenceLost, cause) + } + + if _, err := r.commitOutputFile(ctx, owner, attempt, token, entry.name, entry.path); err != nil { + return err + } + } + + return nil +} + +// commitOutputFile reads one file collectOutputEntries found on disk +// and commits its actual bytes through the artifact service, returning +// the ref it was recorded under. The size and content type the +// resulting artifact row carries both come from what the backend +// actually saw pass through it while committing these exact bytes — +// nothing here is influenced by anything the sandbox itself claimed +// about its outputs. There is no hash: artifact.ObjectInfo carries no +// content hash for Commit to record, and CommitWriter.Commit never +// populates Artifact.ContentHash on this path — that field is only ever +// assigned on the input-staging path (artifact/staging/middleware.go), +// not here. +// +// It checks FindCommitted first and treats a hit as a no-op success +// rather than re-writing anything: a commit failure is classified as +// StatusLaunchFailed (see terminalFor), which — like any other launch +// failure — requeues without advancing RetryCount, so a retry of it +// reuses the identical (owner, attempt) and, once a lease epoch is +// available, the identical fence token. Without this check, a name an +// earlier, partially-failed pass of the SAME attempt already committed +// would collide with itself on every subsequent retry — forever, on a +// store with no lease grants at all, since nothing there ever changes +// the key a retry computes. Recognising it as already-done instead +// makes a retry converge once whatever caused the original failure +// clears, regardless of whether the handler itself is careful enough +// to consult PriorOutputs and skip regenerating it. +func (r *Runner) commitOutputFile( + ctx context.Context, + owner artifact.OwnerRef, + attempt int, + token string, + name, path string, +) (artifact.Ref, error) { + if existing, err := r.artifacts.FindCommitted(ctx, owner, attempt, name, token); err == nil { + return existing, nil + } else if !errors.Is(err, artifact.ErrNotFound) { + return artifact.Ref{}, fmt.Errorf("dispatch/worker: check existing output %q: %w", name, err) + } + + f, err := openRegularNoFollow(path) + if err != nil { + return artifact.Ref{}, fmt.Errorf("dispatch/worker: open output %q: %w", name, err) + } + defer f.Close() + + // A second, TOCTOU-closing layer behind collectOutputEntries' own + // Lstat-based filter (see its doc comment): openRegularNoFollow + // already refuses to follow a symlink at the final path component on + // platforms that support it, and this confirms what was actually + // opened is still a plain regular file even so — catching the entry + // that was one when listed but has since become something else. + info, statErr := f.Stat() + if statErr != nil { + return artifact.Ref{}, fmt.Errorf("dispatch/worker: stat output %q: %w", name, statErr) + } + + if !info.Mode().IsRegular() { + return artifact.Ref{}, fmt.Errorf("dispatch/worker: output %q is no longer a regular file", name) + } + + var w *artifact.CommitWriter + if token != "" { + w, err = r.artifacts.CreateFenced(ctx, owner, attempt, name, token) + } else { + w, err = r.artifacts.Create(ctx, owner, attempt, name) + } + if err != nil { + return artifact.Ref{}, fmt.Errorf("dispatch/worker: create output %q: %w", name, err) + } + + if _, copyErr := io.Copy(w, f); copyErr != nil { + _ = w.Abort() //nolint:errcheck // best-effort cleanup; the write error below is what the caller acts on + + return artifact.Ref{}, fmt.Errorf("dispatch/worker: write output %q: %w", name, copyErr) + } + + ref, err := w.Commit(ctx) + if err != nil { + return artifact.Ref{}, fmt.Errorf("dispatch/worker: commit output %q: %w", name, err) + } + + return ref, nil +} + +// updateJob persists j's terminal state, fenced on the lease this worker +// held at claim time whenever ctx carries one — see withLeaseFence. +// +// A store that does not implement job.LeaseStore, or a Runner driven +// without a Pool, never sees a fence attached, so this falls back to the +// original unfenced UpdateJob: that path must not become a hard +// requirement of using Runner at all. +// +// job.ErrLeaseLost is returned to the caller exactly like any other +// error here — this method does no special-casing of it. Every call +// site does, through abandonLostLease, because ErrLeaseLost is not "the +// write failed, log and propagate," it is "someone else owns this job +// now, stop." +func (r *Runner) updateJob(ctx context.Context, j *job.Job) error { + if fence, ok := leaseFenceFromContext(ctx); ok { + return fence.store.UpdateLeasedJob(ctx, j, fence.workerID, fence.epoch) + } + + return r.store.UpdateJob(ctx, j) +} + +// abandonLostLease is what every fenced terminal write does on +// job.ErrLeaseLost: the lease moved on before this write landed, so the +// job is running under a different worker's epoch now and this attempt +// has no coherent claim left to make about it. +// +// It does not retry, requeue, or DLQ — both of those write, and the +// winner's outcome must stand untouched. The handler's own side effects +// need no cleanup here either: they already commit under attempt-scoped +// ephemeral artifact keys and carry a link, so SweepOrphans — which only +// ever considers link-less artifacts — never sees them. What collects a +// losing attempt's outputs is SweepEphemeral's owner-terminal path, once +// every owner that links them has gone terminal. +// +// The extension registry emit reuses EmitJobFailed rather than adding a +// new event: audit_hook and relay_hook both already implement +// ext.JobFailed, so they observe a lost lease with no new plumbing. +func (r *Runner) abandonLostLease(ctx context.Context, j *job.Job, cause error) error { + r.logger.Warn("lease lost, discarding terminal write", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + ) + + r.extensions.EmitJobFailed(ctx, j, cause) + + return cause +} + +// handleSuccess marks the job as completed and emits the lifecycle event. +func (r *Runner) handleSuccess(ctx context.Context, j *job.Job, now time.Time, elapsed time.Duration) error { + r.forgetLaunchFailures(j.ID.String()) + + j.State = job.StateCompleted + j.CompletedAt = &now + + if updateErr := r.updateJob(ctx, j); updateErr != nil { + if errors.Is(updateErr, job.ErrLeaseLost) { + return r.abandonLostLease(ctx, j, updateErr) + } + + r.logger.Error("failed to update job after success", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + log.String("error", updateErr.Error()), + ) + return updateErr + } + + r.extensions.EmitJobCompleted(ctx, j, elapsed) + return nil +} + +// handleFailure either requeues the job or increments the retry counter and +// retries, depending on whether the failure was the work's fault. +// +// A failure marked dispatch.ErrPermanent skips the remaining attempts. The +// retry schedule exists to outlast a transient fault, and spending it on a +// condition that cannot change wastes worker time proportional to the backoff +// curve: a job whose input was deleted would otherwise rediscover that the +// object is still gone once per attempt, minutes to hours apart, before +// arriving at the same dead letter queue it could have reached immediately. +func (r *Runner) handleFailure(ctx context.Context, j *job.Job, handlerErr error, now time.Time) error { + j.LastError = handlerErr.Error() + + var execErr *exec.Error + isExecErr := errors.As(handlerErr, &execErr) + + // A launch failure means the handler never ran: an image that would + // not pull, an exhausted quota, a missing runtime. Consuming the + // retry budget for it would let one bad node send healthy work to + // the DLQ, so the job is requeued without counting the attempt — + // up to a bound, since nothing else stops a job that can never launch + // from requeueing itself forever. + if isExecErr && !execErr.Status.CountsAgainstRetries() { + if n := r.recordLaunchFailure(j.ID.String(), now); n > maxLaunchAttempts { + capped := fmt.Errorf("%w: job %s failed to launch %d times: %s", + dispatch.ErrPermanent, j.Name, n, j.LastError) + j.LastError = capped.Error() + + return r.sendToDLQ(ctx, j, capped) + } + + return r.requeueAfterLaunchFailure(ctx, j, now) + } + + j.RetryCount++ + + // Permanence reaches here two ways. In-process the handler's own error + // chain survives, so errors.Is finds the sentinel. Out of process it + // cannot, so the rung sets a flag on the Result instead. Both mean the + // retry schedule would only rediscover the same condition. + if errors.Is(handlerErr, dispatch.ErrPermanent) || (isExecErr && execErr.Permanent) { + r.logger.Info("job failed permanently, skipping remaining retries", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + log.Int("retry_count", j.RetryCount), + log.Int("max_retries", j.MaxRetries), + log.String("error", handlerErr.Error()), + ) + + return r.sendToDLQ(ctx, j, handlerErr) + } + + if j.RetryCount <= j.MaxRetries { + return r.scheduleRetry(ctx, j, now) + } + + return r.sendToDLQ(ctx, j, handlerErr) +} + +// recordLaunchFailure counts one launch failure for a job and returns the +// running total. It is safe for concurrent use. +func (r *Runner) recordLaunchFailure(jobID string, now time.Time) int { + r.launchMu.Lock() + defer r.launchMu.Unlock() + + // Expire stale entries first. A job requeued here and then run + // elsewhere never returns for its deletion, and the sweep is what keeps + // the map bounded in that case. It runs only on launch failures, which + // are rare, over a map that holds only jobs currently failing to launch. + for key, attempt := range r.launches { + if now.Sub(attempt.seen) > launchAttemptTTL { + delete(r.launches, key) + } + } + + attempt := r.launches[jobID] + attempt.count++ + attempt.seen = now + r.launches[jobID] = attempt + + return attempt.count +} + +// forgetLaunchFailures drops a job's launch count. Every path that ends an +// attempt for good — success, retry, dead letter — calls it, so the map +// holds only jobs that are currently failing to launch. +func (r *Runner) forgetLaunchFailures(jobID string) { + r.launchMu.Lock() + defer r.launchMu.Unlock() + delete(r.launches, jobID) +} + +// requeueAfterLaunchFailure returns the job to pending with a backoff +// delay derived from the retry count without advancing it. +func (r *Runner) requeueAfterLaunchFailure(ctx context.Context, j *job.Job, now time.Time) error { + delay := r.backoff.Delay(j.RetryCount + 1) + j.RunAt = now.Add(delay) + j.State = job.StatePending + clearJobAssignment(j) + + if updateErr := r.updateJob(ctx, j); updateErr != nil { + if errors.Is(updateErr, job.ErrLeaseLost) { + return r.abandonLostLease(ctx, j, updateErr) + } + + r.logger.Error("failed to requeue job after launch failure", + log.String("job_id", j.ID.String()), + log.String("error", updateErr.Error()), + ) + + return updateErr + } + + r.logger.Warn("sandbox launch failed; requeued without consuming a retry", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + log.String("error", j.LastError), + log.Duration("delay", delay), + ) + + return fmt.Errorf("job %s launch failed: %s", j.Name, j.LastError) +} + +// scheduleRetry sets the job to StateRetrying with a backoff delay. +func (r *Runner) scheduleRetry(ctx context.Context, j *job.Job, now time.Time) error { + // The attempt reached the handler, so whatever launch trouble this job + // had is behind it. + r.forgetLaunchFailures(j.ID.String()) + + delay := r.backoff.Delay(j.RetryCount) + nextRunAt := now.Add(delay) + j.RunAt = nextRunAt + j.State = job.StateRetrying + + if updateErr := r.updateJob(ctx, j); updateErr != nil { + if errors.Is(updateErr, job.ErrLeaseLost) { + return r.abandonLostLease(ctx, j, updateErr) + } + + r.logger.Error("failed to update job for retry", + log.String("job_id", j.ID.String()), + log.String("error", updateErr.Error()), + ) + return updateErr + } + + r.extensions.EmitJobRetrying(ctx, j, j.RetryCount, nextRunAt) + + r.logger.Info("job scheduled for retry", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + log.Int("attempt", j.RetryCount), + log.Int("max_retries", j.MaxRetries), + log.Duration("delay", delay), + ) + + return fmt.Errorf("job %s retry %d/%d: %w", j.Name, j.RetryCount, j.MaxRetries, fmt.Errorf("%s", j.LastError)) +} + +// sendToDLQ marks the job as failed, pushes it to the DLQ, and emits events. +func (r *Runner) sendToDLQ(ctx context.Context, j *job.Job, handlerErr error) error { + r.forgetLaunchFailures(j.ID.String()) + + j.State = job.StateFailed + + if updateErr := r.updateJob(ctx, j); updateErr != nil { + if errors.Is(updateErr, job.ErrLeaseLost) { + return r.abandonLostLease(ctx, j, updateErr) + } + + r.logger.Error("failed to update job as failed", + log.String("job_id", j.ID.String()), + log.String("error", updateErr.Error()), + ) + return updateErr + } + + if r.dlqService != nil { + if dlqErr := r.dlqService.Push(ctx, j, handlerErr); dlqErr != nil { + r.logger.Error("failed to push job to DLQ", + log.String("job_id", j.ID.String()), + log.String("error", dlqErr.Error()), + ) + } + } + + r.extensions.EmitJobFailed(ctx, j, handlerErr) + r.extensions.EmitJobDLQ(ctx, j, handlerErr) + + // Not always "after exhausting retries" any more: a permanent failure + // arrives here on its first attempt. + r.logger.Warn("job moved to DLQ", + log.String("job_id", j.ID.String()), + log.String("job_name", j.Name), + log.Int("retry_count", j.RetryCount), + log.String("error", handlerErr.Error()), + ) + + return handlerErr +} diff --git a/worker/runner_outputs_test.go b/worker/runner_outputs_test.go new file mode 100644 index 0000000..6408091 --- /dev/null +++ b/worker/runner_outputs_test.go @@ -0,0 +1,1250 @@ +package worker_test + +import ( + "context" + "errors" + "fmt" + "io" + "os" + osexec "os/exec" + "path/filepath" + "runtime" + "sort" + "strconv" + "strings" + "testing" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/inproc" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/store/memory" + "github.com/xraph/dispatch/worker" +) + +// scriptedExecutor is an exec.Executor whose Run writes preconfigured +// files into req.OutputDir and returns a preconfigured Result, without +// spawning any real subprocess or shim. It exists so the Runner's +// directory-lifecycle and output-committing plumbing can be tested +// without exec/subprocess or exec/shim in the loop. +type scriptedExecutor struct { + level exec.Level + + // files are written into req.OutputDir before Run returns, keyed by + // name with their content as the value. Keys may contain "/" to + // place a file in a subdirectory, which is created as needed. + files map[string]string + + // symlinks are created in req.OutputDir before Run returns, keyed + // by the link's name with its target as the value. The target is an + // absolute path, standing in for anything on the worker's own + // filesystem a compromised handler might point at — its own + // config, a credential file — since a symlink's target is not + // confined to OutputDir at all. + symlinks map[string]string + + // claim becomes the returned Result.Outputs. It is deliberately + // independent of files, so a test can make the sandbox lie: claim a + // name it never wrote, omit one it did, or misreport a size. + claim []exec.OutputFile + + status exec.Status + + // beforeReturn, when set, runs after files/symlinks are written but + // before Run returns its Result — the point in time a real + // out-of-process rung's Run has already finished but the worker has + // not yet started committing. It exists so a test can simulate the + // pool's heartbeat loop cancelling the job's context in that exact + // window, the race commitOutputs' fence gate exists to catch. + beforeReturn func() + + got *exec.Request +} + +func (e *scriptedExecutor) Name() string { return "scripted" } +func (e *scriptedExecutor) Level() exec.Level { return e.level } + +func (e *scriptedExecutor) Run(_ context.Context, req *exec.Request) (*exec.Result, error) { + e.got = req + + for name, content := range e.files { + full := filepath.Join(req.OutputDir, name) + if err := os.MkdirAll(filepath.Dir(full), 0o750); err != nil { + return nil, err + } + if err := os.WriteFile(full, []byte(content), 0o600); err != nil { + return nil, err + } + } + + for name, target := range e.symlinks { + if err := os.Symlink(target, filepath.Join(req.OutputDir, name)); err != nil { + return nil, err + } + } + + if e.beforeReturn != nil { + e.beforeReturn() + } + + status := e.status + if status == "" { + status = exec.StatusOK + } + + return &exec.Result{Status: status, Outputs: e.claim}, nil +} + +func (e *scriptedExecutor) Reclaim(context.Context, id.WorkerID) error { return nil } +func (e *scriptedExecutor) Close() error { return nil } + +// artifactPlane bundles the store, backend, and service a test wires +// into a Runner via WithArtifacts. +type artifactPlane struct { + store *memory.Store + backend *artifacttest.Backend + svc *artifact.Service +} + +func newArtifactPlane() *artifactPlane { + s := memory.New() + b := artifacttest.NewBackend() + svc := artifact.NewService(s, b, artifact.WithDefaultBucket("dispatch")) + + return &artifactPlane{store: s, backend: b, svc: svc} +} + +// flakyBackend wraps another artifact.Backend, failing Create for any +// key containing failSubstring a configurable number of times before +// behaving normally — a test double standing in for a transient +// backend outage that clears on its own. +type flakyBackend struct { + inner artifact.Backend + failSubstring string + failN int +} + +var _ artifact.Backend = (*flakyBackend)(nil) + +func (b *flakyBackend) Name() string { return b.inner.Name() } + +func (b *flakyBackend) Open(ctx context.Context, ref artifact.Ref) (io.ReadCloser, error) { + return b.inner.Open(ctx, ref) +} + +func (b *flakyBackend) Stat(ctx context.Context, ref artifact.Ref) (artifact.ObjectInfo, error) { + return b.inner.Stat(ctx, ref) +} + +func (b *flakyBackend) Delete(ctx context.Context, ref artifact.Ref) error { + return b.inner.Delete(ctx, ref) +} + +func (b *flakyBackend) Create(ctx context.Context, bucket, key string) (artifact.Writer, error) { + if b.failN > 0 && strings.Contains(key, b.failSubstring) { + b.failN-- + + return nil, errors.New("simulated transient backend outage") + } + + return b.inner.Create(ctx, bucket, key) +} + +// seedPriorOutput records name as though attempt already committed it, +// so a later attempt's Runner.resolvePriorOutputs has something to find. +func (p *artifactPlane) seedPriorOutput(t *testing.T, jobID id.JobID, name string, attempt int) artifact.Ref { + t.Helper() + + a := &artifact.Artifact{ + ID: id.NewArtifactID(), + Backend: p.backend.Name(), + Bucket: "dispatch", + Key: "ephemeral/job/" + jobID.String() + "/" + strconv.Itoa(attempt) + "/" + name, + Size: int64(len(name)), + Lifecycle: artifact.Ephemeral, + CreatedAt: time.Now().UTC(), + } + link := &artifact.Link{ + ArtifactID: a.ID, + OwnerKind: artifact.OwnerJob, + OwnerID: jobID.String(), + Role: artifact.RoleOutput, + Name: name, + Attempt: attempt, + CreatedAt: time.Now().UTC(), + } + + if err := p.store.CreateArtifact(context.Background(), a, link); err != nil { + t.Fatalf("seed prior output %q: %v", name, err) + } + + return a.Ref() +} + +// newOutputsTestRunner builds a Runner wired to executors, with +// WithArtifacts applied only when plane is non-nil. +func newOutputsTestRunner( + t *testing.T, + reg *job.Registry, + executors *exec.Registry, + plane *artifactPlane, +) *worker.Runner { + t.Helper() + + runner := worker.NewRunner( + reg, + ext.NewRegistry(log.NewNoopLogger()), + newFakeJobStore(), + nil, + backoff.NewExponential(time.Second, time.Hour), + executors, + log.NewNoopLogger(), + ) + + if plane != nil { + runner = runner.WithArtifacts(plane.svc, t.TempDir()) + } + + return runner +} + +// isolatedJobRegistry registers "test.job", requiring LevelProcess +// isolation. +func isolatedJobRegistry(t *testing.T) *job.Registry { + t.Helper() + + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + return reg +} + +func TestRunner_OutOfProcessGetsFreshEmptyOutputDir(t *testing.T) { + rec := &scriptedExecutor{level: exec.LevelProcess} + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner := newOutputsTestRunner(t, reg, executors, nil) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + if rec.got == nil { + t.Fatal("executor was not called") + } + if rec.got.OutputDir == "" { + t.Fatal("Request.OutputDir was not set for an out-of-process rung") + } +} + +func TestRunner_InProcessGetsNoOutputDir(t *testing.T) { + // A job with no declared isolation runs in-process and never reaches + // an executor at all; nothing here should try to build it a scratch + // directory. + reg := job.NewRegistry() + job.NewDefinition("plain.job", func(context.Context, struct{}) error { return nil }).Register(reg) + + runner := newOutputsTestRunner(t, reg, exec.NewRegistry(inproc.New(reg)), nil) + + j := &job.Job{ID: id.NewJobID(), Name: "plain.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } +} + +func TestRunner_PriorOutputsPopulatedWhenArtifactsEnabled(t *testing.T) { + rec := &scriptedExecutor{level: exec.LevelProcess} + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + + jobID := id.NewJobID() + wantA := plane.seedPriorOutput(t, jobID, "a.txt", 0) + wantB := plane.seedPriorOutput(t, jobID, "b.txt", 0) + // A second, later attempt re-committing "a.txt" must win over the + // first — the same tie-break FindLinkByName applies for one name. + wantALatest := plane.seedPriorOutput(t, jobID, "a.txt", 1) + + runner := newOutputsTestRunner(t, reg, executors, plane) + + j := &job.Job{ID: jobID, Name: "test.job", RetryCount: 2, MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + if rec.got == nil { + t.Fatal("executor was not called") + } + + got := make(map[string]artifact.Ref, len(rec.got.PriorOutputs)) + for _, po := range rec.got.PriorOutputs { + got[po.Name] = po.Ref + } + + if len(got) != 2 { + t.Fatalf("PriorOutputs has %d entries, want 2: %+v", len(got), rec.got.PriorOutputs) + } + if got["a.txt"] != wantALatest { + t.Errorf("PriorOutputs[a.txt] = %+v, want the attempt-1 ref %+v (not attempt-0 %+v)", + got["a.txt"], wantALatest, wantA) + } + if got["b.txt"] != wantB { + t.Errorf("PriorOutputs[b.txt] = %+v, want %+v", got["b.txt"], wantB) + } +} + +func TestRunner_PriorOutputsEmptyWhenArtifactsDisabled(t *testing.T) { + rec := &scriptedExecutor{level: exec.LevelProcess} + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + // No WithArtifacts call at all: the plane is off. + runner := newOutputsTestRunner(t, reg, executors, nil) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + if rec.got == nil { + t.Fatal("executor was not called") + } + if len(rec.got.PriorOutputs) != 0 { + t.Errorf("PriorOutputs = %+v, want empty when the artifact plane is disabled", rec.got.PriorOutputs) + } +} + +func TestRunner_CommitsWhatIsActuallyOnDiskNotWhatIsClaimed(t *testing.T) { + // The invariant this task exists to protect: the artifact store must + // reflect the sandbox's actual filesystem, never its claims. "real.txt" + // is written but under-claimed (the sandbox lies about its size); + // "ghost.txt" is claimed but never written at all. + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"real.txt": "hello world"}, + claim: []exec.OutputFile{ + {Name: "real.txt", Size: 999999}, + {Name: "ghost.txt", Size: 12}, + }, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", RetryCount: 0, MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + byName := make(map[string]*artifact.Link, len(links)) + for _, l := range links { + byName[l.Name] = l + } + + if _, ok := byName["ghost.txt"]; ok { + t.Error("an artifact row was created for \"ghost.txt\", which the sandbox claimed but never wrote") + } + + link, ok := byName["real.txt"] + if !ok { + t.Fatal("no artifact row was created for \"real.txt\", which the sandbox actually wrote") + } + + a, err := plane.svc.Get(context.Background(), link.ArtifactID) + if err != nil { + t.Fatalf("Get(%s): %v", link.ArtifactID, err) + } + if want := int64(len("hello world")); a.Size != want { + t.Errorf("committed artifact Size = %d, want %d (the real content length, not the claimed 999999)", a.Size, want) + } + if !plane.backend.Has(a.Bucket, a.Key) { + t.Error("the committed artifact's bytes are not present in the backend") + } +} + +func TestRunner_SkipsCommittingWhenArtifactsDisabled(t *testing.T) { + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"real.txt": "hello"}, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + // No WithArtifacts call: nothing should be committed, and Execute + // must not fail because of it. + runner := newOutputsTestRunner(t, reg, executors, nil) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } +} + +func TestRunner_DoesNotCommitOutputsOnFailure(t *testing.T) { + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"partial.txt": "unfinished"}, + status: exec.StatusHandlerError, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() = nil, want a failure") + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + if len(links) != 0 { + t.Errorf("ListLinks = %+v, want none — a failed attempt must not commit anything", links) + } +} + +func TestRunner_RemovesScratchDirOnSuccess(t *testing.T) { + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"out.txt": "done"}, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner := newOutputsTestRunner(t, reg, executors, newArtifactPlane()) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + if rec.got == nil { + t.Fatal("executor was not called") + } + if _, err := os.Stat(rec.got.OutputDir); !os.IsNotExist(err) { + t.Errorf("OutputDir %q still exists after a successful attempt (stat err = %v)", rec.got.OutputDir, err) + } +} + +func TestRunner_RemovesScratchDirOnFailure(t *testing.T) { + rec := &scriptedExecutor{ + level: exec.LevelProcess, + status: exec.StatusHandlerError, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner := newOutputsTestRunner(t, reg, executors, nil) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() = nil, want a failure") + } + + if rec.got == nil { + t.Fatal("executor was not called") + } + if _, err := os.Stat(rec.got.OutputDir); !os.IsNotExist(err) { + t.Errorf("OutputDir %q still exists after a failed attempt (stat err = %v)", rec.got.OutputDir, err) + } +} + +func TestRunner_ScratchDirIsEmptyWhenHandedToTheExecutor(t *testing.T) { + var sawEntries []string + + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + rec := &recordingEmptyDirExecutor{seen: &sawEntries} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner := newOutputsTestRunner(t, reg, executors, nil) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + if len(sawEntries) != 0 { + t.Errorf("OutputDir contained %v when handed to the executor, want empty", sawEntries) + } +} + +// recordingEmptyDirExecutor records the directory entries it finds in +// req.OutputDir at the moment Run is called, into *seen. +type recordingEmptyDirExecutor struct { + seen *[]string +} + +func (e *recordingEmptyDirExecutor) Name() string { return "recording-empty-dir" } +func (e *recordingEmptyDirExecutor) Level() exec.Level { return exec.LevelProcess } + +func (e *recordingEmptyDirExecutor) Run(_ context.Context, req *exec.Request) (*exec.Result, error) { + entries, err := os.ReadDir(req.OutputDir) + if err != nil { + return nil, err + } + + names := make([]string, 0, len(entries)) + for _, entry := range entries { + names = append(names, entry.Name()) + } + sort.Strings(names) + *e.seen = names + + return &exec.Result{Status: exec.StatusOK}, nil +} + +func (e *recordingEmptyDirExecutor) Reclaim(context.Context, id.WorkerID) error { return nil } +func (e *recordingEmptyDirExecutor) Close() error { return nil } + +// TestRunner_SkipsSymlinksInOutputDir is C1: a compromised handler can +// place a symlink anywhere in its own writable OutputDir pointing at +// anything the worker process itself can read — its config, cloud +// credentials, a mounted service-account token. If the worker ever +// opened what that symlink resolves to and committed the bytes, a +// symlink to any file the worker can read becomes an ordinary, +// downloadable job output. This proves it does not: the secret's bytes +// never appear anywhere in the backend, and no artifact row is created +// for the symlink's name. +func TestRunner_SkipsSymlinksInOutputDir(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("creating a symlink requires elevated privilege on windows") + } + + secretPath := filepath.Join(t.TempDir(), "credentials") + const secret = "AKIA-SUPER-SECRET-WORKER-CREDENTIAL" + if err := os.WriteFile(secretPath, []byte(secret), 0o600); err != nil { + t.Fatalf("write secret file: %v", err) + } + + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"real.txt": "legitimate output"}, + symlinks: map[string]string{"innocent.txt": secretPath}, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil — a symlink must be skipped, not fail the attempt", err) + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + for _, l := range links { + if l.Name == "innocent.txt" { + t.Errorf("an artifact row was created for the symlink %q", l.Name) + } + + a, getErr := plane.svc.Get(context.Background(), l.ArtifactID) + if getErr != nil { + t.Fatalf("Get(%s): %v", l.ArtifactID, getErr) + } + + rc, openErr := plane.svc.Open(context.Background(), a.Ref()) + if openErr != nil { + t.Fatalf("Open(%s): %v", l.ArtifactID, openErr) + } + buf := make([]byte, len(secret)) + _, _ = rc.Read(buf) + _ = rc.Close() + + if strings.Contains(string(buf), secret) { + t.Fatalf("the worker's secret leaked into committed artifact %q (name %q)", l.ArtifactID, l.Name) + } + } + + if len(links) != 1 || links[0].Name != "real.txt" { + t.Errorf("committed links = %+v, want exactly [real.txt]", links) + } +} + +// TestRunner_SkipsNonRegularFilesInOutputDir is a cross-platform +// smoke test alongside the genuine hang reproduction for C2: platform- +// portable FIFO creation lives in the unix-only fifo_unix_test.go +// (TestRunner_OpeningAFIFOWouldHangSoItIsNeverOpened), which is where +// C2's actual measured hang is proven fixed. This one just confirms an +// ordinary, unremarkable attempt with a plain file still behaves, +// exercising the same code path this file's other tests changed. +func TestRunner_SkipsNonRegularFilesInOutputDir(t *testing.T) { + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"real.txt": "kept"}, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } +} + +// TestRunner_HiddenDirectoryContentsAreNotCommitted is m7: a +// dot-prefixed entry must be skipped whole, directory contents +// included — returning bare nil for a hidden directory does not stop +// WalkDir from descending into it, only fs.SkipDir does, so a file +// living inside a hidden directory must not slip through and get +// committed under its own (non-hidden) name. +func TestRunner_HiddenDirectoryContentsAreNotCommitted(t *testing.T) { + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{ + "visible.txt": "kept", + ".hidden/leaked.txt": "must not be committed", + ".hidden/sub/deeper.txt": "must not be committed either", + }, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + for _, l := range links { + if l.Name != "visible.txt" { + t.Errorf("committed a link from inside the hidden directory: %+v", l) + } + } + if len(links) != 1 { + t.Errorf("committed links = %+v, want exactly [visible.txt]", links) + } +} + +// TestRunner_RejectsDuplicateOutputNamesWithoutPartialCommits is C4 and +// C5 together: two files at different paths that would both commit as +// the same base name ("report.csv", nested under "us/" and "eu/") must +// fail the attempt cleanly, with zero partial commits left behind, and +// without consuming the job's real retry budget — a deterministic +// naming collision fails identically on every attempt, so retrying it +// normally would only burn the whole schedule for nothing. +func TestRunner_RejectsDuplicateOutputNamesWithoutPartialCommits(t *testing.T) { + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{ + "us/report.csv": "us data", + "eu/report.csv": "eu data", + }, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", RetryCount: 0, MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() = nil, want a failure for a duplicate output name") + } + + if j.RetryCount != 0 { + t.Errorf("RetryCount = %d, want 0 — a structural naming collision must not consume the real retry budget", + j.RetryCount) + } + if j.State != job.StatePending { + t.Errorf("State = %q, want %q (requeued via the bounded launch-failure path)", j.State, job.StatePending) + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + if len(links) != 0 { + t.Errorf("ListLinks = %+v, want none — a rejected commit must leave nothing behind", links) + } +} + +// TestRunner_DoesNotCommitOutputsWhenLeaseFenceIsAlreadyLost is C3's +// gate half. It simulates the pool's heartbeat loop cancelling the +// job's context with job.ErrLeaseLost in the exact window C3 measured: +// after the sandbox has already finished and is about to be reported a +// success, but before the worker has committed anything. A fenced-out +// attempt must not commit outputs as though it still owned the job. +func TestRunner_DoesNotCommitOutputsWhenLeaseFenceIsAlreadyLost(t *testing.T) { + baseCtx, cancel := context.WithCancelCause(context.Background()) + + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"out.txt": "should never be committed"}, + beforeReturn: func() { + // Exactly what Pool.sendHeartbeats does on a lost lease + // (cancelJob), except fired here, mid-attempt, instead of + // from a concurrent heartbeat goroutine — deterministic + // rather than timing-dependent, testing the same race. + cancel(job.ErrLeaseLost) + }, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + leaseStore := &fakeLeaseJobStore{fakeJobStore: newFakeJobStore()} + + runner := worker.NewRunner( + reg, ext.NewRegistry(log.NewNoopLogger()), leaseStore, nil, + backoff.NewExponential(time.Second, time.Hour), executors, log.NewNoopLogger(), + ).WithArtifacts(plane.svc, t.TempDir()) + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", MaxRetries: 3} + ctx := worker.WithLeaseFenceForTest(baseCtx, leaseStore, id.NewWorkerID(), 5) + + if err := runner.Execute(ctx, j); err == nil { + t.Fatal("Execute() = nil, want a failure once the lease fence is gone") + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + if len(links) != 0 { + t.Errorf("ListLinks = %+v, want none — a fenced-out attempt must not commit anything", links) + } +} + +// TestRunner_TwoHoldersSameAttempt_DifferentEpochsDoNotCollide is C3's +// key half. It reproduces the shape of the race directly: two workers +// each believe they hold the same job at the same RetryCount at once — +// a lease reclaim racing a zombie that has not yet noticed its lease +// expired — and both commit an output under the identical name. Before +// the fix this silently overwrote the earlier commit's backend bytes +// out from under its own still-valid artifact row (see the report for +// the mutation-tested proof); after it, distinct lease epochs give each +// holder its own backend object, so both commits survive with their own +// correct, uncorrupted bytes. +func TestRunner_TwoHoldersSameAttempt_DifferentEpochsDoNotCollide(t *testing.T) { + plane := newArtifactPlane() + leaseStore := &fakeLeaseJobStore{fakeJobStore: newFakeJobStore()} + jobID := id.NewJobID() + + run := func(epoch int, content string) { + t.Helper() + + reg := isolatedJobRegistry(t) + rec := &scriptedExecutor{level: exec.LevelProcess, files: map[string]string{"out.txt": content}} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner := worker.NewRunner( + reg, ext.NewRegistry(log.NewNoopLogger()), leaseStore, nil, + backoff.NewExponential(time.Second, time.Hour), executors, log.NewNoopLogger(), + ).WithArtifacts(plane.svc, t.TempDir()) + + j := &job.Job{ID: jobID, Name: "test.job", RetryCount: 0, MaxRetries: 3} + ctx := worker.WithLeaseFenceForTest(context.Background(), leaseStore, id.NewWorkerID(), epoch) + + if err := runner.Execute(ctx, j); err != nil { + t.Fatalf("Execute() (epoch %d) = %v, want nil", epoch, err) + } + } + + // The winner claims a fresh epoch after reclaiming the job; the + // zombie is still running under the epoch it was originally granted + // and finishes afterward, unaware it has already been reclaimed — + // order matches what C3 measured: the loser finishes second. + run(6, "WINNER-BYTES") + run(5, "LOSER-BYTES-THAT-WOULD-OVERWRITE-THE-WINNERS") + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + + outLinks := make([]*artifact.Link, 0, len(links)) + for _, l := range links { + if l.Name == "out.txt" { + outLinks = append(outLinks, l) + } + } + if len(outLinks) != 2 { + t.Fatalf("out.txt links = %d, want 2 — each holder's commit must survive as its own row", len(outLinks)) + } + + gotContents := make(map[string]bool) + seenKeys := make(map[string]bool) + for _, l := range outLinks { + a, getErr := plane.svc.Get(context.Background(), l.ArtifactID) + if getErr != nil { + t.Fatalf("Get(%s): %v", l.ArtifactID, getErr) + } + + if seenKeys[a.Key] { + t.Errorf("two holders resolved to the identical backend key %q", a.Key) + } + seenKeys[a.Key] = true + + rc, openErr := plane.svc.Open(context.Background(), a.Ref()) + if openErr != nil { + t.Fatalf("Open(%s): %v", l.ArtifactID, openErr) + } + buf := make([]byte, 128) + n, _ := rc.Read(buf) + _ = rc.Close() + + gotContents[string(buf[:n])] = true + } + + if !gotContents["WINNER-BYTES"] { + t.Error("the winner's bytes (epoch 6) were not found intact — something overwrote or lost them") + } + if !gotContents["LOSER-BYTES-THAT-WOULD-OVERWRITE-THE-WINNERS"] { + t.Error("the zombie's bytes (epoch 5) were not found intact — its commit did not survive as its own row") + } +} + +// deadPID starts and waits out a trivial child process, returning its +// PID once the OS has reaped it — a PID this test can be confident +// names no running process, unlike a made-up small integer that might +// coincidentally belong to something on the test machine. +func deadPID(t *testing.T) int { + t.Helper() + + if runtime.GOOS == "windows" { + t.Skip("no portable trivial no-op command on windows") + } + + cmd := osexec.CommandContext(context.Background(), "true") + if err := cmd.Run(); err != nil { + t.Fatalf("run trivial child process: %v", err) + } + + return cmd.Process.Pid +} + +// TestRunner_ReclaimSweepsOnlyScratchDirsWithNoLiveOwner is finding 4 +// from the second review round: age alone cannot tell a stale scratch +// directory from a live one, since a long attempt rewriting the files +// inside it never advances the directory's OWN mtime, and a sibling +// worker process sharing the same scratch root looks exactly like a +// dead one to a purely age-based sweep. Ownership — an embedded PID +// this process can ask the OS about — is what actually distinguishes +// them: a directory whose owning PID is provably gone is swept +// regardless of how old or new it looks; one whose owning PID is alive +// is left alone regardless of how old it looks. +func TestRunner_ReclaimSweepsOnlyScratchDirsWithNoLiveOwner(t *testing.T) { + dead := deadPID(t) + root := t.TempDir() + + // Owned by a PID that is definitely gone, and old — the case the + // sweep exists for. + staleDead := filepath.Join(root, fmt.Sprintf("dispatch-out-%d-job-abc-000000", dead)) + if err := os.Mkdir(staleDead, 0o750); err != nil { + t.Fatalf("mkdir staleDead: %v", err) + } + oldTime := time.Now().Add(-2 * time.Hour) + if err := os.Chtimes(staleDead, oldTime, oldTime); err != nil { + t.Fatalf("chtimes staleDead: %v", err) + } + + // Owned by THIS test process's own PID (definitely alive), but + // backdated to look exactly as old as staleDead — proving liveness + // overrides age, not merely that a fresh mtime is spared. + oldButAlive := filepath.Join(root, fmt.Sprintf("dispatch-out-%d-job-def-111111", os.Getpid())) + if err := os.Mkdir(oldButAlive, 0o750); err != nil { + t.Fatalf("mkdir oldButAlive: %v", err) + } + if err := os.Chtimes(oldButAlive, oldTime, oldTime); err != nil { + t.Fatalf("chtimes oldButAlive: %v", err) + } + + // Owned by a dead PID but too fresh to have cleared the courtesy + // grace window — must survive even though its owner is gone. + freshDead := filepath.Join(root, fmt.Sprintf("dispatch-out-%d-job-ghi-222222", dead)) + if err := os.Mkdir(freshDead, 0o750); err != nil { + t.Fatalf("mkdir freshDead: %v", err) + } + + // Does not match the expected name shape at all — never ours, + // regardless of age. + unrelated := filepath.Join(root, "not-ours-at-all") + if err := os.Mkdir(unrelated, 0o750); err != nil { + t.Fatalf("mkdir unrelated: %v", err) + } + if err := os.Chtimes(unrelated, oldTime, oldTime); err != nil { + t.Fatalf("chtimes unrelated: %v", err) + } + + reg := job.NewRegistry() + runner := worker.NewRunner( + reg, ext.NewRegistry(log.NewNoopLogger()), newFakeJobStore(), nil, + backoff.NewExponential(time.Second, time.Hour), exec.NewRegistry(inproc.New(reg)), log.NewNoopLogger(), + ).WithArtifacts(newArtifactPlane().svc, root) + + if err := runner.Reclaim(context.Background(), id.NewWorkerID()); err != nil { + t.Fatalf("Reclaim() = %v, want nil", err) + } + + if _, err := os.Stat(staleDead); !os.IsNotExist(err) { + t.Errorf("dead-owner, stale scratch dir still exists after Reclaim (stat err = %v)", err) + } + if _, err := os.Stat(oldButAlive); err != nil { + t.Errorf("live-owner scratch dir was removed by Reclaim despite an old mtime: %v", err) + } + if _, err := os.Stat(freshDead); err != nil { + t.Errorf("dead-owner scratch dir was removed before clearing the grace window: %v", err) + } + if _, err := os.Stat(unrelated); err != nil { + t.Errorf("an unrelated directory was removed by Reclaim: %v", err) + } +} + +// TestRunner_ReclaimSweepsWithoutAnArtifactPlaneToo pins the corrected +// behaviour from the comment/doc drift sweep: prepareOutputDir creates a +// scratch directory for ANY out-of-process attempt regardless of +// whether this Runner has an artifact plane configured (see +// terminalFor — only PriorOutputs and committing are gated on +// r.artifacts, not scratch-directory creation itself, exactly as +// WithArtifacts' own doc comment says). A Runner with no artifact plane +// therefore leaks the identical "dispatch-out-…" directories on a crash +// as one with a plane configured, and Reclaim's sweep has to run for it +// too or those directories are never reclaimed at all — the exact gap +// an earlier version of this test enshrined as intentional by asserting +// the opposite. +// +// scratchRoot only has one setter — WithArtifacts — so a Runner with no +// artifact plane at all necessarily also has no configured scratchRoot +// and falls back to the process's real os.TempDir(). That is the exact +// case this test exercises: nothing here is a fabricated test-only +// path, it is what sweepStaleScratchDirs actually resolves to when a +// bare Runner (no WithArtifacts) calls Reclaim. TMPDIR is redirected via +// t.Setenv so the test can observe it without touching the real system +// temp directory. +func TestRunner_ReclaimSweepsWithoutAnArtifactPlaneToo(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("os.TempDir() does not read TMPDIR on windows") + } + + dead := deadPID(t) + root := t.TempDir() + t.Setenv("TMPDIR", root) + + if got := os.TempDir(); got != root { + t.Fatalf("os.TempDir() = %q after redirecting TMPDIR, want %q — test setup is not exercising what it thinks", + got, root) + } + + staleDead := filepath.Join(root, fmt.Sprintf("dispatch-out-%d-job-abc-000000", dead)) + if err := os.Mkdir(staleDead, 0o750); err != nil { + t.Fatalf("mkdir staleDead: %v", err) + } + oldTime := time.Now().Add(-2 * time.Hour) + if err := os.Chtimes(staleDead, oldTime, oldTime); err != nil { + t.Fatalf("chtimes staleDead: %v", err) + } + + // A live-owner directory in the same root must still survive even + // with no artifact plane configured: ownership, not the presence of + // an artifact plane, is what protects a sibling's in-use directory. + oldButAlive := filepath.Join(root, fmt.Sprintf("dispatch-out-%d-job-def-111111", os.Getpid())) + if err := os.Mkdir(oldButAlive, 0o750); err != nil { + t.Fatalf("mkdir oldButAlive: %v", err) + } + if err := os.Chtimes(oldButAlive, oldTime, oldTime); err != nil { + t.Fatalf("chtimes oldButAlive: %v", err) + } + + reg := job.NewRegistry() + // No WithArtifacts call at all, so r.scratchRoot is unset and falls + // back to os.TempDir() — exactly the configuration the sweep must + // still reach into now. + runner := worker.NewRunner( + reg, ext.NewRegistry(log.NewNoopLogger()), newFakeJobStore(), nil, + backoff.NewExponential(time.Second, time.Hour), exec.NewRegistry(inproc.New(reg)), log.NewNoopLogger(), + ) + + if err := runner.Reclaim(context.Background(), id.NewWorkerID()); err != nil { + t.Fatalf("Reclaim() = %v, want nil", err) + } + + if _, err := os.Stat(staleDead); !os.IsNotExist(err) { + t.Errorf("a Runner with no artifact plane left a dead-owner scratch dir behind (stat err = %v) — "+ + "it leaks the identical directories a Runner with a plane does and must sweep them too", err) + } + if _, err := os.Stat(oldButAlive); err != nil { + t.Errorf("a live-owner scratch dir was removed despite no artifact plane being configured: %v", err) + } +} + +// TestRunner_MissingOutputDirIsALoudFailureNotEmptySuccess is finding +// 4's third requirement, independent of the sweep itself: a vanished +// OutputDir must never be read as "the handler produced nothing." That +// distinction matters most exactly when something else — a sibling +// process's sweep, however it decides staleness — has removed a still +// in-use OutputDir out from under a running attempt: silently reporting +// success with zero artifacts would be worse than any error. +func TestRunner_MissingOutputDirIsALoudFailureNotEmptySuccess(t *testing.T) { + rec := &scriptedExecutor{level: exec.LevelProcess} + rec.beforeReturn = func() { + // Simulates OutputDir having been removed by something else + // entirely — a sibling process's sweep, or the sandbox itself + // misbehaving — between the sandbox finishing and the worker + // walking it. + if err := os.RemoveAll(rec.got.OutputDir); err != nil { + t.Fatalf("remove OutputDir: %v", err) + } + } + + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() = nil, want a failure — a missing OutputDir must never look like success") + } + + if j.State == job.StateCompleted { + t.Errorf("State = %q, want anything but completed — nothing was actually committed", j.State) + } +} + +// TestRunner_PartialCommitSurvivesForRetryToConverge is findings 2 and +// 3 from the second review round, verified together: a transient +// backend failure partway through committing outputs must not roll +// back what already landed (finding 2 — a rollback would delete bytes +// PriorOutputs still points a retried handler at, actively telling it +// to skip regenerating data that no longer exists), and a retry of the +// SAME attempt — a commit failure is classified StatusLaunchFailed, so +// RetryCount does not advance and the retry reuses the identical +// (owner, attempt) key namespace — must actually converge once the +// failure clears, even when the handler itself is not careful enough +// to consult PriorOutputs and skip regenerating a name it already +// produced (finding 3's "any store without lease grants" case, where +// nothing about a retry changes the key at all). +func TestRunner_PartialCommitSurvivesForRetryToConverge(t *testing.T) { + store := memory.New() + flaky := &flakyBackend{inner: artifacttest.NewBackend(), failSubstring: "b-second.txt", failN: 1} + svc := artifact.NewService(store, flaky, artifact.WithDefaultBucket("dispatch")) + + reg := isolatedJobRegistry(t) + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"a-first.txt": "AAA", "b-second.txt": "BBB"}, + } + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner := worker.NewRunner( + reg, ext.NewRegistry(log.NewNoopLogger()), newFakeJobStore(), nil, + backoff.NewExponential(time.Second, time.Hour), executors, log.NewNoopLogger(), + ).WithArtifacts(svc, t.TempDir()) + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", RetryCount: 0, MaxRetries: 3} + + // Pass 1: "a-first.txt" commits; "b-second.txt" hits the simulated + // outage, failing the whole attempt. It must requeue via the bounded + // launch-failure path — RetryCount must not advance. + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() (pass 1) = nil, want a failure from the simulated outage") + } + if j.RetryCount != 0 { + t.Fatalf("RetryCount after pass 1 = %d, want 0 (a commit failure must not consume the real retry budget)", + j.RetryCount) + } + if j.State != job.StatePending { + t.Fatalf("State after pass 1 = %q, want %q", j.State, job.StatePending) + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + linksAfterPass1, err := store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks after pass 1: %v", err) + } + if len(linksAfterPass1) != 1 || linksAfterPass1[0].Name != "a-first.txt" { + t.Fatalf("links after pass 1 = %+v, want exactly [a-first.txt] (not rolled back)", linksAfterPass1) + } + + // Finding 2's specific check: a-first.txt's bytes must still be + // genuinely openable — a rollback would have deleted them out from + // under the very link a retried handler's PriorOutputs check relies + // on. + firstArtifact, err := svc.Get(context.Background(), linksAfterPass1[0].ArtifactID) + if err != nil { + t.Fatalf("Get(a-first.txt) after pass 1: %v", err) + } + rc, err := svc.Open(context.Background(), firstArtifact.Ref()) + if err != nil { + t.Fatalf("Open(a-first.txt) after pass 1: %v — a rollback would have deleted these bytes", err) + } + buf := make([]byte, 16) + n, _ := rc.Read(buf) + _ = rc.Close() + if got := string(buf[:n]); got != "AAA" { + t.Errorf("a-first.txt bytes after pass 1 = %q, want %q", got, "AAA") + } + firstArtifactID := firstArtifact.ID + + // Pass 2 reuses the same job object (same RetryCount, same attempt). + // The handler is deliberately NOT idempotent here — it rewrites BOTH + // files again, including the one that already succeeded — to prove + // convergence does not depend on the handler consulting PriorOutputs + // on its own. The simulated outage has cleared (failN is exhausted). + rec2 := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"a-first.txt": "AAA", "b-second.txt": "BBB"}, + } + executors.Add(rec2) + + if execErr := runner.Execute(context.Background(), j); execErr != nil { + t.Fatalf("Execute() (pass 2) = %v, want nil — the outage cleared", execErr) + } + if j.State != job.StateCompleted { + t.Errorf("State after pass 2 = %q, want %q", j.State, job.StateCompleted) + } + + // The other half of finding 2: rec2 must have SEEN a-first.txt in + // PriorOutputs, pointing at pass 1's own still-valid artifact. + foundPrior := false + for _, po := range rec2.got.PriorOutputs { + if po.Name != "a-first.txt" { + continue + } + foundPrior = true + if po.Ref.ID != firstArtifactID { + t.Errorf("PriorOutputs[a-first.txt].Ref.ID = %s, want %s (pass 1's own artifact)", + po.Ref.ID, firstArtifactID) + } + } + if !foundPrior { + t.Error("PriorOutputs on pass 2 did not include a-first.txt from pass 1's partial success") + } + + linksAfterPass2, err := store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks after pass 2: %v", err) + } + byName := make(map[string]*artifact.Link, len(linksAfterPass2)) + for _, l := range linksAfterPass2 { + byName[l.Name] = l + } + if len(byName) != 2 { + t.Fatalf("links after pass 2 = %+v, want exactly [a-first.txt, b-second.txt]", linksAfterPass2) + } + + // Finding 3's core assertion: a-first.txt was recognised as + // already-committed and skipped, not re-written under a fresh + // artifact — the SAME artifact ID both times. + if byName["a-first.txt"].ArtifactID != firstArtifactID { + t.Errorf("a-first.txt was re-committed on pass 2 (new artifact %s), want the SAME artifact %s from pass 1", + byName["a-first.txt"].ArtifactID, firstArtifactID) + } +} + +// TestRunner_PlainCancellationDuringCommitIsNotLeaseLoss is finding 1 +// from the second review round: only a definitive job.ErrLeaseLost +// from the pool's heartbeat loop may be classified as lease-fence loss. +// A plain cancellation for an unrelated reason — a graceful shutdown +// mid-commit is the obvious one — must not be reported as "lease fence +// lost," since that label determines which retry path (and therefore +// which store write, fenced or not) the resulting failure takes. +func TestRunner_PlainCancellationDuringCommitIsNotLeaseLoss(t *testing.T) { + baseCtx, cancel := context.WithCancelCause(context.Background()) + + rec := &scriptedExecutor{ + level: exec.LevelProcess, + files: map[string]string{"out.txt": "unrelated cancellation"}, + beforeReturn: func() { + cancel(context.Canceled) // deliberately NOT job.ErrLeaseLost + }, + } + reg := isolatedJobRegistry(t) + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + plane := newArtifactPlane() + runner := newOutputsTestRunner(t, reg, executors, plane) + + jobID := id.NewJobID() + j := &job.Job{ID: jobID, Name: "test.job", MaxRetries: 3} + + if err := runner.Execute(baseCtx, j); err != nil { + t.Fatalf("Execute() = %v, want nil — a plain cancellation must not gate the commit", err) + } + + owner := artifact.OwnerRef{Kind: artifact.OwnerJob, ID: jobID.String()} + links, err := plane.store.ListLinks(context.Background(), owner) + if err != nil { + t.Fatalf("ListLinks: %v", err) + } + if len(links) != 1 || links[0].Name != "out.txt" { + t.Errorf("links = %+v, want exactly [out.txt] — a plain cancellation wrongly gated on as lease loss "+ + "would have skipped committing it", links) + } +} diff --git a/worker/runner_test.go b/worker/runner_test.go new file mode 100644 index 0000000..d3bbc48 --- /dev/null +++ b/worker/runner_test.go @@ -0,0 +1,591 @@ +package worker_test + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + "time" + + log "github.com/xraph/go-utils/log" + + "github.com/xraph/dispatch" + "github.com/xraph/dispatch/backoff" + "github.com/xraph/dispatch/exec" + "github.com/xraph/dispatch/exec/inproc" + "github.com/xraph/dispatch/ext" + "github.com/xraph/dispatch/id" + "github.com/xraph/dispatch/job" + "github.com/xraph/dispatch/resource" + "github.com/xraph/dispatch/worker" +) + +// recordingExecutor captures the Request the runner built. +type recordingExecutor struct { + got *exec.Request + result *exec.Result + err error + reclaimed int + closed int +} + +func (r *recordingExecutor) Name() string { return "recording" } +func (r *recordingExecutor) Level() exec.Level { return exec.LevelProcess } + +func (r *recordingExecutor) Run(_ context.Context, req *exec.Request) (*exec.Result, error) { + r.got = req + if r.err != nil { + return nil, r.err + } + if r.result != nil { + return r.result, nil + } + + return &exec.Result{Status: exec.StatusOK}, nil +} + +func (r *recordingExecutor) Reclaim(context.Context, id.WorkerID) error { + r.reclaimed++ + return nil +} + +func (r *recordingExecutor) Close() error { + r.closed++ + return nil +} + +func newTestRunner(t *testing.T, reg *job.Registry, executors *exec.Registry) (*worker.Runner, *fakeJobStore) { + t.Helper() + + store := newFakeJobStore() + + return worker.NewRunner( + reg, + ext.NewRegistry(log.NewNoopLogger()), + store, + nil, + backoff.NewExponential(time.Second, time.Hour), + executors, + log.NewNoopLogger(), + ), store +} + +func TestRunner_ExecuteBuildsRequestFromJob(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + rec := &recordingExecutor{} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, _ := newTestRunner(t, reg, executors) + + j := &job.Job{ + ID: id.NewJobID(), + Name: "test.job", + Payload: []byte(`{"a":1}`), + RetryCount: 2, + MaxRetries: 3, + ScopeAppID: "app_1", + ScopeOrgID: "org_1", + ResourceLimits: resource.Set{resource.Memory: 256 << 20}, + } + + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + if rec.got == nil { + t.Fatal("executor was not called") + } + if rec.got.Name != "test.job" { + t.Errorf("Request.Name = %q, want %q", rec.got.Name, "test.job") + } + if rec.got.Attempt != 2 { + t.Errorf("Request.Attempt = %d, want 2", rec.got.Attempt) + } + if rec.got.ScopeAppID != "app_1" || rec.got.ScopeOrgID != "org_1" { + t.Errorf("Request scope = (%q, %q), want (app_1, org_1)", rec.got.ScopeAppID, rec.got.ScopeOrgID) + } + if rec.got.Policy.Level != exec.LevelProcess { + t.Errorf("Request.Policy.Level = %v, want %v", rec.got.Policy.Level, exec.LevelProcess) + } + // job.WithResourceLimits' resolved ceiling must cross the execution + // boundary intact — this is the only thing that lets an isolated + // rung enforce a per-job limit rather than a deployment-wide one. + if got := rec.got.ResourceLimits[resource.Memory]; got != 256<<20 { + t.Errorf("Request.ResourceLimits[memory] = %d, want %d", got, 256<<20) + } +} + +func TestRunner_ExecuteRoutesByPolicy(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition("plain.job", func(context.Context, struct{}) error { return nil }).Register(reg) + + rec := &recordingExecutor{} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, _ := newTestRunner(t, reg, executors) + + // No declared isolation, so this must go to the default executor and + // never reach the recording one. + j := &job.Job{ID: id.NewJobID(), Name: "plain.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + if rec.got != nil { + t.Error("a job with no declared isolation was routed to the isolated executor") + } +} + +func TestRunner_LaunchFailureDoesNotConsumeRetries(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + rec := &recordingExecutor{ + result: &exec.Result{Status: exec.StatusLaunchFailed, HandlerErr: "image pull backoff"}, + } + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, store := newTestRunner(t, reg, executors) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + err := runner.Execute(context.Background(), j) + if err == nil { + t.Fatal("Execute() = nil, want a failure") + } + if j.RetryCount != 0 { + t.Errorf("RetryCount = %d, want 0 — a launch failure is infrastructure", j.RetryCount) + } + if j.State != job.StatePending && j.State != job.StateRetrying { + t.Errorf("State = %q, want the job requeued", j.State) + } + if store.updates == 0 { + t.Error("the job was never persisted") + } +} + +func TestRunner_RunErrorIsTreatedAsLaunchFailure(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + rec := &recordingExecutor{err: errors.New("image pull backoff")} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, store := newTestRunner(t, reg, executors) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + err := runner.Execute(context.Background(), j) + if err == nil { + t.Fatal("Execute() = nil, want a failure") + } + if j.RetryCount != 0 { + t.Errorf("RetryCount = %d, want 0 — a raw Run error is a launch failure", j.RetryCount) + } + if j.State != job.StatePending && j.State != job.StateRetrying { + t.Errorf("State = %q, want the job requeued", j.State) + } + if store.updates == 0 { + t.Error("the job was never persisted") + } +} + +func TestRunner_RunErrorWrappingInvalidRequestGoesToDLQ(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + rec := &recordingExecutor{err: fmt.Errorf("bad: %w", exec.ErrInvalidRequest)} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, _ := newTestRunner(t, reg, executors) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + err := runner.Execute(context.Background(), j) + if err == nil { + t.Fatal("Execute() = nil, want a failure") + } + if j.State != job.StateFailed { + t.Errorf("State = %q, want %q — an invalid request must not requeue forever", j.State, job.StateFailed) + } +} + +// TestRunner_UnsatisfiablePolicyReachesDLQ guards the fix for a job whose +// declared execution policy no configured executor can satisfy: before +// this fix, terminalFor's error from exec.Registry.Select came straight +// back out of Execute without ever reaching handleFailure, so the job +// was never marked failed, never retried, and never persisted — it sat +// at StateRunning until its lease expired and got reaped, then re-leased +// to discover the identical, permanently unsatisfiable policy again, in +// an unbounded loop. Mutation check: reverting Execute to `return err` +// for terminalFor's error (bypassing handleFailure) makes this test fail +// with store.updates == 0 and State still StateRunning. +func TestRunner_UnsatisfiablePolicyReachesDLQ(t *testing.T) { + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + // No process-level executor registered: only the in-process default, + // which cannot satisfy exec.LevelProcess. + executors := exec.NewRegistry(inproc.New(reg)) + + runner, store := newTestRunner(t, reg, executors) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", State: job.StateRunning, MaxRetries: 3} + err := runner.Execute(context.Background(), j) + if err == nil { + t.Fatal("Execute() = nil, want a failure") + } + if !errors.Is(err, exec.ErrNoExecutor) { + t.Errorf("Execute() = %v, want it to wrap %v", err, exec.ErrNoExecutor) + } + if !errors.Is(err, dispatch.ErrPermanent) { + t.Errorf("Execute() = %v, want it to wrap %v", err, dispatch.ErrPermanent) + } + if j.State != job.StateFailed { + t.Errorf("State = %q, want %q — an unsatisfiable policy must reach a terminal state, not loop forever", + j.State, job.StateFailed) + } + if j.LastError == "" { + t.Error("LastError was never recorded") + } + if store.updates == 0 { + t.Error("the job was never persisted — it would sit at StateRunning until its lease expired, forever") + } +} + +func TestRunner_HandlerErrorConsumesRetries(t *testing.T) { + sentinel := errors.New("bad file") + + reg := job.NewRegistry() + job.NewDefinition("test.job", func(context.Context, struct{}) error { return sentinel }).Register(reg) + + runner, _ := newTestRunner(t, reg, exec.NewRegistry(inproc.New(reg))) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() = nil, want a failure") + } + if j.RetryCount != 1 { + t.Errorf("RetryCount = %d, want 1", j.RetryCount) + } +} + +func TestRunner_PermanentHandlerErrorSkipsRetriesThroughAnExecutor(t *testing.T) { + // The regression this guards: routing the attempt through an executor + // must not cost the handler's error its identity, or ErrPermanent stops + // being honoured for every engine user. + reg := job.NewRegistry() + job.NewDefinition("test.job", func(context.Context, struct{}) error { + return fmt.Errorf("malformed payload: %w", dispatch.ErrPermanent) + }).Register(reg) + + runner, _ := newTestRunner(t, reg, exec.NewRegistry(inproc.New(reg))) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + err := runner.Execute(context.Background(), j) + if err == nil { + t.Fatal("Execute() = nil, want a failure") + } + if j.State != job.StateFailed { + t.Errorf("State = %q, want %q on the first attempt", j.State, job.StateFailed) + } + if j.RetryCount != 1 { + t.Errorf("RetryCount = %d, want 1 — the backoff schedule must not be spent", j.RetryCount) + } + if got, want := j.LastError, "malformed payload: dispatch: permanent failure"; got != want { + t.Errorf("LastError = %q, want %q", got, want) + } + if !errors.Is(err, dispatch.ErrPermanent) { + t.Errorf("errors.Is(Execute(), ErrPermanent) = false, want true (err = %v)", err) + } +} + +func TestRunner_PermanentResultFlagSkipsRetries(t *testing.T) { + // The out-of-process shape of the same thing: no error chain crossed + // the boundary, only the flag. + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + rec := &recordingExecutor{result: &exec.Result{ + Status: exec.StatusHandlerError, + HandlerErr: "the input object was deleted", + Permanent: true, + }} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, _ := newTestRunner(t, reg, executors) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() = nil, want a failure") + } + if j.State != job.StateFailed { + t.Errorf("State = %q, want %q on the first attempt", j.State, job.StateFailed) + } + if j.RetryCount != 1 { + t.Errorf("RetryCount = %d, want 1", j.RetryCount) + } +} + +func TestRunner_UserErrorIdentitySurvivesTheExecutor(t *testing.T) { + // What extensions receive: EmitJobFailed is handed this error, and an + // extension matching its own error type on it must still succeed. + reg := job.NewRegistry() + job.NewDefinition("test.job", func(context.Context, struct{}) error { + return fmt.Errorf("parse: %w", &userError{code: 42}) + }).Register(reg) + + runner, _ := newTestRunner(t, reg, exec.NewRegistry(inproc.New(reg))) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 0} + err := runner.Execute(context.Background(), j) + if err == nil { + t.Fatal("Execute() = nil, want a failure") + } + + var target *userError + if !errors.As(err, &target) { + t.Fatalf("errors.As(%v, **userError) = false, want true", err) + } + if target.code != 42 { + t.Errorf("target.code = %d, want 42", target.code) + } + if !errors.Is(err, exec.ErrHandler) { + t.Errorf("errors.Is(%v, ErrHandler) = false, want true", err) + } +} + +// userError stands in for an error type a caller or extension owns. +type userError struct{ code int } + +func (e *userError) Error() string { return fmt.Sprintf("user error %d", e.code) } + +func TestRunner_LaunchFailuresAreBounded(t *testing.T) { + // A launch failure does not consume the retry budget, so without a + // bound a job that can never launch requeues about once a second + // forever, costing a store write and a worker slot each cycle. + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + rec := &recordingExecutor{ + result: &exec.Result{Status: exec.StatusLaunchFailed, HandlerErr: "image pull backoff"}, + } + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, _ := newTestRunner(t, reg, executors) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + + const wantRequeues = 5 + for attempt := 1; attempt <= wantRequeues; attempt++ { + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatalf("attempt %d: Execute() = nil, want a failure", attempt) + } + if j.State != job.StatePending { + t.Fatalf("attempt %d: State = %q, want %q", attempt, j.State, job.StatePending) + } + if j.RetryCount != 0 { + t.Fatalf("attempt %d: RetryCount = %d, want 0", attempt, j.RetryCount) + } + } + + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() = nil after the launch cap, want a failure") + } + if j.State != job.StateFailed { + t.Errorf("State = %q, want %q once the launch cap is exceeded", j.State, job.StateFailed) + } + if !strings.Contains(j.LastError, "failed to launch") { + t.Errorf("LastError = %q, want it to explain the launch cap", j.LastError) + } +} + +func TestRunner_LaunchCounterIsForgottenOnceTheJobRuns(t *testing.T) { + // The counter must not accumulate across a job's lifetime, or a job + // that occasionally fails to launch would eventually be dead-lettered + // for it. It also must not outlive the job in memory. + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + + rec := &recordingExecutor{ + result: &exec.Result{Status: exec.StatusLaunchFailed, HandlerErr: "no capacity"}, + } + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, _ := newTestRunner(t, reg, executors) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + for range 4 { + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatal("Execute() = nil, want a launch failure") + } + } + + // The sandbox comes up and the job succeeds, which clears the count. + rec.result = &exec.Result{Status: exec.StatusOK} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + + // It may now fail to launch a full cap's worth again without being + // dead-lettered for the failures that preceded the success. + rec.result = &exec.Result{Status: exec.StatusLaunchFailed, HandlerErr: "no capacity"} + for attempt := range 5 { + if err := runner.Execute(context.Background(), j); err == nil { + t.Fatalf("attempt %d: Execute() = nil, want a launch failure", attempt+1) + } + if j.State != job.StatePending { + t.Fatalf("attempt %d: State = %q, want %q — the count was not reset", + attempt+1, j.State, job.StatePending) + } + } +} + +func TestRunner_RequestCarriesTheHandlerSetFingerprint(t *testing.T) { + // Drift protection is only as good as the fingerprint being populated: + // a rung handed an empty one either rejects every job or skips the + // check, and a stale image then runs an old handler and reports + // success. + reg := job.NewRegistry() + job.NewDefinition("test.job", + func(context.Context, struct{}) error { return nil }, + job.WithExecution(exec.Isolate(exec.LevelProcess)), + ).Register(reg) + job.NewDefinition("other.job", func(context.Context, struct{}) error { return nil }).Register(reg) + + rec := &recordingExecutor{} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, _ := newTestRunner(t, reg, executors) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := runner.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + if rec.got == nil { + t.Fatal("executor was not called") + } + if want := exec.Fingerprint(reg.Names()); rec.got.Fingerprint != want { + t.Errorf("Request.Fingerprint = %q, want %q", rec.got.Fingerprint, want) + } +} + +func TestRunner_ReclaimAndCloseReachEveryExecutor(t *testing.T) { + // Nothing called these before, so a rung with children to kill would + // have leaked them on every restart and shutdown. + reg := job.NewRegistry() + rec := &recordingExecutor{} + executors := exec.NewRegistry(inproc.New(reg)) + executors.Add(rec) + + runner, _ := newTestRunner(t, reg, executors) + + if err := runner.Reclaim(context.Background(), id.NewWorkerID()); err != nil { + t.Errorf("Reclaim() = %v, want nil", err) + } + if rec.reclaimed != 1 { + t.Errorf("executor reclaimed %d times, want 1", rec.reclaimed) + } + if err := runner.Close(); err != nil { + t.Errorf("Close() = %v, want nil", err) + } + if rec.closed != 1 { + t.Errorf("executor closed %d times, want 1", rec.closed) + } +} + +func TestNewExecutor_StillCompilesAndRuns(t *testing.T) { + // The deprecated constructor must keep working for existing callers. + reg := job.NewRegistry() + job.NewDefinition("test.job", func(context.Context, struct{}) error { return nil }).Register(reg) + + e := worker.NewExecutor( + reg, + ext.NewRegistry(log.NewNoopLogger()), + newFakeJobStore(), + nil, + backoff.NewExponential(time.Second, time.Hour), + log.NewNoopLogger(), + ) + + j := &job.Job{ID: id.NewJobID(), Name: "test.job", MaxRetries: 3} + if err := e.Execute(context.Background(), j); err != nil { + t.Fatalf("Execute() = %v, want nil", err) + } + if j.State != job.StateCompleted { + t.Errorf("State = %q, want %q", j.State, job.StateCompleted) + } +} + +// fakeJobStore is a job.Store that records UpdateJob calls. Only the +// method the runner uses does anything. +// +// DequeueJobs takes job.DequeueOpts rather than the brief's original +// (queues []string, limit int) — the Store interface's dequeue contract +// was widened by a concurrent change (see job/store.go) after the brief +// was written; this stub matches the current interface. +type fakeJobStore struct { + updates int +} + +func newFakeJobStore() *fakeJobStore { return &fakeJobStore{} } + +func (f *fakeJobStore) UpdateJob(context.Context, *job.Job) error { + f.updates++ + return nil +} + +func (f *fakeJobStore) EnqueueJob(context.Context, *job.Job) error { return nil } + +func (f *fakeJobStore) DequeueJobs(context.Context, job.DequeueOpts) ([]*job.Job, error) { + return nil, nil +} + +func (f *fakeJobStore) GetJob(context.Context, id.JobID) (*job.Job, error) { return nil, nil } + +func (f *fakeJobStore) DeleteJob(context.Context, id.JobID) error { return nil } + +func (f *fakeJobStore) ListJobsByState( + context.Context, job.State, job.ListOpts, +) ([]*job.Job, error) { + return nil, nil +} + +func (f *fakeJobStore) HeartbeatJob(context.Context, id.JobID, id.WorkerID) error { return nil } + +func (f *fakeJobStore) ReapStaleJobs(context.Context, time.Duration) ([]*job.Job, error) { + return nil, nil +} + +func (f *fakeJobStore) CountJobs(context.Context, job.CountOpts) (int64, error) { return 0, nil } diff --git a/worker/shared_manager_test.go b/worker/shared_manager_test.go new file mode 100644 index 0000000..9c2d413 --- /dev/null +++ b/worker/shared_manager_test.go @@ -0,0 +1,158 @@ +package worker + +import ( + "context" + "testing" + + "github.com/xraph/dispatch/artifact" + "github.com/xraph/dispatch/artifact/artifacttest" + "github.com/xraph/dispatch/artifact/cache" + "github.com/xraph/dispatch/resource" +) + +const mib = int64(1) << 20 + +// stageInto puts an object of size bytes in the backend, stages it +// through c, and returns the stager's release func. +func stageInto(t *testing.T, c *cache.Cache, b *artifacttest.Backend, key string, size int64) func() { + t.Helper() + + b.Put("stage", key, make([]byte, size)) + + _, _, release, err := c.Stage(context.Background(), artifact.Ref{ + Backend: b.Name(), + Bucket: "stage", + Key: key, + Size: size, + }) + if err != nil { + t.Fatalf("stage %s: %v", key, err) + } + + return release +} + +// TestCacheAndPoolShareOneLedger is the integration this whole wiring +// task exists to get right. +// +// The staging cache and the worker pool must hold the SAME +// resource.Manager. The cache takes a lease per cached entry and +// registers itself as the manager's disk reclaimer; the pool offers disk +// at dequeue as free PLUS what that reclaimer could evict. Wire two +// managers instead of one and every part still works in isolation — +// which is why this needs a test rather than a code review. The cache +// accounts perfectly against its private ledger, the pool accounts +// perfectly against its own, Reclaimable() on the pool's side is +// permanently zero, and a worker with a warm cache quietly stops claiming +// disk-hungry work. No error, no log line, just a worker that went quiet. +// +// So this asserts the coupling end to end, on real bytes: stage through +// the cache, then read the pool's budget. +func TestCacheAndPoolShareOneLedger(t *testing.T) { + const ( + capacityBytes = 64 * mib + stagedBytes = 8 * mib + ) + + mgr := resource.NewManager(resource.Set{ + resource.Memory: 4 * gib, + resource.Disk: capacityBytes, + }) + + backend := artifacttest.NewBackend() + + c, err := cache.New(t.TempDir(), backend, cache.WithManager(mgr)) + if err != nil { + t.Fatalf("cache.New: %v", err) + } + + p := &Pool{resources: mgr} + + if got := p.dequeueBudget()[resource.Disk]; got != capacityBytes { + t.Fatalf("empty cache: budget disk = %d, want %d", got, capacityBytes) + } + + release := stageInto(t, c, backend, "model.bin", stagedBytes) + + // Pinned by a live stager. The bytes are spent — they are neither free + // nor evictable — so the budget has to shrink by exactly that much. + // This is the half that proves the cache's lease landed in the + // manager the pool is reading. + if got, want := p.dequeueBudget()[resource.Disk], capacityBytes-stagedBytes; got != want { + t.Fatalf("while pinned: budget disk = %d, want %d", got, want) + } + + if got, want := mgr.Free()[resource.Disk], capacityBytes-stagedBytes; got != want { + t.Fatalf("while pinned: manager free disk = %d, want %d", got, want) + } + + release() + + // Unpinned: still on disk, still leased, but now evictable. The pool + // must offer it back, because admission can redeem it by reclaiming. + // A full cache is a healthy cache. + if got := mgr.Reclaimable()[resource.Disk]; got != stagedBytes { + t.Fatalf("after release: reclaimable disk = %d, want %d", got, stagedBytes) + } + + if got := p.dequeueBudget()[resource.Disk]; got != capacityBytes { + t.Fatalf("after release: budget disk = %d, want %d (free + reclaimable)", + got, capacityBytes) + } + + // And the promise is redeemable: a job asking for the whole volume + // gets it, by evicting. + lease, aerr := mgr.Acquire(context.Background(), "big-job", + resource.Set{resource.Disk: capacityBytes}) + if aerr != nil { + t.Fatalf("acquire the full budget: %v", aerr) + } + + lease.Release() +} + +// TestPrivateCacheManagerIsInvisibleToThePool is the negative control: +// the exact miswiring the test above guards against, asserted to produce +// the exact symptom described. +// +// A cache built without WithManager constructs its own ledger. Everything +// still works — the cache admits, evicts and accounts correctly — but the +// pool's manager never hears about any of it, so staged bytes are neither +// spent nor reclaimable from where admission is looking. +func TestPrivateCacheManagerIsInvisibleToThePool(t *testing.T) { + const ( + capacityBytes = 64 * mib + stagedBytes = 8 * mib + ) + + poolMgr := resource.NewManager(resource.Set{resource.Disk: capacityBytes}) + + backend := artifacttest.NewBackend() + + // No WithManager: the cache builds a private one over its own budget. + c, err := cache.New(t.TempDir(), backend, cache.WithBudget(capacityBytes)) + if err != nil { + t.Fatalf("cache.New: %v", err) + } + + release := stageInto(t, c, backend, "model.bin", stagedBytes) + release() + + p := &Pool{resources: poolMgr} + + if got := p.dequeueBudget()[resource.Disk]; got != capacityBytes { + t.Fatalf("private manager: budget disk = %d, want %d", got, capacityBytes) + } + + // The tell: the cache is holding bytes and the pool's ledger reports + // nothing to reclaim. Free() alone cannot distinguish this from an + // empty cache, which is why the miswiring is silent. + if got := poolMgr.Reclaimable()[resource.Disk]; got != 0 { + t.Fatalf("private manager: reclaimable disk = %d, want 0 — "+ + "the pool's manager cannot see a cache it was not given", got) + } + + if got := c.Used(); got != stagedBytes { + t.Fatalf("the cache did hold the bytes: used = %d, want %d", got, stagedBytes) + } +}