Add more config options for channels

tonic/src/transport/endpoint.rs

@@ -8,6 +8,7 @@ pub struct Endpoint {
     pub(super) uri: Uri,
     pub(super) timeout: Option<Duration>,
     pub(super) concurrency_limit: Option<usize>,
+    pub(super) rate_limit: Option<(u64, Duration)>,
     pub(super) cert: Option<Cert>,
 }
 
@@ -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<u8>, domain: Option<String>) -> &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, super::Error> {
         Channel::builder().connect(self.clone())
     }
@@ -51,6 +59,7 @@ impl From<Uri> for Endpoint {
         Self {
             uri,
             concurrency_limit: None,
+            rate_limit: None,
             timeout: None,
             cert: None,
         }

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<BoxBody>;
-type Response = http::Response<hyper::Body>;
+pub(crate) type Request = http::Request<BoxBody>;
+pub(crate) type Response = http::Response<hyper::Body>;
 
 pub struct Connection {
-    // inner: AddOrigin<Reconnect<HyperConnect<Connector, BoxBody, Uri>, Uri>>,
     inner: BoxService<Request, Response, crate::Error>,
 }
 
@@ -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());

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;

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;