From 459765a6668c09db761be9ed44d48736d1f9873d Mon Sep 17 00:00:00 2001 From: Waldemort Agent Date: Sun, 16 Aug 2026 23:46:15 +0000 Subject: [PATCH] fix(orchestrator): never return an already-expired lock lease fetch_orchestration_item could return an orchestration item whose lock lease had already expired at the moment it was returned. Every ack against that lease failed with 'Invalid lock token', so the turn's work was discarded, the message became visible again immediately, and the cycle repeated indefinitely -- including the max_attempts poison path, which commits through the same ack and so could never terminate it. Migration 0024 rebuilds the procedure on top of 0019: - pg_try_advisory_xact_lock replaces the blocking pg_advisory_xact_lock. Candidate selection is deterministic (ORDER BY visible_at, id LIMIT 1), so every dispatcher converged on the same key and queued behind it. Contended instances are now skipped and excluded from re-selection within the call. Same-instance mutual exclusion -- the deadlock guard -- is preserved; only the waiting is removed. - The lease is derived from an effective clock, v_now_ms, equal to the caller's p_now_ms advanced by time actually elapsed inside the call. Only the elapsed duration comes from the server clock, never an absolute server timestamp, so the Rust-supplied epoch established by 0006/0012 remains the single time reference. - The advisory key is scoped to '.'. PostgreSQL advisory locks share one database-wide space, so the previous unscoped key made two duroxide schemas in the same database serialize against each other on identical instance ids. Blocking acquisition hid this; non-blocking acquisition would surface it as a spurious skip. - The retry loop is bounded; exhausting it returns no item, which the dispatcher already treats as 'nothing available'. Signature, return shape, and all downstream behavior are unchanged. Measured on PostgreSQL 17.11 and 18.6 (contended fetch, 1s lock timeout, 3s competing hold): before: blocked 2.61s, lease returned 1658ms already expired after: blocked 0.05s, contended instance skipped, live lease tests/expired_lease_regression_test.rs fails without this migration ('the advisory lock is convoying') and passes with it. Full suite: 353 passed / 0 failed on both 17.11 and 18.6. Fixes #21 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...4_fix_expired_orchestration_lock_lease.sql | 329 ++++++++++++++++++ tests/expired_lease_regression_test.rs | 189 ++++++++++ 2 files changed, 518 insertions(+) create mode 100644 migrations/0024_fix_expired_orchestration_lock_lease.sql create mode 100644 tests/expired_lease_regression_test.rs diff --git a/migrations/0024_fix_expired_orchestration_lock_lease.sql b/migrations/0024_fix_expired_orchestration_lock_lease.sql new file mode 100644 index 0000000..3a28024 --- /dev/null +++ b/migrations/0024_fix_expired_orchestration_lock_lease.sql @@ -0,0 +1,329 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the MIT License. + +-- Migration 0024: never hand out an already-expired orchestration lock lease +-- +-- NOTE: 0022 is reserved by the concurrent worker_queue dequeue PR +-- (perf/worker-queue-partial-index-dequeue). This migration is independent of it +-- and applies cleanly whether or not 0022 is present. +-- +-- Rebuilds fetch_orchestration_item on top of the 0019 definition. Three related +-- defects allowed a message to enter a state the runtime could not recover from: +-- +-- 1. Every dispatcher selected the same candidate (ORDER BY visible_at, id +-- LIMIT 1 is deterministic) and then blocked on the same advisory key via +-- pg_advisory_xact_lock. The FOR UPDATE ... SKIP LOCKED re-verification +-- could not prevent the convoy because the blocking had already happened. +-- +-- 2. The lease was computed from p_now_ms -- a timestamp minted by the Rust +-- caller BEFORE the call, never refreshed. When the advisory wait exceeded +-- p_lock_timeout_ms, locked_until was written already in the past, so the +-- item was returned with a dead lease and every ack against it failed with +-- 'Invalid lock token'. The same stale p_now_ms was also the liveness +-- reference for candidate visibility and for stealing expired locks, so a +-- long-running call became progressively less able to make progress. +-- +-- 3. The retry loop had no iteration bound, so each failed attempt took +-- another unbounded blocking advisory lock. +-- +-- Fixes: +-- +-- 1. pg_try_advisory_xact_lock instead of pg_advisory_xact_lock. Contended +-- instances are skipped, not waited on, and are excluded from subsequent +-- candidate selection within the call so the loop cannot re-pick them. +-- This PRESERVES the deadlock protection introduced with instance-level +-- advisory locking: same-instance work is still mutually exclusive. Only +-- the *waiting* is removed. +-- +-- The advisory key is also now schema-scoped. PostgreSQL advisory locks +-- live in a single database-wide space, but the key was hashtext() of the +-- instance id alone, so two duroxide schemas in the same database +-- serialized against each other whenever they happened to use the same +-- instance id -- entirely unrelated work contending on one key. While the +-- acquisition blocked this was invisible (the loser waited and then +-- proceeded); with a non-blocking acquisition it would surface as a +-- spurious skip. Hashing '.' confines contention to +-- the schema that owns the instance. +-- +-- Rolling-upgrade note: during a mixed deployment, nodes on the previous +-- migration hash the unscoped key while upgraded nodes hash the scoped +-- one, so the advisory layer does not mutually exclude across that +-- boundary for the duration of the rollout. Correctness does not depend +-- on it -- instance_locks plus FOR UPDATE ... SKIP LOCKED still serialize +-- the claim -- and the provider already retries deadlock (40P01) errors, +-- so the window degrades to the pre-advisory-lock behavior rather than to +-- incorrect results. +-- +-- 2. An effective clock, v_now_ms, replaces p_now_ms everywhere inside the +-- function. It is the caller epoch advanced by the time actually spent +-- inside this call, measured as an INTERVAL from clock_timestamp(): +-- +-- v_now_ms := p_now_ms + elapsed_ms_since_function_entry +-- +-- Only the elapsed *duration* comes from the server clock, never an +-- absolute server timestamp, so the Rust-supplied epoch established by +-- 0006/0012 remains the single time reference and no cross-clock skew is +-- introduced between this function and visible_at/locked_until values +-- written elsewhere. The lease is therefore always p_lock_timeout_ms in +-- the future relative to the moment it is actually written. +-- +-- 3. The loop is bounded. Exhausting the bound returns no item, which the +-- dispatcher already handles as "nothing available"; the next poll starts +-- with a fresh clock. +-- +-- The signature, return shape, and all downstream behavior (message batching, +-- history/metadata fallback, KV snapshot) are unchanged from 0019. + +DO $$ +DECLARE + v_schema_name TEXT := current_schema(); +BEGIN + EXECUTE format($fmt$DROP FUNCTION IF EXISTS %1$I.fetch_orchestration_item(BIGINT, BIGINT, BIGINT, BIGINT);$fmt$, v_schema_name); + + EXECUTE format($fmt$ + CREATE OR REPLACE FUNCTION %1$I.fetch_orchestration_item( + p_now_ms BIGINT, + p_lock_timeout_ms BIGINT, + p_min_version_packed BIGINT DEFAULT NULL, + p_max_version_packed BIGINT DEFAULT NULL + ) + RETURNS TABLE( + out_instance_id TEXT, + out_orchestration_name TEXT, + out_orchestration_version TEXT, + out_execution_id BIGINT, + out_history JSONB, + out_messages JSONB, + out_lock_token TEXT, + out_attempt_count INTEGER, + out_kv_snapshot JSONB + ) AS $fetch_orch$ + DECLARE + v_instance_id TEXT; + v_lock_token TEXT; + v_locked_until BIGINT; + v_orchestration_name TEXT; + v_orchestration_version TEXT; + v_current_execution_id BIGINT; + v_history JSONB; + v_messages JSONB; + v_kv_snapshot JSONB; + v_lock_acquired INTEGER; + v_max_attempt_count INTEGER; + -- Entry instant, used only to measure elapsed time within this call. + v_call_start TIMESTAMPTZ := clock_timestamp(); + -- Caller epoch advanced by elapsed time inside this call. + v_now_ms BIGINT; + -- Instances proven unavailable during this call; never re-picked. + v_skipped TEXT[] := ARRAY[]::TEXT[]; + v_iterations INTEGER := 0; + c_max_iterations CONSTANT INTEGER := 50; + BEGIN + LOOP + v_iterations := v_iterations + 1; + IF v_iterations > c_max_iterations THEN + -- Heavily contended; report nothing available. The next poll + -- re-enters with a fresh clock and a fresh skip list. + RETURN; + END IF; + + v_now_ms := p_now_ms + + (EXTRACT(EPOCH FROM (clock_timestamp() - v_call_start)) * 1000)::BIGINT; + + v_instance_id := NULL; + + -- Phase 1: Find a candidate instance (no FOR UPDATE yet) + -- When version filter is provided, join to instances+executions to filter + IF p_min_version_packed IS NOT NULL THEN + SELECT q.instance_id INTO v_instance_id + FROM %1$I.orchestrator_queue q + LEFT JOIN %1$I.instances i ON i.instance_id = q.instance_id + LEFT JOIN %1$I.executions e ON e.instance_id = i.instance_id + AND e.execution_id = i.current_execution_id + WHERE q.visible_at <= TO_TIMESTAMP(v_now_ms / 1000.0) + AND NOT (q.instance_id = ANY(v_skipped)) + AND NOT EXISTS ( + SELECT 1 FROM %1$I.instance_locks il + WHERE il.instance_id = q.instance_id AND il.locked_until > v_now_ms + ) + AND ( + e.duroxide_version_major IS NULL + OR (e.duroxide_version_major * 1000000 + e.duroxide_version_minor * 1000 + e.duroxide_version_patch) + BETWEEN p_min_version_packed AND p_max_version_packed + ) + ORDER BY q.visible_at, q.id + LIMIT 1; + ELSE + SELECT q.instance_id INTO v_instance_id + FROM %1$I.orchestrator_queue q + WHERE q.visible_at <= TO_TIMESTAMP(v_now_ms / 1000.0) + AND NOT (q.instance_id = ANY(v_skipped)) + AND NOT EXISTS ( + SELECT 1 FROM %1$I.instance_locks il + WHERE il.instance_id = q.instance_id AND il.locked_until > v_now_ms + ) + ORDER BY q.visible_at, q.id + LIMIT 1; + END IF; + + IF NOT FOUND THEN + RETURN; + END IF; + + -- Phase 2: Acquire the instance-level advisory lock WITHOUT waiting. + -- Same-instance mutual exclusion is preserved (this is still the + -- deadlock guard); a contended instance is skipped instead of + -- convoying every dispatcher onto one key. + IF NOT pg_try_advisory_xact_lock(hashtext(%2$L || '.' || v_instance_id)) THEN + v_skipped := array_append(v_skipped, v_instance_id); + CONTINUE; + END IF; + + -- Re-read the clock after the acquisition attempt so the lease and + -- the liveness predicates below reflect time actually elapsed. + v_now_ms := p_now_ms + + (EXTRACT(EPOCH FROM (clock_timestamp() - v_call_start)) * 1000)::BIGINT; + + -- Phase 3: Re-verify the instance is still available with FOR UPDATE. + -- If contention invalidated the candidate, keep searching. + IF p_min_version_packed IS NOT NULL THEN + SELECT q.instance_id INTO v_instance_id + FROM %1$I.orchestrator_queue q + LEFT JOIN %1$I.instances i ON i.instance_id = q.instance_id + LEFT JOIN %1$I.executions e ON e.instance_id = i.instance_id + AND e.execution_id = i.current_execution_id + WHERE q.instance_id = v_instance_id + AND q.visible_at <= TO_TIMESTAMP(v_now_ms / 1000.0) + AND NOT EXISTS ( + SELECT 1 FROM %1$I.instance_locks il + WHERE il.instance_id = q.instance_id AND il.locked_until > v_now_ms + ) + AND ( + e.duroxide_version_major IS NULL + OR (e.duroxide_version_major * 1000000 + e.duroxide_version_minor * 1000 + e.duroxide_version_patch) + BETWEEN p_min_version_packed AND p_max_version_packed + ) + FOR UPDATE OF q SKIP LOCKED; + ELSE + SELECT q.instance_id INTO v_instance_id + FROM %1$I.orchestrator_queue q + WHERE q.instance_id = v_instance_id + AND q.visible_at <= TO_TIMESTAMP(v_now_ms / 1000.0) + AND NOT EXISTS ( + SELECT 1 FROM %1$I.instance_locks il + WHERE il.instance_id = q.instance_id AND il.locked_until > v_now_ms + ) + FOR UPDATE OF q SKIP LOCKED; + END IF; + + IF NOT FOUND THEN + v_skipped := array_append(v_skipped, v_instance_id); + CONTINUE; + END IF; + + -- Step 2: Generate lock token and acquire instance lock. + -- v_locked_until is derived from the refreshed clock, so the lease + -- is always p_lock_timeout_ms in the future when it is written. + v_lock_token := 'lock_' || gen_random_uuid()::TEXT; + v_locked_until := v_now_ms + p_lock_timeout_ms; + + INSERT INTO %1$I.instance_locks (instance_id, lock_token, locked_until, locked_at) + VALUES (v_instance_id, v_lock_token, v_locked_until, v_now_ms) + ON CONFLICT(instance_id) DO UPDATE + SET lock_token = EXCLUDED.lock_token, + locked_until = EXCLUDED.locked_until, + locked_at = EXCLUDED.locked_at + WHERE %1$I.instance_locks.locked_until <= v_now_ms; + + GET DIAGNOSTICS v_lock_acquired = ROW_COUNT; + + IF v_lock_acquired = 0 THEN + v_skipped := array_append(v_skipped, v_instance_id); + CONTINUE; + END IF; + + -- Step 3: Mark all visible messages with our lock and increment attempt_count + UPDATE %1$I.orchestrator_queue q + SET lock_token = v_lock_token, + locked_until = v_locked_until, + attempt_count = q.attempt_count + 1 + WHERE q.instance_id = v_instance_id + AND q.visible_at <= TO_TIMESTAMP(v_now_ms / 1000.0) + AND (q.lock_token IS NULL OR q.locked_until <= v_now_ms); + + -- Step 4: Fetch all locked messages and get max attempt_count + SELECT COALESCE(JSONB_AGG(q.work_item::JSONB ORDER BY q.id), '[]'::JSONB), + COALESCE(MAX(q.attempt_count), 1) + INTO v_messages, v_max_attempt_count + FROM %1$I.orchestrator_queue q + WHERE q.lock_token = v_lock_token; + + -- Step 5: Load instance metadata + SELECT i.orchestration_name, i.orchestration_version, i.current_execution_id + INTO v_orchestration_name, v_orchestration_version, v_current_execution_id + FROM %1$I.instances i + WHERE i.instance_id = v_instance_id; + + -- Step 6: Load history or implement fallback + IF FOUND THEN + SELECT COALESCE(JSONB_AGG(h.event_data::JSONB ORDER BY h.event_id), '[]'::JSONB) + INTO v_history + FROM %1$I.history h + WHERE h.instance_id = v_instance_id AND h.execution_id = v_current_execution_id; + + v_orchestration_version := COALESCE(v_orchestration_version, 'unknown'); + ELSE + SELECT COALESCE(JSONB_AGG(h.event_data::JSONB ORDER BY h.execution_id, h.event_id), '[]'::JSONB) + INTO v_history + FROM %1$I.history h + WHERE h.instance_id = v_instance_id; + + IF JSONB_ARRAY_LENGTH(v_history) > 0 AND v_history->0 ? 'OrchestrationStarted' THEN + v_orchestration_name := v_history->0->'OrchestrationStarted'->>'name'; + v_orchestration_version := v_history->0->'OrchestrationStarted'->>'version'; + v_current_execution_id := 1; + ELSIF JSONB_ARRAY_LENGTH(v_messages) > 0 AND v_messages->0 ? 'StartOrchestration' THEN + v_orchestration_name := v_messages->0->'StartOrchestration'->>'orchestration'; + v_orchestration_version := COALESCE(v_messages->0->'StartOrchestration'->>'version', 'unknown'); + v_current_execution_id := COALESCE((v_messages->0->'StartOrchestration'->>'execution_id')::BIGINT, 1); + ELSIF JSONB_ARRAY_LENGTH(v_messages) > 0 AND v_messages->0 ? 'ContinueAsNew' THEN + v_orchestration_name := v_messages->0->'ContinueAsNew'->>'orchestration'; + v_orchestration_version := COALESCE(v_messages->0->'ContinueAsNew'->>'version', 'unknown'); + v_current_execution_id := 1; + ELSE + v_orchestration_name := 'Unknown'; + v_orchestration_version := 'unknown'; + v_current_execution_id := 1; + END IF; + END IF; + + -- Load KV snapshot for this instance + SELECT COALESCE( + jsonb_object_agg( + ks.key, + jsonb_build_object('value', ks.value, 'last_updated_at_ms', ks.last_updated_at_ms) + ), + '{}'::jsonb + ) + INTO v_kv_snapshot + FROM %1$I.kv_store ks + WHERE ks.instance_id = v_instance_id; + + RETURN QUERY SELECT + v_instance_id, + v_orchestration_name, + v_orchestration_version, + v_current_execution_id, + v_history, + v_messages, + v_lock_token, + v_max_attempt_count, + v_kv_snapshot; + RETURN; + END LOOP; + END; + $fetch_orch$ LANGUAGE plpgsql; +$fmt$, v_schema_name, v_schema_name); + + RAISE NOTICE 'Migration 0024: fetch_orchestration_item never returns an expired lease; advisory lock is non-blocking and the retry loop is bounded'; +END $$; diff --git a/tests/expired_lease_regression_test.rs b/tests/expired_lease_regression_test.rs new file mode 100644 index 0000000..950379c --- /dev/null +++ b/tests/expired_lease_regression_test.rs @@ -0,0 +1,189 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Regression test for the expired orchestration lock lease bug. +//! +//! Reference: +//! +//! # Summary +//! +//! `fetch_orchestration_item` selected its candidate deterministically +//! (`ORDER BY visible_at, id LIMIT 1`) and then took a *blocking* +//! `pg_advisory_xact_lock` on that instance. Two consequences followed: +//! +//! 1. Every dispatcher converged on the same key and queued behind it, even +//! when other instances were free to run (head-of-line blocking). +//! 2. The lock lease was computed from `p_now_ms` — a timestamp minted by the +//! Rust caller *before* the call — which was never refreshed. When the +//! advisory wait exceeded the lock timeout, `locked_until` was written +//! already in the past, so the item was handed back with a dead lease and +//! every subsequent ack failed with `Invalid lock token`. The message was +//! then never poisoned, never backed off, and retried forever. +//! +//! Migration 0024 makes the advisory acquisition non-blocking +//! (`pg_try_advisory_xact_lock`, skipping contended instances) and derives the +//! lease from the caller epoch advanced by time actually elapsed inside the +//! call. +//! +//! # What this test does +//! +//! A competing session holds the instance advisory lock for `HOLD` seconds +//! while a fetch runs with a lock timeout *shorter* than that hold. Against the +//! pre-0024 procedure the fetch blocks for the full hold and then returns the +//! contended instance with a negative remaining lease. Measured on +//! PostgreSQL 15.19: blocked 2.64s, lease -1658ms. + +use duroxide::providers::{Provider, WorkItem}; +use duroxide_pg::PostgresProvider; +use sqlx::postgres::PgPoolOptions; +use std::time::{Duration, Instant}; + +fn get_database_url() -> String { + dotenvy::dotenv().ok(); + std::env::var("DATABASE_URL").expect("DATABASE_URL must be set") +} + +fn unique_schema_name() -> String { + let guid = uuid::Uuid::new_v4().to_string(); + let suffix = &guid[guid.len() - 8..]; + format!("expired_lease_test_{suffix}") +} + +/// The contended instance. Enqueued first, so it is the head-of-queue candidate +/// that every dispatcher selects. +const HOT: &str = "expired-lease-hot"; +/// A second, entirely uncontended instance that is always available to run. +const OTHER: &str = "expired-lease-other"; + +/// How long the competing session holds the advisory lock. +const HOLD: Duration = Duration::from_secs(3); +/// Lock timeout for the fetch under test. Deliberately shorter than `HOLD`: +/// this is what makes a pre-refresh lease expire before it is returned. +const LOCK_TIMEOUT: Duration = Duration::from_secs(1); +/// A fetch that does not convoy returns far below this; one that waits out the +/// competing session takes at least `HOLD`. +const MAX_ACCEPTABLE_FETCH: Duration = Duration::from_secs(2); + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn fetch_orchestration_item_never_returns_an_expired_lease() { + let schema = unique_schema_name(); + let database_url = get_database_url(); + + let provider = PostgresProvider::new_with_schema(&database_url, Some(&schema)) + .await + .expect("Failed to create provider"); + + for (index, instance) in [HOT, OTHER].iter().enumerate() { + provider + .enqueue_for_orchestrator( + WorkItem::StartOrchestration { + instance: (*instance).to_string(), + orchestration: "ExpiredLeaseRegression".to_string(), + input: "{}".to_string(), + version: Some("1.0.0".to_string()), + parent_instance: None, + parent_id: None, + parent_execution_id: None, + execution_id: 1, + }, + None, + ) + .await + .unwrap_or_else(|e| panic!("failed to enqueue instance {index}: {e:?}")); + } + + let pool = PgPoolOptions::new() + .max_connections(2) + .connect(&database_url) + .await + .expect("Failed to connect for lock holding"); + + // Competing session: hold the instance-level advisory lock for HOLD. + // This is exactly the contention the dispatcher fleet creates on a hot + // instance; the lock is taken on the same key the provider uses. + // + // Two keys are held, so this test reproduces contention against the + // procedure both before and after migration 0024: + // + // '' -- the key used before 0024 + // '.' -- the key used from 0024 onward + // + // PostgreSQL advisory locks share one database-wide space, so the unscoped + // key made two duroxide schemas in the same database contend on identical + // instance ids. Scoping the key confines contention to the owning schema. + // Holding both means this test always describes the defect that is + // actually present, instead of silently failing to contend at all. + let legacy_key = HOT.to_string(); + let scoped_key = format!("{schema}.{HOT}"); + let holder_pool = pool.clone(); + let holder = tokio::spawn(async move { + let mut tx = holder_pool.begin().await.expect("begin failed"); + for key in [&legacy_key, &scoped_key] { + sqlx::query("SELECT pg_advisory_xact_lock(hashtext($1))") + .bind(key) + .execute(&mut *tx) + .await + .expect("advisory lock failed"); + } + tokio::time::sleep(HOLD).await; + tx.rollback().await.expect("rollback failed"); + }); + + // Let the competing session take the lock before the fetch starts. + tokio::time::sleep(Duration::from_millis(400)).await; + + let started = Instant::now(); + let fetched = provider + .fetch_orchestration_item(LOCK_TIMEOUT, Duration::ZERO, None) + .await + .expect("fetch_orchestration_item should succeed"); + let blocked_for = started.elapsed(); + + // Defect 1: a contended instance must not stall the dispatcher. Before the + // fix this waited out the whole hold. + assert!( + blocked_for < MAX_ACCEPTABLE_FETCH, + "fetch_orchestration_item blocked for {blocked_for:?} behind a contended \ + instance (limit {MAX_ACCEPTABLE_FETCH:?}); the advisory lock is convoying" + ); + + // Defect 2: whatever comes back must carry a lease that is still live. This + // is the invariant the runtime depends on — an expired lease makes every + // ack fail, including the poison path that is supposed to terminate the + // message. + let (item, lock_token, _attempt) = + fetched.expect("expected the uncontended instance to be returned"); + + let remaining_ms: i64 = sqlx::query_scalar(&format!( + "SELECT locked_until - (EXTRACT(EPOCH FROM now()) * 1000)::BIGINT + FROM {schema}.instance_locks + WHERE lock_token = $1" + )) + .bind(&lock_token) + .fetch_one(&pool) + .await + .expect("lease row should exist for the returned lock token"); + + assert!( + remaining_ms > 0, + "fetch_orchestration_item returned instance {} with a lease that had \ + already expired {}ms earlier; every ack against it will fail with \ + 'Invalid lock token'", + item.instance, + -remaining_ms + ); + + // The uncontended instance was available the whole time, so a dispatcher + // that routes around contention makes progress instead of stalling. + assert_eq!( + item.instance, OTHER, + "expected the uncontended instance to be picked while {HOT} was locked" + ); + + holder.await.expect("lock holder task panicked"); + + sqlx::query(&format!("DROP SCHEMA IF EXISTS {schema} CASCADE")) + .execute(&pool) + .await + .expect("Failed to drop schema"); +}