Extract pipeline schedule logic

Co-Authored-By: Codex <noreply@openai.com>

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

@@ -12,20 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::hash::{DefaultHasher, Hash, Hasher};
-use std::sync::{Arc, LazyLock};
+use std::hash::Hash;
+use std::sync::Arc;
 use std::time::Duration;
 
 use quickwit_proto::indexing::{CpuCapacity, PIPELINE_FULL_CAPACITY, PipelineMetrics};
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use tokio::time::Instant;
 
-/// We allow ourselves to adjust the sleep time by at most `NUDGE_TOLERANCE`
-/// in order to steer a pipeline to its phase.
-const NUDGE_TOLERANCE: Duration = Duration::from_secs(5);
-
-// Origin of time. It is used to compute the phase of the pipeline.
-static ORIGIN_OF_TIME: LazyLock<Instant> = LazyLock::new(Instant::now);
+use super::pipeline_schedule::{NUDGE_TOLERANCE, PipelineSchedule};
 
 /// Cooperative indexing is a mechanism to deal with a large amount of pipelines.
 ///
@@ -69,10 +64,9 @@ static ORIGIN_OF_TIME: LazyLock<Instant> = LazyLock::new(Instant::now);
 /// mailbox to be drained.
 ///
 /// We then allow ourselves to tweak the sleep time one way or another by at
-/// most two seconds to eventually nudge the system toward the desired phase.
+/// most five seconds to eventually nudge the system toward the desired phase.
 pub(crate) struct CooperativeIndexingCycle {
-    target_phase: Duration,
-    commit_timeout: Duration,
+    schedule: PipelineSchedule,
     indexing_permits: Arc<Semaphore>,
 }
 
@@ -84,36 +78,26 @@ impl CooperativeIndexingCycle {
         commit_timeout: Duration,
         indexing_permits: Arc<Semaphore>,
     ) -> CooperativeIndexingCycle {
-        assert!(commit_timeout.as_millis() > 0);
-        let mut hasher = DefaultHasher::new();
-        phase_id.hash(&mut hasher);
-        let target_phase_millis: u64 = hasher.finish() % commit_timeout.as_millis() as u64;
-        Self::new_with_phase(
-            Duration::from_millis(target_phase_millis),
-            commit_timeout,
+        CooperativeIndexingCycle {
+            schedule: PipelineSchedule::new(phase_id, commit_timeout),
             indexing_permits,
-        )
+        }
     }
 
