From 782c0c19db29fc5fa127147c3e710aee318378bc Mon Sep 17 00:00:00 2001 From: Lucio Franco Date: Mon, 2 Sep 2019 16:39:16 -0400 Subject: [PATCH] Add basic load balancing --- tonic-examples/src/helloworld/client.rs | 8 +++- tonic/Cargo.toml | 15 +++++-- tonic/src/service/connect.rs | 52 +++++++++++++++++++++++++ tonic/src/service/connector.rs | 48 +++++++++++++++++++++++ tonic/src/service/discover.rs | 41 +++++++++++++++++++ tonic/src/service/io.rs | 44 +++++++++++++++++++++ tonic/src/service/mod.rs | 7 ++++ tonic/src/service/reconnect.rs | 24 ++++++++++++ tonic/src/service/tls/mod.rs | 1 + tonic/src/transport/channel.rs | 20 +++++++++- 10 files changed, 253 insertions(+), 7 deletions(-) create mode 100644 tonic/src/service/connect.rs create mode 100644 tonic/src/service/connector.rs create mode 100644 tonic/src/service/discover.rs create mode 100644 tonic/src/service/io.rs create mode 100644 tonic/src/service/reconnect.rs create mode 100644 tonic/src/service/tls/mod.rs diff --git a/tonic-examples/src/helloworld/client.rs b/tonic-examples/src/helloworld/client.rs index 8165819..a208f51 100644 --- a/tonic-examples/src/helloworld/client.rs +++ b/tonic-examples/src/helloworld/client.rs @@ -7,9 +7,13 @@ pub mod hello_world { #[tokio::main] async fn main() -> Result<(), Box> { - 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); diff --git a/tonic/Cargo.toml b/tonic/Cargo.toml index 5d284e7..370bdb4 100644 --- a/tonic/Cargo.toml +++ b/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"] diff --git a/tonic/src/service/connect.rs b/tonic/src/service/connect.rs new file mode 100644 index 0000000..e6f4a5d --- /dev/null +++ b/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, 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> for Connection { + type Response = Response; + type Error = crate::Error; + + type Future = + Pin> + Send + 'static>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + Service::poll_ready(&mut self.inner, cx).map_err(Into::into) + } + + fn call(&mut self, req: Request) -> Self::Future { + let fut = self.inner.call(req); + Box::pin(fut) + } +} + +impl Load for Connection { + type Metric = usize; + + fn load(&self) -> Self::Metric { + 0 + } +} diff --git a/tonic/src/service/connector.rs b/tonic/src/service/connector.rs new file mode 100644 index 0000000..23f8c92 --- /dev/null +++ b/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 = >::Future; + +pub struct Connector { + http: HttpConnector, +} + +impl Connector { + pub fn new() -> Self { + Self { + http: HttpConnector::new(), + } + } +} + +impl Service for Connector { + type Response = BoxedIo; + type Error = crate::Error; + + type Future = + Pin> + Send + 'static>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + 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 { + let io = connect.await?; + + // TODO: build tls based on creds and features + + Ok(BoxedIo::new(io)) +} diff --git a/tonic/src/service/discover.rs b/tonic/src/service/discover.rs new file mode 100644 index 0000000..083a2ed --- /dev/null +++ b/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, + i: usize, +} + +impl ServiceList { + pub fn new(list: Vec) -> 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, 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, + } + } +} diff --git a/tonic/src/service/io.rs b/tonic/src/service/io.rs new file mode 100644 index 0000000..290aa18 --- /dev/null +++ b/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 Io for T where T: AsyncRead + AsyncWrite + Send + Unpin + 'static {} + +pub struct BoxedIo(Pin>); + +impl BoxedIo { + pub(super) fn new(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> { + 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> { + Pin::new(&mut self.0).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.0).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.0).poll_shutdown(cx) + } +} diff --git a/tonic/src/service/mod.rs b/tonic/src/service/mod.rs index da899a2..67b5a6c 100644 --- a/tonic/src/service/mod.rs +++ b/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; diff --git a/tonic/src/service/reconnect.rs b/tonic/src/service/reconnect.rs new file mode 100644 index 0000000..086a234 --- /dev/null +++ b/tonic/src/service/reconnect.rs @@ -0,0 +1,24 @@ +use tower_make::MakeService; +use tower_service::Service; + +#[derive(Debug)] +pub struct Reconnect { + inner: M, +} + +impl Service for Reconnect +where + M: MakeService, +{ + type Response = M::Response; + type Error = M::Error; + type Future = M::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Target) -> Self::Future { + unimplmented!() + } +} diff --git a/tonic/src/service/tls/mod.rs b/tonic/src/service/tls/mod.rs new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/tonic/src/service/tls/mod.rs @@ -0,0 +1 @@ + diff --git a/tonic/src/transport/channel.rs b/tonic/src/transport/channel.rs index f9cff8b..49961cc 100644 --- a/tonic/src/transport/channel.rs +++ b/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 for Channel { } #[derive(Debug)] -pub struct Builder { +pub struct Builder { ca: Option>, override_domain: Option, buffer_size: usize, + balance: Option, } 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) -> Result { + 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(&mut self, discover: D) -> &mut Self { + // self.balance = Some(discover); + // self + // } + pub fn build(&self, uri: T) -> Result where Uri: http::HttpTryFrom, @@ -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) };