tower_batch_control/
future.rs

Help
1//! Future types for the `Batch` middleware.
2
3use super::{error::Closed, message};
4use futures_core::ready;
5use pin_project::pin_project;
6use std::{
7    future::Future,
8    pin::Pin,
9    task::{Context, Poll},
10};
11
12/// Future that completes when the batch processing is complete.
13#[pin_project]
14#[derive(Debug)]
15pub struct ResponseFuture<T> {
16    #[pin]
17    state: ResponseState<T>,
18}
19
20#[pin_project(project = ResponseStateProj)]
21#[derive(Debug)]
22enum ResponseState<T> {
23    Failed(Option<crate::BoxError>),
24    Rx(#[pin] message::Rx<T>),
25    Poll(#[pin] T),
26}
27
28impl<T> ResponseFuture<T> {
29    pub(crate) fn new(rx: message::Rx<T>) -> Self {
30        ResponseFuture {
31            state: ResponseState::Rx(rx),
32        }
33    }
34
35    pub(crate) fn failed(err: crate::BoxError) -> Self {
36        ResponseFuture {
37            state: ResponseState::Failed(Some(err)),
38        }
39    }
40}
41
42impl<F, T, E> Future for ResponseFuture<F>
43where
44    F: Future<Output = Result<T, E>>,
45    E: Into<crate::BoxError>,
46{
47    type Output = Result<T, crate::BoxError>;
48
49    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
50        let mut this = self.project();
51
52        // CORRECTNESS
53        //
54        // The current task must be scheduled for wakeup every time we return
55        // `Poll::Pending`.
56        //
57        // This loop ensures that the task is scheduled as required, because it
58        // only returns Pending when another future returns Pending.
59        loop {
60            match this.state.as_mut().project() {
61                ResponseStateProj::Failed(e) => {
62                    return Poll::Ready(Err(e.take().expect("polled after error")));
63                }
64                ResponseStateProj::Rx(rx) => match ready!(rx.poll(cx)) {
65                    Ok(Ok(f)) => this.state.set(ResponseState::Poll(f)),
66                    Ok(Err(e)) => return Poll::Ready(Err(e.into())),
67                    Err(_) => return Poll::Ready(Err(Closed::new().into())),
68                },
69                ResponseStateProj::Poll(fut) => return fut.poll(cx).map_err(Into::into),
70            }
71        }
72    }
73}
tower_batch_control/future.rs

tower_batch_control/
future.rs
