Write out task statistics if NEXT_TURBOPACK_TASK_STATISTICS is set

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

@@ -1,4 +1,4 @@
-use std::{path::PathBuf, sync::Arc, thread, time::Duration};
+use std::{io::Write, path::PathBuf, sync::Arc, thread, time::Duration};
 
 use anyhow::{anyhow, bail, Context, Result};
 use napi::{
@@ -39,9 +39,9 @@
         },
         ecmascript_hmr_protocol::{ClientUpdateInstruction, ResourceIdentifier},
         trace_utils::{
-            exit::ExitGuard,
+            exit::{ExitHandler, ExitReceiver},
             raw_trace::RawTraceLayer,
-            trace_writer::{TraceWriter, TraceWriterGuard},
+            trace_writer::TraceWriter,
         },
     },
 };
@@ -252,8 +252,7 @@
 pub struct ProjectInstance {
     turbo_tasks: Arc<TurboTasks<MemoryBackend>>,
     container: Vc<ProjectContainer>,
-    #[allow(dead_code)]
-    guard: Option<ExitGuard<TraceWriterGuard>>,
+    exit_receiver: tokio::sync::Mutex<Option<ExitReceiver>>,
 }
 
 #[napi(ts_return_type = "{ __napiType: \"Project\" }")]
@@ -264,8 +263,9 @@
     register();
 
     let trace = std::env::var("NEXT_TURBOPACK_TRACING").ok();
+    let (exit, exit_receiver) = ExitHandler::new_receiver();
 
-    let guard = if let Some(mut trace) = trace {
+    if let Some(mut trace) = trace {
         // Trace presets
         match trace.as_str() {
             "overview" | "1" => {
@@ -297,10 +297,12 @@
             .unwrap();
         let trace_file = internal_dir.join("trace.log");
         let trace_writer = std::fs::File::create(trace_file.clone()).unwrap();
-        let (trace_writer, guard) = TraceWriter::new(trace_writer);
+        let (trace_writer, trace_writer_guard) = TraceWriter::new(trace_writer);
         let subscriber = subscriber.with(RawTraceLayer::new(trace_writer));
 
-        let guard = ExitGuard::new(guard).unwrap();
+        exit.on_exit(async move {
+            tokio::task::spawn_blocking(move || drop(trace_writer_guard));
+        });
 
         let trace_server = std::env::var("NEXT_TURBOPACK_TRACE_SERVER").ok();
         if trace_server.is_some() {
@@ -313,18 +315,30 @@
         }
 
         subscriber.init();
-
-        Some(guard)
-    } else {
-        None
-    };
+    }
 
     let turbo_tasks = TurboTasks::new(MemoryBackend::new(
         turbo_engine_options
             .memory_limit
             .map(|m| m as usize)
             .unwrap_or(usize::MAX),
     ));
+    let stats_path = std::env::var_os("NEXT_TURBOPACK_TASK_STATISTICS");
+    if let Some(stats_path) = stats_path {
+        let task_stats = turbo_tasks.backend().task_statistics().enable().clone();
+        exit.on_exit(async move {
+            tokio::task::spawn_blocking(move || {
+                let mut file = std::fs::File::create(&stats_path)
+                    .with_context(|| format!("failed to create or open {stats_path:?}"))?;
+                serde_json::to_writer(&file, &task_stats)
+                    .context("failed to serialize or write task statistics")?;
+                file.flush().context("failed to flush file")
+            })
+            .await
+            .unwrap()
+            .unwrap();
+        });
+    }
     let options = options.into();
     let container = turbo_tasks
         .run_once(async move {
@@ -344,7 +358,7 @@
         ProjectInstance {
             turbo_tasks,
             container,
-            guard,
+            exit_receiver: tokio::sync::Mutex::new(Some(exit_receiver)),
         },
         100,
     ))
@@ -883,7 +897,7 @@
     }
 }
 
 /// Subscribes to lifecycle events of the compilation.
 /// Emits an [UpdateMessage::Start] event when any computation starts.
 /// Emits an [UpdateMessage::End] event when there was no computation for the
 /// specified time (`aggregation_ms`). The [UpdateMessage::End] event contains
@@ -1100,3 +1114,15 @@
 
     Ok(source)
 }
+
+/// Runs exit handlers for the project registered using the [`ExitHandler`] API.
+#[napi]
+pub async fn project_on_exit(
+    #[napi(ts_arg_type = "{ __napiType: \"Project\" }")] project: External<ProjectInstance>,
+) {
+    let exit_receiver = project.exit_receiver.lock().await.take();
+    exit_receiver
+        .expect("`project.onExitSync` must only be called once")
+        .run_exit_handler()
+        .await;
+}

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

