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/mod.rs

@@ -17,6 +17,7 @@ use cling::prelude::*;
 
 #[derive(Run, Subcommand, Clone)]
 #[clap(visible_alias = "dp", alias = "deployment")]
+#[allow(clippy::large_enum_variant)]
 pub enum Deployments {
     /// List the registered deployments
     List(list::List),

cli/src/commands/deployments/register.rs

@@ -18,6 +18,7 @@ use comfy_table::Table;
 use http::{HeaderName, HeaderValue, StatusCode, Uri};
 use indicatif::ProgressBar;
 use indoc::indoc;
+
 use restate_admin_rest_model::deployments::{
     DetailedDeploymentResponse, RegisterDeploymentRequest, RegisterDeploymentResponse,
 };
@@ -55,6 +56,7 @@ 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)]

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

@@ -8,8 +8,7 @@
 // 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::{Uri, Version};
 use restate_serde_util::SerdeableHeaderHashMap;
 use restate_types::identifiers::ServiceRevision;
 use restate_types::identifiers::{DeploymentId, LambdaARN};
@@ -44,7 +43,7 @@ pub enum RegisterDeploymentRequest {
         ///
         /// Additional headers added to every discover/invoke request to the deployment.
         ///
-        /// Fill this field with API Tokens and other authorization headers needed to send 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>,
 
         /// # Metadata

crates/admin/src/schema_registry/mod.rs

@@ -154,8 +154,8 @@ impl SchemaRegistry {
         // Verify first if we have the service. If we do, no need to do anything here.
         if overwrite == updater::Overwrite::No {
             // Verify if we have a service for this endpoint already or not
-            if let Some((deployment, services)) =
-                Metadata::with_current(|m| m.schema_ref()).find_deployment(&deployment_address)
+            if let Some((deployment, services)) = Metadata::with_current(|m| m.schema_ref())
+                .find_deployment(&deployment_address, &additional_headers)
             {
                 return Ok((RegisterDeploymentResult::Unchanged, deployment, services));
             }

crates/invoker-impl/src/lib.rs

@@ -1525,7 +1525,7 @@ mod tests {
     use restate_test_util::{check, let_assert};
     use restate_time_util::FriendlyDuration;
     use restate_types::config::InvokerOptionsBuilder;
-    use restate_types::deployment::DeploymentAddress;
+    use restate_types::deployment::{DeploymentAddress, Headers};
     use restate_types::errors::{InvocationError, codes};
     use restate_types::identifiers::{LeaderEpoch, PartitionId, ServiceRevision};
     use restate_types::invocation::ServiceType;
@@ -1707,6 +1707,7 @@ mod tests {
         fn find_deployment(
             &self,
             _: &DeploymentAddress,
+            _: &Headers,
         ) -> Option<(Deployment, Vec<ServiceMetadata>)> {
             None
         }

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/storage-query-datafusion/src/mocks.rs

@@ -30,7 +30,7 @@ use restate_partition_store::{PartitionStore, PartitionStoreManager};
 use restate_rocksdb::RocksDbManager;
 use restate_types::NodeId;
 use restate_types::config::QueryEngineOptions;
-use restate_types::deployment::DeploymentAddress;
+use restate_types::deployment::{DeploymentAddress, Headers};
 use restate_types::errors::GenericError;
 use restate_types::identifiers::{DeploymentId, PartitionId, PartitionKey, ServiceRevision};
 use restate_types::live::Live;
@@ -74,8 +74,13 @@ impl DeploymentResolver for MockSchemas {
         self.1.resolve_latest_deployment_for_service(service_name)
     }
 
-    fn find_deployment(&self, _: &DeploymentAddress) -> Option<(Deployment, Vec<ServiceMetadata>)> {
-        None
+    fn find_deployment(
+        &self,
+        deployment_address: &DeploymentAddress,
+        additional_headers: &Headers,
+    ) -> Option<(Deployment, Vec<ServiceMetadata>)> {
+        self.1
+            .find_deployment(deployment_address, additional_headers)
     }
 
     fn get_deployment(&self, deployment_id: &DeploymentId) -> Option<Deployment> {

crates/types/src/config/admin.rs

@@ -12,7 +12,7 @@ use std::net::SocketAddr;
 use std::num::NonZeroUsize;
 use std::path::PathBuf;
 
-use http::Uri;
+use http::{HeaderName, Uri};
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
 use tokio::sync::Semaphore;
@@ -41,6 +41,19 @@ pub struct AdminOptions {
     #[cfg_attr(feature = "schemars", schemars(with = "String", url))]
     pub advertised_admin_endpoint: Option<Uri>,
 
+    /// # Deployment routing headers
+    ///
+    /// List of header names considered routing headers.
+    ///
+    /// These will be used during deployment creation to distinguish between an already existing deployment and a new deployment.
+    #[serde(
+        default,
+        skip_serializing_if = "Vec::is_empty",
+        with = "serde_with::As::<Vec<restate_serde_util::HeaderNameSerde>>"
+    )]
+    #[cfg_attr(feature = "schemars", schemars(with = "Vec<String>"))]
+    pub deployment_routing_headers: Vec<HeaderName>,
+
     /// # Concurrency limit
     ///
     /// Concurrency limit for the Admin APIs. Default is unlimited.
@@ -113,6 +126,7 @@ impl Default for AdminOptions {
             bind_address: "0.0.0.0:9070".parse().unwrap(),
             advertised_admin_endpoint: None,
             // max is limited by Tower's LoadShedLayer.
+            deployment_routing_headers: vec![],
             concurrent_api_requests_limit: None,
             query_engine: Default::default(),
             heartbeat_interval: NonZeroFriendlyDuration::from_millis_unchecked(1500),

crates/types/src/schema/deployment.rs

@@ -13,7 +13,10 @@ use std::fmt;
 use std::fmt::{Display, Formatter};
 use std::ops::RangeInclusive;
 
-use crate::deployment::{DeploymentAddress, HttpDeploymentAddress, LambdaDeploymentAddress};
+use crate::config::Configuration;
+use crate::deployment::{
+    DeploymentAddress, Headers, HttpDeploymentAddress, LambdaDeploymentAddress,
+};
 use crate::identifiers::{DeploymentId, LambdaARN, ServiceRevision};
 use crate::schema::service::ServiceMetadata;
 use crate::time::MillisSinceEpoch;
@@ -141,15 +144,24 @@ impl DeploymentMetadata {
         self.created_at
     }
 
-    pub fn semantic_eq_with_address_and_headers(&self, other_addess: &DeploymentAddress) -> bool {
+    pub fn semantic_eq_with_address_and_headers(
+        &self,
+        other_addess: &DeploymentAddress,
+        other_additional_headers: &Headers,
+    ) -> bool {
         match (&self.ty, other_addess) {
             (
                 DeploymentType::Http {
                     address: this_address,
                     ..
                 },
                 DeploymentAddress::Http(HttpDeploymentAddress { uri: other_address }),
-            ) => Self::semantic_eq_http(this_address, other_address),
+            ) => Self::semantic_eq_http(
+                this_address,
+                other_address,
+                &self.delivery_options.additional_headers,
+                other_additional_headers,
+            ),
             (
                 DeploymentType::Lambda { arn: this_arn, .. },
                 DeploymentAddress::Lambda(LambdaDeploymentAddress { arn: other_arn, .. }),
@@ -162,10 +174,21 @@ impl DeploymentMetadata {
         this_arn == other_arn
     }
 
-    pub(crate) fn semantic_eq_http(this_address: &Uri, other_address: &Uri) -> bool {
+    pub(crate) fn semantic_eq_http(
+        this_address: &Uri,
+        other_address: &Uri,
+        this_additional_headers: &Headers,
+        other_additional_headers: &Headers,
+    ) -> bool {
+        let deployment_routing_headers = &Configuration::pinned().admin.deployment_routing_headers;
+
         this_address.authority().expect("Must have authority")
             == other_address.authority().expect("Must have authority")
             && this_address.path() == other_address.path()
+            && deployment_routing_headers.iter().all(|routing_header_key| {
+                this_additional_headers.get(routing_header_key)
+                    == other_additional_headers.get(routing_header_key)
+            })
     }
 }
 
@@ -245,6 +268,7 @@ pub trait DeploymentResolver {
     fn find_deployment(
         &self,
         deployment_address: &DeploymentAddress,
+        additional_headers: &Headers,
     ) -> Option<(Deployment, Vec<ServiceMetadata>)>;
 
     fn get_deployment(&self, deployment_id: &DeploymentId) -> Option<Deployment>;
@@ -445,10 +469,13 @@ pub mod test_util {
         fn find_deployment(
             &self,
             deployment_address: &DeploymentAddress,
+            additional_headers: &Headers,
         ) -> Option<(Deployment, Vec<ServiceMetadata>)> {
             self.deployments
                 .iter()
-                .find(|(_, d)| d.semantic_eq_with_address_and_headers(deployment_address))
+                .find(|(_, d)| {
+                    d.semantic_eq_with_address_and_headers(deployment_address, additional_headers)
+                })
                 .and_then(|(dp_id, _)| self.get_deployment_and_services(dp_id))
         }
 
@@ -504,6 +531,7 @@ pub mod test_util {
         fn find_deployment(
             &self,
             _: &DeploymentAddress,
+            _: &Headers,
         ) -> Option<(Deployment, Vec<ServiceMetadata>)> {
             None
         }