Skip to main content

hermes_tokio_runtime_components/impls/
parallel_task.rs

1use core::task::{Context, Poll};
2
3use futures::stream::{Stream, StreamExt};
4use futures::task::noop_waker;
5use hermes_async_runtime_components::stream::traits::boxed::HasBoxedStreamType;
6use hermes_runtime_components::traits::task::{ConcurrentTaskRunner, Task};
7use tokio::task::JoinSet;
8
9pub struct TokioRunParallelTasks;
10
11impl<Runtime> ConcurrentTaskRunner<Runtime> for TokioRunParallelTasks
12where
13    Runtime: HasBoxedStreamType,
14{
15    async fn run_concurrent_tasks<T>(_runtime: &Runtime, tasks: Vec<T>)
16    where
17        T: Task,
18    {
19        run_parallel_tasks(tasks).await
20    }
21
22    async fn run_concurrent_task_stream<T>(_runtime: &Runtime, tasks: Runtime::Stream<T>)
23    where
24        T: Task,
25    {
26        run_parallel_task_stream(Runtime::to_boxed_stream(tasks)).await
27    }
28}
29
30pub async fn run_parallel_tasks<T>(tasks: Vec<T>)
31where
32    T: Task,
33{
34    let mut join_set = JoinSet::new();
35
36    for task in tasks.into_iter() {
37        join_set.spawn(async move {
38            task.run().await;
39        });
40    }
41
42    while join_set.join_next().await.is_some() {}
43}
44
45pub async fn run_parallel_task_stream<T>(tasks: impl Stream<Item = T>)
46where
47    T: Task,
48{
49    let mut join_set = JoinSet::new();
50
51    let waker = noop_waker();
52
53    tasks
54        .for_each_concurrent(None, |task| {
55            join_set.spawn(async move {
56                task.run().await;
57            });
58
59            let mut context = Context::from_waker(&waker);
60
61            // Remove finished tasks from the JoinSet to avoid accumulation of
62            // tasks from a potentially infinite stream.
63            while let Poll::Ready(Some(_)) = join_set.poll_join_next(&mut context) {}
64
65            async {}
66        })
67        .await;
68
69    // Wait for all tasks to complete once the stream ends.
70    while join_set.join_next().await.is_some() {}
71}