From 069c4b178baba2c7df67d9cc2ff468cf4efed574 Mon Sep 17 00:00:00 2001 From: Scott Robinson Date: Thu, 29 Oct 2026 08:01:56 +0000 Subject: [PATCH] feat(storage): SnapshotEngine, a consistent table snapshot on every backend Adds the one new obligation the engine-owned backup path places on a backend: open a consistent view of one table and stream its items. The trait returns the table definition, a snapshot marker, an ordering flag, and the item stream together, from one operation, so a backup's manifest and data files always describe the same instant. The in-catalog backup code keeps its own copy of this logic for now; the engine-side backup job (next PR) is the first caller. Stacked on #394: SnapshotInfo is the manifest's SnapshotDescriptor (one definition, with the three KIND_* constants), and key_schema_is_orderable is KeyOrder::resolve, so a stream a backend marks ordered is one the manifest parser accepts and a table whose range key has no attribute definition streams unordered and complete instead of paging by the hash key alone. Trait (crates/storage/src/snapshot.rs): open_table_snapshot(account_id, table_name) returns TableSnapshot { description: TableDescription, info: SnapshotInfo, key_ordered: bool, items: BoxStream> }. The description is the storage-level one at the snapshot instant, not the API response (restore_summary may be set, table_throughput_mode_summary is not, the counts are the stats pass's); items are yielded as stored (valid N normalized, invalid raw, TTL-expired unswept items included), and any error means start over, SnapshotExpired being the one worth retrying. The stream owns the backend transaction or session; dropping it releases it. A stream that has not reached its end (including one never polled) must be dropped inside a Tokio runtime: the sqlx backends return its connection through the runtime and panic outside one, MongoDB does not, and the trait doc says so. The first Err ends the stream on every backend. StorageEngine now requires SnapshotEngine; all four in-tree backends implement it. The trait doc names what each backend holds while a stream is open and which errors each returns (Cassandra answers Unsupported before it resolves the table), the idle timeouts that end a stream as Internal, and why the trait has no default body and StorageError stays exhaustive; PostgreSQL and SQLite bound open snapshots per server (below) and refuse over the bound with LimitExceeded. StorageError gains SnapshotExpired for a snapshot the server discarded before the stream drained; the engine maps it to a retryable 503 like Transient. Indexes: the description keeps every index in whatever build state it had and drops only one being deleted (drop_deleting_indexes, for GSIs and vector indexes; no in-tree backend reports a DELETING index today, and unit tests pin the rule); on PostgreSQL a GSI row whose data table is missing is also left out, with the attribute definitions only it used, as the in-catalog backup does (since #401 UpdateTable creates the data table before committing the row, so this is a guard, not an expected state). A restore recreates the indexes in the description and builds them from the items, so a CREATING index with a data table comes back complete; one left out would not come back at all. PostgreSQL, one shared database: the open uses one connection and one REPEATABLE READ READ ONLY transaction whose first statement reads the tables row, so the definition and the items come from one MVCC snapshot and no row lock is taken; a DeleteTable that drops the data table between that snapshot and the open's LOCK TABLE makes the open TableNotFound. Separate data database: the tables row is held FOR SHARE on a catalog connection from the definition read until the data transaction (REPEATABLE READ READ ONLY, ACCESS SHARE on the relation) has pinned its snapshot with SELECT pg_current_snapshot() and the GSI data-table probe has run; the definition is read on that barrier connection (build_table_description_on), so an open holds one catalog connection and the stream one data connection. UpdateTable and DeleteTable take the row FOR UPDATE, so neither commits in the gap; index-status workers do not, so an index may be reported in a slightly older status than the items. The marker is pg_current_snapshot() text on both. An engine-wide permit is held per snapshot from the open until the stream's transaction ends, half the pool on either layout, so open snapshots leave half of the pool their items come from to requests (with a separate data database and the whole pool allowed, ten open snapshots made GetItem wait out the pool timeout); restores' own half-pool slots are separate, so snapshots and restores together can hold the whole pool, and the field doc says so. Items stream through a server-side cursor, FETCH 1000 per round trip, ordered by the primary key. A fetch is bounded by rows, not bytes (up to 27.7 MB of text per row at the default item limit); the byte budget is the backup job PR's. A table whose text key columns do not have a collation named "C" or "POSIX" (created before the key columns were declared with it; "default" on a C database also counts as not bytewise, which streams unordered and complete) is streamed unordered, with a warning: ordering it would turn the read into a sort of the whole table before the first row. SQLSTATE 40001 (recovery conflict on a standby) and 72000 (snapshot_too_old on a primary, PostgreSQL 16 and earlier) map to SnapshotExpired, pinned by a test that raises each code on the live server. The stream holds ACCESS SHARE on the data table until its last fetch, so a DROP TABLE waits for it; DeleteTable takes the asynchronous path (row DELETING, worker drop under a lock timeout) when another session holds ACCESS SHARE on any relation it would drop (the base, GSI, LSI and vector-index data tables: data::table_storage_relations, the one list drop_table_storage also drops), probing only when its answer is used (delay < 1 s) and, on the shared layout, on its own catalog transaction rather than a second connection from the same pool (twenty concurrent DeleteTables on a ten-connection shared pool timed out without that). On both layouts the synchronous path takes every drop's lock before the catalog removal commits. It deletes the table's pending queue rows first, waiting for in-flight claims with no timeout, as #401 intends. It then drops the base, GSI, LSI and vector-index data tables, base first, within one 1 s budget for the group (SYNC_DROP_LOCK_TIMEOUT; before each DROP, lock_timeout is the time left). On the shared layout that runs in the catalog transaction behind a savepoint. With a separate data database it runs in a data transaction opened first, and a task the request awaits commits the catalog and then that transaction, so a cancelled caller cannot stop between the two. A lock wait that fails (55P03, a deadlock victim 40P01, or a statement_timeout's 57014) rolls the drops back and takes the asynchronous path with the row still held. A write transaction, a SHARE holder, an index build or a reader that starts after the probe therefore costs the request at most the 1 s, and no data table is left without its row. With a separate data database a data commit that fails after the catalog commit has the drop retried on a fresh connection for up to 60 s, so a failover, a dropped connection or a network blip no longer orphans the data tables; only a process that dies, or a data database unreachable for that whole window, between the catalog commit and the last retry leaves them without a row, and then they stay until an operator drops the relations the error log names (docs/troubleshooting.md, "Data tables without a catalog row"). The base table goes first because a write holds it before it writes an LSI or a synchronous GSI: with the index tables first, a write in progress waited about 1 s on the drop, and the drop went asynchronous. On the shared layout a DeleteTable probing between an open's snapshot and its LOCK does not see the open; if the LOCK is granted first, the drop times out after 1 s, is rolled back with the removal, and DeleteTable returns DELETING (in 80 unaided open-and-delete pairs, the delete first in half of them, 0 left a data table without a row on either layout; 39 to 40 deletes ran synchronously and 40 to 41 asynchronously). The data transaction lives in a guard from BEGIN on, so every path that ends it after it may have taken ACCESS SHARE wakes the control-plane poller: the last page (before its commit is awaited), a fetch error, a stream dropped before its end, and an open that fails or is cancelled; the in-catalog backup's data transaction is held in the same guard and wakes it too. The poller stays active while a DELETING row's drop is deferred, also after a restart: startup recovery runs the same pass and wakes the poller when it deferred a drop. A drop is deferred when the probe finds another session holding any lock on the table's relations, or when the drop fails in a lock wait (55P03, 40P01 or 57014); any other failure is logged and left to the normal window and the 60 s idle pass. Each pass probes a held table and skips it without waiting, and drops a free one within the same 1 s budget, so every held table is retried about once a second however many are held, and a CreateTable that is due is activated on the same cadence (1.1 s with four drops deferred). A drop that DeleteTable deferred is due at once, not after control_plane_delay_seconds. The active loop waits on the wake as well as on its timer, so a snapshot's or an in-catalog backup's release has the drop retried at once, and any holder's release is followed within about a second: 108 to 166 ms in the tests, with a 10 s test-only active poll for the wake cases. The admin guide's control_plane_delay_seconds row and the troubleshooting page's DELETING entry state the default (0.25 on every backend), the bound and the retry cadence, the same-name CreateTable refusal until then, and a two-step query, over every relation of the table, that finds any session holding a lock on them and its lock mode. A DeleteTable of a table that is already DELETING returns its description and changes nothing. SQLite: one deferred read transaction covers catalog and data; its first read pins the WAL snapshot and no lock is taken, so writers never wait on a snapshot being opened and the WAL grows until the stream ends (a TRUNCATE checkpoint waits out the busy timeout and reports busy; the module doc gives the measured reply). #384's create_backup_from_wal_snapshot holds the writer lock while it pins because it also inserts the CREATING row; this path writes nothing. build_table_description_on reads the definition, including #384's restore_summary from table_restores, on that transaction. The stream holds a backup slot (backup_server_slots, shared with the in-catalog backup) until its read transaction ends (the short page that shows the end, or the first error) or it is dropped, so the pool arithmetic that keeps half the pool free counts snapshots; with no slot free the open returns LimitExceeded at once, and a CreateBackup that waited out its 30 s for a slot says that backups or table snapshots hold them. The pool_size doc gives the arithmetic. Items page through row-value keyset pagination on the primary key index, so key_ordered is true for single-range-key tables; multi-part keys fall back to rowid order. A page is bounded by rows, not bytes, as on PostgreSQL. An in-memory database has one connection, so the open reads every page in its transaction and returns a stream over the copy, holding no connection once it returns; the copy's stream keeps its backup slot until it has yielded its last item or is dropped, so the slots bound how many decoded copies exist, as #384's create_backup_single_connection does for the in-catalog backup. SQLite DeleteTable of a DELETING table (control_plane_delay_seconds of 1 or more) returns its description too. Cassandra: open_table_snapshot returns Unsupported. Cassandra has no read that fixes an instant across a table (a token-range scan observes each range at a different time; LWT serializes single partitions), and its in-catalog logical backup says so (ADR-0019). Returning that copy under this trait's contract would make the manifest's snapshot member false, so the backend refuses and the in-catalog BackupEngine path stays its backup path. A direct test asserts the refusal naming the backend on an existing table and a missing one. MongoDB: a session started with snapshot=true reads the tables document, the indexes documents (doc_to_table_description_in now takes an optional session) and the base collection at one cluster time; the marker is that time. Every read sets readConcern: snapshot explicitly. The production catalog handle carries readConcern: local (as a URI with readConcernLevel= would), and on it a read without the level fails: the first read reports no snapshot time, which the open refuses naming the read concern, and a later read in the pinned session is refused by the server with code 72 (InvalidOptions). A server that is not a replica set member refuses the first read with NotAReplicaSet (123), mapped to Unsupported; each of the two refusals names its own cause. SnapshotTooOld (239) and SnapshotUnavailable (246) map to SnapshotExpired, pinned by a unit test that builds the command error from a reply document; the server keeps the snapshot for at least minSnapshotHistoryWindowInSeconds and may discard it after, and an idle cursor times out after cursorTimeoutMillis (10 minutes, CursorNotFound mapped to Internal). Error messages carry the server's code name and message, not the driver's dump of the reply. A document the stream cannot decode ends it with that error. key_ordered is false: for a composite key _id is a netstring of the key text, which is not the format's order for any key type. Test hooks (extenddb_storage::snapshot::paging): a page-size knob, a round-trip counter, and a pause at two points of the open (AfterDefinition, BeforeItems) that the conformance check holds while it commits writes. PostgreSQL and SQLite record each FETCH or page query; for MongoDB, whose driver hides getMore, the test fixture records, through command monitoring, each find the server receives on the data database at readConcern snapshot (the stream's; the GSI backfill's find reads at majority) and each getMore on a cursor those finds returned, so the count is the snapshot stream's paging as the server saw it. They and the conformance check compile only with the test-hooks feature of extenddb-storage, which the test-hooks features of storage-postgres and storage-sqlite and the new snapshot-test-hooks of storage-mongodb enable (storage-mongodb's existing test-hooks, the integration binary's coordination hooks, is unchanged). The backend crates' tests enable them through a dev-dependency on themselves. Production builds carry neither the hooks nor the harness. PostgreSQL also gains test-hooks entries spawn_control_plane_poller_for_tests and spawn_control_plane_poller_with_active_poll_for_tests (the production poller, the second with its active retry interval as a parameter), recover_control_plane_at_startup_for_tests, control_plane_drop_attempts_for_tests (a per-engine count of drop attempts, a held table the probe skipped counted), arm_delete_table_pause_at (holds DeleteTable at BeforeDrops, between its probe and its drops, or at BetweenCommits, between the catalog commit and the data commit on the separate layout; arm_delete_table_pause is the BeforeDrops form), shorten_data_drop_retry_bound_for_tests (a process-global shorter bound for the retried data drop), and re-exports of the data relation names and of apply_all_migrations, so tests/table_snapshot.rs binds the crate's lists instead of copying them. The two pauses (the open's and DeleteTable's) are one generic type, extenddb_storage::test_pause::Pause, and snapshot_conformance::PageSize is the one page-size guard the three backends' tests use. tests/table_snapshot.rs fails instead of skipping without a connection string when EXTENDDB_TEST_PG_REQUIRED is set, which integration.yml sets. Tests: extenddb_storage::snapshot_conformance::run is one check every backend runs against a live database. It sets the shared page size to 50, opens a snapshot on a 700-item composite-key table with a GSI, reads 51 items and asserts the round-trip counter advanced by at least two since just before the open (so a backend that ignored the knob, or a MongoDB cursor without batch_size, fails here instead of passing vacuously), then a writer deletes 100 items, inserts 100, and overwrites all 600 remaining items twice before the reader drains. Asserts: 700 items, no duplicates, every deleted item present, no inserted item, every version at its pre-open value, the counter advanced after the writer committed (a stream that read every page during the open fails), key_ordered as the backend declares it and strictly ascending keys by KeyOrder::compare (the format's comparator) when ordered, and a second snapshot taken after the writer stops shows the final round. Then, once per pause point, on a fresh table: with the open held, it commits W1, an UpdateTable adding a GSI, and W2 (in that order; on PostgreSQL with a separate data database the UpdateTable waits for the row lock and the open resumes after 2 s), and asserts what one instant implies: a description with the new GSI has W1, items with W2 have the GSI in the description, and W2 is not in the items at all, since the pause is inside the open's read window (a backend whose pauses sat before both reads passed the two implications), and which of W2's commit and the 2 s timeout resumed the open, against the backend's declared Expectations::update_table_waits_at (PostgreSQL with a separate data database at AfterDefinition: the timeout; everywhere else: W2), so a slow UpdateTable cannot make the W2 assertion vacuous. Before the writer starts the snapshot's description must equal the storage layer's describe_table with the running counts and table_throughput_mode_summary cleared; restore_summary is compared, and each backend's suite also compares a table given a restore record (assert_snapshot_describes_the_table). Every drain is under a 120 s deadline naming the stage and the items read, and a test waiting for the hooks lock gives up when one holder has kept it 600 s, naming the wait. READ COMMITTED or no FOR SHARE on PostgreSQL, the data cursor or the index read outside the session on MongoDB, and the pauses moved outside the read window on SQLite or MongoDB, fail it, as do the pauses moved outside the read window with a 2.5 s UpdateTable (the branch assertion), a description built from a second code path (the DescribeTable comparison), and a stream that never ends (the drain deadline, within 120 s instead of hanging the binary). Six more key shapes (hash-only S, N and B; S hash with S, N or B range) run at a page of 3 with keys whose byte, collation and base64 orders differ and range numbers whose text and numeric orders differ (negative, fractional, 1E-130, 9.9E125, and 38- and 34-digit pairs a last digit apart; the check's Expectations say how many digits a backend stores, and MongoDB's refusal of the 38-digit pair at the put, Decimal128's 34 digits, is asserted). Every test that opens a snapshot or touches the hooks holds snapshot_conformance::hooks_lock(), whose guard run and run_refusals take as an argument, so no other test in the binary can take the pause, reset the knob or feed the counter during a check. run_refusals covers a missing and a non-ACTIVE table on the three backends that implement snapshots. The suites pass on PostgreSQL 16 (both layouts), file-backed SQLite and MongoDB 7 (replica set; the standalone refusal against a standalone mongod), and the Cassandra test on Cassandra 4.1. PG: tests/table_snapshot.rs (81 tests and two ignored race measurements: consistency, marker and key order, tables with legacy-collation key columns streaming unordered and complete on a scratch database whose collation is checked not to be C, five held snapshots on a ten-connection pool leaving GetItem and DescribeTable unblocked with the sixth LimitExceeded and an open succeeding after one is dropped, a GSI without its data table left out with its attribute, DeleteTable during a snapshot returning DELETING, with the production poller at a 10 s active poll the table gone within 3 s of a release during the pause after the first drop attempt (for a dropped stream, a drained stream, a fetch error, an open cancelled at BeforeItems and the in-catalog backup), for a plain reader holding ACCESS SHARE that sends no wake, held 10 s past the poller's window, the table gone within 3 s of its release, a reader across a restart (the table gone within 3 s of its release), a drop failing on a dependent view making at most 3 attempts in the 15 s after the window, a restored table's restore_summary in the snapshot description, DeleteTable with only a GSI or vector-index data table locked taking the async path, an unparseable row ending the stream and freeing its permit, twenty concurrent DeleteTables on a ten-connection pool, an open completing with nine of ten connections held, an open held while nine DeleteTables queue, DeleteTable at BeforeItems taking the async path, a fetch error freeing the permit, a drained stream out of its transaction before its end is polled, a ROW EXCLUSIVE and a SHARE holder sending DeleteTable down the asynchronous path within about 1 s with no data table left without its row, a DeleteTable cancelled during that wait leaving the table ACTIVE with its data, write holders on the base and three GSI tables released 0.8 s apart taking it there at about 1 s, a write holding the base and then writing the GSI table finishing first with the table removed in the request, a claimed queue row holding DeleteTable until it ends, a probed reader sending it asynchronous in milliseconds, a statement timeout and a deadlock victim taking the asynchronous path, DeleteTable of a DELETING table returning its description and changing nothing, a deferred drop due at once, four held tables probed with no lock wait (the first released gone within 2 s, a CreateTable ACTIVE within 2 s), all on both pool layouts; on the separate layout, a data commit that fails after the catalog commit (its backend terminated between the two) retried until the drop goes through with the request returning DELETING and no relation left, a data database refusing connections past a shortened bound returning Internal naming every leaked relation with the row gone, and a caller cancelled between the two commits leaving no orphan; on the shared layout, a snapshot open or a GSI-only reader taking its lock after DeleteTable's probe sending it down the asynchronous path with the row and both data tables kept, and a table dropped before the shared-layout LOCK being TableNotFound; the SQLSTATE mapping; the two ignored race measurements, which delay the delete or the open in alternate pairs and assert both outcomes occurred) with a Drop guard that drops the scratch databases on success or panic, added to the direct-storage step in integration.yml; the full storage-postgres suite is 259 tests. SQLite: in-module tests (consistency, key order, the in-memory copy releasing the connection, in-memory copies holding their slot, a damaged attribute definition streaming unordered and complete, the rowid fallback at a page of seven with no duplicates and at least six fetches, slots shared with CreateBackup and its refusal text, the WAL checkpoint reply while a snapshot is held, a restored table's restore_summary compared on a file-backed and an in-memory engine, a drained stream and a stream that hit an error freeing their slot while still held, and an in-memory copy freeing its slot with its last item) with the same guard for the scratch files, created under the test binary's target directory. MongoDB: in-module tests (consistency, an undecodable document ending the stream, reads without the snapshot level failing with code 72 and the error text, the code mapping, a restored table's restore fields compared, the round-trip counter counting only the snapshot stream (a GSI backfill leaves it unchanged), and a standalone server answering Unsupported, gated on EXTENDDB_TEST_MONGODB_STANDALONE_URI, which the CI job sets against a standalone mongod container it starts (and failing rather than skipping when EXTENDDB_TEST_MONGODB_REQUIRED is set without it)) on per-test scratch databases built without the bootstrapper, as #384's live tests do, the catalog handle with the production option set, dropped by a guard on success or panic, so a URI in the environment never makes cargo test touch another instance's databases; gated on EXTENDDB_TEST_MONGODB_URI (failing rather than skipping when EXTENDDB_TEST_MONGODB_REQUIRED is set, which the CI job does), holding the snapshot hooks lock, run by a new integration-mongodb.yml job that starts a replica set through devtools/run-mongodb-tests --mongo-only and a standalone mongod (both on 127.0.0.1 only). tests/test_table_operations.py: DeleteTable of a DELETING table returns its description, on PostgreSQL and SQLite (the state written into the catalog), skipped on MongoDB and Cassandra, which never leave a table DELETING. Signed-off-by: Scott Robinson --- .github/workflows/integration-mongodb.yml | 54 +- .github/workflows/integration.yml | 7 +- Cargo.lock | 3 + crates/engine/src/create_table.rs | 25 + crates/storage-cassandra/src/lib.rs | 1 + .../storage-cassandra/src/snapshot_engine.rs | 62 + .../tests/direct/snapshot_engine.rs | 36 + .../tests/direct_integration.rs | 2 + crates/storage-mongodb/Cargo.toml | 11 +- crates/storage-mongodb/src/lib.rs | 9 +- crates/storage-mongodb/src/snapshot_engine.rs | 996 +++++ crates/storage-mongodb/src/table_engine.rs | 50 +- crates/storage-postgres/Cargo.toml | 8 + crates/storage-postgres/src/backup_engine.rs | 53 +- crates/storage-postgres/src/data/ddl.rs | 279 +- crates/storage-postgres/src/data/mod.rs | 13 +- .../storage-postgres/src/data/vector_index.rs | 10 +- crates/storage-postgres/src/delete_table.rs | 636 ++- crates/storage-postgres/src/lib.rs | 161 +- crates/storage-postgres/src/migrations.rs | 33 +- .../storage-postgres/src/snapshot_engine.rs | 680 ++++ crates/storage-postgres/src/table_helpers.rs | 51 +- crates/storage-postgres/src/worker_store.rs | 176 +- crates/storage-postgres/src/workers.rs | 119 +- .../tests/gsi_online_build.rs | 19 +- .../storage-postgres/tests/table_snapshot.rs | 3407 +++++++++++++++++ crates/storage-sqlite/Cargo.toml | 8 + crates/storage-sqlite/src/backup.rs | 30 +- crates/storage-sqlite/src/config.rs | 20 +- crates/storage-sqlite/src/delete_table.rs | 22 +- crates/storage-sqlite/src/lib.rs | 1 + crates/storage-sqlite/src/snapshot_engine.rs | 1094 ++++++ crates/storage-sqlite/src/store.rs | 7 +- crates/storage-sqlite/src/table_helpers.rs | 35 +- crates/storage/Cargo.toml | 5 + crates/storage/src/error.rs | 6 + crates/storage/src/lib.rs | 21 +- crates/storage/src/snapshot.rs | 458 +++ crates/storage/src/snapshot_conformance.rs | 1526 ++++++++ crates/storage/src/test_pause.rs | 133 + devtools/run-mongodb-tests | 42 +- docs/manuals/05-admin-guide.md | 2 +- docs/troubleshooting.md | 64 +- tests/conftest.py | 29 + tests/test_table_operations.py | 109 +- 45 files changed, 10199 insertions(+), 314 deletions(-) create mode 100644 crates/storage-cassandra/src/snapshot_engine.rs create mode 100644 crates/storage-cassandra/tests/direct/snapshot_engine.rs create mode 100644 crates/storage-mongodb/src/snapshot_engine.rs create mode 100644 crates/storage-postgres/src/snapshot_engine.rs create mode 100644 crates/storage-postgres/tests/table_snapshot.rs create mode 100644 crates/storage-sqlite/src/snapshot_engine.rs create mode 100644 crates/storage/src/snapshot.rs create mode 100644 crates/storage/src/snapshot_conformance.rs create mode 100644 crates/storage/src/test_pause.rs diff --git a/.github/workflows/integration-mongodb.yml b/.github/workflows/integration-mongodb.yml index 3df405a75..03223f967 100644 --- a/.github/workflows/integration-mongodb.yml +++ b/.github/workflows/integration-mongodb.yml @@ -51,6 +51,55 @@ jobs: - name: Run MongoDB pytest suite run: devtools/run-mongodb-tests -- --pytest --comprehensive --parallel + # SnapshotEngine consistency check against a disposable replica set. The + # tests build their own per-test scratch databases (snap_cat and + # _data) and drop them; they run in this job because it is the one that + # exports EXTENDDB_TEST_MONGODB_URI to cargo test. + mongodb-snapshot: + runs-on: extenddb_ubuntu-2404_4-core + timeout-minutes: 30 + steps: + - uses: actions/checkout@v6 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + with: + cache-on-failure: true + - name: Start replica set + # Same bootstrap developers run locally; --mongo-only skips the + # server and prints the URI. No pipe: a bootstrap failure must fail + # this step, not the next one's connection attempt. + run: devtools/run-mongodb-tests --mongo-only --mongo-port 27018 + - name: Start standalone mongod + # No --replSet: a_standalone_server_is_unsupported checks that a + # server which cannot serve snapshot reads answers NotAReplicaSet and + # the open maps it to Unsupported. Published on 127.0.0.1 only. + run: | + docker run -d --name extenddb-snapshot-standalone-27019 \ + -p 127.0.0.1:27019:27017 mongo:7 + for _ in $(seq 1 60); do + if docker exec extenddb-snapshot-standalone-27019 \ + mongosh --quiet --eval 'db.runCommand({ ping: 1 }).ok' 2>/dev/null \ + | grep -qx 1; then + exit 0 + fi + sleep 1 + done + echo "standalone mongod did not answer ping within 60 s" >&2 + docker logs extenddb-snapshot-standalone-27019 >&2 || true + exit 1 + - name: Run snapshot tests + # The tests serialise themselves on the snapshot hooks lock; the + # default parallel runner is what developers get locally and what + # this checks. + env: + EXTENDDB_TEST_MONGODB_URI: mongodb://127.0.0.1:27018/?directConnection=true + EXTENDDB_TEST_MONGODB_STANDALONE_URI: mongodb://127.0.0.1:27019/?directConnection=true + EXTENDDB_TEST_MONGODB_REQUIRED: "1" + run: cargo test -p extenddb-storage-mongodb snapshot_engine + - name: Stop MongoDB containers + if: always() + run: docker rm -f extenddb-runtests-mongo-27018 extenddb-snapshot-standalone-27019 + mongodb-rust: runs-on: extenddb_ubuntu-2404_4-core timeout-minutes: 40 @@ -85,12 +134,13 @@ jobs: mongodb-integration: runs-on: extenddb_ubuntu-2404_4-core timeout-minutes: 40 - needs: [mongodb-pytest, mongodb-rust, mongodb-production-build] + needs: [mongodb-pytest, mongodb-rust, mongodb-production-build, mongodb-snapshot] if: always() steps: - run: | if [ "${{ needs.mongodb-pytest.result }}" != "success" ] || \ [ "${{ needs.mongodb-rust.result }}" != "success" ] || \ - [ "${{ needs.mongodb-production-build.result }}" != "success" ]; then + [ "${{ needs.mongodb-production-build.result }}" != "success" ] || \ + [ "${{ needs.mongodb-snapshot.result }}" != "success" ]; then exit 1 fi diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 8a4818f75..5d9402a2e 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -305,13 +305,18 @@ jobs: # pause writers with a test hook. They build their own throwaway databases; # the connection string is the server, not a database. - name: Run PostgreSQL storage-level tests + # table_snapshot: the SnapshotEngine consistency check; it needs the + # shared page-size hook, compiled only into the crate's own test + # build, and scratch databases it creates and drops itself. env: EXTENDDB_TEST_PG_CONNECTION_STRING: postgresql://postgres:devpass@127.0.0.1:5432 + # table_snapshot fails rather than skips without the string above. + EXTENDDB_TEST_PG_REQUIRED: "1" run: >- cargo test --release -p extenddb-storage-postgres --test vector_control_plane --test put_create_race --test twi_conflict --test stream_shard_retention --test backup_restore - --test gsi_online_build --test gsi_stale_writer + --test gsi_online_build --test gsi_stale_writer --test table_snapshot # The daemonized server logs to syslog; dump it so server-side failures # are diagnosable from the job log. diff --git a/Cargo.lock b/Cargo.lock index f8d11ad7c..ede01e9c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1452,6 +1452,7 @@ dependencies = [ "extenddb-auth", "extenddb-core", "extenddb-storage", + "extenddb-storage-mongodb", "futures", "mongodb", "proptest", @@ -1480,6 +1481,7 @@ dependencies = [ "extenddb-auth", "extenddb-core", "extenddb-storage", + "extenddb-storage-postgres", "futures", "libc", "pgvector", @@ -1509,6 +1511,7 @@ dependencies = [ "extenddb-auth", "extenddb-core", "extenddb-storage", + "extenddb-storage-sqlite", "futures", "libc", "rand 0.9.4", diff --git a/crates/engine/src/create_table.rs b/crates/engine/src/create_table.rs index 424e27b55..e9948c111 100755 --- a/crates/engine/src/create_table.rs +++ b/crates/engine/src/create_table.rs @@ -162,6 +162,13 @@ pub fn storage_err_to_dynamo(e: extenddb_storage::error::StorageError) -> Dynamo tracing::warn!(transient_error = %msg, "transient storage error"); DynamoDbError::ServiceUnavailable("Service is temporarily unavailable".to_owned()) } + // Only backup and export jobs hold snapshots; a request that reaches + // this arm lost one mid-operation. The job owner restarts it; the + // client sees the same retryable 503 as a transient fault. + StorageError::SnapshotExpired(msg) => { + tracing::warn!(snapshot_error = %msg, "storage snapshot expired"); + DynamoDbError::ServiceUnavailable("Service is temporarily unavailable".to_owned()) + } StorageError::DeletionProtected(arn) => DynamoDbError::ValidationException(format!( "Resource '{arn}' cannot be deleted as it is currently protected against deletion. Disable deletion protection first then try again." )), @@ -247,3 +254,21 @@ pub(crate) fn storage_err_to_dynamo_with_ccf( other => storage_err_to_dynamo(other), } } + +#[cfg(test)] +mod storage_error_mapping_tests { + use extenddb_core::error::DynamoDbError; + use extenddb_storage::error::StorageError; + + use super::storage_err_to_dynamo; + + /// A snapshot the backend discarded is retryable, like a transient error. + #[test] + fn snapshot_expired_is_service_unavailable() { + let err = storage_err_to_dynamo(StorageError::SnapshotExpired("gone".to_owned())); + assert!( + matches!(err, DynamoDbError::ServiceUnavailable(_)), + "{err:?}" + ); + } +} diff --git a/crates/storage-cassandra/src/lib.rs b/crates/storage-cassandra/src/lib.rs index 8f49e9664..7dbfe98c4 100644 --- a/crates/storage-cassandra/src/lib.rs +++ b/crates/storage-cassandra/src/lib.rs @@ -27,6 +27,7 @@ mod metadata_engine; pub mod migrations; pub mod operations; pub(crate) mod propagation_hold; +mod snapshot_engine; mod stream_engine; pub mod stream_util; pub mod table_engine; diff --git a/crates/storage-cassandra/src/snapshot_engine.rs b/crates/storage-cassandra/src/snapshot_engine.rs new file mode 100644 index 000000000..caf4c189c --- /dev/null +++ b/crates/storage-cassandra/src/snapshot_engine.rs @@ -0,0 +1,62 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! `SnapshotEngine` for Cassandra: refused. +//! +//! The trait promises a view of one table as of one instant: writes committed +//! after the open are not visible, writes committed before it all are, and the +//! definition describes the same instant as the items. Cassandra has no +//! mechanism that gives that across a token-range scan. Reads see whatever +//! replicas hold at the consistency level of each query; a scan that pages +//! through `token(pk)` ranges observes each range at a different time, and +//! `LOCAL_SERIAL` reads serialize single partitions, not a table. The +//! in-catalog backup this backend ships (`backup_engine.rs`, ADR-0019) is a +//! logical copy that documents exactly this: "concurrent writes can race backup +//! creation ... not advertised as strict point-in-time recovery." +//! +//! Returning that copy under the `SnapshotEngine` contract would make the +//! backup manifest's `snapshot` member a lie, so this backend returns +//! [`StorageError::Unsupported`]. The in-catalog `BackupEngine` path is +//! unaffected. A token-range implementation +//! with a documented non-point-in-time `SnapshotInfo.kind` is a design change +//! the backup design (`docs/design/15-backup-restore.md`, section 3) leaves +//! open; it must not arrive as a silent widening of this method. + +use extenddb_storage::error::StorageError; +use extenddb_storage::{SnapshotEngine, TableSnapshot}; +use futures::future::BoxFuture; + +use crate::CassandraEngine; + +/// The reason every `open_table_snapshot` call returns. +pub const SNAPSHOT_UNSUPPORTED_REASON: &str = "table snapshots are not supported on the Cassandra \ + backend: Cassandra has no point-in-time read across a table, so a backup taken through a \ + snapshot could not promise the instant its manifest records"; + +impl SnapshotEngine for CassandraEngine { + fn open_table_snapshot( + &self, + account_id: &str, + table_name: &str, + ) -> BoxFuture<'_, Result> { + // No query is issued: the answer does not depend on the table. + let _ = (account_id, table_name); + let result = Err(StorageError::Unsupported( + SNAPSHOT_UNSUPPORTED_REASON.to_owned(), + )); + Box::pin(async move { result }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The reason is one sentence an operator can act on: it names the + /// backend and the cause (no point-in-time read across a table). + #[test] + fn reason_names_the_backend_and_the_cause() { + assert!(SNAPSHOT_UNSUPPORTED_REASON.contains("Cassandra")); + assert!(SNAPSHOT_UNSUPPORTED_REASON.contains("point-in-time")); + } +} diff --git a/crates/storage-cassandra/tests/direct/snapshot_engine.rs b/crates/storage-cassandra/tests/direct/snapshot_engine.rs new file mode 100644 index 000000000..26616602f --- /dev/null +++ b/crates/storage-cassandra/tests/direct/snapshot_engine.rs @@ -0,0 +1,36 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! `SnapshotEngine` on Cassandra refuses every table with `Unsupported`: +//! an existing ACTIVE table and a missing one alike, so an implementation +//! that someday consulted the catalog before refusing would still be held to +//! the refusal. That no query is issued is a property of the code (the impl +//! returns before touching the session), not asserted here. + +#[cfg(test)] +mod tests { + use extenddb_storage::SnapshotEngine; + use extenddb_storage::error::StorageError; + + #[tokio::test] + async fn open_table_snapshot_is_unsupported() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = crate::helpers::setup_engine().await; + let table = crate::helpers::TestTable::new(&engine, "snap_refused", true).await; + let account = table.key_info.account_id.clone(); + + for name in [ + table.key_info.table_name.as_str(), + "does_not_exist_for_snapshot", + ] { + let err = engine + .open_table_snapshot(&account, name) + .await + .expect_err("refused"); + assert!(matches!(err, StorageError::Unsupported(_)), "{name}: {err}"); + assert!(err.to_string().contains("Cassandra"), "{err}"); + } + } +} diff --git a/crates/storage-cassandra/tests/direct_integration.rs b/crates/storage-cassandra/tests/direct_integration.rs index 43eb28b81..c85398327 100644 --- a/crates/storage-cassandra/tests/direct_integration.rs +++ b/crates/storage-cassandra/tests/direct_integration.rs @@ -45,6 +45,8 @@ mod roles; mod scan; #[path = "direct/settings_store.rs"] mod settings_store; +#[path = "direct/snapshot_engine.rs"] +mod snapshot_engine; #[path = "direct/streams.rs"] mod streams; #[path = "direct/table_engine.rs"] diff --git a/crates/storage-mongodb/Cargo.toml b/crates/storage-mongodb/Cargo.toml index a8d5380a9..f455971dc 100644 --- a/crates/storage-mongodb/Cargo.toml +++ b/crates/storage-mongodb/Cargo.toml @@ -10,8 +10,14 @@ license.workspace = true [features] # Enables deterministic coordination hooks used only by the MongoDB integration -# test binary. Production builds do not compile these hooks. +# test binary (the mongodb-test-hooks server build). Production builds do not +# compile these hooks. test-hooks = [] +# Compiles the snapshot page-size setter, the open pause and the shared +# snapshot conformance check the crate's snapshot tests use. Separate from +# `test-hooks` so the crate's own tests build the production commit and +# backfill paths, not the coordination hooks. +snapshot-test-hooks = ["extenddb-storage/test-hooks"] [dependencies] # Internal crates @@ -61,3 +67,6 @@ dashmap.workspace = true # Property-based testing for the filter-pushdown parity harness # (crates/storage-mongodb/tests/pushdown_parity.rs). proptest = "1" +# The crate's snapshot tests need `snapshot-test-hooks`; a dev-dependency on +# itself turns it on for test builds only. +extenddb-storage-mongodb = { path = ".", features = ["snapshot-test-hooks"] } diff --git a/crates/storage-mongodb/src/lib.rs b/crates/storage-mongodb/src/lib.rs index c1d28882f..b27cc0baf 100644 --- a/crates/storage-mongodb/src/lib.rs +++ b/crates/storage-mongodb/src/lib.rs @@ -21,6 +21,7 @@ mod management_store; mod metadata_engine; mod operations; pub mod pushdown; +mod snapshot_engine; mod stream_engine; mod table_engine; mod ttl_worker; @@ -536,9 +537,15 @@ pub(crate) fn majority_write_concern( /// concern other than `majority` in the connection string; this does not rely /// on it. fn catalog_database(client: &mongodb::Client) -> mongodb::Database { + catalog_database_named(client, "extenddb_catalog") +} + +/// [`catalog_database`] under another name, for test fixtures that must run +/// the production option set on a scratch database. +pub(crate) fn catalog_database_named(client: &mongodb::Client, name: &str) -> mongodb::Database { use mongodb::options::{DatabaseOptions, ReadConcern, ReadPreference, SelectionCriteria}; client.database_with_options( - "extenddb_catalog", + name, DatabaseOptions::builder() .selection_criteria(SelectionCriteria::ReadPreference(ReadPreference::Primary)) .read_concern(ReadConcern::local()) diff --git a/crates/storage-mongodb/src/snapshot_engine.rs b/crates/storage-mongodb/src/snapshot_engine.rs new file mode 100644 index 000000000..c28c55397 --- /dev/null +++ b/crates/storage-mongodb/src/snapshot_engine.rs @@ -0,0 +1,996 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! `SnapshotEngine` for `MongoDB`. +//! +//! One snapshot session covers catalog and data: a session started with +//! `snapshot: true` runs every read at the cluster time of its first read +//! (`MongoDB` 5.0+ snapshot reads outside a transaction). The `tables` document +//! is read first and pins the time; the index documents and the base +//! collection cursor then read at the same time, so the definition and the +//! items describe one instant without any lock. +//! +//! Every read sets `readConcern: snapshot` explicitly. The driver only +//! inserts that level when neither the operation nor its database or +//! collection carries a read concern, and the production catalog handle is +//! built with `readConcern: local` (`catalog_database`), as a connection +//! string with `readConcernLevel=` would build every handle. Without the +//! explicit level the reads fail, checked against a `MongoDB` 7 replica set +//! (`reads_without_the_snapshot_level_fail`): the first read runs at the +//! handle's level and the session reports no snapshot time, which the open +//! refuses, and a later read in a pinned session goes out with +//! `atClusterTime` and the handle's level, which the server refuses with +//! code 72 (`InvalidOptions`: "atClusterTime field can be set only if level +//! is equal to snapshot"). +//! +//! A server that is not a replica set member (a standalone `mongod`) cannot +//! serve snapshot reads and refuses the first read with `NotAReplicaSet` +//! (code 123); the open maps that to [`StorageError::Unsupported`], with a +//! sentence that names that cause only (a first read that succeeded without +//! a snapshot time is an internal error with a sentence of its own). Checked +//! against a standalone `MongoDB` 7 by `a_standalone_server_is_unsupported`, +//! which CI's `mongodb-snapshot` job runs against a standalone container. +//! +//! Error messages are built from the server's code name, code and message +//! (`describe`), not the driver error's `Display`, which carries the raw BSON +//! reply as a hex dump. +//! +//! The server keeps at least `minSnapshotHistoryWindowInSeconds` (default +//! 300) of snapshot history. A stream that is still open past that window +//! may fail with `SnapshotTooOld` (code 239) once the server prunes the +//! snapshot, and a server that cannot serve the snapshot answers +//! `SnapshotUnavailable` (code 246); both surface as +//! [`StorageError::SnapshotExpired`] and the caller restarts the backup. At +//! typical cursor throughput the window bounds a single backup at tens of +//! millions of documents unless the server parameter is raised; the backup +//! design (#396) records that as an operator setting. The cursor is opened +//! without `noCursorTimeout`, so it has the server's default idle timeout +//! (`cursorTimeoutMillis`, 10 minutes): a stream not polled for that long +//! gets `CursorNotFound` (code 43) on its next batch, which surfaces as +//! [`StorageError::Internal`], and the caller likewise starts over. +//! +//! `key_ordered` is `false`. For a composite key `_id` is a netstring of the +//! key text (`:,:,`), which sorts by length before content +//! and is not the format's key order for any key type; a hash-only table's +//! `_id` is the hash key text and would be in order, but one rule for the +//! backend is simpler than a second ordered path no restore needs. The +//! stream reads in natural order. +//! +//! The stream does not record round trips into the shared paging counter: +//! the driver issues `getMore` inside `cursor.next` and does not expose the +//! batch boundary, so a count here could only echo the batch size the +//! stream asked for. The test fixture counts, through command monitoring, +//! the `find` commands the server receives on the data database at +//! `readConcern: snapshot` (only this stream's `find` on the table's base +//! collection reads at that level; the GSI backfill worker's `find` reads at +//! `majority` and is not counted) and the `getMore` commands on the cursor +//! ids those finds returned, so the conformance check observes the server's +//! paging of the snapshot stream and nothing else. +//! +//! Memory: a batch is at most `batch_size` documents +//! (`paging::DEFAULT_FETCH_BATCH`, 1000) and at most the server's 16 MiB +//! reply limit, whichever is reached first, so unlike the SQL backends' +//! row-bounded fetches one batch here is bounded in bytes too. + +use extenddb_core::types::Item; +use extenddb_storage::error::StorageError; +use extenddb_storage::snapshot::{drop_deleting_indexes, paging}; +use extenddb_storage::{SnapshotEngine, SnapshotInfo, TableSnapshot}; +use futures::future::BoxFuture; +use futures::stream::BoxStream; +use mongodb::bson::{Document, doc}; +use mongodb::options::ReadConcern; +use mongodb::{ClientSession, SessionCursor}; + +use crate::MongoEngine; +use crate::data::{data_collection_name, document_to_item}; + +/// Server error code for `SnapshotTooOld`: the read timestamp is older than +/// the oldest the server still holds. +const SNAPSHOT_TOO_OLD: i32 = 239; +/// Server error code for `SnapshotUnavailable`: the server cannot serve a +/// read at the requested snapshot and asks the client to retry. The remedy +/// for a backup is the same restart. +const SNAPSHOT_UNAVAILABLE: i32 = 246; +/// Server error code for `NotAReplicaSet`: a standalone server refusing a +/// snapshot read. +const NOT_A_REPLICA_SET: i32 = 123; + +/// Why an open cannot proceed on a standalone server, which refuses the +/// first read with `NotAReplicaSet` ([`map_first_read_error`]). +const STANDALONE_SERVER: &str = "MongoDB refused the first read at readConcern snapshot: a \ + standalone server cannot serve snapshot reads, and table snapshots need a replica set \ + member"; + +/// Why an open cannot proceed when the first read succeeded but the session +/// reports no snapshot time. +const NO_SNAPSHOT_TIME: &str = "MongoDB served the first read without a snapshot time: a read \ + concern overrode the snapshot level"; + +impl SnapshotEngine for MongoEngine { + fn open_table_snapshot( + &self, + account_id: &str, + table_name: &str, + ) -> BoxFuture<'_, Result> { + let account_id = account_id.to_owned(); + let table_name = table_name.to_owned(); + Box::pin(async move { + self.open_table_snapshot_impl(&account_id, &table_name) + .await + }) + } +} + +impl MongoEngine { + async fn open_table_snapshot_impl( + &self, + account_id: &str, + table_name: &str, + ) -> Result { + Self::validate_account_id(account_id)?; + + let mut session = self + .client + .start_session() + .snapshot(true) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + + // First read pins the snapshot time for the whole session. + let tables = self.catalog_db.collection::("tables"); + let table_doc = tables + .find_one(doc! { "_id": { "account_id": account_id, "table_name": table_name } }) + .read_concern(ReadConcern::snapshot()) + .session(&mut session) + .await + .map_err(|e| map_first_read_error(&e))? + .ok_or_else(|| StorageError::TableNotFound(table_name.to_owned()))?; + let table_id = table_doc + .get_str("table_id") + .map_err(|_| StorageError::Internal("missing table_id".to_owned()))? + .to_owned(); + + // The first read must have pinned a cluster time; without one the + // server is not serving snapshot reads and the stream would not be + // consistent. The cluster time is the marker: what a later change-log + // archive (the oplog from this time on) anchors against. Nothing + // reopens a snapshot from it. + let marker = session + .snapshot_time() + .map(|t| format!("{}:{}", t.time, t.increment)) + .ok_or_else(|| StorageError::Internal(NO_SNAPSHOT_TIME.to_owned()))?; + + // After the `tables` document, before its index documents (the + // stage's definition for this backend): an index read outside the + // session would see what commits here. + paging::open_pause(paging::OpenStage::AfterDefinition).await; + let mut description = self + .doc_to_table_description_in(&table_doc, Some(&mut session)) + .await?; + // One rule for "is this table ACTIVE": the description's, as + // DescribeTable reports it. + if description.table_status != extenddb_core::types::TableStatus::Active { + return Err(StorageError::TableNotActive(table_name.to_owned())); + } + drop_deleting_indexes(&mut description); + + let collection = self + .data_db + .collection::(&data_collection_name(&table_id)); + let batch = paging::fetch_batch(); + paging::open_pause(paging::OpenStage::BeforeItems).await; + let cursor = collection + .find(doc! {}) + .batch_size(u32::try_from(batch).unwrap_or(u32::MAX)) + .read_concern(ReadConcern::snapshot()) + .session(&mut session) + .await + .map_err(|e| map_mongo_error(&e))?; + + Ok(TableSnapshot { + description, + info: SnapshotInfo { + kind: SnapshotInfo::KIND_MONGODB_SNAPSHOT_READ_CONCERN.to_owned(), + marker, + }, + key_ordered: false, + items: session_stream(session, cursor), + }) + } +} + +/// Stream items from a session cursor. The session and cursor live in the +/// stream's state; dropping the stream ends both. Round trips are counted +/// by the test fixture from the server's command stream (module doc). +fn session_stream( + session: ClientSession, + cursor: SessionCursor, +) -> BoxStream<'static, Result> { + struct State { + session: ClientSession, + cursor: SessionCursor, + done: bool, + } + let state = State { + session, + cursor, + done: false, + }; + Box::pin(futures::stream::unfold(state, |mut state| async move { + if state.done { + return None; + } + // The first error ends the stream, as on every backend: a document + // that does not decode, like a failed getMore, means the stream is + // not the whole table. + match state.cursor.next(&mut state.session).await { + Some(Ok(doc)) => { + let item = document_to_item(&doc); + state.done = item.is_err(); + Some((item, state)) + } + Some(Err(e)) => { + state.done = true; + Some((Err(map_mongo_error(&e)), state)) + } + None => None, + } + })) +} + +pub(crate) fn map_mongo_error(e: &mongodb::error::Error) -> StorageError { + // `find` and `getMore` failures arrive as command errors. + let snapshot_gone = matches!( + &*e.kind, + mongodb::error::ErrorKind::Command(c) + if c.code == SNAPSHOT_TOO_OLD || c.code == SNAPSHOT_UNAVAILABLE + ); + if snapshot_gone { + return StorageError::SnapshotExpired(describe(e)); + } + StorageError::Internal(describe(e)) +} + +/// [`map_mongo_error`] for the open's first read, where a standalone server +/// answers `NotAReplicaSet`: that is a configuration the backend cannot take +/// a snapshot in, not an internal error. +fn map_first_read_error(e: &mongodb::error::Error) -> StorageError { + if matches!(&*e.kind, mongodb::error::ErrorKind::Command(c) if c.code == NOT_A_REPLICA_SET) { + return StorageError::Unsupported(format!("{STANDALONE_SERVER} ({})", describe(e))); + } + map_mongo_error(e) +} + +/// A driver error as text: the server's code name, code and message for a +/// command error, the error kind otherwise. Not `Error`'s `Display`, which +/// appends the raw reply document as a hex dump. +fn describe(e: &mongodb::error::Error) -> String { + match &*e.kind { + mongodb::error::ErrorKind::Command(c) => { + format!("{} (code {}): {}", c.code_name, c.code, c.message) + } + kind => kind.to_string(), + } +} + +#[cfg(test)] +mod error_mapping_tests { + //! Pin the server error codes without a server: a wrong constant here + //! turns the one retryable failure the module promises to classify into + //! an internal error. + + use mongodb::error::{CommandError, Error, ErrorKind}; + + use super::*; + + /// `CommandError` is non-exhaustive with a crate-private field, so it is + /// built the way the driver builds it: from the server's reply document. + fn command_error(code: i32, code_name: &str) -> Error { + let reply = doc! { "code": code, "codeName": code_name, "errmsg": "synthetic" }; + let command: CommandError = + mongodb::bson::from_document(reply).expect("command error document"); + Error::from(ErrorKind::Command(command)) + } + + #[test] + fn snapshot_too_old_and_unavailable_are_snapshot_expired() { + for (code, name) in [(239, "SnapshotTooOld"), (246, "SnapshotUnavailable")] { + let mapped = map_mongo_error(&command_error(code, name)); + assert!( + matches!(mapped, StorageError::SnapshotExpired(_)), + "{name} ({code}): {mapped:?}" + ); + } + } + + #[test] + fn not_a_replica_set_on_the_first_read_is_unsupported() { + let StorageError::Unsupported(message) = + map_first_read_error(&command_error(123, "NotAReplicaSet")) + else { + panic!("NotAReplicaSet on the first read is Unsupported"); + }; + assert!(message.starts_with(STANDALONE_SERVER), "{message}"); + assert!(message.contains("NotAReplicaSet (code 123)"), "{message}"); + // Only the first read maps it; later reads keep the general rule. + assert!(matches!( + map_mongo_error(&command_error(123, "NotAReplicaSet")), + StorageError::Internal(_) + )); + // And the first read keeps the snapshot codes. + assert!(matches!( + map_first_read_error(&command_error(239, "SnapshotTooOld")), + StorageError::SnapshotExpired(_) + )); + } + + #[test] + fn other_command_errors_are_internal() { + for (code, name) in [ + (50, "MaxTimeMSExpired"), + (13, "Unauthorized"), + (0, "unknown"), + ] { + let mapped = map_mongo_error(&command_error(code, name)); + assert!( + matches!(mapped, StorageError::Internal(_)), + "{name} ({code}): {mapped:?}" + ); + } + } +} + +#[cfg(test)] +mod tests { + //! Runs against a live replica set named by `EXTENDDB_TEST_MONGODB_URI` + //! (for example `mongodb://127.0.0.1:27017/?directConnection=true`); + //! skips otherwise. Each test builds an engine on its own + //! `snap_cat` / `snap_data` databases, the way the backup + //! live tests do, and drops them when the fixture drops, so a URI in the + //! environment never makes `cargo test` touch another instance's + //! databases. `devtools/run-mongodb-tests --mongo-only` starts a + //! replica set. + + use std::sync::Arc; + + use extenddb_storage::snapshot_conformance::{self, PageSize}; + use extenddb_storage::{SnapshotEngine, TableEngine}; + + use super::*; + + const ACCOUNT: &str = "123456789012"; + /// `MongoDB`: natural order, number keys stored as Decimal128 (34 digits); + /// the open holds nothing an `UpdateTable` waits on. + const EXPECT: snapshot_conformance::Expectations = snapshot_conformance::Expectations { + key_ordered: false, + numeric_key_digits: 34, + update_table_waits_at: &[], + }; + + /// Every live test here opens a snapshot (or tries to) or sets the shared + /// page size, so each holds `snapshot_conformance::hooks_lock` for its + /// whole body: the knob, the round-trip counter and the open pause are + /// process-global, and one lock covers all three. + async fn hooks() -> snapshot_conformance::HooksGuard { + snapshot_conformance::hooks_lock().await + } + + /// The replica set URI, or `None` to skip. With + /// `EXTENDDB_TEST_MONGODB_REQUIRED` set a missing URI is a failure, so CI + /// cannot pass by skipping. + fn uri() -> Option { + let uri = std::env::var("EXTENDDB_TEST_MONGODB_URI") + .ok() + .filter(|u| !u.trim().is_empty()); + assert!( + uri.is_some() || std::env::var_os("EXTENDDB_TEST_MONGODB_REQUIRED").is_none(), + "EXTENDDB_TEST_MONGODB_REQUIRED is set but EXTENDDB_TEST_MONGODB_URI is not; the \ + MongoDB snapshot tests must run, not skip" + ); + uri + } + + /// An engine on two scratch databases plus the GSI backfill worker that + /// `serve` would run (the check table has a GSI, which only the worker + /// flips to ACTIVE). Dropped with the fixture on success or panic (not + /// when the test binary is killed): the worker is + /// aborted, and the databases are dropped on a thread with its own + /// runtime and a fresh client, since `Drop` cannot await on the test + /// runtime and the test's client belongs to it. + struct Fixture { + engine: Option, + worker: tokio::task::JoinHandle<()>, + uri: String, + names: [String; 2], + } + + impl std::ops::Deref for Fixture { + type Target = MongoEngine; + fn deref(&self) -> &MongoEngine { + self.engine.as_ref().expect("engine present until drop") + } + } + + impl Drop for Fixture { + fn drop(&mut self) { + self.worker.abort(); + drop(self.engine.take()); + let uri = self.uri.clone(); + let names = self.names.clone(); + std::thread::scope(|scope| { + scope.spawn(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("cleanup runtime"); + rt.block_on(async move { + let Ok(client) = mongodb::Client::with_uri_str(&uri).await else { + eprintln!("scratch cleanup: could not connect; {names:?} left behind"); + return; + }; + for name in &names { + if let Err(e) = client.database(name).drop().await { + eprintln!("scratch cleanup: drop {name}: {e}"); + } + } + client.shutdown().await; + }); + }); + }); + } + } + + async fn fixture(uri: &str) -> Fixture { + let stem = format!("snap{}", uuid::Uuid::new_v4().simple()); + let names = [format!("{stem}_cat"), format!("{stem}_data")]; + // The shared paging counter is fed from what the server receives, + // counting only the snapshot stream's round trips: a `find` on the + // data database at `readConcern: snapshot` (the stream's read of the + // table's base collection; the GSI backfill worker also runs `find` + // on base collections, at `majority`, and the engine's writes are + // `findAndModify`, `update` and `delete`), and a `getMore` on a + // cursor id one of those finds returned (read from its reply). + let mut options = mongodb::options::ClientOptions::parse(uri) + .await + .expect("client options"); + options.command_event_handler = Some(snapshot_round_trip_counter(names[1].clone())); + let client = mongodb::Client::with_options(options).expect("client"); + // The production option set (`readConcern: local` among them), so + // the explicit snapshot read concern is tested against what `new` + // builds. + let catalog_db = crate::catalog_database_named(&client, &names[0]); + catalog_db + .collection::("accounts") + .insert_one(doc! { "_id": ACCOUNT, "account_name": "snapshot-test" }) + .await + .expect("account"); + catalog_db + .collection::("settings") + .insert_one(doc! { "_id": "control_plane_delay_seconds", "value": "0" }) + .await + .expect("setting"); + let slots = crate::backup_engine::restore_server_slot_count(10); + let engine = MongoEngine { + client: client.clone(), + catalog_db, + data_db: client.database(&names[1]), + region: "us-east-1".to_owned(), + gsi_cache: Arc::default(), + tx_read_concern: ReadConcern::majority(), + backup_liveness: crate::backup_engine::BackupLiveness::default(), + restore_server_slot_count: slots, + restore_server_slots: Arc::new(tokio::sync::Semaphore::new(slots)), + restore_account_slots: crate::backup_engine::RestoreAccountSlots::default(), + }; + let worker = tokio::spawn(crate::ttl_worker::gsi_backfill_worker(Arc::new( + engine.clone(), + ))); + Fixture { + engine: Some(engine), + worker, + uri: uri.to_owned(), + names, + } + } + + /// Command monitor that records one paging round trip per `find` at + /// `readConcern: snapshot` on `data_name` and per `getMore` on a cursor + /// such a `find` opened. + fn snapshot_round_trip_counter( + data_name: String, + ) -> mongodb::event::EventHandler { + use mongodb::event::command::CommandEvent; + #[derive(Default)] + struct Seen { + /// Request ids of counted finds whose reply has not arrived. + finds: std::collections::HashSet, + /// Cursor ids those finds returned. + cursors: std::collections::HashSet, + } + let seen = std::sync::Mutex::new(Seen::default()); + mongodb::event::EventHandler::callback(move |event: CommandEvent| { + let mut seen = seen + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match event { + CommandEvent::Started(started) if started.db == data_name => { + match started.command_name.as_str() { + "find" + if started + .command + .get_document("readConcern") + .and_then(|rc| rc.get_str("level")) + == Ok("snapshot") => + { + seen.finds.insert(started.request_id); + paging::record_fetch(); + } + "getMore" + if started + .command + .get_i64("getMore") + .is_ok_and(|id| seen.cursors.contains(&id)) => + { + paging::record_fetch(); + } + _ => {} + } + } + CommandEvent::Succeeded(succeeded) if seen.finds.remove(&succeeded.request_id) => { + if let Ok(id) = succeeded + .reply + .get_document("cursor") + .and_then(|c| c.get_i64("id")) + && id != 0 + { + seen.cursors.insert(id); + } + } + CommandEvent::Failed(failed) => { + seen.finds.remove(&failed.request_id); + } + _ => {} + } + }) + } + + async fn set_updating(engine: &MongoEngine, table: &str) { + engine + .catalog_db + .collection::("tables") + .update_one( + doc! { "_id": { "account_id": ACCOUNT, "table_name": table } }, + doc! { "$set": { "table_status": "UPDATING" } }, + ) + .await + .expect("flip status"); + } + + #[tokio::test] + async fn snapshot_is_consistent_under_concurrent_writes() { + let Some(uri) = uri() else { + eprintln!("SKIP snapshot_is_consistent_under_concurrent_writes: no MongoDB"); + return; + }; + let hooks = hooks().await; + let engine = fixture(&uri).await; + snapshot_conformance::run( + &hooks, + &*engine, + ACCOUNT, + "snap_src", + 50, + EXPECT, + || async { + let _ = extenddb_storage::WorkerStore::process_control_plane_transitions(&*engine) + .await; + }, + ) + .await; + snapshot_conformance::run_refusals( + &hooks, + &*engine, + ACCOUNT, + "snap_refused", + || async { + let _ = extenddb_storage::WorkerStore::process_control_plane_transitions(&*engine) + .await; + }, + |table| { + let engine = &*engine; + async move { set_updating(engine, &table).await } + }, + ) + .await; + } + + /// A restored table's snapshot description carries the same + /// `restore_summary` as `describe_table`'s (both read the restore fields + /// of the `tables` document). The fields are set directly, as + /// `RestoreTableFromBackup` leaves them. + #[tokio::test] + async fn a_restored_table_describes_its_restore() { + let Some(uri) = uri() else { + eprintln!("SKIP a_restored_table_describes_its_restore: no MongoDB"); + return; + }; + let hooks = hooks().await; + let engine = fixture(&uri).await; + engine + .create_table(ACCOUNT, snapshot_conformance::table_input("restored")) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "restored", || async { + let _ = + extenddb_storage::WorkerStore::process_control_plane_transitions(&*engine).await; + }) + .await; + snapshot_conformance::populate(&*engine, ACCOUNT, "restored", 5).await; + engine + .catalog_db + .collection::("tables") + .update_one( + doc! { "_id": { "account_id": ACCOUNT, "table_name": "restored" } }, + doc! { "$set": { + "restore_source_backup_arn": + "arn:aws:dynamodb:us-east-1:123456789012:table/src/backup/01759320000250-0a1b2c3d", + "restore_source_table_arn": "arn:aws:dynamodb:us-east-1:123456789012:table/src", + "restore_date_time": mongodb::bson::DateTime::from_millis(1_759_320_000_250), + } }, + ) + .await + .expect("restore fields"); + snapshot_conformance::assert_snapshot_describes_the_table( + &hooks, &*engine, ACCOUNT, "restored", true, + ) + .await; + } + + /// A document the decoder cannot read ends the stream at its error, as on + /// the other backends: nothing after it, so a consumer cannot mistake the + /// rest for the whole table. + #[tokio::test] + async fn an_undecodable_document_ends_the_stream() { + let Some(uri) = uri() else { + eprintln!("SKIP an_undecodable_document_ends_the_stream: no MongoDB"); + return; + }; + let _hooks = hooks().await; + let engine = fixture(&uri).await; + engine + .create_table(ACCOUNT, snapshot_conformance::table_input("bad")) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "bad", || async { + let _ = + extenddb_storage::WorkerStore::process_control_plane_transitions(&*engine).await; + }) + .await; + let table_id = engine + .describe_table( + ACCOUNT, + extenddb_core::types::DescribeTableInput { + table_name: "bad".to_owned(), + }, + ) + .await + .expect("describe") + .table_id; + let collection = engine + .data_db + .collection::(&data_collection_name(&table_id)); + // Five items, the undecodable document, five more: natural order is + // insertion order here, so the error falls mid-stream. + snapshot_conformance::populate(&*engine, ACCOUNT, "bad", 5).await; + collection + .insert_one(doc! { "_id": "undecodable" }) + .await + .expect("insert a document without item_data"); + let ki = extenddb_core::types::TableKeyInfo { + table_name: "bad".to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: table_id.clone(), + key_schema: snapshot_conformance::table_input("bad").key_schema, + base_key_schema: snapshot_conformance::table_input("bad").key_schema, + attribute_definitions: snapshot_conformance::table_input("bad").attribute_definitions, + ..Default::default() + }; + for i in 100..105 { + let item: Item = [ + ( + "pk".to_owned(), + extenddb_core::types::AttributeValue::S("late".to_owned()), + ), + ( + "sk".to_owned(), + extenddb_core::types::AttributeValue::N(i.to_string()), + ), + ] + .into_iter() + .collect(); + extenddb_storage::DataEngine::put_item( + &*engine, + &ki, + item, + false, + None, + &extenddb_core::expression::ExpressionMaps::default(), + None, + ) + .await + .expect("put"); + } + // Where the document falls in natural order, from a plain read: the + // stream must yield exactly the items before it, then the error. + let natural: Vec = + futures::TryStreamExt::try_collect(collection.find(doc! {}).await.expect("find")) + .await + .expect("natural order"); + let position = natural + .iter() + .position(|d| d.get_str("_id") == Ok("undecodable")) + .expect("the document is there"); + assert!( + position + 1 < natural.len(), + "precondition: the undecodable document is not last ({position} of {}), or the \ + test cannot tell a stream that ends from one that continues", + natural.len() + ); + let snapshot = engine + .open_table_snapshot(ACCOUNT, "bad") + .await + .expect("open"); + let results: Vec> = + futures::StreamExt::collect(snapshot.items).await; + let errors = results.iter().filter(|r| r.is_err()).count(); + assert_eq!(errors, 1, "one error for the one undecodable document"); + assert_eq!( + results.len(), + position + 1, + "the {position} items before the document, then the error, then the end" + ); + assert!(results.last().is_some_and(Result::is_err)); + } + + /// Pins the module doc's read-concern sentence and the error text: on a + /// handle carrying `readConcern: local` (the production catalog handle) + /// a snapshot session's first read without the explicit level runs at + /// `local` and leaves no snapshot time, and a later read in a pinned + /// session is refused with code 72. The mapped error names the server's + /// code and message and carries no hex dump of the reply. + #[tokio::test] + async fn reads_without_the_snapshot_level_fail() { + let Some(uri) = uri() else { + eprintln!("SKIP reads_without_the_snapshot_level_fail: no MongoDB"); + return; + }; + let _hooks = hooks().await; + let engine = fixture(&uri).await; + let accounts = engine.catalog_db.collection::("accounts"); + + let mut session = engine + .client + .start_session() + .snapshot(true) + .await + .expect("session"); + accounts + .find_one(doc! { "_id": ACCOUNT }) + .session(&mut session) + .await + .expect("a read at the handle's level succeeds"); + assert!( + session.snapshot_time().is_none(), + "a read at readConcern local pins no snapshot time" + ); + + let mut session = engine + .client + .start_session() + .snapshot(true) + .await + .expect("session"); + accounts + .find_one(doc! { "_id": ACCOUNT }) + .read_concern(ReadConcern::snapshot()) + .session(&mut session) + .await + .expect("pinning read"); + assert!(session.snapshot_time().is_some(), "pinned"); + let err = accounts + .find_one(doc! { "_id": ACCOUNT }) + .session(&mut session) + .await + .expect_err("the server refuses a non-snapshot level with atClusterTime"); + let raw = err.to_string(); + assert!( + raw.contains("RawDocumentBuf"), + "precondition: the driver's Display carries the reply dump: {raw}" + ); + let StorageError::Internal(message) = map_mongo_error(&err) else { + panic!("code 72 is an internal error"); + }; + assert!( + message.starts_with("InvalidOptions (code 72): atClusterTime"), + "{message}" + ); + assert!( + !message.contains("RawDocumentBuf") && !message.contains("server response"), + "{message}" + ); + } + + /// A standalone server refuses the open's first read with `NotAReplicaSet`; + /// the open answers `Unsupported`, not an internal error. Runs only when + /// `EXTENDDB_TEST_MONGODB_STANDALONE_URI` names a standalone `mongod` + /// (for example one started with `docker run -p + /// 127.0.0.1:27118:27017 mongo:7`, no `--replSet`); with + /// `EXTENDDB_TEST_MONGODB_REQUIRED` set it fails rather than skips when + /// that variable is missing, like every other live test here. + #[tokio::test] + async fn a_standalone_server_is_unsupported() { + let uri = std::env::var("EXTENDDB_TEST_MONGODB_STANDALONE_URI") + .ok() + .filter(|u| !u.trim().is_empty()); + assert!( + uri.is_some() || std::env::var_os("EXTENDDB_TEST_MONGODB_REQUIRED").is_none(), + "EXTENDDB_TEST_MONGODB_REQUIRED is set but EXTENDDB_TEST_MONGODB_STANDALONE_URI is \ + not; the standalone test must run, not skip" + ); + let Some(uri) = uri else { + eprintln!( + "SKIP a_standalone_server_is_unsupported: EXTENDDB_TEST_MONGODB_STANDALONE_URI \ + is not set" + ); + return; + }; + let _hooks = hooks().await; + let engine = fixture(&uri).await; + let err = engine + .open_table_snapshot(ACCOUNT, "anything") + .await + .expect_err("a standalone server cannot serve a snapshot"); + let StorageError::Unsupported(message) = err else { + panic!("a standalone server is Unsupported"); + }; + eprintln!("standalone server: {message}"); + assert!(message.contains("NotAReplicaSet (code 123)"), "{message}"); + } + + #[tokio::test] + async fn snapshot_marker_is_the_cluster_time() { + let Some(uri) = uri() else { + eprintln!("SKIP snapshot_marker_is_the_cluster_time: no MongoDB"); + return; + }; + let _hooks = hooks().await; + let engine = fixture(&uri).await; + engine + .create_table(ACCOUNT, snapshot_conformance::table_input("marker")) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "marker", || async { + let _ = + extenddb_storage::WorkerStore::process_control_plane_transitions(&*engine).await; + }) + .await; + let snapshot = engine + .open_table_snapshot(ACCOUNT, "marker") + .await + .expect("open"); + assert_eq!( + snapshot.info.kind, + SnapshotInfo::KIND_MONGODB_SNAPSHOT_READ_CONCERN + ); + assert!(!snapshot.key_ordered); + assert!(!snapshot.info.marker.is_empty()); + let (secs, inc) = snapshot + .info + .marker + .split_once(':') + .expect("marker is :"); + assert!(secs.parse::().unwrap() > 1_600_000_000); + inc.parse::().unwrap(); + } + + /// A connection string that sets a read concern must not silently turn + /// snapshot reads into ordinary reads: the explicit per-operation read + /// concern wins, and the marker is still a cluster time. + #[tokio::test] + async fn uri_read_concern_does_not_defeat_the_snapshot() { + let Some(uri) = uri() else { + eprintln!("SKIP uri_read_concern_does_not_defeat_the_snapshot: no MongoDB"); + return; + }; + let hooks = hooks().await; + let sep = if uri.contains('?') { '&' } else { '?' }; + let uri = format!("{uri}{sep}readConcernLevel=majority"); + let engine = fixture(&uri).await; + engine + .create_table(ACCOUNT, snapshot_conformance::table_input("rc")) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "rc", || async { + let _ = + extenddb_storage::WorkerStore::process_control_plane_transitions(&*engine).await; + }) + .await; + let snapshot = engine + .open_table_snapshot(ACCOUNT, "rc") + .await + .expect("open"); + assert!( + !snapshot.info.marker.is_empty(), + "snapshot time must be pinned" + ); + // Release the session before the consistency run opens its own. + drop(snapshot); + snapshot_conformance::run(&hooks, &*engine, ACCOUNT, "rc_src", 50, EXPECT, || async { + let _ = + extenddb_storage::WorkerStore::process_control_plane_transitions(&*engine).await; + }) + .await; + } + + /// The fixture's round-trip counter counts the snapshot stream and + /// nothing else: a GSI backfill (which scans the base collection with + /// `find`) leaves it where it was, and a stream at a page of three over + /// twelve items moves it by at least four. + #[tokio::test] + async fn the_round_trip_counter_counts_only_the_snapshot_stream() { + let Some(uri) = uri() else { + eprintln!("SKIP the_round_trip_counter_counts_only_the_snapshot_stream: no MongoDB"); + return; + }; + let _hooks = hooks().await; + let engine = fixture(&uri).await; + let pump = || async { + let _ = + extenddb_storage::WorkerStore::process_control_plane_transitions(&*engine).await; + }; + let before = paging::fetches(); + engine + .create_table(ACCOUNT, snapshot_conformance::table_input("counted")) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "counted", pump).await; + snapshot_conformance::populate(&*engine, ACCOUNT, "counted", 12).await; + let update: extenddb_core::types::UpdateTableInput = + serde_json::from_value(serde_json::json!({ + "TableName": "counted", + "AttributeDefinitions": [{"AttributeName": "u", "AttributeType": "S"}], + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "gu", + "KeySchema": [{"AttributeName": "u", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "KEYS_ONLY"} + }}] + })) + .expect("UpdateTable input"); + engine + .update_table(ACCOUNT, update) + .await + .expect("add a GSI"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "counted", pump).await; + assert_eq!( + paging::fetches(), + before, + "the backfill of the new GSI moved the snapshot round-trip counter" + ); + let page = PageSize::set(3); + let snapshot = engine + .open_table_snapshot(ACCOUNT, "counted") + .await + .expect("open"); + let items: Vec> = + futures::StreamExt::collect(snapshot.items).await; + let items: Vec = items.into_iter().map(|r| r.expect("item")).collect(); + drop(page); + assert_eq!(items.len(), 12); + let counted = paging::fetches() - before; + assert!( + counted >= 4, + "{counted} round trips for twelve items at a page of three" + ); + } +} diff --git a/crates/storage-mongodb/src/table_engine.rs b/crates/storage-mongodb/src/table_engine.rs index aa8c5272c..77706c3c7 100644 --- a/crates/storage-mongodb/src/table_engine.rs +++ b/crates/storage-mongodb/src/table_engine.rs @@ -1361,6 +1361,20 @@ impl MongoEngine { async fn doc_to_table_description( &self, doc: &Document, + ) -> Result { + self.doc_to_table_description_in(doc, None).await + } + + /// Build the description from a `tables` document. When `session` is a + /// snapshot session (see `snapshot_engine`), the index documents are read + /// through it so they come from the same snapshot as `doc`. (The body is + /// the former `doc_to_table_description`, long before the session + /// parameter; splitting it is not part of the snapshot change.) + #[allow(clippy::too_many_lines)] + pub(crate) async fn doc_to_table_description_in( + &self, + doc: &Document, + session: Option<&mut mongodb::ClientSession>, ) -> Result { let id_doc = doc .get_document("_id") @@ -1466,14 +1480,34 @@ impl MongoEngine { // Fetch indexes let indexes_coll = self.catalog_db.collection::("indexes"); use futures::TryStreamExt; - let index_cursor = indexes_coll - .find(doc! { "_id.table_id": &table_id }) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - let index_docs: Vec = index_cursor - .try_collect() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + let index_filter = doc! { "_id.table_id": &table_id }; + let index_docs: Vec = match session { + Some(session) => { + let mut cursor = indexes_coll + .find(index_filter) + .read_concern(mongodb::options::ReadConcern::snapshot()) + .session(&mut *session) + .await + .map_err(|e| crate::snapshot_engine::map_mongo_error(&e))?; + let mut docs = Vec::new(); + while let Some(d) = cursor + .next(&mut *session) + .await + .transpose() + .map_err(|e| crate::snapshot_engine::map_mongo_error(&e))? + { + docs.push(d); + } + docs + } + None => indexes_coll + .find(index_filter) + .await + .map_err(|e| StorageError::Internal(e.to_string()))? + .try_collect() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?, + }; let mut gsis = Vec::new(); let mut lsis = Vec::new(); diff --git a/crates/storage-postgres/Cargo.toml b/crates/storage-postgres/Cargo.toml index 604a31176..887647b57 100755 --- a/crates/storage-postgres/Cargo.toml +++ b/crates/storage-postgres/Cargo.toml @@ -7,10 +7,18 @@ edition.workspace = true rust-version.workspace = true license.workspace = true +[features] +# Compiles the snapshot page-size setter the consistency test uses and the +# shared conformance check. Production builds do not compile them. +test-hooks = ["extenddb-storage/test-hooks"] + [dev-dependencies] # The storage-level vector control-plane tests drive the engine directly against # a scratch database, so they need a runtime and the test macro. tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +# The crate's own tests need `test-hooks`; a dev-dependency on itself turns +# the feature on for test builds only. +extenddb-storage-postgres = { path = ".", features = ["test-hooks"] } [dependencies] anyhow = { workspace = true } diff --git a/crates/storage-postgres/src/backup_engine.rs b/crates/storage-postgres/src/backup_engine.rs index 20baee942..3255dbb2e 100755 --- a/crates/storage-postgres/src/backup_engine.rs +++ b/crates/storage-postgres/src/backup_engine.rs @@ -1077,15 +1077,23 @@ impl BackupEngine for PostgresEngine { // REPEATABLE READ snapshot, so the real relation SELECT must run // before meta.commit(). This fixes the data snapshot at the same // instant as the captured definition; the lock then keeps a DROP - // from removing the relation while its rows are read. The - // control-plane drop path uses a short lock timeout and retries a - // blocked DeleteTable rather than stalling its worker pass. + // from removing the relation while its rows are read. A + // DeleteTable that finds this lock leaves its row DELETING, and + // the control-plane poller drops the tables under a short lock + // timeout, retrying while the lock is held; the guard wakes it + // when this transaction ends (commit, error or cancellation), as + // the table snapshot's does, and the active poller waits on that + // wake as well as on its timer, so the drop follows the backup at + // once. let ddb_table = data_table_name(&table_id); - let mut snapshot = self - .data_pool - .begin_with("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY") - .await - .map_err(db_err)?; + let mut held = crate::snapshot_engine::HeldSnapshot::new( + self.data_pool + .begin_with("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY") + .await + .map_err(db_err)?, + Arc::clone(&self.control_plane_notify), + ); + let snapshot = held.open_conn()?; sqlx::query(&format!("LOCK TABLE {ddb_table} IN ACCESS SHARE MODE")) .execute(&mut *snapshot) .await @@ -1096,29 +1104,32 @@ impl BackupEngine for PostgresEngine { .map_err(db_err)?; meta.commit().await.map_err(db_err)?; - // UpdateTable commits a new index's catalog row before it creates - // and fills the index's data table, and removes the row again if - // that fails. An index whose data table is not there yet is - // therefore not part of the table as of this snapshot; leave it - // out rather than record an index that may never exist. + // Since #401 UpdateTable creates a new GSI's data table before it + // commits the catalog row, so a GSI row without a data table is + // not expected: it is what an index whose table was dropped out + // from under the catalog, or a crash in a pre-#401 build, leaves. + // Such an index is not part of the table as of this snapshot; + // leave it out rather than record an index a restore cannot + // rebuild. The probe is the table snapshot's + // (`snapshot_engine::index_data_table_exists`). let gsi_ids_len = gsi_ids.len(); // The ids were read under the catalog barrier, so this probes // exactly the indexes the captured definition describes. The // to_regclass lookup reads pg_class rather than the pinned data // snapshot and can see later committed DDL; that remains consistent: - // UpdateTable creates and fills a GSI in one data transaction, and - // restore rebuilds its rows from the captured base items. + // a GSI's data table exists from before its catalog row commits + // (#401), and restore rebuilds its rows from the captured base + // items. let mut present = Vec::with_capacity(definition.global_secondary_indexes.len()); for (gsi, (name, index_id)) in std::mem::take(&mut definition.global_secondary_indexes) .into_iter() .zip(gsi_ids) { debug_assert_eq!(gsi.index_name, name); - let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") - .bind(index_table_name(&index_id)) - .fetch_one(&mut *snapshot) - .await - .map_err(db_err)?; + let exists = + crate::snapshot_engine::index_data_table_exists(&mut *snapshot, &index_id) + .await + .map_err(db_err)?; if exists { present.push(gsi); } @@ -1226,7 +1237,7 @@ impl BackupEngine for PostgresEngine { } } count += insert_backup_batch(&mut tx, &backup_arn, &mut pks, &mut datas).await?; - snapshot.commit().await.map_err(db_err)?; + held.finish().await.map_err(db_err)?; (count, size) }; sqlx::query( diff --git a/crates/storage-postgres/src/data/ddl.rs b/crates/storage-postgres/src/data/ddl.rs index f608c89e8..16f1a3172 100755 --- a/crates/storage-postgres/src/data/ddl.rs +++ b/crates/storage-postgres/src/data/ddl.rs @@ -120,26 +120,6 @@ impl PostgresEngine { Ok(()) } - /// Drop the per-DynamoDB-table data table. - /// - /// Called when a table deletion transition completes. - /// - /// # Errors - /// - /// Returns [`StorageError::Internal`] if the DDL execution fails. - pub(crate) async fn drop_data_table( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - table_id: &str, - ) -> Result<(), StorageError> { - let ddb_table = data_table_name(table_id); - let ddl = format!("DROP TABLE IF EXISTS {ddb_table}"); - sqlx::query(&ddl) - .execute(&mut **tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - Ok(()) - } - /// Create a GSI/LSI data table in `PostgreSQL`. /// /// GSI tables use the same `(pk, sk_*)` structure as base tables but add @@ -376,20 +356,6 @@ impl PostgresEngine { Ok(()) } - /// Drop a GSI/LSI data table. - pub(crate) async fn drop_index_data_table( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - index_id: &str, - ) -> Result<(), StorageError> { - let idx_table = index_table_name(index_id); - let ddl = format!("DROP TABLE IF EXISTS {idx_table}"); - sqlx::query(&ddl) - .execute(&mut **tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - Ok(()) - } - /// Fetch key schema and attribute definitions for a table from the catalog. /// /// Uses a single query that combines the table row with an LSI existence @@ -643,8 +609,8 @@ impl PostgresEngine { } impl PostgresEngine { - /// Drop a deleted table's data tables and its queue rows, in an order that - /// cannot deadlock with the propagation queue. + /// The control plane's drop of a `DELETING` table's data tables and queue + /// rows, in an order that cannot deadlock with the propagation queue. /// /// The worker claims a `gsi_pending` row, then takes `FOR SHARE` on the base /// row (the state-based apply). A drop that takes the base table first and @@ -661,66 +627,217 @@ impl PostgresEngine { /// could be waiting on them. The hold is removed in the same transaction as /// the drops, so a table id that no longer exists is never left held; a /// crash in between leaves a hold the recovery sweep clears by age (ten - /// minutes for a hold of this kind, which a drop waiting behind a long - /// client transaction can legitimately keep for a while). + /// minutes for a hold of this kind). /// /// What this does not cover is a deadlock that does not go through the /// queue: the drop waits on a client transaction that waits, through another /// table, on a client transaction queued behind the drop. That is the same - /// exposure any DDL has here and was there before. + /// exposure any DDL has here and was there before; the server's deadlock + /// check ends it, and [`lock_wait_failed`] counts a deadlock victim as a + /// drop to retry. /// - /// With `lock_timeout`, every wait in the dropping transaction (the pending - /// rows' claims, a CreateBackup reading the table, a client transaction) is - /// bounded; a timeout fails the call with nothing dropped, and the hold - /// stays until the next attempt or the recovery sweep. The control plane - /// worker uses it so one table cannot stall a whole pass; DeleteTable itself - /// waits. + /// Every wait in the dropping transaction shares one `budget`: the + /// pending-row delete (a claim in flight) runs under it as its lock + /// timeout, and each `DROP` under what is left of it + /// ([`drop_table_relations`]). Running out fails the call with SQLSTATE + /// 55P03 ([`lock_wait_failed`]) and nothing dropped, and the hold stays + /// until the next attempt or the recovery sweep. `DeleteTable` does not + /// call this: it deletes the pending rows with no timeout and runs + /// [`drop_table_relations`] in a transaction of its own + /// (`delete_table.rs`). + /// + /// Errors are returned as `sqlx` errors so a caller can tell a lock wait + /// from a failure that will repeat. pub(crate) async fn drop_table_storage( data: &sqlx::PgPool, table_id: &str, index_ids: &[String], vector_index_ids: &[String], - lock_timeout: Option<&str>, - ) -> Result<(), StorageError> { + budget: std::time::Duration, + ) -> Result<(), sqlx::Error> { sqlx::query( "INSERT INTO vector_index_holds (table_id, index_id) VALUES ($1, $1) \ ON CONFLICT (table_id, index_id) DO NOTHING", ) .bind(table_id) .execute(data) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - let mut tx = data - .begin() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - if let Some(timeout) = lock_timeout { - sqlx::query("SELECT set_config('lock_timeout', $1, true)") - .bind(timeout) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - } - sqlx::query("DELETE FROM gsi_pending WHERE table_id = $1") - .bind(table_id) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - Self::drop_data_table(&mut tx, table_id).await?; - for idx_id in index_ids { - Self::drop_index_data_table(&mut tx, idx_id).await?; - } - for idx_id in vector_index_ids { - Self::drop_vector_data_table(&mut tx, idx_id).await?; - } - sqlx::query("DELETE FROM vector_index_holds WHERE table_id = $1") - .bind(table_id) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - tx.commit() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - Ok(()) + .await?; + let deadline = std::time::Instant::now() + budget; + let mut tx = data.begin().await?; + set_lock_timeout_until(&mut tx, deadline).await?; + delete_pending_queue_rows(&mut tx, table_id).await?; + drop_table_relations(&mut tx, table_id, index_ids, vector_index_ids, deadline).await?; + tx.commit().await + } +} + +/// The first statement of a table's storage drop, on `conn`'s open +/// transaction: delete the table's pending queue rows, which waits for the +/// claims in flight to commit ([`PostgresEngine::drop_table_storage`]). +pub(crate) async fn delete_pending_queue_rows( + conn: &mut sqlx::PgConnection, + table_id: &str, +) -> Result<(), sqlx::Error> { + sqlx::query("DELETE FROM gsi_pending WHERE table_id = $1") + .bind(table_id) + .execute(&mut *conn) + .await?; + Ok(()) +} + +/// The rest of a table's storage drop, on `conn`'s open transaction after +/// [`delete_pending_queue_rows`]: every relation +/// ([`table_storage_relations`], base table first), then the table's queue +/// holds, all by `deadline`. +/// +/// One deadline for the group, not one lock timeout per statement: before +/// each statement `lock_timeout` is set to what is left until `deadline` +/// (at least 1 ms once it has passed, so a statement that would have to wait +/// fails with 55P03 at once and one that need not wait still runs). A lock +/// timeout per `DROP` let a table with many indexes wait up to the timeout +/// once per relation, holding the base table's `ACCESS EXCLUSIVE` (granted +/// first) for all of it. The transaction's `lock_timeout` is reset to its +/// default afterwards, so statements the caller runs next are not bounded +/// by what was left. +/// +/// The caller owns the transaction and whatever stops the queue before the +/// drop starts ([`PostgresEngine::drop_table_storage`]). +pub(crate) async fn drop_table_relations( + conn: &mut sqlx::PgConnection, + table_id: &str, + index_ids: &[String], + vector_index_ids: &[String], + deadline: std::time::Instant, +) -> Result<(), sqlx::Error> { + for relation in table_storage_relations(table_id, index_ids, vector_index_ids) { + set_lock_timeout_until(conn, deadline).await?; + sqlx::query(&format!("DROP TABLE IF EXISTS {relation}")) + .execute(&mut *conn) + .await?; } + set_lock_timeout_until(conn, deadline).await?; + sqlx::query("DELETE FROM vector_index_holds WHERE table_id = $1") + .bind(table_id) + .execute(&mut *conn) + .await?; + sqlx::query("SET LOCAL lock_timeout TO DEFAULT") + .execute(&mut *conn) + .await?; + Ok(()) +} + +/// Set `conn`'s transaction-local `lock_timeout` to the time left until +/// `deadline`, and to 1 ms once it has passed (`0` would disable the +/// timeout). +async fn set_lock_timeout_until( + conn: &mut sqlx::PgConnection, + deadline: std::time::Instant, +) -> Result<(), sqlx::Error> { + let left = deadline.saturating_duration_since(std::time::Instant::now()); + sqlx::query("SELECT set_config('lock_timeout', $1, true)") + .bind(format!("{}ms", left.as_millis().max(1))) + .execute(&mut *conn) + .await?; + Ok(()) +} + +/// SQLSTATE `lock_not_available` (55P03), which `lock_timeout` raises. +const SQLSTATE_LOCK_NOT_AVAILABLE: &str = "55P03"; + +/// Whether `e` is a lock wait that ran out its `lock_timeout` (55P03): a +/// session held a lock the statement needed, which a later attempt can get +/// once that session ends. Every other failure (a dependent object, a +/// missing privilege, a lost connection) is not one. +pub(crate) fn lock_not_available(e: &sqlx::Error) -> bool { + sqlstate_is(e, &[SQLSTATE_LOCK_NOT_AVAILABLE]) +} + +/// The SQLSTATEs that end a drop's lock wait without anything being wrong +/// with the drop itself: `lock_not_available` (55P03, its `lock_timeout` ran +/// out), `deadlock_detected` (40P01, the server chose its transaction as the +/// victim of a wait cycle with another session) and `query_canceled` (57014, +/// a `statement_timeout` set for the role, the database or the session that +/// is shorter than the lock timeout, or an administrator's +/// `pg_cancel_backend`). Each rolls back the drop's transaction (or the +/// savepoint it ran under), and a later attempt succeeds once the other +/// session is done. +const LOCK_WAIT_FAILURES: [&str; 3] = [SQLSTATE_LOCK_NOT_AVAILABLE, "40P01", "57014"]; + +/// Whether a table storage drop failed in one of its lock waits +/// ([`LOCK_WAIT_FAILURES`]), so the table is still in use and the drop is +/// to be retried rather than reported: `DeleteTable` then takes its +/// asynchronous path, and the control plane counts the row as deferred. +/// The one classifier for both. +pub(crate) fn lock_wait_failed(e: &sqlx::Error) -> bool { + sqlstate_is(e, &LOCK_WAIT_FAILURES) +} + +fn sqlstate_is(e: &sqlx::Error, codes: &[&str]) -> bool { + e.as_database_error() + .and_then(sqlx::error::DatabaseError::code) + .is_some_and(|c| codes.contains(&c.as_ref())) +} + +/// Which sessions [`storage_in_use`] looks for. +#[derive(Debug, Clone, Copy)] +pub(crate) enum LockHolders { + /// Sessions holding `ACCESS SHARE` (`AccessShareLock`): the long + /// readers. `DeleteTable` probes for these only (`delete_table.rs`). + Readers, + /// Sessions holding any lock on the relations. The control plane probes + /// for these before it retries a drop it already had to defer + /// (`worker_store.rs`). + Any, +} + +/// Whether a session other than this connection's holds a granted lock of +/// the `holders` kind on one of `relations` ([`table_storage_relations`]) +/// in this database. Runs on `executor`; `pg_locks` is cluster-wide, so the +/// database filter keeps a lock on the same OID in another database from +/// matching. `pid <> pg_backend_pid()` keeps the caller's own locks out. +pub(crate) async fn storage_in_use<'e, E>( + executor: E, + relations: Vec, + holders: LockHolders, +) -> Result +where + E: sqlx::PgExecutor<'e>, +{ + sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_locks \ + WHERE locktype = 'relation' AND granted \ + AND ($2 OR mode = 'AccessShareLock') \ + AND pid <> pg_backend_pid() \ + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) \ + AND relation = ANY (SELECT to_regclass(name)::oid FROM unnest($1::text[]) AS name))", + ) + .bind(relations) + .bind(matches!(holders, LockHolders::Any)) + .fetch_one(executor) + .await +} + +/// Every data relation one table's storage consists of, in drop order: the +/// base data table, then each GSI/LSI data table, then each vector-index data +/// table. The one list of them: the drops ([`drop_table_relations`]) remove +/// exactly these, and `DeleteTable` and the control plane probe exactly these +/// for holders ([`storage_in_use`]), so a relation added to a table's storage +/// is added here and every path sees it. +/// +/// The base table goes first. A client write holds the base table (`ROW +/// EXCLUSIVE`) before it writes an LSI or a synchronous GSI in the same +/// transaction, and a queue worker holds a base row `FOR SHARE` before it +/// writes the index table. A drop that holds the base waits for neither on +/// an index; a drop that took an index first would wait for such a session +/// on the base while it waits for the drop on the index, a wait cycle that +/// ends only at a lock timeout or the server's deadlock check. +pub(crate) fn table_storage_relations( + table_id: &str, + index_ids: &[String], + vector_index_ids: &[String], +) -> Vec { + std::iter::once(data_table_name(table_id)) + .chain(index_ids.iter().map(|id| index_table_name(id))) + .chain(vector_index_ids.iter().map(|id| vector_table_name(id))) + .collect() } diff --git a/crates/storage-postgres/src/data/mod.rs b/crates/storage-postgres/src/data/mod.rs index a9010220d..7f37ea17d 100755 --- a/crates/storage-postgres/src/data/mod.rs +++ b/crates/storage-postgres/src/data/mod.rs @@ -16,12 +16,14 @@ use extenddb_storage::error::StorageError; /// Includes `account_id` for multi-account isolation (Phase 12a). /// Table names are validated at the engine layer (alphanumeric + `_.-`), /// so this is safe for identifier construction. -pub(crate) fn data_table_name(table_id: &str) -> String { +#[must_use] +pub fn data_table_name(table_id: &str) -> String { format!("\"_ddb_{table_id}\"") } /// SQL table name for a GSI/LSI data table. -pub(crate) fn index_table_name(index_id: &str) -> String { +#[must_use] +pub fn index_table_name(index_id: &str) -> String { format!("\"_ddb_{index_id}\"") } @@ -38,7 +40,8 @@ pub(crate) fn index_table_name(index_id: &str) -> String { /// for the same reason. /// /// The id is a server-generated UUID, so no client input reaches this identifier. -pub(crate) fn vector_table_name(index_id: &str) -> String { +#[must_use] +pub fn vector_table_name(index_id: &str) -> String { format!("\"_ddb_vec_{index_id}\"") } @@ -142,4 +145,8 @@ mod tx_helpers; mod update_item; pub mod vector_index; +pub(crate) use ddl::{ + LockHolders, delete_pending_queue_rows, drop_table_relations, lock_not_available, + lock_wait_failed, storage_in_use, table_storage_relations, +}; pub(crate) use index::{insert_index_rows_multi, item_has_index_keys, project_item_for_index}; diff --git a/crates/storage-postgres/src/data/vector_index.rs b/crates/storage-postgres/src/data/vector_index.rs index 4a8c70bef..b48170e24 100644 --- a/crates/storage-postgres/src/data/vector_index.rs +++ b/crates/storage-postgres/src/data/vector_index.rs @@ -865,10 +865,12 @@ async fn sweep_orphan_holds( engine: &crate::PostgresEngine, live: &[String], ) -> Result<(), StorageError> { - // A hold whose index_id is its own table_id is a DeleteTable stopping the - // queue while it drops the table (`drop_table_storage`). The drop can wait - // behind a long client transaction; sweeping its hold meanwhile only makes - // workers claim rows they then time out on, so such a hold gets longer. + // A hold whose index_id is its own table_id is the control plane stopping + // the queue while it drops a DELETING table (`drop_table_storage`). The + // drop runs within a one-second budget, so a hold that is still there + // after one was left by a process that died mid-drop; it gets the longer + // age all the same, since sweeping a live one only makes workers claim + // rows they then time out on. let swept = sqlx::query( "DELETE FROM vector_index_holds \ WHERE NOT (index_id = ANY($1)) \ diff --git a/crates/storage-postgres/src/delete_table.rs b/crates/storage-postgres/src/delete_table.rs index 7475eed6a..f03d2c7f0 100755 --- a/crates/storage-postgres/src/delete_table.rs +++ b/crates/storage-postgres/src/delete_table.rs @@ -2,15 +2,30 @@ // SPDX-License-Identifier: Apache-2.0 //! `delete_table` implementation for `PostgresEngine`. +//! +//! The synchronous path (`control_plane_delay_seconds` below 1 and no reader +//! on the table's relations) removes the catalog row and drops the data +//! relations in the request. On one shared database both go in one commit. +//! With a separate data database the drops are run first and held +//! uncommitted, the catalog removal commits, and the data transaction +//! commits after it ([`drop_on_data_database`]); a data commit that fails +//! then has the drop retried on fresh connections for +//! [`DATA_DROP_RETRY_BOUND`] before the request gives up. What is left +//! behind when it does, and how to remove it, is stated at +//! [`DATA_DROP_RETRY_BOUND`]. + +use std::time::{Duration, Instant}; use extenddb_core::types::{DeleteTableInput, TableDescription, TableStatus}; use extenddb_storage::error::StorageError; use crate::PostgresEngine; +use crate::data::LockHolders; use crate::table_helpers::{IndexRow, TableRow}; impl PostgresEngine { /// Core implementation of `delete_table` (H-5). + #[allow(clippy::too_many_lines)] pub(crate) async fn delete_table_impl( &self, account_id: &str, @@ -23,7 +38,9 @@ impl PostgresEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; - // Lock and fetch the row atomically with SELECT ... FOR UPDATE + // Lock and fetch the row atomically with SELECT ... FOR UPDATE. A + // DELETING row is fetched too: DeleteTable of a table that is already + // being deleted returns its description, as the service does. let row: Option = sqlx::query_as( r"SELECT table_name, key_schema, attribute_definitions, billing_mode, provisioned_throughput, stream_specification, table_status, @@ -31,7 +48,8 @@ impl PostgresEngine { table_size_bytes, item_count, table_arn, table_id, deletion_protection_enabled, stream_label, table_class, sse_specification, on_demand_throughput - FROM tables WHERE account_id = $1 AND table_name = $2 AND table_status IN ('ACTIVE', 'CREATING') + FROM tables WHERE account_id = $1 AND table_name = $2 + AND table_status IN ('ACTIVE', 'CREATING', 'DELETING') FOR UPDATE", ) .bind(account_id) @@ -42,6 +60,31 @@ impl PostgresEngine { let row = row.ok_or_else(|| StorageError::TableNotFound(input.table_name.clone()))?; + // Fetch indexes for the response description. + let index_rows: Vec = sqlx::query_as( + r"SELECT index_name, index_id, index_type, key_schema, projection, + index_status, provisioned_throughput + FROM indexes WHERE table_id = $1", + ) + .bind(&row.table_id) + .fetch_all(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + + if row.table_status == "DELETING" { + // Already being removed (a deferred drop the control plane + // retries, or a delay of a second or more): nothing to change. + // The rollback releases the row lock. + tx.rollback() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let desc = self.build_table_description_from_row(account_id, row, index_rows)?; + return Ok(TableDescription { + table_status: TableStatus::Deleting, + ..desc + }); + } + // REQ: DeletionProtectionEnabled check — real DynamoDB returns ValidationException if row.deletion_protection_enabled { return Err(StorageError::DeletionProtected(row.table_arn.clone())); @@ -70,17 +113,6 @@ impl PostgresEngine { } } - // Fetch indexes for the response description. - let index_rows: Vec = sqlx::query_as( - r"SELECT index_name, index_id, index_type, key_schema, projection, - index_status, provisioned_throughput - FROM indexes WHERE table_id = $1", - ) - .bind(&row.table_id) - .fetch_all(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - // Vector index ids, read while the rows still exist. Deleting the table // row cascades them away, and each id names a data table that has to be // dropped afterwards: the same ordering the GSI path uses, and for the same @@ -106,42 +138,176 @@ impl PostgresEngine { let index_ids: Vec = index_rows.iter().map(|r| r.index_id.clone()).collect(); - if delay_secs < 1.0 { - // Synchronous delete: remove tags, catalog row, and data tables inline. - // Note: deleting the tables row cascades to indexes and stream rows via FK CASCADE. - // Vector index rows cascade the same way, through - // vector_indexes_table_id_fkey, so a table with vector indexes needs - // no extra catalog cleanup. Their data tables do need dropping, and - // the cascade takes away the rows that name them, which is why the - // index ids are collected BEFORE the catalog row goes. - sqlx::query("DELETE FROM tags WHERE resource_arn = $1") - .bind(&row.table_arn) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - sqlx::query("DELETE FROM tables WHERE table_id = $1") - .bind(&row.table_id) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + // Whether another session is reading the table. A `DROP TABLE` + // waits behind any lock on the relation, so the synchronous path is + // taken only when no session holds `ACCESS SHARE` on any relation the + // drop touches (`table_storage_relations`, the list the drops + // remove); otherwise the row is left `DELETING` and the + // control-plane poller drops the tables once the reader is done + // (`workers.rs`). + // + // Matched: `AccessShareLock` holders, which are the long readers: a + // table snapshot (`snapshot_engine.rs`) and the in-catalog backup hold + // it on the base data table for as long as they stream; a Query or + // Scan with `IndexName` holds it on a GSI data table for the + // statement; pg_dump and any plain `SELECT` or `LOCK ... IN ACCESS + // SHARE MODE` hold it too. + // + // Not matched: every other mode, which PostgreSQL takes instead of + // `ACCESS SHARE` for `SELECT ... FOR SHARE`/`FOR UPDATE`, for + // `INSERT`/`UPDATE`/`DELETE`, for `LOCK ... IN SHARE MODE`, for an + // index build or `VACUUM`: a client write transaction, a GSI backfill + // batch (`gsi_build.rs` reads the base `FOR SHARE SKIP LOCKED` and + // writes the index table), a queue worker's claim, a TTL index build. + // Most are short, and the synchronous path waits for them, within + // `SYNC_DROP_LOCK_TIMEOUT` for the drops; one that is not makes the + // drops run out that budget, and the row is left `DELETING` as for a + // reader, with the drops rolled back (below for each layout). The + // queue claim is left unmatched on purpose: the pending-row delete + // waits for it with no timeout. + // + // The probe runs only when its answer is used (the synchronous path + // needs `delay_secs < 1.0`). On one shared pool it runs on the + // connection this transaction already holds: taking a second one from + // the pool while holding the first lets `pool_size` concurrent + // DeleteTables hold every connection and wait out the acquire timeout + // for a second. `pg_locks` is cluster-wide, and `current_database()` + // is then the data database. With a separate data database it runs on + // the data pool; every path acquires catalog then data, so that order + // cannot form a cycle. `pid <> pg_backend_pid()` is defensive: this + // transaction has touched only catalog tables, so it cannot hold a + // lock on a data relation and no test can make it; the clause keeps a + // future statement added above the probe from matching itself. + let data_table_in_use = if delay_secs < 1.0 { + let relations = + crate::data::table_storage_relations(&row.table_id, &index_ids, &vector_index_ids); + if self.shared_data_pool { + crate::data::storage_in_use(&mut *tx, relations, LockHolders::Readers).await + } else { + crate::data::storage_in_use(&self.data_pool, relations, LockHolders::Readers).await + } + .map_err(|e| StorageError::Internal(e.to_string()))? + } else { + false + }; - tx.commit() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + // Between the probe and the drops: the window in which a reader that + // the probe did not see can take its lock. + #[cfg(feature = "test-hooks")] + test_pause::reached(test_pause::DeleteTablePoint::BeforeDrops).await; - // Drop data tables after catalog commit. See `drop_table_storage` for - // why the queue is stopped before anything is dropped. - Self::drop_table_storage( + // Deleting the tables row cascades to indexes and stream rows via FK + // CASCADE, and to vector index rows through + // vector_indexes_table_id_fkey. The data tables need dropping, and + // the cascade takes away the rows that name them, which is why the + // index ids are collected before the catalog row goes. On either + // layout nothing commits before every drop has its lock: a reader + // the probe missed (an open between its snapshot and its `LOCK + // TABLE`, a session that took its lock after the probe) or a holder + // above makes a drop wait, `SYNC_DROP_LOCK_TIMEOUT` bounds the wait + // for the whole group of drops, and a wait that fails + // (`lock_wait_failed`) rolls the drops back and takes the + // asynchronous path below on this transaction, with the row still + // held `FOR UPDATE`. A caller cancelled before the commit leaves both + // transactions to roll back: the table `ACTIVE` with its data. + let mut data_commit: Option> = None; + let removed = if delay_secs < 1.0 && !data_table_in_use { + if self.shared_data_pool { + // One database: the drops run in this transaction, behind a + // savepoint, so the row and its data tables go in one commit + // or not at all. + remove_table_and_storage( + &mut tx, + &row.table_arn, + &row.table_id, + &index_ids, + &vector_index_ids, + ) + .await? + } else if let Some(data_tx) = drop_on_data_database( &self.data_pool, &row.table_id, &index_ids, &vector_index_ids, - None, ) - .await?; + .await? + { + // Separate data database: no transaction covers both. The + // drops hold their locks, uncommitted, in `data_tx`; the + // catalog removal is made here and the two commits follow, + // catalog first, below. + sqlx::query("DELETE FROM tags WHERE resource_arn = $1") + .bind(&row.table_arn) + .execute(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + sqlx::query("DELETE FROM tables WHERE table_id = $1") + .bind(&row.table_id) + .execute(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + data_commit = Some(data_tx); + true + } else { + false + } } else { - // Async delete: set DELETING status with scheduled removal time. + false + }; + + if removed { + if let Some(data_tx) = data_commit { + // Both commits, and the retries a failed data commit needs, + // run in a task of their own, which the request awaits: a + // caller cancelled from here on cannot stop between the + // commits and leave the row gone with the data tables kept, + // and cannot stop the retries either (dropping the request + // detaches the task; it does not cancel it). The drops + // already hold their locks, so the data commit does not + // wait. Once the catalog commit is in, the row is gone + // whatever happens next: a data commit that fails (the + // connection terminated, a failover, a network fault) has + // the drop redone on a fresh connection by + // `drop_after_catalog_commit`, and only a process that dies, + // or a data database unreachable for longer than + // `DATA_DROP_RETRY_BOUND`, between the catalog commit and + // the last retry leaves the data tables without a row. + let data_pool = self.data_pool.clone(); + let table_id = row.table_id.clone(); + tokio::spawn(async move { + tx.commit() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + #[cfg(feature = "test-hooks")] + test_pause::reached(test_pause::DeleteTablePoint::BetweenCommits).await; + match data_tx.commit().await { + Ok(()) => Ok(()), + Err(e) => { + drop_after_catalog_commit( + &data_pool, + &table_id, + &index_ids, + &vector_index_ids, + &e, + ) + .await + } + } + }) + .await + .map_err(|e| StorageError::Internal(e.to_string()))??; + } else { + tx.commit() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + } + } else { + // Async delete: set DELETING status with a scheduled removal time. + // A drop deferred because the table is in use (`delay_secs < 1`) + // is due at once: the table is already being removed, and a + // reader that ends right after this has the poller's next pass + // drop it rather than a pass after the delay. + let schedule_secs = if delay_secs < 1.0 { 0.0 } else { delay_secs }; sqlx::query( r"UPDATE tables SET table_status = 'DELETING', status_transition_at = NOW() + make_interval(secs => $3) @@ -149,15 +315,13 @@ impl PostgresEngine { ) .bind(account_id) .bind(&input.table_name) - .bind(delay_secs) + .bind(schedule_secs) .execute(&mut *tx) .await .map_err(|e| StorageError::Internal(e.to_string()))?; - tx.commit() .await .map_err(|e| StorageError::Internal(e.to_string()))?; - // F-3: Wake the control plane poller so it processes the DELETING → // removed transition without waiting for the idle timeout. self.control_plane_notify.notify_one(); @@ -172,3 +336,387 @@ impl PostgresEngine { }) } } + +/// The budget of the synchronous `DROP TABLE`s of one `DeleteTable`, and of +/// each drop the control plane retries (`worker_store.rs`). One deadline +/// for the whole group of drops (`data::drop_table_relations`), not a lock +/// timeout per `DROP`. +/// +/// `DeleteTable`'s drops run only after the in-use probe found no `ACCESS +/// SHARE` holder, so what they usually wait on is short: a write +/// transaction, a GSI backfill batch, or an open between its snapshot and +/// its `LOCK TABLE`. One second covers one such batch or request at +/// ordinary latencies; a holder that keeps its lock longer (an open write +/// transaction, an index build, `VACUUM`) costs the request this second and +/// sends it down the asynchronous path. The bound that results for a +/// `DeleteTable` is the wait for the pending queue rows' claims, which has no +/// timeout by design (a claim is short, and the drop is meant to wait for +/// it), plus at most this budget for the drops. It is also the longest a +/// drop's queued `ACCESS EXCLUSIVE` request can stall other requests on the +/// table, and, after the claims, the longest this `DeleteTable` holds its +/// connections and the row `FOR UPDATE` before it gives up; giving up costs +/// only the asynchronous path. +pub(crate) const SYNC_DROP_LOCK_TIMEOUT: Duration = Duration::from_secs(1); + +/// Separate data database: how long a `DeleteTable` whose data commit failed +/// after its catalog commit keeps retrying the drop of the table's data +/// relations ([`drop_after_catalog_commit`]). Attempts start until this much +/// time has passed since the failed commit; the last one may run on past it +/// for its own waits (the drop's `SYNC_DROP_LOCK_TIMEOUT`, or the data pool's +/// 30 s acquire timeout when the server does not answer at all). +/// +/// One minute covers the outages a data database recovers from on its own +/// (a terminated or lost connection, a restart, a failover, each of which +/// the pool's own acquire timeout bounds at 30 s), and keeps the request, +/// which the client is still waiting on, from outliving the client's +/// patience by much. An outage longer than this is an incident an operator +/// is looking at, and the error log written then (the table id, the relation +/// names, and the `DROP TABLE` to run) is what they need to finish the +/// delete by hand. +/// +/// What is left behind when the retries are exhausted, or when the process +/// dies between the catalog commit and the last retry: the table's data +/// relations (`_ddb_`, one `_ddb_` per GSI or LSI, one +/// `_ddb_vec_` per vector index) on the data database with no +/// catalog row naming them, and possibly a `vector_index_holds` row for the +/// table id, which the hold sweep clears by age. Nothing in the server drops +/// such relations later: `DescribeTable` reports the table gone, a +/// `CreateTable` of the same name gets a new table id, and the control +/// plane enumerates relations from catalog rows only. They stay until an +/// operator drops them (`docs/troubleshooting.md`, "Data tables without a +/// catalog row"). On one shared database the row and the drops are one +/// commit, and none of this applies. +pub(crate) const DATA_DROP_RETRY_BOUND: Duration = Duration::from_secs(60); + +/// The retry bound in force: [`DATA_DROP_RETRY_BOUND`], or the shorter one a +/// test set (`test-hooks` only). +fn data_drop_retry_bound() -> Duration { + #[cfg(feature = "test-hooks")] + if let Some(bound) = test_pause::shortened_data_drop_retry_bound() { + return bound; + } + DATA_DROP_RETRY_BOUND +} + +/// Separate data database: the catalog commit is in and the data commit +/// failed with `commit_error`, so the table's row is gone while its data +/// relations may still exist (a commit whose reply was lost may also have +/// gone through; nothing here can tell). Redo the drop on fresh connections +/// from `data` until it succeeds or [`DATA_DROP_RETRY_BOUND`] has passed. +/// +/// Each attempt is the control plane's own drop +/// ([`PostgresEngine::drop_table_storage`]: a queue hold, the pending queue +/// rows, every relation, the hold, under `SYNC_DROP_LOCK_TIMEOUT`), whose +/// statements are all idempotent (`DROP TABLE IF EXISTS`, `DELETE ... WHERE +/// table_id = $1`, `INSERT ... ON CONFLICT DO NOTHING`), so repeating a drop +/// that did go through changes nothing. Every failure is retried, after a +/// backoff from 100 ms doubling to 5 s: a lock wait that fails +/// (`lock_wait_failed`) because a session took a lock after the first drop's +/// locks were released with its connection, a connection the pool cannot +/// open, a statement the server did not answer. On success the delete is +/// complete and the caller returns the description as usual. When the bound +/// is exhausted the relation names and the table id are logged at error with +/// the `DROP TABLE` that removes them, and the request fails with +/// `Internal`; the row is gone either way. +async fn drop_after_catalog_commit( + data: &sqlx::PgPool, + table_id: &str, + index_ids: &[String], + vector_index_ids: &[String], + commit_error: &sqlx::Error, +) -> Result<(), StorageError> { + let bound = data_drop_retry_bound(); + let started = Instant::now(); + let relations = crate::data::table_storage_relations(table_id, index_ids, vector_index_ids); + tracing::warn!( + table_id, + ?relations, + "DeleteTable: the catalog row is removed but the data commit failed ({commit_error}); \ + retrying the drop of the data relations for up to {bound:?}" + ); + let mut backoff = Duration::from_millis(100); + let mut attempts: u32 = 0; + let last_error = loop { + attempts += 1; + match PostgresEngine::drop_table_storage( + data, + table_id, + index_ids, + vector_index_ids, + SYNC_DROP_LOCK_TIMEOUT, + ) + .await + { + Ok(()) => { + tracing::info!( + table_id, + "DeleteTable: the data relations were dropped on attempt {attempts}, \ + {:?} after the failed data commit", + started.elapsed() + ); + return Ok(()); + } + Err(e) => { + let elapsed = started.elapsed(); + if elapsed >= bound { + break e; + } + tokio::time::sleep(backoff.min(bound.saturating_sub(elapsed))).await; + backoff = (backoff * 2).min(Duration::from_secs(5)); + } + } + }; + let drop_statement = format!("DROP TABLE IF EXISTS {};", relations.join(", ")); + tracing::error!( + table_id, + ?relations, + "DeleteTable: the catalog row of table {table_id} is removed but its data relations \ + could not be dropped in {attempts} attempt(s) over {:?} ({last_error}); they stay on \ + the data database until dropped by hand: {drop_statement}", + started.elapsed() + ); + Err(StorageError::Internal(format!( + "DeleteTable: table {table_id} is removed from the catalog but its data relations \ + {relations:?} could not be dropped ({last_error})" + ))) +} + +/// Separate data database: begin a data transaction and run the table's +/// storage drop in it, leaving it uncommitted for the caller to commit after +/// the catalog removal. `Ok(Some(tx))`: every drop holds its lock. `Ok(None)`: +/// a lock wait failed (`data::lock_wait_failed`), the data transaction is +/// rolled back, and the caller takes the asynchronous path. Any other failure +/// is returned, and the data transaction rolls back when it is dropped. +/// +/// Nothing commits here, so a failure up to and including the catalog commit +/// leaves the table as it was. After the catalog commit the row is gone, and +/// a data commit that fails for any reason (not only a process death) is +/// followed by [`drop_after_catalog_commit`]'s retries; what a process death +/// or an outage longer than [`DATA_DROP_RETRY_BOUND`] leaves behind is +/// stated there. +/// +/// The order is `remove_table_and_storage`'s: the pending queue rows first, +/// with no timeout (a claim in flight is short and the drop waits for it, as +/// the control plane's drop does), then the relations, base table first, by +/// one deadline `SYNC_DROP_LOCK_TIMEOUT` away. Unlike the control plane's +/// drop (`PostgresEngine::drop_table_storage`), no queue hold is committed +/// first. The hold stops new claims for the table, and none is needed here: +/// a claim taken after the pending rows are deleted holds a base row `FOR +/// SHARE` before it writes an index table, so it either finishes before the +/// base drop is granted (the drop waits for it within the deadline) or waits +/// for the drop, holding nothing the drop needs; it cannot hold an index +/// table the drop is waiting for, because the base table is dropped first. +/// Committing a hold would also be the one write here that outlives a +/// cancelled request, and a hold left behind stops the table's queue until +/// the recovery sweep removes it ten minutes later. +async fn drop_on_data_database( + data: &sqlx::PgPool, + table_id: &str, + index_ids: &[String], + vector_index_ids: &[String], +) -> Result>, StorageError> { + let mut data_tx = data + .begin() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let dropped = async { + crate::data::delete_pending_queue_rows(&mut data_tx, table_id).await?; + let deadline = std::time::Instant::now() + SYNC_DROP_LOCK_TIMEOUT; + crate::data::drop_table_relations( + &mut data_tx, + table_id, + index_ids, + vector_index_ids, + deadline, + ) + .await + } + .await; + match dropped { + Ok(()) => Ok(Some(data_tx)), + Err(e) if crate::data::lock_wait_failed(&e) => { + tracing::info!( + table_id, + "DeleteTable: a session holds the table's storage ({e}); leaving it DELETING \ + for the control plane to drop" + ); + data_tx + .rollback() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + Ok(None) + } + Err(e) => Err(StorageError::Internal(e.to_string())), + } +} + +/// One shared database: drop the table's storage and remove its tags and +/// catalog row in `tx`, behind a savepoint. `Ok(true)`: done, for the caller +/// to commit. `Ok(false)`: a lock wait failed (`data::lock_wait_failed`: +/// the drops ran out `SYNC_DROP_LOCK_TIMEOUT`, were chosen as a deadlock +/// victim, or hit a `statement_timeout`) and everything since the savepoint +/// is rolled back, so the caller can take the asynchronous path on the same +/// transaction with the row still held `FOR UPDATE`. Any other failure is +/// returned, and the caller's rollback leaves the table as it was. The row +/// and the drops are one commit, so on this layout no failure at any point +/// leaves a data relation without its row (the separate layout's exposure is +/// at [`DATA_DROP_RETRY_BOUND`]). +/// +/// The order is `drop_on_data_database`'s, and for the same reason no queue +/// hold is committed first (nothing here commits before the row goes): the +/// pending queue rows first, which waits for the claims in flight to commit +/// with no timeout (a caller cancelled in that wait has committed nothing), +/// then the relations, base table first, by one deadline +/// `SYNC_DROP_LOCK_TIMEOUT` away. +async fn remove_table_and_storage( + tx: &mut sqlx::PgConnection, + table_arn: &str, + table_id: &str, + index_ids: &[String], + vector_index_ids: &[String], +) -> Result { + let internal = |e: sqlx::Error| StorageError::Internal(e.to_string()); + sqlx::query("SAVEPOINT delete_table_storage") + .execute(&mut *tx) + .await + .map_err(internal)?; + match remove_after_savepoint(tx, table_arn, table_id, index_ids, vector_index_ids).await { + Ok(()) => { + sqlx::query("RELEASE SAVEPOINT delete_table_storage") + .execute(&mut *tx) + .await + .map_err(internal)?; + Ok(true) + } + Err(e) if crate::data::lock_wait_failed(&e) => { + tracing::info!( + table_id, + "DeleteTable: a session holds the table's storage ({e}); leaving it DELETING \ + for the control plane to drop" + ); + sqlx::query("ROLLBACK TO SAVEPOINT delete_table_storage") + .execute(&mut *tx) + .await + .map_err(internal)?; + Ok(false) + } + Err(e) => Err(internal(e)), + } +} + +async fn remove_after_savepoint( + tx: &mut sqlx::PgConnection, + table_arn: &str, + table_id: &str, + index_ids: &[String], + vector_index_ids: &[String], +) -> Result<(), sqlx::Error> { + crate::data::delete_pending_queue_rows(tx, table_id).await?; + let deadline = std::time::Instant::now() + SYNC_DROP_LOCK_TIMEOUT; + crate::data::drop_table_relations(tx, table_id, index_ids, vector_index_ids, deadline).await?; + sqlx::query("DELETE FROM tags WHERE resource_arn = $1") + .bind(table_arn) + .execute(&mut *tx) + .await?; + sqlx::query("DELETE FROM tables WHERE table_id = $1") + .bind(table_id) + .execute(&mut *tx) + .await?; + Ok(()) +} + +/// Test hooks: hold the next `DeleteTable` at a chosen point, and shorten +/// the retry bound of a drop redone after the catalog commit. Compiled only +/// with `test-hooks`. +#[cfg(feature = "test-hooks")] +pub(crate) mod test_pause { + use std::sync::atomic::{AtomicU64, Ordering}; + use std::time::Duration; + + use extenddb_storage::test_pause::{Pause, PausePoint}; + + /// Where a `DeleteTable` can be held. + #[doc(hidden)] + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub enum DeleteTablePoint { + /// Between the in-use probe and the drops, so a test can take a lock + /// on the table in the window a reader the probe missed has. + BeforeDrops, + /// Separate data database only: between the catalog commit and the + /// data commit of the synchronous path, so a test can make the data + /// commit fail with the row already gone. + BetweenCommits, + } + + static PAUSE: PausePoint = PausePoint::new(); + + /// A pause armed by [`arm_delete_table_pause`] or + /// [`arm_delete_table_pause_at`]. Dropping it disarms the pause and + /// releases a `DeleteTable` waiting at it. + #[doc(hidden)] + pub type DeleteTablePause = Pause; + + /// Arm a pause for the next `DeleteTable` in this process that reaches + /// the point between its probe and its drops (one call, once), whatever + /// its engine or table. Process-global: a test that arms it holds the + /// snapshot hooks lock, as every test in `tests/table_snapshot.rs` does. + /// + /// # Panics + /// + /// In debug builds, when a pause is already armed ([`PausePoint::arm`]). + #[doc(hidden)] + #[must_use] + pub fn arm_delete_table_pause() -> DeleteTablePause { + arm_delete_table_pause_at(DeleteTablePoint::BeforeDrops) + } + + /// Arm a pause for the next `DeleteTable` in this process that reaches + /// `point`; otherwise as [`arm_delete_table_pause`]. + /// + /// # Panics + /// + /// In debug builds, when a pause is already armed ([`PausePoint::arm`]). + #[doc(hidden)] + #[must_use] + pub fn arm_delete_table_pause_at(point: DeleteTablePoint) -> DeleteTablePause { + PAUSE.arm(point) + } + + pub(crate) async fn reached(point: DeleteTablePoint) { + PAUSE.reached(point).await; + } + + /// Milliseconds; 0 means the production bound. + static DATA_DROP_RETRY_BOUND_MS: AtomicU64 = AtomicU64::new(0); + + /// The shortened bound set by [`shorten_data_drop_retry_bound_for_tests`] + /// while its guard lives. + pub(crate) fn shortened_data_drop_retry_bound() -> Option { + match DATA_DROP_RETRY_BOUND_MS.load(Ordering::Relaxed) { + 0 => None, + ms => Some(Duration::from_millis(ms)), + } + } + + /// Restores the production retry bound when dropped. + #[doc(hidden)] + pub struct DataDropRetryBound(()); + + impl Drop for DataDropRetryBound { + fn drop(&mut self) { + DATA_DROP_RETRY_BOUND_MS.store(0, Ordering::Relaxed); + } + } + + /// Make every `DeleteTable` in this process whose data commit fails + /// after its catalog commit give up after `bound` instead of + /// `DATA_DROP_RETRY_BOUND`, until the guard is dropped. Process-global, + /// under the snapshot hooks lock like the pauses. + #[doc(hidden)] + #[must_use] + pub fn shorten_data_drop_retry_bound_for_tests(bound: Duration) -> DataDropRetryBound { + let ms = u64::try_from(bound.as_millis()).unwrap_or(u64::MAX).max(1); + DATA_DROP_RETRY_BOUND_MS.store(ms, Ordering::Relaxed); + DataDropRetryBound(()) + } +} diff --git a/crates/storage-postgres/src/lib.rs b/crates/storage-postgres/src/lib.rs index e9f1ae2e9..5a9d733c0 100755 --- a/crates/storage-postgres/src/lib.rs +++ b/crates/storage-postgres/src/lib.rs @@ -24,6 +24,9 @@ mod metadata_engine; mod migrations; mod operations; mod pg_util; +mod snapshot_engine; +#[doc(hidden)] +pub use snapshot_engine::map_fetch_error; mod stream_engine; mod table_engine; mod table_helpers; @@ -72,6 +75,25 @@ pub use data::vector_index::rebuild_stuck_vector_indexes; /// Rebuild vector indexes a crash left mid-build. A startup step, exported because /// it is part of bringing a deployment up rather than an internal detail. pub use data::vector_index::reconcile_incomplete_vector_indexes; +/// The relation names a table's storage uses, for tests that pin what a drop +/// removes; the crate owns the format (`data/mod.rs`). +#[cfg(feature = "test-hooks")] +#[doc(hidden)] +pub use data::{data_table_name, index_table_name, vector_table_name}; +/// Hold a `DeleteTable` between its in-use probe and its drops, or (separate +/// data database) between its catalog commit and its data commit; and shorten +/// the bound on the drop retried after a failed data commit. +#[cfg(feature = "test-hooks")] +#[doc(hidden)] +pub use delete_table::test_pause::{ + DataDropRetryBound, DeleteTablePause, DeleteTablePoint, arm_delete_table_pause, + arm_delete_table_pause_at, shorten_data_drop_retry_bound_for_tests, +}; +/// Every migration, as `extenddb migrate --yes` applies them, for the test +/// harnesses in `tests/` so they do not carry their own copy of the lists. +#[cfg(feature = "test-hooks")] +#[doc(hidden)] +pub use migrations::apply_all_migrations; /// The `PostgreSQL` storage backend. /// @@ -246,6 +268,67 @@ pub struct PostgresEngine { /// Restores each account may run at once in this process; see /// `backup_engine::RestoreAccountSlots`. pub(crate) restore_account_slots: backup_engine::RestoreAccountSlots, + /// Bounds how many table snapshots may be open at once (one permit from + /// the open until the stream's transaction ends): half the pool + /// ([`snapshot_slot_count`]) on either layout. + /// + /// On one shared database an open and its stream use one connection and + /// one transaction for catalog and data, so the permits cap snapshots at + /// half the pool's connections and no open ever waits for a second + /// connection while holding one. With a separate data database an open + /// holds a catalog connection (the barrier) while it takes a data + /// connection, and the stream keeps the data connection; the permits cap + /// barriers at half the catalog pool and streams at half the data pool, + /// and the catalog-then-data order every path follows cannot form a + /// cycle. + /// + /// What the permits do not prevent: request load alone exhausting a pool + /// (an open then waits for its connection up to the acquire timeout like + /// any request), and restores, which take their own slots (half the + /// pool, at most 16: `backup_engine::restore_server_slot_count`), so + /// snapshots and restores together can hold the whole pool. On the separate layout requests that block on + /// the `tables` row an open holds `FOR SHARE` (`UpdateTable`, + /// `DeleteTable` of that table) hold catalog connections while they + /// wait; the open needs only a data connection by then, so they delay it + /// no further. + pub(crate) snapshot_open_slots: Arc, + /// Whether `data_pool` is the catalog pool (no separate data database). + /// The snapshot open uses one transaction on that layout, and + /// `DeleteTable` runs its in-use probe and its synchronous drop on the + /// connection it holds rather than taking a second from the same pool. + pub(crate) shared_data_pool: bool, + /// Drop attempts the control plane has made (`control_plane_pass`), a + /// `DELETING` table its probe found held and skipped counted as one, for + /// tests that count them. Shared by clones. + #[cfg(feature = "test-hooks")] + pub(crate) drop_attempts: Arc, +} + +/// Table snapshots one engine may hold open at once: half the pool (at least +/// five, since the pool is at least `MIN_POOL_SIZE`). Each open snapshot +/// holds one connection of the pool its items come from (on one shared +/// database, the one pool), so open snapshots leave at least half of it to +/// requests; see `PostgresEngine::snapshot_open_slots` for what that does and +/// does not prevent on each layout. +fn snapshot_slot_count(pool_size: u32) -> usize { + usize::try_from((pool_size / 2).max(1)).unwrap_or(usize::MAX) +} + +/// Stops the poller started by +/// [`PostgresEngine::spawn_control_plane_poller_for_tests`] when dropped. +#[cfg(feature = "test-hooks")] +#[doc(hidden)] +pub struct ControlPlanePollerGuard { + token: extenddb_storage::CancellationToken, + handle: tokio::task::JoinHandle<()>, +} + +#[cfg(feature = "test-hooks")] +impl Drop for ControlPlanePollerGuard { + fn drop(&mut self) { + self.token.cancel(); + self.handle.abort(); + } } impl PostgresEngine { @@ -278,6 +361,7 @@ impl PostgresEngine { // P54 Bug 1: Read data database connection string from catalog settings. // Falls back to the catalog pool if no separate data database is configured. + let mut shared_data_pool = false; let data_pool = match sqlx::query_as::<_, (String,)>( "SELECT value FROM settings WHERE key = 'data_database_connection_string'", ) @@ -294,7 +378,10 @@ impl PostgresEngine { .map_err(|e| { StorageError::Connection(format!("data database connection failed: {e}")) })?, - _ => pool.clone(), + _ => { + shared_data_pool = true; + pool.clone() + } }; // P119: Read initial GSI propagation delay from settings table. @@ -335,6 +422,12 @@ impl PostgresEngine { restore_server_slots: Arc::new(tokio::sync::Semaphore::new(server_slots)), restore_server_slot_count: server_slots, restore_account_slots: backup_engine::RestoreAccountSlots::default(), + snapshot_open_slots: Arc::new(tokio::sync::Semaphore::new(snapshot_slot_count( + pool_size, + ))), + shared_data_pool, + #[cfg(feature = "test-hooks")] + drop_attempts: Arc::default(), }) } @@ -387,6 +480,57 @@ impl PostgresEngine { Arc::clone(&self.control_plane_notify) } + /// Run the control-plane poller `serve` runs (`spawn_workers`), on a + /// clone of this engine with a settings store on its catalog pool, until + /// the returned guard is dropped. For tests that need the production + /// wake-up path rather than calling `process_control_plane_transitions` + /// themselves. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + #[must_use] + pub fn spawn_control_plane_poller_for_tests(&self) -> ControlPlanePollerGuard { + self.spawn_control_plane_poller_with_active_poll_for_tests(workers::ACTIVE_POLL) + } + + /// [`Self::spawn_control_plane_poller_for_tests`] with the active loop's + /// retry interval set to `active_poll` instead of one second, so a test + /// can tell a pass run by a wake from one run by the timer. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + #[must_use] + pub fn spawn_control_plane_poller_with_active_poll_for_tests( + &self, + active_poll: std::time::Duration, + ) -> ControlPlanePollerGuard { + let token = extenddb_storage::CancellationToken::new(); + let handle = tokio::spawn(workers::poll_control_plane_transitions( + Arc::new(self.clone()), + self.control_plane_notify(), + Arc::new(PostgresCatalogStore::new(self.pool.clone())), + token.clone(), + active_poll, + )); + ControlPlanePollerGuard { token, handle } + } + + /// Run the startup recovery `serve` runs before it starts the poller + /// (`server_components_factory`), on this engine. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub async fn recover_control_plane_at_startup_for_tests(&self) { + workers::recover_control_plane_at_startup(self).await; + } + + /// Drop attempts this engine's control plane has made so far, a table + /// its probe found held and skipped counted as one. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + #[must_use] + pub fn control_plane_drop_attempts_for_tests(&self) -> u64 { + self.drop_attempts + .load(std::sync::atomic::Ordering::Relaxed) + } + /// Defense-in-depth: validate `account_id` before use in SQL identifiers. /// /// `account_id` is interpolated into SQL identifiers via `data_table_name()`. @@ -653,6 +797,7 @@ impl ServerRuntimeHooks for PostgresRuntimeHooks { cp_notify, catalog_store, token, + workers::ACTIVE_POLL, ) .await; }); @@ -784,16 +929,10 @@ fn server_components_factory( Err(e) => tracing::error!("Failed to reconcile incomplete vector indexes: {e}"), } - // Recover control plane transitions (ignore errors) - match engine.process_control_plane_transitions().await { - Ok(ref t) if t.is_empty() => {} - Ok(transitions) => { - for (name, transition) in &transitions { - tracing::info!("Recovered table '{name}': {transition}"); - } - } - Err(e) => tracing::error!("Failed to recover control plane transitions: {e}"), - } + // Recover control plane transitions (errors are logged, not fatal). + // A drop deferred by a session still holding the table wakes the + // poller, which `spawn_workers` starts later. + workers::recover_control_plane_at_startup(&engine).await; // Start GSI workers let engine = engine.start_gsi_workers(); diff --git a/crates/storage-postgres/src/migrations.rs b/crates/storage-postgres/src/migrations.rs index 2873bbc77..952f7920a 100755 --- a/crates/storage-postgres/src/migrations.rs +++ b/crates/storage-postgres/src/migrations.rs @@ -145,9 +145,6 @@ const CATALOG_SHAPE_STEPS: &[ShapeStep] = &[ /// timeout the server refuses to start instead. pub(crate) const SERVE_SHAPE_LOCK_TIMEOUT_SECS: u64 = 5; -/// SQLSTATE `lock_not_available`, which `lock_timeout` raises. -const SQLSTATE_LOCK_NOT_AVAILABLE: &str = "55P03"; - async fn shape_step_needed(pool: &PgPool, step: &ShapeStep) -> OpResult { sqlx::query_scalar(step.needed) .fetch_one(pool) @@ -198,11 +195,7 @@ pub(crate) async fn converge_catalog_shapes_at_serve(pool: &PgPool) -> OpResult< continue; } let fail = |e: sqlx::Error| { - let lock_timeout = e - .as_database_error() - .and_then(|d| d.code()) - .is_some_and(|c| c == SQLSTATE_LOCK_NOT_AVAILABLE); - if lock_timeout { + if crate::data::lock_not_available(&e) { OpError::Internal(format!( "Catalog shape step ({}): could not lock the table within \ {SERVE_SHAPE_LOCK_TIMEOUT_SECS} s, another session holds it", @@ -373,6 +366,30 @@ pub(crate) async fn run_data_code_migrations( Ok(()) } +/// Bring a fresh catalog database and a fresh data database (the same pool +/// twice on the shared layout) to the current schema the way `extenddb +/// migrate --yes` does: every catalog migration, every data migration, then +/// the data code migrations, each recorded in its ledger. For test harnesses +/// in this crate's own `tests/` directory, which otherwise copy the file +/// lists and go stale when a migration is added here. +/// +/// # Errors +/// +/// Returns the first migration's error, formatted. +#[cfg(feature = "test-hooks")] +#[doc(hidden)] +pub async fn apply_all_migrations(catalog: &PgPool, data: &PgPool) -> Result<(), String> { + run_catalog_migrations(catalog) + .await + .map_err(|e| format!("{e:?}"))?; + run_data_migrations(data) + .await + .map_err(|e| format!("{e:?}"))?; + run_data_code_migrations(catalog, data) + .await + .map_err(|e| format!("{e:?}")) +} + /// Create the base-table-key index on every existing GSI/LSI table. /// /// During GSI propagation each index table (`_ddb_`) is looked up back to diff --git a/crates/storage-postgres/src/snapshot_engine.rs b/crates/storage-postgres/src/snapshot_engine.rs new file mode 100644 index 000000000..c9be8e553 --- /dev/null +++ b/crates/storage-postgres/src/snapshot_engine.rs @@ -0,0 +1,680 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! `SnapshotEngine` for `PostgreSQL`. +//! +//! The definition is in the catalog and the items in the data tables. How the +//! open fixes one instant for both depends on the layout. +//! +//! One shared database (no `data_database_connection_string`): catalog and +//! data are one database, so one transaction reads both. The open begins +//! `REPEATABLE READ READ ONLY`, and its first statement reads the `tables` +//! row, which pins the transaction's snapshot; the description, the GSI ids +//! and the items are then all read from that one MVCC snapshot, so no row +//! lock or barrier is needed and `UpdateTable` and `DeleteTable` never wait +//! on an open's row lock. The open then takes `LOCK TABLE ... IN ACCESS SHARE MODE` on +//! the data table. A `DeleteTable` that committed and dropped the data table +//! between the snapshot and the `LOCK` makes the lock fail with +//! `undefined_table` (42P01), which the open reports as `TableNotFound`. A +//! `DeleteTable` whose in-use probe runs in that same window (after the +//! snapshot, before the `LOCK` is granted) does not see the open; on this +//! layout its drops run in its own catalog transaction within a 1 s budget +//! before the catalog removal commits, so if the `LOCK` is granted first the +//! drop times out, is rolled back with the removal, and the `DeleteTable` +//! takes its asynchronous path (`delete_table.rs`). The row and the data +//! tables are removed together or not at all. +//! +//! Separate data database: no one transaction covers both, and the fence is +//! the one the in-catalog backup uses: the `tables` row is held `FOR SHARE` +//! on a catalog transaction from the definition read until the data +//! transaction (`REPEATABLE READ READ ONLY`) holds `ACCESS SHARE` on the data +//! table and has pinned its snapshot. `UpdateTable` and `DeleteTable` take +//! that row `FOR UPDATE`, so neither can commit in the gap, and a +//! `DeleteTable` waiting on the row finds the `ACCESS SHARE` when it probes. +//! The GSI data-table probe also runs before the barrier is released. +//! +//! On either layout, background workers that flip an index status (a vector +//! build finishing, `mark_active`) do not take the row lock, so on the +//! separate layout an index can be reported `CREATING` while the items are a +//! few milliseconds newer; the description keeps it either way (only an index +//! being deleted is dropped), and a restore builds it from the items. +//! +//! The in-catalog backup (`backup_engine.rs`) opens its data snapshot the +//! same way as the separate layout (`REPEATABLE READ READ ONLY`, `ACCESS +//! SHARE` on the relation) on both layouts, reads the definition on its +//! barrier connection, and drops a GSI whose data table is missing, and the +//! attribute definitions only that GSI used, by the same rule (the probe is +//! one function, [`index_data_table_exists`]); this module differs in +//! returning the definition and the stream together, in ordering by the +//! primary key, and in using one transaction on the shared layout. +//! +//! Connections: on one shared database an open and its stream hold one +//! connection, from the open until the stream's transaction ends. With a +//! separate data database the open holds one catalog connection (the +//! barrier, which also reads the definition) until the data snapshot is +//! pinned, and the stream holds one data connection until it ends; every +//! path takes catalog then data, so the two pools cannot deadlock each +//! other. The engine's `snapshot_open_slots` bounds open snapshots (one +//! permit per snapshot, from the open until the stream's transaction ends) +//! at half the pool (`snapshot_slot_count`) on both layouts, so open +//! snapshots leave at least half of the pool their items come from to +//! requests (with the bound at the whole pool, ten open snapshots on a +//! ten-connection data pool made `GetItem` and `DeleteTable` wait out the +//! pool's 30 s acquire timeout). The `snapshot_open_slots` field doc says +//! what the permits do not prevent. An open over the bound is refused with +//! `LimitExceeded` rather than queued. The backup design's +//! `max_concurrent_jobs` is the caller-side bound; the restore path's slots +//! (`restore_server_slot_count`) do not count snapshots. +//! +//! `ACCESS SHARE` on the data table is held until the stream's last fetch +//! (the transaction commits as soon as a short page shows the end), not +//! until the stream is dropped. While it is held a `DROP TABLE` waits; +//! `DeleteTable` therefore takes its asynchronous path for a table another +//! session holds `ACCESS SHARE` on (any relation of its storage, +//! `data::table_storage_relations`), so the catalog row is not removed while +//! its data tables are still in use. The poller stays active while a +//! `DELETING` row's drop is deferred (`workers.rs`): each pass probes the +//! table for holders and skips it while one remains, then the poller pauses +//! for a second, so the drop comes at the next pass after the release. Every +//! path that ends this stream's data transaction after it may have taken +//! `ACCESS SHARE` also wakes the poller, and the active poller waits on that +//! wake as well as on its timer, so the drop is retried at once: the +//! transaction lives in a guard (`HeldSnapshot`) from +//! `BEGIN` on, whose drop wakes the poller (an open that fails or is +//! cancelled at any await, a stream dropped before its end, a stream that +//! hit an error), and the last page wakes it before it awaits the commit. +//! +//! `LOCK TABLE ... IN ACCESS SHARE MODE` keeps a `DROP TABLE` from removing +//! the relation while it is read (ordinary DML is unaffected). `LOCK` is a +//! utility statement and does not pin the snapshot; on the separate layout +//! the `SELECT pg_current_snapshot()` that follows does, and on the shared +//! layout the first catalog read already did. Its text is the manifest +//! marker: the `xmin:xmax:xip` of the snapshot, recorded so a later +//! change-log archive can tell which commits the snapshot includes; nothing +//! reopens a snapshot from it. Items are then read through a server-side +//! cursor in batches. +//! +//! Memory: one `FETCH` is `paging::DEFAULT_FETCH_BATCH` (1000) rows, all of +//! them held as item text at once. At the default 400 KB item limit one +//! row's text can reach 27.7 MB (the measured worst case in #394's line-cap +//! table), so a fetch is bounded by rows, not bytes. A byte budget (keyset +//! paging with a byte cutoff) is required before the stream gets a caller; +//! the backup design assigns it to the backup job PR. +//! +//! Order: the cursor orders by the primary key (`key_order_clause`) only when +//! the text key columns (`pk`, and `sk_s` for an `S` range key) have a +//! collation named "C" or "POSIX", read from `pg_attribute` inside the +//! snapshot transaction. The rule is the name: any other collation, +//! including "default" on a database whose own collation is C, streams +//! unordered, which is complete. A table created before the key columns were +//! declared `COLLATE "C"` keeps the database collation, and ordering it +//! bytewise is a Seq Scan and an external sort of the whole table before the +//! first FETCH, with the lock, the connection and the permit held +//! throughout. Such a table is streamed with a plain `SELECT`, reported +//! `key_ordered = false`, and a warning names it. + +use std::sync::Arc; + +use extenddb_core::types::ScalarAttributeType; +use extenddb_core::types::{Item, TableDescription}; +use extenddb_storage::error::StorageError; +use extenddb_storage::snapshot::{drop_deleting_indexes, key_schema_is_orderable, paging}; +use extenddb_storage::util::sk_column; +use extenddb_storage::{SnapshotEngine, SnapshotInfo, TableSnapshot}; +use futures::future::BoxFuture; +use futures::stream::BoxStream; +use sqlx::{PgConnection, Postgres, Transaction}; +use tokio::sync::{Notify, OwnedSemaphorePermit}; + +use crate::PostgresEngine; +use crate::data::{all_sort_key_info, data_table_name, index_table_name}; + +const CURSOR: &str = "extenddb_snapshot"; + +const BEGIN_SNAPSHOT: &str = "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY"; + +fn db_err(e: &sqlx::Error) -> StorageError { + StorageError::Internal(format!("Database error: {e}")) +} + +impl SnapshotEngine for PostgresEngine { + fn open_table_snapshot( + &self, + account_id: &str, + table_name: &str, + ) -> BoxFuture<'_, Result> { + let account_id = account_id.to_owned(); + let table_name = table_name.to_owned(); + Box::pin(async move { + self.open_table_snapshot_impl(&account_id, &table_name) + .await + }) + } +} + +/// The transaction that holds the snapshot's `ACCESS SHARE`, from `BEGIN` to +/// its end. Every way it ends wakes the control-plane poller: a +/// `DeleteTable` that found the table in use left the row `DELETING`, and +/// the active poller, which waits on the wake as well as on its one-second +/// timer, then retries the drop at once rather than at its next poll (the +/// poller stays active while a drop is deferred, so the wake saves up to a +/// second, not the difference between a second and its 60 s idle pass). +/// Dropping the guard with the transaction still set rolls it back (sqlx +/// queues the rollback and returns the connection through the runtime) and +/// wakes the poller; the poller's drop waits on the lock under a lock +/// timeout, so waking it before the rollback has run is harmless. Waking +/// when no lock was taken yet is harmless too. The in-catalog backup +/// (`backup_engine.rs`) holds its data snapshot in the same guard. +pub(crate) struct HeldSnapshot { + tx: Option>, + wake: Arc, +} + +impl HeldSnapshot { + pub(crate) fn new(tx: Transaction<'static, Postgres>, wake: Arc) -> Self { + Self { tx: Some(tx), wake } + } + + /// The open transaction's connection; `None` once it has ended. + fn conn(&mut self) -> Option<&mut PgConnection> { + self.tx.as_deref_mut() + } + + /// The connection during the open, where the transaction cannot have + /// ended yet. + pub(crate) fn open_conn(&mut self) -> Result<&mut PgConnection, StorageError> { + self.conn().ok_or_else(|| { + StorageError::Internal("table snapshot transaction ended during the open".to_owned()) + }) + } + + /// End the transaction by rolling it back, and wake the poller. + fn release(&mut self) { + if self.tx.take().is_some() { + self.wake.notify_one(); + } + } + + /// End the transaction by committing it after the last row. The poller + /// is woken before the commit is awaited, so a stream dropped during + /// that await has already woken it. The commit's own error is returned + /// for a caller that reports it (the in-catalog backup); the stream + /// ignores it, since a `READ ONLY` transaction's commit has nothing to + /// fail on that the caller could act on. + pub(crate) async fn finish(&mut self) -> Result<(), sqlx::Error> { + if let Some(tx) = self.tx.take() { + self.wake.notify_one(); + tx.commit().await?; + } + Ok(()) + } +} + +impl Drop for HeldSnapshot { + fn drop(&mut self) { + self.release(); + } +} + +/// What the layout-specific half of the open hands to the common half: the +/// held transaction with `ACCESS SHARE` taken and the snapshot pinned, the +/// description as of that snapshot, and the marker. +struct Pinned { + held: HeldSnapshot, + description: TableDescription, + table_id: String, + marker: String, +} + +impl PostgresEngine { + async fn open_table_snapshot_impl( + &self, + account_id: &str, + table_name: &str, + ) -> Result { + // Held for the snapshot's life (module doc, Connections); refused at + // once rather than queued, as the backup design refuses a job over + // its limit. + let slot = Arc::clone(&self.snapshot_open_slots) + .try_acquire_owned() + .map_err(|_| { + StorageError::LimitExceeded( + "too many table snapshots are open on this server; retry when one ends" + .to_owned(), + ) + })?; + let pinned = if self.shared_data_pool { + self.pin_on_one_transaction(account_id, table_name).await? + } else { + self.pin_behind_barrier(account_id, table_name).await? + }; + open_cursor(pinned, slot, table_name).await + } + + /// Shared layout (module doc): one `REPEATABLE READ READ ONLY` + /// transaction whose first statement, the `tables` row read, pins the + /// snapshot for the definition and the items alike. + async fn pin_on_one_transaction( + &self, + account_id: &str, + table_name: &str, + ) -> Result { + let mut held = HeldSnapshot::new( + self.pool + .begin_with(BEGIN_SNAPSHOT) + .await + .map_err(|e| db_err(&e))?, + Arc::clone(&self.control_plane_notify), + ); + let table_id = active_table_id(held.open_conn()?, account_id, table_name, "").await?; + let (mut description, gsi_ids) = self + .definition_on(held.open_conn()?, account_id, table_name, &table_id) + .await?; + // Between the snapshot and the LOCK: the window in which a + // DeleteTable can drop the data table (module doc). + paging::open_pause(paging::OpenStage::AfterDefinition).await; + lock_data_table(held.open_conn()?, &table_id, table_name).await?; + let marker = current_snapshot(held.open_conn()?).await?; + leave_out_gsis_without_data_table(held.open_conn()?, &mut description, &gsi_ids).await?; + Ok(Pinned { + held, + description, + table_id, + marker, + }) + } + + /// Separate layout (module doc): the `tables` row is held `FOR SHARE` on + /// a catalog transaction until the data transaction holds `ACCESS SHARE` + /// and has pinned its snapshot, and until the GSI data-table probe has + /// run. + async fn pin_behind_barrier( + &self, + account_id: &str, + table_name: &str, + ) -> Result { + let mut meta = self.pool.begin().await.map_err(|e| db_err(&e))?; + let table_id = active_table_id(&mut meta, account_id, table_name, " FOR SHARE").await?; + // The description is read on the connection that holds the row, so + // an open needs exactly one catalog connection (the in-catalog backup + // reads its definition on its barrier transaction the same way). A + // concurrent UpdateTable blocks on the FOR UPDATE it needs before it + // can commit, so what is read here is what is in force. + let (mut description, gsi_ids) = self + .definition_on(&mut meta, account_id, table_name, &table_id) + .await?; + let mut held = HeldSnapshot::new( + self.data_pool + .begin_with(BEGIN_SNAPSHOT) + .await + .map_err(|e| db_err(&e))?, + Arc::clone(&self.control_plane_notify), + ); + lock_data_table(held.open_conn()?, &table_id, table_name).await?; + // Directly before the pin, with the barrier held: a barrier released + // any earlier would fall inside this pause, where the conformance + // check's writes would then be visible to the pin. + paging::open_pause(paging::OpenStage::AfterDefinition).await; + let marker = current_snapshot(held.open_conn()?).await?; + // Before the barrier is released: an UpdateTable dropping a GSI's + // data table cannot commit until then, so a GSI that existed at the + // instant is not left out. + leave_out_gsis_without_data_table(held.open_conn()?, &mut description, &gsi_ids).await?; + meta.commit().await.map_err(|e| db_err(&e))?; + Ok(Pinned { + held, + description, + table_id, + marker, + }) + } + + /// The description (minus indexes being deleted) and the GSI ids, read + /// on `conn`. + async fn definition_on( + &self, + conn: &mut PgConnection, + account_id: &str, + table_name: &str, + table_id: &str, + ) -> Result<(TableDescription, Vec<(String, String)>), StorageError> { + let mut description = self + .build_table_description_on(conn, account_id, table_name) + .await?; + drop_deleting_indexes(&mut description); + let gsi_ids: Vec<(String, String)> = sqlx::query_as( + "SELECT index_name, index_id FROM indexes \ + WHERE table_id = $1 AND index_type = 'GSI' ORDER BY index_name", + ) + .bind(table_id) + .fetch_all(&mut *conn) + .await + .map_err(|e| db_err(&e))?; + Ok((description, gsi_ids)) + } +} + +/// Read the `tables` row (with `lock`, `""` or `" FOR SHARE"`) and return the +/// table id of an `ACTIVE` table. +async fn active_table_id( + conn: &mut PgConnection, + account_id: &str, + table_name: &str, + lock: &str, +) -> Result { + let row: Option<(String, String)> = sqlx::query_as(&format!( + "SELECT table_id, table_status FROM tables \ + WHERE account_id = $1 AND table_name = $2{lock}" + )) + .bind(account_id) + .bind(table_name) + .fetch_optional(&mut *conn) + .await + .map_err(|e| db_err(&e))?; + let (table_id, status) = + row.ok_or_else(|| StorageError::TableNotFound(table_name.to_owned()))?; + if status != "ACTIVE" { + return Err(StorageError::TableNotActive(table_name.to_owned())); + } + Ok(table_id) +} + +/// `ACCESS SHARE` on the data table. A table dropped since the snapshot saw +/// its row (`undefined_table`, 42P01) is `TableNotFound`. +async fn lock_data_table( + conn: &mut PgConnection, + table_id: &str, + table_name: &str, +) -> Result<(), StorageError> { + let ddb_table = data_table_name(table_id); + match sqlx::query(&format!("LOCK TABLE {ddb_table} IN ACCESS SHARE MODE")) + .execute(&mut *conn) + .await + { + Ok(_) => Ok(()), + Err(sqlx::Error::Database(db)) if db.code().as_deref() == Some("42P01") => { + Err(StorageError::TableNotFound(table_name.to_owned())) + } + Err(e) => Err(db_err(&e)), + } +} + +/// `pg_current_snapshot()` text: on the separate layout the first +/// non-utility statement of the data transaction, which pins its snapshot. +async fn current_snapshot(conn: &mut PgConnection) -> Result { + sqlx::query_scalar("SELECT pg_current_snapshot()::text") + .fetch_one(&mut *conn) + .await + .map_err(|e| db_err(&e)) +} + +/// Whether a GSI's data table exists. Reads `pg_class`, not the pinned +/// snapshot. Shared with the in-catalog backup, which leaves out a GSI +/// without its data table by the same rule. +pub(crate) async fn index_data_table_exists( + conn: &mut PgConnection, + index_id: &str, +) -> Result { + sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(index_table_name(index_id)) + .fetch_one(&mut *conn) + .await +} + +/// Since #401 `UpdateTable` creates a new GSI's data table before it commits +/// the catalog row, so a GSI row without a data table is not expected; it is +/// what an index whose table was dropped out from under the catalog, or a +/// crash in a pre-#401 build, leaves. Such an index is not part of the table +/// as of this snapshot: leave it out, and the attribute definitions only it +/// used, as the in-catalog backup does, rather than describe an index a +/// restore cannot rebuild or a definition `CreateTable` would refuse. A GSI +/// in the description with no id row is left out the same way. A restore +/// rebuilds index rows from the items. +async fn leave_out_gsis_without_data_table( + conn: &mut PgConnection, + description: &mut TableDescription, + gsi_ids: &[(String, String)], +) -> Result<(), StorageError> { + let mut omitted = false; + if let Some(gsis) = description.global_secondary_indexes.as_mut() { + let mut present = Vec::with_capacity(gsis.len()); + for gsi in std::mem::take(gsis) { + let Some((_, index_id)) = gsi_ids.iter().find(|(name, _)| *name == gsi.index_name) + else { + omitted = true; + continue; + }; + if index_data_table_exists(conn, index_id) + .await + .map_err(|e| db_err(&e))? + { + present.push(gsi); + } else { + omitted = true; + } + } + *gsis = present; + } + if omitted { + prune_unused_attribute_definitions(description); + } + Ok(()) +} + +/// The common half of the open: decide the order, then `DECLARE` the cursor +/// on the held transaction. +async fn open_cursor( + mut pinned: Pinned, + slot: OwnedSemaphorePermit, + table_name: &str, +) -> Result { + let ddb_table = data_table_name(&pinned.table_id); + // Ordered only when the text key columns already sort bytewise: on a + // table whose key columns keep the database collation the ORDER BY would + // be a Seq Scan and a full sort before the first FETCH, with the lock, + // the connection and the permit held throughout (module doc). + let mut order_by = key_order_clause(&pinned.description); + if order_by.is_some() { + let not_bytewise: Vec = sqlx::query_scalar( + "SELECT a.attname::text FROM pg_attribute a \ + LEFT JOIN pg_collation c ON c.oid = a.attcollation \ + WHERE a.attrelid = to_regclass($1) AND a.attname = ANY($2) \ + AND NOT a.attisdropped \ + AND COALESCE(c.collname, '') NOT IN ('C', 'POSIX')", + ) + .bind(&ddb_table) + .bind(key_text_columns(&pinned.description)) + .fetch_all(pinned.held.open_conn()?) + .await + .map_err(|e| db_err(&e))?; + if !not_bytewise.is_empty() { + tracing::warn!( + table = table_name, + table_id = %pinned.table_id, + columns = ?not_bytewise, + "table snapshot: key columns do not have collation \"C\"; streaming the \ + table unordered (key_ordered = false) rather than sorting it" + ); + order_by = None; + } + } + let key_ordered = order_by.is_some(); + let select = match &order_by { + Some(cols) => format!("SELECT item_data::text FROM {ddb_table} ORDER BY {cols}"), + None => format!("SELECT item_data::text FROM {ddb_table}"), + }; + paging::open_pause(paging::OpenStage::BeforeItems).await; + sqlx::query(&format!("DECLARE {CURSOR} NO SCROLL CURSOR FOR {select}")) + .execute(pinned.held.open_conn()?) + .await + .map_err(|e| db_err(&e))?; + + Ok(TableSnapshot { + description: pinned.description, + info: SnapshotInfo { + kind: SnapshotInfo::KIND_POSTGRES_REPEATABLE_READ.to_owned(), + marker: pinned.marker, + }, + key_ordered, + items: cursor_stream(pinned.held, slot), + }) +} + +/// The text key columns whose collation decides the order: `pk`, and `sk_s` +/// for an `S` range key (`sk_n` is `NUMERIC` and `sk_b` is `bytea`). +fn key_text_columns(description: &TableDescription) -> Vec { + let sks = all_sort_key_info(&description.key_schema, &description.attribute_definitions); + let mut columns = vec!["pk".to_owned()]; + if matches!(sks.as_slice(), [(_, ScalarAttributeType::S)]) { + columns.push("sk_s".to_owned()); + } + columns +} + +/// Drop attribute definitions no key schema (table, GSI, LSI) names, after +/// an index has been left out of a description. `CreateTable` refuses a +/// definition no key uses. +fn prune_unused_attribute_definitions(description: &mut TableDescription) { + let mut used: std::collections::HashSet<&str> = description + .key_schema + .iter() + .map(|k| k.attribute_name.as_str()) + .collect(); + for gsi in description.global_secondary_indexes.iter().flatten() { + used.extend(gsi.key_schema.iter().map(|k| k.attribute_name.as_str())); + } + for lsi in description.local_secondary_indexes.iter().flatten() { + used.extend(lsi.key_schema.iter().map(|k| k.attribute_name.as_str())); + } + let used: std::collections::HashSet = used.into_iter().map(str::to_owned).collect(); + description + .attribute_definitions + .retain(|d| used.contains(&d.attribute_name)); +} + +/// `ORDER BY` columns for the format's key order, or `None` when the key +/// shape has none (`key_schema_is_orderable`). The open uses the clause only +/// after checking that the text key columns (`pk`, and `sk_s` for an `S` +/// range key) have collation "C" or "POSIX" (`pg_attribute.attcollation`); +/// a table created before the key columns were declared `TEXT COLLATE "C"` +/// keeps the database collation (`en_US.utf8` in the default setups, where +/// `a < B`) and is streamed unordered instead, because `ORDER BY ... COLLATE +/// "C"` on it is a full sort before the first row. On the tables that reach +/// the clause the explicit `COLLATE "C"` therefore does not change the order; +/// it is kept as defence in depth (no test can tell it from its absence), so +/// the query states the order it relies on. `pk` and `sk_s` sort bytewise and +/// `sk_n` (`NUMERIC`) and `sk_b` (`BYTEA`) sort by value: the order +/// `extenddb_core::backup_format::key_order` defines. +fn key_order_clause(description: &TableDescription) -> Option { + if !key_schema_is_orderable(description) { + return None; + } + let sks = all_sort_key_info(&description.key_schema, &description.attribute_definitions); + match sks.as_slice() { + [] => Some("pk COLLATE \"C\"".to_owned()), + [(_, ScalarAttributeType::S)] => Some("pk COLLATE \"C\", sk_s COLLATE \"C\"".to_owned()), + [(_, sk_type)] => Some(format!("pk COLLATE \"C\", {}", sk_column(*sk_type))), + _ => None, + } +} + +/// Stream items from the open cursor, one `FETCH` per batch. The held +/// transaction lives in the stream's state: the last page commits it (waking +/// the poller first), an error rolls it back, and dropping the stream before +/// its end rolls it back through `HeldSnapshot`'s drop; each wakes the +/// control-plane poller. +fn cursor_stream( + held: HeldSnapshot, + slot: OwnedSemaphorePermit, +) -> BoxStream<'static, Result> { + struct State { + held: HeldSnapshot, + /// Released with the transaction: the slot counts the connection. + slot: Option, + buffer: std::vec::IntoIter, + done: bool, + } + impl State { + /// End the stream after an error: roll back, wake the poller, free + /// the slot. + fn abandon(&mut self) { + self.done = true; + self.held.release(); + self.slot = None; + } + } + let state = State { + held, + slot: Some(slot), + buffer: Vec::new().into_iter(), + done: false, + }; + let batch = paging::fetch_batch(); + Box::pin(futures::stream::unfold( + state, + move |mut state| async move { + loop { + if let Some(item) = state.buffer.next() { + return Some((Ok(item), state)); + } + if state.done { + return None; + } + let conn = state.held.conn()?; + paging::record_fetch(); + let rows: Result, sqlx::Error> = + sqlx::query_as(&format!("FETCH {batch} FROM {CURSOR}")) + .fetch_all(conn) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(e) => { + state.abandon(); + return Some((Err(map_fetch_error(&e)), state)); + } + }; + if rows.len() < batch { + state.done = true; + // Close the transaction now rather than at drop so vacuum is + // released as soon as the last row is out. + let _ = state.held.finish().await; + state.slot = None; + } + let mut items = Vec::with_capacity(rows.len()); + for (data,) in rows { + match serde_json::from_str::(&data) { + Ok(item) => items.push(item), + Err(e) => { + state.abandon(); + return Some(( + Err(StorageError::Internal(format!("Parse item: {e}"))), + state, + )); + } + } + } + state.buffer = items.into_iter(); + } + }, + )) +} + +/// Two SQLSTATEs mean the server discarded the snapshot's rows before the +/// stream drained them: 40001 (`serialization_failure`), which on a +/// `REPEATABLE READ READ ONLY` transaction can only come from a hot standby +/// cancelling the query because its snapshot conflicts with recovery, and +/// 72000 (`snapshot_too_old`), raised on a primary when +/// `old_snapshot_threshold` is set and the snapshot outlives it, on +/// PostgreSQL 16 and earlier (17 removed the setting, so there only 40001 +/// reaches this). Everything +/// else is an ordinary database error. +#[doc(hidden)] +#[must_use] +pub fn map_fetch_error(e: &sqlx::Error) -> StorageError { + if let sqlx::Error::Database(db) = e + && matches!(db.code().as_deref(), Some("40001" | "72000")) + { + return StorageError::SnapshotExpired(db.message().to_owned()); + } + db_err(e) +} diff --git a/crates/storage-postgres/src/table_helpers.rs b/crates/storage-postgres/src/table_helpers.rs index 4dfb7f18f..545f3701e 100755 --- a/crates/storage-postgres/src/table_helpers.rs +++ b/crates/storage-postgres/src/table_helpers.rs @@ -97,14 +97,39 @@ impl PostgresEngine { data::index_table_name(index_id) } - // TODO(fidelity): These two queries are not in a transaction. Under concurrent - // UpdateTable (future phase), the table row and index rows could be read at - // different points in time, producing an inconsistent snapshot. Wrap in a - // transaction or use SELECT ... FOR SHARE when concurrent DDL is supported. + /// `DescribeTable`'s read of the definition, on a connection taken for the + /// four statements and nothing else: the table row and its index rows are + /// not read in one transaction, so a concurrent `UpdateTable` committing + /// between them can show a row with an index set from a moment later. That + /// is what the service's eventually consistent `DescribeTable` also allows. + /// The snapshot engine (`snapshot_engine.rs`) is the caller that needs one + /// instant and gets it: it calls [`Self::build_table_description_on`] on + /// the connection that holds the `tables` row `FOR SHARE` (separate data + /// database) or inside its `REPEATABLE READ` transaction (one shared + /// database). The end of `update_table` calls this one, after its own + /// transaction has committed. pub(crate) async fn build_table_description( &self, account_id: &str, table_name: &str, + ) -> Result { + let mut conn = self + .pool + .acquire() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + self.build_table_description_on(&mut conn, account_id, table_name) + .await + } + + /// Build the description through one catalog connection, so a caller + /// that already holds one (the snapshot engine, with the `tables` row + /// `FOR SHARE`) does not need a second from the pool while it waits. + pub(crate) async fn build_table_description_on( + &self, + conn: &mut sqlx::PgConnection, + account_id: &str, + table_name: &str, ) -> Result { let row: Option = sqlx::query_as( r"SELECT table_name, key_schema, attribute_definitions, billing_mode, @@ -117,7 +142,7 @@ impl PostgresEngine { ) .bind(account_id) .bind(table_name) - .fetch_optional(&self.pool) + .fetch_optional(&mut *conn) .await .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -129,7 +154,7 @@ impl PostgresEngine { FROM indexes WHERE table_id = $1", ) .bind(&row.table_id) - .fetch_all(&self.pool) + .fetch_all(&mut *conn) .await .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -138,14 +163,15 @@ impl PostgresEngine { ORDER BY index_name" )) .bind(&row.table_id) - .fetch_all(&self.pool) + .fetch_all(&mut *conn) .await .map_err(|e| StorageError::Internal(e.to_string()))?; let table_name_owned = row.table_name.clone(); let table_id = row.table_id.clone(); let mut desc = self.build_table_description_from_row(account_id, row, index_rows)?; - desc.restore_summary = self.restore_summary(&table_id, &desc.table_status).await?; + desc.restore_summary = + Self::restore_summary_on(conn, &table_id, &desc.table_status).await?; desc.vector_indexes = vector_index_descriptions( &self.region, account_id, @@ -342,9 +368,10 @@ impl PostgresEngine { /// The RestoreSummary of a table created by RestoreTableFromBackup, or /// `None` for any other table. In progress until the table is ACTIVE; the /// engine drops the summary from DescribeTable once the table is not - /// CREATING. - pub(crate) async fn restore_summary( - &self, + /// CREATING. Takes the connection so a caller holding the catalog barrier + /// reads it there. + pub(crate) async fn restore_summary_on( + conn: &mut sqlx::PgConnection, table_id: &str, status: &extenddb_core::types::TableStatus, ) -> Result, StorageError> { @@ -356,7 +383,7 @@ impl PostgresEngine { FROM table_restores WHERE table_id = $1", ) .bind(table_id) - .fetch_optional(&self.pool) + .fetch_optional(&mut *conn) .await .map_err(|e| StorageError::Internal(e.to_string()))?; Ok(row.map( diff --git a/crates/storage-postgres/src/worker_store.rs b/crates/storage-postgres/src/worker_store.rs index 0f6871145..e24f8f38e 100755 --- a/crates/storage-postgres/src/worker_store.rs +++ b/crates/storage-postgres/src/worker_store.rs @@ -9,6 +9,8 @@ use extenddb_storage::WorkerStore; use extenddb_storage::error::StorageError; use crate::PostgresEngine; +use crate::data::LockHolders; +use crate::delete_table::SYNC_DROP_LOCK_TIMEOUT; impl WorkerStore for PostgresEngine { fn process_control_plane_transitions( @@ -21,6 +23,24 @@ impl WorkerStore for PostgresEngine { } } +/// What one pass of [`PostgresEngine::control_plane_pass`] did. +pub(crate) struct ControlPlanePass { + /// `(table_name, transition)` pairs, for the caller's log. + pub(crate) transitions: Vec<(String, &'static str)>, + /// `DELETING` rows that were due but still in use in this pass: the + /// probe found another session holding a lock on one of the table's + /// relations, or the drop failed in a lock wait + /// (`data::lock_wait_failed`). The drop succeeds once that session ends. + /// While this is non-zero the poller keeps polling rather than wait for a + /// wake, since the holder (a plain transaction, `pg_dump`, a client read + /// or write) may not be one that wakes it. A drop that failed for any + /// other reason (a dependent object, a privilege) is not counted: it + /// would fail again every second, so it is logged and retried at the + /// poller's normal cadence, the next pass after a wake or the 60 s idle + /// pass. + pub(crate) deferred_drops: usize, +} + impl PostgresEngine { /// Process pending control plane transitions (H-5). /// @@ -28,8 +48,10 @@ impl PostgresEngine { /// moved to ACTIVE. Tables in DELETING state whose transition time has /// passed are removed (along with their indexes and tags). /// - /// Called by the background poller in `cmd_serve`. Also called at startup - /// to recover in-flight operations from a previous server instance. + /// Called by tests that drive the control plane themselves. Startup + /// recovery (`workers::recover_control_plane_at_startup`) and the + /// background poller call `control_plane_pass`, which also reports the + /// drops this pass had to defer. /// /// Returns a list of `(table_name, transition)` pairs describing what /// changed, so the caller can log meaningful state-change messages (D-4). @@ -40,18 +62,25 @@ impl PostgresEngine { pub async fn process_control_plane_transitions( &self, ) -> Result, StorageError> { + Ok(self.control_plane_pass().await?.transitions) + } + + /// One pass of the control plane: [`Self::process_control_plane_transitions`] + /// plus the count of `DELETING` rows whose drop was deferred by the lock + /// timeout ([`ControlPlanePass::deferred_drops`]), which the poller uses + /// to decide whether to stay active. + /// + /// # Errors + /// + /// Returns [`StorageError`] if the database is unreachable or a query fails. + pub(crate) async fn control_plane_pass(&self) -> Result { let mut transitions = Vec::new(); + let mut deferred_drops = 0; - // CREATING → ACTIVE - let activated: Vec<(String,)> = sqlx::query_as( - r"UPDATE tables SET table_status = 'ACTIVE', status_transition_at = NULL - WHERE table_status = 'CREATING' AND status_transition_at <= NOW() - RETURNING table_name", - ) - .fetch_all(&self.pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - for (name,) in activated { + // CREATING → ACTIVE, in its own statement ahead of the drops, so a + // table that is due is activated at the start of every pass whatever + // the drops below find. + for name in self.activate_due_tables().await? { transitions.push((name, "CREATING → active")); } @@ -91,26 +120,38 @@ impl PostgresEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; - // See `drop_table_storage` for why the queue is stopped before - // anything is dropped. CreateBackup holds ACCESS SHARE through its - // item read; the lock timeout keeps a DROP waiting for that lock from - // stalling every transition in this worker pass. A failure leaves - // the catalog row DELETING for a later retry. - let dropped = Self::drop_table_storage( - &self.data_pool, - table_id, - &index_ids, - &vector_index_ids, - Some("3s"), - ) - .await; - - if let Err(e) = dropped { - tracing::warn!( - "could not drop data tables for '{name}' ({table_id}); the table remains \ - DELETING and the control plane will retry: {e}" - ); - continue; + let dropped = self + .drop_unless_held(&mut tx, table_id, &index_ids, &vector_index_ids) + .await?; + #[cfg(feature = "test-hooks")] + self.drop_attempts + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + + match dropped { + None => { + tracing::debug!( + "table '{name}' ({table_id}) stays DELETING: a session holds a lock on \ + its data tables" + ); + deferred_drops += 1; + continue; + } + Some(Ok(())) => {} + Some(Err(e)) if crate::data::lock_wait_failed(&e) => { + tracing::info!( + "table '{name}' ({table_id}) stays DELETING: a session holds its data \ + tables ({e}); the control plane retries the drop while it waits" + ); + deferred_drops += 1; + continue; + } + Some(Err(e)) => { + tracing::warn!( + "could not drop data tables for '{name}' ({table_id}); the table remains \ + DELETING and the control plane will retry: {e}" + ); + continue; + } } // Delete tags explicitly (not covered by CASCADE from tables). @@ -153,6 +194,75 @@ impl PostgresEngine { Err(e) => tracing::warn!("abandoned restore sweep failed: {e}"), } - Ok(transitions) + Ok(ControlPlanePass { + transitions, + deferred_drops, + }) + } + + /// Moves every `CREATING` table whose `status_transition_at` has passed + /// to `ACTIVE` and returns their names. + async fn activate_due_tables(&self) -> Result, StorageError> { + let activated: Vec<(String,)> = sqlx::query_as( + r"UPDATE tables SET table_status = 'ACTIVE', status_transition_at = NULL + WHERE table_status = 'CREATING' AND status_transition_at <= NOW() + RETURNING table_name", + ) + .fetch_all(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + Ok(activated.into_iter().map(|(name,)| name).collect()) + } + + /// Drops a `DELETING` table's data tables unless another session holds a + /// lock on one of them. `None` means the probe found one held and nothing + /// was tried; `Some` carries the drop's result. + async fn drop_unless_held( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + table_id: &str, + index_ids: &[String], + vector_index_ids: &[String], + ) -> Result>, StorageError> { + // A table snapshot, the in-catalog backup, a client read or write, + // an index build or `VACUUM` holds a lock on the table's relations + // for as long as it runs, and a `DROP` would wait for it. So the + // drop is tried only when no other session holds any lock on them + // (on the connection this pass holds when catalog and data share + // a pool, as in `DeleteTable`), and a held table is skipped at + // once. A pass therefore costs a probe per held table rather than + // a lock wait per held table, and a table whose holder has ended + // is dropped at the next pass, about `active_poll` after the + // release however many other tables are held. + let relations = crate::data::table_storage_relations(table_id, index_ids, vector_index_ids); + let held = if self.shared_data_pool { + crate::data::storage_in_use(&mut **tx, relations, LockHolders::Any).await + } else { + crate::data::storage_in_use(&self.data_pool, relations, LockHolders::Any).await + } + .map_err(|e| StorageError::Internal(e.to_string()))?; + // A table the probe found free is dropped under the same budget + // as DeleteTable's drops: a holder that took its lock after the + // probe, or a claim in flight on the table's queue rows, costs + // this pass at most `SYNC_DROP_LOCK_TIMEOUT`. See + // `drop_table_storage` for why the queue is stopped before + // anything is dropped. Any failure leaves the catalog row + // DELETING. A lock wait is counted, so the poller keeps retrying + // every `active_poll` until the holder is gone rather than waiting + // for a wake that only a snapshot stream or the backup sends; any + // other failure is not (`ControlPlanePass`). + if held { + return Ok(None); + } + Ok(Some( + Self::drop_table_storage( + &self.data_pool, + table_id, + index_ids, + vector_index_ids, + SYNC_DROP_LOCK_TIMEOUT, + ) + .await, + )) } } diff --git a/crates/storage-postgres/src/workers.rs b/crates/storage-postgres/src/workers.rs index 067679d39..7fcc33339 100644 --- a/crates/storage-postgres/src/workers.rs +++ b/crates/storage-postgres/src/workers.rs @@ -31,53 +31,114 @@ async fn read_index_propagation_delay( .await } +/// How often the control-plane poller retries while active. +pub(crate) const ACTIVE_POLL: Duration = Duration::from_secs(1); + +/// The control-plane poller `serve` runs (`spawn_workers`). +/// +/// Idle, it waits for a wake on `notify` (`CreateTable`, `DeleteTable`, a table +/// snapshot or in-catalog backup ending its read, startup recovery that left +/// a drop deferred) or for its 60 s idle timeout, and then runs an active +/// window of `control_plane_delay_seconds` + 5 s: a pass at once and one +/// every `active_poll` (one second in `serve`). The window lasts past its end +/// for as long as a `DELETING` row's drop is deferred because a session +/// holds a lock on the table's data tables +/// ([`crate::worker_store::ControlPlanePass::deferred_drops`]): only a +/// snapshot stream and the in-catalog backup wake this poller when they +/// release, and a plain transaction, `pg_dump` or a client read or write do +/// not, so going idle while a drop is deferred would leave the row +/// `DELETING` until the idle pass (up to 60 s after the holder is gone). +/// +/// A pass probes each due `DELETING` table for holders and skips a held one +/// without waiting on it, and drops a free one under +/// `delete_table::SYNC_DROP_LOCK_TIMEOUT` (1 s), so a pass costs a probe per +/// held table rather than a lock wait per held table. Staying active +/// therefore bounds the time from a holder's release to its table's drop at +/// about `active_poll` (plus the pass in progress, which is short), whatever +/// the holder and however many other tables are held. A `CreateTable` that is +/// due is activated at the start of the next pass on the same cadence. +/// +/// While active it waits on `notify` as well as on its timer, so a wake runs +/// the next pass at once rather than at the next poll, and restarts the +/// window (as a wake while idle starts one). A release that wakes the +/// poller therefore has the drop retried at once. pub(crate) async fn poll_control_plane_transitions( storage: Arc, notify: Arc, settings: Arc, token: CancellationToken, + active_poll: Duration, ) { - const ACTIVE_POLL: Duration = Duration::from_secs(1); const IDLE_TIMEOUT: Duration = Duration::from_secs(60); const MARGIN_SECS: f64 = 5.0; + let window = |delay_secs: f64| Duration::from_secs_f64(delay_secs + MARGIN_SECS); + // `Some(deadline)` while active. + let mut deadline: Option = None; loop { - // Idle: wait for a wake signal, an idle timeout (defensive sweep), or - // shutdown. - let shutting_down = tokio::select! { - () = token.cancelled() => true, - _ = tokio::time::timeout(IDLE_TIMEOUT, notify.notified()) => false, - }; - if shutting_down { - return; + if deadline.is_some() { + let woken = tokio::select! { + () = token.cancelled() => return, + () = tokio::time::sleep(active_poll) => false, + () = notify.notified() => true, + }; + if woken { + let delay_secs = read_control_plane_delay(&*settings).await; + deadline = Some(tokio::time::Instant::now() + window(delay_secs)); + } + } else { + // Idle: a wake, the idle timeout (defensive sweep), or shutdown. + tokio::select! { + () = token.cancelled() => return, + _ = tokio::time::timeout(IDLE_TIMEOUT, notify.notified()) => {} + } + let delay_secs = read_control_plane_delay(&*settings).await; + deadline = Some(tokio::time::Instant::now() + window(delay_secs)); } - // Read control_plane_delay_seconds from settings to compute active window - let delay_secs = read_control_plane_delay(&*settings).await; - let active_window = Duration::from_secs_f64(delay_secs + MARGIN_SECS); - - // Active: poll every second for active_window - let deadline = tokio::time::Instant::now() + active_window; - loop { - match storage.process_control_plane_transitions().await { - Ok(ref t) if t.is_empty() => {} - Ok(transitions) => { - for (name, transition) in &transitions { - tracing::info!("Table '{name}': {transition}"); - } + match storage.control_plane_pass().await { + Ok(pass) => { + for (name, transition) in &pass.transitions { + tracing::info!("Table '{name}': {transition}"); } - Err(e) => { - tracing::warn!("Control plane transition poll failed: {e}"); - break; + let past_window = deadline.is_some_and(|d| tokio::time::Instant::now() >= d); + if pass.deferred_drops == 0 && past_window { + deadline = None; } } - if tokio::time::Instant::now() >= deadline { - break; + Err(e) => { + tracing::warn!("Control plane transition poll failed: {e}"); + deadline = None; } - if !tick(&token, ACTIVE_POLL).await { - return; + } + } +} + +/// Startup recovery: one control-plane pass for the transitions a previous +/// server process left (a CREATING table past its transition time, a +/// DELETING row whose drop it never finished). When the pass had to defer a +/// drop because a session still holds the table (a reader that outlived the +/// restart: another server, `pg_dump`, a client transaction), the poller is +/// woken so it starts in its active loop and retries every `active_poll`, +/// rather than starting idle and leaving the row `DELETING` until its first +/// idle pass 60 s later. The wake is stored by the `Notify` until the poller +/// first waits. +pub(crate) async fn recover_control_plane_at_startup(engine: &PostgresEngine) { + match engine.control_plane_pass().await { + Ok(pass) => { + for (name, transition) in &pass.transitions { + tracing::info!("Recovered table '{name}': {transition}"); + } + if pass.deferred_drops > 0 { + tracing::info!( + deferred = pass.deferred_drops, + "DELETING table(s) still held by a session at startup; the control plane \ + retries the drop while active" + ); + engine.control_plane_notify.notify_one(); } } + Err(e) => tracing::error!("Failed to recover control plane transitions: {e}"), } } diff --git a/crates/storage-postgres/tests/gsi_online_build.rs b/crates/storage-postgres/tests/gsi_online_build.rs index 969ede4a8..7fdb52b81 100644 --- a/crates/storage-postgres/tests/gsi_online_build.rs +++ b/crates/storage-postgres/tests/gsi_online_build.rs @@ -829,15 +829,16 @@ async fn delete_table_waits_for_a_claimed_queue_row_then_drops_everything() { early.is_err(), "DeleteTable finished while a claim for the table was still held" ); - // With the claim held the queue is stopped and nothing has been dropped. - let held: i64 = sqlx::query_scalar( - "SELECT count(*) FROM vector_index_holds WHERE table_id = $1 AND index_id = $1", - ) - .bind(&table_id) - .fetch_one(&s.db) - .await - .expect("count the drop's hold"); - assert_eq!(held, 1, "the drop stopped the queue for the table first"); + // With the claim held nothing has been dropped and the catalog row is + // still there. On this layout (one database) the drop runs in + // DeleteTable's own transaction, its pending-row delete waiting on the + // claim, so nothing commits (no queue hold either) until the claim ends. + let rows: i64 = sqlx::query_scalar("SELECT count(*) FROM tables WHERE table_id = $1") + .bind(&table_id) + .fetch_one(&s.db) + .await + .expect("count the catalog row"); + assert_eq!(rows, 1, "the catalog row stays until the drop commits"); let base_exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") .bind(format!("\"_ddb_{table_id}\"")) .fetch_one(&s.db) diff --git a/crates/storage-postgres/tests/table_snapshot.rs b/crates/storage-postgres/tests/table_snapshot.rs new file mode 100644 index 000000000..adfad63b2 --- /dev/null +++ b/crates/storage-postgres/tests/table_snapshot.rs @@ -0,0 +1,3407 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! `SnapshotEngine` on `PostgreSQL`: the shared consistency check from +//! `extenddb_storage::snapshot_conformance` against real scratch databases, +//! plus the PostgreSQL-specific marker and key-order checks. +//! +//! Requires `EXTENDDB_TEST_PG_CONNECTION_STRING` (a base URL with no database +//! name); skips otherwise, and fails instead when `EXTENDDB_TEST_PG_REQUIRED` +//! is set. Each test creates its own catalog and data databases and drops +//! them when its `Scratch` is dropped, on success or panic (a killed test +//! binary leaves them behind). Most tests run on both layouts the engine +//! supports (`on_both_layouts!`): a separate data database, and one database +//! shared by catalog and data, where the open uses one transaction. + +use extenddb_storage::snapshot::paging::OpenStage; +use extenddb_storage::snapshot_conformance::{self, PageSize}; +use extenddb_storage::{SnapshotEngine, TableEngine}; +use extenddb_storage_postgres::{ + PostgresConfig, PostgresEngine, apply_all_migrations, data_table_name, index_table_name, + vector_table_name, +}; +use sqlx::postgres::PgPoolOptions; +use sqlx::{Connection, PgConnection, PgPool}; + +const ACCOUNT: &str = "123456789012"; +/// SQL backend: ordered keys, numbers to the service's 38 digits. With a +/// separate data database an open at `AfterDefinition` holds the `tables` row +/// `FOR SHARE`, so an `UpdateTable` issued then waits for it; on one shared +/// database the open holds no row lock and nothing waits. +const fn expectations(shared: bool) -> snapshot_conformance::Expectations { + snapshot_conformance::Expectations { + key_ordered: true, + numeric_key_digits: 38, + update_table_waits_at: if shared { + &[] + } else { + &[OpenStage::AfterDefinition] + }, + } +} +const REGION: &str = "us-east-1"; + +/// Runs `async fn $name(shared: bool)` as two tests, +/// `$name::separate_data_database` and `$name::one_shared_database`, each +/// skipped without `PostgreSQL`. +macro_rules! on_both_layouts { + ($name:ident) => { + mod $name { + #[tokio::test] + async fn separate_data_database() { + if super::base_conn().is_none() { + eprintln!("SKIP {}: no PostgreSQL", module_path!()); + return; + } + super::$name(false).await; + } + + #[tokio::test] + async fn one_shared_database() { + if super::base_conn().is_none() { + eprintln!("SKIP {}: no PostgreSQL", module_path!()); + return; + } + super::$name(true).await; + } + } + }; +} + +struct Scratch { + engine: PostgresEngine, + catalog: PgPool, + base: String, + names: [String; 2], + /// No separate data database: catalog and data in `names[0]`. + shared: bool, + /// Every test here opens a snapshot or touches the paging statics, so + /// each holds the snapshot hooks lock for its whole body. + hooks: snapshot_conformance::HooksGuard, +} + +impl Drop for Scratch { + /// Drops the two scratch databases on success or panic (not when the + /// test binary is killed). + /// Runs on its own thread and runtime with a fresh connection, because + /// the test runtime's pools cannot be awaited from `Drop`; `WITH (FORCE)` + /// terminates the sessions those pools still hold. + fn drop(&mut self) { + let base = self.base.clone(); + let names = self.names.clone(); + std::thread::scope(|scope| { + scope.spawn(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("cleanup runtime"); + rt.block_on(async move { + let Ok(mut conn) = PgConnection::connect(&format!("{base}/postgres")).await + else { + eprintln!("scratch cleanup: could not connect; {names:?} left behind"); + return; + }; + for name in &names { + if let Err(e) = + sqlx::query(&format!("DROP DATABASE IF EXISTS \"{name}\" WITH (FORCE)")) + .execute(&mut conn) + .await + { + eprintln!("scratch cleanup: drop {name}: {e}"); + } + } + let _ = conn.close().await; + }); + }); + }); + } +} + +impl Scratch { + /// The database the data tables live in. + fn data_url(&self) -> String { + let name = if self.shared { + &self.names[0] + } else { + &self.names[1] + }; + format!("{}/{name}", self.base) + } + + /// The data database's name, for `pg_stat_activity` filters. + fn data_name(&self) -> &str { + if self.shared { + &self.names[0] + } else { + &self.names[1] + } + } + + /// A pool on the scratch data database, for probes the engine does not + /// expose. + async fn data_pool(&self) -> PgPool { + PgPoolOptions::new() + .max_connections(1) + .connect(&self.data_url()) + .await + .expect("data pool") + } +} + +/// The server's base URL from `EXTENDDB_TEST_PG_CONNECTION_STRING`, or +/// `None` to skip. With `EXTENDDB_TEST_PG_REQUIRED` set (CI sets it beside +/// the connection string) a missing or empty URL is a failure, so a typo in +/// the variable's name cannot turn the step green by skipping every test. +fn base_conn() -> Option { + let conn = std::env::var("EXTENDDB_TEST_PG_CONNECTION_STRING").unwrap_or_default(); + if conn.trim().is_empty() { + assert!( + std::env::var_os("EXTENDDB_TEST_PG_REQUIRED").is_none(), + "EXTENDDB_TEST_PG_REQUIRED is set but EXTENDDB_TEST_PG_CONNECTION_STRING is not" + ); + return None; + } + Some(conn.trim().trim_end_matches('/').to_owned()) +} + +async fn scratch() -> Scratch { + scratch_with(false).await +} + +async fn scratch_with(shared: bool) -> Scratch { + scratch_opts(shared, false) + .await + .expect("no optional extension requested") +} + +/// `shared`: no separate data database; data migrations go on the catalog +/// database and the engine runs catalog and data on one pool. `pgvector`: +/// install the extension on the data database before the engine probes it; +/// `None` (after printing why) when the server does not offer it. +async fn scratch_opts(shared: bool, pgvector: bool) -> Option { + let hooks = snapshot_conformance::hooks_lock().await; + let base = base_conn().expect("caller checks base_conn() first"); + let stem = format!("eddb_snap_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let names = [format!("{stem}_c"), format!("{stem}_d")]; + if !create_databases(&base, &names, pgvector).await { + return None; + } + let catalog_url = format!("{base}/{}", names[0]); + let data_url = if shared { + catalog_url.clone() + } else { + format!("{base}/{}", names[1]) + }; + let catalog = PgPoolOptions::new() + .max_connections(2) + .connect(&catalog_url) + .await + .expect("catalog pool"); + migrate(&catalog, &data_url, pgvector).await; + configure_catalog(&catalog, (!shared).then_some(&data_url)).await; + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind(ACCOUNT) + .bind(format!("acct-{stem}")) + .execute(&catalog) + .await + .expect("account"); + let engine = PostgresEngine::new( + &PostgresConfig { + connection_string: catalog_url, + pool_size: 8, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("engine"); + Some(Scratch { + engine, + catalog, + base, + names, + shared, + hooks, + }) +} + +/// Create the two scratch databases; `false` (after printing why, with +/// nothing created) when `pgvector` is wanted and the server does not offer +/// it. +async fn create_databases(base: &str, names: &[String; 2], pgvector: bool) -> bool { + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to postgres"); + if pgvector { + let available: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_available_extensions WHERE name = 'vector')", + ) + .fetch_one(&admin) + .await + .expect("list the available extensions"); + if !available { + eprintln!("this PostgreSQL server does not offer the pgvector extension"); + admin.close().await; + return false; + } + } + for name in names { + sqlx::query(&format!("CREATE DATABASE \"{name}\"")) + .execute(&admin) + .await + .expect("create scratch database"); + } + admin.close().await; + true +} + +/// Apply every catalog and data migration the crate ships, as `extenddb +/// migrate --yes` does (the data ones on `data_url`, which is the catalog +/// database on the shared layout), and pgvector if asked. The lists live in +/// the crate (`migrations.rs`), so a migration added there is applied here +/// without this file knowing. +async fn migrate(catalog: &PgPool, data_url: &str, pgvector: bool) { + let data = PgPoolOptions::new() + .max_connections(2) + .connect(data_url) + .await + .expect("data pool"); + apply_all_migrations(catalog, &data) + .await + .expect("apply the crate's migrations"); + if pgvector { + sqlx::query("CREATE EXTENSION IF NOT EXISTS vector") + .execute(&data) + .await + .expect("install pgvector on the data database"); + } + data.close().await; +} + +/// Zero the control-plane and propagation delays, and point the engine at +/// `separate_data_url` when there is one. +async fn configure_catalog(catalog: &PgPool, separate_data_url: Option<&String>) { + for key in ["control_plane_delay_seconds", "index_propagation_delay_ms"] { + sqlx::query("UPDATE settings SET value = '0' WHERE key = $1") + .bind(key) + .execute(catalog) + .await + .expect("setting"); + } + if let Some(data_url) = separate_data_url { + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('data_database_connection_string', $1) \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .bind(data_url) + .execute(catalog) + .await + .expect("data connection setting"); + } +} + +async fn snapshot_is_consistent_under_concurrent_writes(shared: bool) { + let s = scratch_with(shared).await; + snapshot_conformance::run( + &s.hooks, + &s.engine, + ACCOUNT, + "snap_src", + 50, + expectations(shared), + || async { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + }, + ) + .await; + let catalog = s.catalog.clone(); + snapshot_conformance::run_refusals( + &s.hooks, + &s.engine, + ACCOUNT, + "snap_refused", + || async { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + }, + |table| async move { + sqlx::query("UPDATE tables SET table_status = 'UPDATING' WHERE table_name = $1") + .bind(table) + .execute(&catalog) + .await + .expect("flip status"); + }, + ) + .await; +} +on_both_layouts!(snapshot_is_consistent_under_concurrent_writes); + +async fn create_active(s: &Scratch, input: extenddb_core::types::CreateTableInput) { + let name = input.table_name.clone(); + s.engine.create_table(ACCOUNT, input).await.expect("create"); + snapshot_conformance::wait_active(&s.engine, ACCOUNT, &name, || async { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; +} + +async fn key_info(s: &Scratch, table: &str) -> extenddb_core::types::TableKeyInfo { + let desc = s + .engine + .describe_table( + ACCOUNT, + extenddb_core::types::DescribeTableInput { + table_name: table.to_owned(), + }, + ) + .await + .expect("describe"); + extenddb_core::types::TableKeyInfo { + table_name: table.to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: desc.table_id.clone(), + key_schema: desc.key_schema.clone(), + base_key_schema: desc.key_schema.clone(), + attribute_definitions: desc.attribute_definitions.clone(), + ..Default::default() + } +} + +/// A composite-key table with an `S` range key and no index. +fn string_range_input(name: &str) -> extenddb_core::types::CreateTableInput { + serde_json::from_value(serde_json::json!({ + "TableName": name, + "KeySchema": [ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ], + "AttributeDefinitions": [ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"} + ], + "BillingMode": "PAY_PER_REQUEST" + })) + .expect("create input") +} + +/// Strings `en_US` and `C` order differently. +const COLLATION_KEYS: [&str; 10] = ["b", "B", "_x", "a", "ab", "Z", "-", "10", "9", "\u{e9}"]; + +/// Put one item per key in `COLLATION_KEYS`, as the hash key (`pk`, with `sk` +/// `N` 1) or as the range key (`sk` `S`, under one `pk`). +async fn put_collation_keys(s: &Scratch, table: &str, as_range: bool) { + use extenddb_core::expression::ExpressionMaps; + use extenddb_core::types::{AttributeValue, Item}; + use extenddb_storage::DataEngine; + let ki = key_info(s, table).await; + for key in COLLATION_KEYS { + let item: Item = if as_range { + [ + ("pk".to_owned(), AttributeValue::S("p".to_owned())), + ("sk".to_owned(), AttributeValue::S(key.to_owned())), + ] + } else { + [ + ("pk".to_owned(), AttributeValue::S(key.to_owned())), + ("sk".to_owned(), AttributeValue::N("1".to_owned())), + ] + } + .into_iter() + .collect(); + s.engine + .put_item(&ki, item, false, None, &ExpressionMaps::default(), None) + .await + .expect("put"); + } +} + +/// The key text the stream yielded, in stream order: `pk`, or `sk` for a +/// range-key table. +async fn stream_keys( + snapshot: extenddb_storage::TableSnapshot, + attribute: &'static str, +) -> Vec { + use extenddb_core::types::AttributeValue; + use futures::StreamExt; + snapshot + .items + .map(move |r| match r.expect("item").get(attribute) { + Some(AttributeValue::S(v)) => v.clone(), + _ => panic!("{attribute} is not a string"), + }) + .collect() + .await +} + +/// A table created before the key columns were declared `COLLATE "C"` keeps +/// the database collation. Ordering it bytewise would be a full sort before +/// the first row, so it streams unordered and complete, whichever text key +/// column (`pk`, or `sk_s` for an `S` range key) has the database collation. +/// A table with "C" key columns on the same database is still ordered, +/// bytewise. Keys chosen so `en_US` and `C` disagree. +async fn legacy_collation_table_streams_unordered_and_complete(shared: bool) { + let s = scratch_with(shared).await; + let data = s.data_pool().await; + // On a server whose default collation is C or POSIX the recollated + // columns below would still sort bytewise and the test would show + // nothing. + let datcollate: String = sqlx::query_scalar( + "SELECT datcollate::text FROM pg_database WHERE datname = current_database()", + ) + .fetch_one(&data) + .await + .expect("datcollate"); + assert!( + !matches!(datcollate.as_str(), "C" | "POSIX"), + "the scratch data database must have a non-C default collation, got {datcollate}" + ); + + create_active(&s, snapshot_conformance::table_input("legacy_pk")).await; + create_active(&s, string_range_input("legacy_sk")).await; + create_active(&s, string_range_input("current")).await; + put_collation_keys(&s, "legacy_pk", false).await; + put_collation_keys(&s, "legacy_sk", true).await; + put_collation_keys(&s, "current", true).await; + // What a pre-#361 table looks like: a key column on the database + // collation. One table per column the check reads. + for (table, column) in [("legacy_pk", "pk"), ("legacy_sk", "sk_s")] { + let relation = data_table_name(&key_info(&s, table).await.table_id); + sqlx::query(&format!( + "ALTER TABLE {relation} ALTER COLUMN {column} TYPE TEXT COLLATE \"default\"" + )) + .execute(&data) + .await + .expect("recollate"); + } + data.close().await; + + let mut expected: Vec = COLLATION_KEYS.iter().map(|k| (*k).to_owned()).collect(); + expected.sort(); + for (table, attribute) in [("legacy_pk", "pk"), ("legacy_sk", "sk")] { + let snapshot = s + .engine + .open_table_snapshot(ACCOUNT, table) + .await + .expect("open"); + assert!( + !snapshot.key_ordered, + "{table}: a key column without collation \"C\" streams unordered" + ); + let mut streamed = stream_keys(snapshot, attribute).await; + streamed.sort(); + assert_eq!(streamed, expected, "{table}: every item streams"); + } + + let snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "current") + .await + .expect("open"); + assert!(snapshot.key_ordered, "\"C\" key columns stream ordered"); + assert_eq!( + stream_keys(snapshot, "sk").await, + expected, + "stream is not in bytewise order" + ); +} +on_both_layouts!(legacy_collation_table_streams_unordered_and_complete); + +/// The snapshot bound is half the pool, with a separate data database or one +/// shared pool: five snapshots holding their data connections on a +/// ten-connection pool, the sixth open refused at once with `LimitExceeded`, +/// item reads and `DescribeTable` still served, and a dropped snapshot's +/// permit free again. +async fn half_the_pool_holds_snapshots(shared: bool) { + use extenddb_core::types::{AttributeValue, Item}; + use extenddb_storage::DataEngine; + use futures::StreamExt; + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("many")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "many", 20).await; + let ki = key_info(&s, "many").await; + // Each stream reads one item of a five-row page, so its transaction and + // data connection stay open. + let _page = PageSize::set(5); + // The engine's pool is MIN_POOL_SIZE (10) here. + let mut held = Vec::new(); + for i in 0..5 { + let mut snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "many") + .await + .unwrap_or_else(|e| panic!("open {i}: {e}")); + snapshot.items.next().await.expect("an item").expect("ok"); + held.push(snapshot); + } + let err = tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.open_table_snapshot(ACCOUNT, "many"), + ) + .await + .expect("an open over the bound is refused at once, not queued") + .expect_err("the sixth open is over the bound"); + assert!( + matches!(err, extenddb_storage::error::StorageError::LimitExceeded(_)), + "{err}" + ); + let key: Item = [ + ("pk".to_owned(), AttributeValue::S("p000".to_owned())), + ("sk".to_owned(), AttributeValue::N("0".to_owned())), + ] + .into_iter() + .collect(); + tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.get_item(&ki, &key), + ) + .await + .expect("GetItem is not starved by the open snapshots") + .expect("get_item") + .expect("the item exists"); + tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.describe_table( + ACCOUNT, + extenddb_core::types::DescribeTableInput { + table_name: "many".to_owned(), + }, + ), + ) + .await + .expect("DescribeTable is not starved by the open snapshots") + .expect("describe"); + drop(held.pop()); + s.engine + .open_table_snapshot(ACCOUNT, "many") + .await + .expect("a dropped snapshot frees its permit"); +} + +#[tokio::test] +async fn five_snapshots_on_a_separate_data_pool_leave_half_of_it() { + if base_conn().is_none() { + eprintln!("SKIP five_snapshots_on_a_separate_data_pool_leave_half_of_it: no PostgreSQL"); + return; + } + half_the_pool_holds_snapshots(false).await; +} + +#[tokio::test] +async fn five_snapshots_on_one_shared_pool_leave_half_of_it() { + if base_conn().is_none() { + eprintln!("SKIP five_snapshots_on_one_shared_pool_leave_half_of_it: no PostgreSQL"); + return; + } + half_the_pool_holds_snapshots(true).await; +} + +/// Hold `ACCESS SHARE` on one data relation from another session, run +/// `DeleteTable`, and check it took the asynchronous path (it returns at once +/// and the row is `DELETING`, not gone); then release the lock and check the +/// worker drops `relations` and removes the row. +async fn delete_with_relation_locked(s: &Scratch, table: &str, locked: &str, relations: &[String]) { + let mut holder = PgConnection::connect(&s.data_url()).await.expect("holder"); + sqlx::query("BEGIN") + .execute(&mut holder) + .await + .expect("begin"); + sqlx::query(&format!("LOCK TABLE {locked} IN ACCESS SHARE MODE")) + .execute(&mut holder) + .await + .expect("lock the relation"); + let deleted = tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.delete_table( + ACCOUNT, + extenddb_core::types::DeleteTableInput { + table_name: table.to_owned(), + }, + ), + ) + .await + .expect("DeleteTable must not wait behind a reader of an index data table") + .expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + let status: Option = + sqlx::query_scalar("SELECT table_status FROM tables WHERE table_name = $1") + .bind(table) + .fetch_optional(&s.catalog) + .await + .expect("status"); + assert_eq!( + status.as_deref(), + Some("DELETING"), + "with {locked} in use the row stays until the worker's drop" + ); + sqlx::query("ROLLBACK") + .execute(&mut holder) + .await + .expect("release"); + holder.close().await.expect("close holder"); + + let data = s.data_pool().await; + for _ in 0..100 { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + let mut all_gone = true; + for relation in relations { + let gone: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NULL") + .bind(relation) + .fetch_one(&data) + .await + .expect("probe"); + all_gone &= gone; + } + let row: Option = + sqlx::query_scalar("SELECT table_status FROM tables WHERE table_name = $1") + .bind(table) + .fetch_optional(&s.catalog) + .await + .expect("status"); + if all_gone && row.is_none() { + data.close().await; + return; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + panic!("{relations:?} and the catalog row of {table} were not all removed"); +} + +/// `DeleteTable` probes every relation its synchronous drop touches, not only +/// the base data table: a session holding `ACCESS SHARE` on a GSI data table +/// alone (a Query or Scan with `IndexName`, the backfill) sends it down the +/// asynchronous path at once, rather than into a drop that would wait out its +/// lock timeout behind the reader. +async fn delete_table_with_a_gsi_data_table_in_use_takes_the_async_path(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("gsilock")).await; + let table_id = key_info(&s, "gsilock").await.table_id; + let index_id: String = sqlx::query_scalar( + "SELECT i.index_id FROM indexes i JOIN tables t ON t.table_id = i.table_id \ + WHERE t.table_name = 'gsilock' AND i.index_name = 'gi'", + ) + .fetch_one(&s.catalog) + .await + .expect("gi row"); + let gsi_table = index_table_name(&index_id); + delete_with_relation_locked( + &s, + "gsilock", + &gsi_table, + &[data_table_name(&table_id), gsi_table.clone()], + ) + .await; +} +on_both_layouts!(delete_table_with_a_gsi_data_table_in_use_takes_the_async_path); + +/// The same for a vector-index data table, on a server with pgvector. +async fn delete_table_with_a_vector_data_table_in_use_takes_the_async_path(shared: bool) { + let Some(s) = scratch_opts(shared, true).await else { + eprintln!("SKIP: no pgvector on this server, so no vector-index data table"); + return; + }; + assert!( + s.engine.vector_capable(), + "pgvector installed on the data database" + ); + let input: extenddb_core::types::CreateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": "veclock", + "KeySchema": [{"AttributeName": "pk", "KeyType": "HASH"}], + "AttributeDefinitions": [{"AttributeName": "pk", "AttributeType": "S"}], + "BillingMode": "PAY_PER_REQUEST", + "VectorIndexes": [{ + "IndexName": "vi", + "VectorAttribute": {"AttributeName": "emb"}, + "Dimensions": 3, + "DistanceFunction": "COSINE", + "Projection": {"ProjectionType": "ALL"} + }] + })) + .expect("create input"); + create_active(&s, input).await; + let table_id = key_info(&s, "veclock").await.table_id; + let index_id: String = sqlx::query_scalar( + "SELECT v.index_id FROM vector_indexes v JOIN tables t ON t.table_id = v.table_id \ + WHERE t.table_name = 'veclock'", + ) + .fetch_one(&s.catalog) + .await + .expect("vector index row"); + let vector_table = vector_table_name(&index_id); + let data = s.data_pool().await; + let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(&vector_table) + .fetch_one(&data) + .await + .expect("probe"); + data.close().await; + assert!(exists, "the vector index has a data table to lock"); + delete_with_relation_locked( + &s, + "veclock", + &vector_table, + &[data_table_name(&table_id), vector_table.clone()], + ) + .await; +} +on_both_layouts!(delete_table_with_a_vector_data_table_in_use_takes_the_async_path); + +/// A row the stream cannot parse ends it at the error, as on every backend, +/// and gives up the snapshot's permit at the error rather than at the next +/// poll or the drop: with the errored stream still held, a full bound of five +/// opens succeeds. +async fn an_unparseable_row_ends_the_stream_and_frees_its_permit(shared: bool) { + use futures::StreamExt as _; + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("corrupt")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "corrupt", 20).await; + let table_id = key_info(&s, "corrupt").await.table_id; + let data = s.data_pool().await; + // Valid JSON that is not an item, keyed between p005 and p006: the + // seventh row, on the second (full) page of five, so the error is not + // also the end of the table. + sqlx::query(&format!( + "INSERT INTO {} (pk, sk_n, item_data) VALUES ('p005a', 1, '{{\"pk\": 5}}')", + data_table_name(&table_id) + )) + .execute(&data) + .await + .expect("insert an unparseable row"); + data.close().await; + let _page = PageSize::set(5); + let mut snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "corrupt") + .await + .expect("open"); + let mut items = 0; + loop { + match snapshot.items.next().await { + Some(Ok(_)) => items += 1, + Some(Err(_)) => break, + None => panic!("the stream ended without reporting the unparseable row"), + } + } + assert_eq!( + items, 5, + "the first page, then the error for the page holding the unparseable row" + ); + // Before the stream is polled again (its end drops the state, and the + // permit with it), the permit must already be free. + let mut held = Vec::new(); + for i in 0..5 { + held.push( + s.engine + .open_table_snapshot(ACCOUNT, "corrupt") + .await + .unwrap_or_else(|e| panic!("open {i} with the errored stream held: {e}")), + ); + } + assert!( + snapshot.items.next().await.is_none(), + "nothing follows the first error" + ); +} +on_both_layouts!(an_unparseable_row_ends_the_stream_and_frees_its_permit); + +/// A GSI row whose data table is missing is left out of the description, +/// and so are the attribute definitions only it used: `CreateTable` would +/// refuse a definition no key names. +async fn gsi_without_a_data_table_is_left_out_with_its_attributes(shared: bool) { + let s = scratch_with(shared).await; + s.engine + .create_table(ACCOUNT, snapshot_conformance::table_input("half")) + .await + .expect("create"); + snapshot_conformance::wait_active(&s.engine, ACCOUNT, "half", || async { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; + s.engine + .update_table( + ACCOUNT, + serde_json::from_value(serde_json::json!({ + "TableName": "half", + "AttributeDefinitions": [{"AttributeName": "h", "AttributeType": "S"}], + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "gi2", + "KeySchema": [{"AttributeName": "h", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "KEYS_ONLY"}, + "ProvisionedThroughput": {"ReadCapacityUnits": 1, "WriteCapacityUnits": 1} + }}] + })) + .expect("update input"), + ) + .await + .expect("add gi2"); + for _ in 0..50 { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + } + let index_id: String = sqlx::query_scalar( + "SELECT i.index_id FROM indexes i JOIN tables t ON t.table_id = i.table_id \ + WHERE t.table_name = 'half' AND i.index_name = 'gi2'", + ) + .fetch_one(&s.catalog) + .await + .expect("gi2 row"); + let data = s.data_pool().await; + // An index the snapshot cannot see whole: drop its data table. + sqlx::query(&format!("DROP TABLE {}", index_table_name(&index_id))) + .execute(&data) + .await + .expect("drop the GSI data table"); + data.close().await; + + let snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "half") + .await + .expect("open"); + let gsis: Vec<_> = snapshot + .description + .global_secondary_indexes + .iter() + .flatten() + .map(|g| g.index_name.clone()) + .collect(); + assert!(!gsis.contains(&"gi2".to_owned()), "{gsis:?}"); + let attrs: Vec<_> = snapshot + .description + .attribute_definitions + .iter() + .map(|d| d.attribute_name.clone()) + .collect(); + assert!( + !attrs.contains(&"h".to_owned()), + "unused definition left in: {attrs:?}" + ); + drop(snapshot); +} +on_both_layouts!(gsi_without_a_data_table_is_left_out_with_its_attributes); + +/// `DeleteTable` while a snapshot stream holds the data table never leaves +/// the catalog row gone with the data table orphaned: it takes the +/// asynchronous path, returns at once, and the worker finishes once the +/// stream is done, even if the caller gave up. +async fn delete_table_during_a_snapshot_does_not_orphan_the_data_table(shared: bool) { + use futures::StreamExt as _; + let s = scratch_with(shared).await; + s.engine + .create_table(ACCOUNT, snapshot_conformance::table_input("victim")) + .await + .expect("create"); + snapshot_conformance::wait_active(&s.engine, ACCOUNT, "victim", || async { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "victim", 20).await; + let table_id: String = + sqlx::query_scalar("SELECT table_id FROM tables WHERE table_name = 'victim'") + .fetch_one(&s.catalog) + .await + .expect("table id"); + let page = PageSize::set(5); + let mut snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "victim") + .await + .expect("open"); + // Read one page so the cursor is mid-table and the lock is held. + snapshot.items.next().await.expect("an item").expect("ok"); + drop(page); + + let deleted = tokio::time::timeout( + std::time::Duration::from_secs(3), + s.engine.delete_table( + ACCOUNT, + extenddb_core::types::DeleteTableInput { + table_name: "victim".to_owned(), + }, + ), + ) + .await + .expect("DeleteTable must not wait behind the stream") + .expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + let status: Option = + sqlx::query_scalar("SELECT table_status FROM tables WHERE table_name = 'victim'") + .fetch_optional(&s.catalog) + .await + .expect("status"); + assert_eq!( + status.as_deref(), + Some("DELETING"), + "the row stays until the drop" + ); + + // The stream still reads its snapshot to the end. + while let Some(item) = snapshot.items.next().await { + item.expect("stream item"); + } + drop(snapshot); + let data = s.data_pool().await; + for _ in 0..100 { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + let gone: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NULL") + .bind(data_table_name(&table_id)) + .fetch_one(&data) + .await + .expect("probe"); + let row: Option = + sqlx::query_scalar("SELECT table_status FROM tables WHERE table_name = 'victim'") + .fetch_optional(&s.catalog) + .await + .expect("status"); + if gone && row.is_none() { + data.close().await; + return; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + panic!("the data table and the catalog row were not both removed"); +} +on_both_layouts!(delete_table_during_a_snapshot_does_not_orphan_the_data_table); + +/// The two SQLSTATEs that mean the server discarded the snapshot map to +/// `SnapshotExpired`; a third does not. The errors are raised by the server +/// so the test exercises the real `sqlx::Error::Database` shape. +#[tokio::test] +async fn sqlstate_mapping_is_pinned() { + if base_conn().is_none() { + eprintln!("SKIP sqlstate_mapping_is_pinned: no PostgreSQL"); + return; + } + let s = scratch().await; + for (code, expired) in [("40001", true), ("72000", true), ("57014", false)] { + let err = sqlx::query(&format!( + "DO $$ BEGIN RAISE EXCEPTION 'probe' USING ERRCODE = '{code}'; END $$" + )) + .execute(&s.catalog) + .await + .expect_err("raised"); + let mapped = extenddb_storage_postgres::map_fetch_error(&err); + assert_eq!( + matches!( + mapped, + extenddb_storage::error::StorageError::SnapshotExpired(_) + ), + expired, + "{code}: {mapped:?}" + ); + } +} + +async fn snapshot_marker_and_key_order(shared: bool) { + let s = scratch_with(shared).await; + s.engine + .create_table(ACCOUNT, snapshot_conformance::table_input("ordered")) + .await + .expect("create"); + snapshot_conformance::wait_active(&s.engine, ACCOUNT, "ordered", || async { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; + let snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "ordered") + .await + .expect("open"); + assert!(snapshot.key_ordered); + assert_eq!( + snapshot.info.kind, + extenddb_storage::SnapshotInfo::KIND_POSTGRES_REPEATABLE_READ + ); + // pg_current_snapshot() text: "xmin:xmax:xip_list". + assert!( + snapshot.info.marker.split(':').count() == 3, + "marker {:?}", + snapshot.info.marker + ); + drop(snapshot); + + // A table that is not ACTIVE is refused before any data read. + sqlx::query("UPDATE tables SET table_status = 'UPDATING' WHERE table_name = 'ordered'") + .execute(&s.catalog) + .await + .expect("flip status"); + let err = s + .engine + .open_table_snapshot(ACCOUNT, "ordered") + .await + .expect_err("refused"); + assert!( + matches!( + err, + extenddb_storage::error::StorageError::TableNotActive(_) + ), + "{err}" + ); +} +on_both_layouts!(snapshot_marker_and_key_order); + +/// A hash-only table with no index, for tests that need many of them. +fn plain_input(name: &str) -> extenddb_core::types::CreateTableInput { + serde_json::from_value(serde_json::json!({ + "TableName": name, + "KeySchema": [{"AttributeName": "pk", "KeyType": "HASH"}], + "AttributeDefinitions": [{"AttributeName": "pk", "AttributeType": "S"}], + "BillingMode": "PAY_PER_REQUEST" + })) + .expect("create input") +} + +fn delete_input(table: &str) -> extenddb_core::types::DeleteTableInput { + extenddb_core::types::DeleteTableInput { + table_name: table.to_owned(), + } +} + +/// `DeleteTable` under a 5 s timeout, so a delete that waits behind a reader +/// or the pool fails the test instead of hanging it. +async fn delete_within_5s( + s: &Scratch, + table: &str, +) -> Result { + tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.delete_table(ACCOUNT, delete_input(table)), + ) + .await + .unwrap_or_else(|_| panic!("DeleteTable of {table} did not return within 5 s")) +} + +async fn catalog_status(s: &Scratch, table: &str) -> Option { + sqlx::query_scalar("SELECT table_status FROM tables WHERE table_name = $1") + .bind(table) + .fetch_optional(&s.catalog) + .await + .expect("status") +} + +/// Twenty concurrent `DeleteTable`s of twenty tables on a ten-connection +/// engine all succeed within 5 s. Each holds one catalog connection with its +/// row `FOR UPDATE` while it probes for readers; on one shared pool the probe +/// must run on that connection, or `pool_size` deletes hold every +/// connection, each waiting for a second, until the acquire timeout. +async fn concurrent_delete_tables_do_not_exhaust_the_pool(shared: bool) { + let s = scratch_with(shared).await; + let tables: Vec = (0..20).map(|i| format!("many{i:02}")).collect(); + for table in &tables { + create_active(&s, plain_input(table)).await; + } + let started = std::time::Instant::now(); + let results = tokio::time::timeout( + std::time::Duration::from_secs(5), + futures::future::join_all( + tables + .iter() + .map(|t| s.engine.delete_table(ACCOUNT, delete_input(t))), + ), + ) + .await + .expect("twenty concurrent DeleteTables finish within 5 s"); + eprintln!("twenty DeleteTables took {:?}", started.elapsed()); + for (table, result) in tables.iter().zip(results) { + result.unwrap_or_else(|e| panic!("DeleteTable {table}: {e}")); + } +} +on_both_layouts!(concurrent_delete_tables_do_not_exhaust_the_pool); + +/// An open held at `AfterDefinition` while nine `DeleteTable`s of the same +/// table are issued, then resumed: the open and every delete return, none +/// waiting out the pool's acquire timeout. On the separate layout the +/// deletes wait on the row the open holds `FOR SHARE` with catalog +/// connections while the open needs only a data connection, and the open +/// then streams every item; the first delete finds the table in use and +/// leaves it `DELETING`, and the other eight find it `DELETING` and return +/// its description. On the shared layout the open holds no row lock and, at +/// this pause, no lock on the data table yet, so the first delete drops the +/// table at once, the other eight find it gone, and the open reports +/// `TableNotFound`. +async fn an_open_held_with_deletes_queued_completes(shared: bool) { + use futures::StreamExt as _; + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("queued")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "queued", 20).await; + let mut pause = extenddb_storage::snapshot::paging::arm_open_pause( + extenddb_storage::snapshot::paging::OpenStage::AfterDefinition, + ); + let reached = pause.take_reached(); + let open = s.engine.open_table_snapshot(ACCOUNT, "queued"); + let controller = async { + reached.await.expect("the open reached the pause"); + let deletes = futures::future::join_all( + (0..9).map(|_| s.engine.delete_table(ACCOUNT, delete_input("queued"))), + ); + let resume = async { + tokio::time::sleep(std::time::Duration::from_secs(1)).await; + pause.resume(); + }; + let (results, ()) = tokio::join!(deletes, resume); + results + }; + let started = std::time::Instant::now(); + let (opened, deletes) = tokio::time::timeout(std::time::Duration::from_secs(15), async { + tokio::join!(open, controller) + }) + .await + .expect("the open and the deletes finish well inside the 30 s acquire timeout"); + eprintln!("open and nine deletes took {:?}", started.elapsed()); + let mut returned = 0; + for result in deletes { + match result { + Ok(d) => { + assert_eq!(d.table_status, extenddb_core::types::TableStatus::Deleting); + returned += 1; + } + Err(extenddb_storage::error::StorageError::TableNotFound(_)) => {} + Err(e) => panic!("a queued DeleteTable failed: {e}"), + } + } + let (want, status) = if shared { + (1, None) + } else { + (9, Some("DELETING")) + }; + assert_eq!( + returned, want, + "DeleteTables that returned the table's description" + ); + assert_eq!(catalog_status(&s, "queued").await.as_deref(), status); + if shared { + let err = opened.expect_err("the table was dropped before the open's LOCK"); + assert!( + matches!(err, extenddb_storage::error::StorageError::TableNotFound(_)), + "{err}" + ); + } else { + let snapshot = opened.expect("the held open completes"); + let items: Vec<_> = snapshot.items.collect().await; + assert_eq!( + items.len(), + 20, + "the open streams the table as of its snapshot" + ); + assert!(items.iter().all(Result::is_ok)); + } +} +on_both_layouts!(an_open_held_with_deletes_queued_completes); + +/// A restored table's snapshot description carries the same +/// `restore_summary` as `describe_table`'s (both read the `table_restores` +/// row). The record is written directly, as `RestoreTableFromBackup` leaves +/// it. +async fn a_restored_table_describes_its_restore(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("restored")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "restored", 5).await; + sqlx::query( + "INSERT INTO table_restores \ + (table_id, source_backup_arn, source_table_arn, restore_date_time) \ + SELECT table_id, $1, $2, '2026-10-01T12:00:00.250Z' FROM tables \ + WHERE table_name = 'restored'", + ) + .bind("arn:aws:dynamodb:us-east-1:123456789012:table/src/backup/01759320000250-0a1b2c3d") + .bind("arn:aws:dynamodb:us-east-1:123456789012:table/src") + .execute(&s.catalog) + .await + .expect("restore record"); + snapshot_conformance::assert_snapshot_describes_the_table( + &s.hooks, &s.engine, ACCOUNT, "restored", true, + ) + .await; +} +on_both_layouts!(a_restored_table_describes_its_restore); + +/// The pid of the one backend whose last statement was the stream's `FETCH`. +async fn stream_backend(s: &Scratch, data: &PgPool) -> i32 { + sqlx::query_scalar( + "SELECT pid FROM pg_stat_activity WHERE datname = $1 \ + AND state = 'idle in transaction' AND query LIKE 'FETCH%'", + ) + .bind(s.data_name()) + .fetch_one(data) + .await + .expect("the stream's backend is idle in its transaction after a FETCH") +} + +/// End the stream's backend from another session, as a server restart or +/// an administrator would. +async fn terminate_stream_backend(s: &Scratch) { + let data = s.data_pool().await; + let pid = stream_backend(s, &data).await; + let terminated: bool = sqlx::query_scalar("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(&data) + .await + .expect("terminate"); + assert!(terminated, "pg_terminate_backend({pid})"); + data.close().await; +} + +/// How a snapshot's data transaction ends after a `DeleteTable` found the +/// table in use. +#[derive(Debug, Clone, Copy)] +enum Release { + /// The stream is dropped with rows left. + Dropped, + /// The stream is read to its end and kept. + Drained, + /// The stream's backend is terminated and the stream yields the error. + FetchError, + /// The open is cancelled while held at `BeforeItems`. + OpenCancelled, +} + +/// A held snapshot or an open in progress, whichever `Release` needs. +enum Holder<'a> { + Stream(Box), + Open( + futures::future::BoxFuture< + 'a, + Result, + >, + extenddb_storage::snapshot::paging::OpenPause, + ), +} + +/// The active poll of the poller in the wake tests: long enough that a drop +/// retried by the timer instead of a wake comes well after the tests' 3 s +/// bound. +const WAKE_TEST_ACTIVE_POLL: std::time::Duration = std::time::Duration::from_secs(10); + +/// Wait (up to 15 s) until `engine`'s control plane has made more than +/// `after` drop attempts, and return the count. +async fn wait_for_drop_attempt(engine: &PostgresEngine, after: u64) -> u64 { + for _ in 0..150 { + let attempts = engine.control_plane_drop_attempts_for_tests(); + if attempts > after { + return attempts; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + panic!("the control plane made no drop attempt after the first {after} within 15 s"); +} + +/// Wait (up to 3 s from `released_at`) until `relations` are gone and +/// `table` has no catalog row; `what` names the release in the output. +async fn removed_within_3s( + s: &Scratch, + table: &str, + relations: &[String], + released_at: std::time::Instant, + what: &str, +) { + let data = s.data_pool().await; + while released_at.elapsed() < std::time::Duration::from_secs(3) { + let mut all_gone = true; + for relation in relations { + let gone: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NULL") + .bind(relation) + .fetch_one(&data) + .await + .expect("probe"); + all_gone &= gone; + } + if all_gone && catalog_status(s, table).await.is_none() { + eprintln!( + "{what}: table removed {:?} after the release", + released_at.elapsed() + ); + data.close().await; + return; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + panic!("{what}: {table} was not removed within 3 s of the release"); +} + +/// `DeleteTable` while a snapshot holds the table leaves it `DELETING`. The +/// control-plane poller (the one `serve` runs, with its active poll set to +/// `WAKE_TEST_ACTIVE_POLL`) makes its first drop attempt, which finds the +/// table held and skips it, and pauses before the next. The snapshot's +/// transaction then ends, by `how`, during that pause: the wake the release +/// sends has the poller retry at once, and the table is gone within 3 s. +/// Without the wake the next attempt comes about 9.5 s after the release. +/// `Dropped`, `FetchError` and `OpenCancelled` end the transaction through +/// `HeldSnapshot`'s drop, `Drained` through its `finish`. +/// `a_plain_reader_past_the_active_window_delays_the_drop_only_until_it_releases` +/// pins the poller's behaviour for a reader that sends no wake. +async fn release_wakes_the_active_poller(shared: bool, how: Release) { + use futures::StreamExt as _; + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("woken")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "woken", 20).await; + let table_id = key_info(&s, "woken").await.table_id; + let _page = PageSize::set(5); + let holder = if let Release::OpenCancelled = how { + let mut pause = extenddb_storage::snapshot::paging::arm_open_pause( + extenddb_storage::snapshot::paging::OpenStage::BeforeItems, + ); + let reached = pause.take_reached(); + let mut open = s.engine.open_table_snapshot(ACCOUNT, "woken"); + tokio::select! { + r = &mut open => panic!("the open finished while it should be paused: {r:?}"), + r = reached => r.expect("the open reached BeforeItems"), + } + Holder::Open(open, pause) + } else { + let mut snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "woken") + .await + .expect("open"); + snapshot.items.next().await.expect("an item").expect("ok"); + Holder::Stream(Box::new(snapshot)) + }; + + let deleted = delete_within_5s(&s, "woken").await.expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + // Started after DeleteTable, so the wakes CreateTable and DeleteTable + // left are one stored permit: the poller starts its active window at + // once, and nothing else wakes it before the release. + let _poller = s + .engine + .spawn_control_plane_poller_with_active_poll_for_tests(WAKE_TEST_ACTIVE_POLL); + let attempts = wait_for_drop_attempt(&s.engine, 0).await; + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + assert_eq!( + s.engine.control_plane_drop_attempts_for_tests(), + attempts, + "{how:?}: precondition: the poller is in the pause after its first attempt" + ); + assert_eq!( + catalog_status(&s, "woken").await.as_deref(), + Some("DELETING"), + "{how:?}: the poller cannot drop the table while the snapshot holds it" + ); + + // Kept alive until the end, so only the release under test can wake the + // poller (a dropped stream would wake it through its drop). + let mut kept = None; + match (how, holder) { + (Release::Dropped, Holder::Stream(snapshot)) => drop(snapshot), + (Release::Drained, Holder::Stream(mut snapshot)) => { + while let Some(item) = snapshot.items.next().await { + item.expect("stream item"); + } + kept = Some(snapshot); + } + (Release::FetchError, Holder::Stream(mut snapshot)) => { + terminate_stream_backend(&s).await; + loop { + match snapshot.items.next().await { + Some(Ok(_)) => {} + Some(Err(_)) => break, + None => panic!("the stream ended without the FETCH error"), + } + } + kept = Some(snapshot); + } + (Release::OpenCancelled, Holder::Open(open, pause)) => { + drop(open); + drop(pause); + } + (how, _) => unreachable!("holder built for {how:?}"), + } + let released_at = std::time::Instant::now(); + removed_within_3s( + &s, + "woken", + &[data_table_name(&table_id)], + released_at, + &format!("{how:?}"), + ) + .await; + drop(kept); +} + +/// A reader that is not a snapshot stream (a plain transaction holding +/// `ACCESS SHARE`, as `pg_dump` or a client read do) sends the poller no +/// wake when it releases. The poller +/// therefore stays active for as long as a `DELETING` row's drop is deferred +/// because the table is held: with the reader holding the table for 10 s, +/// well past the poller's active window of 5 s (`control_plane_delay_seconds` +/// is 0; each attempt in the window finds the table held and skips it at +/// once, so the window's attempts end with it), the row is `DELETING` +/// throughout and the table is gone within 3 s of the release. A poller that went idle at +/// the end of its window would drop it at its 60 s idle pass instead, and +/// this test's 3 s bound fails it. +/// +/// The poller is started after the `DeleteTable`, so the wakes `CreateTable` +/// and `DeleteTable` left are one stored permit and one window. Started +/// before, the `DeleteTable` wake would arrive during the first window and be +/// stored for a second one right after it, and a poller that went idle would +/// still be active at about 10 s from two chained windows; the hold would +/// have to outlast both to tell the two apart. +async fn a_plain_reader_past_the_active_window_delays_the_drop_only_until_it_releases( + shared: bool, +) { + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("read")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "read", 5).await; + let relation = data_table_name(&key_info(&s, "read").await.table_id); + let mut holder = PgConnection::connect(&s.data_url()).await.expect("holder"); + sqlx::query("BEGIN") + .execute(&mut holder) + .await + .expect("begin"); + sqlx::query(&format!("LOCK TABLE {relation} IN ACCESS SHARE MODE")) + .execute(&mut holder) + .await + .expect("hold the data table as a plain reader"); + + let deleted = delete_within_5s(&s, "read").await.expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting, + "a held table takes the asynchronous path" + ); + let deleted_at = std::time::Instant::now(); + // The production poller, woken at once by the permit the DeleteTable + // (and the CreateTable) left; its active window is 5 s from now. + let _poller = s.engine.spawn_control_plane_poller_for_tests(); + tokio::time::sleep(std::time::Duration::from_secs(10)).await; + assert_eq!( + catalog_status(&s, "read").await.as_deref(), + Some("DELETING"), + "the poller cannot drop the table while the reader holds it" + ); + + sqlx::query("ROLLBACK") + .execute(&mut holder) + .await + .expect("release"); + holder.close().await.expect("close holder"); + let released_at = std::time::Instant::now(); + let data = s.data_pool().await; + while released_at.elapsed() < std::time::Duration::from_secs(3) { + let gone: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NULL") + .bind(&relation) + .fetch_one(&data) + .await + .expect("probe"); + if gone && catalog_status(&s, "read").await.is_none() { + eprintln!( + "plain reader: table removed {:?} after the release ({:?} after DeleteTable)", + released_at.elapsed(), + deleted_at.elapsed() + ); + data.close().await; + return; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + panic!( + "the table was not removed within 3 s of a plain reader's release, {:?} after \ + DeleteTable: the poller went idle while the drop was deferred", + deleted_at.elapsed() + ); +} +on_both_layouts!(a_plain_reader_past_the_active_window_delays_the_drop_only_until_it_releases); + +async fn dropping_a_stream_wakes_the_active_poller(shared: bool) { + release_wakes_the_active_poller(shared, Release::Dropped).await; +} +on_both_layouts!(dropping_a_stream_wakes_the_active_poller); + +async fn draining_a_stream_wakes_the_active_poller(shared: bool) { + release_wakes_the_active_poller(shared, Release::Drained).await; +} +on_both_layouts!(draining_a_stream_wakes_the_active_poller); + +async fn a_fetch_error_wakes_the_active_poller(shared: bool) { + release_wakes_the_active_poller(shared, Release::FetchError).await; +} +on_both_layouts!(a_fetch_error_wakes_the_active_poller); + +async fn cancelling_an_open_wakes_the_active_poller(shared: bool) { + release_wakes_the_active_poller(shared, Release::OpenCancelled).await; +} +on_both_layouts!(cancelling_an_open_wakes_the_active_poller); + +/// The in-catalog backup holds the data table's `ACCESS SHARE` from its +/// `LOCK TABLE` until it has copied the items (`HeldSnapshot::finish`). Held +/// there (its catalog `INSERT INTO backups` waits on a session holding +/// `backups` in `SHARE` mode), a `DeleteTable` leaves the row `DELETING`; +/// the backup is then let go during the active poller's pause after its +/// first drop attempt, and its end wakes the poller, which drops the table +/// within 3 s rather than at its next timed attempt about 9.5 s later. +async fn the_in_catalog_backup_wakes_the_active_poller(shared: bool) { + use extenddb_storage::BackupEngine as _; + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("backed")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "backed", 20).await; + let relation = data_table_name(&key_info(&s, "backed").await.table_id); + let catalog_url = format!("{}/{}", s.base, s.names[0]); + let mut blocker = PgConnection::connect(&catalog_url).await.expect("blocker"); + sqlx::query("BEGIN") + .execute(&mut blocker) + .await + .expect("begin"); + sqlx::query("LOCK TABLE backups IN SHARE MODE") + .execute(&mut blocker) + .await + .expect("hold backups"); + + let backup = s.engine.create_backup(ACCOUNT, "backed", "held"); + let controller = async { + let mut watcher = PgConnection::connect(&format!("{}/postgres", s.base)) + .await + .expect("watcher"); + let mut waiting = 0; + for _ in 0..100 { + waiting = sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity WHERE datname = $1 \ + AND wait_event_type = 'Lock' AND query LIKE 'INSERT INTO backups%'", + ) + .bind(&s.names[0]) + .fetch_one(&mut watcher) + .await + .expect("waiters"); + if waiting == 1 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + watcher.close().await.expect("close watcher"); + assert_eq!( + waiting, 1, + "precondition: the backup holds the data table and waits on backups" + ); + let deleted = delete_within_5s(&s, "backed").await.expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting, + "the backup's lock sends DeleteTable down the asynchronous path" + ); + let poller = s + .engine + .spawn_control_plane_poller_with_active_poll_for_tests(WAKE_TEST_ACTIVE_POLL); + let attempts = wait_for_drop_attempt(&s.engine, 0).await; + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + assert_eq!( + s.engine.control_plane_drop_attempts_for_tests(), + attempts, + "precondition: the poller is in the pause after its first attempt" + ); + assert_eq!( + catalog_status(&s, "backed").await.as_deref(), + Some("DELETING"), + "the poller cannot drop the table while the backup holds it" + ); + sqlx::query("ROLLBACK") + .execute(&mut blocker) + .await + .expect("let the backup go"); + (poller, std::time::Instant::now()) + }; + let (backed_up, (poller, released_at)) = tokio::join!(backup, controller); + backed_up.expect("the backup completes"); + removed_within_3s(&s, "backed", &[relation], released_at, "in-catalog backup").await; + drop(poller); + blocker.close().await.expect("close blocker"); +} +on_both_layouts!(the_in_catalog_backup_wakes_the_active_poller); + +/// A `FETCH` that fails (the backend ended under the stream) ends the stream +/// at the error and frees its permit and connection at once: with the +/// errored stream still held a full bound of five opens succeeds. The +/// `SnapshotExpired` codes take this branch. +async fn a_fetch_error_ends_the_stream_and_frees_its_permit(shared: bool) { + use futures::StreamExt as _; + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("ended")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "ended", 20).await; + let _page = PageSize::set(5); + let mut snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "ended") + .await + .expect("open"); + snapshot.items.next().await.expect("an item").expect("ok"); + terminate_stream_backend(&s).await; + let mut items = 1; + loop { + match snapshot.items.next().await { + Some(Ok(_)) => items += 1, + Some(Err(e)) => { + eprintln!("FETCH error after {items} items: {e}"); + break; + } + None => panic!("the stream ended without reporting the FETCH error"), + } + } + assert_eq!(items, 5, "the first page, then the error"); + let mut held = Vec::new(); + for i in 0..5 { + held.push( + s.engine + .open_table_snapshot(ACCOUNT, "ended") + .await + .unwrap_or_else(|e| panic!("open {i} with the errored stream held: {e}")), + ); + } + assert!( + snapshot.items.next().await.is_none(), + "nothing follows the first error" + ); +} +on_both_layouts!(a_fetch_error_ends_the_stream_and_frees_its_permit); + +/// `DeleteTable` while an open is held at `BeforeItems` (definition read, +/// snapshot pinned, cursor not declared) takes the asynchronous path: the +/// open already holds `ACCESS SHARE` on the data table, so the drop cannot +/// remove it under the open, and the open then streams every item. +async fn delete_table_at_before_items_takes_the_async_path(shared: bool) { + use futures::StreamExt as _; + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("early")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "early", 20).await; + let mut pause = extenddb_storage::snapshot::paging::arm_open_pause( + extenddb_storage::snapshot::paging::OpenStage::BeforeItems, + ); + let reached = pause.take_reached(); + let mut open = s.engine.open_table_snapshot(ACCOUNT, "early"); + tokio::select! { + r = &mut open => panic!("the open finished while it should be paused: {r:?}"), + r = reached => r.expect("the open reached BeforeItems"), + } + let deleted = delete_within_5s(&s, "early").await.expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert_eq!( + catalog_status(&s, "early").await.as_deref(), + Some("DELETING"), + "the row stays until the worker's drop" + ); + pause.resume(); + let snapshot = open.await.expect("the held open completes"); + let items: Vec<_> = snapshot.items.collect().await; + assert_eq!(items.len(), 20, "the open streams every item"); + assert!(items.iter().all(Result::is_ok)); +} +on_both_layouts!(delete_table_at_before_items_takes_the_async_path); + +/// A stream commits its transaction at the fetch that shows the end (a short +/// page), not when the stream next yields `None` or is dropped: with every +/// item read and the stream still held, its backend is idle (out of its +/// transaction, holding no lock on the data table) and a full bound of five +/// opens succeeds. Twenty-two items at a page of five: the fifth fetch +/// returns two rows, so the end is known before the last item is yielded +/// (after `None` the stream's state is dropped, which would release +/// everything whether or not the last page did). +async fn a_drained_stream_ends_its_transaction(shared: bool) { + use futures::StreamExt as _; + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("drained")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "drained", 22).await; + let table_id = key_info(&s, "drained").await.table_id; + let page = PageSize::set(5); + let mut snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "drained") + .await + .expect("open"); + snapshot.items.next().await.expect("an item").expect("ok"); + let data = s.data_pool().await; + let pid = stream_backend(&s, &data).await; + for _ in 1..22 { + snapshot + .items + .next() + .await + .expect("an item") + .expect("stream item"); + } + drop(page); + let (state, locks): (Option, i64) = sqlx::query_as( + "SELECT (SELECT state FROM pg_stat_activity WHERE pid = $1), \ + (SELECT count(*) FROM pg_locks WHERE pid = $1 AND relation = to_regclass($2))", + ) + .bind(pid) + .bind(data_table_name(&table_id)) + .fetch_one(&data) + .await + .expect("backend state"); + data.close().await; + assert_eq!( + state.as_deref(), + Some("idle"), + "the drained stream's backend left its transaction" + ); + assert_eq!(locks, 0, "and holds no lock on the data table"); + let mut held = Vec::new(); + for i in 0..5 { + held.push( + s.engine + .open_table_snapshot(ACCOUNT, "drained") + .await + .unwrap_or_else(|e| panic!("open {i} with the drained stream held: {e}")), + ); + } + assert!( + snapshot.items.next().await.is_none(), + "twenty-two items, then the end" + ); +} +on_both_layouts!(a_drained_stream_ends_its_transaction); + +/// With all but one of the engine's ten connections held (nine +/// `DeleteTable`s of another table, each waiting on that table's row, which +/// an outside session holds `FOR UPDATE`), an open still completes: on the +/// shared layout it needs one connection, and with a separate data database +/// one catalog and one data connection from different pools. An open that +/// took two connections from one pool would wait out the acquire timeout. +async fn an_open_needs_one_connection_per_pool(shared: bool) { + use futures::StreamExt as _; + let s = scratch_with(shared).await; + create_active(&s, plain_input("blocked")).await; + create_active(&s, snapshot_conformance::table_input("opened")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "opened", 20).await; + let mut holder = PgConnection::connect(&format!("{}/{}", s.base, s.names[0])) + .await + .expect("holder"); + sqlx::query("BEGIN") + .execute(&mut holder) + .await + .expect("begin"); + sqlx::query("SELECT 1 FROM tables WHERE table_name = 'blocked' FOR UPDATE") + .execute(&mut holder) + .await + .expect("hold the row"); + let deletes = futures::future::join_all( + (0..9).map(|_| s.engine.delete_table(ACCOUNT, delete_input("blocked"))), + ); + let started = std::time::Instant::now(); + let open_while_held = async { + // Until all nine deletes hold a connection and wait on the row. Read + // from a session outside any transaction: pg_stat_activity is + // snapshotted once per transaction. + let mut watcher = PgConnection::connect(&format!("{}/postgres", s.base)) + .await + .expect("watcher"); + let mut waiting = 0; + for _ in 0..100 { + waiting = sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity WHERE datname = $1 \ + AND wait_event_type = 'Lock'", + ) + .bind(&s.names[0]) + .fetch_one(&mut watcher) + .await + .expect("waiters"); + if waiting == 9 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + watcher.close().await.expect("close watcher"); + assert_eq!( + waiting, 9, + "precondition: nine DeleteTables hold a connection each" + ); + let opened = tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.open_table_snapshot(ACCOUNT, "opened"), + ) + .await; + let opened = opened + .expect("the open gets its connections with nine of ten held") + .expect("open"); + let items: Vec<_> = opened.items.collect().await; + sqlx::query("ROLLBACK") + .execute(&mut holder) + .await + .expect("release the row"); + items + }; + let (results, items) = tokio::join!(deletes, open_while_held); + eprintln!("open and nine waiting deletes took {:?}", started.elapsed()); + assert_eq!(items.len(), 20); + assert!(items.iter().all(Result::is_ok)); + for result in results { + match result { + Ok(_) | Err(extenddb_storage::error::StorageError::TableNotFound(_)) => {} + Err(e) => panic!("a queued DeleteTable failed: {e}"), + } + } +} +on_both_layouts!(an_open_needs_one_connection_per_pool); + +/// On one shared database the open holds no row lock: a `DeleteTable` that +/// runs after the open's snapshot and before its `LOCK TABLE` takes the +/// synchronous path and drops the data table, and the open then reports +/// `TableNotFound` (the `LOCK` fails with `undefined_table`), not an internal +/// error. On the separate layout the barrier makes that delete wait, so the +/// case does not arise there. +#[tokio::test] +async fn a_table_dropped_before_the_lock_is_not_found_on_one_shared_database() { + if base_conn().is_none() { + eprintln!("SKIP a_table_dropped_before_the_lock_is_not_found_on_one_shared_database"); + return; + } + let s = scratch_with(true).await; + create_active(&s, snapshot_conformance::table_input("gone")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "gone", 5).await; + let mut pause = extenddb_storage::snapshot::paging::arm_open_pause( + extenddb_storage::snapshot::paging::OpenStage::AfterDefinition, + ); + let reached = pause.take_reached(); + let mut open = s.engine.open_table_snapshot(ACCOUNT, "gone"); + tokio::select! { + r = &mut open => panic!("the open finished while it should be paused: {r:?}"), + r = reached => r.expect("the open reached AfterDefinition"), + } + delete_within_5s(&s, "gone").await.expect("delete"); + assert_eq!( + catalog_status(&s, "gone").await, + None, + "with no lock on the data table yet, DeleteTable took the synchronous path" + ); + pause.resume(); + let err = open + .await + .expect_err("the data table is gone before the open could lock it"); + assert!( + matches!(err, extenddb_storage::error::StorageError::TableNotFound(_)), + "{err}" + ); +} + +/// The base and GSI data tables of `table` (a `table_input` table). +async fn base_and_gsi_relations(s: &Scratch, table: &str) -> Vec { + let table_id = key_info(s, table).await.table_id; + let index_ids: Vec = + sqlx::query_scalar("SELECT index_id FROM indexes WHERE table_id = $1") + .bind(&table_id) + .fetch_all(&s.catalog) + .await + .expect("index ids"); + assert_eq!(index_ids.len(), 1, "the check table has one GSI"); + std::iter::once(data_table_name(&table_id)) + .chain(index_ids.iter().map(|id| index_table_name(id))) + .collect() +} + +/// Which of `relations` exist on the data database. +async fn present_relations(s: &Scratch, relations: &[String]) -> Vec { + let data = s.data_pool().await; + let mut present = Vec::new(); + for relation in relations { + let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(relation) + .fetch_one(&data) + .await + .expect("probe"); + if exists { + present.push(relation.clone()); + } + } + data.close().await; + present +} + +/// A reader that takes its lock after `DeleteTable`'s in-use probe. +#[derive(Debug, Clone, Copy)] +enum LateReader { + /// A table snapshot opened in the window (it locks the base data table). + Open, + /// A plain session locking only the GSI data table. + GsiSession, +} + +/// On one shared database `DeleteTable` runs its drops in its own catalog +/// transaction, under a 1 s lock timeout, before the catalog removal +/// commits. A reader that takes its lock after the in-use probe (held there +/// by the `DeleteTable` pause) makes a drop time out: the savepoint is +/// rolled back and `DeleteTable` returns `DELETING` with the row and every +/// data table kept, and the control plane drops them once the reader ends. +/// No data table is left without its row at any point. With the drops after +/// the commit, the delete would wait behind the reader with the row already +/// gone. +async fn a_reader_after_the_probe_sends_delete_table_down_the_async_path(how: LateReader) { + use futures::StreamExt as _; + let s = scratch_with(true).await; + create_active(&s, snapshot_conformance::table_input("late")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "late", 5).await; + let relations = base_and_gsi_relations(&s, "late").await; + + let mut pause = extenddb_storage_postgres::arm_delete_table_pause(); + let reached = pause.take_reached(); + let mut delete = Box::pin(tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.delete_table(ACCOUNT, delete_input("late")), + )); + tokio::select! { + r = &mut delete => panic!("DeleteTable finished while it should be paused: {r:?}"), + r = reached => r.expect("DeleteTable reached the pause"), + } + let (mut snapshot, mut session) = (None, None); + match how { + LateReader::Open => { + snapshot = Some( + s.engine + .open_table_snapshot(ACCOUNT, "late") + .await + .expect("the open does not wait on DeleteTable's row lock"), + ); + } + LateReader::GsiSession => { + let mut conn = PgConnection::connect(&s.data_url()).await.expect("reader"); + sqlx::query("BEGIN") + .execute(&mut conn) + .await + .expect("begin"); + sqlx::query(&format!("LOCK TABLE {} IN ACCESS SHARE MODE", relations[1])) + .execute(&mut conn) + .await + .expect("lock the GSI data table"); + session = Some(conn); + } + } + pause.resume(); + let deleted = delete + .await + .unwrap_or_else(|_| { + panic!( + "{how:?}: DeleteTable did not return within 5 s; its drop waited behind the \ + reader instead of giving up after its lock timeout" + ) + }) + .expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert_eq!( + catalog_status(&s, "late").await.as_deref(), + Some("DELETING"), + "{how:?}: the row stays while the reader holds the table" + ); + assert_eq!( + present_relations(&s, &relations).await, + relations, + "{how:?}: the timed-out drop was rolled back with the catalog removal" + ); + + if let Some(snapshot) = snapshot.take() { + let items: Vec<_> = snapshot.items.collect().await; + assert_eq!(items.len(), 5, "the snapshot streams every item"); + assert!(items.iter().all(Result::is_ok)); + } + if let Some(mut conn) = session.take() { + sqlx::query("ROLLBACK") + .execute(&mut conn) + .await + .expect("release"); + conn.close().await.expect("close reader"); + } + for _ in 0..100 { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + let present = present_relations(&s, &relations).await; + let row = catalog_status(&s, "late").await; + assert!( + row.is_some() || present.is_empty(), + "{how:?}: data tables {present:?} without a catalog row" + ); + if row.is_none() { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + panic!("{how:?}: the control plane did not drop the table after the reader ended"); +} + +#[tokio::test] +async fn an_open_after_the_probe_sends_delete_table_down_the_async_path_on_one_shared_database() { + if base_conn().is_none() { + eprintln!("SKIP an_open_after_the_probe_...: no PostgreSQL"); + return; + } + a_reader_after_the_probe_sends_delete_table_down_the_async_path(LateReader::Open).await; +} + +#[tokio::test] +async fn a_gsi_reader_after_the_probe_sends_delete_table_down_the_async_path_on_one_shared_database() + { + if base_conn().is_none() { + eprintln!("SKIP a_gsi_reader_after_the_probe_...: no PostgreSQL"); + return; + } + a_reader_after_the_probe_sends_delete_table_down_the_async_path(LateReader::GsiSession).await; +} + +/// Hold `ACCESS SHARE` on `relation` from a new session and return it. +async fn hold_relation(s: &Scratch, relation: &str) -> PgConnection { + let mut holder = PgConnection::connect(&s.data_url()).await.expect("holder"); + sqlx::query("BEGIN") + .execute(&mut holder) + .await + .expect("begin"); + sqlx::query(&format!("LOCK TABLE {relation} IN ACCESS SHARE MODE")) + .execute(&mut holder) + .await + .expect("hold the relation"); + holder +} + +async fn release_relation(mut holder: PgConnection) { + sqlx::query("ROLLBACK") + .execute(&mut holder) + .await + .expect("release"); + holder.close().await.expect("close holder"); +} + +/// A reader that spans a restart: `DeleteTable` left the row `DELETING` +/// while a session held the data table, and the server restarted (a new +/// engine on the same databases, running the startup recovery `serve` runs) +/// before the reader ended. The startup pass defers the drop and wakes the +/// poller, which starts in its active loop and stays there past its window +/// for as long as the drop is deferred, so the table is gone within 3 s of +/// the release. A poller that started idle would drop it at its first idle +/// pass, 60 s after startup, and one that went idle at the end of its window +/// would drop it at the idle pass after that. +/// +/// The window is `control_plane_delay_seconds` (0) + 5 s from the wake, with +/// a pass at once and one every second (each pass probes the held table and +/// skips it, so it takes milliseconds); the pass that finds the window over, +/// at about 5 s, is the sixth. The reader is released only after the poller +/// has made more attempts than the startup pass and those six, so a poller +/// that went idle at the end of its window never gets there and fails the +/// wait instead. +async fn a_drop_deferred_at_startup_is_retried_while_the_reader_holds(shared: bool) { + /// The startup pass's attempt plus the most the window can start. + const ATTEMPTS_THE_WINDOW_CAN_MAKE: u64 = 1 + 6; + let s = scratch_with(shared).await; + create_active(&s, plain_input("spans")).await; + let relation = data_table_name(&key_info(&s, "spans").await.table_id); + let holder = hold_relation(&s, &relation).await; + let deleted = delete_within_5s(&s, "spans").await.expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + + let restarted = PostgresEngine::new( + &PostgresConfig { + connection_string: format!("{}/{}", s.base, s.names[0]), + pool_size: 8, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("the restarted engine"); + restarted.recover_control_plane_at_startup_for_tests().await; + assert_eq!( + restarted.control_plane_drop_attempts_for_tests(), + 1, + "precondition: the startup pass tried the drop" + ); + let _poller = restarted.spawn_control_plane_poller_for_tests(); + let started = std::time::Instant::now(); + // Past the attempts the window can start, so only the deferred drop can + // have kept the poller active. + let attempts = wait_for_drop_attempt(&restarted, ATTEMPTS_THE_WINDOW_CAN_MAKE).await; + assert_eq!( + catalog_status(&s, "spans").await.as_deref(), + Some("DELETING"), + "the poller cannot drop the table while the reader holds it" + ); + eprintln!( + "{attempts} drop attempts by the restarted engine {:?} after its poller started, \ + before the release", + started.elapsed() + ); + release_relation(holder).await; + removed_within_3s( + &s, + "spans", + &[relation], + std::time::Instant::now(), + "reader across a restart", + ) + .await; +} +on_both_layouts!(a_drop_deferred_at_startup_is_retried_while_the_reader_holds); + +/// A drop that fails for a reason that persists is not a deferred drop. A +/// view on the data table makes `DROP TABLE` fail with 2BP01 (dependent +/// objects) on every attempt; the poller logs it and follows its normal +/// window, a pass every second for `control_plane_delay_seconds` (0) + 5 s +/// and then idle until its 60 s idle pass, rather than retrying every second +/// for as long as the failure lasts. Counted over the 15 s after the window: +/// at most 3 attempts (a poller kept active would make about 15). +async fn a_drop_that_keeps_failing_does_not_keep_the_poller_active(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, plain_input("viewed")).await; + let relation = data_table_name(&key_info(&s, "viewed").await.table_id); + let data = s.data_pool().await; + sqlx::query(&format!( + "CREATE VIEW viewed_dependent AS SELECT * FROM {relation}" + )) + .execute(&data) + .await + .expect("a view on the data table"); + // A reader sends DeleteTable down the asynchronous path, so the row is + // DELETING and the drop is the poller's. + let holder = hold_relation(&s, &relation).await; + let deleted = delete_within_5s(&s, "viewed").await.expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + release_relation(holder).await; + + let _poller = s.engine.spawn_control_plane_poller_for_tests(); + tokio::time::sleep(std::time::Duration::from_secs(7)).await; + let after_window = s.engine.control_plane_drop_attempts_for_tests(); + assert!( + after_window >= 1, + "precondition: the poller attempted the drop during its window" + ); + tokio::time::sleep(std::time::Duration::from_secs(15)).await; + let in_15s = s.engine.control_plane_drop_attempts_for_tests() - after_window; + eprintln!("{after_window} drop attempts in the window, {in_15s} in the 15 s after it"); + assert!( + in_15s <= 3, + "{in_15s} drop attempts in 15 s after the window: a persistent failure kept the \ + poller in its one-second loop" + ); + assert_eq!( + catalog_status(&s, "viewed").await.as_deref(), + Some("DELETING"), + "the drop still fails, so the row stays" + ); + sqlx::query("DROP VIEW viewed_dependent") + .execute(&data) + .await + .expect("drop the view"); + data.close().await; +} +on_both_layouts!(a_drop_that_keeps_failing_does_not_keep_the_poller_active); + +/// What one open-and-delete pair of [`open_and_delete_race`] left. +#[derive(Debug, Default)] +struct RaceTally { + /// `DeleteTable` removed the row in the request. + synchronous: usize, + /// `DeleteTable` returned with the row `DELETING`. + asynchronous: usize, + /// The caller's 2 s timeout cancelled `DeleteTable`. + delete_timed_out: usize, + /// The open found the table gone. + open_not_found: usize, + /// No catalog row, with a data table still present after the stream + /// ended and a control-plane pass ran. + orphaned: Vec, +} + +/// An open and a `DeleteTable` of the same table started together, with no +/// test hook: in even pairs the delete 0 to 5 ms after the open, in odd ones +/// the open 0 to 5 ms after the delete, so both can win; the stream held 3 s, +/// the delete under a 2 s caller timeout, then a control-plane pass. Counts +/// the pairs that left a data table with no catalog row. Timing-dependent, +/// so not part of the suite: run with `--ignored open_and_delete_race`. +async fn open_and_delete_race(shared: bool) -> RaceTally { + let s = scratch_with(shared).await; + let mut tally = RaceTally::default(); + for i in 0..80_u64 { + let table = format!("race{i:02}"); + create_active(&s, snapshot_conformance::table_input(&table)).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, &table, 5).await; + let relations = base_and_gsi_relations(&s, &table).await; + let stagger = std::time::Duration::from_micros((i * 997) % 5000); + let open = async { + if i % 2 == 1 { + tokio::time::sleep(stagger).await; + } + match s.engine.open_table_snapshot(ACCOUNT, &table).await { + Ok(snapshot) => { + tokio::time::sleep(std::time::Duration::from_secs(3)).await; + drop(snapshot); + false + } + Err(extenddb_storage::error::StorageError::TableNotFound(_)) => true, + Err(e) => panic!("open {table}: {e}"), + } + }; + let delete = async { + if i % 2 == 0 { + tokio::time::sleep(stagger).await; + } + tokio::time::timeout( + std::time::Duration::from_secs(2), + s.engine.delete_table(ACCOUNT, delete_input(&table)), + ) + .await + }; + let (not_found, deleted) = tokio::join!(open, delete); + tally.open_not_found += usize::from(not_found); + let returned = match deleted { + Err(_) => { + tally.delete_timed_out += 1; + false + } + Ok(Ok(_)) => true, + Ok(Err(e)) => panic!("DeleteTable {table}: {e}"), + }; + let status_after_delete = catalog_status(&s, &table).await; + if returned { + if status_after_delete.is_some() { + tally.asynchronous += 1; + } else { + tally.synchronous += 1; + } + } + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + if catalog_status(&s, &table).await.is_none() { + for relation in present_relations(&s, &relations).await { + tally.orphaned.push(format!("{table}: {relation}")); + } + } + } + eprintln!("open-and-delete race, shared = {shared}: {tally:?}"); + tally +} + +#[tokio::test] +#[ignore = "timing-dependent race measurement; run by hand"] +async fn open_and_delete_race_on_one_shared_database() { + if base_conn().is_none() { + eprintln!("SKIP open_and_delete_race_on_one_shared_database: no PostgreSQL"); + return; + } + let tally = open_and_delete_race(true).await; + assert!(tally.orphaned.is_empty(), "{tally:?}"); + assert!( + tally.synchronous > 0 && tally.asynchronous > 0, + "the race never had both outcomes, so it measured nothing: {tally:?}" + ); +} + +#[tokio::test] +#[ignore = "timing-dependent race measurement; run by hand"] +async fn open_and_delete_race_on_a_separate_data_database() { + if base_conn().is_none() { + eprintln!("SKIP open_and_delete_race_on_a_separate_data_database: no PostgreSQL"); + return; + } + let tally = open_and_delete_race(false).await; + assert!(tally.orphaned.is_empty(), "{tally:?}"); + assert!( + tally.synchronous > 0 && tally.asynchronous > 0, + "the race never had both outcomes, so it measured nothing: {tally:?}" + ); +} + +/// Hold `relation` in lock `mode` (`ROW EXCLUSIVE` for an open write +/// transaction, `SHARE` for `LOCK ... IN SHARE MODE`, ...) from a new session +/// on the data database and return it. +async fn hold_relation_in(s: &Scratch, relation: &str, mode: &str) -> PgConnection { + let mut holder = PgConnection::connect(&s.data_url()).await.expect("holder"); + sqlx::query("BEGIN") + .execute(&mut holder) + .await + .expect("begin"); + sqlx::query(&format!("LOCK TABLE {relation} IN {mode} MODE")) + .execute(&mut holder) + .await + .unwrap_or_else(|e| panic!("hold {relation} in {mode} mode: {e}")); + holder +} + +/// Panics when `table` has no catalog row while one of `relations` is still +/// present on the data database: data tables no request can reach or remove. +async fn assert_no_orphan(s: &Scratch, table: &str, relations: &[String], when: &str) { + let row = catalog_status(s, table).await; + let present = present_relations(s, relations).await; + assert!( + row.is_some() || present.is_empty(), + "{when}: data tables {present:?} left without a catalog row" + ); +} + +/// Pump the control plane until `table` has no catalog row, asserting at +/// every step that no relation of it is left without one. +async fn pump_until_removed(s: &Scratch, table: &str, relations: &[String]) { + for _ in 0..100 { + s.engine + .process_control_plane_transitions() + .await + .expect("pump"); + assert_no_orphan(s, table, relations, "while the control plane drops it").await; + if catalog_status(s, table).await.is_none() { + assert!( + present_relations(s, relations).await.is_empty(), + "the row and every relation go together" + ); + return; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + panic!("the control plane did not remove {table} after its holder ended"); +} + +/// A session that holds the base data table in a mode other than `ACCESS +/// SHARE` (an open write transaction, `LOCK ... IN SHARE MODE`) is not one +/// the in-use probe matches, so `DeleteTable` tries its synchronous drop and +/// waits on it. That wait is bounded by `SYNC_DROP_LOCK_TIMEOUT` (1 s) on +/// both layouts: `DeleteTable` returns `DELETING` with the row and every data +/// table kept, and the control plane drops them once the holder ends. No +/// data table is left without its row at any point, including while the +/// `DeleteTable` waits. +async fn a_non_reader_holder_sends_delete_table_down_the_async_path(shared: bool, mode: &str) { + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("written")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "written", 5).await; + let relations = base_and_gsi_relations(&s, "written").await; + let holder = hold_relation_in(&s, &relations[0], mode).await; + + let started = std::time::Instant::now(); + let deleted = tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.delete_table(ACCOUNT, delete_input("written")), + ) + .await; + let elapsed = started.elapsed(); + assert_no_orphan(&s, "written", &relations, &format!("{mode} holder")).await; + let deleted = deleted + .unwrap_or_else(|_| panic!("{mode} holder: DeleteTable did not return within 5 s")) + .unwrap_or_else(|e| panic!("{mode} holder: DeleteTable failed: {e}")); + eprintln!("{mode} holder: DeleteTable returned after {elapsed:?}"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert!( + elapsed < std::time::Duration::from_secs(3), + "{mode} holder: DeleteTable took {elapsed:?}, more than about its 1 s lock timeout" + ); + assert_eq!( + catalog_status(&s, "written").await.as_deref(), + Some("DELETING"), + "{mode} holder: the row stays while the holder has the table" + ); + assert_eq!( + present_relations(&s, &relations).await, + relations, + "{mode} holder: the drop that timed out was rolled back" + ); + release_relation(holder).await; + pump_until_removed(&s, "written", &relations).await; +} + +async fn a_write_holder_sends_delete_table_down_the_async_path(shared: bool) { + a_non_reader_holder_sends_delete_table_down_the_async_path(shared, "ROW EXCLUSIVE").await; +} +on_both_layouts!(a_write_holder_sends_delete_table_down_the_async_path); + +async fn a_share_holder_sends_delete_table_down_the_async_path(shared: bool) { + a_non_reader_holder_sends_delete_table_down_the_async_path(shared, "SHARE").await; +} +on_both_layouts!(a_share_holder_sends_delete_table_down_the_async_path); + +/// A `DeleteTable` whose caller gives up while it waits on a write holder (a +/// client timeout shorter than the lock timeout, a dropped connection, a +/// server shutting down) leaves the table either `ACTIVE` with its data or +/// `DELETING` for the control plane, never with its row gone and its data +/// tables left: nothing commits before the drops have their locks. +async fn a_cancelled_delete_table_leaves_no_orphan(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("dropped")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "dropped", 5).await; + let relations = base_and_gsi_relations(&s, "dropped").await; + let holder = hold_relation_in(&s, &relations[0], "ROW EXCLUSIVE").await; + + let cancelled = tokio::time::timeout( + std::time::Duration::from_millis(300), + s.engine.delete_table(ACCOUNT, delete_input("dropped")), + ) + .await; + assert!( + cancelled.is_err(), + "precondition: DeleteTable was still waiting on the holder at 300 ms: {cancelled:?}" + ); + // Long enough for the abandoned transactions to end on the server: the + // drop waits out what is left of its 1 s lock timeout, then rolls back. + tokio::time::sleep(std::time::Duration::from_millis(1500)).await; + assert_no_orphan(&s, "dropped", &relations, "after the cancelled DeleteTable").await; + let status = catalog_status(&s, "dropped").await; + eprintln!("cancelled DeleteTable left the row {status:?}"); + assert!( + matches!(status.as_deref(), Some("ACTIVE" | "DELETING")), + "the row is kept: {status:?}" + ); + assert_eq!( + present_relations(&s, &relations).await, + relations, + "and so is every data table" + ); + release_relation(holder).await; + if status.as_deref() == Some("ACTIVE") { + let deleted = delete_within_5s(&s, "dropped").await.expect("delete again"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + } + pump_until_removed(&s, "dropped", &relations).await; +} +on_both_layouts!(a_cancelled_delete_table_leaves_no_orphan); + +/// Connect to the server's `postgres` database as an admin, for commands +/// against the scratch data database itself (its connection limit). +async fn data_admin_pool(s: &Scratch) -> PgPool { + PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{}/postgres", s.base)) + .await + .expect("admin pool on postgres") +} + +/// Terminate the data-database backend that holds a `DeleteTable`'s +/// uncommitted drops, so its data commit fails. At the `BetweenCommits` +/// pause that backend sits idle in its transaction after the drops; it is +/// the only one on the data database in that state. +async fn terminate_held_data_backend(s: &Scratch) { + let data = s.data_pool().await; + let mut pid = None; + for _ in 0..400 { + pid = sqlx::query_scalar::<_, i32>( + "SELECT pid FROM pg_stat_activity \ + WHERE datname = $1 AND state = 'idle in transaction' \ + AND pid <> pg_backend_pid() \ + ORDER BY state_change DESC LIMIT 1", + ) + .bind(s.data_name()) + .fetch_optional(&data) + .await + .expect("find the held data backend"); + if pid.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + let pid = pid.expect("a data backend idle in its transaction after the drops"); + let terminated: bool = sqlx::query_scalar("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(&data) + .await + .expect("terminate"); + assert!(terminated, "pg_terminate_backend({pid})"); + data.close().await; +} + +/// Set whether the scratch data database accepts connections. `false` blocks +/// even the superuser the engine connects as (a connection limit does not). +async fn set_data_allow_connections(s: &Scratch, allow: bool) { + let admin = data_admin_pool(s).await; + sqlx::query(&format!( + "ALTER DATABASE \"{}\" WITH ALLOW_CONNECTIONS {allow}", + s.data_name() + )) + .execute(&admin) + .await + .expect("set the data database's allow_connections"); + admin.close().await; +} + +/// Make the scratch data database refuse every new connection and end the +/// ones it has, so a `DeleteTable`'s data commit and every retry of its drop +/// fail. +async fn make_data_unreachable(s: &Scratch) { + set_data_allow_connections(s, false).await; + let admin = data_admin_pool(s).await; + sqlx::query( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity \ + WHERE datname = $1 AND pid <> pg_backend_pid()", + ) + .bind(s.data_name()) + .execute(&admin) + .await + .expect("terminate every data backend"); + admin.close().await; +} + +/// Separate data database: when the data commit fails after the catalog +/// commit, the spawned task retries the drop on a fresh connection, so the +/// delete still succeeds and no data table is left without its row. This is +/// the window the PR used to describe as "a process that dies between the +/// two commits"; any failure of the data commit reaches it, with the process +/// up, and the retry closes it. Mutant: return the commit error instead of +/// retrying, and the relations stay. +#[tokio::test] +async fn a_failed_data_commit_after_the_catalog_commit_is_retried() { + if base_conn().is_none() { + eprintln!("SKIP a_failed_data_commit_after_the_catalog_commit_is_retried: no PostgreSQL"); + return; + } + let s = scratch_with(false).await; + create_active(&s, snapshot_conformance::table_input("recommit")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "recommit", 5).await; + let relations = base_and_gsi_relations(&s, "recommit").await; + + let mut pause = extenddb_storage_postgres::arm_delete_table_pause_at( + extenddb_storage_postgres::DeleteTablePoint::BetweenCommits, + ); + let reached = pause.take_reached(); + let mut delete = Box::pin(tokio::time::timeout( + std::time::Duration::from_secs(30), + s.engine.delete_table(ACCOUNT, delete_input("recommit")), + )); + tokio::select! { + r = &mut delete => panic!("DeleteTable finished before the between-commits pause: {r:?}"), + r = reached => r.expect("DeleteTable reached the between-commits pause"), + } + assert_eq!( + catalog_status(&s, "recommit").await.as_deref(), + None, + "the catalog commit is in by the pause" + ); + terminate_held_data_backend(&s).await; + pause.resume(); + + let deleted = delete + .await + .expect("the retry completes well within 30 s") + .expect("DeleteTable returns Ok once the drop is redone"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + for _ in 0..50 { + if present_relations(&s, &relations).await.is_empty() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + assert!( + present_relations(&s, &relations).await.is_empty(), + "the retry dropped every data relation" + ); + assert_no_orphan(&s, "recommit", &relations, "after the retried drop").await; +} + +/// Separate data database: when the data database stays unreachable past the +/// retry bound, the delete gives up, names the leaked relations in its error +/// (and logs them with the `DROP TABLE` to run), and the catalog row is gone +/// regardless. A test-shortened bound keeps it quick. +#[tokio::test] +async fn an_unreachable_data_database_past_the_bound_reports_the_leaked_relations() { + if base_conn().is_none() { + eprintln!( + "SKIP an_unreachable_data_database_past_the_bound_reports_the_leaked_relations: no \ + PostgreSQL" + ); + return; + } + let s = scratch_with(false).await; + let _bound = extenddb_storage_postgres::shorten_data_drop_retry_bound_for_tests( + std::time::Duration::from_millis(500), + ); + create_active(&s, snapshot_conformance::table_input("stranded")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "stranded", 5).await; + let relations = base_and_gsi_relations(&s, "stranded").await; + + let mut pause = extenddb_storage_postgres::arm_delete_table_pause_at( + extenddb_storage_postgres::DeleteTablePoint::BetweenCommits, + ); + let reached = pause.take_reached(); + let mut delete = Box::pin(tokio::time::timeout( + std::time::Duration::from_secs(45), + s.engine.delete_table(ACCOUNT, delete_input("stranded")), + )); + tokio::select! { + r = &mut delete => panic!("DeleteTable finished before the pause: {r:?}"), + r = reached => r.expect("DeleteTable reached the between-commits pause"), + } + make_data_unreachable(&s).await; + pause.resume(); + + let err = delete + .await + .expect("the delete gives up within the shortened bound") + .expect_err("an unreachable data database past the bound is an error"); + let msg = err.to_string(); + for relation in &relations { + let core = relation.trim_matches('"'); + assert!(msg.contains(core), "the error names {core}: {msg}"); + } + assert_eq!( + catalog_status(&s, "stranded").await.as_deref(), + None, + "the catalog row is removed even though the data drop could not be redone" + ); + // Let the engine reach the data database again so the Drop cleanup can + // force-drop it. + set_data_allow_connections(&s, true).await; +} + +/// Separate data database: a caller that gives up after the catalog commit +/// does not stop the data drop. The commits and retries run in a task the +/// request spawns, so dropping the request only detaches it. Mutant: run the +/// body inline instead of spawning, and a dropped request strands the data +/// tables. +#[tokio::test] +async fn a_caller_cancelled_between_the_commits_does_not_strand_the_data() { + if base_conn().is_none() { + eprintln!( + "SKIP a_caller_cancelled_between_the_commits_does_not_strand_the_data: no PostgreSQL" + ); + return; + } + let s = scratch_with(false).await; + create_active(&s, snapshot_conformance::table_input("detached")).await; + snapshot_conformance::populate(&s.engine, ACCOUNT, "detached", 5).await; + let relations = base_and_gsi_relations(&s, "detached").await; + + let mut pause = extenddb_storage_postgres::arm_delete_table_pause_at( + extenddb_storage_postgres::DeleteTablePoint::BetweenCommits, + ); + let reached = pause.take_reached(); + let mut delete = Box::pin(s.engine.delete_table(ACCOUNT, delete_input("detached"))); + tokio::select! { + r = &mut delete => panic!("DeleteTable finished before the pause: {r:?}"), + r = reached => r.expect("DeleteTable reached the between-commits pause"), + } + // The caller gives up now, after the catalog commit. + drop(delete); + // Let the detached task continue past the pause and commit the data drop. + pause.resume(); + + for _ in 0..50 { + if catalog_status(&s, "detached").await.is_none() + && present_relations(&s, &relations).await.is_empty() + { + assert_no_orphan(&s, "detached", &relations, "after the detached drop").await; + return; + } + assert_no_orphan(&s, "detached", &relations, "while the detached task runs").await; + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + panic!("the detached task did not finish the drop after the caller was cancelled"); +} + +/// A hash-only table with three GSIs, for the drop-deadline test. +fn three_gsi_input(name: &str) -> extenddb_core::types::CreateTableInput { + let gsi = |i: usize| { + serde_json::json!({ + "IndexName": format!("g{i}"), + "KeySchema": [{"AttributeName": format!("a{i}"), "KeyType": "HASH"}], + "Projection": {"ProjectionType": "KEYS_ONLY"} + }) + }; + serde_json::from_value(serde_json::json!({ + "TableName": name, + "KeySchema": [{"AttributeName": "pk", "KeyType": "HASH"}], + "AttributeDefinitions": [ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "a0", "AttributeType": "S"}, + {"AttributeName": "a1", "AttributeType": "S"}, + {"AttributeName": "a2", "AttributeType": "S"} + ], + "BillingMode": "PAY_PER_REQUEST", + "GlobalSecondaryIndexes": [gsi(0), gsi(1), gsi(2)] + })) + .expect("create input") +} + +/// The drops of one `DeleteTable` share one deadline. A write holder on the +/// base data table and on each of three GSI data tables, released 0.8 s +/// apart in drop order, keeps every single wait under the 1 s lock timeout; +/// a timeout per `DROP` would have the request wait for all four, about +/// 3.2 s, while it holds the base table's `ACCESS EXCLUSIVE` and stalls every +/// reader of the table. With one deadline the drops give up when the 1 s +/// budget is spent, and `DeleteTable` returns `DELETING` at about 1 s. +async fn the_drops_of_one_delete_table_share_one_deadline(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, three_gsi_input("many_gsis")).await; + let table_id = key_info(&s, "many_gsis").await.table_id; + // The drop order: the base table, then the GSIs in the order the + // engine reads their ids. + let index_ids: Vec = + sqlx::query_scalar("SELECT index_id FROM indexes WHERE table_id = $1") + .bind(&table_id) + .fetch_all(&s.catalog) + .await + .expect("index ids"); + assert_eq!(index_ids.len(), 3, "precondition: three GSIs"); + let relations: Vec = std::iter::once(data_table_name(&table_id)) + .chain(index_ids.iter().map(|id| index_table_name(id))) + .collect(); + let mut holders = Vec::new(); + for relation in &relations { + holders.push(hold_relation_in(&s, relation, "ROW EXCLUSIVE").await); + } + + let started = std::time::Instant::now(); + let release = async { + for holder in holders { + tokio::time::sleep(std::time::Duration::from_millis(800)).await; + release_relation(holder).await; + } + }; + let delete = async { + let deleted = tokio::time::timeout( + std::time::Duration::from_secs(10), + s.engine.delete_table(ACCOUNT, delete_input("many_gsis")), + ) + .await; + (deleted, started.elapsed()) + }; + let ((deleted, elapsed), ()) = tokio::join!(delete, release); + let deleted = deleted + .expect("DeleteTable returned within 10 s") + .expect("DeleteTable"); + eprintln!("staggered holders: DeleteTable returned after {elapsed:?}"); + assert!( + elapsed < std::time::Duration::from_millis(1600), + "DeleteTable took {elapsed:?}: its drops waited once per relation instead of sharing \ + one 1 s budget" + ); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert_eq!( + catalog_status(&s, "many_gsis").await.as_deref(), + Some("DELETING"), + "the budget ran out with holders left, so the drop was deferred" + ); + pump_until_removed(&s, "many_gsis", &relations).await; +} +on_both_layouts!(the_drops_of_one_delete_table_share_one_deadline); + +/// The pids of sessions waiting (not granted) for a lock on one of +/// `relations` on the data database. +async fn lock_waiters(s: &Scratch, relations: &[String]) -> Vec { + let data = s.data_pool().await; + let pids = sqlx::query_scalar( + "SELECT pid FROM pg_locks WHERE locktype = 'relation' AND NOT granted \ + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) \ + AND relation = ANY (SELECT to_regclass(n)::oid FROM unnest($1::text[]) AS n)", + ) + .bind(relations) + .fetch_all(&data) + .await + .expect("lock waiters"); + data.close().await; + pids +} + +/// Four tables left `DELETING` by readers, the production poller running. +struct HeldDeletes { + tables: Vec, + relations: Vec, + holders: Vec, +} + +async fn four_held_deletes(s: &Scratch) -> HeldDeletes { + let tables: Vec = (0..4).map(|i| format!("held{i}")).collect(); + let mut relations = Vec::new(); + let mut holders = Vec::new(); + for table in &tables { + create_active(s, plain_input(table)).await; + let relation = data_table_name(&key_info(s, table).await.table_id); + holders.push(hold_relation(s, &relation).await); + let deleted = delete_within_5s(s, table).await.expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + relations.push(relation); + } + HeldDeletes { + tables, + relations, + holders, + } +} + +/// With four tables `DELETING` behind readers, the poller probes each one +/// for holders and skips it without waiting, so a pass costs a few queries +/// however many tables are held, rather than a lock timeout per table: no +/// drop ever queues for a held table's lock, and a table whose reader ends +/// is gone within about `active_poll` (1 s) of the release. +async fn held_deleting_tables_are_probed_not_waited_on(shared: bool) { + let s = scratch_with(shared).await; + let HeldDeletes { + tables, + relations, + mut holders, + } = four_held_deletes(&s).await; + let _poller = s.engine.spawn_control_plane_poller_for_tests(); + // Two passes over the four tables: the poller is retrying them. + wait_for_drop_attempt(&s.engine, 3).await; + let watch_until = std::time::Instant::now() + std::time::Duration::from_secs(2); + while std::time::Instant::now() < watch_until { + let waiters = lock_waiters(&s, &relations).await; + assert!( + waiters.is_empty(), + "the control plane queued for a held table's lock (pids {waiters:?}) instead of \ + skipping it" + ); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + let attempts = s.engine.control_plane_drop_attempts_for_tests(); + release_relation(holders.remove(0)).await; + let released_at = std::time::Instant::now(); + let data = s.data_pool().await; + loop { + let gone: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NULL") + .bind(&relations[0]) + .fetch_one(&data) + .await + .expect("probe"); + if gone && catalog_status(&s, &tables[0]).await.is_none() { + break; + } + assert!( + released_at.elapsed() < std::time::Duration::from_secs(2), + "{} not removed within 2 s of its reader's release with three other tables held", + tables[0] + ); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + data.close().await; + eprintln!( + "four held tables: {attempts} attempts before the release; the released table was \ + removed {:?} after it", + released_at.elapsed() + ); + for table in &tables[1..] { + assert_eq!( + catalog_status(&s, table).await.as_deref(), + Some("DELETING"), + "{table} is still held" + ); + } + for holder in holders { + release_relation(holder).await; + } +} +on_both_layouts!(held_deleting_tables_are_probed_not_waited_on); + +/// A `CreateTable` reaches `ACTIVE` at its usual pace while other tables' +/// drops are deferred: the passes that retry them do not wait on their +/// holders, so the next pass after the table is due comes about +/// `active_poll` (1 s) later. +async fn create_table_is_not_held_up_by_deferred_drops(shared: bool) { + let s = scratch_with(shared).await; + let held = four_held_deletes(&s).await; + let _poller = s.engine.spawn_control_plane_poller_for_tests(); + wait_for_drop_attempt(&s.engine, 3).await; + sqlx::query("UPDATE settings SET value = '0.25' WHERE key = 'control_plane_delay_seconds'") + .execute(&s.catalog) + .await + .expect("a delay, so CreateTable goes through the control plane"); + let started = std::time::Instant::now(); + let created = s + .engine + .create_table(ACCOUNT, plain_input("fresh")) + .await + .expect("create"); + assert_eq!( + created.table_status, + extenddb_core::types::TableStatus::Creating, + "precondition: the table waits for the control plane" + ); + while catalog_status(&s, "fresh").await.as_deref() != Some("ACTIVE") { + assert!( + started.elapsed() < std::time::Duration::from_secs(2), + "CreateTable was not ACTIVE within 2 s with four drops deferred" + ); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + eprintln!( + "CreateTable ACTIVE after {:?} with four drops deferred", + started.elapsed() + ); + for table in &held.tables { + assert_eq!( + catalog_status(&s, table).await.as_deref(), + Some("DELETING"), + "precondition: {table} is still held" + ); + } + for holder in held.holders { + release_relation(holder).await; + } +} +on_both_layouts!(create_table_is_not_held_up_by_deferred_drops); + +/// `DeleteTable` of a table that is already `DELETING` (a retried or +/// duplicated request while a reader holds it) returns the table's +/// description with `TableStatus: DELETING`, as the service does, and +/// changes nothing: the row and its scheduled removal stay as they were. +async fn delete_table_of_a_deleting_table_returns_its_description(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("twice")).await; + let relations = base_and_gsi_relations(&s, "twice").await; + let holder = hold_relation(&s, &relations[0]).await; + let first = delete_within_5s(&s, "twice").await.expect("first delete"); + assert_eq!( + first.table_status, + extenddb_core::types::TableStatus::Deleting + ); + let scheduled = || async { + sqlx::query_scalar::<_, String>( + "SELECT table_status || ' ' || status_transition_at::text FROM tables \ + WHERE table_name = 'twice'", + ) + .fetch_one(&s.catalog) + .await + .expect("row") + }; + let before = scheduled().await; + assert!(before.starts_with("DELETING "), "precondition: {before}"); + + let second = delete_within_5s(&s, "twice") + .await + .expect("DeleteTable of a DELETING table returns its description"); + assert_eq!( + second.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert_eq!(second.table_arn, first.table_arn); + assert_eq!(second.table_id, first.table_id); + assert_eq!(second.key_schema, first.key_schema); + assert_eq!( + second + .global_secondary_indexes + .iter() + .flatten() + .map(|g| g.index_name.clone()) + .collect::>(), + ["gi"], + "the description lists the indexes" + ); + assert_eq!( + scheduled().await, + before, + "the second DeleteTable changed nothing" + ); + release_relation(holder).await; + pump_until_removed(&s, "twice", &relations).await; +} +on_both_layouts!(delete_table_of_a_deleting_table_returns_its_description); + +/// A drop deferred by a reader is due at once, whatever +/// `control_plane_delay_seconds` is (below 1, so `DeleteTable` tried to drop +/// in the request): the table is already being removed, and a reader that +/// ends right after the `DeleteTable` has the next pass drop it rather than a +/// pass after the delay. Here the reader ends before the poller starts, and +/// the poller's first pass, a few milliseconds later and well inside the +/// 0.9 s delay, removes the table; with the drop scheduled after the delay +/// that pass finds nothing due and the next comes 10 s later. +async fn a_deferred_drop_is_due_at_once(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, plain_input("due")).await; + let relation = data_table_name(&key_info(&s, "due").await.table_id); + sqlx::query("UPDATE settings SET value = '0.9' WHERE key = 'control_plane_delay_seconds'") + .execute(&s.catalog) + .await + .expect("setting"); + let holder = hold_relation(&s, &relation).await; + let deleted = delete_within_5s(&s, "due").await.expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + let due_now: bool = sqlx::query_scalar( + "SELECT status_transition_at <= NOW() FROM tables WHERE table_name = 'due'", + ) + .fetch_one(&s.catalog) + .await + .expect("row"); + assert!( + due_now, + "the deferred drop is scheduled after control_plane_delay_seconds" + ); + release_relation(holder).await; + let _poller = s + .engine + .spawn_control_plane_poller_with_active_poll_for_tests(WAKE_TEST_ACTIVE_POLL); + removed_within_3s( + &s, + "due", + &[relation], + std::time::Instant::now(), + "deferred drop", + ) + .await; +} +on_both_layouts!(a_deferred_drop_is_due_at_once); + +/// An engine on `s`'s databases whose connections take the settings +/// `ALTER DATABASE ... SET` gave the data database after `s` was created. +async fn engine_after_alter(s: &Scratch) -> PostgresEngine { + PostgresEngine::new( + &PostgresConfig { + connection_string: format!("{}/{}", s.base, s.names[0]), + pool_size: 8, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("engine") +} + +/// `ALTER DATABASE SET = `, for the +/// connections opened afterwards. +async fn alter_data_database(s: &Scratch, setting: &str, value: &str) { + let mut admin = PgConnection::connect(&format!("{}/postgres", s.base)) + .await + .expect("admin"); + sqlx::query(&format!( + "ALTER DATABASE \"{}\" SET {setting} = '{value}'", + s.data_name() + )) + .execute(&mut admin) + .await + .expect("alter database"); + admin.close().await.expect("close admin"); +} + +/// A `statement_timeout` shorter than the drop's lock timeout (set here on +/// the data database, as an administrator might set it for the role or the +/// database) ends the drop's wait with 57014 rather than 55P03. That is a +/// lock wait like the other, and `DeleteTable` takes its asynchronous path +/// rather than failing. +async fn a_statement_timeout_sends_delete_table_down_the_async_path(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, plain_input("timed")).await; + let relation = data_table_name(&key_info(&s, "timed").await.table_id); + alter_data_database(&s, "statement_timeout", "300ms").await; + let engine = engine_after_alter(&s).await; + let holder = hold_relation_in(&s, &relation, "ROW EXCLUSIVE").await; + let deleted = tokio::time::timeout( + std::time::Duration::from_secs(5), + engine.delete_table(ACCOUNT, delete_input("timed")), + ) + .await + .expect("DeleteTable returned within 5 s") + .unwrap_or_else(|e| panic!("a statement timeout failed DeleteTable: {e}")); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert_eq!( + catalog_status(&s, "timed").await.as_deref(), + Some("DELETING") + ); + release_relation(holder).await; + pump_until_removed(&s, "timed", std::slice::from_ref(&relation)).await; +} +on_both_layouts!(a_statement_timeout_sends_delete_table_down_the_async_path); + +/// A `DeleteTable` chosen as a deadlock victim (40P01) takes its +/// asynchronous path rather than failing. The cycle: a session holds the +/// base data table `ROW EXCLUSIVE`; the drop deletes the table's pending +/// queue rows (`ROW EXCLUSIVE` on `gsi_pending`) and waits on the base +/// table; the session then asks for `gsi_pending` in `SHARE` mode, which +/// waits on the drop. The drop's connections have a 400 ms +/// `deadlock_timeout` and the session's is 10 s, so the drop runs the +/// deadlock check, well before its 1 s lock timeout, and is the victim. +async fn a_deadlock_victim_delete_table_takes_the_async_path(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, plain_input("victim")).await; + let relation = data_table_name(&key_info(&s, "victim").await.table_id); + alter_data_database(&s, "deadlock_timeout", "400ms").await; + let engine = engine_after_alter(&s).await; + let mut holder = PgConnection::connect(&s.data_url()).await.expect("holder"); + for statement in [ + "SET deadlock_timeout = '10s'".to_owned(), + "BEGIN".to_owned(), + format!("LOCK TABLE {relation} IN ROW EXCLUSIVE MODE"), + ] { + sqlx::query(&statement) + .execute(&mut holder) + .await + .expect("holder"); + } + let delete = tokio::time::timeout( + std::time::Duration::from_secs(5), + engine.delete_table(ACCOUNT, delete_input("victim")), + ); + let close_the_cycle = async { + let mut watcher = PgConnection::connect(&format!("{}/postgres", s.base)) + .await + .expect("watcher"); + let mut waiting = 0; + for _ in 0..200 { + waiting = sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity WHERE datname = $1 \ + AND wait_event_type = 'Lock' AND query LIKE 'DROP TABLE%'", + ) + .bind(s.data_name()) + .fetch_one(&mut watcher) + .await + .expect("waiters"); + if waiting == 1 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + watcher.close().await.expect("close watcher"); + assert_eq!(waiting, 1, "precondition: the drop waits on the holder"); + let started = std::time::Instant::now(); + sqlx::query("LOCK TABLE gsi_pending IN SHARE MODE") + .execute(&mut holder) + .await + .expect("the holder's lock is granted once the drop is gone"); + started.elapsed() + }; + let (deleted, waited) = tokio::join!(delete, close_the_cycle); + let deleted = deleted + .expect("DeleteTable returned within 5 s") + .unwrap_or_else(|e| panic!("a deadlock failed DeleteTable: {e}")); + eprintln!("the holder waited {waited:?} for the drop to give way"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert!( + waited < std::time::Duration::from_millis(900), + "the drop gave way after {waited:?}: it was not the deadlock victim but ran out its \ + lock timeout" + ); + assert_eq!( + catalog_status(&s, "victim").await.as_deref(), + Some("DELETING") + ); + sqlx::query("ROLLBACK") + .execute(&mut holder) + .await + .expect("release"); + holder.close().await.expect("close holder"); + pump_until_removed(&s, "victim", std::slice::from_ref(&relation)).await; +} +on_both_layouts!(a_deadlock_victim_delete_table_takes_the_async_path); + +/// A reader found by the in-use probe sends `DeleteTable` down the +/// asynchronous path without trying the drop, on both layouts: it returns in +/// well under the 1 s lock timeout a drop would wait out, and queues no +/// `ACCESS EXCLUSIVE` request that would stall the table's new readers. +async fn a_probed_reader_sends_delete_table_down_the_async_path_at_once(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, plain_input("probed")).await; + let relation = data_table_name(&key_info(&s, "probed").await.table_id); + let holder = hold_relation(&s, &relation).await; + let started = std::time::Instant::now(); + let deleted = delete_within_5s(&s, "probed").await.expect("delete"); + let elapsed = started.elapsed(); + eprintln!("probed reader: DeleteTable returned after {elapsed:?}"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert!( + elapsed < std::time::Duration::from_millis(500), + "DeleteTable took {elapsed:?}: it tried the drop instead of seeing the reader" + ); + release_relation(holder).await; + pump_until_removed(&s, "probed", std::slice::from_ref(&relation)).await; +} +on_both_layouts!(a_probed_reader_sends_delete_table_down_the_async_path_at_once); + +/// A queue worker's claim on one of the table's pending GSI rows holds +/// `DeleteTable` for as long as the claim lasts, with no timeout, on both +/// layouts: the pending-row delete waits for it before any lock timeout is +/// set, and nothing commits meanwhile, so the row and every data table stay +/// until the claim ends. Then the drops run and the table is removed in the +/// request. The claim is held by hand (the pending row `FOR UPDATE`, the base +/// row `FOR SHARE`, as a worker mid-apply holds them), so the wait is what is +/// tested. +async fn a_claimed_queue_row_holds_delete_table_until_it_ends(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("claimed")).await; + let relations = base_and_gsi_relations(&s, "claimed").await; + let table_id = key_info(&s, "claimed").await.table_id; + sqlx::query("UPDATE settings SET value = '600000' WHERE key = 'index_propagation_delay_ms'") + .execute(&s.catalog) + .await + .expect("hold queue rows back so one can be claimed by hand"); + snapshot_conformance::populate(&s.engine, ACCOUNT, "claimed", 1).await; + let mut claim = PgConnection::connect(&s.data_url()).await.expect("claim"); + sqlx::query("BEGIN") + .execute(&mut claim) + .await + .expect("begin"); + sqlx::query("SELECT id FROM gsi_pending WHERE table_id = $1 ORDER BY id LIMIT 1 FOR UPDATE") + .bind(&table_id) + .fetch_one(&mut claim) + .await + .expect("a pending row to claim"); + sqlx::query(&format!( + "SELECT 1 FROM {} WHERE pk = 'p000' FOR SHARE", + relations[0] + )) + .fetch_one(&mut claim) + .await + .expect("share-lock the base row"); + + let mut delete = Box::pin(s.engine.delete_table(ACCOUNT, delete_input("claimed"))); + let early = tokio::time::timeout(std::time::Duration::from_millis(2500), &mut delete).await; + assert!( + early.is_err(), + "DeleteTable returned while a claim for the table was held: {early:?}" + ); + assert_eq!( + catalog_status(&s, "claimed").await.as_deref(), + Some("ACTIVE"), + "nothing commits while the claim is in flight" + ); + assert_eq!(present_relations(&s, &relations).await, relations); + + sqlx::query("COMMIT") + .execute(&mut claim) + .await + .expect("end the claim"); + let deleted = tokio::time::timeout(std::time::Duration::from_secs(10), &mut delete) + .await + .expect("DeleteTable finished once the claim ended") + .expect("delete"); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert_eq!( + catalog_status(&s, "claimed").await, + None, + "removed in the request" + ); + assert!(present_relations(&s, &relations).await.is_empty()); + claim.close().await.expect("close claim"); +} +on_both_layouts!(a_claimed_queue_row_holds_delete_table_until_it_ends); + +/// `DeleteTable` drops the base data table before the index tables, the +/// order a write takes its locks in (the base row, then an LSI or a +/// synchronous GSI in the same transaction). A write that holds the base +/// when the drops start, and then writes the GSI table, is therefore not +/// held up: the drop waits for it on the base, holding nothing the write +/// needs, and once it commits the table is removed in the request. With the +/// index tables dropped first the drop would hold the GSI table while it +/// waits on the base, the write would wait on the drop, and the two would +/// wait for each other until the drops' 1 s budget ran out. +async fn a_write_in_progress_finishes_before_delete_table_drops(shared: bool) { + let s = scratch_with(shared).await; + create_active(&s, snapshot_conformance::table_input("written2")).await; + let relations = base_and_gsi_relations(&s, "written2").await; + let mut writer = hold_relation_in(&s, &relations[0], "ROW EXCLUSIVE").await; + let delete = async { + let started = std::time::Instant::now(); + let deleted = tokio::time::timeout( + std::time::Duration::from_secs(5), + s.engine.delete_table(ACCOUNT, delete_input("written2")), + ) + .await; + (deleted, started.elapsed()) + }; + let write_the_index = async { + let mut watcher = PgConnection::connect(&format!("{}/postgres", s.base)) + .await + .expect("watcher"); + let mut waiting = 0; + for _ in 0..400 { + waiting = sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity WHERE datname = $1 \ + AND wait_event_type = 'Lock' AND query LIKE 'DROP TABLE%'", + ) + .bind(s.data_name()) + .fetch_one(&mut watcher) + .await + .expect("waiters"); + if waiting == 1 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + watcher.close().await.expect("close watcher"); + assert_eq!(waiting, 1, "precondition: a drop waits on the writer"); + let started = std::time::Instant::now(); + sqlx::query(&format!( + "LOCK TABLE {} IN ROW EXCLUSIVE MODE", + relations[1] + )) + .execute(&mut writer) + .await + .expect("the writer's index write"); + let waited = started.elapsed(); + sqlx::query("COMMIT") + .execute(&mut writer) + .await + .expect("commit"); + waited + }; + let ((deleted, took), waited) = tokio::join!(delete, write_the_index); + let deleted = deleted + .expect("DeleteTable returned within 5 s") + .expect("delete"); + eprintln!("the writer's index write waited {waited:?}; DeleteTable took {took:?}"); + assert!( + waited < std::time::Duration::from_millis(500), + "the writer's index write waited {waited:?} on DeleteTable's drops" + ); + assert_eq!( + deleted.table_status, + extenddb_core::types::TableStatus::Deleting + ); + assert_eq!( + catalog_status(&s, "written2").await, + None, + "the drops waited for the write and removed the table in the request" + ); + assert!(present_relations(&s, &relations).await.is_empty()); + writer.close().await.expect("close writer"); +} +on_both_layouts!(a_write_in_progress_finishes_before_delete_table_drops); diff --git a/crates/storage-sqlite/Cargo.toml b/crates/storage-sqlite/Cargo.toml index a750bf4b8..09447140d 100644 --- a/crates/storage-sqlite/Cargo.toml +++ b/crates/storage-sqlite/Cargo.toml @@ -13,6 +13,9 @@ license.workspace = true # no file on disk). Without this feature the backend is file-backed and the path # must be configured as usual. memory = [] +# Compiles the snapshot page-size setter the consistency test uses and the +# shared conformance check. Production builds do not compile them. +test-hooks = ["extenddb-storage/test-hooks"] [dependencies] extenddb-core = { workspace = true } @@ -36,3 +39,8 @@ bigdecimal = { workspace = true } crc32fast = { workspace = true } zeroize = { workspace = true } libc = { workspace = true } + +[dev-dependencies] +# The crate's own tests need `test-hooks`; a dev-dependency on itself turns +# the feature on for test builds only. +extenddb-storage-sqlite = { path = ".", features = ["test-hooks"] } diff --git a/crates/storage-sqlite/src/backup.rs b/crates/storage-sqlite/src/backup.rs index 4cea564e2..fca131e5c 100644 --- a/crates/storage-sqlite/src/backup.rs +++ b/crates/storage-sqlite/src/backup.rs @@ -46,8 +46,9 @@ const BACKUP_SERVER_SLOTS_MAX: usize = 16; /// `LimitExceeded`. pub(crate) const BACKUP_SLOT_WAIT: std::time::Duration = std::time::Duration::from_secs(30); -/// Backups the whole process may run at once, across accounts: -/// `pool_size / 2 - 1`, clamped to `1..=16`. +/// Backups and table snapshots the whole process may run at once, across +/// accounts. The formula is stated once, on `SqliteConfig::pool_size` +/// (`config.rs`); this computes it. /// /// A file-backed backup holds one pooled connection for its whole duration /// (the WAL snapshot its items are read from) and its batches and final @@ -67,22 +68,28 @@ pub(crate) fn backup_server_slot_count(pool_size: u32) -> usize { /// Take a server slot for a backup, waiting at most `wait`; `LimitExceeded` /// when none frees up in time. -async fn acquire_backup_permit( +pub(crate) async fn acquire_backup_permit( slots: std::sync::Arc, slot_count: usize, wait: std::time::Duration, ) -> Result { tokio::time::timeout(wait, slots.acquire_owned()) .await - .map_err(|_| { - StorageError::LimitExceeded(format!( - "Too many backups are in progress on this server (at most {slot_count} at once); \ - retry later" - )) - })? + .map_err(|_| slots_exhausted(slot_count))? .map_err(|e| StorageError::Internal(format!("backup slot: {e}"))) } +/// The refusal when every server slot is held. Backups and table snapshots +/// (`snapshot_engine.rs`) share the slots, so the message names both: with +/// snapshots holding every slot a `CreateBackup` is refused with no backup +/// running. +pub(crate) fn slots_exhausted(slot_count: usize) -> StorageError { + StorageError::LimitExceeded(format!( + "Too many backups or table snapshots are in progress on this server (at most \ + {slot_count} at once); retry later" + )) +} + /// Read the next batch of `(cursor, item_data)` from `table` after /// `last_cursor`, restricted by `filter` (a fixed SQL condition with one `?` /// bound to `filter_arg`, or `None`). `cursor` is `rowid` for immutable source @@ -3500,10 +3507,7 @@ mod restore_tests { .expect("first"); match acquire_backup_permit(slots.clone(), 1, Duration::from_millis(50)).await { Err(StorageError::LimitExceeded(m)) => { - assert!( - m.contains("backups are in progress on this server (at most 1"), - "{m}" - ); + assert!(m.contains("in progress on this server (at most 1"), "{m}"); } other => panic!("a full bound refuses: {:?}", other.err()), } diff --git a/crates/storage-sqlite/src/config.rs b/crates/storage-sqlite/src/config.rs index 0c7f62454..01f69c0e2 100644 --- a/crates/storage-sqlite/src/config.rs +++ b/crates/storage-sqlite/src/config.rs @@ -16,13 +16,19 @@ use serde::Deserialize; /// two connections so a backup reader can coexist with its batched writer, and /// log at info when they do. A backup holds its reader connection for its whole /// duration, so `pool_size >= 3` keeps a running backup from contending with -/// ordinary requests for connections. The process runs at most -/// `pool_size / 2 - 1` backups at once (between 1 and 16; 4 for the default -/// of 10) and answers `LimitExceededException` to a CreateBackup that waited -/// 30 s for a slot. Running backups hold one connection each for their read -/// snapshot plus one shared writer (`n` backups hold `n + 1`), so from -/// `pool_size` 4 up a flood of backups leaves at least half the pool to other -/// requests; with 2 or 3 the one admitted backup holds two of them. +/// ordinary requests for connections. The process has `pool_size / 2 - 1` +/// server slots (integer division, at least 1, at most 16: 4 for the default +/// of 10, and 1 for every `pool_size` of 5 or less; this is the one statement +/// of the formula, `backup::backup_server_slot_count` computes it), shared by +/// backups and table snapshots: each running backup and each open snapshot +/// stream (file-backed or in-memory) holds one. A snapshot that finds no slot +/// free is refused at once; a `CreateBackup` waits 30 s for one and then +/// answers `LimitExceededException`, with no backup running if snapshots hold +/// every slot. With one slot, one held snapshot blocks every `CreateBackup` +/// for as long as it is open. Running backups hold one connection each for +/// their read snapshot plus one shared writer (`n` backups hold `n + 1`), so +/// from `pool_size` 4 up a flood of backups leaves at least half the pool to +/// other requests; with 2 or 3 the one admitted backup holds two of them. #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] pub struct SqliteConfig { diff --git a/crates/storage-sqlite/src/delete_table.rs b/crates/storage-sqlite/src/delete_table.rs index 72533abea..476177d6c 100644 --- a/crates/storage-sqlite/src/delete_table.rs +++ b/crates/storage-sqlite/src/delete_table.rs @@ -22,9 +22,11 @@ impl SqliteEngine { ) -> Result { Self::validate_account_id(account_id)?; + // A DELETING row is fetched too: DeleteTable of a table that is + // already being deleted returns its description, as the service does. let row: Option = sqlx::query_as(&format!( - "SELECT {TABLE_COLUMNS} FROM tables \ - WHERE account_id = ? AND table_name = ? AND table_status IN ('ACTIVE', 'CREATING')" + "SELECT {TABLE_COLUMNS} FROM tables WHERE account_id = ? AND table_name = ? \ + AND table_status IN ('ACTIVE', 'CREATING', 'DELETING')" )) .bind(account_id) .bind(&input.table_name) @@ -33,6 +35,22 @@ impl SqliteEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; let row = row.ok_or_else(|| StorageError::TableNotFound(input.table_name.clone()))?; + if row.table_status == "DELETING" { + // Already scheduled for removal by the control-plane worker: + // nothing to change. + let index_rows: Vec = sqlx::query_as(&format!( + "SELECT {INDEX_COLUMNS} FROM indexes WHERE table_id = ?" + )) + .bind(&row.table_id) + .fetch_all(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let desc = self.build_table_description_from_row(account_id, row, index_rows)?; + return Ok(TableDescription { + table_status: TableStatus::Deleting, + ..desc + }); + } if row.deletion_protection_enabled { return Err(StorageError::DeletionProtected(row.table_arn.clone())); } diff --git a/crates/storage-sqlite/src/lib.rs b/crates/storage-sqlite/src/lib.rs index 3953fc350..715f755b3 100644 --- a/crates/storage-sqlite/src/lib.rs +++ b/crates/storage-sqlite/src/lib.rs @@ -37,6 +37,7 @@ mod number_key; mod operations; mod schema; mod serve_lock; +mod snapshot_engine; mod sqlite_util; mod store; mod stream; diff --git a/crates/storage-sqlite/src/snapshot_engine.rs b/crates/storage-sqlite/src/snapshot_engine.rs new file mode 100644 index 000000000..f67679874 --- /dev/null +++ b/crates/storage-sqlite/src/snapshot_engine.rs @@ -0,0 +1,1094 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! `SnapshotEngine` for `SQLite`. +//! +//! Catalog and data share one file, so one read transaction covers both. In +//! WAL mode a read transaction sees the database as of its first read for +//! its whole life and never blocks writers; writers continue and the WAL +//! grows (it cannot checkpoint past the reader) until the transaction ends. +//! A checkpoint does not fail while a snapshot is held: `PRAGMA +//! wal_checkpoint(TRUNCATE)` waits out the busy timeout (5 s), then returns +//! a row with `busy = 1` and no frame written after the open checkpointed, +//! and the WAL file keeps every frame written since the open (19.9 MB after +//! 900 puts of small items in `wal_checkpoint_while_a_snapshot_is_held`). +//! Once the stream ends the next checkpoint truncates it. +//! +//! No lock is needed: one read transaction reads the `tables`, `indexes`, +//! `vector_indexes` and data rows, and its first read pins the WAL snapshot +//! for all of them. The engine's writer lock is not taken, so writers never +//! wait on a snapshot being opened, and the snapshot never waits on a pool +//! connection while holding the lock. The in-catalog backup +//! (`backup.rs`, `create_backup_from_wal_snapshot`) does hold the lock while +//! it pins its snapshot, because it also inserts the CREATING catalog row in +//! that step and pages by rowid; this path writes nothing and pages by the +//! primary key, so the lock buys it nothing. +//! +//! The snapshot holds one pool connection and one backup slot +//! (`backup_server_slots`, shared with the in-catalog backup) until the +//! stream ends or is dropped, so the slot arithmetic that keeps half the +//! pool free for requests counts snapshots too; when no slot is free the +//! open returns `LimitExceeded` at once rather than queueing. The slot count +//! is stated once, in `SqliteConfig`'s doc (`config.rs`). The other side of +//! sharing: a `CreateBackup` waits 30 s for a slot, so snapshots holding +//! every slot make every `CreateBackup` wait and then fail with the shared +//! refusal (`backup::slots_exhausted`), which names table snapshots for that +//! reason. +//! +//! In-memory databases are pinned to a single connection (a second connection +//! would open a different database), so a stream that held its transaction +//! would hold the only connection and every other request would wait for it. +//! On such a database the open reads every page in its one read transaction, +//! keeps the items in memory, and returns the stream over that copy: the +//! open holds the connection for the length of the read (other requests wait +//! for it, as they wait for `create_backup_single_connection`, and fail after +//! the pool's acquire timeout if the read takes longer), and no connection is +//! held once it returns. The copy's stream keeps the backup slot until it +//! ends or is dropped, so the slot bounds how many copies exist at once: the +//! copy is the table again in decoded form, so with `k` copies open the +//! process holds the table `k + 1` times, and `k` is at most the slot count. +//! The snapshot kind and the consistency are those of the file-backed path: +//! one read transaction covers catalog and data. +//! +//! Items are read in primary-key order through row-value keyset pagination +//! on the primary key columns (`WHERE (pk, sk_x) > (?, ?) ORDER BY pk, sk_x +//! LIMIT n`), which `SQLite` serves from the primary key index. Multi-part +//! range keys fall back to rowid order and report `key_ordered = false`. +//! +//! Memory: one page is `LIMIT paging::DEFAULT_FETCH_BATCH` (1000) rows, all +//! of them held as item text at once. At the default 400 KB item limit one +//! row's text can reach 27.7 MB (the measured worst case in #394's line-cap +//! table), so a page is bounded by rows, not bytes. A byte budget (keyset +//! paging with a byte cutoff) is required before the stream gets a caller; +//! the backup design assigns it to the backup job PR. + +use extenddb_core::types::{Item, ScalarAttributeType, TableDescription}; +use extenddb_storage::error::StorageError; +use extenddb_storage::snapshot::{drop_deleting_indexes, key_schema_is_orderable, paging}; +use extenddb_storage::util::sk_column; +use extenddb_storage::{SnapshotEngine, SnapshotInfo, TableSnapshot}; +use futures::StreamExt as _; +use futures::future::BoxFuture; +use futures::stream::BoxStream; +use sqlx::{Row, Sqlite, Transaction}; + +use crate::data::{all_sort_key_info, data_table_name}; +use crate::store::SqliteEngine; + +fn db_err(e: &sqlx::Error) -> StorageError { + StorageError::Internal(format!("Database error: {e}")) +} + +impl SnapshotEngine for SqliteEngine { + fn open_table_snapshot( + &self, + account_id: &str, + table_name: &str, + ) -> BoxFuture<'_, Result> { + let account_id = account_id.to_owned(); + let table_name = table_name.to_owned(); + Box::pin(async move { + self.open_table_snapshot_impl(&account_id, &table_name) + .await + }) + } +} + +/// How the stream pages through the table. +enum Paging { + /// Hash-only key: `pk > ?`. + Pk, + /// One range key stored in `sk_col`: `(pk, sk_col) > (?, ?)`. + PkSk { + sk_col: &'static str, + sk_type: ScalarAttributeType, + }, + /// Multi-part range key: rowid order. + Rowid, +} + +impl SqliteEngine { + async fn open_table_snapshot_impl( + &self, + account_id: &str, + table_name: &str, + ) -> Result { + // A snapshot holds one pool connection for the stream's life, the + // same resource an in-catalog backup's reader holds, so it takes a + // backup slot: the slot arithmetic in `backup.rs` keeps at least half + // the pool free for requests, and without a slot enough open snapshots + // would hold every connection and every writer would wait out the + // pool's acquire timeout under the engine write lock. Refuse rather + // than queue; the caller bounds concurrency. + let permit = std::sync::Arc::clone(&self.backup_server_slots) + .try_acquire_owned() + .map_err(|_| crate::backup::slots_exhausted(self.backup_server_slot_count))?; + + let mut tx: Transaction<'static, Sqlite> = + self.pool.begin().await.map_err(|e| db_err(&e))?; + // This first read pins the WAL snapshot; everything below, catalog + // and data, is read from it. + let row: Option<(String, String)> = sqlx::query_as( + "SELECT table_id, table_status FROM tables WHERE account_id = ? AND table_name = ?", + ) + .bind(account_id) + .bind(table_name) + .fetch_optional(&mut *tx) + .await + .map_err(|e| db_err(&e))?; + let (table_id, status) = + row.ok_or_else(|| StorageError::TableNotFound(table_name.to_owned()))?; + if status != "ACTIVE" { + return Err(StorageError::TableNotActive(table_name.to_owned())); + } + let mut description = self + .build_table_description_on(&mut tx, account_id, table_name) + .await?; + drop_deleting_indexes(&mut description); + // Both stages fall here: the first read above pinned the snapshot for + // catalog and data alike, and no page is read until the stream is + // polled. + paging::open_pause(paging::OpenStage::AfterDefinition).await; + let ddb_table = data_table_name(&table_id); + + let mode = paging_for(&description); + let key_ordered = !matches!(mode, Paging::Rowid); + paging::open_pause(paging::OpenStage::BeforeItems).await; + let items = if self.in_memory { + // The only connection must not stay with the stream (module + // doc): read every page now, in the one transaction, and serve + // the stream from the copy. Draining the inner stream ends the + // transaction; the slot moves to the copy's stream, so the slot + // bounds how many copies exist at once. + let mut inner = keyset_stream(tx, None, ddb_table, mode); + let mut copy = Vec::new(); + while let Some(item) = inner.next().await { + copy.push(item?); + } + drop(inner); + copy_stream(copy, permit) + } else { + keyset_stream(tx, Some(permit), ddb_table, mode) + }; + Ok(TableSnapshot { + description, + info: SnapshotInfo { + kind: SnapshotInfo::KIND_SQLITE_WAL_READ_TRANSACTION.to_owned(), + // A WAL read transaction has no addressable position, so + // there is nothing for a later change-log archive to anchor + // on; the marker is for that and for nothing else (no + // backend reopens a snapshot from it). + marker: String::new(), + }, + key_ordered, + items, + }) + } +} + +/// `pk` is TEXT under BINARY collation and the sort column is typed (text, +/// order-preserving number encoding, blob), so paging on the primary key +/// yields the order `extenddb_core::backup_format::key_order` defines. +fn paging_for(description: &TableDescription) -> Paging { + if !key_schema_is_orderable(description) { + return Paging::Rowid; + } + let sks = all_sort_key_info(&description.key_schema, &description.attribute_definitions); + match sks.as_slice() { + [] => Paging::Pk, + [(_, sk_type)] => Paging::PkSk { + sk_col: sk_column(*sk_type), + sk_type: *sk_type, + }, + _ => Paging::Rowid, + } +} + +/// Last-seen key for keyset pagination. +enum Cursor { + Start, + Pk(String), + PkSkText(String, String), + PkSkBlob(String, Vec), + Rowid(i64), +} + +/// The in-memory copy as a stream that owns the backup slot until it has +/// yielded its last item (or is dropped), so a consumer that stops at the +/// item count and keeps the stream does not keep the slot. +fn copy_stream( + copy: Vec, + permit: tokio::sync::OwnedSemaphorePermit, +) -> BoxStream<'static, Result> { + let permit = (!copy.is_empty()).then_some(permit); + Box::pin(futures::stream::unfold( + (copy.into_iter(), permit), + |(mut items, mut permit)| async move { + let item = items.next()?; + if items.len() == 0 { + permit = None; + } + Some((Ok(item), (items, permit))) + }, + )) +} + +/// Page through `table` in one read transaction. `permit` is the backup slot +/// the stream holds until its transaction ends (`None` for the in-memory +/// copy's inner read, whose slot the copy's stream holds): at the short page +/// that shows the end, at the first error, or when the stream is dropped, +/// whichever comes first. A consumer that stops at the last item or at the +/// first `Err` and keeps the stream therefore holds no slot a `CreateBackup` +/// is waiting for. +fn keyset_stream( + tx: Transaction<'static, Sqlite>, + permit: Option, + table: String, + mode: Paging, +) -> BoxStream<'static, Result> { + struct State { + tx: Option>, + /// The backup slot, released with the transaction. + permit: Option, + table: String, + paging: Paging, + cursor: Cursor, + buffer: std::vec::IntoIter, + done: bool, + } + let state = State { + tx: Some(tx), + permit, + table, + paging: mode, + cursor: Cursor::Start, + buffer: Vec::new().into_iter(), + done: false, + }; + // Read the page size once so a change to the test knob during a stream + // cannot end it early or late. + let batch = paging::fetch_batch(); + let limit = i64::try_from(batch).unwrap_or(i64::MAX); + Box::pin(futures::stream::unfold( + state, + move |mut state| async move { + loop { + if let Some(item) = state.buffer.next() { + return Some((Ok(item), state)); + } + if state.done { + return None; + } + let tx = state.tx.as_mut()?; + paging::record_fetch(); + let page = fetch_page(tx, &state.table, &state.paging, &state.cursor, limit).await; + let (items, next_cursor) = match page { + Ok(page) => page, + Err(e) => { + state.done = true; + state.tx = None; + state.permit = None; + return Some((Err(e), state)); + } + }; + if let Some(next) = next_cursor { + state.cursor = next; + } + if items.len() < batch { + state.done = true; + if let Some(tx) = state.tx.take() { + let _ = tx.commit().await; + } + state.permit = None; + } + state.buffer = items.into_iter(); + } + }, + )) +} + +/// Query one page of up to `limit` rows for the paging mode and cursor. +async fn query_page( + tx: &mut Transaction<'static, Sqlite>, + table: &str, + paging: &Paging, + cursor: &Cursor, + limit: i64, +) -> Result, StorageError> { + match (paging, cursor) { + (Paging::Pk, Cursor::Start) => { + sqlx::query(&format!( + "SELECT pk, item_data FROM {table} ORDER BY pk LIMIT ?" + )) + .bind(limit) + .fetch_all(&mut **tx) + .await + } + (Paging::Pk, Cursor::Pk(last)) => { + sqlx::query(&format!( + "SELECT pk, item_data FROM {table} WHERE pk > ? ORDER BY pk LIMIT ?" + )) + .bind(last) + .bind(limit) + .fetch_all(&mut **tx) + .await + } + (Paging::PkSk { sk_col, .. }, Cursor::Start) => { + sqlx::query(&format!( + "SELECT pk, {sk_col}, item_data FROM {table} ORDER BY pk, {sk_col} LIMIT ?" + )) + .bind(limit) + .fetch_all(&mut **tx) + .await + } + (Paging::PkSk { sk_col, .. }, Cursor::PkSkText(pk, sk)) => { + sqlx::query(&format!( + "SELECT pk, {sk_col}, item_data FROM {table} \ + WHERE (pk, {sk_col}) > (?, ?) ORDER BY pk, {sk_col} LIMIT ?" + )) + .bind(pk) + .bind(sk) + .bind(limit) + .fetch_all(&mut **tx) + .await + } + (Paging::PkSk { sk_col, .. }, Cursor::PkSkBlob(pk, sk)) => { + sqlx::query(&format!( + "SELECT pk, {sk_col}, item_data FROM {table} \ + WHERE (pk, {sk_col}) > (?, ?) ORDER BY pk, {sk_col} LIMIT ?" + )) + .bind(pk) + .bind(sk) + .bind(limit) + .fetch_all(&mut **tx) + .await + } + (Paging::Rowid, Cursor::Start) => { + sqlx::query(&format!( + "SELECT rowid, item_data FROM {table} ORDER BY rowid LIMIT ?" + )) + .bind(limit) + .fetch_all(&mut **tx) + .await + } + (Paging::Rowid, Cursor::Rowid(last)) => { + sqlx::query(&format!( + "SELECT rowid, item_data FROM {table} WHERE rowid > ? ORDER BY rowid LIMIT ?" + )) + .bind(last) + .bind(limit) + .fetch_all(&mut **tx) + .await + } + _ => { + return Err(StorageError::Internal( + "snapshot cursor does not match its paging mode".to_owned(), + )); + } + } + .map_err(|e| db_err(&e)) +} + +/// One page: items and the cursor for the next page. Fewer items than +/// `limit` means the table is exhausted. +async fn fetch_page( + tx: &mut Transaction<'static, Sqlite>, + table: &str, + paging: &Paging, + cursor: &Cursor, + limit: i64, +) -> Result<(Vec, Option), StorageError> { + let rows = query_page(tx, table, paging, cursor, limit).await?; + let mut items = Vec::with_capacity(rows.len()); + let mut next = None; + for row in &rows { + let data: String = row.try_get("item_data").map_err(|e| db_err(&e))?; + let item: Item = serde_json::from_str(&data) + .map_err(|e| StorageError::Internal(format!("Parse item: {e}")))?; + items.push(item); + } + if let Some(last) = rows.last() { + next = Some(match paging { + Paging::Pk => Cursor::Pk(last.try_get("pk").map_err(|e| db_err(&e))?), + Paging::PkSk { sk_col, sk_type } => { + let pk: String = last.try_get("pk").map_err(|e| db_err(&e))?; + match sk_type { + ScalarAttributeType::B => Cursor::PkSkBlob( + pk, + last.try_get::, _>(*sk_col) + .map_err(|e| db_err(&e))?, + ), + ScalarAttributeType::S | ScalarAttributeType::N => Cursor::PkSkText( + pk, + last.try_get::(*sk_col).map_err(|e| db_err(&e))?, + ), + } + } + Paging::Rowid => Cursor::Rowid(last.try_get("rowid").map_err(|e| db_err(&e))?), + }); + } + Ok((items, next)) +} + +#[cfg(test)] +mod tests { + use extenddb_storage::snapshot_conformance::{self, PageSize}; + use extenddb_storage::{SnapshotEngine, TableEngine}; + use futures::StreamExt; + + use super::*; + + const ACCOUNT: &str = "123456789012"; + /// SQL backend: ordered keys, numbers to the service's 38 digits; the + /// open's read transaction makes no `UpdateTable` wait. + const EXPECT: snapshot_conformance::Expectations = snapshot_conformance::Expectations { + key_ordered: true, + numeric_key_digits: 38, + update_table_waits_at: &[], + }; + + async fn prepare(engine: &SqliteEngine) { + crate::schema::apply(&engine.pool).await.expect("schema"); + for key in ["control_plane_delay_seconds", "index_propagation_delay_ms"] { + sqlx::query("UPDATE settings SET value = '0' WHERE key = ?") + .bind(key) + .execute(&engine.pool) + .await + .expect("setting"); + } + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES (?, 'default')") + .bind(ACCOUNT) + .execute(&engine.pool) + .await + .expect("account"); + } + + /// A file-backed engine on a scratch database under the cargo target + /// directory. The files are removed when the value is dropped, on success + /// or on a failed assertion. + struct FileEngine { + engine: Option, + path: std::path::PathBuf, + } + + impl std::ops::Deref for FileEngine { + type Target = SqliteEngine; + fn deref(&self) -> &SqliteEngine { + self.engine.as_ref().expect("engine present until drop") + } + } + + impl Drop for FileEngine { + fn drop(&mut self) { + drop(self.engine.take()); + for suffix in ["", "-wal", "-shm"] { + let _ = std::fs::remove_file(format!("{}{suffix}", self.path.display())); + } + } + } + + /// A file-backed engine on a scratch database under the cargo target + /// directory, found from this test binary's own path + /// (`//deps/`), so a relative `CARGO_TARGET_DIR` + /// or a `--target-dir` on the command line, neither of which the test + /// process can resolve from its environment, is honoured too. The files + /// go in a `snapshot-scratch` directory beside the profile directories, + /// which the test creates. + async fn file_engine() -> FileEngine { + let exe = std::env::current_exe().expect("the test binary's path"); + let target = exe + .ancestors() + .nth(3) + .unwrap_or_else(|| panic!("{} is not //deps/", exe.display())); + let scratch = target.join("snapshot-scratch"); + // Owner-only, like the database files the engine creates in it (0600), + // so another local user cannot list the scratch databases either. + let mut dir = std::fs::DirBuilder::new(); + dir.recursive(true); + #[cfg(unix)] + std::os::unix::fs::DirBuilderExt::mode(&mut dir, 0o700); + dir.create(&scratch) + .unwrap_or_else(|e| panic!("create {}: {e}", scratch.display())); + let path = scratch.join(format!("snapshot-test-{}.sqlite", uuid::Uuid::new_v4())); + let engine = SqliteEngine::new(path.to_str().expect("UTF-8 path"), 4, "us-east-1", 409_600) + .await + .expect("file engine"); + prepare(&engine).await; + FileEngine { + engine: Some(engine), + path, + } + } + + /// Every test here opens a snapshot or sets the shared page size, so each + /// holds `snapshot_conformance::hooks_lock` for its whole body: the knob, + /// the round-trip counter and the open pause are process-global, and a + /// test that opened a snapshot outside the lock could take a pause the + /// conformance check armed. + async fn hooks() -> snapshot_conformance::HooksGuard { + snapshot_conformance::hooks_lock().await + } + + #[tokio::test] + async fn snapshot_is_consistent_under_concurrent_writes() { + let hooks = hooks().await; + let engine = file_engine().await; + snapshot_conformance::run( + &hooks, + &*engine, + ACCOUNT, + "snap_src", + 50, + EXPECT, + || async { + engine + .process_control_plane_transitions() + .await + .expect("pump"); + }, + ) + .await; + let pool = engine.pool.clone(); + snapshot_conformance::run_refusals( + &hooks, + &*engine, + ACCOUNT, + "snap_refused", + || async { + engine + .process_control_plane_transitions() + .await + .expect("pump"); + }, + |table| async move { + sqlx::query("UPDATE tables SET table_status = 'UPDATING' WHERE table_name = ?") + .bind(table) + .execute(&pool) + .await + .expect("flip status"); + }, + ) + .await; + } + + /// A range key whose attribute definition is missing (the shape issue + /// #259 left behind) streams unordered and complete, not hash-only and + /// short. It pages by rowid, at a page of seven so the rowid keyset + /// query runs (40 items: six fetches), and no page repeats the previous + /// page's last row. + #[tokio::test] + async fn missing_sort_key_definition_streams_unordered_and_complete() { + use extenddb_core::expression::ExpressionMaps; + use extenddb_core::types::{AttributeValue, DescribeTableInput, TableKeyInfo}; + use extenddb_storage::DataEngine; + let _hooks = hooks().await; + let engine = file_engine().await; + engine + .create_table_impl(ACCOUNT, snapshot_conformance::table_input("damaged"), false) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "damaged", || async { + engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; + let desc = engine + .describe_table( + ACCOUNT, + DescribeTableInput { + table_name: "damaged".to_owned(), + }, + ) + .await + .expect("describe"); + let ki = TableKeyInfo { + table_name: "damaged".to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: desc.table_id.clone(), + key_schema: desc.key_schema.clone(), + base_key_schema: desc.key_schema.clone(), + attribute_definitions: desc.attribute_definitions.clone(), + ..Default::default() + }; + // Four partitions, ten items each: hash-only paging would lose rows. + for i in 0..40u32 { + let item: Item = [ + ("pk".to_owned(), AttributeValue::S(format!("p{}", i % 4))), + ("sk".to_owned(), AttributeValue::N(i.to_string())), + ] + .into_iter() + .collect(); + engine + .put_item(&ki, item, false, None, &ExpressionMaps::default(), None) + .await + .expect("put"); + } + // Damage the catalog the way #259 did: drop the sort key's definition. + let defs: Vec = desc + .attribute_definitions + .iter() + .filter(|d| d.attribute_name != "sk") + .cloned() + .collect(); + sqlx::query("UPDATE tables SET attribute_definitions = ? WHERE table_name = ?") + .bind(serde_json::to_string(&defs).unwrap()) + .bind("damaged") + .execute(&engine.pool) + .await + .expect("damage"); + let page = PageSize::set(7); + let fetches_before = paging::fetches(); + let snapshot = engine + .open_table_snapshot(ACCOUNT, "damaged") + .await + .expect("open"); + assert!( + !snapshot.key_ordered, + "a key without a definition has no format order" + ); + let items: Vec = snapshot.items.map(|r| r.expect("item")).collect().await; + let fetches = paging::fetches() - fetches_before; + drop(page); + let mut keys = std::collections::BTreeSet::new(); + for item in &items { + let key = format!("{:?} {:?}", item.get("pk"), item.get("sk")); + assert!(keys.insert(key.clone()), "duplicate item {key}"); + } + assert!( + fetches >= 6, + "{fetches} fetches: the rowid keyset query did not page at 7" + ); + assert_eq!(items.len(), 40, "every item streams"); + } + + /// Every open takes a backup slot; with none free the open is refused at + /// once, and dropping a stream frees its slot. + #[tokio::test] + async fn snapshots_take_backup_slots() { + let _hooks = hooks().await; + let engine = file_engine().await; + engine + .create_table_impl(ACCOUNT, snapshot_conformance::table_input("slots"), false) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "slots", || async { + engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; + let slots = engine.backup_server_slot_count; + let mut held = Vec::new(); + for _ in 0..slots { + held.push( + engine + .open_table_snapshot(ACCOUNT, "slots") + .await + .expect("open"), + ); + } + let err = engine + .open_table_snapshot(ACCOUNT, "slots") + .await + .expect_err("no slot left"); + assert!(matches!(err, StorageError::LimitExceeded(_)), "{err}"); + // CreateBackup's admission on the same slots, with its 30 s wait + // shortened: with snapshots holding every slot and no backup running, + // the refusal must name snapshots. + match crate::backup::acquire_backup_permit( + std::sync::Arc::clone(&engine.backup_server_slots), + engine.backup_server_slot_count, + std::time::Duration::from_millis(50), + ) + .await + { + Err(StorageError::LimitExceeded(m)) => assert!( + m.contains("backups or table snapshots") && m.contains(&format!("at most {slots}")), + "{m}" + ), + other => panic!( + "CreateBackup is refused while snapshots hold every slot: {:?}", + other.err() + ), + } + drop(held.pop()); + engine + .open_table_snapshot(ACCOUNT, "slots") + .await + .expect("a freed slot is reusable"); + } + + /// What a held snapshot does to checkpoints: `PRAGMA + /// wal_checkpoint(TRUNCATE)` while a stream is open, after writes that + /// followed the open. Pins the reply the module doc describes. + #[tokio::test] + async fn wal_checkpoint_while_a_snapshot_is_held() { + use futures::StreamExt as _; + let _hooks = hooks().await; + let engine = file_engine().await; + engine + .create_table_impl(ACCOUNT, snapshot_conformance::table_input("wal"), false) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "wal", || async { + engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; + snapshot_conformance::populate(&*engine, ACCOUNT, "wal", 20).await; + let wal = format!("{}-wal", engine.path.display()); + let checkpoint = |label: &'static str| { + let pool = engine.pool.clone(); + async move { + let started = std::time::Instant::now(); + let (busy, log, done): (i64, i64, i64) = + sqlx::query_as("PRAGMA wal_checkpoint(TRUNCATE)") + .fetch_one(&pool) + .await + .expect("checkpoint"); + eprintln!( + "{label}: busy={busy} log={log} checkpointed={done} in {:?}", + started.elapsed() + ); + (busy, log, done) + } + }; + checkpoint("before the open").await; + let page = PageSize::set(5); + let mut snapshot = engine + .open_table_snapshot(ACCOUNT, "wal") + .await + .expect("open"); + snapshot.items.next().await.expect("an item").expect("ok"); + drop(page); + snapshot_conformance::populate(&*engine, ACCOUNT, "wal", 900).await; + let held = checkpoint("snapshot held").await; + let held_size = std::fs::metadata(&wal).expect("WAL file").len(); + eprintln!("WAL while held: {held_size} bytes"); + drop(snapshot); + // The pooled connection's rollback and return run on the runtime. + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + let released = checkpoint("snapshot dropped").await; + let released_size = std::fs::metadata(&wal).expect("WAL file").len(); + eprintln!("WAL after release: {released_size} bytes"); + // The statement does not fail while the snapshot is held: it waits + // out the busy timeout and reports busy with nothing checkpointed, + // so the WAL keeps every frame written since the open. + assert_eq!(held.0, 1, "busy while the snapshot is held"); + assert!(held.1 > 0, "the writes after the open are in the WAL"); + assert_eq!(held.2, 0, "no frame after the snapshot can be checkpointed"); + assert!(held_size > 0, "the WAL holds them on disk"); + assert_eq!(released.0, 0, "the checkpoint completes once it is dropped"); + assert_eq!(released_size, 0, "and truncates the WAL"); + } + + #[tokio::test] + async fn snapshot_reports_key_order_for_single_range_key() { + let _hooks = hooks().await; + let engine = file_engine().await; + engine + .create_table_impl(ACCOUNT, snapshot_conformance::table_input("ordered"), false) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, "ordered", || async { + engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; + let snapshot = engine + .open_table_snapshot(ACCOUNT, "ordered") + .await + .expect("open"); + assert!(snapshot.key_ordered); + assert_eq!( + snapshot.info.kind, + SnapshotInfo::KIND_SQLITE_WAL_READ_TRANSACTION + ); + } + + /// A restored table's snapshot description carries the same + /// `restore_summary` as `describe_table`'s, on a file-backed engine and + /// on an in-memory one (whose stream is a copy). The restore record is + /// written directly, as `RestoreTableFromBackup` leaves it. + #[tokio::test] + async fn a_restored_table_describes_its_restore() { + let hooks = hooks().await; + let file = file_engine().await; + let memory = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("memory engine"); + prepare(&memory).await; + for engine in [&*file, &memory] { + engine + .create_table_impl( + ACCOUNT, + snapshot_conformance::table_input("restored"), + false, + ) + .await + .expect("create"); + snapshot_conformance::populate(engine, ACCOUNT, "restored", 5).await; + sqlx::query( + "INSERT INTO table_restores \ + (table_id, source_backup_arn, source_table_arn, restore_date_time) \ + SELECT table_id, ?, ?, '2026-10-01T12:00:00.250Z' FROM tables \ + WHERE table_name = 'restored'", + ) + .bind( + "arn:aws:dynamodb:us-east-1:123456789012:table/src/backup/01759320000250-0a1b2c3d", + ) + .bind("arn:aws:dynamodb:us-east-1:123456789012:table/src") + .execute(&engine.pool) + .await + .expect("restore record"); + snapshot_conformance::assert_snapshot_describes_the_table( + &hooks, engine, ACCOUNT, "restored", true, + ) + .await; + } + } + + #[tokio::test] + async fn in_memory_snapshot_copies_the_table_and_releases_the_connection() { + let _hooks = hooks().await; + let engine = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("memory engine"); + prepare(&engine).await; + engine + .create_table_impl(ACCOUNT, snapshot_conformance::table_input("mem"), false) + .await + .expect("create"); + snapshot_conformance::populate(&engine, ACCOUNT, "mem", 40).await; + let snapshot = engine + .open_table_snapshot(ACCOUNT, "mem") + .await + .expect("open"); + assert!(snapshot.key_ordered); + // The stream holds nothing: on the one connection these writes would + // otherwise wait for it until the pool's acquire timeout. + tokio::time::timeout( + std::time::Duration::from_secs(5), + snapshot_conformance::populate(&engine, ACCOUNT, "mem", 60), + ) + .await + .expect("writes proceed while the stream is held"); + let items: Vec = snapshot.items.map(|r| r.expect("item")).collect().await; + assert_eq!( + items.len(), + 40, + "the items as of the open, not the later 20" + ); + let order = extenddb_core::backup_format::KeyOrder::resolve( + &snapshot.description.key_schema, + &snapshot.description.attribute_definitions, + ) + .expect("key order"); + assert!( + items + .windows(2) + .all(|w| order.compare(&w[0], &w[1]).expect("compare").is_lt()), + "key order" + ); + } + + /// The in-memory copy's stream keeps its backup slot until it has yielded + /// its last item or is dropped, so the slot bounds how many decoded copies + /// of a table exist: with every slot held by a copy a further open is + /// refused, and a drained copy frees its slot. + #[tokio::test] + async fn in_memory_copies_hold_their_backup_slot() { + let _hooks = hooks().await; + let engine = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("memory engine"); + prepare(&engine).await; + engine + .create_table_impl(ACCOUNT, snapshot_conformance::table_input("mem"), false) + .await + .expect("create"); + snapshot_conformance::populate(&engine, ACCOUNT, "mem", 10).await; + let slots = engine.backup_server_slot_count; + let mut held = Vec::new(); + for i in 0..slots { + held.push( + engine + .open_table_snapshot(ACCOUNT, "mem") + .await + .unwrap_or_else(|e| panic!("open {i}: {e}")), + ); + } + let err = engine + .open_table_snapshot(ACCOUNT, "mem") + .await + .expect_err("every slot is held by a copy"); + assert!(matches!(err, StorageError::LimitExceeded(_)), "{err}"); + let drained = held.pop().expect("a held copy"); + let items: Vec = drained.items.map(|r| r.expect("item")).collect().await; + assert_eq!(items.len(), 10); + engine + .open_table_snapshot(ACCOUNT, "mem") + .await + .expect("a drained copy frees its slot"); + } + + /// A file-backed engine with one backup slot (pool 4) and a populated + /// check table, for the tests that hold a finished stream. + async fn one_slot_engine(table: &str, items: usize) -> FileEngine { + let engine = file_engine().await; + engine + .create_table_impl(ACCOUNT, snapshot_conformance::table_input(table), false) + .await + .expect("create"); + snapshot_conformance::wait_active(&*engine, ACCOUNT, table, || async { + engine + .process_control_plane_transitions() + .await + .expect("pump"); + }) + .await; + snapshot_conformance::populate(&*engine, ACCOUNT, table, items).await; + assert_eq!( + engine.backup_server_slot_count, 1, + "precondition: one backup slot, so a second open needs the first one's" + ); + engine + } + + /// A stream ends its read transaction and gives up its backup slot at + /// the fetch that shows the end (a short page), not when it next yields + /// `None` or is dropped: with every item read and the stream still held, + /// a checkpoint is not held back by it and the one slot is free for + /// another open. Twenty-two items at a page of five: the fifth fetch + /// returns two rows, so the end is known before the last item is yielded + /// (after `None` the stream's state is dropped, which would release + /// everything whether or not the last page did). + #[tokio::test] + async fn a_drained_stream_ends_its_transaction() { + let _hooks = hooks().await; + let engine = one_slot_engine("drained", 22).await; + let page = PageSize::set(5); + let mut snapshot = engine + .open_table_snapshot(ACCOUNT, "drained") + .await + .expect("open"); + for _ in 0..22 { + snapshot + .items + .next() + .await + .expect("an item") + .expect("stream item"); + } + drop(page); + let (busy, _, _): (i64, i64, i64) = sqlx::query_as("PRAGMA wal_checkpoint(TRUNCATE)") + .fetch_one(&engine.pool) + .await + .expect("checkpoint"); + assert_eq!(busy, 0, "the drained stream's read transaction has ended"); + let second = engine + .open_table_snapshot(ACCOUNT, "drained") + .await + .unwrap_or_else(|e| panic!("open with the drained stream held: {e}")); + drop(second); + assert!( + snapshot.items.next().await.is_none(), + "twenty-two items, then the end" + ); + } + + /// A page that fails (a row whose item does not parse) ends the stream at + /// the error and frees its backup slot at once: with the errored stream + /// still held, the one slot is free for another open. The unparseable row + /// is the seventh in key order, on the second (full) page of five, so the + /// error is not also the end of the table. + #[tokio::test] + async fn a_fetch_error_ends_the_stream_and_frees_its_permit() { + let _hooks = hooks().await; + let engine = one_slot_engine("ended", 20).await; + let desc = engine + .describe_table( + ACCOUNT, + extenddb_core::types::DescribeTableInput { + table_name: "ended".to_owned(), + }, + ) + .await + .expect("describe"); + sqlx::query(&format!( + "UPDATE {} SET item_data = '{{\"pk\": 5}}' WHERE pk = 'p006'", + data_table_name(&desc.table_id) + )) + .execute(&engine.pool) + .await + .expect("make one row unparseable"); + let page = PageSize::set(5); + let mut snapshot = engine + .open_table_snapshot(ACCOUNT, "ended") + .await + .expect("open"); + let mut items = 0; + loop { + match snapshot.items.next().await { + Some(Ok(_)) => items += 1, + Some(Err(e)) => { + eprintln!("page error after {items} items: {e}"); + break; + } + None => panic!("the stream ended without reporting the unparseable row"), + } + } + drop(page); + assert_eq!(items, 5, "the first page, then the error"); + let second = engine + .open_table_snapshot(ACCOUNT, "ended") + .await + .unwrap_or_else(|e| panic!("open with the errored stream held: {e}")); + drop(second); + assert!( + snapshot.items.next().await.is_none(), + "nothing follows the first error" + ); + } + + /// The in-memory copy's stream gives up its backup slot with its last + /// item, not when it next yields `None` or is dropped: with every item + /// read and the copy still held, a further open gets the slot. + #[tokio::test] + async fn an_in_memory_copy_read_to_its_last_item_frees_its_slot() { + let _hooks = hooks().await; + let engine = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("memory engine"); + prepare(&engine).await; + engine + .create_table_impl(ACCOUNT, snapshot_conformance::table_input("mem"), false) + .await + .expect("create"); + snapshot_conformance::populate(&engine, ACCOUNT, "mem", 10).await; + assert_eq!( + engine.backup_server_slot_count, 1, + "precondition: one backup slot" + ); + let mut copy = engine + .open_table_snapshot(ACCOUNT, "mem") + .await + .expect("open"); + for _ in 0..10 { + copy.items.next().await.expect("an item").expect("item"); + } + let second = engine + .open_table_snapshot(ACCOUNT, "mem") + .await + .unwrap_or_else(|e| panic!("open with the read copy held: {e}")); + drop(second); + assert!(copy.items.next().await.is_none(), "ten items, then the end"); + } +} diff --git a/crates/storage-sqlite/src/store.rs b/crates/storage-sqlite/src/store.rs index d67f721e5..075d25b28 100644 --- a/crates/storage-sqlite/src/store.rs +++ b/crates/storage-sqlite/src/store.rs @@ -115,9 +115,10 @@ impl SqliteEngine { if !in_memory && max_connections != pool_size { tracing::info!( "SQLite pool_size {pool_size} raised to {max_connections} for a file-backed \ - database: a backup holds one reader connection for its whole duration while \ - its batches write through another. Set pool_size >= 3 so a running backup \ - does not contend with ordinary requests for connections" + database: a backup or a table snapshot holds one reader connection for its \ + whole duration while writes go through another. Set pool_size >= 3 so a \ + running backup or snapshot does not contend with ordinary requests for \ + connections" ); } let mut opts = SqlitePoolOptions::new() diff --git a/crates/storage-sqlite/src/table_helpers.rs b/crates/storage-sqlite/src/table_helpers.rs index bf9288436..4827391f2 100644 --- a/crates/storage-sqlite/src/table_helpers.rs +++ b/crates/storage-sqlite/src/table_helpers.rs @@ -139,13 +139,31 @@ impl SqliteEngine { &self, account_id: &str, table_name: &str, + ) -> Result { + let mut conn = self + .pool + .acquire() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + self.build_table_description_on(&mut conn, account_id, table_name) + .await + } + + /// Build the description through one connection, so a caller holding a + /// read transaction (the snapshot engine) reads the catalog rows from the + /// same snapshot as everything else it reads. + pub(crate) async fn build_table_description_on( + &self, + conn: &mut sqlx::SqliteConnection, + account_id: &str, + table_name: &str, ) -> Result { let row: Option = sqlx::query_as(&format!( "SELECT {TABLE_COLUMNS} FROM tables WHERE account_id = ? AND table_name = ?" )) .bind(account_id) .bind(table_name) - .fetch_optional(&self.pool) + .fetch_optional(&mut *conn) .await .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -155,7 +173,7 @@ impl SqliteEngine { "SELECT {INDEX_COLUMNS} FROM indexes WHERE table_id = ?" )) .bind(&row.table_id) - .fetch_all(&self.pool) + .fetch_all(&mut *conn) .await .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -163,14 +181,15 @@ impl SqliteEngine { "SELECT {VECTOR_INDEX_COLUMNS} FROM vector_indexes WHERE table_id = ?" )) .bind(&row.table_id) - .fetch_all(&self.pool) + .fetch_all(&mut *conn) .await .map_err(|e| StorageError::Internal(e.to_string()))?; let table_name_owned = row.table_name.clone(); let table_id = row.table_id.clone(); let mut desc = self.build_table_description_from_row(account_id, row, index_rows)?; - desc.restore_summary = self.restore_summary(&table_id, &desc.table_status).await?; + desc.restore_summary = + Self::restore_summary_on(conn, &table_id, &desc.table_status).await?; let catalog_rows = vector_rows .into_iter() .map(VectorIndexRow::into_catalog_row) @@ -412,8 +431,10 @@ fn zero_throughput() -> ProvisionedThroughputDescription { impl SqliteEngine { /// The RestoreSummary of a table created by RestoreTableFromBackup, or /// `None` for any other table. In progress until the table is ACTIVE. - pub(crate) async fn restore_summary( - &self, + /// Takes the connection so a snapshot reads it inside the transaction + /// that pinned its WAL snapshot. + async fn restore_summary_on( + conn: &mut sqlx::SqliteConnection, table_id: &str, status: &extenddb_core::types::TableStatus, ) -> Result, StorageError> { @@ -422,7 +443,7 @@ impl SqliteEngine { FROM table_restores WHERE table_id = ?", ) .bind(table_id) - .fetch_optional(&self.pool) + .fetch_optional(&mut *conn) .await .map_err(|e| StorageError::Internal(e.to_string()))?; Ok(row.map( diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml index 68b64913b..b604d9265 100755 --- a/crates/storage/Cargo.toml +++ b/crates/storage/Cargo.toml @@ -7,6 +7,11 @@ edition.workspace = true rust-version.workspace = true license.workspace = true +[features] +# Compiles the `snapshot_conformance` check that each backend's tests run +# against a real database. Production builds do not compile it. +test-hooks = [] + [dependencies] aes-gcm = { workspace = true } async-trait = { workspace = true } diff --git a/crates/storage/src/error.rs b/crates/storage/src/error.rs index 34a93c340..8fe3ac287 100755 --- a/crates/storage/src/error.rs +++ b/crates/storage/src/error.rs @@ -71,6 +71,12 @@ pub enum StorageError { /// and the worker dropped claimed rows on transient errors. #[error("{0}")] Transient(String), + /// The backend discarded a consistent snapshot before the reader drained + /// it (`MongoDB` `SnapshotTooOld`, a `PostgreSQL` snapshot conflict on a hot + /// standby). The operation that held it must restart from the beginning; + /// nothing read so far may be combined with a fresh snapshot. + #[error("snapshot expired: {0}")] + SnapshotExpired(String), #[error( "Catalog version mismatch: expected {expected}, found {found}. Run 'extenddb migrate' to update." )] diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs index c760583a6..19cd6bcbe 100755 --- a/crates/storage/src/lib.rs +++ b/crates/storage/src/lib.rs @@ -22,12 +22,20 @@ pub mod operations; pub mod policy_migration; pub mod server_components; pub mod settings_store; +pub mod snapshot; +#[doc(hidden)] +#[cfg(feature = "test-hooks")] +pub mod snapshot_conformance; +#[doc(hidden)] +#[cfg(feature = "test-hooks")] +pub mod test_pause; pub mod transact; pub mod vector_catalog; pub mod vector_lifecycle; pub use backend::{Backend, BackendAlreadySet, backend_name, set_backend, try_backend}; pub use backup_definition::{BackupRestoreSource, RestoreRequest, RestoreTableOverrides}; +pub use snapshot::{SnapshotEngine, SnapshotInfo, TableSnapshot}; pub use transact::{IdempotencyKey, TransactGetOp, TransactWriteOp}; @@ -806,17 +814,26 @@ pub trait BackupEngine: Send + Sync { /// DynamoDB-compatible API. This trait has NO additional methods beyond /// the trait bounds — backend-specific concerns belong in `ServerRuntimeHooks`. pub trait StorageEngine: - TableEngine + DataEngine + MetadataEngine + StreamEngine + BackupEngine + WorkerStore + Send + Sync + TableEngine + + DataEngine + + MetadataEngine + + StreamEngine + + BackupEngine + + SnapshotEngine + + WorkerStore + + Send + + Sync { } -// Blanket implementation: any type implementing all 6 traits is a StorageEngine +// Blanket implementation: any type implementing all 7 traits is a StorageEngine impl StorageEngine for T where T: TableEngine + DataEngine + MetadataEngine + StreamEngine + BackupEngine + + SnapshotEngine + WorkerStore + Send + Sync diff --git a/crates/storage/src/snapshot.rs b/crates/storage/src/snapshot.rs new file mode 100644 index 000000000..766eb7aa8 --- /dev/null +++ b/crates/storage/src/snapshot.rs @@ -0,0 +1,458 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Consistent table snapshots: the one obligation a backend has toward the +//! engine-owned backup path. +//! +//! A backup needs two things from a backend that nothing else needs: the +//! table's definition and every item in it, both as of one instant. The +//! definition lives in the catalog and the items in the data store, and on +//! `PostgreSQL` with a separate data database (the layout `extenddb init` +//! creates) those are different databases, so no single transaction covers +//! both. Each backend therefore has its own way of fixing the instant: +//! `PostgreSQL` with a separate data database holds the catalog row `FOR +//! SHARE` from the definition read until the data transaction's snapshot is +//! pinned; `PostgreSQL` on one database shared by catalog and data reads both +//! in one `REPEATABLE READ` transaction whose first read pins the snapshot, +//! like `SQLite`; `SQLite` reads catalog and data inside one deferred read +//! transaction whose first read pins the WAL snapshot, taking no lock; +//! `MongoDB` reads both through one session with snapshot read concern. +//! Cassandra has no read that fixes an instant across a table and returns +//! [`StorageError::Unsupported`]; its in-catalog logical backup (ADR-0019) +//! stays the backup path for that backend. +//! +//! What that fixes: on `PostgreSQL` with a separate data database, +//! `UpdateTable` and `DeleteTable` take the row `FOR UPDATE`, so neither +//! commits between the definition read and the data pin; background workers +//! that flip an index's status (a backfill finishing) do not take the row +//! lock, so an index may be reported in the status it had at the definition +//! read while the items are a few milliseconds newer. On `SQLite`, `MongoDB` +//! and `PostgreSQL` on one shared database the definition and the items come +//! from one snapshot, so no write of any kind sits between them; on that +//! `PostgreSQL` layout a `DeleteTable` that commits between the snapshot and +//! the open's lock on the data table makes the open report the table as not +//! found ([`paging::OpenStage::AfterDefinition`]). +//! +//! Indexes are reported in whatever lifecycle status they had, except that +//! an index being deleted is dropped from the description +//! ([`drop_deleting_indexes`]), and on `PostgreSQL` a GSI row whose data table +//! is missing (not an expected state: `UpdateTable` creates the data table +//! before it commits the row) is left out as the in-catalog backup leaves it +//! out. A restore recreates the indexes the +//! description lists and builds them from the items, so an index that was +//! still backfilling at the snapshot comes back complete; an index that is +//! not in the description does not come back at all, which is why a +//! `CREATING` index with a data table is kept. +//! +//! The returned [`TableSnapshot`] owns the open transaction, session, or +//! connection for as long as its `items` stream is alive. Dropping the stream +//! releases it. The stream must be dropped inside a Tokio runtime: on the +//! sqlx backends (`SQLite`, `PostgreSQL`) a stream that has not reached its end +//! (including one never polled) returns its pooled connection on drop by +//! spawning onto the current runtime, and dropping it outside one panics +//! ("this functionality requires a Tokio context"); a `MongoDB` stream drops +//! anywhere. A caller tested only on `MongoDB` will not see this, so test drops +//! inside a runtime. The caller must +//! drain or drop the stream promptly: a held `PostgreSQL` snapshot delays +//! vacuum on every table in the data database, a held `SQLite` read +//! transaction stops WAL checkpoints, and `MongoDB` keeps at least +//! `minSnapshotHistoryWindowInSeconds` (default 300 s) of snapshot history, +//! past which the server may discard the snapshot and the stream fails with +//! [`StorageError::SnapshotExpired`]. + +use extenddb_core::backup_format::KeyOrder; +use extenddb_core::types::{IndexStatus, Item, TableDescription}; +use futures::future::BoxFuture; +use futures::stream::BoxStream; + +use crate::error::StorageError; + +/// How a backend took its snapshot: the manifest's `snapshot` member, one +/// definition shared with the backup format. `kind` is one of the `KIND_*` +/// constants on the type; `marker` is `pg_current_snapshot()` text on +/// `PostgreSQL`, the cluster time `:` on `MongoDB`, empty on +/// `SQLite` (a WAL read transaction has no addressable position). The marker +/// is recorded for a later change-log archive to anchor against (which +/// writes committed before the snapshot); it is not a resume point. No +/// backend can reopen a snapshot at its marker, so a stream that fails is +/// started over from a new open, and on `SQLite`, where the marker is empty, +/// a change-log archive has nothing to anchor on. +pub use extenddb_core::backup_format::SnapshotDescriptor as SnapshotInfo; + +/// Page size of a snapshot stream and, under `test-hooks`, the knob and the +/// round-trip counter the conformance check uses. One knob for every backend, +/// so the check can set it and observe its effect without a per-backend +/// setter that a test could forget to call. +pub mod paging { + /// Rows or documents a backend fetches per round trip. Measured on + /// localhost `PostgreSQL` 16 with 318-byte items: 1000 and 10000 read a + /// million rows in the same time; 1000 holds less memory. + /// + /// The bound is rows, not bytes. `PostgreSQL`'s `FETCH` and `SQLite`'s + /// `LIMIT` each hold up to this many rows of item text in memory at + /// once, and at the default 400 KB item limit one row's text can reach + /// 27.7 MB (the measured worst case in #394's line-cap table), so one + /// fetch is bounded by rows, not bytes. A byte budget (keyset paging with a byte cutoff) is + /// required before the stream gets a caller; the backup design assigns + /// it to the backup job PR. `MongoDB`'s batches are bounded by the + /// server's 16 MiB reply limit as well as by this count. + pub const DEFAULT_FETCH_BATCH: usize = 1000; + + #[cfg(feature = "test-hooks")] + static OVERRIDE: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + #[cfg(feature = "test-hooks")] + static FETCHES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + + /// The page size a stream is built with. Backends read it once when they + /// open a stream, so a change to the knob during a stream cannot end it + /// early or late. + #[must_use] + pub fn fetch_batch() -> usize { + #[cfg(feature = "test-hooks")] + { + match OVERRIDE.load(std::sync::atomic::Ordering::Relaxed) { + 0 => DEFAULT_FETCH_BATCH, + n => n, + } + } + #[cfg(not(feature = "test-hooks"))] + { + DEFAULT_FETCH_BATCH + } + } + + /// Record one round trip to the backend for a page of items. A no-op + /// without `test-hooks`. + /// + /// The counter is trusted: the conformance check reads it to confirm + /// the backend paged at the knob's size, so a backend must call this + /// exactly once per round trip (one `FETCH`, one page query, one + /// `find` or `getMore`), never per item. A backend that ignored the + /// knob and called it per item would pass the paging assertions with + /// the whole table read in one trip. `MongoDB`'s driver hides `getMore` + /// from the stream, so its test fixture feeds the counter from command + /// monitoring (what the server received) for that reason, and its + /// engine does not call this. + pub fn record_fetch() { + #[cfg(feature = "test-hooks")] + FETCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + + /// Set the page size for every stream opened afterwards; `0` restores the + /// default. + #[cfg(feature = "test-hooks")] + pub fn set_fetch_batch(rows: usize) { + OVERRIDE.store(rows, std::sync::atomic::Ordering::Relaxed); + } + + /// Round trips recorded by every stream in this process so far. + #[cfg(feature = "test-hooks")] + #[must_use] + pub fn fetches() -> u64 { + FETCHES.load(std::sync::atomic::Ordering::Relaxed) + } + + /// The two points in `open_table_snapshot` at which every backend calls + /// [`open_pause`]. Under `test-hooks` the conformance check holds an open + /// at one of them while it commits writes and an `UpdateTable`, to check + /// that the description and the items come from one instant. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub enum OpenStage { + /// The definition has been read (on `MongoDB`, the `tables` document; + /// its index documents are read after this point, in the same + /// session). On `PostgreSQL` with a separate data database the data + /// transaction has begun and holds `ACCESS SHARE` on the data table, + /// the catalog barrier is still held, and the next statement is the + /// `SELECT pg_current_snapshot()` that pins the data snapshot, so a + /// barrier released any earlier would fall inside the pause. On + /// `SQLite`, `MongoDB` and `PostgreSQL` with one shared database the first + /// catalog read already pinned the snapshot; on the shared `PostgreSQL` + /// layout the pause is before the `LOCK TABLE`, in the window where a + /// `DeleteTable` can still drop the data table. + AfterDefinition, + /// The data snapshot is pinned; the item read (`PostgreSQL`'s `DECLARE + /// ... CURSOR`, `MongoDB`'s `find`, `SQLite`'s first page) has not + /// started. + BeforeItems, + } + + /// Hold the open at `stage` if a test armed a pause for it. Without + /// `test-hooks` this is a ready future and the open never waits (it stays + /// `async` so the backends call it the same way in both builds). + #[cfg_attr(not(feature = "test-hooks"), allow(clippy::unused_async))] + pub async fn open_pause(stage: OpenStage) { + #[cfg(feature = "test-hooks")] + PAUSE.reached(stage).await; + #[cfg(not(feature = "test-hooks"))] + let _ = stage; + } + + #[cfg(feature = "test-hooks")] + static PAUSE: crate::test_pause::PausePoint = crate::test_pause::PausePoint::new(); + + /// A pause armed by [`arm_open_pause`]. Dropping it disarms the pause and + /// releases an open waiting at it, so a failing test cannot leave a pause + /// armed for the next open in the process. + #[cfg(feature = "test-hooks")] + pub type OpenPause = crate::test_pause::Pause; + + /// Arm a pause at `stage` for the next open that reaches it (one open, + /// once), whatever its engine or table. Process-global like the knob and + /// the counter: every test that opens a snapshot or touches these + /// statics holds `snapshot_conformance::hooks_lock`, so no foreign open + /// can take the pause. + /// + /// # Panics + /// + /// In debug builds, when a pause is already armed: a second arm would + /// replace the first silently, and the first handle's drop would then + /// disarm the second. + #[cfg(feature = "test-hooks")] + #[must_use] + pub fn arm_open_pause(stage: OpenStage) -> OpenPause { + PAUSE.arm(stage) + } +} + +/// Drop indexes that were being deleted at the snapshot instant and keep +/// every other one, whatever its build state. One rule for every backend: a +/// restore recreates the indexes in the description and builds them from the +/// items, so a `CREATING` or `UPDATING` index comes back complete, and an +/// index the user had removed must not come back. No in-tree backend +/// currently reports a `DELETING` index (each removes the index row in the +/// operation that deletes it); the rule is here so one that does cannot leak +/// a deleted index into a backup. +pub fn drop_deleting_indexes(description: &mut TableDescription) { + if let Some(gsis) = description.global_secondary_indexes.as_mut() { + gsis.retain(|g| g.index_status != "DELETING"); + } + if let Some(vectors) = description.vector_indexes.as_mut() { + vectors.retain(|v| v.index_status != IndexStatus::Deleting); + } +} + +/// Whether a key schema has the order the backup format defines, by the +/// format's own rule (`KeyOrder::resolve`): one hash key listed first, at +/// most one range key with a name of its own, every key attribute defined +/// exactly once. A table that fails the +/// rule streams unordered and complete (multi-part keys; a range key whose +/// attribute definition is missing or duplicated, which would otherwise make +/// a backend page by the hash key alone and lose rows). Backends use this to +/// decide `key_ordered`, so a stream they mark ordered is one the manifest +/// parser accepts. +#[must_use] +pub fn key_schema_is_orderable(description: &TableDescription) -> bool { + KeyOrder::resolve(&description.key_schema, &description.attribute_definitions).is_ok() +} + +/// A consistent view of one table: its definition, how the snapshot was +/// taken, and every item as of the same instant. +pub struct TableSnapshot { + /// The table as the storage layer's `describe_table` would have reported + /// it at the snapshot instant: key schema, attribute definitions, indexes + /// (minus any being deleted; see [`drop_deleting_indexes`]), billing, + /// throughput, class, SSE, vector indexes, ARN, id and creation time. + /// Streams, TTL and tags are not part of a backup and may be absent. + /// + /// This is the storage-level description, not the API response: the + /// engine's `prepare_for_response` is not applied, so `restore_summary` + /// may be present on an `ACTIVE` table that was restored (the API drops + /// it once the table is not `CREATING`), `table_throughput_mode_summary` + /// is unset (the API derives it from `billing_mode_summary`), and + /// `item_count` and `table_size_bytes` are the catalog row's values, the + /// stats pass's last figures and not the instant's: a caller that needs + /// the count as of the snapshot counts the stream. + pub description: TableDescription, + /// Snapshot mechanism and position. + pub info: SnapshotInfo, + /// True when `items` yields rows in the backup format's key order + /// (`extenddb_core::backup_format::key_order`): hash key by its text + /// encoding, bytewise, so a number hash key sorts as text; range key `S` + /// and `B` bytewise, `N` numeric. The SQL backends produce it from their + /// primary key columns. On `PostgreSQL` the rule is the collation's name: + /// the table is ordered only when its text key columns (`pk`, and `sk_s` + /// for an `S` range key) have a collation named "C" or "POSIX". Any other + /// collation, including "default" on a database whose own collation is + /// C, streams unordered (complete, only not ordered), because ordering a + /// column that does not already sort bytewise would sort the whole table + /// before the first row. A backup records the flag so a restore into a + /// backend that benefits from ordered input can rely on it; a restore + /// must produce the same table either way. + pub key_ordered: bool, + /// Every item in the table as of the snapshot. Holds the snapshot's + /// transaction or session until it ends or is dropped; drop it inside a + /// Tokio runtime (the sqlx backends panic when a stream that has not + /// reached its end is dropped outside one; see the module doc). The first + /// `Err` ends the stream on every backend: nothing follows it, and since + /// no backend can reopen a snapshot at its marker, any error means the + /// caller starts over with a new open ([`StorageError::SnapshotExpired`] + /// is the one worth retrying at once; the others name a fault). + /// + /// Items are yielded as stored, after the wire deserializer: there is no + /// promise that every item has its key attributes with the declared + /// types (the backends do not validate stored rows on the way out), and + /// a valid `N` value arrives normalized (the deserializer rewrites it) + /// while an invalid one stored by an earlier build passes through raw. + /// Items whose TTL has expired but that the sweep has not yet removed + /// are rows and are yielded like any other. + pub items: BoxStream<'static, Result>, +} + +impl std::fmt::Debug for TableSnapshot { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TableSnapshot") + .field("table", &self.description.table_name) + .field("info", &self.info) + .field("key_ordered", &self.key_ordered) + .finish_non_exhaustive() + } +} + +/// Open a consistent snapshot of one table. +/// +/// The one method has no default body on purpose: a backend that cannot +/// take a consistent snapshot must say so (the Cassandra implementation is +/// the refusal, `StorageError::Unsupported` before it resolves the table), +/// and a default that refused would let a new backend compile with the +/// backup path silently absent. Likewise `StorageError` is not +/// `#[non_exhaustive]` on purpose: the engine's exhaustive match over it +/// (`storage_err_to_dynamo`) is what gives [`StorageError::SnapshotExpired`] +/// its wire mapping, and a new variant must be mapped there, not fall into +/// a wildcard. Both are breaking changes for an out-of-tree backend, and +/// both are meant to be. +pub trait SnapshotEngine: Send + Sync { + /// Open a snapshot of `table_name` in `account_id`. + /// + /// The table must be `ACTIVE` (a restore target is `CREATING` and is + /// refused with [`StorageError::TableNotActive`]). The definition and the + /// first item read are fixed at the same instant; writes committed after + /// that instant are not visible in `items`, and items committed before it + /// all are, exactly once. + /// + /// Resources, until the stream ends or is dropped: on `PostgreSQL` one + /// connection and one snapshot permit (with a separate data database the + /// open also holds a catalog connection until the data snapshot is + /// pinned); on file-backed `SQLite` one pool connection and one backup + /// slot; on in-memory `SQLite` one backup slot and no connection (the + /// stream serves a copy read during the open); on `MongoDB` one session + /// and one server cursor. The caller bounds how many snapshots are open + /// at once; `PostgreSQL` and `SQLite` also refuse over their own bound + /// (`LimitExceeded`) rather than queue, `MongoDB` has no bound of its own. + /// + /// The stream ends at its first `Err` on every backend; a caller that + /// sees an error has not seen the whole table and must not treat what it + /// read as complete. No snapshot can be reopened at its marker, so any + /// error from the stream means a new open from the beginning; nothing + /// read from the failed stream may be combined with the new one. + /// Drop the stream inside a Tokio runtime: the sqlx + /// backends return a connection on drop through the runtime and panic + /// outside one when the stream has not reached its end, `MongoDB` does + /// not, so a caller tested only on `MongoDB` will not see the panic. + /// + /// A stream left idle is ended by the server, not by this crate: on + /// `PostgreSQL` an operator's `idle_in_transaction_session_timeout` ends + /// the backend between two fetches and the stream reports `Internal` + /// (not `SnapshotExpired`); on `MongoDB` the cursor has the server's + /// default idle timeout (10 minutes; it is opened without + /// `noCursorTimeout`), after which the next batch reports `Internal` + /// likewise. A caller whose consumer can stall (an upload) must keep + /// polling or expect to start over. + /// + /// # Errors + /// + /// [`StorageError::TableNotFound`] when the table does not exist in the + /// account, [`StorageError::TableNotActive`] (or the backend's + /// equivalent) when it is not `ACTIVE`, [`StorageError::Unsupported`] + /// when the backend cannot take a consistent snapshot at all (Cassandra, + /// which answers `Unsupported` before it resolves the table and so never + /// returns `TableNotFound` or `TableNotActive`) or in this configuration + /// (a `MongoDB` server that is not a replica set member, which cannot + /// serve snapshot reads), + /// [`StorageError::LimitExceeded`] when the backend's bound on open + /// snapshots is reached (`PostgreSQL`, `SQLite`), + /// [`StorageError::Validation`] when the account id is not one the + /// backend accepts (`MongoDB` validates it before the catalog read), and + /// [`StorageError::Internal`] for a database error during the open. + /// + /// From the stream: [`StorageError::SnapshotExpired`] when the backend + /// discarded the snapshot before the stream was drained (retry with a + /// new open), and [`StorageError::Internal`] for a database error during + /// a fetch (a terminated backend, a timed-out cursor, a lost connection) + /// or a stored row that does not parse as an item. + fn open_table_snapshot( + &self, + account_id: &str, + table_name: &str, + ) -> BoxFuture<'_, Result>; +} + +#[cfg(test)] +mod tests { + use super::*; + use extenddb_core::types::GsiDescription; + + #[test] + fn deleting_indexes_are_dropped_and_others_kept() { + let gsi = |name: &str, status: &str| -> GsiDescription { + serde_json::from_value(serde_json::json!({ + "IndexName": name, + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"}, + "IndexStatus": status, + "IndexSizeBytes": 0, + "ItemCount": 0, + "IndexArn": format!("arn:aws:dynamodb:us-east-1:123456789012:table/t/index/{name}") + })) + .expect("gsi description") + }; + let mut description = TableDescription { + global_secondary_indexes: Some(vec![ + gsi("kept_active", "ACTIVE"), + gsi("kept_creating", "CREATING"), + gsi("gone", "DELETING"), + ]), + ..TableDescription::default() + }; + drop_deleting_indexes(&mut description); + let names: Vec<_> = description + .global_secondary_indexes + .unwrap() + .into_iter() + .map(|g| g.index_name) + .collect(); + assert_eq!(names, ["kept_active", "kept_creating"]); + } + + #[test] + fn deleting_vector_indexes_are_dropped_and_others_kept() { + let vector = |name: &str, status: &str| -> extenddb_core::types::VectorIndexDescription { + serde_json::from_value(serde_json::json!({ + "IndexName": name, + "VectorAttribute": {"AttributeName": "emb"}, + "Dimensions": 3, + "DistanceFunction": "COSINE", + "IndexStatus": status, + "IndexSizeBytes": 0, + "ItemCount": 0, + "IndexArn": format!("arn:aws:dynamodb:us-east-1:123456789012:table/t/index/{name}") + })) + .expect("vector index description") + }; + let mut description = TableDescription { + vector_indexes: Some(vec![ + vector("kept_active", "ACTIVE"), + vector("gone", "DELETING"), + vector("kept_creating", "CREATING"), + ]), + ..TableDescription::default() + }; + drop_deleting_indexes(&mut description); + let names: Vec<_> = description + .vector_indexes + .unwrap() + .into_iter() + .map(|v| v.index_name) + .collect(); + assert_eq!(names, ["kept_active", "kept_creating"]); + } +} diff --git a/crates/storage/src/snapshot_conformance.rs b/crates/storage/src/snapshot_conformance.rs new file mode 100644 index 000000000..04b58c9b2 --- /dev/null +++ b/crates/storage/src/snapshot_conformance.rs @@ -0,0 +1,1526 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Backend-independent consistency check for [`SnapshotEngine`], run by each +//! backend's own test suite against a real database. +//! +//! The check opens a snapshot, then has a writer delete pre-existing items, +//! insert new ones, and overwrite every remaining item several times while +//! the snapshot stream is still being read. A consistent snapshot shows the +//! table exactly as it was at open: every deleted item present, no inserted +//! item, every overwritten item at its pre-open version, no duplicates. +//! +//! For the check to mean anything the stream must still be open against the +//! backend while the writer runs. Every backend fetches a page at a time +//! through the shared [`paging`] knob; `run` sets the page size itself before +//! it opens anything, requires the table to span several pages +//! (`BASE_ITEMS + DOOMED_ITEMS > 3 * page_size`), reads one page plus one +//! item, and then asserts that the shared round-trip counter advanced by at +//! least two since just before the open, so a backend that ignored the knob (or a stream +//! that fetched the whole table at once) fails here rather than letting the +//! rest of the check pass vacuously. It then pauses until the writer has +//! committed, then drains. A paged reader that re-queries per page without a +//! snapshot fails the version or count assertion: such mutants, run at a page +//! of 50 on all three backends (SQLite: fresh `BEGIN` per page; PostgreSQL: +//! `OFFSET`/`LIMIT` in READ COMMITTED; MongoDB: a plain majority-read cursor +//! outside the session), each failed. The +//! counter is trusted, not verified: a backend must record one fetch per +//! round trip (PostgreSQL and SQLite record one beside each `FETCH` or page +//! query), and one that recorded per item would pass these assertions with +//! the whole table read in one trip. The MongoDB test fixture therefore +//! feeds the counter from what the server received, each `find` and +//! `getMore` on the data database (command monitoring), since the driver +//! hides `getMore` from the stream and the engine records nothing. +//! +//! Every drain here is under a deadline (`DRAIN_TIMEOUT`, 120 s): a stream that +//! never ends (a keyset page that re-queries an empty page forever) fails +//! the check naming the stage and the items read, instead of hanging it and, +//! through the hooks lock, every other snapshot test in the binary. +//! [`hooks_lock`] itself gives up when one holder has kept the lock for +//! `HOOKS_LOCK_WAIT` (600 s), naming the wait. +//! +//! Before the writer starts, the check compares the snapshot's `description` +//! with the storage layer's `describe_table` of the quiescent table, with the +//! running counts and `table_throughput_mode_summary` cleared on both sides +//! (`comparable_description`): the trait's promise that the description is +//! the table as described, which a backend building it from a second code +//! path could break without any other assertion noticing. +//! `restore_summary` is compared: storage sets it on a restored table, so a +//! backend's suite also runs [`assert_snapshot_describes_the_table`] on a +//! table that has one. +//! +//! The knob, the counter and the open pause are process-global, and the +//! pause goes to the next open that reaches it whatever its engine or table. +//! One lock covers all three: [`hooks_lock`]. Every test in a backend crate +//! that opens a snapshot or touches the [`paging`] statics holds it for its +//! whole body, and [`run`] and [`run_refusals`] take its guard as their first +//! argument, so a check cannot run without it and no other test in the +//! binary can take the pause, reset the knob, or feed the counter while a +//! check is running. +//! +//! Key order is checked with `KeyOrder::compare`, the comparator the +//! manifest parser and the restore use, so the backend's SQL order, the +//! format's rule and this oracle are one rule, not three. +//! +//! Reads after the commit: the counter is read when the writer's rounds +//! have committed, and must have advanced again by the end of the drain. A +//! stream that fetched every page during the open (each page in its own +//! transaction, say) passes every content assertion, because the writer has +//! not started; this assertion fails it. +//! +//! One instant: for each [`paging::OpenStage`] the check arms a pause in the +//! next open (every backend calls [`paging::open_pause`] at both stages) and, +//! while the open waits, commits a put W1, an `UpdateTable` U adding a GSI, +//! and a put W2 after U returns. With W1 < U < W2 in commit order, a +//! description and items of one instant satisfy: U's GSI in the description +//! implies W1 in the items, and W2 in the items implies U's GSI in the +//! description. Both are implications, and both hold when a backend reads +//! the definition and the items after W2, so the check also asserts that W2 +//! is not in the items: W2 commits after U returns, and on every correct +//! backend U cannot commit inside the open's read window (PostgreSQL with a +//! separate data database makes U wait on the row until the data pin; every +//! other backend pinned before the pause). W2 in the items therefore means +//! the pause sat outside the read window and the scenario checked nothing; +//! mutants that moved both pauses to the top of the open (SQLite reading the +//! definition on the pool before `begin`, MongoDB reading the index +//! documents outside the session) passed without this +//! assertion and fail on it. That assertion is live only when W2 commits +//! while the open is held, so the check also asserts which branch resumed +//! the open against [`Expectations::update_table_waits_at`]: where U does +//! not wait on the open, W2's commit must resume it (within the 2 s +//! timeout), and where U waits (PostgreSQL, separate data database, +//! `AfterDefinition`), the timeout must. The sixth round's mutant, the +//! pauses before `begin` plus a 2.5 s delay in `UpdateTable`, passed the W2 +//! assertion with W2 committed after the resume and fails on this one. `AfterDefinition` catches a definition read that is not +//! fenced against the data pin (PostgreSQL without `FOR SHARE`; MongoDB's +//! index documents read outside the session), `BeforeItems` an item read +//! that is not at the pinned snapshot (PostgreSQL at READ COMMITTED, where +//! the cursor takes a fresh snapshot; a MongoDB cursor outside the session; +//! a SQLite stream that reads its pages in fresh transactions). +//! +//! Key shapes: hash-only `S`, `N` and `B` keys and `S` hash with `S`, `N` or +//! `B` range key, at a page of three, with values whose bytewise order +//! differs from collation order, binary values whose base64 order differs +//! from their byte order, and range numbers (negative, fractional, 1E-130, +//! 9.9E125, 38-digit values a last digit apart) whose text order differs +//! from their numeric order. Each table is checked for completeness, duplicates +//! (a keyset page that repeats its last row) and, on a backend that reports +//! `key_ordered`, ascending `KeyOrder::compare` order. +//! +//! Liveness (lesson 5): a second snapshot taken after the writer stops must +//! show the writer's final state, which proves the writer wrote and the first +//! snapshot's view was a snapshot, not a frozen database. + +use std::collections::BTreeSet; +use std::time::Duration; + +use extenddb_core::expression::ExpressionMaps; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, BillingMode, CreateTableInput, DescribeTableInput, + GsiInput, Item, KeySchemaElement, KeyType, Projection, ProjectionType, ScalarAttributeType, + TableKeyInfo, TableStatus, +}; +use futures::StreamExt; + +use crate::error::StorageError; +use crate::snapshot::paging; +use crate::{DataEngine, SnapshotEngine, TableEngine}; + +/// Number of items present before the snapshot opens and never deleted. +pub const BASE_ITEMS: usize = 600; +/// Items present before the snapshot opens that the writer deletes after. +pub const DOOMED_ITEMS: usize = 100; +/// Items the writer inserts after the snapshot opens. +pub const LATE_ITEMS: usize = 100; +/// Overwrite rounds the writer must complete while the stream is open. +pub const ROUNDS_DURING_STREAM: u32 = 2; + +fn s(v: impl Into) -> AttributeValue { + AttributeValue::S(v.into()) +} + +fn n(v: impl std::fmt::Display) -> AttributeValue { + AttributeValue::N(v.to_string()) +} + +fn base_item(i: usize, version: u32) -> Item { + [ + ("pk".to_owned(), s(format!("p{:03}", i % 20))), + ("sk".to_owned(), n(i)), + ("v".to_owned(), n(version)), + ("g".to_owned(), s(format!("g{}", i % 7))), + ] + .into_iter() + .collect() +} + +fn doomed_item(i: usize) -> Item { + [ + ("pk".to_owned(), s(format!("doomed{:02}", i % 5))), + ("sk".to_owned(), n(i)), + ("v".to_owned(), n(0)), + ] + .into_iter() + .collect() +} + +fn late_item(i: usize) -> Item { + [ + ("pk".to_owned(), s(format!("late{:02}", i % 5))), + ("sk".to_owned(), n(i)), + ("v".to_owned(), n(0)), + ] + .into_iter() + .collect() +} + +fn key_of(item: &Item) -> Item { + item.iter() + .filter(|(k, _)| *k == "pk" || *k == "sk") + .map(|(k, v)| (k.clone(), v.clone())) + .collect() +} + +/// `(pk, sk)` as a set key for identity checks. +fn sort_key(item: &Item) -> (String, i64) { + let pk = match item.get("pk") { + Some(AttributeValue::S(v)) => v.clone(), + other => panic!("pk must be a string, got {other:?}"), + }; + let sk = match item.get("sk") { + Some(AttributeValue::N(v)) => v.parse::().expect("integer sort key"), + other => panic!("sk must be a number, got {other:?}"), + }; + (pk, sk) +} + +/// The table the check creates: composite key (`pk` S, `sk` N) with one GSI. +#[must_use] +pub fn table_input(name: &str) -> CreateTableInput { + CreateTableInput { + table_name: name.to_owned(), + key_schema: vec![ + KeySchemaElement { + attribute_name: "pk".to_owned(), + key_type: KeyType::Hash, + }, + KeySchemaElement { + attribute_name: "sk".to_owned(), + key_type: KeyType::Range, + }, + ], + attribute_definitions: vec![ + AttributeDefinition { + attribute_name: "pk".to_owned(), + attribute_type: ScalarAttributeType::S, + }, + AttributeDefinition { + attribute_name: "sk".to_owned(), + attribute_type: ScalarAttributeType::N, + }, + AttributeDefinition { + attribute_name: "g".to_owned(), + attribute_type: ScalarAttributeType::S, + }, + ], + billing_mode: Some(BillingMode::PayPerRequest), + global_secondary_indexes: Some(vec![GsiInput { + index_name: "gi".to_owned(), + key_schema: vec![KeySchemaElement { + attribute_name: "g".to_owned(), + key_type: KeyType::Hash, + }], + projection: Projection { + projection_type: ProjectionType::KeysOnly, + non_key_attributes: None, + }, + provisioned_throughput: None, + }]), + ..Default::default() + } +} + +async fn key_info(engine: &E, account: &str, table: &str) -> TableKeyInfo { + let desc = engine + .describe_table( + account, + DescribeTableInput { + table_name: table.to_owned(), + }, + ) + .await + .expect("describe the check table"); + TableKeyInfo { + table_name: table.to_owned(), + account_id: account.to_owned(), + table_id: desc.table_id.clone(), + key_schema: desc.key_schema.clone(), + base_key_schema: desc.key_schema.clone(), + attribute_definitions: desc.attribute_definitions.clone(), + ..Default::default() + } +} + +/// Put `count` items of the check's shape (`pk`, `sk`, `v`) into an +/// existing table made with [`table_input`], for backend tests that need a +/// populated table outside [`run`]. +pub async fn populate(engine: &E, account: &str, table: &str, count: usize) +where + E: TableEngine + DataEngine + Send + Sync, +{ + let ki = key_info(engine, account, table).await; + for i in 0..count { + put(engine, &ki, base_item(i, 0)).await; + } +} + +/// The lock behind [`hooks_lock`]. +static HARNESS: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +/// How long one holder may keep the hooks lock before a test waiting for it +/// panics naming the wait. A holder that hangs (a stream that never ends, +/// say) would otherwise block every other snapshot test in the binary +/// silently until the CI job's own timeout. Counted from when the current +/// holder took the lock, not from when the waiter queued: the suite runs +/// serially behind the lock, so with as many test threads as tests the last +/// one queues for nearly the whole suite while no holder is slow. +const HOOKS_LOCK_WAIT: Duration = Duration::from_secs(600); + +/// When the current holder of [`HARNESS`] took it; `None` while it is free. +static HELD_SINCE: std::sync::Mutex> = std::sync::Mutex::new(None); + +fn held_since() -> Option { + *HELD_SINCE + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +fn set_held_since(since: Option) { + *HELD_SINCE + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = since; +} + +/// How long a drain in this module waits for the next item or the end of a +/// stream before it panics naming the stage and the items read so far. Well +/// above any drain here (the longest, the main scenario's, takes seconds) +/// and below [`HOOKS_LOCK_WAIT`], so the test that hangs is the one that +/// fails, not its neighbours. +const DRAIN_TIMEOUT: Duration = Duration::from_secs(120); + +/// Proof that the holder has the process-wide lock on the snapshot test +/// hooks (the page-size knob, the round-trip counter and the open pause). +/// [`run`] and [`run_refusals`] take it, so a check cannot run without it. +pub struct HooksGuard { + _guard: tokio::sync::MutexGuard<'static, ()>, +} + +impl Drop for HooksGuard { + fn drop(&mut self) { + set_held_since(None); + } +} + +/// Take the lock on the process-global snapshot test hooks. Hold it for the +/// whole body of any test that opens a snapshot or touches the [`paging`] +/// statics (module doc); it is a Tokio mutex, so a test that panicked while +/// holding it does not poison it for the next. +/// +/// # Panics +/// +/// When one holder has kept the lock for `HOOKS_LOCK_WAIT` (600 s) while this +/// test waits: that holder is the hung test to look at. A wait behind a +/// queue of holders that each finish in time does not panic, however long +/// the queue is. +pub async fn hooks_lock() -> HooksGuard { + loop { + let since = held_since(); + let deadline = since.unwrap_or_else(std::time::Instant::now) + HOOKS_LOCK_WAIT; + if let Ok(guard) = + tokio::time::timeout_at(tokio::time::Instant::from_std(deadline), HARNESS.lock()).await + { + set_held_since(Some(std::time::Instant::now())); + return HooksGuard { _guard: guard }; + } + // Panic only if the holder the deadline was counted from still holds + // the lock; a newer holder gets its own `HOOKS_LOCK_WAIT`. + let now = held_since(); + assert!( + since.is_none() || now != since, + "another test has held the snapshot hooks lock for {} s; it is the one that hangs", + HOOKS_LOCK_WAIT.as_secs() + ); + } +} + +/// The next item of `stream`, or `None` at its end, within +/// [`DRAIN_TIMEOUT`]; panics naming `stage` and `read` (items read so far) +/// when neither arrives, so a stream that never ends fails the check instead +/// of hanging it and every test behind the hooks lock. +async fn next_within(stream: &mut S, stage: &str, read: usize) -> Option +where + S: futures::Stream + Unpin, +{ + tokio::time::timeout(DRAIN_TIMEOUT, stream.next()) + .await + .unwrap_or_else(|_| { + panic!( + "{stage}: the stream yielded neither an item nor its end within {} s after \ + {read} item(s)", + DRAIN_TIMEOUT.as_secs() + ) + }) +} + +/// Drain `stream` into `into`, each item within [`DRAIN_TIMEOUT`] (see +/// [`next_within`]); `stage` names the drain in the panic. +async fn drain_within(stream: &mut S, stage: &str, into: &mut Vec) +where + S: futures::Stream> + Unpin, +{ + while let Some(next) = next_within(stream, stage, into.len()).await { + into.push(next.unwrap_or_else(|e| panic!("{stage}: snapshot stream item: {e}"))); + } +} + +/// The members of a description that are not fixed at the snapshot instant, +/// cleared so two descriptions of one table compare equal: the table's and +/// each index's item count and size are the stats pass's last values, +/// `table_throughput_mode_summary` is set by the API layer's response +/// preparation and not by storage, and an index being deleted is dropped +/// from a snapshot's description by rule. Everything else (key schema, +/// attribute definitions, every index's definition and status, billing, +/// throughput, class, SSE, ARN, id, creation time, streams, TTL, deletion +/// protection, and `restore_summary`, which storage fills from its restore +/// record on a restored table) must agree. +fn comparable_description(mut description: extenddb_core::types::TableDescription) -> String { + crate::snapshot::drop_deleting_indexes(&mut description); + description.item_count = 0; + description.table_size_bytes = 0; + description.table_throughput_mode_summary = None; + for gsi in description.global_secondary_indexes.iter_mut().flatten() { + gsi.item_count = 0; + gsi.index_size_bytes = 0; + } + for lsi in description.local_secondary_indexes.iter_mut().flatten() { + lsi.item_count = 0; + lsi.index_size_bytes = 0; + } + for vector in description.vector_indexes.iter_mut().flatten() { + vector.item_count = 0; + vector.index_size_bytes = 0; + } + serde_json::to_string_pretty(&description).expect("a description serializes") +} + +/// Assert that `snapshot`'s description is `described`, compared by +/// [`comparable_description`]. +fn assert_same_description( + described: extenddb_core::types::TableDescription, + snapshot: extenddb_core::types::TableDescription, +) { + let want = comparable_description(described); + let got = comparable_description(snapshot); + assert!( + want == got, + "the snapshot's description is not DescribeTable's (counts and throughput mode summary \ + excluded); differing members:\n{}", + description_diff(&want, &got) + ); +} + +/// Describe `table` (quiescent, `ACTIVE`), open a snapshot of it, and +/// assert the two descriptions agree as [`run`] compares them. With +/// `restored`, the table must have a `restore_summary` (asserted on +/// `describe_table`'s side, so the comparison of it is not vacuous): a +/// backend's suite calls this on a table it has given a restore record. +/// +/// # Panics +/// +/// When the describe or the open fails, when `restored` and the table has +/// no `restore_summary`, or when the descriptions differ. +pub async fn assert_snapshot_describes_the_table( + hooks: &HooksGuard, + engine: &E, + account: &str, + table: &str, + restored: bool, +) where + E: SnapshotEngine + TableEngine + Send + Sync, +{ + let _ = hooks; + let described = engine + .describe_table( + account, + DescribeTableInput { + table_name: table.to_owned(), + }, + ) + .await + .expect("describe the table"); + assert_eq!( + described.restore_summary.is_some(), + restored, + "precondition: restore_summary on DescribeTable's description: {:?}", + described.restore_summary + ); + let snapshot = engine + .open_table_snapshot(account, table) + .await + .expect("open the snapshot"); + assert_same_description(described, snapshot.description); +} + +/// The lines of `want` (`DescribeTable`'s description) and `got` (the +/// snapshot's) that the other does not have, for an assertion message. Both +/// are pretty-printed JSON of the same type, so a differing member shows as +/// one line on each side. Lines are compared without a trailing `,`: a member +/// present on one side only (`RestoreSummary`, say) changes whether the line +/// before it ends in one, and that line is not a difference. +fn description_diff(want: &str, got: &str) -> String { + let only_in = |a: &str, b: &str| -> Vec { + let other: BTreeSet<&str> = b.lines().map(|l| l.trim().trim_end_matches(',')).collect(); + a.lines() + .map(|l| l.trim().trim_end_matches(',')) + .filter(|l| !other.contains(l)) + .map(str::to_owned) + .collect() + }; + format!( + " only in DescribeTable's: {:?}\n only in the snapshot's: {:?}", + only_in(want, got), + only_in(got, want) + ) +} + +async fn put(engine: &E, key_info: &TableKeyInfo, item: Item) { + engine + .put_item( + key_info, + item, + false, + None, + &ExpressionMaps::default(), + None, + ) + .await + .expect("put_item"); +} + +async fn delete(engine: &E, key_info: &TableKeyInfo, key: &Item) { + engine + .delete_item(key_info, key, false, None, &ExpressionMaps::default(), None) + .await + .expect("delete_item"); +} + +/// Wait until the table is `ACTIVE`, calling `pump` between polls so a test +/// without background workers can drive the control plane itself. +pub async fn wait_active(engine: &E, account: &str, table: &str, mut pump: F) +where + E: TableEngine, + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + for _ in 0..600 { + pump().await; + let desc = engine + .describe_table( + account, + DescribeTableInput { + table_name: table.to_owned(), + }, + ) + .await + .expect("describe"); + let indexes_active = desc + .global_secondary_indexes + .as_ref() + .is_none_or(|g| g.iter().all(|g| g.index_status == "ACTIVE")); + if desc.table_status == TableStatus::Active && indexes_active { + return; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("table {table} did not become ACTIVE"); +} + +/// Sets the shared snapshot page size ([`paging::set_fetch_batch`]) and +/// restores the default when dropped, so a failed assertion cannot leave a +/// small page for the next test. The one helper for every backend's tests +/// and for [`run`]; hold [`hooks_lock`] while one is alive. +pub struct PageSize(()); + +impl PageSize { + /// Set the page size to `rows` until the returned value is dropped. + #[must_use] + pub fn set(rows: usize) -> Self { + paging::set_fetch_batch(rows); + Self(()) + } +} + +impl Drop for PageSize { + fn drop(&mut self) { + paging::set_fetch_batch(0); + } +} + +/// What the check expects of the backend under test. +#[derive(Debug, Clone, Copy)] +pub struct Expectations { + /// Whether the backend reports `key_ordered` for hash-only and + /// single-range-key tables (the SQL backends). + pub key_ordered: bool, + /// The most significant digits a number key may have: 38 (the service's + /// limit) on the SQL backends, 34 on `MongoDB`, which stores number keys as + /// Decimal128. Range numbers with more are expected to be refused with + /// `StorageError::Validation` and are then not part of the table. + pub numeric_key_digits: usize, + /// The one-instant stages at which an `UpdateTable` issued while the open + /// is held at the pause waits for the open to finish, so the pause is + /// resumed by the 2 s timeout rather than by W2's commit: `PostgreSQL` with + /// a separate data database at `AfterDefinition` (U takes the `tables` + /// row `FOR UPDATE`, which the open holds `FOR SHARE` until the data pin); + /// empty on `PostgreSQL` with one shared database, `SQLite` and `MongoDB`, + /// where nothing the open holds makes U wait. The check asserts which + /// branch ran (lesson 22): a backend whose U and W2 take longer than the + /// timeout would otherwise pass the "W2 not in items" assertion with W2 + /// committed after the resume and the scenario checking nothing. + pub update_table_waits_at: &'static [paging::OpenStage], +} + +/// Significant digits of the longest `N` key attribute of `key` (0 when it +/// has none). +fn max_number_digits(key: &Item) -> usize { + key.values() + .filter_map(|v| match v { + AttributeValue::N(n) => Some(n), + _ => None, + }) + .map(|n| { + let mantissa = n + .trim_start_matches('-') + .split(['e', 'E']) + .next() + .unwrap_or_default(); + let digits: String = mantissa.chars().filter(char::is_ascii_digit).collect(); + digits + .trim_start_matches('0') + .trim_end_matches('0') + .len() + .max(1) + }) + .max() + .unwrap_or(0) +} + +/// Run the consistency check against `engine`. `hooks` is the caller's +/// [`hooks_lock`] guard. `table` must not exist, and +/// neither may tables named with it as a prefix (the check creates +/// `_i0`, `
_i1` and `
_k` too); +/// `pump` is called while waiting for a table to become `ACTIVE`; +/// `page_size` is the number of items every snapshot stream opened during +/// the main scenario fetches per round trip, set through +/// [`paging::set_fetch_batch`] here and restored to the default when the +/// check returns; `expect` states what the backend is known to do +/// ([`Expectations`]), so a backend that stops ordering, or that starts +/// refusing a number it accepted, fails instead of skipping an assertion. +/// +/// After the main scenario (writes while a stream is read), the check runs +/// the one-instant scenario at each [`paging::OpenStage`] and the key-shape +/// scenario. +/// +/// # Panics +/// +/// Panics on any consistency violation, with the violation named, when +/// `page_size` is too large for the table to span several pages, and when the +/// backend did not page at `page_size` (fewer than two round trips after one +/// page plus one item). +#[allow(clippy::too_many_lines)] +pub async fn run( + hooks: &HooksGuard, + engine: &E, + account: &str, + table: &str, + page_size: usize, + expect: Expectations, + mut pump: F, +) where + E: SnapshotEngine + TableEngine + DataEngine + Send + Sync, + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + assert!( + BASE_ITEMS + DOOMED_ITEMS > 3 * page_size, + "page_size {page_size} must be well under the table size {} or the first fetch \ + holds the whole table and the check proves nothing", + BASE_ITEMS + DOOMED_ITEMS + ); + let _ = hooks; + let _restore = PageSize::set(page_size); + engine + .create_table(account, table_input(table)) + .await + .expect("create the check table"); + wait_active(engine, account, table, &mut pump).await; + let ki = key_info(engine, account, table).await; + + for i in 0..BASE_ITEMS { + put(engine, &ki, base_item(i, 0)).await; + } + for i in 0..DOOMED_ITEMS { + put(engine, &ki, doomed_item(i)).await; + } + + // The table is quiescent here (its puts are done, the writer has not + // started), so storage's `describe_table` and the snapshot's + // `description` must agree on everything that is not a running count: + // the trait promises the table as DescribeTable would have reported it. + let described = engine + .describe_table( + account, + DescribeTableInput { + table_name: table.to_owned(), + }, + ) + .await + .expect("describe the check table before the open"); + + // Taken before the open: a backend may fetch its first page while + // opening (MongoDB's `find` returns the first batch), and that fetch is + // part of the paging the check relies on. + let fetches_before_open = paging::fetches(); + let snapshot = engine + .open_table_snapshot(account, table) + .await + .expect("open the snapshot"); + assert_eq!(snapshot.description.table_name, table); + assert_eq!( + snapshot.description.key_schema.len(), + 2, + "composite key schema" + ); + let gsis = snapshot + .description + .global_secondary_indexes + .as_deref() + .unwrap_or(&[]); + assert_eq!(gsis.len(), 1, "the ACTIVE GSI is in the description"); + assert_eq!(gsis[0].index_name, "gi"); + assert_same_description(described, snapshot.description.clone()); + assert!( + !snapshot.info.kind.is_empty(), + "snapshot kind must name the mechanism" + ); + let key_ordered = snapshot.key_ordered; + assert_eq!( + key_ordered, expect.key_ordered, + "key_ordered for a single-range-key table" + ); + let snapshot_description = snapshot.description; + let mut items = snapshot.items; + + // Writer: runs only after the snapshot is open. Signals once the deletes, + // inserts and ROUNDS_DURING_STREAM overwrite rounds have committed, then + // keeps overwriting until told to stop. Reader and writer are joined in + // this task; both make progress at their await points. + let (committed_tx, committed_rx) = tokio::sync::oneshot::channel::(); + let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let writer = { + let ki = ki.clone(); + let stop = stop.clone(); + async move { + for i in 0..DOOMED_ITEMS { + delete(engine, &ki, &key_of(&doomed_item(i))).await; + } + for i in 0..LATE_ITEMS { + put(engine, &ki, late_item(i)).await; + } + let mut round = 0u32; + let mut committed_tx = Some(committed_tx); + loop { + round += 1; + for i in 0..BASE_ITEMS { + put(engine, &ki, base_item(i, round)).await; + } + if round == ROUNDS_DURING_STREAM + && let Some(tx) = committed_tx.take() + { + let _ = tx.send(round); + } + if round >= ROUNDS_DURING_STREAM && stop.load(std::sync::atomic::Ordering::SeqCst) { + return round; + } + } + } + }; + + // Reader: one page plus one item (so a second fetch has happened and the + // backend transaction or cursor is provably still open), then hold until + // the writer has committed everything it will commit during the stream, + // then drain. + let reader = { + let stop = stop.clone(); + async move { + let mut seen: Vec = Vec::new(); + let before_pause = page_size + 1; + for _ in 0..before_pause { + match next_within(&mut items, "main scenario, first page", seen.len()).await { + Some(Ok(item)) => seen.push(item), + Some(Err(e)) => panic!("snapshot stream failed: {e}"), + None => panic!("snapshot ended after fewer than {before_pause} items"), + } + } + // The precondition the rest of the check rests on: the backend + // paged at `page_size`, so item page_size + 1 came from a second + // round trip with the snapshot still open. + let fetches = paging::fetches() - fetches_before_open; + assert!( + fetches >= 2, + "{fetches} round trip(s) after reading {before_pause} items: the backend is not \ + paging at {page_size}, so the check would pass against anything" + ); + tokio::time::timeout(Duration::from_secs(120), committed_rx) + .await + .expect("writer committed its rounds within the deadline") + .expect("writer sent its round count"); + // Fetches recorded by the time the writer's rounds committed: the + // drain below must fetch again, or the stream read everything + // before the writes and the check above shows nothing. + let fetches_at_commit = paging::fetches(); + drain_within( + &mut items, + "main scenario, drain after the writer committed", + &mut seen, + ) + .await; + drop(items); + let fetches_after_drain = paging::fetches(); + assert!( + fetches_after_drain > fetches_at_commit, + "no round trip after the writer committed ({fetches_at_commit} before, \ + {fetches_after_drain} after the drain): the stream read the table before \ + the writes, so the snapshot was not exercised" + ); + stop.store(true, std::sync::atomic::Ordering::SeqCst); + seen + } + }; + let (final_round, seen) = tokio::join!(writer, reader); + + // Safety: the stream is the table as of open. + assert_eq!( + seen.len(), + BASE_ITEMS + DOOMED_ITEMS, + "snapshot item count (base + doomed, no late)" + ); + let mut keys = BTreeSet::new(); + for item in &seen { + assert!(keys.insert(sort_key(item)), "duplicate item {item:?}"); + let pk = match item.get("pk") { + Some(AttributeValue::S(v)) => v.as_str(), + other => panic!("pk {other:?}"), + }; + assert!( + !pk.starts_with("late"), + "item inserted after open is visible: {item:?}" + ); + assert_eq!( + item.get("v"), + Some(&n(0)), + "item overwritten after open shows a post-open version: {item:?}" + ); + } + for i in 0..DOOMED_ITEMS { + assert!( + keys.contains(&sort_key(&doomed_item(i))), + "item deleted after open is missing: doomed {i}" + ); + } + if key_ordered { + let order = extenddb_core::backup_format::KeyOrder::resolve( + &snapshot_description.key_schema, + &snapshot_description.attribute_definitions, + ) + .expect("a key_ordered snapshot has an orderable key schema"); + for pair in seen.windows(2) { + assert_eq!( + order.compare(&pair[0], &pair[1]).expect("comparable keys"), + std::cmp::Ordering::Less, + "key_ordered snapshot must yield strictly ascending keys: {:?} then {:?}", + key_of(&pair[0]), + key_of(&pair[1]) + ); + } + } + + // Liveness through the data: a fresh snapshot sees the writer's final + // state, so the writes above were real and the first snapshot's view was + // the snapshot's, not a frozen database. + let mut after = engine + .open_table_snapshot(account, table) + .await + .expect("open a second snapshot"); + let mut after_items: Vec = Vec::new(); + drain_within( + &mut after.items, + "liveness, second snapshot", + &mut after_items, + ) + .await; + drop(after.items); + assert_eq!(after_items.len(), BASE_ITEMS + LATE_ITEMS); + assert!( + after_items + .iter() + .filter(|i| !matches!(i.get("pk"), Some(AttributeValue::S(p)) if p.starts_with("late"))) + .all(|i| i.get("v") == Some(&n(final_round))), + "second snapshot shows the final round {final_round}" + ); + + for (i, stage) in [ + paging::OpenStage::AfterDefinition, + paging::OpenStage::BeforeItems, + ] + .into_iter() + .enumerate() + { + one_instant( + engine, + account, + &format!("{table}_i{i}"), + stage, + expect, + &mut pump, + ) + .await; + } + key_shapes(engine, account, table, expect, &mut pump).await; +} + +/// Items in a one-instant or key-shape table before its snapshot opens, and +/// the page size those scenarios read at: several pages each. +const SMALL_PAGE: usize = 3; +const INSTANT_ITEMS: usize = 12; + +fn named_item(pk: &str) -> Item { + [ + ("pk".to_owned(), s(pk)), + ("sk".to_owned(), n(0)), + ("v".to_owned(), n(0)), + ] + .into_iter() + .collect() +} + +fn has_pk(items: &[Item], pk: &str) -> bool { + items + .iter() + .any(|i| matches!(i.get("pk"), Some(AttributeValue::S(p)) if p == pk)) +} + +/// One instant: hold an open at `stage` and, while it waits, commit W1 (a new +/// item), then `UpdateTable` U (a GSI on a new attribute), then W2 (another +/// new item) once U has returned. Commit order is W1 < U < W2, so a snapshot +/// of one instant satisfies: the description has U's GSI only if the items +/// have W1, and the items have W2 only if the description has U's GSI. W2 +/// must not be in the items at all: the pause is inside the open's read +/// window, so W2 commits after the instant (module doc). +/// +/// The open resumes when W2 has committed or after 2 s, whichever is first. +/// On `PostgreSQL` with a separate data database at `AfterDefinition` U waits +/// on the `tables` row the open holds `FOR SHARE`, so the timeout resumes it +/// and U and W2 commit after the open; elsewhere W2 commits during the +/// pause. Which of the two happened is asserted against +/// `Expectations::update_table_waits_at`, since the "W2 not in items" +/// assertion is live only in the second case. The reader waits for W2 +/// before it drains, so the drain's round trips are reads after the last +/// commit, and the check asserts the stream made at least one. +#[allow(clippy::too_many_lines)] +async fn one_instant( + engine: &E, + account: &str, + table: &str, + stage: paging::OpenStage, + expect: Expectations, + pump: &mut F, +) where + E: SnapshotEngine + TableEngine + DataEngine + Send + Sync, + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + paging::set_fetch_batch(SMALL_PAGE); + engine + .create_table(account, table_input(table)) + .await + .expect("create the one-instant table"); + wait_active(engine, account, table, &mut *pump).await; + let ki = key_info(engine, account, table).await; + for i in 0..INSTANT_ITEMS { + put(engine, &ki, base_item(i, 0)).await; + } + let update: extenddb_core::types::UpdateTableInput = + serde_json::from_value(serde_json::json!({ + "TableName": table, + "AttributeDefinitions": [{"AttributeName": "u", "AttributeType": "S"}], + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "gu", + "KeySchema": [{"AttributeName": "u", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "KEYS_ONLY"} + }}] + })) + .expect("UpdateTable input"); + + let mut pause = paging::arm_open_pause(stage); + let reached = pause.take_reached(); + let (w1_tx, w1_rx) = tokio::sync::oneshot::channel::<()>(); + // Some(fetches recorded when W2's put returned) once W2 has committed. + let (w2_tx, w2_rx) = tokio::sync::watch::channel::>(None); + + let controller = { + let ki = ki.clone(); + let mut w2_seen = w2_rx.clone(); + async move { + tokio::time::timeout(Duration::from_secs(60), reached) + .await + .unwrap_or_else(|_| panic!("the open did not reach the {stage:?} pause")) + .expect("pause dropped before the open reached it"); + put(engine, &ki, named_item("w1")).await; + let _ = w1_tx.send(()); + let writes = async { + engine + .update_table(account, update) + .await + .expect("UpdateTable adding a GSI"); + put(engine, &ki, named_item("w2")).await; + let _ = w2_tx.send(Some(paging::fetches())); + }; + let resume = async { + let by_w2 = + tokio::time::timeout(Duration::from_secs(2), w2_seen.wait_for(Option::is_some)) + .await + .is_ok(); + pause.resume(); + by_w2 + }; + let ((), by_w2) = tokio::join!(writes, resume); + by_w2 + } + }; + let reader = { + let mut w2 = w2_rx.clone(); + async move { + let snapshot = engine + .open_table_snapshot(account, table) + .await + .expect("open the one-instant snapshot"); + assert!(w1_rx.await.is_ok(), "W1 committed during the pause"); + let fetches_at_w2 = { + let committed = + tokio::time::timeout(Duration::from_secs(60), w2.wait_for(Option::is_some)) + .await + .expect("W2 committed within the deadline") + .expect("W2 sender alive"); + (*committed).expect("W2 committed") + }; + let mut items: Vec = Vec::new(); + let mut stream = snapshot.items; + drain_within(&mut stream, &format!("one-instant {stage:?}"), &mut items).await; + drop(stream); + (snapshot.description, items, fetches_at_w2) + } + }; + let (by_w2, (description, items, fetches_at_w2)) = tokio::join!(controller, reader); + let fetches_after_drain = paging::fetches(); + + let has_gsi = description + .global_secondary_indexes + .iter() + .flatten() + .any(|g| g.index_name == "gu"); + let w1_in = has_pk(&items, "w1"); + let w2_in = has_pk(&items, "w2"); + eprintln!( + "one-instant {stage:?} on {table}: resumed by {}, description has U's GSI: {has_gsi}, \ + W1 in items: {w1_in}, W2 in items: {w2_in}", + if by_w2 { + "W2's commit" + } else { + "the 2 s timeout" + } + ); + // Not vacuous: U and W2 completed (the controller expects both), and the + // table as of either side of them is there. + assert!( + w2_rx.borrow().is_some(), + "W2 committed by the end of the scenario" + ); + // Which branch ran is known per backend and stage (`Expectations`), so + // it is asserted: when U does not wait on the open, W2 must have + // committed inside the pause (the "W2 not in items" assertion below is + // live only then); when it does, the timeout must have resumed the open. + // A slow U or W2 (over the 2 s timeout) on a backend where they should + // not wait fails here instead of passing the scenario vacuously. + let u_waits = expect.update_table_waits_at.contains(&stage); + assert_eq!( + by_w2, + !u_waits, + "{stage:?}: the open was resumed by {}, but on this backend UpdateTable {} at this \ + stage, so it should have been resumed by {}; with the wrong branch the W2 \ + assertion below checks nothing", + if by_w2 { + "W2's commit" + } else { + "the 2 s timeout" + }, + if u_waits { + "waits for the held open" + } else { + "does not wait for the held open" + }, + if u_waits { + "the 2 s timeout (U and W2 commit after the open)" + } else { + "W2's commit (U and W2 commit during the pause, within 2 s)" + } + ); + assert_eq!( + items.len(), + INSTANT_ITEMS + usize::from(w1_in) + usize::from(w2_in), + "every pre-open item, once" + ); + for i in 0..INSTANT_ITEMS { + let (pk, sk) = sort_key(&base_item(i, 0)); + assert!( + items + .iter() + .any(|it| it.get("pk") == Some(&s(pk.clone())) && it.get("sk") == Some(&n(sk))), + "pre-open item {i} missing" + ); + } + assert!( + !has_gsi || w1_in, + "{stage:?}: the description has the GSI UpdateTable added after W1 committed, but \ + the items do not have W1: the definition is newer than the items" + ); + assert!( + !w2_in || has_gsi, + "{stage:?}: the items have W2, committed after UpdateTable returned, but the \ + description does not have UpdateTable's GSI: the items are newer than the definition" + ); + assert!( + !w2_in, + "{stage:?}: W2 committed while the open was held at the pause, yet the items have it: \ + the pause is outside the open's read window, so this scenario checked nothing" + ); + assert!( + fetches_after_drain > fetches_at_w2, + "{stage:?}: no round trip after W2 committed ({fetches_at_w2} then, \ + {fetches_after_drain} after the drain): the stream read its items before the writes" + ); +} + +/// Key shapes other than the main scenario's `S` hash with `N` range. +#[derive(Debug, Clone, Copy)] +enum Shape { + HashS, + HashN, + HashB, + HashSRangeS, + HashSRangeN, + HashSRangeB, +} + +/// Strings whose bytewise order differs from collation order: case, `_`, +/// `-`, `~`, a non-ASCII letter, digits as text. +const TEXT_KEYS: [&str; 19] = [ + "a", "B", "_", "\u{e9}", "A", "b", "Z", "z", "-", "10", "9", "ab", "aB", "Ab", "~", "\u{e9}1", + "a_", "a-", "0", +]; +/// Normalized numbers whose text order differs from numeric order. +const NUMBER_KEYS: [&str; 14] = [ + "10", "9", "-1", "100", "2", "-20", "33", "1000", "7", "-3", "0", "15", "1", "22", +]; +/// Range numbers whose text order differs from numeric order, inserted out +/// of order: negatives, fractions, the smallest and largest magnitudes a +/// number may have, and two 38-digit and two 34-digit values a last digit +/// apart (a lossy numeric type or an order-preserving text encoding that +/// drops digits would tie or swap them; `MongoDB` stores 34 digits and refuses +/// the 38-digit pair, see [`Expectations`]). Normalized before the put, as +/// the engine stores them. +const RANGE_NUMBER_KEYS: [&str; 14] = [ + "10", + "-9.5", + "12345678901234567890123456789012345679", + "0", + "1E-130", + "-100", + "9.9E125", + "0.5", + "-0.001", + "12345678901234567890123456789012345678", + "2", + "-1", + "1234567890123456789012345678901235", + "1234567890123456789012345678901234", +]; +/// Byte strings whose base64 order differs from their byte order (`[0xff]` +/// is `/w==`, `[0xfb]` is `+w==`, `[0x00]` is `AA==`). +const BINARY_KEYS: [&[u8]; 13] = [ + &[0x00], + &[0xff], + &[0xfb], + &[0x3e], + &[0x7f], + &[0x80], + &[0x01], + &[0x00, 0x00], + &[0xff, 0x00], + &[0x10], + &[0xf0], + &[0x41], + &[0x61], +]; + +impl Shape { + const ALL: [Self; 6] = [ + Self::HashS, + Self::HashN, + Self::HashB, + Self::HashSRangeS, + Self::HashSRangeN, + Self::HashSRangeB, + ]; + + fn types(self) -> (ScalarAttributeType, Option) { + match self { + Self::HashS => (ScalarAttributeType::S, None), + Self::HashN => (ScalarAttributeType::N, None), + Self::HashB => (ScalarAttributeType::B, None), + Self::HashSRangeS => (ScalarAttributeType::S, Some(ScalarAttributeType::S)), + Self::HashSRangeN => (ScalarAttributeType::S, Some(ScalarAttributeType::N)), + Self::HashSRangeB => (ScalarAttributeType::S, Some(ScalarAttributeType::B)), + } + } + + fn input(self, name: &str) -> CreateTableInput { + let (hash, range) = self.types(); + let mut key_schema = vec![KeySchemaElement { + attribute_name: "pk".to_owned(), + key_type: KeyType::Hash, + }]; + let mut attribute_definitions = vec![AttributeDefinition { + attribute_name: "pk".to_owned(), + attribute_type: hash, + }]; + if let Some(range) = range { + key_schema.push(KeySchemaElement { + attribute_name: "sk".to_owned(), + key_type: KeyType::Range, + }); + attribute_definitions.push(AttributeDefinition { + attribute_name: "sk".to_owned(), + attribute_type: range, + }); + } + CreateTableInput { + table_name: name.to_owned(), + key_schema, + attribute_definitions, + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + } + } + + /// The keys present before the open, in a fixed order that is not key + /// order. + fn keys(self) -> Vec { + let key = |pk: AttributeValue, sk: Option| -> Item { + std::iter::once(("pk".to_owned(), pk)) + .chain(sk.map(|sk| ("sk".to_owned(), sk))) + .collect() + }; + match self { + Self::HashS => TEXT_KEYS.iter().map(|k| key(s(*k), None)).collect(), + Self::HashN => NUMBER_KEYS + .iter() + .map(|k| key(AttributeValue::N((*k).to_owned()), None)) + .collect(), + Self::HashB => BINARY_KEYS + .iter() + .map(|k| key(AttributeValue::B(k.to_vec()), None)) + .collect(), + Self::HashSRangeS => ["a", "B", "\u{e9}"] + .iter() + .flat_map(|pk| { + ["a", "B", "_", "\u{e9}", "b", "Z", "-"] + .iter() + .map(|sk| key(s(*pk), Some(s(*sk)))) + }) + .collect(), + Self::HashSRangeN => ["a", "B"] + .iter() + .flat_map(|pk| { + RANGE_NUMBER_KEYS.iter().map(|sk| { + let sk = + extenddb_core::validation::number::validate_and_normalize_number(sk) + .unwrap_or_else(|e| panic!("range number {sk}: {e:?}")); + key(s(*pk), Some(AttributeValue::N(sk))) + }) + }) + .collect(), + Self::HashSRangeB => ["a", "B"] + .iter() + .flat_map(|pk| { + BINARY_KEYS + .iter() + .map(|sk| key(s(*pk), Some(AttributeValue::B(sk.to_vec())))) + }) + .collect(), + } + } + + /// Keys inserted after the open, none of them in [`Shape::keys`]. + fn late_keys(self) -> Vec { + let (hash, range) = self.types(); + (0..3u8) + .map(|i| { + let value = |ty: ScalarAttributeType| match ty { + ScalarAttributeType::S => s(format!("late{i}")), + ScalarAttributeType::N => AttributeValue::N(format!("5{i}5")), + ScalarAttributeType::B => AttributeValue::B(vec![0xee, i]), + }; + std::iter::once(("pk".to_owned(), value(hash))) + .chain(range.map(|r| ("sk".to_owned(), value(r)))) + .collect() + }) + .collect() + } +} + +/// A key as a set member. +fn key_id(item: &Item) -> String { + format!("{:?} {:?}", item.get("pk"), item.get("sk")) +} + +fn with_version(key: &Item, version: u32) -> Item { + let mut item = key.clone(); + item.insert("v".to_owned(), n(version)); + item +} + +/// For each [`Shape`]: fill a table, open a snapshot at a page of three, read +/// one page plus one item, then delete three items, insert three and +/// overwrite the rest before draining. The stream must hold every pre-open +/// item once at its pre-open version and nothing inserted, and when +/// `key_ordered` (which must equal `expect.key_ordered`) ascend by +/// `KeyOrder::compare`. Duplicates catch a keyset page that repeats its last +/// row; the order catches a text key sorted by collation rather than +/// bytewise, a `B` range key sorted by base64 text where the format wants +/// bytes (the hash key compares as base64 text; the range key bytewise), and +/// an `N` key sorted as text. +#[allow(clippy::too_many_lines)] +async fn key_shapes( + engine: &E, + account: &str, + table: &str, + expect: Expectations, + pump: &mut F, +) where + E: SnapshotEngine + TableEngine + DataEngine + Send + Sync, + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + paging::set_fetch_batch(SMALL_PAGE); + for (i, shape) in Shape::ALL.into_iter().enumerate() { + let name = format!("{table}_k{i}"); + engine + .create_table(account, shape.input(&name)) + .await + .unwrap_or_else(|e| panic!("create the {shape:?} table: {e}")); + wait_active(engine, account, &name, &mut *pump).await; + let ki = key_info(engine, account, &name).await; + // Numbers with more significant digits than the backend stores are + // refused at the put (MongoDB's Decimal128), so they are not part of + // the table; the refusal is asserted, not tolerated. + let (keys, too_long): (Vec, Vec) = shape + .keys() + .into_iter() + .partition(|key| max_number_digits(key) <= expect.numeric_key_digits); + for key in &too_long { + match engine + .put_item( + &ki, + with_version(key, 0), + false, + None, + &ExpressionMaps::default(), + None, + ) + .await + { + Err(StorageError::Validation(_)) => {} + other => panic!( + "{shape:?}: a number key with more than {} significant digits is refused \ + with a validation error, got {other:?} for {}", + expect.numeric_key_digits, + key_id(key) + ), + } + } + if !too_long.is_empty() { + eprintln!( + "{shape:?}: {} key(s) over {} significant digits refused at the put", + too_long.len(), + expect.numeric_key_digits + ); + } + assert!(keys.len() > 3 * SMALL_PAGE, "{shape:?}: several pages"); + for key in &keys { + put(engine, &ki, with_version(key, 0)).await; + } + + let fetches_before_open = paging::fetches(); + let snapshot = engine + .open_table_snapshot(account, &name) + .await + .unwrap_or_else(|e| panic!("open the {shape:?} snapshot: {e}")); + assert_eq!( + snapshot.key_ordered, expect.key_ordered, + "{shape:?}: key_ordered" + ); + let description = snapshot.description; + let mut stream = snapshot.items; + let mut seen = Vec::new(); + for _ in 0..=SMALL_PAGE { + seen.push( + next_within( + &mut stream, + &format!("key shape {shape:?}, first page"), + seen.len(), + ) + .await + .unwrap_or_else(|| panic!("{shape:?}: stream ended early")) + .expect("stream item"), + ); + } + assert!( + paging::fetches() - fetches_before_open >= 2, + "{shape:?}: the backend is not paging at {SMALL_PAGE}" + ); + for key in &keys[..3] { + delete(engine, &ki, key).await; + } + for key in shape.late_keys() { + put(engine, &ki, with_version(&key, 0)).await; + } + for key in &keys[3..] { + put(engine, &ki, with_version(key, 1)).await; + } + let fetches_at_commit = paging::fetches(); + drain_within( + &mut stream, + &format!("key shape {shape:?}, drain"), + &mut seen, + ) + .await; + drop(stream); + assert!( + paging::fetches() > fetches_at_commit, + "{shape:?}: no round trip after the writes committed" + ); + + let expected: BTreeSet = keys.iter().map(key_id).collect(); + let mut got = BTreeSet::new(); + for item in &seen { + let id = key_id(item); + assert!(got.insert(id.clone()), "{shape:?}: duplicate item {id}"); + assert_eq!( + item.get("v"), + Some(&n(0)), + "{shape:?}: item {id} shows a post-open version" + ); + } + assert_eq!( + got, expected, + "{shape:?}: the stream is the table as of open" + ); + if expect.key_ordered { + let order = extenddb_core::backup_format::KeyOrder::resolve( + &description.key_schema, + &description.attribute_definitions, + ) + .expect("orderable key schema"); + for pair in seen.windows(2) { + assert_eq!( + order.compare(&pair[0], &pair[1]).expect("comparable keys"), + std::cmp::Ordering::Less, + "{shape:?}: not in key order: {} then {}", + key_id(&pair[0]), + key_id(&pair[1]) + ); + } + } + } +} + +/// Checks that a snapshot cannot be opened on a missing or non-`ACTIVE` +/// table. `hooks` is the caller's [`hooks_lock`] guard. `table` must not +/// exist; `set_updating` flips that table's catalog +/// status to `UPDATING` through the backend test's own catalog access, so +/// every backend covers the non-`ACTIVE` refusal with the same assertion. +pub async fn run_refusals( + hooks: &HooksGuard, + engine: &E, + account: &str, + table: &str, + pump: F, + set_updating: S, +) where + E: SnapshotEngine + TableEngine + Send + Sync, + F: FnMut() -> Fut, + Fut: std::future::Future, + S: FnOnce(String) -> SFut, + SFut: std::future::Future, +{ + let _ = hooks; + let err = engine + .open_table_snapshot(account, "does_not_exist_for_snapshot") + .await + .expect_err("missing table refused"); + assert!(matches!(err, StorageError::TableNotFound(_)), "{err}"); + + engine + .create_table(account, table_input(table)) + .await + .expect("create the refusal table"); + wait_active(engine, account, table, pump).await; + engine + .open_table_snapshot(account, table) + .await + .expect("an ACTIVE table opens"); + set_updating(table.to_owned()).await; + let err = engine + .open_table_snapshot(account, table) + .await + .expect_err("a non-ACTIVE table is refused"); + assert!(matches!(err, StorageError::TableNotActive(_)), "{err}"); +} + +#[cfg(test)] +mod tests { + use super::description_diff; + + /// A member on one side only is the one difference reported, not also + /// the line before it whose trailing comma it changes. + #[test] + fn description_diff_ignores_a_trailing_comma() { + let want = "{\n \"A\": 1,\n \"B\": false,\n \"RestoreSummary\": {}\n}"; + let got = "{\n \"A\": 1,\n \"B\": false\n}"; + assert_eq!( + description_diff(want, got), + " only in DescribeTable's: [\"\\\"RestoreSummary\\\": {}\"]\n only in the snapshot's: []" + ); + } +} diff --git a/crates/storage/src/test_pause.rs b/crates/storage/src/test_pause.rs new file mode 100644 index 000000000..ee0683595 --- /dev/null +++ b/crates/storage/src/test_pause.rs @@ -0,0 +1,133 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! One pause machinery for the test hooks that hold an operation at a chosen +//! point (`snapshot::paging::arm_open_pause`, the `PostgreSQL` backend's +//! `arm_delete_table_pause`), so a test can act in the window the point +//! opens. Compiled only with `test-hooks`. +//! +//! A [`PausePoint`] is a process-global slot (a `static`) that the operation +//! calls [`PausePoint::reached`] on; a test arms it with [`PausePoint::arm`] +//! and gets a [`Pause`] handle. The next operation that reaches the point +//! with the armed key signals the handle and waits until the test resumes it +//! or drops the handle. Every test that arms a point holds +//! `snapshot_conformance::hooks_lock`, so no foreign operation can take it. + +use std::sync::{Mutex, PoisonError}; + +use tokio::sync::oneshot; + +struct Armed { + key: K, + reached: oneshot::Sender<()>, + resume: oneshot::Receiver<()>, +} + +/// A point an operation can be held at, keyed by `K` (`()` for a point with +/// one place, an enum for one with several). +pub struct PausePoint { + slot: Mutex>>, +} + +impl PausePoint { + /// An unarmed point, for a `static`. + #[must_use] + pub const fn new() -> Self { + Self { + slot: Mutex::new(None), + } + } + + /// Arm the point at `key` for the next operation that reaches it (one + /// operation, once). + /// + /// # Panics + /// + /// In debug builds, when the point is already armed: a second arm would + /// replace the first silently, and the first handle's drop would then + /// disarm the second. + #[must_use] + pub fn arm(&'static self, key: K) -> Pause { + let (reached_tx, reached_rx) = oneshot::channel(); + let (resume_tx, resume_rx) = oneshot::channel(); + let mut slot = self.slot.lock().unwrap_or_else(PoisonError::into_inner); + debug_assert!( + slot.is_none(), + "a pause is already armed at {:?}; drop its handle before arming another", + slot.as_ref().map(|a| a.key) + ); + *slot = Some(Armed { + key, + reached: reached_tx, + resume: resume_rx, + }); + drop(slot); + Pause { + point: self, + reached: Some(reached_rx), + resume: Some(resume_tx), + } + } + + /// Called by the operation at the point: when the point is armed at + /// `key`, disarm it, signal the handle and wait until the test resumes or + /// drops it. Otherwise returns at once. + pub async fn reached(&self, key: K) { + let armed = { + let mut slot = self.slot.lock().unwrap_or_else(PoisonError::into_inner); + if slot.as_ref().is_some_and(|a| a.key == key) { + slot.take() + } else { + None + } + }; + if let Some(armed) = armed { + let _ = armed.reached.send(()); + let _ = armed.resume.await; + } + } + + fn disarm(&self) { + *self.slot.lock().unwrap_or_else(PoisonError::into_inner) = None; + } +} + +impl Default for PausePoint { + fn default() -> Self { + Self::new() + } +} + +/// A pause armed by [`PausePoint::arm`]. Dropping it disarms the point and +/// releases an operation waiting at it, so a failing test cannot leave a +/// pause armed for the next operation in the process. +pub struct Pause { + point: &'static PausePoint, + reached: Option>, + resume: Option>, +} + +impl Pause { + /// The receiver that resolves when the next operation reaches the point. + /// + /// # Panics + /// + /// When called twice. + pub fn take_reached(&mut self) -> oneshot::Receiver<()> { + self.reached.take().expect("take_reached called once") + } + + /// Let the operation continue. + pub fn resume(&mut self) { + if let Some(tx) = self.resume.take() { + let _ = tx.send(()); + } + } +} + +impl Drop for Pause { + fn drop(&mut self) { + self.point.disarm(); + self.resume(); + } +} diff --git a/devtools/run-mongodb-tests b/devtools/run-mongodb-tests index 14ecb9867..eeff98bb8 100755 --- a/devtools/run-mongodb-tests +++ b/devtools/run-mongodb-tests @@ -18,6 +18,11 @@ # (default: /tmp/run-mongodb-tests-) # --keep Do not tear down the mongo container / server # on exit (useful for post-run inspection) +# --mongo-only Start the replica set, print its URI, and exit without +# writing a config or starting extenddb; implies --keep. +# For tests that talk to MongoDB directly (the storage +# crate's snapshot tests). Tear down with +# docker rm -f extenddb-runtests-mongo-. # -h, --help Show this help # # Arguments after `--` are passed to `devtools/run-tests` verbatim. @@ -48,6 +53,7 @@ BIND_PORT=18443 MONGO_PORT=27021 OUTPUT_DIR="" KEEP=false +MONGO_ONLY=false RUN_TESTS_ARGS=() while [[ $# -gt 0 ]]; do @@ -56,8 +62,9 @@ while [[ $# -gt 0 ]]; do --mongo-port) MONGO_PORT="$2"; shift 2 ;; --output) OUTPUT_DIR="$2"; shift 2 ;; --keep) KEEP=true; shift ;; + --mongo-only) MONGO_ONLY=true; KEEP=true; shift ;; --) shift; RUN_TESTS_ARGS=("$@"); break ;; - -h|--help) sed -n '5,35p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; + -h|--help) sed -n '5,40p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; *) echo "error: unknown option: $1" >&2; exit 1 ;; esac done @@ -73,7 +80,10 @@ fi if [[ -z "$OUTPUT_DIR" ]]; then OUTPUT_DIR="/tmp/run-mongodb-tests-${MONGO_PORT}-$(date +%Y%m%d-%H%M%S)-$$" fi -mkdir -p "$OUTPUT_DIR" +# --mongo-only writes nothing, so it creates no directory. +if ! $MONGO_ONLY; then + mkdir -p "$OUTPUT_DIR" +fi CONFIG="$OUTPUT_DIR/extenddb.toml" SERVER_LOG="$OUTPUT_DIR/server.log" @@ -85,7 +95,7 @@ BINARY="./target/release/extenddb" # compiled only by the `mongodb-test-hooks` feature. export EXTENDDB_TEST_MONGODB_TEST_HOOKS=1 -if [[ ! -x "$BINARY" ]]; then +if ! $MONGO_ONLY && [[ ! -x "$BINARY" ]]; then echo "error: $BINARY not found. Build first: cargo build --release --no-default-features --features mongodb-test-hooks" >&2 exit 1 fi @@ -94,7 +104,7 @@ fi # capability before starting Docker or ExtendDB so a plain `mongodb` binary # fails with the build error instead of timing out while waiting for a gate # that was compiled out. -if ! "$BINARY" version 2>/dev/null | grep -q '^test-hooks enabled$'; then +if ! $MONGO_ONLY && ! "$BINARY" version 2>/dev/null | grep -q '^test-hooks enabled$'; then echo "error: $BINARY does not include MongoDB test hooks; rebuild with:" >&2 echo " cargo build --release --no-default-features --features mongodb-test-hooks" >&2 exit 1 @@ -110,6 +120,10 @@ TMP_CANON=$(realpath /tmp) # --------------------------------------------------------------------------- cleanup() { + if $MONGO_ONLY; then + echo " mongo container: $CONTAINER_NAME (port $MONGO_PORT); tear down with: docker rm -f $CONTAINER_NAME" + return + fi if $KEEP; then echo "" echo "=== --keep set; leaving mongo + extenddb running ===" @@ -135,7 +149,11 @@ trap cleanup EXIT echo "=== Starting MongoDB 7 replica set on host port $MONGO_PORT ===" docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true -docker run -d --name "$CONTAINER_NAME" -p "$MONGO_PORT:27017" mongo:7 \ +# Published on the loopback interface only: the replica set has no +# authentication, and everything that connects to it (the printed URI, the +# server config below, the CI job) uses 127.0.0.1. `--bind_ip_all` applies +# inside the container, where Docker's port forwarding arrives. +docker run -d --name "$CONTAINER_NAME" -p "127.0.0.1:$MONGO_PORT:27017" mongo:7 \ mongod --replSet rs0 --bind_ip_all --wiredTigerCacheSizeGB 1 >/dev/null # Wait for mongod to accept pings @@ -161,6 +179,16 @@ for i in $(seq 1 30); do fi sleep 1 done +if [[ "$state" != "PRIMARY" ]]; then + echo "error: replica set did not elect a PRIMARY" >&2 + exit 1 +fi + +if $MONGO_ONLY; then + echo "" + echo "EXTENDDB_TEST_MONGODB_URI=mongodb://127.0.0.1:$MONGO_PORT/?directConnection=true" + exit 0 +fi # --------------------------------------------------------------------------- # Initialize extenddb @@ -176,7 +204,7 @@ cat > "$CONFIG" <', '_ddb_']) AS n); +``` + **Fix:** If a table appears stuck: 1. Check that extenddb is running (`extenddb status`). -2. Wait for the configured `control_plane_delay_seconds` (default: 5) — the poller is woken immediately when a table enters a transitional state. -3. If the server was restarted, transitions are recovered automatically at startup. +2. Wait for the configured `control_plane_delay_seconds` (default: 0.25) — the poller is woken immediately when a table enters a transitional state. +3. On PostgreSQL, look for a session holding a lock on the table (above); the delete finishes within about a second of its end. +4. If the server was restarted, transitions are recovered automatically at startup. + +### Data tables without a catalog row (PostgreSQL, separate data database) + +**Cause:** With a separate data database, a synchronous DeleteTable commits the catalog row removal first, then commits the drop of the table's data tables on the data database. If that data commit fails (the connection is terminated, a failover, a network fault), the server redoes the drop on a fresh connection for up to 60 seconds. Only a process that dies, or a data database unreachable for that whole window, between the catalog commit and the last retry leaves the table's data tables on the data database with no catalog row naming them. Nothing reclaims them afterwards: DescribeTable reports the table gone, a same-name CreateTable gets a new table id, and the control plane only enumerates relations that still have a catalog row. When the retries are exhausted the server logs, at error, the table id, the relation names, and the exact `DROP TABLE` to run. (On one shared database the row and the drops commit together, so this cannot happen there.) + +To find such relations, compare the `_ddb_%` tables on the data database with the table ids the catalog still knows. On the catalog database: + +```sql +SELECT table_id FROM tables +UNION ALL SELECT index_id FROM indexes +UNION ALL SELECT index_id FROM vector_indexes; +``` + +On the data database, every data relation and whether the catalog still names it (replace the `VALUES` list with the ids from the first query): + +```sql +SELECT d.tablename + FROM pg_tables d + LEFT JOIN (VALUES (''), (''), ('')) AS live(id) + ON d.tablename = '_ddb_' || live.id + OR d.tablename = '_ddb_vec_' || live.id + WHERE d.tablename LIKE '\_ddb\_%' AND live.id IS NULL; +``` + +Each row is a data relation with no catalog row. After confirming the table is gone from the catalog (`DescribeTable` returns `ResourceNotFoundException`), drop the listed relations on the data database: + +```sql +DROP TABLE IF EXISTS "_ddb_", "_ddb_", "_ddb_vec_"; +``` + +The error log line names the exact relations for one leaked table and gives this `DROP TABLE` verbatim. + +**Do not** use `extenddb catalog-check --fix` to clean these up: it expects a different relation-naming scheme and, against the current schema, reports every live data table as orphaned and would drop tables still in use. ### CreateTable returns CREATING instead of ACTIVE -**Cause:** extenddb emulates real DynamoDB's async control plane behavior. Tables start in CREATING state and transition to ACTIVE after a configurable delay (default: 5 seconds). +**Cause:** extenddb emulates real DynamoDB's async control plane behavior. Tables start in CREATING state and transition to ACTIVE after a configurable delay (`control_plane_delay_seconds`, default: 0.25 seconds). **Fix:** This is expected behavior matching real DynamoDB. Poll with DescribeTable until `TableStatus` is `ACTIVE` before performing operations on the table. All test code should use a `wait_for_active()` helper after CreateTable. diff --git a/tests/conftest.py b/tests/conftest.py index 1d35d88da..971547dc8 100755 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -72,6 +72,35 @@ def _is_real_dynamodb() -> bool: return not os.environ.get("EXTENDDB_TEST_ENDPOINT", "").strip() +def config_value(section: str, key: str) -> str | None: + """`key` of the `[section]` table in the server's config file + (`EXTENDDB_CONFIG`, else `extenddb.toml`), unquoted, or None when the file + or the key is missing. Reads plain `key = "value"` lines only, which is + what `extenddb init` writes.""" + path = os.environ.get("EXTENDDB_CONFIG", "").strip() or "extenddb.toml" + try: + with open(path, encoding="utf-8") as f: + text = f.read() + except OSError: + return None + current = None + for line in text.splitlines(): + line = line.split("#", 1)[0].strip() + if line.startswith("["): + current = line.strip("[]").strip() + elif current == section and "=" in line and line.split("=", 1)[0].strip() == key: + return line.split("=", 1)[1].strip().strip('"').strip("'") + return None + + +def configured_backend() -> str | None: + """The backend of the server under test: `EXTENDDB_TEST_BACKEND` (exported + by devtools/run-tests), else the `[storage] backend` of its config file, + else None (real DynamoDB, or a server started elsewhere).""" + explicit = os.environ.get("EXTENDDB_TEST_BACKEND", "").strip() + return explicit or config_value("storage", "backend") + + def _poll_interval() -> float: """Return polling interval: 200ms for real DynamoDB, 20ms for ExtendDB.""" return 0.2 if _is_real_dynamodb() else 0.02 diff --git a/tests/test_table_operations.py b/tests/test_table_operations.py index 43c72ec9f..568702c1b 100755 --- a/tests/test_table_operations.py +++ b/tests/test_table_operations.py @@ -9,11 +9,13 @@ from __future__ import annotations +import contextlib +import os import uuid import pytest from botocore.exceptions import ClientError -from conftest import wait_for_active +from conftest import _is_real_dynamodb, config_value, configured_backend, wait_for_active class TestCreateTable: """CreateTable operation tests.""" @@ -175,6 +177,111 @@ def test_delete_protected_table(self, create_and_cleanup_table, dynamodb_client, DeletionProtectionEnabled=False, ) + def test_delete_of_a_deleting_table_returns_its_description( + self, create_and_cleanup_table, dynamodb_client, unique_table_name + ): + """DeleteTable of a table that is already DELETING returns the table's + description with TableStatus DELETING and no error, as the service does, + and leaves the table as it was. + + On ExtendDB the DELETING state is written into the catalog directly: on + PostgreSQL a reader or writer holding the table leaves it there, and on + SQLite a control_plane_delay_seconds of 1 or more, neither of which a + test sharing the server can arrange without affecting other tests. + MongoDB and Cassandra remove a table in the request and never leave it + DELETING, so there the test does not apply. + """ + create_and_cleanup_table(unique_table_name) + if _is_real_dynamodb(): + dynamodb_client.delete_table(TableName=unique_table_name) + second = dynamodb_client.delete_table(TableName=unique_table_name) + assert second["TableDescription"]["TableStatus"] == "DELETING" + return + backend = configured_backend() or ( + "postgres" if config_value("storage.postgres", "connection_string") else None + ) + print(f"backend under test: {backend}") + if backend not in ("postgres", "sqlite"): + pytest.skip(f"backend {backend!r} never leaves a table DELETING") + with _catalog_connection(backend) as catalog: + _set_status(catalog, backend, unique_table_name, "DELETING") + try: + before = _status_row(catalog, backend, unique_table_name) + second = dynamodb_client.delete_table(TableName=unique_table_name) + desc = second["TableDescription"] + assert desc["TableStatus"] == "DELETING" + assert desc["TableName"] == unique_table_name + assert desc["TableArn"].endswith(f"table/{unique_table_name}") + assert desc["KeySchema"] == [{"AttributeName": "pk", "KeyType": "HASH"}] + described = dynamodb_client.describe_table(TableName=unique_table_name) + assert described["Table"]["TableStatus"] == "DELETING" + assert _status_row(catalog, backend, unique_table_name) == before, ( + "the second DeleteTable changed the table" + ) + finally: + # Back to ACTIVE, so the fixture's DeleteTable removes it. + _set_status(catalog, backend, unique_table_name, "ACTIVE") + + +@contextlib.contextmanager +def _catalog_connection(backend: str): + """A DB-API connection on the running server's catalog, from its config.""" + if backend == "postgres": + import psycopg2 + + url = config_value("storage.postgres", "connection_string") + assert url, "no [storage.postgres] connection_string in the server's config" + conn = psycopg2.connect(url) + conn.autocommit = True + else: + import sqlite3 + + path = config_value("storage.sqlite", "path") + assert path and ":memory:" not in path, ( + f"the SQLite catalog must be a file the test can open: {path!r}" + ) + conn = sqlite3.connect(os.path.abspath(path), timeout=30, isolation_level=None) + try: + yield conn + finally: + conn.close() + + +def _placeholder(backend: str) -> str: + return "%s" if backend == "postgres" else "?" + + +def _set_status(conn, backend: str, table_name: str, status: str) -> None: + """Set the table's catalog status: DELETING with a removal an hour away, or + back to ACTIVE with none scheduled.""" + p = _placeholder(backend) + if status == "DELETING": + when = ( + "NOW() + INTERVAL '1 hour'" + if backend == "postgres" + else "strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '+1 hour')" + ) + else: + when = "NULL" + cur = conn.cursor() + cur.execute( + f"UPDATE tables SET table_status = {p}, status_transition_at = {when} " + f"WHERE table_name = {p}", + (status, table_name), + ) + assert cur.rowcount == 1, f"{table_name} has one catalog row" + + +def _status_row(conn, backend: str, table_name: str): + cur = conn.cursor() + p = _placeholder(backend) + cur.execute( + f"SELECT table_status, CAST(status_transition_at AS TEXT) FROM tables " + f"WHERE table_name = {p}", + (table_name,), + ) + return cur.fetchall() + # --------------------------------------------------------------------------- # CreateTable validation additions (covers commits since 6b98234dcf)