[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.
- It also fixes #4456 by making sure:
  - Request stream is closed immediately after the we receive a terminal state
  - Drain the response stream. This also fixes a connection thrashing issue

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/invoker-impl/src/invocation_task/service_protocol_runner_v4.rs

@@ -147,8 +147,6 @@ where
         .try_into()
         .expect("must be able to build a valid invocation path");
 
-        let journal_size = journal_metadata.length;
-
         debug!(
             restate.invocation.id = %self.invocation_task.invocation_id,
             deployment.address = %deployment.address_display(),
@@ -159,15 +157,14 @@ where
 
         // Create an arc of the parent SpanContext.
         // We send this with every journal entry to correctly link new spans generated from journal entries.
-        let service_invocation_span_context = journal_metadata.span_context;
 
         // Prepare the request
-        let (mut http_stream_tx, request) = Self::prepare_request(
+        let (http_stream_tx, request) = Self::prepare_request(
             path,
             deployment,
             self.service_protocol_version,
             &self.invocation_task.invocation_id,
-            &service_invocation_span_context,
+            &journal_metadata.span_context,
         );
 
         // Initialize the response stream state
@@ -183,6 +180,49 @@ where
             .throttle(self.invocation_task.action_token_bucket.take())
         );
 
+        let result = self
+            .run_inner(
+                txn,
+                protocol_type,
+                journal_metadata,
+                keyed_service_id,
+                cached_journal_items,
+                http_stream_tx,
+                &mut decoder_stream,
+            )
+            .await;
+        // Sanity check of the stream decoder
+        if decoder_stream.inner().has_remaining() {
+            warn_it!(
+                InvokerError::WriteAfterEndOfStream,
+                "The read buffer is non empty after the stream has been closed."
+            );
+        }
+
+        let inner_stream = &mut decoder_stream.inner_pin_mut().inner;
+
+        while inner_stream.next().await.is_some() {}
+
+        result
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    async fn run_inner<Txn, S>(
+        &mut self,
+        txn: Txn,
+        protocol_type: ProtocolType,
+        journal_metadata: JournalMetadata,
+        keyed_service_id: Option<ServiceId>,
+        cached_journal_items: Option<Vec<JournalEntry>>,
+        mut http_stream_tx: mpsc::Sender<Result<Frame<Bytes>, Infallible>>,
+        decoder_stream: &mut S,
+    ) -> TerminalLoopState<()>
+    where
+        Txn: InvocationReaderTransaction,
+        S: Stream<Item = Result<DecoderStreamItem, InvokerError>> + Unpin,
+    {
+        let journal_size = journal_metadata.length;
+        let service_invocation_span_context = journal_metadata.span_context;
         // === Replay phase (transaction alive) ===
         {
             // Read state if needed (state is collected for the START message)
@@ -215,7 +255,7 @@ where
                 crate::shortcircuit!(
                     self.replay_loop(
                         &mut http_stream_tx,
-                        &mut decoder_stream,
+                        decoder_stream,
                         journal_stream,
                         journal_metadata.length
                     )
@@ -234,7 +274,7 @@ where
                 crate::shortcircuit!(
                     self.replay_loop(
                         &mut http_stream_tx,
-                        &mut decoder_stream,
+                        decoder_stream,
                         journal_stream,
                         journal_metadata.length
                     )
@@ -255,32 +295,24 @@ where
                 self.bidi_stream_loop(
                     &service_invocation_span_context,
                     http_stream_tx,
-                    &mut decoder_stream
+                    decoder_stream
                 )
                 .await
             );
         } else {
-            trace!("Protocol is in bidi stream mode, will now drop the sender side of the request");
+            trace!(
+                "Protocol is not in bidi stream mode, will now drop the sender side of the request"
+            );
             // Drop the http_stream_tx.
             // This is required in HTTP/1.1 to let the deployment send the headers back
             drop(http_stream_tx)
         }
 
         // We don't have the invoker_rx, so we simply consume the response
         trace!("Sender side of the request has been dropped, now processing the response");
-        let result = self
-            .response_stream_loop(&service_invocation_span_context, &mut decoder_stream)
-            .await;
-
-        // Sanity check of the stream decoder
-        if decoder_stream.inner().has_remaining() {
-            warn_it!(
-                InvokerError::WriteAfterEndOfStream,
-                "The read buffer is non empty after the stream has been closed."
-            );
-        }
 
-        result
+        self.response_stream_loop(&service_invocation_span_context, decoder_stream)
+            .await
     }
 
     fn prepare_request(

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,
         }
     }