[9/N][VQueues] Ensure the scheduler's internal clock wheel is synchronized with eligibility_at

This changeset makes sure that we have a reliable datum point that's shared between the delayed queue internal timer wheel and the scheduler.
Additionally, it ensures that the scheduler can create `Instant` values that are synchronized with the datum point acquired from UniqueTimestamp.

This fixes the issue introduced in changeset 6e65909 where we would become eligible slightly before the actual eligibility time, causing supurious rescheduling.

Author: AhmedSoliman

crates/types/src/clock/unique_timestamp.rs

@@ -96,10 +96,22 @@ impl UniqueTimestamp {
         Ok(Self(nz))
     }
 
+    /// Compare the physical clock of both timestamps and ignore the logical clock.
+    ///
+    /// Note that this has millisecond precision.
+    pub fn cmp_physical(&self, other: &Self) -> std::cmp::Ordering {
+        self.physical_raw().cmp(&other.physical_raw())
+    }
+
     pub fn as_u64(&self) -> u64 {
         self.0.get()
     }
 
+    /// Fails if the resulting timestamp is out of range
+    pub fn add_millis(&self, millis: u64) -> Result<Self, Error> {
+        Self::from_parts(self.physical_raw() + millis, self.logical_raw())
+    }
+
     #[inline(always)]
     fn physical_raw(&self) -> u64 {
         // extract the physical time

crates/vqueues/src/delay_queue.rs

@@ -471,10 +471,10 @@ impl<T> DelayQueue<T> {
     ///
     /// ```ignore
     /// # use restate_vqueues::delay_queue::DelayQueue;
-    /// let delay_queue: DelayQueue<u32> = DelayQueue::new();
+    /// let delay_queue: DelayQueue<u32> = DelayQueue::new(tokio::time::Instant::now());
     /// ```
-    pub fn new() -> DelayQueue<T> {
-        DelayQueue::with_capacity(0)
+    pub fn new(start: Instant) -> DelayQueue<T> {
+        DelayQueue::with_capacity(start, 0)
     }
 
     /// Creates a new, empty, `DelayQueue` with the specified capacity.
@@ -491,7 +491,7 @@ impl<T> DelayQueue<T> {
     ///
     /// # #[tokio::main(flavor = "current_thread")]
     /// # async fn main() {
-    /// let mut delay_queue = DelayQueue::with_capacity(10);
+    /// let mut delay_queue = DelayQueue::with_capacity(tokio::time::Instant::now(), 10);
     ///
     /// // These insertions are done without further allocation
     /// for i in 0..10 {
@@ -502,14 +502,14 @@ impl<T> DelayQueue<T> {
     /// delay_queue.insert(11, Duration::from_secs(11));
     /// # }
     /// ```
-    pub fn with_capacity(capacity: usize) -> DelayQueue<T> {
+    pub fn with_capacity(start: Instant, capacity: usize) -> DelayQueue<T> {
         DelayQueue {
             wheel: Wheel::new(),
             slab: SlabStorage::with_capacity(capacity),
             expired: Stack::default(),
             delay: None,
             wheel_now: 0,
-            start: Instant::now(),
+            start,
             waker: None,
         }
     }
@@ -1215,12 +1215,6 @@ impl<T> DelayQueue<T> {
 // We never put `T` in a `Pin`...
 impl<T> Unpin for DelayQueue<T> {}
 
-impl<T> Default for DelayQueue<T> {
-    fn default() -> DelayQueue<T> {
-        DelayQueue::new()
-    }
-}
-
 impl<T> Stream for DelayQueue<T> {
     // DelayQueue seems much more specific, where a user may care that it
     // has reached capacity, so return those errors instead of panicking.

crates/vqueues/src/scheduler.rs

@@ -20,15 +20,14 @@ use smallvec::SmallVec;
 use restate_futures_util::concurrency::{Concurrency, Permit};
 use restate_storage_api::StorageError;
 use restate_storage_api::vqueue_table::{ScanVQueueTable, VQueueStore};
-use restate_types::clock::UniqueTimestamp;
-use restate_types::time::MillisSinceEpoch;
 use restate_types::vqueue::VQueueId;
 
 use crate::VQueueEvent;
 use crate::{VQueuesMeta, VQueuesMetaMut};
 
 use self::drr::DRRScheduler;
 
+mod clock;
 mod drr;
 mod vqueue_state;
 
@@ -185,12 +184,11 @@ where
 
     pub fn on_inbox_event(
         &mut self,
-        now: UniqueTimestamp,
         vqueues: VQueuesMeta<'_>,
         event: &VQueueEvent<S::Item>,
     ) -> Result<(), StorageError> {
         if let State::Active(ref mut drr_scheduler) = self.state {
-            drr_scheduler.as_mut().on_inbox_event(now, vqueues, event)?;
+            drr_scheduler.as_mut().on_inbox_event(vqueues, event)?;
         }
         Ok(())
     }
@@ -199,21 +197,19 @@ where
         &mut self,
         vqueues: VQueuesMeta<'_>,
     ) -> Result<Decision<S::Item>, StorageError> {
-        let now = UniqueTimestamp::from_unix_millis(MillisSinceEpoch::now()).unwrap();
-        poll_fn(|cx| self.poll_schedule_next(now, cx, vqueues)).await
+        poll_fn(|cx| self.poll_schedule_next(cx, vqueues)).await
     }
 
     pub fn poll_schedule_next(
         &mut self,
-        now: UniqueTimestamp,
         cx: &mut std::task::Context<'_>,
         vqueues: VQueuesMeta<'_>,
     ) -> Poll<Result<Decision<S::Item>, StorageError>> {
         match self.state {
             // if scheduler is disabled, we always return pending.
             State::Disabled => Poll::Pending,
             State::Active(ref mut drr_scheduler) => {
-                drr_scheduler.as_mut().poll_schedule_next(now, cx, vqueues)
+                drr_scheduler.as_mut().poll_schedule_next(cx, vqueues)
             }
         }
     }

crates/vqueues/src/scheduler/clock.rs

@@ -0,0 +1,50 @@
+// Copyright (c) 2023 - 2025 Restate Software, Inc., Restate GmbH.
+// All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+use std::time::Duration;
+
+use restate_types::clock::UniqueTimestamp;
+
+/// A clock that tracks the physical clock of the scheduler that is synchronized
+/// with the UniqueTimestamp physical clock.
+pub struct SchedulerClock {
+    tokio_origin: tokio::time::Instant,
+    ts_origin: UniqueTimestamp,
+}
+
+impl SchedulerClock {
+    /// `now` is the current physical clock that corresponds to `Instant::now()` at the time
+    /// of construction.
+    pub fn new(now: UniqueTimestamp) -> Self {
+        Self {
+            tokio_origin: tokio::time::Instant::now(),
+            ts_origin: now,
+        }
+    }
+
+    pub fn origin_instant(&self) -> tokio::time::Instant {
+        self.tokio_origin
+    }
+
+    pub fn now_ts(&self) -> UniqueTimestamp {
+        let delta = self.tokio_origin.elapsed().as_millis() as u64;
+        self.ts_origin
+            .add_millis(delta)
+            .expect("clock doesn't overflow")
+    }
+
+    /// Calculates a future tokio Instant from the physical clock of `ts`.
+    ///
+    /// Returns the clock datum/origin point if the input `ts` is in the past.
+    pub fn ts_to_future_instant(&self, ts: UniqueTimestamp) -> tokio::time::Instant {
+        let delta = ts.milliseconds_since(self.ts_origin);
+        self.tokio_origin + Duration::from_millis(delta)
+    }
+}

crates/vqueues/src/scheduler/drr.rs

@@ -13,17 +13,20 @@ use std::num::NonZeroU16;
 use std::pin::Pin;
 use std::task::Poll;
 use std::task::Waker;
+use std::time::Duration;
 
 use hashbrown::HashMap;
 use hashbrown::hash_map;
 use pin_project::pin_project;
+use tokio::time::Instant;
 use tracing::{info, trace};
 
 use restate_futures_util::concurrency::Concurrency;
 use restate_futures_util::concurrency::Permit;
 use restate_storage_api::StorageError;
 use restate_storage_api::vqueue_table::VQueueStore;
 use restate_types::clock::UniqueTimestamp;
+use restate_types::time::MillisSinceEpoch;
 use restate_types::vqueue::VQueueId;
 
 use crate::EventDetails;
@@ -34,6 +37,7 @@ use crate::scheduler::Assignments;
 use crate::scheduler::vqueue_state::Eligibility;
 
 use super::Decision;
+use super::clock::SchedulerClock;
 use super::vqueue_state::VQueueState;
 
 /// Capacity to maintain for N vqueues (N=100)
@@ -55,8 +59,9 @@ pub struct DRRScheduler<S: VQueueStore, Token> {
     remaining_in_round: usize,
     /// Waker to be notified when scheduler is potentially able to scheduler more work
     waker: Waker,
+    datum: SchedulerClock,
     /// Time of the last memory reporting and memory compaction
-    last_report: Option<UniqueTimestamp>,
+    last_report: Instant,
     // SAFETY NOTE: **must** Keep this at the end since it needs to outlive all readers.
     storage: S,
 }
@@ -94,17 +99,27 @@ where
             q.len(),
         );
 
+        let datum = SchedulerClock::new(
+            UniqueTimestamp::from_unix_millis(MillisSinceEpoch::now())
+                .expect("clock does not overflow"),
+        );
+        // Makes sure we use the same clock datum for the internal timer wheel and for our
+        // own eligibility checks.
+        let start = datum.origin_instant();
+        let delayed_eligibility = DelayQueue::new(start);
+
         Self {
             limit_qid_per_poll,
             concurrency_limiter,
             q,
             eligible,
             global_sched_round: 0,
             remaining_in_round: 0,
-            delayed_eligibility: DelayQueue::new(),
+            delayed_eligibility,
             unconfirmed_capacity_permits: Permit::new_empty(),
             waker: Waker::noop().clone(),
-            last_report: None,
+            datum,
+            last_report: start,
             storage,
         }
     }
@@ -122,7 +137,6 @@ where
 
     pub fn poll_schedule_next(
         mut self: Pin<&mut Self>,
-        now: UniqueTimestamp,
         cx: &mut std::task::Context<'_>,
         vqueues: VQueuesMeta<'_>,
     ) -> Poll<Result<Decision<S::Item>, StorageError>> {
@@ -131,14 +145,11 @@ where
             Abort,
         }
 
-        if self
-            .last_report
-            .is_none_or(|t| now.milliseconds_since(t) >= 10000)
-        {
+        if self.last_report.elapsed() >= Duration::from_secs(10) {
             vqueues.report();
             self.report();
             // also report vqueues states
-            self.last_report = Some(now);
+            self.last_report = Instant::now();
             // compact memory
             self.q.shrink_to(MIN_VQUEUES_CAPACITY);
             self.delayed_eligibility.compact();
@@ -151,20 +162,6 @@ where
             if self.remaining_in_round == 0 {
                 // Pop all eligible vqueues that were delayed since we are starting a new round
                 // Once we hit pending, the waker will be registered.
-                //
-                // There is currently an issue due to using two different clock sources. The timer
-                // wheel uses tokio's internal clock but our scheduler's design relies on explicit
-                // monotonic timestamps. This will cause the timer to expire slightly before or after
-                // the actual point at which the input (now) would satisfy the head item's
-                // eligibility requirements. As a result, we will see one of three scenarios:
-                //  1. Time aligns. We pop the timer and the head element will be immediately eligible.
-                //  2. We are woken up before `now` satifies the requirement, in this case will
-                //     schedule a new timer to catch up on the difference.
-                //  3. We are worken up after `now` by a few hundres of millis, in this case the
-                //     head will be eligible but we will be a little late.
-                //
-                //  This will be fixed in a later change after we reason about whether we need any
-                //  causal entanglement between the RSM's clock and the scheduler's internal clock.
                 let previous_round = self.global_sched_round;
                 while let Poll::Ready(Some(expired)) = self.delayed_eligibility.poll_expired(cx) {
                     let Some(qstate) = self.q.get_mut(expired.get_ref()) else {
@@ -188,6 +185,7 @@ where
                 }
             }
 
+            let now = self.datum.now_ts();
             while self.remaining_in_round > 0 {
                 // bail if we exhausted coop budget.
                 let coop = match tokio::task::coop::poll_proceed(cx) {
@@ -227,7 +225,10 @@ where
                         }
                         Eligibility::EligibleAt(wake_up_at) => {
                             *this.remaining_in_round -= 1;
-                            qstate.maybe_schedule_wakeup(wake_up_at, this.delayed_eligibility, now);
+                            qstate.maybe_schedule_wakeup(
+                                this.datum.ts_to_future_instant(wake_up_at),
+                                this.delayed_eligibility,
+                            );
                             this.eligible.pop_front();
                             break 'single_vqueue Outcome::ContinueRound;
                         }
@@ -306,7 +307,6 @@ where
     #[tracing::instrument(skip_all)]
     pub fn on_inbox_event(
         &mut self,
-        now: UniqueTimestamp,
         vqueues: VQueuesMeta<'_>,
         event: &VQueueEvent<S::Item>,
     ) -> Result<(), StorageError> {
@@ -326,7 +326,8 @@ where
                     return Ok(());
                 }
 
-                match qstate.check_eligibility(now, meta, config) {
+                let now_ts = self.datum.now_ts();
+                match qstate.check_eligibility(now_ts, meta, config) {
                     Eligibility::Eligible if !self.eligible.contains(&qid) => {
                         // Make eligible immediately.
                         qstate.deficit.set_last_round(self.global_sched_round);
@@ -335,9 +336,8 @@ where
                     }
                     Eligibility::EligibleAt(eligiblility_ts) if !self.eligible.contains(&qid) => {
                         qstate.maybe_schedule_wakeup(
-                            eligiblility_ts,
+                            self.datum.ts_to_future_instant(eligiblility_ts),
                             &mut self.delayed_eligibility,
-                            now,
                         );
                     }
                     _ => { /* do nothing */ }
@@ -354,7 +354,10 @@ where
                     // drop the already acquired permit
                     let _ = self.unconfirmed_capacity_permits.split(1);
 
-                    if qstate.check_eligibility(now, meta, config).is_eligible() {
+                    if qstate
+                        .check_eligibility(self.datum.now_ts(), meta, config)
+                        .is_eligible()
+                    {
                         if !self.eligible.contains(&qid) {
                             self.eligible.push_back(qid);
                             qstate.deficit.set_last_round(self.global_sched_round);
@@ -377,7 +380,7 @@ where
                     // drop the already acquired permit
                     let _ = self.unconfirmed_capacity_permits.split(1);
                 } else if qstate.notify_removed(item_hash) {
-                    match qstate.check_eligibility(now, meta, config) {
+                    match qstate.check_eligibility(self.datum.now_ts(), meta, config) {
                         Eligibility::Eligible if !self.eligible.contains(&qid) => {
                             // Make eligible immediately.
                             qstate.deficit.set_last_round(self.global_sched_round);
@@ -388,9 +391,8 @@ where
                             if !self.eligible.contains(&qid) =>
                         {
                             qstate.maybe_schedule_wakeup(
-                                eligiblility_ts,
+                                self.datum.ts_to_future_instant(eligiblility_ts),
                                 &mut self.delayed_eligibility,
-                                now,
                             );
                         }
                         _ => { /* do nothing */ }