Compute partition key eagerly in ServiceId

src/common/src/types.rs

@@ -59,11 +59,12 @@ impl ServiceInvocationId {
         key: impl Into<Bytes>,
         invocation_id: impl Into<InvocationId>,
     ) -> Self {
+        Self::with_service_id(ServiceId::new(service_name, key), invocation_id)
+    }
+
+    pub fn with_service_id(service_id: ServiceId, invocation_id: impl Into<InvocationId>) -> Self {
         Self {
-            service_id: ServiceId {
-                service_name: service_name.into(),
-                key: key.into(),
-            },
+            service_id,
             invocation_id: invocation_id.into(),
         }
     }
@@ -80,19 +81,33 @@ pub struct ServiceId {
     pub service_name: ByteString,
     /// Identifies the service instance for the given service name
     pub key: Bytes,
+
+    partition_key: PartitionKey,
 }
 
 impl ServiceId {
     pub fn new(service_name: impl Into<ByteString>, key: impl Into<Bytes>) -> Self {
+        let key = key.into();
+        let partition_key = HashPartitioner::compute_partition_key(&key);
+        Self::with_partition_key(partition_key, service_name, key)
+    }
+
+    /// # Important
+    /// The `partition_key` must be hash of the `key` computed via [`HashPartitioner`].
+    pub fn with_partition_key(
+        partition_key: PartitionKey,
+        service_name: impl Into<ByteString>,
+        key: impl Into<Bytes>,
+    ) -> Self {
         Self {
             service_name: service_name.into(),
             key: key.into(),
+            partition_key,
         }
     }
 
     pub fn partition_key(&self) -> PartitionKey {
-        // Todo: Figure out whether to cache this value in ServiceId struct
-        HashPartitioner::compute_partition_key(&self.key)
+        self.partition_key
     }
 }
 

src/storage_api/src/inbox_table/mod.rs

@@ -1,30 +1,12 @@
 use crate::{GetFuture, GetStream, PutFuture};
-use restate_common::types::{InboxEntry, PartitionKey, ServiceId};
+use restate_common::types::{InboxEntry, ServiceId};
 
 pub trait InboxTable {
-    fn put_invocation(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-        inbox_entry: InboxEntry,
-    ) -> PutFuture;
+    fn put_invocation(&mut self, service_id: &ServiceId, inbox_entry: InboxEntry) -> PutFuture;
 
-    fn delete_invocation(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-        sequence_number: u64,
-    ) -> PutFuture;
+    fn delete_invocation(&mut self, service_id: &ServiceId, sequence_number: u64) -> PutFuture;
 
-    fn peek_inbox(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-    ) -> GetFuture<Option<InboxEntry>>;
+    fn peek_inbox(&mut self, service_id: &ServiceId) -> GetFuture<Option<InboxEntry>>;
 
-    fn inbox(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-    ) -> GetStream<InboxEntry>;
+    fn inbox(&mut self, service_id: &ServiceId) -> GetStream<InboxEntry>;
 }

src/storage_api/src/journal_table/mod.rs

@@ -1,33 +1,25 @@
 use crate::{GetFuture, GetStream, PutFuture};
-use restate_common::types::{EntryIndex, JournalEntry, PartitionKey, ServiceId};
+use restate_common::types::{EntryIndex, JournalEntry, ServiceId};
 
 pub trait JournalTable {
     fn put_journal_entry(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
         journal_index: u32,
         journal_entry: JournalEntry,
     ) -> PutFuture;
 
     fn get_journal_entry(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
         journal_index: u32,
     ) -> GetFuture<Option<JournalEntry>>;
 
     fn get_journal(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
         journal_length: EntryIndex,
     ) -> GetStream<JournalEntry>;
 
-    fn delete_journal(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-        journal_length: EntryIndex,
-    ) -> PutFuture;
+    fn delete_journal(&mut self, service_id: &ServiceId, journal_length: EntryIndex) -> PutFuture;
 }

