feat: h1 connection pooling

This is imported from the orogene project, with massive thanks to Kat Marchán (@zkat) for letting us re-use this work.

Specifically:
https://github.com/orogene/orogene/tree/82b5c1e6773ceb3ee5a06c8273d7262d6073ce43/crates/oro-client/src/http_client

Co-Authored-By: Kat Marchán <[email protected]>

Author: Fishrock123

Cargo.toml

@@ -22,8 +22,8 @@ rustdoc-args = ["--cfg", "feature=\"docs\""]
 [features]
 default = ["h1_client"]
 docs = ["h1_client", "curl_client", "wasm_client", "hyper_client"]
-h1_client = ["async-h1", "async-std", "async-native-tls"]
-h1_client_rustls = ["async-h1", "async-std", "async-tls"]
+h1_client = ["async-h1", "async-std", "async-native-tls", "deadpool", "futures"]
+h1_client_rustls = ["async-h1", "async-std", "async-tls", "deadpool", "futures"]
 native_client = ["curl_client", "wasm_client"]
 curl_client = ["isahc", "async-std"]
 wasm_client = ["js-sys", "web-sys", "wasm-bindgen", "wasm-bindgen-futures", "futures"]
@@ -38,6 +38,8 @@ log = "0.4.7"
 async-h1 = { version = "2.0.0", optional = true }
 async-std = { version = "1.6.0", default-features = false, optional = true }
 async-native-tls = { version = "0.3.1", optional = true }
+deadpool = { version = "0.6.0", optional = true }
+futures = { version = "0.3.8", optional = true }
 
 # h1_client_rustls
 async-tls = { version = "0.10.0", optional = true }

src/h1.rs renamed to src/h1/mod.rs

@@ -1,14 +1,45 @@
-//! http-client implementation for async-h1.
+//! http-client implementation for async-h1, with connecton pooling ("Keep-Alive").
 
-use super::{async_trait, Error, HttpClient, Request, Response};
+use std::collections::HashMap;
+use std::net::SocketAddr;
+use std::{fmt::Debug, sync::Arc};
 
 use async_h1::client;
+use async_std::net::TcpStream;
+use async_std::sync::Mutex;
+use deadpool::managed::Pool;
 use http_types::StatusCode;
 
