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

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(())
+}