[PartitionStore] Allow table key prefix builders to use references

- Allows key prefixes and (get) operations no partition store to be encoded without the extra clone into the builder (via the new builder_ref() type).
- Allows table key designers to use Option<T> by implementing their own KeyEncode trait with their own semantic of what "None" means.

Author: AhmedSoliman

crates/partition-store/src/durable_lsn_tracking.rs

@@ -26,7 +26,9 @@ use restate_types::{identifiers::PartitionId, logs::Lsn};
 
 use crate::SharedState;
 use crate::fsm_table::{PartitionStateMachineKey, fsm_variable};
-use crate::keys::{KeyKind, TableKey};
+#[cfg(test)]
+use crate::keys::EncodeTableKey;
+use crate::keys::{DecodeTableKey, KeyKind};
 
 const APPLIED_LSNS_PROPERTY_PREFIX: &str = "p:";
 

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

@@ -21,7 +21,7 @@ use restate_storage_api::protobuf_types::PartitionStoreProtobufValue;
 use restate_storage_api::{Result, StorageError};
 use restate_types::identifiers::{IdempotencyId, PartitionKey, WithPartitionKey};
 
-use crate::keys::{KeyKind, TableKey, define_table_key};
+use crate::keys::{DecodeTableKey, KeyKind, define_table_key};
 use crate::scan::TableScan;
 use crate::{PartitionStore, PartitionStoreTransaction, StorageAccess, TableKind, break_on_err};
 

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

@@ -24,7 +24,7 @@ use restate_types::identifiers::{PartitionKey, ServiceId, WithPartitionKey};
 use restate_types::message::MessageIndex;
 
 use crate::TableKind::Inbox;
-use crate::keys::{KeyKind, TableKey, define_table_key};
+use crate::keys::{DecodeTableKey, KeyKind, define_table_key};
 use crate::{
     PartitionStore, PartitionStoreTransaction, StorageAccess, TableScan,
     TableScanIterationDecision, break_on_err,
@@ -211,7 +211,7 @@ fn decode_inbox_key_value(mut k: &[u8], mut v: &[u8]) -> Result<SequenceNumberIn
 #[cfg(test)]
 mod tests {
     use crate::inbox_table::InboxKey;
-    use crate::keys::TableKeyPrefix;
+    use crate::keys::EncodeTableKeyPrefix;
     use bytes::Bytes;
     use restate_types::identifiers::{ServiceId, WithPartitionKey};
 

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

@@ -24,7 +24,7 @@ use restate_storage_api::{Result, StorageError};
 use restate_types::identifiers::{InvocationId, InvocationUuid, PartitionKey, WithPartitionKey};
 
 use crate::TableScan::FullScanPartitionKeyRange;
-use crate::keys::{KeyKind, TableKey, define_table_key};
+use crate::keys::{DecodeTableKey, KeyKind, define_table_key};
 use crate::scan::TableScan;
 use crate::{PartitionStore, PartitionStoreTransaction, StorageAccess, TableKind, break_on_err};
 
@@ -289,7 +289,7 @@ impl WriteInvocationStatusTable for PartitionStoreTransaction<'_> {
 
 #[cfg(test)]
 mod tests {
-    use crate::keys::TableKeyPrefix;
+    use crate::keys::EncodeTableKeyPrefix;
 
     use super::*;
 

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

@@ -26,7 +26,7 @@ use restate_types::time::MillisSinceEpoch;
 
 use crate::TableKind::{JournalEvent, State};
 use crate::error::break_on_err;
-use crate::keys::{KeyKind, TableKey, define_table_key};
+use crate::keys::{DecodeTableKey, KeyKind, define_table_key};
 use crate::{
     PartitionStore, PartitionStoreTransaction, StorageAccess, TableScan, TableScanIterationDecision,
 };

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

@@ -24,7 +24,7 @@ use restate_types::identifiers::{
 };
 
 use crate::TableKind::Journal;
-use crate::keys::{KeyKind, TableKey, define_table_key};
+use crate::keys::{DecodeTableKey, EncodeTableKey, KeyKind, define_table_key};
 use crate::{PartitionStore, PartitionStoreTransaction, StorageAccess, TableScan, break_on_err};
 
 define_table_key!(
@@ -473,7 +473,7 @@ impl WriteJournalTable for PartitionStoreTransaction<'_> {
 #[cfg(test)]
 mod tests {
     use crate::journal_table::write_journal_entry_key;
-    use crate::keys::TableKeyPrefix;
+    use crate::keys::EncodeTableKeyPrefix;
     use bytes::Bytes;
     use restate_types::identifiers::{InvocationId, InvocationUuid};
 

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

@@ -31,7 +31,7 @@ use restate_types::journal_v2::{CompletionId, EntryMetadata, NotificationId};
 use restate_types::storage::{StoredRawEntry, StoredRawEntryHeader};
 
 use crate::TableKind::Journal;
-use crate::keys::{KeyKind, TableKey, define_table_key};
+use crate::keys::{DecodeTableKey, EncodeTableKey, KeyKind, define_table_key};
 use crate::owned_iter::OwnedIterator;
 use crate::{PartitionStore, PartitionStoreTransaction, StorageAccess, TableScan, break_on_err};
 
@@ -791,7 +791,7 @@ mod tests {
 
     use super::write_journal_entry_key;
 
-    use crate::keys::TableKeyPrefix;
+    use crate::keys::EncodeTableKeyPrefix;
     use bytes::Bytes;
     use restate_types::identifiers::{InvocationId, InvocationUuid};