reth_node_core/
exit.rs

1//! Helper types for waiting for the node to exit.
2
3use futures::{future::BoxFuture, FutureExt};
4use std::{
5    fmt,
6    future::Future,
7    pin::Pin,
8    task::{ready, Context, Poll},
9};
10
11/// A Future which resolves when the node exits
12pub struct NodeExitFuture {
13    /// The consensus engine future.
14    /// This can be polled to wait for the consensus engine to exit.
15    consensus_engine_fut: Option<BoxFuture<'static, eyre::Result<()>>>,
16}
17
18impl fmt::Debug for NodeExitFuture {
19    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
20        f.debug_struct("NodeExitFuture").field("consensus_engine_fut", &"...").finish()
21    }
22}
23
24impl NodeExitFuture {
25    /// Create a new `NodeExitFuture`.
26    pub fn new<F>(consensus_engine_fut: F) -> Self
27    where
28        F: Future<Output = eyre::Result<()>> + 'static + Send,
29    {
30        Self { consensus_engine_fut: Some(Box::pin(consensus_engine_fut)) }
31    }
32}
33
34impl Future for NodeExitFuture {
35    type Output = eyre::Result<()>;
36
37    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
38        let this = self.get_mut();
39        if let Some(rx) = this.consensus_engine_fut.as_mut() {
40            match ready!(rx.poll_unpin(cx)) {
41                Ok(_) => {
42                    this.consensus_engine_fut.take();
43                    Poll::Ready(Ok(()))
44                }
45                Err(err) => Poll::Ready(Err(err)),
46            }
47        } else {
48            Poll::Pending
49        }
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[tokio::test]
58    async fn test_node_exit_future() {
59        let fut = async { Ok(()) };
60
61        let node_exit_future = NodeExitFuture::new(fut);
62
63        let res = node_exit_future.await;
64
65        assert!(res.is_ok());
66    }
67}
reth_node_core/exit.rs

reth_node_core/
exit.rs
