From 046a9b26ed7992bec50b8744260de8e38cff1ee9 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Tue, 3 Nov 2026 03:06:53 +0000 Subject: [PATCH] fix(sqlite): scan a table's size without holding the write lock refresh_table_size took the engine write lock before it read the table id and ran SELECT COUNT(*), COALESCE(SUM(length(item_data)), 0) over the table's data, and held it until the UPDATE of the tables row. Every write needs that lock, so every five minutes, when the size refresh worker visited each active table in turn, writes waited for a full scan of each table. The scan only reads, and D1 serializes writers only, so it now runs without the lock in scan_table_size. One read transaction reads the table's id and status and then counts its data, so all three come from the same snapshot; a table that is not ACTIVE, such as a restore target still being filled, is skipped, and a table deleted after the lookup can no longer make the scan fail. The lock is taken only for the UPDATE, in store_table_size, which also matches the table_id that was scanned: a table deleted, or deleted and created again under the same name, since the scan keeps its counts. On an in-memory database the engine has a single connection, so a write still waits for that connection while a scan runs, now with the pool's acquire timeout instead of the mutex; the change helps file databases, where reads and the writer use separate connections. Unit tests: the refresh stores the item count and size; the scan runs while a writer holds the lock; a table that is not ACTIVE is neither scanned nor given counts; the refresh does not store its counts until the lock is released; and counts for another table_id are not stored. The last three each fail against one mutant: the status check removed, the UPDATE run without the lock, and the UPDATE without the table_id match. Fixes #719 --- crates/storage-sqlite/src/metadata.rs | 284 ++++++++++++++++++++++++-- 1 file changed, 264 insertions(+), 20 deletions(-) diff --git a/crates/storage-sqlite/src/metadata.rs b/crates/storage-sqlite/src/metadata.rs index 93c319873..27153e5fa 100644 --- a/crates/storage-sqlite/src/metadata.rs +++ b/crates/storage-sqlite/src/metadata.rs @@ -33,6 +33,71 @@ fn safe_index_attr(attr: &str) -> bool { } impl SqliteEngine { + /// The id, the item count and the summed item size of the table named + /// `table_name`, read without the engine write lock in one read + /// transaction, so all three come from the same snapshot. `None` when the + /// table is not `ACTIVE`, for example a restore target that is still + /// being filled. + async fn scan_table_size( + &self, + account_id: &str, + table_name: &str, + ) -> Result, StorageError> { + let internal = |e: sqlx::Error| StorageError::Internal(e.to_string()); + let mut tx = self.pool.begin().await.map_err(internal)?; + let row: Option<(String, String)> = sqlx::query_as( + "SELECT table_id, table_status FROM tables WHERE account_id = ? AND table_name = ?", + ) + .bind(account_id) + .bind(table_name) + .fetch_optional(&mut *tx) + .await + .map_err(internal)?; + let Some((table_id, status)) = row else { + return Err(StorageError::TableNotFound(table_name.to_owned())); + }; + if status != "ACTIVE" { + return Ok(None); + } + let data_table = data::data_table_name(&table_id); + let (item_count, table_size): (i64, i64) = sqlx::query_as(&format!( + "SELECT COUNT(*), COALESCE(SUM(length(item_data)), 0) FROM {data_table}" + )) + .fetch_one(&mut *tx) + .await + .map_err(internal)?; + tx.rollback().await.map_err(internal)?; + Ok(Some((table_id, item_count, table_size))) + } + + /// Record the counts of the scan of `table_id`. The caller holds the + /// engine write lock. A table deleted, or deleted and created again under + /// the same name, since the scan has another `table_id` or is no longer + /// `ACTIVE`, and keeps its counts. + async fn store_table_size( + &self, + account_id: &str, + table_name: &str, + table_id: &str, + item_count: i64, + table_size: i64, + ) -> Result<(), StorageError> { + sqlx::query( + "UPDATE tables SET item_count = ?, table_size_bytes = ? \ + WHERE account_id = ? AND table_name = ? AND table_id = ? \ + AND table_status = 'ACTIVE'", + ) + .bind(item_count) + .bind(table_size) + .bind(account_id) + .bind(table_name) + .bind(table_id) + .execute(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + Ok(()) + } + async fn table_id_for( &self, account_id: &str, @@ -359,28 +424,18 @@ impl MetadataEngine for SqliteEngine { let table_name = table_name.to_owned(); Box::pin(async move { Self::validate_account_id(&account_id)?; + // The scan only reads, so it takes no lock (D1). Holding the write + // lock across it made writes wait for a full scan of each table in + // turn. + let Some((table_id, item_count, table_size)) = + self.scan_table_size(&account_id, &table_name).await? + else { + return Ok(()); + }; // D1: every writer holds the engine write lock. let _writer = self.write_lock.lock().await; - let table_id = self.table_id_for(&account_id, &table_name).await?; - let data_table = data::data_table_name(&table_id); - let (item_count, table_size): (i64, i64) = sqlx::query_as(&format!( - "SELECT COUNT(*), COALESCE(SUM(length(item_data)), 0) FROM {data_table}" - )) - .fetch_one(&self.pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - sqlx::query( - "UPDATE tables SET item_count = ?, table_size_bytes = ? \ - WHERE account_id = ? AND table_name = ? AND table_status = 'ACTIVE'", - ) - .bind(item_count) - .bind(table_size) - .bind(&account_id) - .bind(&table_name) - .execute(&self.pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + self.store_table_size(&account_id, &table_name, &table_id, item_count, table_size) + .await?; Ok(()) }) } @@ -415,3 +470,192 @@ impl MetadataEngine for SqliteEngine { }) } } + +#[cfg(test)] +mod size_refresh_tests { + use std::time::Duration; + + use extenddb_core::types::{Item, TableKeyInfo}; + use extenddb_storage::{DataEngine, MetadataEngine}; + use serde_json::json; + + use crate::SqliteEngine; + + const ACCOUNT: &str = "000000000000"; + + async fn engine() -> SqliteEngine { + let engine = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("engine"); + crate::migrations::apply_for_tests(&engine.pool).await; + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES (?, 'default')") + .bind(ACCOUNT) + .execute(&engine.pool) + .await + .expect("account"); + // Zero control-plane delay: tables are ACTIVE at create time. + sqlx::query( + "INSERT OR REPLACE INTO settings (key, value) \ + VALUES ('control_plane_delay_seconds', '0')", + ) + .execute(&engine.pool) + .await + .expect("settings"); + engine + } + + /// A table with a string hash key holding `items` items; returns its id. + async fn table_with_items(engine: &SqliteEngine, name: &str, items: usize) -> String { + let input = serde_json::from_value(json!({ + "TableName": name, + "KeySchema": [{"AttributeName": "pk", "KeyType": "HASH"}], + "AttributeDefinitions": [{"AttributeName": "pk", "AttributeType": "S"}], + "BillingMode": "PAY_PER_REQUEST" + })) + .expect("input"); + let desc = engine + .create_table_impl(ACCOUNT, input, false) + .await + .expect("create"); + let key_info = TableKeyInfo { + table_name: name.to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: desc.table_id.clone(), + key_schema: desc.key_schema.clone(), + base_key_schema: desc.key_schema.clone(), + attribute_definitions: desc.attribute_definitions.clone(), + ..Default::default() + }; + for i in 0..items { + let item: Item = serde_json::from_value(json!({ + "pk": {"S": format!("p{i}")}, + "v": {"S": "x"} + })) + .expect("item"); + engine + .put_item( + &key_info, + item, + false, + None, + &extenddb_core::expression::ExpressionMaps::default(), + None, + ) + .await + .expect("put"); + } + desc.table_id + } + + async fn stored_counts(engine: &SqliteEngine, name: &str) -> (i64, i64) { + sqlx::query_as( + "SELECT item_count, table_size_bytes FROM tables \ + WHERE account_id = ? AND table_name = ?", + ) + .bind(ACCOUNT) + .bind(name) + .fetch_one(&engine.pool) + .await + .expect("counts") + } + + #[tokio::test] + async fn refresh_stores_the_item_count_and_size() { + let engine = engine().await; + let table_id = table_with_items(&engine, "t", 3).await; + MetadataEngine::refresh_table_size(&engine, ACCOUNT, "t") + .await + .expect("refresh"); + let (scanned_id, count, size) = engine + .scan_table_size(ACCOUNT, "t") + .await + .expect("scan") + .expect("active"); + assert_eq!(scanned_id, table_id); + assert_eq!(count, 3); + assert!(size > 0, "{size}"); + assert_eq!(stored_counts(&engine, "t").await, (count, size)); + } + + /// The scan takes no lock, so it runs while a writer holds the engine + /// write lock (#719). + #[tokio::test] + async fn the_scan_runs_while_a_writer_holds_the_lock() { + let engine = engine().await; + table_with_items(&engine, "t", 3).await; + let _writer = engine.write_lock.lock().await; + let scanned = tokio::time::timeout( + Duration::from_secs(10), + engine.scan_table_size(ACCOUNT, "t"), + ) + .await + .expect("the scan waited for the write lock") + .expect("scan") + .expect("active"); + assert_eq!(scanned.1, 3); + } + + /// A table that is not ACTIVE, such as a restore target still being + /// filled, is neither scanned nor given counts. + #[tokio::test] + async fn a_table_that_is_not_active_is_not_counted() { + let engine = engine().await; + table_with_items(&engine, "t", 3).await; + sqlx::query("UPDATE tables SET table_status = 'CREATING' WHERE table_name = 't'") + .execute(&engine.pool) + .await + .expect("status"); + assert_eq!( + engine.scan_table_size(ACCOUNT, "t").await.expect("scan"), + None + ); + MetadataEngine::refresh_table_size(&engine, ACCOUNT, "t") + .await + .expect("refresh"); + assert_eq!(stored_counts(&engine, "t").await, (0, 0)); + } + + /// Storing the counts is a write, so the refresh waits for the lock before + /// it stores them (D1). + #[tokio::test] + async fn refresh_takes_the_write_lock_to_store_the_counts() { + let engine = engine().await; + table_with_items(&engine, "t", 3).await; + let guard = engine.write_lock.lock().await; + let e = engine.clone(); + let task = + tokio::spawn(async move { MetadataEngine::refresh_table_size(&e, ACCOUNT, "t").await }); + tokio::time::sleep(Duration::from_millis(200)).await; + assert!( + !task.is_finished(), + "the refresh stored its counts while the write lock was held" + ); + assert_eq!(stored_counts(&engine, "t").await, (0, 0)); + drop(guard); + tokio::time::timeout(Duration::from_secs(10), task) + .await + .expect("the refresh did not finish after the lock was released") + .expect("task") + .expect("refresh"); + assert_eq!(stored_counts(&engine, "t").await.0, 3); + } + + /// Counts scanned from a table that was then deleted, or deleted and + /// created again under the same name, are not stored on the name's row. + #[tokio::test] + async fn counts_for_another_table_id_are_not_stored() { + let engine = engine().await; + let table_id = table_with_items(&engine, "t", 3).await; + let _writer = engine.write_lock.lock().await; + engine + .store_table_size(ACCOUNT, "t", "another-table-id", 7, 70) + .await + .expect("store"); + assert_eq!(stored_counts(&engine, "t").await, (0, 0)); + engine + .store_table_size(ACCOUNT, "t", &table_id, 7, 70) + .await + .expect("store"); + assert_eq!(stored_counts(&engine, "t").await, (7, 70)); + } +}