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  ffi::c_void,
5  future::Future,
6  sync::{
7    Arc,
8    atomic::{AtomicPtr, AtomicU32},
9  },
10};
11
12use tokio::{
13  task::{JoinHandle, futures::TaskLocalFuture},
14  task_local,
15};
16
17// don't overuse this and put everything here, it's mostly used for store isolated id generator
18#[derive(Debug)]
19pub struct CompilerContext {
20  dependenc_id_generator: AtomicU32,
21  exports_info_artifact_ptr: AtomicPtr<c_void>,
22}
23
24task_local! {
25  // implicit COMPIlER_CONTEXT for current running compiler, every compiler has its own isolated compiler context
26 pub static CURRENT_COMPILER_CONTEXT: Arc<CompilerContext>;
27}
28#[allow(clippy::new_without_default)]
29impl CompilerContext {
30  pub fn new() -> Self {
31    Self {
32      dependenc_id_generator: AtomicU32::new(0),
33      exports_info_artifact_ptr: AtomicPtr::new(std::ptr::null_mut()),
34    }
35  }
36  pub fn fetch_new_dependency_id(&self) -> u32 {
37    self
38      .dependenc_id_generator
39      .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
40  }
41  pub fn dependency_id(&self) -> u32 {
42    self
43      .dependenc_id_generator
44      .load(std::sync::atomic::Ordering::SeqCst)
45  }
46  pub fn set_dependency_id(&self, id: u32) {
47    self
48      .dependenc_id_generator
49      .store(id, std::sync::atomic::Ordering::SeqCst);
50  }
51
52  pub fn exports_info_artifact_ptr(&self) -> Option<*mut c_void> {
53    let ptr = self
54      .exports_info_artifact_ptr
55      .load(std::sync::atomic::Ordering::SeqCst);
56    (!ptr.is_null()).then_some(ptr)
57  }
58
59  pub fn set_exports_info_artifact_ptr(&self, ptr: Option<*mut c_void>) {
60    self.exports_info_artifact_ptr.store(
61      ptr.unwrap_or(std::ptr::null_mut()),
62      std::sync::atomic::Ordering::SeqCst,
63    );
64  }
65}
66
67pub fn fetch_new_dependency_id() -> u32 {
68  CURRENT_COMPILER_CONTEXT.get().fetch_new_dependency_id()
69}
70pub fn get_current_dependency_id() -> u32 {
71  CURRENT_COMPILER_CONTEXT.get().dependency_id()
72}
73pub fn set_current_dependency_id(id: u32) {
74  CURRENT_COMPILER_CONTEXT.get().set_dependency_id(id);
75}
76
77pub fn within_compiler_context<F>(
78  compiler_context: Arc<CompilerContext>,
79  f: F,
80) -> TaskLocalFuture<Arc<CompilerContext>, F>
81where
82  F: Future,
83{
84  CURRENT_COMPILER_CONTEXT.scope(compiler_context, f)
85}
86pub fn within_compiler_context_sync<F, R>(compiler_context: Arc<CompilerContext>, f: F) -> R
87where
88  F: FnOnce() -> R,
89{
90  CURRENT_COMPILER_CONTEXT.sync_scope(compiler_context, f)
91}
92
93// this is only used for testing rust builder api, we need to find better api in the future
94/// For test use only.
95pub fn within_compiler_context_for_testing_sync<F, R>(f: F) -> R
96where
97  F: FnOnce() -> R,
98{
99  CURRENT_COMPILER_CONTEXT.sync_scope(Arc::new(CompilerContext::new()), f)
100}
101/// For test use only.
102pub fn within_compiler_context_for_testing<F>(f: F) -> TaskLocalFuture<Arc<CompilerContext>, F>
103where
104  F: Future,
105{
106  CURRENT_COMPILER_CONTEXT.scope(Arc::new(CompilerContext::new()), f)
107}
108
109pub fn spawn_in_compiler_context<F>(future: F) -> JoinHandle<F::Output>
110where
111  F: Future + Send + 'static,
112  F::Output: Send + 'static,
113{
114  let compiler_context = CURRENT_COMPILER_CONTEXT.get();
115
116  tokio::spawn(CURRENT_COMPILER_CONTEXT.scope(compiler_context, future))
117}
118
119/// Like [`spawn_in_compiler_context`], but falls back to a plain [`tokio::spawn`]
120/// when there is no active compiler context (e.g. in unit tests or utility code
121/// that is not driven by a compiler).
122pub fn spawn_in_context<F>(future: F) -> JoinHandle<F::Output>
123where
124  F: Future + Send + 'static,
125  F::Output: Send + 'static,
126{
127  match CURRENT_COMPILER_CONTEXT.try_get() {
128    Ok(compiler_context) => tokio::spawn(CURRENT_COMPILER_CONTEXT.scope(compiler_context, future)),
129    Err(_) => tokio::spawn(future),
130  }
131}
rspack_tasks/lib.rs

rspack_tasks/
lib.rs
