Skip to content

Commit

Permalink
chore(server): Refactor recover error service (#2100)
Browse files Browse the repository at this point in the history
  • Loading branch information
tottoto authored Dec 14, 2024
1 parent 36bc21c commit 40fb659
Showing 1 changed file with 20 additions and 22 deletions.
42 changes: 20 additions & 22 deletions tonic/src/transport/server/service/recover_error.rs
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
use crate::Status;
use http::Response;
use http_body::Frame;
use pin_project::pin_project;
use std::{
future::Future,
pin::Pin,
task::{ready, Context, Poll},
};
use tower::Service;

use http::Response;
use pin_project::pin_project;
use tower_service::Service;

use crate::Status;

/// Middleware that attempts to recover from service errors by turning them into a response built
/// from the `Status`.
Expand All @@ -22,20 +23,20 @@ impl<S> RecoverError<S> {
}
}

impl<S, R, ResBody> Service<R> for RecoverError<S>
impl<S, Req, ResBody> Service<Req> for RecoverError<S>
where
S: Service<R, Response = Response<ResBody>>,
S: Service<Req, Response = Response<ResBody>>,
S::Error: Into<crate::BoxError>,
{
type Response = Response<MaybeEmptyBody<ResBody>>;
type Response = Response<ResponseBody<ResBody>>;
type Error = crate::BoxError;
type Future = ResponseFuture<S::Future>;

fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx).map_err(Into::into)
}

fn call(&mut self, req: R) -> Self::Future {
fn call(&mut self, req: Req) -> Self::Future {
ResponseFuture {
inner: self.inner.call(req),
}
Expand All @@ -53,21 +54,18 @@ where
F: Future<Output = Result<Response<ResBody>, E>>,
E: Into<crate::BoxError>,
{
type Output = Result<Response<MaybeEmptyBody<ResBody>>, crate::BoxError>;
type Output = Result<Response<ResponseBody<ResBody>>, crate::BoxError>;

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let result: Result<Response<_>, crate::BoxError> =
ready!(self.project().inner.poll(cx)).map_err(Into::into);

match result {
match ready!(self.project().inner.poll(cx)) {
Ok(response) => {
let response = response.map(MaybeEmptyBody::full);
let response = response.map(ResponseBody::full);
Poll::Ready(Ok(response))
}
Err(err) => match Status::try_from_error(err) {
Err(err) => match Status::try_from_error(err.into()) {
Ok(status) => {
let (parts, ()) = status.into_http::<()>().into_parts();
let res = Response::from_parts(parts, MaybeEmptyBody::empty());
let res = Response::from_parts(parts, ResponseBody::empty());
Poll::Ready(Ok(res))
}
Err(err) => Poll::Ready(Err(err)),
Expand All @@ -77,22 +75,22 @@ where
}

#[pin_project]
pub(crate) struct MaybeEmptyBody<B> {
pub(crate) struct ResponseBody<B> {
#[pin]
inner: Option<B>,
}

impl<B> MaybeEmptyBody<B> {
impl<B> ResponseBody<B> {
fn full(inner: B) -> Self {
Self { inner: Some(inner) }
}

fn empty() -> Self {
const fn empty() -> Self {
Self { inner: None }
}
}

impl<B> http_body::Body for MaybeEmptyBody<B>
impl<B> http_body::Body for ResponseBody<B>
where
B: http_body::Body,
{
Expand All @@ -102,7 +100,7 @@ where
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
match self.project().inner.as_pin_mut() {
Some(b) => b.poll_frame(cx),
None => Poll::Ready(None),
Expand Down

0 comments on commit 40fb659

Please sign in to comment.