src/storage_api/src/state_table/mod.rs

@@ -1,33 +1,26 @@
 use crate::{GetFuture, GetStream, PutFuture};
 use bytes::Bytes;
-use restate_common::types::{PartitionKey, ServiceId};
+use restate_common::types::ServiceId;
 
 pub trait StateTable {
     fn put_user_state(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
         state_key: impl AsRef<[u8]>,
         state_value: impl AsRef<[u8]>,
     ) -> PutFuture;
 
     fn delete_user_state(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
         state_key: impl AsRef<[u8]>,
     ) -> PutFuture;
 
     fn get_user_state(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
         state_key: impl AsRef<[u8]>,
     ) -> GetFuture<Option<Bytes>>;
 
-    fn get_all_user_states(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-    ) -> GetStream<(Bytes, Bytes)>;
+    fn get_all_user_states(&mut self, service_id: &ServiceId) -> GetStream<(Bytes, Bytes)>;
 }

src/storage_api/src/status_table/mod.rs

@@ -5,22 +5,16 @@ use std::ops::RangeInclusive;
 pub trait StatusTable {
     fn put_invocation_status(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
         status: InvocationStatus,
     ) -> PutFuture;
 
     fn get_invocation_status(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
     ) -> GetFuture<Option<InvocationStatus>>;
 
-    fn delete_invocation_status(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-    ) -> PutFuture;
+    fn delete_invocation_status(&mut self, service_id: &ServiceId) -> PutFuture;
 
     fn invoked_invocations(
         &mut self,

src/storage_grpc/tests/integration_test.rs

@@ -41,8 +41,7 @@ async fn test_state() {
     for i in 0..101 {
         let svc_key = format!("key-{}", i / 10);
         txn.put_user_state(
-            1337,
-            &ServiceId::new("svc-1", svc_key),
+            &ServiceId::with_partition_key(1337, "svc-1", svc_key),
             &Bytes::from(format!("{i}")),
             &Bytes::from(format!("{i}")),
         )

src/storage_rocksdb/src/inbox_table/mod.rs

@@ -26,15 +26,14 @@ define_table_key!(
 impl InboxTable for RocksDBTransaction {
     fn put_invocation(
         &mut self,
-        partition_key: PartitionKey,
         service_id: &ServiceId,
         InboxEntry {
             inbox_sequence_number,
             service_invocation,
         }: InboxEntry,
     ) -> PutFuture {
         let key = InboxKey::default()
-            .partition_key(partition_key)
+            .partition_key(service_id.partition_key())
             .service_name(service_id.service_name.clone())
             .service_key(service_id.key.clone())
             .sequence_number(inbox_sequence_number);
@@ -46,14 +45,9 @@ impl InboxTable for RocksDBTransaction {
         ready()
     }
 
-    fn delete_invocation(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-        sequence_number: u64,
-    ) -> PutFuture {
+    fn delete_invocation(&mut self, service_id: &ServiceId, sequence_number: u64) -> PutFuture {
         let key = InboxKey::default()
-            .partition_key(partition_key)
+            .partition_key(service_id.partition_key())
             .service_name(service_id.service_name.clone())
             .service_key(service_id.key.clone())
             .sequence_number(sequence_number);
@@ -62,22 +56,14 @@ impl InboxTable for RocksDBTransaction {
         ready()
     }
 
-    fn peek_inbox(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-    ) -> GetFuture<Option<InboxEntry>> {
-        let mut stream = self.inbox(partition_key, service_id);
+    fn peek_inbox(&mut self, service_id: &ServiceId) -> GetFuture<Option<InboxEntry>> {
+        let mut stream = self.inbox(service_id);
         async move { stream.next().await.transpose() }.boxed()
     }
 
-    fn inbox(
-        &mut self,
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-    ) -> GetStream<InboxEntry> {
+    fn inbox(&mut self, service_id: &ServiceId) -> GetStream<InboxEntry> {
         let key = InboxKey::default()
-            .partition_key(partition_key)
+            .partition_key(service_id.partition_key())
             .service_name(service_id.service_name.clone())
             .service_key(service_id.key.clone());
 
@@ -104,15 +90,11 @@ mod tests {
     use crate::inbox_table::InboxKey;
     use crate::keys::TableKey;
     use bytes::{Bytes, BytesMut};
-    use restate_common::types::{PartitionKey, ServiceId};
+    use restate_common::types::ServiceId;
 
-    fn message_key(
-        partition_key: PartitionKey,
-        service_id: &ServiceId,
-        sequence_number: u64,
-    ) -> Bytes {
+    fn message_key(service_id: &ServiceId, sequence_number: u64) -> Bytes {
         let key = InboxKey {
-            partition_key: Some(partition_key),
+            partition_key: Some(service_id.partition_key()),
             service_name: Some(service_id.service_name.clone()),
             service_key: Some(service_id.key.clone()),
             sequence_number: Some(sequence_number),
@@ -122,9 +104,9 @@ mod tests {
         buf.freeze()
     }
 
-    fn inbox_key(partition_key: PartitionKey, service_id: &ServiceId) -> Bytes {
+    fn inbox_key(service_id: &ServiceId) -> Bytes {
         let key = InboxKey {
-            partition_key: Some(partition_key),
+            partition_key: Some(service_id.partition_key()),
             service_name: Some(service_id.service_name.clone()),
             service_key: Some(service_id.key.clone()),
             sequence_number: None,
@@ -135,12 +117,15 @@ mod tests {
 
     #[test]
     fn inbox_key_covers_all_messages_of_a_service() {
-        let prefix_key = inbox_key(1337, &ServiceId::new("svc-1", "key-a"));
+        let prefix_key = inbox_key(&ServiceId::with_partition_key(1337, "svc-1", "key-a"));
 
-        let low_key = message_key(1337, &ServiceId::new("svc-1", "key-a"), 0);
+        let low_key = message_key(&ServiceId::with_partition_key(1337, "svc-1", "key-a"), 0);
         assert!(low_key.starts_with(&prefix_key));
 
-        let high_key = message_key(1337, &ServiceId::new("svc-1", "key-a"), u64::MAX);
+        let high_key = message_key(
+            &ServiceId::with_partition_key(1337, "svc-1", "key-a"),
+            u64::MAX,
+        );
         assert!(high_key.starts_with(&prefix_key));
     }
 
@@ -150,22 +135,24 @@ mod tests {
         // across services
         //
         assert!(
-            message_key(1337, &ServiceId::new("svc-1", ""), 0)
-                < message_key(1337, &ServiceId::new("svc-2", ""), 0)
+            message_key(&ServiceId::with_partition_key(1337, "svc-1", ""), 0)
+                < message_key(&ServiceId::with_partition_key(1337, "svc-2", ""), 0)
         );
         //
         // same service across keys
         //
         assert!(
-            message_key(1337, &ServiceId::new("svc-1", "a"), 0)
-                < message_key(1337, &ServiceId::new("svc-1", "b"), 0)
+            message_key(&ServiceId::with_partition_key(1337, "svc-1", "a"), 0)
+                < message_key(&ServiceId::with_partition_key(1337, "svc-1", "b"), 0)
         );
         //
         // within the same service and key
         //
-        let mut previous_key = message_key(1337, &ServiceId::new("svc-1", "key-a"), 0);
+        let mut previous_key =
+            message_key(&ServiceId::with_partition_key(1337, "svc-1", "key-a"), 0);
         for i in 1..300 {
-            let current_key = message_key(1337, &ServiceId::new("svc-1", "key-a"), i);
+            let current_key =
+                message_key(&ServiceId::with_partition_key(1337, "svc-1", "key-a"), i);
             assert!(previous_key < current_key);
             previous_key = current_key;
         }