Limit tcp connections instead of http requests

Signed-off-by: Ryan Levick <[email protected]>

Author: rylev

crates/factor-outbound-http/src/lib.rs

@@ -56,10 +56,10 @@ impl Factor for OutboundHttpFactor {
         Ok(AppState {
             wasi_http_clients: wasi::HttpClients::new(config.connection_pooling_enabled),
             connection_pooling_enabled: config.connection_pooling_enabled,
-            concurrent_outbound_requests_semaphore: config
-                .max_concurrent_requests
-                // Permit count is the max concurrent requests + 1.
-                // i.e., 0 concurrent requests means 1 total request.
+            concurrent_outbound_connections_semaphore: config
+                .max_concurrent_connections
+                // Permit count is the max concurrent connections + 1.
+                // i.e., 0 concurrent connections means 1 total connection.
                 .map(|n| Arc::new(Semaphore::new(n + 1))),
         })
     }
@@ -82,9 +82,9 @@ impl Factor for OutboundHttpFactor {
             spin_http_client: None,
             wasi_http_clients: ctx.app_state().wasi_http_clients.clone(),
             connection_pooling_enabled: ctx.app_state().connection_pooling_enabled,
-            concurrent_outbound_requests_semaphore: ctx
+            concurrent_outbound_connections_semaphore: ctx
                 .app_state()
-                .concurrent_outbound_requests_semaphore
+                .concurrent_outbound_connections_semaphore
                 .clone(),
         })
     }
