From 4000c42478414aa9addfd4f0df623088ddd7312d Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Sat, 31 Aug 2019 03:18:52 -0400
Subject: [PATCH] Add buffer support

---
 tonic-interop/Cargo.toml        |  2 ++
 tonic-interop/src/bin/client.rs |  3 ++-
 tonic/Cargo.toml                |  3 ++-
 tonic/src/transport/client.rs   | 36 ++++++++++++++++++---------------
 4 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/tonic-interop/Cargo.toml b/tonic-interop/Cargo.toml
index 6f9af94..e0fd778 100644
--- a/tonic-interop/Cargo.toml
+++ b/tonic-interop/Cargo.toml
@@ -25,6 +25,8 @@ hyper = { git = "https://github.com/hyperium/hyper" }
 console = "0.7"
 structopt = "0.2"
 pretty_env_logger = "0.3"
+tracing-fmt = "0.0.1-alpha.1"
+tracing = "0.1"
 
 [build-dependencies]
 tonic-build = { path = "../tonic-build" }
diff --git a/tonic-interop/src/bin/client.rs b/tonic-interop/src/bin/client.rs
index e812923..2854052 100644
--- a/tonic-interop/src/bin/client.rs
+++ b/tonic-interop/src/bin/client.rs
@@ -14,7 +14,8 @@ struct Opts {
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    pretty_env_logger::init();
+    let sub = tracing_fmt::FmtSubscriber::builder().finish();
+    tracing::subscriber::set_global_default(sub).unwrap();
 
     let matches = Opts::from_args();
 
diff --git a/tonic/Cargo.toml b/tonic/Cargo.toml
index 4baf9a3..bb21d4b 100644
--- a/tonic/Cargo.toml
+++ b/tonic/Cargo.toml
@@ -26,7 +26,8 @@ hyper = { git = "https://github.com/hyperium/hyper", optional = true}
 
 # tower
 tower-reconnect = { path = "../../tower/tower-reconnect", optional = true }
+tower-buffer = { path = "../../tower/tower-buffer", optional = true }
 
 [features]
 default = ["transport"]
-transport = ["hyper", "tower-reconnect"]
+transport = ["hyper", "tower-reconnect", "tower-buffer"]
diff --git a/tonic/src/transport/client.rs b/tonic/src/transport/client.rs
index 02f2b87..dd31aa1 100644
--- a/tonic/src/transport/client.rs
+++ b/tonic/src/transport/client.rs
@@ -11,19 +11,23 @@ use hyper::{Request, Response};
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use tower_buffer::{future::ResponseFuture, Buffer};
+use tower_service::Service;
 
 type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
-// #[derive/(Clone)]
+type Inner = Box<
+    dyn Service<
+            Request<BoxBody>,
+            Response = Response<hyper::Body>,
+            Error = crate::Error,
+            Future = BoxFuture<'static, Result<Response<hyper::Body>, crate::Error>>,
+        > + Send
+        + 'static,
+>;
+
+#[derive(Clone)]
 pub struct Client {
-    svc: Box<
-        dyn GrpcService<
-                BoxBody,
-                ResponseBody = hyper::Body,
-                Error = crate::Error,
-                Future = BoxFuture<'static, Result<Response<hyper::Body>, crate::Error>>,
-            > + Send
-            + 'static,
-    >,
+    svc: Buffer<Inner, Request<BoxBody>>,
 }
 
 impl Client {
@@ -35,7 +39,9 @@ impl Client {
         let svc = AddOrigin::new(svc, addr);
         let svc = BoxService::new(svc);
 
-        Ok(Self { svc: Box::new(svc) })
+        let svc = Buffer::new(Box::new(svc) as Inner, 100);
+
+        Ok(Self { svc })
     }
 }
 
@@ -44,19 +50,17 @@ impl GrpcService<BoxBody> for Client {
     type Error = super::Error;
 
     type Future = MapErr<
-        BoxFuture<'static, Result<Response<Self::ResponseBody>, crate::Error>>,
+        ResponseFuture<BoxFuture<'static, Result<Response<Self::ResponseBody>, crate::Error>>>,
         fn(crate::Error) -> super::Error,
     >;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        self.svc
-            .poll_ready(cx)
+        GrpcService::poll_ready(&mut self.svc, cx)
             .map_err(|e| super::Error::from((super::ErrorKind::Client, e)))
     }
 
     fn call(&mut self, request: Request<BoxBody>) -> Self::Future {
-        self.svc
-            .call(request)
+        GrpcService::call(&mut self.svc, request)
             .map_err(|e| super::Error::from((super::ErrorKind::Client, e)))
     }
 }