rspack_parallel/
scope.rs

1use std::{cell::RefCell, future::Future, marker::PhantomData, pin::Pin};
2
3use tokio::task::{JoinError, JoinHandle};
4
5/// Scope Token
6pub struct Token<'scope, 'spawner, O> {
7  list: &'spawner RefCell<Vec<JoinHandle<O>>>,
8  _phantom: PhantomData<&'scope mut &'scope ()>,
9}
10
11/// Scope Spawner
12pub struct Spawner<'scope, 'spawner, T, O> {
13  list: &'spawner RefCell<Vec<JoinHandle<O>>>,
14  used: T,
15  _phantom: PhantomData<&'scope mut &'scope ()>,
16}
17
18/// Async scope helper
19///
20/// This function helps you write unsafe
21/// asynchronous structured concurrent code more easily.
22/// but it is **still unsafe**, so need to be careful when using it.
23///
24/// To use it safely,
25/// the user needs to ensure that the task is done within used reference lifetime.
26/// Due to `std::mem::forget`, the Rust currently cannot guarantee it.
27///
28/// From a practical point of view, the following points need to be note
29///
30/// * `.await` as early as possible
31/// * Don't put task into container unless you know what you are doing
32/// * Don't call `std::mem::forget`
33///
34/// # Example
35///
36/// ```rust
37/// # #[tokio::test]
38/// # async fn foo() {
39/// let list: Vec<u32> = vec![1, 2, 3, 4];
40///
41/// rspack_parallel::scope(|token| {
42///   for i in 0..list.len() {
43///     let s = unsafe { token.used(&list) };
44///
45///     s.spawn(move |list| async move {
46///       &list[i];
47///     });
48///   }
49/// })
50/// .await;
51/// # }
52/// ```
53///
54/// This doesn't compile
55///
56/// ```rust,compile_fail
57/// # async fn foo() {
58/// rspack_parallel::scope(|token| {
59///   let list: Vec<u32> = vec![1, 2, 3, 4];
60///
61///   for i in 0..list.len() {
62///     let s = unsafe { token.used(&list) };
63///
64///     s.spawn(move |list| async move {
65///       &list[i];
66///     });
67///   }
68/// })
69/// .await;
70/// # }
71/// ```
72pub async fn scope<'scope, F, O>(f: F) -> Vec<Result<O, JoinError>>
73where
74  for<'spawner> F: FnOnce(Token<'scope, 'spawner, O>),
75  O: Send + 'static,
76{
77  struct ScopeGuard(());
78
79  impl ScopeGuard {
80    fn forget(self) {
81      #[allow(clippy::disallowed_methods)]
82      std::mem::forget(self);
83    }
84  }
85
86  impl Drop for ScopeGuard {
87    fn drop(&mut self) {
88      // avoid unsound caused by poll interruption
89      std::process::abort();
90    }
91  }
92
93  let guard = ScopeGuard(());
94  let list = RefCell::new(Vec::new());
95
96  let token = Token {
97    list: &list,
98    _phantom: PhantomData,
99  };
100
101  f(token);
102
103  let list = RefCell::into_inner(list);
104  let mut output = Vec::with_capacity(list.len());
105
106  for j in list {
107    output.push(j.await);
108  }
109
110  guard.forget();
111  output
112}
113
114impl<'scope, 'spawner, O> Token<'scope, 'spawner, O> {
115  /// Use references
116  ///
117  /// Specify the reference to use when spawning the task.
118  ///
119  /// # Safety
120  ///
121  /// This is not sound.
122  ///
123  /// the user must ensure that `scope` task is legally consumed,
124  /// and assume that the runtime handles the task correctly.
125  pub unsafe fn used<T: 'scope>(&self, used: T) -> Spawner<'scope, 'spawner, T, O> {
126    Spawner {
127      list: self.list,
128      used,
129      _phantom: PhantomData,
130    }
131  }
132}
133
134impl<'scope, T, O> Spawner<'scope, '_, T, O> {
135  /// Spawn task from used reference
136  pub fn spawn<F, Fut>(self, f: F)
137  where
138    // TODO Use AsyncFnOnce
139    F: FnOnce(T) -> Fut + 'static,
140    Fut: Future<Output = O> + Send + 'scope,
141    T: Send + Sync + 'scope,
142    O: Send + 'static,
143  {
144    let fut = f(self.used);
145    let fut: Pin<Box<dyn Future<Output = O> + Send + 'scope>> = Box::pin(fut);
146
147    // # Safety
148    //
149    // The safety guarantee here comes from `Token::used`.
150    // The user needs to ensure that the task will done within used reference lifetime.
151    let fut: Pin<Box<dyn Future<Output = O> + Send + 'static>> =
152      unsafe { std::mem::transmute(fut) };
153
154    let j = rspack_tasks::spawn_in_compiler_context(fut);
155    self.list.borrow_mut().push(j);
156  }
157}
rspack_parallel/scope.rs

rspack_parallel/
scope.rs
