From 2ea17b2ecfc40a20f4d9608f807b3d099a8f415d Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Mon, 27 Jul 2020 12:44:11 -0400
Subject: [PATCH] fix(transport): Return connection error on `Channel::connect`
 (#413)

Before this fix, if the connect phase of the transport failed before
ever establishing a connection, we would never return the error until
the first call to send a request. This PR changes that behavior to only
forward the error to the call method if we have ever made a connection
before. If we have never established a connection before then
`Reconnect` will return an error on the call to `poll_ready`.

Fixes #403
---
 tests/integration_tests/tests/connection.rs | 55 +++++++++++++++++++++
 tonic/src/transport/service/reconnect.rs    | 15 +++++-
 2 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 tests/integration_tests/tests/connection.rs

diff --git a/tests/integration_tests/tests/connection.rs b/tests/integration_tests/tests/connection.rs
new file mode 100644
index 0000000..788a043
--- /dev/null
+++ b/tests/integration_tests/tests/connection.rs
@@ -0,0 +1,55 @@
+use futures_util::FutureExt;
+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};
+
+#[tokio::test]
+async fn connect_returns_err() {
+    let res = TestClient::connect("http://thisdoesntexist").await;
+
+    assert!(res.is_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));
+
+    let jh = tokio::spawn(async move {
+        Server::builder()
+            .add_service(svc)
+            .serve_with_shutdown("127.0.0.1:1338".parse().unwrap(), rx.map(drop))
+            .await
+            .unwrap();
+    });
+
+    tokio::time::delay_for(Duration::from_millis(100)).await;
+
+    let mut client = TestClient::connect("http://127.0.0.1:1338").await.unwrap();
+
+    // First call should pass, then shutdown the server
+    client.unary_call(Request::new(Input {})).await.unwrap();
+
+    tokio::time::delay_for(Duration::from_millis(100)).await;
+
+    let res = client.unary_call(Request::new(Input {})).await;
+
+    assert!(res.is_err());
+
+    jh.await.unwrap();
+}
diff --git a/tonic/src/transport/service/reconnect.rs b/tonic/src/transport/service/reconnect.rs
index 3e6e701..1d6394a 100644
--- a/tonic/src/transport/service/reconnect.rs
+++ b/tonic/src/transport/service/reconnect.rs
@@ -18,6 +18,7 @@ where
     state: State<M::Future, M::Response>,
     target: Target,
     error: Option<M::Error>,
+    has_been_connected: bool,
 }
 
 #[derive(Debug)]
@@ -37,6 +38,7 @@ where
             state: State::Idle,
             target,
             error: None,
+            has_been_connected: false,
         }
     }
 }
@@ -84,14 +86,23 @@ where
                         }
                         Poll::Ready(Err(e)) => {
                             trace!("poll_ready; error");
+
                             state = State::Idle;
-                            self.error = Some(e.into());
-                            break;
+
+                            if self.has_been_connected {
+                                self.error = Some(e.into());
+                                break;
+                            } else {
+                                return Poll::Ready(Err(e.into()));
+                            }
                         }
                     }
                 }
                 State::Connected(ref mut inner) => {
                     trace!("poll_ready; connected");
+
+                    self.has_been_connected = true;
+
                     match inner.poll_ready(cx) {
                         Poll::Ready(Ok(())) => {
                             trace!("poll_ready; ready");