refactor: wip

Author: BioPhoton

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

@@ -1,3 +1,5 @@
+import path from 'node:path';
+import { performance } from 'node:perf_hooks';
 import type { PerformanceEntry } from 'node:perf_hooks';
 import process from 'node:process';
 import { threadId } from 'node:worker_threads';
@@ -26,7 +28,14 @@ import {
 import { entryToTraceEvents } from './trace-file-utils.js';
 import type { UserTimingTraceEvent } from './trace-file.type.js';
 import { traceEventWalFormat } from './wal-json-trace.js';
-import { ShardedWal, WriteAheadLogFile } from './wal.js';
+import {
+  ShardedWal,
+  WriteAheadLogFile,
+  getShardId,
+  getShardedGroupId,
+  isLeaderWal,
+  setLeaderWal,
+} from './wal.js';
 import type { WalFormat } from './wal.js';
 
 /**
@@ -80,11 +89,6 @@ export class Profiler<T extends ActionTrackConfigs> {
    *
    */
   constructor(options: ProfilerOptions<T>) {
-    // Initialize origin PID early - must happen before user code runs
-    if (!process.env[PROFILER_ORIGIN_PID_ENV_VAR]) {
-      process.env[PROFILER_ORIGIN_PID_ENV_VAR] = String(process.pid);
-    }
-
     const { tracks, prefix, enabled, ...defaults } = options;
     const dataType = 'track-entry';
 
@@ -113,6 +117,13 @@ export class Profiler<T extends ActionTrackConfigs> {
     this.#enabled = enabled;
   }
 
+  /**
+   * Close the profiler. Subclasses should override this to perform cleanup.
+   */
+  close(): void {
+    // Base implementation does nothing
+  }
+
   /**
    * Is profiling enabled?
    *
@@ -235,74 +246,103 @@ export class Profiler<T extends ActionTrackConfigs> {
   }
 }
 
-/**
- * Determines if this process is the leader WAL process using the origin PID heuristic.
- *
- * The leader is the process that first enabled profiling (the one that set CP_PROFILER_ORIGIN_PID).
- * All descendant processes inherit the environment but have different PIDs.
- *
- * @returns true if this is the leader WAL process, false otherwise
- */
-export function isLeaderWal(): boolean {
-  return process.env[PROFILER_ORIGIN_PID_ENV_VAR] === String(process.pid);
-}
-
 export class NodeProfiler<
   TracksConfig extends ActionTrackConfigs = ActionTrackConfigs,
-  CodecOutput extends string | object = UserTimingTraceEvent,
+  CodecOutput extends UserTimingTraceEvent = UserTimingTraceEvent,
 > extends Profiler<TracksConfig> {
   #shard: WriteAheadLogFile<CodecOutput>;
   #perfObserver: PerformanceObserverSink<CodecOutput>;
   #shardWal: ShardedWal<CodecOutput>;
   readonly #format: WalFormat<CodecOutput>;
+  readonly #debug: boolean;
+  #closed: boolean = false;
+
   constructor(
     options: ProfilerOptions<TracksConfig> & {
       directory?: string;
       performanceEntryEncode: (entry: PerformanceEntry) => CodecOutput[];
-      format: WalFormat<CodecOutput>;
+      debug?: boolean;
     },
   ) {
+    // Initialize origin PID early - must happen before user code runs
+    setLeaderWal(PROFILER_ORIGIN_PID_ENV_VAR);
+
     const {
       directory = PROFILER_DIRECTORY,
       performanceEntryEncode,
-      format,
+      debug = false,
+      ...profilerOptions
     } = options;
-    super(options);
-    const shardId = `${process.pid}-${threadId}`;
+    super(profilerOptions);
+    const walGroupId = getShardedGroupId();
+    const shardId = getShardId(process.pid, threadId);
 
-    this.#format = format;
-    this.#shardWal = new ShardedWal(directory, format);
+    this.#format = traceEventWalFormat({ groupId: walGroupId });
+    this.#debug = debug;
+    this.#shardWal = new ShardedWal(
+      path.join(directory, walGroupId),
+      this.#format,
+    );
     this.#shard = this.#shardWal.shard(shardId);
 
     this.#perfObserver = new PerformanceObserverSink({
       sink: this.#shard,
       encode: performanceEntryEncode,
       buffered: true,
-      flushThreshold: 100,
+      flushThreshold: 1, // Lower threshold for immediate flushing
     });
 
+    this.#perfObserver.subscribe();
+
     installExitHandlers({
       onExit: () => {
-        this.#perfObserver.flush();
-        this.#perfObserver.unsubscribe();
-        this.#shard.close();
-        if (isLeaderWal()) {
-          this.#shardWal.finalize();
-          this.#shardWal.cleanup();
-        }
+        this.close();
       },
     });
   }
 
   getFinalPath() {
     return this.#format.finalPath();
   }
