feat: add support for tracing rspack (#1150)

* feat: add support for profiling rspack

* fix: shutdown gracefully

* chore: changeset

* fix: tests

Author: jbroma

.changeset/shy-worlds-type.md

@@ -0,0 +1,5 @@
+---
+"@callstack/repack": minor
+---
+
+Support `RSPACK_PROFILE` env var for obtaining traces from Rspack

packages/repack/package.json

@@ -84,6 +84,7 @@
     "estree-util-is-identifier-name": "^1.1.0",
     "events": "^3.3.0",
     "execa": "^5.0.0",
+    "exit-hook": "^4.0.0",
     "flow-remove-types": "^2.268.0",
     "gradient-string": "^2.0.2",
     "image-size": "^1.1.1",

packages/repack/src/commands/common/__tests__/setupInteractions.test.ts

@@ -23,6 +23,11 @@ describe('setupInteractions', () => {
     } as unknown as Logger;
 
     mockProcess = {
+      on: (event: string) => {
+        if (event === 'SIGINT') {
+          mockProcess.exit();
+        }
+      },
       stdin: {
         setRawMode: jest.fn(),
         on: jest.fn(),

packages/repack/src/commands/common/setupInteractions.ts

@@ -41,12 +41,17 @@ export function setupInteractions(
   readline.emitKeypressEvents(process.stdin);
   process.stdin.setRawMode(true);
 
+  // graceful shutdown
+  process.on('SIGINT', () => {
+    process.exit();
+  });
+
   process.stdin.on('keypress', (_key, data) => {
     const { ctrl, name } = data;
     if (ctrl === true) {
       switch (name) {
         case 'c':
-          process.exit();
+          process.emit('SIGINT', 'SIGINT');
           break;
 
         case 'z':

packages/repack/src/commands/rspack/profile.ts

@@ -0,0 +1,75 @@
+/**
+ * Reference: https://github.com/web-infra-dev/rspack/blob/bad70431988d77d8a95320066f72c77b6cc4c120/packages/rspack-cli/src/utils/profile.ts
+ */
+
+/**
+ * `RSPACK_PROFILE=ALL` // all trace events
+ * `RSPACK_PROFILE=OVERVIEW` // overview trace events
+ * `RSPACK_PROFILE=warn,tokio::net=info` // trace filter from  https://docs.rs/tracing-subscriber/latest/tracing_subscriber/filter/struct.EnvFilter.html#example-syntax
+ */
+import fs from 'node:fs';
+import path from 'node:path';
+import { rspack } from '@rspack/core';
+
+const overviewTraceFilter = 'info';
+const allTraceFilter = 'trace';
+const defaultRustTraceLayer = 'chrome';
+
+function resolveLayer(value: string) {
+  if (value === 'OVERVIEW') {
+    return overviewTraceFilter;
+  }
+  if (value === 'ALL') {
+    return allTraceFilter;
+  }
+  return value;
+}
+
+export async function applyProfile(
+  filterValue: string,
+  traceLayer: string = defaultRustTraceLayer,
+  traceOutput?: string
+) {
+  const { asyncExitHook } = await import('exit-hook');
+
+  if (traceLayer !== 'chrome' && traceLayer !== 'logger') {
+    throw new Error(`unsupported trace layer: ${traceLayer}`);
+  }
+
+  if (!traceOutput) {
+    const timestamp = Date.now();
+    const defaultOutputDir = path.resolve(
+      `.rspack-profile-${timestamp}-${process.pid}`
+    );
+    const defaultRustTraceChromeOutput = path.join(
+      defaultOutputDir,
+      'trace.json'
+    );
+    const defaultRustTraceLoggerOutput = 'stdout';
+
+    const defaultTraceOutput =
+      traceLayer === 'chrome'
+        ? defaultRustTraceChromeOutput
+        : defaultRustTraceLoggerOutput;
+
+    // biome-ignore lint/style/noParameterAssign: setting default value makes sense
+    traceOutput = defaultTraceOutput;
+  }
+
+  const filter = resolveLayer(filterValue);
+
+  await ensureFileDir(traceOutput);
+  await rspack.experiments.globalTrace.register(
+    filter,
+    traceLayer,
+    traceOutput
+  );
+  asyncExitHook(rspack.experiments.globalTrace.cleanup, {
+    wait: 500,
+  });
+}
+
+async function ensureFileDir(outputFilePath: string) {
+  const dir = path.dirname(outputFilePath);
+  await fs.promises.mkdir(dir, { recursive: true });
+}

packages/repack/src/commands/rspack/start.ts

@@ -73,6 +73,15 @@ export async function start(
     });
   }
 
+  if (process.env.RSPACK_PROFILE) {
+    const { applyProfile } = await import('./profile.js');
+    await applyProfile(
+      process.env.RSPACK_PROFILE,
+      process.env.RSPACK_TRACE_LAYER,
+      process.env.RSPACK_TRACE_OUTPUT
+    );
+  }
+
   const compiler = new Compiler(configs, reporter, cliConfig.root);
 
   const { createServer } = await import('@callstack/repack-dev-server');