Spread indexing pipelines behind an environment variable

quickwit/quickwit-actors/src/scheduler.rs

@@ -60,6 +60,11 @@ enum SchedulerMessage {
         callback: Callback,
         timeout: Duration,
     },
+    #[cfg(any(test, feature = "testsuite"))]
+    ShiftTime {
+        duration: Duration,
+        done_tx: oneshot::Sender<()>,
+    },
 }
 
 #[derive(Clone)]
@@ -158,6 +163,16 @@ impl SchedulerClient {
         }
     }
 
+    #[cfg(any(test, feature = "testsuite"))]
+    pub async fn shift_time(&self, duration: Duration) {
+        let (done_tx, done_rx) = oneshot::channel();
+        let _ = self
+            .inner
+            .tx
+            .send(SchedulerMessage::ShiftTime { duration, done_tx });
+        let _ = done_rx.await;
+    }
+
     // Triggers an event, telling the Scheduler to process time,
     // checks whether some scheduled events have timed out, or whether we should
     // jump forward in time.
@@ -207,6 +222,11 @@ pub fn start_scheduler() -> SchedulerClient {
                     SchedulerMessage::Schedule { callback, timeout } => {
                         scheduler.process_schedule(callback, timeout);
                     }
+                    #[cfg(any(test, feature = "testsuite"))]
+                    SchedulerMessage::ShiftTime { duration, done_tx } => {
+                        scheduler.process_shift_time(duration);
+                        let _ = done_tx.send(());
+                    }
                 }
             }
         },
@@ -261,6 +281,12 @@ impl Scheduler {
         self.process_time();
     }
 
+    #[cfg(any(test, feature = "testsuite"))]
+    fn process_shift_time(&mut self, duration: Duration) {
+        self.simulated_time_shift += duration;
+        self.process_time();
+    }
+
     fn scheduler_client(&self) -> Option<SchedulerClient> {
         let scheduler_client = SchedulerClient {
             inner: self.weak_scheduler_client.upgrade()?,
@@ -434,4 +460,25 @@ mod tests {
         assert!(elapsed.as_millis() < 50);
         universe.assert_quit().await;
     }
+
+    #[tokio::test]
+    async fn test_scheduler_shift_time() {
+        let count: Arc<AtomicUsize> = Default::default();
+        let simple_actor = ClockActor {
+            count: count.clone(),
+        };
+        let universe = Universe::new();
+        let (_mailbox, handle) = universe.spawn_builder().spawn(simple_actor);
+        handle.process_pending_and_observe().await;
+        assert_eq!(count.load(Ordering::SeqCst), 1);
+
+        universe.shift_time(Duration::from_millis(500)).await;
+        handle.process_pending_and_observe().await;
+        assert_eq!(count.load(Ordering::SeqCst), 1);
+
+        universe.shift_time(Duration::from_millis(600)).await;
+        handle.process_pending_and_observe().await;
+        assert_eq!(count.load(Ordering::SeqCst), 2);
+        universe.assert_quit().await;
+    }
 }

quickwit/quickwit-actors/src/universe.rs

@@ -47,7 +47,7 @@ impl Universe {
         }
     }
 
-    /// Creates a universe were time is accelerated.
+    /// Creates a universe where time is accelerated.
     ///
     /// Time is accelerated in a way to exhibit a behavior as close as possible
     /// to what would have happened with normal time but faster.
@@ -112,6 +112,11 @@ impl Universe {
         self.spawn_ctx.scheduler_client.sleep(duration).await;
     }
 
+    #[cfg(any(test, feature = "testsuite"))]
+    pub async fn shift_time(&self, duration: Duration) {
+        self.spawn_ctx.scheduler_client.shift_time(duration).await;
+    }
+
     pub fn spawn_builder<A: Actor>(&self) -> SpawnBuilder<A> {
         self.spawn_ctx.spawn_builder()
     }

quickwit/quickwit-indexing/src/actors/cooperative_indexing.rs

@@ -109,6 +109,11 @@ impl CooperativeIndexingCycle {
         Duration::from_millis(initial_sleep_millis)
     }
 
