feat(next-swc): add compile-time heap profiler flag

packages/next-swc/Cargo.toml

@@ -126,3 +126,4 @@ tracing = "0.1.37"
 url = "2.2.2"
 urlencoding = "2.1.2"
 webbrowser = "0.8.7"
+dhat = { version = "0.3.2" }

packages/next-swc/crates/napi/Cargo.toml

@@ -26,10 +26,18 @@ __internal_nextjs_integration_test = [
   "next-dev/serializable"
 ]
 
+# Enable dhat profiling allocator for heap profiling.
+__internal_dhat-heap = ["dhat"]
+# Enable dhat profiling allocator for ad hoc profiling.
+# [Note]: we do not have any ad hoc event in the codebase yet, so enabling this
+# effectively does nothing.
+__internal_dhat-ad-hoc = ["dhat"]
+
 [dependencies]
 anyhow = "1.0.66"
 backtrace = "0.3"
 fxhash = "0.2.1"
+dhat = { workspace = true, optional = true }
 napi = { version = "2", default-features = false, features = [
   "napi3",
   "serde-json",

packages/next-swc/crates/napi/src/lib.rs

@@ -52,10 +52,18 @@ pub mod util;
 
 // don't use turbo malloc (`mimalloc`) on linux-musl-aarch64 because of the
 // compile error
-#[cfg(not(all(target_os = "linux", target_env = "musl", target_arch = "aarch64")))]
+#[cfg(not(any(
+    all(target_os = "linux", target_env = "musl", target_arch = "aarch64"),
+    feature = "__internal_dhat-heap",
+    feature = "__internal_dhat-ad-hoc"
+)))]
 #[global_allocator]
 static ALLOC: turbo_binding::turbo::malloc::TurboMalloc = turbo_binding::turbo::malloc::TurboMalloc;
 
