[cluster migration] run retention during migration (#24504)

GitOrigin-RevId: 9146f965fb2009492d0e06b04e2fcc105e6e840e

Author: ldanilek

crates/database/src/index_worker.rs

@@ -925,18 +925,15 @@ impl<RT: Runtime> IndexWriter<RT> {
         let min_snapshot_ts = self.retention_validator.min_snapshot_ts().await?;
         let all_indexes = btreemap! { index_id => (index_name, indexed_fields) };
         // TODO(lee) add checkpointing.
-        let mut cursor_ts = backfill_begin_ts;
-        while cursor_ts.succ()? < min_snapshot_ts {
-            (cursor_ts, _) = LeaderRetentionManager::delete(
-                min_snapshot_ts,
-                self.persistence.clone(),
-                &self.runtime,
-                cursor_ts,
-                &all_indexes,
-                self.retention_validator.clone(),
-            )
-            .await?;
-        }
+        LeaderRetentionManager::delete_all_no_checkpoint(
+            backfill_begin_ts,
+            min_snapshot_ts,
+            self.persistence.clone(),
+            &self.runtime,
+            &all_indexes,
+            self.retention_validator.clone(),
+        )
+        .await?;
         Ok(())
     }
 }

crates/database/src/lib.rs

@@ -153,6 +153,8 @@ pub use self::{
     retention::{
         latest_retention_min_snapshot_ts,
         FollowerRetentionManager,
+        LeaderRetentionManager,
+        RetentionType,
     },
     snapshot_manager::{
         Snapshot,

crates/database/src/retention.rs

@@ -623,7 +623,7 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
     /// entries is the number of index entries we found were expired, not
     /// necessarily the total we deleted or wanted to delete, though they're
     /// correlated.
-    pub(crate) async fn delete(
+    async fn delete(
         min_snapshot_ts: Timestamp,
         persistence: Arc<dyn Persistence>,
         rt: &RT,
@@ -690,6 +690,32 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
             .map(|timestamp| (timestamp, total_expired_entries))
     }
 
+    pub async fn delete_all_no_checkpoint(
+        mut cursor_ts: Timestamp,
+        min_snapshot_ts: Timestamp,
+        persistence: Arc<dyn Persistence>,
+        rt: &RT,
+        all_indexes: &BTreeMap<IndexId, (GenericIndexName<TableId>, IndexedFields)>,
+        retention_validator: Arc<dyn RetentionValidator>,
+    ) -> anyhow::Result<()> {
+        while cursor_ts.succ()? < min_snapshot_ts {
+            let (new_cursor_ts, _) = Self::delete(
+                min_snapshot_ts,
+                persistence.clone(),
+                rt,
+                cursor_ts,
+                all_indexes,
+                retention_validator.clone(),
+            )
+            .await?;
+            tracing::info!(
+                "custom index retention completed between ts {cursor_ts} and {new_cursor_ts}"
+            );
+            cursor_ts = new_cursor_ts;
+        }
+        Ok(())
+    }
+
     #[try_stream(ok = (Timestamp, InternalDocumentId), error = anyhow::Error)]
     async fn expired_documents(
         rt: &RT,
@@ -1222,9 +1248,8 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
         Ok(())
     }
 
-    async fn get_checkpoint(
+    pub async fn get_checkpoint_no_logging(
         persistence: &dyn PersistenceReader,
-        snapshot_reader: Reader<SnapshotManager>,
         retention_type: RetentionType,
     ) -> anyhow::Result<Timestamp> {
         let key = match retention_type {
@@ -1239,24 +1264,30 @@ impl<RT: Runtime> LeaderRetentionManager<RT> {
             .map(ConvexValue::try_from)
             .transpose()?;
         let checkpoint = match checkpoint_value {
-            Some(ConvexValue::Int64(ts)) => {
-                let checkpoint = Timestamp::try_from(ts)?;
-                match retention_type {
-                    RetentionType::Document => log_document_retention_cursor_age(
-                        (*snapshot_reader.lock().latest_ts()).secs_since_f64(checkpoint),
-                    ),
-                    RetentionType::Index => log_retention_cursor_age(
-                        (*snapshot_reader.lock().latest_ts()).secs_since_f64(checkpoint),
-                    ),
-                }
-                checkpoint
-            },
+            Some(ConvexValue::Int64(ts)) => Timestamp::try_from(ts)?,
             None => Timestamp::MIN,
             _ => anyhow::bail!("invalid retention checkpoint {checkpoint_value:?}"),
         };
         Ok(checkpoint)
     }
 
+    async fn get_checkpoint(
+        persistence: &dyn PersistenceReader,
+        snapshot_reader: Reader<SnapshotManager>,
+        retention_type: RetentionType,
+    ) -> anyhow::Result<Timestamp> {
+        let checkpoint = Self::get_checkpoint_no_logging(persistence, retention_type).await?;
+        match retention_type {
+            RetentionType::Document => log_document_retention_cursor_age(
+                (*snapshot_reader.lock().latest_ts()).secs_since_f64(checkpoint),
+            ),
+            RetentionType::Index => log_retention_cursor_age(
+                (*snapshot_reader.lock().latest_ts()).secs_since_f64(checkpoint),
+            ),
+        }
+        Ok(checkpoint)
+    }
+
     fn accumulate_index_document(
         maybe_doc: Option<ResolvedDocument>,
         all_indexes: &mut BTreeMap<IndexId, (GenericIndexName<TableId>, IndexedFields)>,

crates/indexing/src/index_registry.rs

@@ -12,6 +12,7 @@ use common::{
         database_index::{
             DatabaseIndexState,
             DeveloperDatabaseIndexConfig,
+            IndexedFields,
         },
         DeveloperIndexConfig,
         IndexConfig,
@@ -405,6 +406,23 @@ impl IndexRegistry {
             .map(|index| index.metadata())
     }
 
+    pub fn all_database_index_configs(
+        &self,
+    ) -> BTreeMap<IndexId, (TabletIndexName, IndexedFields)> {
+        self.all_indexes()
+            .filter_map(|index| {
+                let index_id = index.id().internal_id();
+                let index_name = index.name.clone();
+                match &index.config {
+                    IndexConfig::Database {
+                        developer_config, ..
+                    } => Some((index_id, (index_name, developer_config.fields.clone()))),
+                    IndexConfig::Search { .. } | IndexConfig::Vector { .. } => None,
+                }
+            })
+            .collect()
+    }
+
     pub fn all_enabled_indexes(&self) -> Vec<ParsedDocument<TabletIndexMetadata>> {
         self.enabled_indexes
             .values()