Migrate partition key ranges to KeyRange and extract sharding types

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

Replace all ~85 uses of std::ops::RangeInclusive<PartitionKey> with the new
KeyRange type from restate-util-sharding. Move PartitionId and
EqualSizedPartitionPartitioner into restate-util-sharding alongside KeyRange,
re-exporting from restate-types for backwards compatibility.

The migration is mechanical: .start()/.end() now return u64 by value instead
of &u64, .clone() calls are removed (KeyRange is Copy), and construction
switches from a..=b syntax to KeyRange::new(a, b).

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

Author: AhmedSoliman

AGENTS.md

@@ -52,6 +52,7 @@ If these rules conflict with normal behavior, always follow the rules above.
 1. Check crates/cli-util/README.md when doing CLI changes to make sure that you are adhering to the CLI style guide.
 1. New or deprecated config options must have `/// Since vX.Y.Z` in their doc comment.
 1. Use `ByteCount::from(value)` (from `restate_memory`) when displaying byte sizes in errors/logs.
+1. Use `KeyRange` (from `restate_sharding`, re-exported via `restate_types::sharding::KeyRange`) instead of `std::ops::RangeInclusive<PartitionKey>` for partition key ranges. `KeyRange` is `Copy`, 16 bytes (vs 24), and has wire-compatible serde/bilrost encoding.
 
 
 # Validation Before Committing Changes

Cargo.lock

@@ -11,6 +11,7 @@ publish = false
 restate-workspace-hack = { workspace = true }
 
 restate-encoding-derive = { version = "0.1.0", path = "derive" }
+restate-sharding = { workspace = true }
 
 bilrost = { workspace = true }
 bytes = { workspace = true }

crates/encoding/Cargo.toml

@@ -98,6 +98,7 @@ where
 
 impl<V> NetSerde for HashSet<V> where V: NetSerde {}
 impl<Idx> NetSerde for RangeInclusive<Idx> where Idx: NetSerde {}
+impl NetSerde for restate_sharding::PartitionId {}
 impl<T> NetSerde for Arc<T> where T: NetSerde {}
 impl<T> NetSerde for Arc<[T]> where T: NetSerde {}
 impl<T> NetSerde for Box<T> where T: NetSerde {}
@@ -131,4 +132,49 @@ mod test {
 
         assert_eq!(x.id.0, y.id);
     }
+
+    /// Validates that `KeyRange`'s general bilrost encoding produces the same
+    /// wire format as `RangeInclusive<u64>` with `RestateEncoding`. This ensures
+    /// that `KeyRange` fields can use `#[bilrost(N)]` and remain wire-compatible
+    /// with the old `#[bilrost(tag(N), encoding(RestateEncoding))] RangeInclusive<u64>`.
+    #[test]
+    fn key_range_wire_compat_with_range_inclusive() {
+        use restate_sharding::KeyRange;
+
+        use super::RestateEncoding;
+
+        #[derive(Debug, PartialEq, bilrost::Message)]
+        struct WithKeyRange {
+            #[bilrost(1)]
+            range: KeyRange,
+        }
+
+        #[derive(Debug, PartialEq, bilrost::Message)]
+        struct WithRangeInclusive {
+            #[bilrost(tag(1), encoding(RestateEncoding))]
+            range: std::ops::RangeInclusive<u64>,
+        }
+
+        for (start, end) in [(0u64, 0u64), (1, 100), (0, u64::MAX), (42, 42)] {
+            let kr = KeyRange::new(start, end);
+            let ri = start..=end;
+
+            let kr_bytes = WithKeyRange { range: kr }.encode_to_vec();
+            let ri_bytes = WithRangeInclusive { range: ri.clone() }.encode_to_vec();
+
+            assert_eq!(
+                kr_bytes, ri_bytes,
+                "wire format mismatch for range ({start}, {end})"
+            );
+
+            // Cross-decode
+            let decoded: WithRangeInclusive =
+                WithRangeInclusive::decode(&*kr_bytes).expect("cross-decode KeyRange→RI");
+            assert_eq!(decoded.range, ri);
+
+            let decoded: WithKeyRange =
+                WithKeyRange::decode(&*ri_bytes).expect("cross-decode RI→KeyRange");
+            assert_eq!(decoded.range, kr);
+        }
+    }
 }

