feat: Allow providing custom connector to balanced channels by exposing Channel::balance

tonic/src/transport/channel/mod.rs

@@ -10,7 +10,7 @@ pub use endpoint::Endpoint;
 #[cfg(feature = "_tls-any")]
 pub use tls::ClientTlsConfig;
 
-use self::service::{Connection, DynamicServiceStream, Executor, SharedExec};
+use self::service::{Connection, DynamicServiceStream, Executor, MapDiscover, SharedExec};
 use crate::body::Body;
 use bytes::Bytes;
 use http::{
@@ -187,14 +187,22 @@ impl Channel {
         Ok(Channel { svc })
     }
 
-    pub(crate) fn balance<D, E>(discover: D, buffer_size: usize, executor: E) -> Self
+    /// Create a balanced [`Channel`] using a custom [`Discover`] implementation.
+    ///
+    /// This is a lower level API, prefer to use [`Channel::balance_list`] or
+    /// [`Channel::balance_channel`] if you are not using a custom connector.
+    pub fn balance<D, C, E>(discover: D, buffer_size: usize, executor: E) -> Self
     where
-        D: Discover<Service = Connection> + Unpin + Send + 'static,
+        D: Discover<Service = (C, Endpoint)> + Unpin + Send + 'static,
         D::Error: Into<crate::BoxError>,
         D::Key: Hash + Send + Clone,
+        C: Service<Uri> + Send + 'static,
+        C::Response: hyper::rt::Read + hyper::rt::Write + Unpin + Send,
+        C::Error: Into<crate::BoxError> + Send,
+        C::Future: Send,
         E: Executor<BoxFuture<'static, ()>> + Send + Sync + 'static,
     {
-        let svc = Balance::new(discover);
+        let svc = Balance::new(MapDiscover::new(discover));
 
         let svc = BoxService::new(svc);
         let (svc, worker) = Buffer::pair(svc, buffer_size);

tonic/src/transport/channel/service/discover.rs

@@ -1,13 +1,21 @@
-use super::super::{Connection, Endpoint};
+use super::{
+    super::{Connection, Endpoint},
+    Connector,
+};
 
+use http::Uri;
+use hyper_util::client::legacy::connect::HttpConnector;
+use pin_project::pin_project;
 use std::{
-    hash::Hash,
     pin::Pin,
     task::{Context, Poll},
 };
 use tokio::sync::mpsc::Receiver;
 use tokio_stream::Stream;
-use tower::discover::Change as TowerChange;
+use tower::{
+    discover::{Change as TowerChange, Discover},
+    Service,
+};
 
 /// A change in the service set.
 #[derive(Debug, Clone)]
@@ -18,31 +26,69 @@ pub enum Change<K, V> {
     Remove(K),
 }
 
-pub(crate) struct DynamicServiceStream<K: Hash + Eq + Clone> {
+/// Implements [`Discover<Service = Connection>`](Discover) for any
+/// [`Discover<Service = (Connector, Endpoint)>`](Discover)
+#[pin_project]
+pub(crate) struct MapDiscover<D> {
+    #[pin]
+    discover: D,
+}
+
+impl<D> MapDiscover<D> {
+    pub(crate) fn new(discover: D) -> Self {
+        Self { discover }
+    }
+}
+
+impl<D, C> Stream for MapDiscover<D>
+where
+    D: Discover<Service = (C, Endpoint)>,
+    C: Service<Uri> + Send + 'static,
+    C::Response: hyper::rt::Read + hyper::rt::Write + Unpin + Send,
+    C::Error: Into<crate::BoxError> + Send,
+    C::Future: Send,
+{
+    type Item = Result<TowerChange<D::Key, Connection>, D::Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        match this.discover.poll_discover(cx) {
+            Poll::Ready(Some(Ok(change))) => match change {
+                TowerChange::Insert(k, (conn, e)) => {
+                    Poll::Ready(Some(Ok(TowerChange::Insert(k, Connection::lazy(conn, e)))))
+                }
+                TowerChange::Remove(k) => Poll::Ready(Some(Ok(TowerChange::Remove(k)))),
+            },
+            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+/// Implements [`Discover<Service = Connection>`](Discover) for [`Receiver`]
+pub(crate) struct DynamicServiceStream<K> {
     changes: Receiver<Change<K, Endpoint>>,
 }
 
-impl<K: Hash + Eq + Clone> DynamicServiceStream<K> {
+impl<K> DynamicServiceStream<K> {
     pub(crate) fn new(changes: Receiver<Change<K, Endpoint>>) -> Self {
         Self { changes }
     }
 }
 
-impl<K: Hash + Eq + Clone> Stream for DynamicServiceStream<K> {
-    type Item = Result<TowerChange<K, Connection>, crate::BoxError>;
+impl<K> Stream for DynamicServiceStream<K> {
+    type Item = Result<TowerChange<K, (Connector<HttpConnector>, Endpoint)>, crate::BoxError>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         match Pin::new(&mut self.changes).poll_recv(cx) {
             Poll::Pending | Poll::Ready(None) => Poll::Pending,
             Poll::Ready(Some(change)) => match change {
-                Change::Insert(k, endpoint) => {
-                    let connection = Connection::lazy(endpoint.http_connector(), endpoint);
-                    Poll::Ready(Some(Ok(TowerChange::Insert(k, connection))))
+                Change::Insert(k, e) => {
+                    Poll::Ready(Some(Ok(TowerChange::Insert(k, (e.http_connector(), e)))))
                 }
                 Change::Remove(k) => Poll::Ready(Some(Ok(TowerChange::Remove(k)))),
             },
         }
     }
 }
-
-impl<K: Hash + Eq + Clone> Unpin for DynamicServiceStream<K> {}

tonic/src/transport/channel/service/mod.rs

@@ -12,7 +12,7 @@ pub(super) use self::connection::Connection;
 
 mod discover;
 pub use self::discover::Change;
-pub(super) use self::discover::DynamicServiceStream;
+pub(super) use self::discover::{DynamicServiceStream, MapDiscover};
 
 mod io;
 use self::io::BoxedIo;