Add routing header feature to deployment registration.

This lets you use a specific header as "discriminator" when registering a service.

Author: slinkydeveloper

cli/src/commands/deployments/register.rs

@@ -18,7 +18,7 @@ use comfy_table::Table;
 use http::{HeaderName, HeaderValue, StatusCode, Uri};
 use indicatif::ProgressBar;
 
-use restate_admin_rest_model::deployments::RegisterDeploymentRequest;
+use restate_admin_rest_model::deployments::{Header, RegisterDeploymentRequest};
 use restate_cli_util::ui::console::{Styled, StyledTable, confirm_or_exit};
 use restate_cli_util::ui::stylesheet::Style;
 use restate_cli_util::{c_eprintln, c_error, c_indent_table, c_indentln, c_success, c_warn};
@@ -48,11 +48,20 @@ pub struct Register {
     assume_role_arn: Option<String>,
 
     /// Additional header that will be sent to the endpoint during the discovery request.
+    /// You typically want to include here API keys and other tokens required to send requests to deployments.
     ///
     /// Use `--extra-header name=value` format and repeat --extra-header for each additional header.
     #[clap(long="extra-header", value_parser = parse_header, action = clap::ArgAction::Append)]
     extra_headers: Option<Vec<HeaderKeyValue>>,
 
+    /// Header used for routing to a specific deployment.
+    /// If the load balancer between restate-server and your deployments uses a specific header to route,
+    /// you should set this as the routing header, as it will be used to distinguish this deployment with other deployments with the same URL.
+    ///
+    /// Use `--routing-header name=value` format.
+    #[clap(long="routing-header", value_parser = parse_header)]
+    routing_header: Option<HeaderKeyValue>,
+
     /// Attempt discovery using a client that defaults to HTTP1.1 instead of a prior-knowledge HTTP2 client.
     /// This may be necessary if you see `META0014` discovering local dev servers like `wrangler dev`.
     #[clap(long = "use-http1.1")]
@@ -131,6 +140,14 @@ pub async fn run_register(State(env): State<CliEnv>, discover_opts: &Register) -
     let headers = discover_opts.extra_headers.as_ref().map(|headers| {
         HashMap::from_iter(headers.iter().map(|kv| (kv.key.clone(), kv.value.clone())))
     });
+    let routing_header =
+        discover_opts
+            .routing_header
+            .as_ref()
+            .map(|HeaderKeyValue { key, value }| Header {
+                key: key.clone(),
+                value: value.clone(),
+            });
 
     // Preparing the discovery request
     let client = AdminClient::new(&env).await?;
@@ -189,6 +206,7 @@ pub async fn run_register(State(env): State<CliEnv>, discover_opts: &Register) -
     let mk_request_body = |force, dry_run| match &deployment {
         DeploymentEndpoint::Uri(uri) => RegisterDeploymentRequest::Http {
             uri: uri.clone(),
+            routing_header: routing_header.clone(),
             additional_headers: headers.clone().map(Into::into),
             use_http_11: discover_opts.use_http_11,
             force,

crates/admin-rest-model/src/deployments.rs

@@ -8,8 +8,8 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use http::Uri;
 use http::Version;
+use http::{HeaderName, HeaderValue, Uri};
 use restate_serde_util::SerdeableHeaderHashMap;
 use restate_types::identifiers::ServiceRevision;
 use restate_types::identifiers::{DeploymentId, LambdaARN};
@@ -18,6 +18,18 @@ use restate_types::schema::service::ServiceMetadata;
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
 
+#[serde_as]
+#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Header {
+    #[serde_as(as = "restate_serde_util::HeaderNameSerde")]
+    #[cfg_attr(feature = "schema", schemars(with = "String"))]
+    pub key: HeaderName,
+    #[serde_as(as = "restate_serde_util::HeaderValueSerde")]
+    #[cfg_attr(feature = "schema", schemars(with = "String"))]
+    pub value: HeaderValue,
+}
+
 // This enum could be a struct with a nested enum to avoid repeating some fields, but serde(flatten) unfortunately breaks the openapi code generation
 #[serde_as]
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
@@ -39,10 +51,18 @@ pub enum RegisterDeploymentRequest {
         #[cfg_attr(feature = "schema", schemars(with = "String"))]
         uri: Uri,
 
+        /// # Routing header
+        ///
+        /// Header used for routing to a specific deployment.
+        /// If the load balancer between restate-server and your deployments uses a specific header to route,
+        /// you should set this as the routing header, as it will be used to distinguish this deployment with other deployments with the same URL.
+        routing_header: Option<Header>,
+
         /// # Additional headers
         ///
         /// Additional headers added to the discover/invoke requests to the deployment.
         ///
+        /// You typically want to include here API keys and other tokens required to send requests to deployments.
         additional_headers: Option<SerdeableHeaderHashMap>,
 
         /// # Use http1.1

crates/admin/src/rest_api/deployments.rs

@@ -10,6 +10,7 @@
 
 use super::error::*;
 use crate::state::AdminServiceState;
+use std::collections::HashMap;
 use std::time::SystemTime;
 
 use axum::Json;
@@ -30,7 +31,7 @@ use serde::Deserialize;
 
 /// Create deployment and return discovered services.
 #[openapi(
-    summary = "Create deployment",
+    summary = " pubCreate deployment",
     description = "Create deployment. Restate will invoke the endpoint to gather additional information required for registration, such as the services exposed by the deployment. If the deployment is already registered, this method will fail unless `force` is set to `true`.",
     operation_id = "create_deployment",
     tags = "deployment",
@@ -49,10 +50,11 @@ pub async fn create_deployment<V, IC>(
     State(state): State<AdminServiceState<V, IC>>,
     #[request_body(required = true)] Json(payload): Json<RegisterDeploymentRequest>,
 ) -> Result<impl IntoResponse, MetaApiError> {
-    let (discover_endpoint, force, dry_run) = match payload {
+    let (discover_endpoint, routing_header, force, dry_run) = match payload {
         RegisterDeploymentRequest::Http {
             uri,
             additional_headers,
+            routing_header,
             use_http_11,
             force,
             dry_run,
@@ -69,6 +71,13 @@ pub async fn create_deployment<V, IC>(
 
             let is_using_https = uri.scheme().unwrap() == &Scheme::HTTPS;
 
+            // Add it as part of the additional headers
+            let mut additional_headers: HashMap<_, _> =
+                additional_headers.unwrap_or_default().into();
+            if let Some(routing_header) = &routing_header {
+                additional_headers.insert(routing_header.key.clone(), routing_header.value.clone());
+            }
+
             (
                 DiscoverEndpoint::new(
                     Endpoint::Http(
@@ -83,8 +92,9 @@ pub async fn create_deployment<V, IC>(
                             Some(http::Version::HTTP_2)
                         },
                     ),
-                    additional_headers.unwrap_or_default().into(),
+                    additional_headers,
                 ),
+                routing_header.map(|h| (h.key, h.value)),
                 force,
                 dry_run,
             )
@@ -106,6 +116,7 @@ pub async fn create_deployment<V, IC>(
                 ),
                 additional_headers.unwrap_or_default().into(),
             ),
+            None,
             force,
             dry_run,
         ),
@@ -125,7 +136,7 @@ pub async fn create_deployment<V, IC>(
 
     let (deployment, services) = state
         .schema_registry
-        .register_deployment(discover_endpoint, force, apply_mode)
+        .register_deployment(discover_endpoint, routing_header, force, apply_mode)
         .await
         .inspect_err(|e| warn_it!(e))?;
 

crates/admin/src/schema_registry/mod.rs

@@ -13,7 +13,7 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use anyhow::Context;
-use http::{HeaderMap, HeaderValue, Uri, uri::PathAndQuery};
+use http::{HeaderMap, HeaderName, HeaderValue, Uri, uri::PathAndQuery};
 use tracing::subscriber::NoSubscriber;
 use tracing::trace;
 
@@ -90,6 +90,7 @@ impl<V> SchemaRegistry<V> {
     pub async fn register_deployment(
         &self,
         discover_endpoint: DiscoverEndpoint,
+        routing_header: Option<(HeaderName, HeaderValue)>,
         force: updater::Force,
         apply_mode: updater::ApplyMode,
     ) -> Result<(Deployment, Vec<ServiceMetadata>), SchemaRegistryError> {
@@ -128,6 +129,7 @@ impl<V> SchemaRegistry<V> {
             let id = tracing::subscriber::with_default(NoSubscriber::new(), || {
                 updater.add_deployment(
                     deployment_metadata,
+                    routing_header,
                     discovered_metadata.services,
                     force.force_enabled(),
                 )
@@ -153,6 +155,7 @@ impl<V> SchemaRegistry<V> {
 
                     new_deployment_id = Some(updater.add_deployment(
                         deployment_metadata.clone(),
+                        routing_header.clone(),
                         discovered_metadata.services.clone(),
                         force.force_enabled(),
                     )?);

crates/serde-util/src/header_name.rs

@@ -0,0 +1,37 @@
+// Copyright (c) 2023 - 2025 Restate Software, Inc., Restate GmbH.
+// All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+use http::HeaderName;
+use serde::Deserialize;
+use serde_with::{DeserializeAs, SerializeAs};
+use std::str::FromStr;
+
+/// SerializeAs/DeserializeAs to implement ser/de trait for [HeaderName]
+/// Use it with `#[serde(with = "serde_with::As::<HeaderNameSerde>")]`.
+pub struct HeaderNameSerde;
+
+impl SerializeAs<HeaderName> for HeaderNameSerde {
+    fn serialize_as<S>(source: &HeaderName, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serializer.serialize_str(source.as_str())
+    }
+}
+
+impl<'de> DeserializeAs<'de, HeaderName> for HeaderNameSerde {
+    fn deserialize_as<D>(deserializer: D) -> Result<HeaderName, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let buf = String::deserialize(deserializer)?;
+        HeaderName::from_str(&buf).map_err(serde::de::Error::custom)
+    }
+}

crates/serde-util/src/lib.rs

@@ -15,12 +15,14 @@ mod proto;
 
 pub mod authority;
 pub mod default;
+mod header_name;
 pub mod header_value;
 mod map_as_vec;
 mod version;
 
 pub use byte_count::*;
 pub use header_map::SerdeableHeaderHashMap;
+pub use header_name::HeaderNameSerde;
 pub use header_value::HeaderValueSerde;
 pub use map_as_vec::{MapAsVec, MapAsVecItem};
 #[cfg(feature = "proto")]

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

@@ -27,7 +27,7 @@ use crate::schema::invocation_target::{
 use crate::schema::subscriptions::{
     EventInvocationTargetTemplate, Sink, Source, Subscription, SubscriptionValidator,
 };
-use http::{HeaderValue, Uri};
+use http::{HeaderName, HeaderValue, Uri};
 use serde_json::Value;
 use std::borrow::Borrow;
 use std::collections::HashMap;
@@ -345,6 +345,7 @@ impl SchemaUpdater {
     pub fn add_deployment(
         &mut self,
         deployment_metadata: DeploymentMetadata,
+        routing_header: Option<(HeaderName, HeaderValue)>,
         services: Vec<endpoint_manifest::Service>,
         force: bool,
     ) -> Result<DeploymentId, SchemaError> {
@@ -357,6 +358,16 @@ impl SchemaUpdater {
         let mut existing_deployments = self.schema.deployments.iter().filter(|(_, schemas)| {
             schemas.ty.protocol_type() == deployment_metadata.ty.protocol_type()
                 && schemas.ty.normalized_address() == deployment_metadata.ty.normalized_address()
+                && (routing_header.as_ref().is_none()
+                    || routing_header.as_ref().is_some_and(
+                        |(routing_header_key, routing_header_value)| {
+                            schemas
+                                .delivery_options
+                                .additional_headers
+                                .get(routing_header_key)
+                                .is_some_and(|v| v == routing_header_value)
+                        },
+                    ))
         });
 
         let mut services_to_remove = Vec::default();