Change routing header feature.

Accept the key of the routing header via configuration. This makes the behavior more consistent.

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::{Header, RegisterDeploymentRequest};
+use restate_admin_rest_model::deployments::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,20 +48,11 @@ 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")]
@@ -140,14 +131,6 @@ 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?;
@@ -206,7 +189,6 @@ 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,18 +18,6 @@ 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))]
@@ -51,13 +39,6 @@ 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.

crates/admin/src/rest_api/deployments.rs

@@ -10,7 +10,6 @@
 
 use super::error::*;
 use crate::state::AdminServiceState;
-use std::collections::HashMap;
 use std::time::SystemTime;
 
 use axum::Json;
@@ -31,7 +30,7 @@ use serde::Deserialize;
 
 /// Create deployment and return discovered services.
 #[openapi(
-    summary = " pubCreate deployment",
+    summary = "Create 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",
@@ -50,11 +49,10 @@ 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, routing_header, force, dry_run) = match payload {
+    let (discover_endpoint, force, dry_run) = match payload {
         RegisterDeploymentRequest::Http {
             uri,
             additional_headers,
-            routing_header,
             use_http_11,
             force,
             dry_run,
@@ -71,13 +69,6 @@ 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(
@@ -92,9 +83,8 @@ pub async fn create_deployment<V, IC>(
                             Some(http::Version::HTTP_2)
                         },
                     ),
-                    additional_headers,
+                    additional_headers.unwrap_or_default().into(),
                 ),
-                routing_header.map(|h| (h.key, h.value)),
                 force,
                 dry_run,
             )
@@ -116,7 +106,6 @@ pub async fn create_deployment<V, IC>(
                 ),
                 additional_headers.unwrap_or_default().into(),
             ),
-            None,
             force,
             dry_run,
         ),
@@ -136,7 +125,7 @@ pub async fn create_deployment<V, IC>(
 
     let (deployment, services) = state
         .schema_registry
-        .register_deployment(discover_endpoint, routing_header, force, apply_mode)
+        .register_deployment(discover_endpoint, 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, HeaderName, HeaderValue, Uri, uri::PathAndQuery};
+use http::{HeaderMap, HeaderValue, Uri, uri::PathAndQuery};
 use tracing::subscriber::NoSubscriber;
 use tracing::trace;
 
@@ -90,7 +90,6 @@ 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> {
@@ -129,7 +128,6 @@ 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(),
                 )
@@ -155,7 +153,6 @@ 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/types/src/config/admin.rs

@@ -13,7 +13,7 @@ use std::num::NonZeroUsize;
 use std::path::PathBuf;
 use std::time::Duration;
 
-use http::Uri;
+use http::{HeaderName, Uri};
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
 use tokio::sync::Semaphore;
@@ -42,6 +42,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.
@@ -130,6 +143,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/metadata/updater/mod.rs

@@ -27,7 +27,7 @@ use crate::schema::invocation_target::{
 use crate::schema::subscriptions::{
     EventInvocationTargetTemplate, Sink, Source, Subscription, SubscriptionValidator,
 };
-use http::{HeaderName, HeaderValue, Uri};
+use http::{HeaderValue, Uri};
 use serde_json::Value;
 use std::borrow::Borrow;
 use std::collections::HashMap;
@@ -344,8 +344,7 @@ impl SchemaUpdater {
 
     pub fn add_deployment(
         &mut self,
-        mut deployment_metadata: DeploymentMetadata,
-        routing_header: Option<(HeaderName, HeaderValue)>,
+        deployment_metadata: DeploymentMetadata,
         services: Vec<endpoint_manifest::Service>,
         force: bool,
     ) -> Result<DeploymentId, SchemaError> {
@@ -354,6 +353,8 @@ impl SchemaUpdater {
             .map(|c| ServiceName::try_from(c.name.to_string()).map(|name| (name, c)))
             .collect::<Result<HashMap<_, _>, _>>()?;
 
+        let deployment_routing_headers = &Configuration::pinned().admin.deployment_routing_headers;
+
         // Did we find an existing deployment with a conflicting endpoint url?
         let existing_deployment = self
             .schema
@@ -363,15 +364,16 @@ impl SchemaUpdater {
                 schemas.ty.protocol_type() == deployment_metadata.ty.protocol_type()
                     && schemas.ty.normalized_address()
                         == deployment_metadata.ty.normalized_address()
-                    && routing_header.as_ref().is_none_or(
-                        |(routing_header_key, routing_header_value)| {
-                            schemas
+                    && deployment_routing_headers.iter().all(|routing_header_key| {
+                        deployment_metadata
+                            .delivery_options
+                            .additional_headers
+                            .get(routing_header_key)
+                            == schemas
                                 .delivery_options
                                 .additional_headers
                                 .get(routing_header_key)
-                                .is_some_and(|v| v == routing_header_value)
-                        },
-                    )
+                    })
             })
             // There are few situations where we might have multiple deployments for the same endpoint:
             // * If the user specified in the at least two previous registrations the routing-header,
@@ -442,12 +444,6 @@ impl SchemaUpdater {
             computed_services.insert(service_name.to_string(), Arc::new(new_service_revision));
         }
 
-        if let Some((key, value)) = routing_header {
-            deployment_metadata
-                .delivery_options
-                .additional_headers
-                .insert(key, value);
-        }
         self.schema.deployments.insert(
             deployment_id,
             Deployment {