-/// Async-h1 based HTTP Client.
-#[derive(Debug)]
+#[cfg(not(feature = "h1_client_rustls"))]
+use async_native_tls::TlsStream;
+#[cfg(feature = "h1_client_rustls")]
+use async_tls::client::TlsStream;
+
+use super::{async_trait, Error, HttpClient, Request, Response};
+
+mod tcp;
+mod tls;
+
+use tcp::{TcpConnWrapper, TcpConnection};
+use tls::{TlsConnWrapper, TlsConnection};
+
+// TODO: Move this to a parameter. This current number is based on a few
+// random benchmarks and see whatever gave decent perf vs resource use.
+static MAX_CONCURRENT_CONNECTIONS: usize = 50;
+
+type HttpPool = HashMap<SocketAddr, Pool<TcpStream, std::io::Error>>;
+type HttpsPool = HashMap<SocketAddr, Pool<TlsStream<TcpStream>, Error>>;
+
+/// Async-h1 based HTTP Client, with connecton pooling ("Keep-Alive").
 pub struct H1Client {
-    _priv: (),
+    http_pool: Arc<Mutex<HttpPool>>,
+    https_pool: Arc<Mutex<HttpsPool>>,
+}
+
+impl Debug for H1Client {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("H1Client")
+    }
 }
 
 impl Default for H1Client {
@@ -20,13 +51,20 @@ impl Default for H1Client {
 impl H1Client {
     /// Create a new instance.
     pub fn new() -> Self {
-        Self { _priv: () }
+        Self {
+            http_pool: Arc::new(Mutex::new(HashMap::new())),
+            https_pool: Arc::new(Mutex::new(HashMap::new())),
+        }
     }
 }
 
 #[async_trait]
 impl HttpClient for H1Client {
     async fn send(&self, mut req: Request) -> Result<Response, Error> {
+        let http_pool = self.http_pool.clone();
+        let https_pool = self.https_pool.clone();
+        req.insert_header("Connection", "keep-alive");
+
         // Insert host
         let host = req
             .url()
@@ -57,40 +95,62 @@ impl HttpClient for H1Client {
 
         match scheme {
             "http" => {
-                let stream = async_std::net::TcpStream::connect(addr).await?;
+                let mut hash = http_pool.lock().await;
+                let pool = if let Some(pool) = hash.get(&addr) {
+                    pool
+                } else {
+                    let manager = TcpConnection::new(addr);
+                    let pool =
+                        Pool::<TcpStream, std::io::Error>::new(manager, MAX_CONCURRENT_CONNECTIONS);
+                    hash.insert(addr, pool);
+                    hash.get(&addr).unwrap()
+                };
+                let pool = pool.clone();
+                std::mem::drop(hash);
+                let stream = pool.get().await?;
                 req.set_peer_addr(stream.peer_addr().ok());
                 req.set_local_addr(stream.local_addr().ok());
-                client::connect(stream, req).await
+                client::connect(TcpConnWrapper::new(stream), req).await
+
+                // let stream = async_std::net::TcpStream::connect(addr).await?;
+                // req.set_peer_addr(stream.peer_addr().ok());
+                // req.set_local_addr(stream.local_addr().ok());
+                // client::connect(stream, req).await
             }
             "https" => {
-                let raw_stream = async_std::net::TcpStream::connect(addr).await?;
-                req.set_peer_addr(raw_stream.peer_addr().ok());
-                req.set_local_addr(raw_stream.local_addr().ok());
-                let tls_stream = add_tls(host, raw_stream).await?;
-                client::connect(tls_stream, req).await
+                let mut hash = https_pool.lock().await;
+                let pool = if let Some(pool) = hash.get(&addr) {
+                    pool
+                } else {
+                    let manager = TlsConnection::new(host.clone(), addr);
+                    let pool = Pool::<TlsStream<TcpStream>, Error>::new(
+                        manager,
+                        MAX_CONCURRENT_CONNECTIONS,
+                    );
+                    hash.insert(addr, pool);
+                    hash.get(&addr).unwrap()
+                };
+                let pool = pool.clone();
+                std::mem::drop(hash);
+                let stream = pool.get().await.unwrap(); // TODO: remove unwrap
+                req.set_peer_addr(stream.get_ref().peer_addr().ok());
+                req.set_local_addr(stream.get_ref().local_addr().ok());
+
+                client::connect(TlsConnWrapper::new(stream), req).await
+
+                // let raw_stream = async_std::net::TcpStream::connect(addr).await?;
+                // req.set_peer_addr(raw_stream.peer_addr().ok());
+                // req.set_local_addr(raw_stream.local_addr().ok());
+
+                // let stream = async_native_tls::connect(host, raw_stream).await?;
+
+                // client::connect(stream, req).await
             }
             _ => unreachable!(),
         }
     }
 }
 
-#[cfg(not(feature = "h1_client_rustls"))]
-async fn add_tls(
-    host: String,
-    stream: async_std::net::TcpStream,
-) -> Result<async_native_tls::TlsStream<async_std::net::TcpStream>, async_native_tls::Error> {
-    async_native_tls::connect(host, stream).await
-}
-
-#[cfg(feature = "h1_client_rustls")]
-async fn add_tls(
-    host: String,
-    stream: async_std::net::TcpStream,
-) -> std::io::Result<async_tls::client::TlsStream<async_std::net::TcpStream>> {
-    let connector = async_tls::TlsConnector::default();
-    connector.connect(host, stream).await
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

src/h1/tcp.rs

@@ -0,0 +1,68 @@
+use std::fmt::Debug;
+use std::net::SocketAddr;
+use std::pin::Pin;
+
+use async_std::net::TcpStream;
+use async_trait::async_trait;
+use deadpool::managed::{Manager, Object, RecycleResult};
+use futures::io::{AsyncRead, AsyncWrite};
+use futures::task::{Context, Poll};
+
+#[derive(Clone, Debug)]
+pub(crate) struct TcpConnection {
+    addr: SocketAddr,
+}
+impl TcpConnection {
+    pub(crate) fn new(addr: SocketAddr) -> Self {
+        Self { addr }
+    }
+}
+
+pub(crate) struct TcpConnWrapper {
+    conn: Object<TcpStream, std::io::Error>,
+}
+impl TcpConnWrapper {
+    pub(crate) fn new(conn: Object<TcpStream, std::io::Error>) -> Self {
+        Self { conn }
+    }
+}
+
+impl AsyncRead for TcpConnWrapper {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        Pin::new(&mut *self.conn).poll_read(cx, buf)
+    }
+}
+
+impl AsyncWrite for TcpConnWrapper {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        let amt = futures::ready!(Pin::new(&mut *self.conn).poll_write(cx, buf))?;
+        Poll::Ready(Ok(amt))
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut *self.conn).poll_flush(cx)
+    }
+
+    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut *self.conn).poll_close(cx)
+    }
+}
+
+#[async_trait]
+impl Manager<TcpStream, std::io::Error> for TcpConnection {
+    async fn create(&self) -> Result<TcpStream, std::io::Error> {
+        Ok(TcpStream::connect(self.addr).await?)
+    }
+
+    async fn recycle(&self, _conn: &mut TcpStream) -> RecycleResult<std::io::Error> {
+        Ok(())
+    }
+}

