From d77d129572a420de346d2c721c7420da61ffc558 Mon Sep 17 00:00:00 2001
From: Benjamin Woodruff <benjamin.woodruff@vercel.com>
Date: Thu, 13 Jun 2024 22:51:51 -0700
Subject: [PATCH] Fix dhat support for turbo dev and build

---
 crates/napi/src/next_api/project.rs           | 12 +++-
 crates/napi/src/util.rs                       | 62 ++++++++++---------
 packages/next/src/build/index.ts              |  2 -
 packages/next/src/build/output/store.ts       |  5 +-
 .../next/src/build/swc/generated-native.d.ts  |  4 --
 packages/next/src/build/swc/index.ts          | 38 +-----------
 packages/next/src/build/swc/types.ts          |  2 -
 turbopack/crates/turbo-tasks/Cargo.toml       |  1 +
 8 files changed, 45 insertions(+), 81 deletions(-)

diff --git a/crates/napi/src/next_api/project.rs b/crates/napi/src/next_api/project.rs
index 03f979770e60e..788c4605a6028 100644
--- a/crates/napi/src/next_api/project.rs
+++ b/crates/napi/src/next_api/project.rs
@@ -55,7 +55,7 @@ use super::{
         NextTurboTasks, RootTask, 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.
@@ -295,11 +295,19 @@ pub async fn project_new(
     turbo_engine_options: NapiTurboEngineOptions,
 ) -> napi::Result<External<ProjectInstance>> {
     register();
+    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 mut trace = std::env::var("NEXT_TURBOPACK_TRACING")
         .ok()
         .filter(|v| !v.is_empty());
-    let (exit, exit_receiver) = ExitHandler::new_receiver();
 
     if cfg!(feature = "tokio-console") && trace.is_none() {
         // ensure `trace` is set to *something* so that the `tokio-console` feature works, otherwise
diff --git a/crates/napi/src/util.rs b/crates/napi/src/util.rs
index a94825d3e60b0..ec324c6edcbfe 100644
--- a/crates/napi/src/util.rs
+++ b/crates/napi/src/util.rs
@@ -138,45 +138,47 @@ 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.
+///
+/// [`dhat::Profiler`]: https://docs.rs/dhat/latest/dhat/struct.Profiler.html
+#[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))))
-}
-
-#[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]
diff --git a/packages/next/src/build/index.ts b/packages/next/src/build/index.ts
index ef74f3c0f158d..c7bae5774119b 100644
--- a/packages/next/src/build/index.ts
+++ b/packages/next/src/build/index.ts
@@ -148,7 +148,6 @@ import {
   loadBindings,
   lockfilePatchPromise,
   teardownTraceSubscriber,
-  teardownHeapProfiler,
   createDefineEnv,
 } from './swc'
 import { getNamedRouteRegex } from '../shared/lib/router/utils/route-regex'
@@ -3796,7 +3795,6 @@ export default async function build(
     // Ensure all traces are flushed before finishing the command
     await flushAllTraces()
     teardownTraceSubscriber()
-    teardownHeapProfiler()
 
     if (traceUploadUrl && loadedConfig) {
       uploadTrace({
diff --git a/packages/next/src/build/output/store.ts b/packages/next/src/build/output/store.ts
index 7c84d1c602fa4..fb195d1c266f6 100644
--- a/packages/next/src/build/output/store.ts
+++ b/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 { teardownTraceSubscriber } from '../swc'
 import * as Log from './log'
 
 const MAX_LOG_SKIP_DURATION = 500 // 500ms
@@ -152,7 +152,6 @@ store.subscribe((state) => {
     // Ensure traces are flushed after each compile in development mode
     flushAllTraces()
     teardownTraceSubscriber()
-    teardownHeapProfiler()
     return
   }
 
@@ -176,7 +175,6 @@ store.subscribe((state) => {
     // Ensure traces are flushed after each compile in development mode
     flushAllTraces()
     teardownTraceSubscriber()
-    teardownHeapProfiler()
     return
   }
 
@@ -207,5 +205,4 @@ store.subscribe((state) => {
   // Ensure traces are flushed after each compile in development mode
   flushAllTraces()
   teardownTraceSubscriber()
-  teardownHeapProfiler()
 })
diff --git a/packages/next/src/build/swc/generated-native.d.ts b/packages/next/src/build/swc/generated-native.d.ts
index af16dedb1e93b..df4881fda49a3 100644
--- a/packages/next/src/build/swc/generated-native.d.ts
+++ b/packages/next/src/build/swc/generated-native.d.ts
@@ -381,10 +381,6 @@ export interface NapiRewrite {
   missing?: Array<NapiRouteHas>
 }
 export declare function getTargetTriple(): string
-export declare function initHeapProfiler(): ExternalObject<RefCell>
-export declare function teardownHeapProfiler(
-  guardExternal: ExternalObject<RefCell>
-): void
 /**
  * Initialize tracing subscriber to emit traces. This configures subscribers
  * for Trace Event Format <https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview>.
diff --git a/packages/next/src/build/swc/index.ts b/packages/next/src/build/swc/index.ts
index da2f57f3fa1f7..7182a9c889727 100644
--- a/packages/next/src/build/swc/index.ts
+++ b/packages/next/src/build/swc/index.ts
@@ -1225,8 +1225,6 @@ function loadNative(importPath?: string) {
       getTargetTriple: bindings.getTargetTriple,
       initCustomTraceSubscriber: bindings.initCustomTraceSubscriber,
       teardownTraceSubscriber: bindings.teardownTraceSubscriber,
-      initHeapProfiler: bindings.initHeapProfiler,
-      teardownHeapProfiler: bindings.teardownHeapProfiler,
       turbo: {
         createProject: bindingToApi(customBindings ?? bindings, false),
         startTurbopackTraceServer(traceFilePath) {
@@ -1324,30 +1322,13 @@ export function getBinaryMetadata() {
  *
  */
 export function initCustomTraceSubscriber(traceFileName?: string) {
-  if (!swcTraceFlushGuard) {
+  if (swcTraceFlushGuard) {
     // Wasm binary doesn't support trace emission
     let bindings = loadNative()
     swcTraceFlushGuard = bindings.initCustomTraceSubscriber?.(traceFileName)
   }
 }
 
-/**
- * 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 function initHeapProfiler() {
-  try {
-    if (!swcHeapProfilerFlushGuard) {
-      let bindings = loadNative()
-      swcHeapProfilerFlushGuard = bindings.initHeapProfiler?.()
-    }
-  } catch (_) {
-    // Suppress exceptions, this fn allows to fail to load native bindings
-  }
-}
-
 function once(fn: () => void): () => void {
   let executed = false
 
@@ -1360,23 +1341,6 @@ function once(fn: () => void): () => void {
   }
 }
 
-/**
- * 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 = once(() => {
-  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.
  *
diff --git a/packages/next/src/build/swc/types.ts b/packages/next/src/build/swc/types.ts
index d79b05e70412a..67af229ca489b 100644
--- a/packages/next/src/build/swc/types.ts
+++ b/packages/next/src/build/swc/types.ts
@@ -27,8 +27,6 @@ export interface Binding {
 
   initCustomTraceSubscriber?(traceOutFilePath?: string): ExternalObject<RefCell>
   teardownTraceSubscriber?(guardExternal: ExternalObject<RefCell>): void
-  initHeapProfiler?(): ExternalObject<RefCell>
-  teardownHeapProfiler?(guardExternal: ExternalObject<RefCell>): void
   css: {
     lightning: {
       transform(transformOptions: any): Promise<any>
diff --git a/turbopack/crates/turbo-tasks/Cargo.toml b/turbopack/crates/turbo-tasks/Cargo.toml
index fb521d5fb645c..ec74381ed3018 100644
--- a/turbopack/crates/turbo-tasks/Cargo.toml
+++ b/turbopack/crates/turbo-tasks/Cargo.toml
@@ -12,6 +12,7 @@ bench = false
 default = []
 tokio_tracing = ["tokio/tracing"]
 hanging_detection = []
+local_resolution = []
 
 [lints]
 workspace = true