chore(transport): Clean up server and channel (#174)

* chore(transport): Clean up server and channel

* Fix tls feature compilation

tonic/src/transport/channel/endpoint.rs

@@ -1,9 +1,9 @@
-use super::channel::Channel;
+use super::Channel;
 #[cfg(feature = "tls")]
-use super::{
+use super::ClientTlsConfig;
-    service::TlsConnector,
+#[cfg(feature = "tls")]
-    tls::{Certificate, Identity},
+use crate::transport::service::TlsConnector;
-};
+use crate::transport::{Error, ErrorKind};
 use bytes::Bytes;
 use http::uri::{InvalidUri, Uri};
 use std::{
@@ -18,33 +18,33 @@ use std::{
 /// This struct is used to build and configure HTTP/2 channels.
 #[derive(Clone)]
 pub struct Endpoint {
-    pub(super) uri: Uri,
+    pub(crate) uri: Uri,
-    pub(super) timeout: Option<Duration>,
+    pub(crate) timeout: Option<Duration>,
-    pub(super) concurrency_limit: Option<usize>,
+    pub(crate) concurrency_limit: Option<usize>,
-    pub(super) rate_limit: Option<(u64, Duration)>,
+    pub(crate) rate_limit: Option<(u64, Duration)>,
     #[cfg(feature = "tls")]
-    pub(super) tls: Option<TlsConnector>,
+    pub(crate) tls: Option<TlsConnector>,
-    pub(super) buffer_size: Option<usize>,
+    pub(crate) buffer_size: Option<usize>,
-    pub(super) interceptor_headers:
+    pub(crate) interceptor_headers:
         Option<Arc<dyn Fn(&mut http::HeaderMap) + Send + Sync + 'static>>,
-    pub(super) init_stream_window_size: Option<u32>,
+    pub(crate) init_stream_window_size: Option<u32>,
-    pub(super) init_connection_window_size: Option<u32>,
+    pub(crate) init_connection_window_size: Option<u32>,
-    pub(super) tcp_keepalive: Option<Duration>,
+    pub(crate) tcp_keepalive: Option<Duration>,
-    pub(super) tcp_nodelay: bool,
+    pub(crate) tcp_nodelay: bool,
 }
 
 impl Endpoint {
     // FIXME: determine if we want to expose this or not. This is really
     // just used in codegen for a shortcut.
     #[doc(hidden)]
-    pub fn new<D>(dst: D) -> Result<Self, super::Error>
+    pub fn new<D>(dst: D) -> Result<Self, Error>
     where
         D: TryInto<Self>,
         D::Error: Into<crate::Error>,
     {
         let me = dst
             .try_into()
-            .map_err(|e| super::Error::from_source(super::ErrorKind::Client, e.into()))?;
+            .map_err(|e| Error::from_source(ErrorKind::Client, e.into()))?;
         Ok(me)
     }
 
@@ -181,7 +181,7 @@ impl Endpoint {
     }
 
     /// Create a channel from this config.
-    pub async fn connect(&self) -> Result<Channel, super::Error> {
+    pub async fn connect(&self) -> Result<Channel, Error> {
         Channel::connect(self.clone()).await
     }
 }
@@ -245,90 +245,3 @@ impl fmt::Debug for Endpoint {
         f.debug_struct("Endpoint").finish()
     }
 }
-
-/// Configures TLS settings for endpoints.
-#[cfg(feature = "tls")]
-#[derive(Clone)]
-pub struct ClientTlsConfig {
-    domain: Option<String>,
-    cert: Option<Certificate>,
-    identity: Option<Identity>,
-    rustls_raw: Option<tokio_rustls::rustls::ClientConfig>,
-}
-
-#[cfg(feature = "tls")]
-impl fmt::Debug for ClientTlsConfig {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("ClientTlsConfig")
-            .field("domain", &self.domain)
-            .field("cert", &self.cert)
-            .field("identity", &self.identity)
-            .finish()
-    }
-}
-
-#[cfg(feature = "tls")]
-impl ClientTlsConfig {
-    /// Creates a new `ClientTlsConfig` using Rustls.
-    pub fn with_rustls() -> Self {
-        ClientTlsConfig {
-            domain: None,
-            cert: None,
-            identity: None,
-            rustls_raw: None,
-        }
-    }
-
-    /// Sets the domain name against which to verify the server's TLS certificate.
-    ///
-    /// This has no effect if `rustls_client_config` is used to configure Rustls.
-    pub fn domain_name(self, domain_name: impl Into<String>) -> Self {
-        ClientTlsConfig {
-            domain: Some(domain_name.into()),
-            ..self
-        }
-    }
-
-    /// Sets the CA Certificate against which to verify the server's TLS certificate.
-    ///
-    /// This has no effect if `rustls_client_config` is used to configure Rustls.
-    pub fn ca_certificate(self, ca_certificate: Certificate) -> Self {
-        ClientTlsConfig {
-            cert: Some(ca_certificate),
-            ..self
-        }
-    }
-
-    /// Sets the client identity to present to the server.
-    ///
-    /// This has no effect if `rustls_client_config` is used to configure Rustls.
-    pub fn identity(self, identity: Identity) -> Self {
-        ClientTlsConfig {
-            identity: Some(identity),
-            ..self
-        }
-    }
-
-    /// Use options specified by the given `ClientConfig` to configure TLS.
-    ///
-    /// This overrides all other TLS options set via other means.
-    pub fn rustls_client_config(self, config: tokio_rustls::rustls::ClientConfig) -> Self {
-        ClientTlsConfig {
-            rustls_raw: Some(config),
-            ..self
-        }
-    }
-
-    fn tls_connector(&self, uri: Uri) -> Result<TlsConnector, crate::Error> {
-        let domain = match &self.domain {
-            None => uri.to_string(),
-            Some(domain) => domain.clone(),
-        };
-        match &self.rustls_raw {
-            None => {
-                TlsConnector::new_with_rustls_cert(self.cert.clone(), self.identity.clone(), domain)
-            }
-            Some(c) => TlsConnector::new_with_rustls_raw(c.clone(), domain),
-        }
-    }
-}

tonic/src/transport/channel/mod.rs

@@ -1,9 +1,14 @@
 //! Client implementation and builder.
 
-use super::{
+mod endpoint;
-    service::{Connection, ServiceList},
+#[cfg(feature = "tls")]
-    Endpoint,
+mod tls;
-};
+
+pub use endpoint::Endpoint;
+#[cfg(feature = "tls")]
+pub use tls::ClientTlsConfig;
+
+use super::service::{Connection, ServiceList};
 use crate::{body::BoxBody, client::GrpcService};
 use bytes::Bytes;
 use http::{

tonic/src/transport/channel/tls.rs

@@ -0,0 +1,93 @@
+use crate::transport::{
+    service::TlsConnector,
+    tls::{Certificate, Identity},
+};
+use http::Uri;
+use std::fmt;
+
+/// Configures TLS settings for endpoints.
+#[cfg(feature = "tls")]
+#[derive(Clone)]
+pub struct ClientTlsConfig {
+    domain: Option<String>,
+    cert: Option<Certificate>,
+    identity: Option<Identity>,
+    rustls_raw: Option<tokio_rustls::rustls::ClientConfig>,
+}
+
+#[cfg(feature = "tls")]
+impl fmt::Debug for ClientTlsConfig {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ClientTlsConfig")
+            .field("domain", &self.domain)
+            .field("cert", &self.cert)
+            .field("identity", &self.identity)
+            .finish()
+    }
+}
+
+#[cfg(feature = "tls")]
+impl ClientTlsConfig {
+    /// Creates a new `ClientTlsConfig` using Rustls.
+    pub fn with_rustls() -> Self {
+        ClientTlsConfig {
+            domain: None,
+            cert: None,
+            identity: None,
+            rustls_raw: None,
+        }
+    }
+
+    /// Sets the domain name against which to verify the server's TLS certificate.
+    ///
+    /// This has no effect if `rustls_client_config` is used to configure Rustls.
+    pub fn domain_name(self, domain_name: impl Into<String>) -> Self {
+        ClientTlsConfig {
+            domain: Some(domain_name.into()),
+            ..self
+        }
+    }
+
+    /// Sets the CA Certificate against which to verify the server's TLS certificate.
+    ///
+    /// This has no effect if `rustls_client_config` is used to configure Rustls.
+    pub fn ca_certificate(self, ca_certificate: Certificate) -> Self {
+        ClientTlsConfig {
+            cert: Some(ca_certificate),
+            ..self
+        }
+    }
+
+    /// Sets the client identity to present to the server.
+    ///
+    /// This has no effect if `rustls_client_config` is used to configure Rustls.
+    pub fn identity(self, identity: Identity) -> Self {
+        ClientTlsConfig {
+            identity: Some(identity),
+            ..self
+        }
+    }
+
+    /// Use options specified by the given `ClientConfig` to configure TLS.
+    ///
+    /// This overrides all other TLS options set via other means.
+    pub fn rustls_client_config(self, config: tokio_rustls::rustls::ClientConfig) -> Self {
+        ClientTlsConfig {
+            rustls_raw: Some(config),
+            ..self
+        }
+    }
+
+    pub(crate) fn tls_connector(&self, uri: Uri) -> Result<TlsConnector, crate::Error> {
+        let domain = match &self.domain {
+            None => uri.to_string(),
+            Some(domain) => domain.clone(),
+        };
+        match &self.rustls_raw {
+            None => {
+                TlsConnector::new_with_rustls_cert(self.cert.clone(), self.identity.clone(), domain)
+            }
+            Some(c) => TlsConnector::new_with_rustls_raw(c.clone(), domain),
+        }
+    }
+}

tonic/src/transport/mod.rs

@@ -94,14 +94,12 @@
 pub mod channel;
 pub mod server;
 
-mod endpoint;
 mod error;
 mod service;
 mod tls;
 
 #[doc(inline)]
-pub use self::channel::Channel;
+pub use self::channel::{Channel, Endpoint};
-pub use self::endpoint::Endpoint;
 pub use self::error::Error;
 #[doc(inline)]
 pub use self::server::{Server, ServiceName};
@@ -109,7 +107,7 @@ pub use self::tls::{Certificate, Identity};
 pub use hyper::Body;
 
 #[cfg(feature = "tls")]
-pub use self::endpoint::ClientTlsConfig;
+pub use self::channel::ClientTlsConfig;
 #[cfg(feature = "tls")]
 pub use self::server::ServerTlsConfig;
 

tonic/src/transport/server/mod.rs

@@ -1,13 +1,19 @@
 //! Server implementation and builder.
 
-use super::service::{layer_fn, BoxedIo, Or, Routes, ServiceBuilderExt};
 #[cfg(feature = "tls")]
-use super::{service::TlsAcceptor, tls::Identity, Certificate};
+mod tls;
+
+#[cfg(feature = "tls")]
+pub use tls::ServerTlsConfig;
+
+#[cfg(feature = "tls")]
+use super::service::TlsAcceptor;
+
+use super::service::{layer_fn, BoxedIo, Or, Routes, ServiceBuilderExt};
 use crate::body::BoxBody;
-use futures_core::Stream;
 use futures_util::{
-    future::{self, MapErr},
+    future::{self, poll_fn, MapErr},
-    ready, TryFutureExt, TryStreamExt,
+    TryFutureExt,
 };
 use http::{Request, Response};
 use hyper::{
@@ -222,69 +228,11 @@ impl Server {
         Router::new(self.clone(), svc)
     }
 
-    pub(crate) async fn serve<S>(self, addr: SocketAddr, svc: S) -> Result<(), super::Error>
-    where
-        S: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
-        S::Future: Send + 'static,
-        S::Error: Into<crate::Error> + Send,
-    {
-        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::<_, _, crate::Error>(
-            async_stream::try_stream! {
-                let mut tcp = TcpIncoming::bind(addr)?
-                    .set_nodelay(self.tcp_nodelay)
-                    .set_keepalive(self.tcp_keepalive);
-
-                while let Some(stream) = tcp.try_next().await? {
-                    #[cfg(feature = "tls")]
-                    {
-                        if let Some(tls) = &self.tls {
-                            let io = match tls.connect(stream.into_inner()).await {
-                                Ok(io) => io,
-                                Err(error) => {
-                                    error!(message = "Unable to accept incoming connection.", %error);
-                                    continue
-                                },
-                            };
-                            yield BoxedIo::new(io);
-                            continue;
-                        }
-                    }
-
-                    yield BoxedIo::new(stream);
-                }
-            },
-        );
-        let svc = MakeSvc {
-            inner: svc,
-            interceptor,
-            concurrency_limit,
-            // timeout,
-        };
-
-        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)?;
-
-        Ok(())
-    }
-
     pub(crate) async fn serve_with_shutdown<S, F>(
         self,
         addr: SocketAddr,
         svc: S,
-        signal: F,
+        signal: Option<F>,
     ) -> Result<(), super::Error>
     where
         S: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
@@ -301,11 +249,14 @@ impl Server {
 
         let incoming = hyper::server::accept::from_stream::<_, _, crate::Error>(
             async_stream::try_stream! {
-                let mut tcp = TcpIncoming::bind(addr)?
+                let mut incoming = conn::AddrIncoming::bind(&addr)?;
-                    .set_nodelay(self.tcp_nodelay)
-                    .set_keepalive(self.tcp_keepalive);
 
-                while let Some(stream) = tcp.try_next().await? {
+                incoming.set_nodelay(self.tcp_nodelay);
+                incoming.set_keepalive(self.tcp_keepalive);
+
+
+
+                while let Some(stream) = next_accept(&mut incoming).await? {
                     #[cfg(feature = "tls")]
                     {
                         if let Some(tls) = &self.tls {
@@ -332,15 +283,22 @@ impl Server {
             concurrency_limit,
             // timeout,
         };
-        hyper::Server::builder(incoming)
+
+        let 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)
+            .http2_max_concurrent_streams(max_concurrent_streams);
-            .serve(svc)
+
-            .with_graceful_shutdown(signal)
+        if let Some(signal) = signal {
-            .await
+            server
-            .map_err(map_err)?;
+                .serve(svc)
+                .with_graceful_shutdown(signal)
+                .await
+                .map_err(map_err)?
+        } else {
+            server.serve(svc).await.map_err(map_err)?;
+        }
 
         Ok(())
     }
@@ -410,7 +368,9 @@ where
     ///
     /// [`Server`]: struct.Server.html
     pub async fn serve(self, addr: SocketAddr) -> Result<(), super::Error> {
-        self.server.serve(addr, self.routes).await
+        self.server
+            .serve_with_shutdown::<_, future::Ready<()>>(addr, self.routes, None)
+            .await
     }
 
     /// Consume this [`Server`] creating a future that will execute the server
@@ -423,7 +383,9 @@ where
         addr: SocketAddr,
         f: F,
     ) -> Result<(), super::Error> {
-        self.server.serve_with_shutdown(addr, self.routes, f).await
+        self.server
+            .serve_with_shutdown(addr, self.routes, Some(f))
+            .await
     }
 }
 
@@ -437,105 +399,6 @@ impl fmt::Debug for Server {
     }
 }
 
-/// Configures TLS settings for servers.
-#[cfg(feature = "tls")]
-#[derive(Clone)]
-pub struct ServerTlsConfig {
-    identity: Option<Identity>,
-    client_ca_root: Option<Certificate>,
-    rustls_raw: Option<tokio_rustls::rustls::ServerConfig>,
-}
-
-#[cfg(feature = "tls")]
-impl fmt::Debug for ServerTlsConfig {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("ServerTlsConfig").finish()
-    }
-}
-
-#[cfg(feature = "tls")]
-impl ServerTlsConfig {
-    /// Creates a new `ServerTlsConfig`.
-    pub fn with_rustls() -> Self {
-        ServerTlsConfig {
-            identity: None,
-            client_ca_root: None,
-            rustls_raw: None,
-        }
-    }
-
-    /// Sets the [`Identity`] of the server.
-    pub fn identity(self, identity: Identity) -> Self {
-        ServerTlsConfig {
-            identity: Some(identity),
-            ..self
-        }
-    }
-
-    /// Sets a certificate against which to validate client TLS certificates.
-    pub fn client_ca_root(self, cert: Certificate) -> Self {
-        ServerTlsConfig {
-            client_ca_root: Some(cert),
-            ..self
-        }
-    }
-
-    /// Use options specified by the given `ServerConfig` to configure TLS.
-    ///
-    /// This overrides all other TLS options set via other means.
-    pub fn rustls_server_config(
-        &mut self,
-        config: tokio_rustls::rustls::ServerConfig,
-    ) -> &mut Self {
-        self.rustls_raw = Some(config);
-        self
-    }
-
-    fn tls_acceptor(&self) -> Result<TlsAcceptor, crate::Error> {
-        match &self.rustls_raw {
-            None => TlsAcceptor::new_with_rustls_identity(
-                self.identity.clone().unwrap(),
-                self.client_ca_root.clone(),
-            ),
-            Some(config) => TlsAcceptor::new_with_rustls_raw(config.clone()),
-        }
-    }
-}
-
-#[derive(Debug)]
-struct TcpIncoming {
-    inner: conn::AddrIncoming,
-}
-
-impl TcpIncoming {
-    fn bind(addr: SocketAddr) -> Result<Self, crate::Error> {
-        let inner = conn::AddrIncoming::bind(&addr).map_err(Box::new)?;
-        Ok(Self { inner })
-    }
-
-    fn set_nodelay(mut self, enabled: bool) -> Self {
-        self.inner.set_nodelay(enabled);
-        self
-    }
-
-    fn set_keepalive(mut self, tcp_keepalive: Option<Duration>) -> Self {
-        self.inner.set_keepalive(tcp_keepalive);
-        self
-    }
-}
-
-impl Stream for TcpIncoming {
-    type Item = Result<conn::AddrStream, crate::Error>;
-
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        match ready!(Accept::poll_accept(Pin::new(&mut self.inner), cx)) {
-            Some(Ok(s)) => Poll::Ready(Some(Ok(s))),
-            Some(Err(e)) => Poll::Ready(Some(Err(e.into()))),
-            None => Poll::Ready(None),
-        }
-    }
-}
-
 #[derive(Debug)]
 struct Svc<S>(S);
 
@@ -628,3 +491,16 @@ impl Service<Request<Body>> for Unimplemented {
         )
     }
 }
