Add partition-store schema metadata accessors

## Summary
- add a dedicated metadata slot and read/write helpers for storing service schemas in the partition-store FSM table
- extend the storage-api FSM table traits with schema accessors so callers can persist schema metadata

Author: muhamadazmy

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

@@ -17,6 +17,7 @@ use restate_types::SemanticRestateVersion;
 use restate_types::identifiers::PartitionId;
 use restate_types::logs::Lsn;
 use restate_types::message::MessageIndex;
+use restate_types::schema::Schema;
 
 use crate::TableKind::PartitionStateMachine;
 use crate::keys::{KeyKind, define_table_key};
@@ -37,7 +38,10 @@ pub(crate) mod fsm_variable {
     pub(crate) const PARTITION_DURABILITY: u64 = 4;
 
     /// Schema versions are represented as a strictly monotonically increasing number.
+    /// This represent the partition storage schema version, not the user services schema.
     pub(crate) const SCHEMA_VERSION: u64 = 5;
+
+    pub(crate) const SERVICES_SCHEMA_METADATA: u64 = 6;
 }
 
 fn get<T: PartitionStoreProtobufValue, S: StorageAccess>(
@@ -143,6 +147,13 @@ impl ReadFsmTable for PartitionStore {
             fsm_variable::PARTITION_DURABILITY,
         )
     }
+
+    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);
+        self.get_value_storage_codec(key)
+    }
 }
 
 impl WriteFsmTable for PartitionStoreTransaction<'_> {
@@ -190,4 +201,11 @@ impl WriteFsmTable for PartitionStoreTransaction<'_> {
             durability,
         )
     }
+
+    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);
+        self.put_kv_storage_codec(key, schema)
+    }
 }

crates/partition-store/src/partition_store.rs

@@ -38,6 +38,8 @@ use restate_types::identifiers::{PartitionId, PartitionKey, SnapshotId, WithPart
 use restate_types::logs::Lsn;
 use restate_types::partitions::Partition;
 use restate_types::storage::StorageCodec;
+use restate_types::storage::StorageDecode;
+use restate_types::storage::StorageEncode;
 
 use crate::fsm_table::{get_locally_durable_lsn, get_schema_version, put_schema_version};
 use crate::keys::KeyKind;
@@ -1075,17 +1077,25 @@ pub(crate) trait StorageAccess {
         &mut self,
         key: K,
         value: &V,
+    ) -> Result<()> {
+        self.put_kv_storage_codec(
+            key,
+            &ProtobufStorageWrapper::<V::ProtobufType>(value.clone().into()),
+        )
+    }
+
+    #[inline]
+    fn put_kv_storage_codec<K: TableKey, V: StorageEncode + 'static>(
+        &mut self,
+        key: K,
+        value: &V,
     ) -> Result<()> {
         let key_buffer = self.cleared_key_buffer_mut(key.serialized_length());
         key.serialize_to(key_buffer);
         let key_buffer = key_buffer.split();
 
         let value_buffer = self.cleared_value_buffer_mut(0);
-        StorageCodec::encode(
-            &ProtobufStorageWrapper::<V::ProtobufType>(value.clone().into()),
-            value_buffer,
-        )
-        .map_err(|e| StorageError::Generic(e.into()))?;
+        StorageCodec::encode(value, value_buffer).map_err(|e| StorageError::Generic(e.into()))?;
         let value_buffer = value_buffer.split();
 
         self.put_cf(K::TABLE, key_buffer, value_buffer)
@@ -1107,23 +1117,33 @@ pub(crate) trait StorageAccess {
         V: PartitionStoreProtobufValue,
         <<V as PartitionStoreProtobufValue>::ProtobufType as TryInto<V>>::Error:
             Into<anyhow::Error>,
+    {
+        let value: Option<ProtobufStorageWrapper<V::ProtobufType>> =
+            self.get_value_storage_codec(key)?;
+
+        value
+            .map(|v| v.0.try_into())
+            .transpose()
+            .map_err(|err| StorageError::Conversion(err.into()))
+    }
+
+    #[inline]
+    fn get_value_storage_codec<K, V>(&mut self, key: K) -> Result<Option<V>>
+    where
+        K: TableKey,
+        V: StorageDecode,
     {
         let mut buf = self.cleared_key_buffer_mut(key.serialized_length());
         key.serialize_to(&mut buf);
         let buf = buf.split();
 
-        match self.get(K::TABLE, &buf) {
-            Ok(value) => {
-                let slice = value.as_ref().map(|v| v.as_ref());
-
-                if let Some(mut slice) = slice {
-                    Ok(Some(V::decode(&mut slice)?))
-                } else {
-                    Ok(None)
-                }
-            }
-            Err(err) => Err(err),
-        }
+        self.get(K::TABLE, &buf)?
+            .map(|value| {
+                let mut slice = value.as_ref();
+                StorageCodec::decode(&mut slice)
+            })
+            .transpose()
+            .map_err(|err| StorageError::Generic(err.into()))
     }
 
     /// Forces a read from persistent storage, bypassing memtables and block cache.

crates/storage-api/src/fsm_table/mod.rs

@@ -13,6 +13,7 @@ use std::future::Future;
 use restate_types::SemanticRestateVersion;
 use restate_types::logs::Lsn;
 use restate_types::message::MessageIndex;
+use restate_types::schema::Schema;
 use restate_types::time::MillisSinceEpoch;
 
 use crate::Result;
@@ -32,6 +33,8 @@ pub trait ReadFsmTable {
     fn get_partition_durability(
         &mut self,
     ) -> impl Future<Output = Result<Option<PartitionDurability>>> + Send + '_;
+
+    fn get_schema(&mut self) -> impl Future<Output = Result<Option<Schema>>> + Send + '_;
 }
 
 pub trait WriteFsmTable {
@@ -44,6 +47,8 @@ pub trait WriteFsmTable {
     fn put_min_restate_version(&mut self, version: &SemanticRestateVersion) -> Result<()>;
 
     fn put_partition_durability(&mut self, durability: &PartitionDurability) -> Result<()>;
+
+    fn put_schema(&mut self, schema: &Schema) -> Result<()>;
 }
 
 #[derive(Debug, Clone, Copy, derive_more::From, derive_more::Into)]