fix(next-swc): explicitly flush traces when exit process

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

@@ -24,9 +24,11 @@ swc_ecma_loader = { version = "0.29.0", features = ["node", "lru"] }
 swc_ecmascript = { version = "0.140.0", features = ["codegen", "minifier", "optimization", "parser", "react", "transforms", "typescript", "utils", "visit"] }
 swc_node_base = "0.5.1"
 tracing = { version = "0.1.32", features = ["release_max_level_info"] }
-tracing-chrome = "0.5.0"
 tracing-futures = "0.2.5"
 tracing-subscriber = "0.3.9"
 
+# https://github.com/thoren-d/tracing-chrome/pull/10
+tracing-chrome = { git = "https://github.com/kwonoj/tracing-chrome", rev = "0dd9d6c9e0f74f43993b58560f4ac0103e058df8" }
+
 [build-dependencies]
 napi-build = "1"

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

@@ -77,6 +77,7 @@ fn init(mut exports: JsObject) -> napi::Result<()> {
         "initCustomTraceSubscriber",
         util::init_custom_trace_subscriber,
     )?;
+    exports.create_named_method("teardownTraceSubscriber", util::teardown_trace_subscriber)?;
 
     Ok(())
 }

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

@@ -27,16 +27,13 @@ DEALINGS IN THE SOFTWARE.
 */
 
 use anyhow::{Context, Error};
-use napi::{CallContext, Env, JsBuffer, JsString, JsUndefined, JsUnknown, Status};
-use once_cell::sync::OnceCell;
+use napi::{CallContext, Env, JsBuffer, JsExternal, JsString, JsUndefined, JsUnknown, Status};
 use serde::de::DeserializeOwned;
 use std::{any::type_name, convert::TryFrom, path::PathBuf};
-use tracing_chrome::ChromeLayerBuilder;
+use tracing_chrome::{ChromeLayerBuilder, FlushGuard};
 use tracing_subscriber::{filter, prelude::*, util::SubscriberInitExt, Layer};
 
 static TARGET_TRIPLE: &str = include_str!(concat!(env!("OUT_DIR"), "/triple.txt"));
