diff --git a/scripts/bench-streaming-import.mts b/scripts/bench-streaming-import.mts deleted file mode 100644 index 398a50a4..00000000 --- a/scripts/bench-streaming-import.mts +++ /dev/null @@ -1,484 +0,0 @@ -/** - * End-to-end streaming import benchmark. - * - * Each measured run executes in a fresh process and writes fresh databases. - * Fixture generation and result verification are excluded from the timed range. - * - * Usage: - * pnpm exec tsx scripts/bench-streaming-import.mts single 100000 [runs=3] - * pnpm exec tsx scripts/bench-streaming-import.mts batch 100 10000 [runs=3] [concurrency=1] - */ - -import { createHash } from 'node:crypto' -import { once } from 'node:events' -import { copyFileSync, createWriteStream, linkSync, mkdirSync, rmSync, statSync, unlinkSync } from 'node:fs' -import path from 'node:path' -import { fileURLToPath } from 'node:url' -import { spawnSync } from 'node:child_process' -import Database from 'better-sqlite3' -import { CHAT_DB_TABLES } from '@openchatlab/core' -import { getNativeParserStatus, isNativeFormatAvailable, type NativeParserStatus } from '@openchatlab/parser' -import { BetterSqliteAdapter } from '../packages/node-runtime/src/better-sqlite3-adapter' -import { computeAndSetOverviewCache } from '../packages/node-runtime/src/cache/session-cache' -import { runKeyedBatch } from '../packages/node-runtime/src/import/batch-coordinator' -import { - streamingImport, - type ImportLogger, - type ImportStageTimings, - type StreamImportDeps, -} from '../packages/node-runtime/src/import/streaming-importer' -import { createChatLabTempDir } from './chatlab-temp.mjs' - -const scriptPath = fileURLToPath(import.meta.url) -const nativeBinding = path.resolve('apps/cli/native/better_sqlite3.node') -const memberCount = 200 -const sampleTexts = [ - '今天天气不错,我们出去玩吧!', - '哈哈哈哈哈哈这也太好笑了', - '[图片]', - '好的,收到,明天见。', - 'This is a mixed language message with some English words 和中文混排。', - '周末有人一起打球吗?地点老地方,时间下午三点,人齐就开打。', -] - -interface BenchmarkParser { - formatId: 'weflow' - implementation: 'rust-native' - nativeModuleAvailable: true -} - -interface BenchmarkResult { - scenario: 'single' | 'batch' - parser: BenchmarkParser - fileCount: number - concurrency: number - inputMessages: number - inputBytes: number - durationMs: number - sampledPeakRssMb: number - sampledRssDeltaMb: number - databaseBytes: number - messagesWritten: number - membersWritten: number - outputSignature: string - timings: ImportStageTimings -} - -export interface BenchmarkParserMonitor { - logger: ImportLogger - assertRustNativeCompleted(importIndex: number): void -} - -export function resolveBenchmarkParser(status: NativeParserStatus, nativeFormatAvailable: boolean): BenchmarkParser { - if (!status.available || !nativeFormatAvailable) { - const reason = status.disabled - ? 'CHATLAB_DISABLE_NATIVE_PERF=1 disables the native parser' - : status.error - ? `native parser failed to load: ${status.error}` - : 'the native parser does not provide the WeFlow kernel' - throw new Error( - `Streaming import benchmark requires the Rust Native WeFlow parser; ${reason}. ` + - 'Run pnpm build:native before collecting a baseline.' - ) - } - - return { - formatId: 'weflow', - implementation: 'rust-native', - nativeModuleAvailable: true, - } -} - -export function createBenchmarkParserMonitor(): BenchmarkParserMonitor { - let nativeStarted = false - let fallbackDetected = false - const ignoreNonParserLog = () => undefined - - const inspectMessage = (message: string) => { - if (message.includes('[NativeParser] Parsing WeFlow export with Rust kernel')) { - nativeStarted = true - } - if (message.includes('[NativeParser]') && message.includes('falling back to TS parser')) { - fallbackDetected = true - } - } - - return { - logger: { - info: inspectMessage, - error: inspectMessage, - perf: ignoreNonParserLog, - perfDetail: ignoreNonParserLog, - summary: ignoreNonParserLog, - reset() { - nativeStarted = false - fallbackDetected = false - }, - init: ignoreNonParserLog, - getCurrentLogFile() { - return null - }, - }, - assertRustNativeCompleted(importIndex: number) { - if (fallbackDetected) { - throw new Error(`Import ${importIndex} fell back to the TypeScript WeFlow parser; benchmark result rejected`) - } - if (!nativeStarted) { - throw new Error(`Import ${importIndex} did not start the Rust Native WeFlow parser; benchmark result rejected`) - } - }, - } -} - -const zeroTimings = (): ImportStageTimings => ({ - detectionMs: 0, - preprocessingMs: 0, - databaseSetupMs: 0, - parserMs: 0, - metaWriteMs: 0, - memberWriteMs: 0, - messageWriteMs: 0, - nicknameHistoryMs: 0, - indexCreationMs: 0, - ftsMs: 0, - checkpointMs: 0, - sessionIndexMs: 0, - postImportHookMs: 0, - totalMs: 0, -}) - -function addTimings(target: ImportStageTimings, source: ImportStageTimings): void { - for (const key of Object.keys(target) as Array) target[key] += source[key] -} - -async function generateWeflowFixture(filePath: string, count: number): Promise { - const stream = createWriteStream(filePath, { encoding: 'utf-8' }) - const write = async (chunk: string) => { - if (!stream.write(chunk)) await once(stream, 'drain') - } - - await write( - `{"weflow":{"version":"1.0.0","exportedAt":1704164645},` + - `"session":{"wxid":"bench@chatroom","nickname":"性能测试群","displayName":"性能测试群","type":"群聊"},` + - '"avatars":{},"messages":[' - ) - - const batch: string[] = [] - for (let index = 0; index < count; index++) { - const member = index % memberCount - batch.push( - JSON.stringify({ - localId: index + 1, - createTime: 1_704_164_645 + index * 3, - formattedTime: '2024-01-02 03:04:05', - type: '文本消息', - localType: 1, - content: `${sampleTexts[index % sampleTexts.length]} #${index}`, - isSend: member === 0 ? 1 : 0, - senderUsername: `wxid_member_${member}`, - senderDisplayName: `成员${member}号`, - senderAvatarKey: `wxid_member_${member}`, - source: '', - }) - ) - if (batch.length >= 5000) { - await write((index + 1 > batch.length ? ',' : '') + batch.join(',')) - batch.length = 0 - } - } - if (batch.length > 0) await write((count > batch.length ? ',' : '') + batch.join(',')) - await write(']}') - stream.end() - await once(stream, 'finish') -} - -function createImportDeps(dbPath: string, cacheDir: string, logger: ImportLogger): StreamImportDeps { - return { - openDatabase() { - const raw = new Database(dbPath, { nativeBinding }) - raw.pragma('journal_mode = WAL') - raw.pragma('synchronous = NORMAL') - raw.exec(CHAT_DB_TABLES) - return new BetterSqliteAdapter(raw) - }, - deleteDatabase() { - for (const suffix of ['', '-wal', '-shm']) { - try { - unlinkSync(dbPath + suffix) - } catch { - /* ignore cleanup failures in a disposable benchmark workspace */ - } - } - }, - onProgress() { - /* benchmark excludes UI transport */ - }, - logger, - postImportHook(db, sessionId) { - computeAndSetOverviewCache(db, sessionId, cacheDir) - }, - } -} - -export function inspectDatabase(dbPath: string): { - messages: number - members: number - signature: string -} { - const db = new Database(dbPath, { readonly: true, nativeBinding }) - const hash = createHash('sha256') - const hashRows = (label: string, query: string): number => { - hash.update(`table:${label}\n`) - let count = 0 - let chunk = '' - for (const row of db.prepare(query).raw().iterate() as Iterable) { - chunk += `${JSON.stringify(row)}\n` - count++ - if (chunk.length >= 1024 * 1024) { - hash.update(chunk) - chunk = '' - } - } - if (chunk) hash.update(chunk) - return count - } - - try { - hashRows( - 'meta', - `SELECT name, platform, type, group_id, group_avatar, owner_id, schema_version, session_gap_threshold - FROM meta ORDER BY rowid` - ) - const members = hashRows( - 'member', - `SELECT id, platform_id, account_name, group_nickname, aliases, avatar, roles - FROM member ORDER BY id` - ) - hashRows( - 'member_name_history', - `SELECT id, member_id, name_type, name, start_ts, end_ts - FROM member_name_history ORDER BY id` - ) - const messages = hashRows( - 'message', - `SELECT id, sender_id, sender_account_name, sender_group_nickname, ts, type, content, - reply_to_message_id, platform_message_id - FROM message ORDER BY id` - ) - hashRows( - 'segment', - `SELECT id, start_ts, end_ts, message_count, is_manual, summary - FROM segment ORDER BY id` - ) - hashRows( - 'message_context', - `SELECT message_id, segment_id, topic_id - FROM message_context ORDER BY message_id` - ) - hashRows('message_fts', 'SELECT rowid FROM message_fts ORDER BY rowid') - return { messages, members, signature: hash.digest('hex') } - } finally { - db.close() - } -} - -async function runWorker(fileCount: number, messagesPerFile: number, concurrency: number): Promise { - if (typeof global.gc !== 'function') { - throw new Error('Benchmark worker requires Node.js --expose-gc') - } - - const parser = resolveBenchmarkParser(getNativeParserStatus(), isNativeFormatAvailable('weflow')) - const root = createChatLabTempDir('bench', 'streaming-import-') - const fixturePath = path.join(root, 'fixture.json') - const cacheDir = path.join(root, 'cache') - mkdirSync(cacheDir, { recursive: true }) - - try { - await generateWeflowFixture(fixturePath, messagesPerFile) - const inputPaths = [fixturePath] - for (let index = 1; index < fileCount; index++) { - const linkedPath = path.join(root, `fixture-${index}.json`) - try { - linkSync(fixturePath, linkedPath) - } catch { - copyFileSync(fixturePath, linkedPath) - } - inputPaths.push(linkedPath) - } - - global.gc() - const rssStartMb = process.memoryUsage().rss / 1024 / 1024 - const timings = zeroTimings() - let sampledPeakRssMb = rssStartMb - const dbPaths: string[] = [] - const startedAt = performance.now() - - await runKeyedBatch( - inputPaths.map((inputPath, index) => ({ - value: { inputPath, index }, - key: `session:${index}`, - })), - { - concurrency, - async run({ inputPath, index }) { - const dbPath = path.join(root, `session-${index}.db`) - dbPaths[index] = dbPath - const parserMonitor = createBenchmarkParserMonitor() - const result = await streamingImport( - inputPath, - createImportDeps(dbPath, cacheDir, parserMonitor.logger), - { formatId: 'weflow' }, - `bench-${index}` - ) - if (!result.success || !result.diagnostics?.performance) { - throw new Error(`Import ${index} failed: ${result.error ?? 'missing performance diagnostics'}`) - } - parserMonitor.assertRustNativeCompleted(index) - addTimings(timings, result.diagnostics.performance.timings) - sampledPeakRssMb = Math.max(sampledPeakRssMb, result.diagnostics.performance.rssSampledPeakMb) - }, - } - ) - const durationMs = performance.now() - startedAt - - let messagesWritten = 0 - let membersWritten = 0 - let databaseBytes = 0 - const signatures: string[] = [] - for (const dbPath of dbPaths) { - const inspected = inspectDatabase(dbPath) - messagesWritten += inspected.messages - membersWritten += inspected.members - signatures.push(inspected.signature) - databaseBytes += statSync(dbPath).size - } - - return { - scenario: fileCount === 1 ? 'single' : 'batch', - parser, - fileCount, - concurrency, - inputMessages: fileCount * messagesPerFile, - inputBytes: statSync(fixturePath).size * fileCount, - durationMs, - sampledPeakRssMb, - sampledRssDeltaMb: Math.max(0, sampledPeakRssMb - rssStartMb), - databaseBytes, - messagesWritten, - membersWritten, - outputSignature: createHash('sha256').update(signatures.sort().join(':')).digest('hex'), - timings, - } - } finally { - rmSync(root, { recursive: true, force: true }) - } -} - -function runIsolated(fileCount: number, messagesPerFile: number, concurrency: number): BenchmarkResult { - const child = spawnSync( - process.execPath, - [ - '--expose-gc', - ...process.execArgv, - scriptPath, - '--worker', - String(fileCount), - String(messagesPerFile), - String(concurrency), - ], - { - cwd: process.cwd(), - encoding: 'utf-8', - maxBuffer: 10 * 1024 * 1024, - env: { ...process.env }, - } - ) - if (child.status !== 0) { - throw new Error(`Benchmark worker failed (${child.status}):\n${child.stdout}\n${child.stderr}`) - } - const resultLine = child.stdout.split('\n').find((line) => line.startsWith('BENCH_RESULT ')) - if (!resultLine) throw new Error(`Benchmark worker returned no result:\n${child.stdout}\n${child.stderr}`) - return JSON.parse(resultLine.slice('BENCH_RESULT '.length)) as BenchmarkResult -} - -function formatBytes(bytes: number): string { - return `${(bytes / 1024 / 1024).toFixed(1)} MB` -} - -function printResult(result: BenchmarkResult, label: string): void { - console.log( - `${label}: parser ${result.parser.implementation} | ${result.durationMs.toFixed(0)} ms | ` + - `sampled peak RSS ${result.sampledPeakRssMb.toFixed(0)} MB ` + - `(+${result.sampledRssDeltaMb.toFixed(0)} MB) | DB ${formatBytes(result.databaseBytes)} | ` + - `signature ${result.outputSignature}` - ) -} - -async function main(): Promise { - const [mode, first, second, third, fourth] = process.argv.slice(2) - if (mode === '--worker') { - const result = await runWorker(Number(first), Number(second), Number(third)) - console.log(`BENCH_RESULT ${JSON.stringify(result)}`) - return - } - - let fileCount: number - let messagesPerFile: number - let runs: number - let concurrency: number - if (mode === 'single') { - fileCount = 1 - messagesPerFile = Number(first ?? 100_000) - runs = Number(second ?? 3) - concurrency = 1 - } else if (mode === 'batch') { - fileCount = Number(first ?? 100) - messagesPerFile = Number(second ?? 10_000) - runs = Number(third ?? 3) - concurrency = Number(fourth ?? 1) - } else { - throw new Error('Usage: single [runs=3] | batch [runs=3] [concurrency=1]') - } - if (![fileCount, messagesPerFile, runs, concurrency].every((value) => Number.isInteger(value) && value > 0)) { - throw new Error('All benchmark counts must be positive integers') - } - - console.log( - `Benchmarking ${fileCount} file(s), ${messagesPerFile.toLocaleString()} messages/file, concurrency ${concurrency}, ${runs} isolated run(s)` - ) - const results: BenchmarkResult[] = [] - for (let index = 0; index < runs; index++) { - const result = runIsolated(fileCount, messagesPerFile, concurrency) - results.push(result) - printResult(result, `run ${index + 1}`) - } - - const sorted = [...results].sort((left, right) => left.durationMs - right.durationMs) - const median = sorted[Math.floor(sorted.length / 2)] - const parserSignatures = new Set(results.map((result) => JSON.stringify(result.parser))) - if (parserSignatures.size !== 1) { - throw new Error('Parser implementations differ between isolated runs') - } - if (results.some((result) => result.outputSignature !== median.outputSignature)) { - throw new Error('Output signatures differ between isolated runs') - } - if (median.messagesWritten !== fileCount * messagesPerFile) { - throw new Error(`Expected ${fileCount * messagesPerFile} messages, got ${median.messagesWritten}`) - } - - printResult(median, 'median') - const stageRows = Object.entries(median.timings) - .filter(([stage]) => stage !== 'totalMs') - .sort(([, left], [, right]) => right - left) - console.log('Median phase totals:') - for (const [stage, durationMs] of stageRows) { - console.log(` ${stage.padEnd(22)} ${durationMs.toFixed(1).padStart(10)} ms`) - } -} - -if (path.resolve(process.argv[1] ?? '') === scriptPath) { - void main().catch((error: unknown) => { - console.error(error) - process.exitCode = 1 - }) -} diff --git a/scripts/bench-streaming-import.test.ts b/scripts/bench-streaming-import.test.ts deleted file mode 100644 index 68d17ce2..00000000 --- a/scripts/bench-streaming-import.test.ts +++ /dev/null @@ -1,100 +0,0 @@ -import assert from 'node:assert/strict' -import fs from 'node:fs' -import path from 'node:path' -import test from 'node:test' -import Database from 'better-sqlite3' -import { CHAT_DB_TABLES } from '@openchatlab/core' -import { BetterSqliteAdapter } from '../packages/node-runtime/src/better-sqlite3-adapter' -import { buildFtsIndex } from '../packages/node-runtime/src/fts' -import { createChatLabTempDir } from './chatlab-temp.mjs' -import { createBenchmarkParserMonitor, inspectDatabase, resolveBenchmarkParser } from './bench-streaming-import.mts' - -const nativeBinding = path.resolve('apps/cli/native/better_sqlite3.node') - -test('benchmark requires and records the Rust Native WeFlow parser', () => { - assert.deepEqual(resolveBenchmarkParser({ available: true, disabled: false }, true), { - formatId: 'weflow', - implementation: 'rust-native', - nativeModuleAvailable: true, - }) - assert.throws( - () => resolveBenchmarkParser({ available: false, disabled: true }, false), - /CHATLAB_DISABLE_NATIVE_PERF=1/ - ) - assert.throws( - () => resolveBenchmarkParser({ available: true, disabled: false }, false), - /does not provide the WeFlow kernel/ - ) -}) - -test('benchmark rejects imports that do not complete with the Rust Native WeFlow parser', () => { - const missingNative = createBenchmarkParserMonitor() - assert.throws(() => missingNative.assertRustNativeCompleted(0), /did not start the Rust Native WeFlow parser/) - - const native = createBenchmarkParserMonitor() - native.logger.info('[NativeParser] Parsing WeFlow export with Rust kernel, size: 1.00 MB') - assert.doesNotThrow(() => native.assertRustNativeCompleted(1)) - - native.logger.info('[NativeParser] Rust parse failed, falling back to TS parser: invalid fixture') - assert.throws(() => native.assertRustNativeCompleted(1), /fell back to the TypeScript WeFlow parser/) -}) - -function createFixtureDatabase(dbPath: string): void { - const raw = new Database(dbPath, { nativeBinding }) - raw.exec(CHAT_DB_TABLES) - raw - .prepare( - `INSERT INTO meta - (name, platform, type, imported_at, group_id, group_avatar, owner_id) - VALUES (?, ?, ?, ?, ?, ?, ?)` - ) - .run('Signature Test', 'WeChat', 'group', 1_704_164_645, 'group-1', null, 'member-1') - const insertMember = raw.prepare( - `INSERT INTO member - (platform_id, account_name, group_nickname, aliases, avatar, roles) - VALUES (?, ?, ?, ?, ?, ?)` - ) - insertMember.run('member-1', 'Alice', 'Alice', '[]', null, '[]') - insertMember.run('member-2', 'Bob', 'Bob', '[]', null, '[]') - const insertMessage = raw.prepare( - `INSERT INTO message - (sender_id, sender_account_name, sender_group_nickname, ts, type, content) - VALUES (?, ?, ?, ?, ?, ?)` - ) - insertMessage.run(1, 'Alice', 'Alice', 1_704_164_645, 0, 'alpha') - insertMessage.run(2, 'Bob', 'Bob', 1_704_164_648, 0, 'bravo') - buildFtsIndex(new BetterSqliteAdapter(raw)) - raw.close() -} - -test('canonical database signature changes for equal-length content and sender permutations', () => { - const root = createChatLabTempDir('tests', 'bench-signature-') - const dbPath = path.join(root, 'session.db') - try { - createFixtureDatabase(dbPath) - const baseline = inspectDatabase(dbPath) - assert.equal(baseline.messages, 2) - assert.equal(baseline.members, 2) - - const raw = new Database(dbPath, { nativeBinding }) - raw.prepare('UPDATE message SET content = ? WHERE id = 1').run('omega') - raw.close() - const changedContent = inspectDatabase(dbPath) - assert.notEqual(changedContent.signature, baseline.signature) - - const senderDb = new Database(dbPath, { nativeBinding }) - senderDb.prepare('UPDATE message SET content = ? WHERE id = 1').run('alpha') - senderDb.exec( - `UPDATE message - SET sender_id = CASE id - WHEN 1 THEN 2 - WHEN 2 THEN 1 - END` - ) - senderDb.close() - const changedSenders = inspectDatabase(dbPath) - assert.notEqual(changedSenders.signature, baseline.signature) - } finally { - fs.rmSync(root, { recursive: true, force: true }) - } -})