diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a363105..333c42b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -77,4 +77,6 @@ jobs: env: # Use the PostgreSQL service container for CI tests DATABASE_URL: postgres://postgres:postgres@localhost:5432/duroxide_test - run: cargo test --verbose + run: | + cargo test --verbose + DUROXIDE_PG_TEST_LONG_POLL=true cargo test --test postgres_provider_test --verbose diff --git a/CHANGELOG.md b/CHANGELOG.md index 917c21b..00e4320 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,52 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Opt-in long-polling through `ProviderConfig::long_poll`, for both URL and + Entra connections. It uses LISTEN/NOTIFY, bounded timer tracking for both + queues, periodic recovery of missed notifications, and reconnect handling. + Existing constructors remain short-polling by default. +- Follow-up empty fetches and timer catch-up scans are coalesced at 100 ms + intervals; the first notification wake remains immediate. The configuration + type is non-exhaustive to permit compatible additions. +- Acknowledgments notify other orchestration waiters after releasing a claim, + and per-queue catch-up scans wake work that became due between timer wakes. +- Claims, renewals, and session leases announce future expiry deadlines, allowing + waiting dispatchers to recover crashed owners without waiting for the full + poll timeout. Delayed orchestration abandonment also wakes other visible work + immediately when it releases the instance lock. +- Long-poll startup retains a warm fetch connection alongside LISTEN. Read-only + timer scans use 200 ms server-side lock and statement timeouts and return + their rolled-back connections to the pool. A one-second client watchdog + retains any pending scan instead of repeatedly detaching server backends; + degraded listeners keep polling independently of reconnect attempts. +- Migration `0030` adds disabled notification triggers. An `ApplyAll` provider + with long-polling enabled explicitly activates them for its schema; + `VerifyOnly` requires prior activation and performs no DDL. Once activated, + queue and session-lease writes cannot use PostgreSQL `PREPARE TRANSACTION`. + Schemas that have not opted in retain their previous transaction behavior. +- Long-polling reserves one connection within the configured pool and requires + at least two connections. Client-only providers can leave it disabled. + Notification channels support long and mixed-case schema names. This does + not convert the incompatible `duroxide-pg-opt` migration history. + +### Changed + +- Require SQLx 0.8.3 or later within the 0.8 series for explicit listener + reconnect handling. The existing native-TLS feature selection is unchanged. +- Migration and notification-activation transactions use bounded lock waits + and jittered retries for deadlocks and lock contention, with a ten-second + retry budget. Permanent errors remain errors and failed attempts do not + change the migration ledger. Retry warnings include the SQLSTATE and full + error context. +- Add a reproducible sparse/loaded workflow benchmark alongside the idle-query + benchmark, recording concurrency, duration, pool size, and commit settings. +- Add a synchronous-standby contention profile with 16 orchestration and + 16 worker slots, 64 clients, and concurrent lease-renewal measurements. + Benchmarks record sampled notification-commit and replication wait time; + documentation makes the per-claim and per-renewal notification costs explicit. + ## [0.1.38] - 2026-10-26 ### Fixed diff --git a/Cargo.toml b/Cargo.toml index b10b01b..6e99845 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,7 +32,7 @@ exclude = [ duroxide = { version = "0.1.32", features = ["provider-test"] } async-trait = "0.1" tokio = { version = "1", features = ["full"] } -sqlx = { version = "0.8", features = ["runtime-tokio-native-tls", "postgres", "chrono"], default-features = false } +sqlx = { version = "0.8.3", features = ["runtime-tokio-native-tls", "postgres", "chrono"], default-features = false } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" dotenvy = "0.15" diff --git a/README.md b/README.md index 056b834..071b336 100644 --- a/README.md +++ b/README.md @@ -174,9 +174,182 @@ Two test layers cover the Entra integration: |---------------------|-------------|---------| | `DUROXIDE_PG_POOL_MAX` | Maximum connection pool size | `10` | +### Optional long-polling + +Existing constructors use short polling. To reduce idle queries and wake workers +when work arrives, explicitly enable long-polling: + +```rust,no_run +use duroxide_pg::{PostgresProvider, ProviderConfig}; + +# async fn example(database_url: &str) -> anyhow::Result<()> { +let mut config = ProviderConfig::url(database_url); +config.schema_name = Some("my_app".to_string()); +config.long_poll.enabled = true; +let provider = PostgresProvider::new_with_config(config).await?; +# Ok(()) +# } +``` + +The same setting works with `ProviderConfig::entra(...)`. Long-polling reserves +one connection **within** the provider's pool for PostgreSQL `LISTEN`, so the +pool must allow at least two connections. It keeps another connection warm for +fetches rather than leaving the first dispatcher to open one. Leave it disabled +for client-only providers, especially one-connection backend clients. It requires a +session-preserving database connection; transaction-pooling proxies are not a +substitute for a dedicated LISTEN session. + +`LongPollConfig` also provides `notifier_poll_interval` (default 60 seconds) +and `timer_grace_period` (default 100 milliseconds). The provider schedules +future work and orchestration, activity, and session lease expiry, rescans after +a timer fires, and periodically wakes fetches to recover missed notifications. +Claims and renewals announce their lease deadlines without immediately waking +fetches. An expired claim can therefore be recovered even if its owner crashes. +Notifications are hints; queue predicates and locks remain authoritative. + +Read-only timer scans run in a transaction with 200 ms server-side lock and +statement timeouts. PostgreSQL cancels blocked scans, and the rolled-back +connection returns to the pool rather than leaving a blocked server backend +behind. Scan errors retry after about one second. + +A one-second client watchdog keeps fetches checking for work if a scan stalls +beyond those server limits. It does not cancel the scan or start another one; +the same in-flight scan is retained across listener reconnects. Reconnecting +also has an independent one-second polling fallback. Successful reconnection +resubscribes and rechecks work. + +An enabled fetch waits for work until its `poll_timeout`; a zero timeout still +performs a single fetch. The timeout bounds notification waiting, not an +in-flight database operation. Dropping the provider or closing its pool stops +the notification/timer task; connection cleanup has the network-failure +limitation described below. +The first notification wakes a fetch immediately. Further notifications that +find no eligible work, and timer catch-up scans, are coalesced to a 100 ms +cadence to bound query load. This can add up to 100 ms after an unsuccessful +wake; the overall polling deadline is not extended. + +Migration `0030` installs **disabled** queue notification triggers. Constructing +a provider with both `MigrationPolicy::ApplyAll` and `long_poll.enabled = true` +activates them for that schema. `VerifyOnly` workers require that activation +first and perform no DDL themselves. See [Migration Policy](docs/migration-policy.md). + +Once activated, inserts, visibility changes, lease changes, released locks, and +consumption of claimed orchestration messages emit notifications. Acknowledgment +notifications are delivered after the transaction releases the instance lock, so other +workers recheck events that arrived during the turn. Delayed abandonment also +announces the immediate instance-lock release, so other visible messages need +not wait for the abandoned message's delay. This includes writes from +short-polling producers. Turning off a provider's listener does not deactivate +the schema's triggers. Transactions writing those queues or session leases cannot use +`PREPARE TRANSACTION`, because PostgreSQL does +not support preparing transactions that emit notifications. Schemas that have +not opted in retain their existing transaction behavior. + +PostgreSQL serializes transactions that emit notifications through a +server-wide lock held until commit completes. Slow WAL flushes or synchronous +replication waits can therefore delay other notification writers, even in +other schemas or databases on that server. This cost includes **every successful +activity and orchestration claim**, renewals that change their lease deadlines, +and session-lease writes. Empty fetches do not notify. A normal completed +activity has three notifying commits (enqueue, claim, completion acknowledgment); +an orchestration turn has two (claim and acknowledgment). Each lease-renewal +transaction also joins this serialization when it emits a hint. Multiple hints +in one transaction still acquire the commit lock only once. + +These are schema-wide costs, including short-polling callers of an activated +schema. Benchmark concurrent writers with representative commit latency before +activation; idle-query savings alone do not establish loaded throughput. + +Detected connection failures reconnect automatically, but silent network +partitions are not actively probed. Until the transport reports a failure, +work relies on the polling and refresh fallbacks instead of prompt +notifications. SQLx listener cleanup can also wait indefinitely on a wedged +socket, so provider drop does not guarantee bounded connection cleanup and +`pool.close()` has no hard shutdown deadline in that situation. + +Recreate providers if their queue tables are dropped and recreated: notification +channels identify the original tables. Runtime session-capacity decisions are +also made once per fetch; a slot that started polling without session capacity +may need another fetch after its poll timeout before it can claim session work. + +This feature does **not** convert a `duroxide-pg-opt` database into a +`duroxide-pg` database. Their migration histories remain incompatible. + +#### Reproducing long-poll performance + +The ignored benchmarks in [tests/long_poll_tests.rs](tests/long_poll_tests.rs) +compare short and long polling without changing the runtime's lease durations. +Use a disposable, optimized PostgreSQL server with assertions disabled: + +```sh +DUROXIDE_PG_POOL_MAX=10 cargo nextest run --release --test long_poll_tests \ + --run-ignored only --nocapture \ + -E 'test(workflow_latency_and_throughput_with_long_polling)' +``` + +Use otherwise idle, controlled hosts for comparisons. Concurrent workloads can +invalidate both throughput comparisons and wait-time measurements. + +The workflow driver uses two orchestration slots, two worker slots, a 100 ms +minimum poll interval, and a 30 s long-poll timeout. It measures 20 sparse +one-activity workflows with a 200 ms idle gap, then 20 concurrent clients running +one- and five-sequential-activity workflows for 10 seconds, with three loaded +repetitions per mode. Activities return immediately. A completed warm-up workflow +is excluded. Latency includes enqueue and the client's completion-status polling; +throughput includes draining in-flight workflows after submissions stop. + +Each `BENCH` JSON line records the mode, pool limit, concurrency, duration, +PostgreSQL version, commit settings (`fsync`, `full_page_writes`, +`synchronous_commit`, and `synchronous_standby_names`), and replication state. +Loaded runs sample notification-commit lock waits and `SyncRep` waits every +10 ms. The reported backend-seconds are estimates integrated across the +provider's connections, not individual-query latency measurements, and can +exceed wall-clock time. The notification measurement uses PostgreSQL's shared +`pg_database` object-0 heavyweight lock, not the `NotifyQueue` LWLock. +Keep the settings and source revision alongside results; the driver does not +relax durability. +The separate `idle_query_load_is_reduced_with_long_polling` benchmark requires +`pg_stat_statements` and measures 65 seconds per mode. Run it alone in an +otherwise idle disposable database: it counts database-wide top-level statements +except its statistics queries, including timer-scan transaction control. +The `pg-stress` binary continues to use short polling. + +For commit-contention measurements, use a disposable primary with a physical +synchronous standby. The higher-concurrency profile uses 16 orchestration +slots, 16 worker slots, 64 clients, and a warmed pool; it also measures +concurrent orchestration, activity, and session renewals with 60-second leases. +Both workflow and renewal runs use three 10-second repetitions per mode. +The renewal workload deliberately saturates renewal calls rather than using +the runtime's normal renewal intervals. + +For a reproducible 10 ms apply-latency experiment, configure the standby's +connection with `application_name=pr41_bench_standby` and set +`recovery_min_apply_delay = '10ms'` there. On the primary, keep `fsync` and +`full_page_writes` enabled and use: + +```conf +synchronous_standby_names = 'FIRST 1 (pr41_bench_standby)' +synchronous_commit = 'remote_apply' +``` + +Once that standby is streaming, run against the primary: + +```sh +DUROXIDE_PG_POOL_MAX=100 cargo nextest run --release --test long_poll_tests \ + --run-ignored only --nocapture \ + -E 'test(notification_contention_with_synchronous_standby)' +``` + +The test rejects missing synchronous replication and insufficient pool capacity. +It requires observed notification and replication waits in the enabled runs; +disabled-schema controls must have no notification-lock waits. This experiment +models a delayed synchronous apply acknowledgment, not a particular cloud +service's replication topology. + ## Features - Automatic schema migration on startup +- Optional notification-driven long-polling for orchestration and activity work - Connection pooling via sqlx - Custom schema support for multi-tenant isolation - Full implementation of the Duroxide `Provider` and `ProviderAdmin` traits diff --git a/docs/migration-policy.md b/docs/migration-policy.md index e5716d0..5fe834a 100644 --- a/docs/migration-policy.md +++ b/docs/migration-policy.md @@ -33,6 +33,71 @@ separate privileged process has already run migrations. `ApplyAll` takes a PostgreSQL session advisory lock scoped to the target schema before checking or applying migrations. That serializes concurrent provider startup so multiple nodes do not race to apply the same migration. +Each migration transaction uses a 250 ms `lock_timeout`. Deadlocks (`40P01`) and +lock timeouts (`55P03`) roll back and retry the entire transaction with jitter, +within a ten-second contention budget per migration. Other errors return +immediately. A failed attempt does not advance the migration ledger. The +migration connection is closed on exit or cancellation so a failed unlock +cannot leave a session advisory lock in the pool. + +The contention budget is currently fixed, not a provider configuration option. +If it runs out, inspect the logged schema, SQLSTATE, and error chain, stop or +drain competing workers and admin writers, and allow transactions holding the +blocking locks to finish. Then retry `ApplyAll` from the process with DDL +privileges. Earlier successful migrations remain recorded; the failed +transaction is rolled back and will be retried on startup. Do not delete or +edit migration ledger rows to force a retry. Backfill migrations must use their +drained-upgrade procedure: a transaction retry repeats that backfill's work. + +## Long-Poll Notification Activation + +Migration `0030` installs disabled notification triggers. With +`long_poll.enabled = true`, `ApplyAll` also enables the queue and session triggers in a +transaction while holding the migration lock. `VerifyOnly` requires that the +triggers are already active; otherwise construction returns an error explaining +that an `ApplyAll` provider must activate them first. It never performs DDL. +Activation follows acknowledgments' worker-before-orchestration lock order. +Every lock acquisition, including the first queue, uses the same 250 ms +`lock_timeout` and ten-second retry budget as migration transactions. Each +failed attempt rolls back before the jittered delay. This bounds interference +with writers and allows activation to outlast short bursts of reverse-order +admin-delete contention. Exhausting the budget returns the database error +without partially activating triggers. + +Activation is schema-wide and persists after the initializing provider exits. +It allows short-polling clients to notify long-polling workers. Constructing a +provider with long-polling disabled does not deactivate these shared triggers. +Schemas that have not opted in do not emit notifications. + +PostgreSQL cannot prepare a transaction that emits a notification, so +`PREPARE TRANSACTION` is not supported for queue or session-lease writes after activation. +Reverting to two-phase-commit use requires stopping all long-polling workers +and explicitly disabling all six notification triggers: + +- `worker_queue`: `trg_notify_worker_work`, `trg_notify_worker_work_released` +- `orchestrator_queue`: `trg_notify_orch_work`, `trg_notify_orch_work_released`, + `trg_notify_orch_work_consumed` +- `sessions`: `trg_notify_session_lease` + +These are verified and enabled from a single trigger inventory; +changing a process's configuration alone is not sufficient. + +Notification-emitting commits also participate in PostgreSQL's server-wide +notification-writer serialization. WAL or synchronous-replication delays can +block other notification senders, including in different databases. Activation +adds this cost to every successful queue claim, lease-changing orchestration +and activity renewal, and session-lease write, not just enqueue and acknowledgment. +A normal activity therefore has three notifying commits and an orchestration +turn two; renewal transactions also notify. Empty fetches do not. +Evaluate the loaded-write and renewal costs before opting a schema in; +activation affects short-polling callers too. Disabling a process's listener +does not remove the commit-lock cost of the schema's triggers. See the +[synchronous-standby benchmark](../README.md#reproducing-long-poll-performance) +for a higher-concurrency measurement that records notification-lock wait time. + +This is an additive migration within the `duroxide-pg` schema lineage, not a +conversion of `duroxide-pg-opt` state. Older binaries still reject schemas with +unknown migration versions, even when notifications remain disabled. ## Error Messages diff --git a/migrations/0030_add_long_poll_notifications.sql b/migrations/0030_add_long_poll_notifications.sql new file mode 100644 index 0000000..7670ae9 --- /dev/null +++ b/migrations/0030_add_long_poll_notifications.sql @@ -0,0 +1,157 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the MIT License. + +-- Migration 0030: notification hints for optional long-polling. +-- Table OIDs keep channels short and independent of identifier case/truncation. + +CREATE OR REPLACE FUNCTION duroxide_notify_queue_work() +RETURNS TRIGGER +LANGUAGE plpgsql +SET search_path = pg_catalog, pg_temp +AS $notify$ +BEGIN + IF TG_TABLE_NAME = 'sessions' THEN + PERFORM pg_catalog.pg_notify( + 'duroxide_pg_' || pg_catalog.format('%I.worker_queue', TG_TABLE_SCHEMA)::pg_catalog.regclass::pg_catalog.oid::TEXT, + NEW.locked_until::TEXT + ); + RETURN NEW; + ELSIF TG_OP = 'DELETE' THEN + -- Acknowledgment consumes claimed rows and releases the instance lock + -- in one transaction; NOTIFY is delivered only after that commit. + PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); + RETURN OLD; + ELSIF TG_OP = 'UPDATE' AND TG_TABLE_NAME = 'orchestrator_queue' THEN + IF OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL THEN + -- Other visible messages can run even when this message is delayed. + PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); + END IF; + END IF; + + PERFORM pg_catalog.pg_notify( + 'duroxide_pg_' || TG_RELID::TEXT, + CASE WHEN NEW.lock_token IS NOT NULL THEN NEW.locked_until::TEXT + ELSE pg_catalog.ceil(EXTRACT(EPOCH FROM NEW.visible_at) * 1000)::BIGINT::TEXT + END + ); + RETURN NEW; +END; +$notify$; + +-- Match normal acknowledgments' worker-before-orchestration table lock order. +DROP TRIGGER IF EXISTS trg_notify_worker_work ON worker_queue; +CREATE TRIGGER trg_notify_worker_work + AFTER INSERT ON worker_queue + FOR EACH ROW EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE worker_queue DISABLE TRIGGER trg_notify_worker_work; + +DROP TRIGGER IF EXISTS trg_notify_worker_work_released ON worker_queue; +CREATE TRIGGER trg_notify_worker_work_released + AFTER UPDATE OF visible_at, lock_token, locked_until ON worker_queue + FOR EACH ROW + WHEN ( + NEW.visible_at IS DISTINCT FROM OLD.visible_at + OR (OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL) + OR (NEW.lock_token IS NOT NULL AND ( + NEW.lock_token IS DISTINCT FROM OLD.lock_token + OR NEW.locked_until IS DISTINCT FROM OLD.locked_until + )) + ) + EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE worker_queue DISABLE TRIGGER trg_notify_worker_work_released; + +DROP TRIGGER IF EXISTS trg_notify_orch_work ON orchestrator_queue; +CREATE TRIGGER trg_notify_orch_work + AFTER INSERT ON orchestrator_queue + FOR EACH ROW EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work; + +DROP TRIGGER IF EXISTS trg_notify_orch_work_released ON orchestrator_queue; +CREATE TRIGGER trg_notify_orch_work_released + AFTER UPDATE OF visible_at, lock_token, locked_until ON orchestrator_queue + FOR EACH ROW + WHEN ( + NEW.visible_at IS DISTINCT FROM OLD.visible_at + OR (OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL) + OR (NEW.lock_token IS NOT NULL AND ( + NEW.lock_token IS DISTINCT FROM OLD.lock_token + OR NEW.locked_until IS DISTINCT FROM OLD.locked_until + )) + ) + EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work_released; + +DROP TRIGGER IF EXISTS trg_notify_orch_work_consumed ON orchestrator_queue; +CREATE TRIGGER trg_notify_orch_work_consumed + AFTER DELETE ON orchestrator_queue + FOR EACH ROW + WHEN (OLD.lock_token IS NOT NULL) + EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work_consumed; + +DROP TRIGGER IF EXISTS trg_notify_session_lease ON sessions; +CREATE TRIGGER trg_notify_session_lease + AFTER INSERT OR UPDATE OF locked_until ON sessions + FOR EACH ROW EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE sessions DISABLE TRIGGER trg_notify_session_lease; + +DO $$ +DECLARE + v_schema_name TEXT := current_schema(); +BEGIN + EXECUTE format($fmt$ + CREATE OR REPLACE FUNCTION %1$I.cleanup_schema() + RETURNS VOID AS $cleanup$ + BEGIN + DROP TABLE IF EXISTS %1$I.sessions CASCADE; + DROP TABLE IF EXISTS %1$I.kv_delta CASCADE; + DROP TABLE IF EXISTS %1$I.kv_store CASCADE; + DROP TABLE IF EXISTS %1$I.instances CASCADE; + DROP TABLE IF EXISTS %1$I.executions CASCADE; + DROP TABLE IF EXISTS %1$I.history CASCADE; + DROP TABLE IF EXISTS %1$I.orchestrator_queue CASCADE; + DROP TABLE IF EXISTS %1$I.worker_queue CASCADE; + DROP TABLE IF EXISTS %1$I.instance_locks CASCADE; + DROP TABLE IF EXISTS %1$I._duroxide_migrations CASCADE; + + DROP FUNCTION IF EXISTS %1$I.duroxide_notify_queue_work(); + DROP FUNCTION IF EXISTS %1$I.cleanup_schema(); + DROP FUNCTION IF EXISTS %1$I.list_instances(); + DROP FUNCTION IF EXISTS %1$I.list_executions(TEXT); + DROP FUNCTION IF EXISTS %1$I.latest_execution_id(TEXT); + DROP FUNCTION IF EXISTS %1$I.list_instances_by_status(TEXT); + DROP FUNCTION IF EXISTS %1$I.get_instance_info(TEXT); + DROP FUNCTION IF EXISTS %1$I.get_execution_info(TEXT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.get_system_metrics(); + DROP FUNCTION IF EXISTS %1$I.get_queue_depths(BIGINT); + DROP FUNCTION IF EXISTS %1$I.enqueue_worker_work(TEXT, BIGINT, TEXT, BIGINT, BIGINT, TEXT); + DROP FUNCTION IF EXISTS %1$I.enqueue_worker_work(TEXT, BIGINT, TEXT, BIGINT, BIGINT, TEXT, TEXT); + DROP FUNCTION IF EXISTS %1$I.ack_worker(TEXT, TEXT, TEXT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.renew_work_item_lock(TEXT, BIGINT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.fetch_work_item(BIGINT, BIGINT, TEXT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.fetch_work_item(BIGINT, BIGINT, TEXT, BIGINT, TEXT[], TEXT); + DROP FUNCTION IF EXISTS %1$I.abandon_work_item(TEXT, BIGINT, BIGINT, BOOLEAN); + DROP FUNCTION IF EXISTS %1$I.enqueue_orchestrator_work(TEXT, TEXT, TIMESTAMPTZ, TEXT, TEXT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.fetch_orchestration_item(BIGINT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.fetch_orchestration_item(BIGINT, BIGINT, BIGINT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.ack_orchestration_item(TEXT, BIGINT, BIGINT, JSONB, JSONB, JSONB, JSONB, JSONB); + DROP FUNCTION IF EXISTS %1$I.abandon_orchestration_item(TEXT, BIGINT, BIGINT, BOOLEAN); + DROP FUNCTION IF EXISTS %1$I.renew_orchestration_item_lock(TEXT, BIGINT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.fetch_history(TEXT); + DROP FUNCTION IF EXISTS %1$I.fetch_history_with_execution(TEXT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.append_history(TEXT, BIGINT, JSONB); + DROP FUNCTION IF EXISTS %1$I.list_children(TEXT); + DROP FUNCTION IF EXISTS %1$I.get_parent_id(TEXT); + DROP FUNCTION IF EXISTS %1$I.delete_instances_atomic(TEXT[], BOOLEAN); + DROP FUNCTION IF EXISTS %1$I.prune_executions(TEXT, INTEGER, BIGINT); + DROP FUNCTION IF EXISTS %1$I.prune_executions_batch(TEXT, BIGINT, BIGINT, INTEGER); + DROP FUNCTION IF EXISTS %1$I.renew_session_lock(TEXT[], BIGINT, BIGINT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.cleanup_orphaned_sessions(BIGINT); + DROP FUNCTION IF EXISTS %1$I.get_custom_status(TEXT, BIGINT); + DROP FUNCTION IF EXISTS %1$I.get_kv_value(TEXT, TEXT); + DROP FUNCTION IF EXISTS %1$I.get_kv_all_values(TEXT); + DROP FUNCTION IF EXISTS %1$I.get_instance_stats(TEXT); + END; + $cleanup$ LANGUAGE plpgsql; + $fmt$, v_schema_name); +END $$; diff --git a/migrations/0030_diff.md b/migrations/0030_diff.md new file mode 100644 index 0000000..f54c8e9 --- /dev/null +++ b/migrations/0030_diff.md @@ -0,0 +1,181 @@ +# Diff for migration 0030 + +**Migration file:** `0030_add_long_poll_notifications.sql` + +Each changed function is shown **in full** with `+`/`-` markers on changed lines. + +## Table Changes + +None. + +## New Indexes + +None. + +## Trigger Changes + +These triggers are installed disabled. `ApplyAll` construction with +`long_poll.enabled = true` enables them atomically under the migration lock; +short-polling-only schemas retain their previous transaction behavior. +Installation and activation take worker-queue locks before orchestration-queue +locks, matching normal acknowledgments. + +The six-trigger inventory covers both queues and session leases. Claims and +renewals announce future deadlines, not immediate availability. Orchestration +updates and deletes have separate `WHEN` clauses. Releasing a claim announces +both its message's visibility and the immediate instance-lock release; +acknowledgment notifications are delivered after that release commits. + +```sql +CREATE TRIGGER trg_notify_worker_work + AFTER INSERT ON worker_queue + FOR EACH ROW EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE worker_queue DISABLE TRIGGER trg_notify_worker_work; + +CREATE TRIGGER trg_notify_worker_work_released + AFTER UPDATE OF visible_at, lock_token, locked_until ON worker_queue + FOR EACH ROW + WHEN ( + NEW.visible_at IS DISTINCT FROM OLD.visible_at + OR (OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL) + OR (NEW.lock_token IS NOT NULL AND ( + NEW.lock_token IS DISTINCT FROM OLD.lock_token + OR NEW.locked_until IS DISTINCT FROM OLD.locked_until + )) + ) + EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE worker_queue DISABLE TRIGGER trg_notify_worker_work_released; + +CREATE TRIGGER trg_notify_orch_work + AFTER INSERT ON orchestrator_queue + FOR EACH ROW EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work; + +CREATE TRIGGER trg_notify_orch_work_released + AFTER UPDATE OF visible_at, lock_token, locked_until ON orchestrator_queue + FOR EACH ROW + WHEN ( + NEW.visible_at IS DISTINCT FROM OLD.visible_at + OR (OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL) + OR (NEW.lock_token IS NOT NULL AND ( + NEW.lock_token IS DISTINCT FROM OLD.lock_token + OR NEW.locked_until IS DISTINCT FROM OLD.locked_until + )) + ) + EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work_released; + +CREATE TRIGGER trg_notify_orch_work_consumed + AFTER DELETE ON orchestrator_queue + FOR EACH ROW + WHEN (OLD.lock_token IS NOT NULL) + EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work_consumed; + +CREATE TRIGGER trg_notify_session_lease + AFTER INSERT OR UPDATE OF locked_until ON sessions + FOR EACH ROW EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE sessions DISABLE TRIGGER trg_notify_session_lease; +``` + +## Function Changes + +### `cleanup_schema` - body modified (baseline: 0028) + +Full function with diff: +```diff + CREATE OR REPLACE FUNCTION SCHEMA.cleanup_schema() + RETURNS void + LANGUAGE plpgsql + AS $function$ + BEGIN + DROP TABLE IF EXISTS SCHEMA.sessions CASCADE; + DROP TABLE IF EXISTS SCHEMA.kv_delta CASCADE; + DROP TABLE IF EXISTS SCHEMA.kv_store CASCADE; + DROP TABLE IF EXISTS SCHEMA.instances CASCADE; + DROP TABLE IF EXISTS SCHEMA.executions CASCADE; + DROP TABLE IF EXISTS SCHEMA.history CASCADE; + DROP TABLE IF EXISTS SCHEMA.orchestrator_queue CASCADE; + DROP TABLE IF EXISTS SCHEMA.worker_queue CASCADE; + DROP TABLE IF EXISTS SCHEMA.instance_locks CASCADE; + DROP TABLE IF EXISTS SCHEMA._duroxide_migrations CASCADE; + ++ DROP FUNCTION IF EXISTS SCHEMA.duroxide_notify_queue_work(); + DROP FUNCTION IF EXISTS SCHEMA.cleanup_schema(); + DROP FUNCTION IF EXISTS SCHEMA.list_instances(); + DROP FUNCTION IF EXISTS SCHEMA.list_executions(TEXT); + DROP FUNCTION IF EXISTS SCHEMA.latest_execution_id(TEXT); + DROP FUNCTION IF EXISTS SCHEMA.list_instances_by_status(TEXT); + DROP FUNCTION IF EXISTS SCHEMA.get_instance_info(TEXT); + DROP FUNCTION IF EXISTS SCHEMA.get_execution_info(TEXT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.get_system_metrics(); + DROP FUNCTION IF EXISTS SCHEMA.get_queue_depths(BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.enqueue_worker_work(TEXT, BIGINT, TEXT, BIGINT, BIGINT, TEXT); + DROP FUNCTION IF EXISTS SCHEMA.enqueue_worker_work(TEXT, BIGINT, TEXT, BIGINT, BIGINT, TEXT, TEXT); + DROP FUNCTION IF EXISTS SCHEMA.ack_worker(TEXT, TEXT, TEXT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.renew_work_item_lock(TEXT, BIGINT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.fetch_work_item(BIGINT, BIGINT, TEXT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.fetch_work_item(BIGINT, BIGINT, TEXT, BIGINT, TEXT[], TEXT); + DROP FUNCTION IF EXISTS SCHEMA.abandon_work_item(TEXT, BIGINT, BIGINT, BOOLEAN); + DROP FUNCTION IF EXISTS SCHEMA.enqueue_orchestrator_work(TEXT, TEXT, TIMESTAMPTZ, TEXT, TEXT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.fetch_orchestration_item(BIGINT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.fetch_orchestration_item(BIGINT, BIGINT, BIGINT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.ack_orchestration_item(TEXT, BIGINT, BIGINT, JSONB, JSONB, JSONB, JSONB, JSONB); + DROP FUNCTION IF EXISTS SCHEMA.abandon_orchestration_item(TEXT, BIGINT, BIGINT, BOOLEAN); + DROP FUNCTION IF EXISTS SCHEMA.renew_orchestration_item_lock(TEXT, BIGINT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.fetch_history(TEXT); + DROP FUNCTION IF EXISTS SCHEMA.fetch_history_with_execution(TEXT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.append_history(TEXT, BIGINT, JSONB); + DROP FUNCTION IF EXISTS SCHEMA.list_children(TEXT); + DROP FUNCTION IF EXISTS SCHEMA.get_parent_id(TEXT); + DROP FUNCTION IF EXISTS SCHEMA.delete_instances_atomic(TEXT[], BOOLEAN); + DROP FUNCTION IF EXISTS SCHEMA.prune_executions(TEXT, INTEGER, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.prune_executions_batch(TEXT, BIGINT, BIGINT, INTEGER); + DROP FUNCTION IF EXISTS SCHEMA.renew_session_lock(TEXT[], BIGINT, BIGINT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.cleanup_orphaned_sessions(BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.get_custom_status(TEXT, BIGINT); + DROP FUNCTION IF EXISTS SCHEMA.get_kv_value(TEXT, TEXT); + DROP FUNCTION IF EXISTS SCHEMA.get_kv_all_values(TEXT); + DROP FUNCTION IF EXISTS SCHEMA.get_instance_stats(TEXT); + END; + $function$ +``` + +### `duroxide_notify_queue_work` - new function +```diff ++CREATE OR REPLACE FUNCTION SCHEMA.duroxide_notify_queue_work() ++ RETURNS trigger ++ LANGUAGE plpgsql ++ SET search_path TO 'pg_catalog', 'pg_temp' ++AS $function$ ++BEGIN ++ IF TG_TABLE_NAME = 'sessions' THEN ++ PERFORM pg_catalog.pg_notify( ++ 'duroxide_pg_' || pg_catalog.format('%I.worker_queue', TG_TABLE_SCHEMA)::pg_catalog.regclass::pg_catalog.oid::TEXT, ++ NEW.locked_until::TEXT ++ ); ++ RETURN NEW; ++ ELSIF TG_OP = 'DELETE' THEN ++ -- Acknowledgment consumes claimed rows and releases the instance lock ++ -- in one transaction; NOTIFY is delivered only after that commit. ++ PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); ++ RETURN OLD; ++ ELSIF TG_OP = 'UPDATE' AND TG_TABLE_NAME = 'orchestrator_queue' THEN ++ IF OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL THEN ++ -- Other visible messages can run even when this message is delayed. ++ PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); ++ END IF; ++ END IF; ++ ++ PERFORM pg_catalog.pg_notify( ++ 'duroxide_pg_' || TG_RELID::TEXT, ++ CASE WHEN NEW.lock_token IS NOT NULL THEN NEW.locked_until::TEXT ++ ELSE pg_catalog.ceil(EXTRACT(EPOCH FROM NEW.visible_at) * 1000)::BIGINT::TEXT ++ END ++ ); ++ RETURN NEW; ++END; ++$function$ +``` + +--- diff --git a/migrations/README.md b/migrations/README.md index 3114817..721fa13 100644 --- a/migrations/README.md +++ b/migrations/README.md @@ -90,6 +90,51 @@ CREATE INDEX IF NOT EXISTS idx_worker_processed ON worker_queue(processed_at); Migrations are automatically applied when creating a `PostgresProvider`. Each test schema gets its own migration history, so you can test migrations in isolation. +Migration `0030` installs disabled queue-notification triggers. A provider with +`MigrationPolicy::ApplyAll` and `long_poll.enabled = true` activates them for the +schema under the migration lock. `VerifyOnly` requires prior activation. See +[Migration Policy](../docs/migration-policy.md) for compatibility and privilege +requirements. +The six triggers cover both queues and session leases. Claims and renewals send +future lease-expiry hints; releasing an orchestration claim also sends an +immediate hint, even when the abandoned message is delayed. Installation and +activation use bounded transaction retries under lock contention. + +The long-poll migration is reserved as `0030`, after `0027`-`0028` from +[#32](https://github.com/microsoft/duroxide-pg/pull/32) and `0029` from +[#51](https://github.com/microsoft/duroxide-pg/pull/51). Its `cleanup_schema` +definition retains migration `0028`'s batched-pruning cleanup. +Disposable development/test databases that applied an earlier draft of the +long-poll migration as `0027` must be recreated before using this sequence. +Do not delete migration records or apply this draft over such a database that +contains data you need to preserve. + +The provider validation harness can exercise either mode: + +```sh +cargo nextest run --test postgres_provider_test +DUROXIDE_PG_TEST_LONG_POLL=true cargo nextest run --test postgres_provider_test +cargo nextest run --test long_poll_tests +``` + +The ignored two-phase-commit test additionally requires a disposable PostgreSQL +server with `max_prepared_transactions >= 1`: + +```sh +cargo nextest run --test long_poll_tests --run-ignored only \ + -E 'test(two_phase_commit_is_preserved_until_explicit_notification_activation)' +``` + +The ignored idle-query benchmark requires `pg_stat_statements` to be preloaded +and installed in a disposable local database, and an optimized PostgreSQL build +without assertions or memory wiping. It measures 65 seconds per mode, including +the normal long-poll timeouts and notifier refresh: + +```sh +cargo nextest run --release --test long_poll_tests --run-ignored only \ + -E 'test(idle_query_load_is_reduced_with_long_polling)' --success-output final +``` + ## Migration Execution Context - Migrations run inside transactions @@ -112,4 +157,3 @@ Migrations are automatically applied when creating a `PostgresProvider`. Each te ### Schema conflicts - Each schema maintains its own migration history - Migrations can be applied independently to different schemas - diff --git a/src/lib.rs b/src/lib.rs index 7549f3a..abf7c80 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -73,6 +73,27 @@ //! |---------------------|-------------|---------| //! | `DUROXIDE_PG_POOL_MAX` | Maximum connection pool size | `10` | //! +//! ## Optional Long-Polling +//! +//! Long-polling is disabled by default. Enable it on worker providers to wait +//! for PostgreSQL notifications instead of repeatedly querying empty queues: +//! +//! ```rust,no_run +//! use duroxide_pg::{PostgresProvider, ProviderConfig}; +//! +//! # async fn example(database_url: &str) -> anyhow::Result<()> { +//! let mut config = ProviderConfig::url(database_url); +//! config.long_poll.enabled = true; +//! let provider = PostgresProvider::new_with_config(config).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! One pool connection is reserved for LISTEN; at least two are required. +//! Client-only providers can leave this disabled. See [`LongPollConfig`] for +//! timer and recovery settings. A zero fetch timeout still returns immediately +//! after checking the queue. +//! //! ## Features //! //! - Automatic schema migration on startup @@ -82,7 +103,9 @@ pub mod entra; pub mod migrations; +mod notifier; pub mod provider; pub use entra::EntraAuthOptions; +pub use notifier::LongPollConfig; pub use provider::{ConnectionConfig, MigrationPolicy, PostgresProvider, ProviderConfig}; diff --git a/src/migrations.rs b/src/migrations.rs index 6f933f4..3eb3cd6 100644 --- a/src/migrations.rs +++ b/src/migrations.rs @@ -1,15 +1,29 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -use anyhow::Result; +use anyhow::{Context, Result}; use include_dir::{include_dir, Dir}; use sqlx::Connection; use sqlx::PgPool; use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use std::time::Duration; +use tokio::time::Instant; +// Touch this module when adding embedded files so incremental builds see them. Latest: 0030. static MIGRATIONS: Dir = include_dir!("$CARGO_MANIFEST_DIR/migrations"); +const DDL_LOCK_TIMEOUT_SQL: &str = "SET LOCAL lock_timeout = '250ms'"; +const DDL_RETRY_BUDGET: Duration = Duration::from_secs(10); +const NOTIFICATION_TRIGGERS: &[(&str, &str)] = &[ + ("worker_queue", "trg_notify_worker_work"), + ("worker_queue", "trg_notify_worker_work_released"), + ("orchestrator_queue", "trg_notify_orch_work"), + ("orchestrator_queue", "trg_notify_orch_work_released"), + ("orchestrator_queue", "trg_notify_orch_work_consumed"), + ("sessions", "trg_notify_session_lease"), +]; + const BASE_TABLES: &[&str] = &[ "instances", "executions", @@ -80,18 +94,30 @@ impl MigrationRunner { Ok(()) } - async fn unlock_for_migrations(&self, conn: &mut sqlx::postgres::PgConnection) { + async fn unlock_for_migrations(&self, conn: &mut sqlx::postgres::PgConnection) -> Result<()> { let key = self.advisory_lock_key(); - // Best-effort unlock. - let _ = sqlx::query("SELECT pg_advisory_unlock($1)") + let unlocked: bool = sqlx::query_scalar("SELECT pg_advisory_unlock($1)") .bind(key) - .execute(&mut *conn) - .await; + .fetch_one(&mut *conn) + .await + .context("Releasing the migration advisory lock")?; + anyhow::ensure!( + unlocked, + "Migration advisory lock was not held for schema {:?}", + self.schema_name + ); + Ok(()) } /// Run all pending migrations pub async fn migrate(&self) -> Result<()> { + self.migrate_with_notifications(false).await + } + + pub(crate) async fn migrate_with_notifications(&self, notifications: bool) -> Result<()> { let mut conn = self.pool.acquire().await?; + // Cancellation or failed unlocking must not return a session lock to the pool. + conn.close_on_drop(); let conn = &mut *conn; self.lock_for_migrations(conn).await?; @@ -99,13 +125,18 @@ impl MigrationRunner { // older binary cannot rewrite a schema that is ahead of its code. // Short-circuit: do NOT run migrate_inner if unknown migrations are // detected. - let result = match self.check_no_unknown_migrations(conn).await { - Ok(()) => self.migrate_inner(conn).await, - Err(e) => Err(e), - }; - self.unlock_for_migrations(conn).await; - - result + let result: Result<()> = async { + self.check_no_unknown_migrations(conn).await?; + self.migrate_inner(conn).await?; + if notifications { + self.enable_notifications(conn).await?; + } + Ok(()) + } + .await; + let unlock_result = self.unlock_for_migrations(conn).await; + result?; + unlock_result } /// Verify that the migration tracking table exists and that every embedded @@ -119,6 +150,10 @@ impl MigrationRunner { /// backends, where a separately privileged worker is responsible for /// applying schema changes). pub async fn verify(&self) -> Result<()> { + self.verify_with_notifications(false).await + } + + pub(crate) async fn verify_with_notifications(&self, notifications: bool) -> Result<()> { let mut conn = self.pool.acquire().await?; let conn = &mut *conn; @@ -164,9 +199,133 @@ impl MigrationRunner { ); } + if notifications + && self.notification_trigger_counts(conn).await?.1 != NOTIFICATION_TRIGGERS.len() as i64 + { + anyhow::bail!( + "long-poll notifications are not enabled in schema {:?}. \ + Construct a provider with MigrationPolicy::ApplyAll and \ + long_poll.enabled = true from a process with DDL privileges first.", + self.schema_name, + ); + } + Ok(()) } + async fn notification_trigger_counts( + &self, + conn: &mut sqlx::postgres::PgConnection, + ) -> Result<(i64, i64)> { + let (relations, names): (Vec<_>, Vec<_>) = NOTIFICATION_TRIGGERS + .iter() + .map(|(table, trigger)| (format!("{}.{table}", self.schema_name), *trigger)) + .unzip(); + let counts: (i64, i64) = sqlx::query_as( + "SELECT COUNT(*), COUNT(*) FILTER (WHERE tgenabled IN ('O', 'A')) \ + FROM pg_catalog.pg_trigger \ + JOIN unnest($1::TEXT[], $2::TEXT[]) expected(relation, name) \ + ON tgrelid = expected.relation::REGCLASS AND tgname = expected.name", + ) + .bind(relations) + .bind(names) + .fetch_one(conn) + .await?; + anyhow::ensure!( + counts.0 == NOTIFICATION_TRIGGERS.len() as i64, + "long-poll notification triggers are missing in schema {:?}; repair the schema before enabling long-polling", + self.schema_name, + ); + Ok(counts) + } + + async fn enable_notifications(&self, conn: &mut sqlx::postgres::PgConnection) -> Result<()> { + if self.notification_trigger_counts(conn).await?.1 == NOTIFICATION_TRIGGERS.len() as i64 { + return Ok(()); + } + + let deadline = Instant::now() + DDL_RETRY_BUDGET; + loop { + let result = async { + let mut tx = conn.begin().await?; + let result = async { + sqlx::query(DDL_LOCK_TIMEOUT_SQL).execute(&mut *tx).await?; + // Match acknowledgments, but bound every wait because admin + // deletion can acquire the queue locks in the opposite order. + for table in ["worker_queue", "orchestrator_queue", "sessions"] { + sqlx::query(&format!( + "LOCK TABLE {}.{table} IN SHARE ROW EXCLUSIVE MODE", + self.schema_name, + )) + .execute(&mut *tx) + .await?; + } + for (table, trigger) in NOTIFICATION_TRIGGERS { + sqlx::query(&format!( + "ALTER TABLE {}.{table} ENABLE TRIGGER {trigger}", + self.schema_name, + )) + .execute(&mut *tx) + .await?; + } + Ok(()) + } + .await; + Self::finish_ddl(tx, result).await + } + .await; + match result { + Ok(()) => return Ok(()), + Err(error) => { + if !self + .retry_ddl(&error, deadline, "activating long-poll notifications") + .await + { + return Err(error.context(format!( + "Failed to activate long-poll notifications in schema {:?}", + self.schema_name + ))); + } + } + } + } + } + + async fn finish_ddl( + tx: sqlx::Transaction<'_, sqlx::Postgres>, + result: Result<()>, + ) -> Result<()> { + match result { + Ok(()) => tx.commit().await.context("Committing DDL transaction"), + Err(error) => { + tx.rollback() + .await + .with_context(|| format!("Rolling back DDL after {error:#}"))?; + Err(error) + } + } + } + + async fn retry_ddl(&self, error: &anyhow::Error, deadline: Instant, operation: &str) -> bool { + let sqlstate = error + .downcast_ref::() + .and_then(sqlx::Error::as_database_error) + .and_then(|error| error.code()); + if !matches!(sqlstate.as_deref(), Some("40P01" | "55P03")) || Instant::now() >= deadline { + return false; + } + tracing::warn!( + schema = %self.schema_name, + operation, + sqlstate = sqlstate.as_deref(), + error = %format_args!("{error:#}"), + "DDL lock contention; retrying the rolled-back transaction" + ); + let jitter = Duration::from_millis(50 + u64::from(uuid::Uuid::new_v4().as_bytes()[0])); + tokio::time::sleep(jitter.min(deadline.saturating_duration_since(Instant::now()))).await; + Instant::now() < deadline + } + /// Check that the database has no migrations the running binary does not /// recognize. Used by both `migrate()` (to refuse running DDL against a /// schema ahead of the code) and `verify()` (to refuse claiming @@ -638,20 +797,24 @@ impl MigrationRunner { conn: &mut sqlx::postgres::PgConnection, migration: &Migration, ) -> Result<()> { - // Start transaction - let mut tx = conn.begin().await?; - - // Set search_path for this transaction. - // - // `pg_temp` is pinned explicitly at the lowest priority. Without it, - // `pg_temp` keeps its implicit highest-priority position, which would let a - // temporary object shadow the unqualified references this migration relies on - // the search_path to resolve. This is defense-in-depth (CVE-2018-1058 - // guidance); see `migration_search_path_stmt` for the threat model. - sqlx::query(&Self::migration_search_path_stmt(&self.schema_name)) - .execute(&mut *tx) - .await?; + let deadline = Instant::now() + DDL_RETRY_BUDGET; + loop { + match self.apply_migration_once(conn, migration).await { + Ok(()) => return Ok(()), + Err(error) => { + if !self.retry_ddl(&error, deadline, &migration.name).await { + return Err(error); + } + } + } + } + } + async fn apply_migration_once( + &self, + conn: &mut sqlx::postgres::PgConnection, + migration: &Migration, + ) -> Result<()> { // Remove comment lines and split SQL into individual statements let sql = migration.sql.trim(); let cleaned_sql: String = sql @@ -684,41 +847,46 @@ impl MigrationRunner { migration.version ); - for (idx, statement) in statements.iter().enumerate() { - if !statement.trim().is_empty() { - tracing::debug!( - "Executing statement {} of {}: {}...", - idx + 1, - statements.len(), - &statement.chars().take(50).collect::() - ); - sqlx::query(statement) + let mut tx = conn.begin().await?; + let result = + async { + sqlx::query(DDL_LOCK_TIMEOUT_SQL).execute(&mut *tx).await?; + sqlx::query(&Self::migration_search_path_stmt(&self.schema_name)) .execute(&mut *tx) - .await - .map_err(|e| { - anyhow::anyhow!( - "Failed to execute statement {} in migration {}: {}\nStatement: {}", + .await?; + for (idx, statement) in statements.iter().enumerate() { + if !statement.trim().is_empty() { + tracing::debug!( + "Executing statement {} of {}: {}...", idx + 1, - migration.version, - e, - statement + statements.len(), + &statement.chars().take(50).collect::() + ); + sqlx::query(statement) + .execute(&mut *tx) + .await + .with_context(|| { + format!( + "Failed to execute statement {} in migration {}:\nStatement: {}", + idx + 1, migration.version, statement ) - })?; - } - } + })?; + } + } - // Record migration as applied - sqlx::query(&format!( - "INSERT INTO {}._duroxide_migrations (version, name) VALUES ($1, $2)", - self.schema_name - )) - .bind(migration.version) - .bind(&migration.name) - .execute(&mut *tx) - .await?; + sqlx::query(&format!( + "INSERT INTO {}._duroxide_migrations (version, name) VALUES ($1, $2)", + self.schema_name + )) + .bind(migration.version) + .bind(&migration.name) + .execute(&mut *tx) + .await?; + Ok(()) + } + .await; - // Commit transaction - tx.commit().await?; + Self::finish_ddl(tx, result).await?; tracing::info!( "Applied migration {}: {}", @@ -734,6 +902,75 @@ impl MigrationRunner { mod tests { use super::*; + #[tokio::test] + async fn ddl_retry_logs_sqlstate_and_full_error_context() { + #[derive(Clone, Default)] + struct Capture(Arc>>); + + impl std::io::Write for Capture { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + dotenvy::dotenv().ok(); + let Ok(url) = std::env::var("DATABASE_URL") else { + eprintln!("DATABASE_URL not set; skipping DDL logging database test"); + return; + }; + let pool = sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .connect(&url) + .await + .unwrap(); + let runner = MigrationRunner::new(Arc::new(pool.clone()), "diagnostic_test".to_string()); + let captured = Capture::default(); + let writer = captured.clone(); + let subscriber = tracing_subscriber::fmt() + .without_time() + .with_ansi(false) + .with_env_filter("off,duroxide_pg::migrations=warn") + .with_writer(move || writer.clone()) + .finish(); + let _guard = tracing::subscriber::set_default(subscriber); + for (code, message) in [ + ("40P01", "simulated migration deadlock"), + ("55P03", "simulated migration lock timeout"), + ] { + let error = sqlx::query(&format!( + "DO $$ BEGIN RAISE EXCEPTION '{message}' USING ERRCODE = '{code}'; END $$" + )) + .execute(&pool) + .await + .unwrap_err(); + let error = anyhow::Error::new(error).context("outer migration statement context"); + assert!( + runner + .retry_ddl( + &error, + Instant::now() + Duration::from_secs(5), + "test migration" + ) + .await + ); + let output = String::from_utf8(captured.0.lock().unwrap().clone()).unwrap(); + assert!( + output.contains("sqlstate=") && output.contains(code), + "{output}" + ); + assert!( + output.contains("outer migration statement context") && output.contains(message), + "{output}" + ); + } + pool.close().await; + } + #[test] fn migration_search_path_pins_pg_temp_last() { // The security property: pg_temp must be present and last, so temporary diff --git a/src/notifier.rs b/src/notifier.rs new file mode 100644 index 0000000..8ef1dec --- /dev/null +++ b/src/notifier.rs @@ -0,0 +1,1320 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use anyhow::{ensure, Context, Result}; +use duroxide::providers::ProviderError; +use sqlx::postgres::{types::Oid, PgListener, PgNotification}; +use sqlx::{pool::PoolConnection, Connection, PgPool, Postgres}; +use std::future::Future; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Notify; +use tokio::task::JoinHandle; +use tokio::time::{interval_at, sleep, sleep_until, timeout_at, Instant, MissedTickBehavior}; +use tracing::warn; + +use crate::provider::PostgresProvider; + +const MIN_RETRY_INTERVAL: Duration = Duration::from_millis(100); +const DEGRADED_RETRY_INTERVAL: Duration = Duration::from_secs(1); +const TIMER_REFRESH_BACKSTOP: Duration = Duration::from_secs(1); + +/// Optional LISTEN/NOTIFY-based waiting for orchestration and activity work. +/// +/// Disabled by default. When enabled, one connection from the provider's pool +/// is reserved for LISTEN, so the pool must allow at least two connections. +/// +/// `ApplyAll` construction also activates notification triggers for the schema. +/// `VerifyOnly` requires that a privileged `ApplyAll` provider has done this +/// first. Turning this option off does not deactivate a schema's triggers. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct LongPollConfig { + /// Enable notification-driven waiting. Default: `false`. + pub enabled: bool, + /// Rescan future work and wake waiting fetches to recover missed + /// notifications. Must be nonzero when enabled. Default: 60 seconds. + pub notifier_poll_interval: Duration, + /// Delay added to future visibility timestamps to avoid waking early. + /// Default: 100 milliseconds. + pub timer_grace_period: Duration, +} + +impl Default for LongPollConfig { + fn default() -> Self { + Self { + enabled: false, + notifier_poll_interval: Duration::from_secs(60), + timer_grace_period: Duration::from_millis(100), + } + } +} + +impl LongPollConfig { + pub(crate) fn validate(&self, max_connections: u32) -> Result<()> { + if self.enabled { + ensure!( + max_connections >= 2, + "long-polling requires at least two pool connections: one is reserved for LISTEN" + ); + ensure!( + !self.notifier_poll_interval.is_zero(), + "long-polling notifier_poll_interval must be nonzero" + ); + deadline_after(self.notifier_poll_interval)?; + deadline_after(self.timer_grace_period)?; + } + Ok(()) + } +} + +pub(crate) struct LongPoller { + pub(crate) orch_notify: Arc, + pub(crate) worker_notify: Arc, + task: JoinHandle<()>, +} + +impl LongPoller { + pub(crate) async fn start( + pool: PgPool, + schema: &str, + config: LongPollConfig, + ) -> Result> { + if !config.enabled { + return Ok(None); + } + + // Resolve the same unquoted identifiers used by provider SQL. Table OIDs + // avoid identifier folding/truncation and NOTIFY's channel-length limit. + let (orch_oid, worker_oid): (Oid, Oid) = sqlx::query_as( + "SELECT $1::pg_catalog.text::pg_catalog.regclass::pg_catalog.oid, \ + $2::pg_catalog.text::pg_catalog.regclass::pg_catalog.oid", + ) + .bind(format!("{schema}.orchestrator_queue")) + .bind(format!("{schema}.worker_queue")) + .fetch_one(&pool) + .await?; + let channels = [ + format!("duroxide_pg_{}", orch_oid.0), + format!("duroxide_pg_{}", worker_oid.0), + ]; + let listener = subscribe(&pool, &channels).await?; + let orch_notify = Arc::new(Notify::new()); + let worker_notify = Arc::new(Notify::new()); + let notifications = [orch_notify.clone(), worker_notify.clone()]; + let timer_query = format!( + "SELECT LEAST(\ + (SELECT CEIL(EXTRACT(EPOCH FROM visible_at) * 1000)::BIGINT \ + FROM {schema}.orchestrator_queue \ + WHERE visible_at > TO_TIMESTAMP($1::BIGINT / 1000.0) \ + ORDER BY visible_at LIMIT 1), \ + (SELECT locked_until FROM {schema}.instance_locks \ + WHERE locked_until > $1 ORDER BY locked_until LIMIT 1)), \ + LEAST(\ + (SELECT CEIL(EXTRACT(EPOCH FROM visible_at) * 1000)::BIGINT \ + FROM {schema}.worker_queue \ + WHERE visible_at > TO_TIMESTAMP($1::BIGINT / 1000.0) \ + ORDER BY visible_at LIMIT 1), \ + (SELECT MIN(locked_until) FROM {schema}.worker_queue \ + WHERE lock_token IS NOT NULL AND locked_until > $1), \ + (SELECT MIN(locked_until) FROM {schema}.sessions WHERE locked_until > $1))" + ); + let task = tokio::spawn(run( + pool, + listener, + channels, + notifications, + timer_query, + config, + )); + Ok(Some(Self { + orch_notify, + worker_notify, + task, + })) + } + + pub(crate) fn stop(&self) { + self.task.abort(); + } +} + +impl Drop for LongPoller { + fn drop(&mut self) { + self.stop(); + } +} + +async fn subscribe(pool: &PgPool, channels: &[String; 2]) -> Result { + let mut listener = PgListener::connect_with(pool).await?; + // Detect loss before SQLx silently reconnects, so timers and ready work can + // be recovered after resubscribing. + listener.eager_reconnect(false); + listener + .listen_all(channels.iter().map(String::as_str)) + .await?; + // LISTEN must not consume the pool's only warm connection. + drop( + pool.acquire() + .await + .context("Warming a fetch connection after LISTEN")?, + ); + Ok(listener) +} + +async fn run( + pool: PgPool, + mut listener: PgListener, + channels: [String; 2], + notifications: [Arc; 2], + timer_query: String, + config: LongPollConfig, +) { + let mut refresh = TimerRefresh::default(); + loop { + let result = run_connected( + &pool, + &mut listener, + &channels, + ¬ifications, + &timer_query, + &config, + &mut refresh, + ) + .await; + for notify in ¬ifications { + notify.notify_waiters(); + } + if pool.is_closed() { + return; + } + match result { + Ok(()) => warn!( + target: "duroxide::providers::postgres::notifier", + "LISTEN connection lost; reconnecting" + ), + Err(error) => warn!( + target: "duroxide::providers::postgres::notifier", + error = %error, + "Long-poll listener failed; reconnecting" + ), + } + drop(listener); + let reconnect = async { + loop { + sleep(DEGRADED_RETRY_INTERVAL).await; + match subscribe(&pool, &channels).await { + Ok(listener) => break listener, + Err(error) => warn!( + target: "duroxide::providers::postgres::notifier", + error = %error, + "Long-poll reconnect failed; continuing one-second polling fallback" + ), + } + } + }; + tokio::pin!(reconnect); + let mut fallback = interval_at( + Instant::now() + DEGRADED_RETRY_INTERVAL, + DEGRADED_RETRY_INTERVAL, + ); + fallback.set_missed_tick_behavior(MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = pool.close_event() => { + for notify in ¬ifications { + notify.notify_waiters(); + } + return; + } + _ = fallback.tick() => { + for notify in ¬ifications { + notify.notify_waiters(); + } + } + reconnected = &mut reconnect => { + listener = reconnected; + for notify in ¬ifications { + notify.notify_waiters(); + } + break; + } + } + } + } +} + +type TimerDeadlines = (Option, Option); + +#[derive(Default)] +struct TimerRefresh { + task: Option>>, +} + +impl TimerRefresh { + fn start(&mut self, pool: &PgPool, query: &str) { + // Keep the same scan across listener reconnects and client watchdog wakes. + if self.task.is_none() { + let pool = pool.clone(); + let query = query.to_string(); + self.task = Some(tokio::spawn( + async move { refresh_timers(&pool, &query).await }, + )); + } + } + + async fn result(&mut self) -> Result { + let result = self.task.as_mut().expect("timer scan was started").await; + self.task.take(); + result.context("Long-poll timer scan task failed")? + } +} + +impl Drop for TimerRefresh { + fn drop(&mut self) { + if let Some(task) = &self.task { + task.abort(); + } + } +} + +struct RefreshConnection(Option>); + +impl Drop for RefreshConnection { + fn drop(&mut self) { + // SQLx's normal pool return pings the connection, which can wait forever + // for a cancelled query. Discard only unfinished refresh connections. + if let Some(connection) = self.0.take() { + drop(connection.detach()); + } + } +} + +async fn refresh_timers(pool: &PgPool, query: &str) -> Result { + let mut connection = RefreshConnection(Some(pool.acquire().await?)); + let mut tx = connection + .0 + .as_mut() + .expect("refresh owns its connection") + .begin() + .await?; + let result: Result<_> = async { + sqlx::Executor::execute( + &mut *tx, + "SET LOCAL lock_timeout = '200ms'; SET LOCAL statement_timeout = '200ms'", + ) + .await + .context("Setting long-poll timer scan timeouts")?; + sqlx::query_as(query) + .bind(PostgresProvider::now_millis()) + .fetch_one(&mut *tx) + .await + .context("Reading long-poll timer deadlines") + } + .await; + match &result { + Ok(_) => tx + .commit() + .await + .context("Committing long-poll timer scan")?, + Err(error) => tx + .rollback() + .await + .with_context(|| format!("Rolling back long-poll timer scan after {error:#}"))?, + } + drop(connection.0.take()); + result +} + +async fn run_connected( + pool: &PgPool, + listener: &mut PgListener, + channels: &[String; 2], + notifications: &[Arc; 2], + timer_query: &str, + config: &LongPollConfig, + refresh: &mut TimerRefresh, +) -> Result<()> { + let mut timers: [Option; 2] = [None, None]; + let mut next_refresh = Instant::now(); + let mut next_recovery = next_refresh; + let mut refresh_started = next_refresh; + refresh.start(pool, timer_query); + let mut refreshing = true; + let mut refresh_failed = false; + let mut refresh_stalled = false; + let mut next_backstop = refresh_started + TIMER_REFRESH_BACKSTOP; + // A paced scan must recheck work that matured since a queue's timer wake. + // Keep events during an in-flight scan pending for the following scan. + let mut pending_catchup = [false; 2]; + let mut refreshing_catchup = [false; 2]; + + loop { + let next_timer = timers.iter().flatten().copied().min(); + tokio::select! { + _ = pool.close_event() => return Ok(()), + notification = listener.try_recv() => { + let Some(notification) = notification? else { + return Ok(()); + }; + handle_notification(notification, channels, notifications, &mut timers, config); + } + result = refresh.result(), if refreshing => { + refreshing = false; + let recover = refresh_failed || Instant::now() >= next_recovery; + if recover { + next_recovery = deadline_after(config.notifier_poll_interval)?; + } + next_refresh = next_recovery; + match result { + Ok((orch, worker)) => { + refresh_failed = false; + for (index, timestamp) in [orch, worker].into_iter().enumerate() { + // Recovery must also find ready work whose notification + // was lost. Timer-only scans need not wake the other queue. + if recover || refreshing_catchup[index] { + notifications[index].notify_waiters(); + } + if let Some(timestamp) = timestamp { + pending_catchup[index] |= timestamp <= PostgresProvider::now_millis(); + schedule(timestamp, ¬ifications[index], &mut timers[index], config); + } + } + if pending_catchup.iter().any(|pending| *pending) { + next_refresh = next_refresh.min( + (refresh_started + MIN_RETRY_INTERVAL).max(Instant::now()) + ); + } + } + Err(error) => { + refresh_failed = true; + next_refresh = (refresh_started + DEGRADED_RETRY_INTERVAL).max(Instant::now()); + for notify in notifications { + notify.notify_waiters(); + } + for (pending, attempted) in pending_catchup.iter_mut().zip(refreshing_catchup) { + *pending |= attempted; + } + warn!( + target: "duroxide::providers::postgres::notifier", + error = %format_args!("{error:#}"), + "Long-poll timer refresh failed; retrying with one-second polling fallback" + ); + } + } + } + _ = sleep_until(next_backstop), if refreshing => { + if !refresh_stalled { + warn!( + target: "duroxide::providers::postgres::notifier", + "Long-poll timer scan is still running; polling without cancelling or starting another scan" + ); + } + refresh_stalled = true; + refresh_failed = true; + for notify in notifications { + notify.notify_waiters(); + } + next_backstop = Instant::now() + DEGRADED_RETRY_INTERVAL; + } + _ = sleep_until(next_refresh), if !refreshing => { + refresh_started = Instant::now(); + refreshing_catchup = std::mem::take(&mut pending_catchup); + refresh.start(pool, timer_query); + refreshing = true; + refresh_stalled = false; + next_backstop = refresh_started + TIMER_REFRESH_BACKSTOP; + } + _ = sleep_until(next_timer.unwrap_or(next_refresh)), if next_timer.is_some() => { + let now = Instant::now(); + for (index, timer) in timers.iter_mut().enumerate() { + if timer.is_some_and(|deadline| deadline <= now) { + *timer = None; + pending_catchup[index] = true; + notifications[index].notify_waiters(); + } + } + // Only the earliest timer per queue is retained. Rescan after it + // fires to discover the next one without an unbounded timer heap. + next_refresh = next_refresh.min( + (refresh_started + MIN_RETRY_INTERVAL).max(now) + ); + } + } + } +} + +fn handle_notification( + notification: PgNotification, + channels: &[String; 2], + notifications: &[Arc; 2], + timers: &mut [Option; 2], + config: &LongPollConfig, +) { + let Some(index) = channels + .iter() + .position(|channel| channel == notification.channel()) + else { + warn!( + target: "duroxide::providers::postgres::notifier", + "Ignoring a notification from an unexpected channel" + ); + return; + }; + match notification.payload().parse::() { + Ok(timestamp) => schedule(timestamp, ¬ifications[index], &mut timers[index], config), + Err(error) => { + warn!( + target: "duroxide::providers::postgres::notifier", + error = %error, + "Invalid long-poll timestamp; waking fetches to recheck the queue" + ); + notifications[index].notify_waiters(); + } + } +} + +fn schedule(timestamp: i64, notify: &Notify, timer: &mut Option, config: &LongPollConfig) { + let now = PostgresProvider::now_millis(); + if timestamp <= now { + notify.notify_waiters(); + return; + } + let delay = Duration::from_millis(timestamp.saturating_sub(now) as u64) + .checked_add(config.timer_grace_period); + let deadline = delay.and_then(|delay| Instant::now().checked_add(delay)); + match deadline { + Some(deadline) => *timer = Some(timer.map_or(deadline, |current| current.min(deadline))), + None => warn!( + target: "duroxide::providers::postgres::notifier", + "Long-poll timer exceeds the clock range; relying on polling fallback" + ), + } +} + +fn deadline_after(duration: Duration) -> Result { + Instant::now() + .checked_add(duration) + .context("long-poll duration exceeds the supported clock range") +} + +pub(crate) async fn fetch_with_poll( + notify: Option<&Notify>, + poll_timeout: Duration, + operation: &'static str, + mut fetch: F, +) -> Result, ProviderError> +where + F: FnMut() -> Fut, + Fut: Future, ProviderError>>, +{ + let Some(notify) = notify.filter(|_| !poll_timeout.is_zero()) else { + return fetch().await; + }; + let deadline = Instant::now().checked_add(poll_timeout).ok_or_else(|| { + ProviderError::permanent(operation, "poll_timeout exceeds the supported clock range") + })?; + let mut earliest_retry = Instant::now(); + loop { + let notified = notify.notified(); + tokio::pin!(notified); + // Register before reading the queue so a concurrent commit cannot be lost. + notified.as_mut().enable(); + let result = fetch().await?; + if result.is_some() { + return Ok(result); + } + if Instant::now() >= deadline || timeout_at(deadline, notified).await.is_err() { + return Ok(None); + } + // The first wake is immediate; repeated unsuccessful wakes coalesce. + sleep_until(earliest_retry.min(deadline)).await; + if Instant::now() >= deadline { + return Ok(None); + } + earliest_retry = Instant::now() + MIN_RETRY_INTERVAL; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::time::timeout; + + #[test] + fn defaults_preserve_short_polling() { + let config = LongPollConfig::default(); + assert!(!config.enabled); + assert_eq!(config.notifier_poll_interval, Duration::from_secs(60)); + assert_eq!(config.timer_grace_period, Duration::from_millis(100)); + config.validate(1).unwrap(); + } + + #[test] + fn enabled_configuration_requires_usable_pool_and_timers() { + let mut config = LongPollConfig { + enabled: true, + ..Default::default() + }; + assert!(config.validate(1).unwrap_err().to_string().contains("two")); + config.validate(2).unwrap(); + config.notifier_poll_interval = Duration::ZERO; + assert!(config.validate(2).is_err()); + config.notifier_poll_interval = Duration::MAX; + assert!(config.validate(2).is_err()); + config.notifier_poll_interval = Duration::from_secs(60); + config.timer_grace_period = Duration::MAX; + assert!(config.validate(2).is_err()); + } + + #[tokio::test] + async fn empty_long_poll_fetches_once_without_busy_polling() { + let notify = Notify::new(); + let calls = AtomicUsize::new(0); + let start = Instant::now(); + let result = fetch_with_poll(Some(¬ify), Duration::from_millis(50), "test", || async { + calls.fetch_add(1, Ordering::Relaxed); + Ok(None::<()>) + }) + .await + .unwrap(); + assert!(result.is_none()); + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert!(start.elapsed() >= Duration::from_millis(50)); + } + + #[tokio::test] + async fn notification_between_empty_read_and_wait_is_not_lost() { + let notify = Notify::new(); + let calls = AtomicUsize::new(0); + let result = fetch_with_poll( + Some(¬ify), + Duration::from_millis(100), + "test", + || async { + if calls.fetch_add(1, Ordering::Relaxed) == 0 { + notify.notify_waiters(); + Ok(None) + } else { + Ok(Some(7)) + } + }, + ) + .await + .unwrap(); + assert_eq!(result, Some(7)); + assert_eq!(calls.load(Ordering::Relaxed), 2); + } + + #[tokio::test] + async fn repeated_unmatched_notifications_do_not_extend_deadline() { + let notify = Notify::new(); + let calls = AtomicUsize::new(0); + let start = Instant::now(); + let result = fetch_with_poll(Some(¬ify), Duration::from_millis(30), "test", || async { + calls.fetch_add(1, Ordering::Relaxed); + notify.notify_waiters(); + Ok(None::<()>) + }) + .await + .unwrap(); + assert!(result.is_none()); + assert!(calls.load(Ordering::Relaxed) > 1); + assert!(start.elapsed() < Duration::from_millis(300)); + } + + #[tokio::test] + async fn repeated_empty_wakes_are_paced_but_the_first_wake_is_immediate() { + let notify = Notify::new(); + let calls = std::sync::Mutex::new(Vec::new()); + let result = fetch_with_poll( + Some(¬ify), + Duration::from_millis(350), + "test", + || async { + calls.lock().unwrap().push(Instant::now()); + notify.notify_waiters(); + Ok(None::<()>) + }, + ) + .await + .unwrap(); + assert!(result.is_none()); + let calls = calls.lock().unwrap(); + assert!( + (2..=5).contains(&calls.len()), + "fetch count: {}", + calls.len() + ); + assert!(calls[1].duration_since(calls[0]) < MIN_RETRY_INTERVAL); + for pair in calls[1..].windows(2) { + assert!(pair[1].duration_since(pair[0]) >= MIN_RETRY_INTERVAL); + } + } + + #[tokio::test] + async fn timer_scan_timeouts_are_server_errors_and_reuse_the_connection() { + dotenvy::dotenv().ok(); + let Ok(url) = std::env::var("DATABASE_URL") else { + eprintln!("DATABASE_URL not set; skipping timer-scan database test"); + return; + }; + let pool = sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .connect(&url) + .await + .unwrap(); + let control = sqlx::postgres::PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .unwrap(); + let schema = format!("bounded_scan_{}", uuid::Uuid::new_v4().simple()); + sqlx::Executor::execute( + &control, + format!( + "CREATE SCHEMA {schema}; CREATE TABLE {schema}.probe (id INTEGER); + INSERT INTO {schema}.probe VALUES (1)" + ) + .as_str(), + ) + .await + .unwrap(); + let before: (i32, String, String) = sqlx::query_as( + "SELECT pg_backend_pid(), current_setting('lock_timeout'), current_setting('statement_timeout')", + ).fetch_one(&pool).await.unwrap(); + let valid = "SELECT $1::BIGINT, NULL::BIGINT"; + assert!(refresh_timers(&pool, valid).await.unwrap().0.is_some()); + let mut blocker = control.begin().await.unwrap(); + sqlx::query(&format!( + "LOCK TABLE {schema}.probe IN ACCESS EXCLUSIVE MODE" + )) + .execute(&mut *blocker) + .await + .unwrap(); + let mut outcomes = Vec::new(); + for query in [ + "SELECT $1::BIGINT, NULL::BIGINT FROM pg_sleep(2)".to_string(), + format!("SELECT $1::BIGINT, NULL::BIGINT FROM {schema}.probe"), + ] { + let error = timeout(Duration::from_secs(1), refresh_timers(&pool, &query)) + .await + .expect("server timeout should beat the client backstop") + .unwrap_err(); + let after: (i32, String, String) = sqlx::query_as( + "SELECT pg_backend_pid(), current_setting('lock_timeout'), current_setting('statement_timeout')", + ).fetch_one(&pool).await.unwrap(); + outcomes.push((error, after)); + } + blocker.rollback().await.unwrap(); + let rows: i64 = sqlx::query_scalar(&format!("SELECT count(*) FROM {schema}.probe")) + .fetch_one(&control) + .await + .unwrap(); + assert!(refresh_timers(&pool, valid).await.unwrap().0.is_some()); + pool.close().await; + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&control) + .await + .unwrap(); + control.close().await; + assert_eq!(rows, 1, "failed scans must not change the durable state"); + for (error, after) in outcomes { + let code = error + .downcast_ref::() + .and_then(sqlx::Error::as_database_error) + .and_then(|error| error.code()); + assert!( + matches!(code.as_deref(), Some("55P03" | "57014")), + "{error:#}" + ); + assert_eq!( + after, before, + "timeouts must roll back and return the same healthy connection" + ); + } + } + + #[tokio::test] + async fn dense_overdue_refreshes_are_paced_and_do_not_wake_the_other_queue() { + dotenvy::dotenv().ok(); + let Ok(url) = std::env::var("DATABASE_URL") else { + eprintln!("DATABASE_URL not set; skipping timer-refresh database test"); + return; + }; + let pool = sqlx::postgres::PgPoolOptions::new() + .max_connections(3) + .connect(&url) + .await + .unwrap(); + let schema = format!("timer_scan_{}", uuid::Uuid::new_v4().simple()); + sqlx::query(&format!("CREATE SCHEMA {schema}")) + .execute(&pool) + .await + .unwrap(); + sqlx::query(&format!("CREATE SEQUENCE {schema}.calls")) + .execute(&pool) + .await + .unwrap(); + let channels = [format!("{schema}_orch"), format!("{schema}_worker")]; + let mut listener = subscribe(&pool, &channels).await.unwrap(); + let notifications = [Arc::new(Notify::new()), Arc::new(Notify::new())]; + let config = LongPollConfig { + enabled: true, + ..Default::default() + }; + // A future timestamp matures during every query, as with dense timers + // and nontrivial round-trip latency. Count scans without global resets. + let query = format!( + "SELECT $1::BIGINT + 1, NULL::BIGINT FROM pg_catalog.pg_sleep(0.005), \ + pg_catalog.nextval('{schema}.calls')", + ); + let mut worker_wakes = 0; + { + let mut refresh = TimerRefresh::default(); + let run = run_connected( + &pool, + &mut listener, + &channels, + ¬ifications, + &query, + &config, + &mut refresh, + ); + tokio::pin!(run); + let duration = tokio::time::timeout(Duration::from_millis(350), async { + loop { + tokio::select! { + result = &mut run => panic!("listener stopped unexpectedly: {result:?}"), + _ = notifications[1].notified() => worker_wakes += 1, + } + } + }); + assert!(duration.await.is_err()); + } + let calls: i64 = sqlx::query_scalar(&format!("SELECT last_value FROM {schema}.calls")) + .fetch_one(&pool) + .await + .unwrap(); + drop(listener); + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&pool) + .await + .unwrap(); + pool.close().await; + assert!((2..=5).contains(&calls), "refreshes in 350ms: {calls}"); + assert_eq!( + worker_wakes, 1, + "only the initial recovery should wake the empty worker queue" + ); + } + + #[tokio::test] + async fn paced_catchup_wakes_waiters_when_no_future_timestamp_remains() { + dotenvy::dotenv().ok(); + let Ok(url) = std::env::var("DATABASE_URL") else { + eprintln!("DATABASE_URL not set; skipping catch-up database test"); + return; + }; + let pool = sqlx::postgres::PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .unwrap(); + let control = sqlx::postgres::PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .unwrap(); + let schema = format!("catchup_{}", uuid::Uuid::new_v4().simple()); + let gate_key = i64::from_be_bytes(uuid::Uuid::new_v4().as_bytes()[..8].try_into().unwrap()); + sqlx::Executor::execute( + &control, + format!( + "CREATE SCHEMA {schema}; + CREATE SEQUENCE {schema}.calls; + CREATE FUNCTION {schema}.scan(p_now BIGINT) RETURNS TABLE(orch BIGINT, worker BIGINT) + LANGUAGE plpgsql AS $$ + DECLARE n BIGINT := nextval('{schema}.calls'); + BEGIN + IF n = 2 THEN + PERFORM pg_advisory_xact_lock({gate_key}); + END IF; + RETURN QUERY SELECT CASE WHEN n = 1 THEN p_now + 20 END, NULL::BIGINT; + END $$" + ) + .as_str(), + ) + .await + .unwrap(); + let mut gate = control.begin().await.unwrap(); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(gate_key) + .execute(&mut *gate) + .await + .unwrap(); + let channels = [format!("{schema}_orch"), format!("{schema}_worker")]; + let mut listener = subscribe(&pool, &channels).await.unwrap(); + let notifications = [Arc::new(Notify::new()), Arc::new(Notify::new())]; + let task = tokio::spawn({ + let pool = pool.clone(); + let notifications = notifications.clone(); + let query = format!("SELECT * FROM {schema}.scan($1)"); + async move { + let mut refresh = TimerRefresh::default(); + run_connected( + &pool, + &mut listener, + &channels, + ¬ifications, + &query, + &LongPollConfig { + enabled: true, + ..Default::default() + }, + &mut refresh, + ) + .await + } + }); + let pid = timeout(Duration::from_secs(2), async { + loop { + let pid: Option = sqlx::query_scalar( + "SELECT pid FROM pg_stat_activity WHERE wait_event_type = 'Lock' \ + AND strpos(query, $1) > 0", + ) + .bind(format!("{schema}.scan")) + .fetch_optional(&control) + .await + .unwrap(); + if let Some(pid) = pid { + break pid; + } + sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("the catch-up scan must reach the gate"); + // Register after the timer's wake but before its catch-up scan completes. + let notified = notifications[0].notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + gate.rollback().await.unwrap(); + let wake = timeout(Duration::from_secs(1), notified).await; + let scan_succeeded: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE pid = $1 AND state = 'idle')", + ) + .bind(pid) + .fetch_one(&control) + .await + .unwrap(); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + pool.close().await; + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&control) + .await + .unwrap(); + control.close().await; + wake.expect("a catch-up scan must wake work that matured before its snapshot"); + assert!( + scan_succeeded, + "a scan timeout must not substitute for the catch-up wake" + ); + } + + #[tokio::test] + async fn failed_refresh_retries_before_the_recovery_interval() { + dotenvy::dotenv().ok(); + let Ok(url) = std::env::var("DATABASE_URL") else { + eprintln!("DATABASE_URL not set; skipping timer-refresh database test"); + return; + }; + let pool = sqlx::postgres::PgPoolOptions::new() + .max_connections(3) + .connect(&url) + .await + .unwrap(); + let schema = format!("failed_scan_{}", uuid::Uuid::new_v4().simple()); + sqlx::query(&format!("CREATE SCHEMA {schema}")) + .execute(&pool) + .await + .unwrap(); + sqlx::query(&format!("CREATE SEQUENCE {schema}.calls")) + .execute(&pool) + .await + .unwrap(); + let channels = [format!("{schema}_orch"), format!("{schema}_worker")]; + let mut listener = subscribe(&pool, &channels).await.unwrap(); + let notifications = [Arc::new(Notify::new()), Arc::new(Notify::new())]; + let query = format!( + "WITH attempt AS MATERIALIZED (SELECT nextval('{schema}.calls') AS n) \ + SELECT 1 / (n - 1), NULL::BIGINT FROM attempt WHERE $1::BIGINT IS NOT NULL" + ); + let mut refresh = TimerRefresh::default(); + let result = tokio::time::timeout( + Duration::from_millis(1500), + run_connected( + &pool, + &mut listener, + &channels, + ¬ifications, + &query, + &LongPollConfig { + enabled: true, + ..Default::default() + }, + &mut refresh, + ), + ) + .await; + drop(refresh); + assert!( + result.is_err(), + "the listener should stay connected after a scan error" + ); + let calls: i64 = sqlx::query_scalar(&format!("SELECT last_value FROM {schema}.calls")) + .fetch_one(&pool) + .await + .unwrap(); + drop(listener); + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&pool) + .await + .unwrap(); + pool.close().await; + assert!( + calls >= 2, + "a failed scan waited for the 60-second recovery: {calls}" + ); + } + + #[tokio::test] + async fn subscribing_keeps_a_warm_fetch_connection() { + dotenvy::dotenv().ok(); + let Ok(url) = std::env::var("DATABASE_URL") else { + eprintln!("DATABASE_URL not set; skipping listener database test"); + return; + }; + let pool = sqlx::postgres::PgPoolOptions::new() + .min_connections(1) + .max_connections(2) + .connect(&url) + .await + .unwrap(); + let channel = format!("warm_listener_{}", uuid::Uuid::new_v4().simple()); + let listener = subscribe( + &pool, + &[format!("{channel}_orch"), format!("{channel}_worker")], + ) + .await + .unwrap(); + let connections = pool.size(); + drop(listener); + pool.close().await; + assert_eq!(connections, 2, "LISTEN must leave another connection warm"); + } + + #[tokio::test] + async fn watchdog_and_reconnect_keep_the_same_pending_timer_scan() { + use std::sync::atomic::AtomicBool; + + dotenvy::dotenv().ok(); + let Ok(url) = std::env::var("DATABASE_URL") else { + eprintln!("DATABASE_URL not set; skipping pending-scan database test"); + return; + }; + let block_next = Arc::new(AtomicBool::new(false)); + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let pool = sqlx::postgres::PgPoolOptions::new() + .min_connections(3) + .max_connections(3) + .before_acquire({ + let block_next = block_next.clone(); + let entered = entered.clone(); + let release = release.clone(); + move |_, _| { + let block_next = block_next.clone(); + let entered = entered.clone(); + let release = release.clone(); + Box::pin(async move { + if block_next.swap(false, Ordering::Relaxed) { + let released = release.notified(); + tokio::pin!(released); + released.as_mut().enable(); + entered.notify_one(); + released.await; + } + Ok(true) + }) + } + }) + .connect(&url) + .await + .unwrap(); + let control = sqlx::postgres::PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .unwrap(); + let schema = format!("pending_scan_{}", uuid::Uuid::new_v4().simple()); + sqlx::Executor::execute( + &control, + format!("CREATE SCHEMA {schema}; CREATE SEQUENCE {schema}.calls").as_str(), + ) + .await + .unwrap(); + let channels = [format!("{schema}_orch"), format!("{schema}_worker")]; + let listener = subscribe(&pool, &channels).await.unwrap(); + let pid: i32 = sqlx::query_scalar( + "SELECT pid FROM pg_stat_activity WHERE query LIKE 'LISTEN%' AND strpos(query, $1) > 0", + ) + .bind(&channels[0]) + .fetch_one(&control) + .await + .unwrap(); + let notifications = [Arc::new(Notify::new()), Arc::new(Notify::new())]; + let initial_scan = notifications[0].notified(); + tokio::pin!(initial_scan); + initial_scan.as_mut().enable(); + let task = tokio::spawn(run( + pool.clone(), listener, channels.clone(), notifications.clone(), + format!("SELECT NULL::BIGINT, NULL::BIGINT FROM nextval('{schema}.calls') WHERE $1::BIGINT IS NOT NULL"), + LongPollConfig { + enabled: true, + notifier_poll_interval: Duration::from_millis(50), + ..Default::default() + }, + )); + timeout(Duration::from_secs(3), initial_scan).await.unwrap(); + block_next.store(true, Ordering::Relaxed); + timeout(Duration::from_secs(3), entered.notified()) + .await + .unwrap(); + for _ in 0..2 { + timeout(Duration::from_millis(1500), notifications[0].notified()) + .await + .expect("the watchdog must wake fetches while a scan is pending"); + } + let before_reconnect: i64 = + sqlx::query_scalar(&format!("SELECT last_value FROM {schema}.calls")) + .fetch_one(&control) + .await + .unwrap(); + sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .execute(&control) + .await + .unwrap(); + timeout(Duration::from_secs(3), async { + loop { + let reconnected: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE query LIKE 'LISTEN%' \ + AND strpos(query, $1) > 0 AND pid <> $2)", + ) + .bind(&channels[0]) + .bind(pid) + .fetch_one(&control) + .await + .unwrap(); + if reconnected { + break; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("listener must reconnect without replacing the pending scan"); + let after_reconnect: i64 = + sqlx::query_scalar(&format!("SELECT last_value FROM {schema}.calls")) + .fetch_one(&control) + .await + .unwrap(); + release.notify_waiters(); + timeout(Duration::from_secs(3), async { + loop { + let calls: i64 = + sqlx::query_scalar(&format!("SELECT last_value FROM {schema}.calls")) + .fetch_one(&control) + .await + .unwrap(); + if calls >= 2 { + break; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the original scan must resume once connection acquisition finishes"); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + timeout(Duration::from_secs(3), pool.close()).await.unwrap(); + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&control) + .await + .unwrap(); + control.close().await; + assert_eq!( + before_reconnect, 1, + "the watchdog must not replace the blocked scan" + ); + assert_eq!( + after_reconnect, 1, + "reconnecting must not replace the blocked scan" + ); + } + + #[tokio::test] + async fn fallback_wakes_fetches_while_reconnect_is_pending() { + use std::sync::atomic::AtomicBool; + + dotenvy::dotenv().ok(); + let Ok(url) = std::env::var("DATABASE_URL") else { + eprintln!("DATABASE_URL not set; skipping reconnect database test"); + return; + }; + let blocked = Arc::new(AtomicBool::new(false)); + let reconnecting = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let pool = sqlx::postgres::PgPoolOptions::new() + .min_connections(2) + .max_connections(2) + .before_acquire({ + let blocked = blocked.clone(); + let reconnecting = reconnecting.clone(); + let release = release.clone(); + move |_, _| { + let blocked = blocked.clone(); + let reconnecting = reconnecting.clone(); + let release = release.clone(); + Box::pin(async move { + if blocked.load(Ordering::Relaxed) { + let released = release.notified(); + tokio::pin!(released); + released.as_mut().enable(); + reconnecting.notify_one(); + released.await; + } + Ok(true) + }) + } + }) + .connect(&url) + .await + .unwrap(); + let control = sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .connect(&url) + .await + .unwrap(); + let channel = format!("pending_reconnect_{}", uuid::Uuid::new_v4().simple()); + let channels = [format!("{channel}_orch"), format!("{channel}_worker")]; + let listener = subscribe(&pool, &channels).await.unwrap(); + let pid: i32 = sqlx::query_scalar( + "SELECT pid FROM pg_stat_activity WHERE query LIKE 'LISTEN%' AND strpos(query, $1) > 0", + ) + .bind(&channels[0]) + .fetch_one(&control) + .await + .unwrap(); + let notifications = [Arc::new(Notify::new()), Arc::new(Notify::new())]; + let initial_scan = notifications[0].notified(); + tokio::pin!(initial_scan); + initial_scan.as_mut().enable(); + let task = tokio::spawn(run( + pool.clone(), + listener, + channels, + notifications.clone(), + "SELECT NULL::BIGINT, NULL::BIGINT WHERE $1::BIGINT IS NOT NULL".to_string(), + LongPollConfig { + enabled: true, + ..Default::default() + }, + )); + tokio::time::timeout(Duration::from_secs(3), initial_scan) + .await + .expect("initial scan must finish"); + blocked.store(true, Ordering::Relaxed); + sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .execute(&control) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(3), reconnecting.notified()) + .await + .expect("reconnect must enter the blocked acquire hook"); + let result = tokio::time::timeout(Duration::from_millis(2500), async { + for notify in ¬ifications { + let calls = AtomicUsize::new(0); + let item = + fetch_with_poll(Some(notify), Duration::from_secs(10), "test", || async { + Ok((calls.fetch_add(1, Ordering::Relaxed) != 0).then_some("ready")) + }) + .await + .unwrap(); + assert_eq!(item, Some("ready")); + assert_eq!(calls.load(Ordering::Relaxed), 2); + } + }) + .await; + blocked.store(false, Ordering::Relaxed); + release.notify_waiters(); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(3), pool.close()) + .await + .unwrap(); + control.close().await; + result.expect("fallback wakes must continue independently of a pending reconnect"); + } + + #[tokio::test] + async fn zero_timeout_and_disabled_mode_fetch_once() { + let notify = Notify::new(); + for (notify, timeout) in [ + (Some(¬ify), Duration::ZERO), + (None, Duration::from_secs(30)), + ] { + let result = tokio::time::timeout( + Duration::from_millis(100), + fetch_with_poll(notify, timeout, "test", || async { Ok(None::<()>) }), + ) + .await + .expect("short polling must not wait") + .unwrap(); + assert!(result.is_none()); + } + } + + #[tokio::test] + async fn a_claim_completing_after_the_wait_deadline_is_still_returned() { + let notify = Notify::new(); + let result = fetch_with_poll(Some(¬ify), Duration::from_millis(10), "test", || async { + sleep(Duration::from_millis(30)).await; + Ok(Some("claimed")) + }) + .await + .unwrap(); + assert_eq!(result, Some("claimed")); + } + + #[tokio::test] + async fn fetch_errors_and_invalid_timeout_are_explicit() { + let notify = Notify::new(); + let error = fetch_with_poll(Some(¬ify), Duration::from_secs(30), "test", || async { + Err::, _>(ProviderError::retryable("test", "unavailable")) + }) + .await + .unwrap_err(); + assert!(error.is_retryable()); + let error = fetch_with_poll(Some(¬ify), Duration::MAX, "test", || async { + Ok(None::<()>) + }) + .await + .unwrap_err(); + assert!(!error.is_retryable()); + } +} diff --git a/src/provider.rs b/src/provider.rs index 6aa67b6..7aebe5e 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -21,6 +21,7 @@ use tracing::{debug, error, instrument, warn}; use crate::entra::{EntraAuthOptions, TokenSource}; use crate::migrations::MigrationRunner; +use crate::notifier::{fetch_with_poll, LongPollConfig, LongPoller}; /// PostgreSQL-based provider for Duroxide durable orchestrations. /// @@ -166,6 +167,7 @@ pub struct PostgresProvider { /// preserving FR-006 byte-equivalent behavior on the password path. is_entra: bool, _refresh_task: Option, + long_poll: Option, } /// Migration policy applied at [`PostgresProvider`] construction. @@ -217,6 +219,10 @@ pub struct ProviderConfig { pub schema_name: Option, /// Policy for handling embedded migrations at startup. pub migration_policy: MigrationPolicy, + /// Notification-driven waiting for work. Disabled by default. + /// `ApplyAll` also activates schema-wide notification triggers; + /// `VerifyOnly` requires prior activation. + pub long_poll: LongPollConfig, } impl ProviderConfig { @@ -228,6 +234,7 @@ impl ProviderConfig { connection: ConnectionConfig::Url(database_url.into()), schema_name: None, migration_policy: MigrationPolicy::default(), + long_poll: LongPollConfig::default(), } } @@ -252,6 +259,7 @@ impl ProviderConfig { }, schema_name: None, migration_policy: MigrationPolicy::default(), + long_poll: LongPollConfig::default(), } } } @@ -337,7 +345,7 @@ impl PostgresProvider { /// Create a provider from a [`ProviderConfig`]. This is the single /// constructor that fully exposes the configuration surface: - /// connection variant (URL or Entra), schema, and migration policy. + /// connection variant (URL or Entra), schema, migration policy, and long-polling. /// /// All other public constructors delegate to this one. pub async fn new_with_config(config: ProviderConfig) -> Result { @@ -345,6 +353,7 @@ impl PostgresProvider { connection, schema_name, migration_policy, + long_poll, } = config; if let Some(ref s) = schema_name { @@ -353,7 +362,13 @@ impl PostgresProvider { match connection { ConnectionConfig::Url(database_url) => { - Self::new_from_url(&database_url, schema_name.as_deref(), migration_policy).await + Self::new_from_url( + &database_url, + schema_name.as_deref(), + migration_policy, + long_poll, + ) + .await } ConnectionConfig::Entra { host, @@ -362,6 +377,7 @@ impl PostgresProvider { user, options, } => { + long_poll.validate(options.max_connections_value())?; let token_source = options.default_token_source().context( "Entra credential resolution failed: could not build the default credential chain", )?; @@ -375,6 +391,7 @@ impl PostgresProvider { token_source, PgSslMode::VerifyFull, migration_policy, + long_poll, ) .await } @@ -385,33 +402,46 @@ impl PostgresProvider { database_url: &str, schema_name: Option<&str>, migration_policy: MigrationPolicy, + long_poll: LongPollConfig, ) -> Result { let max_connections = std::env::var("DUROXIDE_PG_POOL_MAX") .ok() .and_then(|s| s.parse::().ok()) .unwrap_or(10); + long_poll.validate(max_connections)?; let pool = PgPoolOptions::new() .max_connections(max_connections) - .min_connections(1) + .min_connections(if long_poll.enabled { 2 } else { 1 }) .acquire_timeout(std::time::Duration::from_secs(30)) .connect(database_url) .await?; let schema_name = schema_name.unwrap_or("public").to_string(); - let provider = Self { + let mut provider = Self { pool: Arc::new(pool), schema_name: schema_name.clone(), is_entra: false, _refresh_task: None, + long_poll: None, }; let migration_runner = MigrationRunner::new(provider.pool.clone(), schema_name); match migration_policy { - MigrationPolicy::ApplyAll => migration_runner.migrate().await?, - MigrationPolicy::VerifyOnly => migration_runner.verify().await?, + MigrationPolicy::ApplyAll => { + migration_runner + .migrate_with_notifications(long_poll.enabled) + .await? + } + MigrationPolicy::VerifyOnly => { + migration_runner + .verify_with_notifications(long_poll.enabled) + .await? + } } + provider.long_poll = + LongPoller::start((*provider.pool).clone(), &provider.schema_name, long_poll).await?; Ok(provider) } @@ -500,7 +530,9 @@ impl PostgresProvider { token_source: Arc, ssl_mode: PgSslMode, migration_policy: MigrationPolicy, + long_poll: LongPollConfig, ) -> Result { + long_poll.validate(options.max_connections_value())?; let audience = options.audience_str().to_string(); let token = token_source .fetch_token(&[audience.as_str()]) @@ -513,7 +545,7 @@ impl PostgresProvider { let pool = PgPoolOptions::new() .max_connections(options.max_connections_value()) - .min_connections(1) + .min_connections(if long_poll.enabled { 2 } else { 1 }) .acquire_timeout(options.acquire_timeout_value()) .connect_with(base_options.clone().password(&token.secret)) .await?; @@ -523,9 +555,18 @@ impl PostgresProvider { let migration_runner = MigrationRunner::new(pool.clone(), schema_name.clone()); match migration_policy { - MigrationPolicy::ApplyAll => migration_runner.migrate().await?, - MigrationPolicy::VerifyOnly => migration_runner.verify().await?, + MigrationPolicy::ApplyAll => { + migration_runner + .migrate_with_notifications(long_poll.enabled) + .await? + } + MigrationPolicy::VerifyOnly => { + migration_runner + .verify_with_notifications(long_poll.enabled) + .await? + } } + let long_poll = LongPoller::start((*pool).clone(), &schema_name, long_poll).await?; let refresh_handle = spawn_token_refresh_task( pool.clone(), @@ -541,6 +582,7 @@ impl PostgresProvider { schema_name, is_entra: true, _refresh_task: Some(AbortOnDropHandle(refresh_handle)), + long_poll, }) } @@ -556,7 +598,7 @@ impl PostgresProvider { } /// Get current timestamp in milliseconds (Unix epoch) - fn now_millis() -> i64 { + pub(crate) fn now_millis() -> i64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() @@ -646,6 +688,9 @@ impl PostgresProvider { /// **SAFETY**: Never drops the "public" schema itself, only tables within it. /// Only drops the schema if it's a custom schema (not "public"). pub async fn cleanup_schema(&self) -> Result<()> { + if let Some(long_poll) = &self.long_poll { + long_poll.stop(); + } const MAX_RETRIES: u32 = 5; const BASE_RETRY_DELAY_MS: u64 = 50; @@ -952,21 +997,10 @@ fn compute_next_refresh_sleep( ceiling.min(expiry_driven).max(ENTRA_REFRESH_MIN_INTERVAL) } -#[async_trait::async_trait] -impl Provider for PostgresProvider { - fn name(&self) -> &str { - "duroxide-pg" - } - - fn version(&self) -> &str { - env!("CARGO_PKG_VERSION") - } - - #[instrument(skip(self), target = "duroxide::providers::postgres")] - async fn fetch_orchestration_item( +impl PostgresProvider { + async fn fetch_orchestration_item_once( &self, lock_timeout: Duration, - _poll_timeout: Duration, filter: Option<&DispatcherCapabilityFilter>, ) -> Result, ProviderError> { let start = std::time::Instant::now(); @@ -999,6 +1033,7 @@ impl Provider for PostgresProvider { for attempt in 0..=MAX_RETRIES { let now_ms = Self::now_millis(); + #[allow(clippy::type_complexity)] let result: Result< Option<( String, @@ -1158,6 +1193,10 @@ impl Provider for PostgresProvider { vec![], ) .await?; + // Consuming an orphan does not mean the queue is empty. + if let Some(poller) = &self.long_poll { + poller.orch_notify.notify_waiters(); + } return Ok(None); } @@ -1184,6 +1223,126 @@ impl Provider for PostgresProvider { Ok(None) } + + async fn fetch_work_item_once( + &self, + lock_timeout: Duration, + session: Option<&SessionFetchConfig>, + tag_filter: &TagFilter, + ) -> Result, ProviderError> { + let start = std::time::Instant::now(); + + // Convert Duration to milliseconds + let lock_timeout_ms = lock_timeout.as_millis() as i64; + + // Extract session parameters + let (owner_id, session_lock_timeout_ms): (Option<&str>, Option) = match session { + Some(config) => ( + Some(&config.owner_id), + Some(config.lock_timeout.as_millis() as i64), + ), + None => (None, None), + }; + + // Convert TagFilter to SQL parameters + let (tag_mode, tag_names) = Self::tag_filter_to_sql(tag_filter); + + let row = match sqlx::query_as::<_, (String, String, i32)>(&format!( + "SELECT * FROM {}.fetch_work_item($1, $2, $3, $4, $5, $6)", + self.schema_name + )) + .bind(Self::now_millis()) + .bind(lock_timeout_ms) + .bind(owner_id) + .bind(session_lock_timeout_ms) + .bind(&tag_names) + .bind(tag_mode) + .fetch_optional(&*self.pool) + .await + { + Ok(row) => row, + Err(e) => { + return Err(self.sqlx_to_provider_error("fetch_work_item", e)); + } + }; + + let (work_item_json, lock_token, attempt_count) = match row { + Some(row) => row, + None => return Ok(None), + }; + + let work_item: WorkItem = serde_json::from_str(&work_item_json).map_err(|e| { + ProviderError::permanent( + "fetch_work_item", + format!("Failed to deserialize worker item: {e}"), + ) + })?; + + let duration_ms = start.elapsed().as_millis() as u64; + + // Extract instance for logging - different work item types have different structures + let instance_id = match &work_item { + WorkItem::ActivityExecute { instance, .. } => instance.as_str(), + WorkItem::ActivityCompleted { instance, .. } => instance.as_str(), + WorkItem::ActivityFailed { instance, .. } => instance.as_str(), + WorkItem::StartOrchestration { instance, .. } => instance.as_str(), + WorkItem::TimerFired { instance, .. } => instance.as_str(), + WorkItem::ExternalRaised { instance, .. } => instance.as_str(), + WorkItem::CancelInstance { instance, .. } => instance.as_str(), + WorkItem::ContinueAsNew { instance, .. } => instance.as_str(), + WorkItem::SubOrchCompleted { + parent_instance, .. + } => parent_instance.as_str(), + WorkItem::SubOrchFailed { + parent_instance, .. + } => parent_instance.as_str(), + WorkItem::QueueMessage { instance, .. } => instance.as_str(), + }; + + debug!( + target = "duroxide::providers::postgres", + operation = "fetch_work_item", + instance_id = %instance_id, + attempt_count = attempt_count, + duration_ms = duration_ms, + "Fetched activity work item via stored procedure" + ); + + Ok(Some((work_item, lock_token, attempt_count as u32))) + } +} + +#[async_trait::async_trait] +impl Provider for PostgresProvider { + fn name(&self) -> &str { + "duroxide-pg" + } + + fn version(&self) -> &str { + env!("CARGO_PKG_VERSION") + } + + #[instrument(skip(self), target = "duroxide::providers::postgres")] + async fn fetch_orchestration_item( + &self, + lock_timeout: Duration, + poll_timeout: Duration, + filter: Option<&DispatcherCapabilityFilter>, + ) -> Result, ProviderError> { + if filter.is_some_and(|filter| filter.supported_duroxide_versions.is_empty()) { + return Ok(None); + } + fetch_with_poll( + self.long_poll + .as_ref() + .map(|poller| poller.orch_notify.as_ref()), + poll_timeout, + "fetch_orchestration_item", + || self.fetch_orchestration_item_once(lock_timeout, filter), + ) + .await + } + #[instrument(skip(self), fields(lock_token = %lock_token, execution_id = execution_id), target = "duroxide::providers::postgres")] async fn ack_orchestration_item( &self, @@ -1601,7 +1760,7 @@ impl Provider for PostgresProvider { async fn fetch_work_item( &self, lock_timeout: Duration, - _poll_timeout: Duration, + poll_timeout: Duration, session: Option<&SessionFetchConfig>, tag_filter: &TagFilter, ) -> Result, ProviderError> { @@ -1610,85 +1769,15 @@ impl Provider for PostgresProvider { return Ok(None); } - let start = std::time::Instant::now(); - - // Convert Duration to milliseconds - let lock_timeout_ms = lock_timeout.as_millis() as i64; - - // Extract session parameters - let (owner_id, session_lock_timeout_ms): (Option<&str>, Option) = match session { - Some(config) => ( - Some(&config.owner_id), - Some(config.lock_timeout.as_millis() as i64), - ), - None => (None, None), - }; - - // Convert TagFilter to SQL parameters - let (tag_mode, tag_names) = Self::tag_filter_to_sql(tag_filter); - - let row = match sqlx::query_as::<_, (String, String, i32)>(&format!( - "SELECT * FROM {}.fetch_work_item($1, $2, $3, $4, $5, $6)", - self.schema_name - )) - .bind(Self::now_millis()) - .bind(lock_timeout_ms) - .bind(owner_id) - .bind(session_lock_timeout_ms) - .bind(&tag_names) - .bind(tag_mode) - .fetch_optional(&*self.pool) + fetch_with_poll( + self.long_poll + .as_ref() + .map(|poller| poller.worker_notify.as_ref()), + poll_timeout, + "fetch_work_item", + || self.fetch_work_item_once(lock_timeout, session, tag_filter), + ) .await - { - Ok(row) => row, - Err(e) => { - return Err(self.sqlx_to_provider_error("fetch_work_item", e)); - } - }; - - let (work_item_json, lock_token, attempt_count) = match row { - Some(row) => row, - None => return Ok(None), - }; - - let work_item: WorkItem = serde_json::from_str(&work_item_json).map_err(|e| { - ProviderError::permanent( - "fetch_work_item", - format!("Failed to deserialize worker item: {e}"), - ) - })?; - - let duration_ms = start.elapsed().as_millis() as u64; - - // Extract instance for logging - different work item types have different structures - let instance_id = match &work_item { - WorkItem::ActivityExecute { instance, .. } => instance.as_str(), - WorkItem::ActivityCompleted { instance, .. } => instance.as_str(), - WorkItem::ActivityFailed { instance, .. } => instance.as_str(), - WorkItem::StartOrchestration { instance, .. } => instance.as_str(), - WorkItem::TimerFired { instance, .. } => instance.as_str(), - WorkItem::ExternalRaised { instance, .. } => instance.as_str(), - WorkItem::CancelInstance { instance, .. } => instance.as_str(), - WorkItem::ContinueAsNew { instance, .. } => instance.as_str(), - WorkItem::SubOrchCompleted { - parent_instance, .. - } => parent_instance.as_str(), - WorkItem::SubOrchFailed { - parent_instance, .. - } => parent_instance.as_str(), - WorkItem::QueueMessage { instance, .. } => instance.as_str(), - }; - - debug!( - target = "duroxide::providers::postgres", - operation = "fetch_work_item", - instance_id = %instance_id, - attempt_count = attempt_count, - duration_ms = duration_ms, - "Fetched activity work item via stored procedure" - ); - - Ok(Some((work_item, lock_token, attempt_count as u32))) } #[instrument(skip(self), fields(token = %token), target = "duroxide::providers::postgres")] @@ -3428,6 +3517,7 @@ mod entra_pipeline_tests { token_source, PgSslMode::Disable, MigrationPolicy::ApplyAll, + LongPollConfig::default(), ) .await .expect("Entra pipeline must succeed against local PG with correct token"); @@ -3450,6 +3540,127 @@ mod entra_pipeline_tests { drop_schema(provider.pool(), &schema).await; } + #[tokio::test] + async fn long_poll_rejects_single_connection_before_fetching_a_token() { + let source = RecordingFakeTokenSource::with_tokens(vec![token("unused", 3600)]); + let result = PostgresProvider::new_with_entra_with_token_source( + "localhost", + 1, + "unused", + "unused", + Some("unused"), + EntraAuthOptions::new().max_connections(1), + source.clone(), + PgSslMode::Disable, + MigrationPolicy::ApplyAll, + LongPollConfig { + enabled: true, + ..Default::default() + }, + ) + .await; + let error = match result { + Ok(_) => panic!("a long-poll listener cannot share a one-connection pool"), + Err(error) => error, + }; + assert!(error.to_string().contains("two pool connections")); + assert_eq!(source.call_count(), 0); + } + + #[tokio::test] + async fn pipeline_long_poll_with_two_connections() { + let Some((host, port, db, user, password)) = pg_connection_or_skip() else { + return; + }; + let schema = unique_schema(); + let source = RecordingFakeTokenSource::with_tokens(vec![token(&password, 3600)]); + let provider = PostgresProvider::new_with_entra_with_token_source( + &host, + port, + &db, + &user, + Some(&schema), + EntraAuthOptions::new().max_connections(2), + source, + PgSslMode::Disable, + MigrationPolicy::ApplyAll, + LongPollConfig { + enabled: true, + notifier_poll_interval: Duration::from_millis(50), + ..Default::default() + }, + ) + .await + .expect("Entra long-poll construction should work with two connections"); + let control = PgPoolOptions::new() + .max_connections(2) + .connect_with( + build_entra_connect_options(&host, port, &db, &user, PgSslMode::Disable) + .password(&password), + ) + .await + .unwrap(); + let mut blocker = control.begin().await.unwrap(); + sqlx::query(&format!( + "LOCK TABLE {schema}.orchestrator_queue IN ACCESS EXCLUSIVE MODE" + )) + .execute(&mut *blocker) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(3), async { + loop { + let blocked: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity \ + WHERE wait_event_type = 'Lock' AND strpos(query, $1) > 0)", + ) + .bind(format!("FROM {schema}.orchestrator_queue")) + .fetch_one(&control) + .await + .unwrap(); + if blocked { + break; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("timer scan must reach the blocked queue"); + let fetch = provider.fetch_work_item( + Duration::from_secs(30), + Duration::from_secs(5), + None, + &TagFilter::Any, + ); + let enqueue = async { + sleep(Duration::from_millis(100)).await; + provider + .enqueue_for_worker(WorkItem::ActivityExecute { + instance: schema.clone(), + execution_id: 1, + id: 1, + name: "test".to_string(), + input: String::new(), + session_id: None, + tag: None, + }) + .await + .unwrap(); + }; + let result = tokio::time::timeout(Duration::from_secs(3), async { + tokio::join!(fetch, enqueue) + }) + .await; + blocker.rollback().await.unwrap(); + provider.cleanup_schema().await.unwrap(); + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + control.close().await; + let (result, ()) = + result.expect("a blocked scan must not starve the two-connection Entra provider"); + assert!(result.unwrap().is_some()); + } + #[tokio::test] async fn pipeline_with_wrong_token_fails_before_migrations() { let Some((host, port, db, user, _password)) = pg_connection_or_skip() else { @@ -3477,6 +3688,7 @@ mod entra_pipeline_tests { token_source, PgSslMode::Disable, MigrationPolicy::ApplyAll, + LongPollConfig::default(), ) .await; @@ -3526,6 +3738,7 @@ mod entra_pipeline_tests { token_source, PgSslMode::Disable, MigrationPolicy::ApplyAll, + LongPollConfig::default(), ) .await .expect("default-constructor variant must succeed"); diff --git a/tests/deadlocks/harness.rs b/tests/deadlocks/harness.rs index e196fd6..bd3d8a6 100644 --- a/tests/deadlocks/harness.rs +++ b/tests/deadlocks/harness.rs @@ -60,6 +60,43 @@ pub async fn raw_pool(database_url: &str) -> PgPool { .expect("failed to create raw PostgreSQL pool") } +pub async fn migrate_through(pool: &PgPool, schema: &str, last_version: i64) { + sqlx::query(&format!("CREATE SCHEMA {schema}")) + .execute(pool) + .await + .unwrap(); + let directory = include_dir::include_dir!("$CARGO_MANIFEST_DIR/migrations"); + let mut files: Vec<_> = directory + .files() + .filter(|file| file.path().extension().is_some_and(|ext| ext == "sql")) + .collect(); + files.sort_by_key(|file| file.path()); + for file in files { + let name = file.path().file_name().unwrap().to_str().unwrap(); + let version: i64 = name.split('_').next().unwrap().parse().unwrap(); + if version > last_version { + break; + } + let mut tx = pool.begin().await.unwrap(); + sqlx::query(&format!("SET LOCAL search_path TO {schema}, pg_temp")) + .execute(&mut *tx) + .await + .unwrap(); + sqlx::Executor::execute(&mut *tx, file.contents_utf8().unwrap()) + .await + .unwrap(); + sqlx::query(&format!( + "INSERT INTO {schema}._duroxide_migrations(version, name) VALUES ($1, $2)" + )) + .bind(version) + .bind(name) + .execute(&mut *tx) + .await + .unwrap(); + tx.commit().await.unwrap(); + } +} + pub async fn timed(what: &str, future: impl Future) -> T { tokio::time::timeout(STEP_TIMEOUT, future) .await diff --git a/tests/deadlocks/main.rs b/tests/deadlocks/main.rs index fbaaee6..d3ea21d 100644 --- a/tests/deadlocks/main.rs +++ b/tests/deadlocks/main.rs @@ -31,12 +31,14 @@ //! - `renew_ack_abandon`: renewal against acknowledge and abandon. //! - `victim_retry`: the provider retries a renewal aborted as a deadlock //! victim. +//! - `notification_ddl`: notification upgrades and activation under queue contention. //! - `harness`: shared helpers, also used by //! `tests/orchestration_lock_renewal_tests.rs` and //! `tests/orchestration_ack_lease_tests.rs`. mod ack_fetch; mod harness; +mod notification_ddl; mod renew_ack_abandon; mod renew_fetch; mod victim_retry; @@ -46,6 +48,7 @@ const SOURCES: &[(&str, &str)] = &[ ("main.rs", include_str!("main.rs")), ("ack_fetch.rs", include_str!("ack_fetch.rs")), ("harness.rs", include_str!("harness.rs")), + ("notification_ddl.rs", include_str!("notification_ddl.rs")), ("renew_ack_abandon.rs", include_str!("renew_ack_abandon.rs")), ("renew_fetch.rs", include_str!("renew_fetch.rs")), ("victim_retry.rs", include_str!("victim_retry.rs")), diff --git a/tests/deadlocks/notification_ddl.rs b/tests/deadlocks/notification_ddl.rs new file mode 100644 index 0000000..3731f86 --- /dev/null +++ b/tests/deadlocks/notification_ddl.rs @@ -0,0 +1,323 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::harness::*; +use duroxide_pg::migrations::MigrationRunner; +use duroxide_pg::{PostgresProvider, ProviderConfig}; +use serde_json::Value; +use sqlx::{postgres::PgPoolOptions, PgPool}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +async fn snapshot(pool: &PgPool, schema: &str) -> Value { + sqlx::query_scalar(&format!( + "SELECT jsonb_build_object( + 'ledger', (SELECT jsonb_agg(to_jsonb(m) ORDER BY version) FROM {schema}._duroxide_migrations m), + 'functions', (SELECT jsonb_agg(pg_get_functiondef(p.oid) ORDER BY p.oid) + FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace WHERE n.nspname = $1), + 'triggers', (SELECT jsonb_agg(pg_get_triggerdef(t.oid) ORDER BY t.oid) + FROM pg_trigger t JOIN pg_class c ON c.oid = t.tgrelid + JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1), + 'worker', (SELECT jsonb_agg(to_jsonb(q) ORDER BY id) FROM {schema}.worker_queue q), + 'orchestration', (SELECT jsonb_agg(to_jsonb(q) ORDER BY id) FROM {schema}.orchestrator_queue q))" + )) + .bind(schema) + .fetch_one(pool) + .await + .unwrap() +} + +#[tokio::test] +async fn notification_upgrade_retries_reverse_order_delete_contention() { + with_schema(|schema| async move { + let pool = raw_pool(&database_url()).await; + migrate_through(&pool, &schema, 29).await; + sqlx::Executor::execute( + &pool, + format!( + "INSERT INTO {schema}.instances(instance_id, orchestration_name, created_at, updated_at) + VALUES ('delete-me', 'Test', now(), now()); + INSERT INTO {schema}.orchestrator_queue(instance_id, work_item, visible_at, created_at) + VALUES ('delete-me', '{{}}', now(), now()); + INSERT INTO {schema}.worker_queue(instance_id, work_item, visible_at, created_at) + VALUES ('delete-me', '{{}}', now(), now());" + ) + .as_str(), + ) + .await + .unwrap(); + let before = snapshot(&pool, &schema).await; + let migrator = PgPoolOptions::new() + .max_connections(1) + .after_connect(|conn, _| { + Box::pin(async move { + sqlx::query("SET deadlock_timeout = '500ms'") + .execute(conn) + .await?; + Ok(()) + }) + }) + .connect(&database_url()) + .await + .unwrap(); + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&migrator) + .await + .unwrap(); + let mut delete = pool.begin().await.unwrap(); + sqlx::Executor::execute( + &mut *delete, + format!( + "SET LOCAL deadlock_timeout = '10s'; + SET LOCAL statement_timeout = '5s'; + LOCK TABLE {schema}.orchestrator_queue IN ROW EXCLUSIVE MODE" + ) + .as_str(), + ) + .await + .unwrap(); + let runner = MigrationRunner::new(Arc::new(migrator.clone()), schema.clone()); + let upgrade = tokio::spawn(async move { runner.migrate().await }); + wait_for_lock_waiter(&pool, pid).await; + let deleted: (i64, i64, i64, i64) = timed( + "reverse-order admin delete", + sqlx::query_as(&format!( + "SELECT * FROM {schema}.delete_instances_atomic(ARRAY['delete-me'], true)" + )) + .fetch_one(&mut *delete), + ) + .await + .expect("admin delete must finish when the DDL transaction rolls back"); + delete.commit().await.unwrap(); + assert_eq!(deleted, (1, 0, 0, 2)); + let result = timed("notification upgrade retry", upgrade).await.unwrap(); + let after = snapshot(&pool, &schema).await; + migrator.close().await; + pool.close().await; + result.expect("notification upgrade must retry its aborted transaction"); + let before_ledger = before["ledger"].as_array().unwrap(); + let after_ledger = after["ledger"].as_array().unwrap(); + assert_eq!(&after_ledger[..before_ledger.len()], before_ledger); + assert_eq!(after_ledger.len(), before_ledger.len() + 1); + assert_eq!(after_ledger.last().unwrap()["version"], 30); + assert!(after["worker"].is_null()); + assert!(after["orchestration"].is_null()); + }) + .await; +} + +#[tokio::test] +async fn notification_migration_retries_only_transient_errors_and_rolls_back() { + for code in ["40P01", "55P03", "42501"] { + with_schema(move |schema| async move { + let pool = raw_pool(&database_url()).await; + migrate_through(&pool, &schema, 29).await; + sqlx::Executor::execute(&pool, format!( + "CREATE SEQUENCE {schema}.migration_attempts; + CREATE FUNCTION {schema}.reject_first_migration() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN + IF NEW.version = 30 AND nextval('{schema}.migration_attempts') = 1 THEN + RAISE EXCEPTION 'injected migration failure' USING ERRCODE = '{code}'; + END IF; + RETURN NEW; + END $$; + CREATE TRIGGER reject_first_migration BEFORE INSERT ON {schema}._duroxide_migrations + FOR EACH ROW EXECUTE FUNCTION {schema}.reject_first_migration(); + INSERT INTO {schema}.worker_queue(work_item, visible_at, created_at) + VALUES ('retained work', now(), now());" + ).as_str()).await.unwrap(); + let before = snapshot(&pool, &schema).await; + let runner = MigrationRunner::new(Arc::new(pool.clone()), schema.clone()); + let result = runner.migrate().await; + let attempts: i64 = sqlx::query_scalar(&format!( + "SELECT last_value FROM {schema}.migration_attempts" + )).fetch_one(&pool).await.unwrap(); + let after = snapshot(&pool, &schema).await; + pool.close().await; + if code == "42501" { + let error = result.expect_err("permanent migration failures must be returned"); + let sql_error = error.downcast_ref::().expect("original SQL error must survive"); + assert_eq!(sqlstate(sql_error).as_deref(), Some(code)); + assert_eq!(attempts, 1, "permanent errors must not be retried"); + assert_eq!(after, before, "failed DDL must not change the ledger, functions, triggers, or work"); + } else { + result.expect("transient migration failures must retry"); + assert_eq!(attempts, 2, "the failed transaction must have been retried once"); + assert_eq!(after["worker"], before["worker"]); + assert_eq!(after["orchestration"], before["orchestration"]); + let before_ledger = before["ledger"].as_array().unwrap(); + let after_ledger = after["ledger"].as_array().unwrap(); + assert_eq!(&after_ledger[..before_ledger.len()], before_ledger); + assert_eq!(after_ledger.len(), before_ledger.len() + 1); + assert_eq!(after_ledger.last().unwrap()["version"], 30); + } + }).await; + } +} + +#[tokio::test] +async fn notification_activation_retries_sustained_contention_on_either_queue() { + for table in ["worker_queue", "orchestrator_queue"] { + with_schema(move |schema| async move { + let provider = PostgresProvider::new_with_schema(&database_url(), Some(&schema)) + .await + .unwrap(); + let pool = raw_pool(&database_url()).await; + let before = snapshot(&pool, &schema).await; + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(&format!( + "LOCK TABLE {schema}.{table} IN ROW EXCLUSIVE MODE" + )) + .execute(&mut *blocker) + .await + .unwrap(); + let application = format!("{schema}_activation"); + let mut url = url::Url::parse(&database_url()).unwrap(); + url.query_pairs_mut() + .append_pair("application_name", &application); + let mut config = ProviderConfig::url(url.as_str()); + config.schema_name = Some(schema.clone()); + config.long_poll.enabled = true; + let mut activation = tokio::spawn(PostgresProvider::new_with_config(config)); + let deadline = Instant::now() + WAIT_BUDGET; + let first_query = loop { + let started: Option> = sqlx::query_scalar( + "SELECT query_start FROM pg_stat_activity WHERE application_name = $1 \ + AND strpos(query, $2) > 0", + ) + .bind(&application) + .bind(format!("{table} IN SHARE ROW EXCLUSIVE MODE")) + .fetch_optional(&pool) + .await + .unwrap(); + if let Some(started) = started { + break started; + } + + assert!( + Instant::now() < deadline, + "activation never attempted the queue lock" + ); + tokio::time::sleep(POLL_INTERVAL).await; + }; + tokio::time::sleep(Duration::from_secs(1)).await; + let retry_started: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE application_name = $1 \ + AND query_start > $2)", + ) + .bind(&application) + .bind(first_query) + .fetch_one(&pool) + .await + .unwrap(); + let still_retrying = !activation.is_finished(); + blocker.rollback().await.unwrap(); + let result = tokio::time::timeout(Duration::from_secs(5), &mut activation).await; + let activated = result + .expect("activation should finish after the writer exits") + .unwrap() + .expect("activation needs a retry budget longer than 300ms"); + let active: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_trigger t JOIN pg_class c ON c.oid = t.tgrelid \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND tgname LIKE 'trg_notify_%' AND tgenabled = 'O'", + ) + .bind(&schema) + .fetch_one(&pool) + .await + .unwrap(); + let after = snapshot(&pool, &schema).await; + let listener_pool = activated.pool().clone(); + drop(activated); + listener_pool.close().await; + provider.pool().close().await; + pool.close().await; + assert!( + retry_started, + "{table}: the first lock wait must be bounded" + ); + assert!( + still_retrying, + "{table}: activation exhausted its budget too early" + ); + assert_eq!( + active, 6, + "all notification triggers must be enabled together" + ); + assert_eq!( + after, before, + "activation must preserve the ledger, function definitions, and work" + ); + }) + .await; + } +} + +#[tokio::test] +async fn notification_activation_exhausts_its_budget_without_partial_activation() { + with_schema(|schema| async move { + let provider = PostgresProvider::new_with_schema(&database_url(), Some(&schema)) + .await + .unwrap(); + let pool = raw_pool(&database_url()).await; + let before = snapshot(&pool, &schema).await; + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(&format!( + "LOCK TABLE {schema}.worker_queue IN ROW EXCLUSIVE MODE" + )) + .execute(&mut *blocker) + .await + .unwrap(); + let mut config = ProviderConfig::url(database_url()); + config.schema_name = Some(schema.clone()); + config.long_poll.enabled = true; + let start = Instant::now(); + let result = tokio::time::timeout( + Duration::from_secs(12), + PostgresProvider::new_with_config(config.clone()), + ) + .await; + let elapsed = start.elapsed(); + blocker.rollback().await.unwrap(); + let error = match result.expect("activation's lock retry budget must be bounded") { + Ok(_) => panic!("activation must not succeed while the writer holds its lock"), + Err(error) => error, + }; + assert_eq!( + sqlstate( + error + .downcast_ref::() + .expect("original lock error") + ) + .as_deref(), + Some("55P03"), + ); + assert!( + elapsed >= Duration::from_secs(9), + "retry budget ended after {elapsed:?}" + ); + assert_eq!(snapshot(&pool, &schema).await, before); + let active: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_trigger t JOIN pg_class c ON c.oid = t.tgrelid \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND tgname LIKE 'trg_notify_%' AND tgenabled <> 'D'", + ) + .bind(&schema) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!( + active, 0, + "failed activation must leave every trigger disabled" + ); + let activated = PostgresProvider::new_with_config(config) + .await + .expect("activation must succeed after contention is removed"); + let listener_pool = activated.pool().clone(); + drop(activated); + listener_pool.close().await; + provider.pool().close().await; + pool.close().await; + }) + .await; +} diff --git a/tests/long_poll_tests.rs b/tests/long_poll_tests.rs new file mode 100644 index 0000000..56577c6 --- /dev/null +++ b/tests/long_poll_tests.rs @@ -0,0 +1,2319 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use chrono::{DateTime, Utc}; +use duroxide::provider_validation::long_polling; +use duroxide::providers::{ + DispatcherCapabilityFilter, ExecutionMetadata, Provider, SessionFetchConfig, TagFilter, + WorkItem, +}; +use duroxide_pg::{LongPollConfig, MigrationPolicy, PostgresProvider, ProviderConfig}; +use sqlx::postgres::{types::Oid, PgPoolOptions}; +use sqlx::PgPool; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::task::JoinHandle; + +#[path = "deadlocks/harness.rs"] +mod harness; + +const LOCK: Duration = Duration::from_secs(30); +const POLL: Duration = Duration::from_secs(5); + +fn database_url() -> String { + dotenvy::dotenv().ok(); + std::env::var("DATABASE_URL").expect("DATABASE_URL must be set") +} + +fn schema_name() -> String { + format!("lp_{}", uuid::Uuid::new_v4().simple()) +} + +fn polling_config(enabled: bool) -> LongPollConfig { + let mut config = LongPollConfig::default(); + config.enabled = enabled; + config +} + +fn activity(instance: &str, session: Option<&str>, tag: Option<&str>) -> WorkItem { + WorkItem::ActivityExecute { + instance: instance.to_string(), + execution_id: 1, + id: 1, + name: "Test".to_string(), + input: String::new(), + session_id: session.map(str::to_string), + tag: tag.map(str::to_string), + } +} + +#[derive(Clone, Copy, Debug)] +enum Queue { + Orchestration, + Worker, +} + +#[derive(Debug)] +struct Claim { + token: String, + attempt: u32, +} + +impl Queue { + fn table(self) -> &'static str { + match self { + Self::Orchestration => "orchestrator_queue", + Self::Worker => "worker_queue", + } + } + + fn fetch_name(self) -> &'static str { + match self { + Self::Orchestration => "fetch_orchestration_item", + Self::Worker => "fetch_work_item", + } + } + + async fn enqueue(self, fixture: &Fixture, instance: &str, delay: Option) { + match self { + Self::Orchestration => fixture + .provider + .enqueue_for_orchestrator( + WorkItem::StartOrchestration { + instance: instance.to_string(), + orchestration: "Test".to_string(), + version: Some("1.0.0".to_string()), + input: String::new(), + parent_instance: None, + parent_id: None, + parent_execution_id: None, + execution_id: 1, + }, + delay, + ) + .await + .unwrap(), + Self::Worker => { + if let Some(delay) = delay { + // The worker enqueue API has no initial-delay parameter. + let now = Utc::now(); + sqlx::query(&format!( + "INSERT INTO {}.worker_queue \ + (work_item, visible_at, created_at, instance_id, execution_id, activity_id) \ + VALUES ($1, $2, $3, $4, 1, 1)", + fixture.schema + )) + .bind(serde_json::to_string(&activity(instance, None, None)).unwrap()) + .bind(now + chrono::Duration::from_std(delay).unwrap()) + .bind(now) + .bind(instance) + .execute(&fixture.control) + .await + .unwrap(); + } else { + fixture + .provider + .enqueue_for_worker(activity(instance, None, None)) + .await + .unwrap(); + } + } + } + } + + async fn fetch(self, provider: &PostgresProvider, timeout: Duration) -> Option { + self.fetch_with_lock(provider, LOCK, timeout).await + } + + async fn fetch_with_lock( + self, + provider: &PostgresProvider, + lock_timeout: Duration, + poll_timeout: Duration, + ) -> Option { + match self { + Self::Orchestration => provider + .fetch_orchestration_item(lock_timeout, poll_timeout, None) + .await + .unwrap() + .map(|(_, token, attempt)| Claim { token, attempt }), + Self::Worker => provider + .fetch_work_item(lock_timeout, poll_timeout, None, &TagFilter::Any) + .await + .unwrap() + .map(|(_, token, attempt)| Claim { token, attempt }), + } + } + + async fn abandon(self, provider: &PostgresProvider, token: &str, delay: Option) { + match self { + Self::Orchestration => provider + .abandon_orchestration_item(token, delay, false) + .await + .unwrap(), + Self::Worker => provider + .abandon_work_item(token, delay, false) + .await + .unwrap(), + } + } + + async fn ack(self, provider: &PostgresProvider, token: &str) { + match self { + Self::Orchestration => provider + .ack_orchestration_item( + token, + 1, + vec![], + vec![], + vec![], + ExecutionMetadata::default(), + vec![], + ) + .await + .unwrap(), + Self::Worker => provider.ack_work_item(token, None).await.unwrap(), + } + } +} + +struct Fixture { + provider: Arc, + control: PgPool, + schema: String, +} + +impl Fixture { + async fn new(config: LongPollConfig) -> Self { + Self::with_schema(config, schema_name()).await + } + + async fn enabled() -> Self { + let mut config = polling_config(true); + config.timer_grace_period = Duration::from_millis(20); + Self::new(config).await + } + + async fn with_schema(long_poll: LongPollConfig, schema: String) -> Self { + let enabled = long_poll.enabled; + let application = format!("lp_fixture_{}", uuid::Uuid::new_v4().simple()); + let mut url = url::Url::parse(&database_url()).unwrap(); + url.query_pairs_mut() + .append_pair("application_name", &application); + let mut config = ProviderConfig::url(url.as_str()); + config.schema_name = Some(schema.clone()); + config.long_poll = long_poll; + let provider = Arc::new(PostgresProvider::new_with_config(config).await.unwrap()); + let control = PgPoolOptions::new() + .max_connections(2) + .connect(&database_url()) + .await + .unwrap(); + let fixture = Self { + provider, + control, + schema, + }; + if enabled { + fixture + .wait_for_query("COMMIT", None, Some(&application)) + .await; + tokio::task::yield_now().await; + } + fixture + } + + async fn wait_for_query( + &self, + query: &str, + since: Option>, + application: Option<&str>, + ) { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let seen: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity \ + WHERE datname = current_database() AND state = 'idle' \ + AND strpos(query, $1) > 0 AND ($2::TIMESTAMPTZ IS NULL OR query_start >= $2) \ + AND ($3::TEXT IS NULL OR application_name = $3))", + ) + .bind(query) + .bind(since) + .bind(application) + .fetch_one(&self.control) + .await + .unwrap(); + if seen { + return; + } + assert!(Instant::now() < deadline, "query was not observed: {query}"); + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + async fn stamp(&self) -> DateTime { + sqlx::query_scalar("SELECT clock_timestamp()") + .fetch_one(&self.control) + .await + .unwrap() + } + + async fn wait_for_blocked_refresh(&self) { + let deadline = Instant::now() + Duration::from_secs(3); + loop { + let blocked: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity \ + WHERE wait_event_type = 'Lock' AND strpos(query, $1) > 0)", + ) + .bind(format!("FROM {}.orchestrator_queue", self.schema)) + .fetch_one(&self.control) + .await + .unwrap(); + if blocked { + return; + } + assert!( + Instant::now() < deadline, + "refresh did not block on the test lock" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + async fn waiting_fetch(&self, queue: Queue) -> JoinHandle> { + let since = self.stamp().await; + let provider = self.provider.clone(); + let fetch = tokio::spawn(async move { queue.fetch(&provider, POLL).await }); + self.wait_for_query( + &format!("{}.{}", self.schema, queue.fetch_name()), + Some(since), + None, + ) + .await; + assert!( + !fetch.is_finished(), + "fetch must be waiting before the producer acts" + ); + fetch + } + + async fn channel(&self, queue: Queue) -> String { + let oid: Oid = sqlx::query_scalar("SELECT $1::TEXT::REGCLASS::OID") + .bind(format!("{}.{}", self.schema, queue.table())) + .fetch_one(&self.control) + .await + .unwrap(); + format!("duroxide_pg_{}", oid.0) + } + + async fn listener_pid(&self) -> Option { + let channel = format!("\"{}\"", self.channel(Queue::Worker).await); + sqlx::query_scalar( + "SELECT pid FROM pg_stat_activity WHERE datname = current_database() \ + AND query LIKE 'LISTEN%' AND strpos(query, $1) > 0", + ) + .bind(channel) + .fetch_optional(&self.control) + .await + .unwrap() + } + + async fn cleanup(self) { + let pool = self.provider.pool().clone(); + drop(self.provider); + tokio::time::timeout(Duration::from_secs(5), pool.close()) + .await + .expect("listener and refresh must release pooled connections"); + sqlx::query(&format!("DROP SCHEMA IF EXISTS {} CASCADE", self.schema)) + .execute(&self.control) + .await + .unwrap(); + self.control.close().await; + } +} + +async fn claimed(fetch: JoinHandle>) -> Claim { + tokio::time::timeout(Duration::from_secs(3), fetch) + .await + .expect("notification/timer must beat the five-second polling fallback") + .expect("fetch task") + .expect("work must be found") +} + +#[tokio::test] +async fn default_and_zero_timeout_preserve_short_polling() { + for enabled in [false, true] { + let fixture = Fixture::new(polling_config(enabled)).await; + if !enabled { + long_polling::test_short_poll_returns_immediately( + fixture.provider.as_ref(), + Duration::from_millis(500), + ) + .await; + long_polling::test_short_poll_work_item_returns_immediately( + fixture.provider.as_ref(), + Duration::from_millis(500), + ) + .await; + assert!(fixture.listener_pid().await.is_none()); + } + for queue in [Queue::Orchestration, Queue::Worker] { + let start = Instant::now(); + assert!(queue + .fetch(&fixture.provider, Duration::ZERO) + .await + .is_none()); + assert!(start.elapsed() < Duration::from_secs(1)); + } + fixture.cleanup().await; + } +} + +#[tokio::test] +async fn enabled_provider_satisfies_upstream_long_poll_contract() { + let fixture = Fixture::enabled().await; + long_polling::test_long_poll_waits_for_timeout(fixture.provider.as_ref()).await; + long_polling::test_long_poll_work_item_waits_for_timeout(fixture.provider.as_ref()).await; + long_polling::test_fetch_respects_timeout_upper_bound(fixture.provider.as_ref()).await; + fixture.cleanup().await; +} + +#[tokio::test] +async fn existing_work_returns_without_waiting() { + let fixture = Fixture::enabled().await; + for queue in [Queue::Orchestration, Queue::Worker] { + queue.enqueue(&fixture, &fixture.schema, None).await; + let start = Instant::now(); + let claim = queue.fetch(&fixture.provider, POLL).await.unwrap(); + assert!(start.elapsed() < Duration::from_secs(1)); + queue.ack(&fixture.provider, &claim.token).await; + } + fixture.cleanup().await; +} + +#[tokio::test] +async fn long_poll_construction_leaves_a_warm_fetch_connection() { + let schema = schema_name(); + let mut config = ProviderConfig::url(database_url()); + config.schema_name = Some(schema.clone()); + config.long_poll.enabled = true; + let provider = Arc::new(PostgresProvider::new_with_config(config).await.unwrap()); + let connections = provider.pool().size(); + let control = PgPoolOptions::new() + .max_connections(2) + .connect(&database_url()) + .await + .unwrap(); + Fixture { + provider, + control, + schema, + } + .cleanup() + .await; + assert!( + connections >= 2, + "LISTEN must not consume the only warm connection: {connections}" + ); +} + +#[tokio::test] +async fn expired_claims_wake_fetches_before_the_poll_timeout() { + let mut failures = Vec::new(); + for queue in [Queue::Orchestration, Queue::Worker] { + for claim_before_listener in [false, true] { + let owner = Fixture::new(polling_config(false)).await; + let waiter = if claim_before_listener { + None + } else { + Some(Fixture::with_schema(polling_config(true), owner.schema.clone()).await) + }; + queue.enqueue(&owner, "crashed-owner", None).await; + let first = queue + .fetch_with_lock(&owner.provider, Duration::from_secs(1), Duration::ZERO) + .await + .unwrap(); + let waiter = match waiter { + Some(waiter) => waiter, + None => Fixture::with_schema(polling_config(true), owner.schema.clone()).await, + }; + let pool = owner.provider.pool().clone(); + drop(owner.provider); + pool.close().await; + owner.control.close().await; + let before: (String, i64, i32) = sqlx::query_as(&format!( + "SELECT lock_token, locked_until, attempt_count FROM {}.{}", + waiter.schema, + queue.table(), + )) + .fetch_one(&waiter.control) + .await + .unwrap(); + assert_eq!(before.0, first.token); + assert_eq!(before.2, 1); + let result = tokio::time::timeout(Duration::from_secs(3), async { + loop { + if let Some(claim) = + queue.fetch(&waiter.provider, Duration::from_secs(10)).await + { + break claim; + } + } + }) + .await; + let after: (String, i64, i32) = sqlx::query_as(&format!( + "SELECT lock_token, locked_until, attempt_count FROM {}.{}", + waiter.schema, + queue.table(), + )) + .fetch_one(&waiter.control) + .await + .unwrap(); + waiter.cleanup().await; + match result { + Ok(claim) => { + assert_ne!(claim.token, first.token); + assert_eq!(claim.attempt, 2); + assert_eq!(after.0, claim.token); + assert!(after.1 > before.1); + assert_eq!(after.2, 2); + } + Err(error) => { + assert_eq!(after, before, "a waiting fetch must not alter the claim"); + failures.push((queue, claim_before_listener, error)); + } + } + } + } + assert!( + failures.is_empty(), + "lease recovery was delayed: {failures:?}" + ); +} + +#[tokio::test] +async fn session_expiry_wakes_a_waiting_new_owner() { + let fixture = Fixture::enabled().await; + for instance in ["held-session-item", "waiting-session-item"] { + fixture + .provider + .enqueue_for_worker(activity(instance, Some("expiring-session"), None)) + .await + .unwrap(); + } + let first = fixture + .provider + .fetch_work_item( + LOCK, + Duration::ZERO, + Some(&SessionFetchConfig { + owner_id: "crashed-session-owner".to_string(), + lock_timeout: Duration::from_secs(1), + }), + &TagFilter::Any, + ) + .await + .unwrap() + .unwrap(); + let recovered = tokio::time::timeout( + Duration::from_secs(3), + fixture.provider.fetch_work_item( + LOCK, + Duration::from_secs(10), + Some(&SessionFetchConfig { + owner_id: "new-session-owner".to_string(), + lock_timeout: LOCK, + }), + &TagFilter::Any, + ), + ) + .await; + let owner: String = sqlx::query_scalar(&format!( + "SELECT worker_id FROM {}.sessions WHERE session_id = 'expiring-session'", + fixture.schema, + )) + .fetch_one(&fixture.control) + .await + .unwrap(); + fixture.cleanup().await; + let (item, token, attempt) = recovered + .expect("session expiry must beat the ten-second polling timeout") + .unwrap() + .unwrap(); + assert!(matches!( + item, + WorkItem::ActivityExecute { instance, .. } if instance == "waiting-session-item" + )); + assert_ne!(token, first.1); + assert_eq!(attempt, 1); + assert_eq!(owner, "new-session-owner"); +} + +#[tokio::test] +async fn delayed_abandon_immediately_releases_other_visible_orchestration_messages() { + let fixture = Fixture::enabled().await; + Queue::Orchestration + .enqueue(&fixture, "delayed-abandon", None) + .await; + let first = Queue::Orchestration + .fetch(&fixture.provider, Duration::ZERO) + .await + .unwrap(); + fixture + .provider + .enqueue_for_orchestrator( + WorkItem::ExternalRaised { + instance: "delayed-abandon".to_string(), + name: "ready".to_string(), + data: String::new(), + }, + None, + ) + .await + .unwrap(); + let waiting = fixture.waiting_fetch(Queue::Orchestration).await; + let start = Instant::now(); + Queue::Orchestration + .abandon( + &fixture.provider, + &first.token, + Some(Duration::from_secs(2)), + ) + .await; + let next = claimed(waiting).await; + let elapsed = start.elapsed(); + Queue::Orchestration + .ack(&fixture.provider, &next.token) + .await; + fixture.cleanup().await; + assert!(elapsed < Duration::from_secs(1), "release took {elapsed:?}"); + assert_eq!( + next.attempt, 1, + "only the already-visible message should be claimed" + ); +} + +#[tokio::test] +async fn orphan_cleanup_does_not_put_ready_orchestrations_to_sleep() { + let producer = Fixture::new(LongPollConfig::default()).await; + for instance in ["orphan-a", "orphan-b"] { + producer + .provider + .enqueue_for_orchestrator( + WorkItem::QueueMessage { + instance: instance.to_string(), + name: "signal".to_string(), + data: String::new(), + }, + None, + ) + .await + .unwrap(); + } + Queue::Orchestration.enqueue(&producer, "ready", None).await; + let consumer = Fixture::with_schema(polling_config(true), producer.schema.clone()).await; + let result = tokio::time::timeout( + Duration::from_secs(2), + consumer.provider.fetch_orchestration_item(LOCK, POLL, None), + ) + .await; + consumer.cleanup().await; + producer.cleanup().await; + let (item, _, _) = result + .expect("orphan cleanup must retry ready work rather than wait for another notification") + .unwrap() + .expect("ready orchestration"); + assert_eq!(item.instance, "ready"); +} + +#[tokio::test] +async fn inserts_wake_already_waiting_fetches_for_both_queues() { + let fixture = Fixture::enabled().await; + for queue in [Queue::Orchestration, Queue::Worker] { + let fetch = fixture.waiting_fetch(queue).await; + queue.enqueue(&fixture, &fixture.schema, None).await; + let claim = claimed(fetch).await; + assert_eq!(claim.attempt, 1); + queue.ack(&fixture.provider, &claim.token).await; + } + fixture.cleanup().await; +} + +#[tokio::test] +async fn acknowledgment_wakes_another_provider_after_it_consumed_a_locked_instances_notification() { + let owner = Fixture::new(LongPollConfig::default()).await; + Queue::Orchestration + .enqueue(&owner, "same-instance", None) + .await; + let first = Queue::Orchestration + .fetch(&owner.provider, Duration::ZERO) + .await + .unwrap(); + let waiter = Fixture::with_schema(polling_config(true), owner.schema.clone()).await; + let waiting = waiter.waiting_fetch(Queue::Orchestration).await; + let inserted_after = waiter.stamp().await; + owner + .provider + .enqueue_for_orchestrator( + WorkItem::ExternalRaised { + instance: "same-instance".to_string(), + name: "event".to_string(), + data: "arrived during the turn".to_string(), + }, + None, + ) + .await + .unwrap(); + waiter + .wait_for_query( + &format!("{}.fetch_orchestration_item", waiter.schema), + Some(inserted_after), + None, + ) + .await; + assert!(!waiting.is_finished(), "the instance is still locked"); + owner + .provider + .ack_orchestration_item( + &first.token, + 1, + vec![], + vec![], + vec![], + ExecutionMetadata { + orchestration_name: Some("Test".to_string()), + orchestration_version: Some("1.0.0".to_string()), + ..Default::default() + }, + vec![], + ) + .await + .unwrap(); + let next = claimed(waiting).await; + assert_eq!(next.attempt, 1); + Queue::Orchestration + .ack(&waiter.provider, &next.token) + .await; + waiter.cleanup().await; + owner.cleanup().await; +} + +#[tokio::test] +async fn orchestration_claims_announce_only_expiry_but_acknowledgments_notify() { + let fixture = Fixture::enabled().await; + let mut listener = sqlx::postgres::PgListener::connect_with(&fixture.control) + .await + .unwrap(); + listener + .listen(&fixture.channel(Queue::Orchestration).await) + .await + .unwrap(); + Queue::Orchestration + .enqueue(&fixture, "claim-and-ack", None) + .await; + tokio::time::timeout(Duration::from_secs(3), listener.recv()) + .await + .unwrap() + .unwrap(); + let claim = Queue::Orchestration + .fetch(&fixture.provider, Duration::ZERO) + .await + .unwrap(); + let notification = tokio::time::timeout(Duration::from_secs(3), listener.recv()) + .await + .unwrap() + .unwrap(); + let deadline: i64 = sqlx::query_scalar(&format!( + "SELECT locked_until FROM {}.instance_locks WHERE lock_token = $1", + fixture.schema, + )) + .bind(&claim.token) + .fetch_one(&fixture.control) + .await + .unwrap(); + assert_eq!(notification.payload().parse::().unwrap(), deadline); + assert!(deadline > Utc::now().timestamp_millis()); + assert!( + tokio::time::timeout(Duration::from_millis(100), listener.recv()) + .await + .is_err(), + "claiming work must not emit an availability notification", + ); + Queue::Orchestration + .ack(&fixture.provider, &claim.token) + .await; + let notification = tokio::time::timeout(Duration::from_secs(3), listener.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(notification.payload(), "0"); + drop(listener); + fixture.cleanup().await; +} + +#[tokio::test] +async fn renewed_claims_are_not_recovered_at_the_original_expiry() { + for queue in [Queue::Orchestration, Queue::Worker] { + let fixture = Fixture::enabled().await; + queue.enqueue(&fixture, "renewed", None).await; + let first = queue + .fetch_with_lock(&fixture.provider, Duration::from_secs(1), Duration::ZERO) + .await + .unwrap(); + match queue { + Queue::Orchestration => fixture + .provider + .renew_orchestration_item_lock(&first.token, Duration::from_secs(2)) + .await + .unwrap(), + Queue::Worker => fixture + .provider + .renew_work_item_lock(&first.token, Duration::from_secs(2)) + .await + .unwrap(), + } + let start = Instant::now(); + let second = claimed(fixture.waiting_fetch(queue).await).await; + assert!(start.elapsed() >= Duration::from_millis(1800)); + assert_ne!(first.token, second.token); + assert_eq!(second.attempt, 2); + queue.ack(&fixture.provider, &second.token).await; + fixture.cleanup().await; + } +} + +#[tokio::test] +async fn delayed_inserts_and_successive_timers_wake_both_queues() { + for queue in [Queue::Orchestration, Queue::Worker] { + let fixture = Fixture::enabled().await; + let fetch = fixture.waiting_fetch(queue).await; + let start = Instant::now(); + queue + .enqueue(&fixture, "first", Some(Duration::from_millis(300))) + .await; + queue + .enqueue(&fixture, "second", Some(Duration::from_millis(700))) + .await; + let first = claimed(fetch).await; + assert!(start.elapsed() >= Duration::from_millis(300)); + queue.ack(&fixture.provider, &first.token).await; + let next = fixture.waiting_fetch(queue).await; + let second = claimed(next).await; + assert!(start.elapsed() >= Duration::from_millis(700)); + queue.ack(&fixture.provider, &second.token).await; + fixture.cleanup().await; + } +} + +#[tokio::test] +async fn abandon_wakes_waiters_with_and_without_visibility_changes() { + for queue in [Queue::Orchestration, Queue::Worker] { + for delay in [None, Some(Duration::from_millis(300))] { + let fixture = Fixture::enabled().await; + queue.enqueue(&fixture, &fixture.schema, None).await; + let first = queue + .fetch(&fixture.provider, Duration::ZERO) + .await + .unwrap(); + let waiting = fixture.waiting_fetch(queue).await; + let start = Instant::now(); + queue.abandon(&fixture.provider, &first.token, delay).await; + let second = claimed(waiting).await; + assert_eq!(second.attempt, 2); + if let Some(delay) = delay { + assert!(start.elapsed() >= delay); + } + queue.ack(&fixture.provider, &second.token).await; + fixture.cleanup().await; + } + } +} + +#[tokio::test] +async fn refresh_recovers_ready_and_future_work_without_notifications() { + for queue in [Queue::Orchestration, Queue::Worker] { + for delay in [None, Some(Duration::from_millis(300))] { + let mut config = polling_config(true); + config.notifier_poll_interval = Duration::from_millis(100); + config.timer_grace_period = Duration::from_millis(20); + let fixture = Fixture::new(config).await; + sqlx::query(&format!( + "ALTER TABLE {}.{} DISABLE TRIGGER USER", + fixture.schema, + queue.table() + )) + .execute(&fixture.control) + .await + .unwrap(); + let waiting = fixture.waiting_fetch(queue).await; + let start = Instant::now(); + queue.enqueue(&fixture, &fixture.schema, delay).await; + let claim = claimed(waiting).await; + if let Some(delay) = delay { + assert!(start.elapsed() >= delay); + } + queue.ack(&fixture.provider, &claim.token).await; + fixture.cleanup().await; + } + } +} + +#[tokio::test] +async fn reconnect_recovers_missed_work_and_resubscribes_both_channels() { + for queue in [Queue::Orchestration, Queue::Worker] { + let fixture = Fixture::enabled().await; + let waiting = fixture.waiting_fetch(queue).await; + let old_pid = fixture.listener_pid().await.expect("listener must exist"); + let terminated: bool = sqlx::query_scalar("SELECT pg_terminate_backend($1)") + .bind(old_pid) + .fetch_one(&fixture.control) + .await + .unwrap(); + assert!(terminated); + queue.enqueue(&fixture, "during-disconnect", None).await; + let claim = claimed(waiting).await; + queue.ack(&fixture.provider, &claim.token).await; + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if fixture + .listener_pid() + .await + .is_some_and(|pid| pid != old_pid) + { + break; + } + assert!(Instant::now() < deadline, "listener did not reconnect"); + tokio::time::sleep(Duration::from_millis(20)).await; + } + let waiting = fixture.waiting_fetch(queue).await; + queue.enqueue(&fixture, "after-reconnect", None).await; + let claim = claimed(waiting).await; + queue.ack(&fixture.provider, &claim.token).await; + fixture.cleanup().await; + } +} + +#[tokio::test] +async fn cancelled_fetch_does_not_consume_later_notifications() { + let fixture = Fixture::enabled().await; + for queue in [Queue::Orchestration, Queue::Worker] { + let cancelled = fixture.waiting_fetch(queue).await; + cancelled.abort(); + assert!(cancelled.await.unwrap_err().is_cancelled()); + let waiting = fixture.waiting_fetch(queue).await; + queue.enqueue(&fixture, &fixture.schema, None).await; + let claim = claimed(waiting).await; + queue.ack(&fixture.provider, &claim.token).await; + } + fixture.cleanup().await; +} + +#[tokio::test] +async fn schema_names_do_not_constrain_notification_channel_length_or_case() { + for schema in [ + format!("lp_{}{}", uuid::Uuid::new_v4().simple(), "x".repeat(28)), + format!("LP_{}", uuid::Uuid::new_v4().simple()), + ] { + let fixture = Fixture::with_schema(polling_config(true), schema).await; + for queue in [Queue::Orchestration, Queue::Worker] { + assert!(fixture.channel(queue).await.len() < 63); + let waiting = fixture.waiting_fetch(queue).await; + queue.enqueue(&fixture, &fixture.schema, None).await; + let claim = claimed(waiting).await; + queue.ack(&fixture.provider, &claim.token).await; + } + fixture.cleanup().await; + } +} + +#[tokio::test] +async fn notification_timestamp_does_not_round_before_submillisecond_visibility() { + let fixture = Fixture::enabled().await; + let mut listener = sqlx::postgres::PgListener::connect_with(&fixture.control) + .await + .unwrap(); + listener + .listen(&fixture.channel(Queue::Worker).await) + .await + .unwrap(); + sqlx::query(&format!( + "INSERT INTO {}.worker_queue (work_item, visible_at, created_at) \ + VALUES ($1, '2100-01-01 00:00:00.000001+00', now())", + fixture.schema, + )) + .bind(serde_json::to_string(&activity("fractional", None, None)).unwrap()) + .execute(&fixture.control) + .await + .unwrap(); + let notification = tokio::time::timeout(Duration::from_secs(3), listener.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(notification.payload(), "4102444800001"); + drop(listener); + fixture.cleanup().await; +} + +#[tokio::test] +async fn disabled_producer_wakes_verify_only_consumer() { + let fixture = Fixture::enabled().await; + let mut producer_config = ProviderConfig::url(database_url()); + producer_config.schema_name = Some(fixture.schema.clone()); + producer_config.migration_policy = MigrationPolicy::VerifyOnly; + let producer = PostgresProvider::new_with_config(producer_config) + .await + .unwrap(); + let mut config = ProviderConfig::url(database_url()); + config.schema_name = Some(fixture.schema.clone()); + config.migration_policy = MigrationPolicy::VerifyOnly; + config.long_poll.enabled = true; + let consumer = Arc::new(PostgresProvider::new_with_config(config).await.unwrap()); + let since = fixture.stamp().await; + let waiting = { + let consumer = consumer.clone(); + tokio::spawn(async move { Queue::Worker.fetch(&consumer, POLL).await }) + }; + fixture + .wait_for_query( + &format!("{}.fetch_work_item", fixture.schema), + Some(since), + None, + ) + .await; + producer + .enqueue_for_worker(activity("mixed-config", None, None)) + .await + .unwrap(); + let claim = claimed(waiting).await; + Queue::Worker.ack(&consumer, &claim.token).await; + let pool = consumer.pool().clone(); + drop(consumer); + pool.close().await; + let pool = producer.pool().clone(); + drop(producer); + pool.close().await; + fixture.cleanup().await; +} + +#[tokio::test] +async fn verify_only_listener_initializes_without_ddl() { + let fixture = Fixture::enabled().await; + let mut url = url::Url::parse(&database_url()).unwrap(); + url.query_pairs_mut() + .append_pair("options", "-c default_transaction_read_only=on"); + let mut config = ProviderConfig::url(url.as_str()); + config.schema_name = Some(fixture.schema.clone()); + config.migration_policy = MigrationPolicy::VerifyOnly; + config.long_poll.enabled = true; + let provider = PostgresProvider::new_with_config(config).await.unwrap(); + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + fixture.cleanup().await; +} + +#[tokio::test] +async fn verify_only_requires_explicit_schema_opt_in() { + let fixture = Fixture::new(LongPollConfig::default()).await; + let mut config = ProviderConfig::url(database_url()); + config.schema_name = Some(fixture.schema.clone()); + config.migration_policy = MigrationPolicy::VerifyOnly; + config.long_poll.enabled = true; + let error = match PostgresProvider::new_with_config(config).await { + Ok(_) => panic!("VerifyOnly must not activate notifications"), + Err(error) => error, + }; + assert!(error.to_string().contains("ApplyAll")); + let active: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_trigger WHERE tgrelid IN \ + ($1::TEXT::REGCLASS, $2::TEXT::REGCLASS) AND NOT tgisinternal AND tgenabled <> 'D'", + ) + .bind(format!("{}.orchestrator_queue", fixture.schema)) + .bind(format!("{}.worker_queue", fixture.schema)) + .fetch_one(&fixture.control) + .await + .unwrap(); + assert_eq!(active, 0); + fixture.cleanup().await; +} + +#[tokio::test] +async fn pool_closure_wakes_waiters_and_releases_listener() { + let fixture = Fixture::enabled().await; + let since = fixture.stamp().await; + let provider = fixture.provider.clone(); + let waiting = tokio::spawn(async move { + provider + .fetch_work_item(LOCK, POLL, None, &TagFilter::Any) + .await + }); + fixture + .wait_for_query( + &format!("{}.fetch_work_item", fixture.schema), + Some(since), + None, + ) + .await; + tokio::time::timeout(Duration::from_secs(3), fixture.provider.pool().close()) + .await + .expect("LISTEN must not prevent pool closure"); + let error = tokio::time::timeout(Duration::from_secs(2), waiting) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!(error.is_retryable()); + fixture.cleanup().await; +} + +#[tokio::test] +async fn pool_closure_during_reconnect_does_not_leave_fetch_waiting() { + let fixture = Fixture::enabled().await; + let pid = fixture.listener_pid().await.unwrap(); + let terminated: bool = sqlx::query_scalar("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(&fixture.control) + .await + .unwrap(); + assert!(terminated); + tokio::time::sleep(Duration::from_millis(100)).await; + let since = fixture.stamp().await; + let provider = fixture.provider.clone(); + let waiting = tokio::spawn(async move { + provider + .fetch_work_item(LOCK, POLL, None, &TagFilter::Any) + .await + }); + fixture + .wait_for_query( + &format!("{}.fetch_work_item", fixture.schema), + Some(since), + None, + ) + .await; + tokio::time::timeout(Duration::from_secs(2), fixture.provider.pool().close()) + .await + .unwrap(); + let error = tokio::time::timeout(Duration::from_secs(2), waiting) + .await + .expect("pool closure must wake fetches even during reconnect backoff") + .unwrap() + .unwrap_err(); + assert!(error.is_retryable()); + fixture.cleanup().await; +} + +#[tokio::test] +async fn ineligible_filters_do_not_wait() { + let fixture = Fixture::enabled().await; + let empty = DispatcherCapabilityFilter { + supported_duroxide_versions: vec![], + }; + let started = Instant::now(); + assert!(fixture + .provider + .fetch_orchestration_item(LOCK, POLL, Some(&empty)) + .await + .unwrap() + .is_none()); + assert!(fixture + .provider + .fetch_work_item(LOCK, POLL, None, &TagFilter::None) + .await + .unwrap() + .is_none()); + assert!(started.elapsed() < Duration::from_millis(300)); + fixture.cleanup().await; +} + +#[tokio::test] +async fn notifications_wake_the_matching_tag_waiter() { + let fixture = Fixture::enabled().await; + let since = fixture.stamp().await; + let mut waits = Vec::new(); + for tag in ["red", "blue"] { + let provider = fixture.provider.clone(); + waits.push(tokio::spawn(async move { + provider + .fetch_work_item(LOCK, POLL, None, &TagFilter::tags([tag])) + .await + .unwrap() + })); + } + fixture + .wait_for_query( + &format!("{}.fetch_work_item", fixture.schema), + Some(since), + None, + ) + .await; + fixture + .provider + .enqueue_for_worker(activity("blue", None, Some("blue"))) + .await + .unwrap(); + let blue = tokio::time::timeout(Duration::from_secs(3), waits.pop().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(blue.0, activity("blue", None, Some("blue"))); + assert!(!waits[0].is_finished()); + fixture + .provider + .enqueue_for_worker(activity("red", None, Some("red"))) + .await + .unwrap(); + let red = tokio::time::timeout(Duration::from_secs(3), waits.pop().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(red.0, activity("red", None, Some("red"))); + fixture.cleanup().await; +} + +#[tokio::test] +async fn notifications_wake_the_session_owner_not_just_an_arbitrary_slot() { + let fixture = Fixture::enabled().await; + let owner = SessionFetchConfig { + owner_id: "owner".to_string(), + lock_timeout: LOCK, + }; + fixture + .provider + .enqueue_for_worker(activity("first", Some("session"), None)) + .await + .unwrap(); + let (_, token, _) = fixture + .provider + .fetch_work_item(LOCK, Duration::ZERO, Some(&owner), &TagFilter::Any) + .await + .unwrap() + .unwrap(); + fixture.provider.ack_work_item(&token, None).await.unwrap(); + let since = fixture.stamp().await; + let other = { + let provider = fixture.provider.clone(); + tokio::spawn(async move { + let session = SessionFetchConfig { + owner_id: "other".to_string(), + lock_timeout: LOCK, + }; + provider + .fetch_work_item(LOCK, POLL, Some(&session), &TagFilter::Any) + .await + .unwrap() + }) + }; + let eligible = { + let provider = fixture.provider.clone(); + tokio::spawn(async move { + provider + .fetch_work_item(LOCK, POLL, Some(&owner), &TagFilter::Any) + .await + .unwrap() + }) + }; + fixture + .wait_for_query( + &format!("{}.fetch_work_item", fixture.schema), + Some(since), + None, + ) + .await; + fixture + .provider + .enqueue_for_worker(activity("second", Some("session"), None)) + .await + .unwrap(); + let claim = tokio::time::timeout(Duration::from_secs(3), eligible) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(claim.0, activity("second", Some("session"), None)); + assert!(!other.is_finished()); + other.abort(); + assert!(other.await.unwrap_err().is_cancelled()); + fixture.cleanup().await; +} + +#[tokio::test] +async fn public_style_cleanup_removes_notification_function() { + let fixture = Fixture::new(LongPollConfig::default()).await; + let pruning_function = format!( + "{}.prune_executions_batch(TEXT, BIGINT, BIGINT, INTEGER)", + fixture.schema, + ); + let pruning_exists: bool = sqlx::query_scalar("SELECT to_regprocedure($1) IS NOT NULL") + .bind(&pruning_function) + .fetch_one(&fixture.control) + .await + .unwrap(); + if !pruning_exists { + sqlx::query(&format!( + "CREATE FUNCTION {pruning_function} RETURNS VOID LANGUAGE plpgsql AS $$ BEGIN RETURN; END $$" + )).execute(&fixture.control).await.unwrap(); + } + sqlx::query(&format!("SELECT {}.cleanup_schema()", fixture.schema)) + .execute(&fixture.control) + .await + .unwrap(); + let exists: bool = sqlx::query_scalar("SELECT to_regprocedure($1) IS NOT NULL") + .bind(format!("{}.duroxide_notify_queue_work()", fixture.schema)) + .fetch_one(&fixture.control) + .await + .unwrap(); + let pruning_exists: bool = sqlx::query_scalar("SELECT to_regprocedure($1) IS NOT NULL") + .bind(&pruning_function) + .fetch_one(&fixture.control) + .await + .unwrap(); + fixture.cleanup().await; + assert!(!exists); + assert!( + !pruning_exists, + "cleanup must retain migration 0028's batched-pruning DROP" + ); +} + +#[tokio::test] +async fn notification_migration_preserves_an_application_function_with_the_old_generic_name() { + let control = PgPoolOptions::new() + .max_connections(1) + .connect(&database_url()) + .await + .unwrap(); + let schema = schema_name(); + sqlx::query(&format!("CREATE SCHEMA {schema}")) + .execute(&control) + .await + .unwrap(); + sqlx::query(&format!( + "CREATE FUNCTION {schema}.notify_queue_work() RETURNS TRIGGER \ + LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'application-owned function'; END $$", + )) + .execute(&control) + .await + .unwrap(); + let fixture = Fixture::with_schema(polling_config(true), schema.clone()).await; + Queue::Worker + .enqueue(&fixture, "provider-owned-trigger", None) + .await; + let claimed = Queue::Worker + .fetch(&fixture.provider, Duration::ZERO) + .await + .unwrap(); + Queue::Worker.ack(&fixture.provider, &claimed.token).await; + sqlx::query(&format!("SELECT {schema}.cleanup_schema()")) + .execute(&control) + .await + .unwrap(); + let definition: Option = + sqlx::query_scalar("SELECT prosrc FROM pg_proc WHERE oid = to_regprocedure($1)") + .bind(format!("{schema}.notify_queue_work()")) + .fetch_optional(&control) + .await + .unwrap(); + assert!(definition.is_some_and(|definition| definition.contains("application-owned function"))); + fixture.cleanup().await; + control.close().await; +} + +#[tokio::test] +async fn adv_blocked_timer_scans_do_not_accumulate_server_backends() { + let mut config = polling_config(true); + config.notifier_poll_interval = Duration::from_millis(500); + let fixture = Fixture::new(config).await; + let schema = fixture.schema.clone(); + let admin = PgPoolOptions::new() + .max_connections(2) + .connect(&database_url()) + .await + .unwrap(); + let count_waiting = |admin: PgPool| { + let schema = schema.clone(); + async move { + // Other parallel tests deliberately hold unrelated database locks. + sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity \ + WHERE datname = current_database() AND wait_event_type = 'Lock' \ + AND strpos(query, $1) > 0", + ) + .bind(format!("FROM {schema}.")) + .fetch_one(&admin) + .await + .unwrap() + } + }; + let count_backends = |admin: PgPool| async move { + sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity \ + WHERE datname = current_database() AND backend_type = 'client backend'", + ) + .fetch_one(&admin) + .await + .unwrap() + }; + tokio::time::sleep(Duration::from_millis(1500)).await; + let backends_before = count_backends(admin.clone()).await; + let mut locker = admin.begin().await.unwrap(); + sqlx::query(&format!( + "LOCK TABLE {schema}.worker_queue IN ACCESS EXCLUSIVE MODE" + )) + .execute(&mut *locker) + .await + .unwrap(); + let mut waiting = Vec::new(); + for _ in 0..8 { + tokio::time::sleep(Duration::from_secs(1)).await; + waiting.push(count_waiting(admin.clone()).await); + } + let backends_locked = count_backends(admin.clone()).await; + locker.rollback().await.unwrap(); + tokio::time::sleep(Duration::from_secs(2)).await; + let waiting_after = count_waiting(admin.clone()).await; + let backends_after = count_backends(admin.clone()).await; + println!( + "ADV-RESULT schema={schema} lock=ACCESS_EXCLUSIVE(worker_queue) held_s=8 notifier_poll_ms=500 lock_waiting_backends_each_second={waiting:?} database_client_backends_before={backends_before} at_8s={backends_locked} 2s_after_release={backends_after} lock_waiting_after_release={waiting_after}" + ); + admin.close().await; + fixture.cleanup().await; + assert!( + waiting.iter().all(|n| *n <= 2), + "blocked timer scans must not accumulate waiting server backends: {waiting:?}" + ); +} + +#[tokio::test] +async fn a_blocked_timer_refresh_does_not_block_notifications() { + let mut config = polling_config(true); + config.notifier_poll_interval = Duration::from_millis(50); + let fixture = Fixture::new(config).await; + let mut blocker = fixture.control.begin().await.unwrap(); + sqlx::query(&format!( + "LOCK TABLE {}.orchestrator_queue IN ACCESS EXCLUSIVE MODE", + fixture.schema, + )) + .execute(&mut *blocker) + .await + .unwrap(); + fixture.wait_for_blocked_refresh().await; + let waiting = fixture.waiting_fetch(Queue::Worker).await; + Queue::Worker + .enqueue(&fixture, "refresh-blocked", None) + .await; + let claim = claimed(waiting).await; + Queue::Worker.ack(&fixture.provider, &claim.token).await; + blocker.rollback().await.unwrap(); + fixture.cleanup().await; +} + +#[tokio::test] +async fn cancelling_a_blocked_refresh_does_not_strand_a_pool_connection() { + for reconnect in [false, true] { + let mut config = polling_config(true); + config.notifier_poll_interval = Duration::from_millis(50); + let fixture = Fixture::new(config).await; + let mut blocker = fixture.control.begin().await.unwrap(); + sqlx::query(&format!( + "LOCK TABLE {}.orchestrator_queue IN ACCESS EXCLUSIVE MODE", + fixture.schema, + )) + .execute(&mut *blocker) + .await + .unwrap(); + fixture.wait_for_blocked_refresh().await; + if reconnect { + let pid = fixture.listener_pid().await.unwrap(); + sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .execute(&fixture.control) + .await + .unwrap(); + } + let Fixture { + provider, + control, + schema, + } = fixture; + let pool = provider.pool().clone(); + let mut provider = Some(provider); + if !reconnect { + drop(provider.take()); + } + // Let SQLx begin returning the cancelled query while the pool is still open. + tokio::time::sleep(Duration::from_millis(300)).await; + let closed = tokio::time::timeout(Duration::from_secs(2), pool.close()).await; + blocker.rollback().await.unwrap(); + drop(provider); + pool.close().await; + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&control) + .await + .unwrap(); + control.close().await; + assert!( + closed.is_ok(), + "pool closure waited for a cancelled refresh's database lock; reconnect={reconnect}" + ); + } +} + +#[tokio::test] +async fn timers_expiring_during_refresh_do_not_hide_the_next_timer() { + let fixture = Fixture::enabled().await; + let mut blocker = fixture.control.begin().await.unwrap(); + sqlx::query(&format!( + "LOCK TABLE {}.orchestrator_queue IN ACCESS EXCLUSIVE MODE", + fixture.schema, + )) + .execute(&mut *blocker) + .await + .unwrap(); + let first = fixture.waiting_fetch(Queue::Worker).await; + Queue::Worker + .enqueue(&fixture, "start-refresh", Some(Duration::from_millis(200))) + .await; + let first = claimed(first).await; + Queue::Worker.ack(&fixture.provider, &first.token).await; + fixture.wait_for_blocked_refresh().await; + + let second = fixture.waiting_fetch(Queue::Worker).await; + Queue::Worker + .enqueue(&fixture, "during-refresh", Some(Duration::from_millis(50))) + .await; + let third_enqueued = Instant::now(); + Queue::Worker + .enqueue(&fixture, "after-refresh", Some(Duration::from_millis(300))) + .await; + let second = claimed(second).await; + Queue::Worker.ack(&fixture.provider, &second.token).await; + blocker.rollback().await.unwrap(); + + let third = fixture.waiting_fetch(Queue::Worker).await; + let third = claimed(third).await; + let elapsed = third_enqueued.elapsed(); + Queue::Worker.ack(&fixture.provider, &third.token).await; + fixture.cleanup().await; + assert!( + elapsed >= Duration::from_millis(300), + "timer fired early: {elapsed:?}" + ); + assert!( + elapsed < Duration::from_millis(600), + "the timer, not the one-second polling fallback, must wake this fetch: {elapsed:?}" + ); +} + +#[tokio::test] +async fn multiple_providers_receive_notifications_without_duplicate_claims() { + let fixture = Fixture::enabled().await; + let mut config = ProviderConfig::url(database_url()); + config.schema_name = Some(fixture.schema.clone()); + config.migration_policy = MigrationPolicy::VerifyOnly; + config.long_poll.enabled = true; + let second = Arc::new(PostgresProvider::new_with_config(config).await.unwrap()); + for queue in [Queue::Orchestration, Queue::Worker] { + let first_waiter = fixture.waiting_fetch(queue).await; + let second_waiter = { + let provider = second.clone(); + tokio::spawn(async move { queue.fetch(&provider, POLL).await }) + }; + queue.enqueue(&fixture, "node-a", None).await; + queue.enqueue(&fixture, "node-b", None).await; + let a = claimed(first_waiter).await; + let b = claimed(second_waiter).await; + assert_ne!(a.token, b.token); + assert_eq!((a.attempt, b.attempt), (1, 1)); + queue.ack(&fixture.provider, &a.token).await; + queue.ack(&second, &b.token).await; + } + let pool = second.pool().clone(); + drop(second); + pool.close().await; + fixture.cleanup().await; +} + +#[tokio::test] +async fn concurrent_schema_activation_is_idempotent() { + let fixture = Fixture::new(LongPollConfig::default()).await; + let mut config = ProviderConfig::url(database_url()); + config.schema_name = Some(fixture.schema.clone()); + config.long_poll.enabled = true; + let (a, b) = tokio::join!( + PostgresProvider::new_with_config(config.clone()), + PostgresProvider::new_with_config(config), + ); + for provider in [a.unwrap(), b.unwrap()] { + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + } + let mut config = ProviderConfig::url(database_url()); + config.schema_name = Some(fixture.schema.clone()); + config.long_poll.enabled = true; + config.migration_policy = MigrationPolicy::VerifyOnly; + let provider = PostgresProvider::new_with_config(config).await.unwrap(); + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + fixture.cleanup().await; +} + +#[tokio::test] +async fn notification_installation_and_activation_follow_writer_lock_order() { + let mut failures = Vec::new(); + for install in [false, true] { + let fixture = Fixture::new(LongPollConfig::default()).await; + let mut blocker = fixture.control.begin().await.unwrap(); + sqlx::query("SET LOCAL lock_timeout = '500ms'") + .execute(&mut *blocker) + .await + .unwrap(); + sqlx::query(&format!( + "LOCK TABLE {}.worker_queue IN ROW EXCLUSIVE MODE", + fixture.schema, + )) + .execute(&mut *blocker) + .await + .unwrap(); + let application = format!("{}_activation", fixture.schema); + let mut url = url::Url::parse(&database_url()).unwrap(); + url.query_pairs_mut() + .append_pair("application_name", &application); + let schema = fixture.schema.clone(); + let activation = async move { + if install { + let pool = PgPoolOptions::new() + .max_connections(1) + .connect(url.as_str()) + .await + .unwrap(); + let mut tx = pool.begin().await.unwrap(); + sqlx::query(&format!("SET LOCAL search_path TO {schema}, pg_temp")) + .execute(&mut *tx) + .await + .unwrap(); + sqlx::raw_sql(include_str!( + "../migrations/0030_add_long_poll_notifications.sql" + )) + .execute(&mut *tx) + .await + .unwrap(); + tx.commit().await.unwrap(); + pool.close().await; + } else { + let mut config = ProviderConfig::url(url.as_str()); + config.schema_name = Some(schema); + config.long_poll.enabled = true; + let provider = PostgresProvider::new_with_config(config).await.unwrap(); + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + } + }; + let writer = async { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity a \ + JOIN pg_locks l ON l.pid = a.pid \ + WHERE a.application_name = $1 AND a.wait_event_type = 'Lock' \ + AND NOT l.granted AND l.relation = $2::TEXT::REGCLASS)", + ) + .bind(&application) + .bind(format!("{}.worker_queue", fixture.schema)) + .fetch_one(&fixture.control) + .await + .unwrap(); + if waiting { + break; + } + assert!( + Instant::now() < deadline, + "activation did not reach the worker lock" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + let writer_lock = sqlx::query(&format!( + "LOCK TABLE {}.orchestrator_queue IN ROW EXCLUSIVE MODE", + fixture.schema, + )) + .execute(&mut *blocker) + .await; + blocker.rollback().await.unwrap(); + writer_lock + }; + let ((), writer_lock) = tokio::time::timeout(Duration::from_secs(10), async { + tokio::join!(activation, writer) + }) + .await + .unwrap(); + fixture.cleanup().await; + if let Err(error) = writer_lock { + failures.push((install, error.to_string())); + } + } + assert!( + failures.is_empty(), + "installation/activation must not hold the orchestration queue while waiting for a writer's worker lock: {failures:?}", + ); +} + +#[tokio::test] +async fn activation_releases_worker_lock_before_retrying_reverse_order_contention() { + let fixture = Fixture::new(LongPollConfig::default()).await; + let mut blocker = fixture.control.begin().await.unwrap(); + sqlx::query("SET LOCAL lock_timeout = '500ms'") + .execute(&mut *blocker) + .await + .unwrap(); + sqlx::query(&format!( + "LOCK TABLE {}.orchestrator_queue IN ROW EXCLUSIVE MODE", + fixture.schema, + )) + .execute(&mut *blocker) + .await + .unwrap(); + let application = format!("{}_retry", fixture.schema); + let mut url = url::Url::parse(&database_url()).unwrap(); + url.query_pairs_mut() + .append_pair("application_name", &application); + let mut config = ProviderConfig::url(url.as_str()); + config.schema_name = Some(fixture.schema.clone()); + config.long_poll.enabled = true; + let activation = tokio::spawn(PostgresProvider::new_with_config(config)); + let deadline = Instant::now() + Duration::from_secs(3); + loop { + let tried_second_lock: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE application_name = $1 \ + AND strpos(query, 'orchestrator_queue IN SHARE ROW EXCLUSIVE MODE') > 0)", + ) + .bind(&application) + .fetch_one(&fixture.control) + .await + .unwrap(); + if tried_second_lock { + break; + } + assert!( + Instant::now() < deadline, + "activation did not try the contended lock" + ); + tokio::time::sleep(Duration::from_millis(1)).await; + } + let writer_lock = sqlx::query(&format!( + "LOCK TABLE {}.worker_queue IN ROW EXCLUSIVE MODE", + fixture.schema, + )) + .execute(&mut *blocker) + .await; + blocker.rollback().await.unwrap(); + let provider = tokio::time::timeout(Duration::from_secs(5), activation) + .await + .unwrap() + .unwrap() + .expect("activation should retry after the blocker exits"); + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + fixture.cleanup().await; + writer_lock.expect("activation must release the worker lock before retrying"); +} + +#[tokio::test] +async fn runtime_executes_activities_and_timer_then_shuts_down() { + use duroxide::runtime::{registry::ActivityRegistry, Runtime}; + use duroxide::{ActivityContext, Client, OrchestrationContext, OrchestrationRegistry}; + + let fixture = Fixture::enabled().await; + let activities = ActivityRegistry::builder() + .register("Echo", |_ctx: ActivityContext, input: String| async move { + Ok(input) + }) + .build(); + let orchestrations = OrchestrationRegistry::builder() + .register( + "Sequence", + |ctx: OrchestrationContext, input: String| async move { + let value = ctx.schedule_activity("Echo", input).await?; + ctx.schedule_timer(Duration::from_millis(200)).await; + ctx.schedule_activity_on_session("Echo", value, "session") + .await + }, + ) + .build(); + let runtime = + Runtime::start_with_store(fixture.provider.clone(), activities, orchestrations).await; + let client = Client::new(fixture.provider.clone()); + client + .start_orchestration("runtime-test", "Sequence", "result") + .await + .unwrap(); + let status = client + .wait_for_orchestration("runtime-test", Duration::from_secs(10)) + .await + .unwrap(); + assert!( + matches!(status, duroxide::OrchestrationStatus::Completed { output, .. } if output == "result") + ); + tokio::time::timeout(Duration::from_secs(3), runtime.shutdown(Some(1000))) + .await + .expect("shutdown must cancel waiting dispatchers"); + drop(client); + fixture.cleanup().await; +} + +#[tokio::test] +async fn upgrading_pre_notification_schema_preserves_work_and_defaults_to_disabled() { + let control = PgPoolOptions::new() + .max_connections(2) + .connect(&database_url()) + .await + .unwrap(); + let schema = schema_name(); + harness::migrate_through(&control, &schema, 29).await; + let payload = serde_json::to_string(&activity("retained", None, None)).unwrap(); + sqlx::query(&format!( + "INSERT INTO {schema}.worker_queue \ + (work_item, visible_at, created_at, instance_id, execution_id, activity_id) \ + VALUES ($1, now(), now(), 'retained', 1, 1)", + )) + .bind(&payload) + .execute(&control) + .await + .unwrap(); + let provider = PostgresProvider::new_with_schema(&database_url(), Some(&schema)) + .await + .unwrap(); + let active: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_trigger WHERE tgrelid = $1::TEXT::REGCLASS \ + AND NOT tgisinternal AND tgenabled <> 'D'", + ) + .bind(format!("{schema}.worker_queue")) + .fetch_one(&control) + .await + .unwrap(); + assert_eq!(active, 0); + let retained: String = + sqlx::query_scalar(&format!("SELECT work_item FROM {schema}.worker_queue")) + .fetch_one(&control) + .await + .unwrap(); + assert_eq!(retained, payload); + let fetched = provider + .fetch_work_item(LOCK, Duration::ZERO, None, &TagFilter::Any) + .await + .unwrap(); + assert!(fetched.is_some()); + provider.cleanup_schema().await.unwrap(); + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + control.close().await; +} + +#[tokio::test] +#[ignore = "requires PostgreSQL max_prepared_transactions >= 1"] +async fn two_phase_commit_is_preserved_until_explicit_notification_activation() { + let fixture = Fixture::new(LongPollConfig::default()).await; + let setting: String = sqlx::query_scalar("SHOW max_prepared_transactions") + .fetch_one(&fixture.control) + .await + .unwrap(); + assert!( + setting.parse::().unwrap() > 0, + "enable max_prepared_transactions for this test" + ); + let mut connection = fixture.control.acquire().await.unwrap(); + for enabled in [false, true] { + if enabled { + let mut config = ProviderConfig::url(database_url()); + config.schema_name = Some(fixture.schema.clone()); + config.long_poll.enabled = true; + let provider = PostgresProvider::new_with_config(config).await.unwrap(); + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + } + sqlx::query("BEGIN") + .execute(&mut *connection) + .await + .unwrap(); + sqlx::query(&format!( + "INSERT INTO {}.worker_queue(work_item, visible_at, created_at) VALUES ($1, now(), now())", + fixture.schema, + )) + .bind(serde_json::to_string(&activity("prepared", None, None)).unwrap()) + .execute(&mut *connection).await.unwrap(); + let prepared = sqlx::query(&format!("PREPARE TRANSACTION '{}'", fixture.schema)) + .execute(&mut *connection) + .await; + if enabled { + assert!(prepared.unwrap_err().to_string().contains("NOTIFY")); + sqlx::query("ROLLBACK") + .execute(&mut *connection) + .await + .unwrap(); + } else { + prepared.expect("short-polling-only schemas must still support two-phase commit"); + sqlx::query(&format!("ROLLBACK PREPARED '{}'", fixture.schema)) + .execute(&mut *connection) + .await + .unwrap(); + } + } + drop(connection); + fixture.cleanup().await; +} + +const COMMIT_WAIT_SAMPLE_INTERVAL: Duration = Duration::from_millis(10); + +#[derive(Default, serde::Serialize)] +struct CommitWaits { + samples: u64, + elapsed_seconds: f64, + notification_lock_backend_seconds: f64, + synchronous_replication_backend_seconds: f64, + max_notification_waiters: i64, + max_synchronous_replication_waiters: i64, +} + +async fn start_commit_wait_sampling( + fixture: &Fixture, +) -> (tokio::sync::oneshot::Sender<()>, JoinHandle) { + let application: String = sqlx::query_scalar("SELECT current_setting('application_name')") + .fetch_one(fixture.provider.pool()) + .await + .unwrap(); + let control = fixture.control.clone(); + let (stop, mut stopped) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(async move { + let start = Instant::now(); + let mut previous = start; + let mut totals = CommitWaits::default(); + let mut ticks = tokio::time::interval(COMMIT_WAIT_SAMPLE_INTERVAL); + ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = &mut stopped => break, + _ = ticks.tick() => { + // PreCommit_Notify serializes commits on the shared pg_database/0 + // heavyweight object lock, not the short-lived NotifyQueue LWLock. + let (notify, replication): (i64, i64) = sqlx::query_as( + "WITH notification_waiters AS MATERIALIZED ( + SELECT DISTINCT l.pid FROM pg_locks l + WHERE NOT l.granted AND l.locktype = 'object' + AND l.classid = 'pg_catalog.pg_database'::regclass + AND l.objid = 0 AND l.objsubid = 0 AND l.mode = 'AccessExclusiveLock') + SELECT count(*) FILTER (WHERE n.pid IS NOT NULL), + count(*) FILTER (WHERE a.wait_event = 'SyncRep') + FROM pg_stat_activity a + LEFT JOIN notification_waiters n ON n.pid = a.pid + WHERE a.datname = current_database() AND a.application_name = $1 + AND a.backend_type = 'client backend'", + ).bind(&application).fetch_one(&control).await.unwrap(); + let now = Instant::now(); + let elapsed = now.duration_since(previous).as_secs_f64(); + totals.samples += 1; + totals.notification_lock_backend_seconds += notify as f64 * elapsed; + totals.synchronous_replication_backend_seconds += replication as f64 * elapsed; + totals.max_notification_waiters = totals.max_notification_waiters.max(notify); + totals.max_synchronous_replication_waiters = + totals.max_synchronous_replication_waiters.max(replication); + previous = now; + } + } + } + totals.elapsed_seconds = start.elapsed().as_secs_f64(); + totals + }); + (stop, task) +} + +async fn benchmark_renewals( + fixture: &Fixture, + clients: usize, + duration: Duration, + enabled: bool, + repetition: usize, +) { + let lease = Duration::from_secs(60); + let mut claims = Vec::new(); + for lane in 0..clients { + let owner = format!("bench-renew-{lane}"); + Queue::Orchestration.enqueue(fixture, &owner, None).await; + let (item, orch_token, _) = fixture + .provider + .fetch_orchestration_item(lease, Duration::ZERO, None) + .await + .unwrap() + .unwrap(); + assert_eq!(item.instance, owner); + fixture + .provider + .enqueue_for_worker(activity(&owner, Some(&owner), None)) + .await + .unwrap(); + let (item, worker_token, _) = fixture + .provider + .fetch_work_item( + lease, + Duration::ZERO, + Some(&SessionFetchConfig { + owner_id: owner.clone(), + lock_timeout: lease, + }), + &TagFilter::Any, + ) + .await + .unwrap() + .unwrap(); + assert!(matches!(item, WorkItem::ActivityExecute { instance, .. } if instance == owner)); + claims.push((owner, orch_token, worker_token)); + } + let (stop_sampling, samples) = start_commit_wait_sampling(fixture).await; + let start = Instant::now(); + let deadline = start + duration; + let mut tasks = tokio::task::JoinSet::new(); + for (owner, orch, worker) in claims { + let provider = fixture.provider.clone(); + tasks.spawn(async move { + let mut renewals = 0_u64; + while Instant::now() < deadline { + provider + .renew_orchestration_item_lock(&orch, lease) + .await + .unwrap(); + provider.renew_work_item_lock(&worker, lease).await.unwrap(); + assert_eq!( + provider + .renew_session_lock(&[owner.as_str()], lease, lease) + .await + .unwrap(), + 1, + ); + renewals += 3; + } + renewals + }); + } + let mut renewals = 0; + while let Some(result) = tasks.join_next().await { + renewals += result.unwrap(); + } + let elapsed = start.elapsed().as_secs_f64(); + stop_sampling + .send(()) + .expect("commit-wait sampler must remain active"); + let waits = samples.await.unwrap(); + assert!(renewals > 0); + if enabled { + assert!( + waits.max_notification_waiters > 0, + "renewal contention must be observed" + ); + } else { + assert_eq!(waits.notification_lock_backend_seconds, 0.0); + } + println!( + "BENCH {}", + serde_json::json!({ + "phase": "renewals", "long_poll": enabled, "clients": clients, "repetition": repetition, + "lease_seconds": lease.as_secs(), "renewals": renewals, "elapsed_seconds": elapsed, + "renewals_per_second": renewals as f64 / elapsed, + "sample_interval_ms": COMMIT_WAIT_SAMPLE_INTERVAL.as_millis(), + "commit_waits": waits, + }) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "local release-mode benchmark; requires PostgreSQL without assertions"] +async fn workflow_latency_and_throughput_with_long_polling() { + benchmark_workflows(2, 20, false).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires a synchronous standby and DUROXIDE_PG_POOL_MAX >= 98; run in release mode"] +async fn notification_contention_with_synchronous_standby() { + benchmark_workflows(16, 64, true).await; +} + +async fn benchmark_workflows(concurrency: usize, clients: usize, require_standby: bool) { + use duroxide::runtime::{registry::ActivityRegistry, Runtime, RuntimeOptions}; + use duroxide::{ + ActivityContext, Client, OrchestrationContext, OrchestrationRegistry, OrchestrationStatus, + }; + + const REPETITIONS: usize = 3; + const DURATION: Duration = Duration::from_secs(10); + const SPARSE_SAMPLES: usize = 20; + const SPARSE_GAP: Duration = Duration::from_millis(200); + + async fn complete(client: &Client, instance: &str, activities: usize) -> Duration { + let start = Instant::now(); + let input = activities.to_string(); + client + .start_orchestration(instance, "Bench", &input) + .await + .unwrap(); + let status = client + .wait_for_orchestration(instance, Duration::from_secs(30)) + .await + .unwrap(); + assert!( + matches!(&status, OrchestrationStatus::Completed { output, .. } if output == &input), + "unexpected benchmark outcome: {status:?}" + ); + start.elapsed() + } + + for enabled in [false, true] { + for activities_per_workflow in [1, 5] { + let mut throughputs = Vec::new(); + for repetition in 0..REPETITIONS { + let fixture = Fixture::new(polling_config(enabled)).await; + let settings: serde_json::Value = sqlx::query_scalar( + "SELECT jsonb_build_object( + 'postgres_version', current_setting('server_version'), + 'debug_assertions', current_setting('debug_assertions'), + 'fsync', current_setting('fsync'), + 'full_page_writes', current_setting('full_page_writes'), + 'synchronous_commit', current_setting('synchronous_commit'), + 'synchronous_standby_names', current_setting('synchronous_standby_names'), + 'replicas', (SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'application_name', application_name, 'state', state, 'sync_state', sync_state) + ORDER BY application_name), '[]'::jsonb) FROM pg_stat_replication))", + ) + .fetch_one(&fixture.control) + .await + .unwrap(); + assert_eq!(settings["debug_assertions"], "off"); + if require_standby { + assert_eq!(settings["fsync"], "on"); + assert_eq!(settings["full_page_writes"], "on"); + assert!(matches!( + settings["synchronous_commit"].as_str(), + Some("on" | "remote_apply") + )); + assert!( + settings["replicas"] + .as_array() + .unwrap() + .iter() + .any(|replica| replica["state"] == "streaming" + && matches!( + replica["sync_state"].as_str(), + Some("sync" | "quorum") + )), + "this profile requires an active synchronous standby: {settings}" + ); + assert!( + fixture.provider.pool().options().get_max_connections() as usize + >= clients + 2 * concurrency + 2, + "provide enough connections to measure server locks rather than pool queueing", + ); + let mut warm = Vec::new(); + for _ in 0..clients + 2 * concurrency + 1 { + warm.push(fixture.provider.pool().acquire().await.unwrap()); + } + drop(warm); + } + let activities = ActivityRegistry::builder() + .register("Echo", |_ctx: ActivityContext, input: String| async move { + Ok(input) + }) + .build(); + let orchestrations = OrchestrationRegistry::builder() + .register( + "Bench", + |ctx: OrchestrationContext, input: String| async move { + let count = + input.parse::().map_err(|error| error.to_string())?; + for _ in 0..count { + ctx.schedule_activity("Echo", input.clone()).await?; + } + Ok(input) + }, + ) + .build(); + let options = RuntimeOptions { + orchestration_concurrency: concurrency, + worker_concurrency: concurrency, + dispatcher_min_poll_interval: Duration::from_millis(100), + dispatcher_long_poll_timeout: Duration::from_secs(30), + ..Default::default() + }; + println!( + "BENCH {}", + serde_json::json!({ + "phase": "config", "long_poll": enabled, "repetition": repetition, + "rust_debug_assertions": cfg!(debug_assertions), + "activities_per_workflow": activities_per_workflow, + "orch_concurrency": options.orchestration_concurrency, + "worker_concurrency": options.worker_concurrency, + "pool_max": fixture.provider.pool().options().get_max_connections(), + "clients": clients, "duration_seconds": DURATION.as_secs(), + "postgres": settings, + }) + ); + let runtime = Runtime::start_with_options( + fixture.provider.clone(), + activities, + orchestrations, + options, + ) + .await; + let client = Arc::new(Client::new(fixture.provider.clone())); + complete(&client, "warmup", activities_per_workflow).await; + if activities_per_workflow == 1 && repetition == 0 { + let mut latencies = Vec::new(); + for sample in 0..SPARSE_SAMPLES { + tokio::time::sleep(SPARSE_GAP).await; + latencies.push(complete(&client, &format!("sparse-{sample}"), 1).await); + } + latencies.sort_unstable(); + let median_ms = (latencies[SPARSE_SAMPLES / 2 - 1].as_secs_f64() + + latencies[SPARSE_SAMPLES / 2].as_secs_f64()) + * 500.0; + println!( + "BENCH {}", + serde_json::json!({ + "phase": "sparse", "long_poll": enabled, "samples": SPARSE_SAMPLES, + "idle_gap_ms": SPARSE_GAP.as_millis(), "median_ms": median_ms, + }) + ); + } + let (stop_sampling, samples) = start_commit_wait_sampling(&fixture).await; + let start = Instant::now(); + let deadline = start + DURATION; + let mut tasks = tokio::task::JoinSet::new(); + for lane in 0..clients { + let client = client.clone(); + tasks.spawn(async move { + let mut completed = 0; + while Instant::now() < deadline { + complete( + &client, + &format!("loaded-{lane}-{completed}"), + activities_per_workflow, + ) + .await; + completed += 1; + } + completed + }); + } + let mut completed = 0; + while let Some(result) = tasks.join_next().await { + completed += result.unwrap(); + } + let elapsed = start.elapsed().as_secs_f64(); + stop_sampling + .send(()) + .expect("commit-wait sampler must remain active"); + let waits = samples.await.unwrap(); + if require_standby && enabled { + assert!( + waits.max_notification_waiters > 0, + "notification contention must be observed" + ); + assert!( + waits.max_synchronous_replication_waiters > 0, + "synchronous commit waits must be observed" + ); + } + if !enabled { + assert_eq!(waits.notification_lock_backend_seconds, 0.0); + } + assert!(completed > 0); + let throughput = completed as f64 / elapsed; + throughputs.push(throughput); + println!( + "BENCH {}", + serde_json::json!({ + "phase": "loaded", "long_poll": enabled, "repetition": repetition, + "activities_per_workflow": activities_per_workflow, + "completed": completed, "elapsed_seconds": elapsed, + "workflows_per_second": throughput, + "sample_interval_ms": COMMIT_WAIT_SAMPLE_INTERVAL.as_millis(), + "commit_waits": waits, + }) + ); + runtime.shutdown(Some(1000)).await; + drop(client); + if require_standby && activities_per_workflow == 1 { + benchmark_renewals(&fixture, clients, DURATION, enabled, repetition).await; + } + fixture.cleanup().await; + } + throughputs.sort_by(f64::total_cmp); + println!( + "BENCH {}", + serde_json::json!({ + "phase": "loaded_median", "long_poll": enabled, + "activities_per_workflow": activities_per_workflow, + "workflows_per_second": throughputs[REPETITIONS / 2], + }) + ); + } + } +} + +#[tokio::test] +#[ignore = "slow local benchmark; requires pg_stat_statements and PostgreSQL without assertions"] +async fn idle_query_load_is_reduced_with_long_polling() { + use duroxide::runtime::{registry::ActivityRegistry, Runtime}; + use duroxide::OrchestrationRegistry; + + let mut calls = Vec::new(); + for enabled in [false, true] { + let fixture = Fixture::new(polling_config(enabled)).await; + let assertions: String = sqlx::query_scalar("SHOW debug_assertions") + .fetch_one(&fixture.control) + .await + .unwrap(); + assert_eq!( + assertions, "off", + "use a non-assert PostgreSQL build for benchmarks" + ); + let statement_count = || async { + sqlx::query_scalar::<_, i64>( + "SELECT COALESCE(SUM(calls), 0)::BIGINT FROM pg_stat_statements \ + WHERE toplevel AND dbid = (SELECT oid FROM pg_database WHERE datname = current_database()) \ + AND strpos(query, 'pg_stat_statements') = 0", + ) + .fetch_one(&fixture.control) + .await + .expect("create pg_stat_statements in the benchmark database") + }; + let runtime = Runtime::start_with_store( + fixture.provider.clone(), + ActivityRegistry::builder().build(), + OrchestrationRegistry::builder().build(), + ) + .await; + let before = statement_count().await; + let start = Instant::now(); + tokio::time::sleep(Duration::from_secs(65)).await; + let elapsed = start.elapsed().as_secs_f64(); + let count = statement_count().await - before; + println!( + "long_poll={enabled}: {count} provider SQL calls in {elapsed:.3}s ({:.3}/s)", + count as f64 / elapsed, + ); + calls.push(count); + runtime.shutdown(Some(1000)).await; + fixture.cleanup().await; + } + assert!(calls[0] > 0); + assert!( + calls[1] < calls[0] / 20, + "idle query counts: short={}, long={}", + calls[0], + calls[1] + ); +} diff --git a/tests/postgres_provider_test.rs b/tests/postgres_provider_test.rs index 9bd6947..17e9181 100644 --- a/tests/postgres_provider_test.rs +++ b/tests/postgres_provider_test.rs @@ -10,7 +10,7 @@ use duroxide::provider_validation::{ }; use duroxide::provider_validations::{race_replay, ProviderFactory}; use duroxide::providers::Provider; -use duroxide_pg::PostgresProvider; +use duroxide_pg::{PostgresProvider, ProviderConfig}; use sqlx::{postgres::PgPoolOptions, Executor}; use tracing_subscriber::EnvFilter; @@ -76,6 +76,7 @@ pub struct PostgresProviderFactory { database_url: String, lock_timeout_ms: u64, current_schema_name: std::sync::Mutex>, + long_poll: bool, } impl Default for PostgresProviderFactory { @@ -87,10 +88,19 @@ impl Default for PostgresProviderFactory { impl PostgresProviderFactory { pub fn new() -> Self { init_test_logging(); + let long_poll = std::env::var_os("DUROXIDE_PG_TEST_LONG_POLL") + .map(|value| { + value + .to_str() + .and_then(|value| value.parse::().ok()) + .expect("DUROXIDE_PG_TEST_LONG_POLL must be true or false") + }) + .unwrap_or(false); Self { database_url: get_database_url(), lock_timeout_ms: 30_000, // 30 seconds - must match hardcoded timeout in validation tests current_schema_name: std::sync::Mutex::new(None), + long_poll, } } @@ -101,7 +111,10 @@ impl PostgresProviderFactory { // Store schema name for cleanup *self.current_schema_name.lock().unwrap() = Some(schema_name.clone()); - let provider = PostgresProvider::new_with_schema(&self.database_url, Some(&schema_name)) + let mut config = ProviderConfig::url(&self.database_url); + config.schema_name = Some(schema_name); + config.long_poll.enabled = self.long_poll; + let provider = PostgresProvider::new_with_config(config) .await .expect("Failed to create Postgres provider for validation tests"); @@ -340,7 +353,10 @@ mod long_polling_tests { #[tokio::test] async fn test_short_poll_returns_immediately() { - let factory = PostgresProviderFactory::new(); + let factory = PostgresProviderFactory { + long_poll: false, + ..PostgresProviderFactory::new() + }; let provider = factory.create_provider().await; long_polling::test_short_poll_returns_immediately( &*provider, @@ -352,7 +368,10 @@ mod long_polling_tests { #[tokio::test] async fn test_short_poll_work_item_returns_immediately() { - let factory = PostgresProviderFactory::new(); + let factory = PostgresProviderFactory { + long_poll: false, + ..PostgresProviderFactory::new() + }; let provider = factory.create_provider().await; long_polling::test_short_poll_work_item_returns_immediately( &*provider, @@ -369,10 +388,6 @@ mod long_polling_tests { long_polling::test_fetch_respects_timeout_upper_bound(&*provider).await; factory.cleanup_schema().await; } - - // test_long_poll_waits_for_timeout: Not applicable to short-poll providers. - // This test asserts that fetch blocks for poll_timeout, which short-poll - // providers correctly do NOT do (they return immediately). } mod cancellation_tests {