Push schema information to workers from admin service on schema changes

This commit enables the admin service to push schema updates to the worker
in case of schema changes.

crates/admin/src/options.rs

@@ -8,7 +8,7 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use restate_meta::MetaHandle;
+use restate_meta::{FileMetaReader, MetaHandle};
 use restate_schema_impl::Schemas;
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
@@ -44,7 +44,12 @@ impl Default for Options {
 }
 
 impl Options {
-    pub fn build(self, schemas: Schemas, meta_handle: MetaHandle) -> AdminService {
-        AdminService::new(self, schemas, meta_handle)
+    pub fn build(
+        self,
+        schemas: Schemas,
+        meta_handle: MetaHandle,
+        schema_reader: FileMetaReader,
+    ) -> AdminService {
+        AdminService::new(self, schemas, meta_handle, schema_reader)
     }
 }

crates/admin/src/rest_api/deployments.rs

@@ -19,6 +19,7 @@ use restate_service_client::Endpoint;
 use restate_service_protocol::discovery::DiscoverEndpoint;
 use restate_types::identifiers::InvalidLambdaARN;
 
+use crate::rest_api::notify_worker_about_schema_changes;
 use axum::body::Bytes;
 use axum::extract::{Path, Query, State};
 use axum::http::{header, HeaderValue, StatusCode};
@@ -97,6 +98,8 @@ pub async fn create_deployment<W>(
         .register_deployment(discover_endpoint, force, apply_changes)
         .await?;
 
+    notify_worker_about_schema_changes(state.schema_reader(), state.worker_svc_client()).await?;
+
     let response_body = RegisterDeploymentResponse {
         id: registration_result.deployment,
         services: registration_result.services,
@@ -241,6 +244,8 @@ pub async fn delete_deployment<W>(
 ) -> Result<StatusCode, MetaApiError> {
     if let Some(true) = force {
         state.meta_handle().remove_deployment(deployment_id).await?;
+        notify_worker_about_schema_changes(state.schema_reader(), state.worker_svc_client())
+            .await?;
         Ok(StatusCode::ACCEPTED)
     } else {
         Ok(StatusCode::NOT_IMPLEMENTED)

crates/admin/src/rest_api/error.rs

@@ -43,6 +43,8 @@ pub enum MetaApiError {
     Meta(#[from] MetaError),
     #[error(transparent)]
     Worker(#[from] restate_worker_api::Error),
+    #[error(transparent)]
+    Generic(Error),
 }
 
 /// # Error description response

crates/admin/src/rest_api/mod.rs

@@ -18,8 +18,14 @@ mod methods;
 mod services;
 mod subscriptions;
 
+use crate::rest_api::error::MetaApiError;
 use okapi_operation::axum_integration::{delete, get, patch, post};
 use okapi_operation::*;
+use restate_meta::{FileMetaReader, MetaReader};
+use restate_node_services::worker::worker_svc_client::WorkerSvcClient;
+use restate_node_services::worker::UpdateSchemaRequest;
+use tonic::transport::Channel;
+use tracing::debug;
 
 use crate::state::AdminServiceState;
 
@@ -101,3 +107,30 @@ pub fn create_router<W: restate_worker_api::Handle + Clone + Send + Sync + 'stat
         .expect("Error when building the OpenAPI specification")
         .with_state(state)
 }
+
+/// Notifies the worker about schema changes. This method is best-effort and will not fail if the worker
+/// could not be reached.
+async fn notify_worker_about_schema_changes(
+    schema_reader: &FileMetaReader,
+    mut worker_svc_client: WorkerSvcClient<Channel>,
+) -> Result<(), MetaApiError> {
+    let schema_updates = schema_reader
+        .read()
+        .await
+        .map_err(|err| MetaApiError::Meta(err.into()))?;
+
+    // don't fail if the worker is not reachable
+    let result = worker_svc_client
+        .update_schemas(UpdateSchemaRequest {
+            schema_bin: bincode::serde::encode_to_vec(schema_updates, bincode::config::standard())
+                .map_err(|err| MetaApiError::Generic(err.into()))?
+                .into(),
+        })
+        .await;
+
+    if let Err(err) = result {
+        debug!("Failed notifying worker about schema changes: {err}");
+    }
+
+    Ok(())
+}

crates/admin/src/rest_api/services.rs

@@ -15,6 +15,7 @@ use restate_meta_rest_model::services::*;
 use restate_pb::grpc::reflection::v1::FileDescriptorResponse;
 use restate_schema_api::service::ServiceMetadataResolver;
 
+use crate::rest_api::notify_worker_about_schema_changes;
 use axum::extract::{Path, State};
 use axum::http::{header, HeaderValue};
 use axum::response::{IntoResponse, Response};
@@ -91,6 +92,8 @@ pub async fn modify_service<W>(
         .modify_service(service_name.clone(), public)
         .await?;
 
+    notify_worker_about_schema_changes(state.schema_reader(), state.worker_svc_client()).await?;
+
     state
         .schemas()
         .resolve_latest_service_metadata(&service_name)

crates/admin/src/rest_api/subscriptions.rs

@@ -14,6 +14,7 @@ use crate::state::AdminServiceState;
 use restate_meta_rest_model::subscriptions::*;
 use restate_schema_api::subscription::SubscriptionResolver;
 
+use crate::rest_api::notify_worker_about_schema_changes;
 use axum::extract::Query;
 use axum::extract::{Path, State};
 use axum::http::StatusCode;
@@ -52,6 +53,8 @@ pub async fn create_subscription<W>(
         )
         .await?;
 
+    notify_worker_about_schema_changes(state.schema_reader(), state.worker_svc_client()).await?;
+
     Ok((
         StatusCode::CREATED,
         [(
@@ -168,5 +171,7 @@ pub async fn delete_subscription<W>(
         .delete_subscription(subscription_id)
         .await?;
 
+    notify_worker_about_schema_changes(state.schema_reader(), state.worker_svc_client()).await?;
+
     Ok(StatusCode::ACCEPTED)
 }

crates/admin/src/service.rs

@@ -16,7 +16,7 @@ use http::StatusCode;
 use tonic::transport::Channel;
 use tower::ServiceBuilder;
 
-use restate_meta::MetaHandle;
+use restate_meta::{FileMetaReader, MetaHandle};
 use restate_node_services::worker::worker_svc_client::WorkerSvcClient;
 use restate_schema_impl::Schemas;
 use tracing::info;
@@ -29,36 +29,40 @@ pub struct AdminService {
     opts: Options,
     schemas: Schemas,
     meta_handle: MetaHandle,
+    schema_reader: FileMetaReader,
 }
 
 impl AdminService {
-    pub fn new(opts: Options, schemas: Schemas, meta_handle: MetaHandle) -> Self {
+    pub fn new(
+        opts: Options,
+        schemas: Schemas,
+        meta_handle: MetaHandle,
+        schema_reader: FileMetaReader,
+    ) -> Self {
         Self {
             opts,
             schemas,
             meta_handle,
+            schema_reader,
         }
     }
 
     pub async fn run(
         self,
         drain: drain::Watch,
         worker_handle: impl restate_worker_api::Handle + Clone + Send + Sync + 'static,
-        worker_svc_client: Option<WorkerSvcClient<Channel>>,
+        worker_svc_client: WorkerSvcClient<Channel>,
     ) -> Result<(), Error> {
-        let rest_state =
-            state::AdminServiceState::new(self.meta_handle, self.schemas, worker_handle);
-
-        let router = axum::Router::new();
+        let rest_state = state::AdminServiceState::new(
+            self.meta_handle,
+            self.schemas,
+            worker_handle,
+            worker_svc_client.clone(),
+            self.schema_reader,
+        );
 
-        // Stitch query http endpoint if enabled
-        let router = if let Some(worker_svc_client) = worker_svc_client {
-            let query_state = Arc::new(state::QueryServiceState { worker_svc_client });
-            // Merge storage query router
-            router.merge(storage_query::create_router(query_state))
-        } else {
-            router
-        };
+        let query_state = Arc::new(state::QueryServiceState { worker_svc_client });
+        let router = axum::Router::new().merge(storage_query::create_router(query_state));
 
         let router = router
             // Merge meta API router

crates/admin/src/state.rs

@@ -9,7 +9,7 @@
 // by the Apache License, Version 2.0.
 //
 
-use restate_meta::MetaHandle;
+use restate_meta::{FileMetaReader, MetaHandle};
 use restate_node_services::worker::worker_svc_client::WorkerSvcClient;
 use restate_schema_impl::Schemas;
 use tonic::transport::Channel;
@@ -19,6 +19,8 @@ pub struct AdminServiceState<W> {
     meta_handle: MetaHandle,
     schemas: Schemas,
     worker_handle: W,
+    worker_svc_client: WorkerSvcClient<Channel>,
+    schema_reader: FileMetaReader,
 }
 
 #[derive(Clone)]
@@ -27,11 +29,19 @@ pub struct QueryServiceState {
 }
 
 impl<W> AdminServiceState<W> {
-    pub fn new(meta_handle: MetaHandle, schemas: Schemas, worker_handle: W) -> Self {
+    pub fn new(
+        meta_handle: MetaHandle,
+        schemas: Schemas,
+        worker_handle: W,
+        worker_svc_client: WorkerSvcClient<Channel>,
+        schema_reader: FileMetaReader,
+    ) -> Self {
         Self {
             meta_handle,
             schemas,
             worker_handle,
+            worker_svc_client,
+            schema_reader,
         }
     }
 
@@ -42,6 +52,14 @@ impl<W> AdminServiceState<W> {
     pub fn schemas(&self) -> &Schemas {
         &self.schemas
     }
+
+    pub fn worker_svc_client(&self) -> WorkerSvcClient<Channel> {
+        self.worker_svc_client.clone()
+    }
+
+    pub fn schema_reader(&self) -> &FileMetaReader {
+        &self.schema_reader
+    }
 }
 
 impl<W: Clone> AdminServiceState<W> {

crates/meta/src/error.rs

@@ -11,7 +11,7 @@
 use restate_schema_impl::SchemasUpdateError;
 use restate_service_protocol::discovery::ServiceDiscoveryError;
 
-use crate::storage::MetaStorageError;
+use crate::storage::{MetaReaderError, MetaStorageError};
 
 #[derive(Debug, thiserror::Error, codederror::CodedError)]
 pub enum Error {
@@ -26,6 +26,9 @@ pub enum Error {
     Storage(#[from] MetaStorageError),
     #[error(transparent)]
     #[code(unknown)]
+    Reader(#[from] MetaReaderError),
+    #[error(transparent)]
+    #[code(unknown)]
     SchemaRegistry(#[from] SchemasUpdateError),
     #[error("meta closed")]
     #[code(unknown)]

crates/node-services/proto/worker.proto

@@ -25,6 +25,9 @@ service WorkerSvc {
 
   // Queries the storage of the worker and returns the result as a stream of responses
   rpc QueryStorage(StorageQueryRequest) returns (stream StorageQueryResponse);
+
+  // Updates the schema information on the worker node
+  rpc UpdateSchemas(UpdateSchemaRequest) returns (google.protobuf.Empty);
 }
 
 message BifrostVersion {
@@ -48,4 +51,8 @@ message StorageQueryRequest {
 message StorageQueryResponse {
   bytes header = 1;
   bytes data = 2;
+}
+
+message UpdateSchemaRequest {
+  bytes schema_bin = 1;
 }

crates/node/src/lib.rs

@@ -117,6 +117,8 @@ impl Node {
                     worker.bifrost_handle(),
                     worker.worker_command_tx(),
                     worker.storage_query_context().clone(),
+                    worker.schemas(),
+                    worker.subscription_controller(),
                 )
             }),
             cluster_controller_role.as_ref().map(|cluster_controller| {

crates/node/src/roles/cluster_controller.rs

@@ -114,7 +114,7 @@ impl ClusterControllerRole {
 
         component_set.spawn(
             self.admin
-                .run(inner_shutdown_watch, worker_handle, Some(worker_svc_client))
+                .run(inner_shutdown_watch, worker_handle, worker_svc_client)
                 .map_ok(|_| "admin")
                 .map_err(ClusterControllerRoleError::Admin),
         );
@@ -143,7 +143,9 @@ impl TryFrom<Options> for ClusterControllerRole {
 
     fn try_from(options: Options) -> Result<Self, Self::Error> {
         let meta = options.meta.build(options.worker.kafka.clone())?;
-        let admin = options.admin.build(meta.schemas(), meta.meta_handle());
+        let admin = options
+            .admin
+            .build(meta.schemas(), meta.meta_handle(), meta.schema_reader());
 
         Ok(ClusterControllerRole {
             controller: restate_cluster_controller::Service::new(options.cluster_controller),

crates/node/src/roles/mod.rs

@@ -14,4 +14,4 @@ mod worker;
 pub use cluster_controller::{
     ClusterControllerRole, ClusterControllerRoleBuildError, ClusterControllerRoleError,
 };
-pub use worker::{WorkerRole, WorkerRoleBuildError, WorkerRoleError};
+pub use worker::{update_schemas, WorkerRole, WorkerRoleBuildError, WorkerRoleError};