-    fn new_with_phase(
-        target_phase: Duration,
-        commit_timeout: Duration,
-        indexing_permits: Arc<Semaphore>,
-    ) -> CooperativeIndexingCycle {
-        // Force the initial of the origin of time.
-        let _t0 = *ORIGIN_OF_TIME;
+    #[cfg(test)]
+    fn with_schedule(schedule: PipelineSchedule, indexing_permits: Arc<Semaphore>) -> Self {
         CooperativeIndexingCycle {
-            target_phase,
-            commit_timeout,
+            schedule,
             indexing_permits,
         }
     }
 
     pub fn initial_sleep_duration(&self) -> Duration {
-        let t0 = *ORIGIN_OF_TIME;
-        let commit_timeout_millis = self.commit_timeout.as_millis() as u64;
-        let current_phase_millis: u64 = t0.elapsed().as_millis() as u64 % commit_timeout_millis;
-        let target_phase_millis: u64 = self.target_phase.as_millis() as u64 % commit_timeout_millis;
+        let commit_timeout_millis = self.schedule.commit_timeout().as_millis() as u64;
+        let current_phase_millis: u64 =
+            self.schedule.current_phase(Instant::now()).as_millis() as u64;
+        let target_phase_millis: u64 =
+            self.schedule.target_phase().as_millis() as u64 % commit_timeout_millis;
         let initial_sleep_millis: u64 = (commit_timeout_millis + target_phase_millis
             - current_phase_millis)
             % commit_timeout_millis;
@@ -134,8 +118,7 @@ impl CooperativeIndexingCycle {
         CooperativeIndexingPeriod {
             t_wake,
             t_work_start,
-            commit_timeout: self.commit_timeout,
-            target_phase: self.target_phase,
+            schedule: self.schedule,
             _permit: permit,
         }
     }
@@ -146,8 +129,7 @@ pub(crate) struct CooperativeIndexingPeriod {
     t_wake: Instant,
     // measured after the acquisition of the semaphore.
     t_work_start: Instant,
-    commit_timeout: Duration,
-    target_phase: Duration,
+    schedule: PipelineSchedule,
     _permit: OwnedSemaphorePermit,
 }
 
@@ -160,7 +142,7 @@ impl CooperativeIndexingPeriod {
         let elapsed = end - self.t_work_start;
         let throughput_mb_per_sec: u64 =
             uncompressed_num_bytes / (1u64 + elapsed.as_micros() as u64);
-        let commit_timeout = self.commit_timeout;
+        let commit_timeout = self.schedule.commit_timeout();
         let pipeline_throughput_fraction =
             (elapsed.as_micros() as f32 / commit_timeout.as_micros() as f32).min(1.0f32);
         let cpu_load: CpuCapacity = PIPELINE_FULL_CAPACITY * pipeline_throughput_fraction;
@@ -171,30 +153,10 @@ impl CooperativeIndexingPeriod {
     }
 
     fn compute_sleep_duration(&self, t_work_end: Instant) -> Duration {
-        let commit_timeout_millis = self.commit_timeout.as_millis() as u64;
-        let phase_millis: u64 =
-            ((t_work_end - *ORIGIN_OF_TIME).as_millis() as u64) % commit_timeout_millis;
-        let delta_phase: i64 = phase_millis as i64 - self.target_phase.as_millis() as i64;
-        // delta phase is within (-commit_timeout_millis, commit_timeout_millis)
-        // We fold it back to [-commit_timeout_millis/2, commit_timeout_millis/2)
-        let half_commit_timeout_millis = commit_timeout_millis as i64 / 2;
-        let delta_phase = if delta_phase >= half_commit_timeout_millis {
-            delta_phase - commit_timeout_millis as i64
-        } else if delta_phase < -half_commit_timeout_millis {
-            delta_phase + commit_timeout_millis as i64
-        } else {
-            delta_phase
-        };
-        let nudge_tolerance_millis = NUDGE_TOLERANCE.as_millis() as i64;
-        let nudge_millis: i64 = delta_phase.clamp(-nudge_tolerance_millis, nudge_tolerance_millis);
-        let sleep_duration_millis = self.commit_timeout.as_millis() as i64
-            - (t_work_end - self.t_wake).as_millis() as i64
-            - nudge_millis;
-        if sleep_duration_millis > 0 {
-            Duration::from_millis(sleep_duration_millis as u64)
-        } else {
-            Duration::ZERO
-        }
+        let elapsed = Duration::from_millis((t_work_end - self.t_wake).as_millis() as u64);
+        self.schedule
+            .nudged_commit_timeout(t_work_end)
+            .saturating_sub(elapsed)
     }
 
     /// This drops the indexing permit, allowing another indexer to start indexing.
@@ -210,6 +172,7 @@ impl CooperativeIndexingPeriod {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::actors::pipeline_schedule::ORIGIN_OF_TIME;
 
     #[track_caller]
     fn assert_approx_equal_sleep_time(left: Duration, right: Duration) {
@@ -248,16 +211,17 @@ mod tests {
                 let target_phase = Duration::from_secs(target_phase_secs);
                 let semaphore = Arc::new(Semaphore::new(1));
                 tokio::time::sleep(Duration::from_secs(start_time_secs)).await;
-                let cooperative_indexing = CooperativeIndexingCycle::new_with_phase(
-                    target_phase,
-                    Duration::from_secs(30),
+                let cooperative_indexing = CooperativeIndexingCycle::with_schedule(
+                    PipelineSchedule::new_with_phase(target_phase, Duration::from_secs(30)),
                     semaphore.clone(),
                 );
                 let initial_sleep_duration: Duration =
                     cooperative_indexing.initial_sleep_duration();
                 tokio::time::sleep(initial_sleep_duration).await;
-                let target_phase_millis = cooperative_indexing.target_phase.as_millis() as i64;
-                let commit_timeout_ms = cooperative_indexing.commit_timeout.as_millis() as i64;
+                let target_phase_millis =
+                    cooperative_indexing.schedule.target_phase().as_millis() as i64;
+                let commit_timeout_ms =
+                    cooperative_indexing.schedule.commit_timeout().as_millis() as i64;
                 let phase_millis =
                     (t0.elapsed().as_millis() as i64 - target_phase_millis) % commit_timeout_ms;
                 assert!(phase_millis >= -100, "{phase_millis}");
@@ -336,14 +300,13 @@ mod tests {
         let num_steps = 15;
         let semaphore = Arc::new(Semaphore::new(num_threads));
         let commit_timeout = Duration::from_secs(30);
-        let t0 = Instant::now();
+        let t0 = *ORIGIN_OF_TIME;
         let mut handles = Vec::new();
         for i in 0..num_pipelines {
             let target_phase =
                 Duration::from_millis(commit_timeout.as_millis() as u64 * i / num_pipelines);
-            let cooperative_indexing = CooperativeIndexingCycle::new_with_phase(
-                target_phase,
-                commit_timeout,
+            let cooperative_indexing = CooperativeIndexingCycle::with_schedule(
+                PipelineSchedule::new_with_phase(target_phase, commit_timeout),
                 semaphore.clone(),
             );
             let join_handle = tokio::task::spawn(async move {

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

@@ -24,6 +24,7 @@ mod merge_planner;
 mod merge_scheduler_service;
 mod merge_split_downloader;
 mod packager;
+mod pipeline_schedule;
 mod publisher;
 mod sequencer;
 mod uploader;

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

@@ -0,0 +1,119 @@
+// Copyright 2021-Present Datadog, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::hash::{DefaultHasher, Hash, Hasher};
+use std::sync::LazyLock;
+use std::time::Duration;
+
+use tokio::time::Instant;
+
+/// We allow ourselves to adjust the sleep time or commit timeout by at most `NUDGE_TOLERANCE`
+/// in order to steer a pipeline to its phase.
+pub(super) const NUDGE_TOLERANCE: Duration = Duration::from_secs(5);
+
+// All pipelines in the process must measure their phase from the same origin, including respawns.
+pub(super) static ORIGIN_OF_TIME: LazyLock<Instant> = LazyLock::new(Instant::now);
+
+#[derive(Clone, Copy)]
+pub(super) struct PipelineSchedule {
+    target_phase: Duration,
+    commit_timeout: Duration,
+}
+
+impl PipelineSchedule {
+    pub fn new(phase_id: &(impl Hash + ?Sized), commit_timeout: Duration) -> Self {
+        assert!(commit_timeout.as_millis() > 0);
+        let mut hasher = DefaultHasher::new();
+        phase_id.hash(&mut hasher);
+        let target_phase_millis = u128::from(hasher.finish()) % commit_timeout.as_millis();
+        Self::new_with_phase(
+            Duration::from_millis(target_phase_millis as u64),
+            commit_timeout,
+        )
+    }
+
+    pub(super) fn new_with_phase(target_phase: Duration, commit_timeout: Duration) -> Self {
+        // Force the initialization of the origin of time.
+        let _t0 = *ORIGIN_OF_TIME;
+        Self {
+            target_phase,
+            commit_timeout,
+        }
+    }
+
+    pub fn commit_timeout(&self) -> Duration {
+        self.commit_timeout
+    }
+
+    pub fn target_phase(&self) -> Duration {
+        self.target_phase
+    }
+
+    pub fn current_phase(&self, now: Instant) -> Duration {
+        let t0 = *ORIGIN_OF_TIME;
+        let current_phase_millis = (now - t0).as_millis() % self.commit_timeout.as_millis();
+        Duration::from_millis(current_phase_millis as u64)
+    }
+
+    pub fn nudged_commit_timeout(&self, now: Instant) -> Duration {
+        let commit_timeout_millis = self.commit_timeout.as_millis() as i128;
+        let current_phase_millis = self.current_phase(now).as_millis() as i128;
+        let delta_phase = current_phase_millis - self.target_phase.as_millis() as i128;
+        // Take the shortest correction across the period boundary, then bound it to five seconds.
+        let half_period_millis = commit_timeout_millis / 2;
+        let delta_phase = (delta_phase + half_period_millis).rem_euclid(commit_timeout_millis)
+            - half_period_millis;
+        let nudge_tolerance_millis = NUDGE_TOLERANCE.as_millis() as i128;
+        let nudge_millis = delta_phase.clamp(-nudge_tolerance_millis, nudge_tolerance_millis);
+        if nudge_millis >= 0 {
+            self.commit_timeout - Duration::from_millis(nudge_millis as u64)
+        } else {
+            self.commit_timeout + Duration::from_millis((-nudge_millis) as u64)
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_nudged_commit_timeout() {
+        let t0 = *ORIGIN_OF_TIME;
+        let commit_timeout = Duration::from_secs(30);
+        for (target_phase_secs, elapsed_secs, expected_timeout_secs) in [
+            (10, 10, 30),
+            (10, 12, 28),
+            (10, 8, 32),
+            (0, 29, 31),
+            (29, 0, 29),
+            (10, 16, 25),
+            (10, 4, 35),
+            (10, 24, 25),
+            (10, 25, 35),
+            (10, 42, 28),
+            (10, 38, 32),
+        ] {
+            let schedule = PipelineSchedule::new_with_phase(
+                Duration::from_secs(target_phase_secs),
+                commit_timeout,
+            );
+            assert_eq!(
+                schedule.nudged_commit_timeout(t0 + Duration::from_secs(elapsed_secs)),
+                Duration::from_secs(expected_timeout_secs),
+                "target phase {target_phase_secs}s, elapsed time {elapsed_secs}s"
+            );
+        }
+    }
+}