[PartitionStore] Separate concrete keys from key builder

Better type separation between keys and key prefix builders leads to removal of unnecessary "expect()" calls, increased efficiency (Option adds 8 bytes on each field on the key), and safer code. You can only deserialize full keys, and you can serialize key builders and keys alike.

Author: AhmedSoliman

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

@@ -25,15 +25,24 @@ define_table_key!(
     DeduplicationKey(partition_id: PaddedPartitionId, producer_id: ProducerId)
 );
 
+#[inline]
+fn create_key(
+    partition_id: impl Into<PaddedPartitionId>,
+    producer_id: ProducerId,
+) -> DeduplicationKey {
+    DeduplicationKey {
+        partition_id: partition_id.into(),
+        producer_id,
+    }
+}
+
 fn get_dedup_sequence_number<S: StorageAccess>(
     storage: &mut S,
     partition_id: PartitionId,
     producer_id: &ProducerId,
 ) -> Result<Option<DedupSequenceNumber>> {
     let _x = RocksDbPerfGuard::new("get-dedup-seq");
-    let key = DeduplicationKey::default()
-        .partition_id(partition_id.into())
-        .producer_id(producer_id.clone());
+    let key = create_key(partition_id, producer_id.clone());
 
     storage.get_value_proto(key)
 }
@@ -62,9 +71,7 @@ impl WriteDeduplicationTable for PartitionStoreTransaction<'_> {
         producer_id: ProducerId,
         dedup_sequence_number: &DedupSequenceNumber,
     ) -> Result<()> {
-        let key = DeduplicationKey::default()
-            .partition_id(self.partition_id().into())
-            .producer_id(producer_id);
+        let key = create_key(self.partition_id(), producer_id);
 
         self.put_kv_proto(key, dedup_sequence_number)
     }

crates/partition-store/src/durable_lsn_tracking.rs

@@ -157,10 +157,8 @@ fn extract_partition_applied_lsn(
     }
 
     let fsm_key = PartitionStateMachineKey::deserialize_from(&mut key)?;
