anvil/tasks/
block_listener.rs

1//! A task that listens for new blocks
2
3use crate::shutdown::Shutdown;
4use futures::{FutureExt, Stream, StreamExt};
5use std::{
6    pin::Pin,
7    task::{Context, Poll},
8};
9
10/// A future that executes a given `task` for the latest block available on the stream.
11///
12/// Available blocks are coalesced and the latest block is processed after the active task
13/// completes.
14pub struct BlockListener<St, F, Fut> {
15    stream: St,
16    task_factory: F,
17    task: Option<Pin<Box<Fut>>>,
18    stream_done: bool,
19    on_shutdown: Shutdown,
20}
21
22impl<St, F, Fut> BlockListener<St, F, Fut>
23where
24    St: Stream,
25    F: Fn(<St as Stream>::Item) -> Fut,
26{
27    pub const fn new(on_shutdown: Shutdown, block_stream: St, task_factory: F) -> Self {
28        Self { stream: block_stream, task_factory, task: None, stream_done: false, on_shutdown }
29    }
30}
31
32impl<St, F, Fut> Future for BlockListener<St, F, Fut>
33where
34    St: Stream + Unpin,
35    F: Fn(<St as Stream>::Item) -> Fut + Unpin + Send + Sync + 'static,
36    Fut: Future<Output = ()> + Send,
37{
38    type Output = ();
39
40    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
41        let pin = self.get_mut();
42
43        if pin.on_shutdown.poll_unpin(cx).is_ready() {
44            return Poll::Ready(());
45        }
46
47        if let Some(mut task) = pin.task.take()
48            && task.poll_unpin(cx).is_pending()
49        {
50            pin.task = Some(task);
51            return Poll::Pending;
52        }
53
54        if pin.stream_done {
55            return Poll::Ready(());
56        }
57
58        let mut block = None;
59        // drain the stream
60        loop {
61            match pin.stream.poll_next_unpin(cx) {
62                Poll::Ready(Some(next_block)) => block = Some(next_block),
63                Poll::Ready(None) => {
64                    pin.stream_done = true;
65                    break;
66                }
67                Poll::Pending => break,
68            }
69        }
70
71        if let Some(block) = block {
72            pin.task = Some(Box::pin((pin.task_factory)(block)));
73        }
74
75        if let Some(mut task) = pin.task.take()
76            && task.poll_unpin(cx).is_pending()
77        {
78            pin.task = Some(task);
79        }
80
81        if pin.stream_done && pin.task.is_none() { Poll::Ready(()) } else { Poll::Pending }
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::shutdown;
89    use futures::{
90        channel::{mpsc, oneshot},
91        task::noop_waker,
92    };
93    use parking_lot::Mutex;
94    use std::sync::Arc;
95
96    #[test]
97    fn waits_for_active_task_before_processing_latest_block() {
98        let (_signal, shutdown) = shutdown::signal();
99        let (tx, rx) = mpsc::unbounded();
100        let (release_tx, release_rx) = oneshot::channel();
101        let release_rx = Arc::new(Mutex::new(Some(release_rx)));
102        let processed = Arc::new(Mutex::new(Vec::new()));
103        let task_release_rx = release_rx;
104        let task_processed = processed.clone();
105        let mut listener = Box::pin(BlockListener::new(shutdown, rx, move |block| {
106            let release_rx = task_release_rx.clone();
107            let processed = task_processed.clone();
108            async move {
109                let release_rx = (block == 1).then(|| release_rx.lock().take().unwrap());
110                if let Some(release_rx) = release_rx {
111                    let _ = release_rx.await;
112                }
113                processed.lock().push(block);
114            }
115        }));
116        let waker = noop_waker();
117        let mut cx = Context::from_waker(&waker);
118
119        tx.unbounded_send(1).unwrap();
120        assert!(listener.as_mut().poll(&mut cx).is_pending());
121
122        tx.unbounded_send(2).unwrap();
123        tx.unbounded_send(3).unwrap();
124        assert!(listener.as_mut().poll(&mut cx).is_pending());
125        assert!(processed.lock().is_empty());
126
127        drop(tx);
128        release_tx.send(()).unwrap();
129        assert!(listener.as_mut().poll(&mut cx).is_ready());
130        assert_eq!(*processed.lock(), [1, 3]);
131    }
132}
anvil/tasks/block_listener.rs

anvil/tasks/
block_listener.rs
