use core::{
future::Future,
marker::PhantomData,
pin::Pin,
task::{Context, Poll},
};
use pin_project_lite::pin_project;
use super::{Service, ServiceFactory};
pub struct MapErr<S, Req, F, E> {
service: S,
mapper: F,
_t: PhantomData<fn(Req) -> E>,
}
impl<S, Req, F, E> MapErr<S, Req, F, E> {
pub(crate) fn new(service: S, mapper: F) -> Self
where
S: Service<Req>,
F: Fn(S::Error) -> E,
{
Self {
service,
mapper,
_t: PhantomData,
}
}
}
impl<S, Req, F, E> Clone for MapErr<S, Req, F, E>
where
S: Clone,
F: Clone,
{
fn clone(&self) -> Self {
MapErr {
service: self.service.clone(),
mapper: self.mapper.clone(),
_t: PhantomData,
}
}
}
impl<A, Req, F, E> Service<Req> for MapErr<A, Req, F, E>
where
A: Service<Req>,
F: Fn(A::Error) -> E + Clone,
{
type Response = A::Response;
type Error = E;
type Future = MapErrFuture<A, Req, F, E>;
fn poll_ready(&self, ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.service.poll_ready(ctx).map_err(&self.mapper)
}
fn call(&self, req: Req) -> Self::Future {
MapErrFuture::new(self.service.call(req), self.mapper.clone())
}
}
pin_project! {
pub struct MapErrFuture<A, Req, F, E>
where
A: Service<Req>,
F: Fn(A::Error) -> E,
{
f: F,
#[pin]
fut: A::Future,
}
}
impl<A, Req, F, E> MapErrFuture<A, Req, F, E>
where
A: Service<Req>,
F: Fn(A::Error) -> E,
{
fn new(fut: A::Future, f: F) -> Self {
MapErrFuture { f, fut }
}
}
impl<A, Req, F, E> Future for MapErrFuture<A, Req, F, E>
where
A: Service<Req>,
F: Fn(A::Error) -> E,
{
type Output = Result<A::Response, E>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
this.fut.poll(cx).map_err(this.f)
}
}
pub struct MapErrServiceFactory<SF, Req, F, E>
where
SF: ServiceFactory<Req>,
F: Fn(SF::Error) -> E + Clone,
{
a: SF,
f: F,
e: PhantomData<fn(Req) -> E>,
}
impl<SF, Req, F, E> MapErrServiceFactory<SF, Req, F, E>
where
SF: ServiceFactory<Req>,
F: Fn(SF::Error) -> E + Clone,
{
pub(crate) fn new(a: SF, f: F) -> Self {
Self {
a,
f,
e: PhantomData,
}
}
}
impl<SF, Req, F, E> Clone for MapErrServiceFactory<SF, Req, F, E>
where
SF: ServiceFactory<Req> + Clone,
F: Fn(SF::Error) -> E + Clone,
{
fn clone(&self) -> Self {
Self {
a: self.a.clone(),
f: self.f.clone(),
e: PhantomData,
}
}
}
impl<SF, Req, F, E> ServiceFactory<Req> for MapErrServiceFactory<SF, Req, F, E>
where
SF: ServiceFactory<Req>,
F: Fn(SF::Error) -> E + Clone,
{
type Response = SF::Response;
type Error = E;
type Config = SF::Config;
type Service = MapErr<SF::Service, Req, F, E>;
type InitError = SF::InitError;
type Future = MapErrServiceFuture<SF, Req, F, E>;
fn new_service(&self, cfg: SF::Config) -> Self::Future {
MapErrServiceFuture::new(self.a.new_service(cfg), self.f.clone())
}
}
pin_project! {
pub struct MapErrServiceFuture<SF, Req, F, E>
where
SF: ServiceFactory<Req>,
F: Fn(SF::Error) -> E,
{
#[pin]
fut: SF::Future,
mapper: F,
}
}
impl<SF, Req, F, E> MapErrServiceFuture<SF, Req, F, E>
where
SF: ServiceFactory<Req>,
F: Fn(SF::Error) -> E,
{
fn new(fut: SF::Future, mapper: F) -> Self {
MapErrServiceFuture { fut, mapper }
}
}
impl<SF, Req, F, E> Future for MapErrServiceFuture<SF, Req, F, E>
where
SF: ServiceFactory<Req>,
F: Fn(SF::Error) -> E + Clone,
{
type Output = Result<MapErr<SF::Service, Req, F, E>, SF::InitError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
if let Poll::Ready(svc) = this.fut.poll(cx)? {
Poll::Ready(Ok(MapErr::new(svc, this.mapper.clone())))
} else {
Poll::Pending
}
}
}
#[cfg(test)]
mod tests {
use futures_util::future::lazy;
use super::*;
use crate::{
err, ok, IntoServiceFactory, Ready, Service, ServiceExt, ServiceFactory,
ServiceFactoryExt,
};
struct Srv;
impl Service<()> for Srv {
type Response = ();
type Error = ();
type Future = Ready<Result<(), ()>>;
fn poll_ready(&self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Err(()))
}
fn call(&self, _: ()) -> Self::Future {
err(())
}
}
#[actix_rt::test]
async fn test_poll_ready() {
let srv = Srv.map_err(|_| "error");
let res = lazy(|cx| srv.poll_ready(cx)).await;
assert_eq!(res, Poll::Ready(Err("error")));
}
#[actix_rt::test]
async fn test_call() {
let srv = Srv.map_err(|_| "error");
let res = srv.call(()).await;
assert!(res.is_err());
assert_eq!(res.err().unwrap(), "error");
}
#[actix_rt::test]
async fn test_new_service() {
let new_srv = (|| ok::<_, ()>(Srv)).into_factory().map_err(|_| "error");
let srv = new_srv.new_service(&()).await.unwrap();
let res = srv.call(()).await;
assert!(res.is_err());
assert_eq!(res.err().unwrap(), "error");
}
}