From 1a197d6f377b1a3a73273061ca2fc286a47d5775 Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Tue, 20 Oct 2026 13:54:45 +0000 Subject: [PATCH 1/6] Bring long polling over to duroxide-pg --- .github/workflows/test.yml | 4 +- CHANGELOG.md | 24 + Cargo.toml | 2 +- README.md | 77 + docs/migration-policy.md | 34 + .../0027_add_long_poll_notifications.sql | 114 ++ migrations/0027_diff.md | 134 ++ migrations/README.md | 33 +- src/lib.rs | 23 + src/migrations.rs | 115 +- src/notifier.rs | 649 ++++++++ src/provider.rs | 376 +++-- tests/long_poll_tests.rs | 1479 +++++++++++++++++ tests/postgres_provider_test.rs | 31 +- 14 files changed, 2980 insertions(+), 115 deletions(-) create mode 100644 migrations/0027_add_long_poll_notifications.sql create mode 100644 migrations/0027_diff.md create mode 100644 src/notifier.rs create mode 100644 tests/long_poll_tests.rs 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..cc6a03c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,30 @@ 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. +- Migration `0027` 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-writing transactions 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. + ## [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..f1b3c6e 100644 --- a/README.md +++ b/README.md @@ -174,9 +174,86 @@ 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. 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 in both queues, rescans after a timer fires, and periodically wakes +fetches to recover missed notifications. Reconnection resubscribes and rechecks +work. Notifications are hints; queue predicates and locks remain authoritative. +Lost notifications and eligibility changes such as expired leases are also +recovered by subsequent fetches after the runtime's long-poll timeout. + +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 `0027` 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, and released locks emit +notifications, including writes from short-polling producers. Turning off a +provider's listener does not deactivate the schema's triggers. Transactions +writing those queues 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. Benchmark concurrent writers with +representative commit latency before activating notifications; 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. + ## 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..ddd56df 100644 --- a/docs/migration-policy.md +++ b/docs/migration-policy.md @@ -34,6 +34,40 @@ separate privileged process has already run migrations. before checking or applying migrations. That serializes concurrent provider startup so multiple nodes do not race to apply the same migration. +## Long-Poll Notification Activation + +Migration `0027` installs disabled notification triggers. With +`long_poll.enabled = true`, `ApplyAll` also enables both queues' 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. +If the second queue cannot be locked immediately, it releases the first and +retries, rather than deadlocking with operations that lock the queues in reverse +order. Lock contention and deadlocks retry the entire activation transaction +up to three times before returning an error. + +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 writes after activation. +Reverting to two-phase-commit use requires stopping all long-polling workers +and explicitly disabling the four `trg_notify_*` triggers on the two queues; +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. Evaluate +the loaded-write cost before opting a schema in; activation affects all of +its producers, not only processes using a listener. + +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 `VerifyOnly` fails with messages in these shapes: diff --git a/migrations/0027_add_long_poll_notifications.sql b/migrations/0027_add_long_poll_notifications.sql new file mode 100644 index 0000000..4c3b295 --- /dev/null +++ b/migrations/0027_add_long_poll_notifications.sql @@ -0,0 +1,114 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the MIT License. + +-- Migration 0027: 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 + PERFORM pg_catalog.pg_notify( + 'duroxide_pg_' || TG_RELID::TEXT, + pg_catalog.ceil(EXTRACT(EPOCH FROM NEW.visible_at) * 1000)::BIGINT::TEXT + ); + 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 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) + ) + 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 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) + ) + EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work_released; + +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.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/0027_diff.md b/migrations/0027_diff.md new file mode 100644 index 0000000..408f465 --- /dev/null +++ b/migrations/0027_diff.md @@ -0,0 +1,134 @@ +# Diff for migration 0027 + +**Migration file:** `0027_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. + +```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 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) + ) + 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 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) + ) + EXECUTE FUNCTION duroxide_notify_queue_work(); +ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work_released; +``` + +## Function Changes + +### `cleanup_schema` - body modified (baseline: 0021) + +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.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 ++ PERFORM pg_catalog.pg_notify( ++ 'duroxide_pg_' || TG_RELID::TEXT, ++ pg_catalog.ceil(EXTRACT(EPOCH FROM NEW.visible_at) * 1000)::BIGINT::TEXT ++ ); ++ RETURN NEW; ++END; ++$function$ +``` + +--- diff --git a/migrations/README.md b/migrations/README.md index 3114817..04ea34d 100644 --- a/migrations/README.md +++ b/migrations/README.md @@ -90,6 +90,38 @@ 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 `0027` 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 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 +144,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..93220ad 100644 --- a/src/migrations.rs +++ b/src/migrations.rs @@ -8,6 +8,7 @@ use sqlx::PgPool; use std::collections::{HashMap, HashSet}; use std::sync::Arc; +// Touch this module when adding embedded files so incremental builds see them. Latest: 0027. static MIGRATIONS: Dir = include_dir!("$CARGO_MANIFEST_DIR/migrations"); const BASE_TABLES: &[&str] = &[ @@ -91,6 +92,10 @@ impl MigrationRunner { /// 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?; let conn = &mut *conn; self.lock_for_migrations(conn).await?; @@ -99,10 +104,15 @@ 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), - }; + let result = async { + self.check_no_unknown_migrations(conn).await?; + self.migrate_inner(conn).await?; + if notifications { + self.enable_notifications(conn).await?; + } + Ok(()) + } + .await; self.unlock_for_migrations(conn).await; result @@ -119,6 +129,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 +178,102 @@ impl MigrationRunner { ); } + if notifications && self.notification_trigger_counts(conn).await?.1 != 4 { + 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 counts: (i64, i64) = sqlx::query_as( + "SELECT COUNT(*), COUNT(*) FILTER (WHERE tgenabled IN ('O', 'A')) \ + FROM pg_catalog.pg_trigger \ + WHERE (tgrelid = $1::TEXT::REGCLASS AND tgname IN \ + ('trg_notify_orch_work', 'trg_notify_orch_work_released')) \ + OR (tgrelid = $2::TEXT::REGCLASS AND tgname IN \ + ('trg_notify_worker_work', 'trg_notify_worker_work_released'))", + ) + .bind(format!("{}.orchestrator_queue", self.schema_name)) + .bind(format!("{}.worker_queue", self.schema_name)) + .fetch_one(conn) + .await?; + anyhow::ensure!( + counts.0 == 4, + "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 == 4 { + return Ok(()); + } + + let mut retries = 0; + loop { + let result: Result<(), sqlx::Error> = async { + // Match normal acknowledgments. Other administrative operations + // can reverse this order, so never wait while holding both up. + let mut tx = conn.begin().await?; + sqlx::query(&format!( + "LOCK TABLE {}.worker_queue IN SHARE ROW EXCLUSIVE MODE", + self.schema_name, + )) + .execute(&mut *tx) + .await?; + sqlx::query(&format!( + "LOCK TABLE {}.orchestrator_queue IN SHARE ROW EXCLUSIVE MODE NOWAIT", + self.schema_name, + )) + .execute(&mut *tx) + .await?; + for (table, trigger) in [ + ("worker_queue", "trg_notify_worker_work"), + ("orchestrator_queue", "trg_notify_orch_work"), + ] { + sqlx::query(&format!( + "ALTER TABLE {}.{table} ENABLE TRIGGER {trigger}, \ + ENABLE TRIGGER {trigger}_released", + self.schema_name, + )) + .execute(&mut *tx) + .await?; + } + tx.commit().await + } + .await; + match result { + Ok(()) => return Ok(()), + Err(error) + if retries < 3 + && error.as_database_error().is_some_and(|error| { + matches!(error.code().as_deref(), Some("40P01" | "55P03")) + }) => + { + retries += 1; + tracing::warn!( + schema = %self.schema_name, + attempt = retries, + error = %error, + "Queue lock contention activating long-poll notifications; retrying the transaction" + ); + tokio::time::sleep(std::time::Duration::from_millis(50 * retries)).await; + } + Err(error) => return Err(error.into()), + } + } + } + /// 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 diff --git a/src/notifier.rs b/src/notifier.rs new file mode 100644 index 0000000..2e861ad --- /dev/null +++ b/src/notifier.rs @@ -0,0 +1,649 @@ +// 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, 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::{sleep, sleep_until, timeout_at, Instant}; +use tracing::warn; + +use crate::provider::PostgresProvider; + +const MIN_RETRY_INTERVAL: Duration = Duration::from_millis(100); + +/// 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 \ + (SELECT CEIL(EXTRACT(EPOCH FROM visible_at) * 1000)::BIGINT \ + FROM {schema}.orchestrator_queue \ + WHERE visible_at > TO_TIMESTAMP($1::BIGINT / 1000.0) AND lock_token IS NULL \ + ORDER BY visible_at LIMIT 1), \ + (SELECT CEIL(EXTRACT(EPOCH FROM visible_at) * 1000)::BIGINT \ + FROM {schema}.worker_queue \ + WHERE visible_at > TO_TIMESTAMP($1::BIGINT / 1000.0) AND lock_token IS NULL \ + ORDER BY visible_at LIMIT 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?; + Ok(listener) +} + +async fn run( + pool: PgPool, + mut listener: PgListener, + channels: [String; 2], + notifications: [Arc; 2], + timer_query: String, + config: LongPollConfig, +) { + loop { + let result = run_connected( + &pool, + &mut listener, + &channels, + ¬ifications, + &timer_query, + &config, + ) + .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); + loop { + tokio::select! { + _ = pool.close_event() => { + for notify in ¬ifications { + notify.notify_waiters(); + } + return; + } + _ = sleep(Duration::from_secs(1)) => {} + } + match subscribe(&pool, &channels).await { + Ok(reconnected) => { + listener = reconnected; + for notify in ¬ifications { + notify.notify_waiters(); + } + break; + } + Err(error) => warn!( + target: "duroxide::providers::postgres::notifier", + error = %error, + "Long-poll reconnect failed; fetch timeouts retain polling fallback" + ), + } + } + } +} + +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<(Option, Option)> { + let mut connection = RefreshConnection(Some(pool.acquire().await?)); + let result = sqlx::query_as(query) + .bind(PostgresProvider::now_millis()) + .fetch_one(&mut **connection.0.as_mut().expect("refresh owns its connection")) + .await?; + drop(connection.0.take()); + Ok(result) +} + +async fn run_connected( + pool: &PgPool, + listener: &mut PgListener, + channels: &[String; 2], + notifications: &[Arc; 2], + timer_query: &str, + config: &LongPollConfig, +) -> 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; + let refresh = refresh_timers(pool, timer_query); + tokio::pin!(refresh); + let mut refreshing = true; + // A timer can expire after the in-flight query took its snapshot. + let mut refresh_again = false; + + 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 = &mut refresh, if refreshing => { + refreshing = false; + let recover = Instant::now() >= next_recovery; + if recover { + next_recovery = deadline_after(config.notifier_poll_interval)?; + } + next_refresh = next_recovery; + match result { + Ok((orch, worker)) => { + 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 { + notifications[index].notify_waiters(); + } + if let Some(timestamp) = timestamp { + refresh_again |= timestamp <= PostgresProvider::now_millis(); + schedule(timestamp, ¬ifications[index], &mut timers[index], config); + } + } + if refresh_again { + next_refresh = next_refresh.min( + (refresh_started + MIN_RETRY_INTERVAL).max(Instant::now()) + ); + } + } + Err(error) => warn!( + target: "duroxide::providers::postgres::notifier", + error = %error, + "Long-poll timer refresh failed; fetch timeouts retain polling fallback" + ), + } + refresh_again = false; + } + _ = sleep_until(next_refresh), if !refreshing => { + refresh_started = Instant::now(); + refresh.set(refresh_timers(pool, timer_query)); + refreshing = true; + } + _ = 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; + 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. + refresh_again |= refreshing; + 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}; + + #[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 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 run = run_connected( + &pool, + &mut listener, + &channels, + ¬ifications, + &query, + &config, + ); + 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 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..a684bf1 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,11 +402,13 @@ 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) @@ -400,18 +419,29 @@ impl PostgresProvider { 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()]) @@ -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,90 @@ 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, + ..Default::default() + }, + ) + .await + .expect("Entra long-poll construction should work with two connections"); + 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 + .expect("notification must wake the two-connection Entra provider"); + assert!(result.unwrap().is_some()); + provider.cleanup_schema().await.unwrap(); + let pool = provider.pool().clone(); + drop(provider); + pool.close().await; + } + #[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 +3651,7 @@ mod entra_pipeline_tests { token_source, PgSslMode::Disable, MigrationPolicy::ApplyAll, + LongPollConfig::default(), ) .await; @@ -3526,6 +3701,7 @@ mod entra_pipeline_tests { token_source, PgSslMode::Disable, MigrationPolicy::ApplyAll, + LongPollConfig::default(), ) .await .expect("default-constructor variant must succeed"); diff --git a/tests/long_poll_tests.rs b/tests/long_poll_tests.rs new file mode 100644 index 0000000..5e00a27 --- /dev/null +++ b/tests/long_poll_tests.rs @@ -0,0 +1,1479 @@ +// 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; + +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 { + match self { + Self::Orchestration => provider + .fetch_orchestration_item(LOCK, timeout, None) + .await + .unwrap() + .map(|(_, token, attempt)| Claim { token, attempt }), + Self::Worker => provider + .fetch_work_item(LOCK, 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 mut config = ProviderConfig::url(database_url()); + 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( + &format!("FROM {schema}.worker_queue", schema = fixture.schema), + None, + ) + .await; + tokio::task::yield_now().await; + } + fixture + } + + async fn wait_for_query(&self, query: &str, since: Option>) { + 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))", + ) + .bind(query) + .bind(since) + .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), + ) + .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 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 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)) + .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)) + .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)) + .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)) + .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)) + .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; + 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(); + assert!(!exists); + fixture.cleanup().await; +} + +#[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 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(200))) + .await; + Queue::Worker + .enqueue(&fixture, "after-refresh", Some(Duration::from_millis(700))) + .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; + Queue::Worker.ack(&fixture.provider, &third.token).await; + fixture.cleanup().await; +} + +#[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/0027_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 \ + WHERE application_name = $1 AND wait_event_type = 'Lock' \ + AND strpos(query, 'worker_queue') > 0)", + ) + .bind(&application) + .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 NOWAIT') > 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(); + sqlx::query(&format!("CREATE SCHEMA {schema};",)) + .execute(&control) + .await + .unwrap(); + sqlx::query(&format!( + "CREATE TABLE {schema}._duroxide_migrations \ + (version BIGINT PRIMARY KEY, name TEXT NOT NULL, applied_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP)", + )).execute(&control).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(|extension| extension == "sql") + }) + .filter(|file| file.path().file_name().unwrap().to_str().unwrap() < "0027") + .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(); + let mut tx = control.begin().await.unwrap(); + sqlx::query(&format!("SET LOCAL search_path TO {schema}, pg_temp")) + .execute(&mut *tx) + .await + .unwrap(); + sqlx::raw_sql(file.contents_utf8().unwrap()) + .execute(&mut *tx) + .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(); + } + 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; +} + +#[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 strpos(query, $1) > 0", + ) + .bind(format!("{}.", fixture.schema)) + .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 { From 60a2bb9b1b571ad523d3ee6b8f3a5f1e2897b954 Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Tue, 20 Oct 2026 17:14:38 +0000 Subject: [PATCH 2/6] Fixes requested by review --- CHANGELOG.md | 2 + README.md | 11 +- .../0027_add_long_poll_notifications.sql | 20 ++- migrations/0027_diff.md | 25 +++- src/notifier.rs | 31 ++-- tests/long_poll_tests.rs | 141 +++++++++++++++++- 6 files changed, 201 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cc6a03c..6004449 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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. - Migration `0027` 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, diff --git a/README.md b/README.md index f1b3c6e..38c310e 100644 --- a/README.md +++ b/README.md @@ -221,10 +221,13 @@ 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, and released locks emit -notifications, including writes from short-polling producers. Turning off a -provider's listener does not deactivate the schema's triggers. Transactions -writing those queues cannot use `PREPARE TRANSACTION`, because PostgreSQL does +Once activated, inserts, visibility 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. This includes writes from +short-polling producers. Turning off a provider's listener does not deactivate +the schema's triggers. Transactions writing those queues 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. diff --git a/migrations/0027_add_long_poll_notifications.sql b/migrations/0027_add_long_poll_notifications.sql index 4c3b295..95a64d5 100644 --- a/migrations/0027_add_long_poll_notifications.sql +++ b/migrations/0027_add_long_poll_notifications.sql @@ -10,6 +10,20 @@ LANGUAGE plpgsql SET search_path = pg_catalog, pg_temp AS $notify$ BEGIN + IF TG_OP = 'DELETE' THEN + -- Acknowledgment consumes claimed rows and releases the instance lock + -- in one transaction; NOTIFY is delivered only after that commit. + IF OLD.lock_token IS NOT NULL THEN + PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); + END IF; + RETURN OLD; + ELSIF TG_OP = 'UPDATE' THEN + IF NEW.visible_at IS NOT DISTINCT FROM OLD.visible_at + AND NOT (OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL) THEN + RETURN NEW; + END IF; + END IF; + PERFORM pg_catalog.pg_notify( 'duroxide_pg_' || TG_RELID::TEXT, pg_catalog.ceil(EXTRACT(EPOCH FROM NEW.visible_at) * 1000)::BIGINT::TEXT @@ -44,12 +58,8 @@ 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 ON orchestrator_queue + AFTER DELETE OR UPDATE OF visible_at, lock_token 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) - ) EXECUTE FUNCTION duroxide_notify_queue_work(); ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work_released; diff --git a/migrations/0027_diff.md b/migrations/0027_diff.md index 408f465..bde63b8 100644 --- a/migrations/0027_diff.md +++ b/migrations/0027_diff.md @@ -20,6 +20,11 @@ short-polling-only schemas retain their previous transaction behavior. Installation and activation take worker-queue locks before orchestration-queue locks, matching normal acknowledgments. +The orchestration release trigger also covers deletion of claimed messages. +Its helper filters out claims and other updates that do not release work; +acknowledgment notifications are delivered after the transaction removes the +instance lock. + ```sql CREATE TRIGGER trg_notify_worker_work AFTER INSERT ON worker_queue @@ -42,12 +47,8 @@ CREATE TRIGGER trg_notify_orch_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 ON orchestrator_queue + AFTER DELETE OR UPDATE OF visible_at, lock_token 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) - ) EXECUTE FUNCTION duroxide_notify_queue_work(); ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work_released; ``` @@ -122,6 +123,20 @@ Full function with diff: + SET search_path TO 'pg_catalog', 'pg_temp' +AS $function$ +BEGIN ++ IF TG_OP = 'DELETE' THEN ++ -- Acknowledgment consumes claimed rows and releases the instance lock ++ -- in one transaction; NOTIFY is delivered only after that commit. ++ IF OLD.lock_token IS NOT NULL THEN ++ PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); ++ END IF; ++ RETURN OLD; ++ ELSIF TG_OP = 'UPDATE' THEN ++ IF NEW.visible_at IS NOT DISTINCT FROM OLD.visible_at ++ AND NOT (OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL) THEN ++ RETURN NEW; ++ END IF; ++ END IF; ++ + PERFORM pg_catalog.pg_notify( + 'duroxide_pg_' || TG_RELID::TEXT, + pg_catalog.ceil(EXTRACT(EPOCH FROM NEW.visible_at) * 1000)::BIGINT::TEXT diff --git a/src/notifier.rs b/src/notifier.rs index 2e861ad..342f1b8 100644 --- a/src/notifier.rs +++ b/src/notifier.rs @@ -249,8 +249,10 @@ async fn run_connected( let refresh = refresh_timers(pool, timer_query); tokio::pin!(refresh); let mut refreshing = true; - // A timer can expire after the in-flight query took its snapshot. - let mut refresh_again = false; + // 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(); @@ -274,30 +276,35 @@ async fn run_connected( 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 { + if recover || refreshing_catchup[index] { notifications[index].notify_waiters(); } if let Some(timestamp) = timestamp { - refresh_again |= timestamp <= PostgresProvider::now_millis(); + pending_catchup[index] |= timestamp <= PostgresProvider::now_millis(); schedule(timestamp, ¬ifications[index], &mut timers[index], config); } } - if refresh_again { + if pending_catchup.iter().any(|pending| *pending) { next_refresh = next_refresh.min( (refresh_started + MIN_RETRY_INTERVAL).max(Instant::now()) ); } } - Err(error) => warn!( - target: "duroxide::providers::postgres::notifier", - error = %error, - "Long-poll timer refresh failed; fetch timeouts retain polling fallback" - ), + Err(error) => { + for (pending, attempted) in pending_catchup.iter_mut().zip(refreshing_catchup) { + *pending |= attempted; + } + warn!( + target: "duroxide::providers::postgres::notifier", + error = %error, + "Long-poll timer refresh failed; fetch timeouts retain polling fallback" + ); + } } - refresh_again = false; } _ = sleep_until(next_refresh), if !refreshing => { refresh_started = Instant::now(); + refreshing_catchup = std::mem::take(&mut pending_catchup); refresh.set(refresh_timers(pool, timer_query)); refreshing = true; } @@ -306,12 +313,12 @@ async fn run_connected( 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. - refresh_again |= refreshing; next_refresh = next_refresh.min( (refresh_started + MIN_RETRY_INTERVAL).max(now) ); diff --git a/tests/long_poll_tests.rs b/tests/long_poll_tests.rs index 5e00a27..4d855e2 100644 --- a/tests/long_poll_tests.rs +++ b/tests/long_poll_tests.rs @@ -416,6 +416,139 @@ async fn inserts_wake_already_waiting_fetches_for_both_queues() { 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), + ) + .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_stay_silent_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(); + 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 closely_spaced_timers_do_not_disappear_between_paced_scans() { + for queue in [Queue::Orchestration, Queue::Worker] { + let mut config = polling_config(true); + config.timer_grace_period = Duration::ZERO; + let fixture = Fixture::new(config).await; + let first = fixture.waiting_fetch(queue).await; + queue + .enqueue( + &fixture, + "first-close-timer", + Some(Duration::from_millis(20)), + ) + .await; + queue + .enqueue( + &fixture, + "second-close-timer", + Some(Duration::from_millis(40)), + ) + .await; + let first = claimed(first).await; + // Keep the first item claimed so its acknowledgment cannot supply a + // notification that would hide a missed timer wake. + let second = + tokio::time::timeout(Duration::from_secs(3), queue.fetch(&fixture.provider, POLL)) + .await + .expect("the paced scan must wake work that matured since the first timer") + .expect("second timer"); + assert_ne!(first.token, second.token); + queue.ack(&fixture.provider, &first.token).await; + 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] { @@ -1134,11 +1267,13 @@ async fn notification_installation_and_activation_follow_writer_lock_order() { let deadline = Instant::now() + Duration::from_secs(5); loop { let waiting: bool = sqlx::query_scalar( - "SELECT EXISTS(SELECT 1 FROM pg_stat_activity \ - WHERE application_name = $1 AND wait_event_type = 'Lock' \ - AND strpos(query, 'worker_queue') > 0)", + "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(); From c3ba02a4012496b947175c465937e95fd3280e57 Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Fri, 30 Oct 2026 13:32:45 +0000 Subject: [PATCH 3/6] Respond to review feedback - Future-dated lease hints, including session expiry, plus immediate release hints for delayed abandonment. - Warm fetch connections, bounded timer scans, and independent degraded-mode fallback. - Bounded, jittered migration/activation retries with rollback and deadlock regressions. - Updated trigger inventory, deterministic catch-up coverage, and a reproducible benchmark driver. --- CHANGELOG.md | 15 +- README.md | 62 ++- docs/migration-policy.md | 29 +- .../0027_add_long_poll_notifications.sql | 54 +- migrations/0027_diff.md | 61 ++- migrations/README.md | 4 + src/migrations.rs | 270 ++++++---- src/notifier.rs | 395 +++++++++++++- src/provider.rs | 49 +- tests/deadlocks/harness.rs | 37 ++ tests/deadlocks/main.rs | 3 + tests/deadlocks/notification_ddl.rs | 323 ++++++++++++ tests/long_poll_tests.rs | 492 +++++++++++++++--- 13 files changed, 1542 insertions(+), 252 deletions(-) create mode 100644 tests/deadlocks/notification_ddl.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 6004449..3f6af61 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,10 +18,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 are bounded to 250 ms; degraded listeners keep a one-second + polling fallback independent of reconnect attempts. - Migration `0027` 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-writing transactions cannot use PostgreSQL `PREPARE TRANSACTION`. + 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. @@ -32,6 +39,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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. +- Add a reproducible sparse/loaded workflow benchmark alongside the idle-query + benchmark, recording concurrency, duration, pool size, and commit settings. ## [0.1.38] - 2026-10-26 diff --git a/README.md b/README.md index 38c310e..c3605bd 100644 --- a/README.md +++ b/README.md @@ -193,18 +193,26 @@ let provider = PostgresProvider::new_with_config(config).await?; 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. Leave it disabled for client-only -providers, especially one-connection backend clients. It requires a +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 in both queues, rescans after a timer fires, and periodically wakes -fetches to recover missed notifications. Reconnection resubscribes and rechecks -work. Notifications are hints; queue predicates and locks remain authoritative. -Lost notifications and eligibility changes such as expired leases are also -recovered by subsequent fetches after the runtime's long-poll timeout. +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 have a 250 ms budget, including connection acquisition, +so a scan blocked on one queue cannot indefinitely monopolize a two-connection +pool. Failed or cancelled scans discard their connections. Scan errors retry +after about one second; while reconnecting, an independent one-second wakeup +keeps fetches checking for work even if connection establishment stalls. +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 @@ -221,12 +229,14 @@ 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, 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. This includes writes from +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 cannot use +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. @@ -253,6 +263,34 @@ 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)' +``` + +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, and commit settings (`fsync`, `full_page_writes`, +`synchronous_commit`, and `synchronous_standby_names`). Keep those settings and +the 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. The `pg-stress` binary +continues to use short polling. + ## Features - Automatic schema migration on startup diff --git a/docs/migration-policy.md b/docs/migration-policy.md index ddd56df..bf72f6c 100644 --- a/docs/migration-policy.md +++ b/docs/migration-policy.md @@ -33,19 +33,27 @@ 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. ## Long-Poll Notification Activation Migration `0027` installs disabled notification triggers. With -`long_poll.enabled = true`, `ApplyAll` also enables both queues' triggers in a +`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. -If the second queue cannot be locked immediately, it releases the first and -retries, rather than deadlocking with operations that lock the queues in reverse -order. Lock contention and deadlocks retry the entire activation transaction -up to three times before returning an error. +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 @@ -53,9 +61,16 @@ 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 writes after activation. +`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 the four `trg_notify_*` triggers on the two queues; +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 diff --git a/migrations/0027_add_long_poll_notifications.sql b/migrations/0027_add_long_poll_notifications.sql index 95a64d5..54a716b 100644 --- a/migrations/0027_add_long_poll_notifications.sql +++ b/migrations/0027_add_long_poll_notifications.sql @@ -10,23 +10,29 @@ LANGUAGE plpgsql SET search_path = pg_catalog, pg_temp AS $notify$ BEGIN - IF TG_OP = 'DELETE' THEN + 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. - IF OLD.lock_token IS NOT NULL THEN - PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); - END IF; + PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); RETURN OLD; - ELSIF TG_OP = 'UPDATE' THEN - IF NEW.visible_at IS NOT DISTINCT FROM OLD.visible_at - AND NOT (OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL) THEN - RETURN NEW; + 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, - pg_catalog.ceil(EXTRACT(EPOCH FROM NEW.visible_at) * 1000)::BIGINT::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; @@ -41,11 +47,15 @@ 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 ON worker_queue + 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; @@ -58,11 +68,33 @@ 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 DELETE OR UPDATE OF visible_at, lock_token ON orchestrator_queue + 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(); diff --git a/migrations/0027_diff.md b/migrations/0027_diff.md index bde63b8..0a0a69f 100644 --- a/migrations/0027_diff.md +++ b/migrations/0027_diff.md @@ -20,10 +20,11 @@ short-polling-only schemas retain their previous transaction behavior. Installation and activation take worker-queue locks before orchestration-queue locks, matching normal acknowledgments. -The orchestration release trigger also covers deletion of claimed messages. -Its helper filters out claims and other updates that do not release work; -acknowledgment notifications are delivered after the transaction removes the -instance lock. +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 @@ -32,11 +33,15 @@ CREATE TRIGGER trg_notify_worker_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 ON worker_queue + 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; @@ -47,10 +52,30 @@ CREATE TRIGGER trg_notify_orch_work ALTER TABLE orchestrator_queue DISABLE TRIGGER trg_notify_orch_work; CREATE TRIGGER trg_notify_orch_work_released - AFTER DELETE OR UPDATE OF visible_at, lock_token ON orchestrator_queue + 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 @@ -123,23 +148,29 @@ Full function with diff: + SET search_path TO 'pg_catalog', 'pg_temp' +AS $function$ +BEGIN -+ IF TG_OP = 'DELETE' THEN ++ 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. -+ IF OLD.lock_token IS NOT NULL THEN -+ PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); -+ END IF; ++ PERFORM pg_catalog.pg_notify('duroxide_pg_' || TG_RELID::TEXT, '0'); + RETURN OLD; -+ ELSIF TG_OP = 'UPDATE' THEN -+ IF NEW.visible_at IS NOT DISTINCT FROM OLD.visible_at -+ AND NOT (OLD.lock_token IS NOT NULL AND NEW.lock_token IS NULL) THEN -+ RETURN NEW; ++ 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, -+ pg_catalog.ceil(EXTRACT(EPOCH FROM NEW.visible_at) * 1000)::BIGINT::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; diff --git a/migrations/README.md b/migrations/README.md index 04ea34d..815a557 100644 --- a/migrations/README.md +++ b/migrations/README.md @@ -95,6 +95,10 @@ Migration `0027` installs disabled queue-notification triggers. A provider with 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 provider validation harness can exercise either mode: diff --git a/src/migrations.rs b/src/migrations.rs index 93220ad..f54b7e5 100644 --- a/src/migrations.rs +++ b/src/migrations.rs @@ -1,16 +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: 0027. 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", @@ -81,13 +94,19 @@ 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 @@ -97,6 +116,8 @@ impl MigrationRunner { 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?; @@ -104,7 +125,7 @@ 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 = async { + let result: Result<()> = async { self.check_no_unknown_migrations(conn).await?; self.migrate_inner(conn).await?; if notifications { @@ -113,9 +134,9 @@ impl MigrationRunner { Ok(()) } .await; - self.unlock_for_migrations(conn).await; - - result + let unlock_result = self.unlock_for_migrations(conn).await; + result?; + unlock_result } /// Verify that the migration tracking table exists and that every embedded @@ -178,7 +199,9 @@ impl MigrationRunner { ); } - if notifications && self.notification_trigger_counts(conn).await?.1 != 4 { + 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 \ @@ -194,20 +217,22 @@ impl MigrationRunner { &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 \ - WHERE (tgrelid = $1::TEXT::REGCLASS AND tgname IN \ - ('trg_notify_orch_work', 'trg_notify_orch_work_released')) \ - OR (tgrelid = $2::TEXT::REGCLASS AND tgname IN \ - ('trg_notify_worker_work', 'trg_notify_worker_work_released'))", + JOIN unnest($1::TEXT[], $2::TEXT[]) expected(relation, name) \ + ON tgrelid = expected.relation::REGCLASS AND tgname = expected.name", ) - .bind(format!("{}.orchestrator_queue", self.schema_name)) - .bind(format!("{}.worker_queue", self.schema_name)) + .bind(relations) + .bind(names) .fetch_one(conn) .await?; anyhow::ensure!( - counts.0 == 4, + 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, ); @@ -215,65 +240,91 @@ impl MigrationRunner { } async fn enable_notifications(&self, conn: &mut sqlx::postgres::PgConnection) -> Result<()> { - if self.notification_trigger_counts(conn).await?.1 == 4 { + if self.notification_trigger_counts(conn).await?.1 == NOTIFICATION_TRIGGERS.len() as i64 { return Ok(()); } - let mut retries = 0; + let deadline = Instant::now() + DDL_RETRY_BUDGET; loop { - let result: Result<(), sqlx::Error> = async { - // Match normal acknowledgments. Other administrative operations - // can reverse this order, so never wait while holding both up. + let result = async { let mut tx = conn.begin().await?; - sqlx::query(&format!( - "LOCK TABLE {}.worker_queue IN SHARE ROW EXCLUSIVE MODE", - self.schema_name, - )) - .execute(&mut *tx) - .await?; - sqlx::query(&format!( - "LOCK TABLE {}.orchestrator_queue IN SHARE ROW EXCLUSIVE MODE NOWAIT", - self.schema_name, - )) - .execute(&mut *tx) - .await?; - for (table, trigger) in [ - ("worker_queue", "trg_notify_worker_work"), - ("orchestrator_queue", "trg_notify_orch_work"), - ] { - sqlx::query(&format!( - "ALTER TABLE {}.{table} ENABLE TRIGGER {trigger}, \ - ENABLE TRIGGER {trigger}_released", - self.schema_name, - )) - .execute(&mut *tx) - .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(()) } - tx.commit().await + .await; + Self::finish_ddl(tx, result).await } .await; match result { Ok(()) => return Ok(()), - Err(error) - if retries < 3 - && error.as_database_error().is_some_and(|error| { - matches!(error.code().as_deref(), Some("40P01" | "55P03")) - }) => - { - retries += 1; - tracing::warn!( - schema = %self.schema_name, - attempt = retries, - error = %error, - "Queue lock contention activating long-poll notifications; retrying the transaction" - ); - tokio::time::sleep(std::time::Duration::from_millis(50 * retries)).await; + 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 + ))); + } } - Err(error) => return Err(error.into()), } } } + 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 retryable = error + .downcast_ref::() + .and_then(sqlx::Error::as_database_error) + .is_some_and(|error| matches!(error.code().as_deref(), Some("40P01" | "55P03"))); + if !retryable || Instant::now() >= deadline { + return false; + } + tracing::warn!( + schema = %self.schema_name, + operation, + error = %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 @@ -745,20 +796,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 @@ -791,41 +846,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 {}: {}", diff --git a/src/notifier.rs b/src/notifier.rs index 342f1b8..3a89d68 100644 --- a/src/notifier.rs +++ b/src/notifier.rs @@ -10,12 +10,16 @@ use std::sync::Arc; use std::time::Duration; use tokio::sync::Notify; use tokio::task::JoinHandle; -use tokio::time::{sleep, sleep_until, timeout_at, Instant}; +use tokio::time::{ + interval_at, sleep, sleep_until, timeout, 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_TIMEOUT: Duration = Duration::from_millis(250); /// Optional LISTEN/NOTIFY-based waiting for orchestration and activity work. /// @@ -101,15 +105,21 @@ impl LongPoller { let worker_notify = Arc::new(Notify::new()); let notifications = [orch_notify.clone(), worker_notify.clone()]; let timer_query = format!( - "SELECT \ + "SELECT LEAST(\ (SELECT CEIL(EXTRACT(EPOCH FROM visible_at) * 1000)::BIGINT \ FROM {schema}.orchestrator_queue \ - WHERE visible_at > TO_TIMESTAMP($1::BIGINT / 1000.0) AND lock_token IS NULL \ + 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) AND lock_token IS NULL \ - ORDER BY visible_at LIMIT 1)" + 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, @@ -145,6 +155,12 @@ async fn subscribe(pool: &PgPool, channels: &[String; 2]) -> Result 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) } @@ -184,6 +200,25 @@ async fn run( ), } 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() => { @@ -192,21 +227,18 @@ async fn run( } return; } - _ = sleep(Duration::from_secs(1)) => {} - } - match subscribe(&pool, &channels).await { - Ok(reconnected) => { + _ = fallback.tick() => { + for notify in ¬ifications { + notify.notify_waiters(); + } + } + reconnected = &mut reconnect => { listener = reconnected; for notify in ¬ifications { notify.notify_waiters(); } break; } - Err(error) => warn!( - target: "duroxide::providers::postgres::notifier", - error = %error, - "Long-poll reconnect failed; fetch timeouts retain polling fallback" - ), } } } @@ -225,13 +257,17 @@ impl Drop for RefreshConnection { } async fn refresh_timers(pool: &PgPool, query: &str) -> Result<(Option, Option)> { - let mut connection = RefreshConnection(Some(pool.acquire().await?)); - let result = sqlx::query_as(query) - .bind(PostgresProvider::now_millis()) - .fetch_one(&mut **connection.0.as_mut().expect("refresh owns its connection")) - .await?; - drop(connection.0.take()); - Ok(result) + timeout(TIMER_REFRESH_TIMEOUT, async { + let mut connection = RefreshConnection(Some(pool.acquire().await?)); + let result = sqlx::query_as(query) + .bind(PostgresProvider::now_millis()) + .fetch_one(&mut **connection.0.as_mut().expect("refresh owns its connection")) + .await?; + drop(connection.0.take()); + Ok(result) + }) + .await + .context("Long-poll timer scan exceeded its 250ms read-only query budget")? } async fn run_connected( @@ -249,6 +285,7 @@ async fn run_connected( let refresh = refresh_timers(pool, timer_query); tokio::pin!(refresh); let mut refreshing = true; + let mut refresh_failed = false; // 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]; @@ -266,13 +303,14 @@ async fn run_connected( } result = &mut refresh, if refreshing => { refreshing = false; - let recover = Instant::now() >= next_recovery; + 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. @@ -291,13 +329,18 @@ async fn run_connected( } } 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 = %error, - "Long-poll timer refresh failed; fetch timeouts retain polling fallback" + "Long-poll timer refresh failed; retrying with one-second polling fallback" ); } } @@ -607,6 +650,312 @@ mod tests { ); } + #[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 { + run_connected( + &pool, + &mut listener, + &channels, + ¬ifications, + &query, + &LongPollConfig { + enabled: true, + ..Default::default() + }, + ) + .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 result = tokio::time::timeout( + Duration::from_millis(1500), + run_connected( + &pool, + &mut listener, + &channels, + ¬ifications, + &query, + &LongPollConfig { + enabled: true, + ..Default::default() + }, + ), + ) + .await; + 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 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(); diff --git a/src/provider.rs b/src/provider.rs index a684bf1..7aebe5e 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -412,7 +412,7 @@ impl PostgresProvider { 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?; @@ -545,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?; @@ -3586,11 +3586,45 @@ mod entra_pipeline_tests { 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), @@ -3612,16 +3646,19 @@ mod entra_pipeline_tests { .await .unwrap(); }; - let (result, ()) = tokio::time::timeout(Duration::from_secs(3), async { + let result = tokio::time::timeout(Duration::from_secs(3), async { tokio::join!(fetch, enqueue) }) - .await - .expect("notification must wake the two-connection Entra provider"); - assert!(result.unwrap().is_some()); + .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] 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..31c317a --- /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, 26).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"], 27); + 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, 26).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 = 27 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"], 27); + } + }).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 index 4d855e2..3d2d254 100644 --- a/tests/long_poll_tests.rs +++ b/tests/long_poll_tests.rs @@ -14,6 +14,9 @@ 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); @@ -119,14 +122,23 @@ impl Queue { } 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, None) + .fetch_orchestration_item(lock_timeout, poll_timeout, None) .await .unwrap() .map(|(_, token, attempt)| Claim { token, attempt }), Self::Worker => provider - .fetch_work_item(LOCK, timeout, None, &TagFilter::Any) + .fetch_work_item(lock_timeout, poll_timeout, None, &TagFilter::Any) .await .unwrap() .map(|(_, token, attempt)| Claim { token, attempt }), @@ -370,6 +382,208 @@ async fn existing_work_returns_without_waiting() { 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; @@ -475,7 +689,7 @@ async fn acknowledgment_wakes_another_provider_after_it_consumed_a_locked_instan } #[tokio::test] -async fn orchestration_claims_stay_silent_but_acknowledgments_notify() { +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 @@ -495,6 +709,20 @@ async fn orchestration_claims_stay_silent_but_acknowledgments_notify() { .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 @@ -514,36 +742,31 @@ async fn orchestration_claims_stay_silent_but_acknowledgments_notify() { } #[tokio::test] -async fn closely_spaced_timers_do_not_disappear_between_paced_scans() { +async fn renewed_claims_are_not_recovered_at_the_original_expiry() { for queue in [Queue::Orchestration, Queue::Worker] { - let mut config = polling_config(true); - config.timer_grace_period = Duration::ZERO; - let fixture = Fixture::new(config).await; - let first = fixture.waiting_fetch(queue).await; - queue - .enqueue( - &fixture, - "first-close-timer", - Some(Duration::from_millis(20)), - ) - .await; - queue - .enqueue( - &fixture, - "second-close-timer", - Some(Duration::from_millis(40)), - ) - .await; - let first = claimed(first).await; - // Keep the first item claimed so its acknowledgment cannot supply a - // notification that would hide a missed timer wake. - let second = - tokio::time::timeout(Duration::from_secs(3), queue.fetch(&fixture.provider, POLL)) + 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 - .expect("the paced scan must wake work that matured since the first timer") - .expect("second timer"); + .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); - queue.ack(&fixture.provider, &first.token).await; + assert_eq!(second.attempt, 2); queue.ack(&fixture.provider, &second.token).await; fixture.cleanup().await; } @@ -1338,7 +1561,7 @@ async fn activation_releases_worker_lock_before_retrying_reverse_order_contentio 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 NOWAIT') > 0)", + AND strpos(query, 'orchestrator_queue IN SHARE ROW EXCLUSIVE MODE') > 0)", ) .bind(&application) .fetch_one(&fixture.control) @@ -1423,47 +1646,7 @@ async fn upgrading_pre_notification_schema_preserves_work_and_defaults_to_disabl .await .unwrap(); let schema = schema_name(); - sqlx::query(&format!("CREATE SCHEMA {schema};",)) - .execute(&control) - .await - .unwrap(); - sqlx::query(&format!( - "CREATE TABLE {schema}._duroxide_migrations \ - (version BIGINT PRIMARY KEY, name TEXT NOT NULL, applied_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP)", - )).execute(&control).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(|extension| extension == "sql") - }) - .filter(|file| file.path().file_name().unwrap().to_str().unwrap() < "0027") - .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(); - let mut tx = control.begin().await.unwrap(); - sqlx::query(&format!("SET LOCAL search_path TO {schema}, pg_temp")) - .execute(&mut *tx) - .await - .unwrap(); - sqlx::raw_sql(file.contents_utf8().unwrap()) - .execute(&mut *tx) - .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(); - } + harness::migrate_through(&control, &schema, 26).await; let payload = serde_json::to_string(&activity("retained", None, None)).unwrap(); sqlx::query(&format!( "INSERT INTO {schema}.worker_queue \ @@ -1558,6 +1741,171 @@ async fn two_phase_commit_is_preserved_until_explicit_notification_activation() fixture.cleanup().await; } +#[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() { + use duroxide::runtime::{registry::ActivityRegistry, Runtime, RuntimeOptions}; + use duroxide::{ + ActivityContext, Client, OrchestrationContext, OrchestrationRegistry, OrchestrationStatus, + }; + + const CLIENTS: usize = 20; + 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)); + 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'))", + ) + .fetch_one(&fixture.control) + .await + .unwrap(); + assert_eq!(settings["debug_assertions"], "off"); + 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: 2, + worker_concurrency: 2, + 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 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(); + 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, + }) + ); + runtime.shutdown(Some(1000)).await; + drop(client); + 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() { From 9dbe0cf358f72810078e2f5b2dac90f9aadc9817 Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Fri, 30 Oct 2026 16:36:54 +0000 Subject: [PATCH 4/6] fix: bound timer scans without orphaning server backends Run timer reads in transactions with server-side lock and statement timeouts. Retain one owned scan across client watchdog wakeups and listener reconnects instead of cancelling and replacing an in-flight query. Add the supplied backend-accumulation regression, connection-reuse and single-flight coverage, and a timer wake deadline that cannot pass through the polling fallback. Addresses review items 41-N1 and 41-N5. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 6 +- README.md | 16 +- src/notifier.rs | 357 ++++++++++++++++++++++++++++++++++++--- tests/long_poll_tests.rs | 133 +++++++++++++-- 4 files changed, 469 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3f6af61..eca9c11 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,8 +23,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 are bounded to 250 ms; degraded listeners keep a one-second - polling fallback independent of reconnect attempts. + 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 `0027` 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, diff --git a/README.md b/README.md index c3605bd..cbc4d1b 100644 --- a/README.md +++ b/README.md @@ -207,12 +207,16 @@ 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 have a 250 ms budget, including connection acquisition, -so a scan blocked on one queue cannot indefinitely monopolize a two-connection -pool. Failed or cancelled scans discard their connections. Scan errors retry -after about one second; while reconnecting, an independent one-second wakeup -keeps fetches checking for work even if connection establishment stalls. -Successful reconnection resubscribes and rechecks work. +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 diff --git a/src/notifier.rs b/src/notifier.rs index 3a89d68..8ef1dec 100644 --- a/src/notifier.rs +++ b/src/notifier.rs @@ -4,22 +4,20 @@ use anyhow::{ensure, Context, Result}; use duroxide::providers::ProviderError; use sqlx::postgres::{types::Oid, PgListener, PgNotification}; -use sqlx::{pool::PoolConnection, PgPool, Postgres}; +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, timeout_at, Instant, MissedTickBehavior, -}; +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_TIMEOUT: Duration = Duration::from_millis(250); +const TIMER_REFRESH_BACKSTOP: Duration = Duration::from_secs(1); /// Optional LISTEN/NOTIFY-based waiting for orchestration and activity work. /// @@ -172,6 +170,7 @@ async fn run( timer_query: String, config: LongPollConfig, ) { + let mut refresh = TimerRefresh::default(); loop { let result = run_connected( &pool, @@ -180,6 +179,7 @@ async fn run( ¬ifications, &timer_query, &config, + &mut refresh, ) .await; for notify in ¬ifications { @@ -244,6 +244,40 @@ async fn run( } } +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 { @@ -256,18 +290,40 @@ impl Drop for RefreshConnection { } } -async fn refresh_timers(pool: &PgPool, query: &str) -> Result<(Option, Option)> { - timeout(TIMER_REFRESH_TIMEOUT, async { - let mut connection = RefreshConnection(Some(pool.acquire().await?)); - let result = sqlx::query_as(query) +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 **connection.0.as_mut().expect("refresh owns its connection")) - .await?; - drop(connection.0.take()); - Ok(result) - }) - .await - .context("Long-poll timer scan exceeded its 250ms read-only query budget")? + .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( @@ -277,15 +333,17 @@ async fn run_connected( 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; - let refresh = refresh_timers(pool, timer_query); - tokio::pin!(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]; @@ -301,7 +359,7 @@ async fn run_connected( }; handle_notification(notification, channels, notifications, &mut timers, config); } - result = &mut refresh, if refreshing => { + result = refresh.result(), if refreshing => { refreshing = false; let recover = refresh_failed || Instant::now() >= next_recovery; if recover { @@ -339,17 +397,33 @@ async fn run_connected( } warn!( target: "duroxide::providers::postgres::notifier", - error = %error, + 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.set(refresh_timers(pool, timer_query)); + 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(); @@ -466,6 +540,7 @@ where mod tests { use super::*; use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::time::timeout; #[test] fn defaults_preserve_short_polling() { @@ -578,6 +653,89 @@ mod tests { } } + #[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(); @@ -614,6 +772,7 @@ mod tests { ); let mut worker_wakes = 0; { + let mut refresh = TimerRefresh::default(); let run = run_connected( &pool, &mut listener, @@ -621,6 +780,7 @@ mod tests { ¬ifications, &query, &config, + &mut refresh, ); tokio::pin!(run); let duration = tokio::time::timeout(Duration::from_millis(350), async { @@ -702,6 +862,7 @@ mod tests { let notifications = notifications.clone(); let query = format!("SELECT * FROM {schema}.scan($1)"); async move { + let mut refresh = TimerRefresh::default(); run_connected( &pool, &mut listener, @@ -712,6 +873,7 @@ mod tests { enabled: true, ..Default::default() }, + &mut refresh, ) .await } @@ -790,6 +952,7 @@ mod tests { "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( @@ -802,9 +965,11 @@ mod tests { enabled: true, ..Default::default() }, + &mut refresh, ), ) .await; + drop(refresh); assert!( result.is_err(), "the listener should stay connected after a scan error" @@ -851,6 +1016,156 @@ mod tests { 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; diff --git a/tests/long_poll_tests.rs b/tests/long_poll_tests.rs index 3d2d254..d1baae8 100644 --- a/tests/long_poll_tests.rs +++ b/tests/long_poll_tests.rs @@ -196,7 +196,11 @@ impl Fixture { async fn with_schema(long_poll: LongPollConfig, schema: String) -> Self { let enabled = long_poll.enabled; - let mut config = ProviderConfig::url(database_url()); + 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()); @@ -212,26 +216,30 @@ impl Fixture { }; if enabled { fixture - .wait_for_query( - &format!("FROM {schema}.worker_queue", schema = fixture.schema), - None, - ) + .wait_for_query("COMMIT", None, Some(&application)) .await; tokio::task::yield_now().await; } fixture } - async fn wait_for_query(&self, query: &str, since: Option>) { + 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 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(); @@ -279,6 +287,7 @@ impl Fixture { self.wait_for_query( &format!("{}.{}", self.schema, queue.fetch_name()), Some(since), + None, ) .await; assert!( @@ -659,6 +668,7 @@ async fn acknowledgment_wakes_another_provider_after_it_consumed_a_locked_instan .wait_for_query( &format!("{}.fetch_orchestration_item", waiter.schema), Some(inserted_after), + None, ) .await; assert!(!waiting.is_finished(), "the instance is still locked"); @@ -964,7 +974,11 @@ async fn disabled_producer_wakes_verify_only_consumer() { tokio::spawn(async move { Queue::Worker.fetch(&consumer, POLL).await }) }; fixture - .wait_for_query(&format!("{}.fetch_work_item", fixture.schema), Some(since)) + .wait_for_query( + &format!("{}.fetch_work_item", fixture.schema), + Some(since), + None, + ) .await; producer .enqueue_for_worker(activity("mixed-config", None, None)) @@ -1034,7 +1048,11 @@ async fn pool_closure_wakes_waiters_and_releases_listener() { .await }); fixture - .wait_for_query(&format!("{}.fetch_work_item", fixture.schema), Some(since)) + .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 @@ -1067,7 +1085,11 @@ async fn pool_closure_during_reconnect_does_not_leave_fetch_waiting() { .await }); fixture - .wait_for_query(&format!("{}.fetch_work_item", fixture.schema), Some(since)) + .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 @@ -1119,7 +1141,11 @@ async fn notifications_wake_the_matching_tag_waiter() { })); } fixture - .wait_for_query(&format!("{}.fetch_work_item", fixture.schema), Some(since)) + .wait_for_query( + &format!("{}.fetch_work_item", fixture.schema), + Some(since), + None, + ) .await; fixture .provider @@ -1190,7 +1216,11 @@ async fn notifications_wake_the_session_owner_not_just_an_arbitrary_slot() { }) }; fixture - .wait_for_query(&format!("{}.fetch_work_item", fixture.schema), Some(since)) + .wait_for_query( + &format!("{}.fetch_work_item", fixture.schema), + Some(since), + None, + ) .await; fixture .provider @@ -1268,6 +1298,71 @@ async fn notification_migration_preserves_an_application_function_with_the_old_g 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); @@ -1364,10 +1459,11 @@ async fn timers_expiring_during_refresh_do_not_hide_the_next_timer() { let second = fixture.waiting_fetch(Queue::Worker).await; Queue::Worker - .enqueue(&fixture, "during-refresh", Some(Duration::from_millis(200))) + .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(700))) + .enqueue(&fixture, "after-refresh", Some(Duration::from_millis(300))) .await; let second = claimed(second).await; Queue::Worker.ack(&fixture.provider, &second.token).await; @@ -1375,8 +1471,17 @@ async fn timers_expiring_during_refresh_do_not_hide_the_next_timer() { 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] From ee7b5106469b0a5004545bc0299b4d6d885760cc Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Fri, 30 Oct 2026 16:59:48 +0000 Subject: [PATCH 5/6] fix: reserve migration 0030 and clarify startup failures Place long-poll notification DDL after the pending 0027-0029 migrations, preserving migration 0028's prune_executions_batch cleanup and companion diff baseline. Update upgrade tests and documentation without stacking the prerequisite PRs. Include SQLSTATE and the full context chain in DDL retry warnings, test those diagnostics, and document draining writers and retrying ApplyAll when the fixed contention budget expires. Validated the intended predecessor sequence in an isolated source copy. Addresses 41-C1, 41-N3, and 41-N4. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 5 +- README.md | 2 +- docs/migration-policy.md | 11 ++- ...l => 0030_add_long_poll_notifications.sql} | 3 +- migrations/{0027_diff.md => 0030_diff.md} | 7 +- migrations/README.md | 11 ++- src/migrations.rs | 80 +++++++++++++++++-- tests/deadlocks/notification_ddl.rs | 10 +-- tests/long_poll_tests.rs | 29 ++++++- 9 files changed, 136 insertions(+), 22 deletions(-) rename migrations/{0027_add_long_poll_notifications.sql => 0030_add_long_poll_notifications.sql} (97%) rename migrations/{0027_diff.md => 0030_diff.md} (97%) diff --git a/CHANGELOG.md b/CHANGELOG.md index eca9c11..49b9176 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,7 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 `0027` adds disabled notification triggers. An `ApplyAll` provider +- 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`. @@ -44,7 +44,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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. + 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. diff --git a/README.md b/README.md index cbc4d1b..efe42f8 100644 --- a/README.md +++ b/README.md @@ -228,7 +228,7 @@ 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 `0027` installs **disabled** queue notification triggers. Constructing +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). diff --git a/docs/migration-policy.md b/docs/migration-policy.md index bf72f6c..7999169 100644 --- a/docs/migration-policy.md +++ b/docs/migration-policy.md @@ -40,9 +40,18 @@ 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 `0027` installs disabled notification triggers. With +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 diff --git a/migrations/0027_add_long_poll_notifications.sql b/migrations/0030_add_long_poll_notifications.sql similarity index 97% rename from migrations/0027_add_long_poll_notifications.sql rename to migrations/0030_add_long_poll_notifications.sql index 54a716b..7670ae9 100644 --- a/migrations/0027_add_long_poll_notifications.sql +++ b/migrations/0030_add_long_poll_notifications.sql @@ -1,7 +1,7 @@ -- Copyright (c) Microsoft Corporation. -- Licensed under the MIT License. --- Migration 0027: notification hints for optional long-polling. +-- 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() @@ -144,6 +144,7 @@ BEGIN 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); diff --git a/migrations/0027_diff.md b/migrations/0030_diff.md similarity index 97% rename from migrations/0027_diff.md rename to migrations/0030_diff.md index 0a0a69f..f54c8e9 100644 --- a/migrations/0027_diff.md +++ b/migrations/0030_diff.md @@ -1,6 +1,6 @@ -# Diff for migration 0027 +# Diff for migration 0030 -**Migration file:** `0027_add_long_poll_notifications.sql` +**Migration file:** `0030_add_long_poll_notifications.sql` Each changed function is shown **in full** with `+`/`-` markers on changed lines. @@ -80,7 +80,7 @@ ALTER TABLE sessions DISABLE TRIGGER trg_notify_session_lease; ## Function Changes -### `cleanup_schema` - body modified (baseline: 0021) +### `cleanup_schema` - body modified (baseline: 0028) Full function with diff: ```diff @@ -130,6 +130,7 @@ Full function with diff: 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); diff --git a/migrations/README.md b/migrations/README.md index 815a557..721fa13 100644 --- a/migrations/README.md +++ b/migrations/README.md @@ -90,7 +90,7 @@ 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 `0027` installs disabled queue-notification triggers. A provider with +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 @@ -100,6 +100,15 @@ 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 diff --git a/src/migrations.rs b/src/migrations.rs index f54b7e5..3eb3cd6 100644 --- a/src/migrations.rs +++ b/src/migrations.rs @@ -10,7 +10,7 @@ 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: 0027. +// 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'"; @@ -307,17 +307,18 @@ impl MigrationRunner { } async fn retry_ddl(&self, error: &anyhow::Error, deadline: Instant, operation: &str) -> bool { - let retryable = error + let sqlstate = error .downcast_ref::() .and_then(sqlx::Error::as_database_error) - .is_some_and(|error| matches!(error.code().as_deref(), Some("40P01" | "55P03"))); - if !retryable || Instant::now() >= deadline { + .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, - error = %error, + 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])); @@ -901,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/tests/deadlocks/notification_ddl.rs b/tests/deadlocks/notification_ddl.rs index 31c317a..3731f86 100644 --- a/tests/deadlocks/notification_ddl.rs +++ b/tests/deadlocks/notification_ddl.rs @@ -31,7 +31,7 @@ async fn snapshot(pool: &PgPool, schema: &str) -> Value { 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, 26).await; + migrate_through(&pool, &schema, 29).await; sqlx::Executor::execute( &pool, format!( @@ -99,7 +99,7 @@ async fn notification_upgrade_retries_reverse_order_delete_contention() { 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"], 27); + assert_eq!(after_ledger.last().unwrap()["version"], 30); assert!(after["worker"].is_null()); assert!(after["orchestration"].is_null()); }) @@ -111,12 +111,12 @@ 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, 26).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 = 27 AND nextval('{schema}.migration_attempts') = 1 THEN + IF NEW.version = 30 AND nextval('{schema}.migration_attempts') = 1 THEN RAISE EXCEPTION 'injected migration failure' USING ERRCODE = '{code}'; END IF; RETURN NEW; @@ -149,7 +149,7 @@ async fn notification_migration_retries_only_transient_errors_and_rolls_back() { 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"], 27); + assert_eq!(after_ledger.last().unwrap()["version"], 30); } }).await; } diff --git a/tests/long_poll_tests.rs b/tests/long_poll_tests.rs index d1baae8..5cd7904 100644 --- a/tests/long_poll_tests.rs +++ b/tests/long_poll_tests.rs @@ -1242,6 +1242,20 @@ async fn notifications_wake_the_session_owner_not_just_an_arbitrary_slot() { #[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 @@ -1251,8 +1265,17 @@ async fn public_style_cleanup_removes_notification_function() { .fetch_one(&fixture.control) .await .unwrap(); - assert!(!exists); + 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] @@ -1574,7 +1597,7 @@ async fn notification_installation_and_activation_follow_writer_lock_order() { .await .unwrap(); sqlx::raw_sql(include_str!( - "../migrations/0027_add_long_poll_notifications.sql" + "../migrations/0030_add_long_poll_notifications.sql" )) .execute(&mut *tx) .await @@ -1751,7 +1774,7 @@ async fn upgrading_pre_notification_schema_preserves_work_and_defaults_to_disabl .await .unwrap(); let schema = schema_name(); - harness::migrate_through(&control, &schema, 26).await; + 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 \ From 529bb0a62df5f3c4ab4ddc3388923d0da7e6a811 Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Fri, 30 Oct 2026 17:55:23 +0000 Subject: [PATCH 6/6] test: add a synchronous-standby notification contention benchmark Add a 16+16 dispatcher, 64-client profile with sampled notification-commit and synchronous-replication wait time. Exercise workflow claims and saturated orchestration, activity, and session renewals with verified synchronous-standby and pool-capacity preconditions. Document notification-bearing operations and reproducible benchmark methodology without publishing performance results from uncontrolled runs. Include timer transaction control in idle SQL counts. Addresses review item 41-N2. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 4 + README.md | 67 +++++++++-- docs/migration-policy.md | 13 ++- tests/long_poll_tests.rs | 247 +++++++++++++++++++++++++++++++++++++-- 4 files changed, 311 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 49b9176..00e4320 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 diff --git a/README.md b/README.md index efe42f8..071b336 100644 --- a/README.md +++ b/README.md @@ -248,9 +248,17 @@ 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. Benchmark concurrent writers with -representative commit latency before activating notifications; idle-query -savings alone do not establish loaded throughput. +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, @@ -279,6 +287,9 @@ DUROXIDE_PG_POOL_MAX=10 cargo nextest run --release --test long_poll_tests \ -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 @@ -288,12 +299,52 @@ 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, and commit settings (`fsync`, `full_page_writes`, -`synchronous_commit`, and `synchronous_standby_names`). Keep those settings and -the source revision alongside results; the driver does not relax durability. +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. The `pg-stress` binary -continues to use short polling. +`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 diff --git a/docs/migration-policy.md b/docs/migration-policy.md index 7999169..5fe834a 100644 --- a/docs/migration-policy.md +++ b/docs/migration-policy.md @@ -84,9 +84,16 @@ 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. Evaluate -the loaded-write cost before opting a schema in; activation affects all of -its producers, not only processes using a listener. +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 diff --git a/tests/long_poll_tests.rs b/tests/long_poll_tests.rs index 5cd7904..56577c6 100644 --- a/tests/long_poll_tests.rs +++ b/tests/long_poll_tests.rs @@ -1869,15 +1869,185 @@ async fn two_phase_commit_is_preserved_until_explicit_notification_activation() 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 CLIENTS: usize = 20; const REPETITIONS: usize = 3; const DURATION: Duration = Duration::from_secs(10); const SPARSE_SAMPLES: usize = 20; @@ -1894,7 +2064,10 @@ async fn workflow_latency_and_throughput_with_long_polling() { .wait_for_orchestration(instance, Duration::from_secs(30)) .await .unwrap(); - assert!(matches!(status, OrchestrationStatus::Completed { output, .. } if output == input)); + assert!( + matches!(&status, OrchestrationStatus::Completed { output, .. } if output == &input), + "unexpected benchmark outcome: {status:?}" + ); start.elapsed() } @@ -1910,12 +2083,45 @@ async fn workflow_latency_and_throughput_with_long_polling() { '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'))", + '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) @@ -1935,8 +2141,8 @@ async fn workflow_latency_and_throughput_with_long_polling() { ) .build(); let options = RuntimeOptions { - orchestration_concurrency: 2, - worker_concurrency: 2, + orchestration_concurrency: concurrency, + worker_concurrency: concurrency, dispatcher_min_poll_interval: Duration::from_millis(100), dispatcher_long_poll_timeout: Duration::from_secs(30), ..Default::default() @@ -1950,7 +2156,7 @@ async fn workflow_latency_and_throughput_with_long_polling() { "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(), + "clients": clients, "duration_seconds": DURATION.as_secs(), "postgres": settings, }) ); @@ -1981,10 +2187,11 @@ async fn workflow_latency_and_throughput_with_long_polling() { }) ); } + 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 { + for lane in 0..clients { let client = client.clone(); tasks.spawn(async move { let mut completed = 0; @@ -2005,6 +2212,23 @@ async fn workflow_latency_and_throughput_with_long_polling() { 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); @@ -2015,10 +2239,15 @@ async fn workflow_latency_and_throughput_with_long_polling() { "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); @@ -2054,9 +2283,9 @@ async fn idle_query_load_is_reduced_with_long_polling() { let statement_count = || async { sqlx::query_scalar::<_, i64>( "SELECT COALESCE(SUM(calls), 0)::BIGINT FROM pg_stat_statements \ - WHERE toplevel AND strpos(query, $1) > 0", + WHERE toplevel AND dbid = (SELECT oid FROM pg_database WHERE datname = current_database()) \ + AND strpos(query, 'pg_stat_statements') = 0", ) - .bind(format!("{}.", fixture.schema)) .fetch_one(&fixture.control) .await .expect("create pg_stat_statements in the benchmark database")