talksik/tonic
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat(tonic): make it easier to add tower middleware to servers (#651)

Browse Source
This commit is contained in:
David Pedersen
2021-05-19 09:19:58 +02:00
committed by GitHub
parent 4dda4cbcca
commit 4d2667d1cb
22 changed files with 733 additions and 300 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/src/tower/client.rs
+15 -3
View File
@@ -1,8 +1,9 @@
use hello_world::greeter_client::GreeterClient;
use hello_world::HelloRequest;
use service::AuthSvc;
use tower::ServiceBuilder;
use tonic::transport::Channel;
use tonic::{transport::Channel, Request, Status};
pub mod hello_world {
tonic::include_proto!("helloworld");
@@ -11,9 +12,14 @@ pub mod hello_world {
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let channel = Channel::from_static("http://[::1]:50051").connect().await?;
let auth = AuthSvc::new(channel);
let mut client = GreeterClient::new(auth);
let channel = ServiceBuilder::new()
// Interceptors can be also be applied as middleware
.layer(tonic::service::interceptor_fn(intercept))
.layer_fn(AuthSvc::new)
.service(channel);
let mut client = GreeterClient::new(channel);
let request = tonic::Request::new(HelloRequest {
name: "Tonic".into(),
@@ -26,6 +32,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
Ok(())
}
// An interceptor function.
fn intercept(req: Request<()>) -> Result<Request<()>, Status> {
println!("received {:?}", req);
Ok(req)
}
mod service {
use http::{Request, Response};
use std::future::Future;
examples/src/tower/server.rs
+43 -23
View File
@@ -1,11 +1,10 @@
use hyper::{Body, Request as HyperRequest, Response as HyperResponse};
use std::task::{Context, Poll};
use tonic::{
body::BoxBody,
transport::{NamedService, Server},
Request, Response, Status,
use hyper::Body;
use std::{
task::{Context, Poll},
time::Duration,
};
use tower::Service;
use tonic::{body::BoxBody, transport::Server, Request, Response, Status};
use tower::{Layer, Service};
use hello_world::greeter_server::{Greeter, GreeterServer};
use hello_world::{HelloReply, HelloRequest};
@@ -39,27 +38,52 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("GreeterServer listening on {}", addr);
let svc = InterceptedService {
inner: GreeterServer::new(greeter),
};
let svc = GreeterServer::new(greeter);
Server::builder().add_service(svc).serve(addr).await?;
// The stack of middleware that our service will be wrapped in
let layer = tower::ServiceBuilder::new()
// Apply middleware from tower
.timeout(Duration::from_secs(30))
// Apply our own middleware
.layer(MyMiddlewareLayer::default())
// Interceptors can be also be applied as middleware
.layer(tonic::service::interceptor_fn(intercept))
.into_inner();
Server::builder()
// Wrap all services in the middleware stack
.layer(layer)
.add_service(svc)
.serve(addr)
.await?;
Ok(())
}
// An interceptor function.
fn intercept(req: Request<()>) -> Result<Request<()>, Status> {
Ok(req)
}
#[derive(Debug, Clone, Default)]
struct MyMiddlewareLayer;
impl<S> Layer<S> for MyMiddlewareLayer {
type Service = MyMiddleware<S>;
fn layer(&self, service: S) -> Self::Service {
MyMiddleware { inner: service }
}
}
#[derive(Debug, Clone)]
struct InterceptedService<S> {
struct MyMiddleware<S> {
inner: S,
}
impl<S> Service<HyperRequest<Body>> for InterceptedService<S>
impl<S> Service<hyper::Request<Body>> for MyMiddleware<S>
where
S: Service<HyperRequest<Body>, Response = HyperResponse<BoxBody>>
+ NamedService
+ Clone
+ Send
+ 'static,
S: Service<hyper::Request<Body>, Response = hyper::Response<BoxBody>> + Clone + Send + 'static,
S::Future: Send + 'static,
{
type Response = S::Response;
@@ -70,7 +94,7 @@ where
self.inner.poll_ready(cx)
}
fn call(&mut self, req: HyperRequest<Body>) -> Self::Future {
fn call(&mut self, req: hyper::Request<Body>) -> Self::Future {
// This is necessary because tonic internally uses `tower::buffer::Buffer`.
// See https://github.com/tower-rs/tower/issues/547#issuecomment-767629149
// for details on why this is necessary
@@ -85,7 +109,3 @@ where
})
}
}
impl<S: NamedService> NamedService for InterceptedService<S> {
const NAME: &'static str = S::NAME;
}
tests/integration_tests/Cargo.toml
+3
View File
@@ -20,6 +20,9 @@ tokio-stream = { version = "0.1.5", features = ["net"] }
tower-service = "0.3"
hyper = "0.14"
futures = "0.3"
tower = { version = "0.4", features = [] }
http-body = "0.4"
http = "0.2"
[build-dependencies]
tonic-build = { path = "../../tonic-build" }
tests/integration_tests/tests/complex_tower_middleware.rs
+113
View File
@@ -0,0 +1,113 @@
#![allow(unused_variables, dead_code)]
use http_body::Body;
use integration_tests::pb::{test_server, Input, Output};
use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use tonic::{transport::Server, Request, Response, Status};
use tower::{layer::Layer, BoxError, Service};
// all we care about is that this compiles
async fn complex_tower_layers_work() {
struct Svc;
#[tonic::async_trait]
impl test_server::Test for Svc {
async fn unary_call(&self, req: Request<Input>) -> Result<Response<Output>, Status> {
unimplemented!()
}
}
let svc = test_server::TestServer::new(Svc);
Server::builder()
.layer(MyServiceLayer::new())
.add_service(svc)
.serve("127.0.0.1:1322".parse().unwrap())
.await
.unwrap();
}
#[derive(Debug, Clone)]
struct MyServiceLayer {}
impl MyServiceLayer {
fn new() -> Self {
unimplemented!()
}
}
impl<S> Layer<S> for MyServiceLayer {
type Service = MyService<S>;
fn layer(&self, inner: S) -> Self::Service {
unimplemented!()
}
}
#[derive(Debug, Clone)]
struct MyService<S> {
inner: S,
}
impl<S, R, ResBody> Service<R> for MyService<S>
where
S: Service<R, Response = http::Response<ResBody>>,
{
type Response = http::Response<MyBody<ResBody>>;
type Error = BoxError;
type Future = MyFuture<S::Future, ResBody>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
unimplemented!()
}
fn call(&mut self, req: R) -> Self::Future {
unimplemented!()
}
}
struct MyFuture<F, B> {
inner: F,
body: B,
}
impl<F, E, B> Future for MyFuture<F, B>
where
F: Future<Output = Result<http::Response<B>, E>>,
{
type Output = Result<http::Response<MyBody<B>>, BoxError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
unimplemented!()
}
}
struct MyBody<B> {
inner: B,
}
impl<B> Body for MyBody<B>
where
B: Body,
{
type Data = B::Data;
type Error = BoxError;
fn poll_data(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
unimplemented!()
}
fn poll_trailers(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> {
unimplemented!()
}
}
tonic-build/src/client.rs
+13 -7
View File
@@ -36,18 +36,24 @@ pub fn generate<T: Service>(
#connect
impl<T> #service_ident<T>
where T: tonic::client::GrpcService<tonic::body::BoxBody>,
T::ResponseBody: Body + Send + Sync + 'static,
T::Error: Into<StdError>,
<T::ResponseBody as Body>::Error: Into<StdError> + Send, {
where
T: tonic::client::GrpcService<tonic::body::BoxBody>,
T::ResponseBody: Body + Send + Sync + 'static,
T::Error: Into<StdError>,
<T::ResponseBody as Body>::Error: Into<StdError> + Send,
{
pub fn new(inner: T) -> Self {
let inner = tonic::client::Grpc::new(inner);
Self { inner }
}
pub fn with_interceptor(inner: T, interceptor: impl Into<tonic::Interceptor>) -> Self {
let inner = tonic::client::Grpc::with_interceptor(inner, interceptor);
Self { inner }
pub fn with_interceptor<F>(inner: T, interceptor: F) -> #service_ident<InterceptedService<T, F>>
where
F: FnMut(tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status>,
T: Service<http::Request<tonic::body::BoxBody>, Response = http::Response<T::ResponseBody>>,
<T as Service<http::Request<tonic::body::BoxBody>>>::Error: Into<StdError> + Send + Sync,
{
#service_ident::new(InterceptedService::new(inner, interceptor))
}
#methods
tonic-build/src/server.rs
+12 -31
View File
@@ -50,19 +50,20 @@ pub fn generate<T: Service>(
inner: _Inner<T>,
}
struct _Inner<T>(Arc<T>, Option<tonic::Interceptor>);
struct _Inner<T>(Arc<T>);
impl<T: #server_trait> #server_service<T> {
pub fn new(inner: T) -> Self {
let inner = Arc::new(inner);
let inner = _Inner(inner, None);
let inner = _Inner(inner);
Self { inner }
}
pub fn with_interceptor(inner: T, interceptor: impl Into<tonic::Interceptor>) -> Self {
let inner = Arc::new(inner);
let inner = _Inner(inner, Some(interceptor.into()));
Self { inner }
pub fn with_interceptor<F>(inner: T, interceptor: F) -> InterceptedService<Self, F>
where
F: FnMut(tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status>,
{
InterceptedService::new(Self::new(inner), interceptor)
}
}
@@ -107,7 +108,7 @@ pub fn generate<T: Service>(
impl<T: #server_trait> Clone for _Inner<T> {
fn clone(&self) -> Self {
Self(self.0.clone(), self.1.clone())
Self(self.0.clone())
}
}
@@ -336,16 +337,11 @@ fn generate_unary<T: Method>(
let inner = self.inner.clone();
let fut = async move {
let interceptor = inner.1.clone();
let inner = inner.0;
let method = #service_ident(inner);
let codec = #codec_name::default();
let mut grpc = if let Some(interceptor) = interceptor {
tonic::server::Grpc::with_interceptor(codec, interceptor)
} else {
tonic::server::Grpc::new(codec)
};
let mut grpc = tonic::server::Grpc::new(codec);
let res = grpc.unary(method, req).await;
Ok(res)
@@ -391,16 +387,11 @@ fn generate_server_streaming<T: Method>(
let inner = self.inner.clone();
let fut = async move {
let interceptor = inner.1;
let inner = inner.0;
let method = #service_ident(inner);
let codec = #codec_name::default();
let mut grpc = if let Some(interceptor) = interceptor {
tonic::server::Grpc::with_interceptor(codec, interceptor)
} else {
tonic::server::Grpc::new(codec)
};
let mut grpc = tonic::server::Grpc::new(codec);
let res = grpc.server_streaming(method, req).await;
Ok(res)
@@ -443,16 +434,11 @@ fn generate_client_streaming<T: Method>(
let inner = self.inner.clone();
let fut = async move {
let interceptor = inner.1;
let inner = inner.0;
let method = #service_ident(inner);
let codec = #codec_name::default();
let mut grpc = if let Some(interceptor) = interceptor {
tonic::server::Grpc::with_interceptor(codec, interceptor)
} else {
tonic::server::Grpc::new(codec)
};
let mut grpc = tonic::server::Grpc::new(codec);
let res = grpc.client_streaming(method, req).await;
Ok(res)
@@ -498,16 +484,11 @@ fn generate_streaming<T: Method>(
let inner = self.inner.clone();
let fut = async move {
let interceptor = inner.1;
let inner = inner.0;
let method = #service_ident(inner);
let codec = #codec_name::default();
let mut grpc = if let Some(interceptor) = interceptor {
tonic::server::Grpc::with_interceptor(codec, interceptor)
} else {
tonic::server::Grpc::new(codec)
};
let mut grpc = tonic::server::Grpc::new(codec);
let res = grpc.streaming(method, req).await;
Ok(res)
tonic/Cargo.toml
+2
View File
@@ -52,6 +52,7 @@ base64 = "0.13"
percent-encoding = "2.1"
tower-service = "0.3"
tower-layer = "0.3"
tokio-util = { version = "0.6", features = ["codec"] }
async-stream = "0.3"
http-body = "0.4.2"
@@ -83,6 +84,7 @@ rand = "0.8"
bencher = "0.1.5"
quickcheck = "1.0"
quickcheck_macros = "1.0"
tower = { version = "0.4.7", features = ["full"] }
[package.metadata.docs.rs]
all-features = true
tonic/src/client/grpc.rs
+1 -22
View File
@@ -2,7 +2,6 @@ use crate::{
body::BoxBody,
client::GrpcService,
codec::{encode_client, Codec, Streaming},
interceptor::Interceptor,
Code, Request, Response, Status,
};
use futures_core::Stream;
@@ -29,25 +28,12 @@ use std::fmt;
/// [gRPC protocol definition]: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests
pub struct Grpc<T> {
inner: T,
interceptor: Option<Interceptor>,
}
impl<T> Grpc<T> {
/// Creates a new gRPC client with the provided [`GrpcService`].
pub fn new(inner: T) -> Self {
Self {
inner,
interceptor: None,
}
}
/// Creates a new gRPC client with the provided [`GrpcService`] and will apply
/// the provided interceptor on each request.
pub fn with_interceptor(inner: T, interceptor: impl Into<Interceptor>) -> Self {
Self {
inner,
interceptor: Some(interceptor.into()),
}
Self { inner }
}
/// Check if the inner [`GrpcService`] is able to accept a new request.
@@ -153,12 +139,6 @@ impl<T> Grpc<T> {
M1: Send + Sync + 'static,
M2: Send + Sync + 'static,
{
let request = if let Some(interceptor) = &self.interceptor {
interceptor.call(request)?
} else {
request
};
let mut parts = Parts::default();
parts.path_and_query = Some(path);
@@ -217,7 +197,6 @@ impl<T: Clone> Clone for Grpc<T> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
interceptor: self.interceptor.clone(),
}
}
}
tonic/src/codegen.rs
+1
View File
@@ -10,6 +10,7 @@ pub use std::sync::Arc;
pub use std::task::{Context, Poll};
pub use tower_service::Service;
pub type StdError = Box<dyn std::error::Error + Send + Sync + 'static>;
pub use crate::service::interceptor::InterceptedService;
pub use http_body::Body;
pub type BoxFuture<T, E> = self::Pin<Box<dyn self::Future<Output = Result<T, E>> + Send + 'static>>;
tonic/src/extensions.rs
+2 -2
View File
@@ -2,10 +2,10 @@ use std::fmt;
/// A type map of protocol extensions.
///
/// `Extensions` can be used by [`Interceptor`] and [`Request`] to store extra data derived from
/// `Extensions` can be used by [`interceptor_fn`] and [`Request`] to store extra data derived from
/// the underlying protocol.
///
/// [`Interceptor`]: crate::Interceptor
/// [`interceptor_fn`]: crate::service::interceptor_fn
/// [`Request`]: crate::Request
pub struct Extensions {
inner: http::Extensions,
tonic/src/interceptor.rs
-86
View File
@@ -1,86 +0,0 @@
use crate::{Request, Status};
use std::panic::{RefUnwindSafe, UnwindSafe};
use std::{fmt, sync::Arc};
type InterceptorFn = Arc<
dyn Fn(Request<()>) -> Result<Request<()>, Status>
+ Send
+ Sync
+ UnwindSafe
+ RefUnwindSafe
+ 'static,
>;
/// Represents a gRPC interceptor.
///
/// gRPC interceptors are similar to middleware but have much less
/// flexibility. This interceptor allows you to do two main things,
/// one is to add/remove/check items in the `MetadataMap` of each
/// request. Two, cancel a request with any `Status`.
///
/// An interceptor can be used on both the server and client side through
/// the `tonic-build` crate's generated structs.
///
/// These interceptors do not allow you to modify the `Message` of the request
/// but allow you to check for metadata. If you would like to apply middleware like
/// features to the body of the request, going through the `tower` abstraction is recommended.
#[derive(Clone)]
pub struct Interceptor {
f: InterceptorFn,
}
impl Interceptor {
/// Create a new `Interceptor` from the provided function.
pub fn new(
f: impl Fn(Request<()>) -> Result<Request<()>, Status>
+ Send
+ Sync
+ UnwindSafe
+ RefUnwindSafe
+ 'static,
) -> Self {
Interceptor { f: Arc::new(f) }
}
pub(crate) fn call<T>(&self, req: Request<T>) -> Result<Request<T>, Status> {
let (metadata, ext, message) = req.into_parts();
let temp_req = Request::from_parts(metadata, ext, ());
let (metadata, ext, _) = (self.f)(temp_req)?.into_parts();
Ok(Request::from_parts(metadata, ext, message))
}
}
impl<F> From<F> for Interceptor
where
F: Fn(Request<()>) -> Result<Request<()>, Status>
+ Send
+ Sync
+ UnwindSafe
+ RefUnwindSafe
+ 'static,
{
fn from(f: F) -> Self {
Interceptor::new(f)
}
}
impl fmt::Debug for Interceptor {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Interceptor").finish()
}
}
#[cfg(test)]
mod tests {
#[allow(unused_imports)]
use super::*;
#[test]
fn interceptor_fn_is_unwind_safe() {
fn is_unwind_safe<T: UnwindSafe + RefUnwindSafe>() {}
is_unwind_safe::<InterceptorFn>();
}
}
tonic/src/lib.rs
+2 -2
View File
@@ -83,17 +83,18 @@ pub mod client;
pub mod codec;
pub mod metadata;
pub mod server;
pub mod service;
#[cfg(feature = "transport")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport")))]
pub mod transport;
mod extensions;
mod interceptor;
mod macros;
mod request;
mod response;
mod status;
mod util;
/// A re-export of [`async-trait`](https://docs.rs/async-trait) for use with codegen.
#[cfg(feature = "codegen")]
@@ -103,7 +104,6 @@ pub use async_trait::async_trait;
#[doc(inline)]
pub use codec::Streaming;
pub use extensions::Extensions;
pub use interceptor::Interceptor;
pub use request::{IntoRequest, IntoStreamingRequest, Request};
pub use response::Response;
pub use status::{Code, Status};
tonic/src/request.rs
+2 -2
View File
@@ -267,13 +267,13 @@ impl<T> Request<T> {
/// Extensions can be set in interceptors:
///
/// ```no_run
/// use tonic::{Request, Interceptor};
/// use tonic::{Request, service::interceptor_fn};
///
/// struct MyExtension {
/// some_piece_of_data: String,
/// }
///
/// Interceptor::new(|mut request: Request<()>| {
/// interceptor_fn(|mut request: Request<()>| {
/// request.extensions_mut().insert(MyExtension {
/// some_piece_of_data: "foo".to_string(),
/// });
tonic/src/server/grpc.rs
+1 -42
View File
@@ -1,7 +1,6 @@
use crate::{
body::BoxBody,
codec::{encode_server, Codec, Streaming},
interceptor::Interceptor,
server::{ClientStreamingService, ServerStreamingService, StreamingService, UnaryService},
Code, Request, Status,
};
@@ -10,16 +9,6 @@ use futures_util::{future, stream, TryStreamExt};
use http_body::Body;
use std::fmt;
// A try! type macro for intercepting requests
macro_rules! t {
($expr : expr) => {
match $expr {
Ok(request) => request,
Err(res) => return res,
}
};
}
/// A gRPC Server handler.
///
/// This will wrap some inner [`Codec`] and provide utilities to handle
@@ -31,7 +20,6 @@ macro_rules! t {
/// implements some [`Body`].
pub struct Grpc<T> {
codec: T,
interceptor: Option<Interceptor>,
}
impl<T> Grpc<T>
@@ -41,19 +29,7 @@ where
{
/// Creates a new gRPC server with the provided [`Codec`].
pub fn new(codec: T) -> Self {
Self {
codec,
interceptor: None,
}
}
/// Creates a new gRPC server with the provided [`Codec`] and will apply the provided
/// interceptor on each inbound request.
pub fn with_interceptor(codec: T, interceptor: impl Into<Interceptor>) -> Self {
Self {
codec,
interceptor: Some(interceptor.into()),
}
Self { codec }
}
/// Handle a single unary gRPC request.
@@ -77,8 +53,6 @@ where
}
};
let request = t!(self.intercept_request(request));
let response = service
.call(request)
.await
@@ -106,8 +80,6 @@ where
}
};
let request = t!(self.intercept_request(request));
let response = service.call(request).await;
self.map_response(response)
@@ -125,7 +97,6 @@ where
B::Error: Into<crate::Error> + Send + 'static,
{
let request = self.map_request_streaming(req);
let request = t!(self.intercept_request(request));
let response = service
.call(request)
.await
@@ -146,7 +117,6 @@ where
B::Error: Into<crate::Error> + Send,
{
let request = self.map_request_streaming(req);
let request = t!(self.intercept_request(request));
let response = service.call(request).await;
self.map_response(response)
}
@@ -213,17 +183,6 @@ where
Err(status) => status.to_http(),
}
}
fn intercept_request<A>(&self, req: Request<A>) -> Result<Request<A>, http::Response<BoxBody>> {
if let Some(interceptor) = &self.interceptor {
match interceptor.call(req) {
Ok(req) => Ok(req),
Err(status) => Err(status.to_http()),
}
} else {
Ok(req)
}
}
}
impl<T: fmt::Debug> fmt::Debug for Grpc<T> {
tonic/src/service/interceptor.rs
+183
View File
@@ -0,0 +1,183 @@
//! gRPC interceptors which are a kind of middleware.
use crate::Status;
use pin_project::pin_project;
use std::{
fmt,
future::Future,
pin::Pin,
task::{Context, Poll},
};
use tower_layer::Layer;
use tower_service::Service;
/// Create a new interceptor from a function.
///
/// gRPC interceptors are similar to middleware but have less flexibility. This interceptor allows
/// you to do two main things, one is to add/remove/check items in the `MetadataMap` of each
/// request. Two, cancel a request with any `Status`.
///
/// An interceptor can be used on both the server and client side through the `tonic-build` crate's
/// generated structs.
///
/// These interceptors do not allow you to modify the `Message` of the request but allow you to
/// check for metadata. If you would like to apply middleware like features to the body of the
/// request, going through the [tower] abstraction is recommended.
///
/// Interceptors is not recommend should not be used to add logging to your service. For that a
/// [tower] middleware is more appropriate since it can also act on the response.
///
/// See the [interceptor example][example] for more details.
///
/// [tower]: https://crates.io/crates/tower
/// [example]: https://github.com/hyperium/tonic/tree/master/examples/src/interceptor
// TODO: when tower-http is shipped update the docs to mention its `Trace` middleware which has
// support for gRPC and is an easy to add logging
pub fn interceptor_fn<F>(f: F) -> InterceptorFn<F>
where
F: FnMut(crate::Request<()>) -> Result<crate::Request<()>, Status>,
{
InterceptorFn { f }
}
/// An interceptor created from a function.
///
/// See [`interceptor_fn`] for more details.
#[derive(Debug, Clone, Copy)]
pub struct InterceptorFn<F> {
f: F,
}
impl<S, F> Layer<S> for InterceptorFn<F>
where
F: FnMut(crate::Request<()>) -> Result<crate::Request<()>, Status> + Clone,
{
type Service = InterceptedService<S, F>;
fn layer(&self, service: S) -> Self::Service {
InterceptedService::new(service, self.f.clone())
}
}
/// A service wrapped in an interceptor middleware.
///
/// See [`interceptor_fn`] for more details.
#[derive(Clone, Copy)]
pub struct InterceptedService<S, F> {
inner: S,
f: F,
}
impl<S, F> InterceptedService<S, F> {
/// Create a new `InterceptedService` thats wraps `S` and intercepts each request with the
/// function `F`.
pub fn new(service: S, f: F) -> Self
where
F: FnMut(crate::Request<()>) -> Result<crate::Request<()>, Status>,
{
Self { inner: service, f }
}
}
impl<S, F> fmt::Debug for InterceptedService<S, F>
where
S: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("InterceptedService")
.field("inner", &self.inner)
.field("f", &format_args!("{}", std::any::type_name::<F>()))
.finish()
}
}
impl<S, F, ReqBody, ResBody> Service<http::Request<ReqBody>> for InterceptedService<S, F>
where
F: FnMut(crate::Request<()>) -> Result<crate::Request<()>, Status>,
S: Service<http::Request<ReqBody>, Response = http::Response<ResBody>>,
S::Error: Into<crate::Error>,
{
type Response = http::Response<ResBody>;
type Error = crate::Error;
type Future = ResponseFuture<S::Future>;
#[inline]
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, req: http::Request<ReqBody>) -> Self::Future {
let uri = req.uri().clone();
let req = crate::Request::from_http(req);
let (metadata, extensions, msg) = req.into_parts();
match (self.f)(crate::Request::from_parts(metadata, extensions, ())) {
Ok(req) => {
let (metadata, extensions, _) = req.into_parts();
let req = crate::Request::from_parts(metadata, extensions, msg);
let req = req.into_http(uri);
ResponseFuture::future(self.inner.call(req))
}
Err(status) => ResponseFuture::error(status),
}
}
}
// required to use `InterceptedService` with `Router`
#[cfg(feature = "transport")]
impl<S, F> crate::transport::NamedService for InterceptedService<S, F>
where
S: crate::transport::NamedService,
{
const NAME: &'static str = S::NAME;
}
/// Response future for [`InterceptedService`].
#[pin_project]
#[derive(Debug)]
pub struct ResponseFuture<F> {
#[pin]
kind: Kind<F>,
}
impl<F> ResponseFuture<F> {
fn future(future: F) -> Self {
Self {
kind: Kind::Future(future),
}
}
fn error(status: Status) -> Self {
Self {
kind: Kind::Error(Some(status)),
}
}
}
#[pin_project(project = KindProj)]
#[derive(Debug)]
enum Kind<F> {
Future(#[pin] F),
Error(Option<Status>),
}
impl<F, E, B> Future for ResponseFuture<F>
where
F: Future<Output = Result<http::Response<B>, E>>,
E: Into<crate::Error>,
{
type Output = Result<http::Response<B>, crate::Error>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.project().kind.project() {
KindProj::Future(future) => {
let response = futures_core::ready!(future.poll(cx).map_err(Into::into)?);
Poll::Ready(Ok(response))
}
KindProj::Error(status) => {
let error = status.take().unwrap().into();
Poll::Ready(Err(error))
}
}
}
}
tonic/src/service/mod.rs
+6
View File
@@ -0,0 +1,6 @@
//! Utilities for using Tower services with Tonic.
pub mod interceptor;
#[doc(inline)]
pub use self::interceptor::interceptor_fn;
tonic/src/transport/server/incoming.rs
+4 -4
View File
@@ -15,9 +15,9 @@ use std::{
use tokio::io::{AsyncRead, AsyncWrite};
#[cfg(not(feature = "tls"))]
pub(crate) fn tcp_incoming<IO, IE>(
pub(crate) fn tcp_incoming<IO, IE, L>(
incoming: impl Stream<Item = Result<IO, IE>>,
_server: Server,
_server: Server<L>,
) -> impl Stream<Item = Result<ServerIo, crate::Error>>
where
IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
@@ -35,9 +35,9 @@ where
}
#[cfg(feature = "tls")]
pub(crate) fn tcp_incoming<IO, IE>(
pub(crate) fn tcp_incoming<IO, IE, L>(
incoming: impl Stream<Item = Result<IO, IE>>,
server: Server,
server: Server<L>,
) -> impl Stream<Item = Result<ServerIo, crate::Error>>
where
IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
tonic/src/transport/server/mod.rs
+243 -59
View File
@@ -25,26 +25,33 @@ use crate::transport::Error;
use self::recover_error::RecoverError;
use super::service::{GrpcTimeout, Or, Routes, ServerIo};
use crate::{body::BoxBody, request::ConnectionInfo};
use bytes::Bytes;
use futures_core::Stream;
use futures_util::{
future::{self, Either as FutureEither, MapErr},
TryFutureExt,
future::{self, MapErr},
ready, TryFutureExt,
};
use http::{Request, Response};
use http_body::Body as _;
use hyper::{server::accept, Body};
use pin_project::pin_project;
use std::{
fmt,
future::Future,
net::SocketAddr,
pin::Pin,
sync::Arc,
task::{Context, Poll},
time::Duration,
};
use tokio::io::{AsyncRead, AsyncWrite};
use tower::{limit::concurrency::ConcurrencyLimitLayer, util::Either, Service, ServiceBuilder};
use tracing_futures::{Instrument, Instrumented};
use tower::{
layer::util::Identity, layer::Layer, limit::concurrency::ConcurrencyLimitLayer, util::Either,
Service, ServiceBuilder,
};
type BoxService = tower::util::BoxService<Request<Body>, Response<BoxBody>, crate::Error>;
type BoxHttpBody = http_body::combinators::BoxBody<Bytes, crate::Error>;
type BoxService = tower::util::BoxService<Request<Body>, Response<BoxHttpBody>, crate::Error>;
type TraceInterceptor = Arc<dyn Fn(&http::Request<()>) -> tracing::Span + Send + Sync + 'static>;
const DEFAULT_HTTP2_KEEPALIVE_TIMEOUT_SECS: u64 = 20;
@@ -58,7 +65,7 @@ const DEFAULT_HTTP2_KEEPALIVE_TIMEOUT_SECS: u64 = 20;
/// reference implementation that should be a good starting point for anyone
/// wanting to create a more complex and/or specific implementation.
#[derive(Default, Clone)]
pub struct Server {
pub struct Server<L = Identity> {
trace_interceptor: Option<TraceInterceptor>,
concurrency_limit: Option<usize>,
timeout: Option<Duration>,
@@ -73,12 +80,13 @@ pub struct Server {
http2_keepalive_timeout: Option<Duration>,
max_frame_size: Option<u32>,
accept_http1: bool,
layer: L,
}
/// A stack based `Service` router.
#[derive(Debug)]
pub struct Router<A, B> {
server: Server,
pub struct Router<A, B, L = Identity> {
server: Server<L>,
routes: Routes<A, B, Request<Body>>,
}
@@ -88,35 +96,29 @@ pub struct Router<A, B> {
/// gRPC endpoints and can be consumed with the rest of the `tower`
/// ecosystem.
#[derive(Debug)]
pub struct RouterService<A, B> {
router: Router<A, B>,
pub struct RouterService<S> {
inner: S,
}
impl<A, B> Service<Request<Body>> for RouterService<A, B>
impl<S> Service<Request<Body>> for RouterService<S>
where
A: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
A::Future: Send + 'static,
A::Error: Into<crate::Error> + Send,
B: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
B::Future: Send + 'static,
B::Error: Into<crate::Error> + Send,
S: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
S::Future: Send + 'static,
S::Error: Into<crate::Error> + Send,
{
type Response = Response<BoxBody>;
type Error = crate::Error;
#[allow(clippy::type_complexity)]
type Future = FutureEither<
MapErr<A::Future, fn(A::Error) -> crate::Error>,
MapErr<B::Future, fn(B::Error) -> crate::Error>,
>;
type Future = MapErr<S::Future, fn(S::Error) -> crate::Error>;
#[inline]
fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
#[inline]
fn call(&mut self, req: Request<Body>) -> Self::Future {
self.router.routes.call(req)
self.inner.call(req).map_err(Into::into)
}
}
@@ -144,7 +146,7 @@ impl Server {
}
}
impl Server {
impl<L> Server<L> {
/// Configure TLS for this server.
#[cfg(feature = "tls")]
#[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
@@ -319,7 +321,7 @@ impl Server {
///
/// This will clone the `Server` builder and create a router that will
/// route around different services.
pub fn add_service<S>(&mut self, svc: S) -> Router<S, Unimplemented>
pub fn add_service<S>(&mut self, svc: S) -> Router<S, Unimplemented, L>
where
S: Service<Request<Body>, Response = Response<BoxBody>>
+ NamedService
@@ -328,6 +330,7 @@ impl Server {
+ 'static,
S::Future: Send + 'static,
S::Error: Into<crate::Error> + Send,
L: Clone,
{
Router::new(self.clone(), svc)
}
@@ -343,7 +346,7 @@ impl Server {
pub fn add_optional_service<S>(
&mut self,
svc: Option<S>,
) -> Router<Either<S, Unimplemented>, Unimplemented>
) -> Router<Either<S, Unimplemented>, Unimplemented, L>
where
S: Service<Request<Body>, Response = Response<BoxBody>>
+ NamedService
@@ -352,6 +355,7 @@ impl Server {
+ 'static,
S::Future: Send + 'static,
S::Error: Into<crate::Error> + Send,
L: Clone,
{
let svc = match svc {
Some(some) => Either::A(some),
@@ -360,20 +364,104 @@ impl Server {
Router::new(self.clone(), svc)
}
pub(crate) async fn serve_with_shutdown<S, I, F, IO, IE>(
/// Set the [Tower] [`Layer`] all services will be wrapped in.
///
/// This enables using middleware from the [Tower ecosystem][eco].
///
/// # Example
///
/// ```
/// # use tonic::transport::Server;
/// # use tower_service::Service;
/// use tower::timeout::TimeoutLayer;
/// use std::time::Duration;
///
/// # let mut builder = Server::builder();
/// builder.layer(TimeoutLayer::new(Duration::from_secs(30)));
/// ```
///
/// Note that timeouts should be set using [`Server::timeout`]. `TimeoutLayer` is only used
/// here as an example.
///
/// You can build more complex layers using [`ServiceBuilder`]. Those layers can include
/// [interceptors]:
///
/// ```
/// # use tonic::transport::Server;
/// # use tower_service::Service;
/// use tower::ServiceBuilder;
/// use std::time::Duration;
/// use tonic::{Request, Status, service::interceptor_fn};
///
/// fn auth_interceptor(request: Request<()>) -> Result<Request<()>, Status> {
/// if valid_credentials(&request) {
/// Ok(request)
/// } else {
/// Err(Status::unauthenticated("invalid credentials"))
/// }
/// }
///
/// fn valid_credentials(request: &Request<()>) -> bool {
/// // ...
/// # true
/// }
///
/// fn some_other_interceptor(request: Request<()>) -> Result<Request<()>, Status> {
/// Ok(request)
/// }
///
/// let layer = ServiceBuilder::new()
/// .load_shed()
/// .timeout(Duration::from_secs(30))
/// .layer(interceptor_fn(auth_interceptor))
/// .layer(interceptor_fn(some_other_interceptor))
/// .into_inner();
///
/// Server::builder().layer(layer);
/// ```
///
/// [Tower]: https://github.com/tower-rs/tower
/// [`Layer`]: tower::layer::Layer
/// [eco]: https://github.com/tower-rs
/// [`ServiceBuilder`]: tower::ServiceBuilder
/// [interceptors]: crate::service::interceptor_fn
pub fn layer<NewLayer>(self, new_layer: NewLayer) -> Server<NewLayer> {
Server {
layer: new_layer,
trace_interceptor: self.trace_interceptor,
concurrency_limit: self.concurrency_limit,
timeout: self.timeout,
#[cfg(feature = "tls")]
tls: self.tls,
init_stream_window_size: self.init_stream_window_size,
init_connection_window_size: self.init_connection_window_size,
max_concurrent_streams: self.max_concurrent_streams,
tcp_keepalive: self.tcp_keepalive,
tcp_nodelay: self.tcp_nodelay,
http2_keepalive_interval: self.http2_keepalive_interval,
http2_keepalive_timeout: self.http2_keepalive_timeout,
max_frame_size: self.max_frame_size,
accept_http1: self.accept_http1,
}
}
pub(crate) async fn serve_with_shutdown<S, I, F, IO, IE, ResBody>(
self,
svc: S,
incoming: I,
signal: Option<F>,
) -> Result<(), super::Error>
where
S: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
S::Future: Send + 'static,
S::Error: Into<crate::Error> + Send,
L: Layer<S>,
L::Service: Service<Request<Body>, Response = Response<ResBody>> + Clone + Send + 'static,
<<L as Layer<S>>::Service as Service<Request<Body>>>::Future: Send + 'static,
<<L as Layer<S>>::Service as Service<Request<Body>>>::Error: Into<crate::Error> + Send,
I: Stream<Item = Result<IO, IE>>,
IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
IE: Into<crate::Error>,
F: Future<Output = ()>,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
let trace_interceptor = self.trace_interceptor.clone();
let concurrency_limit = self.concurrency_limit;
@@ -387,7 +475,9 @@ impl Server {
let http2_keepalive_interval = self.http2_keepalive_interval;
let http2_keepalive_timeout = self
.http2_keepalive_timeout
.unwrap_or(Duration::new(DEFAULT_HTTP2_KEEPALIVE_TIMEOUT_SECS, 0));
.unwrap_or_else(|| Duration::new(DEFAULT_HTTP2_KEEPALIVE_TIMEOUT_SECS, 0));
let svc = self.layer.layer(svc);
let tcp = incoming::tcp_incoming(incoming, self);
let incoming = accept::from_stream::<_, _, crate::Error>(tcp);
@@ -422,8 +512,8 @@ impl Server {
}
}
impl<S> Router<S, Unimplemented> {
pub(crate) fn new(server: Server, svc: S) -> Self
impl<S, L> Router<S, Unimplemented, L> {
pub(crate) fn new(server: Server<L>, svc: S) -> Self
where
S: Service<Request<Body>, Response = Response<BoxBody>>
+ NamedService
@@ -447,7 +537,7 @@ impl<S> Router<S, Unimplemented> {
}
}
impl<A, B> Router<A, B>
impl<A, B, L> Router<A, B, L>
where
A: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
A::Future: Send + 'static,
@@ -457,7 +547,7 @@ where
B::Error: Into<crate::Error> + Send,
{
/// Add a new service to this router.
pub fn add_service<S>(self, svc: S) -> Router<S, Or<A, B, Request<Body>>>
pub fn add_service<S>(self, svc: S) -> Router<S, Or<A, B, Request<Body>>, L>
where
S: Service<Request<Body>, Response = Response<BoxBody>>
+ NamedService
@@ -486,10 +576,11 @@ where
/// # Note
/// Even when the argument given is `None` this will capture *all* requests to this service name.
/// As a result, one cannot use this to toggle between two identically named implementations.
#[allow(clippy::type_complexity)]
pub fn add_optional_service<S>(
self,
svc: Option<S>,
) -> Router<Either<S, Unimplemented>, Or<A, B, Request<Body>>>
) -> Router<Either<S, Unimplemented>, Or<A, B, Request<Body>>, L>
where
S: Service<Request<Body>, Response = Response<BoxBody>>
+ NamedService
@@ -518,27 +609,53 @@ where
}
/// Consume this [`Server`] creating a future that will execute the server
/// on [`tokio`]'s default executor.
/// on [tokio]'s default executor.
///
/// [`Server`]: struct.Server.html
pub async fn serve(self, addr: SocketAddr) -> Result<(), super::Error> {
/// [tokio]: https://docs.rs/tokio
pub async fn serve<ResBody>(self, addr: SocketAddr) -> Result<(), super::Error>
where
L: Layer<Routes<A, B, Request<Body>>>,
L::Service: Service<Request<Body>, Response = Response<ResBody>> + Clone + Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Future:
Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Error:
Into<crate::Error> + Send,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
let incoming = TcpIncoming::new(addr, self.server.tcp_nodelay, self.server.tcp_keepalive)
.map_err(super::Error::from_source)?;
self.server
.serve_with_shutdown::<_, _, future::Ready<()>, _, _>(self.routes, incoming, None)
.serve_with_shutdown::<_, _, future::Ready<()>, _, _, ResBody>(
self.routes,
incoming,
None,
)
.await
}
/// Consume this [`Server`] creating a future that will execute the server
/// on [`tokio`]'s default executor. And shutdown when the provided signal
/// on [tokio]'s default executor. And shutdown when the provided signal
/// is received.
///
/// [`Server`]: struct.Server.html
pub async fn serve_with_shutdown<F: Future<Output = ()>>(
/// [tokio]: https://docs.rs/tokio
pub async fn serve_with_shutdown<F: Future<Output = ()>, ResBody>(
self,
addr: SocketAddr,
signal: F,
) -> Result<(), super::Error> {
) -> Result<(), super::Error>
where
L: Layer<Routes<A, B, Request<Body>>>,
L::Service: Service<Request<Body>, Response = Response<ResBody>> + Clone + Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Future:
Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Error:
Into<crate::Error> + Send,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
let incoming = TcpIncoming::new(addr, self.server.tcp_nodelay, self.server.tcp_keepalive)
.map_err(super::Error::from_source)?;
self.server
@@ -550,14 +667,29 @@ where
/// the provided incoming stream of `AsyncRead + AsyncWrite`.
///
/// [`Server`]: struct.Server.html
pub async fn serve_with_incoming<I, IO, IE>(self, incoming: I) -> Result<(), super::Error>
pub async fn serve_with_incoming<I, IO, IE, ResBody>(
self,
incoming: I,
) -> Result<(), super::Error>
where
I: Stream<Item = Result<IO, IE>>,
IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
IE: Into<crate::Error>,
L: Layer<Routes<A, B, Request<Body>>>,
L::Service: Service<Request<Body>, Response = Response<ResBody>> + Clone + Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Future:
Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Error:
Into<crate::Error> + Send,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
self.server
.serve_with_shutdown::<_, _, future::Ready<()>, _, _>(self.routes, incoming, None)
.serve_with_shutdown::<_, _, future::Ready<()>, _, _, ResBody>(
self.routes,
incoming,
None,
)
.await
}
@@ -567,7 +699,7 @@ where
/// gracefully shutdown the server.
///
/// [`Server`]: struct.Server.html
pub async fn serve_with_incoming_shutdown<I, IO, IE, F>(
pub async fn serve_with_incoming_shutdown<I, IO, IE, F, ResBody>(
self,
incoming: I,
signal: F,
@@ -577,6 +709,14 @@ where
IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
IE: Into<crate::Error>,
F: Future<Output = ()>,
L: Layer<Routes<A, B, Request<Body>>>,
L::Service: Service<Request<Body>, Response = Response<ResBody>> + Clone + Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Future:
Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Error:
Into<crate::Error> + Send,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
self.server
.serve_with_shutdown(self.routes, incoming, Some(signal))
@@ -584,12 +724,23 @@ where
}
/// Create a tower service out of a router.
pub fn into_service(self) -> RouterService<A, B> {
RouterService { router: self }
pub fn into_service<ResBody>(self) -> RouterService<L::Service>
where
L: Layer<Routes<A, B, Request<Body>>>,
L::Service: Service<Request<Body>, Response = Response<ResBody>> + Clone + Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Future:
Send + 'static,
<<L as Layer<Routes<A, B, Request<Body>>>>::Service as Service<Request<Body>>>::Error:
Into<crate::Error> + Send,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
let inner = self.server.layer.layer(self.routes);
RouterService { inner }
}
}
impl fmt::Debug for Server {
impl<L> fmt::Debug for Server<L> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Builder").finish()
}
@@ -601,16 +752,16 @@ struct Svc<S> {
conn_info: ConnectionInfo,
}
impl<S> Service<Request<Body>> for Svc<S>
impl<S, ResBody> Service<Request<Body>> for Svc<S>
where
S: Service<Request<Body>, Response = Response<BoxBody>>,
S: Service<Request<Body>, Response = Response<ResBody>>,
S::Error: Into<crate::Error>,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
type Response = Response<BoxBody>;
type Response = Response<BoxHttpBody>;
type Error = crate::Error;
#[allow(clippy::type_complexity)]
type Future = MapErr<Instrumented<S::Future>, fn(S::Error) -> crate::Error>;
type Future = SvcFuture<S::Future>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx).map_err(Into::into)
@@ -633,7 +784,36 @@ where
req.extensions_mut().insert(self.conn_info.clone());
self.inner.call(req).instrument(span).map_err(|e| e.into())
SvcFuture {
inner: self.inner.call(req),
span,
}
}
}
#[pin_project]
struct SvcFuture<F> {
#[pin]
inner: F,
span: tracing::Span,
}
impl<F, E, ResBody> Future for SvcFuture<F>
where
F: Future<Output = Result<Response<ResBody>, E>>,
E: Into<crate::Error>,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
type Output = Result<Response<BoxHttpBody>, crate::Error>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
let _guard = this.span.enter();
let response: Response<ResBody> = ready!(this.inner.poll(cx)).map_err(Into::into)?;
let response = response.map(|body| body.map_err(Into::into).boxed());
Poll::Ready(Ok(response))
}
}
@@ -650,11 +830,13 @@ struct MakeSvc<S> {
trace_interceptor: Option<TraceInterceptor>,
}
impl<S> Service<&ServerIo> for MakeSvc<S>
impl<S, ResBody> Service<&ServerIo> for MakeSvc<S>
where
S: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
S: Service<Request<Body>, Response = Response<ResBody>> + Clone + Send + 'static,
S::Future: Send + 'static,
S::Error: Into<crate::Error> + Send,
ResBody: http_body::Body<Data = Bytes> + Send + Sync + 'static,
ResBody::Error: Into<crate::Error>,
{
type Response = BoxService;
type Error = crate::Error;
@@ -681,11 +863,13 @@ where
.layer_fn(|s| GrpcTimeout::new(s, timeout))
.service(svc);
let svc = BoxService::new(Svc {
let svc = Svc {
inner: svc,
trace_interceptor,
conn_info,
});
};
let svc = BoxService::new(svc);
future::ready(Ok(svc))
}
tonic/src/transport/server/recover_error.rs
+71 -10
View File
@@ -1,4 +1,7 @@
use crate::{body::BoxBody, Status};
use crate::{
util::{OptionPin, OptionPinProj},
Status,
};
use futures_util::ready;
use http::Response;
use pin_project::pin_project;
@@ -22,12 +25,12 @@ impl<S> RecoverError<S> {
}
}
impl<S, R> Service<R> for RecoverError<S>
impl<S, R, ResBody> Service<R> for RecoverError<S>
where
S: Service<R, Response = Response<BoxBody>>,
S: Service<R, Response = Response<ResBody>>,
S::Error: Into<crate::Error>,
{
type Response = Response<BoxBody>;
type Response = Response<MaybeEmptyBody<ResBody>>;
type Error = crate::Error;
type Future = ResponseFuture<S::Future>;
@@ -48,22 +51,25 @@ pub(crate) struct ResponseFuture<F> {
inner: F,
}
impl<F, E> Future for ResponseFuture<F>
impl<F, E, ResBody> Future for ResponseFuture<F>
where
F: Future<Output = Result<Response<BoxBody>, E>>,
F: Future<Output = Result<Response<ResBody>, E>>,
E: Into<crate::Error>,
{
type Output = Result<Response<BoxBody>, crate::Error>;
type Output = Result<Response<MaybeEmptyBody<ResBody>>, crate::Error>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let result: Result<Response<BoxBody>, crate::Error> =
let result: Result<Response<_>, crate::Error> =
ready!(self.project().inner.poll(cx)).map_err(Into::into);
match result {
Ok(res) => Poll::Ready(Ok(res)),
Ok(response) => {
let response = response.map(MaybeEmptyBody::full);
Poll::Ready(Ok(response))
}
Err(err) => {
if let Some(status) = Status::try_from_error(&*err) {
let mut res = Response::new(crate::body::empty_body());
let mut res = Response::new(MaybeEmptyBody::empty());
status.add_header(res.headers_mut()).unwrap();
Poll::Ready(Ok(res))
} else {
@@ -73,3 +79,58 @@ where
}
}
}
#[pin_project]
pub(crate) struct MaybeEmptyBody<B> {
#[pin]
inner: OptionPin<B>,
}
impl<B> MaybeEmptyBody<B> {
fn full(inner: B) -> Self {
Self {
inner: OptionPin::Some(inner),
}
}
fn empty() -> Self {
Self {
inner: OptionPin::None,
}
}
}
impl<B> http_body::Body for MaybeEmptyBody<B>
where
B: http_body::Body + Send,
{
type Data = B::Data;
type Error = B::Error;
fn poll_data(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
match self.project().inner.project() {
OptionPinProj::Some(b) => b.poll_data(cx),
OptionPinProj::None => Poll::Ready(None),
}
}
fn poll_trailers(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> {
match self.project().inner.project() {
OptionPinProj::Some(b) => b.poll_trailers(cx),
OptionPinProj::None => Poll::Ready(Ok(None)),
}
}
fn is_end_stream(&self) -> bool {
match &self.inner {
OptionPin::Some(b) => b.is_end_stream(),
OptionPin::None => true,
}
}
}
tonic/src/transport/service/grpc_timeout.rs
+1 -6
View File
@@ -1,4 +1,5 @@
use crate::metadata::GRPC_TIMEOUT_HEADER;
use crate::util::{OptionPin, OptionPinProj};
use http::{HeaderMap, HeaderValue, Request};
use pin_project::pin_project;
use std::{
@@ -97,12 +98,6 @@ where
}
}
#[pin_project(project = OptionPinProj)]
enum OptionPin<T> {
Some(#[pin] T),
None,
}
const SECONDS_IN_HOUR: u64 = 60 * 60;
const SECONDS_IN_MINUTE: u64 = 60;
tonic/src/transport/service/router.rs
+2 -1
View File
@@ -9,8 +9,9 @@ use std::{
};
use tower_service::Service;
#[doc(hidden)]
#[derive(Debug)]
pub(crate) struct Routes<A, B, Request> {
pub struct Routes<A, B, Request> {
routes: Or<A, B, Request>,
}
tonic/src/util.rs
+13
View File
@@ -0,0 +1,13 @@
//! Various utilities used throughout tonic.
// some combinations of features might cause things here not to be used
#![allow(dead_code)]
use pin_project::pin_project;
/// A pin-project compatible `Option`
#[pin_project(project = OptionPinProj)]
pub(crate) enum OptionPin<T> {
Some(#[pin] T),
None,
}