Introduce restate-util-sharding crate

Part of the effort to decompose the restate-types monolith into focused,
composable utility crates.

Add a new `restate-util-sharding` crate (in `util/sharding/`) following the
`util/string` pattern. It defines the foundational sharding types:

- `KeyRange`: a compact `Copy` newtype over `std::range::RangeInclusive<u64>`.
  16 bytes (vs 24 for std::ops::RangeInclusive), wire-format compatible serde
  and bilrost encoding. Provides iter(), is_overlapping(), and split() helpers.
- `PartitionKey`: type alias for u64
- `WithPartitionKey`: trait for types that carry a partition key

The crate is re-exported from `restate-types::sharding` for ergonomic access.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: AhmedSoliman

Cargo.lock

@@ -83,6 +83,7 @@ restate-service-protocol = { path = "crates/service-protocol" }
 restate-service-protocol-v4 = { path = "crates/service-protocol-v4" }
 restate-storage-api = { path = "crates/storage-api" }
 restate-storage-query-datafusion = { path = "crates/storage-query-datafusion" }
+restate-util-sharding = { path = "util/sharding" }
 restate-util-string = { path = "util/string" }
 restate-test-util = { path = "crates/test-util" }
 restate-time-util = { path = "crates/time-util" }

Cargo.toml

@@ -18,4 +18,5 @@ bytestring = { workspace = true }
 
 [dev-dependencies]
 rand = { workspace = true }
+restate-util-sharding = { workspace = true, features = ["bilrost"] }
 static_assertions = { workspace = true }

crates/encoding/Cargo.toml

@@ -30,6 +30,7 @@ restate-memory = { workspace = true }
 restate-serde-util = { workspace = true }
 restate-test-util = { workspace = true, optional = true }
 restate-time-util = { workspace = true, features = ["serde", "serde_with"] }
+restate-util-sharding = { workspace = true, features = ["serde", "bilrost"] }
 restate-util-string = { workspace = true, features = ["serde", "bilrost"] }
 restate-utoipa = { workspace = true }
 

crates/types/Cargo.toml

@@ -13,6 +13,7 @@
 mod partitioned;
 
 pub use partitioned::PartitionedResourceId;
+pub use restate_util_sharding::{PartitionKey, WithPartitionKey};
 
 use std::cell::RefCell;
 use std::fmt::{self, Display, Formatter};
@@ -155,10 +156,6 @@ pub type PartitionLeaderEpoch = (PartitionId, LeaderEpoch);
 // Just an alias
 pub type EntryIndex = u32;
 
