From 0218d58c282d6de6f300229677c99369d3ea20ed Mon Sep 17 00:00:00 2001
From: John Doneth <doneth7@gmail.com>
Date: Fri, 4 Oct 2019 10:08:37 -0400
Subject: [PATCH] feat(transport): Expose http/2 settings (#28)

* expose http2 settings

* add client http2 options
---
 tonic/src/transport/endpoint.rs           | 23 +++++++++++++
 tonic/src/transport/server.rs             | 39 +++++++++++++++++++++++
 tonic/src/transport/service/connection.rs |  6 +++-
 3 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/tonic/src/transport/endpoint.rs b/tonic/src/transport/endpoint.rs
index abd0e1e68..ab11b9922 100644
--- a/tonic/src/transport/endpoint.rs
+++ b/tonic/src/transport/endpoint.rs
@@ -24,6 +24,8 @@ pub struct Endpoint {
     pub(super) buffer_size: Option<usize>,
     pub(super) interceptor_headers:
         Option<Arc<dyn Fn(&mut http::HeaderMap) + Send + Sync + 'static>>,
+    pub(super) init_stream_window_size: Option<u32>,
+    pub(super) init_connection_window_size: Option<u32>,
 }
 
 impl Endpoint {
@@ -101,6 +103,25 @@ impl Endpoint {
         self
     }
 
+    /// Sets the [`SETTINGS_INITIAL_WINDOW_SIZE`][spec] option for HTTP2
+    /// stream-level flow control.
+    ///
+    /// Default is 65,535
+    ///
+    /// [spec]: https://http2.github.io/http2-spec/#SETTINGS_INITIAL_WINDOW_SIZE
+    pub fn initial_stream_window_size(&mut self, sz: impl Into<Option<u32>>) -> &mut Self {
+        self.init_stream_window_size = sz.into();
+        self
+    }
+
+    /// Sets the max connection-level flow control for HTTP2
+    ///
+    /// Default is 65,535
+    pub fn initial_connection_window_size(&mut self, sz: impl Into<Option<u32>>) -> &mut Self {
+        self.init_connection_window_size = sz.into();
+        self
+    }
+
     /// Enable TLS and apply the CA as the root certificate.
     ///
     /// Providing an optional domain to override. If `None` is passed to this
@@ -185,6 +206,8 @@ impl From<Uri> for Endpoint {
             tls: None,
             buffer_size: None,
             interceptor_headers: None,
+            init_stream_window_size: None,
+            init_connection_window_size: None,
         }
     }
 }
diff --git a/tonic/src/transport/server.rs b/tonic/src/transport/server.rs
index b746631e3..34cd32711 100644
--- a/tonic/src/transport/server.rs
+++ b/tonic/src/transport/server.rs
@@ -47,6 +47,9 @@ pub struct Server {
     // timeout: Option<Duration>,
     #[cfg(feature = "tls")]
     tls: Option<TlsAcceptor>,
+    init_stream_window_size: Option<u32>,
+    init_connection_window_size: Option<u32>,
+    max_concurrent_streams: Option<u32>,
 }
 
 impl Server {
@@ -123,6 +126,36 @@ impl Server {
     //     self
     // }
 
+    /// Sets the [`SETTINGS_INITIAL_WINDOW_SIZE`][spec] option for HTTP2
+    /// stream-level flow control.
+    ///
+    /// Default is 65,535
+    ///
+    /// [spec]: https://http2.github.io/http2-spec/#SETTINGS_INITIAL_WINDOW_SIZE
+    pub fn initial_stream_window_size(&mut self, sz: impl Into<Option<u32>>) -> &mut Self {
+        self.init_stream_window_size = sz.into();
+        self
+    }
+
+    /// Sets the max connection-level flow control for HTTP2
+    ///
+    /// Default is 65,535
+    pub fn initial_connection_window_size(&mut self, sz: impl Into<Option<u32>>) -> &mut Self {
+        self.init_connection_window_size = sz.into();
+        self
+    }
+
+    /// Sets the [`SETTINGS_MAX_CONCURRENT_STREAMS`][spec] option for HTTP2
+    /// connections.
+    ///
+    /// Default is no limit (`None`).
+    ///
+    /// [spec]: https://http2.github.io/http2-spec/#SETTINGS_MAX_CONCURRENT_STREAMS
+    pub fn max_concurrent_streams(&mut self, max: impl Into<Option<u32>>) -> &mut Self {
+        self.max_concurrent_streams = max.into();
+        self
+    }
+
     /// Intercept the execution of gRPC methods.
     ///
     /// ```
@@ -162,6 +195,9 @@ impl Server {
     {
         let interceptor = self.interceptor.clone();
         let concurrency_limit = self.concurrency_limit;
+        let init_connection_window_size = self.init_connection_window_size;
+        let init_stream_window_size = self.init_stream_window_size;
+        let max_concurrent_streams = self.max_concurrent_streams;
         // let timeout = self.timeout.clone();
 
         let incoming = hyper::server::accept::from_stream(async_stream::try_stream! {
@@ -190,6 +226,9 @@ impl Server {
 
         hyper::Server::builder(incoming)
             .http2_only(true)
+            .http2_initial_connection_window_size(init_connection_window_size)
+            .http2_initial_stream_window_size(init_stream_window_size)
+            .http2_max_concurrent_streams(max_concurrent_streams)
             .serve(svc)
             .await
             .map_err(map_err)?;
diff --git a/tonic/src/transport/service/connection.rs b/tonic/src/transport/service/connection.rs
index 74bf62eac..c30888839 100644
--- a/tonic/src/transport/service/connection.rs
+++ b/tonic/src/transport/service/connection.rs
@@ -34,7 +34,11 @@ impl Connection {
         #[cfg(not(feature = "tls"))]
         let connector = connector();
 
-        let settings = Builder::new().http2_only(true).clone();
+        let settings = Builder::new()
+            .http2_initial_stream_window_size(endpoint.init_stream_window_size)
+            .http2_initial_connection_window_size(endpoint.init_connection_window_size)
+            .http2_only(true)
+            .clone();
 
         let stack = ServiceBuilder::new()
             .layer_fn(|s| AddOrigin::new(s, endpoint.uri.clone()))