@@ -110,8 +110,8 @@ pub struct InstanceState {
     wasi_http_clients: wasi::HttpClients,
     /// Whether connection pooling is enabled for this instance.
     connection_pooling_enabled: bool,
-    /// A semaphore to limit the number of concurrent outbound requests.
-    concurrent_outbound_requests_semaphore: Option<Arc<Semaphore>>,
+    /// A semaphore to limit the number of concurrent outbound connections.
+    concurrent_outbound_connections_semaphore: Option<Arc<Semaphore>>,
 }
 
 impl InstanceState {
@@ -195,6 +195,6 @@ pub struct AppState {
     wasi_http_clients: wasi::HttpClients,
     /// Whether connection pooling is enabled for this app.
     connection_pooling_enabled: bool,
-    /// A semaphore to limit the number of concurrent outbound requests.
-    concurrent_outbound_requests_semaphore: Option<Arc<Semaphore>>,
+    /// A semaphore to limit the number of concurrent outbound connections.
+    concurrent_outbound_connections_semaphore: Option<Arc<Semaphore>>,
 }

crates/factor-outbound-http/src/runtime_config.rs

@@ -6,15 +6,15 @@ pub mod spin;
 pub struct RuntimeConfig {
     /// If true, enable connection pooling and reuse.
     pub connection_pooling_enabled: bool,
-    /// If set, limits the number of concurrent outbound requests.
-    pub max_concurrent_requests: Option<usize>,
+    /// If set, limits the number of concurrent outbound connections.
+    pub max_concurrent_connections: Option<usize>,
 }
 
 impl Default for RuntimeConfig {
     fn default() -> Self {
         Self {
             connection_pooling_enabled: true,
-            max_concurrent_requests: None,
+            max_concurrent_connections: None,
         }
     }
 }

crates/factor-outbound-http/src/runtime_config/spin.rs

@@ -16,7 +16,7 @@ pub fn config_from_table(
         let outbound_http_toml = outbound_http.clone().try_into::<OutboundHttpToml>()?;
         Ok(Some(super::RuntimeConfig {
             connection_pooling_enabled: outbound_http_toml.connection_pooling,
-            max_concurrent_requests: outbound_http_toml.max_concurrent_requests,
+            max_concurrent_connections: outbound_http_toml.max_concurrent_requests,
         }))
     } else {
         Ok(None)

crates/factor-outbound-http/src/spin.rs

@@ -98,7 +98,7 @@ impl spin_http::Host for crate::InstanceState {
         });
 
         // If we're limiting concurrent outbound requests, acquire a permit
-        let permit = match &self.concurrent_outbound_requests_semaphore {
+        let permit = match &self.concurrent_outbound_connections_semaphore {
             Some(s) => s.acquire().await.ok(),
             None => None,
         };

crates/factor-outbound-http/src/wasi.rs

@@ -26,7 +26,7 @@ use spin_factors::{wasmtime::component::ResourceTable, RuntimeFactorsInstanceSta
 use tokio::{
     io::{AsyncRead, AsyncWrite, ReadBuf},
     net::TcpStream,
-    sync::Semaphore,
+    sync::{OwnedSemaphorePermit, Semaphore},
     time::timeout,
 };
 use tokio_rustls::client::TlsStream;
@@ -129,9 +129,9 @@ impl WasiHttpView for WasiHttpImplInner<'_> {
             self_request_origin: self.state.self_request_origin.clone(),
             blocked_networks: self.state.blocked_networks.clone(),
             http_clients: self.state.wasi_http_clients.clone(),
-            concurrent_outbound_requests_semaphore: self
+            concurrent_outbound_connections_semaphore: self
                 .state
-                .concurrent_outbound_requests_semaphore
+                .concurrent_outbound_connections_semaphore
                 .clone(),
         };
         Ok(HostFutureIncomingResponse::Pending(
@@ -158,7 +158,7 @@ struct RequestSender {
     self_request_origin: Option<SelfRequestOrigin>,
     request_interceptor: Option<Arc<dyn OutboundHttpInterceptor>>,
     http_clients: HttpClients,
-    concurrent_outbound_requests_semaphore: Option<Arc<Semaphore>>,
+    concurrent_outbound_connections_semaphore: Option<Arc<Semaphore>>,
 }
 
 impl RequestSender {
@@ -300,12 +300,18 @@ impl RequestSender {
             None
         };
 
+        // If we're limiting concurrent outbound requests, acquire a permit
+        let permit = match &self.concurrent_outbound_connections_semaphore {
+            Some(s) => s.clone().acquire_owned().await.ok().map(Arc::new),
+            None => None,
+        };
         let resp = CONNECT_OPTIONS.scope(
             ConnectOptions {
                 blocked_networks: self.blocked_networks,
                 connect_timeout,
                 tls_client_config,
                 override_connect_addr,
+                permit,
             },
             async move {
                 if use_tls {
@@ -326,17 +332,11 @@ impl RequestSender {
             },
         );
 
-        // If we're limiting concurrent outbound requests, acquire a permit
-        let permit = match &self.concurrent_outbound_requests_semaphore {
-            Some(s) => s.acquire().await.ok(),
-            None => None,
-        };
         let resp = timeout(first_byte_timeout, resp)
             .await
             .map_err(|_| ErrorCode::ConnectionReadTimeout)?
             .map_err(hyper_legacy_request_error)?
             .map(|body| body.map_err(hyper_request_error).boxed());
-        drop(permit);
 
         tracing::Span::current().record("http.response.status_code", resp.status().as_u16());
 
@@ -378,24 +378,40 @@ impl HttpClients {
     }
 }
 
-// We must use task-local variables for these config options when using
-// `hyper_util::client::legacy::Client::request` because there's no way to plumb
-// them through as parameters.  Moreover, if there's already a pooled connection
-// ready, we'll reuse that and ignore these options anyway.
 tokio::task_local! {
+    /// The options used when establishing a new connection.
+    ///
+    /// We must use task-local variables for these config options when using
+    /// `hyper_util::client::legacy::Client::request` because there's no way to plumb
+    /// them through as parameters.  Moreover, if there's already a pooled connection
+    /// ready, we'll reuse that and ignore these options anyway. After each connection
+    /// is established, the options are dropped.
     static CONNECT_OPTIONS: ConnectOptions;
 }
 
 #[derive(Clone)]
 struct ConnectOptions {
+    /// The blocked networks configuration.
     blocked_networks: BlockedNetworks,
+    /// Timeout for establishing a TCP connection.
     connect_timeout: Duration,
+    /// TLS client configuration to use, if any.
     tls_client_config: Option<TlsClientConfig>,
+    /// If set, override the address to connect to instead of using the given `uri`'s authority.
     override_connect_addr: Option<SocketAddr>,
+    /// A permit for this connection
+    ///
+    /// If there is a permit, it should be dropped when the connection is closed.
+    permit: Option<Arc<OwnedSemaphorePermit>>,
 }
 
 impl ConnectOptions {
-    async fn connect_tcp(&self, uri: &Uri, default_port: u16) -> Result<TcpStream, ErrorCode> {
+    /// Establish a TCP connection to the given URI and default port.
+    async fn connect_tcp(
+        &self,
+        uri: &Uri,
+        default_port: u16,
+    ) -> Result<PermittedTcpStream, ErrorCode> {
         let mut socket_addrs = match self.override_connect_addr {
             Some(override_connect_addr) => vec![override_connect_addr],
             None => {
@@ -430,22 +446,27 @@ impl ConnectOptions {
             return Err(ErrorCode::DestinationIpProhibited);
         }
 
-        timeout(self.connect_timeout, TcpStream::connect(&*socket_addrs))
+        let stream = timeout(self.connect_timeout, TcpStream::connect(&*socket_addrs))
             .await
             .map_err(|_| ErrorCode::ConnectionTimeout)?
             .map_err(|err| match err.kind() {
                 std::io::ErrorKind::AddrNotAvailable => {
                     dns_error("address not available".into(), 0)
                 }
                 _ => ErrorCode::ConnectionRefused,
-            })
+            })?;
+        Ok(PermittedTcpStream {
+            inner: stream,
+            _permit: self.permit.clone(),
+        })
     }
 
+    /// Establish a TLS connection to the given URI and default port.
     async fn connect_tls(
         &self,
         uri: &Uri,
         default_port: u16,
-    ) -> Result<TlsStream<TcpStream>, ErrorCode> {
+    ) -> Result<TlsStream<PermittedTcpStream>, ErrorCode> {
         let tcp_stream = self.connect_tcp(uri, default_port).await?;
 
         let mut tls_client_config = self.tls_client_config.as_deref().unwrap().clone();
@@ -465,20 +486,22 @@ impl ConnectOptions {
     }
 }
 
+/// A connector the uses `ConnectOptions`
 #[derive(Clone)]
 struct HttpConnector;
 
 impl HttpConnector {
-    async fn connect(uri: Uri) -> Result<TokioIo<TcpStream>, ErrorCode> {
+    async fn connect(uri: Uri) -> Result<TokioIo<PermittedTcpStream>, ErrorCode> {
         let stream = CONNECT_OPTIONS.get().connect_tcp(&uri, 80).await?;
         Ok(TokioIo::new(stream))
     }
 }
 
 impl Service<Uri> for HttpConnector {
-    type Response = TokioIo<TcpStream>;
+    type Response = TokioIo<PermittedTcpStream>;
     type Error = ErrorCode;
-    type Future = Pin<Box<dyn Future<Output = Result<TokioIo<TcpStream>, ErrorCode>> + Send>>;
+    type Future =
+        Pin<Box<dyn Future<Output = Result<TokioIo<PermittedTcpStream>, ErrorCode>> + Send>>;
 
     fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         Poll::Ready(Ok(()))
@@ -489,6 +512,7 @@ impl Service<Uri> for HttpConnector {
     }
 }
 
+/// A connector that establishes TLS connections using `rustls` and `ConnectOptions`.
 #[derive(Clone)]
 struct HttpsConnector;
 
@@ -513,7 +537,7 @@ impl Service<Uri> for HttpsConnector {
     }
 }
 
-struct RustlsStream(TlsStream<TcpStream>);
+struct RustlsStream(TlsStream<PermittedTcpStream>);
 
 impl Connection for RustlsStream {
     fn connected(&self) -> Connected {
@@ -568,6 +592,55 @@ impl AsyncWrite for RustlsStream {
     }
 }
 
+/// A TCP stream that holds an optional permit indicating that it is allowed to exist.
+struct PermittedTcpStream {
+    inner: TcpStream,
+    _permit: Option<Arc<OwnedSemaphorePermit>>,
+}
+
+impl PermittedTcpStream {
+    fn connected(&self) -> Connected {
+        self.inner.connected()
+    }
+}
+
+impl Connection for PermittedTcpStream {
+    fn connected(&self) -> Connected {
+        self.inner.connected()
+    }
+}
+
+impl AsyncRead for PermittedTcpStream {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
+    }
+}
+
+impl AsyncWrite for PermittedTcpStream {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
+        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
+    }
+}
+
 /// Translate a [`hyper::Error`] to a wasi-http `ErrorCode` in the context of a request.
 fn hyper_request_error(err: hyper::Error) -> ErrorCode {
     // If there's a source, we might be able to extract a wasi-http error from it.