crates/encoding/src/lib.rs

@@ -440,7 +440,7 @@ mod test {
             let partition = pt.get(&partition_id).unwrap();
             client
                 .ingest(
-                    *partition.key_range.start(),
+                    partition.key_range.start(),
                     InputRecord::from_str(format!("partition {p}")),
                 )
                 .await

crates/ingestion-client/src/client.rs

@@ -8,17 +8,16 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use std::ops::RangeInclusive;
-
-use restate_types::vqueues::VQueueId;
 use tokio::sync::mpsc;
 
 use restate_errors::NotRunningError;
 use restate_futures_util::concurrency::Permit;
 use restate_memory::MemoryLease;
-use restate_types::identifiers::{EntryIndex, InvocationId, PartitionKey, PartitionLeaderEpoch};
+use restate_types::identifiers::{EntryIndex, InvocationId, PartitionLeaderEpoch};
 use restate_types::invocation::InvocationTarget;
 use restate_types::journal_v2::{CommandIndex, NotificationId};
+use restate_types::sharding::KeyRange;
+use restate_types::vqueues::VQueueId;
 
 use super::Effect;
 
@@ -89,7 +88,7 @@ pub trait InvokerHandle<SR> {
     fn register_partition(
         &mut self,
         partition: PartitionLeaderEpoch,
-        partition_key_range: RangeInclusive<PartitionKey>,
+        partition_key_range: KeyRange,
         storage_reader: SR,
         sender: mpsc::Sender<Box<Effect>>,
     ) -> Result<(), NotRunningError>;

crates/invoker-api/src/handle.rs

@@ -26,7 +26,6 @@ pub mod test_util {
     use super::*;
     use std::convert::Infallible;
     use std::marker::PhantomData;
-    use std::ops::RangeInclusive;
 
     use bytes::Bytes;
     use tokio::sync::mpsc::Sender;
@@ -36,10 +35,9 @@ pub mod test_util {
     use restate_memory::{
         IgnorePinnableMemoryStream, LocalMemoryLease, LocalMemoryPool, MemoryLease,
     };
-    use restate_types::identifiers::{
-        EntryIndex, InvocationId, PartitionKey, PartitionLeaderEpoch, ServiceId,
-    };
+    use restate_types::identifiers::{EntryIndex, InvocationId, PartitionLeaderEpoch, ServiceId};
     use restate_types::invocation::{InvocationTarget, ServiceInvocationSpanContext};
+    use restate_types::sharding::KeyRange;
     use restate_types::time::MillisSinceEpoch;
     use restate_types::vqueues::VQueueId;
 
@@ -238,7 +236,7 @@ pub mod test_util {
         fn register_partition(
             &mut self,
             _partition: PartitionLeaderEpoch,
-            _partition_key_range: RangeInclusive<PartitionKey>,
+            _partition_key_range: KeyRange,
             _storage_reader: SR,
             _sender: Sender<Box<Effect>>,
         ) -> Result<(), NotRunningError> {

crates/invoker-api/src/lib.rs

@@ -8,15 +8,17 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::future::Future;
+use std::time::SystemTime;
+
 use codederror::Code;
+
 use restate_types::errors::InvocationError;
-use restate_types::identifiers::{DeploymentId, InvocationId, PartitionKey};
+use restate_types::identifiers::{DeploymentId, InvocationId};
 use restate_types::identifiers::{LeaderEpoch, PartitionId, PartitionLeaderEpoch};
 use restate_types::journal::{EntryIndex, EntryType};
 use restate_types::service_protocol::ServiceProtocolVersion;
-use std::future::Future;
-use std::ops::RangeInclusive;
-use std::time::SystemTime;
+use restate_types::sharding::KeyRange;
 
 // -- Status data structure
 
@@ -125,10 +127,7 @@ pub trait StatusHandle {
     /// filtered by the partition key range
     ///
     /// The data returned by this method is eventually consistent.
-    fn read_status(
-        &self,
-        keys: RangeInclusive<PartitionKey>,
-    ) -> impl Future<Output = Self::Iterator> + Send;
+    fn read_status(&self, keys: KeyRange) -> impl Future<Output = Self::Iterator> + Send;
 }
 
 #[cfg(any(test, feature = "test-util"))]
@@ -148,7 +147,7 @@ pub mod test_util {
     impl StatusHandle for MockStatusHandle {
         type Iterator = std::vec::IntoIter<InvocationStatusReport>;
 
-        async fn read_status(&self, _keys: RangeInclusive<PartitionKey>) -> Self::Iterator {
+        async fn read_status(&self, _keys: KeyRange) -> Self::Iterator {
             self.0.clone().into_iter()
         }
     }

crates/invoker-api/src/status_handle.rs

@@ -8,16 +8,15 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use std::ops::RangeInclusive;
-
 use tokio::sync::mpsc;
 
 use restate_errors::NotRunningError;
 use restate_futures_util::concurrency::Permit;
 use restate_memory::MemoryLease;
-use restate_types::identifiers::{EntryIndex, InvocationId, PartitionKey, PartitionLeaderEpoch};
+use restate_types::identifiers::{EntryIndex, InvocationId, PartitionLeaderEpoch};
 use restate_types::invocation::InvocationTarget;
 use restate_types::journal_v2::{CommandIndex, NotificationId};
+use restate_types::sharding::KeyRange;
 use restate_types::vqueues::VQueueId;
 
 use restate_invoker_api::{Effect, InvocationStatusReport, StatusHandle};
@@ -93,7 +92,7 @@ pub(crate) enum InputCommand<SR> {
     // needed for dynamic registration at Invoker
     RegisterPartition {
         partition: PartitionLeaderEpoch,
-        partition_key_range: RangeInclusive<PartitionKey>,
+        partition_key_range: KeyRange,
         storage_reader: SR,
         sender: mpsc::Sender<Box<Effect>>,
     },
@@ -241,7 +240,7 @@ impl<SR: Send> restate_invoker_api::InvokerHandle<SR> for InvokerHandle<SR> {
     fn register_partition(
         &mut self,
         partition: PartitionLeaderEpoch,
-        partition_key_range: RangeInclusive<PartitionKey>,
+        partition_key_range: KeyRange,
         storage_reader: SR,
         sender: mpsc::Sender<Box<Effect>>,
     ) -> Result<(), NotRunningError> {
@@ -259,10 +258,7 @@ impl<SR: Send> restate_invoker_api::InvokerHandle<SR> for InvokerHandle<SR> {
 #[derive(Debug, Clone)]
 pub struct ChannelStatusReader(
     pub(super)  mpsc::UnboundedSender<
-        restate_futures_util::command::Command<
-            RangeInclusive<PartitionKey>,
-            Vec<InvocationStatusReport>,
-        >,
+        restate_futures_util::command::Command<KeyRange, Vec<InvocationStatusReport>>,
     >,
 );
 
@@ -272,7 +268,7 @@ impl StatusHandle for ChannelStatusReader {
         std::vec::IntoIter<InvocationStatusReport>,
     >;
 
-    async fn read_status(&self, keys: RangeInclusive<PartitionKey>) -> Self::Iterator {
+    async fn read_status(&self, keys: KeyRange) -> Self::Iterator {
         let (cmd, rx) = restate_futures_util::command::Command::prepare(keys);
         if self.0.send(cmd).is_err() {
             return itertools::Either::Left(std::iter::empty::<InvocationStatusReport>());

crates/invoker-impl/src/input_command.rs

@@ -20,7 +20,7 @@ mod status_store;
 
 use std::collections::{HashMap, HashSet};
 use std::io::ErrorKind;
-use std::ops::RangeInclusive;
+use std::ops::RangeBounds;
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::time::{Duration, Instant, SystemTime};
@@ -49,7 +49,7 @@ use restate_service_client::{AssumeRoleCacheMode, ServiceClient};
 use restate_time_util::DurationExt;
 use restate_types::config::{Configuration, InvokerOptions, ServiceClientOptions};
 use restate_types::deployment::PinnedDeployment;
-use restate_types::identifiers::{DeploymentId, InvocationId, PartitionKey, WithPartitionKey};
+use restate_types::identifiers::{DeploymentId, InvocationId, WithPartitionKey};
 use restate_types::identifiers::{PartitionId, PartitionLeaderEpoch};
 use restate_types::invocation::InvocationTarget;
 use restate_types::journal::EntryIndex;
@@ -61,6 +61,7 @@ use restate_types::journal_v2::{CommandIndex, EntryMetadata, NotificationId};
 use restate_types::live::{Live, LiveLoad};
 use restate_types::schema::deployment::DeploymentResolver;
 use restate_types::schema::invocation_target::InvocationTargetResolver;
+use restate_types::sharding::KeyRange;
 use tokio_util::time::DelayQueue;
 use tokio_util::time::delay_queue::Key as RetryTimerKey;
 
@@ -203,10 +204,7 @@ pub struct Service<StorageReader, EntryEnricher, Schemas> {
     // Used for constructing the invoker sender and status reader
     input_tx: mpsc::UnboundedSender<InputCommand<StorageReader>>,
     status_tx: mpsc::UnboundedSender<
-        restate_futures_util::command::Command<
-            RangeInclusive<PartitionKey>,
-            Vec<InvocationStatusReport>,
-        >,
+        restate_futures_util::command::Command<KeyRange, Vec<InvocationStatusReport>>,
     >,
     // For the segment queue
     tmp_dir: PathBuf,
@@ -386,10 +384,7 @@ where
 struct ServiceInner<InvocationTaskRunner, Schemas, StorageReader> {
     input_rx: mpsc::UnboundedReceiver<InputCommand<StorageReader>>,
     status_rx: mpsc::UnboundedReceiver<
-        restate_futures_util::command::Command<
-            RangeInclusive<PartitionKey>,
-            Vec<InvocationStatusReport>,
-        >,
+        restate_futures_util::command::Command<KeyRange, Vec<InvocationStatusReport>>,
     >,
 
     // Channel to communicate with invocation tasks
@@ -449,10 +444,10 @@ where
 
         tokio::select! {
             Some(cmd) = self.status_rx.recv() => {
-                let keys = cmd.payload();
+                let keys = *cmd.payload();
                 let statuses = self
                     .invocation_state_machine_manager
-                    .registered_partitions_with_keys(keys.clone())
+                    .registered_partitions_with_keys(keys)
                     .flat_map(|partition| self.status_store.status_for_partition(partition))
                     .filter(|status| keys.contains(&status.invocation_id().partition_key()))
                     .collect();
@@ -607,7 +602,7 @@ where
     fn handle_register_partition(
         &mut self,
         partition: PartitionLeaderEpoch,
-        partition_key_range: RangeInclusive<PartitionKey>,
+        partition_key_range: KeyRange,
         storage_reader: IR,
         sender: mpsc::Sender<Box<Effect>>,
     ) {
@@ -1850,10 +1845,7 @@ mod tests {
         ) -> (
             mpsc::UnboundedSender<InputCommand<IR>>,
             mpsc::UnboundedSender<
-                restate_futures_util::command::Command<
-                    RangeInclusive<PartitionKey>,
-                    Vec<InvocationStatusReport>,
-                >,
+                restate_futures_util::command::Command<KeyRange, Vec<InvocationStatusReport>>,
             >,
             Self,
         ) {
@@ -1887,7 +1879,7 @@ mod tests {
             let (partition_tx, partition_rx) = mpsc::channel(1024);
             self.handle_register_partition(
                 MOCK_PARTITION,
-                RangeInclusive::new(0, 0),
+                KeyRange::new(0, 0),
                 storage_reader,
                 partition_tx,
             );
@@ -2139,7 +2131,7 @@ mod tests {
         handle
             .register_partition(
                 partition_leader_epoch,
-                RangeInclusive::new(0, 0),
+                KeyRange::new(0, 0),
                 EmptyStorageReader,
                 output_tx,
             )