Add basic load balancing

tonic-examples/src/helloworld/client.rs

@@ -7,9 +7,13 @@ pub mod hello_world {
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let origin = http::Uri::from_static("http://[::1]:50051");
+    let origin = vec![
+        http::Uri::from_static("http://[::1]:50051"),
+        http::Uri::from_static("http://[::1]:50051"),
+        http::Uri::from_static("http://[::1]:50051"),
+    ];
 
-    let svc = Channel::builder().build(origin)?;
+    let svc = Channel::builder().balance_list(origin)?;
 
     let mut client = hello_world::GreeterClient::new(svc);
 

tonic/Cargo.toml

@@ -18,7 +18,7 @@ tower-service = "=0.3.0-alpha.1"
 tokio-codec = "=0.2.0-alpha.4"
 async-stream = { git = "https://github.com/tokio-rs/async-stream" }
 http-body = "0.2.0-alpha.1"
-pin-project = "0.4.0-alpha.2"
+pin-project = "0.4.0-alpha.7"
 
 # optional
 hyper = { git = "https://github.com/hyperium/hyper", optional = true}
@@ -26,6 +26,9 @@ tokio = { version = "=0.2.0-alpha.4", default-features = false, features = ["tcp
 tower-make = "=0.1.0-alpha.2"
 tower-reconnect = { git = "https://github.com/tower-rs/tower", branch = "lucio/update-reconnect-buffer", optional = true }
 tower-buffer = { git = "https://github.com/tower-rs/tower", branch = "lucio/update-reconnect-buffer", optional = true }
+tower-balance = { git = "https://github.com/tower-rs/tower", branch = "lucio/update-balance", optional = true }
+tower-load = { git = "https://github.com/tower-rs/tower", branch = "lucio/update-balance", optional = true }
+tower-discover = { git = "https://github.com/tower-rs/tower", branch = "lucio/update-balance", optional = true }
 
 # openssl
 tokio-openssl = { version = "=0.4.0-alpha.4", optional = true }
@@ -38,10 +41,16 @@ tokio-rustls = { version = "0.12.0-alpha.2", optional = true }
 default = ["transport"]
 transport = [
     "hyper",
-    "tower-reconnect",
-    "tower-buffer",
+    "tower",
     "tokio",
     "openssl-1",
 ]
+tower = [
+    "tower-reconnect",
+    "tower-buffer",
+    "tower-balance",
+    "tower-load",
+    "tower-discover"
+]
 openssl-1 = ["openssl", "tokio-openssl"]
 rustls = ["tokio-rustls"]

tonic/src/service/connect.rs

@@ -0,0 +1,52 @@
+use super::{add_origin::AddOrigin, connector::Connector};
+use crate::body::BoxBody;
+use http::{Request, Response, Uri};
+use hyper::client::conn::Builder;
+use hyper::client::service::Connect as HyperConnect;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tower_load::Load;
+use tower_reconnect::Reconnect;
+use tower_service::Service;
+
+pub struct Connection {
+    inner: AddOrigin<Reconnect<HyperConnect<Connector, BoxBody, Uri>, Uri>>,
+}
+
+impl Connection {
+    pub fn new(uri: Uri) -> Self {
+        let connector = Connector::new();
+        let settings = Builder::new().http2_only(true).clone();
+        let connect = HyperConnect::new(connector, settings);
+        let reconnect = Reconnect::new(connect, uri.clone());
+        let inner = AddOrigin::new(reconnect, uri);
+
+        Self { inner }
+    }
+}
+
+impl Service<Request<BoxBody>> for Connection {
+    type Response = Response<hyper::Body>;
+    type Error = crate::Error;
+
+    type Future =
+        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Service::poll_ready(&mut self.inner, cx).map_err(Into::into)
+    }
+
+    fn call(&mut self, req: Request<BoxBody>) -> Self::Future {
+        let fut = self.inner.call(req);
+        Box::pin(fut)
+    }
+}
+
+impl Load for Connection {
+    type Metric = usize;
+
+    fn load(&self) -> Self::Metric {
+        0
+    }
+}

tonic/src/service/connector.rs

@@ -0,0 +1,48 @@
+use super::io::BoxedIo;
+use http::Uri;
+use hyper::client::connect::HttpConnector;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tower_make::MakeConnection;
+use tower_service::Service;
+
+type ConnectFuture = <HttpConnector as MakeConnection<Uri>>::Future;
+
+pub struct Connector {
+    http: HttpConnector,
+}
+
+impl Connector {
+    pub fn new() -> Self {
+        Self {
+            http: HttpConnector::new(),
+        }
+    }
+}
+
+impl Service<Uri> for Connector {
+    type Response = BoxedIo;
+    type Error = crate::Error;
+
+    type Future =
+        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        MakeConnection::poll_ready(&mut self.http, cx).map_err(Into::into)
+    }
+
+    fn call(&mut self, uri: Uri) -> Self::Future {
+        let connect_fut = MakeConnection::make_connection(&mut self.http, uri);
+
+        Box::pin(connect(connect_fut))
+    }
+}
+
+async fn connect(connect: ConnectFuture) -> Result<BoxedIo, crate::Error> {
+    let io = connect.await?;
+
+    // TODO: build tls based on creds and features
+
+    Ok(BoxedIo::new(io))
+}

tonic/src/service/discover.rs

@@ -0,0 +1,41 @@
+use super::connect::Connection;
+use http::Uri;
+use std::collections::VecDeque;
+use std::task::{Context, Poll};
+use tower_discover::{Change, Discover};
+
+#[derive(Debug)]
+pub struct ServiceList {
+    list: VecDeque<Uri>,
+    i: usize,
+}
+
+impl ServiceList {
+    pub fn new(list: Vec<Uri>) -> Self {
+        Self {
+            list: list.into(),
+            i: 0,
+        }
+    }
+}
+
+impl Discover for ServiceList {
+    type Key = usize;
+    type Service = Connection;
+    type Error = hyper::Error;
+
+    fn poll(
+        &mut self,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Result<Change<Self::Key, Self::Service>, Self::Error>> {
+        match self.list.pop_front() {
+            Some(uri) => {
+                let i = self.i;
+                self.i += 1;
+                let service = Connection::new(uri);
+                Poll::Ready(Ok(Change::Insert(i, service)))
+            }
+            None => Poll::Pending,
+        }
+    }
+}

tonic/src/service/io.rs

@@ -0,0 +1,44 @@
+use std::io;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tokio::io::{AsyncRead, AsyncWrite};
+
+pub(super) trait Io: AsyncRead + AsyncWrite + Send + Unpin + 'static {}
+
+impl<T> Io for T where T: AsyncRead + AsyncWrite + Send + Unpin + 'static {}
+
+pub struct BoxedIo(Pin<Box<dyn Io>>);
+
+impl BoxedIo {
+    pub(super) fn new<I: Io>(io: I) -> Self {
+        BoxedIo(Box::pin(io))
+    }
+}
+
+impl AsyncRead for BoxedIo {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut self.0).poll_read(cx, buf)
+    }
+}
+
+impl AsyncWrite for BoxedIo {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut self.0).poll_write(cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_shutdown(cx)
+    }
+}

