♻️ Add ConcurrencyLimitLayer middleware

Cargo.toml

@@ -36,7 +36,8 @@ async-trait = "0.1.53"
 serde = { version = "1.0.136", features = ["derive"] }
 serde_json = "1.0.79"
 serde_yaml = "0.8.24"
-tower-http = { version = "0.2.3", features = ["trace", "set-header"] }
+tower = { version= "0.4.12", features = ["limit"] }
+tower-http = { version = "0.3.4", features = ["trace", "set-header"] }
 clap = { version = "3.1.0", features = ["derive", "env"] }
 chrono = "0.4.19"
 

src/app/config.rs

@@ -58,6 +58,9 @@ pub struct Config {
 
     #[clap(long = "distributed-wait-duration", default_value_t = 300)]
     pub distributed_wait_duration: u16,
+
+    #[clap(long = "num-concurrent-requests")]
+    pub num_concurrent_requests: Option<u16>,
 }
 
 impl Config {
@@ -86,6 +89,17 @@ impl Config {
         }
         Ok(None)
     }
+
+    pub fn num_concurrent_requests(&self) -> u16 {
+        if let Some(n) = self.num_concurrent_requests {
+            n
+        } else {
+            match self.mode() {
+                Mode::Local => 5,
+                Mode::Distributed => 100,
+            }
+        }
+    }
 }
 
 #[derive(Clone, Debug)]

src/app/web/router.rs

@@ -3,18 +3,21 @@ use crate::app::web::handler;
 use axum::http::{header, HeaderValue};
 use axum::routing::{get, post};
 use axum::{AddExtensionLayer, Router};
+use tower::limit::concurrency::ConcurrencyLimitLayer;
 use tower_http::set_header::SetRequestHeaderLayer;
 use tower_http::trace::TraceLayer;
 
 pub fn routes(deps: DynDependencyProvider) -> axum::Router {
+    let num_concurrent_requests = deps.get_config().num_concurrent_requests();
     let template_execute = post(handler::execute_template)
-        .layer(TraceLayer::new_for_http())
         // TODO remove SetRequestHeaderLayer once Argo sends correct Content-Type header
+        .layer(AddExtensionLayer::new(deps))
         .layer(SetRequestHeaderLayer::if_not_present(
             header::CONTENT_TYPE,
             HeaderValue::from_static("application/json"),
         ))
-        .layer(AddExtensionLayer::new(deps));
+        .layer(ConcurrencyLimitLayer::new(num_concurrent_requests.into()))
+        .layer(TraceLayer::new_for_http());
 
     Router::new()
         .route("/healthz", get(|| async { "ok\n" }))