add support of rate limit and timeout (#1347)

Signed-off-by: Benjamin Coenen <[email protected]>

NEXT_CHANGELOG.md

@@ -29,6 +29,29 @@ By [@USERNAME](https://github.com/USERNAME) in https://github.com/apollographql/
 
 ## 🚀 Features
 
+### Add support of global rate limit and timeout. [PR #1347](https://github.com/apollographql/router/pull/1347)
+
+Additions to the traffic shaping plugin:
+- **Global rate limit** - If you want to rate limit requests to subgraphs or to the router itself.
+- **Timeout**: - Set a timeout to subgraphs and router requests.
+
+```yaml
+traffic_shaping:
+  router: # Rules applied to requests from clients to the router
+    global_rate_limit: # Accept a maximum of 10 requests per 5 secs. Excess requests must be rejected.
+      capacity: 10
+      interval: 5s # Must not be greater than 18_446_744_073_709_551_615 milliseconds and not less than 0 milliseconds
+    timeout: 50s # If a request to the router takes more than 50secs then cancel the request (30 sec by default)
+  subgraphs: # Rules applied to requests from the router to individual subgraphs
+    products:
+      global_rate_limit: # Accept a maximum of 10 requests per 5 secs from the router. Excess requests must be rejected.
+        capacity: 10
+        interval: 5s # Must not be greater than 18_446_744_073_709_551_615 milliseconds and not less than 0 milliseconds
+      timeout: 50s # If a request to the subgraph 'products' takes more than 50secs then cancel the request (30 sec by default)
+```
+
+By [@bnjjj](https://github.com/bnjjj) in https://github.com/apollographql/router/pull/1347
+
 ## 🐛 Fixes
 
 ## 🛠 Maintenance

apollo-router/Cargo.toml

@@ -133,6 +133,7 @@ tower-http = { version = "0.3.4", features = [
     "decompression-br",
     "decompression-deflate",
     "decompression-gzip",
+    "timeout"
 ] }
 tower-service = "0.3.2"
 tower-test = "0.4.0"
@@ -144,7 +145,7 @@ tracing-subscriber = { version = "0.3.11", features = ["env-filter", "json"] }
 url = { version = "2.2.2", features = ["serde"] }
 urlencoding = "2.1.0"
 yaml-rust = "0.4.5"
-
+pin-project-lite = "0.2.9"
 
 [target.'cfg(macos)'.dependencies]
 uname = "0.1.1"

apollo-router/src/axum_http_server_factory.rs

@@ -62,6 +62,8 @@ use crate::http_server_factory::HttpServerHandle;
 use crate::http_server_factory::Listener;
 use crate::http_server_factory::NetworkStream;
 use crate::plugin::Handler;
+use crate::plugins::traffic_shaping::Elapsed;
+use crate::plugins::traffic_shaping::RateLimited;
 use crate::router::ApolloRouterError;
 use crate::router_factory::RouterServiceFactory;
 
@@ -504,6 +506,14 @@ where
 
             match service.call(Request::from_parts(head, body).into()).await {
                 Err(e) => {
+                    if let Some(source_err) = e.source() {
+                        if source_err.is::<RateLimited>() {
+                            return RateLimited::new().into_response();
+                        }
+                        if source_err.is::<Elapsed>() {
+                            return Elapsed::new().into_response();
+                        }
+                    }
                     tracing::error!("router service call failed: {}", e);
                     (
                         StatusCode::INTERNAL_SERVER_ERROR,
@@ -559,6 +569,15 @@ where
         }
         Err(e) => {
             tracing::error!("router service is not available to process request: {}", e);
+            if let Some(source_err) = e.source() {
+                if source_err.is::<RateLimited>() {
+                    return RateLimited::new().into_response();
+                }
+                if source_err.is::<Elapsed>() {
+                    return Elapsed::new().into_response();
+                }
+            }
+
             (
                 StatusCode::SERVICE_UNAVAILABLE,
                 "router service is not available to process request",

apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap

@@ -1,5 +1,6 @@
 ---
 source: apollo-router/src/configuration/mod.rs
+assertion_line: 914
 expression: "&schema"
 ---
 {
@@ -1887,6 +1888,7 @@ expression: "&schema"
       "type": "object",
       "properties": {
         "all": {
+          "description": "Applied on all subgraphs",
           "type": "object",
           "properties": {
             "compression": {
@@ -1899,16 +1901,79 @@ expression: "&schema"
               ],
               "nullable": true
             },
+            "global_rate_limit": {
+              "description": "Enable global rate limiting",
+              "type": "object",
+              "required": [
+                "capacity",
+                "interval"
+              ],
+              "properties": {
+                "capacity": {
+                  "description": "Number of requests allowed",
+                  "type": "integer",
+                  "format": "uint64",
+                  "minimum": 1.0
+                },
+                "interval": {
+                  "description": "Per interval",
+                  "type": "string"
+                }
+              },
+              "additionalProperties": false,
+              "nullable": true
+            },
             "query_deduplication": {
               "description": "Enable query deduplication",
               "type": "boolean",
               "nullable": true
+            },
+            "timeout": {
+              "description": "Enable timeout for incoming requests",
+              "default": null,
+              "type": "string"
+            }
+          },
+          "additionalProperties": false,
+          "nullable": true
+        },
+        "router": {
+          "description": "Applied at the router level",
+          "type": "object",
+          "properties": {
+            "global_rate_limit": {
+              "description": "Enable global rate limiting",
+              "type": "object",
+              "required": [
+                "capacity",
+                "interval"
+              ],
+              "properties": {
+                "capacity": {
+                  "description": "Number of requests allowed",
+                  "type": "integer",
+                  "format": "uint64",
+                  "minimum": 1.0
+                },
+                "interval": {
+                  "description": "Per interval",
+                  "type": "string"
+                }
+              },
+              "additionalProperties": false,
+              "nullable": true
+            },
+            "timeout": {
+              "description": "Enable timeout for incoming requests",
+              "default": null,
+              "type": "string"
             }
           },
           "additionalProperties": false,
           "nullable": true
         },
         "subgraphs": {
+          "description": "Applied on specific subgraphs",
           "type": "object",
           "additionalProperties": {
             "type": "object",
@@ -1923,10 +1988,37 @@ expression: "&schema"
                 ],
                 "nullable": true
               },
+              "global_rate_limit": {
+                "description": "Enable global rate limiting",
+                "type": "object",
+                "required": [
+                  "capacity",
+                  "interval"
+                ],
+                "properties": {
+                  "capacity": {
+                    "description": "Number of requests allowed",
+                    "type": "integer",
+                    "format": "uint64",
+                    "minimum": 1.0
+                  },
+                  "interval": {
+                    "description": "Per interval",
+                    "type": "string"
+                  }
+                },
+                "additionalProperties": false,
+                "nullable": true
+              },
               "query_deduplication": {
                 "description": "Enable query deduplication",
                 "type": "boolean",
                 "nullable": true
+              },
+              "timeout": {
+                "description": "Enable timeout for incoming requests",
+                "default": null,
+                "type": "string"
               }
             },
             "additionalProperties": false

apollo-router/src/plugins/telemetry/mod.rs

@@ -225,8 +225,6 @@ impl Plugin for Telemetry {
         // The trace provider will not be shut down if drop is not called and it will result in a hang.
         // Don't add anything fallible after the tracer provider has been created.
         let tracer_provider = Self::create_tracer_provider(&config)?;
-        //
-        // let metrics_response_queries = Self::create_metrics_queries(&config)?;
 
         let plugin = Ok(Telemetry {
             spaceport_shutdown: shutdown_tx,
@@ -236,7 +234,6 @@ impl Plugin for Telemetry {
             meter_provider: builder.meter_provider(),
             apollo_metrics_sender: builder.apollo_metrics_provider(),
             config,
-            // metrics_response_queries,
         });
 
         // We're safe now for shutdown.