diff --git a/.env.example b/.env.example index 0556cda..11ee54c 100644 --- a/.env.example +++ b/.env.example @@ -13,3 +13,7 @@ LEASE_DURATION=30s STALE_RUNNING_AFTER=1m RETRY_BACKOFF=5s WORKER_CONCURRENCY=1 +# How often the worker upserts its workers row (minimum 1s) +WORKER_HEARTBEAT_INTERVAL=30s +# Stable worker identity; defaults to a random UUID per process +# WORKER_ID=worker-1 diff --git a/INSTRUCTIONS.md b/INSTRUCTIONS.md index cc74715..35e1853 100644 --- a/INSTRUCTIONS.md +++ b/INSTRUCTIONS.md @@ -59,7 +59,7 @@ flowchart LR W2 --> CRDB ``` -**Invariant:** If Redis is flushed or inconsistent, the system remains correct by reconciling from CRDB: **(implemented)** the orchestrator periodically selects **`queued`** tasks with `scheduled_at <= now()` and LPUSHes task IDs after a Redis **pending** marker (`SET NX`) to avoid duplicate enqueue storms; it also **reclaims** **`running`** tasks that are older than a configurable threshold **and** have **no Redis lease key** (worker heartbeats stopped), closing open `task_runs` and returning the task to `queued`. Redis is an optimization, not the system of record. +**Invariant:** If Redis is flushed or inconsistent, the system remains correct by reconciling from CRDB: **(implemented)** the orchestrator periodically selects **`queued`** tasks with `scheduled_at <= now()` and LPUSHes task IDs after a Redis **pending** marker (`SET NX`) to avoid duplicate enqueue storms; it also **reclaims** **`running`** tasks that are older than a configurable threshold **and** have **no Redis lease key** (worker heartbeats stopped), closing open `task_runs` as **`dead`** and returning the task to `queued` without consuming a retry. Redis is an optimization, not the system of record. --- @@ -104,6 +104,7 @@ distributed_task_queue/ - Timestamps: **`TIMESTAMPTZ`** everywhere. - Status fields: use **`STRING`** with enumerated values documented below (or `ENUM`-like check constraints if preferred). - **Transaction boundaries:** creating a job with all tasks and dependency rows MUST occur in **one transaction**. Transitioning a task and inserting a `task_runs` row for a new attempt SHOULD be one transaction where practical. +- **`attempt` is a retry-budget counter, not a run key.** `tasks.attempt` counts how much of `max_attempts` has been consumed. Reclaiming a stale `running` task rolls it back by one — deliberately, so losing a worker does not cost the task a retry — and the next claim re-issues the same number. `(task_id, attempt_number)` in `task_runs` is therefore **not unique** and MUST NOT be treated as one: a reclaimed-and-retried task holds a `dead` row and a later row sharing that number. The run key is `task_runs.id`; history is ordered by `started_at` (the index is built for exactly this). Distinguish an abandoned run from its successor by `status = 'dead'`, not by the `error` text and not by attempt number. ### Enumerated values @@ -111,7 +112,9 @@ distributed_task_queue/ **`tasks.status`:** `pending`, `queued`, `running`, `completed`, `failed`, `cancelled` -**`task_runs.status`:** `running`, `succeeded`, `failed` (implementation; a `dead` / stale-run distinction may be added later) +**`task_runs.status`:** `running`, `succeeded`, `failed`, `dead` + +`failed` means the handler ran and reported an error. `dead` means the run was abandoned — the worker stopped heartbeating and the orchestrator reclaimed the task — so no result was ever reported. Reclaim also writes the reason into `error`, but the status is the field to branch on. ### DDL @@ -158,6 +161,7 @@ CREATE TABLE task_dependencies ( CREATE INDEX idx_task_dependencies_depends ON task_dependencies (depends_on_task_id); +-- attempt_number is intentionally not unique per task_id (see the invariant under Conventions). CREATE TABLE task_runs ( id UUID NOT NULL PRIMARY KEY DEFAULT gen_random_uuid(), task_id UUID NOT NULL REFERENCES tasks (id) ON DELETE CASCADE, @@ -211,11 +215,21 @@ Per claimed task, store lease metadata so other workers do not claim it until ex Fields: `worker_id`, `deadline_ms` (or `deadline` as Unix ms string), optional `run_id` (UUID of `task_runs` row). - **TTL:** set **EXPIRE** on the hash to slightly exceed lease duration so orphaned keys disappear; correctness still comes from CRDB reconciliation if a worker dies without releasing. -**Claim flow (implemented):** worker `BRPOP` → CRDB `ClaimQueuedTask` (queued→running) → `InsertTaskRun` → `SetLease` on the hash → handler runs → `DeleteLease` on success path; heartbeat extends key TTL until completion. +**Claim flow (implemented):** worker `BRPOP` → CRDB `ClaimQueuedTask` (queued→running) → release the pending marker → `InsertTaskRun` → `SetLease` on the hash → handler runs → `DeleteLease` on success path; heartbeat extends key TTL until completion. If the claim does **not** stick (another worker won it, or the DB call failed), the marker is released anyway — `BRPOP` already removed the id from the ready list, so this worker owns the delivery and must hand it back. ### Pending enqueue deduplication -**Implemented:** `{prefix}queue:pending:{task_id}` — short TTL via `SET NX` so the reconciler and other producers do not LPUSH the same task repeatedly while it is already queued or in flight. Released when a worker successfully claims the task. +**Implemented:** `{prefix}queue:pending:{task_id}` — set via `SET NX` with a TTL so the reconciler and other producers do not LPUSH the same task repeatedly while it is already queued or in flight. + +**The marker is a lock on the delivery, and whoever ends a delivery MUST release it.** While it is set, `EnqueueDueTaskID` will not push the task, so a marker that outlives its delivery leaves the task `queued`, due, and unreachable — a deadlock that only the TTL breaks. Release is required on **all three** exit paths, not just the successful one: + +| Path | Released by | +|------|-------------| +| Worker claims the task | `Runtime.processTask` after `ClaimQueuedTask` succeeds | +| Worker pops the id but the claim fails or is lost | `Runtime.processTask` on the error path | +| Orchestrator reclaims a stale `running` task, abandoning the in-flight delivery | `ReclaimStaleRunningOnce` before it re-enqueues | + +The TTL is a **backstop** for a release that never happened (the process died between `LPUSH` and claim), not the primary mechanism. It does not need to exceed the worst-case `LPUSH`→claim latency: expiring while the task is still on the ready list only lets a producer push a duplicate id, and duplicates are harmless because `ClaimQueuedTask` is a single conditional `UPDATE` — the second worker to pop it gets `ErrTaskNotClaimable`. Prefer a **short** TTL: expiring early costs a wasted `BRPOP`, expiring late stalls a task for the whole TTL. ### Delayed / scheduled tasks @@ -239,6 +253,8 @@ Per claimed task, store lease metadata so other workers do not claim it until ex **CockroachDB always wins** for `tasks.status` and attempts. Redis entries that disagree with CRDB are harmless at claim time (`ClaimQueuedTask` gates execution). Recovery: reconciler re-enqueues due **`queued`** rows; separate path **reclaims** stale **`running`** rows when the Redis lease key is absent (see §2 invariant). +One asymmetry is worth stating, because it is the way this design can actually stall: Redis state that says **"there is more work than there is"** self-corrects at claim time, but Redis state that says **"a delivery is already in flight"** — a pending marker — *suppresses* recovery instead of triggering it. The reconciler treats the marker as authoritative and skips the task, so a leaked marker is not harmless the way a duplicate ready-list entry is. That is why the release paths above are mandatory rather than best-effort, and why the marker carries a TTL. Applies only to the pending marker; the lease hash fails the safe way, since a missing lease is what *causes* reclaim. + --- ## 6. Worker interface (Go) @@ -289,7 +305,7 @@ Implementations in `internal/worker` construct `Runtime` with Redis + DB clients | Topic | Behavior | |--------|----------| -| **Delivery** | **At-least-once.** The same logical attempt may be redelivered after crash, lease expiry, or network partition. | +| **Delivery** | **At-least-once.** The same logical attempt may be redelivered after crash, lease expiry, or network partition. Redelivery reuses the attempt number, so `task_runs` holds one `dead` row and one live row sharing it — see the `attempt` invariant in §4. | | **Success** | Handler returns `nil` → runtime **acks**: clear lease in Redis, update CRDB task to `completed`, close `task_runs` as `succeeded`, orchestrator may enqueue dependents. | | **Failure** | Handler returns error → runtime records error, increments attempt if under `max_attempts`, applies **backoff** to `scheduled_at`, sets status to `pending` or `queued` per policy, may re-enqueue to Redis after delay. | | **Lease / heartbeat** | While `Handler` runs, periodically extend Redis lease (and optionally refresh `task_runs.started_at` semantics); if extension fails, cancel handler `ctx` so shutdown is cooperative. | @@ -316,6 +332,7 @@ The control plane is **HTTP only** (no gRPC in this repo). Bind address: `ORCHES | `GET` | `/v1/tasks/{id}` | Task row JSON (status, attempts, timestamps, payload); **404** if unknown; **400** if `id` is not a UUID | | `POST` | `/v1/jobs` | JSON body `{"tasks":[...]}` — DAG job with `name`, `kind`, `payload`, optional `depends_on` (names); **201** + `job_id` and `tasks` name→id map; **400** on validation (cycles, unknown deps, etc.) | | `GET` | `/v1/jobs/{id}` | Job metadata + all tasks in the job; **404** / **400** as above | +| `GET` | `/v1/workers` | Workers that heartbeat recently, from the `workers` table. Optional `?active_since=` duration (default `2m`); **400** if not a positive duration | --- @@ -333,6 +350,7 @@ Loaded from the environment by both `cmd/orchestrator` and `cmd/worker` (`intern | `STALE_RUNNING_AFTER` | `2 × LEASE_DURATION` | Minimum time a task may stay `running` before reclaim is considered (still requires Redis lease to be absent) | | `WORKER_ID` | random UUID | Stable worker identity for `task_runs.worker_id` | | `WORKER_CONCURRENCY` | `1` | Parallel BRPOP loops in the worker | +| `WORKER_HEARTBEAT_INTERVAL` | `30s` | How often the worker upserts its `workers` row; must be ≥ `1s`. Registry only — the Redis **lease** heartbeat is `pkg/worker.Options.HeartbeatEvery` (5s default, code-configured, no env var) | | `LEASE_DURATION` | `30s` | Logical lease window; Redis key TTL adds a buffer | | `RETRY_BACKOFF` | `5s` | Delay before a failed task is re-queued when attempts remain | diff --git a/README.md b/README.md index 4dde894..23d6fa9 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,35 @@ powershell -ExecutionPolicy Bypass -File .\scripts\migrate.ps1 Copy [`.env.example`](.env.example) to `.env` and load it in your shell if you use a tool that reads `.env` automatically; otherwise set `CRDB_DSN` / `REDIS_ADDR` as in the table below. +## Three-node cluster (failover testing) + +[`docker-compose.multinode.yml`](docker-compose.multinode.yml) runs three CockroachDB nodes plus Redis, so you can kill a node and watch the system recover. The single-node compose above is enough for normal development. + +```bash +make cluster-up +make cluster-migrate +make cluster-status +``` + +Point both binaries at **all three** nodes so `pgx` fails over when one dies: + +```bash +export CRDB_DSN='postgresql://root@127.0.0.1:26257,127.0.0.1:26258,127.0.0.1:26259/defaultdb?sslmode=disable' +``` + +- **SQL:** `26257` / `26258` / `26259` (nodes 1–3) — **Admin UI:** `8089` / `8090` / `8091` — **Redis:** `6379` +- Tear down with `make cluster-down` (this also removes volumes). + +To exercise the reclaim and re-delivery paths, submit load while killing a node. The cluster keeps quorum with two of three nodes, so writes continue while in-flight workers fail mid-commit: + +```bash +docker kill dtq-multinode-crdb1-1 +``` + +Restart it with `docker start dtq-multinode-crdb1-1`. Tasks whose worker failed to commit are picked up by `ReclaimStaleRunningOnce` once `STALE_RUNNING_AFTER` passes with no lease. Lower `RECONCILE_INTERVAL` and `LEASE_DURATION` (e.g. `5s` and `10s`) to see it happen in seconds rather than minutes. + +> The integration tests share one database and `ReconcileOnce` scans it globally, so a running orchestrator or worker will compete with the test suite for tasks and cause spurious failures. Stop both before running `go test`. + ## Database migrations (any cluster) Apply the initial schema: @@ -104,9 +133,15 @@ Loaded by both binaries via [`internal/config`](internal/config/config.go). `CRD | `STALE_RUNNING_AFTER` | `2 × LEASE_DURATION` | Min time a task stays `running` before the reconciler may reclaim it if the Redis lease key is missing | | `WORKER_ID` | random UUID | Stable worker identity if set | | `WORKER_CONCURRENCY` | `1` | Parallel BRPOP worker loops | +| `WORKER_HEARTBEAT_INTERVAL` | `30s` | How often the worker upserts its `workers` row (minimum `1s`). **Not** the Redis lease heartbeat — see below | | `LEASE_DURATION` | `30s` | Logical lease window; Redis key TTL adds a buffer | | `RETRY_BACKOFF` | `5s` | Delay before a failed task is re-queued (when attempts remain) | +There are **two** unrelated heartbeats, and only one is configurable by environment: + +- **Worker registry heartbeat** — `WORKER_HEARTBEAT_INTERVAL`, one per process, upserts the `workers` row for dashboards. Nothing in the claim path depends on it. +- **Lease heartbeat** — one per *running task*, extends the Redis lease TTL so the orchestrator does not reclaim work that is still running. It is `pkg/worker.Options.HeartbeatEvery` (**5s** default), set in code when constructing the runtime; `cmd/worker` passes the zero value, so **no environment variable changes it**. This is the one that matters for reclaim: if it stops, the lease expires and `ReclaimStaleRunningOnce` takes the task back. + ## Run (end-to-end) 1. Apply [migrations](migrations/001_initial.sql) to your CockroachDB cluster. @@ -149,6 +184,12 @@ curl -sS "http://127.0.0.1:8080/v1/tasks/" curl -sS "http://127.0.0.1:8080/v1/jobs/" ``` +`GET /v1/workers` lists workers that have heartbeated recently (see `WORKER_HEARTBEAT_INTERVAL`). The window defaults to 2 minutes; `?active_since=` takes any positive duration. + +```bash +curl -sS "http://127.0.0.1:8080/v1/workers?active_since=5m" +``` + Integration tests that hit the DB and Redis run when `CRDB_DSN` is set (e.g. after `docker compose up`): ```bash @@ -156,12 +197,15 @@ set CRDB_DSN=postgresql://root@127.0.0.1:26257/defaultdb?sslmode=disable go test ./internal/orchestrator/ -count=1 -v ``` -Jobs move **`pending` → `running` → `completed`** (or **`failed`**) as tasks finish; the orchestrator **reconciler** periodically re-enqueues **`queued`** rows that are due (`scheduled_at <= now()`), using Redis **pending** markers to avoid spamming duplicate LPUSHes. If a worker dies after claiming a task, the Redis lease **TTL** expires (heartbeats stop); the reconciler also **reclaims** long-running `running` rows with **no lease**—closing the open `task_run`, re-queuing the task (same retry budget), and LPUSHing again. +Jobs move **`pending` → `running` → `completed`** (or **`failed`**) as tasks finish; the orchestrator **reconciler** periodically re-enqueues **`queued`** rows that are due (`scheduled_at <= now()`), using Redis **pending** markers to avoid spamming duplicate LPUSHes. If a worker dies after claiming a task, the Redis lease **TTL** expires (heartbeats stop); the reconciler also **reclaims** long-running `running` rows with **no lease**—closing the open `task_run` as **`dead`**, re-queuing the task (same retry budget), releasing the pending marker left by the abandoned delivery, and LPUSHing again. + +A reclaimed task keeps its attempt number, so `task_runs` holds a `dead` row and its retry sharing one `attempt_number`. That is intended: `(task_id, attempt_number)` is **not** unique — see the `attempt` invariant in [INSTRUCTIONS.md §4](INSTRUCTIONS.md). ## Layout - [`docker-compose.yml`](docker-compose.yml) — local CockroachDB + Redis -- [`Makefile`](Makefile) — `make compose-up`, `make migrate`, … +- [`docker-compose.multinode.yml`](docker-compose.multinode.yml) — three-node CockroachDB + Redis for failover testing +- [`Makefile`](Makefile) — `make compose-up`, `make migrate`, `make cluster-up`, … - `scripts/` — `migrate.sh` / `migrate.ps1` / `migrate.cmd`, `verify.cmd` (same checks as CI / `make verify`) - `cmd/orchestrator` — HTTP API, DB ping, task submission - `cmd/worker` — worker process (`echo` demo handler) diff --git a/internal/db/jobs.go b/internal/db/jobs.go index 77f0efa..9d73c19 100644 --- a/internal/db/jobs.go +++ b/internal/db/jobs.go @@ -127,40 +127,46 @@ INSERT INTO task_dependencies (task_id, depends_on_task_id) VALUES ($1, $2)`, ti // RefreshJobStatus sets jobs.status from tasks: any failed → failed; all completed → completed; // all cancelled → cancelled; otherwise running. +// +// The counts and the write are one statement on purpose. Splitting them into a SELECT and a +// later UPDATE loses updates: two workers finishing tasks of the same job can both read, and +// the one that read the earlier (less complete) snapshot can commit its status last, leaving a +// finished job stuck at 'running' with nothing left to trigger another refresh. As a single +// statement this is an implicit transaction, so the read of tasks and the write to jobs share +// one timestamp and CockroachDB retries it internally when a concurrent refresh conflicts. func RefreshJobStatus(ctx context.Context, pool *pgxpool.Pool, jobID uuid.UUID) error { + // n_total = 0 leaves the row untouched rather than deriving a status from no tasks; + // it is returned so the caller can tell "job has no tasks" from "job does not exist". const q = ` -SELECT - count(*) FILTER (WHERE status = 'failed') AS n_failed, - count(*) FILTER (WHERE status = 'completed') AS n_done, - count(*) FILTER (WHERE status = 'cancelled') AS n_cancelled, - count(*) AS n_total -FROM tasks WHERE job_id = $1` - var nFailed, nDone, nCancel, nTotal int - if err := pool.QueryRow(ctx, q, jobID).Scan(&nFailed, &nDone, &nCancel, &nTotal); err != nil { +WITH agg AS ( + SELECT + count(*) FILTER (WHERE status = 'failed') AS n_failed, + count(*) FILTER (WHERE status = 'completed') AS n_done, + count(*) FILTER (WHERE status = 'cancelled') AS n_cancelled, + count(*) AS n_total + FROM tasks WHERE job_id = $1 +) +UPDATE jobs SET + status = CASE + WHEN agg.n_total = 0 THEN jobs.status + WHEN agg.n_failed > 0 THEN 'failed' + WHEN agg.n_cancelled = agg.n_total THEN 'cancelled' + WHEN agg.n_done + agg.n_cancelled = agg.n_total THEN 'completed' + ELSE 'running' + END, + updated_at = CASE WHEN agg.n_total = 0 THEN jobs.updated_at ELSE now() END +FROM agg +WHERE jobs.id = $1 +RETURNING agg.n_total` + var nTotal int + if err := pool.QueryRow(ctx, q, jobID).Scan(&nTotal); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return fmt.Errorf("db: refresh job: no job row for %s", jobID) + } return err } if nTotal == 0 { return fmt.Errorf("db: refresh job: no tasks for job %s", jobID) } - var status string - switch { - case nFailed > 0: - status = "failed" - case nDone+nCancel == nTotal: - if nCancel == nTotal { - status = "cancelled" - } else { - status = "completed" - } - default: - status = "running" - } - tag, err := pool.Exec(ctx, `UPDATE jobs SET status = $2, updated_at = now() WHERE id = $1`, jobID, status) - if err != nil { - return err - } - if tag.RowsAffected() == 0 { - return fmt.Errorf("db: refresh job: no job row for %s", jobID) - } return nil } diff --git a/internal/db/jobs_concurrency_test.go b/internal/db/jobs_concurrency_test.go new file mode 100644 index 0000000..fdad515 --- /dev/null +++ b/internal/db/jobs_concurrency_test.go @@ -0,0 +1,274 @@ +package db_test + +import ( + "context" + "fmt" + "os" + "runtime" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/distributed_task_queue/distributed_task_queue/internal/db" +) + +// TestRefreshJobStatus_ParallelTaskCompletions completes every task of one job from parallel +// workers, each refreshing the job afterwards as internal/worker does, and asserts the job +// ends 'completed'. Once all tasks are terminal nothing refreshes the job again, so whatever +// status the last refresh commits is permanent. +func TestRefreshJobStatus_ParallelTaskCompletions(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + const ( + rounds = 5 + tasksPerRound = 16 + ) + + for round := 0; round < rounds; round++ { + t.Run(fmt.Sprintf("round_%d", round), func(t *testing.T) { + jobID, ids := createJobWithTasks(ctx, t, pool, tasksPerRound) + + start := make(chan struct{}) + errs := make(chan error, tasksPerRound) + var wg sync.WaitGroup + for _, taskID := range ids { + wg.Add(1) + go func(taskID uuid.UUID) { + defer wg.Done() + <-start + if err := completeTask(ctx, pool, taskID); err != nil { + errs <- fmt.Errorf("complete %s: %w", taskID, err) + return + } + if err := db.RefreshJobStatus(ctx, pool, jobID); err != nil { + errs <- fmt.Errorf("refresh after %s: %w", taskID, err) + } + }(taskID) + } + close(start) + wg.Wait() + close(errs) + for err := range errs { + t.Error(err) + } + if t.Failed() { + t.FailNow() + } + + job, err := db.GetJobByID(ctx, pool, jobID) + if err != nil { + t.Fatal(err) + } + if job.Status != "completed" { + t.Fatalf("job status = %q, want %q; task statuses: %v", + job.Status, "completed", taskStatuses(ctx, t, pool, jobID)) + } + }) + } +} + +// TestRefreshJobStatus_StaleReadCannotWinTheLastWrite pins the reason RefreshJobStatus has to +// aggregate and write in one statement. +// +// A refresh split into a SELECT of task counts and a later UPDATE loses updates: a worker that +// read while tasks were outstanding computes 'running' and can commit that after the worker +// which read the final state committed 'completed', leaving the job stuck at 'running' forever. +// Parallel completions alone rarely expose this — writes to the contended jobs row are served +// in arrival order, so readers usually write in the order they read — so this test forces the +// reordering: W1's refresh runs on a pool with a single connection, and the test queues behind +// that connection to catch W1 between its read and its write. W2 then finishes the remaining +// tasks and refreshes to 'completed' before W1 is allowed to continue. +// +// An implementation that computed the status from its earlier read fails here every run; one +// that aggregates and writes atomically has no window to stall in. +func TestRefreshJobStatus_StaleReadCannotWinTheLastWrite(t *testing.T) { + fast := testPool(t) + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + // A second pool limited to one connection, so the test can take that connection away + // between statements of a refresh running on it. + slow, err := db.NewPool(ctx, withQueryParam(os.Getenv("CRDB_DSN"), "pool_max_conns=1")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(slow.Close) + + jobID, ids := createJobWithTasks(ctx, t, fast, 3) + + // W1 finishes the first of three tasks and starts its refresh, which reads 1 of 3 done. + if err := completeTask(ctx, fast, ids[0]); err != nil { + t.Fatal(err) + } + w1 := make(chan error, 1) + go func() { w1 <- db.RefreshJobStatus(ctx, slow, jobID) }() + + // Wait until W1's refresh holds the pool's only connection, then queue behind it. The + // Acquire returns as soon as W1's first statement releases the connection, which for a + // split read-then-write refresh is exactly the moment between its read and its write. + for slow.Stat().AcquiredConns() == 0 { + if ctx.Err() != nil { + t.Fatal("refresh never acquired a connection") + } + runtime.Gosched() + } + held, err := slow.Acquire(ctx) + if err != nil { + t.Fatal(err) + } + + // W2 finishes the remaining tasks and refreshes: 3 of 3 done, so it writes 'completed'. + for _, taskID := range ids[1:] { + if err := completeTask(ctx, fast, taskID); err != nil { + t.Fatal(err) + } + } + if err := db.RefreshJobStatus(ctx, fast, jobID); err != nil { + t.Fatal(err) + } + + // W1 continues, now with every task terminal. Its write is the last one the job gets. + held.Release() + if err := <-w1; err != nil { + t.Fatal(err) + } + + job, err := db.GetJobByID(ctx, fast, jobID) + if err != nil { + t.Fatal(err) + } + if job.Status != "completed" { + t.Fatalf("job status = %q, want %q: a refresh overwrote the final status with one "+ + "derived from a stale read, and nothing will refresh this job again", job.Status, "completed") + } +} + +// TestRefreshJobStatus_TerminalStates covers the status mapping itself, so a rewrite of the +// aggregation cannot silently change what a mix of task states means. +func TestRefreshJobStatus_TerminalStates(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + tests := []struct { + name string + statuses []string + want string + }{ + {"all completed", []string{"completed", "completed"}, "completed"}, + {"one failed", []string{"completed", "failed"}, "failed"}, + {"all cancelled", []string{"cancelled", "cancelled"}, "cancelled"}, + {"completed and cancelled", []string{"completed", "cancelled"}, "completed"}, + {"still running", []string{"completed", "running"}, "running"}, + {"failed beats outstanding", []string{"failed", "queued"}, "failed"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + jobID, ids := createJobWithTasks(ctx, t, pool, len(tc.statuses)) + for i, st := range tc.statuses { + if _, err := pool.Exec(ctx, `UPDATE tasks SET status = $2 WHERE id = $1`, ids[i], st); err != nil { + t.Fatal(err) + } + } + if err := db.RefreshJobStatus(ctx, pool, jobID); err != nil { + t.Fatal(err) + } + job, err := db.GetJobByID(ctx, pool, jobID) + if err != nil { + t.Fatal(err) + } + if job.Status != tc.want { + t.Fatalf("job status = %q, want %q", job.Status, tc.want) + } + }) + } +} + +// TestRefreshJobStatus_UnknownJob covers the missing-row path, which the single-statement form +// reports from zero returned rows rather than from RowsAffected. +func TestRefreshJobStatus_UnknownJob(t *testing.T) { + pool := testPool(t) + if err := db.RefreshJobStatus(context.Background(), pool, uuid.New()); err == nil { + t.Fatal("RefreshJobStatus on an unknown job returned nil, want an error") + } +} + +func testPool(t *testing.T) *pgxpool.Pool { + t.Helper() + dsn := os.Getenv("CRDB_DSN") + if dsn == "" { + t.Skip("CRDB_DSN not set; start CockroachDB and export CRDB_DSN to run this test") + } + ctx := context.Background() + pool, err := db.NewPool(ctx, dsn) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + if err := pool.Ping(ctx); err != nil { + t.Skipf("database unavailable: %v", err) + } + return pool +} + +// withQueryParam appends a query parameter to a DSN that may or may not already have a query. +func withQueryParam(dsn, param string) string { + if strings.Contains(dsn, "?") { + return dsn + "&" + param + } + return dsn + "?" + param +} + +// createJobWithTasks creates a job whose n tasks all start queued, returning them in spec order. +func createJobWithTasks(ctx context.Context, t *testing.T, pool *pgxpool.Pool, n int) (uuid.UUID, []uuid.UUID) { + t.Helper() + specs := make([]db.TaskSpec, n) + for i := range specs { + specs[i] = db.TaskSpec{Name: fmt.Sprintf("t%d", i), Kind: "test.noop", Payload: []byte(`{}`)} + } + jobID, nameToID, queued, err := db.CreateJobWithTasks(ctx, pool, "refresh-job-status", 1, specs) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, jobID) + }) + if len(queued) != n { + t.Fatalf("queued = %d, want %d (tasks without dependencies must start queued)", len(queued), n) + } + ids := make([]uuid.UUID, n) + for i := range ids { + ids[i] = nameToID[fmt.Sprintf("t%d", i)] + } + return jobID, ids +} + +// completeTask runs the claim → run → complete sequence a worker performs for one task. +func completeTask(ctx context.Context, pool *pgxpool.Pool, taskID uuid.UUID) error { + task, err := db.ClaimQueuedTask(ctx, pool, taskID) + if err != nil { + return err + } + runID, err := db.InsertTaskRun(ctx, pool, taskID, task.Attempt, "test-worker") + if err != nil { + return err + } + return db.CompleteTaskSuccess(ctx, pool, runID, taskID) +} + +func taskStatuses(ctx context.Context, t *testing.T, pool *pgxpool.Pool, jobID uuid.UUID) map[string]int { + t.Helper() + tasks, err := db.ListTasksByJobID(ctx, pool, jobID) + if err != nil { + t.Fatal(err) + } + counts := make(map[string]int) + for _, task := range tasks { + counts[task.Status]++ + } + return counts +} diff --git a/internal/db/stale.go b/internal/db/stale.go index 1d8120d..5fb3036 100644 --- a/internal/db/stale.go +++ b/internal/db/stale.go @@ -42,12 +42,21 @@ LIMIT $2` return out, rows.Err() } -// ReclaimStaleRunningTask marks open task_runs as failed and returns the task to queued with one attempt -// rolled back so the next claim retries the same logical attempt budget. +// ReclaimStaleRunningTask closes the abandoned task_runs row as 'dead' and returns the task to +// queued with one attempt rolled back, so a reclaim does not consume the task's retry budget. +// +// The rollback means the next ClaimQueuedTask re-issues the same attempt number, so the retried +// run writes a second task_runs row carrying the same (task_id, attempt_number) as the abandoned +// one. That is intended, not a collision: attempt_number is a retry-budget counter, not a run key. +// Runs are identified by task_runs.id and ordered by started_at, and the abandoned predecessor is +// the row with status 'dead'. See the task_runs invariant in INSTRUCTIONS.md section 4. func ReclaimStaleRunningTask(ctx context.Context, pool *pgxpool.Pool, taskID uuid.UUID, reason string) (jobID uuid.UUID, err error) { err = InTx(ctx, pool, "reclaim_stale_running_task", func(ctx context.Context, tx pgx.Tx) error { + // 'dead' rather than 'failed': the handler never reported a result, so this run was + // abandoned, not attempted-and-failed. Keeping the two apart lets history be read + // without matching on the error string. if _, err := tx.Exec(ctx, ` -UPDATE task_runs SET status = 'failed', finished_at = now(), error = $2 +UPDATE task_runs SET status = 'dead', finished_at = now(), error = $2 WHERE task_id = $1 AND status = 'running' AND finished_at IS NULL`, taskID, reason); err != nil { return err } diff --git a/internal/db/tasks.go b/internal/db/tasks.go index 0db9b86..be0401c 100644 --- a/internal/db/tasks.go +++ b/internal/db/tasks.go @@ -93,6 +93,9 @@ ORDER BY name ASC, id ASC` // ClaimQueuedTask atomically increments attempt, sets status=running, and timestamps for a queued task. // If the task is not queued, returns ErrTaskNotClaimable. +// +// attempt counts consumed retry budget, not deliveries: ReclaimStaleRunningTask rolls it back by +// one, so claiming a reclaimed task re-issues the attempt number its abandoned run already used. func ClaimQueuedTask(ctx context.Context, pool *pgxpool.Pool, id uuid.UUID) (Task, error) { const q = ` UPDATE tasks SET @@ -117,6 +120,10 @@ RETURNING id, job_id, name, kind, status, payload, max_attempts, attempt, schedu } // InsertTaskRun records the start of an attempt; status is running. +// +// attemptNumber is deliberately not unique per task: a task reclaimed after a lost lease retries +// the same attempt number, so its history holds one 'dead' row and one live row that share it. +// The returned run id is the run key; order history by started_at, never by attempt_number. func InsertTaskRun(ctx context.Context, pool *pgxpool.Pool, taskID uuid.UUID, attemptNumber int, workerID string) (runID uuid.UUID, err error) { const q = ` INSERT INTO task_runs (task_id, attempt_number, worker_id, status) diff --git a/internal/orchestrator/integration_api_test.go b/internal/orchestrator/integration_api_test.go index d3ecddf..3d8ab44 100644 --- a/internal/orchestrator/integration_api_test.go +++ b/internal/orchestrator/integration_api_test.go @@ -12,6 +12,8 @@ import ( "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + goredis "github.com/redis/go-redis/v9" "github.com/distributed_task_queue/distributed_task_queue/internal/config" "github.com/distributed_task_queue/distributed_task_queue/internal/db" @@ -69,6 +71,7 @@ func TestIntegration_SubmitAndReadTaskAndJob(t *testing.T) { if sub.TaskID == "" { t.Fatal("empty task_id") } + cleanupSubmittedTask(t, pool, rdb, prefix, sub.TaskID) getTask, err := http.Get(ts.URL + "/v1/tasks/" + sub.TaskID) if err != nil { @@ -117,6 +120,35 @@ func TestIntegration_SubmitAndReadTaskAndJob(t *testing.T) { } } +// cleanupSubmittedTask removes what submitting a task leaves behind: the job row, which cascades +// to the task and its runs, and the Redis keys under this test's prefix. +// +// The database and Redis are shared with the rest of the suite and across runs, so skipping this +// leaves a queued, due task that every later reconcile - in any test - keeps enqueueing, plus a +// ready list that has no TTL and so never goes away. That accumulation is what makes assertions +// on global counts (ReconcileOnce's return, ready list length) unstable, so tests that submit +// must clean up even though nothing in a single run depends on it. +func cleanupSubmittedTask(t *testing.T, pool *pgxpool.Pool, rdb *goredis.Client, prefix, taskID string) { + t.Helper() + t.Cleanup(func() { + // A fresh context: the test's may already be cancelled by the time cleanup runs. + ctx := context.Background() + var jobID uuid.UUID + if err := pool.QueryRow(ctx, `SELECT job_id FROM tasks WHERE id = $1`, taskID).Scan(&jobID); err != nil { + t.Logf("cleanup: look up job for task %s: %v", taskID, err) + } else if _, err := pool.Exec(ctx, `DELETE FROM jobs WHERE id = $1`, jobID); err != nil { + t.Logf("cleanup: delete job %s: %v", jobID, err) + } + if err := rdb.Del(ctx, + appredis.ReadyList(prefix, appredis.DefaultPriority), + appredis.ScheduledZSet(prefix), + appredis.PendingEnqueueKey(prefix, taskID), + ).Err(); err != nil { + t.Logf("cleanup: delete redis keys for %s: %v", prefix, err) + } + }) +} + func TestIntegration_GetTask_notFound(t *testing.T) { dsn := os.Getenv("CRDB_DSN") if dsn == "" { @@ -281,6 +313,8 @@ func TestIntegration_SubmitDelayedTask(t *testing.T) { t.Fatal(err) } + cleanupSubmittedTask(t, pool, rdb, prefix, sub.TaskID) + readyKey := appredis.ReadyList(prefix, appredis.DefaultPriority) readyLen, err := appredis.ReadyLen(ctx, rdb, readyKey) if err != nil { @@ -298,19 +332,46 @@ func TestIntegration_SubmitDelayedTask(t *testing.T) { t.Fatalf("zset score: got %v want %d", score, runAt.UnixMilli()) } - time.Sleep(3 * time.Second) - n, err := orchestrator.ReconcileOnce(ctx, pool, rdb, cfg) - if err != nil { - t.Fatal(err) - } - if n < 1 { - t.Fatalf("reconcile enqueued: got %d want >= 1", n) + // Poll instead of sleeping a fixed 3s and reconciling exactly once. run_at is set from this + // process's clock but "due" is decided by the database's, so a single reconcile fired at a + // hardcoded instant makes the test a race against clock offset and scheduling delay. + // + // Assert on this task id, never on ReconcileOnce's count or the ready list's length: both + // aggregate over every due task in the database, which is shared with the rest of the suite + // and retains rows from earlier runs. A count assertion passes when unrelated leftovers are + // enqueued and fails when they are absent, neither of which says anything about this task. + deadline := time.Now().Add(20 * time.Second) + for { + if _, err := orchestrator.ReconcileOnce(ctx, pool, rdb, cfg); err != nil { + t.Fatal(err) + } + if readyListContains(t, ctx, rdb, readyKey, sub.TaskID) { + break + } + if time.Now().After(deadline) { + // Say why it never arrived, so a real regression is not read as a slow machine. + var status string + var schedAt, dbNow time.Time + if err := pool.QueryRow(ctx, ` +SELECT status, scheduled_at, now() FROM tasks WHERE id = $1`, sub.TaskID).Scan(&status, &schedAt, &dbNow); err != nil { + t.Fatal(err) + } + t.Fatalf("task %s never reached the ready list: status=%q scheduled_at=%s db now()=%s due=%v", + sub.TaskID, status, schedAt.Format(time.RFC3339Nano), dbNow.Format(time.RFC3339Nano), + !schedAt.After(dbNow)) + } + time.Sleep(200 * time.Millisecond) } - readyLen, err = appredis.ReadyLen(ctx, rdb, readyKey) - if err != nil { + + // Exactly once: the ZSET entry must be removed as part of enqueueing, so that repeated + // reconcile passes over a still-undelivered task do not push it again. + if _, err := orchestrator.ReconcileOnce(ctx, pool, rdb, cfg); err != nil { t.Fatal(err) } - if readyLen < 1 { - t.Fatalf("ready list after reconcile: got %d want >= 1", readyLen) + if got := countInReadyList(t, ctx, rdb, readyKey, sub.TaskID); got != 1 { + t.Errorf("task %s appears %d times on the ready list after a second reconcile, want exactly 1", sub.TaskID, got) + } + if err := rdb.ZScore(ctx, zKey, sub.TaskID).Err(); err != goredis.Nil { + t.Errorf("scheduled zset still holds %s after enqueue: %v", sub.TaskID, err) } } diff --git a/internal/orchestrator/reclaim.go b/internal/orchestrator/reclaim.go index 2b09113..0dfcbbc 100644 --- a/internal/orchestrator/reclaim.go +++ b/internal/orchestrator/reclaim.go @@ -14,7 +14,8 @@ import ( const reclaimBatch = 200 -// ReclaimStaleRunningOnce finds long-running tasks with no Redis lease, returns them to queued, and LPUSHes after TryReservePending. +// ReclaimStaleRunningOnce finds long-running tasks with no Redis lease, returns them to queued, +// drops the stale pending marker left by the abandoned delivery, and LPUSHes after TryReservePending. func ReclaimStaleRunningOnce(ctx context.Context, pool *pgxpool.Pool, rdb *goredis.Client, cfg config.Config) (reclaimed int, err error) { ids, err := db.ListStaleRunningCandidates(ctx, pool, cfg.StaleRunningAfter, reclaimBatch) if err != nil { @@ -45,6 +46,15 @@ func ReclaimStaleRunningOnce(ctx context.Context, pool *pgxpool.Pool, rdb *gored } _ = db.RefreshJobStatus(ctx, pool, jobID) + // The delivery that put this task in flight has been abandoned, so its pending + // marker is stale. Without this the marker outlives the reclaim, EnqueueDueTaskID + // reads it as "already enqueued" and skips the LPUSH, and only the worker's + // post-claim ReleasePending would clear it - which cannot run because the task is + // never delivered. Release before enqueueing so the reservation below can succeed. + if err := appredis.ReleasePending(ctx, rdb, cfg.RedisKeyPrefix, id.String()); err != nil { + return reclaimed, err + } + enqueued, err := EnqueueDueTaskID(ctx, pool, rdb, cfg, id) if err != nil { return reclaimed, err diff --git a/internal/orchestrator/reclaim_dead_run_test.go b/internal/orchestrator/reclaim_dead_run_test.go new file mode 100644 index 0000000..fdb3adc --- /dev/null +++ b/internal/orchestrator/reclaim_dead_run_test.go @@ -0,0 +1,145 @@ +package orchestrator_test + +import ( + "context" + "os" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/distributed_task_queue/distributed_task_queue/internal/config" + "github.com/distributed_task_queue/distributed_task_queue/internal/db" + "github.com/distributed_task_queue/distributed_task_queue/internal/orchestrator" + appredis "github.com/distributed_task_queue/distributed_task_queue/internal/redis" +) + +// TestIntegration_ReclaimMarksRunDeadAndReusesAttempt pins the task_runs invariant documented in +// INSTRUCTIONS.md section 4. Reclaim rolls tasks.attempt back so losing a worker does not cost the +// task a retry, which means the retried run reuses the abandoned run's attempt_number. That is +// intended, so this test asserts both halves: the number IS reused (no unique constraint may be +// added), and the abandoned predecessor is still identifiable - by status 'dead', not by parsing +// the error text and not by attempt number. +func TestIntegration_ReclaimMarksRunDeadAndReusesAttempt(t *testing.T) { + dsn := os.Getenv("CRDB_DSN") + if dsn == "" { + t.Skip("CRDB_DSN not set; start CockroachDB and export CRDB_DSN to run this test") + } + redisAddr := os.Getenv("REDIS_ADDR") + if redisAddr == "" { + redisAddr = "127.0.0.1:6379" + } + ctx := context.Background() + + pool, err := db.NewPool(ctx, dsn) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + if err := pool.Ping(ctx); err != nil { + t.Skipf("database unavailable: %v", err) + } + + rdb := appredis.New(redisAddr) + t.Cleanup(func() { _ = rdb.Close() }) + if err := rdb.Ping(ctx).Err(); err != nil { + t.Skipf("redis unavailable: %v", err) + } + + prefix := "dto:itest:" + uuid.NewString() + ":" + cfg := config.Config{ + RedisKeyPrefix: prefix, + LeaseDuration: 30 * time.Second, + StaleRunningAfter: time.Second, + } + t.Cleanup(func() { + _ = rdb.Del(context.Background(), appredis.ReadyList(prefix, appredis.DefaultPriority)).Err() + }) + + // A task in flight on attempt 1, with the open task_runs row its worker wrote before dying. + taskID := insertRunningTaskForReclaim(t, ctx, pool, 5*time.Minute) + abandonedRunID, err := db.InsertTaskRun(ctx, pool, taskID, 1, "worker-that-died") + if err != nil { + t.Fatal(err) + } + + if _, err := orchestrator.ReclaimStaleRunningOnce(ctx, pool, rdb, cfg); err != nil { + t.Fatal(err) + } + + // The abandoned run is closed as 'dead', not 'failed': no handler ever reported a result. + var gotStatus string + var finishedAt *time.Time + var runErr *string + if err := pool.QueryRow(ctx, ` +SELECT status, finished_at, error FROM task_runs WHERE id = $1`, abandonedRunID). + Scan(&gotStatus, &finishedAt, &runErr); err != nil { + t.Fatal(err) + } + if gotStatus != "dead" { + t.Errorf("abandoned run status = %q, want %q", gotStatus, "dead") + } + if finishedAt == nil { + t.Error("abandoned run finished_at is NULL, want it closed") + } + if runErr == nil || *runErr == "" { + t.Error("abandoned run error is empty, want the reclaim reason") + } + + // The retry budget must not have been consumed by the reclaim. + task, err := db.GetTaskByID(ctx, pool, taskID) + if err != nil { + t.Fatal(err) + } + if task.Attempt != 0 { + t.Fatalf("attempt after reclaim = %d, want 0 (reclaim must not consume a retry)", task.Attempt) + } + + // Replay what a worker does next. The claim re-issues attempt 1, and inserting its run must + // succeed even though attempt_number 1 already exists for this task. + claimed, err := db.ClaimQueuedTask(ctx, pool, taskID) + if err != nil { + t.Fatal(err) + } + if claimed.Attempt != 1 { + t.Fatalf("attempt after re-claim = %d, want 1 (the reclaimed attempt is retried)", claimed.Attempt) + } + retryRunID, err := db.InsertTaskRun(ctx, pool, taskID, claimed.Attempt, "worker-that-took-over") + if err != nil { + t.Fatalf("insert retry run reusing attempt_number %d: %v (a UNIQUE (task_id, attempt_number) "+ + "constraint would break reclaim - attempt_number is a budget counter, not a run key)", + claimed.Attempt, err) + } + if retryRunID == abandonedRunID { + t.Fatal("retry reused the abandoned run id; each run must be its own row") + } + + // Both rows share attempt_number and are told apart by status alone. + rows, err := pool.Query(ctx, ` +SELECT id, status FROM task_runs WHERE task_id = $1 AND attempt_number = 1 ORDER BY started_at ASC`, taskID) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + byID := map[uuid.UUID]string{} + for rows.Next() { + var id uuid.UUID + var status string + if err := rows.Scan(&id, &status); err != nil { + t.Fatal(err) + } + byID[id] = status + } + if err := rows.Err(); err != nil { + t.Fatal(err) + } + if len(byID) != 2 { + t.Fatalf("attempt_number 1 has %d task_runs rows, want 2 (the dead run and its retry)", len(byID)) + } + if got := byID[abandonedRunID]; got != "dead" { + t.Errorf("abandoned run status = %q, want %q", got, "dead") + } + if got := byID[retryRunID]; got != "running" { + t.Errorf("retry run status = %q, want %q", got, "running") + } +} diff --git a/internal/orchestrator/reclaim_pending_test.go b/internal/orchestrator/reclaim_pending_test.go new file mode 100644 index 0000000..c1a0050 --- /dev/null +++ b/internal/orchestrator/reclaim_pending_test.go @@ -0,0 +1,156 @@ +package orchestrator_test + +import ( + "context" + "os" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + goredis "github.com/redis/go-redis/v9" + + "github.com/distributed_task_queue/distributed_task_queue/internal/config" + "github.com/distributed_task_queue/distributed_task_queue/internal/db" + "github.com/distributed_task_queue/distributed_task_queue/internal/orchestrator" + appredis "github.com/distributed_task_queue/distributed_task_queue/internal/redis" +) + +// TestIntegration_ReclaimReleasesStalePendingMarker covers the node-kill strand: a task whose +// worker died mid-flight is reclaimed to 'queued', but the pending marker from its original +// enqueue survives. If reclaim leaves that marker in place, EnqueueDueTaskID reads it as +// "already enqueued" and never LPUSHes, and the worker's post-claim ReleasePending can never +// run because the task is never delivered - the task stays queued and due but unreachable. +func TestIntegration_ReclaimReleasesStalePendingMarker(t *testing.T) { + dsn := os.Getenv("CRDB_DSN") + if dsn == "" { + t.Skip("CRDB_DSN not set; start CockroachDB and export CRDB_DSN to run this test") + } + redisAddr := os.Getenv("REDIS_ADDR") + if redisAddr == "" { + redisAddr = "127.0.0.1:6379" + } + ctx := context.Background() + + pool, err := db.NewPool(ctx, dsn) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + if err := pool.Ping(ctx); err != nil { + t.Skipf("database unavailable: %v", err) + } + + rdb := appredis.New(redisAddr) + t.Cleanup(func() { _ = rdb.Close() }) + if err := rdb.Ping(ctx).Err(); err != nil { + t.Skipf("redis unavailable: %v", err) + } + + prefix := "dto:itest:" + uuid.NewString() + ":" + cfg := config.Config{ + RedisKeyPrefix: prefix, + LeaseDuration: 30 * time.Second, + StaleRunningAfter: time.Second, + } + readyKey := appredis.ReadyList(prefix, appredis.DefaultPriority) + t.Cleanup(func() { + _ = rdb.Del(context.Background(), readyKey).Err() + }) + + // A task that a worker picked up and never finished: status running, started_at old + // enough to be a stale candidate, no Redis lease. + taskID := insertRunningTaskForReclaim(t, ctx, pool, 5*time.Minute) + + // The marker left behind by the delivery that is about to be abandoned. + reserved, err := appredis.TryReservePending(ctx, rdb, prefix, taskID.String(), 5*time.Minute) + if err != nil { + t.Fatal(err) + } + if !reserved { + t.Fatalf("precondition: expected to reserve pending marker for fresh task %s", taskID) + } + t.Cleanup(func() { + _ = appredis.ReleasePending(context.Background(), rdb, prefix, taskID.String()) + }) + + if _, err := orchestrator.ReclaimStaleRunningOnce(ctx, pool, rdb, cfg); err != nil { + t.Fatal(err) + } + + // Reclaim must have returned the task to queued... + status, due, found, err := db.TaskScheduleState(ctx, pool, taskID) + if err != nil { + t.Fatal(err) + } + if !found || status != "queued" || !due { + t.Fatalf("after reclaim: got found=%v status=%q due=%v, want found=true status=queued due=true", found, status, due) + } + + // ...and the task must actually be back on the ready list. This is the assertion that + // fails on the unfixed code: the stale marker makes EnqueueDueTaskID's reservation fail, + // so it returns without pushing and the task is queued, due, and unreachable forever. + // Note we do NOT assert the marker is gone - reclaim's own enqueue immediately takes a + // fresh reservation, which is correct; a marker paired with a real LPUSH is the healthy + // state, a marker with nothing on the ready list is the bug. + if !readyListContains(t, ctx, rdb, readyKey, taskID.String()) { + t.Fatalf("task %s is queued and due but absent from %s: stale pending marker suppressed the enqueue", taskID, readyKey) + } + + // The marker must still do its job afterwards, so that a naive "always release" fix that + // breaks duplicate suppression does not pass this test: a reconcile pass over the same + // still-undelivered task must not push it a second time. + if _, err := orchestrator.ReconcileOnce(ctx, pool, rdb, cfg); err != nil { + t.Fatal(err) + } + if n := countInReadyList(t, ctx, rdb, readyKey, taskID.String()); n != 1 { + t.Errorf("task %s appears %d times on the ready list after a reconcile pass, want exactly 1", taskID, n) + } +} + +// insertRunningTaskForReclaim creates a job and a task already in flight, with started_at +// pushed back by age so ListStaleRunningCandidates sees it. +func insertRunningTaskForReclaim(t *testing.T, ctx context.Context, pool *pgxpool.Pool, age time.Duration) uuid.UUID { + t.Helper() + var jobID uuid.UUID + if err := pool.QueryRow(ctx, ` +INSERT INTO jobs (name, status) VALUES ('reclaim-pending-regression', 'running') +RETURNING id`).Scan(&jobID); err != nil { + t.Fatal(err) + } + var taskID uuid.UUID + if err := pool.QueryRow(ctx, ` +INSERT INTO tasks (job_id, name, kind, status, payload, max_attempts, attempt, scheduled_at, started_at) +VALUES ($1, 'stranded', 'echo', 'running', '{}'::JSONB, 3, 1, now(), now() - ($2::BIGINT * '1 microsecond'::interval)) +RETURNING id`, jobID, age.Microseconds()).Scan(&taskID); err != nil { + t.Fatal(err) + } + // The task is deliberately left queued and due, so without this it would linger in a + // shared dev database and look exactly like a real stranded task to anyone inspecting it. + t.Cleanup(func() { + if _, err := pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, jobID); err != nil { + t.Logf("cleanup job %s: %v", jobID, err) + } + }) + return taskID +} + +func countInReadyList(t *testing.T, ctx context.Context, rdb *goredis.Client, key, want string) int { + t.Helper() + ids, err := rdb.LRange(ctx, key, 0, -1).Result() + if err != nil { + t.Fatal(err) + } + n := 0 + for _, id := range ids { + if id == want { + n++ + } + } + return n +} + +func readyListContains(t *testing.T, ctx context.Context, rdb *goredis.Client, key, want string) bool { + t.Helper() + return countInReadyList(t, ctx, rdb, key, want) > 0 +} diff --git a/internal/redis/pending.go b/internal/redis/pending.go index d71c99f..1eb86fc 100644 --- a/internal/redis/pending.go +++ b/internal/redis/pending.go @@ -8,19 +8,49 @@ import ( goredis "github.com/redis/go-redis/v9" ) -// PendingEnqueueKey marks that a task id was recently placed on the ready list (or is in flight) to reduce duplicate LPUSH from the reconciler. +// The pending marker records that a task id has an in-flight delivery, so the reconciler does +// not LPUSH a task that is already on the ready list or already being claimed. +// +// The marker is a lock on the delivery, and whoever ends a delivery MUST release it. While it is +// set, EnqueueDueTaskID will not push the task, so a marker that outlives its delivery makes the +// task queued, due, and unreachable until the TTL expires. Releasing is required on every path +// that ends a delivery, not just the happy one: +// +// - worker claims the task successfully -> internal/worker.Runtime.processTask +// - worker pops the id but the claim fails or is lost -> internal/worker.Runtime.processTask +// - orchestrator reclaims a stale running task, abandoning the delivery that put it in flight +// -> internal/orchestrator.ReclaimStaleRunningOnce +// +// The TTL passed to TryReservePending is a backstop for a release that never happened (the +// process died between LPUSH and claim), not the primary mechanism - see its doc comment. + +// PendingEnqueueKey names the marker for one task id. // Format: {prefix}queue:pending:{task_id} func PendingEnqueueKey(prefix, taskID string) string { return fmt.Sprintf("%squeue:pending:%s", prefix, taskID) } -// TryReservePending sets a short-lived key with SET NX so only one producer enqueues at a time. -// Returns true if this caller reserved the key (should enqueue to Redis). +// TryReservePending claims the right to enqueue taskID with SET NX, so concurrent producers do +// not all LPUSH it. Returns true if this caller reserved the marker and should enqueue. +// +// ttl bounds how long a leaked marker can suppress delivery, so it is an upper bound on the +// stall a missed release can cause. It does not need to exceed the worst-case time between LPUSH +// and claim: expiring while the task is still on the ready list only lets a producer push a +// duplicate id, and a duplicate is harmless because ClaimQueuedTask is a single conditional +// UPDATE - the second worker to pop it gets ErrTaskNotClaimable and drops it. Prefer a short ttl: +// the cost of expiring early is a wasted BRPOP, the cost of expiring late is a stalled task. func TryReservePending(ctx context.Context, rdb *goredis.Client, prefix, taskID string, ttl time.Duration) (bool, error) { return rdb.SetNX(ctx, PendingEnqueueKey(prefix, taskID), "1", ttl).Result() } -// ReleasePending deletes the pending marker (e.g. after a successful claim). +// ReleasePending deletes the marker, ending the delivery and letting the task be enqueued again. +// Call it on every path that ends a delivery (see the lifecycle note above), including failure +// paths. It is idempotent, and safe to call after losing a claim race: the winner releases its +// own marker, and EnqueueDueTaskID re-checks task status before pushing, so releasing cannot +// cause a running task to be re-enqueued. +// +// Callers on a failure path should pass a context not already cancelled by the failure itself, +// otherwise the release is dropped exactly when it matters most. func ReleasePending(ctx context.Context, rdb *goredis.Client, prefix, taskID string) error { return rdb.Del(ctx, PendingEnqueueKey(prefix, taskID)).Err() } diff --git a/internal/redis/queue.go b/internal/redis/queue.go index 507224d..9bb7b15 100644 --- a/internal/redis/queue.go +++ b/internal/redis/queue.go @@ -34,7 +34,13 @@ func ReadyLen(ctx context.Context, rdb *goredis.Client, key string) (int64, erro return rdb.LLen(ctx, key).Result() } -// EnqueueTaskIDs LPUSHes each task ID on the default ready list after TryReservePending. +// EnqueueTaskIDs LPUSHes each task ID on the ready list for priority, taking a pending marker for +// each so the reconciler does not immediately push it again. +// +// This is the submit path, for ids that are new and therefore cannot already be in flight, so it +// pushes unconditionally and ignores whether the reservation was won. Do not copy that pattern +// into a re-delivery path: there, an existing marker means a delivery is already in flight and +// pushing anyway would duplicate it. Use EnqueueDueTaskID, which honours the reservation. func EnqueueTaskIDs(ctx context.Context, rdb *goredis.Client, prefix, priority string, ids []string) error { key := ReadyList(prefix, priority) for _, id := range ids { diff --git a/internal/worker/runtime.go b/internal/worker/runtime.go index e6425a7..4cf551c 100644 --- a/internal/worker/runtime.go +++ b/internal/worker/runtime.go @@ -117,6 +117,16 @@ func (r *Runtime) runLoop(ctx context.Context) { } } +// releasePending clears the enqueue marker for id on a context independent of the caller's, so +// the release still lands when the task context is already cancelled or its database call failed. +func (r *Runtime) releasePending(id uuid.UUID) { + cctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := appredis.ReleasePending(cctx, r.rdb, r.cfg.RedisKeyPrefix, id.String()); err != nil { + log.Printf("worker: release pending %s: %v", id, err) + } +} + func (r *Runtime) processTask(ctx context.Context, id uuid.UUID) { leaseKey := appredis.LeaseTask(r.cfg.RedisKeyPrefix, id.String()) defer func() { @@ -127,13 +137,21 @@ func (r *Runtime) processTask(ctx context.Context, id uuid.UUID) { task, err := db.ClaimQueuedTask(ctx, r.pool, id) if err != nil { + // BRPOP already took this id off the ready list, so this worker owns the delivery. + // The claim did not stick, so the delivery has to be handed back: leaving the marker + // set makes the task queued, due, and off the ready list, with EnqueueDueTaskID + // reading the marker as "already enqueued" and never re-pushing it. Releasing is safe + // even when another worker won the claim - it has already released the marker itself, + // and EnqueueDueTaskID re-checks status before pushing, so a running task is not + // re-enqueued. Use a fresh context: ctx is typically cancelled on shutdown. + r.releasePending(id) if errors.Is(err, db.ErrTaskNotClaimable) { return } log.Printf("worker: claim %s: %v", id, err) return } - _ = appredis.ReleasePending(ctx, r.rdb, r.cfg.RedisKeyPrefix, task.ID.String()) + r.releasePending(task.ID) r.refreshJob(context.Background(), task.JobID) runID, err := db.InsertTaskRun(ctx, r.pool, task.ID, task.Attempt, r.cfg.WorkerID) diff --git a/internal/worker/runtime_pending_test.go b/internal/worker/runtime_pending_test.go new file mode 100644 index 0000000..f6779dc --- /dev/null +++ b/internal/worker/runtime_pending_test.go @@ -0,0 +1,144 @@ +package worker_test + +import ( + "context" + "os" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/distributed_task_queue/distributed_task_queue/internal/config" + "github.com/distributed_task_queue/distributed_task_queue/internal/db" + appredis "github.com/distributed_task_queue/distributed_task_queue/internal/redis" + wruntime "github.com/distributed_task_queue/distributed_task_queue/internal/worker" + pkgworker "github.com/distributed_task_queue/distributed_task_queue/pkg/worker" +) + +// TestIntegration_FailedClaimReleasesPendingMarker covers the second strand path seen in the +// node-kill experiment: BRPOP takes a task id off the ready list, the claim then fails (during a +// CockroachDB outage the worker logs "claim : unexpected EOF"), and the worker returns without +// releasing the pending marker. The task is left queued, due, absent from the ready list, and +// EnqueueDueTaskID reads the surviving marker as "already enqueued" and never re-pushes it. +// Reclaim cannot rescue it either, because the task never reached status 'running'. +// +// The failed claim is simulated with a task that is not claimable (ErrTaskNotClaimable), which +// reaches the same release path as a transport error without needing to kill a database node. +func TestIntegration_FailedClaimReleasesPendingMarker(t *testing.T) { + dsn := os.Getenv("CRDB_DSN") + if dsn == "" { + t.Skip("CRDB_DSN not set; start CockroachDB and export CRDB_DSN to run this test") + } + redisAddr := os.Getenv("REDIS_ADDR") + if redisAddr == "" { + redisAddr = "127.0.0.1:6379" + } + ctx := context.Background() + + pool, err := db.NewPool(ctx, dsn) + if err != nil { + t.Fatal(err) + } + t.Cleanup(pool.Close) + if err := pool.Ping(ctx); err != nil { + t.Skipf("database unavailable: %v", err) + } + + rdb := appredis.New(redisAddr) + t.Cleanup(func() { _ = rdb.Close() }) + if err := rdb.Ping(ctx).Err(); err != nil { + t.Skipf("redis unavailable: %v", err) + } + + prefix := "dto:wtest:" + uuid.NewString() + ":" + cfg := config.Config{ + RedisKeyPrefix: prefix, + RedisAddr: redisAddr, + WorkerID: "test-worker", + WorkerConcurrency: 1, + LeaseDuration: 30 * time.Second, + // Run starts a heartbeat ticker unconditionally; a zero interval panics. + WorkerHeartbeatInterval: 5 * time.Second, + } + readyKey := appredis.ReadyList(prefix, appredis.DefaultPriority) + t.Cleanup(func() { _ = rdb.Del(context.Background(), readyKey).Err() }) + + // A task the worker cannot claim, standing in for a claim that fails to commit. + taskID := insertUnclaimableTask(t, ctx, pool) + + // The marker set by the enqueue that put this id on the ready list. + reserved, err := appredis.TryReservePending(ctx, rdb, prefix, taskID.String(), 5*time.Minute) + if err != nil { + t.Fatal(err) + } + if !reserved { + t.Fatalf("precondition: expected to reserve pending marker for fresh task %s", taskID) + } + t.Cleanup(func() { + _ = appredis.ReleasePending(context.Background(), rdb, prefix, taskID.String()) + }) + + if err := appredis.EnqueueReady(ctx, rdb, readyKey, taskID.String()); err != nil { + t.Fatal(err) + } + + runCtx, cancel := context.WithCancel(ctx) + defer cancel() + rt := wruntime.NewRuntime(pkgworker.Options{}, cfg, pool, rdb) + rt.Register("echo", func(context.Context, pkgworker.Task) error { return nil }) + done := make(chan struct{}) + go func() { + defer close(done) + _ = rt.Run(runCtx) + }() + + // The worker should pop the id, fail to claim, and hand the delivery back by clearing the + // marker. Without the release this poll runs to timeout and the task is unreachable until + // the marker's TTL expires. + key := appredis.PendingEnqueueKey(prefix, taskID.String()) + deadline := time.Now().Add(15 * time.Second) + released := false + for time.Now().Before(deadline) { + n, err := rdb.Exists(ctx, key).Result() + if err != nil { + t.Fatal(err) + } + if n == 0 { + released = true + break + } + time.Sleep(100 * time.Millisecond) + } + cancel() + <-done + + if !released { + t.Fatalf("pending marker for %s survived a failed claim: the task is queued, off the ready list, and cannot be re-enqueued", taskID) + } +} + +// insertUnclaimableTask creates a task already past the queued state, so ClaimQueuedTask returns +// ErrTaskNotClaimable. +func insertUnclaimableTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool) uuid.UUID { + t.Helper() + var jobID uuid.UUID + if err := pool.QueryRow(ctx, ` +INSERT INTO jobs (name, status) VALUES ('worker-pending-regression', 'running') +RETURNING id`).Scan(&jobID); err != nil { + t.Fatal(err) + } + var taskID uuid.UUID + if err := pool.QueryRow(ctx, ` +INSERT INTO tasks (job_id, name, kind, status, payload, max_attempts, attempt, scheduled_at) +VALUES ($1, 'unclaimable', 'echo', 'completed', '{}'::JSONB, 3, 1, now()) +RETURNING id`, jobID).Scan(&taskID); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if _, err := pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, jobID); err != nil { + t.Logf("cleanup job %s: %v", jobID, err) + } + }) + return taskID +} diff --git a/migrations/001_initial.sql b/migrations/001_initial.sql index 36554e1..3586431 100644 --- a/migrations/001_initial.sql +++ b/migrations/001_initial.sql @@ -44,6 +44,10 @@ CREATE TABLE IF NOT EXISTS task_dependencies ( CREATE INDEX IF NOT EXISTS idx_task_dependencies_depends ON task_dependencies (depends_on_task_id); +-- attempt_number is NOT unique per task_id and must not be constrained to be: reclaiming a stale +-- running task rolls tasks.attempt back so the reclaim does not consume a retry, and the retried +-- run reuses the number its abandoned ('dead') run already wrote. The run key is id; order history +-- by started_at. See the task_runs invariant in INSTRUCTIONS.md section 4. CREATE TABLE IF NOT EXISTS task_runs ( id UUID NOT NULL PRIMARY KEY DEFAULT gen_random_uuid(), task_id UUID NOT NULL REFERENCES tasks (id) ON DELETE CASCADE, diff --git a/pkg/worker/worker.go b/pkg/worker/worker.go index 207c736..e5feffe 100644 --- a/pkg/worker/worker.go +++ b/pkg/worker/worker.go @@ -26,6 +26,11 @@ type Options struct { // DefaultTimeout caps handler execution when the task does not specify its own deadline. DefaultTimeout time.Duration // HeartbeatEvery is the interval between Redis lease extensions while a handler runs. + // Keep it well under the lease duration: if extensions stop, the lease key expires and the + // orchestrator reclaims the task as stale while the handler is still running. + // + // This is not the WORKER_HEARTBEAT_INTERVAL environment variable, which is the separate + // per-process `workers` table registry heartbeat. HeartbeatEvery is set in code only. HeartbeatEvery time.Duration }