feat: make channel creation more flexible (#98)

In the interest of enabling caller-managed channels, and in order to
reduce feature gate combinatorial explosion, this patch adds a layer
of indirection into the etcd client's Channel abstraction. This
unifies the existing ways of creating a native Tonic channel, an
OpenSSL channel, and future support for caller-supplied channels.

The core abstractions introduced here are the `Channel` enum, which
can represent multiple ways of passing a channel to the generated
clients, and the `MakeBalancedChannel` trait, which allows for the
caller to customize the way in which a balanced channel is created.

The `Channel` enum contains a `Custom` variant, which stores any
Tower service that can process an `http::Request<tonic::body::BoxBody>`
and return an `http::Response<tonic::body::BoxBody>`. This allows for
maximum customization of the channel from the caller, supporting
custom load-balancing policies, custom discovery, and custom transport
layers.

Co-authored-by: Nathan Moos <nmoos@cisco.com>

Author: nmooscisco

src/channel.rs

@@ -1,16 +1,168 @@
-#[cfg(not(feature = "tls-openssl"))]
-pub use tonic::transport::Channel;
+use std::{future::Future, pin::Pin, task::ready};
 
+use http::Uri;
+use tokio::sync::mpsc::Sender;
+use tonic::transport::Endpoint;
+use tower::{discover::Change, util::BoxCloneService, Service};
+
+/// A type alias to make the below types easier to represent.
+pub type EndpointUpdater = Sender<Change<Uri, Endpoint>>;
+
+/// Creates a balanced channel.
+pub trait BalancedChannelBuilder {
+    type Error;
+
+    /// Makes a new balanced channel, given the provided options.
+    fn balanced_channel(
+        self,
+        buffer_size: usize,
+    ) -> Result<(Channel, EndpointUpdater), Self::Error>;
+}
+
+/// Create a simple Tonic channel.
+pub struct Tonic;
+
+impl BalancedChannelBuilder for Tonic {
+    type Error = tonic::transport::Error;
+
+    #[inline]
+    fn balanced_channel(
+        self,
+        buffer_size: usize,
+    ) -> Result<(Channel, EndpointUpdater), Self::Error> {
+        let (chan, tx) = tonic::transport::Channel::balance_channel(buffer_size);
+        Ok((Channel::Tonic(chan), tx))
+    }
+}
+
+/// Create an Openssl-backed channel.
 #[cfg(feature = "tls-openssl")]
-pub use self::openssl::Channel;
+pub struct Openssl {
+    pub(crate) conn: crate::openssl_tls::OpenSslConnector,
+}
 
 #[cfg(feature = "tls-openssl")]
-mod openssl {
-    use crate::openssl_tls;
-
-    /// Because we cannot create `Channel` by the balanced, cached channels,
-    /// we cannot create clients (which explicitly requires `Channel` as argument) directly.
-    ///
-    /// This type alias would be useful to 'batch replace' the signature of `Client::new`.
-    pub type Channel = openssl_tls::OpenSslChannel;
+impl BalancedChannelBuilder for Openssl {
+    type Error = crate::error::Error;
+
+    #[inline]
+    fn balanced_channel(self, _: usize) -> Result<(Channel, EndpointUpdater), Self::Error> {
+        let (chan, tx) = crate::openssl_tls::balanced_channel(self.conn)?;
+        Ok((Channel::Openssl(chan), tx))
+    }
+}
+
+type TonicRequest = http::Request<tonic::body::BoxBody>;
+type TonicResponse = http::Response<tonic::body::BoxBody>;
+pub type CustomChannel = BoxCloneService<TonicRequest, TonicResponse, tower::BoxError>;
+
+/// Represents a channel that can be created by a BalancedChannelBuilder
+/// or may be initialized externally and passed into the client.
+#[derive(Clone)]
+pub enum Channel {
+    /// A standard tonic channel.
+    Tonic(tonic::transport::Channel),
+
+    /// An OpenSSL channel.
+    #[cfg(feature = "tls-openssl")]
+    Openssl(crate::openssl_tls::OpenSslChannel),
+
+    /// A custom Service impl, inside a Box.
+    Custom(CustomChannel),
+}
+
+impl std::fmt::Debug for Channel {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Channel").finish_non_exhaustive()
+    }
+}
+
+pub enum ChannelFuture {
+    Tonic(<tonic::transport::Channel as Service<TonicRequest>>::Future),
+    #[cfg(feature = "tls-openssl")]
+    Openssl(<crate::openssl_tls::OpenSslChannel as Service<TonicRequest>>::Future),
+    Custom(<CustomChannel as Service<TonicRequest>>::Future),
+}
+
+impl std::future::Future for ChannelFuture {
+    type Output = Result<TonicResponse, tower::BoxError>;
+
+    #[inline]
+    fn poll(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        // Safety: trivial projection
+        unsafe {
+            let this = self.get_unchecked_mut();
+            match this {
+                ChannelFuture::Tonic(fut) => {
+                    let fut = Pin::new_unchecked(fut);
+                    let result = ready!(Future::poll(fut, cx));
+                    result.map_err(|e| Box::new(e) as tower::BoxError).into()
+                }
+                #[cfg(feature = "tls-openssl")]
+                ChannelFuture::Openssl(fut) => {
+                    let fut = Pin::new_unchecked(fut);
+                    Future::poll(fut, cx)
+                }
+                ChannelFuture::Custom(fut) => {
+                    let fut = Pin::new_unchecked(fut);
+                    Future::poll(fut, cx)
+                }
+            }
+        }
+    }
+}
+
+impl ChannelFuture {
+    #[inline]
+    fn from_tonic(value: <tonic::transport::Channel as Service<TonicRequest>>::Future) -> Self {
+        Self::Tonic(value)
+    }
+
+    #[cfg(feature = "tls-openssl")]
+    #[inline]
+    fn from_openssl(
+        value: <crate::openssl_tls::OpenSslChannel as Service<TonicRequest>>::Future,
+    ) -> Self {
+        Self::Openssl(value)
+    }
+
+    #[inline]
+    fn from_custom(value: <CustomChannel as Service<TonicRequest>>::Future) -> Self {
+        Self::Custom(value)
+    }
+}
+
+impl Service<TonicRequest> for Channel {
+    type Response = TonicResponse;
+    type Error = tower::BoxError;
+    type Future = ChannelFuture;
+
+    #[inline]
+    fn poll_ready(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), Self::Error>> {
+        match self {
+            Channel::Tonic(channel) => {
+                let result = ready!(channel.poll_ready(cx));
+                result.map_err(|e| Box::new(e) as tower::BoxError).into()
+            }
+            #[cfg(feature = "tls-openssl")]
+            Channel::Openssl(openssl) => openssl.poll_ready(cx),
+            Channel::Custom(custom) => custom.poll_ready(cx),
+        }
+    }
+
+    #[inline]
+    fn call(&mut self, req: TonicRequest) -> Self::Future {
+        match self {
+            Channel::Tonic(channel) => ChannelFuture::from_tonic(channel.call(req)),
+            #[cfg(feature = "tls-openssl")]
+            Channel::Openssl(openssl) => ChannelFuture::from_openssl(openssl.call(req)),
+            Channel::Custom(custom) => ChannelFuture::from_custom(custom.call(req)),
+        }
+    }
 }

src/client.rs

@@ -1,12 +1,10 @@
 //! Asynchronous client & synchronous client.
 
-#[cfg(not(feature = "tls-openssl"))]
-use crate::channel::Channel;
 use crate::error::{Error, Result};
 use crate::intercept::{InterceptedChannel, Interceptor};
 use crate::lock::RwLockExt;
 #[cfg(feature = "tls-openssl")]
-use crate::openssl_tls::{self, OpenSslClientConfig, OpenSslConnector};
+use crate::openssl_tls::{OpenSslClientConfig, OpenSslConnector};
 use crate::rpc::auth::Permission;
 use crate::rpc::auth::{AuthClient, AuthDisableResponse, AuthEnableResponse};
 use crate::rpc::auth::{
@@ -77,6 +75,28 @@ impl Client {
         endpoints: S,
         options: Option<ConnectOptions>,
     ) -> Result<Self> {
+        #[cfg(not(feature = "tls-openssl"))]
+        let make_balanced_channel = crate::channel::Tonic;
+        #[cfg(feature = "tls-openssl")]
+        let make_balanced_channel = crate::channel::Openssl {
+            conn: options
+                .clone()
+                .and_then(|o| o.otls)
+                .unwrap_or_else(OpenSslConnector::create_default)?,
+        };
+        Self::connect_with_balanced_channel(endpoints, options, make_balanced_channel).await
+    }
+
+    /// Connect to `etcd` servers from given `endpoints` and a balanced channel.
+    pub async fn connect_with_balanced_channel<E: AsRef<str>, S: AsRef<[E]>, MBC>(
+        endpoints: S,
+        options: Option<ConnectOptions>,
+        make_balanced_channel: MBC,
+    ) -> Result<Self>
+    where
+        MBC: crate::channel::BalancedChannelBuilder,
+        crate::error::Error: From<MBC::Error>,
+    {
         let endpoints = {
             let mut eps = Vec::new();
             for e in endpoints.as_ref() {
@@ -91,15 +111,7 @@ impl Client {
         }
 
         // Always use balance strategy even if there is only one endpoint.
-        #[cfg(not(feature = "tls-openssl"))]
-        let (channel, tx) = Channel::balance_channel(64);
-        #[cfg(feature = "tls-openssl")]
-        let (channel, tx) = openssl_tls::balanced_channel(
-            options
-                .clone()
-                .and_then(|o| o.otls)
-                .unwrap_or_else(OpenSslConnector::create_default)?,
-        )?;
+        let (channel, tx) = make_balanced_channel.balanced_channel(64)?;
         let channel = InterceptedChannel::new(
             channel,
             Interceptor {
@@ -122,8 +134,7 @@ impl Client {
     }
 
     fn build_endpoint(url: &str, options: &Option<ConnectOptions>) -> Result<Endpoint> {
-        #[cfg(feature = "tls-openssl")]
-        use tonic::transport::Channel;
+        use tonic::transport::Channel as TonicChannel;
         let mut endpoint = if url.starts_with(HTTP_PREFIX) {
             #[cfg(feature = "tls")]
             if let Some(connect_options) = options {
@@ -134,7 +145,7 @@ impl Client {
                 }
             }
 
-            Channel::builder(url.parse()?)
+            TonicChannel::builder(url.parse()?)
         } else if url.starts_with(HTTPS_PREFIX) {
             #[cfg(not(any(feature = "tls", feature = "tls-openssl")))]
             return Err(Error::InvalidArgs(String::from(
@@ -143,7 +154,7 @@ impl Client {
 
             #[cfg(all(feature = "tls-openssl", not(feature = "tls")))]
             {
-                Channel::builder(url.parse()?)
+                TonicChannel::builder(url.parse()?)
             }
 
             #[cfg(feature = "tls")]
@@ -155,7 +166,7 @@ impl Client {
                 }
                 .unwrap_or_else(TlsOptions::new);
 
-                Channel::builder(url.parse()?).tls_config(tls)?
+                TonicChannel::builder(url.parse()?).tls_config(tls)?
             }
         } else {
             #[cfg(feature = "tls")]
@@ -169,11 +180,11 @@ impl Client {
                 match tls {
                     Some(tls) => {
                         let e = HTTPS_PREFIX.to_owned() + url;
-                        Channel::builder(e.parse()?).tls_config(tls)?
+                        TonicChannel::builder(e.parse()?).tls_config(tls)?
                     }
                     None => {
                         let e = HTTP_PREFIX.to_owned() + url;
-                        Channel::builder(e.parse()?)
+                        TonicChannel::builder(e.parse()?)
                     }
                 }
             }
@@ -186,13 +197,13 @@ impl Client {
                     HTTP_PREFIX
                 };
                 let e = pfx.to_owned() + url;
-                Channel::builder(e.parse()?)
+                TonicChannel::builder(e.parse()?)
             }
 
             #[cfg(all(not(feature = "tls"), not(feature = "tls-openssl")))]
             {
                 let e = HTTP_PREFIX.to_owned() + url;
-                Channel::builder(e.parse()?)
+                TonicChannel::builder(e.parse()?)
             }
         };
 

src/lib.rs

@@ -68,6 +68,7 @@ mod openssl_tls;
 mod rpc;
 mod vec;
 
+pub use crate::channel::{BalancedChannelBuilder, Channel};
 pub use crate::client::{Client, ConnectOptions};
 pub use crate::error::Error;
 pub use crate::namespace::{KvClientPrefix, LeaseClientPrefix};