Go

Author: ardatan

docs/README.md

@@ -9,7 +9,7 @@
 |[**log**](#log)|`object`|The router logger configuration.<br/>Default: `{"filter":null,"format":"json","level":"info"}`<br/>||
 |[**query\_planner**](#query_planner)|`object`|Query planning configuration.<br/>Default: `{"allow_expose":false,"timeout":"10s"}`<br/>||
 |[**supergraph**](#supergraph)|`object`|Configuration for the Federation supergraph source. By default, the router will use a local file-based supergraph source (`./supergraph.graphql`).<br/>Default: `{"path":"supergraph.graphql","source":"file"}`<br/>||
-|[**traffic\_shaping**](#traffic_shaping)|`object`|Configuration for the traffic-shaper executor. Use these configurations to control how requests are being executed to subgraphs.<br/>Default: `{"all":{"dedupe_enabled":true,"max_connections_per_host":100,"pool_idle_timeout_seconds":50}}`<br/>||
+|[**traffic\_shaping**](#traffic_shaping)|`object`|Configuration for the traffic-shaper executor. Use these configurations to control how requests are being executed to subgraphs.<br/>Default: `{"all":{"dedupe_enabled":true,"pool_idle_timeout_seconds":50},"max_connections_per_host":100}`<br/>||
 
 **Additional Properties:** not allowed  
 **Example**
@@ -48,8 +48,8 @@ supergraph:
 traffic_shaping:
   all:
     dedupe_enabled: true
-    max_connections_per_host: 100
     pool_idle_timeout_seconds: 50
+  max_connections_per_host: 100
 
 ```
 
@@ -1125,16 +1125,17 @@ Configuration for the traffic-shaper executor. Use these configurations to contr
 
 |Name|Type|Description|Required|
 |----|----|-----------|--------|
-|[**all**](#traffic_shapingall)|`object`|The default configuration that will be applied to all subgraphs, unless overridden by a specific subgraph configuration.<br/>Default: `{"dedupe_enabled":true,"max_connections_per_host":100,"pool_idle_timeout_seconds":50}`<br/>||
+|[**all**](#traffic_shapingall)|`object`|The default configuration that will be applied to all subgraphs, unless overridden by a specific subgraph configuration.<br/>Default: `{"dedupe_enabled":true,"pool_idle_timeout_seconds":50}`<br/>||
+|**max\_connections\_per\_host**|`integer`|Limits the concurrent amount of requests/connections per host/subgraph.<br/>Default: `100`<br/>Format: `"uint"`<br/>Minimum: `0`<br/>||
 |[**subgraphs**](#traffic_shapingsubgraphs)|`object`|Optional per-subgraph configurations that will override the default configuration for specific subgraphs.<br/>||
 
 **Example**
 
 ```yaml
 all:
   dedupe_enabled: true
-  max_connections_per_host: 100
   pool_idle_timeout_seconds: 50
+max_connections_per_host: 100
 
 ```
 
@@ -1149,15 +1150,13 @@ The default configuration that will be applied to all subgraphs, unless overridd
 |Name|Type|Description|Required|
 |----|----|-----------|--------|
 |**dedupe\_enabled**|`boolean`|Enables/disables request deduplication to subgraphs.<br/><br/>When requests exactly matches the hashing mechanism (e.g., subgraph name, URL, headers, query, variables), and are executed at the same time, they will<br/>be deduplicated by sharing the response of other in-flight requests.<br/>Default: `true`<br/>||
-|**max\_connections\_per\_host**|`integer`|Limits the concurrent amount of requests/connections per host/subgraph.<br/>Default: `100`<br/>Format: `"uint"`<br/>Minimum: `0`<br/>||
 |**pool\_idle\_timeout\_seconds**|`integer`|Timeout for idle sockets being kept-alive.<br/>Default: `50`<br/>Format: `"uint64"`<br/>Minimum: `0`<br/>||
 |**timeout**||Optional timeout configuration for requests to subgraphs.<br/><br/>Example with a fixed duration:<br/>```yaml<br/>  timeout:<br/>    duration: 5s<br/>```<br/><br/>Or with a VRL expression that can return a duration based on the operation kind:<br/>```yaml<br/>  timeout:<br/>    expression: \|<br/>     if (.request.operation.type == "mutation") {<br/>       10000<br/>     } else {<br/>       5000<br/>     }<br/>```<br/>||
 
 **Example**
 
 ```yaml
 dedupe_enabled: true
-max_connections_per_host: 100
 pool_idle_timeout_seconds: 50
 
 ```
@@ -1182,15 +1181,13 @@ Optional per-subgraph configurations that will override the default configuratio
 |Name|Type|Description|Required|
 |----|----|-----------|--------|
 |**dedupe\_enabled**|`boolean`|Enables/disables request deduplication to subgraphs.<br/><br/>When requests exactly matches the hashing mechanism (e.g., subgraph name, URL, headers, query, variables), and are executed at the same time, they will<br/>be deduplicated by sharing the response of other in-flight requests.<br/>Default: `true`<br/>||
-|**max\_connections\_per\_host**|`integer`|Limits the concurrent amount of requests/connections per host/subgraph.<br/>Default: `100`<br/>Format: `"uint"`<br/>Minimum: `0`<br/>||
 |**pool\_idle\_timeout\_seconds**|`integer`|Timeout for idle sockets being kept-alive.<br/>Default: `50`<br/>Format: `"uint64"`<br/>Minimum: `0`<br/>||
 |**timeout**||Optional timeout configuration for requests to subgraphs.<br/><br/>Example with a fixed duration:<br/>```yaml<br/>  timeout:<br/>    duration: 5s<br/>```<br/><br/>Or with a VRL expression that can return a duration based on the operation kind:<br/>```yaml<br/>  timeout:<br/>    expression: \|<br/>     if (.request.operation.type == "mutation") {<br/>       10000<br/>     } else {<br/>       5000<br/>     }<br/>```<br/>||
 
 **Example**
 
 ```yaml
 dedupe_enabled: true
-max_connections_per_host: 100
 pool_idle_timeout_seconds: 50
 
 ```

lib/executor/src/executors/map.rs

@@ -2,7 +2,9 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
 
 use bytes::{BufMut, Bytes, BytesMut};
 use dashmap::DashMap;
-use hive_router_config::{traffic_shaping::TrafficShapingExecutorConfig, TrafficShapingConfig};
+use hive_router_config::{
+    traffic_shaping::TrafficShapingConfig, traffic_shaping::TrafficShapingExecutorConfig,
+};
 use http::Uri;
 use http_body_util::Full;
 use hyper_tls::HttpsConnector;
@@ -78,8 +80,10 @@ impl SubgraphExecutorMap {
         subgraph_endpoint_map: HashMap<String, String>,
         config: TrafficShapingConfig,
     ) -> Result<Self, SubgraphExecutorError> {
-        let global_client_arc = from_traffic_shaping_config_to_client(&config.all);
-        let global_semaphores_by_origin: DashMap<String, Arc<Semaphore>> = DashMap::new();
+        let max_connections_per_host = config.max_connections_per_host;
+        let global_client_arc =
+            from_traffic_shaping_config_to_client(&config.all, max_connections_per_host);
+        let semaphores_by_origin: DashMap<String, Arc<Semaphore>> = DashMap::new();
         let global_config_arc = Arc::new(config.all);
         let global_in_flight_requests: Arc<
             DashMap<u64, Arc<OnceCell<SharedResponse>>, ABuildHasher>,
@@ -107,19 +111,11 @@ impl SubgraphExecutorMap {
 
                 let subgraph_config = config.subgraphs.get(&subgraph_name);
 
-                let semaphore = get_semaphore_for_subgraph(
-                    &origin,
-                    &global_semaphores_by_origin,
-                    subgraph_config
-                        .map(|cfg| cfg.max_connections_per_host)
-                        .unwrap_or(global_config_arc.max_connections_per_host),
-                    global_config_arc.max_connections_per_host,
-                );
-
                 let http_client = get_http_client_for_subgraph(
                     subgraph_config,
                     &global_config_arc,
                     &global_client_arc,
+                    max_connections_per_host,
                 );
 
                 // TODO: Maybe reuse the in-flight requests map in some cases ???
@@ -131,6 +127,11 @@ impl SubgraphExecutorMap {
                     .map(|cfg| Arc::new(cfg.clone()))
                     .unwrap_or_else(|| global_config_arc.clone());
 
+                let semaphore = semaphores_by_origin
+                    .entry(origin.to_string())
+                    .or_insert_with(|| Arc::new(Semaphore::new(max_connections_per_host)))
+                    .clone();
+
                 let mut executor = HTTPSubgraphExecutor::new(
                     endpoint_uri.clone(),
                     http_client,
@@ -158,12 +159,13 @@ impl SubgraphExecutorMap {
 // Create a new hyper client based on the traffic shaping config
 pub fn from_traffic_shaping_config_to_client(
     config: &TrafficShapingExecutorConfig,
+    max_connections_per_host: usize,
 ) -> Arc<Client<HttpsConnector<HttpConnector>, Full<Bytes>>> {
     Arc::new(
         Client::builder(TokioExecutor::new())
             .pool_timer(TokioTimer::new())
             .pool_idle_timeout(Duration::from_secs(config.pool_idle_timeout_seconds))
-            .pool_max_idle_per_host(config.max_connections_per_host)
+            .pool_max_idle_per_host(max_connections_per_host)
             .build(HttpsConnector::new()),
     )
 }
@@ -174,35 +176,16 @@ fn get_http_client_for_subgraph(
     subgraph_config: Option<&TrafficShapingExecutorConfig>,
     global_config: &TrafficShapingExecutorConfig,
     global_client: &Arc<Client<HttpsConnector<HttpConnector>, Full<Bytes>>>,
+    max_connections_per_host: usize,
 ) -> Arc<Client<HttpsConnector<HttpConnector>, Full<Bytes>>> {
     match subgraph_config {
         Some(cfg) => {
-            if global_config.max_connections_per_host == cfg.max_connections_per_host
-                && global_config.pool_idle_timeout_seconds == cfg.pool_idle_timeout_seconds
-            {
+            if global_config.pool_idle_timeout_seconds == cfg.pool_idle_timeout_seconds {
                 global_client.clone()
             } else {
-                from_traffic_shaping_config_to_client(cfg)
+                from_traffic_shaping_config_to_client(cfg, max_connections_per_host)
             }
         }
         None => global_client.clone(),
     }
 }
-
-// If the subgraph has a specific max_connections_per_host, create a new semaphore for it.
-// Otherwise, reuse the global semaphore for that origin.
-fn get_semaphore_for_subgraph(
-    origin: &str,
-    semaphores_by_origin: &DashMap<String, Arc<Semaphore>>,
-    max_connections_per_host: usize,
-    global_max_connections_per_host: usize,
-) -> Arc<Semaphore> {
-    if max_connections_per_host == global_max_connections_per_host {
-        semaphores_by_origin
-            .entry(origin.to_string())
-            .or_insert_with(|| Arc::new(Semaphore::new(global_max_connections_per_host)))
-            .clone()
-    } else {
-        Arc::new(Semaphore::new(max_connections_per_host))
-    }
-}

lib/executor/src/executors/timeout.rs

@@ -364,7 +364,7 @@ mod tests {
         "#;
 
         let config = hive_router_config::parse_yaml_config(config.to_string()).unwrap();
-        let http_client = from_traffic_shaping_config_to_client(&config.traffic_shaping.all);
+        let http_client = from_traffic_shaping_config_to_client(&config.traffic_shaping.all, 10);
         let http_executor = crate::executors::http::HTTPSubgraphExecutor::new(
             endpoint.clone(),
             http_client,

lib/router-config/src/lib.rs

@@ -6,15 +6,13 @@ pub mod query_planner;
 pub mod supergraph;
 pub mod traffic_shaping;
 
-use std::collections::HashMap;
-
 use config::{Config, Environment, File, FileFormat, FileSourceFile};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 use crate::{
     http_server::HttpServerConfig, log::LoggingConfig, query_planner::QueryPlannerConfig,
-    supergraph::SupergraphSource, traffic_shaping::TrafficShapingExecutorConfig,
+    supergraph::SupergraphSource, traffic_shaping::TrafficShapingConfig,
 };
 
 #[derive(Deserialize, Serialize, JsonSchema)]
@@ -49,16 +47,6 @@ pub struct HiveRouterConfig {
     pub headers: headers::HeadersConfig,
 }
 
-#[derive(Clone, Deserialize, Serialize, JsonSchema, Default)]
-pub struct TrafficShapingConfig {
-    /// The default configuration that will be applied to all subgraphs, unless overridden by a specific subgraph configuration.
-    #[serde(default)]
-    pub all: TrafficShapingExecutorConfig,
-    /// Optional per-subgraph configurations that will override the default configuration for specific subgraphs.
-    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
-    pub subgraphs: HashMap<String, TrafficShapingExecutorConfig>,
-}
-
 #[derive(Debug, thiserror::Error)]
 pub enum RouterConfigError {
     #[error("Failed to load configuration: {0}")]

lib/router-config/src/traffic_shaping.rs

@@ -3,12 +3,23 @@ use std::time::Duration;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
-pub struct TrafficShapingExecutorConfig {
+use std::collections::HashMap;
+
+#[derive(Clone, Deserialize, Serialize, JsonSchema)]
+pub struct TrafficShapingConfig {
+    /// The default configuration that will be applied to all subgraphs, unless overridden by a specific subgraph configuration.
+    #[serde(default)]
+    pub all: TrafficShapingExecutorConfig,
+    /// Optional per-subgraph configurations that will override the default configuration for specific subgraphs.
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub subgraphs: HashMap<String, TrafficShapingExecutorConfig>,
     /// Limits the concurrent amount of requests/connections per host/subgraph.
     #[serde(default = "default_max_connections_per_host")]
     pub max_connections_per_host: usize,
+}
 
+#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
+pub struct TrafficShapingExecutorConfig {
     /// Timeout for idle sockets being kept-alive.
     #[serde(default = "default_pool_idle_timeout_seconds")]
     pub pool_idle_timeout_seconds: u64,
@@ -60,14 +71,23 @@ where
 impl Default for TrafficShapingExecutorConfig {
     fn default() -> Self {
         Self {
-            max_connections_per_host: default_max_connections_per_host(),
             pool_idle_timeout_seconds: default_pool_idle_timeout_seconds(),
             dedupe_enabled: default_dedupe_enabled(),
             timeout: None,
         }
     }
 }
 
+impl Default for TrafficShapingConfig {
+    fn default() -> Self {
+        Self {
+            all: TrafficShapingExecutorConfig::default(),
+            subgraphs: HashMap::new(),
+            max_connections_per_host: default_max_connections_per_host(),
+        }
+    }
+}
+
 fn default_max_connections_per_host() -> usize {
     100
 }