diff --git a/tonic/src/transport/endpoint.rs b/tonic/src/transport/endpoint.rs index d6c20c6..2eabd31 100644 --- a/tonic/src/transport/endpoint.rs +++ b/tonic/src/transport/endpoint.rs @@ -8,6 +8,7 @@ pub struct Endpoint { pub(super) uri: Uri, pub(super) timeout: Option, pub(super) concurrency_limit: Option, + pub(super) rate_limit: Option<(u64, Duration)>, pub(super) cert: Option, } @@ -32,6 +33,11 @@ impl Endpoint { self } + pub fn rate_limit(&mut self, limit: u64, duration: Duration) -> &mut Self { + self.rate_limit = Some((limit, duration)); + self + } + pub fn tls_cert(&mut self, ca: Vec, domain: Option) -> &mut Self { self.cert = Some(Cert { ca, @@ -41,6 +47,8 @@ impl Endpoint { self } + // pub fn metadata_interceptor(f: impl Fn(MetadataMap) ->) + pub fn channel(&self) -> Result { Channel::builder().connect(self.clone()) } @@ -51,6 +59,7 @@ impl From for Endpoint { Self { uri, concurrency_limit: None, + rate_limit: None, timeout: None, cert: None, } diff --git a/tonic/src/transport/service/connect.rs b/tonic/src/transport/service/connection.rs similarity index 85% rename from tonic/src/transport/service/connect.rs rename to tonic/src/transport/service/connection.rs index abc450b..d3a514b 100644 --- a/tonic/src/transport/service/connect.rs +++ b/tonic/src/transport/service/connection.rs @@ -9,18 +9,20 @@ use std::{ task::{Context, Poll}, }; use tower::{ - layer::Layer, limit::concurrency::ConcurrencyLimitLayer, timeout::TimeoutLayer, - util::BoxService, ServiceBuilder, + layer::Layer, + limit::{concurrency::ConcurrencyLimitLayer, rate::RateLimitLayer}, + timeout::TimeoutLayer, + util::BoxService, + ServiceBuilder, }; use tower_load::Load; use tower_reconnect::Reconnect; use tower_service::Service; -type Request = http::Request; -type Response = http::Response; +pub(crate) type Request = http::Request; +pub(crate) type Response = http::Response; pub struct Connection { - // inner: AddOrigin, Uri>>, inner: BoxService, } @@ -38,6 +40,7 @@ impl Connection { .concurrency_limit .map(|l| ConcurrencyLimitLayer::new(l)), ) + .optional_layer(endpoint.rate_limit.map(|(l, d)| RateLimitLayer::new(l, d))) .into_inner(); let conn = Reconnect::new(HyperConnect::new(connector, settings), endpoint.uri.clone()); diff --git a/tonic/src/transport/service/discover.rs b/tonic/src/transport/service/discover.rs index d3f77e1..154ef59 100644 --- a/tonic/src/transport/service/discover.rs +++ b/tonic/src/transport/service/discover.rs @@ -1,4 +1,4 @@ -use super::connect::Connection; +use super::connection::Connection; use crate::transport::Endpoint; use std::collections::VecDeque; use std::pin::Pin; diff --git a/tonic/src/transport/service/mod.rs b/tonic/src/transport/service/mod.rs index fa89055..f6b6f5d 100644 --- a/tonic/src/transport/service/mod.rs +++ b/tonic/src/transport/service/mod.rs @@ -1,6 +1,6 @@ mod add_origin; mod boxed; -mod connect; +mod connection; mod connector; mod discover; mod io; @@ -8,7 +8,7 @@ mod layer; pub(crate) use self::add_origin::AddOrigin; pub(crate) use self::boxed::BoxService; -pub(crate) use self::connect::Connection; +pub(crate) use self::connection::Connection; pub(crate) use self::connector::Connector; pub(crate) use self::discover::ServiceList; pub(crate) use self::io::BoxedIo;