From d1604c567aa424be5b8dff6b3013b4e98c9aad48 Mon Sep 17 00:00:00 2001 From: Lee Hannigan Date: Fri, 7 Aug 2026 14:58:22 +0000 Subject: [PATCH] fix(postgres): restore must not report ACTIVE before the data copy completes DescribeTable's contract is that the first ACTIVE observation on a restore target implies the restored data is fully present. The Postgres restore path created the target through the normal create-table flow, which schedules a CREATING->ACTIVE transition on the creation-delay timer. The row-by-row copy of a large backup outlasts that timer, so a client that waits-for-ACTIVE can observe an empty or partial table (reproduced live: ACTIVE at 515/40000 items on baseline). Fix: create_table_impl gains a defer_active flag. The restore path passes true, writing the row CREATING with no scheduled transition; the existing explicit ACTIVE update at the end of the copy is now the only flip. The public TableEngine::create_table path passes false and is unchanged. Verification (live, Postgres 16 at 127.0.0.1:5433): - New test restore_active_completeness: observer spawns before restore, polls DescribeTable, counts items via paginated Select=COUNT scan at first ACTIVE. Transient describe errors are retried, and scan errors fail the test rather than truncating the count (avoids the oracle bug that makes the equivalent upstream test flaky). - Negative control on baseline: FAILED with ACTIVE at 515/40000 (112s). - With fix: single test PASS (129s); full integration suite --test-threads=1: 418/421 passed, 0 filtered. The 3 reds are capacity_throttling tests requiring 'settings set throttling_enabled true' which the ad-hoc harness omitted; CI sets it (integration.yml). --- crates/storage-postgres/src/backup_engine.rs | 23 +-- crates/storage-postgres/src/create_table.rs | 14 +- crates/storage-postgres/src/table_engine.rs | 2 +- tests/rust/src/main.rs | 5 +- tests/rust/src/restore_active_completeness.rs | 151 ++++++++++++++++++ 5 files changed, 180 insertions(+), 15 deletions(-) create mode 100644 tests/rust/src/restore_active_completeness.rs diff --git a/crates/storage-postgres/src/backup_engine.rs b/crates/storage-postgres/src/backup_engine.rs index 6486f522..ee02a17c 100755 --- a/crates/storage-postgres/src/backup_engine.rs +++ b/crates/storage-postgres/src/backup_engine.rs @@ -8,7 +8,6 @@ use extenddb_core::types::{ PointInTimeRecoveryDescription, SourceTableDetails, TableDescription, }; use extenddb_storage::BackupEngine; -use extenddb_storage::TableEngine; use extenddb_storage::error::StorageError; use futures::future::BoxFuture; @@ -451,16 +450,20 @@ impl BackupEngine for PostgresEngine { on_demand_throughput: None, }; - let desc = self.create_table(&account_id, create_input).await?; + // Create the table with the ACTIVE transition deferred: it enters + // CREATING with no scheduled flip, so the control-plane worker + // cannot mark it ACTIVE while the item copy below is still + // running. The explicit ACTIVE update after the copy is the only + // way this table leaves CREATING, so ACTIVE implies the restored + // data is fully present. + let desc = self + .create_table_impl(&account_id, create_input, true) + .await?; let new_table_id = &desc.table_id; let ddb_table = data_table_name(new_table_id); let ddb_table_unquoted = ddb_table.trim_matches('"'); - // Do NOT force ACTIVE — let the control plane handle the transition - // (steering rule D-2: tests run with control_plane_delay_seconds > 0). - // The table starts in CREATING and transitions to ACTIVE after the delay. - let has_sk: bool = sqlx::query_scalar( "SELECT EXISTS(SELECT 1 FROM information_schema.columns \ WHERE table_name = $1 AND column_name = 'sk')", @@ -513,10 +516,10 @@ impl BackupEngine for PostgresEngine { .await .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - // Mark the restored table ACTIVE immediately — the data is fully - // populated and the table is ready to serve requests. This matches - // real DynamoDB behavior where restored tables become ACTIVE once - // the restore completes (the CREATING status is transient). + // Mark the restored table ACTIVE now that the copy has fully + // drained. The table was created with the transition deferred, so + // this is the first point it can become ACTIVE, by ordering rather + // than by timing. sqlx::query( "UPDATE tables SET table_status = 'ACTIVE', status_transition_at = NULL \ WHERE account_id = $1 AND table_name = $2", diff --git a/crates/storage-postgres/src/create_table.rs b/crates/storage-postgres/src/create_table.rs index b808a1fa..9ef2e39d 100755 --- a/crates/storage-postgres/src/create_table.rs +++ b/crates/storage-postgres/src/create_table.rs @@ -14,10 +14,17 @@ use crate::PostgresEngine; impl PostgresEngine { /// Core implementation of `create_table` (Fix #4: wrapped in a transaction). + /// Create a table. When `defer_active` is set (the restore path), the row + /// is written `CREATING` with **no** scheduled transition, so the + /// background control-plane worker cannot flip it to `ACTIVE` while the + /// caller is still populating it; the caller sets `ACTIVE` itself once the + /// data copy completes. Normal `CreateTable` passes `false` and gets the + /// usual timed transition. pub(crate) async fn create_table_impl( &self, account_id: &str, input: CreateTableInput, + defer_active: bool, ) -> Result { Self::validate_account_id(account_id)?; let table_id = uuid::Uuid::new_v4().to_string(); @@ -77,10 +84,12 @@ impl PostgresEngine { creation_date_time, table_arn, table_id, deletion_protection_enabled, status_transition_at, table_class, sse_specification, on_demand_throughput) VALUES ($1, $2, $3, $4, $5, $6, $7, - CASE WHEN (SELECT secs FROM delay) = 0 + CASE WHEN $14 THEN 'CREATING' + WHEN (SELECT secs FROM delay) = 0 THEN 'ACTIVE' ELSE 'CREATING' END, NOW(), $8, $9, $10, - CASE WHEN (SELECT secs FROM delay) = 0 + CASE WHEN $14 THEN NULL + WHEN (SELECT secs FROM delay) = 0 THEN NULL ELSE NOW() + make_interval(secs => (SELECT secs FROM delay)) END, @@ -100,6 +109,7 @@ impl PostgresEngine { .bind(&input.table_class) .bind(&sse_spec_json) .bind(&on_demand_json) + .bind(defer_active) .fetch_one(&mut *tx) .await .map_err(|e| match &e { diff --git a/crates/storage-postgres/src/table_engine.rs b/crates/storage-postgres/src/table_engine.rs index a3a4e1d4..09c871bc 100755 --- a/crates/storage-postgres/src/table_engine.rs +++ b/crates/storage-postgres/src/table_engine.rs @@ -22,7 +22,7 @@ impl TableEngine for PostgresEngine { input: CreateTableInput, ) -> BoxFuture<'_, Result> { let account_id = account_id.to_string(); - Box::pin(async move { self.create_table_impl(&account_id, input).await }) + Box::pin(async move { self.create_table_impl(&account_id, input, false).await }) } // H-5: Set status to DELETING with a scheduled transition to removal, diff --git a/tests/rust/src/main.rs b/tests/rust/src/main.rs index 1befdc6f..9287f713 100755 --- a/tests/rust/src/main.rs +++ b/tests/rust/src/main.rs @@ -7,6 +7,7 @@ //! Requires a running extenddb instance with `EXTENDDB_TEST_ENDPOINT` and credentials set. mod helpers; +mod restore_active_completeness; mod test_base; #[cfg(test)] @@ -18,10 +19,10 @@ mod backup_restore; #[cfg(test)] mod batch_get_item; #[cfg(test)] -mod batch_write_item; -#[cfg(test)] mod batch_transact_authz; #[cfg(test)] +mod batch_write_item; +#[cfg(test)] mod binary_sort_key; #[cfg(test)] mod capacity_throttling; diff --git a/tests/rust/src/restore_active_completeness.rs b/tests/rust/src/restore_active_completeness.rs new file mode 100644 index 00000000..cc4ea758 --- /dev/null +++ b/tests/rust/src/restore_active_completeness.rs @@ -0,0 +1,151 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! A restored table must not report ACTIVE before its data copy completes. +//! +//! DynamoDB's contract: when `DescribeTable` on a restore target first +//! returns ACTIVE, the restored data is fully present. A restore that flips +//! ACTIVE on a control-plane timer decoupled from the copy exposes an empty +//! or partial table to a client that waits-for-ACTIVE. + +use crate::test_base::*; +use aws_sdk_dynamodb::types::{ + AttributeDefinition, AttributeValue, BillingMode, KeySchemaElement, KeyType, PutRequest, + ScalarAttributeType, Select, WriteRequest, +}; + +const ITEMS: usize = 40000; + +#[tokio::test] +async fn restored_table_has_all_items_when_first_active() { + let c = client(); + let src = format!("RestoreRaceSrc_{}", ts()); + c.create_table() + .table_name(&src) + .key_schema( + KeySchemaElement::builder() + .attribute_name("pk") + .key_type(KeyType::Hash) + .build() + .unwrap(), + ) + .attribute_definitions( + AttributeDefinition::builder() + .attribute_name("pk") + .attribute_type(ScalarAttributeType::S) + .build() + .unwrap(), + ) + .billing_mode(BillingMode::PayPerRequest) + .send() + .await + .unwrap(); + wait_for_active(&c, &src).await; + + // Enough data that the restore copy takes longer than the control-plane + // transition delay, so a timer-driven ACTIVE flip would win the race. + let pad = "x".repeat(2000); + for chunk in (0..ITEMS).collect::>().chunks(25) { + let reqs: Vec = chunk + .iter() + .map(|i| { + WriteRequest::builder() + .put_request( + PutRequest::builder() + .item("pk", AttributeValue::S(format!("k{i:06}"))) + .item("d", AttributeValue::S(pad.clone())) + .build() + .unwrap(), + ) + .build() + }) + .collect(); + c.batch_write_item() + .request_items(&src, reqs) + .send() + .await + .unwrap(); + } + + let backup = c + .create_backup() + .table_name(&src) + .backup_name("restore-race-probe") + .send() + .await + .unwrap(); + let arn = backup.backup_details().unwrap().backup_arn().to_string(); + // Wait until the backup is AVAILABLE. + for _ in 0..240 { + let d = c.describe_backup().backup_arn(&arn).send().await.unwrap(); + if d.backup_description() + .and_then(|b| b.backup_details()) + .map(|b| b.backup_status().as_str() == "AVAILABLE") + .unwrap_or(false) + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + } + + let dst = format!("RestoreRaceDst_{}", ts()); + // This backend's RestoreTableFromBackup blocks until the copy completes + // (real DynamoDB returns immediately), so the ACTIVE-before-data window + // is only observable by a CONCURRENT client. Race an observer task + // against the in-flight restore: the moment it sees ACTIVE, it counts. + let observer = { + let c2 = client().clone(); + let dst2 = dst.clone(); + tokio::spawn(async move { + // Wait for the table to exist, then for first ACTIVE. + loop { + match c2.describe_table().table_name(&dst2).send().await { + Ok(out) => { + let st = out.table().unwrap().table_status().unwrap().clone(); + if st.as_str() == "ACTIVE" { + break; + } + } + Err(_) => {} + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + // First ACTIVE observation: count items immediately. + let mut count = 0usize; + let mut start_key = None; + loop { + let mut req = c2.scan().table_name(&dst2).select(Select::Count); + if let Some(k) = start_key.take() { + req = req.set_exclusive_start_key(Some(k)); + } + let resp = req.send().await.unwrap(); + count += resp.count() as usize; + match resp.last_evaluated_key() { + Some(k) if !k.is_empty() => start_key = Some(k.clone()), + _ => break, + } + } + count + }) + }; + + c.restore_table_from_backup() + .target_table_name(&dst) + .backup_arn(&arn) + .send() + .await + .unwrap(); + + // The concurrent observer counted at first-ACTIVE while the restore call + // above was still in flight (or just after, if the copy was fast). + let count = observer.await.unwrap(); + + c.delete_table().table_name(&src).send().await.ok(); + c.delete_table().table_name(&dst).send().await.ok(); + + assert_eq!( + count, ITEMS, + "restored table reported ACTIVE with {count}/{ITEMS} items present — \ + ACTIVE must imply the restore copy is complete" + ); +}