fix(sqlite): scan a table's size without holding the write lock

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<Option<(String, i64, i64)>, 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));
+    }
+}