perf(db): accelerate parallel trace ingestion / 加速并行 trace 摄取

packages/db/src/etl/async-semaphore.ts

@@ -0,0 +1,38 @@
+/** Minimal FIFO semaphore for bounding concurrent asynchronous operations. */
+export class AsyncSemaphore {
+  private active = 0;
+  private readonly waiters: (() => void)[] = [];
+  private readonly limit: number;
+
+  constructor(limit: number) {
+    if (!Number.isSafeInteger(limit) || limit < 1) {
+      throw new Error(`Semaphore limit must be a positive integer, received ${limit}`);
+    }
+    this.limit = limit;
+  }
+
+  async run<T>(operation: () => Promise<T>): Promise<T> {
+    await this.acquire();
+    try {
+      return await operation();
+    } finally {
+      this.release();
+    }
+  }
+
+  private async acquire(): Promise<void> {
+    if (this.active < this.limit) {
+      this.active += 1;
+      return;
+    }
+    await new Promise<void>((resolve) => {
+      this.waiters.push(resolve);
+    });
+    this.active += 1;
+  }
+
+  private release(): void {
+    this.active -= 1;
+    this.waiters.shift()?.();
+  }
+}

packages/db/src/etl/gzip-json-stream.test.ts

@@ -103,6 +103,54 @@ describe('collectMetricPhases', () => {
     });
   });
 
