Add delete_tablet_documents method to Persistence trait (#43255)

Adds `delete_tablet_documents` method to `Persistence` trait. Soon, we'll call this from the retention worker to delete documents in deleted tables.

GitOrigin-RevId: 9d63c72e4a09d61e5c5c0200f741930d9a121e7c

Author: emmaling27

crates/common/src/knobs.rs

@@ -513,6 +513,10 @@ pub static DOCUMENT_RETENTION_RATE_LIMIT: LazyLock<NonZeroU32> = LazyLock::new(|
 pub static DOCUMENT_RETENTION_MAX_SCANNED_DOCUMENTS: LazyLock<usize> =
     LazyLock::new(|| env_config("DOCUMENT_RETENTION_MAX_SCANNED_DOCUMENTS", 10000));
 
+/// Chunk size for SQL queries deleting documents from Deleting tables.
+pub static DELETE_TABLET_CHUNK_SIZE: LazyLock<u16> =
+    LazyLock::new(|| env_config("DELETE_TABLET_CHUNK_SIZE", 128));
+
 /// Size at which a search index will be queued for snapshotting.
 pub static SEARCH_INDEX_SIZE_SOFT_LIMIT: LazyLock<usize> =
     LazyLock::new(|| env_config("SEARCH_INDEX_SIZE_SOFT_LIMIT", 10 * (1 << 20))); // 10 MiB

crates/common/src/persistence.rs

@@ -237,6 +237,12 @@ pub trait Persistence: Sync + Send + 'static {
         documents: Vec<(Timestamp, InternalDocumentId)>,
     ) -> anyhow::Result<usize>;
 
+    async fn delete_tablet_documents(
+        &self,
+        tablet_id: TabletId,
+        chunk_size: usize,
+    ) -> anyhow::Result<usize>;
+
     // No-op by default. Persistence implementation can override.
     async fn shutdown(&self) -> anyhow::Result<()> {
         Ok(())

crates/common/src/testing/persistence_test_suite.rs

@@ -53,6 +53,7 @@ use crate::{
         Interval,
         StartIncluded,
     },
+    knobs::DELETE_TABLET_CHUNK_SIZE,
     persistence::{
         fake_retention_validator::FakeRetentionValidator,
         ConflictStrategy,
@@ -226,6 +227,14 @@ macro_rules! run_persistence_test_suite {
                 .await
         }
 
+        #[tokio::test]
+        async fn test_persistence_delete_tablet_documents() -> anyhow::Result<()> {
+            let $db = $create_db;
+            let p = $create_persistence;
+            persistence_test_suite::persistence_delete_tablet_documents(::std::sync::Arc::new(p))
+                .await
+        }
+
         #[tokio::test]
         async fn test_persistence_previous_revisions_of_documents() -> anyhow::Result<()> {
             let $db = $create_db;
@@ -1778,6 +1787,38 @@ pub async fn persistence_delete_many_documents<P: Persistence>(p: Arc<P>) -> any
     Ok(())
 }
 
+pub async fn persistence_delete_tablet_documents<P: Persistence>(p: Arc<P>) -> anyhow::Result<()> {
+    let mut id_generator = TestIdGenerator::new();
+    let table: TableName = str::parse("table")?;
+
+    const NUM_DOCS: usize = 256;
+    let ids: Vec<_> = (0..NUM_DOCS)
+        .map(|_| id_generator.user_generate(&table))
+        .collect();
+    let tablet_id = ids[0].tablet_id;
+    let documents: Vec<_> = ids
+        .iter()
+        .enumerate()
+        .map(|(i, &id)| doc(id, i as i32 + 1, Some(i as i64), None).unwrap())
+        .collect();
+
+    p.write(&documents, &[], ConflictStrategy::Error).await?;
+
+    let num_deleted = p
+        .delete_tablet_documents(tablet_id, *DELETE_TABLET_CHUNK_SIZE as usize)
+        .await?;
+    assert_eq!(num_deleted, *DELETE_TABLET_CHUNK_SIZE as usize);
+
+    let reader = p.reader();
+    let stream = reader.load_all_documents();
+    let all_docs = stream.try_collect::<Vec<_>>().await?;
+    assert_eq!(
+        all_docs.len(),
+        NUM_DOCS - *DELETE_TABLET_CHUNK_SIZE as usize
+    );
+    Ok(())
+}
+
 pub async fn persistence_previous_revisions_of_documents<P: Persistence>(
     p: Arc<P>,
 ) -> anyhow::Result<()> {

crates/common/src/testing/test_persistence.rs

@@ -204,6 +204,25 @@ impl Persistence for TestPersistence {
         }
         Ok(total_deleted)
     }
+
+    async fn delete_tablet_documents(
+        &self,
+        tablet_id: TabletId,
+        chunk_size: usize,
+    ) -> anyhow::Result<usize> {
+        let mut inner = self.inner.lock();
+        let log = &mut inner.log;
+        let log_len = log.len();
+        let mut num_deleted = 0;
+        log.retain(|(_ts, id), _| {
+            if num_deleted >= chunk_size {
+                return true;
+            }
+            num_deleted += 1;
+            id.table() != tablet_id
+        });
+        Ok(log_len - log.len())
+    }
 }
 
 #[async_trait]

crates/mysql/src/lib.rs

@@ -598,6 +598,32 @@ impl<RT: Runtime> Persistence for MySqlPersistence<RT> {
             })
             .await
     }
