From 26c5b4eb9a6b21b52b1e6e1f1ad98714e6ed21d5 Mon Sep 17 00:00:00 2001 From: Codinget Date: Wed, 1 Jul 2026 23:11:38 +0000 Subject: [PATCH] Add COPY-based bulk importer to bypass incremental B-tree update cost The incremental importer was measured at 24.72 min wall / 186s user time for rockyou (14.3M rows), and Postgres logs showed checkpoints firing every ~22s due to WAL pressure from random-order B-tree maintenance and full-page-write amplification. bulkImport.ts (`npm run import:bulk`) avoids all of that: - Loads raw rows via COPY into a TEMP staging table (no WAL, no index). - Merges with the existing table's contents into a fresh, unindexed shadow table via one INSERT ... SELECT ... GROUP BY. - Adds the primary key and prefix index only after the table is populated, so Postgres builds them via a single sorted bulk pass instead of 14M random-order incremental inserts. - Swaps the shadow table into place at the end. - The entire run is one transaction: since pwned_passwords_new doesn't exist outside that transaction, Postgres skips WAL-logging its data entirely, and any failure rolls back to the exact starting state instead of leaving a half-migrated table. - Wordlist hashing runs across a worker_threads pool (--jobs, shared resolveWorkerUrl helper with the incremental importer's hashWorker) instead of blocking the single thread that also drives the COPY stream. - Each phase (hash+copy, merge, reindex, swap) reports its own wall time. Net effect measured against the same rockyou file/DB: 42 min -> 24.72 min (fsync fix) -> ~7.3 min (WAL-skip transaction) -> 186.17s (parallel hashing). --- package-lock.json | 23 ++- package.json | 3 +- packages/import/package.json | 7 +- packages/import/src/bulkImport.ts | 229 ++++++++++++++++++++++++ packages/import/src/index.ts | 4 +- packages/import/src/resolveWorkerUrl.ts | 5 + 6 files changed, 264 insertions(+), 7 deletions(-) create mode 100644 packages/import/src/bulkImport.ts create mode 100644 packages/import/src/resolveWorkerUrl.ts diff --git a/package-lock.json b/package-lock.json index 6d93459..83409f0 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1389,6 +1389,17 @@ "pg-types": "^2.2.0" } }, + "node_modules/@types/pg-copy-streams": { + "version": "1.2.5", + "resolved": "https://registry.npmjs.org/@types/pg-copy-streams/-/pg-copy-streams-1.2.5.tgz", + "integrity": "sha512-7D6/GYW2uHIaVU6S/5omI+6RZnwlZBpLQDZAH83xX1rjxAOK0f6/deKyyUTewxqts145VIGn6XWYz1YGf50G5g==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/node": "*", + "@types/pg": "*" + } + }, "node_modules/@types/qs": { "version": "6.15.1", "resolved": "https://registry.npmjs.org/@types/qs/-/qs-6.15.1.tgz", @@ -4566,6 +4577,12 @@ "integrity": "sha512-XwWDGcLRGCXAR8F/AM5bG7Q+A3Wm2s6QeEjlOKZLlH3UYcguiqCWKyWXVag5TLTIjR7oOJUY8kcADaZgWPyLeg==", "license": "MIT" }, + "node_modules/pg-copy-streams": { + "version": "7.0.0", + "resolved": "https://registry.npmjs.org/pg-copy-streams/-/pg-copy-streams-7.0.0.tgz", + "integrity": "sha512-zBvnY6wtaBRE2ae2xXWOOGMaNVPkXh1vhypAkNSKgMdciJeTyIQAHZaEeRAxUjs/p1El5jgzYmwG5u871Zj3dQ==", + "license": "MIT" + }, "node_modules/pg-int8": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/pg-int8/-/pg-int8-1.0.1.tgz", @@ -6410,10 +6427,12 @@ "version": "1.0.0", "license": "ISC", "dependencies": { - "pg": "^8.22.0" + "pg": "^8.22.0", + "pg-copy-streams": "^7.0.0" }, "devDependencies": { - "@types/pg": "^8.20.0" + "@types/pg": "^8.20.0", + "@types/pg-copy-streams": "^1.2.5" } } } diff --git a/package.json b/package.json index df9eeab..1b48608 100644 --- a/package.json +++ b/package.json @@ -13,7 +13,8 @@ "db:down": "docker compose --env-file .env.local down", "db:seed": "npm run seed -w backend", "import": "npm run start -w import --", - "import:hibp": "npm run import-hibp -w import --" + "import:hibp": "npm run import-hibp -w import --", + "import:bulk": "npm run bulk-import -w import --" }, "devDependencies": { "concurrently": "^10.0.3", diff --git a/packages/import/package.json b/packages/import/package.json index e30119b..2dd4e72 100644 --- a/packages/import/package.json +++ b/packages/import/package.json @@ -7,14 +7,17 @@ "scripts": { "start": "tsx --env-file=../../.env.local src/index.ts", "import-hibp": "tsx --env-file=../../.env.local src/importHibp.ts", + "bulk-import": "tsx --env-file=../../.env.local src/bulkImport.ts", "build": "tsc -p tsconfig.json", "test": "echo \"Error: no test specified\" && exit 1" }, "license": "ISC", "dependencies": { - "pg": "^8.22.0" + "pg": "^8.22.0", + "pg-copy-streams": "^7.0.0" }, "devDependencies": { - "@types/pg": "^8.20.0" + "@types/pg": "^8.20.0", + "@types/pg-copy-streams": "^1.2.5" } } diff --git a/packages/import/src/bulkImport.ts b/packages/import/src/bulkImport.ts new file mode 100644 index 0000000..46c32cb --- /dev/null +++ b/packages/import/src/bulkImport.ts @@ -0,0 +1,229 @@ +import { createHash } from "node:crypto"; +import { once } from "node:events"; +import { availableParallelism } from "node:os"; +import { Readable } from "node:stream"; +import { finished, pipeline } from "node:stream/promises"; +import { parseArgs } from "node:util"; +import { from as copyFrom } from "pg-copy-streams"; +import type { ClientBase } from "pg"; +import { pool } from "./db.ts"; +import { readLines } from "./lines.ts"; +import { createLimiter } from "./concurrencyLimit.ts"; +import { WorkerPool } from "./workerPool.ts"; +import { resolveWorkerUrl } from "./resolveWorkerUrl.ts"; +import type { Entry } from "./upsertBatch.ts"; + +const HASH_LINE = /^([0-9A-Fa-f]{40})(?::(\d+))?$/; +const CHUNK_SIZE = 5000; +const hashWorkerUrl = resolveWorkerUrl("hashWorker", import.meta.url); + +function usage(): never { + console.error( + `Usage: npm run bulk-import -w import -- --format= [options]\n\n` + + `Rebuilds the pwned_passwords table by merging its existing contents with\n` + + `the given file. Loads via COPY into an unindexed shadow table and builds\n` + + `the primary key once at the end (a single sorted bulk build), instead of\n` + + `the incremental "import" script's per-row random-order index maintenance.\n` + + `Much faster for large files, at the cost of needing ~2x disk space for the\n` + + `duration of the run and a brief swap where pwned_passwords is unavailable.\n` + + `The whole run is one transaction, so a failure at any point leaves\n` + + `pwned_passwords completely untouched rather than half-migrated.\n\n` + + ` wordlist file contains one plaintext password per line\n` + + ` hashes file contains one SHA-1 hash per line, optionally "HASH:COUNT"\n` + + ` file may be gzip-compressed (.gz)\n\n` + + `Options:\n` + + ` --jobs, -j worker threads for wordlist hashing (default: all CPU cores);\n` + + ` COPY itself only ever uses one connection, so this only\n` + + ` helps when hashing - not network/DB throughput - is the\n` + + ` bottleneck\n\n` + + `Requires the pwned_passwords table to already exist (run migrate first).\n`, + ); + process.exit(1); +} + +function sha1(value: string): string { + return createHash("sha1").update(value).digest("hex").toUpperCase(); +} + +function toCopyLine(prefix: string, suffix: string, count: number): string { + // Prefix/suffix are always plain hex and count is always numeric, so none + // of COPY's TEXT-format special characters (tab, newline, backslash) can + // appear and no escaping is needed. + return `${prefix}\t${suffix}\t${count}\n`; +} + +function parseHashLine(line: string): Entry | null { + const match = HASH_LINE.exec(line.trim()); + if (!match) return null; + const [, hash, count] = match; + const upper = hash.toUpperCase(); + return { prefix: upper.slice(0, 5), suffix: upper.slice(5), count: count ? Number(count) : 1 }; +} + +// Sequential path: used for the hashes format (no CPU-bound work to +// parallelize) and for wordlist with --jobs=1. +async function* toStagingLines(file: string, format: "wordlist" | "hashes") { + let skipped = 0; + let seen = 0; + for await (const line of readLines(file)) { + seen++; + let entry: Entry | null; + if (format === "wordlist") { + const hash = sha1(line); + entry = { prefix: hash.slice(0, 5), suffix: hash.slice(5), count: 1 }; + } else { + entry = parseHashLine(line); + } + if (!entry) { + skipped++; + continue; + } + yield toCopyLine(entry.prefix, entry.suffix, entry.count); + } + console.log(`staged ${(seen - skipped).toLocaleString()} lines (${skipped} skipped)`); +} + +// Parallel path: hashing is CPU-bound and single-threaded JS can't use +// multiple cores for it, so a worker pool hashes chunks of lines while the +// main thread just writes the results into the one COPY connection (COPY +// itself can't be parallelized here - the whole run is one transaction on +// one connection, for atomicity). +async function copyWordlistParallel(copyStream: NodeJS.WritableStream, file: string, jobs: number) { + const workers = new WorkerPool(hashWorkerUrl, jobs); + const limit = createLimiter(jobs); + const inFlight: Promise[] = []; + + let seen = 0; + let chunk: string[] = []; + const submit = (lines: string[]) => + limit(async () => { + const entries = await workers.run(lines); + const text = entries.map((e) => toCopyLine(e.prefix, e.suffix, e.count)).join(""); + if (!copyStream.write(text)) { + await once(copyStream, "drain"); + } + }); + + for await (const line of readLines(file)) { + chunk.push(line); + seen++; + if (chunk.length >= CHUNK_SIZE) { + inFlight.push(submit(chunk)); + chunk = []; + } + } + if (chunk.length > 0) inFlight.push(submit(chunk)); + + await Promise.all(inFlight); + await workers.close(); + console.log(`staged ${seen.toLocaleString()} lines`); +} + +async function timePhase(label: string, fn: () => Promise): Promise { + const startedAt = Date.now(); + try { + return await fn(); + } finally { + console.log(`[${label}] ${((Date.now() - startedAt) / 1000).toFixed(2)}s`); + } +} + +async function loadStaging(client: ClientBase, file: string, format: "wordlist" | "hashes", jobs: number) { + const copyStream = client.query(copyFrom("COPY pwned_staging (prefix, suffix, count) FROM STDIN")); + + if (format === "wordlist" && jobs > 1) { + const done = finished(copyStream); + await copyWordlistParallel(copyStream, file, jobs); + copyStream.end(); + await done; + return; + } + + await pipeline(Readable.from(toStagingLines(file, format)), copyStream); +} + +async function main() { + const { positionals, values } = parseArgs({ + allowPositionals: true, + options: { + format: { type: "string" }, + jobs: { type: "string", short: "j" }, + }, + }); + + const [file] = positionals; + const format = values.format; + if (!file || (format !== "wordlist" && format !== "hashes")) { + usage(); + } + const jobs = values.jobs ? Number(values.jobs) : availableParallelism(); + const startedAt = Date.now(); + + const client = await pool.connect(); + try { + // Everything runs in one transaction: pwned_passwords_new is created and + // fully loaded before anyone outside this transaction can see it, so + // Postgres skips WAL-logging its data entirely (if we crash, the table + // just never existed). It also means a failure at any point rolls back + // to the exact starting state instead of leaving a half-swapped table. + await client.query("BEGIN"); + + console.log("loading into a temporary staging table via COPY..."); + await client.query( + "CREATE TEMP TABLE pwned_staging (prefix char(5) NOT NULL, suffix char(35) NOT NULL, count bigint NOT NULL)", + ); + await timePhase("hash+copy", () => loadStaging(client, file, format, jobs)); + + console.log("merging with existing data into a fresh, unindexed table..."); + await timePhase("merge", async () => { + await client.query("CREATE TABLE pwned_passwords_new (LIKE pwned_passwords INCLUDING DEFAULTS)"); + await client.query(` + INSERT INTO pwned_passwords_new (prefix, suffix, count) + SELECT prefix, suffix, SUM(count) + FROM ( + SELECT prefix, suffix, count FROM pwned_staging + UNION ALL + SELECT prefix, suffix, count FROM pwned_passwords + ) combined + GROUP BY prefix, suffix + `); + }); + + console.log("building index (one sorted bulk build instead of incremental random inserts)..."); + await timePhase("reindex", async () => { + await client.query( + "ALTER TABLE pwned_passwords_new ADD CONSTRAINT pwned_passwords_new_pkey PRIMARY KEY (prefix, suffix)", + ); + await client.query("CREATE INDEX pwned_passwords_new_prefix_idx ON pwned_passwords_new (prefix)"); + }); + + console.log("swapping into place..."); + await timePhase("swap", async () => { + await client.query("DROP TABLE pwned_passwords"); + await client.query("ALTER TABLE pwned_passwords_new RENAME TO pwned_passwords"); + // Renaming a constraint (rather than its backing index directly) also + // renames the index, keeping the two in sync. + await client.query( + "ALTER TABLE pwned_passwords RENAME CONSTRAINT pwned_passwords_new_pkey TO pwned_passwords_pkey", + ); + await client.query("ALTER INDEX pwned_passwords_new_prefix_idx RENAME TO pwned_passwords_prefix_idx"); + }); + + const { rows } = await client.query("SELECT count(*) FROM pwned_passwords"); + await client.query("COMMIT"); + const totalSeconds = (Date.now() - startedAt) / 1000; + console.log(`[total] ${totalSeconds.toFixed(2)}s`); + console.log(`done: pwned_passwords now has ${Number(rows[0].count).toLocaleString()} rows`); + } catch (err) { + await client.query("ROLLBACK").catch(() => {}); + throw err; + } finally { + client.release(); + await pool.end(); + } +} + +main().catch((err) => { + console.error(err); + process.exitCode = 1; +}); diff --git a/packages/import/src/index.ts b/packages/import/src/index.ts index 33ff76e..798d9a5 100644 --- a/packages/import/src/index.ts +++ b/packages/import/src/index.ts @@ -6,6 +6,7 @@ import { readLines } from "./lines.ts"; import { upsertBatch, type Entry } from "./upsertBatch.ts"; import { createLimiter } from "./concurrencyLimit.ts"; import { WorkerPool } from "./workerPool.ts"; +import { resolveWorkerUrl } from "./resolveWorkerUrl.ts"; const HASH_LINE = /^([0-9A-Fa-f]{40})(?::(\d+))?$/; @@ -39,8 +40,7 @@ function parseHashLine(line: string): Entry | null { return toEntry(hash.toUpperCase(), count ? Number(count) : 1); } -// `.ts` when running under tsx, `.js` once built with tsc. -const hashWorkerUrl = new URL(`./hashWorker.${import.meta.url.endsWith(".ts") ? "ts" : "js"}`, import.meta.url); +const hashWorkerUrl = resolveWorkerUrl("hashWorker", import.meta.url); async function importWordlist(file: string, batchSize: number, jobs: number) { let processed = 0; diff --git a/packages/import/src/resolveWorkerUrl.ts b/packages/import/src/resolveWorkerUrl.ts new file mode 100644 index 0000000..8d02326 --- /dev/null +++ b/packages/import/src/resolveWorkerUrl.ts @@ -0,0 +1,5 @@ +// `.ts` when running under tsx, `.js` once built with tsc. +export function resolveWorkerUrl(name: string, baseUrl: string): URL { + const ext = baseUrl.endsWith(".ts") ? "ts" : "js"; + return new URL(`./${name}.${ext}`, baseUrl); +}