+    // with cooperative indexing, the nudging is done only when the actor mailbox is empty
+    pub fn commit_timeout(&self) -> Duration {
+        self.schedule.commit_timeout()
+    }
+
     pub async fn cooperative_indexing_period(&self) -> CooperativeIndexingPeriod {
         let t_wake = Instant::now();
         let permit = Semaphore::acquire_owned(self.indexing_permits.clone())

quickwit/quickwit-indexing/src/actors/indexer.rs

@@ -16,6 +16,7 @@ use std::collections::hash_map::Entry;
 use std::num::NonZeroU32;
 use std::ops::RangeInclusive;
 use std::sync::Arc;
+use std::time::Duration;
 
 use anyhow::Context;
 use async_trait::async_trait;
@@ -47,11 +48,13 @@ use tantivy::tokenizer::TokenizerManager;
 use tantivy::{DateTime, DocId, IndexBuilder, IndexSettings};
 use tokio::runtime::Handle;
 use tokio::sync::Semaphore;
+use tokio::time::Instant;
 use tracing::{Span, debug, info_span, warn};
 use ulid::Ulid;
 
 use super::IndexSerializer;
 use super::cooperative_indexing::{CooperativeIndexingCycle, CooperativeIndexingPeriod};
+use super::pipeline_schedule::PipelineSchedule;
 use crate::docs_clustering::{DocIdClusterer, Fingerprinter};
 use crate::metrics::SPLIT_BUILDERS;
 use crate::models::{
@@ -62,11 +65,50 @@ use crate::models::{
 // Random partition ID used to gather partitions exceeding the maximum number of partitions.
 pub(crate) const OTHER_PARTITION_ID: u64 = 3264326757911759461u64;
 
+const ENABLE_SPREAD_INDEXING_PIPELINES_ENV_KEY: &str = "QW_ENABLE_SPREAD_INDEXING_PIPELINES";
+
 #[derive(Debug)]
 struct CommitTimeout {
     workbench_id: Ulid,
 }
 
+enum IndexingMode {
+    Normal(Duration),
+    Cooperative(CooperativeIndexingCycle),
+    Spread(PipelineSchedule),
+}
+
+impl IndexingMode {
+    fn new(
+        pipeline_id: &IndexingPipelineId,
+        commit_timeout: Duration,
+        cooperative_indexing_permits_opt: Option<Arc<Semaphore>>,
+        enable_spread_indexing_pipelines: bool,
+    ) -> Self {
+        if let Some(permits) = cooperative_indexing_permits_opt {
+            return Self::Cooperative(CooperativeIndexingCycle::new(
+                pipeline_id,
+                commit_timeout,
+                permits,
+            ));
+        }
+
+        if enable_spread_indexing_pipelines {
+            return Self::Spread(PipelineSchedule::new(pipeline_id, commit_timeout));
+        }
+
+        Self::Normal(commit_timeout)
+    }
+
+    fn commit_timeout(&self, now: Instant) -> Duration {
+        match self {
+            Self::Normal(commit_timeout) => *commit_timeout,
+            Self::Cooperative(cycle) => cycle.commit_timeout(),
+            Self::Spread(schedule) => schedule.nudged_commit_timeout(now),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Default, Eq, PartialEq, Serialize)]
 pub struct IndexerCounters {
     /// Number of splits that were emitted by the indexer.
@@ -100,7 +142,7 @@ struct IndexerState {
     tokenizer_manager: TokenizerManager,
     max_num_partitions: NonZeroU32,
     index_settings: IndexSettings,
-    cooperative_indexing_opt: Option<CooperativeIndexingCycle>,
+    indexing_mode: IndexingMode,
     indexing_io_throughput_limiter_opt: Option<Limiter>,
 }
 
@@ -202,7 +244,7 @@ impl IndexerState {
         );
         let indexing_span = info_span!(parent: batch_parent_span.id(), "indexer");
         let cooperative_indexing_period =
-            if let Some(cooperative_indexing) = &self.cooperative_indexing_opt {
+            if let IndexingMode::Cooperative(cooperative_indexing) = &self.indexing_mode {
                 Some(
                     ctx.protect_future(cooperative_indexing.cooperative_indexing_period())
                         .await,
@@ -261,10 +303,8 @@ impl IndexerState {
             let commit_timeout_message = CommitTimeout {
                 workbench_id: indexing_workbench.workbench_id,
             };
-            ctx.schedule_self_msg(
-                self.indexing_settings.commit_timeout(),
-                commit_timeout_message,
-            );
+            let commit_timeout = self.indexing_mode.commit_timeout(Instant::now());
+            ctx.schedule_self_msg(commit_timeout, commit_timeout_message);
             *indexing_workbench_opt = Some(indexing_workbench);
         }
         let current_indexing_workbench = indexing_workbench_opt.as_mut().context(
@@ -404,11 +444,15 @@ impl Actor for Indexer {
     }
 
     async fn initialize(&mut self, ctx: &ActorContext<Self>) -> Result<(), ActorExitStatus> {
-        if let Some(cooperative_indexing_cycle) = &self.indexer_state.cooperative_indexing_opt {
-            let initial_sleep_duration = cooperative_indexing_cycle.initial_sleep_duration();
-            ctx.pause();
-            ctx.schedule_self_msg(initial_sleep_duration, Command::Resume);
-        }
+        let initial_sleep_duration = match &self.indexer_state.indexing_mode {
+            // when spreading indexing pipelines is enabled, start indexing immediately and nudge
+            // pipelines on further commit timeout. This way, we are using all available CPU
+            // right away instead of idling on the first indexing cycle
+            IndexingMode::Normal(_) | IndexingMode::Spread(_) => return Ok(()),
+            IndexingMode::Cooperative(cycle) => cycle.initial_sleep_duration(),
+        };
+        ctx.pause();
+        ctx.schedule_self_msg(initial_sleep_duration, Command::Resume);
         Ok(())
     }
 
@@ -551,14 +595,15 @@ impl Indexer {
             // A configured fingerprinter supplies the mapping when the split is finalized.
             manual_doc_id_mapping: fingerprinter_opt.is_some(),
         };
-        let cooperative_indexing_opt: Option<CooperativeIndexingCycle> =
-            cooperative_indexing_permits_opt.map(|cooperative_indexing_permits| {
-                CooperativeIndexingCycle::new(
-                    &pipeline_id,
-                    indexing_settings.commit_timeout(),
-                    cooperative_indexing_permits,
-                )
-            });
+        let indexing_mode = IndexingMode::new(
+            &pipeline_id,
+            indexing_settings.commit_timeout(),
+            cooperative_indexing_permits_opt,
+            quickwit_common::get_bool_from_env_cached!(
+                ENABLE_SPREAD_INDEXING_PIPELINES_ENV_KEY,
+                false,
+            ),
+        );
         Self {
             indexer_state: IndexerState {
                 pipeline_id,
@@ -572,7 +617,7 @@ impl Indexer {
                 tokenizer_manager: tokenizer_manager.tantivy_manager().clone(),
                 index_settings,
                 max_num_partitions: doc_mapper.max_num_partitions(),
-                cooperative_indexing_opt,
+                indexing_mode,
                 indexing_io_throughput_limiter_opt,
             },
             index_serializer_mailbox,
@@ -715,6 +760,7 @@ mod tests {
     use tantivy::{DateTime, DocAddress, TantivyDocument, doc};
 
     use super::{IndexerCounters, record_timestamp, *};
+    use crate::actors::pipeline_schedule::NUDGE_TOLERANCE;
     use crate::docs_clustering::Fingerprint;
 
     #[allow(clippy::too_many_arguments)]
@@ -1831,4 +1877,93 @@ mod tests {
         universe.assert_quit().await;
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_indexer_spread_mode_nudges_commit_timeout() {
+        let universe = Universe::new();
+        let pipeline_id = IndexingPipelineId {
+            index_uid: IndexUid::new_with_random_ulid("test-index"),
+            source_id: "test-source".to_string(),
+            node_id: NodeId::from_str("test-node"),
+            pipeline_uid: PipelineUid::default(),
+        };
+        let doc_mapper = Arc::new(default_doc_mapper_for_test());
+        let schema = doc_mapper.schema();
+        let body_field = schema.get_field("body").unwrap();
+        let timestamp_field = schema.get_field("timestamp").unwrap();
+        let commit_timeout = Duration::from_secs(30);
+        let mut indexing_settings = IndexingSettings::for_test();
+        indexing_settings.commit_timeout_secs = commit_timeout.as_secs() as usize;
+        let (index_serializer_mailbox, index_serializer_inbox) = universe.create_test_mailbox();
+        let mut mock_metastore = MockMetastoreService::new();
+        mock_metastore
+            .expect_last_delete_opstamp()
+            .returning(|_| Ok(LastDeleteOpstampResponse::new(10)));
+        unsafe { std::env::set_var(ENABLE_SPREAD_INDEXING_PIPELINES_ENV_KEY, "true") };
+        let mut indexer = create_indexer_for_test(
+            pipeline_id,
+            doc_mapper,
+            MetastoreServiceClient::from_mock(mock_metastore),
+            TempDirectory::for_test(),
+            indexing_settings,
+            None,
+            index_serializer_mailbox,
+            None,
+        );
+        assert!(matches!(
+            indexer.indexer_state.indexing_mode,
+            IndexingMode::Spread(_)
+        ));
+
+        // Make the pipeline more than `NUDGE_TOLERANCE` behind its target phase.
+        let current_phase = PipelineSchedule::new_with_phase(Duration::default(), commit_timeout)
+            .current_phase(Instant::now());
+        indexer.indexer_state.indexing_mode = IndexingMode::Spread(
+            PipelineSchedule::new_with_phase(current_phase + NUDGE_TOLERANCE * 2, commit_timeout),
+        );
+        let (indexer_mailbox, indexer_handle) = universe.spawn_builder().spawn(indexer);
+        indexer_mailbox
+            .send_message(ProcessedDocBatch::new(
+                vec![ProcessedDoc {
+                    doc: doc!(
+                        body_field=>"this is a test document",
+                        timestamp_field=>DateTime::from_timestamp_secs(1_662_529_435)
+                    ),
+                    fingerprint_opt: None,
+                    timestamp_opt: Some(DateTime::from_timestamp_secs(1_662_529_435)),
+                    partition: 1,
+                    num_bytes: 30,
+                }],
+                SourceCheckpointDelta::from_range(0..1),
+                false,
+            ))
+            .await
+            .unwrap();
+        let indexer_counters = indexer_handle.process_pending_and_observe().await.state;
+        assert_eq!(indexer_counters.num_docs_in_workbench, 1);
+
+        // The pipeline is more than `NUDGE_TOLERANCE` behind its target phase, so the commit
+        // timeout fires after `commit_timeout`
+        universe.shift_time(commit_timeout).await;
+        let indexer_counters = indexer_handle.process_pending_and_observe().await.state;
+        // Check that no split is emitted yet
+        assert_eq!(indexer_counters.num_splits_emitted, 0);
+        assert_eq!(indexer_counters.num_docs_in_workbench, 1);
+
+        // Move to after `commit_timeout + NUDGE_TOLERANCE`
+        universe.shift_time(NUDGE_TOLERANCE * 2).await;
+        let indexer_counters = indexer_handle.process_pending_and_observe().await.state;
+        // Now the split has been emitted
+        assert_eq!(indexer_counters.num_splits_emitted, 1);
+        assert_eq!(indexer_counters.num_docs_in_workbench, 0);
+
+        let indexed_split_batches: Vec<IndexedSplitBatchBuilder> =
+            index_serializer_inbox.drain_for_test_typed();
+        assert_eq!(indexed_split_batches.len(), 1);
+        assert_eq!(
+            indexed_split_batches[0].commit_trigger,
+            CommitTrigger::Timeout
+        );
+        universe.assert_quit().await;
+    }
 }