+
+// Implement try_next for `Accept::poll_accept`.
+async fn next_accept(
+    incoming: &mut conn::AddrIncoming,
+) -> Result<Option<conn::AddrStream>, crate::Error> {
+    let res = poll_fn(|cx| Pin::new(&mut *incoming).poll_accept(cx)).await;
+
+    if let Some(res) = res {
+        Ok(Some(res?))
+    } else {
+        return Ok(None);
+    }
+}

tonic/src/transport/server/tls.rs

@@ -0,0 +1,70 @@
+use crate::transport::{
+    service::TlsAcceptor,
+    tls::{Certificate, Identity},
+};
+use std::fmt;
+
+/// Configures TLS settings for servers.
+#[cfg(feature = "tls")]
+#[derive(Clone)]
+pub struct ServerTlsConfig {
+    identity: Option<Identity>,
+    client_ca_root: Option<Certificate>,
+    rustls_raw: Option<tokio_rustls::rustls::ServerConfig>,
+}
+
+#[cfg(feature = "tls")]
+impl fmt::Debug for ServerTlsConfig {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ServerTlsConfig").finish()
+    }
+}
+
+#[cfg(feature = "tls")]
+impl ServerTlsConfig {
+    /// Creates a new `ServerTlsConfig`.
+    pub fn with_rustls() -> Self {
+        ServerTlsConfig {
+            identity: None,
+            client_ca_root: None,
+            rustls_raw: None,
+        }
+    }
+
+    /// Sets the [`Identity`] of the server.
+    pub fn identity(self, identity: Identity) -> Self {
+        ServerTlsConfig {
+            identity: Some(identity),
+            ..self
+        }
+    }
+
+    /// Sets a certificate against which to validate client TLS certificates.
+    pub fn client_ca_root(self, cert: Certificate) -> Self {
+        ServerTlsConfig {
+            client_ca_root: Some(cert),
+            ..self
+        }
+    }
+
+    /// Use options specified by the given `ServerConfig` to configure TLS.
+    ///
+    /// This overrides all other TLS options set via other means.
+    pub fn rustls_server_config(
+        &mut self,
+        config: tokio_rustls::rustls::ServerConfig,
+    ) -> &mut Self {
+        self.rustls_raw = Some(config);
+        self
+    }
+
+    pub(crate) fn tls_acceptor(&self) -> Result<TlsAcceptor, crate::Error> {
+        match &self.rustls_raw {
+            None => TlsAcceptor::new_with_rustls_identity(
+                self.identity.clone().unwrap(),
+                self.client_ca_root.clone(),
+            ),
+            Some(config) => TlsAcceptor::new_with_rustls_raw(config.clone()),
+        }
+    }
+}