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

examples/Cargo.toml

@@ -42,6 +42,14 @@ path = "src/load_balance/client.rs"
 name = "load-balance-server"
 path = "src/load_balance/server.rs"
 
+[[bin]]
+name = "dynamic-load-balance-client"
+path = "src/dynamic_load_balance/client.rs"
+
+[[bin]]
+name = "dynamic-load-balance-server"
+path = "src/dynamic_load_balance/server.rs"
+
 [[bin]]
 name = "tls-client"
 path = "src/tls/client.rs"
@@ -123,7 +131,7 @@ serde_json = "1.0"
 rand = "0.7"
 # Tracing
 tracing = "0.1"
-tracing-subscriber = { version = "0.2.0-alpha", features = ["tracing-log"] }
+tracing-subscriber = { version = "0.2", features = ["tracing-log"] }
 tracing-attributes = "0.1"
 tracing-futures = "0.2"
 # Required for wellknown types

examples/README.md

@@ -44,7 +44,7 @@ $ cargo run --bin authentication-client
 $ cargo run --bin authentication-server
 ```
 
-## Load balance
+## Load Balance
 
 ### Client
 
@@ -58,6 +58,20 @@ $ cargo run --bin load-balance-client
 $ cargo run --bin load-balance-server
 ```
 
+## Dyanmic Load Balance
+
+### Client
+
+```bash
+$ cargo run --bin dynamic-load-balance-client
+```
+
+### Server
+
+```bash
+$ cargo run --bin dynamic-load-balance-server
+```
+
 ## TLS (rustls)
 
 ### Client

examples/src/dynamic_load_balance/client.rs

@@ -0,0 +1,81 @@
+pub mod pb {
+ tonic::include_proto!("grpc.examples.echo");
+}
+
+use pb::{echo_client::EchoClient, EchoRequest};
+use tonic::transport::Channel;
+
+use tonic::transport::Endpoint;
+
+use std::sync::Arc;
+
+use std::sync::atomic::{AtomicBool, Ordering::SeqCst};
+use tokio::time::timeout;
+use tower::discover::Change;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+ let e1 = Endpoint::from_static("http://[::1]:50051");
+ let e2 = Endpoint::from_static("http://[::1]:50052");
+
+ let (channel, mut rx) = Channel::balance_channel(10);
+ let mut client = EchoClient::new(channel);
+
+ let done = Arc::new(AtomicBool::new(false));
+ let demo_done = done.clone();
+ tokio::spawn(async move {
+ tokio::time::delay_for(tokio::time::Duration::from_secs(5)).await;
+ println!("Added first endpoint");
+ let change = Change::Insert("1", e1);
+ let res = rx.send(change).await;
+ println!("{:?}", res);
+ tokio::time::delay_for(tokio::time::Duration::from_secs(5)).await;
+ println!("Added second endpoint");
+ let change = Change::Insert("2", e2);
+ let res = rx.send(change).await;
+ println!("{:?}", res);
+ tokio::time::delay_for(tokio::time::Duration::from_secs(5)).await;
+ println!("Removed first endpoint");
+ let change = Change::Remove("1");
+ let res = rx.send(change).await;
+ println!("{:?}", res);
+
+ tokio::time::delay_for(tokio::time::Duration::from_secs(5)).await;
+ println!("Removed second endpoint");
+ let change = Change::Remove("2");
+ let res = rx.send(change).await;
+ println!("{:?}", res);
+
+ tokio::time::delay_for(tokio::time::Duration::from_secs(5)).await;
+ println!("Added third endpoint");
+ let e3 = Endpoint::from_static("http://[::1]:50051");
+ let change = Change::Insert("3", e3);
+ let res = rx.send(change).await;
+ println!("{:?}", res);
+
+ tokio::time::delay_for(tokio::time::Duration::from_secs(5)).await;
+ println!("Removed third endpoint");
+ let change = Change::Remove("3");
+ let res = rx.send(change).await;
+ println!("{:?}", res);
+ demo_done.swap(true, SeqCst);
+ });
+
+ while !done.load(SeqCst) {
+ tokio::time::delay_for(tokio::time::Duration::from_millis(500)).await;
+ let request = tonic::Request::new(EchoRequest {
+ message: "hello".into(),
+ });
+
+ let rx = client.unary_echo(request);
+ if let Ok(resp) = timeout(tokio::time::Duration::from_secs(10), rx).await {
+ println!("RESPONSE={:?}", resp);
+ } else {
+ println!("did not receive value within 10 secs");
+ }
+ }
+
+ println!("... Bye");
+
+ Ok(())
+}

examples/src/dynamic_load_balance/server.rs

@@ -0,0 +1,82 @@
+pub mod pb {
+ tonic::include_proto!("grpc.examples.echo");
+}
+
+use futures::Stream;
+use std::net::SocketAddr;
+use std::pin::Pin;
+use tokio::sync::mpsc;
+use tonic::{transport::Server, Request, Response, Status, Streaming};
+
+use pb::{EchoRequest, EchoResponse};
+
+type EchoResult<T> = Result<Response<T>, Status>;
+type ResponseStream = Pin<Box<dyn Stream<Item = Result<EchoResponse, Status>> + Send + Sync>>;
+
+#[derive(Debug)]
+pub struct EchoServer {
+ addr: SocketAddr,
+}
+
+#[tonic::async_trait]
+impl pb::echo_server::Echo for EchoServer {
+ async fn unary_echo(&self, request: Request<EchoRequest>) -> EchoResult<EchoResponse> {
+ let message = format!("{} (from {})", request.into_inner().message, self.addr);
+
+ Ok(Response::new(EchoResponse { message }))
+ }
+
+ type ServerStreamingEchoStream = ResponseStream;
+
+ async fn server_streaming_echo(
+ &self,
+ _: Request<EchoRequest>,
+ ) -> EchoResult<Self::ServerStreamingEchoStream> {
+ Err(Status::unimplemented("not implemented"))
+ }
+
+ async fn client_streaming_echo(
+ &self,
+ _: Request<Streaming<EchoRequest>>,
+ ) -> EchoResult<EchoResponse> {
+ Err(Status::unimplemented("not implemented"))
+ }
+
+ type BidirectionalStreamingEchoStream = ResponseStream;
+
+ async fn bidirectional_streaming_echo(
+ &self,
+ _: Request<Streaming<EchoRequest>>,
+ ) -> EchoResult<Self::BidirectionalStreamingEchoStream> {
+ Err(Status::unimplemented("not implemented"))
+ }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+ let addrs = ["[::1]:50051", "[::1]:50052"];
+
+ let (tx, mut rx) = mpsc::unbounded_channel();
+
+ for addr in &addrs {
+ let addr = addr.parse()?;
+ let tx = tx.clone();
+
+ let server = EchoServer { addr };
+ let serve = Server::builder()
+ .add_service(pb::echo_server::EchoServer::new(server))
+ .serve(addr);
+
+ tokio::spawn(async move {
+ if let Err(e) = serve.await {
+ eprintln!("Error = {:?}", e);
+ }
+
+ tx.send(()).unwrap();
+ });
+ }
+
+ rx.recv().await;
+
+ Ok(())
+}

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);
-
- let discover = ServiceList::new(list);
+ channel
+ }
 
- 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>