diff --git a/packages/shared/src/nodeSqliteClient.test.ts b/packages/shared/src/nodeSqliteClient.test.ts index 92ae7134017e9..ac65e7f12daa7 100644 --- a/packages/shared/src/nodeSqliteClient.test.ts +++ b/packages/shared/src/nodeSqliteClient.test.ts @@ -174,6 +174,72 @@ const makeTempDatabase = Effect.gen(function* () { return { filename, other }; }); +it.effect.each([false, true])("batches worker results without losing rows (reader=%s)", (reader) => + Effect.gen(function* () { + const { filename } = yield* makeTempDatabase; + yield* Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + const route = (effect: Effect.Effect) => + reader ? effect.pipe(SqliteClient.readOnly) : effect; + yield* route(sql`SELECT 1`); + const worker = startedWorkers + .filter((candidate) => !checkpointWorkers.includes(candidate)) + .at(-1)!; + const batches: Array> = []; + const record = (reply: { rows?: ReadonlyArray }) => { + if (Array.isArray(reply.rows) && reply.rows.length > 0) batches.push(reply.rows); + }; + worker.prependListener("message", record); + yield* Effect.gen(function* () { + const text = "x".repeat(16 * 1024); + const query = sql<{ n: number; text: string }>` + WITH RECURSIVE sequence(n) AS ( + SELECT 1 UNION ALL SELECT n + 1 FROM sequence WHERE n < 300 + ) SELECT n, ${text} AS text FROM sequence + `; + const rows = yield* reader ? route(query.values) : query; + assert.equal(rows.length, 300); + for (const [index, row] of rows.entries()) { + assert.deepEqual(row, reader ? [index + 1, text] : { n: index + 1, text }); + } + assert.isAbove(batches.length, 1); + for (const batch of batches) { + assert.isAtMost(batch.length, 16_384); + assert.isAtMost(batch.length * text.length * 2, 8 * 1024 * 1024); + } + batches.length = 0; + const narrow = sql<{ n: number }>` + WITH RECURSIVE sequence(n) AS ( + SELECT 1 UNION ALL SELECT n + 1 FROM sequence WHERE n < 10000 + ) SELECT n FROM sequence + `; + const narrowRows = yield* reader ? route(narrow.values) : narrow; + assert.equal(narrowRows.length, 10_000); + assert.deepEqual(narrowRows.at(-1), reader ? [10_000] : { n: 10_000 }); + assert.equal(batches.length, 1); + }).pipe(Effect.ensuring(Effect.sync(() => worker.off("message", record)))); + }).pipe( + Effect.provide(SqliteClient.layer({ filename, writerWorker: true, readerWorker: true })), + ); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), +); + +it.effect("reuses a worker query after an iterator constraint failure", () => + Effect.gen(function* () { + const { filename } = yield* makeTempDatabase; + yield* Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + yield* sql`CREATE TABLE unique_entries(name TEXT UNIQUE)`; + const insert = (name: string) => + sql<{ name: string }>`INSERT INTO unique_entries VALUES (${name}) RETURNING name`; + yield* insert("first"); + const error = yield* insert("first").pipe(Effect.flip); + assert.equal(error.reason._tag, "UniqueViolation"); + assert.deepEqual(yield* insert("second"), [{ name: "second" }]); + }).pipe(Effect.provide(SqliteClient.layer({ filename, writerWorker: true }))); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), +); + it.effect("keeps another connection from committing between a transaction's read and write", () => Effect.gen(function* () { const { filename, other } = yield* makeTempDatabase; diff --git a/packages/shared/src/nodeSqliteClient.ts b/packages/shared/src/nodeSqliteClient.ts index 73872a8805b0f..fda878a114d9a 100644 --- a/packages/shared/src/nodeSqliteClient.ts +++ b/packages/shared/src/nodeSqliteClient.ts @@ -242,7 +242,7 @@ parentPort.on("message", (request) => { parentPort.close(); return; } - const { id, sql, params, cached, values, raw, safeIntegers } = request; + const { id, sql, params, cached, values, raw, safeIntegers, received } = request; let operation = "open"; try { if (db === undefined) throw new Error("database connection is unavailable"); @@ -252,12 +252,37 @@ parentPort.on("message", (request) => { statement.setReadBigInts(safeIntegers); if (statement.columns().length > 0) { statement.setReturnArrays(values); - parentPort.postMessage({ id, rows: statement.all(...params) }); + let rows = []; + let bytes = 0; + let pendingBatch = false; + for (const row of statement.iterate(...params)) { + let rowBytes = 0; + for (const key in row) { + const value = row[key]; + rowBytes += typeof value === "string" ? value.length * 2 + : ArrayBuffer.isView(value) ? value.byteLength : 8; + } + if (rows.length > 0 && (rows.length >= 16384 || bytes + rowBytes > 8 * 1024 * 1024)) { + // Build the next batch during receipt, but allow only one unacknowledged message. + if (pendingBatch) Atomics.wait(received, 0, 0); + Atomics.store(received, 0, 0); + parentPort.postMessage({ id, rows, more: true }); + rows = []; + bytes = 0; + pendingBatch = true; + } + rows.push(row); + bytes += rowBytes; + } + if (pendingBatch) Atomics.wait(received, 0, 0); + parentPort.postMessage({ id, rows }); } else { const result = statement.run(...params); parentPort.postMessage({ id, rows: raw ? result : [] }); } } catch (error) { + // A failed iterator can report its previous SQLite error when the statement is reused. + statements.delete(sql); parentPort.postMessage({ id, error: { operation, ...describe(error) } }); } }); @@ -283,7 +308,7 @@ type WorkerStartup = | ({ readonly type: "unavailable" } & WorkerFailure); type WorkerReply = - | { readonly id: number; readonly rows: ReadonlyArray } + | { readonly id: number; readonly rows: ReadonlyArray; readonly more?: boolean } | { readonly id: number; readonly error: WorkerFailure & { readonly operation: "open" | "prepare" | "execute" }; @@ -353,7 +378,12 @@ const makeWorkerConnection = Effect.fnUntraced(function* (options: { let stoppedMessage: string | undefined; const pending = new Map< number, - { readonly endsTransaction: boolean; readonly resume: (reply: WorkerReply) => void } + { + readonly endsTransaction: boolean; + readonly received: Int32Array; + readonly rows: Array; + readonly resume: (reply: WorkerReply) => void; + } >(); const stop = (message: string) => { stoppedMessage = message; @@ -377,10 +407,21 @@ const makeWorkerConnection = Effect.fnUntraced(function* (options: { } const request = pending.get(reply.id); if (request === undefined) return; + if ("rows" in reply && reply.more) { + for (const row of reply.rows) request.rows.push(row); + Atomics.store(request.received, 0, 1); + Atomics.notify(request.received, 0); + return; + } pending.delete(reply.id); if (pending.size === 0) worker.unref(); options.onReply?.(); - request.resume(reply); + if ("rows" in reply && request.rows.length > 0) { + for (const row of reply.rows) request.rows.push(row); + request.resume({ id: reply.id, rows: request.rows }); + } else { + request.resume(reply); + } }); worker.on("error", (cause) => stop(`${label} failed: ${cause.message}`)); worker.on("exit", (code) => stop(`${label} exited with code ${code}`)); @@ -399,8 +440,11 @@ const makeWorkerConnection = Effect.fnUntraced(function* (options: { const safeIntegers = Boolean(Context.get(fiber.context, Client.SafeIntegers)); const reply = Effect.callback, SqlError>((resume) => { const id = nextRequestId++; + const received = new Int32Array(new SharedArrayBuffer(Int32Array.BYTES_PER_ELEMENT)); pending.set(id, { endsTransaction, + received, + rows: [], resume: (reply) => resume( "rows" in reply @@ -411,7 +455,7 @@ const makeWorkerConnection = Effect.fnUntraced(function* (options: { if (pending.size === 1) worker.ref(); const cached = input.cached && !hasLargeBinding(input.params); // oxlint-disable-next-line unicorn/require-post-message-target-origin -- Node workers do not accept a target origin. - worker.postMessage({ id, safeIntegers, ...input, cached }); + worker.postMessage({ id, safeIntegers, ...input, cached, received }); }); // A posted write runs to the end, as it did on the event loop; the caller and its lock // must not move on while the worker is still applying it. Reads stay cancellable.