Consistent schema across partitions

Summary:
This PR make sure that PPs have a consistent service schema by writing
and upsert-schema record to bifrost. The leader processor makes sure
to write this record to bifrost when schema changes are detected

Author: muhamadazmy

crates/types/src/schema/metadata/mod.rs

@@ -38,8 +38,9 @@ use crate::schema::{deployment, service};
 use crate::time::MillisSinceEpoch;
 use crate::{Version, Versioned, identifiers};
 
-#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[derive(derive_more::Debug, Clone, serde::Serialize, serde::Deserialize)]
 #[serde(from = "serde_hacks::Schema", into = "serde_hacks::Schema")]
+#[debug("Schema(version: {version})")]
 pub struct Schema {
     /// This gets bumped on each update.
     version: Version,

crates/wal-protocol/src/lib.rs

@@ -18,6 +18,7 @@ use restate_types::invocation::{
 use restate_types::logs;
 use restate_types::logs::{HasRecordKeys, Keys, MatchKeyQuery};
 use restate_types::message::MessageIndex;
+use restate_types::schema::Schema;
 use restate_types::state_mut::ExternalStateMutation;
 
 use crate::control::{AnnounceLeader, VersionBarrier};
@@ -29,7 +30,7 @@ pub mod control;
 pub mod timer;
 
 /// The primary envelope for all messages in the system.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub struct Envelope {
     pub header: Header,
@@ -125,7 +126,7 @@ pub enum Destination {
 }
 
 /// State machine input commands
-#[derive(Debug, Clone, PartialEq, Eq, strum::EnumDiscriminants, strum::VariantNames)]
+#[derive(Debug, Clone, strum::EnumDiscriminants, strum::VariantNames)]
 #[strum_discriminants(derive(strum::IntoStaticStr))]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub enum Command {
@@ -183,6 +184,9 @@ pub enum Command {
     NotifyGetInvocationOutputResponse(GetInvocationOutputResponse),
     /// Notify a signal.
     NotifySignal(NotifySignalRequest),
+
+    /// Upsert schema for consistent schema across replicas
+    UpsertSchema(Schema),
 }
 
 impl Command {
@@ -233,6 +237,7 @@ impl HasRecordKeys for Envelope {
             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::UpsertSchema(_) => Keys::Single(self.partition_key()),
         }
     }
 }
@@ -337,6 +342,7 @@ mod envelope {
         UpdatePartitionDurability = 17,         // bilrost
         ResumeInvocation = 18,                  // flexbuffers
         RestartAsNewInvocation = 19,            // flexbuffers
+        UpsertSchema = 20,                      // flexbuffers
     }
 
     #[derive(bilrost::Message)]
@@ -524,6 +530,10 @@ mod envelope {
                 let value = protobuf::outbox_message::NotifySignal::from(value.clone());
                 (CommandKind::NotifySignal, Field::encode_protobuf(&value))
             }
+            Command::UpsertSchema(value) => (
+                CommandKind::UpsertSchema,
+                Field::encode_serde(StorageCodecKind::FlexbuffersSerde, value),
+            ),
         };
 
         let dto = Envelope {
@@ -631,6 +641,10 @@ mod envelope {
 
                 Command::NotifySignal(value.try_into()?)
             }
+            CommandKind::UpsertSchema => {
+                codec_or_error!(envelope.command, StorageCodecKind::FlexbuffersSerde);
+                Command::UpsertSchema(envelope.command.decode_serde()?)
+            }
         };
 
         Ok(super::Envelope { header, command })

crates/worker/src/partition/leadership/leader_state.rs

@@ -20,12 +20,12 @@ use futures::future::OptionFuture;
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt, stream};
 use metrics::counter;
-use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::wrappers::{ReceiverStream, WatchStream};
 use tracing::{debug, trace};
 
 use restate_bifrost::CommitToken;
 use restate_core::network::{Oneshot, Reciprocal};