+
+    async fn delete_tablet_documents(
+        &self,
+        tablet_id: TabletId,
+        chunk_size: usize,
+    ) -> anyhow::Result<usize> {
+        let multitenant = self.multitenant;
+        let instance_name = mysql_async::Value::from(&self.instance_name.raw);
+        self.lease
+            .transact(async move |tx| {
+                let mut deleted_count = 0;
+                let mut params =
+                    Vec::with_capacity(sql::DELETE_TABLE_COLUMN_COUNT + (multitenant as usize));
+                let tablet_id: Vec<u8> = tablet_id.0.into();
+                params.push(tablet_id.into());
+                if multitenant {
+                    params.push(instance_name.clone());
+                }
+                params.push(chunk_size.into());
+                deleted_count += tx
+                    .exec_iter(sql::delete_tablet_chunk(multitenant), params)
+                    .await?;
+                Ok(deleted_count as usize)
+            })
+            .await
+    }
 }
 
 #[derive(Clone)]

crates/mysql/src/sql.rs

@@ -541,6 +541,28 @@ pub fn delete_document_chunk(chunk_size: usize, multitenant: bool) -> &'static s
         .unwrap()
 }
 
+pub const DELETE_TABLE_COLUMN_COUNT: usize = 2;
+static DELETE_TABLET_CHUNK_QUERIES: LazyLock<HashMap<bool, String>> = LazyLock::new(|| {
+    [false, true]
+        .into_iter()
+        .map(move |multitenant| {
+            let where_clause = if multitenant {
+                "(table_id = ? AND instance_name = ?)"
+            } else {
+                "(table_id = ?)"
+            };
+            (
+                multitenant,
+                format!("DELETE FROM @db_name.documents WHERE {where_clause} LIMIT ?",),
+            )
+        })
+        .collect()
+});
+
+pub fn delete_tablet_chunk(multitenant: bool) -> &'static str {
+    DELETE_TABLET_CHUNK_QUERIES.get(&multitenant).unwrap()
+}
+
 pub const fn write_persistence_global(multitenant: bool) -> &'static str {
     tableify!(
         multitenant,

crates/postgres/src/lib.rs

@@ -724,6 +724,29 @@ impl Persistence for PostgresPersistence {
             .await
     }
 
+    async fn delete_tablet_documents(
+        &self,
+        tablet_id: TabletId,
+        chunk_size: usize,
+    ) -> anyhow::Result<usize> {
+        let multitenant = self.multitenant;
+        let instance_name = self.instance_name.clone();
+        self.lease
+            .transact(async move |tx| {
+                let delete_doc = tx
+                    .prepare_cached(sql::delete_tablet_chunk(multitenant))
+                    .await?;
+                let mut params = vec![Param::Bytes(tablet_id.0.into())];
+                if multitenant {
+                    params.push(Param::Text(instance_name.to_string()));
+                }
+                params.push(Param::Limit(chunk_size as i64));
+                let deleted_count = tx.execute_raw(&delete_doc, params).await?;
+                Ok(deleted_count as usize)
+            })
+            .await
+    }
+
     async fn import_documents_batch(
         &self,
         mut documents: BoxStream<'_, Vec<DocumentLogEntry>>,

crates/postgres/src/sql.rs

@@ -629,6 +629,26 @@ DELETE FROM @db_name.documents WHERE
     )
 }
 
+pub const fn delete_tablet_chunk(multitenant: bool) -> &'static str {
+    tableify!(multitenant, {
+        const LIMIT_PARAM: i32 = if multitenant { 3 } else { 2 };
+        formatcp!(
+            r#"DELETE FROM @db_name.documents WHERE table_id = $1 AND id IN (
+            SELECT id FROM @db_name.documents
+            WHERE table_id = $1{instance_clause}
+            LIMIT ${limit_param}
+        ){instance_clause}
+"#,
+            instance_clause = if multitenant {
+                " AND instance_name = $2"
+            } else {
+                ""
+            },
+            limit_param = LIMIT_PARAM,
+        )
+    })
+}
+
 pub const fn write_persistence_global(multitenant: bool) -> &'static str {
     tableify!(
         multitenant,

crates/sqlite/src/lib.rs

@@ -421,6 +421,24 @@ impl Persistence for SqlitePersistence {
         tx.commit()?;
         Ok(count_deleted)
     }
+
+    async fn delete_tablet_documents(
+        &self,
+        tablet_id: TabletId,
+        chunk_size: usize,
+    ) -> anyhow::Result<usize> {
+        let mut inner = self.inner.lock();
+        let tx = inner.connection.transaction()?;
+        let mut delete_table_documents_query = tx.prepare_cached(DELETE_TABLE_DOCUMENTS)?;
+        let count_deleted = delete_table_documents_query.execute(params![
+            &tablet_id.0[..],
+            &tablet_id.0[..],
+            chunk_size,
+        ])?;
+        drop(delete_table_documents_query);
+        tx.commit()?;
+        Ok(count_deleted)
+    }
 }
 
 #[async_trait]
@@ -682,6 +700,9 @@ const DELETE_INDEX: &str = "DELETE FROM indexes WHERE index_id = ? AND ts <= ? A
 
 const DELETE_DOCUMENT: &str = "DELETE FROM documents WHERE table_id = ? AND id = ? AND ts <= ?";
 
+const DELETE_TABLE_DOCUMENTS: &str = "DELETE FROM documents WHERE table_id = ? AND id IN (SELECT \
+                                      id FROM documents WHERE table_id = ? LIMIT ?)";
+
 const PREV_REV_QUERY: &str = r#"
 SELECT id, ts, table_id, json_value, deleted, prev_ts
 FROM documents