diff --git a/packages/node-runtime/src/import/streaming-importer.test.ts b/packages/node-runtime/src/import/streaming-importer.test.ts index b54eac05..92b76e1e 100644 --- a/packages/node-runtime/src/import/streaming-importer.test.ts +++ b/packages/node-runtime/src/import/streaming-importer.test.ts @@ -4,7 +4,7 @@ import os from 'node:os' import path from 'node:path' import test from 'node:test' import Database from 'better-sqlite3' -import { CHAT_DB_TABLES } from '@openchatlab/core' +import { CHAT_DB_TABLES, type DatabaseAdapter } from '@openchatlab/core' import { isNativeFormatAvailable } from '@openchatlab/parser' import { BetterSqliteAdapter } from '../better-sqlite3-adapter' import { computeAndSetOverviewCache } from '../cache/session-cache' @@ -25,12 +25,15 @@ function makeTempDir(): string { return fs.mkdtempSync(path.join(baseDir, 'chatlab-streaming-import-')) } -function createImportDeps(dbPath: string): StreamImportDeps { +function createImportDeps( + dbPath: string, + wrapDatabase: (db: DatabaseAdapter) => DatabaseAdapter = (db) => db +): StreamImportDeps { return { openDatabase() { const db = new Database(dbPath, { nativeBinding }) db.exec(CHAT_DB_TABLES) - return new BetterSqliteAdapter(db) + return wrapDatabase(new BetterSqliteAdapter(db)) }, deleteDatabase() { for (const suffix of ['', '-wal', '-shm']) { @@ -47,6 +50,38 @@ function createImportDeps(dbPath: string): StreamImportDeps { } } +function createDelayedTransactionAdapter(db: DatabaseAdapter, delayMs: number): DatabaseAdapter { + let delayedMessageCommit = false + + return { + get readonly() { + return db.readonly + }, + exec(sql) { + if (!delayedMessageCommit && sql.trim().toUpperCase() === 'COMMIT') { + delayedMessageCommit = true + Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, delayMs) + } + db.exec(sql) + }, + prepare(sql) { + return db.prepare(sql) + }, + transaction(fn) { + return db.transaction(fn) + }, + pragma(pragma) { + if (pragma === 'wal_checkpoint(TRUNCATE)') { + Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, delayMs) + } + return db.pragma(pragma) + }, + close() { + db.close() + }, + } +} + function createCollectingLogger(messages: string[]): ImportLogger { return { info(message) { @@ -570,6 +605,23 @@ test('streamingImport applies incremental-equivalent deduplication on first impo assert.equal(row.count, 3) }) +test('streamingImport attributes transaction and checkpoint work to their import phases', async (t) => { + const root = makeTempDir() + t.after(() => fs.rmSync(root, { recursive: true, force: true })) + const filePath = writeDuplicateChatLabExport(root) + const dbPath = path.join(root, 'transaction-timing-test.db') + const delayMs = 50 + const deps = createImportDeps(dbPath, (db) => createDelayedTransactionAdapter(db, delayMs)) + + const result = await streamingImport(filePath, deps, { formatId: 'chatlab' }, 'transaction-timing-test') + + assert.equal(result.success, true) + const timings = result.diagnostics?.performance.timings + assert.ok(timings) + assert.ok(timings.messageWriteMs >= delayMs * 0.8, 'message writes should include the final transaction commit') + assert.ok(timings.checkpointMs >= delayMs * 0.8, 'checkpoint timing should include the final WAL checkpoint') +}) + test('streamingImport canonicalizes reserved SYSTEM senders and excludes them from overview counts', async (t) => { const root = makeTempDir() t.after(() => fs.rmSync(root, { recursive: true, force: true })) diff --git a/packages/node-runtime/src/import/streaming-importer.ts b/packages/node-runtime/src/import/streaming-importer.ts index 3cb48d44..03c41ed4 100644 --- a/packages/node-runtime/src/import/streaming-importer.ts +++ b/packages/node-runtime/src/import/streaming-importer.ts @@ -421,12 +421,30 @@ async function streamImportSingle( let lastCheckpointCount = 0 let inTransaction = false let messageTransactionCount = 0 + let messageTransactionMs = 0 const beginTransaction = () => { if (!inTransaction) { - db.exec('BEGIN TRANSACTION') - inTransaction = true - messageTransactionCount++ + const startedAt = performance.now() + try { + db.exec('BEGIN TRANSACTION') + inTransaction = true + messageTransactionCount++ + } finally { + messageTransactionMs += elapsedMs(startedAt) + } + } + } + + const commitTransaction = () => { + if (inTransaction) { + const startedAt = performance.now() + try { + db.exec('COMMIT') + inTransaction = false + } finally { + messageTransactionMs += elapsedMs(startedAt) + } } } @@ -438,14 +456,22 @@ async function streamImportSingle( } } + const measureCheckpoint = () => { + const startedAt = performance.now() + try { + doCheckpoint() + } finally { + timings.checkpointMs += elapsedMs(startedAt) + } + } + const commitAndBeginNew = () => { if (inTransaction) { - db.exec('COMMIT') - inTransaction = false + commitTransaction() logger?.perf('Commit transaction', totalMessageCount, BATCH_COMMIT_SIZE) if (totalMessageCount - lastCheckpointCount >= CHECKPOINT_INTERVAL) { - doCheckpoint() + measureCheckpoint() logger?.perf('WAL checkpoint', totalMessageCount) lastCheckpointCount = totalMessageCount } @@ -486,7 +512,8 @@ async function streamImportSingle( const finishParserTiming = () => { if (parserTimingFinished) return const parsePipelineMs = elapsedMs(parsePipelineStartedAt) - const writeCallbackMs = timings.metaWriteMs + timings.memberWriteMs + timings.messageWriteMs + const writeCallbackMs = + timings.metaWriteMs + timings.memberWriteMs + timings.messageWriteMs + messageTransactionMs + timings.checkpointMs timings.parserMs = Math.max(0, parsePipelineMs - writeCallbackMs) parserTimingFinished = true sampleRss() @@ -566,6 +593,8 @@ async function streamImportSingle( onMessageBatch: (messages: ParsedMessage[]) => { const startedAt = performance.now() + const transactionMsBefore = messageTransactionMs + const checkpointMsBefore = timings.checkpointMs try { callbackStats.onMessageBatchCalls++ callbackStats.totalMessagesReceived += messages.length @@ -635,7 +664,9 @@ async function streamImportSingle( } } } finally { - timings.messageWriteMs += elapsedMs(startedAt) + const nestedStageMs = + messageTransactionMs - transactionMsBefore + (timings.checkpointMs - checkpointMsBefore) + timings.messageWriteMs += Math.max(0, elapsedMs(startedAt) - nestedStageMs) sampleRss() } }, @@ -645,9 +676,9 @@ async function streamImportSingle( finishParserTiming() if (inTransaction) { - db.exec('COMMIT') - inTransaction = false + commitTransaction() } + timings.messageWriteMs += messageTransactionMs // Flush nickname history in batch onProgress({ @@ -719,9 +750,7 @@ async function streamImportSingle( message: '', }) await yieldToEventLoop() - const checkpointStartedAt = performance.now() - doCheckpoint() - timings.checkpointMs = elapsedMs(checkpointStartedAt) + measureCheckpoint() sampleRss() logger?.perf('WAL checkpoint done', totalMessageCount)