Implement request/response buffering properly, add CLI args

Author: blind-oracle

Cargo.lock

@@ -17,7 +17,7 @@ hostname = "0.4.0"
 http = "1.3.1"
 http-body-util = "0.1.3"
 humantime = "2.2.0"
-ic-bn-lib = { git = "https://github.com/dfinity/ic-bn-lib", rev = "ca9b455682d8f2c855d07eeac0566c7273be8b1d", features = [
+ic-bn-lib = { git = "https://github.com/dfinity/ic-bn-lib", rev = "61c8e29bb5a38fed927e620ca319eab8d5aaa09d", features = [
     "acme-alpn",
     "cert-providers",
     "clients-hyper",

Cargo.toml

@@ -1,4 +1,4 @@
-use std::{fmt::Display, path::PathBuf, sync::Arc, time::Duration};
+use std::{cell::RefCell, fmt::Display, path::PathBuf, sync::Arc, time::Duration};
 
 use anyhow::{Context, Error, anyhow, bail};
 use arc_swap::ArcSwapOption;
@@ -22,6 +22,7 @@ use tokio::{
     fs, select,
     signal::unix::{SignalKind, signal},
     sync::{Mutex, MutexGuard},
+    task_local,
     time::timeout,
 };
 use tokio_util::sync::CancellationToken;
@@ -30,6 +31,16 @@ use url::Url;
 
 pub type LBBackendRouter = BackendRouter<Arc<Backend>, Request, Response, ic_bn_lib::http::Error>;
 
+task_local! {
+    pub static REQUEST_CONTEXT: RefCell<RequestContext>;
+}
+
+/// Request context information
+#[derive(Debug, Clone, Default)]
+pub struct RequestContext {
+    pub backend: Option<Arc<Backend>>,
+}
+
 #[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)]
 pub struct Config {
     strategy: Strategy,
@@ -419,6 +430,11 @@ impl ExecutesRequest<Arc<Backend>> for RequestExecutor {
         backend: &Arc<Backend>,
         mut req: Self::Request,
     ) -> Result<Self::Response, Self::Error> {
+        // Store the selected backend in the request context
+        let _ = REQUEST_CONTEXT.try_with(|x| {
+            x.borrow_mut().backend = Some(backend.clone());
+        });
+
         let uri = match Uri::builder()
             .scheme(backend.url.scheme())
             .authority(backend.url.authority())
@@ -452,16 +468,6 @@ impl ExecutesRequest<Arc<Backend>> for RequestExecutor {
         strip_connection_headers(req.headers_mut());
 
         // Execute it
-        let result = self.client.execute(req).await.map(|mut x| {
-            // Insert the backend into the response for observability
-            x.extensions_mut().insert(backend.clone());
-            x
-        });
-
-        if let Err(e) = &result {
-            warn!("Unable to execute request: {e:#}");
-        }
-
-        result
+        self.client.execute(req).await
     }
 }

src/backend.rs

@@ -90,6 +90,15 @@ pub struct Network {
     /// *** Dangerous *** - use only for testing.
     #[clap(env, long)]
     pub network_http_client_insecure_bypass_tls_verification: bool,
+
+    /// Whether to buffer request body from the client before sending it to the backend.
+    /// If `retry_attempts` is >1 then this is implicitly enabled.
+    #[clap(env, long)]
+    pub network_request_body_buffer: bool,
+
+    /// Whether to buffer response body from the backend before sending it to the client.
+    #[clap(env, long)]
+    pub network_response_body_buffer: bool,
 }
 
 #[derive(Args)]
@@ -177,8 +186,8 @@ pub struct Config {
 pub struct Retry {
     /// Number of request attempts to do.
     /// Only network errors are retried, not HTTP codes.
-    /// If the number of attempts is 1 then we don't retry and don't buffer the request body.
-    #[clap(env, long, default_value = "3", value_parser = clap::value_parser!(u8).range(1..))]
+    /// If the number of attempts is 1 then we don't retry.
+    #[clap(env, long, default_value = "1", value_parser = clap::value_parser!(u8).range(1..))]
     pub retry_attempts: u8,
 
     /// Initial retry interval, with each retry it is doubled.
@@ -190,13 +199,21 @@ pub struct Retry {
 
 #[derive(Args)]
 pub struct Limits {
-    /// Maximum request body size.
+    /// Maximum request body size if buffering.
     #[clap(env, long, default_value = "10MB", value_parser = parse_size_usize)]
     pub limits_request_body_size: usize,
 
-    /// Maximum time allowed to send the request body.
+    /// Maximum time allowed to buffer the request body.
     #[clap(env, long, default_value = "30s", value_parser = parse_duration)]
     pub limits_request_body_timeout: Duration,
+
+    /// Maximum response body size if buffering.
+    #[clap(env, long, default_value = "30MB", value_parser = parse_size_usize)]
+    pub limits_response_body_size: usize,
+
+    /// Maximum time allowed to buffer the response body.
+    #[clap(env, long, default_value = "60s", value_parser = parse_duration)]
+    pub limits_response_body_timeout: Duration,
 }
 
 #[derive(Args)]

src/cli.rs

@@ -1,11 +1,11 @@
 use std::sync::{Arc, OnceLock};
 
 use anyhow::{Context, Error};
-use axum::Router;
+use axum::{Router, body::Body};
 use ic_bn_lib::{
     http::{
-        self as bnhttp, HyperClientLeastLoaded, ReqwestClient, ServerBuilder, dns,
-        redirect_to_https,
+        self as bnhttp, ClientHttp, HyperClient, HyperClientLeastLoaded, ReqwestClient,
+        ServerBuilder, dns, redirect_to_https,
     },
     rustls,
     tasks::TaskManager,
@@ -69,12 +69,16 @@ pub async fn main(
             .context("unable to setup HTTP client")?,
     );
 
-    let http_client = Arc::new(HyperClientLeastLoaded::new(
-        http_client_opts,
-        resolver,
-        cli.network.network_http_client_count as usize,
-        Some(&registry),
-    ));
+    let http_client: Arc<dyn ClientHttp<Body>> = if cli.network.network_http_client_count > 1 {
+        Arc::new(HyperClientLeastLoaded::new(
+            http_client_opts,
+            resolver,
+            cli.network.network_http_client_count as usize,
+            Some(&registry),
+        ))
+    } else {
+        Arc::new(HyperClient::new(http_client_opts, resolver))
+    };
 
     // Setup Vector
     let vector = cli.log.vector.log_vector_url.as_ref().map(|_| {

src/core.rs

@@ -1,4 +1,5 @@
 use std::{
+    cell::RefCell,
     sync::Arc,
     time::{Duration, Instant},
 };
@@ -28,7 +29,7 @@ use tokio::sync::oneshot;
 use tracing::info;
 
 use crate::{
-    backend::Backend,
+    backend::{REQUEST_CONTEXT, RequestContext},
     core::{ENV, HOSTNAME},
     middleware::request_id::RequestId,
     routing::Retries,
@@ -210,14 +211,19 @@ pub async fn middleware(
 
     // Execute the request
     let start = Instant::now();
-    let mut response = next.run(request).await;
+    let (mut response, ctx) = REQUEST_CONTEXT
+        .scope(RefCell::new(RequestContext::default()), async {
+            let r = next.run(request).await;
+            let ctx = REQUEST_CONTEXT.try_with(|x| x.clone()).unwrap_or_default();
+            (r, ctx.into_inner())
+        })
+        .await;
     let duration = start.elapsed().as_secs_f64();
 
-    let backend = response
-        .extensions_mut()
-        .remove::<Arc<Backend>>()
+    let backend = ctx
+        .backend
         .map(|x| x.name.clone())
-        .unwrap_or_default();
+        .unwrap_or_else(|| "unknown".into());
     let response_size = response
         .body()
         .size_hint()