use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use super::{Service, ServiceFactory};
use crate::cell::Cell;
/// Service for the `and_then` combinator, chaining a computation onto the end
/// of another service which completes successfully.
///
/// This is created by the `ServiceExt::and_then` method.
pub struct AndThenService(Cell, Cell);
impl AndThenService {
/// Create new `AndThen` combinator
pub(crate) fn new(a: A, b: B) -> Self
where
A: Service,
B: Service,
{
Self(Cell::new(a), Cell::new(b))
}
}
impl Clone for AndThenService {
fn clone(&self) -> Self {
AndThenService(self.0.clone(), self.1.clone())
}
}
impl Service for AndThenService
where
A: Service,
B: Service,
{
type Request = A::Request;
type Response = B::Response;
type Error = A::Error;
type Future = AndThenServiceResponse;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> {
let not_ready = { !self.0.get_mut().poll_ready(cx)?.is_ready() };
if !self.1.poll_ready(cx)?.is_ready() || not_ready {
Poll::Pending
} else {
Poll::Ready(Ok(()))
}
}
fn call(&mut self, req: A::Request) -> Self::Future {
AndThenServiceResponse {
state: State::A(self.0.get_mut().call(req), Some(self.1.clone())),
}
}
}
#[pin_project::pin_project]
pub struct AndThenServiceResponse
where
A: Service,
B: Service,
{
#[pin]
state: State,
}
#[pin_project::pin_project]
enum State
where
A: Service,
B: Service,
{
A(#[pin] A::Future, Option>),
B(#[pin] B::Future),
Empty,
}
impl Future for AndThenServiceResponse
where
A: Service,
B: Service,
{
type Output = Result;
#[pin_project::project]
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll {
let mut this = self.as_mut().project();
#[project]
match this.state.as_mut().project() {
State::A(fut, b) => match fut.poll(cx)? {
Poll::Ready(res) => {
let mut b = b.take().unwrap();
this.state.set(State::Empty); // drop fut A
let fut = b.get_mut().call(res);
this.state.set(State::B(fut));
self.poll(cx)
}
Poll::Pending => Poll::Pending,
},
State::B(fut) => fut.poll(cx).map(|r| {
this.state.set(State::Empty);
r
}),
State::Empty => panic!("future must not be polled after it returned `Poll::Ready`"),
}
}
}
/// `.and_then()` service factory combinator
pub struct AndThenServiceFactory
where
A: ServiceFactory,
A::Config: Clone,
B: ServiceFactory<
Config = A::Config,
Request = A::Response,
Error = A::Error,
InitError = A::InitError,
>,
{
a: A,
b: B,
}
impl AndThenServiceFactory
where
A: ServiceFactory,
A::Config: Clone,
B: ServiceFactory<
Config = A::Config,
Request = A::Response,
Error = A::Error,
InitError = A::InitError,
>,
{
/// Create new `AndThenFactory` combinator
pub(crate) fn new(a: A, b: B) -> Self {
Self { a, b }
}
}
impl ServiceFactory for AndThenServiceFactory
where
A: ServiceFactory,
A::Config: Clone,
B: ServiceFactory<
Config = A::Config,
Request = A::Response,
Error = A::Error,
InitError = A::InitError,
>,
{
type Request = A::Request;
type Response = B::Response;
type Error = A::Error;
type Config = A::Config;
type Service = AndThenService;
type InitError = A::InitError;
type Future = AndThenServiceFactoryResponse;
fn new_service(&self, cfg: A::Config) -> Self::Future {
AndThenServiceFactoryResponse::new(
self.a.new_service(cfg.clone()),
self.b.new_service(cfg),
)
}
}
impl Clone for AndThenServiceFactory
where
A: ServiceFactory + Clone,
A::Config: Clone,
B: ServiceFactory<
Config = A::Config,
Request = A::Response,
Error = A::Error,
InitError = A::InitError,
> + Clone,
{
fn clone(&self) -> Self {
Self {
a: self.a.clone(),
b: self.b.clone(),
}
}
}
#[pin_project::pin_project]
pub struct AndThenServiceFactoryResponse
where
A: ServiceFactory,
B: ServiceFactory,
{
#[pin]
fut_a: A::Future,
#[pin]
fut_b: B::Future,
a: Option,
b: Option,
}
impl AndThenServiceFactoryResponse
where
A: ServiceFactory,
B: ServiceFactory,
{
fn new(fut_a: A::Future, fut_b: B::Future) -> Self {
AndThenServiceFactoryResponse {
fut_a,
fut_b,
a: None,
b: None,
}
}
}
impl Future for AndThenServiceFactoryResponse
where
A: ServiceFactory,
B: ServiceFactory,
{
type Output = Result, A::InitError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll {
let this = self.project();
if this.a.is_none() {
if let Poll::Ready(service) = this.fut_a.poll(cx)? {
*this.a = Some(service);
}
}
if this.b.is_none() {
if let Poll::Ready(service) = this.fut_b.poll(cx)? {
*this.b = Some(service);
}
}
if this.a.is_some() && this.b.is_some() {
Poll::Ready(Ok(AndThenService::new(
this.a.take().unwrap(),
this.b.take().unwrap(),
)))
} else {
Poll::Pending
}
}
}
#[cfg(test)]
mod tests {
use std::cell::Cell;
use std::rc::Rc;
use std::task::{Context, Poll};
use futures_util::future::{lazy, ok, ready, Ready};
use crate::{fn_factory, pipeline, pipeline_factory, Service, ServiceFactory};
struct Srv1(Rc>);
impl Service for Srv1 {
type Request = &'static str;
type Response = &'static str;
type Error = ();
type Future = Ready>;
fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> {
self.0.set(self.0.get() + 1);
Poll::Ready(Ok(()))
}
fn call(&mut self, req: &'static str) -> Self::Future {
ok(req)
}
}
#[derive(Clone)]
struct Srv2(Rc>);
impl Service for Srv2 {
type Request = &'static str;
type Response = (&'static str, &'static str);
type Error = ();
type Future = Ready>;
fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> {
self.0.set(self.0.get() + 1);
Poll::Ready(Ok(()))
}
fn call(&mut self, req: &'static str) -> Self::Future {
ok((req, "srv2"))
}
}
#[actix_rt::test]
async fn test_poll_ready() {
let cnt = Rc::new(Cell::new(0));
let mut srv = pipeline(Srv1(cnt.clone())).and_then(Srv2(cnt.clone()));
let res = lazy(|cx| srv.poll_ready(cx)).await;
assert_eq!(res, Poll::Ready(Ok(())));
assert_eq!(cnt.get(), 2);
}
#[actix_rt::test]
async fn test_call() {
let cnt = Rc::new(Cell::new(0));
let mut srv = pipeline(Srv1(cnt.clone())).and_then(Srv2(cnt));
let res = srv.call("srv1").await;
assert!(res.is_ok());
assert_eq!(res.unwrap(), (("srv1", "srv2")));
}
#[actix_rt::test]
async fn test_new_service() {
let cnt = Rc::new(Cell::new(0));
let cnt2 = cnt.clone();
let new_srv =
pipeline_factory(fn_factory(move || ready(Ok::<_, ()>(Srv1(cnt2.clone())))))
.and_then(move || ready(Ok(Srv2(cnt.clone()))));
let mut srv = new_srv.new_service(()).await.unwrap();
let res = srv.call("srv1").await;
assert!(res.is_ok());
assert_eq!(res.unwrap(), ("srv1", "srv2"));
}
}
| | |