+
+  /**
+   * Close the profiler and finalize files if this is the leader process.
+   * This method can be called manually to ensure proper cleanup.
+   */
+  close(): void {
+    if (this.#closed) {
+      return;
+    }
+
+    this.#closed = true;
+
+    try {
+      if (!this.#perfObserver || !this.#shard || !this.#shardWal) {
+        console.warn('Warning: Profiler not fully initialized during close');
+        return;
+      }
+
+      this.#perfObserver.flush();
+      this.#perfObserver.unsubscribe();
+
+      this.#shard.close();
+
+      if (isLeaderWal(PROFILER_ORIGIN_PID_ENV_VAR)) {
+        this.#shardWal.finalize();
+        if (!this.#debug) {
+          this.#shardWal.cleanup();
+        }
+      }
+    } catch (error) {
+      console.warn('Warning: Error during profiler close:', error);
+    }
+  }
 }
 
 export const profiler = new NodeProfiler({
   prefix: 'cp',
   track: 'CLI',
   trackGroup: 'Code Pushup',
   performanceEntryEncode: entryToTraceEvents,
-  format: traceEventWalFormat(),
+  debug: process.env.CP_PROFILER_DEBUG === 'true',
 });

packages/utils/src/lib/profiler/trace-file-utils.ts

@@ -263,7 +263,9 @@ export const measureToSpanEvents = (
  * @param entry - Performance entry
  * @returns UserTimingTraceEvent[]
  */
-export function entryToTraceEvents(entry: PerformanceEntry) {
+export function entryToTraceEvents(
+  entry: PerformanceEntry,
+): UserTimingTraceEvent[] {
   if (entry.entryType === 'mark') {
     return [markToInstantEvent(entry as PerformanceMark)];
   }

packages/utils/src/lib/profiler/wal-json-trace.ts

@@ -71,12 +71,14 @@ export function generateTraceContent(
  */
 export const traceEventWalFormat = <
   T extends UserTimingTraceEvent = UserTimingTraceEvent,
->(_opt?: {
+>(opt?: {
   dir?: string;
+  groupId?: string;
 }) => {
   const baseName = 'trace';
   const walExtension = '.jsonl';
   const finalExtension = '.json';
+  const groupId = opt?.groupId || 'default';
   return {
     baseName,
     walExtension,
@@ -85,8 +87,8 @@ export const traceEventWalFormat = <
       encode: event => JSON.stringify(encodeTraceEvent(event)),
       decode: (json: string) => decodeTraceEvent(JSON.parse(json)) as T,
     },
-    shardPath: (id: string) => `${baseName}.${id}${walExtension}`,
-    finalPath: () => `${baseName}${finalExtension}`,
+    shardPath: (id: string) => `${baseName}.${groupId}.${id}${walExtension}`,
+    finalPath: () => `${baseName}.${groupId}${finalExtension}`,
     finalizer: (records, metadata) => generateTraceContent(records, metadata),
   } satisfies WalFormat<T>;
 };

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

@@ -1,5 +1,7 @@
 import * as fs from 'node:fs';
 import path from 'node:path';
+import process from 'node:process';
+import { PROFILER_ORIGIN_PID_ENV_VAR } from './constants';
 
 /**
  * Codec for encoding/decoding values to/from strings for WAL storage.
@@ -153,13 +155,15 @@ export class WriteAheadLogFile<T> {
   /**
    * Append a record to the WAL.
    * @param v - Record to append
-   * @throws Error if WAL is not opened
+   * @throws Error if WAL cannot be opened
    */
   append = (v: T) => {
     if (!this.#fd) {
-      throw new Error('WAL not opened');
+      this.open();
+    }
+    if (this.#fd) {
+      fs.writeSync(this.#fd, `${this.#encode(v)}\n`);
     }
-    fs.writeSync(this.#fd, `${this.#encode(v)}\n`);
   };
 
   /** Close the WAL file */
@@ -278,6 +282,29 @@ export function parseWalFormat<T extends object | string = object>(
   } satisfies WalFormat<T>;
 }
 
+/**
+ * Determines if this process is the leader WAL process using the origin PID heuristic.
+ *
+ * The leader is the process that first enabled profiling (the one that set CP_PROFILER_ORIGIN_PID).
+ * All descendant processes inherit the environment but have different PIDs.
+ *
+ * @returns true if this is the leader WAL process, false otherwise
+ */
+export function isLeaderWal(envVarName: string): boolean {
+  return process.env[envVarName] === String(process.pid);
+}
+
+/**
+ * Initialize the origin PID environment variable if not already set.
+ * This must be done as early as possible before any user code runs.
+ * Set's PROFILER_ORIGIN_PID_ENV_VAR to the current process PID if not already defined.
+ */
+export function setLeaderWal(PROFILER_ORIGIN_PID_ENV_VAR: string): void {
+  if (!process.env[PROFILER_ORIGIN_PID_ENV_VAR]) {
+    process.env[PROFILER_ORIGIN_PID_ENV_VAR] = String(process.pid);
+  }
+}
+
 /**
  * Sharded Write-Ahead Log manager for coordinating multiple WAL shards.
  * Handles distributed logging across multiple processes/files with atomic finalization.
@@ -304,16 +331,21 @@ export class ShardedWal<T extends object | string = object> {
 
   /** Get all shard file paths matching this WAL's base name */
   private shardFiles() {
-    return fs.existsSync(this.#dir)
-      ? fs
-          .readdirSync(this.#dir)
-          .filter(
-            f =>
-              f.startsWith(`${this.#format.baseName}.`) &&
-              f.endsWith(this.#format.walExtension),
-          )
-          .map(f => path.join(this.#dir, f))
-      : [];
+    if (!fs.existsSync(this.#dir)) {
+      return [];
+    }
+
+    const files: string[] = [];
+    const entries = fs.readdirSync(this.#dir);
+
+    for (const entry of entries) {
+      // Look for files matching the pattern: anything ending with .jsonl
+      if (entry.endsWith(this.#format.walExtension)) {
+        files.push(path.join(this.#dir, entry));
+      }
+    }
+
+    return files;
   }
 
   /**
@@ -351,6 +383,32 @@ export class ShardedWal<T extends object | string = object> {
   }
 
   cleanup() {
-    this.shardFiles().forEach(f => fs.unlinkSync(f));
+    this.shardFiles().forEach(f => {
+      // Remove the shard file
+      fs.unlinkSync(f);
+      // Remove the parent directory (shard group directory)
+      const shardDir = path.dirname(f);
+      try {
+        fs.rmdirSync(shardDir);
+      } catch (error) {
+        // Directory might not be empty or already removed, ignore
+      }
+    });
   }
 }
+
+/**
+ * Generates a shard ID.
+ * This is idempotent since PID and TID are fixed for the process/thread.
+ */
+export function getShardId(pid: number, tid: number = 0): string {
+  return `${pid}-${tid}`;
+}
+
+/**
+ * Generates a sharded group ID based on performance.timeOrigin.
+ * This is idempotent per process since timeOrigin is fixed within a process and its worker.
+ */
+export function getShardedGroupId(): string {
+  return Math.floor(performance.timeOrigin).toString();
+}