Use pruning interval as control plane cooldown

Author: rdettai

quickwit/quickwit-control-plane/src/control_plane.rs

@@ -73,7 +73,7 @@ pub(crate) const CONTROL_PLAN_LOOP_INTERVAL: Duration = if cfg!(any(test, featur
 };
 
 /// Minimum period between two identical shard pruning operations.
-const PRUNE_SHARDS_COOLDOWN_PERIOD: Duration = Duration::from_secs(120);
+const PRUNE_SHARDS_DEFAULT_COOLDOWN_PERIOD: Duration = Duration::from_secs(120);
 
 /// Minimum period between two rebuild plan operations.
 const REBUILD_PLAN_COOLDOWN_PERIOD: Duration = Duration::from_secs(2);
@@ -796,19 +796,27 @@ impl Handler<DebouncedPruneShardsRequest> for ControlPlane {
             return Ok(Ok(EmptyResponse {}));
         }
 
-        self.prune_shards_debouncers
+        let debouncer = self
+            .prune_shards_debouncers
             .entry((
                 metastore_request.index_uid().index_id.clone(),
                 metastore_request.source_id.clone(),
             ))
-            .or_insert_with(|| Debouncer::new(PRUNE_SHARDS_COOLDOWN_PERIOD))
-            .self_send_with_cooldown(
-                ctx,
-                DebouncedPruneShardsRequest {
-                    request: Some(metastore_request),
-                    execute_now: true,
-                },
-            );
+            .or_insert_with(|| Debouncer::new(PRUNE_SHARDS_DEFAULT_COOLDOWN_PERIOD));
+
+        // Update the cooldown period in case the config changes. If multiple
+        // indexers run with different configs, the control plane ends up
+        // running one of them in a seemingly arbitrary way.
+        if let Some(interval) = metastore_request.interval {
+            debouncer.set_next_cooldown_period(Duration::from_secs(interval as u64));
+        }
+        debouncer.self_send_with_cooldown(
+            ctx,
+            DebouncedPruneShardsRequest {
+                request: Some(metastore_request),
+                execute_now: true,
+            },
+        );
         Ok(Ok(EmptyResponse {}))
     }
 }

quickwit/quickwit-control-plane/src/debouncer.rs

@@ -132,6 +132,10 @@ impl Debouncer {
             DebouncerState::CooldownNotScheduled | DebouncerState::CooldownScheduled => {}
         }
     }
+
+    pub fn set_next_cooldown_period(&mut self, period: Duration) {
+        self.cooldown_period = period;
+    }
 }
 
 #[cfg(test)]
