use std::future::Future; use std::marker::PhantomData; use std::pin::Pin; use std::task::{Context, Poll}; use super::{Service, ServiceFactory}; /// Service for the `map` combinator, changing the type of a service's response. /// /// This is created by the `ServiceExt::map` method. pub struct Map { service: A, f: F, _t: PhantomData, } impl Map { /// Create new `Map` combinator pub(crate) fn new(service: A, f: F) -> Self where A: Service, F: FnMut(A::Response) -> Response, { Self { service, f, _t: PhantomData, } } } impl Clone for Map where A: Clone, F: Clone, { fn clone(&self) -> Self { Map { service: self.service.clone(), f: self.f.clone(), _t: PhantomData, } } } impl Service for Map where A: Service, F: FnMut(A::Response) -> Response + Clone, { type Request = A::Request; type Response = Response; type Error = A::Error; type Future = MapFuture; fn poll_ready(&mut self, ctx: &mut Context<'_>) -> Poll> { self.service.poll_ready(ctx) } fn call(&mut self, req: A::Request) -> Self::Future { MapFuture::new(self.service.call(req), self.f.clone()) } } #[pin_project::pin_project] pub struct MapFuture where A: Service, F: FnMut(A::Response) -> Response, { f: F, #[pin] fut: A::Future, } impl MapFuture where A: Service, F: FnMut(A::Response) -> Response, { fn new(fut: A::Future, f: F) -> Self { MapFuture { f, fut } } } impl Future for MapFuture where A: Service, F: FnMut(A::Response) -> Response, { type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); match this.fut.poll(cx) { Poll::Ready(Ok(resp)) => Poll::Ready(Ok((this.f)(resp))), Poll::Ready(Err(e)) => Poll::Ready(Err(e)), Poll::Pending => Poll::Pending, } } } /// `MapNewService` new service combinator pub struct MapServiceFactory { a: A, f: F, r: PhantomData, } impl MapServiceFactory { /// Create new `Map` new service instance pub(crate) fn new(a: A, f: F) -> Self where A: ServiceFactory, F: FnMut(A::Response) -> Res, { Self { a, f, r: PhantomData, } } } impl Clone for MapServiceFactory where A: Clone, F: Clone, { fn clone(&self) -> Self { Self { a: self.a.clone(), f: self.f.clone(), r: PhantomData, } } } impl ServiceFactory for MapServiceFactory where A: ServiceFactory, F: FnMut(A::Response) -> Res + Clone, { type Request = A::Request; type Response = Res; type Error = A::Error; type Config = A::Config; type Service = Map; type InitError = A::InitError; type Future = MapServiceFuture; fn new_service(&self, cfg: A::Config) -> Self::Future { MapServiceFuture::new(self.a.new_service(cfg), self.f.clone()) } } #[pin_project::pin_project] pub struct MapServiceFuture where A: ServiceFactory, F: FnMut(A::Response) -> Res, { #[pin] fut: A::Future, f: Option, } impl MapServiceFuture where A: ServiceFactory, F: FnMut(A::Response) -> Res, { fn new(fut: A::Future, f: F) -> Self { MapServiceFuture { f: Some(f), fut } } } impl Future for MapServiceFuture where A: ServiceFactory, F: FnMut(A::Response) -> Res, { type Output = Result, A::InitError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); if let Poll::Ready(svc) = this.fut.poll(cx)? { Poll::Ready(Ok(Map::new(svc, this.f.take().unwrap()))) } else { Poll::Pending } } } #[cfg(test)] mod tests { use futures::future::{lazy, ok, Ready}; use super::*; use crate::{IntoServiceFactory, Service, ServiceFactory}; struct Srv; impl Service for Srv { type Request = (); type Response = (); type Error = (); type Future = Ready>; fn poll_ready(&mut self, _: &mut Context) -> Poll> { Poll::Ready(Ok(())) } fn call(&mut self, _: ()) -> Self::Future { ok(()) } } #[actix_rt::test] async fn test_poll_ready() { let mut srv = Srv.map(|_| "ok"); let res = lazy(|cx| srv.poll_ready(cx)).await; assert_eq!(res, Poll::Ready(Ok(()))); } #[actix_rt::test] async fn test_call() { let mut srv = Srv.map(|_| "ok"); let res = srv.call(()).await; assert!(res.is_ok()); assert_eq!(res.unwrap(), "ok"); } #[actix_rt::test] async fn test_new_service() { let new_srv = (|| ok::<_, ()>(Srv)).into_factory().map(|_| "ok"); let mut srv = new_srv.new_service(&()).await.unwrap(); let res = srv.call(()).await; assert!(res.is_ok()); assert_eq!(res.unwrap(), ("ok")); } }