-    if fsm_key.state_id == Some(fsm_variable::APPLIED_LSN)
-        && let Some(padded_partition_id) = fsm_key.partition_id
-    {
-        let partition_id = PartitionId::from(padded_partition_id);
+    if fsm_key.state_id == fsm_variable::APPLIED_LSN {
+        let partition_id = PartitionId::from(fsm_key.partition_id);
         let applied_lsn = SequenceNumber::decode(&mut value).map(u64::from)?.into();
         return Ok(Some((partition_id, applied_lsn)));
     }
@@ -184,8 +182,8 @@ mod tests {
         let lsn = Lsn::from(12345);
 
         let applied_lsn_key = PartitionStateMachineKey {
-            partition_id: Some(partition_id.into()),
-            state_id: Some(fsm_variable::APPLIED_LSN),
+            partition_id: partition_id.into(),
+            state_id: fsm_variable::APPLIED_LSN,
         };
 
         let mut key_buf = BytesMut::new();
@@ -224,8 +222,8 @@ mod tests {
         let lsn = Lsn::from(12345);
 
         let other_fsm_key = PartitionStateMachineKey {
-            partition_id: Some(partition_id.into()),
-            state_id: Some(fsm_variable::INBOX_SEQ_NUMBER),
+            partition_id: partition_id.into(),
+            state_id: fsm_variable::INBOX_SEQ_NUMBER,
         };
 
         let mut key_buf = BytesMut::new();

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

@@ -29,6 +29,17 @@ define_table_key!(
     PartitionStateMachineKey(partition_id: PaddedPartitionId, state_id: u64)
 );
 
+#[inline]
+fn create_key(
+    partition_id: impl Into<PaddedPartitionId>,
+    state_id: u64,
+) -> PartitionStateMachineKey {
+    PartitionStateMachineKey {
+        partition_id: partition_id.into(),
+        state_id,
+    }
+}
+
 pub(crate) mod fsm_variable {
     pub(crate) const INBOX_SEQ_NUMBER: u64 = 0;
     pub(crate) const OUTBOX_SEQ_NUMBER: u64 = 1;
@@ -52,10 +63,7 @@ fn get<T: PartitionStoreProtobufValue, S: StorageAccess>(
 where
     <<T as PartitionStoreProtobufValue>::ProtobufType as TryInto<T>>::Error: Into<anyhow::Error>,
 {
-    let key = PartitionStateMachineKey::default()
-        .partition_id(partition_id.into())
-        .state_id(state_id);
-    storage.get_value_proto(key)
+    storage.get_value_proto(create_key(partition_id, state_id))
 }
 
 /// Forces a read from persistent storage, bypassing memtables and block cache.
@@ -67,10 +75,7 @@ fn get_durable<T: PartitionStoreProtobufValue, S: StorageAccess>(
 where
     <<T as PartitionStoreProtobufValue>::ProtobufType as TryInto<T>>::Error: Into<anyhow::Error>,
 {
-    let key = PartitionStateMachineKey::default()
-        .partition_id(partition_id.into())
-        .state_id(state_id);
-    storage.get_durable_value(key)
+    storage.get_durable_value(create_key(partition_id, state_id))
 }
 
 fn put<S: StorageAccess>(
@@ -79,9 +84,10 @@ fn put<S: StorageAccess>(
     state_id: u64,
     state_value: &(impl PartitionStoreProtobufValue + Clone + 'static),
 ) -> Result<()> {
-    let key = PartitionStateMachineKey::default()
-        .partition_id(partition_id.into())
-        .state_id(state_id);
+    let key = PartitionStateMachineKey {
+        partition_id: partition_id.into(),
+        state_id,
+    };
     storage.put_kv_proto(key, state_value)
 }
 
@@ -149,9 +155,7 @@ impl ReadFsmTable for PartitionStore {
     }
 
     async fn get_schema(&mut self) -> Result<Option<Schema>> {
-        let key = PartitionStateMachineKey::default()
-            .partition_id(self.partition_id().into())
-            .state_id(fsm_variable::SERVICES_SCHEMA_METADATA);
+        let key = create_key(self.partition_id(), fsm_variable::SERVICES_SCHEMA_METADATA);
         self.get_value_storage_codec(key)
     }
 }
@@ -203,9 +207,7 @@ impl WriteFsmTable for PartitionStoreTransaction<'_> {
     }
 
     fn put_schema(&mut self, schema: &Schema) -> Result<()> {
-        let key = PartitionStateMachineKey::default()
-            .partition_id(self.partition_id().into())
-            .state_id(fsm_variable::SERVICES_SCHEMA_METADATA);
+        let key = create_key(self.partition_id(), fsm_variable::SERVICES_SCHEMA_METADATA);
         self.put_kv_storage_codec(key, schema)
     }
 }

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

@@ -37,20 +37,21 @@ define_table_key!(
     )
 );
 
+#[inline]
 fn create_key(idempotency_id: &IdempotencyId) -> IdempotencyKey {
-    IdempotencyKey::default()
-        .partition_key(idempotency_id.partition_key())
-        .service_name(idempotency_id.service_name.clone())
-        .service_key(
-            idempotency_id
-                .service_key
-                .as_ref()
-                .cloned()
-                .unwrap_or_default()
-                .into_bytes(),
-        )
-        .service_handler(idempotency_id.service_handler.clone())
-        .idempotency_key(idempotency_id.idempotency_key.clone())
+    IdempotencyKey {
+        partition_key: idempotency_id.partition_key(),
+        service_name: idempotency_id.service_name.clone(),
+        service_key: idempotency_id
+            .service_key
+            .as_ref()
+            .cloned()
+            .unwrap_or_default()
+            .into_bytes(),
+
+        service_handler: idempotency_id.service_handler.clone(),
+        idempotency_key: idempotency_id.idempotency_key.clone(),
+    }
 }
 
 fn get_idempotency_metadata<S: StorageAccess>(
@@ -103,17 +104,13 @@ impl ScanIdempotencyTable for PartitionStore {
                 let idempotency_metadata = break_on_err(IdempotencyMetadata::decode(&mut value))?;
 
                 let idempotency_id = IdempotencyId::new(
-                    break_on_err(key.service_name_ok_or())?.clone(),
-                    break_on_err(
-                        key.service_key
-                            .clone()
-                            .map(|b| {
-                                ByteString::try_from(b).map_err(|e| StorageError::Generic(e.into()))
-                            })
-                            .transpose(),
-                    )?,
-                    break_on_err(key.service_handler_ok_or())?.clone(),
-                    break_on_err(key.idempotency_key_ok_or())?.clone(),
+                    key.service_name,
+                    Some(break_on_err(
+                        ByteString::try_from(key.service_key)
+                            .map_err(|e| StorageError::Generic(e.into())),
+                    )?),
+                    key.service_handler,
+                    key.idempotency_key,
                 );
 
                 f((idempotency_id, idempotency_metadata)).map_break(Ok)

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

@@ -8,7 +8,6 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use std::io::Cursor;
 use std::ops::ControlFlow;
 
 use bytestring::ByteString;
@@ -46,7 +45,7 @@ fn peek_inbox<S: StorageAccess>(
     storage: &mut S,
     service_id: &ServiceId,
 ) -> Result<Option<SequenceNumberInboxEntry>> {
-    let key = InboxKey::default()
+    let key = InboxKey::builder()
         .partition_key(service_id.partition_key())
         .service_name(service_id.service_name.clone())
         .service_key(service_id.key.clone());
@@ -67,7 +66,7 @@ fn inbox<S: StorageAccess>(
     storage: &mut S,
     service_id: &ServiceId,
 ) -> Result<impl Stream<Item = Result<SequenceNumberInboxEntry>> + Send> {
-    let key = InboxKey::default()
+    let key = InboxKey::builder()
         .partition_key(service_id.partition_key())
         .service_name(service_id.service_name.clone())
         .service_key(service_id.key.clone());
@@ -113,10 +112,13 @@ impl ScanInboxTable for PartitionStore {
             TableScan::FullScanPartitionKeyRange::<InboxKey>(range),
             move |(mut key, mut value)| {
                 let key = break_on_err(InboxKey::deserialize_from(&mut key))?;
-                let sequence_number = *break_on_err(key.sequence_number_ok_or())?;
                 let inbox_entry = break_on_err(InboxEntry::decode(&mut value))?;
 
-                f(SequenceNumberInboxEntry::new(sequence_number, inbox_entry)).map_break(Ok)
+                f(SequenceNumberInboxEntry::new(
+                    key.sequence_number,
+                    inbox_entry,
+                ))
+                .map_break(Ok)
             },
         )
         .map_err(|_| StorageError::OperationalError)
@@ -150,11 +152,12 @@ impl WriteInboxTable for PartitionStoreTransaction<'_> {
         let service_id = inbox_entry.service_id();
         self.assert_partition_key(service_id)?;
 
-        let key = InboxKey::default()
-            .partition_key(service_id.partition_key())
-            .service_name(service_id.service_name.clone())
-            .service_key(service_id.key.clone())
-            .sequence_number(inbox_sequence_number);
+        let key = InboxKey {
+            partition_key: service_id.partition_key(),
+            service_name: service_id.service_name.clone(),
+            service_key: service_id.key.clone(),
+            sequence_number: inbox_sequence_number,
+        };
 
         self.put_kv_proto(key, inbox_entry)
     }
@@ -185,50 +188,48 @@ fn delete_inbox_entry(
     service_id: &ServiceId,
     sequence_number: u64,
 ) -> Result<()> {
-    let key = InboxKey::default()
-        .partition_key(service_id.partition_key())
-        .service_name(service_id.service_name.clone())
-        .service_key(service_id.key.clone())
-        .sequence_number(sequence_number);
+    let key = InboxKey {
+        partition_key: service_id.partition_key(),
+        service_name: service_id.service_name.clone(),
+        service_key: service_id.key.clone(),
+        sequence_number,
+    };
 
     txn.delete_key(&key)
 }
 
-fn decode_inbox_key_value(k: &[u8], mut v: &[u8]) -> Result<SequenceNumberInboxEntry> {
-    let key = InboxKey::deserialize_from(&mut Cursor::new(k))?;
-    let sequence_number = *key.sequence_number_ok_or()?;
-
+fn decode_inbox_key_value(mut k: &[u8], mut v: &[u8]) -> Result<SequenceNumberInboxEntry> {
+    let key = InboxKey::deserialize_from(&mut k)?;
     let inbox_entry = InboxEntry::decode(&mut v)?;
 
-    Ok(SequenceNumberInboxEntry::new(sequence_number, inbox_entry))
+    Ok(SequenceNumberInboxEntry::new(
+        key.sequence_number,
+        inbox_entry,
+    ))
 }
 
 #[cfg(test)]
 mod tests {
     use crate::inbox_table::InboxKey;
-    use crate::keys::TableKey;
-    use bytes::{Bytes, BytesMut};
+    use crate::keys::TableKeyPrefix;
+    use bytes::Bytes;
     use restate_types::identifiers::{ServiceId, WithPartitionKey};
 
     fn message_key(service_id: &ServiceId, sequence_number: u64) -> Bytes {
         let key = InboxKey {
-            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),
+            partition_key: service_id.partition_key(),
+            service_name: service_id.service_name.clone(),
+            service_key: service_id.key.clone(),
+            sequence_number,
         };
-        let mut buf = BytesMut::new();
-        key.serialize_to(&mut buf);
-        buf.freeze()
+        key.serialize().freeze()
     }
 
     fn inbox_key(service_id: &ServiceId) -> Bytes {
-        let key = InboxKey {
-            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,
-        };
+        let key = InboxKey::builder()
+            .partition_key(service_id.partition_key())
+            .service_name(service_id.service_name.clone())
+            .service_key(service_id.key.clone());
 
         key.serialize().freeze()
     }