import { isCI } from 'ci-info'; import { noop } from 'foxts/noop'; import { basename, extname } from 'node:path'; import process from 'node:process'; import picocolors from 'picocolors'; export const SPAN_STATUS_START = 0; export const SPAN_STATUS_END = 1; const spanTag = Symbol('span'); export interface TraceResult { name: string, start: number, end: number, children: TraceResult[] } /** Pure data object — safe to transfer across Worker Thread boundaries. */ export interface RawSpan { traceResult: TraceResult, status: typeof SPAN_STATUS_START | typeof SPAN_STATUS_END } export interface Span { [spanTag]: true, readonly rawSpan: RawSpan, readonly stop: (time?: number) => void, readonly traceChild: (name: string) => Span, readonly traceSyncFn: (fn: (span: Span) => T) => T, readonly traceAsyncFn: (fn: (span: Span) => T | Promise) => Promise, readonly tracePromise: (promise: Promise) => Promise, readonly traceChildSync: (name: string, fn: (span: Span) => T) => T, readonly traceChildAsync: (name: string, fn: (span: Span) => Promise) => Promise, readonly traceChildPromise: (name: string, promise: Promise) => Promise, readonly traceWorkerChild: (name: string, factory: (rawSpan: RawSpan) => Promise>) => Promise, readonly traceResult: TraceResult } /** * Wraps a serializable {@link RawSpan} with all span methods. * Use this on a worker thread after receiving a {@link RawSpan} (or {@link TraceResult}) * transferred from another thread. */ export function makeSpan(rawSpan: RawSpan): Span { const { traceResult } = rawSpan; const stop = (time?: number) => { if (rawSpan.status === SPAN_STATUS_END) { throw new Error(`span already stopped: ${traceResult.name}`); } traceResult.end = time ?? performance.now(); rawSpan.status = SPAN_STATUS_END; }; const traceChild = (name: string) => createSpan(name, traceResult); const span: Span = { [spanTag]: true, rawSpan, stop, traceChild, traceSyncFn(fn: (span: Span) => T) { const res = fn(span); span.stop(); return res; }, async traceAsyncFn(fn: (span: Span) => T | Promise): Promise { const res = await fn(span); span.stop(); return res; }, traceResult, async tracePromise(promise: Promise): Promise { const res = await promise; span.stop(); return res; }, traceChildSync: (name: string, fn: (span: Span) => T): T => traceChild(name).traceSyncFn(fn), traceChildAsync: (name: string, fn: (span: Span) => T | Promise): Promise => traceChild(name).traceAsyncFn(fn), traceChildPromise: (name: string, promise: Promise): Promise => traceChild(name).tracePromise(promise), async traceWorkerChild(name: string, factory: (rawSpan: RawSpan) => Promise>): Promise { const childSpan = traceChild(name); const { result, traceResult, workerTimeOrigin } = await factory(childSpan.rawSpan); mergeWorkerTrace(childSpan, traceResult, workerTimeOrigin); childSpan.stop(); return result; } }; // eslint-disable-next-line sukka/no-redundant-variable -- self reference return span; } export function createSpan(name: string, parentTraceResult?: TraceResult): Span { const rawSpan: RawSpan = { traceResult: { name, start: performance.now(), end: 0, children: [] }, status: SPAN_STATUS_START }; parentTraceResult?.children.push(rawSpan.traceResult); return makeSpan(rawSpan); } export const dummySpan = createSpan('dummy'); export function task(importMetaMain: boolean, importMetaPath: string) { return (fn: (span: Span, onCleanup: (cb: () => Promise | void) => void) => Promise, customName?: string) => { const taskName = customName ?? basename(importMetaPath, extname(importMetaPath)); let cleanup: () => Promise | void = noop; const onCleanup = (cb: () => void) => { cleanup = cb; }; if (importMetaMain) { const innerSpan = createSpan(taskName); process.on('uncaughtException', (error) => { console.error('Uncaught exception:', error); process.exit(1); }); process.on('unhandledRejection', (reason) => { console.error('Unhandled rejection:', reason); process.exit(1); }); innerSpan.traceChildAsync('dummy', (childSpan) => fn(childSpan, onCleanup)).finally(() => { innerSpan.stop(); printTraceResult(innerSpan.traceResult); process.nextTick(whyIsNodeRunning); process.nextTick(() => process.exit(0)); }); } let runSpan: Span; async function run(parentSpan?: Span | null): Promise { if (parentSpan) { runSpan = parentSpan.traceChild(taskName); } else { runSpan = createSpan(taskName); } try { await fn(runSpan, onCleanup); } finally { runSpan.stop(); cleanup(); } return runSpan.traceResult; } return Object.assign(run, { getInternalTraceResult: () => runSpan.traceResult }); }; } export async function whyIsNodeRunning() { if (isCI && process.env.RUNNER_DEBUG === '1') { const mod = await import('why-is-node-running'); return mod.default(); } } // const isSpan = (obj: any): obj is Span => { // return typeof obj === 'object' && obj && spanTag in obj; // }; // export const universalify = (taskname: string, fn: (this: void, ...args: A) => R) => { // return (...args: A) => { // const lastArg = args[args.length - 1]; // if (isSpan(lastArg)) { // return lastArg.traceChild(taskname).traceSyncFn(() => fn(...args)); // } // return fn(...args); // }; // }; function adjustTraceTimestamps(trace: TraceResult, offset: number): TraceResult { return { name: trace.name, start: trace.start + offset, end: trace.end + offset, children: trace.children.map(child => adjustTraceTimestamps(child, offset)) }; } function mergeWorkerTrace( parentSpan: Span, workerTraceResult: TraceResult, workerTimeOrigin: number ): void { const offset = workerTimeOrigin - performance.timeOrigin; for (const child of workerTraceResult.children) { parentSpan.traceResult.children.push(adjustTraceTimestamps(child, offset)); } } /** The envelope that a worker function returns so the main thread can recover both the result and the trace. */ export interface WorkerJobResult { result: T, traceResult: TraceResult, workerTimeOrigin: number } /** * Worker-side wrapper. Call this instead of manually constructing spans. * * - When `rawSpan` is provided (normal worker invocation from the main thread), * it is wrapped with {@link makeSpan} so all child spans are attached to the * caller's trace tree and can be recovered after the job finishes. * - When `rawSpan` is `undefined` (standalone / CLI invocation), a fresh * child span of {@link dummySpan} is used instead. * * The impl function receives a full {@link Span} and returns its result * normally; the wrapper packages everything into a {@link WorkerJobResult}. */ export async function workerJob( rawSpan: RawSpan | undefined, impl: (span: Span) => Promise ): Promise> { const span = rawSpan == null ? dummySpan.traceChild('worker-standalone') : makeSpan(rawSpan); const result = await impl(span); return { result, traceResult: span.traceResult, workerTimeOrigin: performance.timeOrigin }; } export function printTraceResult(traceResult: TraceResult) { printTree( traceResult, node => { if (node.end - node.start < 0) { return node.name; } return `${node.name} ${picocolors.bold(`${(node.end - node.start).toFixed(3)}ms`)}`; } ); } function printTree(initialTree: TraceResult, printNode: (node: TraceResult, branch: string) => string) { function printBranch(tree: TraceResult, branch: string, isGraphHead: boolean, isChildOfLastBranch: boolean) { const children = tree.children; let branchHead = ''; if (!isGraphHead) { branchHead = children.length > 0 ? '┬ ' : '─ '; } const toPrint = printNode(tree, `${branch}${branchHead}`); if (typeof toPrint === 'string') { console.log(`${branch}${branchHead}${toPrint}`); } let baseBranch = branch; if (!isGraphHead) { baseBranch = branch.slice(0, -2) + (isChildOfLastBranch ? ' ' : '│ '); } const nextBranch = `${baseBranch}├─`; const lastBranch = `${baseBranch}└─`; children.forEach((child, index) => { const last = children.length - 1 === index; printBranch(child, last ? lastBranch : nextBranch, false, last); }); } printBranch(initialTree, '', true, false); } export function printStats(stats: TraceResult[]): void { const longestTaskName = Math.max(...stats.map(i => i.name.length)); const realStart = Math.min(...stats.map(i => i.start)); const realEnd = Math.max(...stats.map(i => i.end)); const statsStep = ((realEnd - realStart) / 120) | 0; stats .sort((a, b) => a.start - b.start) .forEach(stat => { console.log( `[${stat.name}]${' '.repeat(longestTaskName - stat.name.length)}`, ' '.repeat(((stat.start - realStart) / statsStep) | 0), '='.repeat(Math.max(((stat.end - stat.start) / statsStep) | 0, 1)) ); }); }