Fix dhat support for turbo dev and build

packages/next-swc/crates/napi/src/next_api/project.rs

@@ -54,7 +54,7 @@ use super::{
         TurbopackResult, VcArc,
     },
 };
-use crate::register;
+use crate::{register, util::DhatProfilerGuard};
 
 /// Used by [`benchmark_file_io`]. This is a noisy benchmark, so set the
 /// threshold high.
@@ -261,10 +261,17 @@ pub async fn project_new(
     turbo_engine_options: NapiTurboEngineOptions,
 ) -> napi::Result<External<ProjectInstance>> {
     register();
-
-    let trace = std::env::var("NEXT_TURBOPACK_TRACING").ok();
     let (exit, exit_receiver) = ExitHandler::new_receiver();
 
+    if let Some(dhat_profiler) = DhatProfilerGuard::try_init() {
+        exit.on_exit(async move {
+            tokio::task::spawn_blocking(move || drop(dhat_profiler))
+                .await
+                .unwrap()
+        });
+    }
+
+    let trace = std::env::var("NEXT_TURBOPACK_TRACING").ok();
     if let Some(mut trace) = trace {
         // Trace presets
         match trace.as_str() {

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

@@ -47,44 +47,55 @@ pub trait MapErr<T>: Into<Result<T, anyhow::Error>> {
 
 impl<T> MapErr<T> for Result<T, anyhow::Error> {}
 
+/// An opaque type potentially wrapping a [`dhat::Profiler`] instance. If we
+/// were not compiled with dhat support, this is an empty struct.
 #[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))));
-    }
+#[non_exhaustive]
+pub struct DhatProfilerGuard(dhat::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))));
+/// An opaque type potentially wrapping a [`dhat::Profiler`] instance. If we
+/// were not compiled with dhat support, this is an empty struct.
+#[cfg(not(any(feature = "__internal_dhat-heap", feature = "__internal_dhat-ad-hoc")))]
+#[non_exhaustive]
+pub struct DhatProfilerGuard;
+
+impl DhatProfilerGuard {
+    /// Constructs an instance if we were compiled with dhat support.
+    pub fn try_init() -> Option<Self> {
+        #[cfg(feature = "__internal_dhat-heap")]
+        {
+            println!("[dhat-heap]: Initializing heap profiler");
+            Some(Self(dhat::Profiler::new_heap()))
+        }
+        #[cfg(feature = "__internal_dhat-ad-hoc")]
+        {
+            println!("[dhat-ad-hoc]: Initializing ad-hoc profiler");
+            Some(Self(dhat::Profiler::new_ad_hoc()))
+        }
+        #[cfg(not(any(feature = "__internal_dhat-heap", feature = "__internal_dhat-ad-hoc")))]
+        {
+            None
+        }
     }
 }
 
-#[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() {
+impl Drop for DhatProfilerGuard {
+    fn drop(&mut self) {
+        #[cfg(any(feature = "__internal_dhat-heap", feature = "__internal_dhat-ad-hoc"))]
         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))))
+#[napi(js_name = "initDhatProfiler")]
+pub fn init_dhat_profiler_js() -> External<RefCell<Option<DhatProfilerGuard>>> {
+    External::new(RefCell::new(DhatProfilerGuard::try_init()))
 }
 
-#[cfg(not(any(feature = "__internal_dhat-heap", feature = "__internal_dhat-ad-hoc")))]
-#[napi]
-pub fn teardown_heap_profiler(_guard_external: External<RefCell<Option<u32>>>) {}
+#[napi(js_name = "teardownDhatProfiler")]
+pub fn teardown_dhat_profiler_js(guard_external: External<RefCell<Option<DhatProfilerGuard>>>) {
+    let guard_cell = &*guard_external;
+    drop(guard_cell.take());
+}
 
 /// Initialize tracing subscriber to emit traces. This configures subscribers
 /// for Trace Event Format <https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview>.

packages/next/src/build/index.ts

@@ -131,7 +131,7 @@ import {
   loadBindings,
   lockfilePatchPromise,
   teardownTraceSubscriber,
-  teardownHeapProfiler,
+  teardownDhatProfiler,
   createDefineEnv,
 } from './swc'
 import { getNamedRouteRegex } from '../shared/lib/router/utils/route-regex'
