linera_sdk/
util.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Module with helper types and functions used by the SDK.
5
6use std::{
7    future::Future,
8    pin::pin,
9    task::{Context, Poll},
10};
11
12use futures::task;
13
14/// An extension trait to block on a [`Future`] until it completes.
15pub trait BlockingWait {
16    /// The type returned by the [`Future`].
17    type Output;
18
19    /// Waits for the [`Future`] to complete in a blocking manner.
20    ///
21    /// Effectively polls the [`Future`] repeatedly until it returns [`Poll::Ready`].
22    fn blocking_wait(self) -> Self::Output;
23}
24
25impl<AnyFuture> BlockingWait for AnyFuture
26where
27    AnyFuture: Future,
28{
29    type Output = AnyFuture::Output;
30
31    fn blocking_wait(self) -> Self::Output {
32        let mut context = Context::from_waker(task::noop_waker_ref());
33        let mut future = pin!(self);
34
35        loop {
36            match future.as_mut().poll(&mut context) {
37                Poll::Pending => continue,
38                Poll::Ready(output) => return output,
39            }
40        }
41    }
42}
43
44/// Unit tests for the helpers defined in the `util` module.
45#[cfg(test)]
46mod tests {
47    /// Tests the behavior of the [`BlockingWait`] extension.
48    #[test]
49    fn blocking_wait_blocks_until_future_is_ready() {
50        use std::task::Poll;
51
52        use super::BlockingWait as _;
53
54        let mut remaining_polls = 100;
55
56        let future = futures::future::poll_fn(|_context| {
57            if remaining_polls == 0 {
58                Poll::Ready(())
59            } else {
60                remaining_polls -= 1;
61                Poll::Pending
62            }
63        });
64
65        future.blocking_wait();
66
67        assert_eq!(remaining_polls, 0);
68    }
69}