src/h1/tls.rs

@@ -0,0 +1,93 @@
+use std::fmt::Debug;
+use std::net::SocketAddr;
+use std::pin::Pin;
+
+use async_std::net::TcpStream;
+use async_trait::async_trait;
+use deadpool::managed::{Manager, Object, RecycleResult};
+use futures::io::{AsyncRead, AsyncWrite};
+use futures::task::{Context, Poll};
+
+#[cfg(not(feature = "h1_client_rustls"))]
+use async_native_tls::TlsStream;
+#[cfg(feature = "h1_client_rustls")]
+use async_tls::client::TlsStream;
+
+use crate::Error;
+
+#[derive(Clone, Debug)]
+pub(crate) struct TlsConnection {
+    host: String,
+    addr: SocketAddr,
+}
+impl TlsConnection {
+    pub(crate) fn new(host: String, addr: SocketAddr) -> Self {
+        Self { host, addr }
+    }
+}
+
+pub(crate) struct TlsConnWrapper {
+    conn: Object<TlsStream<TcpStream>, Error>,
+}
+impl TlsConnWrapper {
+    pub(crate) fn new(conn: Object<TlsStream<TcpStream>, Error>) -> Self {
+        Self { conn }
+    }
+}
+
+impl AsyncRead for TlsConnWrapper {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        Pin::new(&mut *self.conn).poll_read(cx, buf)
+    }
+}
+
+impl AsyncWrite for TlsConnWrapper {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        let amt = futures::ready!(Pin::new(&mut *self.conn).poll_write(cx, buf))?;
+        Poll::Ready(Ok(amt))
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut *self.conn).poll_flush(cx)
+    }
+
+    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Pin::new(&mut *self.conn).poll_close(cx)
+    }
+}
+
+#[async_trait]
+impl Manager<TlsStream<TcpStream>, Error> for TlsConnection {
+    async fn create(&self) -> Result<TlsStream<TcpStream>, Error> {
+        log::trace!("Creating new socket to {:?}", self.addr);
+        let raw_stream = async_std::net::TcpStream::connect(self.addr).await?;
+        let tls_stream = add_tls(&self.host, raw_stream).await?;
+        Ok(tls_stream)
+    }
+
+    async fn recycle(&self, _conn: &mut TlsStream<TcpStream>) -> RecycleResult<Error> {
+        Ok(())
+    }
+}
+
+#[cfg(not(feature = "h1_client_rustls"))]
+async fn add_tls(
+    host: &str,
+    stream: TcpStream,
+) -> Result<async_native_tls::TlsStream<TcpStream>, async_native_tls::Error> {
+    async_native_tls::connect(host, stream).await
+}
+
+#[cfg(feature = "h1_client_rustls")]
+async fn add_tls(host: &str, stream: TcpStream) -> Result<TlsStream<TcpStream>, std::io::Error> {
+    let connector = async_tls::TlsConnector::default();
+    connector.connect(host, stream).await
+}