@@ -163,6 +167,11 @@ mod tests {
     #[derive(Debug)]
     struct DebouncedIncrement;
 
+    #[derive(Debug, Default, Clone)]
+    struct UpdateCooldownPeriod {
+        cooldown_period: Duration,
+    }
+
     #[async_trait]
     impl Actor for DebouncingActor {
         type ObservableState = usize;
@@ -200,10 +209,25 @@ mod tests {
         }
     }
 
+    #[async_trait]
+    impl Handler<UpdateCooldownPeriod> for DebouncingActor {
+        type Reply = ();
+
+        async fn handle(
+            &mut self,
+            message: UpdateCooldownPeriod,
+            _ctx: &ActorContext<Self>,
+        ) -> Result<Self::Reply, ActorExitStatus> {
+            self.debouncer
+                .set_next_cooldown_period(message.cooldown_period);
+            Ok(())
+        }
+    }
+
     #[tokio::test]
     async fn test_debouncer() {
         let universe = Universe::default();
-        let cooldown_period = Duration::from_millis(1_000);
+        let cooldown_period = Duration::from_millis(100);
         let debouncer = DebouncingActor::new(cooldown_period);
         let (debouncer_mailbox, debouncer_handle) = universe.spawn_builder().spawn(debouncer);
         {
@@ -237,4 +261,51 @@ mod tests {
         }
         universe.assert_quit().await;
     }
+
+    #[tokio::test]
+    async fn test_debouncer_cooldown_update() {
+        let universe = Universe::default();
+        let cooldown_period = Duration::from_millis(100);
+        let debouncer = DebouncingActor::new(cooldown_period);
+        let (debouncer_mailbox, debouncer_handle) = universe.spawn_builder().spawn(debouncer);
+        for _ in 0..10 {
+            debouncer_mailbox.ask(DebouncedIncrement).await.unwrap();
+            let count = *debouncer_handle.process_pending_and_observe().await;
+            assert_eq!(count, 1);
+        }
+        {
+            universe.sleep(cooldown_period.mul_f32(1.2)).await;
+            let count = *debouncer_handle.process_pending_and_observe().await;
+            assert_eq!(count, 2);
+        }
+        // at this point, we are still cooling down
+        debouncer_mailbox
+            .ask(UpdateCooldownPeriod {
+                cooldown_period: cooldown_period.mul_f32(2.0),
+            })
+            .await
+            .unwrap();
+        {
+            universe.sleep(cooldown_period.mul_f32(1.2)).await;
+            let count = *debouncer_handle.process_pending_and_observe().await;
+            assert_eq!(count, 2);
+        }
+        // the previous cooldown now expired, the new one will apply
+        for _ in 0..10 {
+            debouncer_mailbox.ask(DebouncedIncrement).await.unwrap();
+            let count = *debouncer_handle.process_pending_and_observe().await;
+            assert_eq!(count, 3);
+        }
+        {
+            universe.sleep(cooldown_period.mul_f32(1.2)).await;
+            let count = *debouncer_handle.process_pending_and_observe().await;
+            assert_eq!(count, 3);
+        }
+        {
+            universe.sleep(cooldown_period.mul_f32(1.2)).await;
+            let count = *debouncer_handle.process_pending_and_observe().await;
+            assert_eq!(count, 4);
+        }
+        universe.assert_quit().await;
+    }
 }

quickwit/quickwit-indexing/src/source/queue_sources/shared_state.rs

@@ -100,6 +100,7 @@ impl QueueSharedState {
                         source_id: source_id.clone(),
                         max_age,
                         max_count,
+                        interval: Some(pruning_interval.as_secs() as u32),
                     })
                     .await;
                 if let Err(err) = result {

quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs

@@ -641,6 +641,7 @@ mod tests {
             source_id: source_id.clone(),
             max_age: None,
             max_count: None,
+            interval: None,
         };
         let MutationOccurred::No(()) = shards.prune_shards(request).unwrap() else {
             panic!("expected `MutationOccurred::No`");
@@ -651,6 +652,7 @@ mod tests {
             source_id: source_id.clone(),
             max_age: Some(50),
             max_count: None,
+            interval: None,
         };
         let MutationOccurred::No(()) = shards.prune_shards(request).unwrap() else {
             panic!("expected `MutationOccurred::No`");
@@ -687,6 +689,7 @@ mod tests {
             source_id: source_id.clone(),
             max_age: Some(150),
             max_count: None,
+            interval: None,
         };
         let MutationOccurred::Yes(()) = shards.prune_shards(request).unwrap() else {
             panic!("expected `MutationOccurred::Yes`");
@@ -697,6 +700,7 @@ mod tests {
             source_id: source_id.clone(),
             max_age: Some(150),
             max_count: None,
+            interval: None,
         };
         let MutationOccurred::No(()) = shards.prune_shards(request).unwrap() else {
             panic!("expected `MutationOccurred::No`");

quickwit/quickwit-metastore/src/tests/shard.rs

@@ -654,6 +654,7 @@ pub async fn test_metastore_prune_shards<
             source_id: test_index.source_id.clone(),
             max_age: None,
             max_count: None,
+            interval: None,
         };
         metastore.prune_shards(prune_index_request).await.unwrap();
         let all_shards = metastore
@@ -669,6 +670,7 @@ pub async fn test_metastore_prune_shards<
             source_id: test_index.source_id.clone(),
             max_age: Some(oldest_shard_age - 350),
             max_count: None,
+            interval: None,
         };
         metastore.prune_shards(prune_index_request).await.unwrap();
 
@@ -688,6 +690,7 @@ pub async fn test_metastore_prune_shards<
             source_id: test_index.source_id.clone(),
             max_age: None,
             max_count: Some(90),
+            interval: None,
         };
         metastore.prune_shards(prune_index_request).await.unwrap();
         let mut all_shards = metastore
@@ -705,6 +708,7 @@ pub async fn test_metastore_prune_shards<
         source_id: test_index.source_id.clone(),
         max_age: Some(oldest_shard_age - 2950),
         max_count: Some(80),
+        interval: None,
     };
     metastore.prune_shards(prune_index_request).await.unwrap();
     let all_shards = metastore
@@ -718,6 +722,7 @@ pub async fn test_metastore_prune_shards<
         source_id: test_index.source_id.clone(),
         max_age: Some(oldest_shard_age - 4000),
         max_count: Some(50),
+        interval: None,
     };
     metastore.prune_shards(prune_index_request).await.unwrap();
     let all_shards = metastore

quickwit/quickwit-proto/protos/quickwit/metastore.proto

@@ -458,9 +458,11 @@ message PruneShardsRequest {
   quickwit.common.IndexUid index_uid = 1;
   string source_id = 2;
   // The maximum age of the shards to keep, in seconds.
-  optional uint32 max_age = 5;
+  optional uint32 max_age = 3;
   // The maximum number of the shards to keep. Delete older shards first.
-  optional uint32 max_count = 6;
+  optional uint32 max_count = 4;
+  // The interval between two pruning operations, in seconds.
+  optional uint32 interval = 5;
 }
 
 message ListShardsRequest {