refactor: wip

Author: BioPhoton

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

@@ -4,3 +4,5 @@ export const PROFILER_ORIGIN_PID_ENV_VAR = 'CP_PROFILER_ORIGIN_PID';
 export const PROFILER_DIRECTORY_ENV_VAR = 'CP_PROFILER_DIR';
 export const PROFILER_BASE_NAME = 'trace';
 export const PROFILER_DIRECTORY = './tmp/profiles';
+export const SHARDED_WAL_COORDINATOR_ID_ENV_VAR =
+  'CP_SHARDED_WAL_COORDINATOR_ID';

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

@@ -17,8 +17,8 @@ import type {
   DevToolsColor,
   EntryMeta,
 } from '../user-timing-extensibility-api.type.js';
+import { ShardedWal, type WalFormat, WriteAheadLogFile } from '../wal';
 import { PROFILER_ENABLED_ENV_VAR } from './constants.js';
-import { type TraceEvent } from './trace-file.type.js';
 
 /**
  * Generates a unique profiler ID based on performance time origin, process ID, thread ID, and instance count.
@@ -241,37 +241,28 @@ export class Profiler<T extends ActionTrackConfigs> {
   }
 }
 
-// @TODO implement ShardedWAL
-type WalSink = {
-  append: (event: TraceEvent) => void;
-  open: () => void;
-  close: () => void;
-  isClosed: () => boolean;
+export type ProfilerPersistOptions<DomainObject extends string | object> = {
+  format: WalFormat<DomainObject>;
 };
+export type NodeJsProfilerOptions<
+  T extends ActionTrackConfigs,
+  DomainObject extends string | object,
+> = ProfilerOptions<T> & ProfilerPersistOptions<DomainObject>;
 
-export type NodeJsProfilerOptions<T extends ActionTrackConfigs> =
-  ProfilerOptions<T> & {
-    // @TODO implement WALFormat
-    format: {
-      encode: (v: string | object) => string;
-    };
-  };
-
-export class NodeJsProfiler<T extends ActionTrackConfigs> extends Profiler<T> {
+export class NodeJsProfiler<
+  T extends ActionTrackConfigs,
+  DomainObject extends string | object,
+> extends Profiler<T> {
   #exitHandlerSubscription: null | (() => void) = null;
-  protected sink: WalSink | null = null;
+  #shardedWal: ShardedWal<DomainObject>;
+  #sink: WriteAheadLogFile<DomainObject> | null = null;
 
-  constructor(options: NodeJsProfilerOptions<T>) {
-    super(options);
-    // Temporary dummy sink; replaced by real WAL implementation
-    this.sink = {
-      append: event => {
-        options.format.encode(event);
-      },
-      open: () => void 0,
-      close: () => void 0,
-      isClosed: () => false,
-    };
+  constructor(options: NodeJsProfilerOptions<T, DomainObject>) {
+    const { format, ...profilerOptons } = options;
+    super(profilerOptons);
+
+    this.#shardedWal = new ShardedWal({ format });
+    this.#sink = this.#shardedWal.shard();
     this.#exitHandlerSubscription = this.subscribeProcessExit();
   }
 
@@ -317,12 +308,11 @@ export class NodeJsProfiler<T extends ActionTrackConfigs> extends Profiler<T> {
    * data is flushed and the WAL sink is properly closed.
    */
   close(): void {
-    if (!this.isEnabled()) {
-      return;
-    }
-    this.setEnabled(false);
     this.#exitHandlerSubscription?.();
     this.#exitHandlerSubscription = null;
-    this.sink?.close();
+    this.setEnabled(false);
+    this.#sink?.close();
+    this.#sink = null;
+    this.#shardedWal.finalizeIfCoordinator();
   }
 }

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

@@ -574,4 +574,18 @@ describe('NodeJsProfiler', () => {
     expect(setEnabledSpy).toHaveBeenCalledTimes(1);
     expect(setEnabledSpy).toHaveBeenCalledWith(false);
   });