+  it('ignores nested phase lookalikes and preserves selected scalar and compound values', async () => {
+    const adversarial = gzipSync(
+      JSON.stringify({
+        metadata: {
+          metrics: { wanted: { nested: 'must not be selected' } },
+          warmup_metrics: { wanted: { nested: 'must not be selected' } },
+        },
+        metrics: {
+          wanted: [1, { nested: true }, null],
+          scalar: 42,
+          ignored: { very: { deeply: ['nested', 'value'] } },
+        },
+        warmup_metrics: {
+          wanted: false,
+          scalar: 'warmup',
+          ignored: [1, 2, 3],
+        },
+      }),
+    );
+
+    await expect(
+      collectMetricPhases(adversarial, new Set(['wanted', 'scalar']), 1),
+    ).resolves.toEqual({
+      metrics: {
+        wanted: [1, { nested: true }, null],
+        scalar: 42,
+      },
+      warmupMetrics: {
+        wanted: false,
+        scalar: 'warmup',
+      },
+      complete: false,
+    });
+  });
+
+  it('matches JSON semantics for duplicate phase and metric keys', async () => {
+    const duplicateKeys = gzipSync(
+      '{"metrics":{"wanted":1,"wanted":2},"metrics":{"wanted":3},' +
+        '"warmup_metrics":{"wanted":4,"wanted":5}}',
+    );
+
+    await expect(collectMetricPhases(duplicateKeys, new Set(['wanted']), 1)).resolves.toEqual({
+      metrics: { wanted: 3 },
+      warmupMetrics: { wanted: 5 },
+      complete: false,
+    });
+  });
+
   it('rejects malformed gzip input on both paths', async () => {
     await expect(
       collectMetricPhases(Buffer.from('not gzip'), new Set(['wanted']), 1),

packages/db/src/etl/gzip-json-stream.ts

@@ -7,13 +7,14 @@
  * stream-json pipeline collects only the top-level subtrees callers need.
  */
 
-import { PassThrough, Readable } from 'node:stream';
-import { pipeline } from 'node:stream/promises';
+import { Readable } from 'node:stream';
 import { createGunzip, gunzipSync } from 'node:zlib';
 
 import { chain } from 'stream-chain';
 
 import { parser } from 'stream-json';
+import Assembler from 'stream-json/assembler.js';
+import type { Token } from 'stream-json/parser.js';
 import { pick } from 'stream-json/filters/pick.js';
 import { streamObject } from 'stream-json/streamers/stream-object.js';
 
@@ -82,24 +83,122 @@ export interface MetricPhaseMaps<T> {
   complete: boolean;
 }
 
-async function collectTokenBranch<T>(
-  input: PassThrough,
-  filter: 'metrics' | 'warmup_metrics',
+function isValueStart(token: Token): boolean {
+  return (
+    token.name === 'startObject' ||
+    token.name === 'startArray' ||
+    token.name === 'stringValue' ||
+    token.name === 'numberValue' ||
+    token.name === 'nullValue' ||
+    token.name === 'trueValue' ||
+    token.name === 'falseValue'
+  );
+}
+
+function updateDepth(depth: number, token: Token): number {
+  if (token.name === 'startObject' || token.name === 'startArray') return depth + 1;
+  if (token.name === 'endObject' || token.name === 'endArray') return depth - 1;
+  return depth;
+}
+
+/**
+ * Consume the parser token stream once and assemble only selected direct
+ * children of the two metric phase objects. Unselected metric subtrees still
+ * have to be tokenized so the JSON can be validated, but they are never copied
+ * to secondary streams or materialized as JavaScript objects.
+ */
+async function streamCollectMetricPhases<T>(
+  buffer: Buffer,
   wanted: ReadonlySet<string>,
-): Promise<Record<string, T>> {
-  const collected: Record<string, T> = {};
-  const output = chain([input, pick({ filter }), streamObject()]);
-  for await (const chunk of output) {
-    const { key, value } = chunk as { key: string; value: T };
-    if (wanted.has(key)) collected[key] = value;
+): Promise<MetricPhaseMaps<T>> {
+  let metrics: Record<string, T> = {};
+  let warmupMetrics: Record<string, T> = {};
+  let depth = 0;
+  let phase: 'metrics' | 'warmup_metrics' | null = null;
+  let topLevelKey: string | null = null;
+  let metricKey: string | null = null;
+  let valueAssembler: Assembler<T> | null = null;
+  let valueTarget: Record<string, T> | null = null;
+  let valueKey: string | null = null;
+
+  // Packed-only values avoid emitting start/chunk/end tokens in addition to
+  // their complete value token. This materially reduces the token count for
+  // metric-heavy multi-GiB documents.
+  const tokens = chain([
+    Readable.from(buffer),
+    createGunzip(),
+    parser({
+      packKeys: true,
+      packStrings: true,
+      packNumbers: true,
+      streamKeys: false,
+      streamStrings: false,
+      streamNumbers: false,
+    }),
+  ]);
+
+  for await (const rawToken of tokens) {
+    const token = rawToken as Token;
+    const depthBefore = depth;
+
+    if (valueAssembler) {
+      valueAssembler.consume(token);
+      depth = updateDepth(depth, token);
+      if (valueAssembler.done) {
+        valueTarget![valueKey!] = valueAssembler.current as T;
+        valueAssembler = null;
+        valueTarget = null;
+        valueKey = null;
+      }
+      continue;
+    }
+
+    if (token.name === 'keyValue') {
+      if (depthBefore === 1) {
+        topLevelKey = token.value;
+      } else if (depthBefore === 2 && phase) {
+        metricKey = token.value;
+      }
+    } else if (isValueStart(token)) {
+      if (depthBefore === 1 && topLevelKey !== null) {
+        if (token.name === 'startObject' && topLevelKey === 'metrics') {
+          metrics = {};
+          phase = 'metrics';
+        } else if (token.name === 'startObject' && topLevelKey === 'warmup_metrics') {
+          warmupMetrics = {};
+          phase = 'warmup_metrics';
+        }
+        topLevelKey = null;
+      } else if (depthBefore === 2 && phase && metricKey !== null) {
+        if (wanted.has(metricKey)) {
+          valueTarget = phase === 'metrics' ? metrics : warmupMetrics;
+          valueKey = metricKey;
+          valueAssembler = new Assembler<T>();
+          valueAssembler.consume(token);
+          if (valueAssembler.done) {
+            valueTarget[valueKey] = valueAssembler.current as T;
+            valueAssembler = null;
+            valueTarget = null;
+            valueKey = null;
+          }
+        }
+        metricKey = null;
+      }
+    }
+
+    depth = updateDepth(depth, token);
+    if (phase && depth === 1 && (token.name === 'endObject' || token.name === 'endArray')) {
+      phase = null;
+      metricKey = null;
+    }
   }
-  return collected;
+
+  return { metrics, warmupMetrics, complete: false };
 }
 
 /**
- * Gunzip and parse both server-metric phase blocks once. Large documents fan
- * the parser's token stream out to two lightweight selectors, avoiding one
- * complete decompression + JSON tokenization pass per phase.
+ * Gunzip and parse both server-metric phase blocks once. Large documents use a
+ * single token consumer which materializes only the selected metric values.
  */
 export async function collectMetricPhases<T>(
   buffer: Buffer,
@@ -119,30 +218,5 @@ export async function collectMetricPhases<T>(
     };
   }
 
-  // Attach every branch before starting the source pipeline so no parser
-  // tokens can be missed. PassThrough backpressure keeps the two consumers in
-  // lockstep without buffering the full document.
-  const profilingInput = new PassThrough({ objectMode: true });
-  const warmupInput = new PassThrough({ objectMode: true });
-  const tokenTee = new PassThrough({ objectMode: true });
-  tokenTee.pipe(profilingInput);
-  tokenTee.pipe(warmupInput);
-
-  const profiling = collectTokenBranch<T>(profilingInput, 'metrics', wanted);
-  const warmup = collectTokenBranch<T>(warmupInput, 'warmup_metrics', wanted);
-  const tokens = chain([Readable.from(buffer), createGunzip(), parser()]);
-
-  try {
-    const [, metrics, warmupMetrics] = await Promise.all([
-      pipeline(tokens, tokenTee),
-      profiling,
-      warmup,
-    ]);
-    return { metrics, warmupMetrics, complete: false };
-  } catch (error) {
-    tokenTee.destroy();
-    profilingInput.destroy();
-    warmupInput.destroy();
-    throw error;
-  }
+  return await streamCollectMetricPhases<T>(buffer, wanted);
 }

packages/db/src/etl/trace-replay-ingest.test.ts

@@ -9,6 +9,8 @@ import { describe, expect, it } from 'vitest';
 import {
   TRACE_REPLAY_UPLOAD_CHUNK_BYTES,
   gzipTraceReplayInput,
+  persistPreparedTraceReplay,
+  type PreparedTraceReplay,
   uploadTraceReplayPayloadChunks,
 } from './trace-replay-ingest';
 
@@ -26,6 +28,45 @@ function mockTransactionSql(): { sql: postgres.TransactionSql; calls: SqlCall[]
   return { sql, calls };
 }
 
+function preparedFixture(): PreparedTraceReplay {
+  return {
+    profileGz: Buffer.from('profile'),
+    profileSize: 70,
+    serverMetricsCsv: Buffer.from('csv'),
+    serverMetricsCsvSize: 30,
+    serverMetricsJsonGz: Buffer.from('metrics'),
+    serverMetricsJsonSize: 700,
+    aggregateStatsJson: Buffer.from('{"version":1}'),
+    chartSeriesJson: Buffer.from('{"version":12}'),
+    requestTimelineJson: Buffer.from('{"version":1}'),
+    chartWindows: 2,
+    timelineRequests: 3,
+    compressionMs: 10,
+    computeMs: 20,
+    cacheHitRates: null,
+  };
+}
+
+function mockSqlWithTransaction(lockedRows: { id: number }[]): {
+  sql: Parameters<typeof persistPreparedTraceReplay>[0];
+  calls: SqlCall[];
+} {
+  const calls: SqlCall[] = [];
+  const execute = (strings: TemplateStringsArray, ...values: unknown[]) => {
+    const text = strings.join('?');
+    calls.push({ text, values });
+    if (text.includes('for update')) return Promise.resolve(lockedRows);
+    if (text.includes('insert into agentic_trace_replay')) return Promise.resolve([{ id: 123 }]);
+    return Promise.resolve([]);
+  };
+  const tx = Object.assign(execute, { array: (values: unknown[]) => values });
+  const sql = Object.assign(execute, {
+    array: (values: unknown[]) => values,
+    begin: (operation: (transaction: typeof tx) => Promise<void>) => operation(tx),
+  }) as unknown as Parameters<typeof persistPreparedTraceReplay>[0];
+  return { sql, calls };
+}
+
 describe('uploadTraceReplayPayloadChunks', () => {
   it('bounds every Bind payload for the measured 90 MiB staging row', async () => {
     // Exact payload sizes from InferenceX run 29181694248, item 8/9.
@@ -80,3 +121,33 @@ describe('gzipTraceReplayInput', () => {
     }
   });
 });
+
+describe('persistPreparedTraceReplay', () => {
+  it('rechecks links under a row lock and avoids creating an orphan after a concurrent ingest', async () => {
+    const { sql, calls } = mockSqlWithTransaction([]);
+
+    await expect(persistPreparedTraceReplay(sql, [41], preparedFixture())).resolves.toBe(0);
+    expect(calls).toHaveLength(1);
+    expect(calls[0].text).toContain('for update');
+    expect(calls.some((call) => call.text.includes('insert into agentic_trace_replay'))).toBe(
+      false,
+    );
+  });
+
+  it('uploads every payload and links exactly the rows locked by the transaction', async () => {
+    const { sql, calls } = mockSqlWithTransaction([{ id: 41 }]);
+
+    await expect(persistPreparedTraceReplay(sql, [41], preparedFixture())).resolves.toBe(1);
+    const lockIndex = calls.findIndex((call) => call.text.includes('for update'));
+    const blobIndex = calls.findIndex((call) =>
+      call.text.includes('insert into agentic_trace_replay'),
+    );
+    const linkCall = calls.find((call) => call.text.includes('set trace_replay_id'));
+    expect(lockIndex).toBe(0);
+    expect(blobIndex).toBeGreaterThan(lockIndex);
+    expect(
+      calls.filter((call) => call.text.includes('trace_replay_upload_parts (field, part, data)')),
+    ).toHaveLength(6);
+    expect(linkCall?.values.some((value) => Array.isArray(value) && value.includes(41))).toBe(true);
+  });
+});