-/// Identifying to which partition a key belongs. This is unlike the [`PartitionId`]
-/// which identifies a consecutive range of partition keys.
-pub type PartitionKey = u64;
-
 /// Returns the partition key computed from either the service_key, or idempotency_key, if possible
 fn deterministic_partition_key(
     service_key: Option<&str>,
@@ -169,12 +166,6 @@ fn deterministic_partition_key(
         .or_else(|| idempotency_key.map(partitioner::HashPartitioner::compute_partition_key))
 }
 
-/// Trait for data structures that have a partition key
-pub trait WithPartitionKey {
-    /// Returns the partition key
-    fn partition_key(&self) -> PartitionKey;
-}
-
 /// A family of resource identifiers that tracks the timestamp of its creation.
 pub trait TimestampAwareId {
     /// The timestamp when this ID was created.
@@ -607,12 +598,6 @@ impl WithPartitionKey for InvocationId {
     }
 }
 
-impl<T: WithInvocationId> WithPartitionKey for T {
-    fn partition_key(&self) -> PartitionKey {
-        self.invocation_id().partition_key
-    }
-}
-
 impl Display for InvocationId {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         // encode the id such that it is possible to do a string prefix search for a
@@ -767,6 +752,12 @@ impl From<(InvocationId, EntryIndex)> for JournalEntryId {
     }
 }
 
+impl WithPartitionKey for JournalEntryId {
+    fn partition_key(&self) -> PartitionKey {
+        self.invocation_id.partition_key()
+    }
+}
+
 impl WithInvocationId for JournalEntryId {
     fn invocation_id(&self) -> InvocationId {
         self.invocation_id

crates/types/src/identifiers.rs

@@ -417,6 +417,12 @@ impl InvocationRequest {
     }
 }
 
+impl WithPartitionKey for InvocationRequest {
+    fn partition_key(&self) -> PartitionKey {
+        self.header.invocation_id().partition_key()
+    }
+}
+
 impl WithInvocationId for InvocationRequest {
     fn invocation_id(&self) -> InvocationId {
         self.header.invocation_id()
@@ -580,6 +586,12 @@ impl JournalCompletionTarget {
     }
 }
 
+impl WithPartitionKey for JournalCompletionTarget {
+    fn partition_key(&self) -> PartitionKey {
+        self.caller_id.partition_key()
+    }
+}
+
 impl WithInvocationId for JournalCompletionTarget {
     fn invocation_id(&self) -> InvocationId {
         self.caller_id
@@ -597,6 +609,12 @@ pub struct InvocationResponse {
     pub result: ResponseResult,
 }
 
+impl WithPartitionKey for InvocationResponse {
+    fn partition_key(&self) -> PartitionKey {
+        self.target.invocation_id().partition_key()
+    }
+}
+
 impl WithInvocationId for InvocationResponse {
     fn invocation_id(&self) -> InvocationId {
         self.target.invocation_id()
@@ -1351,6 +1369,12 @@ pub struct NotifySignalRequest {
     pub signal: Signal,
 }
 
+impl WithPartitionKey for NotifySignalRequest {
+    fn partition_key(&self) -> PartitionKey {
+        self.invocation_id.partition_key()
+    }
+}
+
 impl WithInvocationId for NotifySignalRequest {
     fn invocation_id(&self) -> InvocationId {
         self.invocation_id

crates/types/src/invocation/mod.rs

@@ -81,6 +81,11 @@ pub mod memory {
     pub use restate_memory::*;
 }
 
+// Re-export restate-util-sharding crate for key range utilities.
+pub mod sharding {
+    pub use restate_util_sharding::*;
+}
+
 // Re-export metrics' SharedString (Space-efficient Cow + RefCounted variant)
 pub type SharedString = metrics::SharedString;
 

crates/types/src/lib.rs

@@ -242,15 +242,17 @@ impl HasRecordKeys for Envelope {
             Command::TruncateOutbox(_) => Keys::Single(self.partition_key()),
             Command::ProxyThrough(_) => Keys::Single(self.partition_key()),
             Command::AttachInvocation(_) => Keys::Single(self.partition_key()),
-            Command::ResumeInvocation(req) => Keys::Single(req.partition_key()),
-            Command::RestartAsNewInvocation(req) => Keys::Single(req.partition_key()),
+            Command::ResumeInvocation(req) => Keys::Single(req.invocation_id.partition_key()),
+            Command::RestartAsNewInvocation(req) => Keys::Single(req.invocation_id.partition_key()),
             // todo: Handle journal entries that request cross-partition invocations
             Command::InvokerEffect(effect) => Keys::Single(effect.invocation_id.partition_key()),
             Command::Timer(timer) => Keys::Single(timer.invocation_id().partition_key()),
             Command::ScheduleTimer(timer) => Keys::Single(timer.invocation_id().partition_key()),
             Command::InvocationResponse(response) => Keys::Single(response.partition_key()),
             Command::NotifySignal(sig) => Keys::Single(sig.partition_key()),
-            Command::NotifyGetInvocationOutputResponse(res) => Keys::Single(res.partition_key()),
+            Command::NotifyGetInvocationOutputResponse(res) => {
+                Keys::Single(res.target.partition_key())
+            }
             Command::UpsertSchema(schema) => schema.partition_key_range.clone(),
             Command::VQWaitingToRunning(_) => Keys::Single(self.partition_key()),
             Command::VQYieldRunning(_) => Keys::Single(self.partition_key()),

crates/wal-protocol/src/lib.rs

@@ -0,0 +1,22 @@
+[package]
+name = "restate-util-sharding"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+publish = false
+
+[features]
+default = []
+serde = ["dep:serde"]
+bilrost = ["dep:bilrost"]
+
+[dependencies]
+restate-workspace-hack = { workspace = true }
+
+bilrost = { workspace = true, optional = true }
+serde = { workspace = true, optional = true }
+
+[dev-dependencies]
+restate-util-sharding = { path = ".", default-features = false, features = ["serde", "bilrost"] }