-static CUSTOM_TRACE_SUBSCRIBER: OnceCell<bool> = OnceCell::new();
-
 #[contextless_function]
 pub fn get_target_triple(env: Env) -> napi::ContextlessResult<JsString> {
     env.create_string(TARGET_TRIPLE).map(Some)
@@ -97,7 +94,7 @@ where
 /// Initialize tracing subscriber to emit traces. This configures subscribers
 /// for Trace Event Format (https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview).
 #[js_function(1)]
-pub fn init_custom_trace_subscriber(cx: CallContext) -> napi::Result<JsUndefined> {
+pub fn init_custom_trace_subscriber(cx: CallContext) -> napi::Result<JsExternal> {
     let optional_trace_out_file_path = cx.get::<JsUnknown>(0)?;
     let trace_out_file_path = match optional_trace_out_file_path.get_type()? {
         napi::ValueType::String => Some(PathBuf::from(
@@ -109,29 +106,30 @@ pub fn init_custom_trace_subscriber(cx: CallContext) -> napi::Result<JsUndefined
         _ => None,
     };
 
-    CUSTOM_TRACE_SUBSCRIBER.get_or_init(|| {
-        let mut layer = ChromeLayerBuilder::new().include_args(true);
-        if let Some(trace_out_file) = trace_out_file_path {
-            layer = layer.file(trace_out_file);
-        }
-
-        let (chrome_layer, guard) = layer.build();
-        tracing_subscriber::registry()
-            .with(chrome_layer.with_filter(filter::filter_fn(|metadata| {
-                !metadata.target().contains("cranelift") && !metadata.name().contains("log ")
-            })))
-            .try_init()
-            .expect("Failed to register tracing subscriber");
-
-        cx.env
-            .add_env_cleanup_hook(guard, |flush_guard| {
-                flush_guard.flush();
-                drop(flush_guard);
-            })
-            .expect("Should able to register custom trace subscriber cleanup");
+    let mut layer = ChromeLayerBuilder::new().include_args(true);
+    if let Some(trace_out_file) = trace_out_file_path {
+        layer = layer.file(trace_out_file);
+    }
+
+    let (chrome_layer, guard) = layer.build();
+    tracing_subscriber::registry()
+        .with(chrome_layer.with_filter(filter::filter_fn(|metadata| {
+            !metadata.target().contains("cranelift") && !metadata.name().contains("log ")
+        })))
+        .try_init()
+        .expect("Failed to register tracing subscriber");
 
-        true
-    });
+    cx.env.create_external(guard, None)
+}
+
+/// Teardown currently running tracing subscriber to flush out remaining traces.
+/// This should be called when parent node.js process exits, otherwise generated
+/// trace will missing traces in the buffer.
+#[js_function(1)]
+pub fn teardown_trace_subscriber(cx: CallContext) -> napi::Result<JsUndefined> {
+    let guard_external = cx.get::<JsExternal>(0)?;
+    let guard = &*cx.env.get_value_external::<FlushGuard>(&guard_external)?;
 
+    guard.close();
     cx.env.get_undefined()
 }

packages/next/build/swc/index.js

@@ -193,6 +193,7 @@ function loadNative() {
 
       getTargetTriple: bindings.getTargetTriple,
       initCustomTraceSubscriber: bindings.initCustomTraceSubscriber,
+      teardownTraceSubscriber: bindings.teardownTraceSubscriber,
     }
     return nativeBindings
   }
@@ -253,8 +254,34 @@ export function getBinaryMetadata() {
   }
 }
 
-export function initCustomTraceSubscriber(filename) {
-  // Wasm binary doesn't support trace emission
-  let bindings = loadNative()
-  return bindings.initCustomTraceSubscriber(filename)
-}
+/**
+ * Initialize trace subscriber to emit traces.
+ *
+ * Returns an internal object to guard async flush emission if subscriber is initialized, caller should manually
+ * tear it down via `teardownTraceSubscriber`.
+ */
+export const initCustomTraceSubscriber = (() => {
+  let guard
+
+  return (filename) => {
+    if (!guard) {
+      // Wasm binary doesn't support trace emission
+      let bindings = loadNative()
+      guard = bindings.initCustomTraceSubscriber(filename)
+    }
+
+    return guard
+  }
+})()
+
+export const teardownTraceSubscriber = (() => {
+  let bindings
+
+  return (guard) => {
+    if (!bindings && !!guard) {
+      // Wasm binary doesn't support trace emission
+      bindings = loadNative()
+      return bindings.teardownTraceSubscriber(guard)
+    }
+  }
+})()

packages/next/build/webpack-config.ts

@@ -88,6 +88,22 @@ const devtoolRevertWarning = execOnce(
 
 let loggedSwcDisabled = false
 let loggedIgnoredCompilerOptions = false
+let swcTraceFlushGuard: unknown = null
+
+/**
+ * Teardown swc's trace subscriber if there's an initialized flush guard exists.
+ *
+ * This is workaround to amend behavior with process.exit
+ * (https://github.com/vercel/next.js/blob/4db8c49cc31e4fc182391fae6903fb5ef4e8c66e/packages/next/bin/next.ts#L134=)
+ * seems preventing napi's cleanup hook execution (https://github.com/swc-project/swc/blob/main/crates/node/src/util.rs#L48-L51=),
+ *
+ * instead parent process manually drops guard when process gets signal to exit.
+ */
+process.on('exit', () => {
+  if (swcTraceFlushGuard) {
+    require('./swc')?.teardownTraceSubscriber?.(swcTraceFlushGuard)
+  }
+})
 
 function getOptimizedAliases(): { [pkg: string]: string } {
   const stubWindowFetch = path.join(__dirname, 'polyfills', 'fetch', 'index.js')
@@ -439,6 +455,20 @@ export default async function getBaseWebpackConfig(
   }
 
   const getBabelOrSwcLoader = (isMiddleware: boolean) => {
+    if (
+      useSWCLoader &&
+      config?.experimental?.swcTrace?.enabled &&
+      !swcTraceFlushGuard
+    ) {
+      // This will init subscribers once only in a single process lifecycle,
+      // even though it can be called multiple times.
+      // Subscriber need to be initialized _before_ any actual swc's call (transform, etcs)
+      // to collect correct trace spans when they are called.
+      swcTraceFlushGuard = require('./swc')?.initCustomTraceSubscriber?.(
+        config?.experimental?.swcTrace?.traceFileName
+      )
+    }
+
     return useSWCLoader
       ? {
           loader: 'next-swc-loader',

packages/next/build/webpack/loaders/next-swc-loader.js

@@ -26,7 +26,7 @@ IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 DEALINGS IN THE SOFTWARE.
 */
 
-import { isWasm, transform, initCustomTraceSubscriber } from '../../swc'
+import { isWasm, transform } from '../../swc'
 import { getLoaderSWCOptions } from '../../swc/options'
 import { isAbsolute } from 'path'
 
@@ -51,14 +51,6 @@ async function loaderTransform(parentTrace, source, inputSourceMap) {
     jsConfig,
   })
 
-  if (nextConfig?.experimental?.swcTrace?.enabled) {
-    // This will init subscribers once only in a single process lifecycle,
-    // even though it can be called multiple times.
-    // Subscriber need to be initialized _before_ any actual swc's call (transform, etcs)
-    // to collect correct trace spans when they are called.
-    initCustomTraceSubscriber(nextConfig?.experimental?.swcTrace?.traceFileName)
-  }
-
   const programmaticOptions = {
     ...swcOptions,
     filename,