refactor: impl unsub

Author: BioPhoton

packages/utils/src/lib/exit-process.ts

@@ -38,7 +38,19 @@ export type ExitHandlerOptions = {
   fatalExitCode?: number;
 };
 
-export function installExitHandlers(options: ExitHandlerOptions = {}): void {
+/**
+ *
+ * @param options - Options for the exit handler
+ * @param options.onExit - Callback to be called when the process exits
+ * @param options.onError - Callback to be called when an error occurs
+ * @param options.exitOnFatal - Whether to exit the process on fatal errors
+ * @param options.exitOnSignal - Whether to exit the process on signals
+ * @param options.fatalExitCode - The exit code to use for fatal errors
+ * @returns A function to unsubscribe from the exit handlers
+ */
+export function subscribeProcessExit(
+  options: ExitHandlerOptions = {},
+): () => void {
   // eslint-disable-next-line functional/no-let
   let closedReason: CloseReason | undefined;
   const {
@@ -57,40 +69,57 @@ export function installExitHandlers(options: ExitHandlerOptions = {}): void {
     onExit?.(code, reason);
   };
 
-  process.on('uncaughtException', err => {
+  const uncaughtExceptionHandler = (err: unknown) => {
     onError?.(err, 'uncaughtException');
     if (exitOnFatal) {
       close(fatalExitCode, {
         kind: 'fatal',
         fatal: 'uncaughtException',
       });
     }
-  });
+  };
 
-  process.on('unhandledRejection', reason => {
+  const unhandledRejectionHandler = (reason: unknown) => {
     onError?.(reason, 'unhandledRejection');
     if (exitOnFatal) {
       close(fatalExitCode, {
         kind: 'fatal',
         fatal: 'unhandledRejection',
       });
     }
-  });
+  };
 
-  (['SIGINT', 'SIGTERM', 'SIGQUIT'] as const).forEach(signal => {
-    process.on(signal, () => {
-      close(SIGNAL_EXIT_CODES()[signal], { kind: 'signal', signal });
-      if (exitOnSignal) {
-        // eslint-disable-next-line n/no-process-exit
-        process.exit(SIGNAL_EXIT_CODES()[signal]);
-      }
-    });
-  });
+  const signalHandlers = (['SIGINT', 'SIGTERM', 'SIGQUIT'] as const).map(
+    signal => {
+      const handler = () => {
+        close(SIGNAL_EXIT_CODES()[signal], { kind: 'signal', signal });
+        if (exitOnSignal) {
+          // eslint-disable-next-line n/no-process-exit
+          process.exit(SIGNAL_EXIT_CODES()[signal]);
+        }
+      };
+      process.on(signal, handler);
+      return { signal, handler };
+    },
+  );
 
-  process.on('exit', code => {
+  const exitHandler = (code: number) => {
     if (closedReason) {
       return;
     }
     close(code, { kind: 'exit' });
-  });
+  };
+
+  process.on('uncaughtException', uncaughtExceptionHandler);
+  process.on('unhandledRejection', unhandledRejectionHandler);
+  process.on('exit', exitHandler);
+
+  return () => {
+    process.removeListener('uncaughtException', uncaughtExceptionHandler);
+    process.removeListener('unhandledRejection', unhandledRejectionHandler);
+    process.removeListener('exit', exitHandler);
+    signalHandlers.forEach(({ signal, handler }) => {
+      process.removeListener(signal, handler);
+    });
+  };
 }

packages/utils/src/lib/profiler/profiler.ts

@@ -1,6 +1,6 @@
 import process from 'node:process';
 import { isEnvVarEnabled } from '../env.js';
-import { installExitHandlers } from '../exit-process';
+import { subscribeProcessExit } from '../exit-process';
 import type { TraceEvent } from '../trace-file.type';
 import {
   type ActionTrackConfigs,
@@ -16,7 +16,6 @@ import type {
   ActionTrackEntryPayload,
   DevToolsColor,
   EntryMeta,
-  UserTimingDetail,
 } from '../user-timing-extensibility-api.type.js';
 import { PROFILER_ENABLED_ENV_VAR } from './constants.js';
 
@@ -236,6 +235,7 @@ type WalSink = {
   append(event: TraceEvent): void;
   open(): void;
   close(): void;
+  isClosed(): boolean;
 };
 
 export type NodeJsProfilerOptions<T extends ActionTrackConfigs> =
@@ -247,6 +247,7 @@ export type NodeJsProfilerOptions<T extends ActionTrackConfigs> =
   };
 
 export class NodeJsProfiler<T extends ActionTrackConfigs> extends Profiler<T> {
+  #exitHandlerSubscribscription: null | (() => void) = null;
   protected sink: WalSink | null = null;
 
   constructor(options: NodeJsProfilerOptions<T>) {
@@ -258,8 +259,9 @@ export class NodeJsProfiler<T extends ActionTrackConfigs> extends Profiler<T> {
       },
       open: () => void 0,
       close: () => void 0,
+      isClosed: () => false,
     };
-    this.installExitHandlers();
+    this.#exitHandlerSubscribscription = this.subscribeProcessExit();
   }
 
   /**
@@ -270,27 +272,36 @@ export class NodeJsProfiler<T extends ActionTrackConfigs> extends Profiler<T> {
    *
    * @protected
    */
-  protected installExitHandlers(): void {
-    installExitHandlers({
+  protected subscribeProcessExit(): () => void {
+    return subscribeProcessExit({
       onError: (err, kind) => {
-        if (!this.isEnabled()) {
+        if (!this.isRunning()) {
           return;
         }
         this.marker('Fatal Error', {
           ...errorToMarkerPayload(err),
           tooltipText: `${kind} caused fatal error`,
         });
-        this.shutdown();
+        this.close();
       },
-      onExit: () => {
-        if (!this.isEnabled()) {
+      onExit: (code, reason) => {
+        if (!this.isRunning()) {
           return;
         }
-        this.shutdown();
+        this.marker('Process Exit', {
+          ...(code !== 0 ? { color: 'warning' } : {}),
+          properties: [['reason', JSON.stringify(reason)]],
+          tooltipText: `Process exited with code ${code}`,
+        });
+        this.close();
       },
     });
   }
 
+  isRunning(): boolean {
+    return this.isEnabled() && !this.sink?.isClosed();
+  }
+
   override setEnabled(enabled: boolean): void {
     super.setEnabled(enabled);
     enabled ? this.sink?.open() : this.sink?.close();
@@ -304,7 +315,10 @@ export class NodeJsProfiler<T extends ActionTrackConfigs> extends Profiler<T> {
    * data is flushed and the WAL sink is properly closed.
    */
   close(): void {
-    this.shutdown();
+    if (!this.isEnabled()) return;
+    this.flush();
+    this.setEnabled(false);
+    this.#exitHandlerSubscribscription?.();
   }
 
   /**
@@ -313,18 +327,4 @@ export class NodeJsProfiler<T extends ActionTrackConfigs> extends Profiler<T> {
   flush(): void {
     // @TODO implement WAL flush, currently all entries are buffered in memory
   }
-
-  /**
-   * Performs internal cleanup of profiling resources.
-   *
-   * Flushes any remaining buffered data and closes the WAL sink.
-   * This method is called automatically on process exit or error.
-   *
-   * @protected
-   */
-  protected shutdown(): void {
-    if (!this.isEnabled()) return;
-    this.flush();
-    this.setEnabled(false);
-  }
 }

packages/utils/src/lib/profiler/profiler.unit.test.ts

@@ -1,10 +1,10 @@
 import { performance } from 'node:perf_hooks';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
-import { installExitHandlers } from '../exit-process.js';
+import { subscribeProcessExit } from '../exit-process.js';
 import type { ActionTrackEntryPayload } from '../user-timing-extensibility-api.type.js';
 import { NodeJsProfiler, Profiler, type ProfilerOptions } from './profiler.js';
 
-// Spy on installExitHandlers to capture handlers
+// Spy on subscribeProcessExit to capture handlers
 vi.mock('../exit-process.js');
 
 describe('Profiler', () => {
@@ -429,7 +429,7 @@ describe('Profiler', () => {
   });
 });
 describe('NodeJsProfiler', () => {
-  const mockInstallExitHandlers = vi.mocked(installExitHandlers);
+  const mockSubscribeProcessExit = vi.mocked(subscribeProcessExit);
 
   let capturedOnError:
     | ((
@@ -456,9 +456,10 @@ describe('NodeJsProfiler', () => {
     capturedOnError = undefined;
     capturedOnExit = undefined;
 
-    mockInstallExitHandlers.mockImplementation(options => {
+    mockSubscribeProcessExit.mockImplementation(options => {
       capturedOnError = options?.onError;
       capturedOnExit = options?.onExit;
+      return vi.fn();
     });
 
     performance.clearMarks();
@@ -470,7 +471,7 @@ describe('NodeJsProfiler', () => {
   it('installs exit handlers on construction', () => {
     expect(() => createProfiler()).not.toThrow();
 
-    expect(mockInstallExitHandlers).toHaveBeenCalledWith({
+    expect(mockSubscribeProcessExit).toHaveBeenCalledWith({
       onError: expect.any(Function),
       onExit: expect.any(Function),
     });