fix(transport): reconnect lazy connections after first failure (#458)

* fix(transport): reconnect lazy connections after first failure

Channels created with lazy connections never try to reconnect if the
first connection attempt fails. This is because `Reconnect` returns
`Poll::Ready(Err)` on poll_ready and the service is considered dead.

This change passes a flag to Reconnect to signal if the connection
is intended to be lazy, in which case reconnect returns the error on
the next call.

fixes #452

tests/integration_tests/tests/connection.rs

@@ -3,7 +3,22 @@ use integration_tests::pb::{test_client::TestClient, test_server, Input, Output}
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use tokio::sync::oneshot;
-use tonic::{transport::Server, Request, Response, Status};
+use tonic::{
+    transport::{Endpoint, Server},
+    Request, Response, Status,
+};
+
+struct Svc(Arc<Mutex<Option<oneshot::Sender<()>>>>);
+
+#[tonic::async_trait]
+impl test_server::Test for Svc {
+    async fn unary_call(&self, _: Request<Input>) -> Result<Response<Output>, Status> {
+        let mut l = self.0.lock().unwrap();
+        l.take().unwrap().send(()).unwrap();
+
+        Ok(Response::new(Output {}))
+    }
+}
 
 #[tokio::test]
 async fn connect_returns_err() {
@@ -14,18 +29,6 @@ async fn connect_returns_err() {
 
 #[tokio::test]
 async fn connect_returns_err_via_call_after_connected() {
-    struct Svc(Arc<Mutex<Option<oneshot::Sender<()>>>>);
-
-    #[tonic::async_trait]
-    impl test_server::Test for Svc {
-        async fn unary_call(&self, _: Request<Input>) -> Result<Response<Output>, Status> {
-            let mut l = self.0.lock().unwrap();
-            l.take().unwrap().send(()).unwrap();
-
-            Ok(Response::new(Output {}))
-        }
-    }
-
     let (tx, rx) = oneshot::channel();
     let sender = Arc::new(Mutex::new(Some(tx)));
     let svc = test_server::TestServer::new(Svc(sender));
@@ -53,3 +56,37 @@ async fn connect_returns_err_via_call_after_connected() {
 
     jh.await.unwrap();
 }
+
+#[tokio::test]
+async fn connect_lazy_reconnects_after_first_failure() {
+    let (tx, rx) = oneshot::channel();
+    let sender = Arc::new(Mutex::new(Some(tx)));
+    let svc = test_server::TestServer::new(Svc(sender));
+
+    let channel = Endpoint::from_static("http://127.0.0.1:1339")
+        .connect_lazy()
+        .unwrap();
+
+    let mut client = TestClient::new(channel);
+
+    // First call should fail, the server is not running
+    client.unary_call(Request::new(Input {})).await.unwrap_err();
+
+    // Start the server now, second call should succeed
+    let jh = tokio::spawn(async move {
+        Server::builder()
+            .add_service(svc)
+            .serve_with_shutdown("127.0.0.1:1339".parse().unwrap(), rx.map(drop))
+            .await
+            .unwrap();
+    });
+
+    tokio::time::delay_for(Duration::from_millis(100)).await;
+    client.unary_call(Request::new(Input {})).await.unwrap();
+
+    // The server shut down, third call should fail
+    tokio::time::delay_for(Duration::from_millis(100)).await;
+    client.unary_call(Request::new(Input {})).await.unwrap_err();
+
+    jh.await.unwrap();
+}

tonic/src/transport/channel/endpoint.rs

@@ -234,7 +234,7 @@ impl Endpoint {
         #[cfg(not(feature = "tls"))]
         let connector = service::connector(http);
 
-        Channel::new(connector, self.clone())
+        Ok(Channel::new(connector, self.clone()))
     }
 
     /// Connect with a custom connector.

tonic/src/transport/channel/mod.rs

@@ -130,7 +130,7 @@ impl Channel {
         (Self::balance(list, DEFAULT_BUFFER_SIZE), tx)
     }
 
-    pub(crate) fn new<C>(connector: C, endpoint: Endpoint) -> Result<Self, super::Error>
+    pub(crate) fn new<C>(connector: C, endpoint: Endpoint) -> Self
     where
         C: Service<Uri> + Send + 'static,
         C::Error: Into<crate::Error> + Send,
@@ -139,10 +139,10 @@ impl Channel {
     {
         let buffer_size = endpoint.buffer_size.clone().unwrap_or(DEFAULT_BUFFER_SIZE);
 
-        let svc = Connection::new(connector, endpoint).map_err(super::Error::from_source)?;
+        let svc = Connection::lazy(connector, endpoint);
         let svc = Buffer::new(Either::A(svc), buffer_size);
 
-        Ok(Channel { svc })
+        Channel { svc }
     }
 
     pub(crate) async fn connect<C>(connector: C, endpoint: Endpoint) -> Result<Self, super::Error>

tonic/src/transport/service/connection.rs

@@ -29,7 +29,7 @@ pub(crate) struct Connection {
 }
 
 impl Connection {
-    pub(crate) fn new<C>(connector: C, endpoint: Endpoint) -> Result<Self, crate::Error>
+    fn new<C>(connector: C, endpoint: Endpoint, is_lazy: bool) -> Self
     where
         C: Service<Uri> + Send + 'static,
         C::Error: Into<crate::Error> + Send,
@@ -61,13 +61,13 @@ impl Connection {
             .into_inner();
 
         let connector = HyperConnect::new(connector, settings);
-        let conn = Reconnect::new(connector, endpoint.uri.clone());
+        let conn = Reconnect::new(connector, endpoint.uri.clone(), is_lazy);
 
         let inner = stack.layer(conn);
 
-        Ok(Self {
+        Self {
             inner: BoxService::new(inner),
-        })
+        }
     }
 
     pub(crate) async fn connect<C>(connector: C, endpoint: Endpoint) -> Result<Self, crate::Error>
@@ -77,7 +77,17 @@ impl Connection {
         C::Future: Unpin + Send,
         C::Response: AsyncRead + AsyncWrite + HyperConnection + Unpin + Send + 'static,
     {
-        Self::new(connector, endpoint)?.ready_oneshot().await
+        Self::new(connector, endpoint, false).ready_oneshot().await
+    }
+
+    pub(crate) fn lazy<C>(connector: C, endpoint: Endpoint) -> Self
+    where
+        C: Service<Uri> + Send + 'static,
+        C::Error: Into<crate::Error> + Send,
+        C::Future: Unpin + Send,
+        C::Response: AsyncRead + AsyncWrite + HyperConnection + Unpin + Send + 'static,
+    {
+        Self::new(connector, endpoint, true)
     }
 }
 

tonic/src/transport/service/reconnect.rs

@@ -19,6 +19,7 @@ where
     target: Target,
     error: Option<M::Error>,
     has_been_connected: bool,
+    is_lazy: bool,
 }
 
 #[derive(Debug)]
@@ -32,13 +33,14 @@ impl<M, Target> Reconnect<M, Target>
 where
     M: Service<Target>,
 {
-    pub(crate) fn new(mk_service: M, target: Target) -> Self {
+    pub(crate) fn new(mk_service: M, target: Target, is_lazy: bool) -> Self {
         Reconnect {
             mk_service,
             state: State::Idle,
             target,
             error: None,
             has_been_connected: false,
+            is_lazy,
         }
     }
 }
@@ -89,11 +91,11 @@ where
 
                             state = State::Idle;
 
-                            if self.has_been_connected {
+                            if !(self.has_been_connected || self.is_lazy) {
+                                return Poll::Ready(Err(e.into()));
+                            } else {
                                 self.error = Some(e.into());
                                 break;
-                            } else {
-                                return Poll::Ready(Err(e.into()));
                             }
                         }
                     }