CoolFace
Datasetpublic

basant307/AI_Governance_Project

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
0likes48downloads
cpuProfiler.ts441 linesDownload Raw Back to utils
1/**2 * @license3 * Copyright 2025 Qwen4 * SPDX-License-Identifier: Apache-2.05 */6 7/**8 * CPU profiling utility that generates .cpuprofile files for Chrome DevTools.9 *10 * Three trigger modes:11 * 1. Environment variable: QWEN_CODE_CPU_PROFILE=1 — records from process start to exit12 * 2. Signal toggle: SIGUSR1 — first signal starts, second stops and writes13 * 3. Command: /doctor cpu-profile [--duration N] — records for N seconds14 *15 * Output: ~/.qwen/cpu-profiles/qwen-code-cpu-<pid>-<timestamp>.cpuprofile16 * Zero overhead when disabled (single env var check at init).17 */18 19import * as fs from 'node:fs';20import * as os from 'node:os';21import * as path from 'node:path';22 23import { registerCleanup } from './cleanup.js';24 25// ---------------------------------------------------------------------------26// Types27// ---------------------------------------------------------------------------28 29type ProfilerState = 'idle' | 'recording' | 'stopping';30 31export type CpuProfileStartResult = { ok: true } | { ok: false; error: string };32 33export type CpuProfileStopResult =34  | { ok: true; filePath: string }35  | { ok: false; error: string };36 37// Custom interface rather than importing from node:inspector/promises because38// the official Session.post() generic overload returns Promise<void>, making39// dynamic method dispatch (Profiler.start/stop) cumbersome without per-call casts.40interface InspectorSession {41  connect(): void;42  disconnect(): void;43  post(method: string, params?: Record<string, unknown>): Promise<unknown>;44}45 46// ---------------------------------------------------------------------------47// Constants48// ---------------------------------------------------------------------------49 50const MAX_PROFILES = 5;51const RATE_LIMIT_MS = 30_000;52const MIN_FREE_BYTES_AFTER_WRITE = 256 * 1024 * 1024;53const DEFAULT_SAMPLING_INTERVAL_US = 1000; // 1ms54const ESTIMATED_PROFILE_BYTES = 10 * 1024 * 1024; // 10 MiB conservative estimate55 56// ---------------------------------------------------------------------------57// Module state58// ---------------------------------------------------------------------------59 60let state: ProfilerState = 'idle';61let session: InspectorSession | null = null;62let initialized = false;63let signalHandlerRegistered = false;64const lastWriteByDir = new Map<string, number>();65 66// ---------------------------------------------------------------------------67// Public API68// ---------------------------------------------------------------------------69 70/**71 * Initialize CPU profiler. Call once at process start.72 * Always registers SIGUSR1 handler (for ad-hoc profiling).73 * When QWEN_CODE_CPU_PROFILE=1, also starts recording immediately.74 */75export function initCpuProfiler(): void {76  if (initialized) return;77  initialized = true;78 79  // Always register signal handler for ad-hoc profiling (non-Windows)80  registerSignalHandler();81 82  // Always register cleanup to flush any in-progress profile on exit83  registerCleanup(async () => {84    if (state === 'recording') {85      const result = await stopCpuProfile();86      if (result.ok) {87        process.stderr.write(88          `[cpu-profiler] Profile written: ${result.filePath}\n`,89        );90      }91    }92  });93 94  const enabled = process.env['QWEN_CODE_CPU_PROFILE'] === '1';95  if (!enabled) return;96 97  // Start recording immediately in env-var mode98  void startCpuProfile().then((result) => {99    if (!result.ok) {100      process.stderr.write(`[cpu-profiler] Failed to start: ${result.error}\n`);101    }102  });103}104 105/**106 * Start CPU profiling.107 * @param opts.samplingInterval - Sampling interval in microseconds (default 1000 = 1ms)108 */109export async function startCpuProfile(opts?: {110  samplingInterval?: number;111}): Promise<CpuProfileStartResult> {112  if (state !== 'idle') {113    return {114      ok: false,115      error:116        state === 'recording'117          ? 'CPU profiling is already in progress.'118          : 'CPU profiler is currently stopping. Please wait a moment and try again.',119    };120  }121 122  // Set state eagerly before the first await to prevent concurrent callers123  // (e.g., rapid SIGUSR1 signals) from both passing the idle guard.124  state = 'recording';125 126  try {127    const inspectorSession = await getOrCreateSession();128    await inspectorSession.post('Profiler.enable');129    await inspectorSession.post('Profiler.setSamplingInterval', {130      interval: opts?.samplingInterval ?? DEFAULT_SAMPLING_INTERVAL_US,131    });132    await inspectorSession.post('Profiler.start');133    return { ok: true };134  } catch (error) {135    state = 'idle';136    disconnectSession();137    return { ok: false, error: formatError(error) };138  }139}140 141/**142 * Stop CPU profiling and write the .cpuprofile file.143 * @returns File path on success.144 */145export async function stopCpuProfile(options?: {146  outputDir?: string;147  now?: Date;148  rateLimitMs?: number;149  maxProfiles?: number;150}): Promise<CpuProfileStopResult> {151  if (state !== 'recording') {152    return {153      ok: false,154      error:155        state === 'idle'156          ? 'CPU profiler is not recording.'157          : 'CPU profiler is already stopping.',158    };159  }160 161  const outputDir = options?.outputDir ?? defaultOutputDir();162  const now = options?.now ?? new Date();163  const rateLimitMs = options?.rateLimitMs ?? RATE_LIMIT_MS;164  const maxProfiles = options?.maxProfiles ?? MAX_PROFILES;165 166  // Check rate limit BEFORE writing to avoid excessive output.167  // If rate-limited, tear down the V8 profiler (data is discarded) and reset168  // state to 'idle' so the user can start a fresh recording later.169  try {170    enforceRateLimit(outputDir, now, rateLimitMs);171  } catch (error) {172    state = 'idle';173    if (session) {174      session.post('Profiler.stop').catch(() => {});175      session.post('Profiler.disable').catch(() => {});176    }177    disconnectSession();178    return { ok: false, error: formatError(error) };179  }180 181  state = 'stopping';182 183  try {184    if (!session) {185      throw new Error(186        'Inspector session lost unexpectedly during Profiler.stop; the profile data could not be retrieved.',187      );188    }189 190    const result = (await session.post('Profiler.stop')) as {191      profile: unknown;192    };193    if (!result.profile) {194      throw new Error(195        'V8 Profiler.stop returned an empty profile; recording may have been interrupted.',196      );197    }198    await session.post('Profiler.disable');199 200    fs.mkdirSync(outputDir, { recursive: true, mode: 0o700 });201    try {202      fs.chmodSync(outputDir, 0o700);203    } catch {204      // Best-effort hardening on filesystems without POSIX chmod.205    }206 207    checkDiskSpace(outputDir);208 209    const filePath = path.join(210      outputDir,211      `qwen-code-cpu-${process.pid}-${formatTimestamp(now)}.cpuprofile`,212    );213 214    try {215      fs.writeFileSync(filePath, JSON.stringify(result.profile), {216        mode: 0o600,217      });218    } catch (writeError) {219      try {220        fs.rmSync(filePath, { force: true });221      } catch {222        // Best-effort cleanup of partial file.223      }224      throw writeError;225    }226 227    recordWrite(outputDir, now);228    cleanupOldProfiles(outputDir, maxProfiles);229 230    state = 'idle';231    return { ok: true, filePath };232  } catch (error) {233    state = 'idle';234    disconnectSession();235    return { ok: false, error: formatError(error) };236  }237}238 239/**240 * Whether the profiler is currently recording.241 */242export function isCpuProfileRecording(): boolean {243  return state === 'recording';244}245 246/**247 * Register SIGUSR1 signal handler for toggle mode.248 * Safe to call multiple times; only registers once.249 * No-op on Windows (SIGUSR1 does not exist).250 */251export function registerSignalHandler(): void {252  if (signalHandlerRegistered) return;253  if (process.platform === 'win32') return;254 255  signalHandlerRegistered = true;256  process.on('SIGUSR1', handleSigusr1);257}258 259// ---------------------------------------------------------------------------260// Test helpers261// ---------------------------------------------------------------------------262 263/** Reset all module state. Test-only. */264export function _resetCpuProfilerForTest(): void {265  state = 'idle';266  initialized = false;267  signalHandlerRegistered = false;268  disconnectSession();269  lastWriteByDir.clear();270}271 272/** Clear rate limit state. Test-only. */273export function clearCpuProfileRateLimit(): void {274  lastWriteByDir.clear();275}276 277// ---------------------------------------------------------------------------278// Internal helpers279// ---------------------------------------------------------------------------280 281function defaultOutputDir(): string {282  return path.join(os.homedir(), '.qwen', 'cpu-profiles');283}284 285function formatTimestamp(now: Date): string {286  return now.toISOString().replace(/[:.]/g, '-');287}288 289function formatError(error: unknown): string {290  return error instanceof Error ? error.message : String(error);291}292 293// Overridable factory for testing (avoids mocking ESM dynamic imports)294let sessionFactory: (() => Promise<InspectorSession>) | null = null;295 296/** Override session creation for testing. */297export function _setSessionFactoryForTest(298  factory: (() => Promise<InspectorSession>) | null,299): void {300  sessionFactory = factory;301}302 303async function getOrCreateSession(): Promise<InspectorSession> {304  if (session) return session;305 306  if (sessionFactory) {307    session = await sessionFactory();308    return session;309  }310 311  // Dynamic import to avoid any overhead when profiling is disabled312  const inspectorModule = await import('node:inspector/promises');313  const newSession =314    new inspectorModule.Session() as unknown as InspectorSession;315  newSession.connect();316  session = newSession;317  return session;318}319 320function disconnectSession(): void {321  if (session) {322    try {323      session.disconnect();324    } catch {325      // Ignore disconnect errors during cleanup.326    }327    session = null;328  }329}330 331function handleSigusr1(): void {332  if (state === 'idle') {333    void startCpuProfile().then((result) => {334      if (result.ok) {335        process.stderr.write(336          `[cpu-profiler] Recording started (PID ${process.pid}). Send SIGUSR1 again to stop.\n`,337        );338      } else {339        process.stderr.write(340          `[cpu-profiler] Failed to start: ${result.error}\n`,341        );342      }343    });344  } else if (state === 'recording') {345    void stopCpuProfile().then((result) => {346      if (result.ok) {347        process.stderr.write(348          `[cpu-profiler] Profile written: ${result.filePath}\n`,349        );350      } else {351        process.stderr.write(352          `[cpu-profiler] Failed to stop: ${result.error}\n`,353        );354      }355    });356  }357  // state === 'stopping': ignore, already in progress358}359 360function enforceRateLimit(361  outputDir: string,362  now: Date,363  rateLimitMs: number,364): void {365  if (rateLimitMs <= 0) return;366 367  const key = path.resolve(outputDir);368  const nowMs = now.getTime();369  const lastWriteMs = lastWriteByDir.get(key);370  if (lastWriteMs !== undefined && nowMs - lastWriteMs < rateLimitMs) {371    const waitSeconds = Math.ceil((rateLimitMs - (nowMs - lastWriteMs)) / 1000);372    throw new Error(373      `CPU profile rate limit: wait ${waitSeconds}s before writing another profile.`,374    );375  }376}377 378function recordWrite(outputDir: string, now: Date): void {379  lastWriteByDir.set(path.resolve(outputDir), now.getTime());380}381 382function checkDiskSpace(outputDir: string): void {383  try {384    const stats = fs.statfsSync(outputDir);385    const available = stats.bavail * stats.bsize;386    if (available - ESTIMATED_PROFILE_BYTES < MIN_FREE_BYTES_AFTER_WRITE) {387      throw new Error(388        'Insufficient free disk space for CPU profile; skipping to avoid filling the disk.',389      );390    }391  } catch (error) {392    if (393      error instanceof Error &&394      error.message.includes('Insufficient free disk')395    ) {396      throw error;397    }398    // statfsSync is not available on all platforms (e.g. Windows).399    // Log a warning so it's not completely silent, but proceed anyway.400    process.stderr.write(401      '[cpu-profiler] Disk space check unavailable on this platform; skipping.\n',402    );403  }404}405 406function cleanupOldProfiles(outputDir: string, maxProfiles: number): void {407  if (maxProfiles < 1) return;408 409  let profiles: string[];410  try {411    profiles = fs412      .readdirSync(outputDir)413      .filter(414        (name) =>415          name.startsWith('qwen-code-cpu-') && name.endsWith('.cpuprofile'),416      )417      .map((name) => path.join(outputDir, name))418      .sort((a, b) => {419        try {420          return (421            fs.lstatSync(b).mtimeMs - fs.lstatSync(a).mtimeMs ||422            path.basename(b).localeCompare(path.basename(a))423          );424        } catch {425          // Fall back to filename comparison if stat fails426          return path.basename(b).localeCompare(path.basename(a));427        }428      });429  } catch {430    return;431  }432 433  for (const filePath of profiles.slice(maxProfiles)) {434    try {435      fs.rmSync(filePath, { force: true });436    } catch {437      // Cleanup is best effort.438    }439  }440}441 
basant307/AI_Governance_Project · CoolFace