+
+  it('close() unsubscribes from exit handlers even when disabled', () => {
+    const unsubscribeFn = vi.fn();
+    mockSubscribeProcessExit.mockReturnValue(unsubscribeFn);
+
+    profiler = createProfiler({ enabled: false });
+    expect(profiler.isEnabled()).toBe(false);
+    expect(mockSubscribeProcessExit).toHaveBeenCalled();
+
+    profiler.close();
+
+    expect(unsubscribeFn).toHaveBeenCalledTimes(1);
+    expect(profiler.isEnabled()).toBe(false);
+  });
 });

packages/utils/src/lib/wal.ts

@@ -3,6 +3,7 @@ import * as fs from 'node:fs';
 import path from 'node:path';
 import process from 'node:process';
 import { threadId } from 'node:worker_threads';
+import { SHARDED_WAL_COORDINATOR_ID_ENV_VAR } from './profiler/constants.js';
 
 /**
  * Codec for encoding/decoding values to/from strings for WAL storage.
@@ -356,6 +357,14 @@ export function setLeaderWal(envVarName: string, profilerID: string): void {
 // eslint-disable-next-line functional/no-let
 let shardCount = 0;
 
+/**
+ * Generates a unique sharded WAL ID based on performance time origin, process ID, thread ID, and instance count.
+ */
+function getShardedWalId() {
+  // eslint-disable-next-line functional/immutable-data
+  return `${Math.round(performance.timeOrigin)}.${process.pid}.${threadId}.${++ShardedWal.instanceCount}`;
+}
+
 /**
  * Generates a human-readable shard ID.
  * This ID is unique per process/thread/shard combination and used in the file name.
@@ -462,24 +471,52 @@ export function getShardedFinalPath<T extends object | string = object>(opt: {
  */
 
 export class ShardedWal<T extends object | string = object> {
+  static instanceCount = 0;
+  readonly #id: string = getShardedWalId();
   readonly groupId = getShardedGroupId();
   readonly #format: WalFormat<T>;
   readonly #dir: string = process.cwd();
+  readonly #isCoordinator: boolean;
+  #finalized = false;
 
   /**
    * Create a sharded WAL manager.
+   *
+   * @param opt.dir - Base directory to store shard files (defaults to process.cwd())
+   * @param opt.format - WAL format configuration
+   * @param opt.groupId - Group ID for sharding (defaults to generated group ID)
+   * @param opt.coordinatorIdEnvVar - Environment variable name for storing coordinator ID (defaults to CP_SHARDED_WAL_COORDINATOR_ID)
    */
   constructor(opt: {
     dir?: string;
     format: Partial<WalFormat<T>>;
     groupId?: string;
+    coordinatorIdEnvVar?: string;
   }) {
-    const { dir, format, groupId } = opt;
+    const {
+      dir,
+      format,
+      groupId,
+      coordinatorIdEnvVar = SHARDED_WAL_COORDINATOR_ID_ENV_VAR,
+    } = opt;
     this.groupId = groupId ?? getShardedGroupId();
     if (dir) {
       this.#dir = dir;
     }
     this.#format = parseWalFormat<T>(format);
+    this.#isCoordinator = isLeaderWal(coordinatorIdEnvVar, this.#id);
+  }
+
+  /**
+   * Is this instance the coordinator?
+   *
+   * Coordinator status is determined from the coordinatorIdEnvVar environment variable.
+   * The coordinator handles finalization and cleanup of shard files.
+   *
+   * @returns true if this instance is the coordinator, false otherwise
+   */
+  isCoordinator(): boolean {
+    return this.#isCoordinator;
   }
 
   shard(shardId: string = getShardId()) {
@@ -570,4 +607,14 @@ export class ShardedWal<T extends object | string = object> {
       // Directory might not be empty or already removed, ignore
     }
   }
+
+  finalizeIfCoordinator(): void {
+    if (this.#finalized) return;
+    this.#finalized = true;
+
+    if (!this.isCoordinator()) return;
+
+    this.finalize();
+    this.cleanup();
+  }
 }