rspack_tasks/
lib.rs

1// borrow the ideas from turbo_tasks https://github.com/vercel/next.js/blob/678ef8b5650871a730ca14c480c762ca53716575/turbopack/crates/turbo-tasks/src/manager.rs#L1
2// which creates a implicit compiler context to support isolated parallel compiler state
3use std::{
4  future::Future,
5  sync::{Arc, atomic::AtomicU32},
6};
7
8use tokio::{
9  task::{JoinHandle, futures::TaskLocalFuture},
10  task_local,
11};
12
13// don't overuse this and put everything here, it's mostly used for store isolated id generator
14#[derive(Debug)]
15pub struct CompilerContext {
16  dependenc_id_generator: AtomicU32,
17}
18
19task_local! {
20  // implicit COMPIlER_CONTEXT for current running compiler, every compiler has its own isolated compiler context
21 pub static CURRENT_COMPILER_CONTEXT: Arc<CompilerContext>;
22}
23#[allow(clippy::new_without_default)]
24impl CompilerContext {
25  pub fn new() -> Self {
26    Self {
27      dependenc_id_generator: AtomicU32::new(0),
28    }
29  }
30  pub fn fetch_new_dependency_id(&self) -> u32 {
31    self
32      .dependenc_id_generator
33      .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
34  }
35  pub fn dependency_id(&self) -> u32 {
36    self
37      .dependenc_id_generator
38      .load(std::sync::atomic::Ordering::SeqCst)
39  }
40  pub fn set_dependency_id(&self, id: u32) {
41    self
42      .dependenc_id_generator
43      .store(id, std::sync::atomic::Ordering::SeqCst);
44  }
45}
46
47pub fn fetch_new_dependency_id() -> u32 {
48  CURRENT_COMPILER_CONTEXT.get().fetch_new_dependency_id()
49}
50pub fn get_current_dependency_id() -> u32 {
51  CURRENT_COMPILER_CONTEXT.get().dependency_id()
52}
53pub fn set_current_dependency_id(id: u32) {
54  CURRENT_COMPILER_CONTEXT.get().set_dependency_id(id);
55}
56
57pub fn within_compiler_context<F>(
58  compiler_context: Arc<CompilerContext>,
59  f: F,
60) -> TaskLocalFuture<Arc<CompilerContext>, F>
61where
62  F: Future,
63{
64  CURRENT_COMPILER_CONTEXT.scope(compiler_context, f)
65}
66pub fn within_compiler_context_sync<F, R>(compiler_context: Arc<CompilerContext>, f: F) -> R
67where
68  F: FnOnce() -> R,
69{
70  CURRENT_COMPILER_CONTEXT.sync_scope(compiler_context, f)
71}
72
73// this is only used for testing rust builder api, we need to find better api in the future
74/// For test use only.
75pub fn within_compiler_context_for_testing_sync<F, R>(f: F) -> R
76where
77  F: FnOnce() -> R,
78{
79  CURRENT_COMPILER_CONTEXT.sync_scope(Arc::new(CompilerContext::new()), f)
80}
81/// For test use only.
82pub fn within_compiler_context_for_testing<F>(f: F) -> TaskLocalFuture<Arc<CompilerContext>, F>
83where
84  F: Future,
85{
86  CURRENT_COMPILER_CONTEXT.scope(Arc::new(CompilerContext::new()), f)
87}
88
89pub fn spawn_in_compiler_context<F>(future: F) -> JoinHandle<F::Output>
90where
91  F: Future + Send + 'static,
92  F::Output: Send + 'static,
93{
94  let compiler_context = CURRENT_COMPILER_CONTEXT.get().clone();
95
96  tokio::spawn(CURRENT_COMPILER_CONTEXT.scope(compiler_context, future))
97}
rspack_tasks/lib.rs

rspack_tasks/
lib.rs
