xet_runtime/utils/
limited_joinset.rs

1use std::future::Future;
2use std::sync::Arc;
3use std::task::{Context, Poll};
4
5use tokio::sync::Semaphore;
6use tokio::task::{AbortHandle, JoinError, JoinSet as TokioJoinSet};
7
8pub struct LimitedJoinSet<T> {
9    inner: TokioJoinSet<T>,
10    semaphore: Arc<Semaphore>,
11}
12
13impl<T: 'static> LimitedJoinSet<T> {
14    pub fn new(max_concurrent: usize) -> Self {
15        Self {
16            inner: TokioJoinSet::new(),
17            semaphore: Arc::new(Semaphore::new(max_concurrent)),
18        }
19    }
20
21    pub fn spawn<F>(&mut self, task: F) -> AbortHandle
22    where
23        F: Future<Output = T>,
24        F: Send + 'static,
25        T: Send,
26    {
27        let semaphore = self.semaphore.clone();
28        self.inner.spawn(async move {
29            let _permit = semaphore.acquire().await;
30            task.await
31        })
32    }
33
34    pub fn try_join_next(&mut self) -> Option<Result<T, JoinError>> {
35        self.inner.try_join_next()
36    }
37
38    pub async fn join_next(&mut self) -> Option<Result<T, JoinError>> {
39        self.inner.join_next().await
40    }
41
42    pub async fn join_all(self) -> Vec<T> {
43        self.inner.join_all().await
44    }
45
46    pub fn poll_join_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<T, JoinError>>> {
47        self.inner.poll_join_next(cx)
48    }
49
50    pub fn len(&self) -> usize {
51        self.inner.len()
52    }
53
54    pub fn is_empty(&self) -> bool {
55        self.inner.is_empty()
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use std::time::Duration;
62
63    use super::*;
64
65    #[tokio::test]
66    async fn test_joinset() {
67        let mut join_set = LimitedJoinSet::new(3);
68
69        for i in 0..4 {
70            join_set.spawn(async move {
71                tokio::time::sleep(Duration::from_millis(10 - i)).await;
72                i
73            });
74        }
75
76        let mut outs = Vec::new();
77        while let Some(Ok(value)) = join_set.join_next().await {
78            outs.push(value);
79        }
80
81        // expect that the task returning 3 was spawned after at least 1 other task finished
82        assert_eq!(outs.len(), 4);
83        for (i, out) in outs.into_iter().enumerate() {
84            if out == 3 {
85                assert!(i > 0);
86            }
87        }
88    }
89}
xet_runtime/utils/limited_joinset.rs

xet_runtime/utils/
limited_joinset.rs
