From d414206736070726bcbaecc51919f1931e9f17cf Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Wed, 14 Apr 2021 09:26:49 +0200
Subject: [PATCH] Use new tower utilities (#560)

* Use new tower utilities

Tower recently introduced `layer_fn` and `ServiceBuilder::option_layer`.
Some very similar things existed in Tonic. This replaces those with what
Tower provides.

* Also use `ServiceBuilder::layer_fn`
---
 tonic/src/transport/server/mod.rs         |  6 +-
 tonic/src/transport/service/connection.rs |  8 +--
 tonic/src/transport/service/layer.rs      | 80 -----------------------
 tonic/src/transport/service/mod.rs        |  2 -
 4 files changed, 7 insertions(+), 89 deletions(-)
 delete mode 100644 tonic/src/transport/service/layer.rs

diff --git a/tonic/src/transport/server/mod.rs b/tonic/src/transport/server/mod.rs
index f13b46b..6f2729d 100644
--- a/tonic/src/transport/server/mod.rs
+++ b/tonic/src/transport/server/mod.rs
@@ -22,7 +22,7 @@ pub(crate) use tokio_rustls::server::TlsStream;
 use crate::transport::Error;
 
 use super::{
-    service::{Or, Routes, ServerIo, ServiceBuilderExt},
+    service::{Or, Routes, ServerIo},
     BoxFuture,
 };
 use crate::{body::BoxBody, request::ConnectionInfo};
@@ -655,8 +655,8 @@ where
 
         Box::pin(async move {
             let svc = ServiceBuilder::new()
-                .optional_layer(concurrency_limit.map(ConcurrencyLimitLayer::new))
-                .optional_layer(timeout.map(TimeoutLayer::new))
+                .option_layer(concurrency_limit.map(ConcurrencyLimitLayer::new))
+                .option_layer(timeout.map(TimeoutLayer::new))
                 .service(svc);
 
             let svc = BoxService::new(Svc {
diff --git a/tonic/src/transport/service/connection.rs b/tonic/src/transport/service/connection.rs
index ed23aac..db3f17e 100644
--- a/tonic/src/transport/service/connection.rs
+++ b/tonic/src/transport/service/connection.rs
@@ -1,5 +1,5 @@
 use super::super::BoxFuture;
-use super::{layer::ServiceBuilderExt, reconnect::Reconnect, AddOrigin, UserAgent};
+use super::{reconnect::Reconnect, AddOrigin, UserAgent};
 use crate::{body::BoxBody, transport::Endpoint};
 use http::Uri;
 use hyper::client::conn::Builder;
@@ -53,9 +53,9 @@ impl Connection {
         let stack = ServiceBuilder::new()
             .layer_fn(|s| AddOrigin::new(s, endpoint.uri.clone()))
             .layer_fn(|s| UserAgent::new(s, endpoint.user_agent.clone()))
-            .optional_layer(endpoint.timeout.map(TimeoutLayer::new))
-            .optional_layer(endpoint.concurrency_limit.map(ConcurrencyLimitLayer::new))
-            .optional_layer(endpoint.rate_limit.map(|(l, d)| RateLimitLayer::new(l, d)))
+            .option_layer(endpoint.timeout.map(TimeoutLayer::new))
+            .option_layer(endpoint.concurrency_limit.map(ConcurrencyLimitLayer::new))
+            .option_layer(endpoint.rate_limit.map(|(l, d)| RateLimitLayer::new(l, d)))
             .into_inner();
 
         let connector = HyperConnect::new(connector, settings);
diff --git a/tonic/src/transport/service/layer.rs b/tonic/src/transport/service/layer.rs
deleted file mode 100644
index b2b1819..0000000
--- a/tonic/src/transport/service/layer.rs
+++ /dev/null
@@ -1,80 +0,0 @@
-use tower::layer::util::Stack;
-use tower::{layer::Layer, util::Either, ServiceBuilder};
-
-pub(crate) trait ServiceBuilderExt<L> {
-    fn layer_fn<F: Fn(S) -> Out, S, Out>(self, f: F) -> ServiceBuilder<Stack<LayerFn<F>, L>>;
-
-    fn optional_layer_fn<F: Fn(S) -> Out, S, Out>(
-        self,
-        f: Option<F>,
-    ) -> ServiceBuilder<Stack<OptionalLayer<LayerFn<F>>, L>>;
-
-    fn optional_layer<T>(self, l: Option<T>) -> ServiceBuilder<Stack<OptionalLayer<T>, L>>;
-}
-
-impl<L> ServiceBuilderExt<L> for ServiceBuilder<L> {
-    fn layer_fn<F, S, Out>(self, f: F) -> ServiceBuilder<Stack<LayerFn<F>, L>>
-    where
-        F: Fn(S) -> Out,
-    {
-        self.layer(LayerFn(f))
-    }
-
-    fn optional_layer_fn<F, S, Out>(
-        self,
-        f: Option<F>,
-    ) -> ServiceBuilder<Stack<OptionalLayer<LayerFn<F>>, L>>
-    where
-        F: Fn(S) -> Out,
-    {
-        let layer = OptionalLayer {
-            inner: f.map(LayerFn),
-        };
-
-        self.layer(layer)
-    }
-
-    fn optional_layer<T>(self, inner: Option<T>) -> ServiceBuilder<Stack<OptionalLayer<T>, L>> {
-        self.layer(OptionalLayer { inner })
-    }
-}
-
-// TODO: figure out why this is causing a warning even though its used in optional_layer_fn
-#[allow(dead_code)]
-pub(crate) fn layer_fn<F>(f: F) -> LayerFn<F> {
-    LayerFn(f)
-}
-
-#[derive(Clone, Copy, Debug)]
-pub(crate) struct LayerFn<F>(F);
-
-impl<F, S, Out> Layer<S> for LayerFn<F>
-where
-    F: Fn(S) -> Out,
-{
-    type Service = Out;
-
-    fn layer(&self, inner: S) -> Self::Service {
-        (self.0)(inner)
-    }
-}
-
-#[derive(Clone, Debug)]
-pub(crate) struct OptionalLayer<L> {
-    inner: Option<L>,
-}
-
-impl<S, L> Layer<S> for OptionalLayer<L>
-where
-    L: Layer<S>,
-{
-    type Service = Either<L::Service, S>;
-
-    fn layer(&self, s: S) -> Self::Service {
-        if let Some(inner) = &self.inner {
-            Either::A(inner.layer(s))
-        } else {
-            Either::B(s)
-        }
-    }
-}
diff --git a/tonic/src/transport/service/mod.rs b/tonic/src/transport/service/mod.rs
index eab3b40..e398830 100644
--- a/tonic/src/transport/service/mod.rs
+++ b/tonic/src/transport/service/mod.rs
@@ -3,7 +3,6 @@ mod connection;
 mod connector;
 mod discover;
 mod io;
-mod layer;
 mod reconnect;
 mod router;
 #[cfg(feature = "tls")]
@@ -15,7 +14,6 @@ pub(crate) use self::connection::Connection;
 pub(crate) use self::connector::connector;
 pub(crate) use self::discover::DynamicServiceStream;
 pub(crate) use self::io::ServerIo;
-pub(crate) use self::layer::ServiceBuilderExt;
 pub(crate) use self::router::{Or, Routes};
 #[cfg(feature = "tls")]
 pub(crate) use self::tls::{TlsAcceptor, TlsConnector};