From 8c317624dcfa0afc7bd7199c7c846388b9964487 Mon Sep 17 00:00:00 2001 From: Scott Robinson Date: Tue, 20 Oct 2026 03:56:56 +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 on the table is refused with LimitExceededException. DeleteTable stops the propagation queue for the table before dropping anything, which removes a deadlock with the queue worker that left orphan data tables; the worker bounds its wait on a held base row to 2 s. 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 | 7 +- Cargo.lock | 1 + crates/core/src/types/mod.rs | 2 +- crates/core/src/types/table.rs | 8 + crates/core/src/validation/mod.rs | 4 +- crates/storage-postgres/Cargo.toml | 1 + crates/storage-postgres/src/data/ddl.rs | 90 +- .../storage-postgres/src/data/delete_item.rs | 29 +- crates/storage-postgres/src/data/gsi_build.rs | 1439 +++++++++++++++++ crates/storage-postgres/src/data/index.rs | 188 ++- crates/storage-postgres/src/data/mod.rs | 5 +- crates/storage-postgres/src/data/put_item.rs | 33 +- .../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 | 45 +- crates/storage-postgres/src/gsi_queue.rs | 79 +- crates/storage-postgres/src/lib.rs | 17 + crates/storage-postgres/src/table_helpers.rs | 103 +- crates/storage-postgres/src/update_table.rs | 278 +++- crates/storage-postgres/src/worker_store.rs | 40 +- .../tests/gsi_online_build.rs | 933 +++++++++++ .../tests/gsi_stale_writer.rs | 342 ++++ 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 | 198 +++ docs/design/04-component-storage.md | 42 +- tests/test_gsi_online_backfill.py | 262 +++ 31 files changed, 3935 insertions(+), 390 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 cfdfdf2ee..d8435a72b 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -300,12 +300,13 @@ jobs: # storage layer directly against this job's PostgreSQL. The PutItem create # race tests hold a create open in an outside transaction, and the # write-transaction conflict tests force a deadlock; both also need direct - # access. They build their own throwaway databases; the connection string is - # the server, not a database. + # 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 - run: cargo test --release -p extenddb-storage-postgres --test vector_control_plane --test put_create_race --test twi_conflict --test stream_shard_retention + run: cargo test --release -p extenddb-storage-postgres --test vector_control_plane --test put_create_race --test twi_conflict --test stream_shard_retention --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 a2797101c..660b0c628 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1475,6 +1475,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 a7ed110e3..16bc72f1c 100755 --- a/crates/core/src/types/mod.rs +++ b/crates/core/src/types/mod.rs @@ -55,7 +55,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 531c4094a..d901fe88a 100755 --- a/crates/core/src/types/table.rs +++ b/crates/core/src/types/table.rs @@ -589,6 +589,14 @@ 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 adds or deletes a GSI while another +/// 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 5b780fd6a..c74a56311 100755 --- a/crates/core/src/validation/mod.rs +++ b/crates/core/src/validation/mod.rs @@ -492,9 +492,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/ddl.rs b/crates/storage-postgres/src/data/ddl.rs index 6edb222f3..811dfb8f3 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,71 @@ 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. + pub(crate) async fn drop_table_storage( + data: &sqlx::PgPool, + table_id: &str, + index_ids: &[String], + vector_index_ids: &[String], + ) -> 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()))?; + 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..e9953665f 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()))?; @@ -254,6 +261,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 +321,6 @@ impl PostgresEngine { sync_indexes( &mut tx, &key_info.key_schema, - &key_info.attribute_definitions, &indexes, oi.as_ref(), None, @@ -371,6 +378,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()))?; @@ -393,6 +403,7 @@ impl PostgresEngine { .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..074356167 --- /dev/null +++ b/crates/storage-postgres/src/data/gsi_build.rs @@ -0,0 +1,1439 @@ +// 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()) +} + +/// Test hook: how long a write sleeps between its catalog read and its data work. +/// Zero outside tests. +static PAUSE_AFTER_INDEX_READ_MS: std::sync::atomic::AtomicU64 = + std::sync::atomic::AtomicU64::new(0); + +/// Make every write on this process pause for `pause` after reading its index +/// set, 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`] exists for, which no timing in a test can otherwise +/// hit reliably. Hidden because no deployment path calls it. +#[doc(hidden)] +pub fn pause_writes_after_index_read_for_tests(pause: Duration) { + let ms = u64::try_from(pause.as_millis()).unwrap_or(u64::MAX); + PAUSE_AFTER_INDEX_READ_MS.store(ms, std::sync::atomic::Ordering::Relaxed); +} + +pub(crate) async fn pause_after_index_read() { + let ms = PAUSE_AFTER_INDEX_READ_MS.load(std::sync::atomic::Ordering::Relaxed); + if ms > 0 { + tokio::time::sleep(Duration::from_millis(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, 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"] + .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 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( + "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 551861dbf..f27fff5bc 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(), @@ -371,10 +502,29 @@ pub(crate) async fn insert_index_row_multi( cols.push("item_data".to_owned()); let placeholders: Vec = (1..=cols.len()).map(|i| format!("${i}")).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. + 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(", "), - placeholders.join(", ") + placeholders.join(", "), + conflict_cols.join(", "), + updates.join(", ") ); let mut query = sqlx::query(&sql).bind(idx_pk_text); diff --git a/crates/storage-postgres/src/data/mod.rs b/crates/storage-postgres/src/data/mod.rs index e16ee3b65..a847890a3 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; @@ -140,7 +141,3 @@ mod transactions; mod tx_helpers; mod update_item; pub mod vector_index; - -pub(crate) use index::{ - delete_index_row_multi, insert_index_row_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..02fab343b 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()))?; @@ -262,6 +270,7 @@ impl PostgresEngine { &item_json, &self.data_pool )?; + self.fix_up_autocommit(stamp, key_info, &item).await; Ok(None) } } else { @@ -347,7 +356,6 @@ impl PostgresEngine { sync_indexes( &mut tx, &key_info.key_schema, - &key_info.attribute_definitions, &indexes, oi.as_ref(), Some(&item), @@ -408,6 +416,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()))?; @@ -434,6 +446,7 @@ impl PostgresEngine { .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 6259160cb..25d7fd49e 100644 --- a/crates/storage-postgres/src/data/vector_index.rs +++ b/crates/storage-postgres/src/data/vector_index.rs @@ -838,9 +838,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) @@ -955,7 +962,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; } }; @@ -971,7 +978,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 5a55231fe..d596c03eb 100755 --- a/crates/storage-postgres/src/delete_table.rs +++ b/crates/storage-postgres/src/delete_table.rs @@ -107,42 +107,15 @@ 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, + ) + .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..c6738868b 100644 --- a/crates/storage-postgres/src/gsi_queue.rs +++ b/crates/storage-postgres/src/gsi_queue.rs @@ -30,11 +30,6 @@ 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` @@ -50,6 +45,24 @@ 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 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 row is retried; nothing is lost. +const ROW_LOCK_TIMEOUT: &str = "2s"; + +/// Pause before re-claiming after a lock timeout, so a held row is not polled +/// at full speed. +const LOCK_RETRY_PAUSE: std::time::Duration = std::time::Duration::from_millis(200); + +/// 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 @@ -289,6 +302,13 @@ async fn worker(worker_id: u64, q: Arc) { tokio::time::timeout(wait, q.notify.notified()).await.ok(); } Ok(_) => continue, + Err(e) if is_lock_timeout(&e) => { + // A writer is holding the base row this partition's next row is + // about to resync. The claim rolled back and the row waits; the + // pause keeps this from spinning on the same held row. + tracing::debug!("GSI worker {worker_id}: row held by a writer, retrying: {e}"); + tokio::time::sleep(LOCK_RETRY_PAUSE).await; + } Err(e) => { tracing::error!("GSI worker {worker_id}: batch error: {e}"); tokio::time::sleep(MAX_IDLE).await; @@ -389,6 +409,13 @@ async fn process_batch(worker_id: u64, q: &GsiQueue) -> Result { apply_index( tx, + table_id, &c.index, old_item.as_ref(), new_item.as_ref(), @@ -499,39 +527,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 0d7a51c0d..3387d3576 100755 --- a/crates/storage-postgres/src/lib.rs +++ b/crates/storage-postgres/src/lib.rs @@ -39,6 +39,14 @@ pub use catalog_store::PostgresCatalogStore; pub use config::PostgresStorageConfig; pub use config::parse_connection_string; pub use credential_store::DbCredentialStore; +/// Make writes pause between their catalog read and their data work, 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::pause_writes_after_index_read_for_tests; +/// 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; /// Apply one queued vector row from its own context. /// /// Reachable so an integration test can drive the classification the propagation @@ -583,6 +591,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(); @@ -607,6 +623,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 d75049ba0..6012c9dc4 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..24da1e68f 100755 --- a/crates/storage-postgres/src/update_table.rs +++ b/crates/storage-postgres/src/update_table.rs @@ -10,6 +10,9 @@ use extenddb_storage::error::StorageError; use extenddb_storage::util::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. /// @@ -413,6 +416,26 @@ 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 { + if updates.iter().any(|u| u.create.is_some()) { + // One index build at a time, as the service has it. Read under the + // table row lock taken above, which every UpdateTable takes, so two + // requests cannot both pass. Measured for a second Create; a Delete + // (including of the index being built, which is how a stuck one is + // cleared by hand) and a throughput Update are left alone. + let building: Option<(i32,)> = sqlx::query_as( + "SELECT 1 FROM indexes WHERE table_id = $1 AND index_status = 'CREATING' \ + LIMIT 1", + ) + .bind(&table_id) + .fetch_optional(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if building.is_some() { + return Err(StorageError::LimitExceeded( + extenddb_core::types::ONLINE_INDEX_LIMIT.to_owned(), + )); + } + } for update in updates { if let Some(create) = &update.create { // Across families, not just this one: see @@ -435,7 +458,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 +470,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 +771,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 +867,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 +1006,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 372fbf49d..bf6c88f35 100755 --- a/crates/storage-postgres/src/worker_store.rs +++ b/crates/storage-postgres/src/worker_store.rs @@ -127,42 +127,12 @@ impl PostgresEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; - // P54 Bug 1: Drop data tables on the data pool after catalog commit. + // Drop data tables on the data pool after catalog commit. See + // `drop_table_storage` for why the queue is stopped before anything is + // dropped. for (table_id, index_ids, vector_index_ids) in &drop_info { - 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 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 now-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()))?; - - // And any vector build hold: the table is gone, so nothing will ever - // release it, and a stale hold blocks claims for that 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()))?; + Self::drop_table_storage(&self.data_pool, table_id, index_ids, vector_index_ids) + .await?; } Ok(transitions) 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..1bbcdbcc0 --- /dev/null +++ b/crates/storage-postgres/tests/gsi_online_build.rs @@ -0,0 +1,933 @@ +// 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 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 is refused with the service's +/// LimitExceeded sentence, a Query or Scan naming the index is told +/// the table has no such index (which is what the engine puts on the wire for +/// `IndexNotFound`), and DeleteTable with queue rows in flight for the table +/// neither deadlocks nor leaves data tables behind. +#[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"); + 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 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..d05fd6146 --- /dev/null +++ b/crates/storage-postgres/tests/gsi_stale_writer.rs @@ -0,0 +1,342 @@ +// 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. This file uses a test hook that makes every write +//! pause after its catalog read, 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::{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_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() +} + +async fn drive(range_cond: 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; + } + + extenddb_storage_postgres::pause_writes_after_index_read_for_tests(Duration::from_millis( + 3_000, + )); + // Writes that read "no index" now and write after the GSI is ACTIVE: a new + // item (an autocommit insert), an overwrite that changes the index key, and + // a delete, the last two through a transaction because of their conditions. + 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, + range_cond.then_some(&cond), + &maps, + None, + ); + let mut key2 = item(2, None, None); + key2.remove("g"); + let c = s + .engine + .delete_item(&ki, &key2, false, range_cond.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; + extenddb_storage_postgres::pause_writes_after_index_read_for_tests(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" + ); + 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); + } + // Unconditional writes on a table with no index take the autocommit path; + // conditional ones take a transaction. One test for both, run in sequence, + // because the pause is process-wide. + drive(false).await; + drive(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 d527c3e6a..41721d5da 100644 --- a/crates/storage-sqlite/src/update_table.rs +++ b/crates/storage-sqlite/src/update_table.rs @@ -315,6 +315,27 @@ 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 { + if updates.iter().any(|u| u.create.is_some()) { + // One index build at a time, as the service has it. The build here + // runs under the engine write lock, so a CREATING row is only seen + // from another process on the same file, or after a crash left it; + // refused all the same, for one answer on both backends. Measured + // for a second Create; a Delete (including of the index being + // built) and a throughput Update are left alone. + let building: Option<(i32,)> = sqlx::query_as( + "SELECT 1 FROM indexes WHERE table_id = ? AND index_status = 'CREATING' \ + LIMIT 1", + ) + .bind(&table_id) + .fetch_optional(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if building.is_some() { + return Err(StorageError::LimitExceeded( + extenddb_core::types::ONLINE_INDEX_LIMIT.to_owned(), + )); + } + } for update in updates { if let Some(create) = &update.create { // Across families, not just this one: see @@ -1302,6 +1323,183 @@ 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` (another process on the same file, or a + /// crash before the flip): a second Create is refused 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"); + let refused = engine.update_table_impl(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!( + 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/docs/design/04-component-storage.md b/docs/design/04-component-storage.md index 5587b8ccc..78aead390 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 second online index change on the table is refused with +LimitExceededException. 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)