Connect lazily in the load balanced channel.

tonic/src/transport/service/discover.rs

@@ -1,9 +1,8 @@
-use super::super::{service, BoxFuture};
+use super::super::{service};
 use super::connection::Connection;
 use crate::transport::Endpoint;
 
 use std::{
-    future::Future,
     hash::Hash,
     pin::Pin,
     task::{Context, Poll},
@@ -16,14 +15,12 @@ type DiscoverResult<K, S, E> = Result<Change<K, S>, E>;
 
 pub(crate) struct DynamicServiceStream<K: Hash + Eq + Clone> {
     changes: Receiver<Change<K, Endpoint>>,
-    connecting: Option<(K, BoxFuture<Connection, crate::Error>)>,
 }
 
 impl<K: Hash + Eq + Clone> DynamicServiceStream<K> {
     pub(crate) fn new(changes: Receiver<Change<K, Endpoint>>) -> Self {
         Self {
             changes,
-            connecting: None,
         }
     }
 }
@@ -38,36 +35,25 @@ impl<K: Hash + Eq + Clone> Discover for DynamicServiceStream<K> {
         cx: &mut Context<'_>,
     ) -> Poll<DiscoverResult<Self::Key, Self::Service, Self::Error>> {
         loop {
-            if let Some((key, connecting)) = &mut self.connecting {
-                let svc = futures_core::ready!(Pin::new(connecting).poll(cx))?;
-                let key = key.to_owned();
-                self.connecting = None;
-                let change = Ok(Change::Insert(key, svc));
-                return Poll::Ready(change);
-            };
-
             let c = &mut self.changes;
-            match Pin::new(&mut *c).poll_next(cx) {
-                Poll::Pending => return Poll::Pending,
-                Poll::Ready(None) => {
-                    return Poll::Pending;
-                }
+            return match Pin::new(&mut *c).poll_next(cx) {
+                Poll::Pending | Poll::Ready(None) => Poll::Pending,
                 Poll::Ready(Some(change)) => match change {
                     Change::Insert(k, endpoint) => {
                         let mut http = hyper::client::connect::HttpConnector::new();
                         http.set_nodelay(endpoint.tcp_nodelay);
                         http.set_keepalive(endpoint.tcp_keepalive);
                         http.enforce_http(false);
                         #[cfg(feature = "tls")]
-                        let connector = service::connector(http, endpoint.tls.clone());
+                            let connector = service::connector(http, endpoint.tls.clone());
 
                         #[cfg(not(feature = "tls"))]
-                        let connector = service::connector(http);
-                        let fut = Connection::connect(connector, endpoint);
-                        self.connecting = Some((k, Box::pin(fut)));
-                        continue;
+                            let connector = service::connector(http);
+                        let connection = Connection::lazy(connector, endpoint);
+                        let change = Ok(Change::Insert(k, connection));
+                        Poll::Ready(change)
                     }
-                    Change::Remove(k) => return Poll::Ready(Ok(Change::Remove(k))),
+                    Change::Remove(k) => Poll::Ready(Ok(Change::Remove(k))),
                 },
             }
         }