Limit concurrent HTTP requests to 32 by default

Author: Ralith

app/buck2_http/src/client.rs

@@ -26,6 +26,7 @@ use hyper::Response;
 use hyper::client::ResponseFuture;
 use hyper::client::connect::Connect;
 use tokio::io::AsyncReadExt;
+use tokio::sync::Semaphore;
 use tokio_util::io::StreamReader;
 
 use crate::HttpError;
@@ -49,6 +50,9 @@ pub struct HttpClient {
     supports_vpnless: bool,
     http2: bool,
     stats: HttpNetworkStats,
+    // tokio::sync::Semaphore doesn't impl Allocative
+    #[allocative(skip)]
+    concurrent_requests_budget: Arc<Semaphore>,
 }
 
 impl HttpClient {
@@ -124,6 +128,7 @@ impl HttpClient {
             );
             change_scheme_to_http(&mut request)?;
         }
+        let semaphore_guard = self.concurrent_requests_budget.acquire().await.unwrap();
         let resp = self.inner.request(request).await.map_err(|e| {
             if is_hyper_error_due_to_timeout(&e) {
                 HttpError::Timeout {
@@ -134,11 +139,14 @@ impl HttpClient {
                 HttpError::SendRequest { uri, source: e }
             }
         })?;
-        Ok(
-            resp.map(|body| {
-                CountingStream::new(body, self.stats.downloaded_bytes().dupe()).boxed()
-            }),
-        )
+        Ok(resp.map(|body| {
+            CountingStream::new(body, self.stats.downloaded_bytes().dupe())
+                .inspect(|_| {
+                    // Ensure we keep a concurrent request permit alive until the stream is consumed
+                    _ = semaphore_guard;
+                })
+                .boxed()
+        }))
     }
 
     /// Send a generic request.

app/buck2_http/src/client/builder.rs

@@ -28,6 +28,7 @@ use hyper_timeout::TimeoutConnector;
 use rustls::ClientConfig;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
+use tokio::sync::Semaphore;
 use tokio_rustls::TlsConnector;
 
 use super::HttpClient;
@@ -66,6 +67,7 @@ pub struct HttpClientBuilder {
     supports_vpnless: bool,
     http2: bool,
     timeout_config: Option<TimeoutConfig>,
+    max_concurrent_requests: usize,
 }
 
 impl HttpClientBuilder {
@@ -104,6 +106,10 @@ impl HttpClientBuilder {
             supports_vpnless: false,
             http2: true,
             timeout_config: None,
+            // Semi-arbitrary sensible default: big enough to ensure good utilization of a typical
+            // internet link, small enough that we won't run out of FDs or make a server think we're
+            // DoSing them.
+            max_concurrent_requests: 32,
         })
     }
 
@@ -214,6 +220,11 @@ impl HttpClientBuilder {
         self.supports_vpnless
     }
 
+    pub fn with_max_concurrent_requests(&mut self, max_concurrent_requests: usize) -> &mut Self {
+        self.max_concurrent_requests = max_concurrent_requests;
+        self
+    }
+
     fn build_inner(&self) -> Arc<dyn RequestClient> {
         match (self.proxies.as_slice(), &self.timeout_config) {
             // Construct x2p unix socket client.
@@ -294,6 +305,7 @@ impl HttpClientBuilder {
             supports_vpnless: self.supports_vpnless,
             http2: self.http2,
             stats: HttpNetworkStats::new(),
+            concurrent_requests_budget: Arc::new(Semaphore::new(self.max_concurrent_requests)),
         }
     }
 }