-use restate_core::{TaskCenter, TaskHandle, TaskId};
+use restate_core::{Metadata, MetadataKind, TaskCenter, TaskHandle, TaskId};
 use restate_partition_store::PartitionStore;
 use restate_types::identifiers::{
     InvocationId, LeaderEpoch, PartitionId, PartitionKey, PartitionProcessorRpcRequestId,
@@ -36,6 +36,7 @@ use restate_types::net::partition_processor::{
     PartitionProcessorRpcError, PartitionProcessorRpcResponse,
 };
 use restate_types::time::MillisSinceEpoch;
+use restate_types::{Version, Versioned};
 use restate_wal_protocol::Command;
 use restate_wal_protocol::timer::TimerKeyValue;
 
@@ -45,7 +46,7 @@ use crate::partition::leadership::self_proposer::SelfProposer;
 use crate::partition::leadership::{ActionEffect, Error, InvokerStream, TimerService};
 use crate::partition::shuffle;
 use crate::partition::shuffle::HintSender;
-use crate::partition::state_machine::Action;
+use crate::partition::state_machine::{Action, StateMachineRef};
 
 use super::durability_tracker::DurabilityTracker;
 
@@ -73,6 +74,7 @@ pub struct LeaderState {
 
     invoker_stream: InvokerStream,
     shuffle_stream: ReceiverStream<shuffle::OutboxTruncation>,
+    schema_stream: WatchStream<Version>,
     pub pending_cleanup_timers_to_schedule: VecDeque<(InvocationId, Duration)>,
     cleaner_task_id: TaskId,
     trimmer_task_id: TaskId,
@@ -103,6 +105,9 @@ impl LeaderState {
             cleaner_task_id,
             trimmer_task_id,
             shuffle_hint_tx,
+            schema_stream: Metadata::with_current(|m| {
+                WatchStream::new(m.watch(MetadataKind::Schema))
+            }),
             timer_service: Box::pin(timer_service),
             self_proposer,
             awaiting_rpc_actions: Default::default(),
@@ -119,7 +124,10 @@ impl LeaderState {
     ///
     /// Important: The future needs to be cancellation safe since it is polled as a tokio::select
     /// arm!
-    pub async fn run(&mut self) -> Result<Vec<ActionEffect>, Error> {
+    pub async fn run(
+        &mut self,
+        state_machine: StateMachineRef<'_>,
+    ) -> Result<Vec<ActionEffect>, Error> {
         let timer_stream = std::pin::pin!(stream::unfold(
             &mut self.timer_service,
             |timer_service| async {
@@ -128,6 +136,21 @@ impl LeaderState {
             }
         ));
 
+        let schema_stream = (&mut self.schema_stream)
+            .filter(|version| {
+                // only upsert schema iff version is newer that
+                futures::future::ready(
+                    state_machine
+                        .schema
+                        .as_ref()
+                        .map(|schema| schema.version() < *version)
+                        .unwrap_or(true),
+                )
+            })
+            .map(|_| {
+                let schema = Metadata::with_current(|m| m.schema().clone());
+                ActionEffect::UpsertSchema(schema)
+            });
         let invoker_stream = (&mut self.invoker_stream).map(ActionEffect::Invoker);
         let shuffle_stream = (&mut self.shuffle_stream).map(ActionEffect::Shuffle);
         let dur_tracker_stream =
@@ -155,7 +178,8 @@ impl LeaderState {
             timer_stream,
             action_effects_stream,
             awaiting_rpc_self_propose_stream,
-            dur_tracker_stream
+            dur_tracker_stream,
+            schema_stream
         );
         let mut all_streams = all_streams.ready_chunks(BATCH_READY_UP_TO);
 
@@ -284,6 +308,16 @@ impl LeaderState {
                         )
                         .await?;
                 }
+                ActionEffect::UpsertSchema(schema) => {
+                    debug!(
+                        "Self purposing {schema:?} for partition key {}",
+                        self.own_partition_key
+                    );
+
+                    self.self_proposer
+                        .propose(self.own_partition_key, Command::UpsertSchema(schema))
+                        .await?;
+                }
                 ActionEffect::AwaitingRpcSelfProposeDone => {
                     // Nothing to do here
                 }

crates/worker/src/partition/leadership/mod.rs

@@ -52,6 +52,7 @@ use restate_types::net::partition_processor::{
 use restate_types::partitions::Partition;
 use restate_types::partitions::state::PartitionReplicaSetStates;
 use restate_types::retries::with_jitter;
+use restate_types::schema::Schema;
 use restate_types::storage::StorageEncodeError;
 use restate_wal_protocol::Command;
 use restate_wal_protocol::control::{AnnounceLeader, PartitionDurability};
@@ -63,7 +64,7 @@ use crate::partition::leadership::leader_state::LeaderState;
 use crate::partition::leadership::self_proposer::SelfProposer;
 use crate::partition::shuffle;
 use crate::partition::shuffle::{OutboxReaderError, Shuffle, ShuffleMetadata};
-use crate::partition::state_machine::Action;
+use crate::partition::state_machine::{Action, StateMachineRef};
 use crate::partition::types::InvokerEffect;
 
 use self::durability_tracker::DurabilityTracker;
@@ -124,6 +125,7 @@ pub(crate) enum ActionEffect {
     Timer(TimerKeyValue),
     ScheduleCleanupTimer(InvocationId, Duration),
     PartitionMaintenance(PartitionDurability),
+    UpsertSchema(Schema),
     AwaitingRpcSelfProposeDone,
 }
 enum State {
@@ -518,7 +520,10 @@ where
     /// * Follower: Nothing to do
     /// * Candidate: Monitor appender task
     /// * Leader: Await action effects and monitor appender task
-    pub async fn run(&mut self) -> Result<Vec<ActionEffect>, Error> {
+    pub async fn run(
+        &mut self,
+        state_machine: StateMachineRef<'_>,
+    ) -> Result<Vec<ActionEffect>, Error> {
         match &mut self.state {
             State::Follower => Ok(futures::future::pending::<Vec<_>>().await),
             State::Candidate { self_proposer, .. } => Err(self_proposer
@@ -527,7 +532,7 @@ where
                 .join_on_err()
                 .await
                 .expect_err("never should never be returned")),
-            State::Leader(leader_state) => leader_state.run().await,
+            State::Leader(leader_state) => leader_state.run(state_machine).await,
         }
     }
 

crates/worker/src/partition/mod.rs

@@ -181,6 +181,7 @@ where
         let outbox_seq_number = partition_store.get_outbox_seq_number().await?;
         let outbox_head_seq_number = partition_store.get_outbox_head_seq_number().await?;
         let min_restate_version = partition_store.get_min_restate_version().await?;
+        let schema = partition_store.get_schema().await?;
 
         if !SemanticRestateVersion::current().is_equal_or_newer_than(&min_restate_version) {
             gauge!(PARTITION_BLOCKED_FLARE, PARTITION_LABEL =>
@@ -199,6 +200,7 @@ where
             partition_store.partition_key_range().clone(),
             min_restate_version,
             EnumSet::empty(),
+            schema,
         );
 
         Ok(state_machine)
@@ -561,7 +563,7 @@ where
                     transaction.commit().await?;
                     self.leadership_state.handle_actions(action_collector.drain(..))?;
                 },
-                result = self.leadership_state.run() => {
+                result = self.leadership_state.run(self.state_machine.state_machine_ref()) => {
                     let action_effects = result?;
                     // We process the action_effects not directly in the run future because it
                     // requires the run future to be cancellation safe. In the future this could be

crates/worker/src/partition/state_machine/lifecycle/version_barrier.rs

@@ -64,6 +64,7 @@ mod tests {
             PartitionKey::MIN..=PartitionKey::MAX,
             SemanticRestateVersion::unknown().clone(),
             Default::default(),
+            None,
         );
         // this is fine as we are always above the unknown version (current > 0.0.0)
         let mut test_env = TestEnv::create_with_state_machine(state_machine).await;
@@ -106,6 +107,7 @@ mod tests {
             PartitionKey::MIN..=PartitionKey::MAX,
             SemanticRestateVersion::unknown().clone(),
             Default::default(),
+            None,
         );
         // this is fine as we are always above the unknown version (current > 0.0.0)
         let mut test_env = TestEnv::create_with_state_machine(state_machine).await;