Skip to main content

persistant_blob/
concurrency.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Runtime-agnostic helpers for sampled-concurrency contract tests, shared with `persistant-kv`.
3//!
4//! No spawn and no runtime: tasks are polled round-robin on the calling task, so interleaving is
5//! at await points and real parallelism comes from the backend's own threads.
6
7use std::future::Future;
8use std::pin::Pin;
9use std::task::Poll;
10
11/// A boxed future borrowed from the test's stack.
12pub type Task<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
13
14/// Let other tasks polled by the same [`join_all`] run once.
15pub async fn yield_now() {
16    let mut yielded = false;
17    std::future::poll_fn(|cx| {
18        if yielded {
19            Poll::Ready(())
20        } else {
21            yielded = true;
22            cx.waker().wake_by_ref();
23            Poll::Pending
24        }
25    })
26    .await;
27}
28
29/// Poll every future on this task until all are done; results are in input order.
30pub async fn join_all<T>(tasks: Vec<Task<'_, T>>) -> Vec<T> {
31    let mut slots: Vec<Option<T>> = tasks.iter().map(|_| None).collect();
32    let mut tasks: Vec<Option<Task<'_, T>>> = tasks.into_iter().map(Some).collect();
33    std::future::poll_fn(|cx| {
34        let mut pending = false;
35        for (i, t) in tasks.iter_mut().enumerate() {
36            let ready = match t.as_mut() {
37                Some(fut) => match fut.as_mut().poll(cx) {
38                    Poll::Ready(v) => Some(v),
39                    Poll::Pending => {
40                        pending = true;
41                        None
42                    }
43                },
44                None => None,
45            };
46            if let Some(v) = ready {
47                slots[i] = Some(v);
48                *t = None;
49            }
50        }
51        if pending {
52            Poll::Pending
53        } else {
54            Poll::Ready(())
55        }
56    })
57    .await;
58    slots.into_iter().flatten().collect()
59}