+#[cfg(feature = "__internal_dhat-heap")]
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
+
 static COMPILER: Lazy<Arc<Compiler>> = Lazy::new(|| {
     let cm = Arc::new(SourceMap::new(FilePathMapping::empty()));
 

packages/next-swc/crates/napi/src/util.rs

@@ -53,6 +53,45 @@ pub trait MapErr<T>: Into<Result<T, anyhow::Error>> {
 
 impl<T> MapErr<T> for Result<T, anyhow::Error> {}
 
+#[cfg(any(feature = "__internal_dhat-heap", feature = "__internal_dhat-ad-hoc"))]
+#[napi]
+pub fn init_heap_profiler() -> napi::Result<External<RefCell<Option<dhat::Profiler>>>> {
+    #[cfg(feature = "__internal_dhat-heap")]
+    {
+        println!("[dhat-heap]: Initializing heap profiler");
+        let _profiler = dhat::Profiler::new_heap();
+        return Ok(External::new(RefCell::new(Some(_profiler))));
+    }
+
+    #[cfg(feature = "__internal_dhat-ad-hoc")]
+    {
+        println!("[dhat-ad-hoc]: Initializing ad-hoc profiler");
+        let _profiler = dhat::Profiler::new_ad_hoc();
+        return Ok(External::new(RefCell::new(Some(_profiler))));
+    }
+}
+
+#[cfg(any(feature = "__internal_dhat-heap", feature = "__internal_dhat-ad-hoc"))]
+#[napi]
+pub fn teardown_heap_profiler(guard_external: External<RefCell<Option<dhat::Profiler>>>) {
+    let guard_cell = &*guard_external;
+
+    if let Some(guard) = guard_cell.take() {
+        println!("[dhat]: Teardown profiler");
+        drop(guard);
+    }
+}
+
+#[cfg(not(any(feature = "__internal_dhat-heap", feature = "__internal_dhat-ad-hoc")))]
+#[napi]
+pub fn init_heap_profiler() -> napi::Result<External<RefCell<Option<u32>>>> {
+    Ok(External::new(RefCell::new(Some(0))))
+}
+
+#[cfg(not(any(feature = "__internal_dhat-heap", feature = "__internal_dhat-ad-hoc")))]
+#[napi]
+pub fn teardown_heap_profiler(_guard_external: External<RefCell<Option<u32>>>) {}
+
 /// Initialize tracing subscriber to emit traces. This configures subscribers
 /// for Trace Event Format (https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview).
 #[napi]

packages/next/src/build/index.ts

@@ -119,6 +119,7 @@ import {
   teardownTraceSubscriber,
   teardownCrashReporter,
   loadBindings,
+  teardownHeapProfiler,
 } from './swc'
 import { getNamedRouteRegex } from '../shared/lib/router/utils/route-regex'
 import { flatReaddir } from '../lib/flat-readdir'
@@ -3144,6 +3145,7 @@ export default async function build(
     // Ensure all traces are flushed before finishing the command
     await flushAllTraces()
     teardownTraceSubscriber()
+    teardownHeapProfiler()
     teardownCrashReporter()
   }
 }

packages/next/src/build/output/store.ts

@@ -1,7 +1,11 @@
 import createStore from 'next/dist/compiled/unistore'
 import stripAnsi from 'next/dist/compiled/strip-ansi'
 import { flushAllTraces } from '../../trace'
-import { teardownCrashReporter, teardownTraceSubscriber } from '../swc'
+import {
+  teardownCrashReporter,
+  teardownHeapProfiler,
+  teardownTraceSubscriber,
+} from '../swc'
 import * as Log from './log'
 
 export type OutputState =
@@ -92,6 +96,7 @@ store.subscribe((state) => {
     // Ensure traces are flushed after each compile in development mode
     flushAllTraces()
     teardownTraceSubscriber()
+    teardownHeapProfiler()
     teardownCrashReporter()
     return
   }
@@ -120,6 +125,7 @@ store.subscribe((state) => {
     // Ensure traces are flushed after each compile in development mode
     flushAllTraces()
     teardownTraceSubscriber()
+    teardownHeapProfiler()
     teardownCrashReporter()
     return
   }
@@ -137,5 +143,6 @@ store.subscribe((state) => {
   // Ensure traces are flushed after each compile in development mode
   flushAllTraces()
   teardownTraceSubscriber()
+  teardownHeapProfiler()
   teardownCrashReporter()
 })

packages/next/src/build/swc/index.ts

@@ -76,6 +76,7 @@ let wasmBindings: any
 let downloadWasmPromise: any
 let pendingBindings: any
 let swcTraceFlushGuard: any
+let swcHeapProfilerFlushGuard: any
 let swcCrashReporterFlushGuard: any
 export const lockfilePatchPromise: { cur?: Promise<void> } = {}
 
@@ -425,9 +426,13 @@ function loadNative(isCustomTurbopack = false) {
       getTargetTriple: bindings.getTargetTriple,
       initCustomTraceSubscriber: bindings.initCustomTraceSubscriber,
       teardownTraceSubscriber: bindings.teardownTraceSubscriber,
+      initHeapProfiler: bindings.initHeapProfiler,
+      teardownHeapProfiler: bindings.teardownHeapProfiler,
       teardownCrashReporter: bindings.teardownCrashReporter,
       turbo: {
         startDev: (options: any) => {
+          initHeapProfiler()
+
           const devOptions = {
             ...options,
             noOpen: options.noOpen ?? true,
@@ -504,13 +509,19 @@ function loadNative(isCustomTurbopack = false) {
           }
         },
         nextBuild: (options: unknown) => {
-          return bindings.nextBuild(options)
+          initHeapProfiler()
+          const ret = bindings.nextBuild(options)
+
+          return ret
         },
-        startTrace: (options = {}, turboTasks: unknown) =>
-          bindings.runTurboTracing(
+        startTrace: (options = {}, turboTasks: unknown) => {
+          initHeapProfiler()
+          const ret = bindings.runTurboTracing(
             toBuffer({ exact: true, ...options }),
             turboTasks
-          ),
+          )
+          return ret
+        },
         createTurboTasks: (memoryLimit?: number): unknown =>
           bindings.createTurboTasks(memoryLimit),
       },
@@ -589,6 +600,43 @@ export const initCustomTraceSubscriber = (traceFileName?: string): void => {
   }
 }
 
+/**
+ * Initialize heap profiler, if possible.
+ * Note this is not available in release build of next-swc by default,
+ * only available by manually building next-swc with specific flags.
+ * Calling in release build will not do anything.
+ */
+export const initHeapProfiler = () => {
+  if (!swcHeapProfilerFlushGuard) {
+    // Wasm binary doesn't support profiler
+    let bindings = loadNative()
+    swcHeapProfilerFlushGuard = bindings.initHeapProfiler()
+  }
+}
+
+/**
+ * Teardown heap profiler, if possible.
+ *
+ * Same as initialization, this is not available in release build of next-swc by default
+ * and calling it will not do anything.
+ */
+export const teardownHeapProfiler = (() => {
+  let flushed = false
+  return (): void => {
+    if (!flushed) {
+      flushed = true
+      try {
+        let bindings = loadNative()
+        if (swcHeapProfilerFlushGuard) {
+          bindings.teardownHeapProfiler(swcHeapProfilerFlushGuard)
+        }
+      } catch (e) {
+        // Suppress exceptions, this fn allows to fail to load native bindings
+      }
+    }
+  }
+})()
+
 /**
  * Teardown swc's trace subscriber if there's an initialized flush guard exists.
  *