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, Pin},
9    task::{Context, Poll},
10};
11
12use futures::task;
13
14/// Yields the current asynchronous task so that other tasks may progress if possible.
15///
16/// After other tasks progress, this task resumes as soon as possible. More explicitly, it is
17/// scheduled to be woken up as soon as possible.
18pub fn yield_once() -> YieldOnce {
19    YieldOnce::default()
20}
21
22/// A [`Future`] that returns [`Poll::Pending`] once and immediately schedules itself to wake up.
23#[derive(Default)]
24pub struct YieldOnce {
25    yielded: bool,
26}
27
28impl Future for YieldOnce {
29    type Output = ();
30
31    fn poll(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
32        let mut this = self.as_mut();
33
34        if this.yielded {
35            Poll::Ready(())
36        } else {
37            this.yielded = true;
38            context.waker().wake_by_ref();
39            Poll::Pending
40        }
41    }
42}
43
44/// An extension trait to block on a [`Future`] until it completes.
45pub trait BlockingWait {
46    /// The type returned by the [`Future`].
47    type Output;
48
49    /// Waits for the [`Future`] to complete in a blocking manner.
50    ///
51    /// Effectively polls the [`Future`] repeatedly until it returns [`Poll::Ready`].
52    fn blocking_wait(self) -> Self::Output;
53}
54
55impl<AnyFuture> BlockingWait for AnyFuture
56where
57    AnyFuture: Future,
58{
59    type Output = AnyFuture::Output;
60
61    fn blocking_wait(mut self) -> Self::Output {
62        let waker = task::noop_waker();
63        let mut task_context = Context::from_waker(&waker);
64        let mut future = pin!(self);
65
66        loop {
67            match future.as_mut().poll(&mut task_context) {
68                Poll::Pending => continue,
69                Poll::Ready(output) => return output,
70            }
71        }
72    }
73}
74
75/// Unit tests for the helpers defined in the `util` module.
76#[cfg(test)]
77mod tests {
78    use std::task::{Context, Poll};
79
80    use futures::{future::poll_fn, task::noop_waker, FutureExt as _};
81
82    use super::{yield_once, BlockingWait};
83
84    /// Tests the behavior of the [`YieldOnce`] future.
85    ///
86    /// Checks the internal state before and after the first and second polls, and ensures that
87    /// only the first poll returns [`Poll::Pending`].
88    #[test]
89    #[expect(clippy::bool_assert_comparison)]
90    fn yield_once_returns_pending_only_on_first_call() {
91        let mut future = yield_once();
92
93        let waker = noop_waker();
94        let mut context = Context::from_waker(&waker);
95
96        assert_eq!(future.yielded, false);
97        assert!(future.poll_unpin(&mut context).is_pending());
98        assert_eq!(future.yielded, true);
99        assert!(future.poll_unpin(&mut context).is_ready());
100        assert_eq!(future.yielded, true);
101    }
102
103    /// Tests the behavior of the [`BlockingWait`] extension.
104    #[test]
105    fn blocking_wait_blocks_until_future_is_ready() {
106        let mut remaining_polls = 100;
107
108        let future = poll_fn(|_context| {
109            if remaining_polls == 0 {
110                Poll::Ready(())
111            } else {
112                remaining_polls -= 1;
113                Poll::Pending
114            }
115        });
116
117        future.blocking_wait();
118
119        assert_eq!(remaining_polls, 0);
120    }
121}
linera_sdk/util.rs

linera_sdk/
util.rs
