From f721d7421db7e5a8d2a8fb38a8244e0b96a4097b Mon Sep 17 00:00:00 2001 From: Scott Robinson Date: Fri, 23 Oct 2026 09:38:59 +0000 Subject: [PATCH] fix(postgres): build a GSI added by UpdateTable correctly while writes continue Adding a GSI to a PostgreSQL table that was taking writes left the index with missing, extra and stale rows, or failed UpdateTable with a duplicate key. The index row was committed ACTIVE before its data table existed, the backfill paged with OFFSET under concurrent writes, and queue rows applied the image they carried, so a late row overwrote newer state. UpdateTable now creates the data table first, publishes the index CREATING, waits 1 s, takes a brief SHARE lock on the base table to wait out in-flight writes, backfills in keyset pages under FOR SHARE SKIP LOCKED, sweeps orphan rows, and flips the index ACTIVE. Writes maintain a CREATING index like an ACTIVE one. The synchronous write path, the queue and the backfill all write an item's current state, read under its row lock. A write that took more than 500 ms from its catalog read re-reads the catalog and queues a repair for any GSI it did not know about; that and the wait cover writes that read the catalog just before the index was created, which no transaction can see because the catalog and the data are separate databases. A build that fails for a transient reason stays CREATING and a background worker finishes it; one that would fail every time removes the index. On both PostgreSQL and SQLite the write path takes index attribute definitions from its own catalog read instead of the cached key info, which predated a new GSI's key attributes until UpdateTable returned, so a GSI with a sort key on a new attribute got rows without that sort key. While a GSI is CREATING, Query and Scan naming it answer as the service does (no such index) rather than reading a partial table, and a second GSI Create, or a Delete of another GSI, on the table is refused with LimitExceededException; on SQLite that check also runs before the engine write lock, which a build holds until it publishes. DeleteTable stops the propagation queue for the table before dropping anything, which removes a deadlock with the queue worker that left orphan data tables, and a request (item call, transaction, Query or Scan) that meets a data table dropped under it gets ResourceNotFoundException instead of a 500. The worker bounds its wait on a held base row to 2 s and passes the item over for 2 s, every queued row for it included, so the rest of its partition keeps moving. UpdateItem on PostgreSQL validates GSI key types on the updated item as PutItem does. The new pytest is xfail on MongoDB (concurrent creates of one item return HTTP 500 there regardless of indexes) and skipped on Cassandra (contended deletes return HTTP 500 and the build did not finish in CI). Signed-off-by: Scott Robinson --- .github/workflows/integration.yml | 6 +- Cargo.lock | 1 + crates/core/src/types/mod.rs | 2 +- crates/core/src/types/table.rs | 9 + crates/core/src/validation/mod.rs | 4 +- crates/storage-postgres/Cargo.toml | 1 + .../storage-postgres/src/data/data_engine.rs | 144 +- crates/storage-postgres/src/data/ddl.rs | 105 +- .../storage-postgres/src/data/delete_item.rs | 31 +- crates/storage-postgres/src/data/gsi_build.rs | 1476 +++++++++++++++++ crates/storage-postgres/src/data/index.rs | 191 ++- crates/storage-postgres/src/data/mod.rs | 6 +- crates/storage-postgres/src/data/put_item.rs | 35 +- .../storage-postgres/src/data/transactions.rs | 58 +- .../storage-postgres/src/data/update_item.rs | 37 +- .../storage-postgres/src/data/vector_index.rs | 13 +- crates/storage-postgres/src/delete_table.rs | 46 +- crates/storage-postgres/src/gsi_queue.rs | 313 +++- crates/storage-postgres/src/lib.rs | 16 + crates/storage-postgres/src/table_helpers.rs | 103 +- crates/storage-postgres/src/update_table.rs | 292 +++- crates/storage-postgres/src/worker_store.rs | 57 +- .../tests/gsi_online_build.rs | 1170 +++++++++++++ .../tests/gsi_stale_writer.rs | 353 ++++ crates/storage-sqlite/src/data/ddl.rs | 7 +- crates/storage-sqlite/src/data/index.rs | 34 +- crates/storage-sqlite/src/data/put_item.rs | 2 +- .../storage-sqlite/src/data/transactions.rs | 20 +- crates/storage-sqlite/src/data/update_item.rs | 8 +- crates/storage-sqlite/src/update_table.rs | 282 +++- crates/storage/src/util/mod.rs | 56 + docs/design/04-component-storage.md | 42 +- tests/test_gsi_online_backfill.py | 262 +++ 33 files changed, 4732 insertions(+), 450 deletions(-) create mode 100644 crates/storage-postgres/src/data/gsi_build.rs create mode 100644 crates/storage-postgres/tests/gsi_online_build.rs create mode 100644 crates/storage-postgres/tests/gsi_stale_writer.rs create mode 100644 tests/test_gsi_online_backfill.py diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 6c5e5cf09..8a4818f75 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -301,8 +301,9 @@ jobs: # race tests hold a create open in an outside transaction, the # write-transaction conflict tests force a deadlock, and the backup restore # cases need catalog rows the current binary would not write; all need - # direct access. They build their own throwaway databases; the connection - # string is the server, not a database. + # direct access. The online GSI build tests reorder queue rows by hand and + # 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 env: EXTENDDB_TEST_PG_CONNECTION_STRING: postgresql://postgres:devpass@127.0.0.1:5432 @@ -310,6 +311,7 @@ jobs: 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 # 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 19bb60c51..003569679 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1478,6 +1478,7 @@ dependencies = [ "extenddb-core", "extenddb-storage", "futures", + "libc", "pgvector", "rand 0.9.4", "serde", diff --git a/crates/core/src/types/mod.rs b/crates/core/src/types/mod.rs index caf90fd19..e21f75110 100755 --- a/crates/core/src/types/mod.rs +++ b/crates/core/src/types/mod.rs @@ -56,7 +56,7 @@ pub use table::{ DescribeTimeToLiveOutput, DistanceFunction, GlobalSecondaryIndexUpdate, GsiDescription, GsiInput, IndexStatus, ListTablesInput, ListTablesOutput, ListTagsOfResourceInput, ListTagsOfResourceOutput, LsiDescription, LsiInput, MAX_VECTOR_INDEXES_PER_TABLE, - OnDemandThroughput, Projection, ProjectionType, ProvisionedThroughput, + ONLINE_INDEX_LIMIT, OnDemandThroughput, Projection, ProjectionType, ProvisionedThroughput, ProvisionedThroughputDescription, RestoreSummary, SearchSchemaElement, SearchSchemaElementType, SseDescription, SseType, StreamSpecification, StreamViewType, TableDescription, TableStatus, Tag, TagResourceInput, TimeToLiveDescription, TimeToLiveSpecification, diff --git a/crates/core/src/types/table.rs b/crates/core/src/types/table.rs index a81b37173..fdebcab30 100755 --- a/crates/core/src/types/table.rs +++ b/crates/core/src/types/table.rs @@ -615,6 +615,15 @@ pub const VECTOR_TABLE_REQUIRES_PAY_PER_REQUEST_MODE: &str = "One or more parame pub const VECTOR_INDEX_COUNT_LIMIT_CREATE: &str = "One or more parameter values were invalid: VectorIndex count exceeds the per-table limit of 5"; +/// A second online GSI change while one is still being built. +/// +/// The service answers an UpdateTable that creates a GSI, or deletes a GSI other +/// than the one building, while a GSI on the table is `CREATING` with this +/// LimitExceededException; the same sentence it uses for two index actions in +/// one request. +pub const ONLINE_INDEX_LIMIT: &str = "Subscriber limit exceeded: Only 1 online index can be \ + created or deleted simultaneously per table"; + /// Per-table vector index limit exceeded on `UpdateTable`. /// /// Measured 2026-08-13: reported as a `LimitExceededException`, not a diff --git a/crates/core/src/validation/mod.rs b/crates/core/src/validation/mod.rs index c144b46b6..3decb4666 100755 --- a/crates/core/src/validation/mod.rs +++ b/crates/core/src/validation/mod.rs @@ -579,9 +579,7 @@ pub fn validate_vector_index_updates( .sum::(); if actions > 1 { return Err(DynamoDbError::LimitExceededException( - "Subscriber limit exceeded: Only 1 online index can be created or deleted \ - simultaneously per table" - .to_owned(), + crate::types::ONLINE_INDEX_LIMIT.to_owned(), )); } Ok(()) diff --git a/crates/storage-postgres/Cargo.toml b/crates/storage-postgres/Cargo.toml index 316f9333e..604a31176 100755 --- a/crates/storage-postgres/Cargo.toml +++ b/crates/storage-postgres/Cargo.toml @@ -35,4 +35,5 @@ aes-gcm = { workspace = true } async-trait = { workspace = true } zeroize = { workspace = true } urlencoding = { workspace = true } +libc = { workspace = true } diff --git a/crates/storage-postgres/src/data/data_engine.rs b/crates/storage-postgres/src/data/data_engine.rs index 65144a6ae..88dd03794 100755 --- a/crates/storage-postgres/src/data/data_engine.rs +++ b/crates/storage-postgres/src/data/data_engine.rs @@ -54,6 +54,7 @@ impl DataEngine for PostgresEngine { stream.as_ref(), ) .await + .map_err(|e| table_gone(e, &key_info.table_name)) }) } @@ -64,7 +65,11 @@ impl DataEngine for PostgresEngine { ) -> BoxFuture<'_, Result, StorageError>> { let key_info = key_info.clone(); let key = key.clone(); - Box::pin(async move { self.get_item_impl(&key_info, &key).await }) + Box::pin(async move { + self.get_item_impl(&key_info, &key) + .await + .map_err(|e| table_gone(e, &key_info.table_name)) + }) } fn delete_item( @@ -91,6 +96,7 @@ impl DataEngine for PostgresEngine { stream.as_ref(), ) .await + .map_err(|e| table_gone(e, &key_info.table_name)) }) } @@ -123,6 +129,7 @@ impl DataEngine for PostgresEngine { stream.as_ref(), ) .await + .map_err(|e| table_gone(e, &key_info.table_name)) }) } @@ -152,6 +159,7 @@ impl DataEngine for PostgresEngine { index_name.as_deref(), ) .await + .map_err(|e| read_target_gone(e, &key_info.table_name, index_name.as_deref())) }) } @@ -177,6 +185,7 @@ impl DataEngine for PostgresEngine { index_name.as_deref(), ) .await + .map_err(|e| read_target_gone(e, &key_info.table_name, index_name.as_deref())) }) } @@ -227,7 +236,13 @@ impl DataEngine for PostgresEngine { .iter() .map(|(key_info, key)| TransactGetOp { key_info, key }) .collect(); - self.transact_get_items_impl(&borrowed_ops).await + let tables: Vec<(&str, &str)> = owned_ops + .iter() + .map(|(k, _)| (k.table_id.as_str(), k.table_name.as_str())) + .collect(); + self.transact_get_items_impl(&borrowed_ops) + .await + .map_err(|e| table_gone_among(e, &tables)) }) } @@ -376,6 +391,17 @@ impl DataEngine for PostgresEngine { }, ) .collect(); + let tables: Vec<(&str, &str)> = borrowed_ops + .iter() + .map(|op| match op { + TransactWriteOp::Put { key_info, .. } + | TransactWriteOp::Delete { key_info, .. } + | TransactWriteOp::Update { key_info, .. } + | TransactWriteOp::ConditionCheck { key_info, .. } => { + (key_info.table_id.as_str(), key_info.table_name.as_str()) + } + }) + .collect(); self.transact_write_items_impl( &borrowed_ops, idempotency @@ -387,6 +413,7 @@ impl DataEngine for PostgresEngine { }), ) .await + .map_err(|e| table_gone_among(e, &tables)) }) } @@ -400,3 +427,116 @@ impl DataEngine for PostgresEngine { }) } } + +/// A data table that is gone mid-request is a table that is gone. +/// +/// A request that resolved the table just before a DeleteTable committed reaches +/// the data database after the drop, and PostgreSQL answers 42P01. The service +/// answers ResourceNotFoundException for a table being deleted, which is what +/// `TableNotFound` becomes on the wire; passing the 42P01 through was a 500. +fn table_gone(e: StorageError, table_name: &str) -> StorageError { + if data_table_missing(&e) { + StorageError::TableNotFound(table_name.to_owned()) + } else { + e + } +} + +/// [`table_gone`] for a transaction over several tables: the table named is the +/// one whose data table the error names, or the first one if the error names an +/// index table (dropped with its table) or is not specific. +fn table_gone_among(e: StorageError, tables: &[(&str, &str)]) -> StorageError { + if !data_table_missing(&e) { + return e; + } + let StorageError::Internal(msg) = &e else { + return e; + }; + let name = tables + .iter() + .find(|(id, _)| msg.contains(&format!("_ddb_{id}\""))) + .or(tables.first()) + .map(|(_, name)| (*name).to_owned()); + match name { + Some(name) => StorageError::TableNotFound(name), + None => e, + } +} + +/// [`table_gone`] for a Query or Scan, which may name an index: a 42P01 there can +/// be the index's table, dropped by a GSI delete that committed after the index +/// was resolved, and the service answers that as an index it does not have. +fn read_target_gone(e: StorageError, table_name: &str, index_name: Option<&str>) -> StorageError { + match index_name { + Some(index) if data_table_missing(&e) => StorageError::IndexNotFound(index.to_owned()), + _ => table_gone(e, table_name), + } +} + +/// Whether an error is PostgreSQL's undefined_table for one of this backend's +/// data tables. Many write paths render the error with `to_string()`, which keeps +/// the message and drops the SQLSTATE, so the message is matched as well: it +/// names the relation, and every data table's name starts `_ddb_`. +fn data_table_missing(e: &StorageError) -> bool { + match e { + StorageError::Internal(msg) => { + crate::gsi_queue::is_undefined_table(e) + || (msg.contains("relation \"_ddb_") && msg.contains("does not exist")) + } + _ => false, + } +} + +#[cfg(test)] +mod table_gone_tests { + use super::*; + + #[test] + fn a_missing_data_table_is_a_missing_table_with_or_without_the_sqlstate() { + for msg in [ + "SQLSTATE 42P01: error returned from database: relation \"_ddb_abc\" does not exist", + "error returned from database: relation \"_ddb_abc\" does not exist", + ] { + let e = table_gone(StorageError::Internal(msg.to_owned()), "t"); + assert!( + matches!(e, StorageError::TableNotFound(ref n) if n == "t"), + "{msg}" + ); + } + } + + #[test] + fn other_errors_pass_through() { + let e = table_gone(StorageError::Internal("deadlock detected".to_owned()), "t"); + assert!(matches!(e, StorageError::Internal(_))); + let e = table_gone( + StorageError::Internal("relation \"settings\" does not exist".to_owned()), + "t", + ); + assert!(matches!(e, StorageError::Internal(_)), "not a data table"); + } + + #[test] + fn a_transaction_names_the_table_whose_data_table_is_gone() { + let tables = [("id-a", "a"), ("id-b", "b")]; + let e = table_gone_among( + StorageError::Internal("relation \"_ddb_id-b\" does not exist".to_owned()), + &tables, + ); + assert!(matches!(e, StorageError::TableNotFound(ref n) if n == "b")); + let e = table_gone_among( + StorageError::Internal("relation \"_ddb_some-index\" does not exist".to_owned()), + &tables, + ); + assert!(matches!(e, StorageError::TableNotFound(ref n) if n == "a")); + let e = table_gone_among(StorageError::Internal("deadlock".to_owned()), &tables); + assert!(matches!(e, StorageError::Internal(_))); + } + + #[test] + fn a_missing_index_table_on_a_read_of_an_index_is_a_missing_index() { + let msg = "relation \"_ddb_idx\" does not exist".to_owned(); + let e = read_target_gone(StorageError::Internal(msg), "t", Some("gsi1")); + assert!(matches!(e, StorageError::IndexNotFound(ref n) if n == "gsi1")); + } +} diff --git a/crates/storage-postgres/src/data/ddl.rs b/crates/storage-postgres/src/data/ddl.rs index 6edb222f3..f608c89e8 100755 --- a/crates/storage-postgres/src/data/ddl.rs +++ b/crates/storage-postgres/src/data/ddl.rs @@ -599,14 +599,20 @@ impl PostgresEngine { table_id: &str, index_name: &str, ) -> Result { - let idx_row: Option<(String, String, serde_json::Value, serde_json::Value)> = sqlx::query_as( - "SELECT index_type, index_id, key_schema, projection FROM indexes WHERE table_id = $1 AND index_name = $2", - ) - .bind(table_id) - .bind(index_name) - .fetch_optional(&self.pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + // A GSI still being built is not readable: its table is partial, or not + // there yet. The service answers a Query or Scan naming one with "The + // table does not have the specified index", the same as for no index at + // all, which is what `IndexNotFound` becomes on the wire. + let idx_row: Option<(String, String, serde_json::Value, serde_json::Value)> = + sqlx::query_as( + "SELECT index_type, index_id, key_schema, projection FROM indexes \ + WHERE table_id = $1 AND index_name = $2 AND index_status <> 'CREATING'", + ) + .bind(table_id) + .bind(index_name) + .fetch_optional(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; let (idx_type_str, idx_id, ks_json, proj_json) = idx_row.ok_or_else(|| StorageError::IndexNotFound(index_name.to_owned()))?; @@ -635,3 +641,86 @@ 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 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 + /// then deletes the pending rows waits on the worker's claimed row while the + /// worker waits on the table: a cycle, and when this side is the victim the + /// catalog rows are already gone and the data tables stay behind. + /// + /// So the queue is stopped first. A `vector_index_holds` row for the table, + /// committed on its own so the claim query can see it, keeps the worker from + /// taking any more of the table's rows (the claim already honours holds, for + /// vector builds). Then the pending rows are deleted, which waits for the + /// claims already in flight to commit, each of which holds the base row only + /// until then. Only after that are the tables dropped, with nothing left that + /// 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). + /// + /// 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. + /// + /// 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. + 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> { + 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(()) + } +} diff --git a/crates/storage-postgres/src/data/delete_item.rs b/crates/storage-postgres/src/data/delete_item.rs index 1ea9dd1e8..9a6dfda93 100755 --- a/crates/storage-postgres/src/data/delete_item.rs +++ b/crates/storage-postgres/src/data/delete_item.rs @@ -9,6 +9,7 @@ use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; use extenddb_storage::util::{SortKeyValue, parse_sk, pk_to_text, sk_column, sk_info}; +use super::gsi_build::fix_up_stale_write; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; use super::query::check_condition; use super::tx_helpers::write_stream_record_in_tx; @@ -47,12 +48,12 @@ impl PostgresEngine { // cached membership gate, and that is gone: an index takes effect the moment // its catalog row commits. What remains is a window between this read and // the data transaction's commit, in which an index created concurrently is - // missed. That window cannot be closed here, because the catalog and the - // data tables are different databases and no transaction spans them. It is - // also exactly the window the secondary indexes have, for the same reason - // and with the same read: parity with a GSI is the bar, and the backfill - // that publishes a new index is what covers writes older than it. - let (indexes, vector_metas) = + // missed. No transaction can close it, because the catalog and the data + // tables are different databases. For a GSI it is closed by time instead: + // the read's stamp lets a write that was slow to commit queue a repair + // for any GSI it did not see (`gsi_build::fix_up_stale_write`), and a fast + // write is one the backfill sees. Vector indexes still have the window. + let (indexes, vector_metas, stamp) = fetch_write_path_indexes(&key_info.table_id, &self.pool).await?; // Read whenever anything can propagate, secondary or vector. Gating this on // the secondary set alone made a vector-only table ignore the configured @@ -65,11 +66,15 @@ impl PostgresEngine { self.index_propagation_delay().await }; + // A write that is already slow takes the transaction path even with no + // index to maintain, so that a repair for a GSI created meanwhile commits + // with the write (`gsi_build::fix_up_stale_write`) instead of after it. let needs_tx = condition.is_some() || return_old || !indexes.is_empty() || !vector_metas.is_empty() - || stream.is_some(); + || stream.is_some() + || stamp.elapsed() > super::gsi_build::GRACE_WRITE; if let Some((sk_name, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) @@ -155,7 +160,6 @@ impl PostgresEngine { sync_indexes( &mut tx, &key_info.key_schema, - &key_info.attribute_definitions, &indexes, oi.as_ref(), None, @@ -213,6 +217,9 @@ impl PostgresEngine { ) .await?; + async_enqueued += + fix_up_stale_write(&mut tx, &self.pool, stamp, key_info, &indexes, key).await?; + tx.commit() .await .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -230,6 +237,7 @@ impl PostgresEngine { } } else { let delete_sql = format!("DELETE FROM {ddb_table} WHERE pk = $1 AND {sk_col} = $2"); + super::gsi_build::pause_before_autocommit_statement().await; match &sk { SortKeyValue::S(s) => { sqlx::query(&delete_sql) @@ -254,6 +262,7 @@ impl PostgresEngine { } } .map_err(|e| StorageError::Internal(e.to_string()))?; + self.fix_up_autocommit(stamp, key_info, key).await; Ok(None) } } else { @@ -313,7 +322,6 @@ impl PostgresEngine { sync_indexes( &mut tx, &key_info.key_schema, - &key_info.attribute_definitions, &indexes, oi.as_ref(), None, @@ -371,6 +379,9 @@ impl PostgresEngine { ) .await?; + async_enqueued += + fix_up_stale_write(&mut tx, &self.pool, stamp, key_info, &indexes, key).await?; + tx.commit() .await .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -388,11 +399,13 @@ impl PostgresEngine { } } else { let delete_sql = format!("DELETE FROM {ddb_table} WHERE pk = $1"); + super::gsi_build::pause_before_autocommit_statement().await; sqlx::query(&delete_sql) .bind(pk_text.as_ref()) .execute(&self.data_pool) .await .map_err(|e| StorageError::Internal(e.to_string()))?; + self.fix_up_autocommit(stamp, key_info, key).await; Ok(None) } } diff --git a/crates/storage-postgres/src/data/gsi_build.rs b/crates/storage-postgres/src/data/gsi_build.rs new file mode 100644 index 000000000..22ee5bf6f --- /dev/null +++ b/crates/storage-postgres/src/data/gsi_build.rs @@ -0,0 +1,1476 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Building a global secondary index on a table that is taking writes. +//! +//! Three pieces, and the order they run in is the design: +//! +//! 1. `UpdateTable` creates the index data table, then commits the catalog row as +//! `CREATING`. From that commit on, every write that reads the catalog maintains +//! the index the way it maintains an `ACTIVE` one, inline or through the queue. +//! 2. It waits [`GRACE_PUBLISH`], takes a [`write_barrier`], and runs +//! [`backfill`], which writes the current state of every base item into the +//! index under that item's row lock. +//! 3. It flips the row to `ACTIVE`. A build that fails or is cut short leaves the +//! index `CREATING`, and [`recovery_worker`] finishes it. +//! +//! Every write to the index, from the backfill, the queue or a write's own inline +//! maintenance, leaves the rows for one base key equal to that item's state at the +//! time, and does so while holding the base row's lock. So the three can run in any +//! order against each other and the index still ends at the base table's state. +//! +//! What that does not cover is a write that read the catalog just before the +//! `CREATING` row committed. It does not know the index exists, so it maintains +//! nothing, and the catalog and the data tables are different databases, so no +//! transaction can tell it. [`WriteStamp`] and [`fix_up_stale_write`] close that +//! gap with a time bound: +//! +//! - A write whose base table statements had all returned within [`GRACE_WRITE`] +//! of its catalog read holds the table's `ROW EXCLUSIVE` lock from the first of +//! them until it commits. The barrier comes [`GRACE_PUBLISH`] (longer) after the +//! catalog commit, so it finds that lock held, waits for the commit, and the +//! backfill, which starts after it, sees the write. (A `SELECT ... FOR UPDATE` +//! alone takes only `ROW SHARE`, which the barrier does not wait for; the check +//! runs after the write's `INSERT`, `UPDATE` or `DELETE`, which is what counts.) +//! - A slower write re-reads the catalog before it commits and queues a repair for +//! any index it did not know about. If the re-read does not show the index, the +//! catalog commit came after it, so that write's lock was again taken before the +//! barrier. +//! +//! Measurement error only ever makes a write look slower or the wait longer, which +//! is the safe direction: the clock is one that keeps counting while the host is +//! suspended, the writer measures to a point after its write, and `UpdateTable` +//! measures from after its commit returned. +//! +//! One gap is left. A write that needs no transaction (a plain put or delete on +//! a table it believed had no index) commits with its one statement, so if it +//! only becomes slow after that decision, its repair goes in a separate +//! transaction. A process that dies between the two, for a write that was over +//! [`GRACE_WRITE`] and raced an index creation, leaves that item out of the +//! index (or, for a delete, leaves its old row in) until the item is next +//! written. A write that is already slow when it decides takes the transaction +//! path instead, which closes most of this window. + +use std::time::Duration; + +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, Item, KeySchemaElement, Projection, ScalarAttributeType, + TableKeyInfo, +}; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::{composite_pk_to_text, parse_sk, sk_column_n}; +use sqlx::Row as _; + +use super::index::{ + IndexMeta, db_error, delete_index_row_multi, fetch_indexes_for_table, insert_index_row_multi, + project_item_for_index, +}; +use super::{all_sort_key_info, data_table_name, index_table_name, json_to_item}; +use crate::gsi_queue::{GsiApplyContext, GsiIndexDef, PendingApplyContext, enqueue_gsi_pending}; + +/// How long a write may take, from just before its catalog read to just before +/// its commit, and still rely on the backfill to have seen it. +pub(crate) const GRACE_WRITE: Duration = Duration::from_millis(500); + +/// How long `UpdateTable` waits after committing a `CREATING` index before the +/// barrier and backfill. Must exceed [`GRACE_WRITE`]; the margin covers the two +/// clocks (a writer's and this call's) disagreeing on how long half a second is, +/// which is the one error the monotonic clock does not remove. +pub(crate) const GRACE_PUBLISH: Duration = Duration::from_millis(1_000); + +/// Base rows read and locked per backfill transaction. +/// +/// A writer to any row in the batch waits for the batch to commit, so this bounds +/// that stall. Small next to the vector backfill's batch because the lock is the +/// point here. +const BACKFILL_BATCH: i64 = 256; + +/// A point in time on a clock that keeps counting while the host is suspended. +/// +/// `std::time::Instant` is `CLOCK_MONOTONIC` on Linux, which stops during suspend, +/// so a write could measure itself as fast when it was not. +#[derive(Clone, Copy, Debug)] +pub(crate) struct WriteStamp(Option); + +impl WriteStamp { + pub(crate) fn now() -> Self { + Self(clock_now()) + } + + /// Time since [`WriteStamp::now`]. An unreadable clock reports as long ago, + /// which sends the write down the repair path rather than trusting it. + pub(crate) fn elapsed(self) -> Duration { + match (self.0, clock_now()) { + (Some(start), Some(now)) => now.saturating_sub(start), + _ => Duration::MAX, + } + } +} + +#[cfg(unix)] +fn clock_now() -> Option { + #[cfg(any(target_os = "linux", target_os = "android"))] + const CLOCK: libc::clockid_t = libc::CLOCK_BOOTTIME; + // CLOCK_MONOTONIC counts across sleep on macOS and the BSDs. + #[cfg(not(any(target_os = "linux", target_os = "android")))] + const CLOCK: libc::clockid_t = libc::CLOCK_MONOTONIC; + + let mut ts = libc::timespec { + tv_sec: 0, + tv_nsec: 0, + }; + // SAFETY: `ts` is a valid, writable timespec for the duration of the call. + let rc = unsafe { libc::clock_gettime(CLOCK, &mut ts) }; + if rc != 0 { + return None; + } + Some(Duration::new( + u64::try_from(ts.tv_sec).ok()?, + u32::try_from(ts.tv_nsec).ok()?, + )) +} + +#[cfg(not(unix))] +fn clock_now() -> Option { + static START: std::sync::OnceLock = std::sync::OnceLock::new(); + Some(START.get_or_init(std::time::Instant::now).elapsed()) +} + +/// Where a write pauses under [`pause_writes_for_tests`]. +#[doc(hidden)] +#[derive(Debug, Clone, Copy)] +pub enum PauseWrites { + /// Between the catalog read and the data work. A write paused here for longer + /// than [`GRACE_WRITE`] takes the transaction path and repairs in-transaction + /// (`fix_up_stale_write`). + AfterIndexRead, + /// Between the decision to write without a transaction and the statement. The + /// write stays on the autocommit path and repairs afterwards + /// (`fix_up_autocommit`), the one repair that is not atomic with the write. + BeforeAutocommitStatement, +} + +/// Test hooks: how long a write sleeps at each [`PauseWrites`] point. Zero +/// outside tests. +static PAUSE_AFTER_INDEX_READ_MS: std::sync::atomic::AtomicU64 = + std::sync::atomic::AtomicU64::new(0); +static PAUSE_BEFORE_AUTOCOMMIT_MS: std::sync::atomic::AtomicU64 = + std::sync::atomic::AtomicU64::new(0); + +/// Make every write on this process pause for `pause` at `at`, so a test can +/// deterministically produce a write that read the catalog before a GSI was +/// created and writes after its backfill finished: the case +/// [`fix_up_stale_write`] and `fix_up_autocommit` exist for, which no timing in a +/// test can otherwise hit reliably. Hidden because no deployment path calls it. +#[doc(hidden)] +pub fn pause_writes_for_tests(at: PauseWrites, pause: Duration) { + let ms = u64::try_from(pause.as_millis()).unwrap_or(u64::MAX); + let slot = match at { + PauseWrites::AfterIndexRead => &PAUSE_AFTER_INDEX_READ_MS, + PauseWrites::BeforeAutocommitStatement => &PAUSE_BEFORE_AUTOCOMMIT_MS, + }; + slot.store(ms, std::sync::atomic::Ordering::Relaxed); +} + +async fn pause(slot: &std::sync::atomic::AtomicU64) { + let ms = slot.load(std::sync::atomic::Ordering::Relaxed); + if ms > 0 { + tokio::time::sleep(Duration::from_millis(ms)).await; + } +} + +pub(crate) async fn pause_after_index_read() { + pause(&PAUSE_AFTER_INDEX_READ_MS).await; +} + +pub(crate) async fn pause_before_autocommit_statement() { + pause(&PAUSE_BEFORE_AUTOCOMMIT_MS).await; +} + +/// One typed key value, as stored in a base or index table's key column. +#[derive(Debug, Clone)] +enum KeyVal { + S(String), + N(sqlx::types::BigDecimal), + B(Vec), +} + +/// A base table primary key, in column order. +#[derive(Debug, Clone)] +struct BaseKey { + pk: String, + sks: Vec, +} + +/// The base table's key columns, `pk` then the typed sort key columns. +fn base_key_columns(base_sks: &[(&str, ScalarAttributeType)]) -> Vec { + let mut cols = vec!["pk".to_owned()]; + for (i, &(_, t)) in base_sks.iter().enumerate() { + cols.push(sk_column_n(i, t)); + } + cols +} + +fn read_key( + row: &sqlx::postgres::PgRow, + base_sks: &[(&str, ScalarAttributeType)], +) -> Result { + let pk: String = row.try_get("pk").map_err(db_error)?; + let mut sks = Vec::with_capacity(base_sks.len()); + for (i, &(_, t)) in base_sks.iter().enumerate() { + let col = sk_column_n(i, t); + let v = match t { + ScalarAttributeType::S => row.try_get::(col.as_str()).map(KeyVal::S), + ScalarAttributeType::N => row + .try_get::(col.as_str()) + .map(KeyVal::N), + ScalarAttributeType::B => row.try_get::, _>(col.as_str()).map(KeyVal::B), + } + .map_err(db_error)?; + sks.push(v); + } + Ok(BaseKey { pk, sks }) +} + +fn key_from_item( + item: &Item, + base_ks: &[KeySchemaElement], + base_sks: &[(&str, ScalarAttributeType)], +) -> Result { + let pk = composite_pk_to_text(item, base_ks)?; + let mut sks = Vec::with_capacity(base_sks.len()); + for &(name, t) in base_sks { + let v = item + .get(name) + .ok_or_else(|| StorageError::Internal(format!("missing sort key {name}")))?; + sks.push(match parse_sk(v, t)? { + extenddb_storage::util::SortKeyValue::S(s) => KeyVal::S(s), + extenddb_storage::util::SortKeyValue::N(n) => KeyVal::N(n), + extenddb_storage::util::SortKeyValue::B(b) => KeyVal::B(b), + }); + } + Ok(BaseKey { pk, sks }) +} + +fn bind_key<'q>( + mut q: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>, + key: &BaseKey, +) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> { + q = q.bind(key.pk.clone()); + for v in &key.sks { + q = match v { + KeyVal::S(s) => q.bind(s.clone()), + KeyVal::N(n) => q.bind(n.clone()), + KeyVal::B(b) => q.bind(b.clone()), + }; + } + q +} + +/// `(c1, c2, ...)` over the given column names, with an optional prefix. +fn tuple(cols: &[String], prefix: &str) -> String { + let parts: Vec = cols.iter().map(|c| format!("{prefix}{c}")).collect(); + format!("({})", parts.join(", ")) +} + +/// `($start, $start+1, ...)` for `n` parameters. +fn params(start: usize, n: usize) -> String { + let parts: Vec = (start..start + n).map(|i| format!("${i}")).collect(); + format!("({})", parts.join(", ")) +} + +/// Whether an item belongs in the index: every key attribute present, of the +/// declared type, and not empty. +/// +/// An item that fails is left out of the index rather than failing the build or +/// stalling a queue partition. It can only be an item written before the index +/// existed, or by a writer that had not yet seen the index's attribute +/// definitions; the service skips such items in a backfill too. +pub(crate) fn indexable( + item: &Item, + index_ks: &[KeySchemaElement], + attr_defs: &[AttributeDefinition], +) -> bool { + index_ks.iter().all(|ks| { + let Some(value) = item.get(&ks.attribute_name) else { + return false; + }; + let declared = attr_defs + .iter() + .find(|ad| ad.attribute_name == ks.attribute_name) + .map(|ad| ad.attribute_type); + match (declared, value) { + (Some(ScalarAttributeType::S), AttributeValue::S(s)) => !s.is_empty(), + (Some(ScalarAttributeType::N), AttributeValue::N(n)) => { + n.parse::().is_ok() + } + (Some(ScalarAttributeType::B), AttributeValue::B(b)) => !b.is_empty(), + _ => false, + } + }) +} + +/// Make one index's rows for one base item equal to `current`, the item's state +/// read under its row lock (or `None` if it has none). +#[allow(clippy::too_many_arguments)] +pub(crate) async fn write_state( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + idx_table: &str, + index_ks: &[KeySchemaElement], + projection: &Projection, + base_ks: &[KeySchemaElement], + attr_defs: &[AttributeDefinition], + base_sks: &[(&str, ScalarAttributeType)], + current: &Item, +) -> Result { + // By base key, whatever index key the stale row carries. + delete_index_row_multi(tx, idx_table, current, base_ks, base_sks).await?; + if !indexable(current, index_ks, attr_defs) { + return Ok(false); + } + let idx_sks = all_sort_key_info(index_ks, attr_defs); + let projected = project_item_for_index(current, index_ks, base_ks, projection); + insert_index_row_multi( + tx, idx_table, current, &projected, index_ks, base_ks, &idx_sks, base_sks, + ) + .await?; + Ok(true) +} + +/// Delete one base key's index rows, but only if the base row is absent. +/// +/// Two statements, and the split is the point. A single `DELETE ... WHERE NOT +/// EXISTS (base row)` that has to wait on an index row a writer is updating +/// re-checks only that row once the writer commits, and evaluates `NOT EXISTS` +/// against its original snapshot: a writer that re-created the item and updated +/// the stale index row in place would have its fresh row deleted. So the rows are +/// locked first, which is where any wait happens, and the delete then runs with a +/// new snapshot over exactly the rows that are locked, by `ctid`. Nobody else can +/// change a locked row, so its `ctid` is stable and the delete waits on nothing; +/// a row some writer added after the lock is not among them and is left alone. +/// +/// No writer in this crate updates an index row in place today: every GSI write +/// goes through `write_state`, which deletes by base key before it upserts, so +/// the single-statement form would not misbehave against them. The guard is for +/// the next writer that does not, which this function cannot see; it costs one +/// statement on a path taken only for an item that is gone. +async fn delete_if_absent( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + base_table: &str, + idx_table: &str, + base_cols: &[String], + key: &BaseKey, +) -> Result<(), StorageError> { + let idx_where: Vec = base_cols + .iter() + .enumerate() + .map(|(i, c)| format!("base_{c} = ${}", i + 1)) + .collect(); + let base_where: Vec = base_cols + .iter() + .enumerate() + .map(|(i, c)| format!("b.{c} = ${}", i + 1)) + .collect(); + let lock = format!( + "SELECT ctid::text FROM {idx_table} WHERE {} FOR UPDATE", + idx_where.join(" AND ") + ); + let locked: Vec = bind_key(sqlx::query(&lock), key) + .fetch_all(&mut **tx) + .await + .map_err(db_error)? + .iter() + .map(|r| r.try_get::(0)) + .collect::>() + .map_err(db_error)?; + if locked.is_empty() { + return Ok(()); + } + let sql = format!( + "DELETE FROM {idx_table} WHERE ctid = ANY(${}::text[]::tid[]) \ + AND NOT EXISTS (SELECT 1 FROM {base_table} b WHERE {})", + base_cols.len() + 1, + base_where.join(" AND ") + ); + bind_key(sqlx::query(&sql), key) + .bind(locked) + .execute(&mut **tx) + .await + .map_err(db_error)?; + Ok(()) +} + +/// Bring one index's rows for one base key in line with the base table. +/// +/// The queue's apply step: a queued row says which item changed, and this writes +/// the item's state now rather than the image the row carried. So a row applied +/// late, out of order, or twice cannot leave the index behind the base table, and +/// a repair row for a write that never told the index about itself needs no image +/// at all. +/// +/// `key_item` only has to carry the base key attributes. +pub(crate) async fn resync_base_key( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + table_id: &str, + idx: &GsiIndexDef, + base_ks: &[KeySchemaElement], + attr_defs: &[AttributeDefinition], + key_item: &Item, +) -> Result<(), StorageError> { + let base_sks = all_sort_key_info(base_ks, attr_defs); + let base_table = data_table_name(table_id); + let idx_table = index_table_name(&idx.index_id); + let key = key_from_item(key_item, base_ks, &base_sks)?; + resync_key( + tx, + &base_table, + &idx_table, + idx, + base_ks, + attr_defs, + &base_sks, + &key, + ) + .await +} + +/// Wait out every write that already holds a lock on the base table. +/// +/// `SHARE` conflicts with the `ROW EXCLUSIVE` lock every insert, update and +/// delete takes, so this returns once each of them has committed or rolled back. +/// New writes queue behind it while it waits, each holding a data pool +/// connection, so the wait is bounded: something that holds the table for long (a +/// long transaction, or a VACUUM, which takes a conflicting lock and need not +/// yield to this one) would otherwise stop every write to the table, and then +/// starve other tables of connections, for as long. On a timeout the lock request +/// is withdrawn, the queued writes proceed, and the barrier is tried again after +/// a growing pause. With the default `deadlock_timeout` the attempts add up to +/// about 12 s of stalls spread over 20 s before the build gives up. +pub(crate) async fn write_barrier(data: &sqlx::PgPool, table_id: &str) -> Result<(), StorageError> { + // The timeout is a second longer than the server's `deadlock_timeout`, on + // purpose: a waiter cancels a conflicting autovacuum only when it runs its + // deadlock check, which happens after that timeout, so a shorter wait could + // never get past one. An anti-wraparound or manual vacuum does not yield at + // all; if one outlasts every attempt the build fails as transient, and the + // index stays CREATING for the recovery worker. + const ATTEMPTS: u32 = 6; + let base_table = data_table_name(table_id); + let mut attempt = 0u32; + loop { + attempt += 1; + let result = async { + let mut tx = data.begin().await.map_err(db_error)?; + sqlx::query( + "SELECT set_config('lock_timeout', \ + ((extract(epoch FROM current_setting('deadlock_timeout')::interval) * 1000)::int \ + + 1000)::text || 'ms', true)", + ) + .execute(&mut *tx) + .await + .map_err(db_error)?; + sqlx::query(&format!("LOCK TABLE {base_table} IN SHARE MODE")) + .execute(&mut *tx) + .await + .map_err(db_error)?; + tx.commit().await.map_err(db_error) + } + .await; + match result { + Ok(()) => return Ok(()), + Err(e) if attempt < ATTEMPTS && e.to_string().contains("SQLSTATE 55P03") => { + // Writers queued behind the attempt run in the gap, and the gap + // grows so a long holder is not met with a steady stream of stalls. + tokio::time::sleep(Duration::from_millis(500 * u64::from(attempt))).await; + } + Err(e) => return Err(e), + } + } +} + +/// Whether a build error is worth leaving to the recovery worker: a conflict, +/// a timeout, a lost or refused connection, a shutdown, a server-side I/O or +/// internal error, or a connection pool that had nothing to give. Anything else +/// (bad data, a row that does not decode, a key too large for the index) would +/// fail the same way every time, so the default is "not transient": a wrong +/// answer here is an index that stays CREATING for good, rebuilt every hour, +/// with nothing to say why. +pub(crate) fn is_transient(e: &StorageError) -> bool { + let msg = e.to_string(); + if let Some(at) = msg.find("SQLSTATE ") { + let code = &msg[at + "SQLSTATE ".len()..]; + return ["40", "08", "53", "55P03", "57", "58", "XX"] + .iter() + .any(|prefix| code.starts_with(prefix)); + } + // sqlx's own wording for errors that never reached the server. + [ + "pool timed out", + "closed pool", + "error communicating with database", + "connection reset", + "broken pipe", + ] + .iter() + .any(|needle| msg.contains(needle)) +} + +/// What a backfill did, for the log line. +#[derive(Debug, Default, Clone, Copy)] +pub(crate) struct BackfillCounts { + pub(crate) indexed: usize, + pub(crate) skipped: usize, +} + +/// Write the current state of every base item into the index. +/// +/// Keyset pages over the full base key, up to the last key present when the +/// backfill starts. A key added later was added by a write that saw the index (it +/// came after the barrier), and that write maintains the index itself; stopping +/// at the starting maximum is what lets the backfill end on a table whose keys +/// keep growing. Each page: +/// +/// 1. Reads the page's keys without locking, which fixes the page's range. +/// 2. In one transaction, share-locks the rows in that range with `SKIP LOCKED` +/// and writes each one's state into the index. Skipping rather than waiting +/// keeps this out of deadlocks with a writer that holds one row of the page and +/// waits on another, because a page that waited would hold its earlier rows +/// while it did. A page that is still chosen as a deadlock victim (or hits a +/// serialization failure under a stricter isolation level) is retried. +/// 3. After that commits, takes each key from step 1 that step 2 skipped, one at +/// a time with a waiting lock. A transaction holding one row lock cannot be +/// part of a deadlock cycle. +/// +/// Then one pass deletes index rows whose base item no longer exists. Only a +/// write that read the catalog before the index existed can leave one, by +/// deleting an item without touching its index row. +pub(crate) async fn backfill( + data: &sqlx::PgPool, + table_id: &str, + idx: &GsiIndexDef, + base_ks: &[KeySchemaElement], + attr_defs: &[AttributeDefinition], +) -> Result { + let base_sks = all_sort_key_info(base_ks, attr_defs); + let base_table = data_table_name(table_id); + let idx_table = index_table_name(&idx.index_id); + let cols = base_key_columns(&base_sks); + let n = cols.len(); + let key_tuple = tuple(&cols, ""); + let selected = cols.join(", "); + let descending: Vec = cols.iter().map(|c| format!("{c} DESC")).collect(); + + let mut counts = BackfillCounts::default(); + + let last_sql = format!( + "SELECT {selected} FROM {base_table} ORDER BY {} LIMIT 1", + descending.join(", ") + ); + let last = sqlx::query(&last_sql) + .fetch_optional(data) + .await + .map_err(db_error)? + .map(|r| read_key(&r, &base_sks)) + .transpose()?; + + let mut cursor: Option = None; + while let Some(end) = &last { + // 1. The page's keys. + let mut conds = Vec::new(); + let mut next = 1; + if cursor.is_some() { + conds.push(format!("{key_tuple} > {}", params(next, n))); + next += n; + } + conds.push(format!("{key_tuple} <= {}", params(next, n))); + let keys_sql = format!( + "SELECT {selected} FROM {base_table} WHERE {} ORDER BY {selected} \ + LIMIT {BACKFILL_BATCH}", + conds.join(" AND ") + ); + let mut q = sqlx::query(&keys_sql); + if let Some(c) = &cursor { + q = bind_key(q, c); + } + q = bind_key(q, end); + let key_rows = q.fetch_all(data).await.map_err(db_error)?; + let page: Vec = key_rows + .iter() + .map(|r| read_key(r, &base_sks)) + .collect::>()?; + let Some(upper) = page.last().cloned() else { + break; + }; + + // 2. Lock what is free in the range and write its state. + let (done, page_counts) = with_retries(|| { + write_page( + data, + &base_table, + &idx_table, + idx, + base_ks, + attr_defs, + &base_sks, + cursor.as_ref(), + &upper, + ) + }) + .await?; + counts.indexed += page_counts.indexed; + counts.skipped += page_counts.skipped; + + // 3. The keys someone else held. + for key in &page { + if done.iter().any(|d| same_key(d, key)) { + continue; + } + with_retries(|| async { + let mut tx = data.begin().await.map_err(db_error)?; + resync_key( + &mut tx, + &base_table, + &idx_table, + idx, + base_ks, + attr_defs, + &base_sks, + key, + ) + .await?; + tx.commit().await.map_err(db_error) + }) + .await?; + } + + if same_key(&upper, end) { + break; + } + cursor = Some(upper); + } + + sweep_orphans(data, &base_table, &idx_table, &base_sks).await?; + Ok(counts) +} + +/// Attempts at one backfill transaction before its error is the build's error. +const PAGE_ATTEMPTS: u32 = 5; + +/// Run one backfill transaction, again if it was a deadlock victim or hit a +/// serialization failure. Anything else is returned as it is. +async fn with_retries(mut run: F) -> Result +where + F: FnMut() -> Fut, + Fut: std::future::Future>, +{ + let mut attempt = 0u32; + loop { + attempt += 1; + match run().await { + Ok(v) => return Ok(v), + Err(e) if attempt < PAGE_ATTEMPTS && is_retryable(&e) => { + tracing::debug!(attempt, "GSI backfill transaction retried: {e}"); + tokio::time::sleep(Duration::from_millis(10 * u64::from(attempt))).await; + } + Err(e) => return Err(e), + } + } +} + +/// A deadlock victim or a serialization failure: the page can simply run again. +fn is_retryable(e: &StorageError) -> bool { + let msg = e.to_string(); + msg.contains("SQLSTATE 40P01") || msg.contains("SQLSTATE 40001") +} + +/// Step 2 of [`backfill`]: one page's transaction. Returns the keys it wrote. +#[allow(clippy::too_many_arguments)] +async fn write_page( + data: &sqlx::PgPool, + base_table: &str, + idx_table: &str, + idx: &GsiIndexDef, + base_ks: &[KeySchemaElement], + attr_defs: &[AttributeDefinition], + base_sks: &[(&str, ScalarAttributeType)], + lower: Option<&BaseKey>, + upper: &BaseKey, +) -> Result<(Vec, BackfillCounts), StorageError> { + let cols = base_key_columns(base_sks); + let n = cols.len(); + let key_tuple = tuple(&cols, ""); + let selected = cols.join(", "); + let (lower_clause, upper_first) = match lower { + Some(_) => (format!("{key_tuple} > {} AND ", params(1, n)), n + 1), + None => (String::new(), 1), + }; + let lock_sql = format!( + "SELECT {selected}, item_data FROM {base_table} \ + WHERE {lower_clause}{key_tuple} <= {} \ + ORDER BY {selected} FOR SHARE SKIP LOCKED", + params(upper_first, n) + ); + + let mut tx = data.begin().await.map_err(db_error)?; + let mut q = sqlx::query(&lock_sql); + if let Some(c) = lower { + q = bind_key(q, c); + } + q = bind_key(q, upper); + let locked = q.fetch_all(&mut *tx).await.map_err(db_error)?; + + let mut counts = BackfillCounts::default(); + let mut done: Vec = Vec::with_capacity(locked.len()); + for row in &locked { + let key = read_key(row, base_sks)?; + let json: serde_json::Value = row.try_get("item_data").map_err(db_error)?; + let item = json_to_item(json)?; + if write_state( + &mut tx, + idx_table, + &idx.key_schema, + &idx.projection, + base_ks, + attr_defs, + base_sks, + &item, + ) + .await? + { + counts.indexed += 1; + } else { + counts.skipped += 1; + } + done.push(key); + } + tx.commit().await.map_err(db_error)?; + Ok((done, counts)) +} + +fn same_key(a: &BaseKey, b: &BaseKey) -> bool { + a.pk == b.pk + && a.sks.len() == b.sks.len() + && a.sks.iter().zip(&b.sks).all(|(x, y)| match (x, y) { + (KeyVal::S(x), KeyVal::S(y)) => x == y, + (KeyVal::N(x), KeyVal::N(y)) => x == y, + (KeyVal::B(x), KeyVal::B(y)) => x == y, + _ => false, + }) +} + +/// [`resync_base_key`] for a key already in column form. +/// +/// The row lock is taken before any index row is touched, which is the order a +/// writer takes them in, so the two cannot deadlock. +#[allow(clippy::too_many_arguments)] +async fn resync_key( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + base_table: &str, + idx_table: &str, + idx: &GsiIndexDef, + base_ks: &[KeySchemaElement], + attr_defs: &[AttributeDefinition], + base_sks: &[(&str, ScalarAttributeType)], + key: &BaseKey, +) -> Result<(), StorageError> { + let cols = base_key_columns(base_sks); + let where_parts: Vec = cols + .iter() + .enumerate() + .map(|(i, c)| format!("{c} = ${}", i + 1)) + .collect(); + let sql = format!( + "SELECT item_data FROM {base_table} WHERE {} FOR SHARE", + where_parts.join(" AND ") + ); + let current = bind_key(sqlx::query(&sql), key) + .fetch_optional(&mut **tx) + .await + .map_err(db_error)?; + match current { + Some(row) => { + let json: serde_json::Value = row.try_get("item_data").map_err(db_error)?; + let item = json_to_item(json)?; + write_state( + tx, + idx_table, + &idx.key_schema, + &idx.projection, + base_ks, + attr_defs, + base_sks, + &item, + ) + .await?; + } + None => delete_if_absent(tx, base_table, idx_table, &cols, key).await?, + } + Ok(()) +} + +/// Delete index rows whose base item no longer exists. +/// +/// The candidates are found without locking, and each is then deleted through +/// [`delete_if_absent`] in its own transaction, which holds no base row lock: a +/// writer re-creating the item at the same moment holds that item's base row and +/// may be waiting on another, and an index row lock taken while holding base +/// locks could close a cycle with it. +async fn sweep_orphans( + data: &sqlx::PgPool, + base_table: &str, + idx_table: &str, + base_sks: &[(&str, ScalarAttributeType)], +) -> Result<(), StorageError> { + let cols = base_key_columns(base_sks); + let idx_cols: Vec = cols.iter().map(|c| format!("i.base_{c} AS {c}")).collect(); + let join: Vec = cols.iter().map(|c| format!("b.{c} = i.base_{c}")).collect(); + let sql = format!( + "SELECT DISTINCT {} FROM {idx_table} i \ + WHERE NOT EXISTS (SELECT 1 FROM {base_table} b WHERE {})", + idx_cols.join(", "), + join.join(" AND ") + ); + let rows = sqlx::query(&sql).fetch_all(data).await.map_err(db_error)?; + for row in &rows { + let key = read_key(row, base_sks)?; + with_retries(|| async { + let mut tx = data.begin().await.map_err(db_error)?; + delete_if_absent(&mut tx, base_table, idx_table, &cols, &key).await?; + tx.commit().await.map_err(db_error) + }) + .await?; + } + Ok(()) +} + +/// The base key attributes of `item`, and nothing else. +fn key_only(item: &Item, base_ks: &[KeySchemaElement]) -> Item { + let mut out = Item::new(); + for ks in base_ks { + if let Some(v) = item.get(&ks.attribute_name) { + out.insert(ks.attribute_name.clone(), v.clone()); + } + } + out +} + +/// Queue a repair for any GSI this write did not know about, if the write was +/// slow enough that the backfill may have missed it. +/// +/// Call inside the write's transaction, after its base table write and before its +/// commit, so the repair commits or rolls back with the write. For a write with no +/// transaction, call it in a new one right after the write. +/// +/// `known` is the index set the write read at `stamp`. Returns the number of rows +/// queued, so the caller can wake the workers. +pub(crate) async fn fix_up_stale_write( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + catalog: &sqlx::PgPool, + stamp: WriteStamp, + key_info: &TableKeyInfo, + known: &[IndexMeta], + item: &Item, +) -> Result { + if stamp.elapsed() <= GRACE_WRITE { + return Ok(0); + } + let missed = missed_gsis(catalog, &key_info.table_id, known).await?; + queue_repairs(tx, key_info, &missed, item).await +} + +/// The GSIs on a table now that are not in `known`. +pub(crate) async fn missed_gsis( + catalog: &sqlx::PgPool, + table_id: &str, + known: &[IndexMeta], +) -> Result, StorageError> { + let fresh = fetch_indexes_for_table(table_id, catalog).await?; + Ok(fresh + .into_iter() + .filter(|idx| idx.index_type == "GSI" && !known.iter().any(|k| k.index_id == idx.index_id)) + .collect()) +} + +/// Queue one repair row per index in `missed` for the item's base key. +pub(crate) async fn queue_repairs( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + key_info: &TableKeyInfo, + missed: &[IndexMeta], + item: &Item, +) -> Result { + let base_ks = &key_info.key_schema; + let key = key_only(item, base_ks); + for idx in missed { + let context = PendingApplyContext::Gsi(GsiApplyContext { + base_key_schema: base_ks.clone(), + attribute_definitions: idx.attribute_definitions.clone(), + index: GsiIndexDef { + index_id: idx.index_id.clone(), + key_schema: idx.key_schema.clone(), + projection: idx.projection.clone(), + }, + }); + enqueue_gsi_pending(tx, &key_info.table_id, Some(&key), None, 0, &context).await?; + } + if !missed.is_empty() { + tracing::debug!( + table_id = %key_info.table_id, + queued = missed.len(), + "queued GSI repair for a write that predates the index" + ); + } + Ok(missed.len()) +} + +/// One GSI an `UpdateTable` (or a recovery sweep) is building. +pub(crate) struct GsiBuild { + pub(crate) index_name: String, + pub(crate) def: GsiIndexDef, + /// Held for the whole build so a peer's sweep leaves it alone; released on drop. + pub(crate) _owner: Option, +} + +/// Drop the data tables of builds that are being abandoned. Logged, not +/// propagated: the caller is already returning an error. +pub(crate) async fn drop_gsi_tables(data: &sqlx::PgPool, builds: &[GsiBuild]) { + for b in builds { + let t = index_table_name(&b.def.index_id); + if let Err(e) = sqlx::query(&format!("DROP TABLE IF EXISTS {t}")) + .execute(data) + .await + { + tracing::warn!( + index_id = %b.def.index_id, + "could not drop the data table of an abandoned GSI build: {e}" + ); + } + } +} + +/// Remove a GSI whose build fails deterministically: the catalog row, then its +/// data table, but only if this call removed the row. The delete is conditional +/// on the row still being `CREATING`: two builds of one index can run at once (a +/// recovery pass next to an UpdateTable whose ownership session could not be +/// opened), and if the other one has published the index since, a failure here +/// is this build's problem and must not take an ACTIVE index away. With the row +/// gone or published, a table writes are maintaining must not vanish under them. +/// +/// The other way round, a second build failing for good while the first is still +/// scanning, removes the row and the table under the first; the first then finds +/// no row to flip and reports the index as deleted while it was being built. The +/// failure would have met the first build too, so the outcome is the same, minus +/// the error to the caller. +/// +/// The table is dropped [`GRACE_PUBLISH`] after the row, so a write that read the +/// catalog just before the delete has usually finished with the table. The +/// synchronous write path also treats a table that is gone as an index that is +/// gone, so the delay is for the common case, not the guarantee. +pub(crate) async fn remove_failed_build( + catalog: &sqlx::PgPool, + data: &sqlx::PgPool, + table_id: &str, + build: &GsiBuild, +) { + // By id, not name: the name may already belong to a newer index. + let deleted = sqlx::query( + "DELETE FROM indexes WHERE table_id = $1 AND index_id = $2 \ + AND index_status = 'CREATING'", + ) + .bind(table_id) + .bind(&build.def.index_id) + .execute(catalog) + .await; + match deleted { + Ok(r) if r.rows_affected() == 1 => { + tokio::time::sleep(GRACE_PUBLISH).await; + drop_gsi_tables(data, std::slice::from_ref(build)).await; + } + Ok(_) => {} + Err(e) => tracing::error!( + index_id = %build.def.index_id, + "could not remove the catalog row of a failed GSI build: {e}" + ), + } +} + +/// Whether the catalog still has this index. A build that fails because the index +/// or its table was deleted under it has nothing left to clean up or report. +async fn index_row_exists( + catalog: &sqlx::PgPool, + table_id: &str, + index_id: &str, +) -> Result { + let found: Option<(i32,)> = + sqlx::query_as("SELECT 1 FROM indexes WHERE table_id = $1 AND index_id = $2") + .bind(table_id) + .bind(index_id) + .fetch_optional(catalog) + .await + .map_err(db_error)?; + Ok(found.is_some()) +} + +/// Wait, barrier, backfill and publish each build, in order. +/// +/// Returns how many this call published, or the position of the first build +/// that failed with its error; builds before it are ACTIVE, and the caller +/// abandons it and the ones after it. A build whose index was deleted while it +/// ran, or that someone else published first, is not a failure. +pub(crate) async fn build_gsis( + catalog: &sqlx::PgPool, + data: &sqlx::PgPool, + table_id: &str, + base_ks: &[KeySchemaElement], + attr_defs: &[AttributeDefinition], + builds: &[GsiBuild], +) -> Result { + if builds.is_empty() { + return Ok(0); + } + let mut published_here = 0usize; + tokio::time::sleep(GRACE_PUBLISH).await; + if let Err(e) = write_barrier(data, table_id).await { + // A table deleted during the wait is not a failed build. + let mut gone = true; + for b in builds { + match index_row_exists(catalog, table_id, &b.def.index_id).await { + Ok(false) => {} + _ => gone = false, + } + } + if gone { + drop_gsi_tables(data, builds).await; + return Ok(0); + } + return Err((0, e)); + } + for (i, b) in builds.iter().enumerate() { + let filled = backfill(data, table_id, &b.def, base_ks, attr_defs).await; + let published = match filled { + Ok(counts) => { + tracing::debug!( + index_id = %b.def.index_id, + indexed = counts.indexed, + skipped = counts.skipped, + "GSI backfill complete" + ); + sqlx::query( + "UPDATE indexes SET index_status = 'ACTIVE' \ + WHERE table_id = $1 AND index_id = $2 AND index_status = 'CREATING'", + ) + .bind(table_id) + .bind(&b.def.index_id) + .execute(catalog) + .await + .map(|r| r.rows_affected() > 0) + .map_err(db_error) + } + Err(e) => Err(e), + }; + match published { + Ok(true) => published_here += 1, + outcome => { + // Either the flip found no CREATING row or the build failed. If the + // row is gone the index was deleted (or its table was) while this + // ran: drop what the delete may have raced past, and move on. + match index_row_exists(catalog, table_id, &b.def.index_id).await { + Ok(false) => { + drop_gsi_tables(data, std::slice::from_ref(b)).await; + tracing::info!( + index_id = %b.def.index_id, + "GSI was deleted while it was being built" + ); + } + Ok(true) => { + let e = match outcome { + Err(e) => e, + // Present but not CREATING: someone else published it. + Ok(_) => continue, + }; + return Err((i, e)); + } + Err(check) => { + let e = match outcome { + Err(e) => e, + Ok(_) => check, + }; + return Err((i, e)); + } + } + } + } + } + Ok(published_here) +} + +/// Finish every GSI a crash left `CREATING`, once. +/// +/// An index whose build is still running holds its ownership lock and is +/// skipped. The data table is created if the crash came before it, and the +/// backfill is re-run from the start: it writes each item's current state, so +/// what the dead build already wrote does not need clearing. Only indexes on +/// ACTIVE tables are taken; a table being deleted takes its indexes with it. +/// +/// Returns the number of indexes this call published. +pub async fn reconcile_creating_gsis( + engine: &crate::PostgresEngine, +) -> Result { + reconcile_pass(engine, None).await +} + +/// When a recovery sweep may next try an index whose rebuild failed. +/// +/// A rebuild that fails every time (a key too large for the index, say) would +/// otherwise rescan the table and take its write barrier every minute forever. +#[derive(Default)] +pub(crate) struct RecoveryBackoff { + failures: std::collections::HashMap, +} + +impl RecoveryBackoff { + const FIRST: Duration = Duration::from_secs(60); + const MAX: Duration = Duration::from_secs(3_600); + + fn due(&self, index_id: &str) -> bool { + self.failures + .get(index_id) + .is_none_or(|(_, next)| std::time::Instant::now() >= *next) + } + + fn failed(&mut self, index_id: &str) { + let n = self.failures.get(index_id).map_or(0, |(n, _)| *n) + 1; + let wait = Self::FIRST + .saturating_mul(2u32.saturating_pow(n.saturating_sub(1))) + .min(Self::MAX); + self.failures + .insert(index_id.to_owned(), (n, std::time::Instant::now() + wait)); + } + + fn succeeded(&mut self, index_id: &str) { + self.failures.remove(index_id); + } + + /// Forget indexes that are no longer CREATING (published or deleted). + fn retain<'a>(&mut self, live: impl Iterator) { + let live: std::collections::HashSet<&str> = live.collect(); + self.failures.retain(|id, _| live.contains(id.as_str())); + } +} + +/// The periodic recovery sweep: runs once immediately, then every minute. +/// +/// Its own task rather than a step at startup, so a large rebuild does not hold +/// the server from serving (writes keep a CREATING index current in the meantime) +/// and does not delay the vector index sweep. +pub(crate) async fn recovery_worker( + engine: std::sync::Arc, + token: extenddb_storage::CancellationToken, +) { + const POLL_INTERVAL: Duration = Duration::from_secs(60); + let mut backoff = RecoveryBackoff::default(); + loop { + match reconcile_pass(&engine, Some(&mut backoff)).await { + Ok(0) => {} + Ok(n) => tracing::info!(finished = n, "finished GSI build(s) left incomplete"), + Err(e) => tracing::error!("incomplete GSI build sweep failed: {e}"), + } + if !extenddb_storage::sleep_or_shutdown(&token, POLL_INTERVAL).await { + return; + } + } +} + +async fn reconcile_pass( + engine: &crate::PostgresEngine, + mut backoff: Option<&mut RecoveryBackoff>, +) -> Result { + type Row = ( + String, + String, + serde_json::Value, + serde_json::Value, + serde_json::Value, + serde_json::Value, + String, + String, + ); + let rows: Vec = sqlx::query_as( + "SELECT i.table_id, i.index_id, i.key_schema, i.projection, \ + t.key_schema, t.attribute_definitions, t.table_name, i.index_name \ + FROM indexes i JOIN tables t ON t.table_id = i.table_id \ + WHERE i.index_type = 'GSI' AND i.index_status = 'CREATING' \ + AND t.table_status = 'ACTIVE' \ + ORDER BY i.table_id, i.index_id", + ) + .fetch_all(&engine.pool) + .await + .map_err(db_error)?; + + if let Some(b) = backoff.as_mut() { + b.retain(rows.iter().map(|r| r.1.as_str())); + } + let mut finished = 0usize; + for (table_id, index_id, idx_ks, idx_proj, base_ks, attr_defs, table_name, index_name) in rows { + if backoff.as_ref().is_some_and(|b| !b.due(&index_id)) { + continue; + } + let Some(owner) = + crate::data::vector_index::build_ownership(&engine.data_pool, &index_id).await + else { + continue; + }; + // Re-read under ownership: the list above is a snapshot, and the build + // that held this index a moment ago may have published or abandoned it. + let still_building: Result, _> = sqlx::query_as( + "SELECT 1 FROM indexes i JOIN tables t ON t.table_id = i.table_id \ + WHERE i.index_id = $1 AND i.index_status = 'CREATING' AND t.table_status = 'ACTIVE'", + ) + .bind(&index_id) + .fetch_optional(&engine.pool) + .await; + match still_building { + Ok(Some(_)) => {} + Ok(None) => continue, + Err(e) => { + tracing::warn!(index_id = %index_id, "could not re-check a CREATING GSI: {e}"); + continue; + } + } + let parse = || -> Result<_, serde_json::Error> { + Ok(( + serde_json::from_value::>(idx_ks.clone())?, + serde_json::from_value::(idx_proj.clone())?, + serde_json::from_value::>(base_ks.clone())?, + serde_json::from_value::>(attr_defs.clone())?, + )) + }; + let (key_schema, projection, base_ks, attr_defs) = match parse() { + Ok(parsed) => parsed, + Err(e) => { + tracing::error!(index_id = %index_id, "unreadable CREATING GSI definition: {e}"); + if let Some(b) = backoff.as_mut() { + b.failed(&index_id); + } + continue; + } + }; + + let made = async { + let mut tx = engine.data_pool.begin().await.map_err(db_error)?; + let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(index_table_name(&index_id)) + .fetch_one(&mut *tx) + .await + .map_err(db_error)?; + if !exists { + crate::PostgresEngine::create_index_data_table( + &mut tx, + &index_id, + &key_schema, + &attr_defs, + &base_ks, + &attr_defs, + ) + .await?; + } + tx.commit().await.map_err(db_error) + } + .await; + if let Err(e) = made { + tracing::error!( + index_id = %index_id, + "could not prepare a CREATING GSI for rebuild: {e}" + ); + if let Some(b) = backoff.as_mut() { + b.failed(&index_id); + } + continue; + } + + let build = GsiBuild { + index_name: index_name.clone(), + def: GsiIndexDef { + index_id: index_id.clone(), + key_schema, + projection, + }, + _owner: Some(owner), + }; + match build_gsis( + &engine.pool, + &engine.data_pool, + &table_id, + &base_ks, + &attr_defs, + std::slice::from_ref(&build), + ) + .await + { + Ok(n) => { + finished += n; + if let Some(b) = backoff.as_mut() { + b.succeeded(&index_id); + } + if n > 0 { + tracing::info!( + table = %table_name, + index = %index_name, + "finished a GSI build left incomplete" + ); + } + } + Err((_, e)) if is_transient(&e) => { + tracing::warn!( + table = %table_name, + index = %index_name, + "could not finish an incomplete GSI build; leaving it CREATING: {e}" + ); + if let Some(b) = backoff.as_mut() { + b.failed(&index_id); + } + } + Err((_, e)) => { + // The same answer UpdateTable gives: an error that would recur on + // every attempt removes the index rather than rebuilding it hourly. + tracing::error!( + table = %table_name, + index = %index_name, + index_id = %index_id, + "GSI build fails on every attempt; removing the index: {e}" + ); + remove_failed_build(&engine.pool, &engine.data_pool, &table_id, &build).await; + if let Some(b) = backoff.as_mut() { + b.succeeded(&index_id); + } + } + } + } + Ok(finished) +} + +impl crate::PostgresEngine { + /// [`fix_up_stale_write`] for a write that ran without a transaction. + /// + /// That path is only taken when the write's catalog read found no index at + /// all, so every GSI the re-read finds is one it missed. The write has already + /// committed, so the repair goes in its own transaction, and a failure is + /// logged rather than returned: reporting an error for a write that happened + /// would invite a retry of something that already took effect. + pub(crate) async fn fix_up_autocommit( + &self, + stamp: WriteStamp, + key_info: &TableKeyInfo, + item: &Item, + ) { + if stamp.elapsed() <= GRACE_WRITE { + return; + } + let repaired = async { + // Asked before a transaction is opened, so a slow write on a table + // that still has no index costs one catalog read and nothing more. + let missed = missed_gsis(&self.pool, &key_info.table_id, &[]).await?; + if missed.is_empty() { + return Ok(0); + } + let mut tx = self.data_pool.begin().await.map_err(db_error)?; + let n = queue_repairs(&mut tx, key_info, &missed, item).await?; + tx.commit().await.map_err(db_error)?; + Ok::(n) + } + .await; + match repaired { + Ok(0) => {} + Ok(_) => { + if let Some(q) = &self.gsi_queue { + q.notify_workers(); + } + } + Err(e) => tracing::error!( + table_id = %key_info.table_id, + "could not queue a GSI repair for a slow write; the index may miss this item \ + until it is written again: {e}" + ), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use extenddb_core::types::KeyType; + + fn ks(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } + } + + fn ad(name: &str, t: ScalarAttributeType) -> AttributeDefinition { + AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type: t, + } + } + + #[test] + fn indexable_requires_every_key_of_the_declared_type() { + let index_ks = vec![ks("g", KeyType::Hash), ks("n", KeyType::Range)]; + let defs = vec![ + ad("g", ScalarAttributeType::S), + ad("n", ScalarAttributeType::N), + ]; + let mut item = Item::new(); + item.insert("g".into(), AttributeValue::S("a".into())); + assert!(!indexable(&item, &index_ks, &defs), "missing range key"); + item.insert("n".into(), AttributeValue::S("1".into())); + assert!(!indexable(&item, &index_ks, &defs), "wrong type"); + item.insert("n".into(), AttributeValue::N("1".into())); + assert!(indexable(&item, &index_ks, &defs)); + item.insert("g".into(), AttributeValue::S(String::new())); + assert!(!indexable(&item, &index_ks, &defs), "empty string key"); + } + + #[test] + fn indexable_is_false_without_a_definition() { + let index_ks = vec![ks("g", KeyType::Hash)]; + let mut item = Item::new(); + item.insert("g".into(), AttributeValue::S("a".into())); + assert!(!indexable(&item, &index_ks, &[])); + } + + #[test] + fn stamp_elapsed_grows_and_is_not_max() { + let s = WriteStamp::now(); + std::thread::sleep(Duration::from_millis(5)); + let e = s.elapsed(); + assert!(e >= Duration::from_millis(5), "{e:?}"); + assert!(e < Duration::from_secs(60), "{e:?}"); + } + + #[test] + fn grace_publish_exceeds_grace_write() { + assert!(GRACE_PUBLISH > GRACE_WRITE); + } + + #[test] + fn transient_errors_are_left_to_recovery() { + let e = |m: &str| StorageError::Internal(m.to_owned()); + assert!(is_transient(&e("SQLSTATE 40P01: deadlock detected"))); + assert!(is_transient(&e("SQLSTATE 55P03: lock timeout"))); + assert!(is_transient(&e("SQLSTATE 08006: connection failure"))); + assert!(is_transient(&e("SQLSTATE 58030: could not read block"))); + assert!(is_transient(&e("SQLSTATE XX000: internal error"))); + assert!(is_transient(&e( + "pool timed out while waiting for an open connection" + ))); + assert!(!is_transient(&e( + "SQLSTATE 54000: index row size exceeds maximum" + ))); + assert!(!is_transient(&e("SQLSTATE 22P02: invalid input syntax"))); + assert!(is_transient(&e( + "error communicating with database: broken pipe" + ))); + assert!(!is_transient(&e("invalid type: null, expected a string"))); + assert!(!is_transient(&e("missing sort key sk"))); + } +} diff --git a/crates/storage-postgres/src/data/index.rs b/crates/storage-postgres/src/data/index.rs index 7340c66da..09c5627d7 100644 --- a/crates/storage-postgres/src/data/index.rs +++ b/crates/storage-postgres/src/data/index.rs @@ -50,6 +50,26 @@ pub(crate) struct IndexMeta { /// Per-GSI propagation delay in milliseconds. `None` means use system /// default. `Some(0)` means synchronous. pub(super) propagation_delay_ms: Option, + /// The table's attribute definitions, read in the same statement as the index + /// rows rather than taken from the cached `TableKeyInfo`. + /// + /// An UpdateTable that adds a GSI also adds the definitions of its key + /// attributes, and the engine drops its cached key info only when the + /// UpdateTable returns, which is after the backfill. A write in between that + /// used the cached set would build this index's row without the columns of a + /// sort key it does not know the type of, and would skip the type check on + /// the new key attributes. + pub(super) attribute_definitions: Vec, +} + +/// The table's attribute definitions as of the write's catalog read, falling back +/// to the cached set when the table has no secondary index to have carried them. +/// See [`IndexMeta::attribute_definitions`]. +pub(crate) fn fresh_attr_defs<'a>( + indexes: &'a [IndexMeta], + cached: &'a [AttributeDefinition], +) -> &'a [AttributeDefinition] { + indexes.first().map_or(cached, |i| &i.attribute_definitions) } /// Both index families a write has to maintain, read in one catalog visit. @@ -58,6 +78,11 @@ pub(crate) struct IndexMeta { /// pool, because every write takes this path and most tables have no vector index /// at all. Also the single place the two reads are ordered together, so the six /// write sites cannot end up disagreeing about which membership they saw. +/// +/// Also returns the moment of the read, for `gsi_build::fix_up_stale_write`. +/// Taken after the connection is acquired, not before: time spent waiting for the +/// catalog pool is time before the read's snapshot, and counting it would send +/// every write under a saturated pool back to that same pool for a second read. pub(crate) async fn fetch_write_path_indexes( table_id: &str, pool: &sqlx::PgPool, @@ -65,6 +90,7 @@ pub(crate) async fn fetch_write_path_indexes( ( Vec, Vec<(extenddb_storage::vector_lifecycle::VectorIndexMeta, String)>, + super::gsi_build::WriteStamp, ), StorageError, > { @@ -72,10 +98,13 @@ pub(crate) async fn fetch_write_path_indexes( .acquire() .await .map_err(|e| StorageError::Internal(e.to_string()))?; + let stamp = super::gsi_build::WriteStamp::now(); let indexes = fetch_indexes_for_table(table_id, &mut *conn).await?; let vector_metas = crate::data::vector_index::fetch_vector_indexes_for_table(&mut *conn, table_id).await?; - Ok((indexes, vector_metas)) + drop(conn); + super::gsi_build::pause_after_index_read().await; + Ok((indexes, vector_metas, stamp)) } /// Fetch all index metadata for a table from the catalog. @@ -86,16 +115,34 @@ pub(crate) async fn fetch_indexes_for_table<'e, E>( where E: sqlx::PgExecutor<'e>, { - let rows: Vec<(String, String, String, serde_json::Value, serde_json::Value, Option)> = sqlx::query_as( - "SELECT index_id, index_name, index_type, key_schema, projection, propagation_delay_ms FROM indexes WHERE table_id = $1", + type Row = ( + String, + String, + String, + serde_json::Value, + serde_json::Value, + Option, + serde_json::Value, + ); + let rows: Vec = sqlx::query_as( + "SELECT i.index_id, i.index_name, i.index_type, i.key_schema, i.projection, \ + i.propagation_delay_ms, t.attribute_definitions \ + FROM indexes i JOIN tables t ON t.table_id = i.table_id \ + WHERE i.table_id = $1", ) .bind(table_id) .fetch_all(executor) .await .map_err(|e| StorageError::Internal(e.to_string()))?; + // Every row carries the same table definitions; parse them once. + let attribute_definitions: Vec = match rows.first() { + Some(row) => serde_json::from_value(row.6.clone()) + .map_err(|e| StorageError::Internal(e.to_string()))?, + None => Vec::new(), + }; rows.into_iter() - .map(|(id, index_name, idx_type, ks_json, proj_json, delay)| { + .map(|(id, index_name, idx_type, ks_json, proj_json, delay, _)| { let key_schema: Vec = serde_json::from_value(ks_json) .map_err(|e| StorageError::Internal(e.to_string()))?; let projection: Projection = serde_json::from_value(proj_json) @@ -107,6 +154,7 @@ where key_schema, projection, propagation_delay_ms: delay, + attribute_definitions: attribute_definitions.clone(), }) }) .collect() @@ -182,31 +230,27 @@ pub(super) fn effective_delay(idx: &IndexMeta, system_default: u64) -> u64 { pub(crate) async fn sync_indexes( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, base_key_schema: &[KeySchemaElement], - attr_defs: &[AttributeDefinition], indexes: &[IndexMeta], old_item: Option<&Item>, new_item: Option<&Item>, system_default_delay: u64, ) -> Result<(), StorageError> { - for idx in indexes { - if idx.index_type != "LSI" && effective_delay(idx, system_default_delay) != 0 { - continue; // Async — handled after commit. LSIs are always synchronous. - } + // LSIs first. An LSI exists from CreateTable, so no write ever missed one + // and none can vanish under a write: the old image says exactly which row to + // remove, and no savepoint is needed. + for idx in indexes.iter().filter(|idx| idx.index_type == "LSI") { + let attr_defs = &idx.attribute_definitions; let idx_table = index_table_name(&idx.index_id); - let idx_sks = all_sort_key_info(&idx.key_schema, attr_defs); let base_sks = all_sort_key_info(base_key_schema, attr_defs); - - // Delete old index row if the old item had index keys if let Some(old) = old_item && item_has_index_keys(old, &idx.key_schema) { delete_index_row_multi(tx, &idx_table, old, base_key_schema, &base_sks).await?; } - - // Insert new index row if the new item has index keys if let Some(new) = new_item - && item_has_index_keys(new, &idx.key_schema) + && super::gsi_build::indexable(new, &idx.key_schema, attr_defs) { + let idx_sks = all_sort_key_info(&idx.key_schema, attr_defs); let projected = project_item_for_index(new, &idx.key_schema, base_key_schema, &idx.projection); insert_index_row_multi( @@ -222,7 +266,94 @@ pub(crate) async fn sync_indexes( .await?; } } - Ok(()) + + let sync_gsis: Vec<&IndexMeta> = indexes + .iter() + .filter(|idx| idx.index_type != "LSI" && effective_delay(idx, system_default_delay) == 0) + .collect(); + if sync_gsis.is_empty() { + return Ok(()); + } + + // A GSI is written the way the online backfill and the queue write it: the + // item's current state, replacing whatever rows the base key had. A write + // that read the catalog before the index existed can leave a row under an + // item it deleted or whose key it removed, and a later write that finds no + // old image, or one without the key, would otherwise insert next to it. + // + // Under a savepoint, because an index may be gone by the time this runs + // (deleted, or a failed build removed it) and that is the index's problem, + // not this write's; the queue tolerates the same race. One savepoint for all + // of the write's GSIs rather than one each, which saves two round trips per + // index past the first. The price is paid only on the rare path: rolling back + // to it undoes every GSI written in that pass, so the pass runs again without + // the index that is gone. Each GSI's work is state-based, so redoing it gives + // the same rows. + let mut gone: Vec<&str> = Vec::new(); + loop { + sqlx::query("SAVEPOINT gsi_sync_apply") + .execute(&mut **tx) + .await + .map_err(db_error)?; + let mut failed: Option<(&str, StorageError)> = None; + for idx in sync_gsis + .iter() + .filter(|idx| !gone.contains(&idx.index_id.as_str())) + { + let attr_defs = &idx.attribute_definitions; + let idx_table = index_table_name(&idx.index_id); + let base_sks = all_sort_key_info(base_key_schema, attr_defs); + let applied = match new_item { + Some(new) => super::gsi_build::write_state( + tx, + &idx_table, + &idx.key_schema, + &idx.projection, + base_key_schema, + attr_defs, + &base_sks, + new, + ) + .await + .map(|_| ()), + None => match old_item { + Some(old) => { + delete_index_row_multi(tx, &idx_table, old, base_key_schema, &base_sks) + .await + } + None => Ok(()), + }, + }; + if let Err(e) = applied { + failed = Some((idx.index_id.as_str(), e)); + break; + } + } + match failed { + None => { + // Released rather than left open: PostgreSQL keeps every savepoint + // of a transaction until commit, and TransactWriteItems runs this + // for each of its items. + sqlx::query("RELEASE SAVEPOINT gsi_sync_apply") + .execute(&mut **tx) + .await + .map_err(db_error)?; + return Ok(()); + } + Some((index_id, e)) if crate::gsi_queue::is_undefined_table(&e) => { + sqlx::query("ROLLBACK TO SAVEPOINT gsi_sync_apply") + .execute(&mut **tx) + .await + .map_err(db_error)?; + tracing::debug!( + index_id, + "GSI data table gone during a synchronous write, skipping" + ); + gone.push(index_id); + } + Some((_, e)) => return Err(e), + } + } } /// Enqueue async GSI propagation for a write, **one `gsi_pending` row per @@ -258,7 +389,7 @@ pub(crate) async fn enqueue_async_indexes( } let context = GsiApplyContext { base_key_schema: key_info.key_schema.clone(), - attribute_definitions: key_info.attribute_definitions.clone(), + attribute_definitions: idx.attribute_definitions.clone(), index: GsiIndexDef { index_id: idx.index_id.clone(), key_schema: idx.key_schema.clone(), @@ -360,7 +491,7 @@ pub(crate) async fn insert_index_row_multi( .await } -/// Insert rows into one index table in a single multi-row INSERT. +/// Write rows into one index table in a single multi-row upsert. /// /// Each row is `(item, projected)`: the item supplies the index and base key /// columns, the projection is stored as `item_data`. The one implementation @@ -404,10 +535,30 @@ pub(crate) async fn insert_index_rows_multi( format!("({})", ps.join(", ")) }) .collect(); + // An upsert on the table's primary key, (pk, base_pk, base sort keys). The + // online backfill and a write can both reach the same index row; each holds + // the base row lock while it writes, so both carry the same values and the + // later one may simply win. A restore copies into a table nothing else + // writes, so for it this is a plain insert. + let mut conflict_cols = vec!["pk".to_owned(), "base_pk".to_owned()]; + for (i, &(_, sk_type)) in base_sks.iter().enumerate() { + conflict_cols.push(if i == 0 { + format!("base_{}", sk_column(sk_type)) + } else { + format!("base_{}", sk_column_n(i, sk_type)) + }); + } + let updates: Vec = cols + .iter() + .filter(|c| !conflict_cols.contains(c)) + .map(|c| format!("{c} = EXCLUDED.{c}")) + .collect(); let sql = format!( - "INSERT INTO {idx_table} ({}) VALUES {}", + "INSERT INTO {idx_table} ({}) VALUES {} ON CONFLICT ({}) DO UPDATE SET {}", cols.join(", "), - tuples.join(", ") + tuples.join(", "), + conflict_cols.join(", "), + updates.join(", ") ); let mut query = sqlx::query(&sql); diff --git a/crates/storage-postgres/src/data/mod.rs b/crates/storage-postgres/src/data/mod.rs index f4da8958b..a9010220d 100755 --- a/crates/storage-postgres/src/data/mod.rs +++ b/crates/storage-postgres/src/data/mod.rs @@ -132,6 +132,7 @@ macro_rules! bind_sk_execute { mod data_engine; mod ddl; mod delete_item; +pub(crate) mod gsi_build; pub(crate) mod index; mod put_item; mod query; @@ -141,7 +142,4 @@ mod tx_helpers; mod update_item; pub mod vector_index; -pub(crate) use index::{ - delete_index_row_multi, insert_index_row_multi, insert_index_rows_multi, item_has_index_keys, - project_item_for_index, -}; +pub(crate) use index::{insert_index_rows_multi, item_has_index_keys, project_item_for_index}; diff --git a/crates/storage-postgres/src/data/put_item.rs b/crates/storage-postgres/src/data/put_item.rs index 35e5b339a..2fbc3102b 100755 --- a/crates/storage-postgres/src/data/put_item.rs +++ b/crates/storage-postgres/src/data/put_item.rs @@ -9,6 +9,7 @@ use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; use extenddb_storage::util::{composite_pk_to_text, parse_sk, pk_to_text, sk_column, sk_info}; +use super::gsi_build::fix_up_stale_write; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; use super::query::check_condition; use super::tx_helpers::write_stream_record_in_tx; @@ -46,12 +47,12 @@ impl PostgresEngine { // cached membership gate, and that is gone: an index takes effect the moment // its catalog row commits. What remains is a window between this read and // the data transaction's commit, in which an index created concurrently is - // missed. That window cannot be closed here, because the catalog and the - // data tables are different databases and no transaction spans them. It is - // also exactly the window the secondary indexes have, for the same reason - // and with the same read: parity with a GSI is the bar, and the backfill - // that publishes a new index is what covers writes older than it. - let (indexes, vector_metas) = + // missed. No transaction can close it, because the catalog and the data + // tables are different databases. For a GSI it is closed by time instead: + // the read's stamp lets a write that was slow to commit queue a repair + // for any GSI it did not see (`gsi_build::fix_up_stale_write`), and a fast + // write is one the backfill sees. Vector indexes still have the window. + let (indexes, vector_metas, stamp) = fetch_write_path_indexes(&key_info.table_id, &self.pool).await?; // Index key attributes present in the item must match their declared @@ -69,7 +70,7 @@ impl PostgresEngine { extenddb_core::validation::validate_index_keys( &item, &index_refs, - &key_info.attribute_definitions, + super::index::fresh_attr_defs(&indexes, &key_info.attribute_definitions), ) .map_err(|e| StorageError::Validation(e.to_string()))?; } @@ -86,11 +87,15 @@ impl PostgresEngine { }; // When there's a condition, return_old, indexes, or stream capture, we need a transaction + // A write that is already slow takes the transaction path even with no + // index to maintain, so that a repair for a GSI created meanwhile commits + // with the write (`gsi_build::fix_up_stale_write`) instead of after it. let needs_tx = condition.is_some() || return_old || !indexes.is_empty() || !vector_metas.is_empty() - || stream.is_some(); + || stream.is_some() + || stamp.elapsed() > super::gsi_build::GRACE_WRITE; if let Some((sk_name, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) @@ -174,7 +179,6 @@ impl PostgresEngine { sync_indexes( &mut tx, &key_info.key_schema, - &key_info.attribute_definitions, &indexes, oi.as_ref(), Some(&item), @@ -235,6 +239,10 @@ impl PostgresEngine { ) .await?; + async_enqueued += + fix_up_stale_write(&mut tx, &self.pool, stamp, key_info, &indexes, &item) + .await?; + tx.commit() .await .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -255,6 +263,7 @@ impl PostgresEngine { "INSERT INTO {ddb_table} (pk, {sk_col}, item_data) VALUES ($1, $2, $3) \ ON CONFLICT (pk, {sk_col}) DO UPDATE SET item_data = EXCLUDED.item_data" ); + super::gsi_build::pause_before_autocommit_statement().await; bind_sk_execute!( &upsert_sql, pk_text.as_str(), @@ -262,6 +271,7 @@ impl PostgresEngine { &item_json, &self.data_pool )?; + self.fix_up_autocommit(stamp, key_info, &item).await; Ok(None) } } else { @@ -347,7 +357,6 @@ impl PostgresEngine { sync_indexes( &mut tx, &key_info.key_schema, - &key_info.attribute_definitions, &indexes, oi.as_ref(), Some(&item), @@ -408,6 +417,10 @@ impl PostgresEngine { ) .await?; + async_enqueued += + fix_up_stale_write(&mut tx, &self.pool, stamp, key_info, &indexes, &item) + .await?; + tx.commit() .await .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -428,12 +441,14 @@ impl PostgresEngine { "INSERT INTO {ddb_table} (pk, item_data) VALUES ($1, $2) \ ON CONFLICT (pk) DO UPDATE SET item_data = EXCLUDED.item_data" ); + super::gsi_build::pause_before_autocommit_statement().await; sqlx::query(&upsert_sql) .bind(pk_text.as_str()) .bind(&item_json) .execute(&self.data_pool) .await .map_err(|e| StorageError::Internal(e.to_string()))?; + self.fix_up_autocommit(stamp, key_info, &item).await; Ok(None) } } diff --git a/crates/storage-postgres/src/data/transactions.rs b/crates/storage-postgres/src/data/transactions.rs index a6cd9f177..e86c00a90 100644 --- a/crates/storage-postgres/src/data/transactions.rs +++ b/crates/storage-postgres/src/data/transactions.rs @@ -101,6 +101,10 @@ impl PostgresEngine { // set would make a transaction skip an index another request has just // created. One read per table also keeps a multi-op transaction from // re-asking for the same answer. + // + // The stamp is for GSIs created between this read and the commit; see + // `gsi_build::fix_up_stale_write`. + let mut stamp: Option = None; let mut table_indexes: HashMap> = HashMap::new(); let mut table_vector_metas: HashMap< String, @@ -110,7 +114,10 @@ impl PostgresEngine { let name = transact_op_table_name(op); if !table_indexes.contains_key(name) { let tid = transact_op_table_id(op); - let (indexes, vector_metas) = fetch_write_path_indexes(tid, &self.pool).await?; + let (indexes, vector_metas, read_at) = + fetch_write_path_indexes(tid, &self.pool).await?; + // The earliest read bounds every table's staleness. + stamp.get_or_insert(read_at); table_indexes.insert(name.to_owned(), indexes); table_vector_metas.insert(name.to_owned(), vector_metas); } @@ -242,6 +249,25 @@ impl PostgresEngine { } } + // GSIs created since the index read, if this transaction took long enough + // that a backfill may have passed its items. Read once per table. + let mut missed: HashMap> = HashMap::new(); + if stamp.is_some_and(|s| s.elapsed() > super::gsi_build::GRACE_WRITE) { + for op in ops { + let name = transact_op_table_name(op); + if !missed.contains_key(name) { + let found = super::gsi_build::missed_gsis( + &self.pool, + transact_op_table_id(op), + &table_indexes[name], + ) + .await + .map_err(cancel_all)?; + missed.insert(name.to_owned(), found); + } + } + } + // Persist async GSI work inside the transaction. let mut needs_notify = false; for (op, (old_item, new_item)) in ops.iter().zip(op_items.iter()) { @@ -285,7 +311,18 @@ impl PostgresEngine { ) .await .map_err(cancel_all)?; - if n > 0 || vector_n > 0 { + let repair_n = match ( + missed.get(transact_op_table_name(op)), + new_item.as_ref().or(old_item.as_ref()), + ) { + (Some(m), Some(item)) if !m.is_empty() => { + super::gsi_build::queue_repairs(&mut tx, key_info, m, item) + .await + .map_err(cancel_all)? + } + _ => 0, + }; + if n > 0 || vector_n > 0 || repair_n > 0 { needs_notify = true; } } @@ -383,7 +420,7 @@ fn transact_op_table_id<'a>(op: &'a TransactWriteOp<'_>) -> &'a str { /// cancellation reasons (BLOCKER #3 fix). /// Build [`validation::IndexKeyRef`] views over the table's indexes for /// secondary-index key validation. -fn index_key_refs(indexes: &[IndexMeta]) -> Vec> { +pub(super) fn index_key_refs(indexes: &[IndexMeta]) -> Vec> { indexes .iter() .map(|idx| validation::IndexKeyRef { @@ -444,10 +481,12 @@ async fn execute_transact_write_op( // per-item cancellation reason; an empty index key is up-front input // validation (a top-level ValidationException). let idx_refs = index_key_refs(indexes); - validation::validate_index_key_types(item, &idx_refs, &key_info.attribute_definitions) - .map_err(|e| { - TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())) - })?; + validation::validate_index_key_types( + item, + &idx_refs, + super::index::fresh_attr_defs(indexes, &key_info.attribute_definitions), + ) + .map_err(|e| TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())))?; validation::validate_index_key_not_empty( item, &idx_refs, @@ -528,7 +567,6 @@ async fn execute_transact_write_op( sync_indexes( tx, &key_info.key_schema, - &key_info.attribute_definitions, indexes, existing.as_ref(), Some(item), @@ -578,7 +616,6 @@ async fn execute_transact_write_op( sync_indexes( tx, &key_info.key_schema, - &key_info.attribute_definitions, indexes, existing.as_ref(), None, @@ -635,7 +672,7 @@ async fn execute_transact_write_op( validation::validate_index_key_types( &item, &idx_refs, - &key_info.attribute_definitions, + super::index::fresh_attr_defs(indexes, &key_info.attribute_definitions), ) .map_err(|e| { TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())) @@ -712,7 +749,6 @@ async fn execute_transact_write_op( sync_indexes( tx, &key_info.key_schema, - &key_info.attribute_definitions, indexes, existing.as_ref(), Some(&item), diff --git a/crates/storage-postgres/src/data/update_item.rs b/crates/storage-postgres/src/data/update_item.rs index 72ce6edb9..80ab9a515 100755 --- a/crates/storage-postgres/src/data/update_item.rs +++ b/crates/storage-postgres/src/data/update_item.rs @@ -10,6 +10,7 @@ use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; use extenddb_storage::util::{parse_sk, pk_to_text, sk_column, sk_info}; +use super::gsi_build::fix_up_stale_write; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; use super::query::check_condition; use super::tx_helpers::write_stream_record_in_tx; @@ -58,12 +59,12 @@ impl PostgresEngine { // cached membership gate, and that is gone: an index takes effect the moment // its catalog row commits. What remains is a window between this read and // the data transaction's commit, in which an index created concurrently is - // missed. That window cannot be closed here, because the catalog and the - // data tables are different databases and no transaction spans them. It is - // also exactly the window the secondary indexes have, for the same reason - // and with the same read: parity with a GSI is the bar, and the backfill - // that publishes a new index is what covers writes older than it. - let (indexes, vector_metas) = + // missed. No transaction can close it, because the catalog and the data + // tables are different databases. For a GSI it is closed by time instead: + // the read's stamp lets a write that was slow to commit queue a repair + // for any GSI it did not see (`gsi_build::fix_up_stale_write`), and a fast + // write is one the backfill sees. Vector indexes still have the window. + let (indexes, vector_metas, stamp) = fetch_write_path_indexes(&key_info.table_id, &self.pool).await?; // Read whenever anything can propagate, secondary or vector. Gating this on // the secondary set alone made a vector-only table ignore the configured @@ -190,6 +191,26 @@ impl PostgresEngine { validation::validate_item_size(&item, self.max_item_size_bytes) .map_err(|e| StorageError::Validation(e.to_string()))?; + // Secondary-index key validation on the post-update item, as the + // TransactWriteItems update path and SQLite already do: a wrong-typed + // index key or an index key set to an empty value is a + // ValidationException, not an item the index then leaves out. + if !indexes.is_empty() { + let idx_refs = super::transactions::index_key_refs(&indexes); + validation::validate_index_key_types( + &item, + &idx_refs, + super::index::fresh_attr_defs(&indexes, &key_info.attribute_definitions), + ) + .map_err(|e| StorageError::Validation(e.to_string()))?; + validation::validate_index_key_not_empty( + &item, + &idx_refs, + validation::SecondaryIndexEmptyContext::UpdateExpression, + ) + .map_err(|e| StorageError::Validation(e.to_string()))?; + } + let new_item = if return_new { Some(item.clone()) } else { None }; let item_json = @@ -263,7 +284,6 @@ impl PostgresEngine { sync_indexes( &mut tx, &key_info.key_schema, - &key_info.attribute_definitions, &indexes, pre_mutation_item.as_ref(), Some(&item), @@ -310,6 +330,9 @@ impl PostgresEngine { ) .await?; + async_enqueued += + fix_up_stale_write(&mut tx, &self.pool, stamp, key_info, &indexes, key).await?; + tx.commit() .await .map_err(|e| StorageError::Internal(e.to_string()))?; diff --git a/crates/storage-postgres/src/data/vector_index.rs b/crates/storage-postgres/src/data/vector_index.rs index bb2852aa3..4a8c70bef 100644 --- a/crates/storage-postgres/src/data/vector_index.rs +++ b/crates/storage-postgres/src/data/vector_index.rs @@ -865,9 +865,16 @@ 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. let swept = sqlx::query( "DELETE FROM vector_index_holds \ - WHERE NOT (index_id = ANY($1)) AND created_at < NOW() - INTERVAL '1 minute'", + WHERE NOT (index_id = ANY($1)) \ + AND created_at < NOW() - CASE WHEN index_id = table_id \ + THEN INTERVAL '10 minutes' \ + ELSE INTERVAL '1 minute' END", ) .bind(live) .execute(&engine.data_pool) @@ -982,7 +989,7 @@ pub async fn build_ownership(data: &sqlx::PgPool, index_id: &str) -> Option::connect_with(&options).await { Ok(conn) => conn, Err(e) => { - tracing::warn!("could not open a session for vector build ownership: {e}"); + tracing::warn!("could not open a session for index build ownership: {e}"); return None; } }; @@ -998,7 +1005,7 @@ pub async fn build_ownership(data: &sqlx::PgPool, index_id: &str) -> Option Some(BuildOwner { _conn: conn }), Ok(false) => None, Err(e) => { - tracing::warn!("vector build ownership probe failed: {e}"); + tracing::warn!("index build ownership probe failed: {e}"); None } } diff --git a/crates/storage-postgres/src/delete_table.rs b/crates/storage-postgres/src/delete_table.rs index 2f56c318a..7475eed6a 100755 --- a/crates/storage-postgres/src/delete_table.rs +++ b/crates/storage-postgres/src/delete_table.rs @@ -130,42 +130,16 @@ impl PostgresEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; - // Drop data tables after catalog commit. - let mut data_tx = self - .data_pool - .begin() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - for idx_id in &index_ids { - Self::drop_index_data_table(&mut data_tx, idx_id).await?; - } - for index_id in &vector_index_ids { - Self::drop_vector_data_table(&mut data_tx, index_id).await?; - } - // Holds go with the table. A hold for a table that no longer exists - // would block claims for a table id that can never be released. - sqlx::query("DELETE FROM vector_index_holds WHERE table_id = $1") - .bind(&row.table_id) - .execute(&mut *data_tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - Self::drop_data_table(&mut data_tx, &row.table_id).await?; - - // Drop any still-pending GSI propagation rows for this table in the - // same transaction that drops the tables. Workers tolerate orphans - // (savepoint-skip on 42P01), but removing them here avoids the - // wasted claim→deserialize→attempt→skip cycle and the log noise, - // especially when many items were written just before deletion. - sqlx::query("DELETE FROM gsi_pending WHERE table_id = $1") - .bind(&row.table_id) - .execute(&mut *data_tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - data_tx - .commit() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + // Drop data tables after catalog commit. See `drop_table_storage` for + // why the queue is stopped before anything is dropped. + Self::drop_table_storage( + &self.data_pool, + &row.table_id, + &index_ids, + &vector_index_ids, + None, + ) + .await?; } else { // Async delete: set DELETING status with scheduled removal time. sqlx::query( diff --git a/crates/storage-postgres/src/gsi_queue.rs b/crates/storage-postgres/src/gsi_queue.rs index d9b551fc3..f068089ed 100644 --- a/crates/storage-postgres/src/gsi_queue.rs +++ b/crates/storage-postgres/src/gsi_queue.rs @@ -17,29 +17,30 @@ //! pending row reappears and is reprocessed — nothing is lost once enqueued. //! //! Rows are partitioned by a stable hash of the base table key, and each worker -//! owns one partition, so all updates to a given base item are applied in `id` -//! order by a single worker (**per-key FIFO**) while distinct keys propagate -//! concurrently. +//! owns one partition, so all rows for a given base item are claimed by a single +//! worker in `id` order while distinct keys propagate concurrently. Vector rows +//! apply the images they carry and rely on that order. Secondary index rows do +//! not: their apply writes the item's state at apply time (`resync_base_key`), +//! which is what lets the worker pass over an item whose base row a writer is +//! holding (`PARK_FOR`) and apply the rows behind it first. use std::sync::Arc; -use extenddb_core::types::{AttributeDefinition, Item, KeySchemaElement, Projection}; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, Item, KeySchemaElement, Projection, +}; use extenddb_storage::error::StorageError; use extenddb_storage::util::composite_pk_to_text; use serde::{Deserialize, Serialize}; use sqlx::PgPool; use tokio::sync::Notify; -use crate::data::{ - all_sort_key_info, delete_index_row_multi, index_table_name, insert_index_row_multi, - item_has_index_keys, project_item_for_index, -}; - /// Number of worker tasks consuming from the persistent queue. /// /// Fixed invariant: a row's `worker_partition` is `hash(base_key) % NUM_WORKERS` -/// and worker `i` consumes partition `i`, so all updates for a base item are -/// applied in `id` order by a single worker (per-key FIFO). Changing this while +/// and worker `i` consumes partition `i`, so all rows for a base item are +/// claimed in `id` order by a single worker (see the module doc for what that +/// does and does not guarantee). Changing this while /// the queue is non-empty could split a key's pending rows across workers; a /// rolling restart drains in-flight rows first, so changing it between deploys /// is safe. @@ -50,6 +51,101 @@ const NUM_WORKERS: u64 = 4; /// holding a long-lived transaction. const BATCH_SIZE: usize = 100; +/// How long a worker waits for a writer's row lock before giving the claim back. +/// +/// The state-based apply of a secondary index row reads the base row `FOR +/// SHARE`, so a client that holds an item's row open (a slow conditional write, +/// a long TransactWriteItems) would otherwise hold this worker's whole partition +/// for as long. The claim rolls back on the timeout and the item is passed over +/// for [`PARK_FOR`]; nothing is lost. +/// +/// Set for secondary index rows only. A vector row applies the images it +/// carries, so rows for one item have to apply in `id` order, and a vector row +/// waits for the base row as it always has. +const ROW_LOCK_TIMEOUT: &str = "2s"; + +/// How long an item whose base row was held is passed over after a lock wait +/// on it timed out. The claim takes the lowest due id, so without this the held +/// item's row is claimed again almost at once and every row behind it in the +/// partition waits for the whole hold. Passed over, the rest of the partition +/// drains; the item is tried again when this runs out, so it is retried about +/// every `PARK_FOR` + [`ROW_LOCK_TIMEOUT`], and each retry stalls the partition +/// for up to [`ROW_LOCK_TIMEOUT`]. +/// +/// The pass-over is by item, not by row: a held item usually has several rows +/// queued (one per asynchronous index for each write to it), and if each of them +/// had to time out on its own the partition would never move, since by the time +/// the second one gave up the first would be due again. A later row for a parked +/// item is given back as soon as it is claimed, without waiting on the lock: a +/// claim and a rollback per such row per cycle. +/// +/// Skipping a row is safe because the apply is state-based: a row writes the +/// item's state at apply time, not the image it carried, so applying a later +/// row for another key first changes nothing, and a later row for the same key +/// writes the same state once the hold ends. +const PARK_FOR: std::time::Duration = std::time::Duration::from_secs(2); + +/// A base item whose row a writer holds: the data table id and the item's key +/// attribute values in key schema order. +type HeldItem = (String, Vec); + +/// The key values that identify `item` in its table, for [`HeldItem`]. +fn held_item(table_id: &str, item: &Item, base_key_schema: &[KeySchemaElement]) -> HeldItem { + let key = base_key_schema + .iter() + .filter_map(|ks| item.get(&ks.attribute_name).cloned()) + .collect(); + (table_id.to_owned(), key) +} + +/// What this worker is passing over because a base row was held: the items, and +/// the queue rows for them that have been claimed and given back, each with the +/// time it may be tried again. Per worker, in memory: a restart simply tries +/// them at once, which is what happens without it. +#[derive(Default)] +struct Parked { + items: Vec<(HeldItem, std::time::Instant)>, + ids: Vec<(i64, std::time::Instant)>, +} + +impl Parked { + /// Ids still being passed over; drops the ones whose time is up. + fn ids(&mut self) -> Vec { + let now = std::time::Instant::now(); + self.ids.retain(|(_, until)| *until > now); + self.ids.iter().map(|(id, _)| *id).collect() + } + + /// When the item is being passed over, until when. + fn item_until(&mut self, item: &HeldItem) -> Option { + let now = std::time::Instant::now(); + self.items.retain(|(_, until)| *until > now); + self.items + .iter() + .find(|(held, _)| held == item) + .map(|(_, until)| *until) + } + + /// Pass the item and this row over for `PARK_FOR`, after a lock wait on its + /// base row timed out. + fn park(&mut self, id: i64, item: HeldItem) { + let until = std::time::Instant::now() + PARK_FOR; + self.items.push((item, until)); + self.ids.push((id, until)); + } + + /// Pass a later row for an item already parked over, until the item is due. + fn park_row(&mut self, id: i64, until: std::time::Instant) { + self.ids.push((id, until)); + } +} + +/// Whether a `StorageError` is PostgreSQL's lock_not_available (55P03), the +/// error `lock_timeout` raises. +fn is_lock_timeout(err: &StorageError) -> bool { + matches!(err, StorageError::Internal(msg) if msg.contains("SQLSTATE 55P03")) +} + /// Marker for PostgreSQL's "undefined_table" SQLSTATE (42P01) as embedded in /// `StorageError` messages by `db_error`. Matching the full `SQLSTATE 42P01` /// prefix (rather than the bare code) avoids false positives from item or @@ -57,7 +153,8 @@ const BATCH_SIZE: usize = 100; const PG_UNDEFINED_TABLE: &str = "SQLSTATE 42P01"; /// Stable partition for a base-table key (FNV-1a, mapped to a worker). Routes -/// all updates for one base item to a single worker for in-order apply. A local +/// all updates for one base item to a single worker, which claims them in `id` +/// order. A local /// hash (not `std`'s, which is not guaranteed stable across builds) keeps the /// mapping fixed for a given key forever. fn partition_for(base_pk_text: &str) -> i32 { @@ -99,7 +196,7 @@ pub(crate) struct GsiIndexDef { /// its own propagation delay (encoded in the row's `ready_at`) and propagates /// independently — mirroring the per-index queue items of the original /// in-memory design. Rows for the same base item share a `worker_partition` -/// and are applied in `id` order, preserving per-key FIFO. +/// and are claimed in `id` order. #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct GsiApplyContext { pub(crate) base_key_schema: Vec, @@ -212,7 +309,9 @@ fn jitter_delay_ms(delay_ms: u64) -> u64 { /// smaller jitter could become eligible first and be applied before an earlier /// write — leaving the index reflecting a stale value. The clamp guarantees a /// row's `ready_at` never precedes that of a lower-`id` row in the same -/// partition, preserving per-key FIFO convergence while still adding jitter. +/// partition, so rows for one item are claimed in write order while still adding +/// jitter. Vector rows, which apply the images they carry, rely on that order; +/// secondary index rows converge whatever the order (see the module doc). /// This matches the original design, where a partition's single consumer slept /// sequentially and so could never reorder updates. pub(crate) async fn enqueue_gsi_pending( @@ -234,7 +333,7 @@ pub(crate) async fn enqueue_gsi_pending( let context_json = serde_json::to_value(context).map_err(|e| StorageError::Internal(e.to_string()))?; - // Route all updates for a given base item to one worker (per-key FIFO). + // Route all updates for a given base item to one worker. // The base key is immutable over an item's lifetime; `new_item` carries it // for puts/updates, `old_item` for deletes. let worker_partition = match new_item.or(old_item) { @@ -276,13 +375,16 @@ async fn worker(worker_id: u64, q: Arc) { tracing::debug!("GSI worker {worker_id} started"); + let mut parked = Parked::default(); loop { - match process_batch(worker_id, &q).await { + match process_batch(worker_id, &q, &mut parked).await { Ok(0) => { // Nothing ready in this partition. Sleep until the earliest // not-yet-due row becomes eligible (honoring the propagation - // delay to the millisecond) or until a write wakes us. - let wait = next_ready_wait(&q.data_pool, worker_id) + // delay to the millisecond) or until a write wakes us. Parked + // rows are left out, or a partition whose only due row is held + // would spin on one millisecond waits until the park expires. + let wait = next_ready_wait(&q.data_pool, worker_id, &parked.ids()) .await .unwrap_or(MAX_IDLE) .min(MAX_IDLE); @@ -301,7 +403,11 @@ async fn worker(worker_id: u64, q: Arc) { /// eligible. `None` when the partition has no pending rows (caller waits its /// backstop interval). A row already due maps to a near-zero wait so the loop /// re-claims promptly. -async fn next_ready_wait(pool: &PgPool, worker_id: u64) -> Option { +async fn next_ready_wait( + pool: &PgPool, + worker_id: u64, + parked: &[i64], +) -> Option { // The ::float8 cast is load-bearing. Since PostgreSQL 14, EXTRACT returns // `numeric`, which sqlx refuses to decode into f64. Without the cast this // query fails on every partition that has a pending row, and the error was @@ -318,10 +424,11 @@ async fn next_ready_wait(pool: &PgPool, worker_id: u64) -> Option = sqlx::query_scalar( "SELECT EXTRACT(EPOCH FROM (MIN(ready_at) - NOW()))::float8 FROM gsi_pending p \ - WHERE worker_partition = $1 \ + WHERE worker_partition = $1 AND NOT (id = ANY($2)) \ AND NOT EXISTS (SELECT 1 FROM vector_index_holds h WHERE h.table_id = p.table_id)", ) .bind(worker_id as i32) + .bind(parked) .fetch_one(pool) .await .map_err(|e| tracing::error!("GSI worker {worker_id}: next_ready_wait failed: {e}")) @@ -340,20 +447,29 @@ async fn next_ready_wait(pool: &PgPool, worker_id: u64) -> Option Result { +async fn process_batch( + worker_id: u64, + q: &GsiQueue, + parked: &mut Parked, +) -> Result { let mut processed = 0usize; for _ in 0..BATCH_SIZE { + let skip = parked.ids(); let mut tx = q .data_pool .begin() @@ -370,7 +486,7 @@ async fn process_batch(worker_id: u64, q: &GsiQueue) -> Result = sqlx::query_as( "SELECT id, table_id, old_item, new_item, index_context FROM gsi_pending \ - WHERE worker_partition = $1 AND ready_at <= NOW() \ + WHERE worker_partition = $1 AND ready_at <= NOW() AND NOT (id = ANY($2)) \ AND NOT EXISTS ( \ SELECT 1 FROM vector_index_holds h WHERE h.table_id = gsi_pending.table_id \ ) \ @@ -379,6 +495,7 @@ async fn process_batch(worker_id: u64, q: &GsiQueue) -> Result Result {} + Err(e) if held.is_some() && is_lock_timeout(&e) => { + // A writer is holding the base row. Give the claim back and pass + // the item over for a while, so the rows behind it are not held + // up for the whole hold; see `PARK_FOR`. Only a secondary index + // row set the timeout; a vector row that meets one (a server-wide + // `lock_timeout`) is an error like any other and is retried. + let _ = tx.rollback().await; + tracing::debug!( + "GSI worker {worker_id}: base row held by a writer, passing over id={id}" + ); + if let Some(item) = held { + parked.park(id, item); + } + continue; + } + Err(e) => return Err(e), + } sqlx::query("DELETE FROM gsi_pending WHERE id = $1") .bind(id) @@ -410,22 +578,12 @@ async fn process_batch(worker_id: u64, q: &GsiQueue) -> Result, - id: i64, - table_id: &str, +/// Decode a claimed row's images and apply context. +fn decode_row( old_json: Option, new_json: Option, ctx_json: serde_json::Value, -) -> Result<(), StorageError> { +) -> Result<(Option, Option, PendingApplyContext), StorageError> { let old_item: Option = old_json .map(serde_json::from_value) .transpose() @@ -436,7 +594,24 @@ async fn apply_claimed_row( .map_err(|e| StorageError::Internal(e.to_string()))?; let context: PendingApplyContext = serde_json::from_value(ctx_json).map_err(|e| StorageError::Internal(e.to_string()))?; + Ok((old_item, new_item, context)) +} +/// Apply the GSI update for one claimed row, within the caller's transaction. +/// +/// The row's single target index is applied under a `SAVEPOINT`: if the index +/// table has been dropped (table deleted; SQLSTATE 42P01) the index is skipped +/// and the row is still consumed. Any other error propagates, rolling back the +/// whole transaction so the row is retried. +async fn apply_claimed_row( + worker_id: u64, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + id: i64, + table_id: &str, + old_item: Option<&Item>, + new_item: Option<&Item>, + context: &PendingApplyContext, +) -> Result<(), StorageError> { // One index per row. Guard the apply with a savepoint so a dropped-index // race (42P01) can be recovered and the row still consumed; aborting the // subtransaction otherwise poisons the outer transaction and blocks the @@ -446,26 +621,21 @@ async fn apply_claimed_row( .await .map_err(|e| StorageError::Internal(e.to_string()))?; - let applied = match &context { + let applied = match context { PendingApplyContext::Gsi(c) => { apply_index( tx, + table_id, &c.index, - old_item.as_ref(), - new_item.as_ref(), + old_item, + new_item, &c.base_key_schema, &c.attribute_definitions, ) .await } PendingApplyContext::Vector(c) => { - crate::data::vector_index::apply_claimed_vector_row( - tx, - c, - old_item.as_ref(), - new_item.as_ref(), - ) - .await + crate::data::vector_index::apply_claimed_vector_row(tx, c, old_item, new_item).await } }; @@ -499,39 +669,22 @@ async fn apply_claimed_row( /// Apply a single index's delete-old / insert-new within the transaction. async fn apply_index( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + table_id: &str, idx: &GsiIndexDef, old_item: Option<&Item>, new_item: Option<&Item>, base_key_schema: &[KeySchemaElement], attr_defs: &[AttributeDefinition], ) -> Result<(), StorageError> { - let idx_table = index_table_name(&idx.index_id); - let idx_sks = all_sort_key_info(&idx.key_schema, attr_defs); - let base_sks = all_sort_key_info(base_key_schema, attr_defs); - - if let Some(old) = old_item - && item_has_index_keys(old, &idx.key_schema) - { - delete_index_row_multi(tx, &idx_table, old, base_key_schema, &base_sks).await?; - } - - if let Some(new) = new_item - && item_has_index_keys(new, &idx.key_schema) - { - let projected = - project_item_for_index(new, &idx.key_schema, base_key_schema, &idx.projection); - insert_index_row_multi( - tx, - &idx_table, - new, - &projected, - &idx.key_schema, - base_key_schema, - &idx_sks, - &base_sks, - ) - .await?; - } - - Ok(()) + // The row says which item changed; what goes into the index is the item's + // state now, read under its row lock. Applying the image the row carried + // instead is only right if this row is the last word on the item, and three + // things break that: an online backfill that already wrote a newer state, + // a write that predated the index and told it nothing, and a second server + // process whose worker claims a later row for the same item first. + let Some(key_item) = new_item.or(old_item) else { + return Ok(()); + }; + crate::data::gsi_build::resync_base_key(tx, table_id, idx, base_key_schema, attr_defs, key_item) + .await } diff --git a/crates/storage-postgres/src/lib.rs b/crates/storage-postgres/src/lib.rs index 4313a48ea..e9f1ae2e9 100755 --- a/crates/storage-postgres/src/lib.rs +++ b/crates/storage-postgres/src/lib.rs @@ -39,6 +39,13 @@ pub use catalog_store::PostgresCatalogStore; pub use config::PostgresStorageConfig; pub use config::parse_connection_string; pub use credential_store::DbCredentialStore; +/// Finish GSI builds left `CREATING`, once. The server runs it from a background +/// worker; exported for its test. +pub use data::gsi_build::reconcile_creating_gsis; +/// Make writes pause at a chosen point, so a test can produce a write that missed +/// a GSI created in between. Hidden because no deployment path calls it. +#[doc(hidden)] +pub use data::gsi_build::{PauseWrites, pause_writes_for_tests}; /// Apply one queued vector row from its own context. /// /// Reachable so an integration test can drive the classification the propagation @@ -688,6 +695,14 @@ impl ServerRuntimeHooks for PostgresRuntimeHooks { workers::vector_stuck_build_worker(storage_for_builds, token).await; }); + // 6b. GSI builds a crash left CREATING. Not a startup step: a rebuild + // scans the table, and writes keep a CREATING index current meanwhile. + let storage_for_gsi_builds = self.engine.clone(); + let token = ctx.shutdown.clone(); + let gsi_builds = tokio::spawn(async move { + crate::data::gsi_build::recovery_worker(storage_for_gsi_builds, token).await; + }); + // 7. Pool metrics worker - needs both catalog and data pools let catalog_pool = self.engine.pool.clone(); let data_pool = self.engine.data_pool().clone(); @@ -712,6 +727,7 @@ impl ServerRuntimeHooks for PostgresRuntimeHooks { idempotency_cleanup, ttl, vector_builds, + gsi_builds, pool_metrics, gsi_poller, ] diff --git a/crates/storage-postgres/src/table_helpers.rs b/crates/storage-postgres/src/table_helpers.rs index 7da4ef419..4dfb7f18f 100755 --- a/crates/storage-postgres/src/table_helpers.rs +++ b/crates/storage-postgres/src/table_helpers.rs @@ -4,9 +4,8 @@ //! Helper types and methods for `TableEngine` operations. use extenddb_core::types::{ - AttributeDefinition, BillingMode, BillingModeSummary, GsiDescription, KeySchemaElement, - LsiDescription, Projection, ProvisionedThroughputDescription, SseDescription, SseType, - TableDescription, TableStatus, + BillingMode, BillingModeSummary, GsiDescription, LsiDescription, + ProvisionedThroughputDescription, SseDescription, SseType, TableDescription, TableStatus, }; use extenddb_storage::error::StorageError; use extenddb_storage::util::{index_arn, stream_arn}; @@ -98,104 +97,6 @@ impl PostgresEngine { data::index_table_name(index_id) } - /// Backfill existing items from the base table into a newly created GSI. - /// - /// Uses batched reads with OFFSET/LIMIT to avoid loading all items into - /// memory at once (CB-19). - // S2: Parameters mirror the SQL schema dimensions (account, table, index, - // key schemas, attribute defs, projection). A wrapper struct would obscure - // the call site without adding clarity. - #[allow(clippy::too_many_arguments)] - pub(crate) async fn backfill_gsi( - tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - table_id: &str, - index_id: &str, - index_key_schema: &[KeySchemaElement], - attr_defs: &[AttributeDefinition], - base_key_schema: &[KeySchemaElement], - base_attr_defs: &[AttributeDefinition], - projection: &Projection, - ) -> Result<(), StorageError> { - const BATCH_SIZE: i64 = 500; - - let base_table = data::data_table_name(table_id); - let idx_table = data::index_table_name(index_id); - - let idx_sks = data::all_sort_key_info(index_key_schema, attr_defs); - let base_sks = data::all_sort_key_info(base_key_schema, base_attr_defs); - - // Order by whatever key columns the base table actually has. A - // hash-only base table has no sort-key columns at all, so naming them - // unconditionally makes the backfill fail with "column sk_s does not - // exist" and the whole GSI add is rolled back. - let mut order_cols = vec!["pk".to_owned()]; - for (i, _) in base_sks.iter().enumerate() { - if i == 0 { - order_cols.extend(["sk_s".to_owned(), "sk_n".to_owned(), "sk_b".to_owned()]); - } else { - let n = i + 1; - order_cols.extend([format!("sk{n}_s"), format!("sk{n}_n"), format!("sk{n}_b")]); - } - } - let sql = format!( - "SELECT item_data FROM {base_table} ORDER BY {} LIMIT $1 OFFSET $2", - order_cols.join(", ") - ); - let mut offset: i64 = 0; - loop { - let rows: Vec<(serde_json::Value,)> = sqlx::query_as(&sql) - .bind(BATCH_SIZE) - .bind(offset) - .fetch_all(&mut **tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - if rows.is_empty() { - break; - } - - // BATCH_SIZE=500, always fits in i64. - let batch_len = i64::from(u16::try_from(rows.len()).unwrap_or(u16::MAX)); - - for (item_json,) in rows { - let item = data::json_to_item(item_json)?; - - let has_all_keys = index_key_schema - .iter() - .all(|ks| item.contains_key(&ks.attribute_name)); - if !has_all_keys { - continue; - } - - let projected = data::project_item_for_index( - &item, - index_key_schema, - base_key_schema, - projection, - ); - - data::insert_index_row_multi( - tx, - &idx_table, - &item, - &projected, - index_key_schema, - base_key_schema, - &idx_sks, - &base_sks, - ) - .await?; - } - - if batch_len < BATCH_SIZE { - break; - } - offset += batch_len; - } - - Ok(()) - } - // 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 diff --git a/crates/storage-postgres/src/update_table.rs b/crates/storage-postgres/src/update_table.rs index 6e9aa58b9..dec99a93b 100755 --- a/crates/storage-postgres/src/update_table.rs +++ b/crates/storage-postgres/src/update_table.rs @@ -7,9 +7,12 @@ use extenddb_core::types::{ AttributeDefinition, BillingMode, KeySchemaElement, TableDescription, UpdateTableInput, }; use extenddb_storage::error::StorageError; -use extenddb_storage::util::effective_attribute_definitions; +use extenddb_storage::util::{building_refuses, effective_attribute_definitions}; use crate::PostgresEngine; +use crate::data::gsi_build::{ + GsiBuild, build_gsis, drop_gsi_tables, is_transient, remove_failed_build, +}; /// Give back the build holds a failed `UpdateTable` took. /// @@ -36,6 +39,28 @@ async fn release_taken_holds( } } +/// Refuse a GSI Create, or a Delete of a GSI other than the one building, while +/// a GSI on the table is `CREATING`. +async fn refuse_while_building( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + table_id: &str, + updates: &[extenddb_core::types::GlobalSecondaryIndexUpdate], +) -> Result<(), StorageError> { + let building: Vec = sqlx::query_scalar( + "SELECT index_name FROM indexes WHERE table_id = $1 AND index_status = 'CREATING'", + ) + .bind(table_id) + .fetch_all(&mut **tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if building_refuses(&building, updates) { + return Err(StorageError::LimitExceeded( + extenddb_core::types::ONLINE_INDEX_LIMIT.to_owned(), + )); + } + Ok(()) +} + /// Refuse a new index whose name is already taken on this table, whatever index /// family holds it. /// @@ -413,6 +438,16 @@ impl PostgresEngine { // index DDL builds its columns from the same set the catalog now holds. let mut merged_attr_defs_for_ddl: Option> = None; if let Some(updates) = &input.global_secondary_index_updates { + { + // One index change at a time while a GSI builds, as the service + // has it: a Create, or a Delete of any other GSI, is refused with + // LimitExceededException (measured for both, in both of the + // service's CREATING phases). A Delete of the index that is + // building is a different case and is left to the delete path. + // Read under the table row lock taken above, which every + // UpdateTable takes, so two requests cannot both pass. + refuse_while_building(&mut tx, &table_id, updates).await?; + } for update in updates { if let Some(create) = &update.create { // Across families, not just this one: see @@ -435,7 +470,7 @@ impl PostgresEngine { r"INSERT INTO indexes (table_id, index_name, index_id, index_type, key_schema, projection, index_status, provisioned_throughput) - VALUES ($1, $2, $3, 'GSI', $4, $5, 'ACTIVE', $6)", + VALUES ($1, $2, $3, 'GSI', $4, $5, 'CREATING', $6)", ) .bind(&table_id) .bind(&create.index_name) @@ -447,9 +482,9 @@ impl PostgresEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; created_index_ids.push(index_id); - - // Create the index data table on the data pool (P54 Bug 1). - // Catalog metadata is committed first; data DDL follows. + // The data table is created before this transaction commits, and + // the index is published CREATING; see the block before the + // commit and `data::gsi_build`. } if let Some(delete) = &update.delete { @@ -748,6 +783,92 @@ impl PostgresEngine { return Err(e); } + // New GSIs get their data table BEFORE the catalog row commits. Every write + // reads the catalog fresh, so from the commit on, a write maintains this + // index; with the table created afterwards, writes in between failed on the + // missing table or queued rows the worker then dropped. + // + // Build ownership is taken here too, before the row is visible, so a peer's + // recovery sweep never sees a CREATING row with nobody holding it. It is an + // efficiency measure only: a second, concurrent build is correct, because a + // full backfill pass is enough on its own to publish the index. + let base_key_schema: Vec = serde_json::from_value(ks_json.clone()) + .map_err(|e| StorageError::Internal(e.to_string()))?; + let base_attr_defs: Vec = serde_json::from_value(ad_json.clone()) + .map_err(|e| StorageError::Internal(e.to_string()))?; + // Built from the merged and pruned set rather than the request's subset: a + // new index may key on an attribute the base table already defined, and on + // a conflicting redeclaration the stored type is the one that counts. + let effective_attr_defs: Vec = merged_attr_defs_for_ddl + .clone() + .unwrap_or_else(|| base_attr_defs.clone()); + let mut gsi_builds: Vec = Vec::new(); + if let Some(updates) = &input.global_secondary_index_updates { + let creates = updates.iter().filter_map(|u| u.create.as_ref()); + for (create, index_id) in creates.zip(&created_index_ids) { + let owner = + crate::data::vector_index::build_ownership(&self.data_pool, index_id).await; + if owner.is_none() { + // The build goes ahead without it: a second build of the same index + // by the recovery worker is harmless to the index's contents, and + // removal of a failed build is conditional on the row still being + // CREATING, so a published index cannot be taken away. + tracing::warn!( + "could not take build ownership for GSI '{}' on '{}'; the recovery \ + worker may build it a second time", + create.index_name, + input.table_name, + ); + } + let made = async { + let mut data_tx = self + .data_pool + .begin() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + Self::create_index_data_table( + &mut data_tx, + index_id, + &create.key_schema, + &effective_attr_defs, + &base_key_schema, + &base_attr_defs, + ) + .await?; + data_tx + .commit() + .await + .map_err(|e| StorageError::Internal(e.to_string())) + } + .await; + if let Err(e) = made { + tracing::error!( + "Failed to create data table for GSI '{}' on '{}': {e}", + create.index_name, + input.table_name, + ); + drop_gsi_tables(&self.data_pool, &gsi_builds).await; + release_taken_holds( + &self.data_pool, + &table_id, + &taken_holds, + &input.table_name, + ) + .await; + return Err(e); + } + gsi_builds.push(GsiBuild { + index_name: create.index_name.clone(), + def: crate::gsi_queue::GsiIndexDef { + index_id: index_id.clone(), + key_schema: create.key_schema.clone(), + projection: create.projection.clone(), + }, + _owner: owner, + }); + } + } + // The commit is inside the same recovery, because it is the one remaining path // that can fail after a hold has been taken. A failing commit is also the // moment the infrastructure is already unhappy, which is the worst time to @@ -758,86 +879,32 @@ impl PostgresEngine { .await .map_err(|e| StorageError::Internal(e.to_string())) { + // A commit can fail on the way back after it took effect. Drop a new + // GSI's data table only if its catalog row really is absent: a CREATING + // row with no table would fail every synchronous write to the table + // until recovery re-created it. + let mut uncommitted = Vec::new(); + for b in gsi_builds { + let present: Option<(i32,)> = + sqlx::query_as("SELECT 1 FROM indexes WHERE table_id = $1 AND index_id = $2") + .bind(&table_id) + .bind(&b.def.index_id) + .fetch_optional(&self.pool) + .await + .unwrap_or(Some((1,))); + if present.is_none() { + uncommitted.push(b); + } + } + drop_gsi_tables(&self.data_pool, &uncommitted).await; release_taken_holds(&self.data_pool, &table_id, &taken_holds, &input.table_name).await; return Err(e); } - // P54 Bug 1: Execute data DDL on the data pool after catalog commit. + // Drop the data tables of deleted GSIs, after the catalog commit. if let Some(updates) = &input.global_secondary_index_updates { - let base_key_schema: Vec = serde_json::from_value(ks_json.clone()) - .map_err(|e| StorageError::Internal(e.to_string()))?; - let base_attr_defs: Vec = serde_json::from_value(ad_json.clone()) - .map_err(|e| StorageError::Internal(e.to_string()))?; - // Build index columns from the merged and pruned set rather than the - // request's subset: a new index may key on an attribute the base - // table already defined and the request need not re-declare it, and - // on a conflicting redeclaration the stored type is the one the data - // table must be built from. - let effective_attr_defs = merged_attr_defs_for_ddl - .as_deref() - .unwrap_or(&base_attr_defs); - - let mut create_idx = 0usize; let mut delete_idx = 0usize; for update in updates { - if let Some(create) = &update.create { - let idx_id = &created_index_ids[create_idx]; - create_idx += 1; - let data_result = async { - let mut data_tx = self - .data_pool - .begin() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - Self::create_index_data_table( - &mut data_tx, - idx_id, - &create.key_schema, - effective_attr_defs, - &base_key_schema, - &base_attr_defs, - ) - .await?; - - Self::backfill_gsi( - &mut data_tx, - &table_id, - idx_id, - &create.key_schema, - effective_attr_defs, - &base_key_schema, - &base_attr_defs, - &create.projection, - ) - .await?; - - data_tx - .commit() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - Ok::<(), StorageError>(()) - } - .await; - - if let Err(e) = data_result { - tracing::error!( - "Failed to create data table for GSI '{}' on '{}', \ - cleaning up catalog: {e}", - create.index_name, - input.table_name, - ); - let _ = sqlx::query( - "DELETE FROM indexes WHERE table_id = $1 AND index_name = $2", - ) - .bind(&table_id) - .bind(&create.index_name) - .execute(&self.pool) - .await; - return Err(e); - } - } - if update.delete.is_some() { let idx_id = &deleted_index_ids[delete_idx]; delete_idx += 1; @@ -951,6 +1018,75 @@ impl PostgresEngine { } } + // Fill and publish the new GSIs, last, so a failed build cannot skip the + // rest of this request's post-commit work above. See `data::gsi_build` for + // why the wait and the barrier make the backfill see every write that did + // not see the index. + // + // A build that fails for a passing reason (a conflict, a timeout, a lost + // connection) leaves its index CREATING, with its data table, for the + // recovery worker to finish, and the request still succeeds, reporting the + // index as CREATING: the index exists, writes are maintaining it, and what + // is missing is a scan that can be re-run. One that would fail the same way + // every time is removed and reported, as before. + // + // In a task of its own, awaited here, so a client that disconnects cannot + // cancel the build halfway: that would leave the barrier's lock request or a + // page's row locks to be released whenever the pooled connection is next + // used. + if !gsi_builds.is_empty() { + let catalog = self.pool.clone(); + let data = self.data_pool.clone(); + let build_table_id = table_id.clone(); + let table_name = input.table_name.clone(); + let base_ks = base_key_schema.clone(); + let attr_defs = effective_attr_defs.clone(); + let task = tokio::spawn(async move { + let built = build_gsis( + &catalog, + &data, + &build_table_id, + &base_ks, + &attr_defs, + &gsi_builds, + ) + .await; + let Err((first_failed, e)) = built else { + return Ok(()); + }; + let unfinished = &gsi_builds[first_failed..]; + if is_transient(&e) { + for b in unfinished { + tracing::warn!( + "Could not finish building GSI '{}' on '{table_name}'; it stays \ + CREATING for the recovery worker to finish: {e}", + b.index_name, + ); + } + return Ok(()); + } + // A failure that would recur on every attempt: remove the index and + // report it, as before. The engine allows one index action per + // UpdateTable, so `unfinished` is the one build in practice; the + // loop is for the shape of the data, not for a second index. + for b in unfinished { + tracing::error!( + "Failed to build GSI '{}' on '{table_name}', removing it: {e}", + b.index_name, + ); + remove_failed_build(&catalog, &data, &build_table_id, b).await; + } + Err(e) + }); + match task.await { + Ok(result) => result?, + Err(e) => tracing::error!( + "GSI build task on '{}' ended abnormally; recovery will finish it: {e}", + input.table_name + ), + } + } + self.build_table_description(account_id, &input.table_name) .await } diff --git a/crates/storage-postgres/src/worker_store.rs b/crates/storage-postgres/src/worker_store.rs index 52ff9ced1..0f6871145 100755 --- a/crates/storage-postgres/src/worker_store.rs +++ b/crates/storage-postgres/src/worker_store.rs @@ -91,51 +91,18 @@ impl PostgresEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; - let dropped = async { - let mut data_tx = self - .data_pool - .begin() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - // CreateBackup holds ACCESS SHARE through its item read. Do not - // let a DROP waiting for that lock stall every transition in - // this worker pass; a timeout rolls this transaction back and - // leaves the catalog row DELETING for a later retry. - sqlx::query("SET LOCAL lock_timeout = '3s'") - .execute(&mut *data_tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - for idx_id in &index_ids { - Self::drop_index_data_table(&mut data_tx, idx_id).await?; - } - for idx_id in &vector_index_ids { - Self::drop_vector_data_table(&mut data_tx, idx_id).await?; - } - Self::drop_data_table(&mut data_tx, table_id).await?; - - // Drop any still-pending GSI propagation rows for this table in - // the same transaction that drops the tables, so workers don't - // waste a claim→deserialize→attempt→skip cycle on orphaned rows. - sqlx::query("DELETE FROM gsi_pending WHERE table_id = $1") - .bind(table_id) - .execute(&mut *data_tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - // The table is gone, so no vector build can release this hold; - // leaving it would block later claims for the table id. - sqlx::query("DELETE FROM vector_index_holds WHERE table_id = $1") - .bind(table_id) - .execute(&mut *data_tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - data_tx - .commit() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - Ok::<(), StorageError>(()) - } + // 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 { diff --git a/crates/storage-postgres/tests/gsi_online_build.rs b/crates/storage-postgres/tests/gsi_online_build.rs new file mode 100644 index 000000000..969ede4a8 --- /dev/null +++ b/crates/storage-postgres/tests/gsi_online_build.rs @@ -0,0 +1,1170 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 +//! Storage-level tests for adding a GSI to a table that is taking writes. +//! +//! The wire test (`tests/test_gsi_online_backfill.py`) shows the end result +//! under real concurrency. These pin down the parts that concurrency only hits +//! sometimes: +//! +//! - a queued index update applied after a newer one for the same item, which a +//! second server process, or a backfill that wrote the newer state first, can +//! both cause; +//! - a GSI left `CREATING` by a crash, which the recovery sweep has to finish. +//! +//! Each test builds its own throwaway database, applies the shipped migrations +//! to it, and drops it when it passes. A failing test leaves its database behind +//! on purpose, named `eddb_gsib_*`, so the state that failed can be inspected. +//! +//! Requires `EXTENDDB_TEST_PG_CONNECTION_STRING`, a base URL with no database +//! component (for example `postgresql://postgres@127.0.0.1:5432`), pointing at a +//! server whose role may create and drop databases. Without it every test here +//! reports a skip and passes, the same convention the wire suites use. + +use std::collections::{BTreeMap, BTreeSet}; +use std::time::Duration; + +use extenddb_core::expression::ExpressionMaps; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, BillingMode, CreateTableInput, Item, KeySchemaElement, + KeyType, ScalarAttributeType, TableKeyInfo, UpdateTableInput, +}; +use extenddb_storage::{DataEngine, TableEngine}; +use extenddb_storage_postgres::{PostgresConfig, PostgresEngine}; +use sqlx::PgPool; +use sqlx::postgres::PgPoolOptions; + +const ACCOUNT: &str = "123456789012"; +const REGION: &str = "us-east-1"; +const TABLE: &str = "t_gsi_online_build"; + +struct Scratch { + engine: PostgresEngine, + db: PgPool, + admin: PgPool, + db_name: String, +} + +impl Scratch { + async fn cleanup(self) { + let Scratch { + engine, + db, + admin, + db_name, + } = self; + drop(engine); + db.close().await; + sqlx::query(&format!( + "DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)" + )) + .execute(&admin) + .await + .expect("drop the scratch database"); + admin.close().await; + } +} + +fn base_conn() -> Option { + let conn = std::env::var("EXTENDDB_TEST_PG_CONNECTION_STRING").ok()?; + (!conn.trim().is_empty()).then(|| conn.trim_end_matches('/').to_owned()) +} + +fn shipped_migrations() -> Vec { + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); + let mut out = Vec::new(); + for dir in ["migrations", "data_migrations"] { + let mut files: Vec<_> = std::fs::read_dir(root.join(dir)) + .expect("read the migrations directory") + .map(|e| e.expect("directory entry").path()) + .filter(|p| p.extension().is_some_and(|e| e == "sql")) + .collect(); + files.sort(); + for f in files { + out.push(std::fs::read_to_string(&f).expect("read a migration file")); + } + } + out +} + +fn skip(test: &str) { + eprintln!( + "SKIP {test}: EXTENDDB_TEST_PG_CONNECTION_STRING is not set, so there is no PostgreSQL \ + to build a scratch catalog in." + ); +} + +async fn scratch() -> Scratch { + let base = base_conn().expect("caller checks base_conn() first"); + let db_name = format!("eddb_gsib_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to the postgres maintenance database"); + sqlx::query(&format!("CREATE DATABASE \"{db_name}\"")) + .execute(&admin) + .await + .expect("create the scratch database"); + let url = format!("{base}/{db_name}"); + let db = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .expect("connect to the scratch database"); + // One scratch database stands in for both the catalog and the data database. + for sql in shipped_migrations() { + sqlx::raw_sql(&sql) + .execute(&db) + .await + .expect("apply a shipped migration"); + } + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'control_plane_delay_seconds'") + .execute(&db) + .await + .expect("pin the control-plane delay to zero"); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind(ACCOUNT) + .bind(format!("acct-{db_name}")) + .execute(&db) + .await + .expect("seed the account row"); + let engine = PostgresEngine::new( + &PostgresConfig { + connection_string: url, + pool_size: 10, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("open a PostgresEngine on the scratch database") + .start_gsi_workers(); + Scratch { + engine, + db, + admin, + db_name, + } +} + +async fn create_table(s: &Scratch) -> TableKeyInfo { + s.engine + .create_table( + ACCOUNT, + CreateTableInput { + table_name: TABLE.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, + }, + ], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + }, + ) + .await + .expect("create the table"); + s.engine + .table_key_info(ACCOUNT, TABLE) + .await + .expect("read the key info") +} + +/// Add GSI `gidx` keyed on `g` (S) and, when `range`, `r` (N). +async fn add_gsi(s: &Scratch, range: bool) { + let mut defs = vec![serde_json::json!({"AttributeName": "g", "AttributeType": "S"})]; + let mut schema = vec![serde_json::json!({"AttributeName": "g", "KeyType": "HASH"})]; + if range { + defs.push(serde_json::json!({"AttributeName": "r", "AttributeType": "N"})); + schema.push(serde_json::json!({"AttributeName": "r", "KeyType": "RANGE"})); + } + let input: UpdateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": TABLE, + "AttributeDefinitions": defs, + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "gidx", + "KeySchema": schema, + "Projection": {"ProjectionType": "ALL"}, + }}], + })) + .expect("build the UpdateTable input"); + s.engine + .update_table(ACCOUNT, input) + .await + .expect("add the GSI"); +} + +fn item(n: u32, g: Option<&str>, r: Option) -> Item { + let mut it = BTreeMap::from([ + ("pk".to_owned(), AttributeValue::S(format!("p{}", n % 5))), + ("sk".to_owned(), AttributeValue::N(n.to_string())), + ]); + if let Some(g) = g { + it.insert("g".to_owned(), AttributeValue::S(g.to_owned())); + } + if let Some(r) = r { + it.insert("r".to_owned(), AttributeValue::N(r.to_string())); + } + it +} + +async fn put(s: &Scratch, key_info: &TableKeyInfo, it: Item) { + s.engine + .put_item(key_info, it, false, None, &ExpressionMaps::default(), None) + .await + .expect("put"); +} + +async fn ids(db: &PgPool) -> (String, String, String) { + sqlx::query_as( + "SELECT t.table_id, i.index_id, i.index_status FROM tables t \ + JOIN indexes i ON i.table_id = t.table_id \ + WHERE t.table_name = $1 AND i.index_name = 'gidx'", + ) + .bind(TABLE) + .fetch_one(db) + .await + .expect("look up the ids") +} + +/// `(base_pk, base_sk, g, item_data)` for every row of the index: the key +/// columns and the projected item, so a row with the right key but a stale +/// projection shows up as a difference. +async fn index_rows(db: &PgPool, index_id: &str) -> BTreeSet<(String, String, String, String)> { + let rows: Vec<(String, String, String, String)> = sqlx::query_as(&format!( + "SELECT base_pk, base_sk_n::text, pk, item_data::text FROM \"_ddb_{index_id}\"" + )) + .fetch_all(db) + .await + .expect("read the index table"); + rows.into_iter().collect() +} + +/// `(pk, sk, g, item_data)` for every base item carrying `g`; the GSIs in this +/// file project ALL, so the index row's item is the base item. +async fn expected_rows(db: &PgPool, table_id: &str) -> BTreeSet<(String, String, String, String)> { + let rows: Vec<(String, String, String, String)> = sqlx::query_as(&format!( + "SELECT pk, sk_n::text, item_data->'g'->>'S', item_data::text FROM \"_ddb_{table_id}\" \ + WHERE item_data ? 'g'" + )) + .fetch_all(db) + .await + .expect("read the base table"); + rows.into_iter().collect() +} + +/// Wait for the GSI to become ACTIVE. UpdateTable publishes it before returning +/// unless the build hit a transient failure, in which case the recovery worker +/// (not running in these tests) would; a short wait tells the two apart. +async fn wait_active(db: &PgPool) -> String { + let mut status = String::new(); + for _ in 0..100 { + status = ids(db).await.2; + if status == "ACTIVE" { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + status +} + +async fn data_table_id(db: &PgPool) -> String { + sqlx::query_scalar("SELECT table_id FROM tables WHERE table_name = $1") + .bind(TABLE) + .fetch_one(db) + .await + .expect("look up the table id") +} + +async fn wait_queue_empty(db: &PgPool) { + for _ in 0..500 { + let n: i64 = sqlx::query_scalar("SELECT count(*) FROM gsi_pending") + .fetch_one(db) + .await + .expect("count pending rows"); + if n == 0 { + return; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("the propagation queue did not drain"); +} + +/// Two queued updates for one item, applied newest first. The index must end at +/// the item's state, not at the last row applied. Replaying the older row's +/// image left that image in the index: a second row for the item, or the old +/// value in place of the new one. +#[tokio::test] +async fn queue_rows_applied_out_of_order_converge_to_the_item() { + let test = "queue_rows_applied_out_of_order_converge_to_the_item"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + add_gsi(&s, false).await; + let key_info = s + .engine + .table_key_info(ACCOUNT, TABLE) + .await + .unwrap_or(key_info); + let (table_id, index_id, _) = ids(&s.db).await; + + // Hold both rows back so their order can be chosen. + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', '600000') \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .execute(&s.db) + .await + .expect("set a long propagation delay"); + put(&s, &key_info, item(1, Some("old"), None)).await; + put(&s, &key_info, item(1, Some("new"), None)).await; + let pending: Vec = sqlx::query_scalar("SELECT id FROM gsi_pending ORDER BY id") + .fetch_all(&s.db) + .await + .expect("read pending ids"); + assert_eq!(pending.len(), 2, "both writes queued"); + + // The newer row first, as a peer process's worker that skipped the locked + // older row would apply it. + sqlx::query("UPDATE gsi_pending SET ready_at = NOW() WHERE id = $1") + .bind(pending[1]) + .execute(&s.db) + .await + .expect("release the newer row"); + let mut left = 2; + for _ in 0..500 { + left = sqlx::query_scalar::<_, i64>("SELECT count(*) FROM gsi_pending") + .fetch_one(&s.db) + .await + .expect("count"); + if left == 1 { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert_eq!(left, 1, "the newer row applied on its own first"); + sqlx::query("UPDATE gsi_pending SET ready_at = NOW()") + .execute(&s.db) + .await + .expect("release the older row"); + wait_queue_empty(&s.db).await; + + assert_eq!( + index_rows(&s.db, &index_id).await, + expected_rows(&s.db, &table_id).await, + "the index holds the item's current state only" + ); + s.cleanup().await; +} + +/// A queued update that lands after the backfill already wrote the same index +/// row. Applying the row's image with a plain INSERT failed with a duplicate key, +/// which rolled the queue row back to be retried forever and stalled every row +/// behind it in its partition. +#[tokio::test] +async fn queue_row_for_a_row_the_backfill_wrote_is_applied() { + let test = "queue_row_for_a_row_the_backfill_wrote_is_applied"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', '600000') \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .execute(&s.db) + .await + .expect("set a long propagation delay"); + add_gsi(&s, false).await; + let key_info = s + .engine + .table_key_info(ACCOUNT, TABLE) + .await + .unwrap_or(key_info); + let (table_id, index_id, _) = ids(&s.db).await; + + // Queue an update, then write its row into the index directly, as a backfill + // reading the item after the write would. + put(&s, &key_info, item(2, Some("v"), None)).await; + sqlx::query(&format!( + "INSERT INTO \"_ddb_{index_id}\" (pk, base_pk, base_sk_n, item_data) \ + SELECT item_data->'g'->>'S', pk, sk_n, item_data FROM \"_ddb_{table_id}\"" + )) + .execute(&s.db) + .await + .expect("write the row the backfill would have"); + sqlx::query("UPDATE gsi_pending SET ready_at = NOW()") + .execute(&s.db) + .await + .expect("release the row"); + wait_queue_empty(&s.db).await; + assert_eq!( + index_rows(&s.db, &index_id).await, + expected_rows(&s.db, &table_id).await + ); + s.cleanup().await; +} + +/// A synchronous GSI write that finds index rows for its item already there, +/// with no old item to say so. Two ways to get there: the backfill wrote the row +/// while the put waited for the item's lock (same index key; a plain INSERT +/// failed the put with a duplicate key), and a write that predated the index +/// deleted the item without removing its row (a different index key; the put +/// inserted next to it and the item was indexed twice). +#[tokio::test] +async fn sync_write_over_an_existing_index_row_succeeds() { + let test = "sync_write_over_an_existing_index_row_succeeds"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', '0') \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .execute(&s.db) + .await + .expect("make index propagation synchronous"); + add_gsi(&s, false).await; + let (table_id, index_id, _) = ids(&s.db).await; + sqlx::query(&format!( + "INSERT INTO \"_ddb_{index_id}\" (pk, base_pk, base_sk_n, item_data) \ + VALUES ('v', 'p3', 3, '{{}}'::jsonb), ('gone', 'p4', 4, '{{}}'::jsonb)" + )) + .execute(&s.db) + .await + .expect("write the rows a backfill and a stale delete would have left"); + put(&s, &key_info, item(3, Some("v"), None)).await; + put(&s, &key_info, item(4, Some("new"), None)).await; + assert_eq!( + index_rows(&s.db, &index_id).await, + expected_rows(&s.db, &table_id).await + ); + s.cleanup().await; +} + +/// Puts, deletes and conditional puts from several writers while the GSI is +/// built, with propagation synchronous and then asynchronous. +async fn concurrent_build(delay_ms: u32) { + let s = scratch().await; + let key_info = create_table(&s).await; + // Several backfill pages' worth, so pages run while writers hold their rows. + for n in 0..1200 { + put( + &s, + &key_info, + item(n, Some(&format!("g{}", n % 4)), Some(n)), + ) + .await; + } + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', $1) \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .bind(delay_ms.to_string()) + .execute(&s.db) + .await + .expect("set the propagation delay"); + + let stop = std::sync::atomic::AtomicBool::new(false); + let maps = ExpressionMaps::default(); + let exists = extenddb_core::expression::parse_condition( + &extenddb_core::expression::tokenize("attribute_exists(pk)").expect("tokenize"), + ) + .expect("parse"); + let writer = |seed: u32| { + let (s, key_info, stop, maps, exists) = (&s, &key_info, &stop, &maps, &exists); + async move { + let mut i = 0u32; + while !stop.load(std::sync::atomic::Ordering::Relaxed) { + i += 1; + let n = (seed * 7919 + i * 31) % 1300; + let result = match i % 4 { + 0 => s + .engine + .delete_item(key_info, &item(n, None, None), false, None, maps, None) + .await + .map(|_| ()), + 1 => s + .engine + .put_item(key_info, item(n, None, Some(i)), false, None, maps, None) + .await + .map(|_| ()), + 2 => match s + .engine + .put_item( + key_info, + item(n, Some(&format!("c{}", i % 3)), Some(i)), + false, + Some(exists), + maps, + None, + ) + .await + { + Err(extenddb_storage::error::StorageError::ConditionFailed(_)) => Ok(()), + other => other.map(|_| ()), + }, + _ => s + .engine + .put_item( + key_info, + item(n, Some(&format!("w{}", i % 5)), Some(i)), + false, + None, + maps, + None, + ) + .await + .map(|_| ()), + }; + result.expect("a write during the build"); + } + } + }; + let build = async { + tokio::time::sleep(Duration::from_millis(100)).await; + add_gsi(&s, false).await; + tokio::time::sleep(Duration::from_millis(300)).await; + stop.store(true, std::sync::atomic::Ordering::Relaxed); + }; + tokio::join!(writer(1), writer(2), writer(3), writer(4), build); + + let (table_id, index_id, _) = ids(&s.db).await; + assert_eq!(wait_active(&s.db).await, "ACTIVE"); + wait_queue_empty(&s.db).await; + assert_eq!( + index_rows(&s.db, &index_id).await, + expected_rows(&s.db, &table_id).await, + "delay {delay_ms} ms" + ); + s.cleanup().await; +} + +#[tokio::test] +async fn concurrent_writes_during_a_build_leave_the_index_equal_to_the_table() { + let test = "concurrent_writes_during_a_build_leave_the_index_equal_to_the_table"; + if base_conn().is_none() { + return skip(test); + } + concurrent_build(0).await; + concurrent_build(10).await; +} + +/// A write carrying the key info cached before UpdateTable added a GSI with a +/// sort key on a new attribute. The engine drops that cache only when the +/// UpdateTable returns, so writes during the build carry definitions that do not +/// know the attribute; the index row was written without its sort key column. +#[tokio::test] +async fn write_with_key_info_older_than_a_new_gsi_indexes_its_sort_key() { + let test = "write_with_key_info_older_than_a_new_gsi_indexes_its_sort_key"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let stale = create_table(&s).await; + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', '0') \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .execute(&s.db) + .await + .expect("make index propagation synchronous"); + add_gsi(&s, true).await; + let (_, index_id, _) = ids(&s.db).await; + put(&s, &stale, item(9, Some("x"), Some(7))).await; + let (rows, missing): (i64, i64) = sqlx::query_as(&format!( + "SELECT COUNT(*), COUNT(*) FILTER (WHERE sk_n IS NULL) FROM \"_ddb_{index_id}\"" + )) + .fetch_one(&s.db) + .await + .expect("count index rows"); + assert_eq!((rows, missing), (1, 0), "one index row, with its sort key"); + s.cleanup().await; +} + +/// A synchronous write to a table with two synchronous GSIs, one of whose data +/// tables is gone (a GSI delete, or a failed build's removal, that the write's +/// catalog read predates). The write succeeds and the surviving index is right. +/// The two GSIs share one savepoint, so the first one's rows are rolled back +/// with the second's failure and have to be written again. +#[tokio::test] +async fn sync_write_survives_a_gsi_whose_table_is_gone() { + let test = "sync_write_survives_a_gsi_whose_table_is_gone"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', '0') \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .execute(&s.db) + .await + .expect("make index propagation synchronous"); + add_gsi(&s, false).await; + let (table_id, kept_id, _) = ids(&s.db).await; + // A second GSI on the same attribute, whose data table then disappears + // without its catalog row: what a write that read the catalog just before a + // GSI delete committed sees. + let input: UpdateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": TABLE, + "AttributeDefinitions": [{"AttributeName": "g", "AttributeType": "S"}], + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "gone", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"}, + }}], + })) + .expect("build the UpdateTable input"); + s.engine + .update_table(ACCOUNT, input) + .await + .expect("add the second GSI"); + let gone_id: String = + sqlx::query_scalar("SELECT index_id FROM indexes WHERE index_name = 'gone'") + .fetch_one(&s.db) + .await + .expect("second index id"); + sqlx::query(&format!("DROP TABLE \"_ddb_{gone_id}\"")) + .execute(&s.db) + .await + .expect("drop the second index's table behind the catalog's back"); + + // The catalog returns the two GSIs in no fixed order. Whichever runs first, + // the final rows must be right; when the kept one runs first its rows are + // rolled back with the other's failure and written again. + put(&s, &key_info, item(1, Some("a"), None)).await; + put(&s, &key_info, item(1, Some("b"), None)).await; + put(&s, &key_info, item(2, Some("c"), None)).await; + assert_eq!( + index_rows(&s.db, &kept_id).await, + expected_rows(&s.db, &table_id).await, + "the surviving index holds the items' current state" + ); + s.cleanup().await; +} + +/// While a GSI is CREATING: a second Create, and a Delete of another GSI, are +/// refused with the service's LimitExceeded sentence; a Delete of the building +/// GSI is not refused that way; and a Query or Scan naming the building index is +/// told the table has no such index (which is what the engine puts on the wire +/// for `IndexNotFound`), and resolves again once the index is ACTIVE. +#[tokio::test] +async fn a_creating_gsi_refuses_a_second_change_and_reads() { + let test = "a_creating_gsi_refuses_a_second_change_and_reads"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + create_table(&s).await; + add_gsi(&s, false).await; + let (table_id, index_id, _) = ids(&s.db).await; + // What a concurrent DescribeTable sees mid-build; the build has finished, + // so the row is put back the way it looked. + sqlx::query("UPDATE indexes SET index_status = 'CREATING' WHERE index_id = $1") + .bind(&index_id) + .execute(&s.db) + .await + .expect("mark CREATING"); + + let second: UpdateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": TABLE, + "AttributeDefinitions": [{"AttributeName": "g", "AttributeType": "S"}], + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "second", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"}, + }}], + })) + .expect("build the UpdateTable input"); + let refused = s.engine.update_table(ACCOUNT, second).await; + let Err(extenddb_storage::error::StorageError::LimitExceeded(msg)) = refused else { + panic!("a second create during a build was not refused with LimitExceeded"); + }; + assert_eq!(msg, extenddb_core::types::ONLINE_INDEX_LIMIT); + assert!( + matches!( + s.engine.index_info_by_table_id(&table_id, "gidx").await, + Err(extenddb_storage::error::StorageError::IndexNotFound(_)) + ), + "a read of a CREATING index resolves as no such index" + ); + sqlx::query("UPDATE indexes SET index_status = 'ACTIVE' WHERE index_id = $1") + .bind(&index_id) + .execute(&s.db) + .await + .expect("mark ACTIVE"); + s.engine + .index_info_by_table_id(&table_id, "gidx") + .await + .expect("an ACTIVE index resolves"); + + // A Delete of another GSI while one builds is refused the same way; a + // Delete of the building index itself is not (the service accepts it once + // the backfill is under way). + let other: UpdateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": TABLE, + "AttributeDefinitions": [{"AttributeName": "g", "AttributeType": "S"}], + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "other", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"}, + }}], + })) + .expect("build the UpdateTable input"); + s.engine + .update_table(ACCOUNT, other) + .await + .expect("add a second GSI while none builds"); + sqlx::query("UPDATE indexes SET index_status = 'CREATING' WHERE index_id = $1") + .bind(&index_id) + .execute(&s.db) + .await + .expect("mark CREATING again"); + let delete_other: UpdateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": TABLE, + "GlobalSecondaryIndexUpdates": [{"Delete": {"IndexName": "other"}}], + })) + .expect("build the UpdateTable input"); + let refused = s.engine.update_table(ACCOUNT, delete_other).await; + let Err(extenddb_storage::error::StorageError::LimitExceeded(msg)) = refused else { + panic!("deleting another GSI during a build was not refused with LimitExceeded"); + }; + assert_eq!(msg, extenddb_core::types::ONLINE_INDEX_LIMIT); + let other_left: i64 = + sqlx::query_scalar("SELECT count(*) FROM indexes WHERE index_name = 'other'") + .fetch_one(&s.db) + .await + .expect("count the other index"); + assert_eq!( + other_left, 1, + "the refused Delete left the other index alone" + ); + let delete_own: UpdateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": TABLE, + "GlobalSecondaryIndexUpdates": [{"Delete": {"IndexName": "gidx"}}], + })) + .expect("build the UpdateTable input"); + let own = s.engine.update_table(ACCOUNT, delete_own).await; + assert!( + !matches!( + own, + Err(extenddb_storage::error::StorageError::LimitExceeded(_)) + ), + "deleting the index that is building is not refused as a second change" + ); + s.cleanup().await; +} + +/// DeleteTable while the propagation queue holds a claimed row for the table. +/// The worker holds its claimed queue row and takes the base row `FOR SHARE`; a +/// drop that took the tables first and the queue rows second deadlocked with it +/// and, as the victim, left the data tables behind once the catalog rows were +/// gone. The drop now stops the queue and waits for the claim before it touches a +/// table. The claim is held by hand here, so the wait is the thing tested rather +/// than a race: the delete must block on it, then finish once it is released. +#[tokio::test] +async fn delete_table_waits_for_a_claimed_queue_row_then_drops_everything() { + let test = "delete_table_waits_for_a_claimed_queue_row_then_drops_everything"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', '600000') \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .execute(&s.db) + .await + .expect("hold queue rows back so one can be claimed by hand"); + add_gsi(&s, false).await; + let (table_id, index_id, _) = ids(&s.db).await; + put(&s, &key_info, item(1, Some("g"), None)).await; + + // A worker mid-apply: the claim lock on the pending row, as `FOR UPDATE SKIP + // LOCKED` leaves it, and the base row share lock it takes next. + let mut claim = s.db.begin().await.expect("begin the fake claim"); + let claimed: i64 = sqlx::query_scalar( + "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("claim the pending row"); + sqlx::query(&format!( + "SELECT 1 FROM \"_ddb_{table_id}\" WHERE pk = 'p1' FOR SHARE" + )) + .execute(&mut *claim) + .await + .expect("share-lock the base row"); + + let input: extenddb_core::types::DeleteTableInput = + serde_json::from_value(serde_json::json!({"TableName": TABLE})).expect("input"); + let engine = &s.engine; + let mut delete = Box::pin(engine.delete_table(ACCOUNT, input)); + let early = tokio::time::timeout(Duration::from_millis(1_500), &mut delete).await; + assert!( + 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"); + let base_exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(format!("\"_ddb_{table_id}\"")) + .fetch_one(&s.db) + .await + .expect("check the base table"); + assert!(base_exists, "nothing is dropped while a claim is in flight"); + + claim.commit().await.expect("release the claim"); + tokio::time::timeout(Duration::from_secs(30), &mut delete) + .await + .expect("DeleteTable finished once the claim was released") + .expect("DeleteTable"); + drop(delete); + let _ = claimed; + for name in [ + format!("\"_ddb_{table_id}\""), + format!("\"_ddb_{index_id}\""), + ] { + let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(&name) + .fetch_one(&s.db) + .await + .expect("check the table"); + assert!(!exists, "{name} was left behind"); + } + let (holds, pending): (i64, i64) = sqlx::query_as( + "SELECT (SELECT count(*) FROM vector_index_holds WHERE table_id = $1), \ + (SELECT count(*) FROM gsi_pending WHERE table_id = $1)", + ) + .bind(&table_id) + .fetch_one(&s.db) + .await + .expect("count leftovers"); + assert_eq!( + (holds, pending), + (0, 0), + "no hold and no queue row left behind" + ); + s.cleanup().await; +} + +/// UpdateItem that sets a GSI key attribute to the wrong type is refused, as the +/// service, PutItem, and TransactWriteItems already refuse it. +#[tokio::test] +async fn update_item_rejects_a_wrong_typed_gsi_key() { + let test = "update_item_rejects_a_wrong_typed_gsi_key"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + add_gsi(&s, false).await; + put(&s, &key_info, item(1, Some("a"), None)).await; + let actions = extenddb_core::expression::parse_update( + &extenddb_core::expression::tokenize("SET g = :v").expect("tokenize"), + ) + .expect("parse"); + let maps = ExpressionMaps::new( + std::collections::HashMap::new(), + std::collections::HashMap::from([("v".to_owned(), AttributeValue::N("9".to_owned()))]), + ); + let result = s + .engine + .update_item( + &key_info, + &item(1, None, None), + &actions, + false, + false, + None, + &maps, + None, + ) + .await; + let Err(extenddb_storage::error::StorageError::Validation(msg)) = result else { + panic!("a wrong-typed index key on UpdateItem was not a ValidationException"); + }; + assert!(msg.contains("Type mismatch for Index Key g"), "{msg}"); + s.cleanup().await; +} + +/// A base row held open by a client does not hold up the rest of its queue +/// partition. The worker claims the lowest due id; with the held item's rows +/// first in line, every row behind them used to wait for the whole hold, since +/// the timed-out row was claimed again at once. The held item has three rows +/// queued, as a table with several asynchronous indexes or a few writes inside +/// the propagation delay produces: passing over a single row at a time would +/// not move the partition either, since by the time the second row had timed out +/// the first would be due again. Every item here has the same partition key, so +/// they share a partition whatever the hash. +#[tokio::test] +async fn a_held_base_row_does_not_hold_up_its_queue_partition() { + let test = "a_held_base_row_does_not_hold_up_its_queue_partition"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + add_gsi(&s, false).await; + let (table_id, index_id, _) = ids(&s.db).await; + put(&s, &key_info, item(0, Some("before"), None)).await; + let db = &s.db; + let set_delay = |ms: &'static str| async move { + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', $1) \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .bind(ms) + .execute(db) + .await + .expect("set the propagation delay"); + }; + set_delay("500").await; + wait_queue_empty(&s.db).await; + + // The held item's updates go in first, due no sooner than 750 ms from now so + // the worker cannot reach them before the row is held; then the base row is + // held open. + set_delay("1500").await; + for g in ["held-1", "held-2", "held"] { + put(&s, &key_info, item(0, Some(g), None)).await; + } + set_delay("500").await; + let queued: i64 = sqlx::query_scalar("SELECT count(*) FROM gsi_pending") + .fetch_one(&s.db) + .await + .expect("count the held item's rows"); + assert_eq!(queued, 3, "three rows queued for the held item"); + let mut holder = s.db.begin().await.expect("begin the holder"); + sqlx::query(&format!( + "SELECT 1 FROM \"_ddb_{table_id}\" WHERE pk = 'p0' AND sk_n = 0 FOR UPDATE" + )) + .execute(&mut *holder) + .await + .expect("hold the base row"); + // Rows behind it in the same partition (pk p0, so the same partition). + let behind: Vec = (1..=10).map(|i| i * 5).collect(); + for n in &behind { + put(&s, &key_info, item(*n, Some("behind"), None)).await; + } + + // The hold lasts 10 s; the rows behind it must be in the index well before. + let start = std::time::Instant::now(); + let mut done = false; + while start.elapsed() < Duration::from_secs(6) { + let n: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM \"_ddb_{index_id}\" WHERE pk = 'behind'" + )) + .fetch_one(&s.db) + .await + .expect("count indexed rows"); + if usize::try_from(n).unwrap_or(0) == behind.len() { + done = true; + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + let reached = start.elapsed(); + assert!( + done, + "the rows behind the held one did not reach the index while it was held" + ); + let held_row: String = sqlx::query_scalar(&format!( + "SELECT pk FROM \"_ddb_{index_id}\" WHERE base_sk_n = 0" + )) + .fetch_one(&s.db) + .await + .expect("read the held item's index row"); + assert_eq!(held_row, "before", "the held item has not moved yet"); + + tokio::time::sleep(Duration::from_secs(10).saturating_sub(reached)).await; + holder.commit().await.expect("release the hold"); + wait_queue_empty(&s.db).await; + assert_eq!( + index_rows(&s.db, &index_id).await, + expected_rows(&s.db, &table_id).await, + "once released, the held item's update applies too" + ); + s.cleanup().await; +} + +/// A request that reaches the data database after a DeleteTable dropped the +/// table is told the table is not there, as the service tells it, rather than +/// getting the 500 that PostgreSQL's undefined_table used to become: PutItem, +/// GetItem, TransactWriteItems and TransactGetItems. +#[tokio::test] +async fn a_write_to_a_dropped_data_table_is_table_not_found() { + let test = "a_write_to_a_dropped_data_table_is_table_not_found"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + let table_id = data_table_id(&s.db).await; + sqlx::query(&format!("DROP TABLE \"_ddb_{table_id}\"")) + .execute(&s.db) + .await + .expect("drop the data table behind the catalog's back"); + let put = s + .engine + .put_item( + &key_info, + item(1, Some("x"), None), + false, + None, + &ExpressionMaps::default(), + None, + ) + .await; + assert!( + matches!(put, Err(extenddb_storage::error::StorageError::TableNotFound(ref n)) if n == TABLE), + "a put to a dropped data table: {}", + put.as_ref() + .map(|_| "ok") + .unwrap_or("an error other than TableNotFound") + ); + let get = s.engine.get_item(&key_info, &item(1, None, None)).await; + assert!( + matches!( + get, + Err(extenddb_storage::error::StorageError::TableNotFound(_)) + ), + "a get from a dropped data table is TableNotFound too" + ); + let maps = ExpressionMaps::default(); + let new_item = item(2, Some("x"), None); + let ops = [extenddb_storage::TransactWriteOp::Put { + key_info: &key_info, + item: &new_item, + condition: None, + maps: &maps, + return_values_on_ccf: extenddb_core::types::ReturnValuesOnConditionCheckFailure::None, + stream: None, + }]; + let twi = s.engine.transact_write_items(&ops, None).await; + assert!( + matches!(twi, Err(extenddb_storage::error::StorageError::TableNotFound(ref n)) if n == TABLE), + "TransactWriteItems against a dropped data table is TableNotFound" + ); + let key = item(1, None, None); + let gets = [extenddb_storage::TransactGetOp { + key_info: &key_info, + key: &key, + }]; + let tgi = s.engine.transact_get_items(&gets).await; + assert!( + matches!( + tgi, + Err(extenddb_storage::error::StorageError::TableNotFound(_)) + ), + "TransactGetItems against a dropped data table is TableNotFound" + ); + s.cleanup().await; +} + +/// A GSI a crash left `CREATING`, with part of its rows written and some of them +/// stale, is finished by the recovery sweep: complete, current, and `ACTIVE`. +#[tokio::test] +async fn recovery_finishes_a_gsi_left_creating() { + let test = "recovery_finishes_a_gsi_left_creating"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let key_info = create_table(&s).await; + for n in 0..40 { + put( + &s, + &key_info, + item(n, Some(&format!("g{}", n % 3)), Some(n)), + ) + .await; + } + add_gsi(&s, true).await; + let (table_id, index_id, status) = ids(&s.db).await; + assert_eq!(status, "ACTIVE"); + + // What a build that died partway leaves: CREATING, some rows missing, one + // stale, and one for an item that no longer exists. + sqlx::query("UPDATE indexes SET index_status = 'CREATING' WHERE index_id = $1") + .bind(&index_id) + .execute(&s.db) + .await + .expect("mark CREATING"); + let idx = format!("\"_ddb_{index_id}\""); + let base = format!("\"_ddb_{table_id}\""); + sqlx::query(&format!("DELETE FROM {idx} WHERE base_sk_n >= 20")) + .execute(&s.db) + .await + .expect("drop half the rows"); + sqlx::query(&format!( + "UPDATE {idx} SET pk = 'stale', item_data = '{{}}'::jsonb WHERE base_sk_n = 3" + )) + .execute(&s.db) + .await + .expect("make one row stale"); + sqlx::query(&format!("DELETE FROM {base} WHERE sk_n = 5")) + .execute(&s.db) + .await + .expect("remove one base item behind the index's back"); + + // The build that ran inside UpdateTable held the ownership lock on a + // connection of its own, which the server releases when it notices that + // connection closed. Until then the sweep skips the index, as it would for a + // build still running, so give it a few passes. + let mut finished = 0; + for _ in 0..50 { + finished = extenddb_storage_postgres::reconcile_creating_gsis(&s.engine) + .await + .expect("run the recovery sweep"); + if finished > 0 { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert_eq!(finished, 1); + let (_, _, status) = ids(&s.db).await; + assert_eq!(status, "ACTIVE"); + assert_eq!( + index_rows(&s.db, &index_id).await, + expected_rows(&s.db, &table_id).await + ); + let range_ok: i64 = + sqlx::query_scalar(&format!("SELECT count(*) FROM {idx} WHERE sk_n IS NULL")) + .fetch_one(&s.db) + .await + .expect("check the sort key column"); + assert_eq!(range_ok, 0, "every row carries the index sort key"); + s.cleanup().await; +} diff --git a/crates/storage-postgres/tests/gsi_stale_writer.rs b/crates/storage-postgres/tests/gsi_stale_writer.rs new file mode 100644 index 000000000..c79fe8cab --- /dev/null +++ b/crates/storage-postgres/tests/gsi_stale_writer.rs @@ -0,0 +1,353 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 +//! A write that read the catalog before a GSI existed and wrote its item after +//! the GSI's backfill had passed it. +//! +//! The catalog and the data tables are separate databases, so nothing in the +//! write's transaction can tell it about the new index. The write notices that it +//! was slow and queues a repair: inside its transaction when it has one, after +//! the statement when it wrote without one. This file uses a test hook that makes +//! every write pause at a chosen point, which is process-wide, so it is a test +//! binary of its own. +//! +//! Each test builds its own throwaway database, applies the shipped migrations +//! to it, and drops it when it passes. A failing test leaves its database behind +//! on purpose, named `eddb_gsis_*`, so the state that failed can be inspected. +//! +//! Requires `EXTENDDB_TEST_PG_CONNECTION_STRING`, a base URL with no database +//! component (for example `postgresql://postgres@127.0.0.1:5432`), pointing at a +//! server whose role may create and drop databases. Without it every test here +//! reports a skip and passes, the same convention the wire suites use. + +use std::collections::{BTreeMap, BTreeSet}; +use std::time::Duration; + +use extenddb_core::expression::ExpressionMaps; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, BillingMode, CreateTableInput, Item, KeySchemaElement, + KeyType, ScalarAttributeType, TableKeyInfo, UpdateTableInput, +}; +use extenddb_storage::{DataEngine, TableEngine}; +use extenddb_storage_postgres::{ + PauseWrites, PostgresConfig, PostgresEngine, pause_writes_for_tests, +}; +use sqlx::PgPool; +use sqlx::postgres::PgPoolOptions; + +const ACCOUNT: &str = "123456789012"; +const REGION: &str = "us-east-1"; +const TABLE: &str = "t_gsi_online_build"; + +struct Scratch { + engine: PostgresEngine, + db: PgPool, + admin: PgPool, + db_name: String, +} + +impl Scratch { + async fn cleanup(self) { + let Scratch { + engine, + db, + admin, + db_name, + } = self; + drop(engine); + db.close().await; + sqlx::query(&format!( + "DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)" + )) + .execute(&admin) + .await + .expect("drop the scratch database"); + admin.close().await; + } +} + +fn base_conn() -> Option { + let conn = std::env::var("EXTENDDB_TEST_PG_CONNECTION_STRING").ok()?; + (!conn.trim().is_empty()).then(|| conn.trim_end_matches('/').to_owned()) +} + +fn shipped_migrations() -> Vec { + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); + let mut out = Vec::new(); + for dir in ["migrations", "data_migrations"] { + let mut files: Vec<_> = std::fs::read_dir(root.join(dir)) + .expect("read the migrations directory") + .map(|e| e.expect("directory entry").path()) + .filter(|p| p.extension().is_some_and(|e| e == "sql")) + .collect(); + files.sort(); + for f in files { + out.push(std::fs::read_to_string(&f).expect("read a migration file")); + } + } + out +} + +fn skip(test: &str) { + eprintln!( + "SKIP {test}: EXTENDDB_TEST_PG_CONNECTION_STRING is not set, so there is no PostgreSQL \ + to build a scratch catalog in." + ); +} + +async fn scratch() -> Scratch { + let base = base_conn().expect("caller checks base_conn() first"); + let db_name = format!("eddb_gsis_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to the postgres maintenance database"); + sqlx::query(&format!("CREATE DATABASE \"{db_name}\"")) + .execute(&admin) + .await + .expect("create the scratch database"); + let url = format!("{base}/{db_name}"); + let db = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .expect("connect to the scratch database"); + // One scratch database stands in for both the catalog and the data database. + for sql in shipped_migrations() { + sqlx::raw_sql(&sql) + .execute(&db) + .await + .expect("apply a shipped migration"); + } + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'control_plane_delay_seconds'") + .execute(&db) + .await + .expect("pin the control-plane delay to zero"); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind(ACCOUNT) + .bind(format!("acct-{db_name}")) + .execute(&db) + .await + .expect("seed the account row"); + let engine = PostgresEngine::new( + &PostgresConfig { + connection_string: url, + pool_size: 10, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("open a PostgresEngine on the scratch database") + .start_gsi_workers(); + Scratch { + engine, + db, + admin, + db_name, + } +} + +async fn create_table(s: &Scratch) -> TableKeyInfo { + s.engine + .create_table( + ACCOUNT, + CreateTableInput { + table_name: TABLE.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, + }, + ], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + }, + ) + .await + .expect("create the table"); + s.engine + .table_key_info(ACCOUNT, TABLE) + .await + .expect("read the key info") +} + +/// Add GSI `gidx` keyed on `g` (S) and, when `range`, `r` (N). +async fn add_gsi(s: &Scratch, range: bool) { + let mut defs = vec![serde_json::json!({"AttributeName": "g", "AttributeType": "S"})]; + let mut schema = vec![serde_json::json!({"AttributeName": "g", "KeyType": "HASH"})]; + if range { + defs.push(serde_json::json!({"AttributeName": "r", "AttributeType": "N"})); + schema.push(serde_json::json!({"AttributeName": "r", "KeyType": "RANGE"})); + } + let input: UpdateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": TABLE, + "AttributeDefinitions": defs, + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "gidx", + "KeySchema": schema, + "Projection": {"ProjectionType": "ALL"}, + }}], + })) + .expect("build the UpdateTable input"); + s.engine + .update_table(ACCOUNT, input) + .await + .expect("add the GSI"); +} + +fn item(n: u32, g: Option<&str>, r: Option) -> Item { + let mut it = BTreeMap::from([ + ("pk".to_owned(), AttributeValue::S(format!("p{}", n % 5))), + ("sk".to_owned(), AttributeValue::N(n.to_string())), + ]); + if let Some(g) = g { + it.insert("g".to_owned(), AttributeValue::S(g.to_owned())); + } + if let Some(r) = r { + it.insert("r".to_owned(), AttributeValue::N(r.to_string())); + } + it +} + +async fn put(s: &Scratch, key_info: &TableKeyInfo, it: Item) { + s.engine + .put_item(key_info, it, false, None, &ExpressionMaps::default(), None) + .await + .expect("put"); +} + +async fn ids(db: &PgPool) -> (String, String, String) { + sqlx::query_as( + "SELECT t.table_id, i.index_id, i.index_status FROM tables t \ + JOIN indexes i ON i.table_id = t.table_id \ + WHERE t.table_name = $1 AND i.index_name = 'gidx'", + ) + .bind(TABLE) + .fetch_one(db) + .await + .expect("look up the ids") +} + +/// `(base_pk, base_sk, g)` for every row of the index. +async fn index_rows(db: &PgPool, index_id: &str) -> BTreeSet<(String, String, String)> { + let rows: Vec<(String, String, String)> = sqlx::query_as(&format!( + "SELECT base_pk, base_sk_n::text, pk FROM \"_ddb_{index_id}\"" + )) + .fetch_all(db) + .await + .expect("read the index table"); + rows.into_iter().collect() +} + +/// `(pk, sk, g)` for every base item carrying `g`. +async fn expected_rows(db: &PgPool, table_id: &str) -> BTreeSet<(String, String, String)> { + let rows: Vec<(String, String, String)> = sqlx::query_as(&format!( + "SELECT pk, sk_n::text, item_data->'g'->>'S' FROM \"_ddb_{table_id}\" \ + WHERE item_data ? 'g'" + )) + .fetch_all(db) + .await + .expect("read the base table"); + rows.into_iter().collect() +} + +/// Three writes that read "no index", then write after the GSI is ACTIVE: a new +/// item, an overwrite that changes the index key, and a delete. Paused at `at`; +/// with `conditional` the last two carry a condition, which puts them on the +/// transaction path whatever the pause. +async fn drive(at: PauseWrites, conditional: bool) { + let s = scratch().await; + let key_info = create_table(&s).await; + for n in 0..10 { + put(&s, &key_info, item(n, Some("before"), None)).await; + } + + pause_writes_for_tests(at, Duration::from_millis(3_000)); + let ki = key_info.clone(); + let maps = ExpressionMaps::default(); + let cond = extenddb_core::expression::parse_condition( + &extenddb_core::expression::tokenize("attribute_exists(pk)").expect("tokenize"), + ) + .expect("parse"); + let writes = async { + let a = s + .engine + .put_item(&ki, item(100, Some("late"), None), false, None, &maps, None); + let b = s.engine.put_item( + &ki, + item(1, Some("changed"), None), + false, + conditional.then_some(&cond), + &maps, + None, + ); + let mut key2 = item(2, None, None); + key2.remove("g"); + let c = s + .engine + .delete_item(&ki, &key2, false, conditional.then_some(&cond), &maps, None); + let (a, b, c) = tokio::join!(a, b, c); + a.expect("late put"); + b.expect("late overwrite"); + c.expect("late delete"); + }; + let build = async { + tokio::time::sleep(Duration::from_millis(200)).await; + add_gsi(&s, false).await; + pause_writes_for_tests(at, Duration::ZERO); + }; + tokio::join!(writes, build); + + let (table_id, index_id, status) = ids(&s.db).await; + assert_eq!(status, "ACTIVE"); + // Give the repair rows time to apply. + for _ in 0..100 { + if index_rows(&s.db, &index_id).await == expected_rows(&s.db, &table_id).await { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert_eq!( + index_rows(&s.db, &index_id).await, + expected_rows(&s.db, &table_id).await, + "the writes that missed the index were repaired ({at:?}, conditional: {conditional})" + ); + s.cleanup().await; +} + +#[tokio::test] +async fn writes_that_missed_a_new_gsi_are_repaired() { + let test = "writes_that_missed_a_new_gsi_are_repaired"; + if base_conn().is_none() { + return skip(test); + } + // One test for the three paths, run in sequence, because the pause is + // process-wide. + // + // Unconditional writes on a table with no index write without a transaction. + // Paused between that decision and the statement, they stay on that path: the + // statement lands after the backfill, and the repair runs after the statement, + // on its own (`fix_up_autocommit`). + drive(PauseWrites::BeforeAutocommitStatement, false).await; + // Paused after the catalog read instead, the same writes are already slow + // when the decision is made and move to the transaction path, where the + // repair commits with the write (`fix_up_stale_write`). + drive(PauseWrites::AfterIndexRead, false).await; + // Conditional writes take a transaction regardless. + drive(PauseWrites::AfterIndexRead, true).await; +} diff --git a/crates/storage-sqlite/src/data/ddl.rs b/crates/storage-sqlite/src/data/ddl.rs index 4052c973e..739e1310e 100644 --- a/crates/storage-sqlite/src/data/ddl.rs +++ b/crates/storage-sqlite/src/data/ddl.rs @@ -524,9 +524,14 @@ impl SqliteEngine { table_id: &str, index_name: &str, ) -> Result { + // A GSI still being built is not readable: its table is partial, or not + // there yet (a Query hit it with a 500 on main). The service answers a + // Query or Scan naming one with "The table does not have the specified + // index", the same as for no index at all, which is what `IndexNotFound` + // becomes on the wire. let row: Option<(String, String, String, String)> = sqlx::query_as( "SELECT index_type, index_id, key_schema, projection \ - FROM indexes WHERE table_id = ? AND index_name = ?", + FROM indexes WHERE table_id = ? AND index_name = ? AND index_status <> 'CREATING'", ) .bind(table_id) .bind(index_name) diff --git a/crates/storage-sqlite/src/data/index.rs b/crates/storage-sqlite/src/data/index.rs index 21e89c944..fa02deeb1 100644 --- a/crates/storage-sqlite/src/data/index.rs +++ b/crates/storage-sqlite/src/data/index.rs @@ -146,6 +146,22 @@ pub(crate) struct IndexMeta { /// Per-GSI propagation delay (ms). `None` = use system default; `Some(0)` = /// synchronous. pub(super) propagation_delay_ms: Option, + /// The table's attribute definitions, read with the index rows rather than + /// taken from the cached `TableKeyInfo`. An UpdateTable that adds a GSI adds + /// its key attributes' definitions too, and the engine drops its cached key + /// info only when the UpdateTable returns, so a write in between would build + /// the index row without the columns of a sort key it does not know the type + /// of. + pub(super) attribute_definitions: Vec, +} + +/// The table's attribute definitions as of the write's index read, falling back +/// to the cached set when the table has no secondary index to have carried them. +pub(crate) fn fresh_attr_defs<'a>( + indexes: &'a [IndexMeta], + cached: &'a [AttributeDefinition], +) -> &'a [AttributeDefinition] { + indexes.first().map_or(cached, |i| &i.attribute_definitions) } /// Fetch all index metadata for a table from the catalog. @@ -153,9 +169,11 @@ pub(crate) async fn fetch_indexes_for_table( table_id: &str, pool: &sqlx::SqlitePool, ) -> Result, StorageError> { - let rows: Vec<(String, String, String, String, String, Option)> = sqlx::query_as( - "SELECT index_id, index_name, index_type, key_schema, projection, propagation_delay_ms \ - FROM indexes WHERE table_id = ?", + type Row = (String, String, String, String, String, Option, String); + let rows: Vec = sqlx::query_as( + "SELECT i.index_id, i.index_name, i.index_type, i.key_schema, i.projection, \ + i.propagation_delay_ms, t.attribute_definitions \ + FROM indexes i JOIN tables t ON t.table_id = i.table_id WHERE i.table_id = ?", ) .bind(table_id) .fetch_all(pool) @@ -163,7 +181,7 @@ pub(crate) async fn fetch_indexes_for_table( .map_err(crate::sqlite_util::map_sqlx_err)?; rows.into_iter() - .map(|(index_id, index_name, index_type, ks, proj, delay)| { + .map(|(index_id, index_name, index_type, ks, proj, delay, ad)| { Ok(IndexMeta { index_id, index_name, @@ -173,6 +191,8 @@ pub(crate) async fn fetch_indexes_for_table( projection: serde_json::from_str(&proj) .map_err(|e| StorageError::Internal(e.to_string()))?, propagation_delay_ms: delay, + attribute_definitions: serde_json::from_str(&ad) + .map_err(|e| StorageError::Internal(e.to_string()))?, }) }) .collect() @@ -213,7 +233,7 @@ pub(crate) async fn enqueue_async_indexes( } let context = GsiApplyContext { base_key_schema: key_info.key_schema.clone(), - attribute_definitions: key_info.attribute_definitions.clone(), + attribute_definitions: idx.attribute_definitions.clone(), index: GsiIndexDef { index_id: idx.index_id.clone(), key_schema: idx.key_schema.clone(), @@ -364,7 +384,9 @@ pub(crate) async fn sync_indexes( continue; // Async GSI — applied later by the propagation worker. } let idx_table = index_table_name(&idx.index_id); - let idx_sks = all_sort_key_info(&idx.key_schema, attr_defs); + // The index's key types come from the write's own catalog read; see + // `IndexMeta::attribute_definitions`. + let idx_sks = all_sort_key_info(&idx.key_schema, &idx.attribute_definitions); if let Some(old) = old_item && item_has_index_keys(old, &idx.key_schema) diff --git a/crates/storage-sqlite/src/data/put_item.rs b/crates/storage-sqlite/src/data/put_item.rs index 4efbaeb86..7ad7a9c21 100644 --- a/crates/storage-sqlite/src/data/put_item.rs +++ b/crates/storage-sqlite/src/data/put_item.rs @@ -56,7 +56,7 @@ impl SqliteEngine { extenddb_core::validation::validate_index_keys( &item, &index_refs, - &key_info.attribute_definitions, + super::index::fresh_attr_defs(&indexes, &key_info.attribute_definitions), ) .map_err(|e| StorageError::Validation(e.to_string()))?; } diff --git a/crates/storage-sqlite/src/data/transactions.rs b/crates/storage-sqlite/src/data/transactions.rs index 93e278f12..e9b5f8f5b 100644 --- a/crates/storage-sqlite/src/data/transactions.rs +++ b/crates/storage-sqlite/src/data/transactions.rs @@ -298,10 +298,12 @@ async fn execute_transact_write_op( // per-item cancellation reason; an empty index key is up-front input // validation (a top-level ValidationException). let idx_refs = index_key_refs(indexes); - validation::validate_index_key_types(item, &idx_refs, &key_info.attribute_definitions) - .map_err(|e| { - TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())) - })?; + validation::validate_index_key_types( + item, + &idx_refs, + super::index::fresh_attr_defs(indexes, &key_info.attribute_definitions), + ) + .map_err(|e| TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())))?; validation::validate_index_key_not_empty( item, &idx_refs, @@ -422,10 +424,12 @@ async fn execute_transact_write_op( // mismatch is a cancellation reason; setting an index key to an // empty value is a top-level ValidationException. let idx_refs = index_key_refs(indexes); - validation::validate_index_key_types(&item, &idx_refs, &key_info.attribute_definitions) - .map_err(|e| { - TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())) - })?; + validation::validate_index_key_types( + &item, + &idx_refs, + super::index::fresh_attr_defs(indexes, &key_info.attribute_definitions), + ) + .map_err(|e| TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())))?; validation::validate_index_key_not_empty( &item, &idx_refs, diff --git a/crates/storage-sqlite/src/data/update_item.rs b/crates/storage-sqlite/src/data/update_item.rs index c0875a673..ae18ae1e9 100644 --- a/crates/storage-sqlite/src/data/update_item.rs +++ b/crates/storage-sqlite/src/data/update_item.rs @@ -80,8 +80,12 @@ impl SqliteEngine { // rather than silently producing a malformed / unmatchable index row. if !indexes.is_empty() { let idx_refs = index_key_refs(&indexes); - validation::validate_index_key_types(&item, &idx_refs, &key_info.attribute_definitions) - .map_err(|e| StorageError::Validation(e.to_string()))?; + validation::validate_index_key_types( + &item, + &idx_refs, + super::index::fresh_attr_defs(&indexes, &key_info.attribute_definitions), + ) + .map_err(|e| StorageError::Validation(e.to_string()))?; validation::validate_index_key_not_empty( &item, &idx_refs, diff --git a/crates/storage-sqlite/src/update_table.rs b/crates/storage-sqlite/src/update_table.rs index ed1e6828c..8adff6173 100644 --- a/crates/storage-sqlite/src/update_table.rs +++ b/crates/storage-sqlite/src/update_table.rs @@ -18,10 +18,32 @@ use extenddb_core::types::{ TableDescription, UpdateTableInput, VectorIndexSpecification, }; use extenddb_storage::error::StorageError; -use extenddb_storage::util::effective_attribute_definitions; +use extenddb_storage::util::{building_refuses, effective_attribute_definitions}; use crate::store::SqliteEngine; +/// Refuse a GSI Create, or a Delete of a GSI other than the one building, while +/// a GSI on the table is `CREATING`. +async fn refuse_while_building( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + table_id: &str, + updates: &[extenddb_core::types::GlobalSecondaryIndexUpdate], +) -> Result<(), StorageError> { + let building: Vec = sqlx::query_scalar( + "SELECT index_name FROM indexes WHERE table_id = ? AND index_status = 'CREATING'", + ) + .bind(table_id) + .fetch_all(&mut **tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if building_refuses(&building, updates) { + return Err(StorageError::LimitExceeded( + extenddb_core::types::ONLINE_INDEX_LIMIT.to_owned(), + )); + } + Ok(()) +} + /// Refuse a new index whose name is already taken on this table, whatever index /// family holds it. /// @@ -64,6 +86,34 @@ impl SqliteEngine { ) -> Result { Self::validate_account_id(account_id)?; + // A GSI build in this process runs under the write lock below and keeps + // it until the index is ACTIVE, so a check taken under the lock would + // only ever see the build finished. Read before taking it, from a WAL + // snapshot that needs no lock, so a second index change sent during a + // build is refused at once, as the service refuses it, instead of + // waiting for the build and then succeeding. It sees only a committed + // CREATING row: a build that holds the lock but is still in its catalog + // transaction is not seen, and the request then waits for that build + // and succeeds, as it did before this check. + if let Some(updates) = &input.global_secondary_index_updates + && !updates.is_empty() + { + let building: Vec = sqlx::query_scalar( + "SELECT i.index_name FROM indexes i JOIN tables t ON t.table_id = i.table_id \ + WHERE t.account_id = ? AND t.table_name = ? AND i.index_status = 'CREATING'", + ) + .bind(account_id) + .bind(&input.table_name) + .fetch_all(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if building_refuses(&building, updates) { + return Err(StorageError::LimitExceeded( + extenddb_core::types::ONLINE_INDEX_LIMIT.to_owned(), + )); + } + } + let _writer = self.write_lock.lock().await; let mut tx = self .pool @@ -315,6 +365,18 @@ impl SqliteEngine { // columns from the same set the catalog now holds. let mut merged_attr_defs_for_ddl: Option> = None; if let Some(updates) = &input.global_secondary_index_updates { + { + // One index change at a time while a GSI builds, as the service + // has it: a Create, or a Delete of any other GSI, is refused with + // LimitExceededException (measured for both, in both of the + // service's CREATING phases). A Delete of the index that is + // building is a different case and is left to the delete path. + // Inside the write lock this sees a CREATING row only from + // another process on the same file or one a crash left; a build + // in this process is caught before the lock is taken, see the + // top of this function. + refuse_while_building(&mut tx, &table_id, updates).await?; + } for update in updates { if let Some(create) = &update.create { // Across families, not just this one: see @@ -1438,6 +1500,224 @@ mod reconciler_tests { ); } + /// A write that carries the key info cached before an UpdateTable added a GSI + /// with a sort key on a new attribute must still give the index row its sort + /// key. The engine drops the cached key info only when UpdateTable returns, so + /// writes in between carry definitions that do not know the new attribute; + /// with them the row was written without the sort key column, and a Query on + /// the full index key did not find the item. + #[tokio::test] + async fn write_with_key_info_older_than_a_new_gsi_indexes_its_sort_key() { + use extenddb_storage::{DataEngine, TableEngine}; + + let engine = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("engine"); + crate::schema::apply(&engine.pool).await.expect("schema"); + let account = "000000000000"; + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES (?, 'default')") + .bind(account) + .execute(&engine.pool) + .await + .expect("account"); + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('index_propagation_delay_ms', '0') \ + ON CONFLICT (key) DO UPDATE SET value = excluded.value", + ) + .execute(&engine.pool) + .await + .expect("synchronous propagation"); + + let input: extenddb_core::types::CreateTableInput = serde_json::from_value(json!({ + "TableName": "t", + "KeySchema": [{"AttributeName": "pk", "KeyType": "HASH"}], + "AttributeDefinitions": [{"AttributeName": "pk", "AttributeType": "S"}], + "BillingMode": "PAY_PER_REQUEST" + })) + .expect("input"); + engine + .create_table_impl(account, input, false) + .await + .expect("create table"); + sqlx::query("UPDATE tables SET table_status = 'ACTIVE' WHERE table_name = 't'") + .execute(&engine.pool) + .await + .expect("activate the table"); + let stale = engine.table_key_info(account, "t").await.expect("key info"); + + let update: extenddb_core::types::UpdateTableInput = serde_json::from_value(json!({ + "TableName": "t", + "AttributeDefinitions": [ + {"AttributeName": "g", "AttributeType": "S"}, + {"AttributeName": "r", "AttributeType": "N"} + ], + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "gidx", + "KeySchema": [ + {"AttributeName": "g", "KeyType": "HASH"}, + {"AttributeName": "r", "KeyType": "RANGE"} + ], + "Projection": {"ProjectionType": "ALL"} + }}] + })) + .expect("update input"); + engine + .update_table_impl(account, update) + .await + .expect("add the GSI"); + + let item: extenddb_core::types::Item = serde_json::from_value(json!({ + "pk": {"S": "a"}, "g": {"S": "x"}, "r": {"N": "7"} + })) + .expect("item"); + engine + .put_item( + &stale, + item, + false, + None, + &extenddb_core::expression::ExpressionMaps::default(), + None, + ) + .await + .expect("put with the older key info"); + + let (index_id,): (String,) = + sqlx::query_as("SELECT index_id FROM indexes WHERE index_name = 'gidx'") + .fetch_one(&engine.pool) + .await + .expect("index id"); + let missing: i64 = sqlx::query_scalar(&format!( + "SELECT COUNT(*) FROM \"_ddb_{index_id}\" WHERE sk_n IS NULL" + )) + .fetch_one(&engine.pool) + .await + .expect("count rows without a sort key"); + let rows: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM \"_ddb_{index_id}\"")) + .fetch_one(&engine.pool) + .await + .expect("count rows"); + assert_eq!((rows, missing), (1, 0), "one index row, with its sort key"); + } + + /// While a GSI row is `CREATING`, with the engine write lock held as a build + /// in this process holds it: a second Create, and a Delete of another GSI, + /// are refused within 500 ms, without waiting for the lock, and with the + /// service's LimitExceeded sentence, and a read naming the index is told the + /// table has no such index, as the service does. Main returned a 500 from + /// the missing or partial data table instead. + #[tokio::test] + async fn a_creating_gsi_refuses_a_second_change_and_reads() { + use extenddb_storage::TableEngine; + + let engine = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("engine"); + crate::schema::apply(&engine.pool).await.expect("schema"); + let account = "000000000000"; + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES (?, 'default')") + .bind(account) + .execute(&engine.pool) + .await + .expect("account"); + let input: extenddb_core::types::CreateTableInput = serde_json::from_value(json!({ + "TableName": "t", + "KeySchema": [{"AttributeName": "pk", "KeyType": "HASH"}], + "AttributeDefinitions": [ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "g", "AttributeType": "S"} + ], + "BillingMode": "PAY_PER_REQUEST" + })) + .expect("input"); + engine + .create_table_impl(account, input, false) + .await + .expect("create table"); + sqlx::query("UPDATE tables SET table_status = 'ACTIVE' WHERE table_name = 't'") + .execute(&engine.pool) + .await + .expect("activate the table"); + let (table_id,): (String,) = + sqlx::query_as("SELECT table_id FROM tables WHERE table_name = 't'") + .fetch_one(&engine.pool) + .await + .expect("table_id"); + sqlx::query( + "INSERT INTO indexes \ + (table_id, index_id, index_name, index_type, key_schema, projection, index_status) \ + VALUES (?, 'idx-1', 'gsi1', 'GSI', ?, ?, 'CREATING')", + ) + .bind(&table_id) + .bind(json!([{"AttributeName": "g", "KeyType": "HASH"}]).to_string()) + .bind(json!({"ProjectionType": "ALL"}).to_string()) + .execute(&engine.pool) + .await + .expect("plant a CREATING index"); + + let second: extenddb_core::types::UpdateTableInput = serde_json::from_value(json!({ + "TableName": "t", + "GlobalSecondaryIndexUpdates": [{"Create": { + "IndexName": "gsi2", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"} + }}] + })) + .expect("update input"); + // A build in this process holds the write lock until it publishes. With + // it held, the refusal must come at once rather than after the build. + let building = engine.write_lock.lock().await; + let refused = tokio::time::timeout( + std::time::Duration::from_millis(500), + engine.update_table_impl(account, second), + ) + .await + .expect("refused without waiting for the build's lock"); + let Err(extenddb_storage::error::StorageError::LimitExceeded(msg)) = refused else { + panic!("a second create during a build was not refused with LimitExceeded"); + }; + assert_eq!(msg, extenddb_core::types::ONLINE_INDEX_LIMIT); + + // A Delete of another GSI is refused the same way, also without waiting. + sqlx::query( + "INSERT INTO indexes \ + (table_id, index_id, index_name, index_type, key_schema, projection, index_status) \ + VALUES (?, 'idx-0', 'gsi0', 'GSI', ?, ?, 'ACTIVE')", + ) + .bind(&table_id) + .bind(json!([{"AttributeName": "g", "KeyType": "HASH"}]).to_string()) + .bind(json!({"ProjectionType": "ALL"}).to_string()) + .execute(&engine.pool) + .await + .expect("plant an ACTIVE index"); + let delete_other: extenddb_core::types::UpdateTableInput = serde_json::from_value(json!({ + "TableName": "t", + "GlobalSecondaryIndexUpdates": [{"Delete": {"IndexName": "gsi0"}}] + })) + .expect("update input"); + let refused = tokio::time::timeout( + std::time::Duration::from_millis(500), + engine.update_table_impl(account, delete_other), + ) + .await + .expect("refused without waiting for the build's lock"); + assert!( + matches!( + refused, + Err(extenddb_storage::error::StorageError::LimitExceeded(_)) + ), + "deleting another GSI during a build was not refused with LimitExceeded" + ); + drop(building); + assert!( + matches!( + engine.index_info_by_table_id(&table_id, "gsi1").await, + Err(extenddb_storage::error::StorageError::IndexNotFound(_)) + ), + "a read of a CREATING index resolves as no such index" + ); + } + /// A vector index left in `CREATING` by a crash must be rebuilt on startup and /// flipped to `ACTIVE`, with `Backfilling` cleared. /// diff --git a/crates/storage/src/util/mod.rs b/crates/storage/src/util/mod.rs index b9c68a0f3..a762738e7 100755 --- a/crates/storage/src/util/mod.rs +++ b/crates/storage/src/util/mod.rs @@ -27,3 +27,59 @@ pub use key::{ merge_attribute_definitions, parse_sk, pk_to_text, recover_sort_key_definitions, sk_column, sk_column_n, sk_info, }; + +/// Whether an UpdateTable's GSI actions must be refused because the GSIs named +/// in `building` are `CREATING`: any Create, or a Delete of a GSI that is not +/// one of them. The service answers both with LimitExceededException ("Only 1 +/// online index can be created or deleted simultaneously per table"). +#[must_use] +pub fn building_refuses( + building: &[String], + updates: &[extenddb_core::types::GlobalSecondaryIndexUpdate], +) -> bool { + if building.is_empty() { + return false; + } + updates.iter().any(|u| { + u.create.is_some() + || u.delete + .as_ref() + .is_some_and(|d| !building.contains(&d.index_name)) + }) +} + +#[cfg(test)] +mod building_refuses_tests { + use super::building_refuses; + use extenddb_core::types::GlobalSecondaryIndexUpdate; + + fn updates(json: serde_json::Value) -> Vec { + serde_json::from_value(json).expect("updates") + } + + #[test] + fn nothing_building_refuses_nothing() { + let u = updates(serde_json::json!([{"Delete": {"IndexName": "a"}}])); + assert!(!building_refuses(&[], &u)); + } + + #[test] + fn a_create_or_a_delete_of_another_index_is_refused() { + let building = vec!["b".to_owned()]; + let create = updates(serde_json::json!([{"Create": { + "IndexName": "c", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"} + }}])); + assert!(building_refuses(&building, &create)); + let other = updates(serde_json::json!([{"Delete": {"IndexName": "a"}}])); + assert!(building_refuses(&building, &other)); + } + + #[test] + fn a_delete_of_the_building_index_is_not_refused_here() { + let building = vec!["b".to_owned()]; + let own = updates(serde_json::json!([{"Delete": {"IndexName": "b"}}])); + assert!(!building_refuses(&building, &own)); + } +} diff --git a/docs/design/04-component-storage.md b/docs/design/04-component-storage.md index 5587b8ccc..429979b59 100755 --- a/docs/design/04-component-storage.md +++ b/docs/design/04-component-storage.md @@ -668,17 +668,37 @@ This column-on-tables approach is sufficient while control plane operations map Until those requirements arise, the single-column approach avoids the complexity of a separate job queue while providing full crash recovery. -### 6.2 GSI Backfill on CreateIndex - -When `UpdateTable` adds a new GSI to a table with existing data: - -1. Set the new index status to `CREATING` in `indexes` -2. Spawn a background task that scans the base table in batches (configurable batch size, default 1000) -3. For each batch, INSERT the projected attributes into the new GSI table -4. On completion, set index status to `ACTIVE` -5. During backfill, writes to the base table also write to the new GSI table (the write path checks index status and -includes `CREATING` indexes) -6. Queries against a `CREATING` index return `ResourceNotFoundException` (matching DynamoDB behavior) +### 6.2 GSI backfill when UpdateTable adds an index + +When `UpdateTable` adds a new GSI to a table with existing data, the PostgreSQL backend +(`crates/storage-postgres/src/data/gsi_build.rs`, which has the full ordering argument): + +1. Creates the GSI's data table, then commits its `indexes` row as `CREATING`. Every write reads the catalog, + so from that commit on writes maintain the index the same way as an `ACTIVE` one, inline or through the + propagation queue. +2. Waits one second, then takes `LOCK TABLE ... IN SHARE MODE` on the base table and releases it at once. The + lock waits for writes already in flight; writes arriving meanwhile queue behind it, each holding a data pool + connection, so the request carries a lock timeout (one second over `deadlock_timeout`) and is retried a few + times with growing pauses. The wait and a time check on the write side (a write that took more than 500 ms + from its catalog read re-reads it and queues a repair for any GSI it missed; a write with no transaction + queues that repair after its own statement, and loses it if the process dies in between) together cover + writes that read the catalog just before the index row committed: the catalog and the data tables are + separate databases, so no transaction can tell those writes the index exists. +3. Scans the base table in keyset pages (256 rows, `FOR SHARE SKIP LOCKED`, skipped rows taken one at a time + afterwards), writing each item's current state into the index, then deletes index rows whose item is gone. +4. Sets the status to `ACTIVE`. A build that fails for a transient reason leaves the index `CREATING`, and a + background worker finishes it; one that would fail every time removes the index and fails the request. + +Queued GSI updates apply the item's state at apply time rather than the image the row carried, so a queue row +applied late, twice, or after the backfill wrote a newer state cannot leave the index behind the table. The +worker reads the base row under `FOR SHARE` for that, so a queued update and a concurrent write to the same item +now wait for each other where before they did not meet. + +`UpdateTable` returns after the build. While a GSI is `CREATING`, a Query or Scan naming it is answered as if the +table had no such index (the service's response), and a GSI Create, or a Delete of another GSI, on the table is +refused with LimitExceededException. A Delete of the building GSI and a throughput Update are not refused. DeleteTable stops the propagation queue for the table (a `vector_index_holds` row, which the +claim query honours) and deletes its pending rows before it drops any data table, so a worker holding a claimed row +and waiting on the base table cannot deadlock with the drop. ## 7. Pagination Token Encoding diff --git a/tests/test_gsi_online_backfill.py b/tests/test_gsi_online_backfill.py new file mode 100644 index 000000000..34e17885e --- /dev/null +++ b/tests/test_gsi_online_backfill.py @@ -0,0 +1,262 @@ +# Copyright 2026 ExtendDB contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Adding a GSI to a table that is taking writes. + +DynamoDB builds a new GSI online: writes keep flowing while it backfills, and +once the index is ACTIVE (and propagation has caught up) it holds exactly the +base items that carry its key attributes, with their latest values. These +tests drive puts, updates that change or remove the index key, deletes and +TransactWriteItems from several threads while UpdateTable adds the index, then +compare a Scan of the index with a Scan of the table. +""" + +from __future__ import annotations + +import os +import re +import threading +import time + +import boto3 +import pytest +from botocore.config import Config + +from conftest import wait_for_active, wait_for_deleted + +ITEMS = 1500 +WRITERS = 4 +PARTITIONS = 40 + +def _configured_backend() -> str: + """The `[storage] backend` of the server's config, or "" if unknown.""" + path = os.environ.get("EXTENDDB_CONFIG", "").strip() + try: + with open(path, encoding="utf-8") as f: + text = f.read() + except OSError: + return "" + match = re.search(r'(?ms)^\[storage\]\s*$.*?^\s*backend\s*=\s*"([^"]+)"', text) + return match.group(1) if match else "" + + +# On Cassandra these writers get HTTP 500 from DeleteItem under contention on one +# item, and the online GSI build had not finished after 10 minutes of them in CI. +# Both are Cassandra backend issues outside this test's subject, and the second +# would cost the job 10 minutes per case, so the test is skipped there. +SKIP_ON_CASSANDRA = pytest.mark.skipif( + _configured_backend() == "cassandra", + reason="Cassandra fails contended writes with HTTP 500 and did not finish the build", +) + +# On MongoDB, concurrent UpdateItem and TransactWriteItems calls that create the +# same item return HTTP 500 (an ImmutableField error on `_id`). That happens with +# these writers on a table with no index at all, so it is not about the GSI build; +# it is the MongoDB write-race fix that test_put_item_create_race.py also waits +# for. Not strict, because the race does not fire in every run. The MongoDB test +# runner sets EXTENDDB_TEST_MONGODB_CONTAINER (the Cassandra runner sets nothing +# of the kind, hence the config file read above). +XFAIL_UNTIL_MONGODB_FIX = pytest.mark.xfail( + bool(os.environ.get("EXTENDDB_TEST_MONGODB_CONTAINER", "").strip()), + reason="MongoDB returns HTTP 500 for concurrent creates of one item", + raises=AssertionError, + strict=False, +) + + +def _client(endpoint_url): + """One client per thread: a boto3 client is not safe for concurrent use.""" + kwargs = { + "service_name": "dynamodb", + "region_name": os.environ.get("AWS_DEFAULT_REGION", "us-east-1"), + "config": Config(retries={"max_attempts": 1}, read_timeout=300), + } + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + if endpoint_url.startswith("https://"): + kwargs["verify"] = False + return boto3.client(**kwargs) + + +def _key(n): + return {"pk": {"S": f"p{n % PARTITIONS}"}, "sk": {"N": str(n)}} + + +def _writer(endpoint_url, table, seed, stop, errors): + ddb = _client(endpoint_url) + i = 0 + while not stop.is_set(): + i += 1 + n = (seed * 7919 + i * 31) % (ITEMS + 300) + try: + op = i % 6 + if op == 0: + ddb.delete_item(TableName=table, Key=_key(n)) + elif op == 1: + ddb.update_item( + TableName=table, Key=_key(n), + UpdateExpression="SET g = :g, r = :r", + ExpressionAttributeValues={":g": {"S": f"u{i % 3}"}, ":r": {"N": str(i)}}, + ) + elif op == 2: + # The item leaves the index. + ddb.update_item(TableName=table, Key=_key(n), UpdateExpression="REMOVE g") + elif op == 3: + m = (n + 1) % (ITEMS + 300) + ddb.transact_write_items(TransactItems=[ + {"Put": {"TableName": table, "Item": { + **_key(n), "g": {"S": f"t{i % 4}"}, "r": {"N": str(i)}}}}, + {"Update": { + "TableName": table, "Key": _key(m), + "UpdateExpression": "SET g = :g, r = :r", + "ExpressionAttributeValues": {":g": {"S": f"x{i % 2}"}, ":r": {"N": str(i)}}, + }}, + ]) + else: + ddb.put_item(TableName=table, Item={ + **_key(n), "g": {"S": f"w{i % 5}"}, "r": {"N": str(i)}}) + except ddb.exceptions.TransactionCanceledException as e: + # Two writers can contend for one item; the service cancels one. + if "TransactionConflict" not in str(e): + errors.append(repr(e)[:300]) + except Exception as e: # noqa: BLE001 - every failure is reported + errors.append(repr(e)[:300]) + + +def _scan(ddb, table, index=None): + """Map of base key to (g, r), and the base keys seen more than once.""" + kwargs = {"TableName": table} + if index: + kwargs["IndexName"] = index + out = {} + repeated = [] + while True: + resp = ddb.scan(**kwargs) + for it in resp["Items"]: + key = (it["pk"]["S"], it["sk"]["N"]) + if key in out: + repeated.append(key) + out[key] = (it.get("g", {}).get("S"), it.get("r", {}).get("N")) + if "LastEvaluatedKey" not in resp: + return out, repeated + kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + + +def _wait_index_active(ddb, table, index, timeout=600.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + gsis = ddb.describe_table(TableName=table)["Table"].get("GlobalSecondaryIndexes", []) + if any(g["IndexName"] == index and g["IndexStatus"] == "ACTIVE" for g in gsis): + return + time.sleep(0.2) + raise TimeoutError(f"{index} on {table} did not become ACTIVE") + + +@pytest.fixture() +def loaded_table(dynamodb_client, unique_table_name): + table = unique_table_name + dynamodb_client.create_table( + TableName=table, + AttributeDefinitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "N"}, + ], + KeySchema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + wait_for_active(dynamodb_client, table) + for start in range(0, ITEMS, 25): + dynamodb_client.batch_write_item(RequestItems={table: [ + {"PutRequest": {"Item": { + **_key(n), "g": {"S": f"g{n % 7}"}, "r": {"N": str(n)}}}} + for n in range(start, min(start + 25, ITEMS)) + ]}) + yield table + try: + dynamodb_client.delete_table(TableName=table) + except Exception: # noqa: BLE001 + pass + else: + wait_for_deleted(dynamodb_client, table) + + +@SKIP_ON_CASSANDRA +@XFAIL_UNTIL_MONGODB_FIX +@pytest.mark.parametrize("with_range", [False, True], ids=["hash-only-gsi", "hash-range-gsi"]) +def test_gsi_added_under_concurrent_writes_matches_table( + dynamodb_client, endpoint_url, loaded_table, with_range, +): + table = loaded_table + stop = threading.Event() + errors: list[str] = [] + threads = [ + threading.Thread(target=_writer, args=(endpoint_url, table, s, stop, errors)) + for s in range(WRITERS) + ] + for t in threads: + t.start() + try: + time.sleep(0.3) + defs = [{"AttributeName": "g", "AttributeType": "S"}] + schema = [{"AttributeName": "g", "KeyType": "HASH"}] + if with_range: + # A sort key on an attribute the table had not defined before. + defs.append({"AttributeName": "r", "AttributeType": "N"}) + schema.append({"AttributeName": "r", "KeyType": "RANGE"}) + # UpdateTable returns when the build is done. The session client's default + # 60 s read timeout and retries would turn a slow build into a retried call + # that fails with ResourceInUse, so the control call uses a writer client. + _client(endpoint_url).update_table( + TableName=table, + AttributeDefinitions=defs, + GlobalSecondaryIndexUpdates=[{"Create": { + "IndexName": "online-gsi", "KeySchema": schema, + "Projection": {"ProjectionType": "ALL"}, + }}], + ) + _wait_index_active(dynamodb_client, table, "online-gsi") + time.sleep(0.5) + finally: + stop.set() + for t in threads: + t.join() + assert errors == [], errors[:5] + + base, _ = _scan(dynamodb_client, table) + if with_range: + want = {k: v for k, v in base.items() if v[0] is not None and v[1] is not None} + else: + want = {k: v for k, v in base.items() if v[0] is not None} + + # The index is eventually consistent, so allow propagation to finish. + got, repeated = {}, [] + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + got, repeated = _scan(dynamodb_client, table, "online-gsi") + if got == want and not repeated: + break + time.sleep(0.5) + # One index entry per base item: an entry left behind under the item's old + # index key would otherwise hide behind the newer one in the map. + assert repeated == [], f"{len(repeated)} items indexed more than once: {repeated[:5]}" + missing = sorted(k for k in want if k not in got) + extra = sorted(k for k in got if k not in want) + stale = sorted(k for k in want if k in got and got[k] != want[k]) + assert (missing, extra, stale) == ([], [], []), ( + f"index vs table: {len(missing)} missing {missing[:5]}, {len(extra)} extra {extra[:5]}, " + f"{len(stale)} stale {[(k, got[k], want[k]) for k in stale[:5]]}" + ) + + if with_range: + # Each indexed item is found by a Query on its full index key, which + # needs the sort key column the index row was written with. + for k, (g, r) in list(want.items())[:100]: + items = dynamodb_client.query( + TableName=table, IndexName="online-gsi", + KeyConditionExpression="g = :g AND r = :r", + ExpressionAttributeValues={":g": {"S": g}, ":r": {"N": r}}, + )["Items"] + assert any((it["pk"]["S"], it["sk"]["N"]) == k for it in items), (k, g, r)