DB index backfill checkpointing (#42012)

Adds checkpointing to database index backfills so they are resumable (restarting the backend won't erase backfill progress).

We store the `BackfillCursor` in the index backfills table to avoid OCCs in the `_index` table, which only allows one write at a time. This is not consistent with the way search index checkpoints are stored, but we should migrate those out of the `_index` table eventually for the same reason.

GitOrigin-RevId: 1a0b21ac46b3edd1dd52da8a37e1ef173538d235

Author: emmaling27

crates/database/src/bootstrap_model/index_backfills/mod.rs

@@ -11,6 +11,7 @@ use common::{
     runtime::Runtime,
     types::IndexId,
 };
+use sync_types::Timestamp;
 use value::{
     DeveloperDocumentId,
     FieldPath,
@@ -21,7 +22,10 @@ use value::{
 };
 
 use crate::{
-    bootstrap_model::index_backfills::types::IndexBackfillMetadata,
+    bootstrap_model::index_backfills::types::{
+        BackfillCursor,
+        IndexBackfillMetadata,
+    },
     system_tables::{
         SystemIndex,
         SystemTable,
@@ -88,14 +92,33 @@ impl<'a, RT: Runtime> IndexBackfillModel<'a, RT> {
             .await
     }
 
+    pub async fn initialize_search_index_backfill(
+        &mut self,
+        index_id: IndexId,
+        total_docs: Option<u64>,
+    ) -> anyhow::Result<ResolvedDocumentId> {
+        self.initialize_backfill(index_id, total_docs, None).await
+    }
+
+    pub async fn initialize_database_index_backfill(
+        &mut self,
+        index_id: IndexId,
+        total_docs: Option<u64>,
+        snapshot_ts: Timestamp,
+    ) -> anyhow::Result<ResolvedDocumentId> {
+        self.initialize_backfill(index_id, total_docs, Some(snapshot_ts))
+            .await
+    }
+
     /// Creates a new index backfill entry or reset existing index backfill
     /// entry with 0 progress and the total number of documents, if available.
     /// total_docs may not be available if table summaries have not yet
     /// bootstrapped. We're ok to update it later (which will be approximate).
-    pub async fn initialize_backfill(
+    async fn initialize_backfill(
         &mut self,
         index_id: IndexId,
         total_docs: Option<u64>,
+        snapshot_ts: Option<Timestamp>,
     ) -> anyhow::Result<ResolvedDocumentId> {
         let index_id = self.index_id_as_developer_id(index_id);
         tracing::info!(
@@ -108,6 +131,10 @@ impl<'a, RT: Runtime> IndexBackfillModel<'a, RT> {
             index_id,
             num_docs_indexed: 0,
             total_docs,
+            cursor: snapshot_ts.map(|ts| BackfillCursor {
+                snapshot_ts: ts,
+                cursor: None,
+            }),
         };
         if let Some(existing_backfill_metadata) = maybe_existing_backfill_metadata {
             system_model
@@ -124,28 +151,63 @@ impl<'a, RT: Runtime> IndexBackfillModel<'a, RT> {
         }
     }
 
+    pub async fn update_search_index_backfill_progress(
+        &mut self,
+        index_id: IndexId,
+        tablet_id: TabletId,
+        num_docs_indexed: u64,
+    ) -> anyhow::Result<()> {
+        self.update_index_backfill_progress(index_id, tablet_id, num_docs_indexed, None)
+            .await
+    }
+
+    pub async fn update_database_index_backfill_progress(
+        &mut self,
+        index_id: IndexId,
+        tablet_id: TabletId,
+        num_docs_indexed: u64,
+        cursor: ResolvedDocumentId,
+    ) -> anyhow::Result<()> {
+        self.update_index_backfill_progress(
+            index_id,
+            tablet_id,
+            num_docs_indexed,
+            Some(cursor.developer_id),
+        )
+        .await
+    }
+
     /// Upserts progress on index backfills. Only call this during the phase of
     /// the backfill where we walk a snapshot of a table, not the catching up
     /// phase where we walk the revision stream. These metrics don't make sense
     /// in the context of the revision stream.
     /// num_docs_indexed is the number of additional documents indexed since the
     /// last call.
-    pub async fn update_index_backfill_progress(
+    async fn update_index_backfill_progress(
         &mut self,
         index_id: IndexId,
         tablet_id: TabletId,
         num_docs_indexed: u64,
+        cursor: Option<DeveloperDocumentId>,
     ) -> anyhow::Result<()> {
         let index_id = self.index_id_as_developer_id(index_id);
         let maybe_existing_backfill_metadata = self.existing_backfill_metadata(index_id).await?;
         let Some(existing_backfill_metadata) = maybe_existing_backfill_metadata else {
             anyhow::bail!("Index backfill not found for index {}", index_id);
         };
+        let cursor = existing_backfill_metadata
+            .cursor
+            .as_ref()
+            .map(|c| BackfillCursor {
+                snapshot_ts: c.snapshot_ts,
+                cursor,
+            });
         if let Some(total_docs) = existing_backfill_metadata.total_docs {
             let new_backfill_metadata = IndexBackfillMetadata {
                 index_id,
                 num_docs_indexed: existing_backfill_metadata.num_docs_indexed + num_docs_indexed,
                 total_docs: Some(total_docs),
+                cursor,
             };
             SystemMetadataModel::new_global(self.tx)
                 .replace(
@@ -162,8 +224,10 @@ impl<'a, RT: Runtime> IndexBackfillModel<'a, RT> {
                 // still bootstrapping when the backfill began.
                 let new_backfill_metadata = IndexBackfillMetadata {
                     index_id,
-                    num_docs_indexed,
+                    num_docs_indexed: existing_backfill_metadata.num_docs_indexed
+                        + num_docs_indexed,
                     total_docs: Some(count),
+                    cursor,
                 };
                 SystemMetadataModel::new_global(self.tx)
                     .replace(

crates/database/src/bootstrap_model/index_backfills/tests.rs

@@ -42,7 +42,7 @@ async fn test_initialize_backfill_creates_new_entry(rt: TestRuntime) -> anyhow::
     let total_docs = Some(1000u64);
 
     let backfill_id = IndexBackfillModel::new(&mut tx)
-        .initialize_backfill(index_id, total_docs)
+        .initialize_backfill(index_id, total_docs, None)
         .await?;
 
     // Verify the backfill was created
@@ -65,7 +65,7 @@ async fn test_initialize_backfill_with_none_total_docs(rt: TestRuntime) -> anyho
     let total_docs = None;
 
     let backfill_id = IndexBackfillModel::new(&mut tx)
-        .initialize_backfill(index_id, total_docs)
+        .initialize_backfill(index_id, total_docs, None)
         .await?;
 
     // Verify the backfill was created with None total_docs
@@ -88,12 +88,12 @@ async fn test_initialize_backfill_resets_existing_entry(rt: TestRuntime) -> anyh
 
     // Create initial backfill
     let first_backfill_id = IndexBackfillModel::new(&mut tx)
-        .initialize_backfill(index_id, Some(500))
+        .initialize_backfill(index_id, Some(500), None)
         .await?;
 
     // Update progress
     IndexBackfillModel::new(&mut tx)
-        .update_index_backfill_progress(index_id, create_test_tablet_id(), 100)
+        .update_index_backfill_progress(index_id, create_test_tablet_id(), 100, None)
         .await?;
 
     // Verify progress was updated
@@ -104,7 +104,7 @@ async fn test_initialize_backfill_resets_existing_entry(rt: TestRuntime) -> anyh
 
     // Initialize again with different total_docs - should reset progress
     let second_backfill_id = IndexBackfillModel::new(&mut tx)
-        .initialize_backfill(index_id, Some(1000))
+        .initialize_backfill(index_id, Some(1000), None)
         .await?;
 
     // Should return the same ID
@@ -132,12 +132,12 @@ async fn test_update_index_backfill_progress_with_total_docs(
 
     // Initialize backfill
     let backfill_id = IndexBackfillModel::new(&mut tx)
-        .initialize_backfill(index_id, total_docs)
+        .initialize_backfill(index_id, total_docs, None)
         .await?;
 
     // Update progress
     IndexBackfillModel::new(&mut tx)
-        .update_index_backfill_progress(index_id, tablet_id, 250)
+        .update_index_backfill_progress(index_id, tablet_id, 250, None)
         .await?;
 
     // Verify progress was updated
@@ -149,7 +149,7 @@ async fn test_update_index_backfill_progress_with_total_docs(
 
     // Update progress again
     IndexBackfillModel::new(&mut tx)
-        .update_index_backfill_progress(index_id, tablet_id, 150)
+        .update_index_backfill_progress(index_id, tablet_id, 150, None)
         .await?;
 
     // Verify progress was accumulated
@@ -173,7 +173,7 @@ async fn test_update_index_backfill_progress_nonexistent_backfill(
 
     // Try to update progress for non-existent backfill
     let result = IndexBackfillModel::new(&mut tx)
-        .update_index_backfill_progress(index_id, tablet_id, 100)
+        .update_index_backfill_progress(index_id, tablet_id, 100, None)
         .await;
 
     // Should return an error
@@ -200,13 +200,13 @@ async fn test_update_index_backfill_progress_with_none_total_docs(
         .tablet_id;
     // Initialize backfill without total_docs
     let backfill_id = IndexBackfillModel::new(&mut tx)
-        .initialize_backfill(index_id, None)
+        .initialize_backfill(index_id, None, None)
         .await?;
 
     // In a test environment, table mapping for a non-existent tablet will fail,
     // so this tests that the method handles the error gracefully
     IndexBackfillModel::new(&mut tx)
-        .update_index_backfill_progress(index_id, tablet_id, 100)
+        .update_index_backfill_progress(index_id, tablet_id, 100, None)
         .await?;
 
     let backfill_doc = tx.get(backfill_id).await?;
@@ -226,7 +226,7 @@ async fn test_delete_index_backfill_existing(rt: TestRuntime) -> anyhow::Result<
 
     // Initialize backfill
     let backfill_id = IndexBackfillModel::new(&mut tx)
-        .initialize_backfill(index_id, total_docs)
+        .initialize_backfill(index_id, total_docs, None)
         .await?;
 
     // Verify it exists
@@ -279,15 +279,19 @@ async fn test_multiple_backfills_different_indexes(rt: TestRuntime) -> anyhow::R
 
     // Initialize backfills for different indexes
     let mut model = IndexBackfillModel::new(&mut tx);
-    let backfill_id1 = model.initialize_backfill(index_id1, Some(1000)).await?;
-    let backfill_id2 = model.initialize_backfill(index_id2, Some(2000)).await?;
+    let backfill_id1 = model
+        .initialize_backfill(index_id1, Some(1000), None)
+        .await?;
+    let backfill_id2 = model
+        .initialize_backfill(index_id2, Some(2000), None)
+        .await?;
 
     // Update progress for both
     model
-        .update_index_backfill_progress(index_id1, tablet_id, 100)
+        .update_index_backfill_progress(index_id1, tablet_id, 100, None)
         .await?;
     model
-        .update_index_backfill_progress(index_id2, tablet_id, 200)
+        .update_index_backfill_progress(index_id2, tablet_id, 200, None)
         .await?;
 
     // Verify both backfills exist with correct progress

crates/database/src/bootstrap_model/index_backfills/types.rs

@@ -2,11 +2,20 @@ use serde::{
     Deserialize,
     Serialize,
 };
+use sync_types::Timestamp;
 use value::{
     codegen_convex_serialization,
     DeveloperDocumentId,
 };
 
+/// Cursor for a database index that has an in-progress backfill
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+#[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
+pub struct BackfillCursor {
+    pub snapshot_ts: Timestamp,
+    pub cursor: Option<DeveloperDocumentId>,
+}
+
 /// Metadata for tracking index backfill progress.
 ///
 /// This structure stores the progress of an index backfill operation,
@@ -29,6 +38,36 @@ pub struct IndexBackfillMetadata {
     /// (does not include documents written since the backfill began)
     /// This field is None if there is no table summary available.
     pub total_docs: Option<u64>,
+    /// We only track the backfill cursor for database indexes because search
+    /// index backfill state is stored in the _index table.
+    pub cursor: Option<BackfillCursor>,
+}
+
+#[derive(Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SerializedBackfillCursor {
+    snapshot_ts: i64,
+    cursor: Option<String>,
+}
+
+impl From<BackfillCursor> for SerializedBackfillCursor {
+    fn from(cursor: BackfillCursor) -> Self {
+        SerializedBackfillCursor {
+            snapshot_ts: cursor.snapshot_ts.into(),
+            cursor: cursor.cursor.map(|id| id.to_string()),
+        }
+    }
+}
+
+impl TryFrom<SerializedBackfillCursor> for BackfillCursor {
+    type Error = anyhow::Error;
+
+    fn try_from(cursor: SerializedBackfillCursor) -> Result<Self, Self::Error> {
+        Ok(BackfillCursor {
+            snapshot_ts: cursor.snapshot_ts.try_into()?,
+            cursor: cursor.cursor.map(|id| id.parse()).transpose()?,
+        })
+    }
 }
 
 #[derive(Serialize, Deserialize)]
@@ -37,6 +76,7 @@ pub struct SerializedIndexBackfillMetadata {
     index_id: String,
     num_docs_indexed: i64,
     total_docs: Option<i64>,
+    cursor: Option<SerializedBackfillCursor>,
 }
 
 impl From<IndexBackfillMetadata> for SerializedIndexBackfillMetadata {
@@ -45,6 +85,7 @@ impl From<IndexBackfillMetadata> for SerializedIndexBackfillMetadata {
             index_id: metadata.index_id.to_string(),
             num_docs_indexed: metadata.num_docs_indexed as i64,
             total_docs: metadata.total_docs.map(|v| v as i64),
+            cursor: metadata.cursor.map(|cursor| cursor.into()),
         }
     }
 }
@@ -57,6 +98,7 @@ impl TryFrom<SerializedIndexBackfillMetadata> for IndexBackfillMetadata {
             index_id: serialized.index_id.parse()?,
             num_docs_indexed: serialized.num_docs_indexed as u64,
             total_docs: serialized.total_docs.map(|v| v as u64),
+            cursor: serialized.cursor.map(|c| c.try_into()).transpose()?,
         })
     }
 }

crates/database/src/database.rs

@@ -1707,6 +1707,10 @@ impl<RT: Runtime> Database<RT> {
         Ok(snapshot)
     }
 
+    pub fn latest_ts_and_snapshot(&self) -> anyhow::Result<(RepeatableTimestamp, Snapshot)> {
+        Ok(self.snapshot_manager.lock().latest())
+    }
+
     pub fn latest_database_snapshot(&self) -> anyhow::Result<DatabaseSnapshot<RT>> {
         let (ts, snapshot) = self.snapshot_manager.lock().latest();
         let repeatable_persistence =