feat: introduce allocative

Author: quininer

Cargo.lock

@@ -113,6 +113,17 @@ wasmparser          = { version = "0.222.0", default-features = false }
 winnow              = { version = "0.7.12", default-features = false, features = ["std", "simd"] }
 xxhash-rust         = { version = "0.8.14", default-features = false }
 
+allocative       = { package = "rspack-allocative", version = "0.3.5", default-features = false, features = [
+  "camino",
+  "dashmap",
+  "hashbrown",
+  "indexmap",
+  "tokio",
+  "once_cell",
+  "ustr",
+  "atomic_refcell",
+] }
+
 # Pinned
 napi        = { version = "3.1.6", default-features = false }
 napi-build  = { version = "2.2.3", default-features = false }
@@ -418,6 +429,10 @@ unused_import_braces                   = "warn"
 unused_lifetimes                       = "warn"
 unused_macro_rules                     = "warn"
 
+[workspace.lints.rust.unexpected_cfgs]
+level = "warn"
+check-cfg = ['cfg(allocative)']
+
 [workspace.lints.clippy]
 cargo_common_metadata   = "allow"
 empty_docs              = "allow" # there are some false positives inside biome_wasm

Cargo.toml

@@ -74,6 +74,9 @@ async function build() {
 			features.push("sftrace-setup");
 			rustflags.push("-Zinstrument-xray=always");
 		}
+		if (process.env.ALLOCATIVE) {
+			rustflags.push("--cfg=allocative");
+		}
 		if (values.profile === "release") {
 			features.push("info-level");
 		}

crates/node_binding/scripts/build.js

@@ -10,6 +10,7 @@ version.workspace = true
 noop = []
 
 [dependencies]
+allocative      = { workspace = true }
 camino          = { workspace = true }
 dashmap         = { workspace = true }
 hashlink        = { workspace = true }
@@ -27,3 +28,7 @@ smol_str        = { workspace = true }
 swc_core        = { workspace = true, features = ["ecma_ast"] }
 ustr            = { workspace = true }
 xxhash-rust     = { workspace = true, features = ["const_xxh64"] }
+
+[lints.rust.unexpected_cfgs]
+level = "warn"
+check-cfg = ['cfg(allocative)']

crates/rspack_cacheable/Cargo.toml

@@ -69,7 +69,9 @@ use std::sync::{Arc, LazyLock};
 
 use dashmap::DashSet;
 
+#[cfg_attr(allocative, allocative::root)]
 pub static CACHED_ARC_STR: LazyLock<DashSet<Arc<str>>> = LazyLock::new(Default::default);
+
 impl AsRefStrConverter for Arc<str> {
   fn as_str(&self) -> &str {
     self.as_ref()

crates/rspack_cacheable/src/with/as_ref_str.rs

@@ -16,3 +16,9 @@ rayon            = { workspace = true }
 rspack_cacheable = { workspace = true }
 serde            = { workspace = true, features = ["derive"] }
 ustr             = { workspace = true, features = ["serde"] }
+allocative       = { workspace = true }
+
+[lints.rust.unexpected_cfgs]
+level = "warn"
+check-cfg = ['cfg(allocative)']
+

crates/rspack_collections/Cargo.toml

@@ -36,6 +36,7 @@ pub type IdentifierLinkedSet = LinkedHashSet<Identifier, BuildHasherDefault<Iden
 
 #[cacheable(hashable)]
 #[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize)]
+#[cfg_attr(allocative, derive(allocative::Allocative))]
 pub struct Identifier(#[cacheable(with=AsPreset)] Ustr);
 
 impl Deref for Identifier {

crates/rspack_collections/src/identifier.rs

@@ -35,6 +35,7 @@ pub trait ItemUkey {
 /// Ukey stands for Unique key
 #[rspack_cacheable::cacheable]
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+#[cfg_attr(allocative, derive(allocative::Allocative))]
 pub struct Ukey(u32);
 
 impl Ukey {

crates/rspack_collections/src/ukey.rs

@@ -27,6 +27,8 @@ use rspack_hook::define_hook;
 use rspack_paths::{ArcPath, ArcPathIndexSet, ArcPathSet};
 use rspack_sources::{BoxSource, CachedSource, SourceExt};
 use rspack_tasks::CompilerContext;
+#[cfg(allocative)]
+use rspack_util::allocative;
 use rspack_util::{itoa, tracing_preset::TRACING_BENCH_TARGET};
 use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet, FxHasher};
 use tracing::instrument;
@@ -132,6 +134,7 @@ pub struct CompilationHooks {
 }
 
 #[derive(Debug, Clone, Copy, Hash, Eq, PartialEq, Ord, PartialOrd)]
+#[cfg_attr(allocative, derive(allocative::Allocative))]
 pub struct CompilationId(pub u32);
 
 impl CompilationId {

crates/rspack_core/src/compiler/compilation.rs

@@ -190,6 +190,32 @@ impl Compiler {
     self.compile_done().await?;
     self.cache.after_compile(&self.compilation).await;
 
+    #[cfg(allocative)]
+    {
+      use std::{
+        path::PathBuf,
+        sync::{
+          LazyLock,
+          atomic::{self, AtomicUsize},
+        },
+      };
+
+      use rspack_util::allocative;
+
+      static ENABLE: LazyLock<Option<PathBuf>> =
+        LazyLock::new(|| std::env::var_os("RSPACK_ALLOCATIVE_DIR").map(Into::into));
+      static COUNT: AtomicUsize = AtomicUsize::new(0);
+
+      if let Some(dir) = ENABLE.as_deref() {
+        let mut builder = allocative::FlameGraphBuilder::default();
+        builder.visit_global_roots();
+        let buf = builder.finish_and_write_flame_graph();
+        let count = COUNT.fetch_add(1, atomic::Ordering::Relaxed);
+        let path = dir.join(format!("allocative.{}.flamegraph", count));
+        std::fs::write(path, buf).expect("allocative write failed");
+      }
+    }
+
     Ok(())
   }
 }