tonic/src/service/mod.rs

@@ -1,7 +1,14 @@
 mod add_origin;
 mod boxed;
 mod grpc;
+// mod reconnect;
+mod connect;
+mod connector;
+mod discover;
+mod io;
+mod tls;
 
 pub use self::add_origin::AddOrigin;
 pub use self::boxed::BoxService;
+pub use self::discover::ServiceList;
 pub use self::grpc::GrpcService;

tonic/src/service/reconnect.rs

@@ -0,0 +1,24 @@
+use tower_make::MakeService;
+use tower_service::Service;
+
+#[derive(Debug)]
+pub struct Reconnect<M> {
+    inner: M,
+}
+
+impl<M, Target, Request> Service<Target> for Reconnect<M>
+where
+    M: MakeService<Target, Request>,
+{
+    type Response = M::Response;
+    type Error = M::Error;
+    type Future = M::Future;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: Target) -> Self::Future {
+        unimplmented!()
+    }
+}

tonic/src/service/tls/mod.rs

@@ -0,0 +1 @@
+

tonic/src/transport/channel.rs

@@ -1,6 +1,6 @@
 use crate::{
     body::BoxBody,
-    service::{AddOrigin, BoxService, GrpcService},
+    service::{AddOrigin, BoxService, GrpcService, ServiceList},
 };
 use futures_util::try_future::{MapErr, TryFutureExt};
 use http::Uri;
@@ -57,10 +57,11 @@ impl GrpcService<BoxBody> for Channel {
 }
 
 #[derive(Debug)]
-pub struct Builder {
+pub struct Builder<D = ServiceList> {
     ca: Option<Vec<u8>>,
     override_domain: Option<String>,
     buffer_size: usize,
+    balance: Option<D>,
 }
 
 impl Builder {
@@ -69,6 +70,7 @@ impl Builder {
             ca: None,
             override_domain: None,
             buffer_size: 1024,
+            balance: None,
         }
     }
 
@@ -89,6 +91,19 @@ impl Builder {
         self
     }
 
+    pub fn balance_list(&mut self, list: Vec<Uri>) -> Result<Channel, super::Error> {
+        let discover = ServiceList::new(list);
+        let svc = tower_balance::p2c::Balance::from_entropy(discover);
+        let svc = BoxService::new(svc);
+        let svc = Buffer::new(Box::new(svc) as Inner, 100);
+        Ok(Channel { svc })
+    }
+
+    // pub fn balance<D: Discover>(&mut self, discover: D) -> &mut Self<D> {
+    //     self.balance = Some(discover);
+    //     self
+    // }
+
     pub fn build<T>(&self, uri: T) -> Result<Channel, super::Error>
     where
         Uri: http::HttpTryFrom<T>,
@@ -128,6 +143,7 @@ impl Builder {
             let svc = tower_reconnect::Reconnect::new(maker, uri.clone());
 
             let svc = AddOrigin::new(svc, uri);
+
             let svc = BoxService::new(svc);
             Buffer::new(Box::new(svc) as Inner, 100)
         };