diff --git a/CHANGELOG.md b/CHANGELOG.md index 3221202b..5677e7e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,34 @@ 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. + 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. + +> **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 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 5233784e..e9a5fe70 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) --- @@ -2011,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. @@ -2545,6 +2547,12 @@ 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, 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. Registration is not synchronous with `df.start()`; unused installations are not @@ -3003,6 +3011,99 @@ 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 not yet terminal in the durable engine, including queued, sleeping, looping, and signal-waiting 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 +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 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 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 +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 +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 +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 +`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. +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`. +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 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. + +--- + ## 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..b747a36f 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,30 @@ 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 | 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` | + +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. + +--- + ### 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..73e62319 100644 --- a/docs/upgrade-testing.md +++ b/docs/upgrade-testing.md @@ -356,6 +356,67 @@ 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.9 → v0.2.10 (unreleased) + +#### `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 + 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 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 + `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 + 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, 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 + `'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. + 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 The post-tag changes in #379, #388, #389, #390, and #380 belong to the 0.2.9 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 c0e77ccc..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) + 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/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/scripts/test-upgrade.sh b/scripts/test-upgrade.sh index 2544483f..5eb91683 100755 --- a/scripts/test-upgrade.sh +++ b/scripts/test-upgrade.sh @@ -952,6 +952,49 @@ 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'); + UPDATE df.instances SET status = 'completed'; + 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 \$\$; + 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 +} + test_b1_result() { assert_sql_contains "SELECT df.result('${B1_INSTANCE_ID}');" "test_value" } @@ -1146,6 +1189,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/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 c21af802..ae062fa3 100644 --- a/src/client.rs +++ b/src/client.rs @@ -10,10 +10,11 @@ //! 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; +use duroxide::{providers::WorkItem, Client}; +use pgrx::pg_sys::panic::ErrorReport; use pgrx::prelude::*; use sqlx::Connection; use tokio::runtime::Runtime; @@ -26,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) }; } @@ -98,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) { @@ -115,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; } }); @@ -130,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 @@ -141,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) @@ -171,6 +182,96 @@ where result } +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)) + .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}"))?; + + // Legacy df rows still consume retained capacity; deduplicate + // starts present in both local metadata and protected accounting. + let owned = format!( + "WITH requested AS ( + 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, + ); + // 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} + 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 ( + 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", + 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(counts) => { + tx.commit() + .await + .map_err(|e| format!("Failed to finish df.start() quota lookup: {e}"))?; + Ok(counts) + } + 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(); @@ -223,32 +324,134 @@ 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, -) -> Result<(), String> { - log!( - "pg_durable: start_durable_function for instance {}", - instance_id - ); +pub(crate) struct PendingSubmission { + transaction: Option>, + schema: &'static str, + instance_id: String, +} - let fn_name = function_name.to_string(); - let inst_id = crate::origin::backend_engine_id(instance_id)?; - let inp = input.to_string(); +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}")) + } - with_duroxide_client(|client, rt| { - rt.block_on(async { - client - .start_orchestration(&inst_id, &fn_name, &inp) - .await - .map_err(|e| format!("Failed to start durable function: {e:?}"))?; - Ok(()) + 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 + .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) + } + } + }) + } +} + +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. /// @@ -278,12 +481,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 +517,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 +528,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 +542,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 +562,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 +592,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 +603,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; @@ -401,9 +640,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:?}"))?; @@ -418,8 +658,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 @@ -444,8 +685,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 5d8190c3..e2e32890 100644 --- a/src/dsl.rs +++ b/src/dsl.rs @@ -5,13 +5,14 @@ use cron::Schedule as CronSchedule; use pgrx::datum::DatumWithOid; +use pgrx::pg_sys::panic::ErrorReport; use pgrx::prelude::*; 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, @@ -949,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, @@ -1062,6 +1063,113 @@ fn require_login_privilege(oid: pgrx::pg_sys::Oid, name: &str, caller: &str) { } } +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_ids(user_oid: pg_sys::Oid) -> Vec { + Spi::connect(|_| unsafe { + let mut argtypes = [pg_sys::OIDOID]; + let plan = pg_sys::SPI_prepare( + 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(), + 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 mut is_null = false; + let datum = pg_sys::SPI_getbinval(*table.vals, table.tupdesc, 1, &mut is_null); + Vec::::from_datum(datum, is_null) + .unwrap_or_else(|| pgrx::error!("df.start() quota query returned NULL instance IDs")) + }) +} + +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. + acquire_instance_quota_lock(user_oid); + + 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, + format!( + "df.start rejected: user \"{user_name}\" has {count} {kind} instances (limit {limit})" + ), + function_name!(), + ) + .set_hint(hint) + .report(PgLogLevel::ERROR); + }; + if instance_limit > 0 && total >= i64::from(instance_limit) { + reject( + total, + "retained", + instance_limit, + "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 && 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.", + ); + } +} + /// `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 +1351,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 +1392,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( @@ -1396,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* @@ -1408,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::(), @@ -1449,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, ) { @@ -1519,11 +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, - ) { + 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/lib.rs b/src/lib.rs index fba8c1ef..617763b2 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` @@ -93,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); @@ -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 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, + GucContext::Suset, + GucFlags::default(), + ); + + GucRegistry::define_int_guc( + c"df.max_instances_per_user", + 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, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( c"pg_durable.retention_days", c"Days a terminal instance is retained before reconciliation removes it", @@ -449,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/src/worker.rs b/src/worker.rs index e2c287b7..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 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 +1312,7 @@ struct MaintenanceProgress { origins: bool, engine: bool, registrations: bool, + submissions: bool, } impl MaintenanceProgress { @@ -1288,11 +1323,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 +2112,161 @@ 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.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 + ORDER BY instance_id LIMIT $2" + )) + .bind(cursor.as_str()) + .bind(i64::from(RECLAIM_BATCH)) + .fetch_all(&mut *tx) + .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(&removable) + .execute(&mut *tx) + .await + .map_err(|error| error.to_string())?; + tx.commit().await.map_err(|error| error.to_string())?; + } + 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) +} + +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 { @@ -2285,6 +2481,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 +2499,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 +2969,7 @@ mod tests { origins: false, engine: false, registrations: false, + submissions: false, }; assert!(!progress.finish_slice(&mut cursor)); assert!(cursor.enter(origin)); @@ -2889,6 +3115,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/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 2ef03e23..6761cc6d 100644 --- a/tests/e2e/sql/89_schema_qualification.sql +++ b/tests/e2e/sql/89_schema_qualification.sql @@ -63,6 +63,21 @@ CREATE FUNCTION schema_qualification.to_char(timestamptz, text) RETURNS text LAN BEGIN RAISE EXCEPTION 'schema qualification function 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.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 new file mode 100644 index 00000000..73362320 --- /dev/null +++ b/tests/e2e/sql/90_start_limits.sql @@ -0,0 +1,249 @@ +-- 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; +-- Older binaries could commit metadata after an enqueue failure. +BEGIN; +INSERT INTO df.instances (id, root_node, submitted_by) +VALUES ('1e9ac001', '1e9ac001', 'start_limits_user'::regrole); +INSERT INTO df.nodes (id, instance_id, node_type, query, submitted_by) +VALUES ('1e9ac001', '1e9ac001', 'SQL', 'SELECT 1', 'start_limits_user'::regrole); +COMMIT; +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; + +DO $$ +DECLARE + removed INTEGER; +BEGIN + -- Simulate a live legacy start without protected accounting. + EXECUTE format('DELETE FROM %I._submissions WHERE instance_id = $1', + df.duroxide_schema()) + USING (SELECT id FROM start_limits_test.ids WHERE kind = 'initial' ORDER BY id LIMIT 1); + GET DIAGNOSTICS removed = ROW_COUNT; + IF removed <> 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; +DO $$ +BEGIN + 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); +RESET ROLE; +DO $$ +DECLARE + actual BIGINT; +BEGIN + 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 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 + 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 $$; +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 $$; +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'); +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'; +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 $$ +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 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; + +-- 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 = 3) 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 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..5048855a --- /dev/null +++ b/tests/e2e/sql/91_start_limits_concurrency.sql @@ -0,0 +1,176 @@ +-- 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 = 3 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_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); + + 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'); + PERFORM public._e2e_wait_for_engine_terminal(id) FROM df.instances; + 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'); + 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', + '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); + -- 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'; + 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'); + PERFORM public._e2e_wait_for_engine_terminal(id) FROM df.instances; + 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'); + PERFORM dblink_disconnect('quota_rollback_fence'); +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; 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..d7aa590c --- /dev/null +++ b/tests/e2e/sql/92_start_limits_trust.sql @@ -0,0 +1,175 @@ +-- 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_fence', 'SET lock_timeout = ''10s'''); + 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); + -- 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; + 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 + 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; + 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)); + 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; + 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 '%FROM requested%' + ) 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 * 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); + 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; + 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; 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..1bb74aaa --- /dev/null +++ b/tests/e2e/sql/94_start_limits_rollback.sql @@ -0,0 +1,184 @@ +-- 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; + 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'); + 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; 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;