@@ -3409,7 +3409,7 @@ export default async function build(
     // Ensure all traces are flushed before finishing the command
     await flushAllTraces()
     teardownTraceSubscriber()
-    teardownHeapProfiler()
+    teardownDhatProfiler()
 
     if (traceUploadUrl && loadedConfig) {
       uploadTrace({

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

@@ -1,7 +1,7 @@
 import createStore from 'next/dist/compiled/unistore'
 import stripAnsi from 'next/dist/compiled/strip-ansi'
 import { type Span, flushAllTraces, trace } from '../../trace'
-import { teardownHeapProfiler, teardownTraceSubscriber } from '../swc'
+import { teardownDhatProfiler, teardownTraceSubscriber } from '../swc'
 import * as Log from './log'
 
 const MAX_LOG_SKIP_DURATION = 500 // 500ms
@@ -125,7 +125,7 @@ store.subscribe((state) => {
     // Ensure traces are flushed after each compile in development mode
     flushAllTraces()
     teardownTraceSubscriber()
-    teardownHeapProfiler()
+    teardownDhatProfiler()
     return
   }
 
@@ -149,7 +149,7 @@ store.subscribe((state) => {
     // Ensure traces are flushed after each compile in development mode
     flushAllTraces()
     teardownTraceSubscriber()
-    teardownHeapProfiler()
+    teardownDhatProfiler()
     return
   }
 
@@ -180,5 +180,5 @@ store.subscribe((state) => {
   // Ensure traces are flushed after each compile in development mode
   flushAllTraces()
   teardownTraceSubscriber()
-  teardownHeapProfiler()
+  teardownDhatProfiler()
 })

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

@@ -146,9 +146,11 @@ let wasmBindings: any
 let downloadWasmPromise: any
 let pendingBindings: any
 let swcTraceFlushGuard: any
-let swcHeapProfilerFlushGuard: any
 let downloadNativeBindingsPromise: Promise<void> | undefined = undefined
 
+type DhatProfiler = { __napiType: 'DhatProfilerGuard' }
+let dhatProfiler: DhatProfiler | undefined = undefined
+
 export const lockfilePatchPromise: { cur?: Promise<void> } = {}
 
 export interface Binding {
@@ -181,8 +183,8 @@ export interface Binding {
 
   initCustomTraceSubscriber?: any
   teardownTraceSubscriber?: any
-  initHeapProfiler?: any
-  teardownHeapProfiler?: any
+  initDhatProfiler: () => DhatProfiler
+  teardownDhatProfiler: (guard: DhatProfiler) => void
   css: {
     lightning: {
       transform(transformOptions: any): Promise<any>
@@ -1535,11 +1537,11 @@ function loadNative(importPath?: string) {
       getTargetTriple: bindings.getTargetTriple,
       initCustomTraceSubscriber: bindings.initCustomTraceSubscriber,
       teardownTraceSubscriber: bindings.teardownTraceSubscriber,
-      initHeapProfiler: bindings.initHeapProfiler,
-      teardownHeapProfiler: bindings.teardownHeapProfiler,
+      initDhatProfiler: bindings.initDhatProfiler,
+      teardownDhatProfiler: bindings.teardownDhatProfiler,
       turbo: {
         startTrace: (options = {}, turboTasks: unknown) => {
-          initHeapProfiler()
+          initDhatProfiler()
           return (customBindings ?? bindings).runTurboTracing(
             toBuffer({ exact: true, ...options }),
             turboTasks
@@ -1685,11 +1687,11 @@ export const initCustomTraceSubscriber = (traceFileName?: string): void => {
  * only available by manually building next-swc with specific flags.
  * Calling in release build will not do anything.
  */
-export const initHeapProfiler = () => {
+export const initDhatProfiler = () => {
   try {
-    if (!swcHeapProfilerFlushGuard) {
+    if (dhatProfiler == null) {
       let bindings = loadNative()
-      swcHeapProfilerFlushGuard = bindings.initHeapProfiler()
+      dhatProfiler = bindings.initDhatProfiler()
     }
   } catch (_) {
     // Suppress exceptions, this fn allows to fail to load native bindings
@@ -1702,22 +1704,16 @@ export const initHeapProfiler = () => {
  * 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
-      }
+export const teardownDhatProfiler = () => {
+  try {
+    if (dhatProfiler != null) {
+      let bindings = loadNative()
+      bindings.teardownDhatProfiler(dhatProfiler)
     }
+  } catch (_) {
+    // Suppress exceptions, this fn allows to fail to load native bindings
   }
-})()
+}
 
 /**
  * Teardown swc's trace subscriber if there's an initialized flush guard exists.