[Service-Client] H2 Pool integation and usage in service client

- This PR finally enable using of the new H2 pool in the http service client.

Author: muhamadazmy

crates/invoker-impl/src/invocation_task/mod.rs

@@ -486,6 +486,8 @@ impl ResponseStream {
         // This task::spawn won't be required by hyper 1.0, as the connection will be driven by a task
         // spawned somewhere else (perhaps in the connection pool).
         // See: https://github.com/restatedev/restate/issues/96 and https://github.com/restatedev/restate/issues/76
+
+        //todo: this is a temp clone to test
         Self::WaitingHeaders {
             join_handle: AbortOnDropHandle::new(tokio::task::spawn(client.call(req))),
         }

crates/service-client/src/http.rs

@@ -10,14 +10,17 @@
 
 use super::proxy::ProxyConnector;
 
+use crate::pool::conn::PermittedRecvStream;
+use crate::pool::tls::TlsConnector;
+use crate::pool::{self, Pool, TcpConnector};
 use crate::utils::ErrorExt;
 
 use bytes::Bytes;
 use futures::FutureExt;
-use futures::future::Either;
+use futures::future::{self, Either};
 use http::Version;
-use http_body_util::BodyExt;
-use hyper::body::Body;
+use http_body_util::{BodyExt, Either as EitherBody};
+use hyper::body::{Body, Incoming};
 use hyper::http::HeaderValue;
 use hyper::http::uri::PathAndQuery;
 use hyper::{HeaderMap, Method, Request, Response, Uri};
@@ -26,10 +29,13 @@ use hyper_util::client::legacy::connect::HttpConnector;
 use restate_types::config::HttpOptions;
 use rustls::{ClientConfig, KeyLogFile};
 use std::error::Error;
+use std::fmt;
 use std::fmt::Debug;
-use std::future::Future;
+use std::num::NonZeroU32;
+use std::pin::Pin;
 use std::sync::{Arc, LazyLock};
-use std::{fmt, future};
+use std::task::{Context, Poll, ready};
+use tower::Layer;
 
 type ProxiedHttpsConnector = ProxyConnector<HttpsConnector<HttpConnector>>;
 
@@ -55,7 +61,7 @@ static TLS_CLIENT_CONFIG: LazyLock<ClientConfig> = LazyLock::new(|| {
 type BoxError = Box<dyn Error + Send + Sync + 'static>;
 type BoxBody = http_body_util::combinators::BoxBody<Bytes, BoxError>;
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct HttpClient {
     /// Client used for HTTPS as long as HTTP1.1 or HTTP2 was not specifically requested.
     /// All HTTP versions are possible.
@@ -68,7 +74,7 @@ pub struct HttpClient {
     /// Client when HTTP2 was specifically requested - for cleartext, we use h2c,
     /// and for HTTPS, we will fail unless the ALPN supports h2.
     /// In practice, at discovery time we never force h2 for HTTPS.
-    h2_client: hyper_util::client::legacy::Client<ProxiedHttpsConnector, BoxBody>,
+    h2_pool: Pool<ProxyConnector<TlsConnector<TcpConnector>>>,
 }
 
 impl HttpClient {
@@ -78,7 +84,7 @@ impl HttpClient {
         builder.timer(hyper_util::rt::TokioTimer::default());
 
         builder
-            .http2_initial_max_send_streams(options.initial_max_send_streams)
+            .http2_initial_max_send_streams(options.initial_max_send_streams.map(|v| v as usize))
             .http2_adaptive_window(true)
             .http2_keep_alive_timeout(options.http_keep_alive_options.timeout.into())
             .http2_keep_alive_interval(Some(options.http_keep_alive_options.interval.into()));
@@ -101,11 +107,25 @@ impl HttpClient {
             .enable_http1()
             .wrap_connector(http_connector.clone());
 
-        let https_h2_connector = hyper_rustls::HttpsConnectorBuilder::new()
-            .with_tls_config(TLS_CLIENT_CONFIG.clone())
-            .https_or_http()
-            .enable_http2()
-            .wrap_connector(http_connector.clone());
+        let h2_pool = {
+            let connector = pool::tls::TlsConnectorLayer::new(TLS_CLIENT_CONFIG.clone())
+                .layer(pool::TcpConnector::new(options.connect_timeout.into()));
+            let connector = ProxyConnector::new(
+                options.http_proxy.clone(),
+                options.no_proxy.clone(),
+                connector,
+            );
+
+            let builder =
+                pool::PoolBuilder::default().max_connections(options.max_http2_connections);
+
+            let builder = match options.initial_max_send_streams.and_then(NonZeroU32::new) {
+                Some(value) => builder.init_max_send_streams(value),
+                None => builder,
+            };
+
+            builder.build(connector)
+        };
 
         HttpClient {
             alpn_client: builder.clone().build::<_, BoxBody>(ProxyConnector::new(
@@ -118,14 +138,7 @@ impl HttpClient {
                 options.no_proxy.clone(),
                 https_h1_connector,
             )),
-            h2_client: {
-                builder.http2_only(true);
-                builder.build::<_, BoxBody>(ProxyConnector::new(
-                    options.http_proxy.clone(),
-                    options.no_proxy.clone(),
-                    https_h2_connector,
-                ))
-            },
+            h2_pool,
         }
     }
 
@@ -186,10 +199,10 @@ impl HttpClient {
         body: B,
         path: PathAndQuery,
         headers: HeaderMap<HeaderValue>,
-    ) -> impl Future<Output = Result<Response<hyper::body::Incoming>, HttpError>> + Send + 'static
+    ) -> impl Future<Output = Result<Response<ResponseBody>, HttpError>> + Send + 'static
     where
         B: Body<Data = Bytes> + Send + Sync + Unpin + Sized + 'static,
-        <B as Body>::Error: Error + Send + Sync + 'static,
+        B::Error: std::error::Error + Send + Sync + 'static,
     {
         let request = match Self::build_request(uri, version, body, method, path, headers) {
             Ok(request) => request,
@@ -198,21 +211,98 @@ impl HttpClient {
 
         let fut = match version {
             // version is set to http1.1 when use_http1.1 is set
-            Some(Version::HTTP_11) => self.h1_client.request(request),
+            Some(Version::HTTP_11) => ResponseMapper {
+                fut: self.h1_client.request(request),
+            }
+            .left_future(),
             // version is set to http2 for cleartext urls when use_http1.1 is not set
-            Some(Version::HTTP_2) => self.h2_client.request(request),
+            Some(Version::HTTP_2) => ResponseMapper {
+                fut: self.h2_pool.request(request),
+            }
+            .right_future(),
             // version is currently set to none for https urls when use_http1.1 is not set
-            None => self.alpn_client.request(request),
+            None => ResponseMapper {
+                fut: self.alpn_client.request(request),
+            }
+            .left_future(),
             // nothing currently sets a different version, but the alpn client is a sensible default
-            Some(_) => self.alpn_client.request(request),
+            Some(_) => ResponseMapper {
+                fut: self.alpn_client.request(request),
+            }
+            .left_future(),
         };
 
-        Either::Left(async move {
-            match fut.await {
-                Ok(res) => Ok(res),
-                Err(err) => Err(err.into()),
-            }
-        })
+        Either::Left(fut)
+    }
+}
+
+#[pin_project::pin_project]
+struct ResponseMapper<F, B, E>
+where
+    F: Future<Output = Result<Response<B>, E>>,
+    E: Into<HttpError>,
+    B: Into<ResponseBody>,
+{
+    #[pin]
+    fut: F,
+}
+
+impl<F, B, E> Future for ResponseMapper<F, B, E>
+where
+    F: Future<Output = Result<Response<B>, E>>,
+    E: Into<HttpError>,
+    B: Into<ResponseBody>,
+{
+    type Output = Result<Response<ResponseBody>, HttpError>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let result = ready!(self.project().fut.poll(cx))
+            .map_err(Into::into)
+            .map(|response| response.map(Into::into));
+
+        Poll::Ready(result)
+    }
+}
+
+/// A wrapper around [`http_body_util::Either`] to hide
+/// type complexity for higher layer
+#[pin_project::pin_project]
+pub struct ResponseBody {
+    #[pin]
+    inner: EitherBody<Incoming, PermittedRecvStream>,
+}
+
+impl From<Incoming> for ResponseBody {
+    fn from(value: Incoming) -> Self {
+        Self {
+            inner: EitherBody::Left(value),
+        }
+    }
+}
+
+impl From<PermittedRecvStream> for ResponseBody {
+    fn from(value: PermittedRecvStream) -> Self {
+        Self {
+            inner: EitherBody::Right(value),
+        }
+    }
+}
+
+impl Body for ResponseBody {
+    type Data = Bytes;
+    type Error = Box<dyn std::error::Error + Send + Sync>;
+
+    fn is_end_stream(&self) -> bool {
+        self.inner.is_end_stream()
+    }
+    fn poll_frame(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
+        self.project().inner.poll_frame(cx)
+    }
+    fn size_hint(&self) -> http_body::SizeHint {
+        self.inner.size_hint()
     }
 }
 
@@ -228,6 +318,8 @@ pub enum HttpError {
     Connect(#[source] hyper_util::client::legacy::Error),
     #[error("{}", FormatHyperError(.0))]
     Hyper(#[source] hyper_util::client::legacy::Error),
+    #[error("h2 pool connection error: {0}")]
+    PoolError(#[from] pool::ConnectionError),
 }
 
 impl HttpError {
@@ -240,6 +332,7 @@ impl HttpError {
             HttpError::PossibleHTTP11Only(_) => false,
             HttpError::PossibleHTTP2Only(_) => false,
             HttpError::Connect(_) => true,
+            HttpError::PoolError(_) => true,
         }
     }
 

crates/service-client/src/lib.rs

@@ -19,7 +19,7 @@ use arc_swap::ArcSwapOption;
 use bytes::Bytes;
 use bytestring::ByteString;
 use core::fmt;
-use futures::FutureExt;
+use futures::{FutureExt, future};
 use http_body_util::Full;
 use hyper::body::Body;
 use hyper::http::uri::PathAndQuery;
@@ -30,8 +30,6 @@ use restate_types::schema::deployment::EndpointLambdaCompression;
 use std::collections::HashMap;
 use std::error::Error;
 use std::fmt::Formatter;
-use std::future;
-use std::future::Future;
 use std::sync::Arc;
 
 mod http;
@@ -41,9 +39,9 @@ mod proxy;
 mod request_identity;
 mod utils;
 
-pub type ResponseBody = http_body_util::Either<hyper::body::Incoming, Full<Bytes>>;
+pub type ResponseBody = http_body_util::Either<http::ResponseBody, Full<Bytes>>;
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct ServiceClient {
     // TODO a single client uses the pooling provided by hyper, but this is not enough.
     //  See https://github.com/restatedev/restate/issues/76 for more background on the topic.

crates/service-protocol/src/discovery.rs

@@ -173,7 +173,7 @@ impl DiscoveryError {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct ServiceDiscovery {
     retry_policy: RetryPolicy,
     client: ServiceClient,

crates/types/src/config/http.rs

@@ -8,6 +8,8 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::num::NonZeroUsize;
+
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
 
@@ -59,7 +61,15 @@ pub struct HttpOptions {
     ///
     /// **NOTE**: Setting this value to None (default) users the default
     /// recommended value from HTTP2 specs
-    pub initial_max_send_streams: Option<usize>,
+    pub initial_max_send_streams: Option<u32>,
+
+    /// # Max HTTP2 Connections
+    ///
+    /// Sets the maximum number of open HTTP/2 connections per
+    /// client for a single host.
+    ///
+    /// Default: 20
+    pub max_http2_connections: NonZeroUsize,
 }
 
 impl Default for HttpOptions {
@@ -70,6 +80,7 @@ impl Default for HttpOptions {
             no_proxy: None,
             connect_timeout: NonZeroFriendlyDuration::from_secs_unchecked(10),
             initial_max_send_streams: None,
+            max_http2_connections: NonZeroUsize::new(20).unwrap(),
         }
     }
 }