feat(transport): Dynamic load balancing (#341)

tonic/src/transport/channel/mod.rs

@@ -9,7 +9,7 @@ pub use endpoint::Endpoint;
 #[cfg(feature = "tls")]
 pub use tls::ClientTlsConfig;
 
-use super::service::{Connection, ServiceList};
+use super::service::{Connection, DynamicServiceStream};
 use crate::{body::BoxBody, client::GrpcService};
 use bytes::Bytes;
 use http::{
@@ -20,13 +20,18 @@ use hyper::client::connect::Connection as HyperConnection;
 use std::{
     fmt,
     future::Future,
+    hash::Hash,
     pin::Pin,
     task::{Context, Poll},
 };
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::{
+    io::{AsyncRead, AsyncWrite},
+    sync::mpsc::{channel, Sender},
+};
+
 use tower::{
     buffer::{self, Buffer},
-    discover::Discover,
+    discover::{Change, Discover},
     util::{BoxService, Either},
     Service,
 };
@@ -104,17 +109,25 @@ impl Channel {
     /// This creates a [`Channel`] that will load balance accross all the
     /// provided endpoints.
     pub fn balance_list(list: impl Iterator<Item = Endpoint>) -> Self {
-        let list = list.collect::<Vec<_>>();
+        let (channel, mut tx) = Self::balance_channel(DEFAULT_BUFFER_SIZE);
+        list.for_each(|endpoint| {
+            tx.try_send(Change::Insert(endpoint.uri.clone(), endpoint))
+                .unwrap();
+        });
 
-        let buffer_size = list
-            .iter()
-            .next()
-            .and_then(|e| e.buffer_size)
-            .unwrap_or(DEFAULT_BUFFER_SIZE);
+        channel
+    }
 
-        let discover = ServiceList::new(list);
-
-        Self::balance(discover, buffer_size)
+    /// Balance a list of [`Endpoint`]'s.
+    ///
+    /// This creates a [`Channel`] that will listen to a stream of change events and will add or remove provided endpoints.
+    pub fn balance_channel<K>(capacity: usize) -> (Self, Sender<Change<K, Endpoint>>)
+    where
+        K: Hash + Eq + Send + Clone + 'static,
+    {
+        let (tx, rx) = channel(capacity);
+        let list = DynamicServiceStream::new(rx);
+        (Self::balance(list, DEFAULT_BUFFER_SIZE), tx)
     }
 
     pub(crate) async fn connect<C>(connector: C, endpoint: Endpoint) -> Result<Self, super::Error>

tonic/src/transport/service/discover.rs

@@ -1,34 +1,36 @@
 use super::super::service;
 use super::connection::Connection;
 use crate::transport::Endpoint;
+
 use std::{
-    collections::VecDeque,
-    fmt,
     future::Future,
+    hash::Hash,
     pin::Pin,
     task::{Context, Poll},
 };
+use tokio::{stream::Stream, sync::mpsc::Receiver};
+
 use tower::discover::{Change, Discover};
 
-pub(crate) struct ServiceList {
-    list: VecDeque<Endpoint>,
-    connecting:
-        Option<Pin<Box<dyn Future<Output = Result<Connection, crate::Error>> + Send + 'static>>>,
-    i: usize,
+pub(crate) struct DynamicServiceStream<K: Hash + Eq + Clone> {
+    changes: Receiver<Change<K, Endpoint>>,
+    connecting: Option<(
+        K,
+        Pin<Box<dyn Future<Output = Result<Connection, crate::Error>> + Send + 'static>>,
+    )>,
 }
 
-impl ServiceList {
-    pub(crate) fn new(list: Vec<Endpoint>) -> Self {
+impl<K: Hash + Eq + Clone> DynamicServiceStream<K> {
+    pub(crate) fn new(changes: Receiver<Change<K, Endpoint>>) -> Self {
         Self {
-            list: list.into(),
+            changes,
             connecting: None,
-            i: 0,
         }
     }
 }
 
-impl Discover for ServiceList {
-    type Key = usize;
+impl<K: Hash + Eq + Clone> Discover for DynamicServiceStream<K> {
+    type Key = K;
     type Service = Connection;
     type Error = crate::Error;
 
@@ -37,43 +39,40 @@ impl Discover for ServiceList {
         cx: &mut Context<'_>,
     ) -> Poll<Result<Change<Self::Key, Self::Service>, Self::Error>> {
         loop {
-            if let Some(connecting) = &mut self.connecting {
+            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 i = self.i;
-                self.i += 1;
-
-                let change = Ok(Change::Insert(i, svc));
-
+                let change = Ok(Change::Insert(key, svc));
                 return Poll::Ready(change);
-            }
+            };
 
-            if let Some(endpoint) = self.list.pop_front() {
-                let mut http = hyper::client::connect::HttpConnector::new();
-                http.set_nodelay(endpoint.tcp_nodelay);
-                http.set_keepalive(endpoint.tcp_keepalive);
-                http.enforce_http(false);
+            let c = &mut self.changes;
+            match Pin::new(&mut *c).poll_next(cx) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(None) => {
+                    return 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());
 
-                #[cfg(feature = "tls")]
-                let connector = service::connector(http, endpoint.tls.clone());
-
-                #[cfg(not(feature = "tls"))]
-                let connector = service::connector(http);
-
-                let fut = Connection::new(connector, endpoint);
-                self.connecting = Some(Box::pin(fut));
-            } else {
-                return Poll::Pending;
+                        #[cfg(not(feature = "tls"))]
+                        let connector = service::connector(http);
+                        let fut = Connection::new(connector, endpoint);
+                        self.connecting = Some((k, Box::pin(fut)));
+                        continue;
+                    }
+                    Change::Remove(k) => return Poll::Ready(Ok(Change::Remove(k))),
+                },
             }
         }
     }
 }
 
-impl fmt::Debug for ServiceList {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("ServiceList")
-            .field("list", &self.list)
-            .finish()
-    }
-}
+impl<K: Hash + Eq + Clone> Unpin for DynamicServiceStream<K> {}

tonic/src/transport/service/mod.rs

@@ -12,7 +12,7 @@ mod tls;
 pub(crate) use self::add_origin::AddOrigin;
 pub(crate) use self::connection::Connection;
 pub(crate) use self::connector::connector;
-pub(crate) use self::discover::ServiceList;
+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};