puzz-middleware 0.1.0

Middleware for puzz.
Documentation
use std::convert::Infallible;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use futures_core::ready;
use pin_project_lite::pin_project;
use puzz_core::service::{Service, Wrap};

pub fn handle_error<F>(f: F) -> HandleErrorWrap<F> {
    HandleErrorWrap::new(f)
}

#[derive(Clone, Copy)]
pub struct HandleErrorWrap<F> {
    f: F,
}

impl<F> HandleErrorWrap<F> {
    pub fn new(f: F) -> Self {
        Self { f }
    }
}

impl<S, F> Wrap<S> for HandleErrorWrap<F> {
    type Service = HandleError<S, F>;

    fn wrap(self, service: S) -> Self::Service {
        HandleError {
            inner: service,
            f: self.f,
        }
    }
}

impl<F> fmt::Debug for HandleErrorWrap<F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("HandleErrorWrap")
            .field("f", &core::any::type_name::<F>())
            .finish()
    }
}

#[derive(Clone, Copy)]
pub struct HandleError<S, F> {
    inner: S,
    f: F,
}

impl<S, F, Req> Service<Req> for HandleError<S, F>
where
    S: Service<Req>,
    F: FnOnce(S::Error) -> S::Response + Clone,
{
    type Response = S::Response;
    type Error = Infallible;
    type Future = HandleErrorFuture<S::Future, F>;

    fn call(&self, request: Req) -> Self::Future {
        HandleErrorFuture::Incomplete {
            fut: self.inner.call(request),
            f: self.f.clone(),
        }
    }
}

impl<S, F> fmt::Debug for HandleError<S, F>
where
    S: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("HandleError")
            .field("inner", &self.inner)
            .field("f", &core::any::type_name::<F>())
            .finish()
    }
}

pin_project! {
    #[project = HandleErrorFutureProj]
    #[project_replace = HandleErrorFutureProjReplace]
    pub enum HandleErrorFuture<Fut, F> {
        Incomplete {
            #[pin]
            fut: Fut,
            f: F,
        },
        Complete,
    }
}

impl<Fut, F, Res, Err> Future for HandleErrorFuture<Fut, F>
where
    Fut: Future<Output = Result<Res, Err>>,
    F: FnOnce(Err) -> Res,
{
    type Output = Result<Res, Infallible>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.as_mut().project() {
            HandleErrorFutureProj::Incomplete { fut, .. } => {
                let output = ready!(fut.poll(cx));
                match self.project_replace(HandleErrorFuture::Complete) {
                    HandleErrorFutureProjReplace::Incomplete { f, .. } => match output {
                        Ok(res) => Poll::Ready(Ok(res)),
                        Err(err) => Poll::Ready(Ok(f(err))),
                    },
                    HandleErrorFutureProjReplace::Complete => unreachable!(),
                }
            }
            HandleErrorFutureProj::Complete => {
                panic!("polled after completion")
            }
        }
    }
}