@@ -677,6 +677,8 @@ export interface Project {
   updateInfoSubscribe(
     aggregationMs: number
   ): AsyncIterableIterator<TurbopackResult<UpdateMessage>>
+
+  onExit(): Promise<void>
 }
 
 export type Route =
@@ -1088,6 +1090,10 @@ function bindingToApi(
       )
       return subscription
     }
+
+    onExit(): Promise<void> {
+      return binding.projectOnExit(this._nativeProject)
+    }
   }
 
   class EndpointImpl implements Endpoint {

packages/next/src/server/dev/hot-reloader-turbopack.ts

@@ -153,6 +153,7 @@ export async function createHotReloaderTurbopack(
       memoryLimit: opts.nextConfig.experimental.turbo?.memoryLimit,
     }
   )
+  opts.onCleanup(() => project.onExit())
   const entrypointsSubscription = project.entrypointsSubscribe()
 
   const currentEntrypoints: Entrypoints = {

packages/next/src/server/lib/router-server.ts

@@ -65,6 +65,7 @@ export async function initialize(opts: {
   dir: string
   port: number
   dev: boolean
+  onCleanup: (listener: () => Promise<void>) => void
   server?: import('http').Server
   minimalMode?: boolean
   hostname?: string
@@ -132,6 +133,7 @@ export async function initialize(opts: {
         isCustomServer: opts.customServer,
         turbo: !!process.env.TURBOPACK,
         port: opts.port,
+        onCleanup: opts.onCleanup,
       })
     )
 

packages/next/src/server/lib/router-utils/setup-dev-bundler.ts

@@ -96,6 +96,7 @@ export type SetupOpts = {
   >
   nextConfig: NextConfigComplete
   port: number
+  onCleanup: (listener: () => Promise<void>) => void
 }
 
 export type ServerFields = {

packages/next/src/server/lib/start-server.ts

@@ -50,6 +50,7 @@ export async function getRequestHandlers({
   dir,
   port,
   isDev,
+  onCleanup,
   server,
   hostname,
   minimalMode,
@@ -60,6 +61,7 @@ export async function getRequestHandlers({
   dir: string
   port: number
   isDev: boolean
+  onCleanup: (listener: () => Promise<void>) => void
   server?: import('http').Server
   hostname?: string
   minimalMode?: boolean
@@ -71,6 +73,7 @@ export async function getRequestHandlers({
     dir,
     port,
     hostname,
+    onCleanup,
     dev: isDev,
     minimalMode,
     server,
@@ -266,9 +269,13 @@ export async function startServer(
       Log.event(`Starting...`)
 
       try {
+        const cleanupListeners = [() => new Promise((res) => server.close(res))]
         const cleanup = () => {
           debug('start-server process cleanup')
-          server.close(() => process.exit(0))
+          ;(async () => {
+            await Promise.all(cleanupListeners.map((f) => f()))
+            process.exit(0)
+          })()
         }
         const exception = (err: Error) => {
           if (isPostpone(err)) {
@@ -298,6 +305,7 @@ export async function startServer(
           dir,
           port,
           isDev,
+          onCleanup: (listener) => cleanupListeners.push(listener),
           server,
           hostname,
           minimalMode,

packages/next/src/server/next.ts

@@ -62,6 +62,7 @@ export class NextServer {
   private reqHandlerPromise?: Promise<NodeRequestHandler>
   private preparedAssetPrefix?: string
 
+  protected cleanupListeners: (() => Promise<void>)[] = []
   protected standaloneMode?: boolean
 
   public options: NextServerOptions
@@ -155,8 +156,15 @@ export class NextServer {
   }
 
   async close() {
-    const server = await this.getServer()
-    return (server as any).close()
+    await Promise.all(
+      [
+        async () => {
+          const server = await this.getServer()
+          await (server as any).close()
+        },
+        ...this.cleanupListeners,
+      ].map((f) => f())
+    )
   }
 
   private async createServer(
@@ -278,6 +286,7 @@ class NextCustomServer extends NextServer {
       dir: this.options.dir!,
       port: this.options.port || 3000,
       isDev: !!this.options.dev,
+      onCleanup: (listener) => this.cleanupListeners.push(listener),
       hostname: this.options.hostname || 'localhost',
       minimalMode: this.options.minimalMode,
       quiet: this.options.quiet,