Enough functionality to implement adaptive load shedding

Cargo.lock

@@ -297,6 +297,7 @@ dependencies = [
  "libc",
  "libtest-mimic",
  "linkme",
+ "little-loadshedder",
  "lru",
  "maplit",
  "mediatype",
@@ -3929,6 +3930,16 @@ version = "0.4.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89"
 
+[[package]]
+name = "little-loadshedder"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1edea14a875dba8659cb864480c27e8eb2d099542b657729a7567625a2c1d65a"
+dependencies = [
+ "tokio",
+ "tower",
+]
+
 [[package]]
 name = "lock_api"
 version = "0.4.12"

apollo-router/Cargo.toml

@@ -116,6 +116,8 @@ jsonwebtoken = "9.3.0"
 lazy_static = "1.4.0"
 libc = "0.2.155"
 linkme = "0.3.27"
+# little-loadshedder = { path = "../../little-loadshedder" }
+little-loadshedder = { version = "0.2.0" }
 lru = "0.12.3"
 maplit = "1.0.2"
 mediatype = "0.19.18"

apollo-router/src/axum_factory/axum_http_server_factory.rs

@@ -1,5 +1,4 @@
 //! Axum http server factory. Axum provides routing capability on top of Hyper HTTP.
-use std::fmt::Display;
 use std::pin::Pin;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::AtomicU64;
@@ -33,6 +32,7 @@ use serde_json::json;
 use tokio::net::UnixListener;
 use tokio::sync::mpsc;
 use tokio_rustls::TlsAcceptor;
+use tower::load_shed::error::Overloaded;
 use tower::service_fn;
 use tower::BoxError;
 use tower::ServiceBuilder;
@@ -429,7 +429,7 @@ pub(crate) fn span_mode(configuration: &Configuration) -> SpanMode {
         .unwrap_or_default()
 }
 
-async fn decompression_error(_error: BoxError) -> axum::response::Response {
+async fn decompression_error(_error: BoxError) -> impl IntoResponse {
     (StatusCode::BAD_REQUEST, "cannot decompress request body").into_response()
 }
 
@@ -657,10 +657,10 @@ async fn handle_graphql(
     let processing_seconds = dur.as_secs_f64();
 
     f64_histogram!(
-        "apollo.router.processing.time",
-        "Time spent by the router actually working on the request, not waiting for its network calls or other queries being processed",
-        processing_seconds
-    );
+            "apollo.router.processing.time",
+            "Time spent by the router actually working on the request, not waiting for its network calls or other queries being processed",
+            processing_seconds
+        );
 
     match res {
         Err(err) => internal_server_error(err),
@@ -689,23 +689,32 @@ async fn handle_graphql(
 
 fn internal_server_error<T>(err: T) -> Response
 where
-    T: Display,
+    T: Into<BoxError>,
 {
+    let err: BoxError = err.into();
+
+    let code = if err.is::<Overloaded>() {
+        StatusCode::SERVICE_UNAVAILABLE
+    } else {
+        StatusCode::INTERNAL_SERVER_ERROR
+    };
+
     tracing::error!(
-        code = "INTERNAL_SERVER_ERROR",
+        code = code.to_string(),
         %err,
     );
 
     // This intentionally doesn't include an error message as this could represent leakage of internal information.
     // The error message is logged above.
     let error = graphql::Error::builder()
-        .message("internal server error")
+        .message(code.to_string())
+        // Note: Decide exactly what this extension_code should be for SERVICE_UNAVAILABLE
         .extension_code("INTERNAL_SERVER_ERROR")
         .build();
 
     let response = graphql::Response::builder().error(error).build();
 
-    (StatusCode::INTERNAL_SERVER_ERROR, Json(json!(response))).into_response()
+    (code, Json(json!(response))).into_response()
 }
 
 struct CancelHandler<'a> {

apollo-router/src/plugin/test/service.rs

@@ -49,7 +49,7 @@ macro_rules! mock_service {
                     std::task::Poll::Ready(Ok(()))
                 }
                 fn call(&mut self, req: $request_type) -> Self::Future {
-                    let r  = self.call(req);
+                    let r = self.call(req);
                     Box::pin(async move { r })
                 }
             }

apollo-router/src/plugins/limits/layer.rs

@@ -23,20 +23,29 @@ struct BodyLimitControlInner {
     current: AtomicUsize,
 }
 
+impl Clone for BodyLimitControlInner {
+    fn clone(&self) -> Self {
+        Self {
+            limit: AtomicUsize::new(self.limit.load(std::sync::atomic::Ordering::SeqCst)),
+            current: AtomicUsize::new(0),
+        }
+    }
+}
+
 /// This structure allows the body limit to be updated dynamically.
 /// It also allows the error message to be updated
 #[derive(Clone)]
 pub(crate) struct BodyLimitControl {
-    inner: Arc<BodyLimitControlInner>,
+    inner: BodyLimitControlInner,
 }
 
 impl BodyLimitControl {
     pub(crate) fn new(limit: usize) -> Self {
         Self {
-            inner: Arc::new(BodyLimitControlInner {
+            inner: BodyLimitControlInner {
                 limit: AtomicUsize::new(limit),
                 current: AtomicUsize::new(0),
-            }),
+            },
         }
     }
 

apollo-router/src/plugins/test.rs

@@ -8,6 +8,7 @@ use apollo_compiler::validation::Valid;
 use serde_json::Value;
 use tower::BoxError;
 use tower::ServiceBuilder;
+use tower::ServiceExt;
 use tower_service::Service;
 
 use crate::introspection::IntrospectionCache;
@@ -149,7 +150,7 @@ impl<T: Plugin> PluginTestHarness<T> {
                 .service_fn(move |req: router::Request| async move { (response_fn)(req).await }),
         );
 
-        self.plugin.router_service(service).call(request).await
+        self.plugin.router_service(service).oneshot(request).await
     }
 
     pub(crate) async fn call_supergraph(
@@ -162,7 +163,10 @@ impl<T: Plugin> PluginTestHarness<T> {
                 .service_fn(move |req: supergraph::Request| async move { Ok((response_fn)(req)) }),
         );
 
-        self.plugin.supergraph_service(service).call(request).await
+        self.plugin
+            .supergraph_service(service)
+            .oneshot(request)
+            .await
     }
 
     #[allow(dead_code)]
@@ -176,7 +180,10 @@ impl<T: Plugin> PluginTestHarness<T> {
                 .service_fn(move |req: execution::Request| async move { Ok((response_fn)(req)) }),
         );
 
-        self.plugin.execution_service(service).call(request).await
+        self.plugin
+            .execution_service(service)
+            .oneshot(request)
+            .await
     }
 
     #[allow(dead_code)]

apollo-router/src/plugins/traffic_shaping/deduplication.rs

@@ -182,19 +182,20 @@ where
     type Error = BoxError;
     type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
-    fn poll_ready(&mut self, _cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
-        Poll::Ready(Ok(()))
+    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.service.poll_ready(cx)
     }
 
     fn call(&mut self, request: SubgraphRequest) -> Self::Future {
         let service = self.service.clone();
+        let inner = std::mem::replace(&mut self.service, service);
 
         if request.operation_kind == OperationKind::Query {
             let wait_map = self.wait_map.clone();
 
-            Box::pin(async move { Self::dedup(service, wait_map, request).await })
+            Box::pin(async move { Self::dedup(inner, wait_map, request).await })
         } else {
-            Box::pin(async move { service.oneshot(request).await })
+            Box::pin(async move { inner.oneshot(request).await })
         }
     }
 }