From 28215dd7d66c6dc31856dc194e3b2f620380f174 Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Tue, 20 Oct 2026 18:16:27 +0000 Subject: [PATCH 1/5] Rate-limit df.start --- CHANGELOG.md | 7 + USER_GUIDE.md | 67 +++++- docs/api-reference.md | 26 +++ docs/upgrade-testing.md | 21 ++ scripts/run-pgspot.sh | 1 + scripts/test-upgrade.sh | 25 +++ src/client.rs | 69 ++++-- src/dsl.rs | 136 ++++++++++-- src/lib.rs | 24 ++ tests/e2e/sql/89_schema_qualification.sql | 16 ++ tests/e2e/sql/90_start_limits.sql | 209 ++++++++++++++++++ tests/e2e/sql/91_start_limits_concurrency.sql | 162 ++++++++++++++ 12 files changed, 728 insertions(+), 35 deletions(-) create mode 100644 tests/e2e/sql/90_start_limits.sql create mode 100644 tests/e2e/sql/91_start_limits_concurrency.sql diff --git a/CHANGELOG.md b/CHANGELOG.md index 3221202b..db8a2728 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -71,6 +71,13 @@ Pre-1.0 note: while `pg_durable` is in major version `0`, minor releases may inc ### Security +- **Start admission quotas (#139):** `df.max_concurrent_per_user` (default + `100`) and `df.max_instances_per_user` (default `10000`) limit active and + retained instances per submitting role and origin database. Both are + superuser-settable; `0` disables a quota and superuser identities are exempt. + Starts serialize per role until the submitting transaction ends and reject + over-quota work before insert/enqueue, including on older installed schemas. + See [Instance Limits](USER_GUIDE.md#instance-limits) before upgrading. - **SQL name-resolution hardening (#418):** extension-owned runtime SQL explicitly qualifies catalog objects, functions, types and operators to resist shadowing through caller or database `search_path` settings. User-authored SQL diff --git a/USER_GUIDE.md b/USER_GUIDE.md index 5233784e..6d649c8f 100644 --- a/USER_GUIDE.md +++ b/USER_GUIDE.md @@ -26,10 +26,11 @@ pg_durable is a PostgreSQL extension that brings durable, fault-tolerant functio 16. [Connection Host](#connection-host) 17. [HTTP Security](#http-security) 18. [HTTP Allowed Domains](#http-allowed-domains) -19. [Connection Limits](#connection-limits) -20. [Troubleshooting](#troubleshooting) -21. [Quick Reference Card](#quick-reference-card) -22. [Appendix: Test Data Setup](#appendix-test-data-setup) +19. [Instance Limits](#instance-limits) +20. [Connection Limits](#connection-limits) +21. [Troubleshooting](#troubleshooting) +22. [Quick Reference Card](#quick-reference-card) +23. [Appendix: Test Data Setup](#appendix-test-data-setup) --- @@ -3003,6 +3004,64 @@ for the full syntax and default domain list. --- +## Instance Limits + +`df.start()` checks two quotas before creating an instance or enqueueing work: + +| Setting | Default | Counts | +|---------|---------|--------| +| `df.max_concurrent_per_user` | `100` | Instances with status `pending` or `running`, including sleeping, looping, and signal-waiting workflows | +| `df.max_instances_per_user` | `10000` | All retained instances, including completed, failed, and cancelled workflows | + +Each quota applies to the **submitting role (`current_user`) in the origin +database**. `SET ROLE` and `SECURITY DEFINER` affect that identity just as they +affect workflow execution. An explicit SQL target database does not move the +quota to that database. Superuser submission identities bypass both quotas, +but still require `pg_durable.enable_superuser_instances = on`. + +Both settings accept integers from `0` to `2147483647`; `0` disables that +individual quota. They are superuser-settable (`SUSET`) and can be changed +without a restart. For example, an administrator can set defaults for new +connections by an application role: + +```sql +ALTER ROLE app_user IN DATABASE appdb SET df.max_concurrent_per_user = 200; +ALTER ROLE app_user IN DATABASE appdb SET df.max_instances_per_user = 25000; +``` + +Server settings in `postgresql.conf` or `ALTER SYSTEM` take effect on reload. +Session `SET` overrides affect `'caller'` starts. **`transaction_mode => 'new'` +uses the separate session's role/database/server defaults**, not the caller's +session overrides. Configure those defaults to apply the same policy to both +modes; `SET ROLE` alone does not load a role's connection defaults. + +At a quota, the call fails with SQLSTATE `53400` (`configuration_limit_exceeded`) +and a `HINT` identifying the relevant setting. Finishing or cancelling work +frees concurrent capacity, but not retained-instance capacity. Retained capacity +is freed when [automatic reconciliation](#automatic-reconciliation) removes +terminal instances. If reconciliation is disabled, or retained work has not +aged out, an administrator must adjust retention or raise the quota. + +Concurrent starts by the same role in the same database wait for the preceding +**submitting transaction** to commit or roll back before checking capacity. +They do not wait for workflows to finish. Keep submitting transactions short +and use PostgreSQL `lock_timeout` to bound admission waits. Counts include +current committed rows and the caller's own changes even under `REPEATABLE READ` +or `SERIALIZABLE`. Rollback, including rollback to a savepoint, releases the +corresponding reservations. + +Do not follow a `'caller'` start with a `'new'` start for the same role and +database while the caller transaction remains open: the separate session waits +on the caller's admission lock and eventually hits its 30-second timeout. +Commit or roll back the caller transaction first. + +These are instance quotas, not calls-per-second or byte limits. They do not +bound an individual graph's size or history. Changing a quota does not cancel +or delete existing work; it only controls admission of new instances. The +defaults also apply after a binary-only upgrade, without `ALTER EXTENSION UPDATE`. + +--- + ## Connection Limits pg_durable uses separate limits for control management, provider, origin metadata, diff --git a/docs/api-reference.md b/docs/api-reference.md index b4c477d6..d0345e05 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -527,6 +527,13 @@ database privileges. Install in `pg_durable.database` first and wait for control readiness before creating satellites; see [Multi-Database Support](../USER_GUIDE.md#multi-database-support). +Non-superuser starts are subject to `df.max_concurrent_per_user` (default `100`) +and `df.max_instances_per_user` (default `10000`) in their origin database. +Rejection raises SQLSTATE `53400` with a tuning hint before any insert or enqueue. +Concurrent starts by the same role wait until the earlier submitting transaction +ends. See [Instance Limits](../USER_GUIDE.md#instance-limits) for quota scope, +retention, transaction interactions, and configuration. + #### transaction_mode Selects which transaction the *start itself* runs in. It changes nothing about @@ -893,6 +900,25 @@ These settings are configured via `ALTER SYSTEM SET` or `postgresql.conf`. See e --- +### df.max_concurrent_per_user / df.max_instances_per_user + +Admission quotas for `df.start()`, per submitting role and origin database. + +| Property | Concurrent quota | Retained-instance quota | +|----------|------------------|-------------------------| +| Default | `100` | `10000` | +| Counts | `pending` and `running` instances | All `df.instances` rows owned by the role | +| Range | `0` to `2147483647`; `0` disables | `0` to `2147483647`; `0` disables | +| Context | `SUSET` | `SUSET` | + +Superuser submission identities are exempt. Administrators can change either +setting at runtime, including per-role/database connection defaults. +`transaction_mode => 'new'` uses its separate session's defaults and does not +inherit the caller's session `SET` overrides. See +[Instance Limits](../USER_GUIDE.md#instance-limits) for tuning and admission waits. + +--- + ### pg_durable.http_security Server-wide outbound HTTP policy for `df.http()` and `df.http_multipart()`. diff --git a/docs/upgrade-testing.md b/docs/upgrade-testing.md index f969dd47..c68bdac7 100644 --- a/docs/upgrade-testing.md +++ b/docs/upgrade-testing.md @@ -363,6 +363,27 @@ development cycle, not the published v0.2.8 release. Multi-database installation identity and HTTP endpoint/secret DDL also belong in this migration; the shipped 0.2.7 to 0.2.8 script must remain byte-identical. +#### `df.start()` admission quotas (#139) + +- **DDL/B1:** No new schema objects, grants, or migration DDL. Both start + wrapper symbols use the same admission check against existing `df.instances` + columns, including on 0.2.2+ schemas without `ALTER EXTENSION UPDATE`. + B1 exercises a non-superuser reaching its quota on each supported schema. +- **Behavior change:** New superuser-settable GUCs + `df.max_concurrent_per_user` (default `100`) and `df.max_instances_per_user` + (default `10000`) apply immediately after loading the new binary. + Existing work remains intact; starts at or above either quota fail with + SQLSTATE `53400`. Set a quota to `0` to disable it. +- **Concurrency:** Admission uses a database-local, per-role transaction + advisory lock and a current snapshot, including under repeatable-read and + serializable isolation. E2E tests cover concurrent commit, rollback, + savepoint rollback, and lock timeout. +- **B2/replay:** No data migration, runtime schema detection, durable payload, + or orchestration changes. Workflows that call `df.start()` from SQL are + subject to the new quotas when that SQL activity executes. The separate + `'new'` session uses its own configuration defaults, not caller session + overrides. See [Instance Limits](../USER_GUIDE.md#instance-limits). + #### Multi-database installation - **Local DDL:** [0.2.8 to 0.2.9](../sql/pg_durable--0.2.8--0.2.9.sql) diff --git a/scripts/run-pgspot.sh b/scripts/run-pgspot.sh index 2f241401..cfe19230 100755 --- a/scripts/run-pgspot.sh +++ b/scripts/run-pgspot.sh @@ -41,6 +41,7 @@ PGSPOT_ALLOW=( # pg_durable's DSL intentionally exposes unqualified custom operators (for # example, `df.sql(...) ~> df.sql(...)`) so users do not need df in search_path. # pgspot reports the generated CREATE OPERATOR name as an unqualified object. + # Match all seven DSL operators explicitly; other unqualified objects still fail. '^PS017: Unqualified object reference: ~> at line [0-9]+$' '^PS017: Unqualified object reference: \|=> at line [0-9]+$' '^PS017: Unqualified object reference: & at line [0-9]+$' diff --git a/scripts/test-upgrade.sh b/scripts/test-upgrade.sh index 2544483f..5d7ebb47 100755 --- a/scripts/test-upgrade.sh +++ b/scripts/test-upgrade.sh @@ -952,6 +952,30 @@ test_b1_status_instance() { assert_sql_equals "SELECT df.status('${B1_INSTANCE_ID}');" "completed" } +test_b1_start_limits() { + run_sql_capture " + BEGIN; + CREATE ROLE upgrade_start_limits LOGIN; + SELECT df.grant_usage('upgrade_start_limits'); + SET LOCAL df.max_concurrent_per_user = 1; + SET LOCAL df.max_instances_per_user = 0; + SET LOCAL ROLE upgrade_start_limits; + DO \$\$ + BEGIN + PERFORM df.start('SELECT 1'); + BEGIN + PERFORM df.start('SELECT 1'); + RAISE EXCEPTION 'legacy df.start bypassed concurrency quota'; + EXCEPTION WHEN configuration_limit_exceeded THEN + IF SQLERRM NOT LIKE '%has 1 active instances (limit 1)%' THEN + RAISE; + END IF; + END; + END \$\$; + ROLLBACK; + " >/dev/null +} + test_b1_result() { assert_sql_contains "SELECT df.result('${B1_INSTANCE_ID}');" "test_value" } @@ -1146,6 +1170,7 @@ else run_test "B1 [v${B1_VERSION}]: DSL chain (~>)" test_b1_dsl_chain run_test "B1 [v${B1_VERSION}]: conditional operators (?>/!>)" test_b1_conditional_operators run_test "B1 [v${B1_VERSION}]: df.start()/wait_for_completion()" test_b1_start_and_complete + run_test "B1 [v${B1_VERSION}]: df.start() admission limits" test_b1_start_limits run_test "B1 [v${B1_VERSION}]: df.status() on real instance" test_b1_status_instance run_test "B1 [v${B1_VERSION}]: df.result()" test_b1_result run_test "B1 [v${B1_VERSION}]: df.list_instances()" test_b1_list_instances diff --git a/src/client.rs b/src/client.rs index c21af802..ed6da7d6 100644 --- a/src/client.rs +++ b/src/client.rs @@ -14,6 +14,7 @@ use std::sync::OnceLock; use std::time::Duration; use duroxide::Client; +use pgrx::pg_sys::panic::ErrorReport; use pgrx::prelude::*; use sqlx::Connection; use tokio::runtime::Runtime; @@ -278,12 +279,34 @@ fn format_new_transaction_start_error(sqlstate: Option<&str>, message: &str) -> } } -fn classify_new_transaction_start_error(err: sqlx::Error) -> String { +fn new_transaction_start_error(message: String) -> Box { + Box::new(ErrorReport::new( + PgSqlErrorCode::ERRCODE_INTERNAL_ERROR, + message, + function_name!(), + )) +} + +fn classify_new_transaction_start_error(err: sqlx::Error) -> Box { let message = err.to_string(); - let sqlstate = err + let database_error = err .as_database_error() - .and_then(|db_err| db_err.code().map(|code| code.to_string())); - format_new_transaction_start_error(sqlstate.as_deref(), &message) + .and_then(|e| e.try_downcast_ref::()); + let sqlstate = database_error.map(|e| e.code()); + let code = if sqlstate == Some("53400") { + PgSqlErrorCode::ERRCODE_CONFIGURATION_LIMIT_EXCEEDED + } else { + PgSqlErrorCode::ERRCODE_INTERNAL_ERROR + }; + let mut report = ErrorReport::new( + code, + format_new_transaction_start_error(sqlstate, &message), + function_name!(), + ); + if let Some(hint) = database_error.and_then(|e| e.hint()) { + report = report.set_hint(hint); + } + Box::new(report) } /// Run `df.start()` on an already-established separate session. @@ -292,7 +315,7 @@ async fn start_on_new_session( fut: &str, label: &Option, database: &Option, -) -> Result { +) -> Result> { use sqlx::Row; // `connect_as_user` marks its connections as workflow-execution sessions @@ -303,7 +326,11 @@ async fn start_on_new_session( sqlx::query("SET df.in_workflow = 'false'") .execute(&mut *conn) .await - .map_err(|e| format!("failed to clear df.in_workflow on new-transaction session: {e}"))?; + .map_err(|e| { + new_transaction_start_error(format!( + "failed to clear df.in_workflow on new-transaction session: {e}" + )) + })?; // Bound both the lock wait and the statement itself. See // NEW_TRANSACTION_START_STATEMENT_TIMEOUT_MS for why this is not optional. Two @@ -313,10 +340,11 @@ async fn start_on_new_session( format!("SET lock_timeout = {NEW_TRANSACTION_START_STATEMENT_TIMEOUT_MS}"), format!("SET statement_timeout = {NEW_TRANSACTION_START_STATEMENT_TIMEOUT_MS}"), ] { - sqlx::query(&stmt) - .execute(&mut *conn) - .await - .map_err(|e| format!("failed to set timeouts on new-transaction session: {e}"))?; + sqlx::query(&stmt).execute(&mut *conn).await.map_err(|e| { + new_transaction_start_error(format!( + "failed to set timeouts on new-transaction session: {e}" + )) + })?; } // Three positional arguments on purpose: this must also resolve on schemas @@ -332,8 +360,11 @@ async fn start_on_new_session( .await .map_err(classify_new_transaction_start_error)?; - row.try_get("id") - .map_err(|e| format!("df.start() on new transaction returned no instance id: {e}")) + row.try_get("id").map_err(|e| { + new_transaction_start_error(format!( + "df.start() on new transaction returned no instance id: {e}" + )) + }) } /// Start a durable function **in its own transaction**. @@ -359,7 +390,7 @@ pub fn start_in_new_transaction( label: Option<&str>, database: Option<&str>, user: &str, -) -> Result { +) -> Result> { use sqlx::Connection; let rt = get_client_runtime(); @@ -370,11 +401,17 @@ pub fn start_in_new_transaction( let user = user.to_string(); let origin_database = Spi::get_one::("SELECT pg_catalog.current_database()::pg_catalog.text") - .map_err(|e| format!("Failed to resolve caller database: {e}"))? - .ok_or_else(|| "Failed to resolve caller database".to_string())?; + .map_err(|e| { + new_transaction_start_error(format!("Failed to resolve caller database: {e}")) + })? + .ok_or_else(|| { + new_transaction_start_error("Failed to resolve caller database".into()) + })?; rt.block_on(async { - let mut conn = connect_as_user_for_new_transaction(&user, &origin_database).await?; + let mut conn = connect_as_user_for_new_transaction(&user, &origin_database) + .await + .map_err(new_transaction_start_error)?; let result = start_on_new_session(&mut conn, &fut, &label, &database).await; diff --git a/src/dsl.rs b/src/dsl.rs index 5d8190c3..2bb522d2 100644 --- a/src/dsl.rs +++ b/src/dsl.rs @@ -5,6 +5,7 @@ use cron::Schedule as CronSchedule; use pgrx::datum::DatumWithOid; +use pgrx::pg_sys::panic::ErrorReport; use pgrx::prelude::*; use std::str::FromStr; @@ -1062,6 +1063,107 @@ fn require_login_privilege(oid: pgrx::pg_sys::Oid, name: &str, caller: &str) { } } +// "PGDQ": separate from the loopback-launch slots. Advisory locks are +// database-local; the second key preserves all bits of the submitting role OID. +const INSTANCE_QUOTA_LOCK_CLASS_ID: i32 = 0x5047_4451; + +fn current_instance_counts(user_oid: pg_sys::Oid) -> (i64, i64) { + Spi::connect(|_| unsafe { + let mut argtypes = [pg_sys::OIDOID]; + let plan = pg_sys::SPI_prepare( + c"SELECT pg_catalog.count(*) FILTER ( + WHERE pg_catalog.lower(status) OPERATOR(pg_catalog.=) + ANY (ARRAY['pending', 'running']::pg_catalog.text[])), + pg_catalog.count(*) + FROM df.instances + WHERE submitted_by OPERATOR(pg_catalog.=) $1::pg_catalog.regrole" + .as_ptr(), + 1, + argtypes.as_mut_ptr(), + ); + if plan.is_null() { + let status = pg_sys::SPI_result; + pgrx::error!("failed to prepare df.start() quota query: SPI result {status}"); + } + + // Even REPEATABLE READ must see starts committed before we acquired the + // admission lock, plus this transaction's own earlier starts. Ordinary + // SPI would reuse the transaction snapshot and allow quota bypasses. + pg_sys::CommandCounterIncrement(); + let mut values = [pg_sys::Datum::from(user_oid)]; + let status = pg_sys::SPI_execute_snapshot( + plan, + values.as_mut_ptr(), + std::ptr::null(), + pg_sys::GetLatestSnapshot(), + std::ptr::null_mut(), + true, + false, + 1, + ); + if status != pg_sys::SPI_OK_SELECT as i32 || pg_sys::SPI_processed != 1 { + pgrx::error!("failed to count df.start() instances: SPI result {status}"); + } + let table = pg_sys::SPI_tuptable + .as_ref() + .unwrap_or_else(|| pgrx::error!("df.start() quota query returned no tuple table")); + let count = |column| { + let mut is_null = false; + let datum = pg_sys::SPI_getbinval(*table.vals, table.tupdesc, column, &mut is_null); + i64::from_datum(datum, is_null) + .unwrap_or_else(|| pgrx::error!("df.start() quota query returned a NULL count")) + }; + (count(1), count(2)) + }) +} + +fn enforce_instance_limits(user_oid: pg_sys::Oid, user_name: &str) { + let concurrent_limit = crate::MAX_CONCURRENT_PER_USER.get(); + let instance_limit = crate::MAX_INSTANCES_PER_USER.get(); + if concurrent_limit == 0 && instance_limit == 0 { + return; + } + + // Keep the lock until commit/rollback: another backend cannot count our + // uncommitted instances. A function-scoped lock would leave a race window. + Spi::run_with_args( + "SELECT pg_catalog.pg_advisory_xact_lock($1::pg_catalog.int4, $2::pg_catalog.int4)", + &[ + INSTANCE_QUOTA_LOCK_CLASS_ID.into(), + (u32::from(user_oid) as i32).into(), + ], + ) + .unwrap_or_else(|e| pgrx::error!("failed to acquire df.start() admission lock: {e}")); + + let (active, total) = current_instance_counts(user_oid); + for (limit, count, kind, hint) in [ + ( + concurrent_limit, + active, + "active", + "Wait for in-flight instances to finish or cancel them, or ask an administrator to raise df.max_concurrent_per_user.", + ), + ( + instance_limit, + total, + "retained", + "Wait for automatic reconciliation to remove terminal instances, or ask an administrator to adjust retention or raise df.max_instances_per_user.", + ), + ] { + if limit > 0 && count >= i64::from(limit) { + ErrorReport::new( + PgSqlErrorCode::ERRCODE_CONFIGURATION_LIMIT_EXCEEDED, + format!( + "df.start rejected: user \"{user_name}\" has {count} {kind} instances (limit {limit})" + ), + function_name!(), + ) + .set_hint(hint) + .report(PgLogLevel::ERROR); + } + } +} + /// `transaction_mode`: the start participates in the caller's transaction and /// is rolled back with it. This is the default and the historical behaviour. const TXN_MODE_CALLER: &str = "caller"; @@ -1243,7 +1345,10 @@ fn start_in_new_transaction(fut: &str, label: Option<&str>, database: Option<&st match crate::client::start_in_new_transaction(fut, label, database, &user_name) { Ok(id) => id, - Err(e) => pgrx::error!("{}", e), + Err(e) => { + e.report(PgLogLevel::ERROR); + unreachable!() + } } } @@ -1281,20 +1386,21 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< require_login_privilege(current_user_oid, ¤t_user_name, "df.start()"); // Reject superuser submission identities unless explicitly enabled. - if !crate::types::superuser_instances_enabled() { - let is_super = match crate::types::is_role_superuser_oid(current_user_oid) { - Ok(v) => v, - Err(e) => pgrx::error!("pg_durable: superuser check failed: {}", e), - }; - if is_super { - pgrx::error!( - "pg_durable: superuser instances are disabled. \ - current_user \"{}\" is a superuser, but \ - pg_durable.enable_superuser_instances is off. \ - Set pg_durable.enable_superuser_instances = on to allow this.", - current_user_name - ); - } + let is_super = match crate::types::is_role_superuser_oid(current_user_oid) { + Ok(v) => v, + Err(e) => pgrx::error!("pg_durable: superuser check failed: {}", e), + }; + if is_super && !crate::types::superuser_instances_enabled() { + pgrx::error!( + "pg_durable: superuser instances are disabled. \ + current_user \"{}\" is a superuser, but \ + pg_durable.enable_superuser_instances is off. \ + Set pg_durable.enable_superuser_instances = on to allow this.", + current_user_name + ); + } + if !is_super { + enforce_instance_limits(current_user_oid, ¤t_user_name); } fn insert_nodes( diff --git a/src/lib.rs b/src/lib.rs index fba8c1ef..e234d970 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -33,6 +33,8 @@ pub static MAX_ORIGIN_CONNECTIONS: GucSetting = GucSetting::::new(12); pub static MAX_DUROXIDE_CONNECTIONS: GucSetting = GucSetting::::new(10); pub static MAX_USER_CONNECTIONS: GucSetting = GucSetting::::new(10); pub static MAX_NEW_TRANSACTION_STARTS: GucSetting = GucSetting::::new(2); +pub static MAX_CONCURRENT_PER_USER: GucSetting = GucSetting::::new(100); +pub static MAX_INSTANCES_PER_USER: GucSetting = GucSetting::::new(10_000); pub static EXECUTION_ACQUIRE_TIMEOUT: GucSetting = GucSetting::::new(30); pub static NEW_TRANSACTION_START_TIMEOUT: GucSetting = GucSetting::::new(5); /// When `false` (default), pg_durable rejects any instance whose `submitted_by` @@ -288,6 +290,28 @@ pub extern "C-unwind" fn _PG_init() { GucFlags::default(), ); + GucRegistry::define_int_guc( + c"df.max_concurrent_per_user", + c"Maximum pending and running instances per submitting role in each database", + c"df.start() rejects new instances at this limit. 0 disables the limit. Superuser submission identities are exempt.", + &MAX_CONCURRENT_PER_USER, + 0, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); + + GucRegistry::define_int_guc( + c"df.max_instances_per_user", + c"Maximum retained instances per submitting role in each database", + c"Includes terminal instances until reconciliation removes them. df.start() rejects new instances at this limit. 0 disables the limit. Superuser submission identities are exempt.", + &MAX_INSTANCES_PER_USER, + 0, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( c"pg_durable.retention_days", c"Days a terminal instance is retained before reconciliation removes it", diff --git a/tests/e2e/sql/89_schema_qualification.sql b/tests/e2e/sql/89_schema_qualification.sql index 2ef03e23..2dc07657 100644 --- a/tests/e2e/sql/89_schema_qualification.sql +++ b/tests/e2e/sql/89_schema_qualification.sql @@ -63,6 +63,22 @@ CREATE FUNCTION schema_qualification.to_char(timestamptz, text) RETURNS text LAN BEGIN RAISE EXCEPTION 'schema qualification function trap'; END $$; +CREATE FUNCTION schema_qualification.lower(text) RETURNS text LANGUAGE plpgsql AS $$ +BEGIN + RAISE EXCEPTION 'schema qualification quota trap'; +END $$; +CREATE FUNCTION schema_qualification.pg_advisory_xact_lock(int4, int4) +RETURNS void LANGUAGE plpgsql AS $$ +BEGIN + RAISE EXCEPTION 'schema qualification quota lock trap'; +END $$; +CREATE FUNCTION schema_qualification.count_trap(bigint) RETURNS bigint LANGUAGE plpgsql AS $$ +BEGIN + RAISE EXCEPTION 'schema qualification aggregate trap'; +END $$; +CREATE AGGREGATE schema_qualification.count(*) ( + SFUNC = schema_qualification.count_trap, STYPE = bigint, INITCOND = '0' +); CREATE FUNCTION schema_qualification.user_value() RETURNS integer LANGUAGE sql AS $$ SELECT 42 $$; -- These also shadow casts in the separate transaction_mode => 'new' backend. diff --git a/tests/e2e/sql/90_start_limits.sql b/tests/e2e/sql/90_start_limits.sql new file mode 100644 index 00000000..d2b63f8b --- /dev/null +++ b/tests/e2e/sql/90_start_limits.sql @@ -0,0 +1,209 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +DROP SCHEMA IF EXISTS start_limits_test CASCADE; +DO $$ +DECLARE + role_name TEXT; +BEGIN + FOREACH role_name IN ARRAY ARRAY['start_limits_user', 'start_limits_new_user'] LOOP + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = role_name) THEN + EXECUTE format('DROP OWNED BY %I', role_name); + EXECUTE format('DROP ROLE %I', role_name); + END IF; + END LOOP; +END $$; +CREATE ROLE start_limits_user LOGIN; +CREATE ROLE start_limits_new_user LOGIN; +SELECT df.grant_usage('start_limits_user'); +SELECT df.grant_usage('start_limits_new_user'); +CREATE SCHEMA start_limits_test; +GRANT USAGE ON SCHEMA start_limits_test TO start_limits_user, start_limits_new_user; +CREATE TABLE start_limits_test.ids (kind TEXT, id TEXT); +GRANT SELECT, INSERT ON start_limits_test.ids TO start_limits_user, start_limits_new_user; + +CREATE FUNCTION start_limits_test.expect_rejected(kind TEXT, n BIGINT, mode TEXT DEFAULT 'caller') +RETURNS VOID LANGUAGE plpgsql AS $$ +DECLARE + hint TEXT; + instance_count BIGINT := (SELECT count(*) FROM df.instances); + node_count BIGINT := (SELECT count(*) FROM df.nodes); + guc TEXT := CASE WHEN kind = 'active' THEN 'df.max_concurrent_per_user' + ELSE 'df.max_instances_per_user' END; +BEGIN + BEGIN + PERFORM df.start('SELECT 1', 'start-limits-rejected', transaction_mode => mode); + RAISE EXCEPTION 'TEST FAILED: over-quota start succeeded'; + EXCEPTION WHEN configuration_limit_exceeded THEN + GET STACKED DIAGNOSTICS hint = PG_EXCEPTION_HINT; + IF position(format('user "%s" has %s %s instances (limit %s)', current_user, n, kind, n) + IN SQLERRM) = 0 OR position(guc IN hint) = 0 OR hint IS NULL THEN + RAISE EXCEPTION 'TEST FAILED: incorrect quota diagnostics: %, hint %', SQLERRM, hint; + END IF; + END; + IF (SELECT count(*) FROM df.instances) <> instance_count OR + (SELECT count(*) FROM df.nodes) <> node_count THEN + RAISE EXCEPTION 'TEST FAILED: rejected start persisted metadata'; + END IF; +END $$; + +DO $$ +DECLARE + setting RECORD; +BEGIN + FOR setting IN SELECT name, boot_val, context, min_val, max_val FROM pg_settings + WHERE name IN ('df.max_concurrent_per_user', 'df.max_instances_per_user') + LOOP + IF setting.context <> 'superuser' OR setting.min_val <> '0' OR + setting.max_val <> '2147483647' OR setting.boot_val <> + (CASE setting.name WHEN 'df.max_concurrent_per_user' THEN '100' ELSE '10000' END) THEN + RAISE EXCEPTION 'TEST FAILED: incorrect quota GUC definition: %', setting; + END IF; + BEGIN + PERFORM set_config(setting.name, '-1', true); + RAISE EXCEPTION 'TEST FAILED: negative quota accepted'; + EXCEPTION WHEN invalid_parameter_value THEN NULL; + END; + END LOOP; + IF (SELECT count(*) FROM pg_settings + WHERE name IN ('df.max_concurrent_per_user', 'df.max_instances_per_user')) <> 2 THEN + RAISE EXCEPTION 'TEST FAILED: quota GUC missing'; + END IF; +END $$; + +SET df.max_concurrent_per_user = 2; +SET df.max_instances_per_user = 0; +SET ROLE start_limits_user; +DO $$ +DECLARE + guc TEXT; +BEGIN + FOREACH guc IN ARRAY ARRAY['df.max_concurrent_per_user', 'df.max_instances_per_user'] LOOP + BEGIN + PERFORM set_config(guc, '0', false); + RAISE EXCEPTION 'TEST FAILED: unprivileged role changed %', guc; + EXCEPTION WHEN insufficient_privilege THEN NULL; + END; + END LOOP; +END $$; + +-- Both calls share a transaction, so the second must count the first's rows. +INSERT INTO start_limits_test.ids +SELECT 'initial', df.start(df.sleep(300), 'start-limits-initial') FROM generate_series(1, 2); +RESET ROLE; + +CREATE TEMP TABLE _quota_engine_count (n BIGINT); +DO $$ +BEGIN + EXECUTE format('INSERT INTO _quota_engine_count SELECT count(*) FROM %I.instances', + df.duroxide_schema()); +END $$; +SET ROLE start_limits_user; +SELECT start_limits_test.expect_rejected('active', 2); +RESET ROLE; +DO $$ +DECLARE + actual BIGINT; +BEGIN + EXECUTE format('SELECT count(*) FROM %I.instances', df.duroxide_schema()) INTO actual; + IF actual <> (SELECT n FROM _quota_engine_count) THEN + RAISE EXCEPTION 'TEST FAILED: rejected start created engine work'; + END IF; +END $$; + +SET ROLE start_limits_user; +DO $$ +BEGIN + PERFORM df.cancel((SELECT id FROM start_limits_test.ids WHERE kind = 'initial' LIMIT 1)); + INSERT INTO start_limits_test.ids + SELECT 'completed', df.start('SELECT 1', 'start-limits-replacement'); +END $$; +DO $$ +BEGIN + IF df.await_instance((SELECT id FROM start_limits_test.ids WHERE kind = 'completed'), 30) + IS DISTINCT FROM 'completed' THEN + RAISE EXCEPTION 'TEST FAILED: replacement did not complete'; + END IF; +END $$; +INSERT INTO start_limits_test.ids SELECT 'after-completion', df.start(df.sleep(300)); +SELECT start_limits_test.expect_rejected('active', 2); +SELECT df.cancel(id) FROM df.instances WHERE status IN ('pending', 'running'); +RESET ROLE; + +-- All four retained rows are terminal; completing/cancelling does not free storage quota. +SET df.max_concurrent_per_user = 0; +SET df.max_instances_per_user = 4; +SET ROLE start_limits_user; +SELECT start_limits_test.expect_rejected('retained', 4); +RESET ROLE; + +CREATE FUNCTION start_limits_test.start_as_owner() RETURNS TEXT +LANGUAGE sql SECURITY DEFINER SET search_path = pg_catalog, pg_temp +AS $$ SELECT df.start('SELECT 1') $$; +ALTER FUNCTION start_limits_test.start_as_owner() OWNER TO start_limits_user; +SET ROLE start_limits_new_user; +DO $$ +BEGIN + BEGIN + PERFORM start_limits_test.start_as_owner(); + RAISE EXCEPTION 'TEST FAILED: SECURITY DEFINER bypassed owner quota'; + EXCEPTION WHEN configuration_limit_exceeded THEN + IF SQLERRM NOT LIKE '%user "start_limits_user" has 4 retained instances%' THEN + RAISE; + END IF; + END; +END $$; +RESET ROLE; + +-- The loopback session uses role/database defaults, not these caller SET values. +SELECT format('ALTER ROLE start_limits_new_user IN DATABASE %I SET df.max_concurrent_per_user = 1', + current_database()) \gexec +SELECT format('ALTER ROLE start_limits_new_user IN DATABASE %I SET df.max_instances_per_user = 2', + current_database()) \gexec +SET df.max_concurrent_per_user = 0; +SET df.max_instances_per_user = 0; +SET ROLE start_limits_new_user; +INSERT INTO start_limits_test.ids +SELECT 'new-sleep', df.start(df.sleep(300), transaction_mode => 'new'); +SELECT start_limits_test.expect_rejected('active', 1, 'new'); +SELECT df.cancel(id) FROM start_limits_test.ids WHERE kind = 'new-sleep'; +INSERT INTO start_limits_test.ids +SELECT 'new-completed', df.start('SELECT 1', transaction_mode => 'new'); +DO $$ +BEGIN + IF df.await_instance((SELECT id FROM start_limits_test.ids WHERE kind = 'new-completed'), 30) + IS DISTINCT FROM 'completed' THEN + RAISE EXCEPTION 'TEST FAILED: new-transaction instance did not complete'; + END IF; +END $$; +SELECT start_limits_test.expect_rejected('retained', 2, 'new'); +RESET ROLE; + +-- Zero disables both limits, including admission locking. +SET ROLE start_limits_user; +BEGIN; +INSERT INTO start_limits_test.ids +SELECT 'unlimited', df.start(df.sleep(300)) FROM generate_series(1, 3); +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_locks WHERE pid = pg_backend_pid() + AND locktype = 'advisory' AND classid = x'50474451'::int AND objsubid = 2) THEN + RAISE EXCEPTION 'TEST FAILED: unlimited starts acquired a quota lock'; + END IF; +END $$; +COMMIT; +RESET ROLE; + +SET df.max_concurrent_per_user = 1; +SET df.max_instances_per_user = 1; +INSERT INTO start_limits_test.ids +SELECT 'superuser', df.start(df.sleep(300)) FROM generate_series(1, 3); + +SELECT df.cancel(id) FROM start_limits_test.ids; +DROP TABLE _quota_engine_count; +DROP SCHEMA start_limits_test CASCADE; +DROP OWNED BY start_limits_user, start_limits_new_user; +DROP ROLE start_limits_user, start_limits_new_user; +RESET df.max_concurrent_per_user; +RESET df.max_instances_per_user; +SELECT 'TEST PASSED' AS result; diff --git a/tests/e2e/sql/91_start_limits_concurrency.sql b/tests/e2e/sql/91_start_limits_concurrency.sql new file mode 100644 index 00000000..2ee9044d --- /dev/null +++ b/tests/e2e/sql/91_start_limits_concurrency.sql @@ -0,0 +1,162 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +DO $$ +DECLARE + role_name TEXT; +BEGIN + FOREACH role_name IN ARRAY ARRAY['start_limits_race', 'start_limits_other'] LOOP + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = role_name) THEN + EXECUTE format('DROP OWNED BY %I', role_name); + EXECUTE format('DROP ROLE %I', role_name); + END IF; + END LOOP; +END $$; +CREATE ROLE start_limits_race LOGIN; +CREATE ROLE start_limits_other LOGIN; +SELECT df.grant_usage('start_limits_race'); +SELECT df.grant_usage('start_limits_other'); +SELECT format('ALTER ROLE start_limits_race IN DATABASE %I SET df.max_concurrent_per_user = 1', + current_database()) \gexec +SELECT format('ALTER ROLE start_limits_race IN DATABASE %I SET df.max_instances_per_user = 0', + current_database()) \gexec + +CREATE FUNCTION pg_temp.quota_wait_for_lock(waiter INTEGER) RETURNS VOID LANGUAGE plpgsql AS $$ +BEGIN + FOR attempt IN 1..300 LOOP + IF EXISTS (SELECT 1 FROM pg_locks WHERE pid = waiter AND locktype = 'advisory' + AND classid = x'50474451'::int AND objsubid = 2 AND NOT granted) THEN + RETURN; + END IF; + PERFORM pg_sleep(0.1); + END LOOP; + RAISE EXCEPTION 'TEST FAILED: concurrent start never waited for admission'; +END $$; + +CREATE FUNCTION pg_temp.quota_finish(conn TEXT, rejected BOOLEAN) RETURNS TEXT LANGUAGE plpgsql AS $$ +DECLARE + instance TEXT; + message TEXT; +BEGIN + FOR attempt IN 1..300 LOOP + EXIT WHEN dblink_is_busy(conn) = 0; + PERFORM pg_sleep(0.1); + END LOOP; + IF dblink_is_busy(conn) <> 0 THEN + RAISE EXCEPTION 'TEST FAILED: concurrent start never finished'; + END IF; + SELECT id INTO instance FROM dblink_get_result(conn, false) AS r(id TEXT); + message := dblink_error_message(conn); + PERFORM * FROM dblink_get_result(conn, false) AS r(id TEXT); + IF rejected THEN + IF message NOT LIKE '%has 1 active instances (limit 1)%' OR instance IS NOT NULL THEN + RAISE EXCEPTION 'TEST FAILED: concurrent start did not reject: %', message; + END IF; + ELSIF message <> 'OK' OR instance IS NULL OR length(instance) <> 8 THEN + RAISE EXCEPTION 'TEST FAILED: concurrent start did not succeed: %', message; + END IF; + RETURN instance; +END $$; + +DO $$ +DECLARE + connstr TEXT := format('host=localhost dbname=%L port=%s user=start_limits_race', + current_database(), current_setting('port')); + waiter INTEGER; + isolation TEXT; + rollback_command TEXT; + instance TEXT; + message TEXT; +BEGIN + PERFORM dblink_connect('quota_holder', connstr); + PERFORM dblink_connect('quota_waiter', connstr); + PERFORM dblink_exec('quota_holder', 'SET statement_timeout = ''15s'''); + PERFORM dblink_exec('quota_waiter', 'SET statement_timeout = ''15s'''); + SELECT pid INTO waiter FROM dblink('quota_waiter', 'SELECT pg_backend_pid()') AS r(pid INTEGER); + + FOREACH isolation IN ARRAY ARRAY['READ COMMITTED', 'REPEATABLE READ', 'SERIALIZABLE'] LOOP + PERFORM dblink_exec('quota_holder', $sql$ + DO $cancel$ BEGIN + PERFORM df.cancel(id) FROM df.instances WHERE status IN ('pending', 'running'); + END $cancel$; + $sql$); + PERFORM dblink_exec('quota_waiter', 'BEGIN ISOLATION LEVEL ' || isolation); + -- Establish the waiter's snapshot before the other transaction inserts. + PERFORM * FROM dblink('quota_waiter', 'SELECT count(*) FROM df.instances') AS r(n BIGINT); + PERFORM dblink_exec('quota_holder', 'BEGIN'); + PERFORM * FROM dblink('quota_holder', + $sql$SELECT df.start(df.sleep(300), 'start-limits-holder')$sql$) AS r(id TEXT); + IF dblink_send_query('quota_waiter', + $sql$SELECT df.start(df.sleep(300), 'start-limits-waiter')$sql$) <> 1 THEN + RAISE EXCEPTION 'TEST FAILED: could not send concurrent start'; + END IF; + PERFORM pg_temp.quota_wait_for_lock(waiter); + + IF isolation = 'READ COMMITTED' THEN + PERFORM dblink_connect('quota_other', format( + 'host=localhost dbname=%L port=%s user=start_limits_other', + current_database(), current_setting('port'))); + PERFORM dblink_exec('quota_other', 'SET statement_timeout = ''5s'''); + PERFORM * FROM dblink('quota_other', + 'SELECT df.start(df.sleep(300))') AS r(id TEXT); + PERFORM dblink_disconnect('quota_other'); + END IF; + + PERFORM dblink_exec('quota_holder', 'COMMIT'); + PERFORM pg_temp.quota_finish('quota_waiter', true); + PERFORM dblink_exec('quota_waiter', 'ROLLBACK'); + IF (SELECT count(*) FROM df.instances WHERE submitted_by = 'start_limits_race'::regrole + AND status IN ('pending', 'running')) <> 1 THEN + RAISE EXCEPTION 'TEST FAILED: concurrent starts exceeded the cap under %', isolation; + END IF; + END LOOP; + + FOREACH rollback_command IN ARRAY ARRAY['ROLLBACK', 'ROLLBACK TO SAVEPOINT quota_start'] LOOP + PERFORM dblink_exec('quota_holder', $sql$ + DO $cancel$ BEGIN + PERFORM df.cancel(id) FROM df.instances WHERE status IN ('pending', 'running'); + END $cancel$; + $sql$); + PERFORM dblink_exec('quota_holder', 'BEGIN'); + PERFORM dblink_exec('quota_holder', 'SAVEPOINT quota_start'); + SELECT id INTO instance FROM dblink('quota_holder', + 'SELECT df.start(df.sleep(300))') AS r(id TEXT); + IF dblink_send_query('quota_waiter', 'SELECT df.start(df.sleep(300))') <> 1 THEN + RAISE EXCEPTION 'TEST FAILED: could not send rollback waiter'; + END IF; + PERFORM pg_temp.quota_wait_for_lock(waiter); + PERFORM dblink_exec('quota_holder', rollback_command); + PERFORM pg_temp.quota_finish('quota_waiter', false); + IF rollback_command <> 'ROLLBACK' THEN + PERFORM dblink_exec('quota_holder', 'COMMIT'); + END IF; + IF EXISTS (SELECT 1 FROM df.instances WHERE id = instance) OR + EXISTS (SELECT 1 FROM df.nodes WHERE instance_id = instance) THEN + RAISE EXCEPTION 'TEST FAILED: rolled-back start retained metadata'; + END IF; + END LOOP; + + PERFORM dblink_exec('quota_holder', $sql$ + DO $cancel$ BEGIN + PERFORM df.cancel(id) FROM df.instances WHERE status IN ('pending', 'running'); + END $cancel$; + $sql$); + PERFORM dblink_exec('quota_holder', 'BEGIN'); + PERFORM * FROM dblink('quota_holder', 'SELECT df.start(df.sleep(300))') AS r(id TEXT); + PERFORM dblink_exec('quota_waiter', 'SET lock_timeout = ''100ms'''); + PERFORM * FROM dblink('quota_waiter', 'SELECT df.start(df.sleep(300))', false) AS r(id TEXT); + message := dblink_error_message('quota_waiter'); + IF message NOT LIKE '%lock timeout%' THEN + RAISE EXCEPTION 'TEST FAILED: admission ignored lock_timeout: %', message; + END IF; + PERFORM dblink_exec('quota_holder', 'ROLLBACK'); + PERFORM dblink_disconnect('quota_holder'); + PERFORM dblink_disconnect('quota_waiter'); +END $$; + +SELECT df.cancel(id) FROM df.instances +WHERE submitted_by IN ('start_limits_race'::regrole, 'start_limits_other'::regrole); +DROP FUNCTION pg_temp.quota_wait_for_lock(INTEGER), pg_temp.quota_finish(TEXT, BOOLEAN); +DROP OWNED BY start_limits_race, start_limits_other; +DROP ROLE start_limits_race, start_limits_other; +SELECT 'TEST PASSED' AS result; From 98b89885dfc466b0470a8f051f49e6ef834967aa Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Wed, 28 Oct 2026 14:20:08 +0000 Subject: [PATCH 2/5] Protect engine accounting, and other fixes from review --- CHANGELOG.md | 32 +++- Cargo.lock | 2 +- Cargo.toml | 2 +- USER_GUIDE.md | 36 +++- docs/api-reference.md | 4 +- docs/upgrade-testing.md | 43 +++-- scripts/run-pgspot.sh | 1 - scripts/test-upgrade.sh | 1 + sql/pg_durable--0.2.9--0.2.10.sql | 4 + src/client.rs | 114 ++++++++++-- src/dsl.rs | 103 ++++++----- src/lib.rs | 7 +- tests/e2e/sql/00_setup_playground.sql | 18 +- tests/e2e/sql/89_schema_qualification.sql | 17 +- tests/e2e/sql/90_start_limits.sql | 25 ++- tests/e2e/sql/91_start_limits_concurrency.sql | 5 +- tests/e2e/sql/92_start_limits_trust.sql | 166 ++++++++++++++++++ .../e2e/sql/93_start_limits_serializable.sql | 102 +++++++++++ 18 files changed, 572 insertions(+), 110 deletions(-) create mode 100644 sql/pg_durable--0.2.9--0.2.10.sql create mode 100644 tests/e2e/sql/92_start_limits_trust.sql create mode 100644 tests/e2e/sql/93_start_limits_serializable.sql diff --git a/CHANGELOG.md b/CHANGELOG.md index db8a2728..03c53708 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,31 @@ All notable changes to this project are documented in this file. The format is b Pre-1.0 note: while `pg_durable` is in major version `0`, minor releases may include breaking changes. +## [Unreleased] + +### Security + +- **Start admission quotas (#139):** `df.max_concurrent_per_user` (default + `100`) and `df.max_instances_per_user` (default `10000`) limit active and + retained instances per submitting role and origin database. Both are + superuser-settable; `0` disables a quota and superuser identities are exempt. + Active accounting uses protected engine state, not caller-writable display + status. Cancellation frees active capacity after the engine processes it. + Starts serialize per role until the submitting transaction ends and reject + over-quota work before insert/enqueue, including on older installed schemas. + +> **Upgrade:** These defaults apply after a binary-only upgrade. Existing +> workflows whose SQL starts more workflows also encounter the quotas, and an +> over-quota SQL step can fail its enclosing workflow. Completed and cancelled +> instances still consume retained capacity. With the default 30-day retention, +> a role holding 10,000 recent instances may remain blocked until history ages +> out or an administrator changes policy; a reconciliation pass alone need not +> free capacity. See [Instance Limits](USER_GUIDE.md#instance-limits). +> Run `ALTER EXTENSION pg_durable UPDATE TO '0.2.10'` to add the owner-leading +> index used by quota checks. The new binary still enforces quotas on older +> schemas; without the update their local quota queries can scan other roles' +> history. + ## [0.2.9] - 2026-10-07 > **Upgrade:** Install/restart the new binary and wait for the control runtime @@ -71,13 +96,6 @@ Pre-1.0 note: while `pg_durable` is in major version `0`, minor releases may inc ### Security -- **Start admission quotas (#139):** `df.max_concurrent_per_user` (default - `100`) and `df.max_instances_per_user` (default `10000`) limit active and - retained instances per submitting role and origin database. Both are - superuser-settable; `0` disables a quota and superuser identities are exempt. - Starts serialize per role until the submitting transaction ends and reject - over-quota work before insert/enqueue, including on older installed schemas. - See [Instance Limits](USER_GUIDE.md#instance-limits) before upgrading. - **SQL name-resolution hardening (#418):** extension-owned runtime SQL explicitly qualifies catalog objects, functions, types and operators to resist shadowing through caller or database `search_path` settings. User-authored SQL diff --git a/Cargo.lock b/Cargo.lock index c6bd4d84..e821bc7a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1924,7 +1924,7 @@ dependencies = [ [[package]] name = "pg_durable" -version = "0.2.9" +version = "0.2.10" dependencies = [ "base64 0.23.1", "bigdecimal", diff --git a/Cargo.toml b/Cargo.toml index ba9cd10a..c6a7eee5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pg_durable" -version = "0.2.9" +version = "0.2.10" edition = "2021" license = "PostgreSQL" repository = "https://github.com/microsoft/pg_durable" diff --git a/USER_GUIDE.md b/USER_GUIDE.md index 6d649c8f..8511e3b0 100644 --- a/USER_GUIDE.md +++ b/USER_GUIDE.md @@ -3010,7 +3010,7 @@ for the full syntax and default domain list. | Setting | Default | Counts | |---------|---------|--------| -| `df.max_concurrent_per_user` | `100` | Instances with status `pending` or `running`, including sleeping, looping, and signal-waiting workflows | +| `df.max_concurrent_per_user` | `100` | Instances not yet terminal in the durable engine, including queued, sleeping, looping, and signal-waiting workflows | | `df.max_instances_per_user` | `10000` | All retained instances, including completed, failed, and cancelled workflows | Each quota applies to the **submitting role (`current_user`) in the origin @@ -3036,12 +3036,22 @@ session overrides. Configure those defaults to apply the same policy to both modes; `SET ROLE` alone does not load a role's connection defaults. At a quota, the call fails with SQLSTATE `53400` (`configuration_limit_exceeded`) -and a `HINT` identifying the relevant setting. Finishing or cancelling work -frees concurrent capacity, but not retained-instance capacity. Retained capacity -is freed when [automatic reconciliation](#automatic-reconciliation) removes +and a `HINT` identifying the relevant setting. Finishing work frees concurrent +capacity once the engine records the terminal state. Cancellation releases it +after the engine processes the request, not immediately when `df.cancel()` +returns. Changing the displayed instance status does not release capacity. +Neither completion nor cancellation frees retained-instance capacity. Retained +capacity is freed when [automatic reconciliation](#automatic-reconciliation) removes terminal instances. If reconciliation is disabled, or retained work has not aged out, an administrator must adjust retention or raise the quota. +The default retained quota equals reconciliation's 10,000-terminal-instance +retention cap. A role with 10,000 recent instances can therefore remain blocked +until its history ages out under `pg_durable.retention_days` (30 days by default); +running reconciliation alone need not free a slot. Tune the quota for the +expected submission volume and retention window before deploying. Retention +settings require a restart; the quota itself can be changed at runtime. + Concurrent starts by the same role in the same database wait for the preceding **submitting transaction** to commit or roll back before checking capacity. They do not wait for workflows to finish. Keep submitting transactions short @@ -3050,6 +3060,13 @@ current committed rows and the caller's own changes even under `REPEATABLE READ` or `SERIALIZABLE`. Rollback, including rollback to a savepoint, releases the corresponding reservations. +Perform business-table updates before starting workflows when possible, and +avoid leaving submitting transactions idle. Consider PostgreSQL +`idle_in_transaction_session_timeout` as well as `lock_timeout`. Applications +must retry entire transactions on deadlock or serialization failures; admission +adds another lock/read to transactions, including workflow SQL that starts +other workflows. Both quotas must be `0` to disable admission locking. + Do not follow a `'caller'` start with a `'new'` start for the same role and database while the caller transaction remains open: the separate session waits on the caller's admission lock and eventually hits its 30-second timeout. @@ -3059,6 +3076,17 @@ These are instance quotas, not calls-per-second or byte limits. They do not bound an individual graph's size or history. Changing a quota does not cancel or delete existing work; it only controls admission of new instances. The defaults also apply after a binary-only upgrade, without `ALTER EXTENSION UPDATE`. +Existing workflows whose SQL calls `df.start()` also encounter the new limits; +an over-quota SQL step can fail its enclosing workflow. + +Fresh 0.2.10 installations include an owner-leading index for quota checks. +Run `ALTER EXTENSION pg_durable UPDATE TO '0.2.10'` to add it on an existing +installation. Binary-only upgrades still enforce quotas, but the local count +can scan other roles' history until the index is installed. Active accounting +also checks the role's retained IDs against the engine, so larger histories +increase admission cost even with the index. Measure representative workloads +before substantially raising retention or quotas. Disabling the active quota +skips the engine lookup; disabling both quotas skips admission entirely. --- diff --git a/docs/api-reference.md b/docs/api-reference.md index d0345e05..990317a4 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -907,12 +907,14 @@ Admission quotas for `df.start()`, per submitting role and origin database. | Property | Concurrent quota | Retained-instance quota | |----------|------------------|-------------------------| | Default | `100` | `10000` | -| Counts | `pending` and `running` instances | All `df.instances` rows owned by the role | +| Counts | Queued instances and executions not terminal in the protected engine state | All `df.instances` rows owned by the role | | Range | `0` to `2147483647`; `0` disables | `0` to `2147483647`; `0` disables | | Context | `SUSET` | `SUSET` | Superuser submission identities are exempt. Administrators can change either setting at runtime, including per-role/database connection defaults. +Cancellation frees active capacity after the engine processes the request; +changing display status alone does not release it. `transaction_mode => 'new'` uses its separate session's defaults and does not inherit the caller's session `SET` overrides. See [Instance Limits](../USER_GUIDE.md#instance-limits) for tuning and admission waits. diff --git a/docs/upgrade-testing.md b/docs/upgrade-testing.md index c68bdac7..0c788c7b 100644 --- a/docs/upgrade-testing.md +++ b/docs/upgrade-testing.md @@ -356,33 +356,52 @@ gate, so they never need to be added to the exclude list. Each schema-changing PR should add a section here documenting what changed, what the upgrade script handles, and any backward compatibility considerations. -### v0.2.8 → v0.2.9 - -The post-tag changes in #379, #388, #389, #390, and #380 belong to the 0.2.9 -development cycle, not the published v0.2.8 release. Multi-database installation -identity and HTTP endpoint/secret DDL also belong in this migration; -the shipped 0.2.7 to 0.2.8 script must remain byte-identical. +### v0.2.9 → v0.2.10 (unreleased) #### `df.start()` admission quotas (#139) -- **DDL/B1:** No new schema objects, grants, or migration DDL. Both start - wrapper symbols use the same admission check against existing `df.instances` - columns, including on 0.2.2+ schemas without `ALTER EXTENSION UPDATE`. +- **DDL/B1:** The upgrade adds `idx_instances_submitted_by` on + `df.instances(submitted_by, id)`, matching fresh installation. No new tables + or grants are needed. Both start wrapper symbols use the same admission + check against existing `df.instances` columns and protected engine state, + including on 0.2.2+ schemas without + `ALTER EXTENSION UPDATE`. No status grants are revoked or used as an + accounting authority; queued or not-yet-materialized engine instances count + as active. Cancellation releases capacity only after the engine processes it. B1 exercises a non-superuser reaching its quota on each supported schema. + Old schemas remain correct without the index but can scan other roles' + history. The normal index build blocks concurrent writes; set `lock_timeout` + and schedule the upgrade appropriately on large installations. - **Behavior change:** New superuser-settable GUCs `df.max_concurrent_per_user` (default `100`) and `df.max_instances_per_user` (default `10000`) apply immediately after loading the new binary. Existing work remains intact; starts at or above either quota fail with SQLSTATE `53400`. Set a quota to `0` to disable it. - **Concurrency:** Admission uses a database-local, per-role transaction - advisory lock and a current snapshot, including under repeatable-read and - serializable isolation. E2E tests cover concurrent commit, rollback, - savepoint rollback, and lock timeout. + lock with a tag that SQL advisory-lock functions cannot construct, and a + current snapshot, including under repeatable-read and serializable isolation. + E2E tests cover concurrent commit, rollback, savepoint rollback, lock timeout, + caller-writable status, and attempts to squat on public advisory-lock keys. - **B2/replay:** No data migration, runtime schema detection, durable payload, or orchestration changes. Workflows that call `df.start()` from SQL are subject to the new quotas when that SQL activity executes. The separate `'new'` session uses its own configuration defaults, not caller session overrides. See [Instance Limits](../USER_GUIDE.md#instance-limits). +- **Scale validation:** Warm-cache PG17 measurements on a development server, + with 500,000 synthetic retained instances distributed across 100 roles, + reduced the median local ID/count query from about 35 ms to 19 ms with the + owner-leading index; retained-only counting fell from about 30 ms to 11 ms. + The protected-engine batch for 5,000 owned IDs took about 244 ms in that run. + These are query measurements, not throughput guarantees. The index reduces + local cross-role scanning, not the per-owner engine-lookup cost. Only retained + counting runs when the active quota is disabled; both disabled skips admission. + +### v0.2.8 → v0.2.9 + +The post-tag changes in #379, #388, #389, #390, and #380 belong to the 0.2.9 +development cycle, not the published v0.2.8 release. Multi-database installation +identity and HTTP endpoint/secret DDL also belong in this migration; +the shipped 0.2.7 to 0.2.8 script must remain byte-identical. #### Multi-database installation diff --git a/scripts/run-pgspot.sh b/scripts/run-pgspot.sh index cfe19230..2f241401 100755 --- a/scripts/run-pgspot.sh +++ b/scripts/run-pgspot.sh @@ -41,7 +41,6 @@ PGSPOT_ALLOW=( # pg_durable's DSL intentionally exposes unqualified custom operators (for # example, `df.sql(...) ~> df.sql(...)`) so users do not need df in search_path. # pgspot reports the generated CREATE OPERATOR name as an unqualified object. - # Match all seven DSL operators explicitly; other unqualified objects still fail. '^PS017: Unqualified object reference: ~> at line [0-9]+$' '^PS017: Unqualified object reference: \|=> at line [0-9]+$' '^PS017: Unqualified object reference: & at line [0-9]+$' diff --git a/scripts/test-upgrade.sh b/scripts/test-upgrade.sh index 5d7ebb47..27a562c1 100755 --- a/scripts/test-upgrade.sh +++ b/scripts/test-upgrade.sh @@ -963,6 +963,7 @@ test_b1_start_limits() { DO \$\$ BEGIN PERFORM df.start('SELECT 1'); + UPDATE df.instances SET status = 'completed'; BEGIN PERFORM df.start('SELECT 1'); RAISE EXCEPTION 'legacy df.start bypassed concurrency quota'; diff --git a/sql/pg_durable--0.2.9--0.2.10.sql b/sql/pg_durable--0.2.9--0.2.10.sql new file mode 100644 index 00000000..e6bc733c --- /dev/null +++ b/sql/pg_durable--0.2.9--0.2.10.sql @@ -0,0 +1,4 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +CREATE INDEX idx_instances_submitted_by ON df.instances(submitted_by, id); diff --git a/src/client.rs b/src/client.rs index ed6da7d6..5839bbe0 100644 --- a/src/client.rs +++ b/src/client.rs @@ -10,7 +10,7 @@ //! recover from connection failures by re-creating the pool on next call. use std::cell::RefCell; -use std::sync::OnceLock; +use std::sync::{Arc, OnceLock}; use std::time::Duration; use duroxide::Client; @@ -27,12 +27,19 @@ use crate::types::{ /// Cached tokio runtime for client operations. static CLIENT_RUNTIME: OnceLock = OnceLock::new(); +struct BackendClient { + connection_string: String, + control: BackendControlState, + provider: Arc, + client: Client, +} + // Per-backend cached Duroxide client. Uses thread_local + RefCell because // PostgreSQL backends are single-threaded forked processes. This allows // the client to be reset on connection failures (unlike OnceLock which // is permanent). thread_local! { - static DUROXIDE_CLIENT: RefCell> = const { RefCell::new(None) }; + static DUROXIDE_CLIENT: RefCell> = const { RefCell::new(None) }; static CONTROL_CONNECTION: RefCell> = const { RefCell::new(None) }; } @@ -99,9 +106,9 @@ fn get_client_runtime() -> &'static Runtime { /// If the client doesn't exist yet, creates it. If `f` returns an error /// that looks like a connection failure, resets the client so the next /// call will re-initialize. -fn with_duroxide_client(f: F) -> Result +fn with_backend_client(f: F) -> Result where - F: FnOnce(&Client, &Runtime) -> Result, + F: FnOnce(&BackendClient, &Runtime) -> Result, { let pg_conn_str = postgres_connection_string(); let control = match backend_control_state(&pg_conn_str) { @@ -116,10 +123,9 @@ where }; DUROXIDE_CLIENT.with(|cell| { let mut cached = cell.borrow_mut(); - if cached - .as_ref() - .is_some_and(|(url, previous, _)| url != &pg_conn_str || previous != &control) - { + if cached.as_ref().is_some_and(|cached| { + cached.connection_string != pg_conn_str || cached.control != control + }) { *cached = None; } }); @@ -131,9 +137,9 @@ where if !has_client { // Need to create a new client let schema = control.schema; - let client = rt.block_on(async { + let provider = rt.block_on(async { // Limit backend provider to 1 connection — backends need minimal - // duroxide access (start/cancel/signal only). + // duroxide access (start/cancel/signal and quota checks). // // SAFETY: Each PostgreSQL backend is a separate process (fork model). // This code runs in a single-threaded tokio runtime with no worker @@ -142,19 +148,23 @@ where std::env::set_var("DUROXIDE_PG_POOL_MAX", "1"); } - let store = new_backend_provider(&pg_conn_str, schema).await?; - Ok::(Client::new(store)) + new_backend_provider(&pg_conn_str, schema).await })?; DUROXIDE_CLIENT.with(|cell| { - *cell.borrow_mut() = Some((pg_conn_str, control, client)); + *cell.borrow_mut() = Some(BackendClient { + connection_string: pg_conn_str, + control, + client: Client::new(provider.clone()), + provider, + }); }); } // Execute the operation with the client let result = DUROXIDE_CLIENT.with(|cell| { let borrow = cell.borrow(); - let (_, _, client) = borrow + let client = borrow .as_ref() .ok_or_else(|| "Client unexpectedly missing".to_string())?; f(client, rt) @@ -172,6 +182,69 @@ where result } +pub(crate) fn count_active_instances(ids: &[String]) -> Result { + if ids.is_empty() { + return Ok(0); + } + let origin = crate::origin::backend_origin()?; + let engine_ids: Vec = ids + .iter() + .map(|id| crate::origin::engine_id_for_origin(origin.as_ref(), id)) + .collect(); + + with_backend_client(|backend, rt| { + rt.block_on(async { + let mut tx = backend + .provider + .pool() + .begin() + .await + .map_err(|e| format!("Failed to begin df.start() quota lookup: {e}"))?; + let result = async { + crate::origin::configure_metadata_transaction(&mut tx) + .await + .map_err(|e| format!("Failed to configure df.start() quota lookup: {e}"))?; + + // Missing engine rows include queued starts and uncommitted graphs. + // Only the engine's protected terminal state can release capacity. + sqlx::query_scalar(&format!( + "SELECT pg_catalog.count(*) + FROM pg_catalog.unnest($1::pg_catalog.text[]) AS requested(id) + LEFT JOIN {schema}.instances AS i + ON i.instance_id OPERATOR(pg_catalog.=) requested.id + LEFT JOIN {schema}.executions AS e + ON e.instance_id OPERATOR(pg_catalog.=) i.instance_id + AND e.execution_id OPERATOR(pg_catalog.=) i.current_execution_id + WHERE e.status IS NULL OR e.status OPERATOR(pg_catalog.<>) ALL ( + ARRAY['Completed', 'Failed']::pg_catalog.text[])", + schema = backend.control.schema, + )) + .bind(&engine_ids) + .fetch_one(&mut *tx) + .await + .map_err(|e| format!("Failed to read df.start() active instances: {e}")) + } + .await; + match result { + Ok(active) => { + tx.commit() + .await + .map_err(|e| format!("Failed to finish df.start() quota lookup: {e}"))?; + Ok(active) + } + Err(error) => { + // This backend's runtime stops polling when block_on returns. + // Finish rollback now instead of leaving it to async Drop. + tx.rollback() + .await + .map_err(|e| format!("{error}; quota lookup rollback failed: {e}"))?; + Err(error) + } + } + }) + }) +} + /// Heuristic to detect connection-level errors that warrant client reset. fn is_connection_error(err: &str) -> bool { let lower = err.to_lowercase(); @@ -239,9 +312,10 @@ pub fn start_durable_function( let inst_id = crate::origin::backend_engine_id(instance_id)?; let inp = input.to_string(); - with_duroxide_client(|client, rt| { + with_backend_client(|backend, rt| { rt.block_on(async { - client + backend + .client .start_orchestration(&inst_id, &fn_name, &inp) .await .map_err(|e| format!("Failed to start durable function: {e:?}"))?; @@ -438,9 +512,10 @@ pub fn cancel_durable_function(instance_id: &str, reason: &str) -> Result<(), St let inst_id = crate::origin::backend_engine_id(instance_id)?; let rsn = reason.to_string(); - with_duroxide_client(|client, rt| { + with_backend_client(|backend, rt| { rt.block_on(async { - client + backend + .client .cancel_instance(&inst_id, &rsn) .await .map_err(|e| format!("Failed to cancel durable function: {e:?}"))?; @@ -455,8 +530,9 @@ pub fn raise_external_event(instance_id: &str, event_name: &str, data: &str) -> let evt_name = event_name.to_string(); let evt_data = data.to_string(); - with_duroxide_client(|client, rt| { + with_backend_client(|backend, rt| { rt.block_on(async { + let client = &backend.client; client .raise_event(&inst_id, &evt_name, &evt_data) .await diff --git a/src/dsl.rs b/src/dsl.rs index 2bb522d2..aa72660b 100644 --- a/src/dsl.rs +++ b/src/dsl.rs @@ -1063,22 +1063,37 @@ fn require_login_privilege(oid: pgrx::pg_sys::Oid, name: &str, caller: &str) { } } -// "PGDQ": separate from the loopback-launch slots. Advisory locks are -// database-local; the second key preserves all bits of the submitting role OID. -const INSTANCE_QUOTA_LOCK_CLASS_ID: i32 = 0x5047_4451; +const INSTANCE_QUOTA_LOCK_CLASS_ID: u32 = 0x5047_4451; + +fn acquire_instance_quota_lock(user_oid: pg_sys::Oid) { + // SQL advisory-lock functions can only construct sub-IDs 1 and 2. Using 3 + // keeps other roles from locking a victim's admission key directly. + let tag = pg_sys::LOCKTAG { + locktag_field1: unsafe { pg_sys::MyDatabaseId.to_u32() }, + locktag_field2: INSTANCE_QUOTA_LOCK_CLASS_ID, + locktag_field3: user_oid.to_u32(), + locktag_field4: 3, + locktag_type: pg_sys::LockTagType::LOCKTAG_ADVISORY as u8, + locktag_lockmethodid: pg_sys::USER_LOCKMETHOD as u8, + }; + let result = unsafe { pg_sys::LockAcquire(&tag, pg_sys::ExclusiveLock as i32, false, false) }; + if result == pg_sys::LockAcquireResult::LOCKACQUIRE_NOT_AVAIL { + pgrx::error!("failed to acquire df.start() admission lock"); + } +} -fn current_instance_counts(user_oid: pg_sys::Oid) -> (i64, i64) { +fn current_instance_quota_rows(user_oid: pg_sys::Oid, include_ids: bool) -> (i64, Vec) { Spi::connect(|_| unsafe { - let mut argtypes = [pg_sys::OIDOID]; + let mut argtypes = [pg_sys::OIDOID, pg_sys::BOOLOID]; let plan = pg_sys::SPI_prepare( - c"SELECT pg_catalog.count(*) FILTER ( - WHERE pg_catalog.lower(status) OPERATOR(pg_catalog.=) - ANY (ARRAY['pending', 'running']::pg_catalog.text[])), - pg_catalog.count(*) + c"SELECT pg_catalog.count(*), + CASE WHEN $2 THEN COALESCE(pg_catalog.array_agg(id::pg_catalog.text), + ARRAY[]::pg_catalog.text[]) + ELSE ARRAY[]::pg_catalog.text[] END FROM df.instances WHERE submitted_by OPERATOR(pg_catalog.=) $1::pg_catalog.regrole" .as_ptr(), - 1, + 2, argtypes.as_mut_ptr(), ); if plan.is_null() { @@ -1090,7 +1105,10 @@ fn current_instance_counts(user_oid: pg_sys::Oid) -> (i64, i64) { // admission lock, plus this transaction's own earlier starts. Ordinary // SPI would reuse the transaction snapshot and allow quota bypasses. pg_sys::CommandCounterIncrement(); - let mut values = [pg_sys::Datum::from(user_oid)]; + let mut values = [ + pg_sys::Datum::from(user_oid), + pg_sys::Datum::from(include_ids), + ]; let status = pg_sys::SPI_execute_snapshot( plan, values.as_mut_ptr(), @@ -1107,13 +1125,14 @@ fn current_instance_counts(user_oid: pg_sys::Oid) -> (i64, i64) { let table = pg_sys::SPI_tuptable .as_ref() .unwrap_or_else(|| pgrx::error!("df.start() quota query returned no tuple table")); - let count = |column| { - let mut is_null = false; - let datum = pg_sys::SPI_getbinval(*table.vals, table.tupdesc, column, &mut is_null); - i64::from_datum(datum, is_null) - .unwrap_or_else(|| pgrx::error!("df.start() quota query returned a NULL count")) - }; - (count(1), count(2)) + let mut is_null = false; + let datum = pg_sys::SPI_getbinval(*table.vals, table.tupdesc, 1, &mut is_null); + let total = i64::from_datum(datum, is_null) + .unwrap_or_else(|| pgrx::error!("df.start() quota query returned a NULL count")); + let datum = pg_sys::SPI_getbinval(*table.vals, table.tupdesc, 2, &mut is_null); + let ids = Vec::::from_datum(datum, is_null) + .unwrap_or_else(|| pgrx::error!("df.start() quota query returned NULL instance IDs")); + (total, ids) }) } @@ -1126,32 +1145,11 @@ fn enforce_instance_limits(user_oid: pg_sys::Oid, user_name: &str) { // Keep the lock until commit/rollback: another backend cannot count our // uncommitted instances. A function-scoped lock would leave a race window. - Spi::run_with_args( - "SELECT pg_catalog.pg_advisory_xact_lock($1::pg_catalog.int4, $2::pg_catalog.int4)", - &[ - INSTANCE_QUOTA_LOCK_CLASS_ID.into(), - (u32::from(user_oid) as i32).into(), - ], - ) - .unwrap_or_else(|e| pgrx::error!("failed to acquire df.start() admission lock: {e}")); + acquire_instance_quota_lock(user_oid); - let (active, total) = current_instance_counts(user_oid); - for (limit, count, kind, hint) in [ - ( - concurrent_limit, - active, - "active", - "Wait for in-flight instances to finish or cancel them, or ask an administrator to raise df.max_concurrent_per_user.", - ), - ( - instance_limit, - total, - "retained", - "Wait for automatic reconciliation to remove terminal instances, or ask an administrator to adjust retention or raise df.max_instances_per_user.", - ), - ] { - if limit > 0 && count >= i64::from(limit) { - ErrorReport::new( + let (total, ids) = current_instance_quota_rows(user_oid, concurrent_limit > 0); + let reject = |count, kind, limit, hint| { + ErrorReport::new( PgSqlErrorCode::ERRCODE_CONFIGURATION_LIMIT_EXCEEDED, format!( "df.start rejected: user \"{user_name}\" has {count} {kind} instances (limit {limit})" @@ -1160,6 +1158,25 @@ fn enforce_instance_limits(user_oid: pg_sys::Oid, user_name: &str) { ) .set_hint(hint) .report(PgLogLevel::ERROR); + }; + if instance_limit > 0 && total >= i64::from(instance_limit) { + reject( + total, + "retained", + instance_limit, + "Wait for terminal history to age out and be reconciled, or ask an administrator to adjust retention or raise df.max_instances_per_user. Completion or cancellation alone does not free retained capacity.", + ); + } + if concurrent_limit > 0 { + let active = crate::client::count_active_instances(&ids) + .unwrap_or_else(|e| pgrx::error!("df.start() quota check failed: {e}")); + if active >= i64::from(concurrent_limit) { + reject( + active, + "active", + concurrent_limit, + "Wait for in-flight instances to finish or for cancellation to be processed, or ask an administrator to raise df.max_concurrent_per_user.", + ); } } } diff --git a/src/lib.rs b/src/lib.rs index e234d970..2811d144 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -292,8 +292,8 @@ pub extern "C-unwind" fn _PG_init() { GucRegistry::define_int_guc( c"df.max_concurrent_per_user", - c"Maximum pending and running instances per submitting role in each database", - c"df.start() rejects new instances at this limit. 0 disables the limit. Superuser submission identities are exempt.", + c"Maximum nonterminal durable instances per submitting role in each database", + c"df.start() checks protected engine state, including queued starts. Cancellation frees capacity after the engine processes it. 0 disables the limit. Superuser submission identities are exempt.", &MAX_CONCURRENT_PER_USER, 0, i32::MAX, @@ -473,6 +473,9 @@ CREATE INDEX idx_instances_created_at ON df.instances(created_at DESC, id); -- pg_get_indexdef() across the fresh-install and upgrade paths. CREATE INDEX idx_instances_label ON df.instances(label, created_at DESC, id) WHERE label IS NOT NULL; +-- Admission reads each role's retained IDs, independently of display status. +CREATE INDEX idx_instances_submitted_by ON df.instances(submitted_by, id); + -- Index for finding nodes by instance CREATE INDEX idx_nodes_instance ON df.nodes(instance_id); diff --git a/tests/e2e/sql/00_setup_playground.sql b/tests/e2e/sql/00_setup_playground.sql index fb893505..6d990ad7 100644 --- a/tests/e2e/sql/00_setup_playground.sql +++ b/tests/e2e/sql/00_setup_playground.sql @@ -95,6 +95,23 @@ BEGIN END LOOP; END $$; +CREATE OR REPLACE FUNCTION public._e2e_wait_for_engine_terminal(p_instance_id TEXT) +RETURNS TEXT LANGUAGE plpgsql SECURITY DEFINER +SET search_path = pg_catalog, pg_temp AS $$ +DECLARE + engine_status TEXT; +BEGIN + FOR attempt IN 1..300 LOOP + EXECUTE format('SELECT status FROM %I.get_instance_info($1)', df.duroxide_schema()) + INTO engine_status USING p_instance_id; + IF engine_status IN ('Completed', 'Failed') THEN + RETURN engine_status; + END IF; + PERFORM pg_sleep(0.1); + END LOOP; + RAISE EXCEPTION 'Engine instance % did not terminate (status %)', p_instance_id, engine_status; +END $$; + -- Install extensions needed by tests (requires superuser) CREATE EXTENSION IF NOT EXISTS dblink; @@ -220,4 +237,3 @@ SELECT 'Playground schema setup complete' AS status; SELECT public._e2e_wait_for_worker_ready(60); SELECT 'TEST PASSED' AS result; - diff --git a/tests/e2e/sql/89_schema_qualification.sql b/tests/e2e/sql/89_schema_qualification.sql index 2dc07657..6761cc6d 100644 --- a/tests/e2e/sql/89_schema_qualification.sql +++ b/tests/e2e/sql/89_schema_qualification.sql @@ -63,15 +63,6 @@ CREATE FUNCTION schema_qualification.to_char(timestamptz, text) RETURNS text LAN BEGIN RAISE EXCEPTION 'schema qualification function trap'; END $$; -CREATE FUNCTION schema_qualification.lower(text) RETURNS text LANGUAGE plpgsql AS $$ -BEGIN - RAISE EXCEPTION 'schema qualification quota trap'; -END $$; -CREATE FUNCTION schema_qualification.pg_advisory_xact_lock(int4, int4) -RETURNS void LANGUAGE plpgsql AS $$ -BEGIN - RAISE EXCEPTION 'schema qualification quota lock trap'; -END $$; CREATE FUNCTION schema_qualification.count_trap(bigint) RETURNS bigint LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'schema qualification aggregate trap'; @@ -79,6 +70,14 @@ END $$; CREATE AGGREGATE schema_qualification.count(*) ( SFUNC = schema_qualification.count_trap, STYPE = bigint, INITCOND = '0' ); +CREATE FUNCTION schema_qualification.array_agg_trap(text[], text) RETURNS text[] +LANGUAGE plpgsql AS $$ +BEGIN + RAISE EXCEPTION 'schema qualification array aggregate trap'; +END $$; +CREATE AGGREGATE schema_qualification.array_agg(text) ( + SFUNC = schema_qualification.array_agg_trap, STYPE = text[] +); CREATE FUNCTION schema_qualification.user_value() RETURNS integer LANGUAGE sql AS $$ SELECT 42 $$; -- These also shadow casts in the separate transaction_mode => 'new' backend. diff --git a/tests/e2e/sql/90_start_limits.sql b/tests/e2e/sql/90_start_limits.sql index d2b63f8b..07393b3f 100644 --- a/tests/e2e/sql/90_start_limits.sql +++ b/tests/e2e/sql/90_start_limits.sql @@ -92,11 +92,12 @@ INSERT INTO start_limits_test.ids SELECT 'initial', df.start(df.sleep(300), 'start-limits-initial') FROM generate_series(1, 2); RESET ROLE; -CREATE TEMP TABLE _quota_engine_count (n BIGINT); +-- Enqueue remains possible, but acknowledgment cannot create an engine instance +-- and consume its queue item while this lock is held. +BEGIN; DO $$ BEGIN - EXECUTE format('INSERT INTO _quota_engine_count SELECT count(*) FROM %I.instances', - df.duroxide_schema()); + EXECUTE format('LOCK TABLE %I.instances IN SHARE MODE', df.duroxide_schema()); END $$; SET ROLE start_limits_user; SELECT start_limits_test.expect_rejected('active', 2); @@ -105,16 +106,22 @@ DO $$ DECLARE actual BIGINT; BEGIN - EXECUTE format('SELECT count(*) FROM %I.instances', df.duroxide_schema()) INTO actual; - IF actual <> (SELECT n FROM _quota_engine_count) THEN - RAISE EXCEPTION 'TEST FAILED: rejected start created engine work'; + EXECUTE format( + 'SELECT count(*) FROM %I.orchestrator_queue WHERE work_item LIKE $1', + df.duroxide_schema() + ) INTO actual USING '%start-limits-rejected%'; + IF actual <> 0 THEN + RAISE EXCEPTION 'TEST FAILED: rejected start enqueued engine work'; END IF; END $$; +COMMIT; SET ROLE start_limits_user; DO $$ BEGIN PERFORM df.cancel((SELECT id FROM start_limits_test.ids WHERE kind = 'initial' LIMIT 1)); + PERFORM public._e2e_wait_for_engine_terminal( + (SELECT id FROM start_limits_test.ids WHERE kind = 'initial' LIMIT 1)); INSERT INTO start_limits_test.ids SELECT 'completed', df.start('SELECT 1', 'start-limits-replacement'); END $$; @@ -125,6 +132,7 @@ BEGIN RAISE EXCEPTION 'TEST FAILED: replacement did not complete'; END IF; END $$; +SELECT public._e2e_wait_for_engine_terminal(id) FROM start_limits_test.ids WHERE kind = 'completed'; INSERT INTO start_limits_test.ids SELECT 'after-completion', df.start(df.sleep(300)); SELECT start_limits_test.expect_rejected('active', 2); SELECT df.cancel(id) FROM df.instances WHERE status IN ('pending', 'running'); @@ -167,6 +175,7 @@ INSERT INTO start_limits_test.ids SELECT 'new-sleep', df.start(df.sleep(300), transaction_mode => 'new'); SELECT start_limits_test.expect_rejected('active', 1, 'new'); SELECT df.cancel(id) FROM start_limits_test.ids WHERE kind = 'new-sleep'; +SELECT public._e2e_wait_for_engine_terminal(id) FROM start_limits_test.ids WHERE kind = 'new-sleep'; INSERT INTO start_limits_test.ids SELECT 'new-completed', df.start('SELECT 1', transaction_mode => 'new'); DO $$ @@ -176,6 +185,7 @@ BEGIN RAISE EXCEPTION 'TEST FAILED: new-transaction instance did not complete'; END IF; END $$; +SELECT public._e2e_wait_for_engine_terminal(id) FROM start_limits_test.ids WHERE kind = 'new-completed'; SELECT start_limits_test.expect_rejected('retained', 2, 'new'); RESET ROLE; @@ -187,7 +197,7 @@ SELECT 'unlimited', df.start(df.sleep(300)) FROM generate_series(1, 3); DO $$ BEGIN IF EXISTS (SELECT 1 FROM pg_locks WHERE pid = pg_backend_pid() - AND locktype = 'advisory' AND classid = x'50474451'::int AND objsubid = 2) THEN + AND locktype = 'advisory' AND classid = x'50474451'::int AND objsubid = 3) THEN RAISE EXCEPTION 'TEST FAILED: unlimited starts acquired a quota lock'; END IF; END $$; @@ -200,7 +210,6 @@ INSERT INTO start_limits_test.ids SELECT 'superuser', df.start(df.sleep(300)) FROM generate_series(1, 3); SELECT df.cancel(id) FROM start_limits_test.ids; -DROP TABLE _quota_engine_count; DROP SCHEMA start_limits_test CASCADE; DROP OWNED BY start_limits_user, start_limits_new_user; DROP ROLE start_limits_user, start_limits_new_user; diff --git a/tests/e2e/sql/91_start_limits_concurrency.sql b/tests/e2e/sql/91_start_limits_concurrency.sql index 2ee9044d..537816d3 100644 --- a/tests/e2e/sql/91_start_limits_concurrency.sql +++ b/tests/e2e/sql/91_start_limits_concurrency.sql @@ -25,7 +25,7 @@ CREATE FUNCTION pg_temp.quota_wait_for_lock(waiter INTEGER) RETURNS VOID LANGUAG BEGIN FOR attempt IN 1..300 LOOP IF EXISTS (SELECT 1 FROM pg_locks WHERE pid = waiter AND locktype = 'advisory' - AND classid = x'50474451'::int AND objsubid = 2 AND NOT granted) THEN + AND classid = x'50474451'::int AND objsubid = 3 AND NOT granted) THEN RETURN; END IF; PERFORM pg_sleep(0.1); @@ -78,6 +78,7 @@ BEGIN PERFORM dblink_exec('quota_holder', $sql$ DO $cancel$ BEGIN PERFORM df.cancel(id) FROM df.instances WHERE status IN ('pending', 'running'); + PERFORM public._e2e_wait_for_engine_terminal(id) FROM df.instances; END $cancel$; $sql$); PERFORM dblink_exec('quota_waiter', 'BEGIN ISOLATION LEVEL ' || isolation); @@ -115,6 +116,7 @@ BEGIN PERFORM dblink_exec('quota_holder', $sql$ DO $cancel$ BEGIN PERFORM df.cancel(id) FROM df.instances WHERE status IN ('pending', 'running'); + PERFORM public._e2e_wait_for_engine_terminal(id) FROM df.instances; END $cancel$; $sql$); PERFORM dblink_exec('quota_holder', 'BEGIN'); @@ -139,6 +141,7 @@ BEGIN PERFORM dblink_exec('quota_holder', $sql$ DO $cancel$ BEGIN PERFORM df.cancel(id) FROM df.instances WHERE status IN ('pending', 'running'); + PERFORM public._e2e_wait_for_engine_terminal(id) FROM df.instances; END $cancel$; $sql$); PERFORM dblink_exec('quota_holder', 'BEGIN'); diff --git a/tests/e2e/sql/92_start_limits_trust.sql b/tests/e2e/sql/92_start_limits_trust.sql new file mode 100644 index 00000000..346a4595 --- /dev/null +++ b/tests/e2e/sql/92_start_limits_trust.sql @@ -0,0 +1,166 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +DO $$ +DECLARE + role_name TEXT; +BEGIN + FOREACH role_name IN ARRAY ARRAY['quota_trust_user', 'quota_trust_attacker'] LOOP + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = role_name) THEN + EXECUTE format('DROP OWNED BY %I', role_name); + EXECUTE format('DROP ROLE %I', role_name); + END IF; + END LOOP; +END $$; +CREATE ROLE quota_trust_user LOGIN; +CREATE ROLE quota_trust_attacker LOGIN; +SELECT df.grant_usage('quota_trust_user'); +SELECT format('ALTER ROLE quota_trust_user IN DATABASE %I SET df.max_concurrent_per_user = 1', + current_database()) \gexec +SELECT format('ALTER ROLE quota_trust_user IN DATABASE %I SET df.max_instances_per_user = 0', + current_database()) \gexec +SELECT format('ALTER ROLE quota_trust_user IN DATABASE %I SET lock_timeout = ''500ms''', + current_database()) \gexec + +DO $$ +DECLARE + base_conn TEXT := format('host=localhost dbname=%L port=%s', + current_database(), current_setting('port')); + instance TEXT; + replacement TEXT; + mode TEXT; + forged_status TEXT; + engine_status TEXT; + victim_oid OID := 'quota_trust_user'::regrole; +BEGIN + PERFORM dblink_connect('quota_trust_user', base_conn || ' user=quota_trust_user'); + PERFORM dblink_connect('quota_trust_attacker', base_conn || ' user=quota_trust_attacker'); + PERFORM dblink_connect('quota_trust_fence', base_conn || ' user=postgres'); + PERFORM dblink_exec('quota_trust_attacker', format($remote$ + DO $lock$ + BEGIN + IF has_schema_privilege(current_user, 'df', 'USAGE') THEN + RAISE EXCEPTION 'TEST SETUP ERROR: attacker must not have df usage'; + END IF; + PERFORM pg_advisory_lock(x'50474451'::int, %1$s::int); + PERFORM pg_advisory_lock((x'50474451'::bigint << 32) | %1$s::bigint); + END $lock$; + $remote$, victim_oid)); + + FOREACH mode IN ARRAY ARRAY['caller', 'new'] LOOP + PERFORM dblink_exec('quota_trust_fence', format( + 'BEGIN; LOCK TABLE %I.instances IN SHARE MODE', df.duroxide_schema() + )); + SELECT id INTO instance FROM dblink('quota_trust_user', format( + 'SELECT df.start(df.sleep(300), ''quota-trust-live'', NULL, %L)', mode + )) AS started(id TEXT); + EXECUTE format('SELECT status FROM %I.get_instance_info($1)', df.duroxide_schema()) + INTO engine_status USING instance; + IF engine_status IS NOT NULL THEN + RAISE EXCEPTION 'TEST FAILED: acknowledgment fence did not hold queued start'; + END IF; + PERFORM dblink_exec('quota_trust_user', format($remote$ + DO $reject$ + BEGIN + BEGIN + PERFORM df.start('SELECT 1', 'quota-trust-queued', NULL, %L); + RAISE EXCEPTION 'TEST FAILED: queued instance did not consume capacity'; + EXCEPTION WHEN configuration_limit_exceeded THEN NULL; + END; + END $reject$; + $remote$, mode)); + PERFORM dblink_exec('quota_trust_fence', 'COMMIT'); + FOR attempt IN 1..300 LOOP + EXECUTE format('SELECT status FROM %I.get_instance_info($1)', df.duroxide_schema()) + INTO engine_status USING instance; + EXIT WHEN engine_status = 'Running'; + PERFORM pg_sleep(0.1); + END LOOP; + IF engine_status IS DISTINCT FROM 'Running' THEN + RAISE EXCEPTION 'TEST FAILED: % start blocked or did not run: %', mode, engine_status; + END IF; + + FOREACH forged_status IN ARRAY ARRAY['completed', 'failed', 'cancelled'] LOOP + PERFORM dblink_exec('quota_trust_user', format( + 'UPDATE df.instances SET status = %L WHERE id = %L', forged_status, instance + )); + PERFORM dblink_exec('quota_trust_user', format($remote$ + DO $reject$ + BEGIN + BEGIN + PERFORM df.start('SELECT 1', 'quota-trust-rejected', NULL, %L); + RAISE EXCEPTION 'TEST FAILED: forged status bypassed active quota'; + EXCEPTION WHEN configuration_limit_exceeded THEN + IF SQLERRM NOT LIKE '%%has 1 active instances (limit 1)%%' THEN + RAISE; + END IF; + END; + END $reject$; + $remote$, mode)); + EXECUTE format('SELECT status FROM %I.get_instance_info($1)', df.duroxide_schema()) + INTO engine_status USING instance; + IF engine_status IS DISTINCT FROM 'Running' THEN + RAISE EXCEPTION 'TEST FAILED: forged display status changed live engine state'; + END IF; + END LOOP; + + PERFORM dblink_exec('quota_trust_fence', format( + 'BEGIN; LOCK TABLE %I.executions IN ACCESS EXCLUSIVE MODE', df.duroxide_schema() + )); + PERFORM dblink_exec('quota_trust_user', format($remote$ + DO $unavailable$ + BEGIN + BEGIN + PERFORM df.start('SELECT 1', 'quota-trust-unavailable', NULL, %L); + RAISE EXCEPTION 'TEST FAILED: unavailable engine accounting admitted work'; + EXCEPTION WHEN internal_error THEN + IF SQLERRM NOT LIKE '%%quota check failed:%%lock timeout%%' THEN + RAISE; + END IF; + END; + END $unavailable$; + $remote$, mode)); + PERFORM pg_stat_clear_snapshot(); + IF EXISTS ( + SELECT 1 FROM pg_stat_activity + WHERE state = 'idle in transaction (aborted)' AND query LIKE '%AS requested(id)%' + ) THEN + RAISE EXCEPTION 'TEST FAILED: failed quota lookup left an aborted provider transaction'; + END IF; + PERFORM dblink_exec('quota_trust_fence', 'COMMIT'); + + -- Hold engine acknowledgment, not enqueue, so cancellation is definitely + -- requested but not yet terminal when admission is checked. + PERFORM dblink_exec('quota_trust_fence', format( + 'BEGIN; LOCK TABLE %I.executions IN SHARE MODE', df.duroxide_schema() + )); + PERFORM * FROM dblink('quota_trust_user', + format('SELECT df.cancel(%L)', instance)) AS cancelled(message TEXT); + PERFORM dblink_exec('quota_trust_user', format($remote$ + DO $reject$ + BEGIN + BEGIN + PERFORM df.start('SELECT 1', 'quota-trust-cancelling', NULL, %L); + RAISE EXCEPTION 'TEST FAILED: unprocessed cancellation freed capacity'; + EXCEPTION WHEN configuration_limit_exceeded THEN NULL; + END; + END $reject$; + $remote$, mode)); + PERFORM dblink_exec('quota_trust_fence', 'COMMIT'); + PERFORM public._e2e_wait_for_engine_terminal(instance); + + SELECT id INTO replacement FROM dblink('quota_trust_user', format( + 'SELECT df.start(''SELECT 1'', ''quota-trust-replacement'', NULL, %L)', mode + )) AS started(id TEXT); + IF public._e2e_wait_for_engine_terminal(replacement) <> 'Completed' THEN + RAISE EXCEPTION 'TEST FAILED: capacity did not recover after cancellation'; + END IF; + END LOOP; + PERFORM dblink_disconnect('quota_trust_user'); + PERFORM dblink_disconnect('quota_trust_attacker'); + PERFORM dblink_disconnect('quota_trust_fence'); +END $$; + +DROP OWNED BY quota_trust_user, quota_trust_attacker; +DROP ROLE quota_trust_user, quota_trust_attacker; +SELECT 'TEST PASSED' AS result; diff --git a/tests/e2e/sql/93_start_limits_serializable.sql b/tests/e2e/sql/93_start_limits_serializable.sql new file mode 100644 index 00000000..60e3a95a --- /dev/null +++ b/tests/e2e/sql/93_start_limits_serializable.sql @@ -0,0 +1,102 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +DO $$ +DECLARE + role_name TEXT; +BEGIN + FOREACH role_name IN ARRAY ARRAY['quota_ssi_a', 'quota_ssi_b'] LOOP + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = role_name) THEN + EXECUTE format('DROP OWNED BY %I', role_name); + EXECUTE format('DROP ROLE %I', role_name); + END IF; + END LOOP; +END $$; +CREATE ROLE quota_ssi_a LOGIN; +CREATE ROLE quota_ssi_b LOGIN; +SELECT df.grant_usage('quota_ssi_a'); +SELECT df.grant_usage('quota_ssi_b'); +CREATE TEMP TABLE _quota_ssi (conn TEXT PRIMARY KEY, id TEXT, state TEXT); +INSERT INTO _quota_ssi (conn) VALUES ('quota_ssi_a'), ('quota_ssi_b'); + +DO $$ +DECLARE + session RECORD; + commit_result TEXT; + retry_count INTEGER := 0; +BEGIN + FOR session IN SELECT * FROM _quota_ssi ORDER BY conn LOOP + PERFORM dblink_connect(session.conn, format( + 'host=localhost dbname=%L port=%s user=%s', + current_database(), current_setting('port'), session.conn)); + PERFORM dblink_exec(session.conn, $remote$ + SET statement_timeout = '15s'; + -- Exercise the scan fallback used by binary-only upgraded schemas. + SET enable_indexscan = off; + SET enable_indexonlyscan = off; + SET enable_bitmapscan = off; + CREATE FUNCTION pg_temp.quota_start() + RETURNS TABLE (id TEXT, state TEXT) LANGUAGE plpgsql AS $start$ + BEGIN + RETURN QUERY SELECT df.start(df.sleep(300)), '00000'::text; + EXCEPTION WHEN serialization_failure THEN + RETURN QUERY SELECT NULL::text, SQLSTATE; + END $start$; + $remote$); + PERFORM dblink_exec(session.conn, 'BEGIN ISOLATION LEVEL SERIALIZABLE'); + END LOOP; + + -- Both transactions remain open through both starts. Their advisory keys + -- differ, but PostgreSQL predicate locking can still require a retry. + FOR session IN SELECT * FROM _quota_ssi ORDER BY conn LOOP + UPDATE _quota_ssi + SET id = started.id, state = started.state + FROM dblink(session.conn, 'SELECT * FROM pg_temp.quota_start()') + AS started(id TEXT, state TEXT) + WHERE conn = session.conn; + END LOOP; + FOR session IN SELECT * FROM _quota_ssi ORDER BY conn LOOP + IF session.state = '40001' THEN + PERFORM dblink_exec(session.conn, 'ROLLBACK'); + ELSIF session.state = '00000' THEN + commit_result := dblink_exec(session.conn, 'COMMIT', false); + IF commit_result = 'ERROR' THEN + IF dblink_error_message(session.conn) NOT LIKE '%could not serialize access%' THEN + RAISE EXCEPTION 'TEST FAILED: unexpected commit error: %', + dblink_error_message(session.conn); + END IF; + UPDATE _quota_ssi SET state = '40001' WHERE conn = session.conn; + PERFORM dblink_exec(session.conn, 'ROLLBACK'); + ELSIF commit_result <> 'COMMIT' THEN + RAISE EXCEPTION 'TEST FAILED: expected COMMIT, got %', commit_result; + END IF; + ELSE + RAISE EXCEPTION 'TEST FAILED: unexpected start result: %', session.state; + END IF; + END LOOP; + FOR session IN SELECT * FROM _quota_ssi WHERE state = '40001' ORDER BY conn LOOP + retry_count := retry_count + 1; + PERFORM dblink_exec(session.conn, 'BEGIN ISOLATION LEVEL SERIALIZABLE'); + UPDATE _quota_ssi + SET id = started.id, state = started.state + FROM dblink(session.conn, 'SELECT * FROM pg_temp.quota_start()') + AS started(id TEXT, state TEXT) + WHERE conn = session.conn; + PERFORM dblink_exec(session.conn, 'COMMIT'); + END LOOP; + FOR session IN SELECT * FROM _quota_ssi ORDER BY conn LOOP + IF session.state IS DISTINCT FROM '00000' OR session.id IS NULL OR + (SELECT count(*) FROM df.instances WHERE submitted_by = session.conn::regrole) <> 1 THEN + RAISE EXCEPTION 'TEST FAILED: serializable retry left incorrect instance state: %', session; + END IF; + PERFORM df.cancel(session.id); + PERFORM public._e2e_wait_for_engine_terminal(session.id); + PERFORM dblink_disconnect(session.conn); + END LOOP; + RAISE NOTICE 'Two overlapping serializable starts completed with % transaction retries', retry_count; +END $$; + +DROP TABLE _quota_ssi; +DROP OWNED BY quota_ssi_a, quota_ssi_b; +DROP ROLE quota_ssi_a, quota_ssi_b; +SELECT 'TEST PASSED' AS result; From a7f9ece69824f6236dcb1acf789480c87c34f3e4 Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Wed, 28 Oct 2026 16:25:10 +0000 Subject: [PATCH 3/5] Fix issue identified by review --- CHANGELOG.md | 3 + USER_GUIDE.md | 30 +++- docs/api-reference.md | 5 +- docs/upgrade-testing.md | 26 ++- scripts/test-e2e-local.sh | 2 +- scripts/test-upgrade.sh | 18 ++ src/client.rs | 135 +++++++++++--- src/dsl.rs | 50 +++--- src/lib.rs | 6 +- src/worker.rs | 127 ++++++++++++- tests/e2e/sql/91_start_limits_concurrency.sql | 13 +- tests/e2e/sql/92_start_limits_trust.sql | 27 ++- tests/e2e/sql/94_start_limits_rollback.sql | 167 ++++++++++++++++++ .../sql/95_start_limits_rollback_cleanup.sql | 85 +++++++++ 14 files changed, 608 insertions(+), 86 deletions(-) create mode 100644 tests/e2e/sql/94_start_limits_rollback.sql create mode 100644 tests/e2e/sql/95_start_limits_rollback_cleanup.sql diff --git a/CHANGELOG.md b/CHANGELOG.md index 03c53708..5677e7e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,9 @@ Pre-1.0 note: while `pg_durable` is in major version `0`, minor releases may inc superuser-settable; `0` disables a quota and superuser identities are exempt. Active accounting uses protected engine state, not caller-writable display status. Cancellation frees active capacity after the engine processes it. + A protected ledger commits atomically with engine enqueue, so transaction + and savepoint rollbacks cannot bypass quotas. Rolled-back submissions remain + charged until the engine processes them and reconciliation removes their history. Starts serialize per role until the submitting transaction ends and reject over-quota work before insert/enqueue, including on older installed schemas. diff --git a/USER_GUIDE.md b/USER_GUIDE.md index 8511e3b0..28d928d6 100644 --- a/USER_GUIDE.md +++ b/USER_GUIDE.md @@ -2012,8 +2012,9 @@ runtime. Distinguish three roles: does not need the extension unless the SQL itself calls `df` APIs. Satellite installation requires a compatible control worker with readiness schema -version `2` (including its `_origins` registry). Provider migrations run only in -control through the worker's `ApplyAll`; satellite creation never creates control +version `3` (including its origin registry and protected submission accounting). +Provider migrations run only in control through the worker's `ApplyAll`; +satellite creation never creates control automatically. The provider namespace is extension-owned, its contents are created by `pg_durable.worker_role`, and satellite objects belong to the local installer. @@ -2546,6 +2547,10 @@ worker runs a **best-effort reconciliation pass** that does two things: age past `retention_days`. Anything the engine is still tracking with a live `df` row is left untouched. +After the corresponding engine work and history are gone, reconciliation removes +the protected submission record and frees its retained quota capacity. This +includes submissions whose caller transaction or savepoint rolled back. + Satellite maintenance visits `_origins` registrations made by activity routing for submitted work, in bounded batches under the shared origin connection budget. Registration is not synchronous with `df.start()`; unused installations are not @@ -3011,7 +3016,7 @@ for the full syntax and default domain list. | Setting | Default | Counts | |---------|---------|--------| | `df.max_concurrent_per_user` | `100` | Instances not yet terminal in the durable engine, including queued, sleeping, looping, and signal-waiting workflows | -| `df.max_instances_per_user` | `10000` | All retained instances, including completed, failed, and cancelled workflows | +| `df.max_instances_per_user` | `10000` | Retained submissions, including rolled-back starts and completed, failed, and cancelled workflows | Each quota applies to the **submitting role (`current_user`) in the origin database**. `SET ROLE` and `SECURITY DEFINER` affect that identity just as they @@ -3040,10 +3045,11 @@ and a `HINT` identifying the relevant setting. Finishing work frees concurrent capacity once the engine records the terminal state. Cancellation releases it after the engine processes the request, not immediately when `df.cancel()` returns. Changing the displayed instance status does not release capacity. -Neither completion nor cancellation frees retained-instance capacity. Retained -capacity is freed when [automatic reconciliation](#automatic-reconciliation) removes -terminal instances. If reconciliation is disabled, or retained work has not -aged out, an administrator must adjust retention or raise the quota. +Neither rollback, completion, nor cancellation frees retained-instance capacity. +Capacity is freed when [automatic reconciliation](#automatic-reconciliation) +removes the engine state and protected submission accounting. If reconciliation +is disabled, or retained work has not aged out, an administrator must adjust +retention or raise the quota. The default retained quota equals reconciliation's 10,000-terminal-instance retention cap. A role with 10,000 recent instances can therefore remain blocked @@ -3058,7 +3064,11 @@ They do not wait for workflows to finish. Keep submitting transactions short and use PostgreSQL `lock_timeout` to bound admission waits. Counts include current committed rows and the caller's own changes even under `REPEATABLE READ` or `SERIALIZABLE`. Rollback, including rollback to a savepoint, releases the -corresponding reservations. +transaction lock and discards the graph, but does not erase submission accounting: +engine enqueue commits independently. Rolled-back starts consume active capacity +until the engine records their failure and retained capacity until reconciliation +removes their engine state and accounting. Quota counts can therefore exceed the +number of visible `df.instances` rows. Perform business-table updates before starting workflows when possible, and avoid leaving submitting transactions idle. Consider PostgreSQL @@ -3086,7 +3096,9 @@ can scan other roles' history until the index is installed. Active accounting also checks the role's retained IDs against the engine, so larger histories increase admission cost even with the index. Measure representative workloads before substantially raising retention or quotas. Disabling the active quota -skips the engine lookup; disabling both quotas skips admission entirely. +skips execution-status lookups, but retained accounting still reads the protected +submission ledger. Disabling both quotas skips admission checks. Successful +enqueues are recorded even while limits are disabled. --- diff --git a/docs/api-reference.md b/docs/api-reference.md index 990317a4..b747a36f 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -907,7 +907,7 @@ Admission quotas for `df.start()`, per submitting role and origin database. | Property | Concurrent quota | Retained-instance quota | |----------|------------------|-------------------------| | Default | `100` | `10000` | -| Counts | Queued instances and executions not terminal in the protected engine state | All `df.instances` rows owned by the role | +| Counts | Queued instances and executions not terminal in the protected engine state | Retained submissions, including starts whose caller transaction or savepoint rolled back | | Range | `0` to `2147483647`; `0` disables | `0` to `2147483647`; `0` disables | | Context | `SUSET` | `SUSET` | @@ -915,6 +915,9 @@ Superuser submission identities are exempt. Administrators can change either setting at runtime, including per-role/database connection defaults. Cancellation frees active capacity after the engine processes the request; changing display status alone does not release it. +Rollback removes the graph, not the protected submission charge. It remains +active until the engine records failure and retained until reconciliation removes +the engine record and accounting. `transaction_mode => 'new'` uses its separate session's defaults and does not inherit the caller's session `SET` overrides. See [Instance Limits](../USER_GUIDE.md#instance-limits) for tuning and admission waits. diff --git a/docs/upgrade-testing.md b/docs/upgrade-testing.md index 0c788c7b..9f52d907 100644 --- a/docs/upgrade-testing.md +++ b/docs/upgrade-testing.md @@ -361,17 +361,28 @@ what the upgrade script handles, and any backward compatibility considerations. #### `df.start()` admission quotas (#139) - **DDL/B1:** The upgrade adds `idx_instances_submitted_by` on - `df.instances(submitted_by, id)`, matching fresh installation. No new tables - or grants are needed. Both start wrapper symbols use the same admission - check against existing `df.instances` columns and protected engine state, - including on 0.2.2+ schemas without - `ALTER EXTENSION UPDATE`. No status grants are revoked or used as an + `df.instances(submitted_by, id)`, matching fresh installation. No new + extension tables or grants are needed. Both start wrapper symbols use the same + admission check against existing `df.instances` columns and protected engine + state, including on 0.2.2+ schemas without `ALTER EXTENSION UPDATE`. + No status grants are revoked or used as an accounting authority; queued or not-yet-materialized engine instances count as active. Cancellation releases capacity only after the engine processes it. B1 exercises a non-superuser reaching its quota on each supported schema. Old schemas remain correct without the index but can scan other roles' history. The normal index build blocks concurrent writes; set `lock_timeout` and schedule the upgrade appropriately on large installations. +- **Worker-managed accounting:** Worker schema version `3` adds protected + `_submissions` in the control provider schema, indexed by origin and submitting + role. The worker creates it before publishing readiness, including against old + extension schemas. Each new root's ledger row and provider enqueue commit in + one transaction; neither caller nor savepoint rollback removes the charge. + Admission unions these records with visible legacy instance IDs, avoiding + double counting. Bounded reconciliation removes ledger entries only after + no corresponding engine instance, execution, history, queue item, or lock + remains. The enqueue WorkItem and orchestration input are unchanged. + Pre-upgrade orphans that already lost their `df` identity cannot be attributed + retroactively; existing orphan cleanup remains responsible for them. - **Behavior change:** New superuser-settable GUCs `df.max_concurrent_per_user` (default `100`) and `df.max_instances_per_user` (default `10000`) apply immediately after loading the new binary. @@ -381,7 +392,8 @@ what the upgrade script handles, and any backward compatibility considerations. lock with a tag that SQL advisory-lock functions cannot construct, and a current snapshot, including under repeatable-read and serializable isolation. E2E tests cover concurrent commit, rollback, savepoint rollback, lock timeout, - caller-writable status, and attempts to squat on public advisory-lock keys. + caller-writable status, attempts to squat on public advisory-lock keys, + repeated transaction/savepoint rollbacks, and quota recovery after cleanup. - **B2/replay:** No data migration, runtime schema detection, durable payload, or orchestration changes. Workflows that call `df.start()` from SQL are subject to the new quotas when that SQL activity executes. The separate @@ -395,6 +407,8 @@ what the upgrade script handles, and any backward compatibility considerations. These are query measurements, not throughput guarantees. The index reduces local cross-role scanning, not the per-owner engine-lookup cost. Only retained counting runs when the active quota is disabled; both disabled skips admission. + These pre-ledger measurements do not include the subsequent protected-ledger + union; ledger lookups use the origin/owner-leading index. ### v0.2.8 → v0.2.9 diff --git a/scripts/test-e2e-local.sh b/scripts/test-e2e-local.sh index c0e77ccc..3ba4d0cb 100755 --- a/scripts/test-e2e-local.sh +++ b/scripts/test-e2e-local.sh @@ -221,7 +221,7 @@ phase_for_test() { 46_connection_limit_startup_validation) echo "connlimit-startup" ;; - 54_reconcile_orphans|76_multi_database_reconcile) + 54_reconcile_orphans|76_multi_database_reconcile|95_start_limits_rollback_cleanup) echo "reconcile" ;; 88_multi_database_maintenance_backlog) diff --git a/scripts/test-upgrade.sh b/scripts/test-upgrade.sh index 27a562c1..5eb91683 100755 --- a/scripts/test-upgrade.sh +++ b/scripts/test-upgrade.sh @@ -973,6 +973,24 @@ test_b1_start_limits() { END IF; END; END \$\$; + RESET ROLE; + SET LOCAL df.max_concurrent_per_user = 0; + SET LOCAL df.max_instances_per_user = 2; + SET LOCAL ROLE upgrade_start_limits; + SAVEPOINT quota_rollback; + SELECT df.start('SELECT 1'); + ROLLBACK TO quota_rollback; + DO \$\$ + BEGIN + BEGIN + PERFORM df.start('SELECT 1'); + RAISE EXCEPTION 'legacy df.start bypassed retained quota after savepoint rollback'; + EXCEPTION WHEN configuration_limit_exceeded THEN + IF SQLERRM NOT LIKE '%has 2 retained instances (limit 2)%' THEN + RAISE; + END IF; + END; + END \$\$; ROLLBACK; " >/dev/null } diff --git a/src/client.rs b/src/client.rs index 5839bbe0..88022b2b 100644 --- a/src/client.rs +++ b/src/client.rs @@ -13,7 +13,7 @@ use std::cell::RefCell; use std::sync::{Arc, OnceLock}; use std::time::Duration; -use duroxide::Client; +use duroxide::{providers::WorkItem, Client}; use pgrx::pg_sys::panic::ErrorReport; use pgrx::prelude::*; use sqlx::Connection; @@ -182,11 +182,13 @@ where result } -pub(crate) fn count_active_instances(ids: &[String]) -> Result { - if ids.is_empty() { - return Ok(0); - } +pub(crate) fn submission_counts( + ids: &[String], + submitted_by: pg_sys::Oid, + include_active: bool, +) -> Result<(i64, i64), String> { let origin = crate::origin::backend_origin()?; + let origin_key = crate::origin::engine_id_for_origin(origin.as_ref(), ""); let engine_ids: Vec = ids .iter() .map(|id| crate::origin::engine_id_for_origin(origin.as_ref(), id)) @@ -205,32 +207,60 @@ pub(crate) fn count_active_instances(ids: &[String]) -> Result { .await .map_err(|e| format!("Failed to configure df.start() quota lookup: {e}"))?; - // Missing engine rows include queued starts and uncommitted graphs. - // Only the engine's protected terminal state can release capacity. - sqlx::query_scalar(&format!( - "SELECT pg_catalog.count(*) - FROM pg_catalog.unnest($1::pg_catalog.text[]) AS requested(id) + // Old binaries left no submission ledger. Their visible df rows + // still count; grouping charges new committed starts only once. + let owned = format!( + "WITH requested AS ( + SELECT id, pg_catalog.bool_or(recorded) AS recorded + FROM ( + SELECT pg_catalog.unnest($1::pg_catalog.text[]) AS id, false AS recorded + UNION ALL + SELECT instance_id, true FROM {schema}._submissions + WHERE origin OPERATOR(pg_catalog.=) $2 + AND submitted_by OPERATOR(pg_catalog.=) $3::pg_catalog.int8::pg_catalog.oid + ) AS owned GROUP BY id + )", + schema = backend.control.schema, + ); + // A recorded start is queued before it materializes. Once both + // are gone, ledger cleanup must not make retired work active again. + let query = if include_active { + format!( + "{owned} + SELECT pg_catalog.count(*), + pg_catalog.count(*) FILTER ( + WHERE e.status OPERATOR(pg_catalog.<>) ALL ( + ARRAY['Completed', 'Failed']::pg_catalog.text[]) + OR (e.status IS NULL AND ( + NOT requested.recorded OR i.instance_id IS NOT NULL OR EXISTS ( + SELECT 1 FROM {schema}.orchestrator_queue AS queued + WHERE queued.instance_id OPERATOR(pg_catalog.=) requested.id)))) + FROM requested LEFT JOIN {schema}.instances AS i ON i.instance_id OPERATOR(pg_catalog.=) requested.id LEFT JOIN {schema}.executions AS e ON e.instance_id OPERATOR(pg_catalog.=) i.instance_id - AND e.execution_id OPERATOR(pg_catalog.=) i.current_execution_id - WHERE e.status IS NULL OR e.status OPERATOR(pg_catalog.<>) ALL ( - ARRAY['Completed', 'Failed']::pg_catalog.text[])", - schema = backend.control.schema, - )) + AND e.execution_id OPERATOR(pg_catalog.=) i.current_execution_id", + schema = backend.control.schema, + ) + } else { + format!("{owned} SELECT pg_catalog.count(*), 0::pg_catalog.int8 FROM requested") + }; + sqlx::query_as(&query) .bind(&engine_ids) + .bind(&origin_key) + .bind(i64::from(submitted_by.to_u32())) .fetch_one(&mut *tx) .await .map_err(|e| format!("Failed to read df.start() active instances: {e}")) } .await; match result { - Ok(active) => { + Ok(counts) => { tx.commit() .await .map_err(|e| format!("Failed to finish df.start() quota lookup: {e}"))?; - Ok(active) + Ok(counts) } Err(error) => { // This backend's runtime stops polling when block_on returns. @@ -302,24 +332,79 @@ pub fn start_durable_function( function_name: &str, instance_id: &str, input: &str, + submitted_by: pg_sys::Oid, ) -> Result<(), String> { log!( "pg_durable: start_durable_function for instance {}", instance_id ); - let fn_name = function_name.to_string(); - let inst_id = crate::origin::backend_engine_id(instance_id)?; - let inp = input.to_string(); + let origin = crate::origin::backend_origin()?; + let origin_key = crate::origin::engine_id_for_origin(origin.as_ref(), ""); + let inst_id = crate::origin::engine_id_for_origin(origin.as_ref(), instance_id); + let work_item = serde_json::to_string(&WorkItem::StartOrchestration { + instance: inst_id.clone(), + orchestration: function_name.to_string(), + input: input.to_string(), + version: None, + parent_instance: None, + parent_id: None, + parent_execution_id: None, + execution_id: duroxide::INITIAL_EXECUTION_ID, + }) + .map_err(|e| format!("Failed to serialize durable start: {e}"))?; with_backend_client(|backend, rt| { rt.block_on(async { - backend - .client - .start_orchestration(&inst_id, &fn_name, &inp) + let mut tx = backend + .provider + .pool() + .begin() .await - .map_err(|e| format!("Failed to start durable function: {e:?}"))?; - Ok(()) + .map_err(|e| format!("Failed to begin durable submission: {e}"))?; + let result = async { + crate::origin::configure_metadata_transaction(&mut tx) + .await + .map_err(|e| format!("Failed to configure durable submission: {e}"))?; + sqlx::query(&format!( + "INSERT INTO {schema}._submissions (instance_id, origin, submitted_by) + VALUES ($1, $2, $3::pg_catalog.int8::pg_catalog.oid)", + schema = backend.control.schema, + )) + .bind(&inst_id) + .bind(&origin_key) + .bind(i64::from(submitted_by.to_u32())) + .execute(&mut *tx) + .await + .map_err(|e| format!("Failed to record durable submission: {e}"))?; + + // This is the same WorkItem and provider function used by + // Client::start_orchestration, but shares the ledger transaction. + sqlx::query(&format!( + "SELECT {schema}.enqueue_orchestrator_work( + $1, $2, pg_catalog.clock_timestamp(), NULL, NULL, NULL)", + schema = backend.control.schema, + )) + .bind(&inst_id) + .bind(&work_item) + .execute(&mut *tx) + .await + .map_err(|e| format!("Failed to enqueue durable submission: {e}"))?; + Ok(()) + } + .await; + match result { + Ok(()) => tx + .commit() + .await + .map_err(|e| format!("Failed to commit durable submission: {e}")), + Err(error) => { + tx.rollback() + .await + .map_err(|e| format!("{error}; submission rollback failed: {e}"))?; + Err(error) + } + } }) }) } diff --git a/src/dsl.rs b/src/dsl.rs index aa72660b..9ad27ea1 100644 --- a/src/dsl.rs +++ b/src/dsl.rs @@ -1082,18 +1082,16 @@ fn acquire_instance_quota_lock(user_oid: pg_sys::Oid) { } } -fn current_instance_quota_rows(user_oid: pg_sys::Oid, include_ids: bool) -> (i64, Vec) { +fn current_instance_ids(user_oid: pg_sys::Oid) -> Vec { Spi::connect(|_| unsafe { - let mut argtypes = [pg_sys::OIDOID, pg_sys::BOOLOID]; + let mut argtypes = [pg_sys::OIDOID]; let plan = pg_sys::SPI_prepare( - c"SELECT pg_catalog.count(*), - CASE WHEN $2 THEN COALESCE(pg_catalog.array_agg(id::pg_catalog.text), - ARRAY[]::pg_catalog.text[]) - ELSE ARRAY[]::pg_catalog.text[] END + c"SELECT COALESCE(pg_catalog.array_agg(id::pg_catalog.text), + ARRAY[]::pg_catalog.text[]) FROM df.instances WHERE submitted_by OPERATOR(pg_catalog.=) $1::pg_catalog.regrole" .as_ptr(), - 2, + 1, argtypes.as_mut_ptr(), ); if plan.is_null() { @@ -1105,10 +1103,7 @@ fn current_instance_quota_rows(user_oid: pg_sys::Oid, include_ids: bool) -> (i64 // admission lock, plus this transaction's own earlier starts. Ordinary // SPI would reuse the transaction snapshot and allow quota bypasses. pg_sys::CommandCounterIncrement(); - let mut values = [ - pg_sys::Datum::from(user_oid), - pg_sys::Datum::from(include_ids), - ]; + let mut values = [pg_sys::Datum::from(user_oid)]; let status = pg_sys::SPI_execute_snapshot( plan, values.as_mut_ptr(), @@ -1127,12 +1122,8 @@ fn current_instance_quota_rows(user_oid: pg_sys::Oid, include_ids: bool) -> (i64 .unwrap_or_else(|| pgrx::error!("df.start() quota query returned no tuple table")); let mut is_null = false; let datum = pg_sys::SPI_getbinval(*table.vals, table.tupdesc, 1, &mut is_null); - let total = i64::from_datum(datum, is_null) - .unwrap_or_else(|| pgrx::error!("df.start() quota query returned a NULL count")); - let datum = pg_sys::SPI_getbinval(*table.vals, table.tupdesc, 2, &mut is_null); - let ids = Vec::::from_datum(datum, is_null) - .unwrap_or_else(|| pgrx::error!("df.start() quota query returned NULL instance IDs")); - (total, ids) + Vec::::from_datum(datum, is_null) + .unwrap_or_else(|| pgrx::error!("df.start() quota query returned NULL instance IDs")) }) } @@ -1147,7 +1138,9 @@ fn enforce_instance_limits(user_oid: pg_sys::Oid, user_name: &str) { // uncommitted instances. A function-scoped lock would leave a race window. acquire_instance_quota_lock(user_oid); - let (total, ids) = current_instance_quota_rows(user_oid, concurrent_limit > 0); + let ids = current_instance_ids(user_oid); + let (total, active) = crate::client::submission_counts(&ids, user_oid, concurrent_limit > 0) + .unwrap_or_else(|e| pgrx::error!("df.start() quota check failed: {e}")); let reject = |count, kind, limit, hint| { ErrorReport::new( PgSqlErrorCode::ERRCODE_CONFIGURATION_LIMIT_EXCEEDED, @@ -1164,20 +1157,16 @@ fn enforce_instance_limits(user_oid: pg_sys::Oid, user_name: &str) { total, "retained", instance_limit, - "Wait for terminal history to age out and be reconciled, or ask an administrator to adjust retention or raise df.max_instances_per_user. Completion or cancellation alone does not free retained capacity.", + "Wait for terminal history, including rolled-back submissions, to be reconciled, or ask an administrator to adjust retention or raise df.max_instances_per_user. Rollback, completion, or cancellation alone does not free retained capacity.", ); } - if concurrent_limit > 0 { - let active = crate::client::count_active_instances(&ids) - .unwrap_or_else(|e| pgrx::error!("df.start() quota check failed: {e}")); - if active >= i64::from(concurrent_limit) { - reject( - active, - "active", - concurrent_limit, - "Wait for in-flight instances to finish or for cancellation to be processed, or ask an administrator to raise df.max_concurrent_per_user.", - ); - } + if concurrent_limit > 0 && active >= i64::from(concurrent_limit) { + reject( + active, + "active", + concurrent_limit, + "Wait for in-flight instances to finish or for cancellation or rollback to be processed, or ask an administrator to raise df.max_concurrent_per_user.", + ); } } @@ -1646,6 +1635,7 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< crate::orchestrations::execute_function_graph::NAME, &instance_id, &input_json, + current_user_oid, ) { // Fail fast: the durable engine could not accept the start, so abort the // whole df.start() transaction rather than committing an instance row diff --git a/src/lib.rs b/src/lib.rs index 2811d144..617763b2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -95,7 +95,7 @@ pub use types::Durofut; /// by the background worker after successful initialization. Increment whenever /// a new binary introduces new duroxide-pg migration scripts or any other /// BGW-applied duroxide schema change. -pub const WORKER_SCHEMA_VERSION: i32 = 2; +pub const WORKER_SCHEMA_VERSION: i32 = 3; ::pgrx::pg_module_magic!(name, version); @@ -303,8 +303,8 @@ pub extern "C-unwind" fn _PG_init() { GucRegistry::define_int_guc( c"df.max_instances_per_user", - c"Maximum retained instances per submitting role in each database", - c"Includes terminal instances until reconciliation removes them. df.start() rejects new instances at this limit. 0 disables the limit. Superuser submission identities are exempt.", + c"Maximum retained submissions per submitting role in each database", + c"Includes rolled-back submissions and terminal instances until reconciliation removes their engine state and accounting. df.start() rejects new instances at this limit. 0 disables the limit. Superuser submission identities are exempt.", &MAX_INSTANCES_PER_USER, 0, i32::MAX, diff --git a/src/worker.rs b/src/worker.rs index e2c287b7..051def57 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -856,6 +856,27 @@ async fn write_worker_ready( .execute(&mut *transaction) .await?; + sqlx::query(&format!( + "CREATE TABLE IF NOT EXISTS {schema_name}._submissions ( + instance_id pg_catalog.text PRIMARY KEY, + origin pg_catalog.text NOT NULL, + submitted_by pg_catalog.oid NOT NULL + )" + )) + .execute(&mut *transaction) + .await?; + sqlx::query(&format!( + "REVOKE ALL ON {schema_name}._submissions FROM PUBLIC" + )) + .execute(&mut *transaction) + .await?; + sqlx::query(&format!( + "CREATE INDEX IF NOT EXISTS idx_submissions_owner + ON {schema_name}._submissions (origin, submitted_by, instance_id)" + )) + .execute(&mut *transaction) + .await?; + sqlx::query(&format!( "CREATE TABLE IF NOT EXISTS {schema}._worker_ready ( sentinel pg_catalog.bool PRIMARY KEY DEFAULT TRUE, @@ -1127,6 +1148,7 @@ async fn run_until_extension_dropped_or_shutdown( let mut engine_cursor = String::new(); let mut legacy_cursor = String::new(); let mut registration_cursor = (0i64, uuid::Uuid::nil()); + let mut submission_cursor = String::new(); let mut pending = MaintenanceProgress::new(); let mut drop_check = tokio::time::interval(drop_poll_interval); @@ -1237,6 +1259,16 @@ async fn run_until_extension_dropped_or_shutdown( Ok(more) => pending.registrations = more, Err(error) => { pending.registrations = registration_cursor.0 != 0; log!("pg_durable: registration cleanup deferred: {error}"); }, } } + if pending.submissions { match cleanup_submissions( + maintenance_pool, &schema, &mut submission_cursor, + ).await { + Ok(more) => pending.submissions = more, + Err(error) => { + pending.submissions = false; + submission_cursor.clear(); + log!("pg_durable: submission accounting cleanup deferred: {error}"); + }, + } } pending.finish_slice(&mut origin_cursor) }; tokio::pin!(maintenance); @@ -1278,6 +1310,7 @@ struct MaintenanceProgress { origins: bool, engine: bool, registrations: bool, + submissions: bool, } impl MaintenanceProgress { @@ -1288,11 +1321,17 @@ impl MaintenanceProgress { origins: true, engine: true, registrations: true, + submissions: true, } } fn has_work(&self) -> bool { - self.control || self.legacy || self.origins || self.engine || self.registrations + self.control + || self.legacy + || self.origins + || self.engine + || self.registrations + || self.submissions } fn finish_slice(&mut self, origins: &mut OriginRetentionCursor) -> bool { @@ -2071,6 +2110,61 @@ async fn cleanup_origin_registrations( Ok(true) } +async fn cleanup_submissions( + pool: &sqlx::PgPool, + schema: &str, + cursor: &mut String, +) -> Result { + let mut tx = pool.begin().await?; + crate::origin::configure_metadata_transaction(&mut tx).await?; + let ids: Vec = sqlx::query_scalar(&format!( + "SELECT instance_id FROM {schema}._submissions + WHERE instance_id OPERATOR(pg_catalog.>) $1 + ORDER BY instance_id LIMIT $2" + )) + .bind(cursor.as_str()) + .bind(i64::from(RECLAIM_BATCH)) + .fetch_all(&mut *tx) + .await?; + if !ids.is_empty() { + // Enqueue and the ledger row commit together; acknowledgment atomically + // replaces queued work with engine state. Neither gap is a cleanup window. + let no_work = [ + "instances", + "executions", + "history", + "orchestrator_queue", + "worker_queue", + "instance_locks", + ] + .iter() + .map(|table| { + format!( + "NOT EXISTS (SELECT 1 FROM {schema}.{table} AS work + WHERE work.instance_id OPERATOR(pg_catalog.=) submission.instance_id)" + ) + }) + .collect::>() + .join(" AND "); + sqlx::query(&format!( + "DELETE FROM {schema}._submissions AS submission + WHERE submission.instance_id OPERATOR(pg_catalog.=) ANY($1::pg_catalog.text[]) + AND {no_work}" + )) + .bind(&ids) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + let more = ids.len() == RECLAIM_BATCH as usize; + *cursor = if more { + ids.last().cloned().expect("a full page has a last ID") + } else { + String::new() + }; + Ok(more) +} + /// Return only roots whose engine state was deleted or is already absent. /// Bulk deletion cannot establish this: it silently skips running instances. async fn retire_engine_records(client: &Client, ids: &[String]) -> Vec { @@ -2285,6 +2379,7 @@ mod tests { CREATE TABLE maintenance_candidate_test.worker_queue(instance_id text); CREATE TABLE maintenance_candidate_test.orchestrator_queue(instance_id text); CREATE TABLE maintenance_candidate_test.instance_locks(instance_id text); + CREATE TABLE maintenance_candidate_test._submissions(instance_id text PRIMARY KEY); INSERT INTO maintenance_candidate_test.instances SELECT 'pgdf-42-00000000000000000000000000000001-' || lpad(to_hex(n),8,'0'),1,NULL FROM generate_series(1,3500) n; INSERT INTO maintenance_candidate_test.instances @@ -2302,6 +2397,34 @@ mod tests { assert!(legacy_orphan_candidates(&pool, "maintenance_candidate_test", &mut cursor).await.unwrap().is_empty()); assert!(cursor.is_empty()); + sqlx::query("INSERT INTO maintenance_candidate_test._submissions + SELECT 'quota-' || lpad(n::text, 8, '0') FROM generate_series(1,2500) n") + .execute(&mut connection).await.unwrap(); + let mut submission_cursor = String::new(); + let mut pages = 1; + while cleanup_submissions(&pool, "maintenance_candidate_test", &mut submission_cursor).await.unwrap() { + pages += 1; + assert!(pages <= 3); + } + assert_eq!(pages, 3); + assert!(submission_cursor.is_empty()); + for table in ["instances","executions","history","orchestrator_queue","worker_queue","instance_locks"] { + sqlx::query("INSERT INTO maintenance_candidate_test._submissions VALUES ('quota-guard')") + .execute(&mut connection).await.unwrap(); + sqlx::query(&format!("INSERT INTO maintenance_candidate_test.{table}(instance_id) VALUES ('quota-guard')")) + .execute(&mut connection).await.unwrap(); + assert!(!cleanup_submissions(&pool, "maintenance_candidate_test", &mut submission_cursor).await.unwrap()); + let count: i64 = sqlx::query_scalar("SELECT count(*) FROM maintenance_candidate_test._submissions") + .fetch_one(&mut connection).await.unwrap(); + assert_eq!(count, 1, "retain submission while {table} has work"); + sqlx::query(&format!("DELETE FROM maintenance_candidate_test.{table} WHERE instance_id = 'quota-guard'")) + .execute(&mut connection).await.unwrap(); + assert!(!cleanup_submissions(&pool, "maintenance_candidate_test", &mut submission_cursor).await.unwrap()); + let count: i64 = sqlx::query_scalar("SELECT count(*) FROM maintenance_candidate_test._submissions") + .fetch_one(&mut connection).await.unwrap(); + assert_eq!(count, 0); + } + let removed = Origin { database_oid: u32::MAX, installation_id: uuid::Uuid::from_u128(9) }; sqlx::query("INSERT INTO maintenance_candidate_test._origins VALUES ($1,$2)") .bind(i64::from(removed.database_oid)).bind(removed.installation_id) @@ -2744,6 +2867,7 @@ mod tests { origins: false, engine: false, registrations: false, + submissions: false, }; assert!(!progress.finish_slice(&mut cursor)); assert!(cursor.enter(origin)); @@ -2889,6 +3013,7 @@ mod tests { origins: false, engine: true, registrations: false, + submissions: false, }; assert!(progress.finish_slice(&mut cursor)); for origin in &origins { diff --git a/tests/e2e/sql/91_start_limits_concurrency.sql b/tests/e2e/sql/91_start_limits_concurrency.sql index 537816d3..5048855a 100644 --- a/tests/e2e/sql/91_start_limits_concurrency.sql +++ b/tests/e2e/sql/91_start_limits_concurrency.sql @@ -70,6 +70,9 @@ DECLARE BEGIN PERFORM dblink_connect('quota_holder', connstr); PERFORM dblink_connect('quota_waiter', connstr); + PERFORM dblink_connect('quota_rollback_fence', format( + 'host=localhost dbname=%L port=%s user=postgres', + current_database(), current_setting('port'))); PERFORM dblink_exec('quota_holder', 'SET statement_timeout = ''15s'''); PERFORM dblink_exec('quota_waiter', 'SET statement_timeout = ''15s'''); SELECT pid INTO waiter FROM dblink('quota_waiter', 'SELECT pg_backend_pid()') AS r(pid INTEGER); @@ -119,6 +122,8 @@ BEGIN PERFORM public._e2e_wait_for_engine_terminal(id) FROM df.instances; END $cancel$; $sql$); + PERFORM dblink_exec('quota_rollback_fence', format( + 'BEGIN; LOCK TABLE %I.instances IN SHARE MODE', df.duroxide_schema())); PERFORM dblink_exec('quota_holder', 'BEGIN'); PERFORM dblink_exec('quota_holder', 'SAVEPOINT quota_start'); SELECT id INTO instance FROM dblink('quota_holder', @@ -128,10 +133,15 @@ BEGIN END IF; PERFORM pg_temp.quota_wait_for_lock(waiter); PERFORM dblink_exec('quota_holder', rollback_command); - PERFORM pg_temp.quota_finish('quota_waiter', false); + -- The transaction lock is released, but the rolled-back engine work + -- still consumes active capacity until its rollback is processed. + PERFORM pg_temp.quota_finish('quota_waiter', true); IF rollback_command <> 'ROLLBACK' THEN PERFORM dblink_exec('quota_holder', 'COMMIT'); END IF; + PERFORM dblink_exec('quota_rollback_fence', 'COMMIT'); + PERFORM public._e2e_wait_for_engine_terminal(instance); + PERFORM * FROM dblink('quota_waiter', 'SELECT df.start(df.sleep(300))') AS r(id TEXT); IF EXISTS (SELECT 1 FROM df.instances WHERE id = instance) OR EXISTS (SELECT 1 FROM df.nodes WHERE instance_id = instance) THEN RAISE EXCEPTION 'TEST FAILED: rolled-back start retained metadata'; @@ -155,6 +165,7 @@ BEGIN PERFORM dblink_exec('quota_holder', 'ROLLBACK'); PERFORM dblink_disconnect('quota_holder'); PERFORM dblink_disconnect('quota_waiter'); + PERFORM dblink_disconnect('quota_rollback_fence'); END $$; SELECT df.cancel(id) FROM df.instances diff --git a/tests/e2e/sql/92_start_limits_trust.sql b/tests/e2e/sql/92_start_limits_trust.sql index 346a4595..d7aa590c 100644 --- a/tests/e2e/sql/92_start_limits_trust.sql +++ b/tests/e2e/sql/92_start_limits_trust.sql @@ -36,6 +36,7 @@ BEGIN PERFORM dblink_connect('quota_trust_user', base_conn || ' user=quota_trust_user'); PERFORM dblink_connect('quota_trust_attacker', base_conn || ' user=quota_trust_attacker'); PERFORM dblink_connect('quota_trust_fence', base_conn || ' user=postgres'); + PERFORM dblink_exec('quota_trust_fence', 'SET lock_timeout = ''10s'''); PERFORM dblink_exec('quota_trust_attacker', format($remote$ DO $lock$ BEGIN @@ -54,8 +55,11 @@ BEGIN SELECT id INTO instance FROM dblink('quota_trust_user', format( 'SELECT df.start(df.sleep(300), ''quota-trust-live'', NULL, %L)', mode )) AS started(id TEXT); - EXECUTE format('SELECT status FROM %I.get_instance_info($1)', df.duroxide_schema()) - INTO engine_status USING instance; + -- Keep engine observations on the fence connection: the driver's DO + -- transaction must not retain AccessShareLock and block the later fence. + SELECT observed.status INTO engine_status FROM dblink('quota_trust_fence', + format('SELECT status FROM %I.get_instance_info(%L)', df.duroxide_schema(), instance) + ) AS observed(status TEXT); IF engine_status IS NOT NULL THEN RAISE EXCEPTION 'TEST FAILED: acknowledgment fence did not hold queued start'; END IF; @@ -71,8 +75,9 @@ BEGIN $remote$, mode)); PERFORM dblink_exec('quota_trust_fence', 'COMMIT'); FOR attempt IN 1..300 LOOP - EXECUTE format('SELECT status FROM %I.get_instance_info($1)', df.duroxide_schema()) - INTO engine_status USING instance; + SELECT observed.status INTO engine_status FROM dblink('quota_trust_fence', + format('SELECT status FROM %I.get_instance_info(%L)', df.duroxide_schema(), instance) + ) AS observed(status TEXT); EXIT WHEN engine_status = 'Running'; PERFORM pg_sleep(0.1); END LOOP; @@ -97,8 +102,9 @@ BEGIN END; END $reject$; $remote$, mode)); - EXECUTE format('SELECT status FROM %I.get_instance_info($1)', df.duroxide_schema()) - INTO engine_status USING instance; + SELECT observed.status INTO engine_status FROM dblink('quota_trust_fence', + format('SELECT status FROM %I.get_instance_info(%L)', df.duroxide_schema(), instance) + ) AS observed(status TEXT); IF engine_status IS DISTINCT FROM 'Running' THEN RAISE EXCEPTION 'TEST FAILED: forged display status changed live engine state'; END IF; @@ -123,7 +129,7 @@ BEGIN PERFORM pg_stat_clear_snapshot(); IF EXISTS ( SELECT 1 FROM pg_stat_activity - WHERE state = 'idle in transaction (aborted)' AND query LIKE '%AS requested(id)%' + WHERE state = 'idle in transaction (aborted)' AND query LIKE '%FROM requested%' ) THEN RAISE EXCEPTION 'TEST FAILED: failed quota lookup left an aborted provider transaction'; END IF; @@ -147,12 +153,15 @@ BEGIN END $reject$; $remote$, mode)); PERFORM dblink_exec('quota_trust_fence', 'COMMIT'); - PERFORM public._e2e_wait_for_engine_terminal(instance); + PERFORM * FROM dblink('quota_trust_fence', + format('SELECT public._e2e_wait_for_engine_terminal(%L)', instance)) AS terminal(status TEXT); SELECT id INTO replacement FROM dblink('quota_trust_user', format( 'SELECT df.start(''SELECT 1'', ''quota-trust-replacement'', NULL, %L)', mode )) AS started(id TEXT); - IF public._e2e_wait_for_engine_terminal(replacement) <> 'Completed' THEN + SELECT terminal.status INTO engine_status FROM dblink('quota_trust_fence', + format('SELECT public._e2e_wait_for_engine_terminal(%L)', replacement)) AS terminal(status TEXT); + IF engine_status <> 'Completed' THEN RAISE EXCEPTION 'TEST FAILED: capacity did not recover after cancellation'; END IF; END LOOP; diff --git a/tests/e2e/sql/94_start_limits_rollback.sql b/tests/e2e/sql/94_start_limits_rollback.sql new file mode 100644 index 00000000..d92f1aa7 --- /dev/null +++ b/tests/e2e/sql/94_start_limits_rollback.sql @@ -0,0 +1,167 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +DO $$ +DECLARE + role_name TEXT; +BEGIN + FOREACH role_name IN ARRAY ARRAY['quota_rollback_tx', 'quota_rollback_sp', 'quota_rollback_atomic'] LOOP + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = role_name) THEN + EXECUTE format('DROP OWNED BY %I', role_name); + EXECUTE format('DROP ROLE %I', role_name); + END IF; + END LOOP; +END $$; +CREATE ROLE quota_rollback_tx LOGIN; +CREATE ROLE quota_rollback_sp LOGIN; +CREATE ROLE quota_rollback_atomic LOGIN; +SELECT df.grant_usage(role_name) +FROM (VALUES ('quota_rollback_tx'), ('quota_rollback_sp'), ('quota_rollback_atomic')) AS roles(role_name); +SELECT format('ALTER ROLE %I IN DATABASE %I SET df.max_concurrent_per_user = 2', + role_name, current_database()) +FROM (VALUES ('quota_rollback_tx'), ('quota_rollback_sp'), ('quota_rollback_atomic')) AS roles(role_name) \gexec +SELECT format('ALTER ROLE %I IN DATABASE %I SET df.max_instances_per_user = 3', + role_name, current_database()) +FROM (VALUES ('quota_rollback_tx'), ('quota_rollback_sp'), ('quota_rollback_atomic')) AS roles(role_name) \gexec + +CREATE TEMP TABLE _rollback_ids (owner_name TEXT, id TEXT); + +CREATE FUNCTION pg_temp.expect_rollback_quota(conn TEXT, kind TEXT, n INT, mode TEXT DEFAULT 'caller') +RETURNS VOID LANGUAGE plpgsql AS $$ +BEGIN + PERFORM dblink_exec(conn, format($remote$ + DO $reject$ + DECLARE + hint TEXT; + BEGIN + BEGIN + PERFORM df.start('SELECT 1', 'quota-rollback-rejected', NULL, %1$L); + RAISE EXCEPTION 'TEST FAILED: rollback bypassed %2$s quota'; + EXCEPTION WHEN configuration_limit_exceeded THEN + GET STACKED DIAGNOSTICS hint = PG_EXCEPTION_HINT; + IF SQLERRM NOT LIKE '%%has %3$s %2$s instances (limit %3$s)%%' + OR hint IS NULL OR hint = '' THEN + RAISE EXCEPTION 'TEST FAILED: unexpected quota diagnostics: %%, %%', SQLERRM, hint; + END IF; + END; + END $reject$; + $remote$, mode, kind, n)); +END $$; + +DO $$ +DECLARE + role_name TEXT; + inst TEXT; + mode TEXT; + count_rows BIGINT; + is_savepoint BOOLEAN; + connbase TEXT := format('host=localhost dbname=%L port=%s', + current_database(), current_setting('port')); +BEGIN + PERFORM dblink_connect('quota_rollback_fence', connbase || ' user=postgres'); + FOREACH role_name IN ARRAY ARRAY['quota_rollback_tx', 'quota_rollback_sp'] LOOP + is_savepoint := role_name = 'quota_rollback_sp'; + PERFORM dblink_connect(role_name, connbase || ' user=' || role_name); + PERFORM dblink_exec(role_name, 'SET statement_timeout = ''15s'''); + PERFORM dblink_exec('quota_rollback_fence', format( + 'BEGIN; LOCK TABLE %I.instances IN SHARE MODE', df.duroxide_schema())); + IF is_savepoint THEN + PERFORM dblink_exec(role_name, 'BEGIN'); + END IF; + FOR attempt IN 1..2 LOOP + PERFORM dblink_exec(role_name, CASE WHEN is_savepoint THEN 'SAVEPOINT submission' ELSE 'BEGIN' END); + SELECT id INTO inst FROM dblink(role_name, + 'SELECT df.start(''SELECT 1'', ''quota-rollback-accepted'')') AS started(id TEXT); + INSERT INTO _rollback_ids VALUES (role_name, inst); + PERFORM dblink_exec(role_name, CASE WHEN is_savepoint + THEN 'ROLLBACK TO submission; RELEASE submission' ELSE 'ROLLBACK' END); + END LOOP; + + FOR attempt IN 1..3 LOOP + PERFORM pg_temp.expect_rollback_quota(role_name, 'active', 2); + END LOOP; + IF is_savepoint THEN + PERFORM dblink_exec(role_name, 'COMMIT'); + END IF; + IF EXISTS (SELECT 1 FROM df.instances WHERE submitted_by = role_name::regrole) THEN + RAISE EXCEPTION 'TEST FAILED: rollback left visible df rows'; + END IF; + EXECUTE format('SELECT count(*) FROM %I._submissions WHERE submitted_by = $1', df.duroxide_schema()) + INTO count_rows USING role_name::regrole::oid; + IF count_rows <> 2 THEN + RAISE EXCEPTION 'TEST FAILED: expected two protected rollback charges, got %', count_rows; + END IF; + EXECUTE format('SELECT count(*) FROM %I.orchestrator_queue WHERE work_item LIKE $1', df.duroxide_schema()) + INTO count_rows USING '%quota-rollback-rejected%'; + IF count_rows <> 0 THEN + RAISE EXCEPTION 'TEST FAILED: over-quota rollback attempt enqueued work'; + END IF; + PERFORM dblink_exec('quota_rollback_fence', 'COMMIT'); + PERFORM public._e2e_wait_for_engine_terminal(id) FROM _rollback_ids AS saved WHERE saved.owner_name = role_name; + + -- Terminal orphan histories no longer use active capacity but still + -- consume the retained quota, in both caller and separate-session mode. + PERFORM dblink_exec(role_name, 'BEGIN'); + SELECT id INTO inst FROM dblink(role_name, + 'SELECT df.start(''SELECT 1'', ''quota-rollback-retained'')') AS started(id TEXT); + INSERT INTO _rollback_ids VALUES (role_name, inst); + PERFORM dblink_exec(role_name, 'ROLLBACK'); + PERFORM public._e2e_wait_for_engine_terminal(inst); + FOREACH mode IN ARRAY ARRAY['caller', 'new'] LOOP + FOR attempt IN 1..3 LOOP + PERFORM pg_temp.expect_rollback_quota(role_name, 'retained', 3, mode); + END LOOP; + END LOOP; + EXECUTE format('SELECT count(*) FROM %I._submissions WHERE submitted_by = $1', df.duroxide_schema()) + INTO count_rows USING role_name::regrole::oid; + IF count_rows <> 3 THEN + RAISE EXCEPTION 'TEST FAILED: rejected attempts altered protected accounting'; + END IF; + PERFORM dblink_disconnect(role_name); + END LOOP; + + -- An enqueue failure must roll back its ledger write on the same connection. + PERFORM dblink_connect('quota_rollback_atomic', connbase || ' user=quota_rollback_atomic'); + PERFORM dblink_exec('quota_rollback_fence', format( + 'BEGIN; LOCK TABLE %I.orchestrator_queue IN SHARE MODE', df.duroxide_schema())); + PERFORM dblink_exec('quota_rollback_atomic', $remote$ + DO $failure$ + BEGIN + BEGIN + PERFORM df.start('SELECT 1', 'quota-rollback-enqueue-failed'); + RAISE EXCEPTION 'TEST FAILED: enqueue fence did not block'; + EXCEPTION WHEN internal_error THEN + IF SQLERRM NOT LIKE '%Failed to enqueue durable submission:%lock timeout%' THEN + RAISE; + END IF; + END; + END $failure$; + $remote$); + PERFORM dblink_exec('quota_rollback_fence', 'COMMIT'); + EXECUTE format('SELECT count(*) FROM %I._submissions WHERE submitted_by = $1', df.duroxide_schema()) + INTO count_rows USING 'quota_rollback_atomic'::regrole::oid; + IF count_rows <> 0 OR EXISTS (SELECT 1 FROM df.instances WHERE label = 'quota-rollback-enqueue-failed') THEN + RAISE EXCEPTION 'TEST FAILED: failed enqueue retained accounting or metadata'; + END IF; + SELECT id INTO inst FROM dblink('quota_rollback_atomic', 'SELECT df.start(''SELECT 1'')') AS started(id TEXT); + PERFORM public._e2e_wait_for_engine_terminal(inst); + PERFORM dblink_disconnect('quota_rollback_atomic'); + PERFORM dblink_disconnect('quota_rollback_fence'); +END $$; + +DO $$ +DECLARE + role_name TEXT; +BEGIN + FOREACH role_name IN ARRAY ARRAY['quota_rollback_tx', 'quota_rollback_sp', 'quota_rollback_atomic'] LOOP + IF has_table_privilege(role_name, format('%I._submissions', df.duroxide_schema()), + 'SELECT, INSERT, UPDATE, DELETE, TRUNCATE') THEN + RAISE EXCEPTION 'TEST FAILED: ordinary user can access protected submission accounting'; + END IF; + END LOOP; +END $$; +DROP TABLE _rollback_ids; +DROP FUNCTION pg_temp.expect_rollback_quota(TEXT, TEXT, INT, TEXT); +DROP OWNED BY quota_rollback_tx, quota_rollback_sp, quota_rollback_atomic; +DROP ROLE quota_rollback_tx, quota_rollback_sp, quota_rollback_atomic; +SELECT 'TEST PASSED' AS result; diff --git a/tests/e2e/sql/95_start_limits_rollback_cleanup.sql b/tests/e2e/sql/95_start_limits_rollback_cleanup.sql new file mode 100644 index 00000000..7c48c05e --- /dev/null +++ b/tests/e2e/sql/95_start_limits_rollback_cleanup.sql @@ -0,0 +1,85 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +-- Runs with retention_days = 0 and reconcile_interval = 2. +DROP DATABASE IF EXISTS _quota_cleanup_origin; +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'quota_cleanup_user') THEN + DROP OWNED BY quota_cleanup_user; + DROP ROLE quota_cleanup_user; + END IF; +END $$; +CREATE ROLE quota_cleanup_user LOGIN; +SELECT df.grant_usage('quota_cleanup_user'); +ALTER ROLE quota_cleanup_user SET df.max_concurrent_per_user = 0; +ALTER ROLE quota_cleanup_user SET df.max_instances_per_user = 1; +CREATE DATABASE _quota_cleanup_origin; + +DO $$ +DECLARE + connbase TEXT := format('host=localhost port=%s', current_setting('port')); + connection TEXT; + local_id TEXT; + count_rows BIGINT; +BEGIN + PERFORM dblink_connect('quota_cleanup_control', connbase || format(' dbname=%L user=quota_cleanup_user', current_database())); + PERFORM dblink_connect('quota_cleanup_origin', connbase || ' dbname=_quota_cleanup_origin user=postgres'); + PERFORM dblink_exec('quota_cleanup_origin', $remote$ + CREATE EXTENSION pg_durable; + DO $grant$ BEGIN PERFORM df.grant_usage('quota_cleanup_user'); END $grant$; + $remote$); + PERFORM dblink_disconnect('quota_cleanup_origin'); + PERFORM dblink_connect('quota_cleanup_origin', connbase || ' dbname=_quota_cleanup_origin user=quota_cleanup_user'); + PERFORM dblink_connect('quota_cleanup_fence', connbase || format(' dbname=%L user=postgres', current_database())); + PERFORM dblink_exec('quota_cleanup_fence', format( + 'BEGIN; LOCK TABLE %I.instances IN SHARE MODE', df.duroxide_schema())); + + -- One role may submit in both origins, but cannot discard either origin's + -- accounting by rolling back its local metadata. + FOREACH connection IN ARRAY ARRAY['quota_cleanup_control', 'quota_cleanup_origin'] LOOP + PERFORM dblink_exec(connection, 'BEGIN'); + SELECT id INTO local_id FROM dblink(connection, + 'SELECT df.start(''SELECT 1'', ''quota-cleanup-rollback'')') AS started(id TEXT); + PERFORM dblink_exec(connection, 'ROLLBACK'); + PERFORM dblink_exec(connection, $remote$ + DO $reject$ + BEGIN + BEGIN + PERFORM df.start('SELECT 1'); + RAISE EXCEPTION 'TEST FAILED: rollback released retained quota'; + EXCEPTION WHEN configuration_limit_exceeded THEN + IF SQLERRM NOT LIKE '%has 1 retained instances (limit 1)%' THEN RAISE; END IF; + END; + END $reject$; + $remote$); + END LOOP; + PERFORM pg_sleep(3); + EXECUTE format('SELECT count(*) FROM %I._submissions WHERE submitted_by = $1', df.duroxide_schema()) + INTO count_rows USING 'quota_cleanup_user'::regrole::oid; + IF count_rows <> 2 THEN + RAISE EXCEPTION 'TEST FAILED: reconciliation removed queued accounting: %', count_rows; + END IF; + PERFORM dblink_exec('quota_cleanup_fence', 'COMMIT'); + + FOR attempt IN 1..600 LOOP + EXECUTE format('SELECT count(*) FROM %I._submissions WHERE submitted_by = $1', df.duroxide_schema()) + INTO count_rows USING 'quota_cleanup_user'::regrole::oid; + EXIT WHEN count_rows = 0; + PERFORM pg_sleep(0.1); + END LOOP; + IF count_rows <> 0 THEN + RAISE EXCEPTION 'TEST FAILED: reconciled rollback accounting still consumes quota: %', count_rows; + END IF; + FOREACH connection IN ARRAY ARRAY['quota_cleanup_control', 'quota_cleanup_origin'] LOOP + SELECT id INTO local_id FROM dblink(connection, + 'SELECT df.start(df.sleep(300), ''quota-cleanup-replacement'', transaction_mode => ''new'')') AS started(id TEXT); + PERFORM * FROM dblink(connection, format('SELECT df.cancel(%L)', local_id)) AS cancelled(message TEXT); + PERFORM dblink_disconnect(connection); + END LOOP; + PERFORM dblink_disconnect('quota_cleanup_fence'); +END $$; +DROP DATABASE _quota_cleanup_origin; +DROP OWNED BY quota_cleanup_user; +DROP ROLE quota_cleanup_user; +SELECT 'TEST PASSED' AS result; From 30bcbd009749f4af6d70c651c19103e845a6f63f Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Wed, 28 Oct 2026 17:52:20 +0000 Subject: [PATCH 4/5] Review fixes --- USER_GUIDE.md | 8 +- docs/upgrade-testing.md | 9 +- scripts/test-e2e-docker.sh | 4 +- scripts/test-e2e-local.sh | 2 +- src/client.rs | 338 ++++++++++++++---- src/dsl.rs | 74 +++- src/worker.rs | 154 ++++++-- tests/e2e/sql/94_start_limits_rollback.sql | 17 + ...6_start_limits_failed_metadata_cleanup.sql | 136 +++++++ 9 files changed, 627 insertions(+), 115 deletions(-) create mode 100644 tests/e2e/sql/96_start_limits_failed_metadata_cleanup.sql diff --git a/USER_GUIDE.md b/USER_GUIDE.md index 28d928d6..e9a5fe70 100644 --- a/USER_GUIDE.md +++ b/USER_GUIDE.md @@ -2547,9 +2547,11 @@ worker runs a **best-effort reconciliation pass** that does two things: age past `retention_days`. Anything the engine is still tracking with a live `df` row is left untouched. -After the corresponding engine work and history are gone, reconciliation removes -the protected submission record and frees its retained quota capacity. This -includes submissions whose caller transaction or savepoint rolled back. +After the corresponding engine work, history, and origin metadata are gone, +reconciliation removes the protected submission record and frees its retained +quota capacity. This includes submissions whose caller transaction or savepoint +rolled back. Failed metadata cleanup delays retained-capacity release, but does +not make retired work consume active capacity again. Satellite maintenance visits `_origins` registrations made by activity routing for submitted work, in bounded batches under the shared origin connection budget. diff --git a/docs/upgrade-testing.md b/docs/upgrade-testing.md index 9f52d907..73e62319 100644 --- a/docs/upgrade-testing.md +++ b/docs/upgrade-testing.md @@ -380,7 +380,14 @@ what the upgrade script handles, and any backward compatibility considerations. Admission unions these records with visible legacy instance IDs, avoiding double counting. Bounded reconciliation removes ledger entries only after no corresponding engine instance, execution, history, queue item, or lock - remains. The enqueue WorkItem and orchestration input are unchanged. + remains and origin metadata is confirmed absent. If metadata deletion fails + or the origin cannot be inspected, the ledger preserves trusted retirement + evidence rather than reclassifying stale metadata as active legacy work. + The enqueue WorkItem and orchestration input are unchanged. + ID allocation reserves the ledger key before the local instance row and retries + collisions in either store, including IDs retained after caller rollback. + The ledger transaction stays uncommitted until enqueue; local persistence + failures release the reservation without leaving a charge or queue item. Pre-upgrade orphans that already lost their `df` identity cannot be attributed retroactively; existing orphan cleanup remains responsible for them. - **Behavior change:** New superuser-settable GUCs diff --git a/scripts/test-e2e-docker.sh b/scripts/test-e2e-docker.sh index 18873c29..2c3d0901 100755 --- a/scripts/test-e2e-docker.sh +++ b/scripts/test-e2e-docker.sh @@ -44,11 +44,13 @@ SKIP_TESTS=( "47_http_dsl_disabled" "48_http_allow_all" # Needs the "reconcile" phase GUCs (reconcile_interval=2, retention_days=0) - # so a pass acts within the 30s test window. The single fixed-config Docker + # so a pass acts within the test window. The single fixed-config Docker # container keeps the production defaults (reconcile_interval=3600, # retention_days=30), under which the orphan is never reclaimed in time. "54_reconcile_orphans" "76_multi_database_reconcile" + "95_start_limits_rollback_cleanup" + "96_start_limits_failed_metadata_cleanup" "88_multi_database_maintenance_backlog" # Require the force-drop phase (one execution slot, reconciliation disabled). "78_multi_database_force_drop" diff --git a/scripts/test-e2e-local.sh b/scripts/test-e2e-local.sh index 3ba4d0cb..7f3515e1 100755 --- a/scripts/test-e2e-local.sh +++ b/scripts/test-e2e-local.sh @@ -221,7 +221,7 @@ phase_for_test() { 46_connection_limit_startup_validation) echo "connlimit-startup" ;; - 54_reconcile_orphans|76_multi_database_reconcile|95_start_limits_rollback_cleanup) + 54_reconcile_orphans|76_multi_database_reconcile|95_start_limits_rollback_cleanup|96_start_limits_failed_metadata_cleanup) echo "reconcile" ;; 88_multi_database_maintenance_backlog) diff --git a/src/client.rs b/src/client.rs index 88022b2b..42d88b67 100644 --- a/src/client.rs +++ b/src/client.rs @@ -327,74 +327,48 @@ async fn list_running_descendants(client: &Client, root_instance_id: &str) -> Ve descendants } -/// Start a durable function via the shared PostgreSQL store. -pub fn start_durable_function( - function_name: &str, - instance_id: &str, - input: &str, - submitted_by: pg_sys::Oid, -) -> Result<(), String> { - log!( - "pg_durable: start_durable_function for instance {}", - instance_id - ); - - let origin = crate::origin::backend_origin()?; - let origin_key = crate::origin::engine_id_for_origin(origin.as_ref(), ""); - let inst_id = crate::origin::engine_id_for_origin(origin.as_ref(), instance_id); - let work_item = serde_json::to_string(&WorkItem::StartOrchestration { - instance: inst_id.clone(), - orchestration: function_name.to_string(), - input: input.to_string(), - version: None, - parent_instance: None, - parent_id: None, - parent_execution_id: None, - execution_id: duroxide::INITIAL_EXECUTION_ID, - }) - .map_err(|e| format!("Failed to serialize durable start: {e}"))?; +pub(crate) struct PendingSubmission { + transaction: Option>, + schema: &'static str, + instance_id: String, +} - with_backend_client(|backend, rt| { - rt.block_on(async { - let mut tx = backend - .provider - .pool() - .begin() - .await - .map_err(|e| format!("Failed to begin durable submission: {e}"))?; - let result = async { - crate::origin::configure_metadata_transaction(&mut tx) - .await - .map_err(|e| format!("Failed to configure durable submission: {e}"))?; - sqlx::query(&format!( - "INSERT INTO {schema}._submissions (instance_id, origin, submitted_by) - VALUES ($1, $2, $3::pg_catalog.int8::pg_catalog.oid)", - schema = backend.control.schema, - )) - .bind(&inst_id) - .bind(&origin_key) - .bind(i64::from(submitted_by.to_u32())) - .execute(&mut *tx) - .await - .map_err(|e| format!("Failed to record durable submission: {e}"))?; +impl PendingSubmission { + pub(crate) fn rollback(mut self) -> Result<(), String> { + let tx = self.transaction.take().expect("submission transaction"); + get_client_runtime() + .block_on(tx.rollback()) + .map_err(|error| format!("Failed to release submission reservation: {error}")) + } - // This is the same WorkItem and provider function used by - // Client::start_orchestration, but shares the ledger transaction. - sqlx::query(&format!( - "SELECT {schema}.enqueue_orchestrator_work( - $1, $2, pg_catalog.clock_timestamp(), NULL, NULL, NULL)", - schema = backend.control.schema, - )) - .bind(&inst_id) - .bind(&work_item) - .execute(&mut *tx) - .await - .map_err(|e| format!("Failed to enqueue durable submission: {e}"))?; - Ok(()) - } - .await; + pub(crate) fn enqueue(mut self, function_name: &str, input: &str) -> Result<(), String> { + let work_item = serde_json::to_string(&WorkItem::StartOrchestration { + instance: self.instance_id.clone(), + orchestration: function_name.to_string(), + input: input.to_string(), + version: None, + parent_instance: None, + parent_id: None, + parent_execution_id: None, + execution_id: duroxide::INITIAL_EXECUTION_ID, + }) + .map_err(|e| format!("Failed to serialize durable start: {e}"))?; + let mut tx = self.transaction.take().expect("submission transaction"); + get_client_runtime().block_on(async { + // This is the same WorkItem and provider function used by + // Client::start_orchestration, but shares the reservation transaction. + let result = sqlx::query(&format!( + "SELECT {schema}.enqueue_orchestrator_work( + $1, $2, pg_catalog.clock_timestamp(), NULL, NULL, NULL)", + schema = self.schema, + )) + .bind(&self.instance_id) + .bind(&work_item) + .execute(&mut *tx) + .await + .map_err(|e| format!("Failed to enqueue durable submission: {e}")); match result { - Ok(()) => tx + Ok(_) => tx .commit() .await .map_err(|e| format!("Failed to commit durable submission: {e}")), @@ -406,9 +380,81 @@ pub fn start_durable_function( } } }) + } +} + +impl Drop for PendingSubmission { + fn drop(&mut self) { + if let Some(tx) = self.transaction.take() { + // SPI errors unwind through this guard. Finish rollback while the + // backend runtime is polling, rather than leaving it to async Drop. + if let Err(error) = get_client_runtime().block_on(tx.rollback()) { + warning!("pg_durable: submission reservation rollback failed: {error}"); + } + } + } +} + +pub(crate) fn reserve_submission_id( + instance_id: &str, + submitted_by: pg_sys::Oid, +) -> Result, String> { + let origin = crate::origin::backend_origin()?; + let origin_key = crate::origin::engine_id_for_origin(origin.as_ref(), ""); + let instance_id = crate::origin::engine_id_for_origin(origin.as_ref(), instance_id); + with_backend_client(|backend, rt| { + let transaction = rt.block_on(reserve_submission_transaction( + backend.provider.pool(), + backend.control.schema, + &instance_id, + &origin_key, + submitted_by.to_u32(), + ))?; + Ok(transaction.map(|transaction| PendingSubmission { + transaction: Some(transaction), + schema: backend.control.schema, + instance_id, + })) }) } +async fn reserve_submission_transaction( + pool: &sqlx::PgPool, + schema: &str, + instance_id: &str, + origin: &str, + submitted_by: u32, +) -> Result>, String> { + let mut tx = pool + .begin() + .await + .map_err(|e| format!("Failed to begin durable submission: {e}"))?; + let result = async { + crate::origin::configure_metadata_transaction(&mut tx) + .await + .map_err(|e| format!("Failed to configure durable submission: {e}"))?; + sqlx::query_scalar::<_, String>(&format!( + "INSERT INTO {schema}._submissions (instance_id, origin, submitted_by) + VALUES ($1, $2, $3::pg_catalog.int8::pg_catalog.oid) + ON CONFLICT (instance_id) DO NOTHING RETURNING instance_id" + )) + .bind(instance_id) + .bind(origin) + .bind(i64::from(submitted_by)) + .fetch_optional(&mut *tx) + .await + .map_err(|e| format!("Failed to reserve durable submission: {e}")) + } + .await; + if matches!(&result, Ok(Some(_))) { + return Ok(Some(tx)); + } + tx.rollback().await.map_err(|error| { + format!("Submission reservation rollback failed: {error}; reservation result: {result:?}") + })?; + result.map(|_| None) +} + /// Upper bound (milliseconds) on the `transaction_mode => 'new'` `df.start()` /// statement. /// @@ -642,8 +688,168 @@ pub fn raise_external_event(instance_id: &str, event_name: &str, data: &str) -> }) } -#[cfg(test)] +#[cfg(any(test, feature = "pg_test"))] +#[pg_schema] mod tests { + use super::*; + use crate::types::worker_provider_config; + + #[pg_test] + fn rolled_back_submission_id_is_retried() { + const SCHEMA: &str = "quota_reservation_test"; + const ROLLED_BACK: &str = "c4270001"; + const LOCAL_COLLISION: &str = "c4270002"; + const AVAILABLE: &str = "c4270003"; + let admin = Spi::get_one::("SELECT current_user::text") + .unwrap() + .unwrap(); + let database = Spi::get_one::("SELECT current_database()::text") + .unwrap() + .unwrap(); + let owner = unsafe { pg_sys::GetUserId().to_u32() }; + let rt = get_client_runtime(); + let mut url = url::Url::parse(&postgres_connection_string()).unwrap(); + url.set_username(&admin).unwrap(); + url.set_path(&database); + let provider = rt + .block_on(duroxide_pg::PostgresProvider::new_with_config( + worker_provider_config(url.as_str(), SCHEMA), + )) + .unwrap(); + let pool = provider.pool(); + rt.block_on( + sqlx::query( + "CREATE TABLE quota_reservation_test._submissions ( + instance_id text PRIMARY KEY, origin text NOT NULL, submitted_by oid NOT NULL)", + ) + .execute(pool), + ) + .unwrap(); + let mut caller = rt.block_on(crate::origin::metadata_test_connection(&admin, &database)); + + let insert = "INSERT INTO df.instances (id,root_node,submitted_by) + VALUES ($1,$1,$2::bigint::oid::regrole) + ON CONFLICT (id) DO NOTHING RETURNING id"; + let mut caller_tx = rt.block_on(caller.begin()).unwrap(); + rt.block_on( + sqlx::query(insert) + .bind(ROLLED_BACK) + .bind(i64::from(owner)) + .execute(&mut *caller_tx), + ) + .unwrap(); + let transaction = rt + .block_on(reserve_submission_transaction( + pool, + SCHEMA, + ROLLED_BACK, + "", + owner, + )) + .unwrap() + .unwrap(); + PendingSubmission { + transaction: Some(transaction), + schema: SCHEMA, + instance_id: ROLLED_BACK.into(), + } + .enqueue("quota-test", "original") + .unwrap(); + rt.block_on(caller_tx.rollback()).unwrap(); + let original: String = rt.block_on(sqlx::query_scalar( + "SELECT work_item FROM quota_reservation_test.orchestrator_queue WHERE instance_id = $1", + ).bind(ROLLED_BACK).fetch_one(pool)).unwrap(); + + let mut caller_tx = rt.block_on(caller.begin()).unwrap(); + rt.block_on( + sqlx::query(insert) + .bind(LOCAL_COLLISION) + .bind(i64::from(owner)) + .execute(&mut *caller_tx), + ) + .unwrap(); + let mut candidates = [ROLLED_BACK, LOCAL_COLLISION, AVAILABLE].into_iter(); + let mut reservation = None; + let mut attempted = Vec::new(); + let id = crate::dsl::pick_id_with_retry( + || candidates.next().expect("bounded candidates").to_string(), + |candidate| { + attempted.push(candidate.to_string()); + let Some(transaction) = rt.block_on(reserve_submission_transaction( + pool, SCHEMA, candidate, "", owner, + ))? + else { + return Ok(false); + }; + let pending = PendingSubmission { + transaction: Some(transaction), + schema: SCHEMA, + instance_id: candidate.into(), + }; + let claimed: Option = rt + .block_on( + sqlx::query_scalar(insert) + .bind(candidate) + .bind(i64::from(owner)) + .fetch_optional(&mut *caller_tx), + ) + .map_err(|error| error.to_string())?; + if claimed.is_some() { + reservation = Some(pending); + Ok(true) + } else { + pending.rollback()?; + Ok(false) + } + }, + 3, + ) + .unwrap(); + assert_eq!(id, AVAILABLE); + assert_eq!(attempted, [ROLLED_BACK, LOCAL_COLLISION, AVAILABLE]); + reservation + .unwrap() + .enqueue("quota-test", "replacement") + .unwrap(); + rt.block_on(caller_tx.rollback()).unwrap(); + + let transaction = rt + .block_on(reserve_submission_transaction( + pool, SCHEMA, "c4270004", "", owner, + )) + .unwrap() + .unwrap(); + drop(PendingSubmission { + transaction: Some(transaction), + schema: SCHEMA, + instance_id: "c4270004".into(), + }); + rt.block_on(async { + let ids: Vec = sqlx::query_scalar( + "SELECT instance_id FROM quota_reservation_test._submissions ORDER BY instance_id", + ).fetch_all(pool).await.unwrap(); + assert_eq!(ids, [ROLLED_BACK, AVAILABLE]); + let unchanged: String = sqlx::query_scalar( + "SELECT work_item FROM quota_reservation_test.orchestrator_queue WHERE instance_id = $1", + ).bind(ROLLED_BACK).fetch_one(pool).await.unwrap(); + assert_eq!(unchanged, original); + let queue_count: i64 = sqlx::query_scalar( + "SELECT count(*) FROM quota_reservation_test.orchestrator_queue", + ).fetch_one(pool).await.unwrap(); + assert_eq!(queue_count, 2); + let local_count: i64 = sqlx::query_scalar( + "SELECT count(*) FROM df.instances WHERE id = ANY($1)", + ).bind(&[ROLLED_BACK, LOCAL_COLLISION, AVAILABLE]).fetch_one(&mut caller).await.unwrap(); + assert_eq!(local_count, 0); + sqlx::query("DROP SCHEMA quota_reservation_test CASCADE").execute(&mut caller).await.unwrap(); + pool.close().await; + caller.close().await.unwrap(); + }); + } +} + +#[cfg(test)] +mod unit_tests { use super::{ backend_control_connection_options, format_new_transaction_start_error, is_connection_error, refresh_backend_control_state, CONTROL_LOOKUP_TIMEOUT, diff --git a/src/dsl.rs b/src/dsl.rs index 9ad27ea1..e2e32890 100644 --- a/src/dsl.rs +++ b/src/dsl.rs @@ -12,7 +12,7 @@ use std::str::FromStr; use std::cell::RefCell; use std::time::{Duration, Instant}; -use crate::client::start_durable_function; +use crate::client::reserve_submission_id; use crate::types::{ flatten_graph, get_max_new_transaction_starts, get_new_transaction_start_timeout, mark_non_future_helper_call, short_id, validate_result_name, Durofut, FunctionInput, @@ -950,7 +950,7 @@ const NODE_INSERT_BATCH_SIZE: usize = 1_000; /// The claim — not the generation — is the loop tail, so this only ever returns /// an ID that `try_claim` confirmed was inserted. On exhaustion it returns an /// `Err` rather than a last, unverified candidate (review finding C1). -fn pick_id_with_retry( +pub(crate) fn pick_id_with_retry( mut generate: impl FnMut() -> String, mut try_claim: impl FnMut(&str) -> Result, max_attempts: usize, @@ -1508,8 +1508,14 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< } } - // Reserve the instance ID before inserting nodes so node rows can reference - // it. Collisions on the 8-hex ID space are rare, but we reserve via + // Take local relation locks before opening the independent reservation + // transaction, so extension DROP cannot hold them while waiting on our ledger. + Spi::run("LOCK TABLE ONLY df.instances, ONLY df.nodes IN ACCESS SHARE MODE") + .unwrap_or_else(|error| pgrx::error!("Failed to lock instance metadata: {error}")); + + // Reserve the protected ledger ID first, retaining its transaction through + // enqueue. Then reserve the local ID before nodes can reference it. + // Collisions on the 8-hex ID space are rare, but we reserve via // INSERT ... ON CONFLICT (id) DO NOTHING RETURNING id and re-roll on // collision so a raw primary-key error never reaches the caller (issue // #129). ON CONFLICT arbitration runs against the global id index *below* @@ -1520,9 +1526,27 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< // referenced root node row is inserted (with that ID) before commit — no // post-insert UPDATE is needed (df callers aren't granted UPDATE on // root_node). + let mut submission = None; + #[cfg(any(test, feature = "pg_test"))] + let mut reservation_error = None; let instance_id = match pick_id_with_retry( short_id, |candidate| { + let reserved = match reserve_submission_id(candidate, current_user_oid) { + Ok(Some(reservation)) => Some(reservation), + Ok(None) => return Ok(false), + Err(error) => { + #[cfg(not(any(test, feature = "pg_test")))] + return Err(error); + // The unit-test server has no worker; preserve the existing + // graph-construction tests and report the enqueue error below. + #[cfg(any(test, feature = "pg_test"))] + { + reservation_error = Some(error); + None + } + } + }; let label_arg: DatumWithOid = match label { Some(l) => l.into(), None => DatumWithOid::null::(), @@ -1561,12 +1585,19 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< ], ) }; - Spi::connect_mut( - |client| match client.update(inst_sql, Some(1), &inst_args) { - Ok(table) => Ok(!table.is_empty()), - Err(e) => Err(format!("{e:?}")), - }, - ) + let claimed = + Spi::connect_mut( + |client| match client.update(inst_sql, Some(1), &inst_args) { + Ok(table) => Ok(!table.is_empty()), + Err(e) => Err(format!("{e:?}")), + }, + )?; + if claimed { + submission = reserved; + } else if let Some(reserved) = reserved { + reserved.rollback()?; + } + Ok(claimed) }, MAX_ID_ATTEMPTS, ) { @@ -1631,12 +1662,23 @@ fn start_in_caller_transaction(fut: &str, label: Option<&str>, database: Option< }; let input_json = serde_json::to_string(&input).unwrap_or(instance_id.clone()); - if let Err(e) = start_durable_function( - crate::orchestrations::execute_function_graph::NAME, - &instance_id, - &input_json, - current_user_oid, - ) { + let start_result = match submission { + Some(submission) => submission.enqueue( + crate::orchestrations::execute_function_graph::NAME, + &input_json, + ), + None => { + #[cfg(not(any(test, feature = "pg_test")))] + { + Err("Missing protected submission reservation".into()) + } + #[cfg(any(test, feature = "pg_test"))] + { + Err(reservation_error.unwrap_or_else(|| "Missing submission reservation".into())) + } + } + }; + if let Err(e) = start_result { // Fail fast: the durable engine could not accept the start, so abort the // whole df.start() transaction rather than committing an instance row // that would never run (a stuck instance nothing recovers). The df rows diff --git a/src/worker.rs b/src/worker.rs index 051def57..274323ac 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -2114,9 +2114,11 @@ async fn cleanup_submissions( pool: &sqlx::PgPool, schema: &str, cursor: &mut String, -) -> Result { - let mut tx = pool.begin().await?; - crate::origin::configure_metadata_transaction(&mut tx).await?; +) -> Result { + let mut tx = pool.begin().await.map_err(|error| error.to_string())?; + crate::origin::configure_metadata_transaction(&mut tx) + .await + .map_err(|error| error.to_string())?; let ids: Vec = sqlx::query_scalar(&format!( "SELECT instance_id FROM {schema}._submissions WHERE instance_id OPERATOR(pg_catalog.>) $1 @@ -2125,37 +2127,96 @@ async fn cleanup_submissions( .bind(cursor.as_str()) .bind(i64::from(RECLAIM_BATCH)) .fetch_all(&mut *tx) - .await?; - if !ids.is_empty() { - // Enqueue and the ledger row commit together; acknowledgment atomically - // replaces queued work with engine state. Neither gap is a cleanup window. - let no_work = [ - "instances", - "executions", - "history", - "orchestrator_queue", - "worker_queue", - "instance_locks", - ] - .iter() - .map(|table| { - format!( - "NOT EXISTS (SELECT 1 FROM {schema}.{table} AS work - WHERE work.instance_id OPERATOR(pg_catalog.=) submission.instance_id)" - ) - }) - .collect::>() - .join(" AND "); + .await + .map_err(|error| error.to_string())?; + // Enqueue and the ledger row commit together; acknowledgment atomically + // replaces queued work with engine state. Neither gap is a cleanup window. + let no_work = [ + "instances", + "executions", + "history", + "orchestrator_queue", + "worker_queue", + "instance_locks", + ] + .iter() + .map(|table| { + format!( + "NOT EXISTS (SELECT 1 FROM {schema}.{table} AS work + WHERE work.instance_id OPERATOR(pg_catalog.=) submission.instance_id)" + ) + }) + .collect::>() + .join(" AND "); + let retired: Vec = sqlx::query_scalar(&format!( + "SELECT submission.instance_id FROM {schema}._submissions AS submission + WHERE submission.instance_id OPERATOR(pg_catalog.=) ANY($1::pg_catalog.text[]) + AND {no_work}" + )) + .bind(&ids) + .fetch_all(&mut *tx) + .await + .map_err(|error| error.to_string())?; + tx.commit().await.map_err(|error| error.to_string())?; + + let mut by_origin = BTreeMap::, Vec>::new(); + for id in retired { + let origin = Origin::from_engine_id(&id)? + .map(|origin| (origin.database_oid, origin.installation_id)); + by_origin.entry(origin).or_default().push(id); + } + let router = Router::new(Arc::new(pool.clone())); + let mut removable = Vec::new(); + for (origin, candidates) in by_origin { + let result = if let Some((database_oid, installation_id)) = origin { + let origin = Origin { + database_oid, + installation_id, + }; + tokio::time::timeout(ORIGIN_OPERATION_TIMEOUT, async { + match router.connect(&origin).await { + Ok(route) => { + let result = + submissions_without_metadata(&route.pool, Some(&origin), &candidates) + .await; + route.close().await; + result + } + Err(error) => { + if origin_is_removed(pool, &origin).await? { + Ok(candidates.clone()) + } else { + Err(error) + } + } + } + }) + .await + .unwrap_or_else(|_| Err("Submission metadata lookup timed out".into())) + } else { + submissions_without_metadata(pool, None, &candidates).await + }; + match result { + Ok(ids) => removable.extend(ids), + Err(error) => log!("pg_durable: retaining submission accounting: {error}"), + } + } + if !removable.is_empty() { + let mut tx = pool.begin().await.map_err(|error| error.to_string())?; + crate::origin::configure_metadata_transaction(&mut tx) + .await + .map_err(|error| error.to_string())?; sqlx::query(&format!( "DELETE FROM {schema}._submissions AS submission WHERE submission.instance_id OPERATOR(pg_catalog.=) ANY($1::pg_catalog.text[]) AND {no_work}" )) - .bind(&ids) + .bind(&removable) .execute(&mut *tx) - .await?; + .await + .map_err(|error| error.to_string())?; + tx.commit().await.map_err(|error| error.to_string())?; } - tx.commit().await?; let more = ids.len() == RECLAIM_BATCH as usize; *cursor = if more { ids.last().cloned().expect("a full page has a last ID") @@ -2165,6 +2226,45 @@ async fn cleanup_submissions( Ok(more) } +async fn submissions_without_metadata( + pool: &sqlx::PgPool, + origin: Option<&Origin>, + ids: &[String], +) -> Result, String> { + let local_ids: Vec<&str> = ids + .iter() + .map(|id| match origin { + Some(origin) => origin_local_id(origin, id) + .ok_or_else(|| format!("Invalid submission instance ID: {id}")), + None => Ok(id.as_str()), + }) + .collect::>()?; + let mut tx = crate::origin::begin_metadata(pool, origin) + .await + .map_err(|error| error.to_string())?; + crate::origin::configure_metadata_transaction(&mut tx) + .await + .map_err(|error| error.to_string())?; + let present: HashSet = sqlx::query_scalar( + "SELECT id::pg_catalog.text FROM df.instances + WHERE id OPERATOR(pg_catalog.=) ANY($1::pg_catalog.text[])", + ) + .bind(&local_ids) + .fetch_all(&mut *tx) + .await + .map_err(|error| error.to_string())? + .into_iter() + .collect(); + tx.commit().await.map_err(|error| error.to_string())?; + // A surviving df row still needs the ledger's trusted retirement evidence. + Ok(ids + .iter() + .zip(local_ids) + .filter(|(_, local_id)| !present.contains(*local_id)) + .map(|(id, _)| id.clone()) + .collect()) +} + /// Return only roots whose engine state was deleted or is already absent. /// Bulk deletion cannot establish this: it silently skips running instances. async fn retire_engine_records(client: &Client, ids: &[String]) -> Vec { diff --git a/tests/e2e/sql/94_start_limits_rollback.sql b/tests/e2e/sql/94_start_limits_rollback.sql index d92f1aa7..1bb74aaa 100644 --- a/tests/e2e/sql/94_start_limits_rollback.sql +++ b/tests/e2e/sql/94_start_limits_rollback.sql @@ -143,6 +143,23 @@ BEGIN IF count_rows <> 0 OR EXISTS (SELECT 1 FROM df.instances WHERE label = 'quota-rollback-enqueue-failed') THEN RAISE EXCEPTION 'TEST FAILED: failed enqueue retained accounting or metadata'; END IF; + PERFORM dblink_exec('quota_rollback_fence', 'REVOKE INSERT (id) ON df.nodes FROM quota_rollback_atomic'); + PERFORM dblink_exec('quota_rollback_atomic', $remote$ + DO $failure$ + BEGIN + BEGIN + PERFORM df.start('SELECT 1', 'quota-reservation-persist-failed'); + RAISE EXCEPTION 'TEST FAILED: node persistence unexpectedly succeeded'; + EXCEPTION WHEN insufficient_privilege THEN NULL; + END; + END $failure$; + $remote$); + PERFORM dblink_exec('quota_rollback_fence', 'GRANT INSERT (id) ON df.nodes TO quota_rollback_atomic'); + EXECUTE format('SELECT count(*) FROM %I._submissions WHERE submitted_by = $1', df.duroxide_schema()) + INTO count_rows USING 'quota_rollback_atomic'::regrole::oid; + IF count_rows <> 0 OR EXISTS (SELECT 1 FROM df.instances WHERE label = 'quota-reservation-persist-failed') THEN + RAISE EXCEPTION 'TEST FAILED: SPI error retained a submission reservation or metadata'; + END IF; SELECT id INTO inst FROM dblink('quota_rollback_atomic', 'SELECT df.start(''SELECT 1'')') AS started(id TEXT); PERFORM public._e2e_wait_for_engine_terminal(inst); PERFORM dblink_disconnect('quota_rollback_atomic'); diff --git a/tests/e2e/sql/96_start_limits_failed_metadata_cleanup.sql b/tests/e2e/sql/96_start_limits_failed_metadata_cleanup.sql new file mode 100644 index 00000000..dc0ba9ff --- /dev/null +++ b/tests/e2e/sql/96_start_limits_failed_metadata_cleanup.sql @@ -0,0 +1,136 @@ +-- Copyright (c) Microsoft Corporation. +-- Licensed under the PostgreSQL License. + +-- Runs with retention_days = 0 and reconcile_interval = 2. +DROP DATABASE IF EXISTS _quota_metadata_origin; +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'quota_metadata_user') THEN + DROP OWNED BY quota_metadata_user; + DROP ROLE quota_metadata_user; + END IF; +END $$; +CREATE ROLE quota_metadata_user LOGIN; +CREATE DATABASE _quota_metadata_origin; + +DO $$ +DECLARE + conn TEXT; + database_name TEXT; + origin_key TEXT; + instance TEXT; + canary TEXT; + replacement TEXT; + attempted BOOLEAN; + present BOOLEAN; + base_conn TEXT := format('host=localhost port=%s user=postgres', current_setting('port')); +BEGIN + FOREACH database_name IN ARRAY ARRAY[current_database(), '_quota_metadata_origin'] LOOP + conn := 'quota_metadata_probe'; + PERFORM dblink_connect(conn, base_conn || format(' dbname=%L', database_name)); + IF database_name <> current_database() THEN + PERFORM dblink_exec(conn, 'CREATE EXTENSION pg_durable'); + END IF; + PERFORM dblink_exec(conn, $remote$ + DO $grant$ BEGIN PERFORM df.grant_usage('quota_metadata_user'); END $grant$; + SET df.max_concurrent_per_user = 1; + SET df.max_instances_per_user = 0; + SET lock_timeout = '10s'; + DROP TRIGGER IF EXISTS quota_metadata_block ON df.instances; + DROP FUNCTION IF EXISTS public.quota_metadata_block_delete(); + DROP SEQUENCE IF EXISTS public.quota_metadata_delete_attempt; + CREATE SEQUENCE public.quota_metadata_delete_attempt; + CREATE FUNCTION public.quota_metadata_block_delete() RETURNS trigger + LANGUAGE plpgsql AS $trigger$ + BEGIN + PERFORM nextval('public.quota_metadata_delete_attempt'); + RAISE EXCEPTION 'quota test: metadata deletion blocked'; + END $trigger$; + CREATE TRIGGER quota_metadata_block + BEFORE DELETE ON df.instances FOR EACH ROW + WHEN (OLD.label = 'quota-metadata-preserve') + EXECUTE FUNCTION public.quota_metadata_block_delete(); + SET ROLE quota_metadata_user; + $remote$); + SELECT prefix INTO origin_key FROM dblink(conn, $remote$ + SELECT CASE WHEN current_database() = current_setting('pg_durable.database') THEN '' + ELSE format('pgdf-%s-%s-', + (SELECT oid FROM pg_database WHERE datname = current_database()), + replace(id::text, '-', '')) END + FROM df._installation + $remote$) AS origin(prefix TEXT); + SELECT id INTO instance FROM dblink(conn, + 'SELECT df.start(''SELECT 1'', ''quota-metadata-preserve'')') AS started(id TEXT); + PERFORM dblink_exec(conn, 'RESET ROLE'); + FOR attempt IN 1..600 LOOP + SELECT is_called INTO attempted FROM dblink(conn, + 'SELECT is_called FROM public.quota_metadata_delete_attempt') AS observed(is_called BOOLEAN); + EXIT WHEN attempted; + PERFORM pg_sleep(0.1); + END LOOP; + IF NOT attempted THEN + RAISE EXCEPTION 'TEST FAILED: metadata deletion was not attempted in %', database_name; + END IF; + + -- Removing a later orphan's ledger entry proves accounting cleanup ran + -- after the failed metadata deletion, without relying on a fixed sleep. + PERFORM dblink_exec(conn, 'BEGIN'); + SELECT id INTO canary FROM dblink(conn, + 'SELECT df.start(''SELECT 1'', ''quota-metadata-canary'')') AS started(id TEXT); + PERFORM dblink_exec(conn, 'ROLLBACK'); + FOR attempt IN 1..600 LOOP + EXECUTE format('SELECT EXISTS (SELECT 1 FROM %I._submissions WHERE instance_id = $1)', + df.duroxide_schema()) INTO present USING origin_key || canary; + EXIT WHEN NOT present; + PERFORM pg_sleep(0.1); + END LOOP; + IF present THEN + RAISE EXCEPTION 'TEST FAILED: accounting cleanup did not process the canary'; + END IF; + EXECUTE format('SELECT EXISTS (SELECT 1 FROM %I.instances WHERE instance_id = $1)', + df.duroxide_schema()) INTO present USING origin_key || instance; + IF present THEN + RAISE EXCEPTION 'TEST FAILED: engine retirement did not precede failed metadata cleanup'; + END IF; + EXECUTE format('SELECT EXISTS (SELECT 1 FROM %I._submissions WHERE instance_id = $1)', + df.duroxide_schema()) INTO present USING origin_key || instance; + IF NOT present THEN + RAISE EXCEPTION 'TEST FAILED: failed metadata cleanup lost trusted retirement accounting'; + END IF; + SELECT observed.found INTO present FROM dblink(conn, format( + 'SELECT EXISTS (SELECT 1 FROM df.instances WHERE id = %L)', instance)) AS observed(found BOOLEAN); + IF NOT present THEN + RAISE EXCEPTION 'TEST FAILED: metadata deletion unexpectedly succeeded'; + END IF; + + PERFORM dblink_exec(conn, 'SET ROLE quota_metadata_user'); + SELECT id INTO replacement FROM dblink(conn, + 'SELECT df.start(df.sleep(300), ''quota-metadata-replacement'')') AS started(id TEXT); + PERFORM * FROM dblink(conn, format('SELECT df.cancel(%L)', replacement)) AS cancelled(message TEXT); + PERFORM dblink_exec(conn, $remote$ + RESET ROLE; + DROP TRIGGER quota_metadata_block ON df.instances; + DROP FUNCTION public.quota_metadata_block_delete(); + DROP SEQUENCE public.quota_metadata_delete_attempt; + $remote$); + FOR attempt IN 1..600 LOOP + EXECUTE format('SELECT EXISTS (SELECT 1 FROM %I._submissions WHERE instance_id = $1)', + df.duroxide_schema()) INTO present USING origin_key || instance; + EXIT WHEN NOT present; + PERFORM pg_sleep(0.1); + END LOOP; + IF present THEN + RAISE EXCEPTION 'TEST FAILED: accounting was not released after metadata cleanup recovered'; + END IF; + SELECT observed.found INTO present FROM dblink(conn, format( + 'SELECT EXISTS (SELECT 1 FROM df.instances WHERE id = %L)', instance)) AS observed(found BOOLEAN); + IF present THEN + RAISE EXCEPTION 'TEST FAILED: accounting was released before metadata removal'; + END IF; + PERFORM dblink_disconnect(conn); + END LOOP; +END $$; +DROP DATABASE _quota_metadata_origin; +DROP OWNED BY quota_metadata_user; +DROP ROLE quota_metadata_user; +SELECT 'TEST PASSED' AS result; From e362bdf4114bcf25d92225241b5f60919e935b95 Mon Sep 17 00:00:00 2001 From: Thom Chiovoloni Date: Thu, 29 Oct 2026 15:19:00 +0000 Subject: [PATCH 5/5] Address a few issues --- scripts/test-epoch-race.sh | 99 ++++++++++++++++++++++++++++++- src/client.rs | 23 ++++--- src/worker.rs | 46 +++++++------- tests/e2e/sql/90_start_limits.sql | 31 ++++++++++ 4 files changed, 161 insertions(+), 38 deletions(-) diff --git a/scripts/test-epoch-race.sh b/scripts/test-epoch-race.sh index 046e48c7..4ea1d724 100755 --- a/scripts/test-epoch-race.sh +++ b/scripts/test-epoch-race.sh @@ -14,13 +14,14 @@ # This test makes the race deterministic with a test hook # (PG_DURABLE_TEST_PAUSE_BEFORE_READY_MS, compiled in only under the `test-hooks` # cargo feature) that pauses the worker between epoch capture / runtime -# initialization and readiness publication. The test waits for the hook's own log +# initialization and epoch revalidation. The test waits for the hook's own log # line before acting, so the DROP/CREATE is always issued inside the window, then # verifies: # 1. The worker tears down the stale runtime and reinitializes (log evidence). # 2. The worker eventually becomes ready. # 3. `_worker_ready` describes the CURRENT provider epoch (a durable function runs). # 4. No stale runtime keeps polling provider objects that were dropped. +# 5. Resumed SQL can start a child before final epoch revalidation on upgrade. # # Usage: ./scripts/test-epoch-race.sh [options] # @@ -263,7 +264,7 @@ fi ok "Baseline worker is ready" # 2. Restart with the pause enabled. On restart the worker captures the current -# epoch, initializes a runtime, then pauses before publishing readiness. +# epoch, initializes a runtime, then pauses before revalidating the epoch. info "Restarting server with a ${PAUSE_MS}ms init->ready pause..." stop_server # Record where the current log ends so we only inspect messages from this restart. @@ -276,7 +277,7 @@ fi # 3. Replace the extension epoch while the worker sits in the pause. Waiting for # the hook's own log line — rather than sleeping a guessed amount — is what # makes the race deterministic: the DROP is issued only once the worker has -# provably initialized a runtime and not yet published readiness. +# provably initialized a runtime and not yet revalidated its epoch. if ! wait_for_log "TEST hook — pausing" "$LOG_MARK" 60; then fail "Worker never reached the init->ready pause (is the test-hooks feature enabled?)" stop_server @@ -345,6 +346,98 @@ fi stop_server +# --------------------------------------------------------------------------- +# Scenario: persisted SQL starts a child while an upgraded worker initializes +# --------------------------------------------------------------------------- + +info "=== Upgrade readiness: resume a workflow that starts a child ===" +if ! start_server ""; then + fail "Could not start PostgreSQL to prepare the persisted parent" + exit 1 +fi + +DX_SCHEMA=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -Atqc \ + "SELECT df.duroxide_schema();") +READY_VERSION=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -Atqc \ + "SELECT schema_version FROM ${DX_SCHEMA}._worker_ready;") +NESTED_INSTANCE=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres \ + -v ON_ERROR_STOP=1 -qAt <<'SQL' +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'epoch_race_nested_user') THEN + CREATE ROLE epoch_race_nested_user LOGIN; + END IF; +END $$; +SELECT df.grant_usage('epoch_race_nested_user'); +SET ROLE epoch_race_nested_user; +SELECT df.start( + df.sleep(10) ~> 'SELECT df.start(''SELECT 42'')', + 'epoch-race-nested-start' +); +SQL +) +NESTED_INSTANCE=$(printf '%s' "$NESTED_INSTANCE" | tr -d ' \n') +if [ -z "$NESTED_INSTANCE" ]; then + fail "Could not persist the parent workflow" + exit 1 +fi +TIMER_WAITING="f" +for attempt in $(seq 1 50); do + TIMER_WAITING=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -Atqc \ + "SELECT EXISTS (SELECT 1 FROM ${DX_SCHEMA}.history + WHERE instance_id = '$NESTED_INSTANCE' + AND event_data::jsonb->>'type' = 'TimerCreated') + AND NOT EXISTS (SELECT 1 FROM ${DX_SCHEMA}.history + WHERE instance_id = '$NESTED_INSTANCE' + AND event_data::jsonb->>'type' = 'TimerFired');") + [ "$TIMER_WAITING" = "t" ] && break + sleep 0.1 +done +if [ "$TIMER_WAITING" != "t" ]; then + fail "Parent did not persist its timer before shutdown" + exit 1 +fi + +# The next boot must replace the old readiness version before resuming SQL. +"$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -v ON_ERROR_STOP=1 \ + -qc "UPDATE ${DX_SCHEMA}._worker_ready SET schema_version = schema_version - 1;" +stop_server +sleep 10 +LOG_MARK=$(wc -l < "$LOG_FILE") +if ! start_server "$PAUSE_MS"; then + fail "Could not restart PostgreSQL for the upgrade readiness scenario" + exit 1 +fi +if ! wait_for_log "TEST hook — pausing" "$LOG_MARK" 60; then + fail "Upgraded worker never reached the post-dispatch test pause" + exit 1 +fi +CURRENT_VERSION=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -Atqc \ + "SELECT schema_version FROM ${DX_SCHEMA}._worker_ready;") +if [ "$CURRENT_VERSION" = "$READY_VERSION" ]; then + ok "Upgraded worker published compatible readiness before dispatch" +else + fail "Expected readiness $READY_VERSION, found $CURRENT_VERSION" +fi + +NESTED_STATUS="" +for attempt in $(seq 1 300); do + NESTED_STATUS=$("$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -Atqc \ + "SELECT lower(status) FROM df.instances WHERE id = '$NESTED_INSTANCE';") + case "$NESTED_STATUS" in + completed|failed|cancelled) break ;; + esac + sleep 0.1 +done +if [ "$NESTED_STATUS" = "completed" ]; then + ok "Persisted parent started its child during upgraded worker initialization" +else + fail "Persisted nested-start parent ended with status '$NESTED_STATUS'" +fi +"$PSQL" -h localhost -p "$PG_PORT" -U postgres -d postgres -v ON_ERROR_STOP=1 \ + -qc "DROP OWNED BY epoch_race_nested_user; DROP ROLE epoch_race_nested_user;" +stop_server + # --------------------------------------------------------------------------- # Summary # --------------------------------------------------------------------------- diff --git a/src/client.rs b/src/client.rs index 42d88b67..ae062fa3 100644 --- a/src/client.rs +++ b/src/client.rs @@ -207,23 +207,20 @@ pub(crate) fn submission_counts( .await .map_err(|e| format!("Failed to configure df.start() quota lookup: {e}"))?; - // Old binaries left no submission ledger. Their visible df rows - // still count; grouping charges new committed starts only once. + // Legacy df rows still consume retained capacity; deduplicate + // starts present in both local metadata and protected accounting. let owned = format!( "WITH requested AS ( - SELECT id, pg_catalog.bool_or(recorded) AS recorded - FROM ( - SELECT pg_catalog.unnest($1::pg_catalog.text[]) AS id, false AS recorded - UNION ALL - SELECT instance_id, true FROM {schema}._submissions - WHERE origin OPERATOR(pg_catalog.=) $2 - AND submitted_by OPERATOR(pg_catalog.=) $3::pg_catalog.int8::pg_catalog.oid - ) AS owned GROUP BY id + SELECT pg_catalog.unnest($1::pg_catalog.text[]) AS id + UNION + SELECT instance_id FROM {schema}._submissions + WHERE origin OPERATOR(pg_catalog.=) $2 + AND submitted_by OPERATOR(pg_catalog.=) $3::pg_catalog.int8::pg_catalog.oid )", schema = backend.control.schema, ); - // A recorded start is queued before it materializes. Once both - // are gone, ledger cleanup must not make retired work active again. + // Queued starts may not have an engine instance yet. Once both + // are gone, lingering metadata or accounting is retained-only. let query = if include_active { format!( "{owned} @@ -232,7 +229,7 @@ pub(crate) fn submission_counts( WHERE e.status OPERATOR(pg_catalog.<>) ALL ( ARRAY['Completed', 'Failed']::pg_catalog.text[]) OR (e.status IS NULL AND ( - NOT requested.recorded OR i.instance_id IS NOT NULL OR EXISTS ( + i.instance_id IS NOT NULL OR EXISTS ( SELECT 1 FROM {schema}.orchestrator_queue AS queued WHERE queued.instance_id OPERATOR(pg_catalog.=) requested.id)))) FROM requested diff --git a/src/worker.rs b/src/worker.rs index 274323ac..4bdbca18 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -303,23 +303,24 @@ async fn run_duroxide_runtime(http_policy: Arc) { INIT_RETRY_INTERVAL, &mgmt_pool, &duroxide_schema, + epoch_oid, &http_policy, ) .await else { - // Shutdown requested or extension dropped while initializing. + // Shutdown or the installation changed while initializing. continue; }; // Test-only pause that widens the window between runtime initialization - // and readiness publication so a regression test can drop/recreate the + // and epoch revalidation so a regression test can drop/recreate the // extension deterministically. Compiled out unless `test-hooks` is on. test_pause_before_ready().await; - // Revalidate the epoch AFTER initialization and BEFORE publishing - // readiness. If the extension was dropped/recreated during init, tear down - // this now-stale runtime and retry against the current epoch, so a runtime - // bound to provider objects that no longer exist is never certified. + // Revalidate the epoch AFTER initialization. If the extension was + // dropped/recreated during init, tear down this now-stale runtime and + // retry against the current epoch, so a runtime bound to provider + // objects that no longer exist is never certified. if extension_epoch_replaced(&poll_pool, epoch_oid).await { log!( "pg_durable: extension replaced during initialization — \ @@ -332,18 +333,6 @@ async fn run_duroxide_runtime(http_policy: Arc) { continue; } - // Write the worker readiness record so backend sessions know the - // duroxide schema is fully initialized for this schema version. - // Skipped if the row already has the current WORKER_SCHEMA_VERSION. - if let Err(e) = write_worker_ready(&mgmt_pool, &duroxide_schema, epoch_oid).await { - log!("pg_durable: failed to write worker readiness record: {}", e); - teardown_runtime(duroxide_runtime, duroxide_store).await; - if !sleep_or_shutdown(STALE_RUNTIME_RETRY_INTERVAL).await { - break; - } - continue; - } - // Write a sentinel so we can detect drop+recreate even if the // extension is always present in pg_extension between polls. let epoch_id = match write_epoch_sentinel(&poll_pool).await { @@ -483,7 +472,7 @@ async fn sleep_or_shutdown(duration: Duration) -> bool { /// Test-only hook: when the `PG_DURABLE_TEST_PAUSE_BEFORE_READY_MS` environment /// variable is set to a positive integer, sleep that many milliseconds between -/// runtime initialization and readiness publication. This creates a deterministic +/// runtime initialization and epoch revalidation. This creates a deterministic /// window for `scripts/test-epoch-race.sh` to DROP/CREATE the extension /// mid-initialization and exercise the stale-runtime detection path. /// @@ -495,7 +484,7 @@ async fn test_pause_before_ready() { if let Ok(ms) = val.parse::() { if ms > 0 { log!( - "pg_durable: TEST hook — pausing {}ms before readiness publication", + "pg_durable: TEST hook — pausing {}ms before epoch revalidation", ms ); tokio::time::sleep(Duration::from_millis(ms)).await; @@ -692,6 +681,7 @@ async fn initialize_duroxide_runtime( retry_interval: Duration, mgmt_pool: &sqlx::PgPool, schema_name: &str, + epoch_oid: i64, http_policy: &Arc, ) -> Option<(Arc, Arc)> { log!("pg_durable: initializing duroxide runtime..."); @@ -778,6 +768,18 @@ async fn initialize_duroxide_runtime( } }; + // Resumed workflow SQL can call df.start(), which requires both the + // submission ledger and current readiness before dispatch begins. + if let Err(error) = write_worker_ready(mgmt_pool, schema_name, epoch_oid).await { + log!( + "pg_durable: failed to write worker readiness record: {}", + error + ); + store.pool().close().await; + sleep_or_shutdown(retry_interval).await; + return None; + } + // Reuse the management pool for activities (graph loading, status updates). // The former dedicated activity pool with its df.in_workflow hook is no // longer needed — connect_as_user() sets that flag independently. @@ -812,8 +814,8 @@ async fn write_epoch_sentinel(pool: &sqlx::PgPool) -> Result 1 THEN + RAISE EXCEPTION 'TEST FAILED: could not prepare a live legacy instance'; + END IF; +END $$; + -- Enqueue remains possible, but acknowledgment cannot create an engine instance -- and consume its queue item while this lock is held. BEGIN; @@ -116,6 +137,16 @@ BEGIN END $$; COMMIT; +SET df.max_instances_per_user = 3; +SET ROLE start_limits_user; +SELECT start_limits_test.expect_rejected('retained', 3); +RESET ROLE; +SET df.max_instances_per_user = 0; +BEGIN; +DELETE FROM df.nodes WHERE instance_id = '1e9ac001'; +DELETE FROM df.instances WHERE id = '1e9ac001'; +COMMIT; + SET ROLE start_limits_user; DO $$ BEGIN