[VQueues] Major redesign of vqueues internal design

Author: AhmedSoliman

Cargo.lock

@@ -65,7 +65,7 @@ fn entry_value(id: u8, original_run_at: u64, num_attempts: u32) -> EntryValue {
         stats.latest_attempt_at = Some(created_at);
     }
     let status = if num_attempts > 0 {
-        Status::Running
+        Status::Started
     } else if stats.first_runnable_at > created_at.to_unix_millis() {
         Status::Scheduled
     } else {

crates/partition-store/src/tests/vqueue_table_test/mod.rs

@@ -34,8 +34,6 @@ define_table_key!(
     )
 );
 
-static_assertions::const_assert_eq!(EntryKind::serialized_length_fixed(), 1);
-
 impl EntryStatusKey {
     pub const fn serialized_length_fixed() -> usize {
         KeyKind::SERIALIZED_LENGTH
@@ -61,8 +59,6 @@ impl From<&InvocationId> for EntryStatusKey {
 pub struct StatusHeaderRawRef<'a> {
     #[bilrost(tag(1))]
     pub(super) qid: VQueueIdRef<'a>,
-    /// Unknown is an invalid state, this will be set to None when the invocation
-    /// leaves the queue.
     #[bilrost(tag(2))]
     pub(super) stage: Stage,
     #[bilrost(tag(3))]

crates/partition-store/src/vqueue_table/entry.rs

@@ -25,7 +25,7 @@ use super::key_codec::HasLock;
 // 'qP' | QID | HAS_LOCK(1B) | RUN_AT(8B) | SEQ(8B) | ENTRY_ID(17B)
 // 'qF' | QID | HAS_LOCK(1B) | RUN_AT(8B) | SEQ(8B) | ENTRY_ID(17B)
 
-// Inbox, Running, Suspended, Paused, and have the same key design.
+// Inbox, Running, Suspended, Paused, and Finished and have the same key design.
 macro_rules! define_stage_keys {
     ( $(Stage::$stage: ident => $key_name: ident),+ $(,)? ) => {
         paste::paste! {

crates/partition-store/src/vqueue_table/inbox.rs

@@ -15,6 +15,9 @@ use crate::keys::{KeyKind, define_table_key};
 
 // Input payloads stored for vqueue items.
 // Vqueue items are stored under the qid they belong to and their creation order
+// The sequence number acts as an epoch to distinguish between different incarinations
+// of the same EntryId (e.g. deterministic invocation re-invoked after expiration).
+//
 // 'qi' | QID | SEQ | ENTRY_ID
 define_table_key!(
     VQueue,

…artition-store/src/vqueue_table/items.rs …artition-store/src/vqueue_table/input.rscrates/partition-store/src/vqueue_table/items.rs renamed to crates/partition-store/src/vqueue_table/input.rs crates/partition-store/src/vqueue_table/items.rs renamed to crates/partition-store/src/vqueue_table/input.rs

@@ -8,7 +8,6 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use anyhow::Context;
 use bytes::{Buf, BufMut};
 
 use restate_clock::RoughTimestamp;
@@ -130,11 +129,14 @@ impl KeyDecode for EntryId {
     }
 }
 
-// RoughTimestamp is encoded as a u64 in big-endian order to enable forward compatibility
-// with potential future higher-precision restate-epoch based timestamp.
+// The on-disk value is milliseconds-since-restate-epoch, quantized to whole
+// seconds by the current seconds-precision `RoughTimestamp`. Storing in the
+// millisecond domain keeps the format forward-compatible with a future
+// higher-precision ms timestamp using the same u64 slot — no migration, and
+// sort order is preserved across the boundary.
 impl KeyEncode for RoughTimestamp {
     fn encode<B: BufMut>(&self, target: &mut B) {
-        target.put_u64(self.as_u32() as u64);
+        target.put_u64(self.as_u32() as u64 * 1_000);
     }
 
     fn serialized_length(&self) -> usize {
@@ -144,11 +146,12 @@ impl KeyEncode for RoughTimestamp {
 
 impl KeyDecode for RoughTimestamp {
     fn decode<B: Buf>(source: &mut B) -> crate::Result<Self> {
-        let raw = source.get_u64();
-        Ok(Self::new(
-            raw.try_into()
-                .context("RoughTimestamp needs to fit into u32")?,
-        ))
+        let raw_ms = source.get_u64();
+        // Floor to seconds. Clamp so that forward-written ms values beyond
+        // `RoughTimestamp`'s range saturate at MAX instead of truncating
+        // on the `as u32` cast.
+        let secs = (raw_ms / 1_000).min(u32::MAX as u64) as u32;
+        Ok(Self::new(secs))
     }
 }
 
@@ -190,3 +193,55 @@ impl KeyEncode for EntryKey {
         Self::serialized_length_fixed()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn rough_timestamp_ms_codec() {
+        // Round-trip preserves seconds across the representable range, and
+        // encoded bytes are always an exact multiple of 1_000 (i.e. the
+        // encoder only ever emits whole-second ms values, so the `*1_000` /
+        // `/1_000` pair is lossless for every `RoughTimestamp`).
+        for &secs in &[0u32, 1, 100, 1_000, u32::MAX / 2, u32::MAX - 1] {
+            let mut buf = Vec::new();
+            RoughTimestamp::new(secs).encode(&mut buf);
+
+            let raw = u64::from_be_bytes(buf.as_slice().try_into().unwrap());
+            assert_eq!(
+                raw % 1_000,
+                0,
+                "encoded value must be a multiple of 1_000 (got {raw} for {secs}s)"
+            );
+            assert_eq!(
+                raw,
+                secs as u64 * 1_000,
+                "encoded ms must equal seconds * 1_000"
+            );
+
+            let mut slice = buf.as_slice();
+            let decoded = RoughTimestamp::decode(&mut slice).expect("decode ok");
+            assert_eq!(decoded.as_u32(), secs, "round-trip failed for {secs}s");
+        }
+
+        // Forward-compat: sub-second ms values floor to the correct second.
+        let mut sub_second = Vec::new();
+        sub_second.put_u64(5_999);
+        let mut slice = sub_second.as_slice();
+        assert_eq!(
+            RoughTimestamp::decode(&mut slice).unwrap().as_u32(),
+            5,
+            "ms value should floor to whole seconds on decode"
+        );
+
+        // Forward-compat: ms values beyond RoughTimestamp::MAX clamp to MAX.
+        let mut over_max = Vec::new();
+        over_max.put_u64(u64::MAX);
+        let mut slice = over_max.as_slice();
+        assert_eq!(
+            RoughTimestamp::decode(&mut slice).unwrap(),
+            RoughTimestamp::MAX,
+        );
+    }
+}

crates/partition-store/src/vqueue_table/key_codec.rs

@@ -10,7 +10,7 @@
 
 mod entry;
 mod inbox;
-mod items;
+mod input;
 mod key_codec;
 mod metadata;
 mod reader;
@@ -21,7 +21,7 @@ use std::io::Cursor;
 
 pub use entry::{EntryStatusKey, StatusHeaderRaw};
 pub use inbox::InboxKey;
-pub use items::InputPayloadKey;
+pub use input::InputPayloadKey;
 pub use metadata::*;
 
 use anyhow::Context;
@@ -298,21 +298,21 @@ impl WriteVQueueTable for PartitionStoreTransaction<'_> {
     ) where
         E: Message,
     {
-        let key_buf = {
-            let seq = seq.into();
-            let key = InputPayloadKey::builder_ref().qid(qid).seq(&seq).id(id);
-            let key_buf = self.cleared_key_buffer_mut(key.serialized_length());
-            key.serialize_to(key_buf);
-            key_buf.split()
-        };
+        let seq = seq.into();
+        let mut key_buffer = [0u8; InputPayloadKey::serialized_length_fixed()];
+        InputPayloadKey::builder_ref()
+            .qid(qid)
+            .seq(&seq)
+            .id(id)
+            .serialize_to(&mut key_buffer.as_mut());
 
         let value_buffer = self.cleared_value_buffer_mut(item.encoded_len());
 
         item.encode(value_buffer)
             .expect("enough space to encode item");
         let value = value_buffer.split();
 
-        self.raw_put_cf(KeyKind::VQueueInput, key_buf, value);
+        self.raw_put_cf(KeyKind::VQueueInput, key_buffer, value);
     }
 
     fn delete_vqueue_input_payload(&mut self, qid: &VQueueId, seq: impl Into<Seq>, id: &EntryId) {

crates/partition-store/src/vqueue_table/mod.rs

@@ -157,7 +157,7 @@ impl EntryKey {
 
 #[derive(Debug, Clone, bilrost::Message)]
 pub struct EntryValue {
-    /// Status is copied over from the entry's state when the last transition
+    /// Status is copied over from the entry status table when the last transition
     /// happened.
     #[bilrost(tag(1))]
     pub status: Status,

crates/storage-api/src/vqueue_table/entry.rs

@@ -23,61 +23,32 @@ pub enum Status {
     New,
     #[bilrost(2)]
     Scheduled,
-    // -- Statuses for an invocation that has already started (attempted at least once)
+    /// Invocation has started running with at least one attempt.
     #[bilrost(3)]
-    Running,
-    #[bilrost(4)]
-    Suspended,
+    Started,
     /// Invocation has previously started but has been placed back on the waiting inbox
     /// due to an attempt error.
-    #[bilrost(5)]
+    #[bilrost(4)]
     BackingOff,
     /// Invocation has previously started but has been placed back on the waiting inbox.
     /// This does not mean that the invocation attempt has failed, it just means that
     /// it has been evicted from the run queue and will be resumed later.
-    #[bilrost(6)]
+    #[bilrost(5)]
     Yielded,
-    /// Inovocation that was suspended and is now waiting for its turn
-    /// to run.
-    #[bilrost(7)]
-    WakingUp,
     ///
     /// -- Terminal states, invocation cannot transition back to any of the previous
     /// statuses
     ///
-    #[bilrost(8)]
+    #[bilrost(6)]
     Killed,
-    #[bilrost(9)]
+    #[bilrost(7)]
     Cancelled,
-    #[bilrost(10)]
+    #[bilrost(8)]
     Failed,
-    #[bilrost(11)]
+    #[bilrost(9)]
     Succeeded,
 }
 
-impl Status {
-    #[inline]
-    pub fn is_terminal(&self) -> bool {
-        matches!(
-            self,
-            Self::Unknown | Self::Killed | Self::Cancelled | Self::Failed | Self::Succeeded
-        )
-    }
-
-    #[inline]
-    pub fn can_move_to_run(&self) -> bool {
-        matches!(
-            self,
-            Self::New
-                | Self::Scheduled
-                | Self::WakingUp
-                | Self::Yielded
-                | Self::BackingOff
-                | Self::Running
-        )
-    }
-}
-
 pub trait EntryStatusHeader: std::fmt::Debug {
     fn vqueue_id(&self) -> &VQueueId;
     fn status(&self) -> Status;
@@ -91,6 +62,18 @@ pub trait EntryStatusHeader: std::fmt::Debug {
     fn seq(&self) -> Seq;
     fn stats(&self) -> &EntryStatistics;
     fn display_entry_id(&self) -> impl std::fmt::Display + '_;
+    /// Returns new if this entry has not started yet.
+    fn has_started(&self) -> bool {
+        self.stats().num_attempts > 0
+    }
+    /// Returns true if this entry is in the terminal state and cannot transition
+    /// out of it.
+    fn is_terminal(&self) -> bool {
+        if matches!(self.stage(), Stage::Finished) {
+            return true;
+        }
+        false
+    }
 }
 
 /// For future support for extra state storage for entries.