From b1daa6509e133dda6d3b3ec102e5ba9debbfbd58 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 06:58:14 +0000 Subject: [PATCH 1/6] perf: pipeline kernel result batch fetching Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 1 + lib/kernel/KernelResultsProvider.ts | 24 ++- .../unit/kernel/KernelResultsProvider.test.ts | 187 ++++++++++++++++++ 3 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 tests/unit/kernel/KernelResultsProvider.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 77d44b38..0e605562 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Pipeline kernel result fetching with JavaScript batch processing to improve throughput. - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: (last error: )`. `` is `HTTP : ` (or `HTTP ` when the status text is empty) for HTTP failures, and `network error `, `network error ()` (e.g. `aborted`, when there is no `code`) or a bare `network error` for network failures (the raw network-error message, which can contain the presigned URL, is not copied). If no cause can be determined, the ` (last error: …)` suffix is omitted. The `RetryError` instance, `errorCode` and `payload` (including the node-fetch `Request` / `Response`) are unchanged. (#471) - Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470) - Reject non-scalar query parameter values at bind time with a `ParameterError` on both the Thrift and kernel backends, instead of stringifying them (which could leave a Thrift operation stuck before FINISHED). Newly rejected on the Thrift backend: arrays, plain objects, boxed primitives (e.g. `new String('x')`, `new Number(1)`), functions, symbols, and a `DBSQLParameter` that wraps another `DBSQLParameter`. Primitives, `Date`, `Int64`, and `DBSQLParameter` are still accepted. diff --git a/lib/kernel/KernelResultsProvider.ts b/lib/kernel/KernelResultsProvider.ts index b84fdca4..80fb67b5 100644 --- a/lib/kernel/KernelResultsProvider.ts +++ b/lib/kernel/KernelResultsProvider.ts @@ -26,6 +26,8 @@ export interface KernelFetchHandle { fetchNextBatch(): Promise<{ ipcBytes: Buffer } | null>; } +type Settled = { ok: true; value: { ipcBytes: Buffer } | null } | { ok: false; error: unknown }; + /** * `IResultsProvider` that pulls Arrow IPC batches from the * kernel via the napi `Statement` handle and adapts them onto the @@ -57,6 +59,8 @@ export default class KernelResultsProvider implements IResultsProvider; + constructor(statement: KernelFetchHandle) { this.statement = statement; } @@ -84,6 +88,16 @@ export default class KernelResultsProvider implements IResultsProvider ({ ok: true, value }), + (error): Settled => ({ ok: false, error }), + ); + } + } + // Pull the next batch from the kernel and stash it in `prefetched`, // or mark the stream exhausted. Used by both `hasMore` and `fetchNext` // to keep one batch buffered ahead so `hasMore` is accurate without @@ -93,12 +107,20 @@ export default class KernelResultsProvider implements IResultsProvider>; + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (error: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +} + +function makeBatch(rows: number[]) { + return { ipcBytes: Buffer.from(tableToIPC(tableFromArrays({ value: Int32Array.from(rows) }), 'stream')) }; +} + +function nextTurn(): Promise { + return new Promise((resolve) => { + setImmediate(resolve); + }); +} + +describe('KernelResultsProvider', () => { + const options = { limit: 1000 }; + + it('starts the next fetch before the consumer receives the current batch', async () => { + const first = deferred(); + const second = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>(); + fetchNextBatch.onFirstCall().returns(first.promise); + fetchNextBatch.onSecondCall().returns(second.promise); + const provider = new KernelResultsProvider({ fetchNextBatch }); + + const received = provider.fetchNext(options); + expect(fetchNextBatch.callCount).to.equal(1); + + const batch = makeBatch([1, 2]); + first.resolve(batch); + expect(await received).to.deep.equal({ batches: [batch.ipcBytes], rowCount: 2 }); + expect(fetchNextBatch.callCount).to.equal(2); + + second.resolve(null); + expect(await provider.hasMore()).to.be.false; + }); + + it('keeps at most one fetch in flight and does not fetch beyond the buffered batch', async () => { + const requests = [deferred(), deferred(), deferred()]; + let active = 0; + let maxActive = 0; + let index = 0; + const fetchNextBatch = sinon.stub<[], Promise>().callsFake(() => { + const request = requests[index]; + index += 1; + active += 1; + maxActive = Math.max(maxActive, active); + return request.promise.finally(() => { + active -= 1; + }); + }); + const provider = new KernelResultsProvider({ fetchNextBatch }); + + const hasMore = provider.hasMore(); + expect(active).to.equal(1); + requests[0].resolve(makeBatch([1])); + expect(await hasMore).to.be.true; + expect(active).to.equal(1); + expect(fetchNextBatch.callCount).to.equal(2); + + expect(await provider.hasMore()).to.be.true; + expect(await provider.hasMore()).to.be.true; + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + expect(fetchNextBatch.callCount).to.equal(2); + + requests[1].resolve(makeBatch([2, 3])); + await nextTurn(); + expect(active).to.equal(0); + expect(fetchNextBatch.callCount).to.equal(2); + + expect(await provider.hasMore()).to.be.true; + expect(active).to.equal(1); + expect((await provider.fetchNext(options)).rowCount).to.equal(2); + + requests[2].resolve(null); + expect(await provider.hasMore()).to.be.false; + expect(active).to.equal(0); + expect(maxActive).to.equal(1); + expect(fetchNextBatch.callCount).to.equal(3); + }); + + for (const prefetch of [false, true]) { + const source = prefetch ? 'prefetch' : 'initial fetch'; + it(`propagates ${source} errors without fetching again`, async () => { + const failed = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>().returns(failed.promise); + if (prefetch) { + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + } + const provider = new KernelResultsProvider({ fetchNextBatch }); + if (prefetch) { + await provider.fetchNext(options); + } + + const expected = new Error('fetch failed'); + const received = provider.fetchNext(options).then( + () => expect.fail('Expected a fetch error'), + (error: unknown) => error, + ); + failed.reject(expected); + expect(await received).to.equal(expected); + await nextTurn(); + expect(fetchNextBatch.callCount).to.equal(prefetch ? 2 : 1); + }); + + it(`stops fetching after ${prefetch ? 'a prefetched' : 'an initial'} null`, async () => { + const fetchNextBatch = sinon.stub<[], Promise>().resolves(null); + if (prefetch) { + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + } + const provider = new KernelResultsProvider({ fetchNextBatch }); + if (prefetch) { + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + } + + expect(await provider.hasMore()).to.be.false; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); + expect(await provider.hasMore()).to.be.false; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); + expect(fetchNextBatch.callCount).to.equal(prefetch ? 2 : 1); + }); + } + + it('handles a rejected prefetch after the consumer stops', async () => { + const abandoned = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>(); + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + fetchNextBatch.onSecondCall().returns(abandoned.promise); + const provider = new KernelResultsProvider({ fetchNextBatch }); + const unhandled: unknown[] = []; + const onUnhandled = (reason: unknown) => { + unhandled.push(reason); + }; + process.on('unhandledRejection', onUnhandled); + try { + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + abandoned.reject(new Error('abandoned prefetch failed')); + await nextTurn(); + await nextTurn(); + expect(unhandled).to.be.empty; + expect(fetchNextBatch.callCount).to.equal(2); + } finally { + process.removeListener('unhandledRejection', onUnhandled); + } + }); + + it('skips empty batches before data and before exhaustion', async () => { + const batch = makeBatch([1, 2, 3]); + const batches = [makeBatch([]), makeBatch([]), batch, makeBatch([]), null]; + const fetchNextBatch = sinon.stub<[], Promise>(); + batches.forEach((value, index) => fetchNextBatch.onCall(index).resolves(value)); + const provider = new KernelResultsProvider({ fetchNextBatch }); + + expect(await provider.hasMore()).to.be.true; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [batch.ipcBytes], rowCount: 3 }); + expect(await provider.hasMore()).to.be.false; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); + expect(fetchNextBatch.callCount).to.equal(batches.length); + }); +}); From f32bcaf26b6cec209c20d078aacbd59304ba9dea Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 16:24:54 +0000 Subject: [PATCH 2/6] fix: restrict kernel prefetch to async result handles Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 2 +- lib/kernel/KernelOperationBackend.ts | 5 +- lib/kernel/KernelResultsProvider.ts | 10 +- .../kernel/KernelOperationBackend.test.ts | 119 ++++++++++++++++++ .../unit/kernel/KernelResultsProvider.test.ts | 31 ++++- 5 files changed, 154 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e605562..87a3f97e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Unreleased -- Pipeline kernel result fetching with JavaScript batch processing to improve throughput. +- Prefetch one kernel result batch when backed by `AsyncStatement`; terminal `Statement` handles keep fetching on demand so close/cancel do not wait behind an added prefetch. Live throughput validation is pending. - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: (last error: )`. `` is `HTTP : ` (or `HTTP ` when the status text is empty) for HTTP failures, and `network error `, `network error ()` (e.g. `aborted`, when there is no `code`) or a bare `network error` for network failures (the raw network-error message, which can contain the presigned URL, is not copied). If no cause can be determined, the ` (last error: …)` suffix is omitted. The `RetryError` instance, `errorCode` and `payload` (including the node-fetch `Request` / `Response`) are unchanged. (#471) - Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470) - Reject non-scalar query parameter values at bind time with a `ParameterError` on both the Thrift and kernel backends, instead of stringifying them (which could leave a Thrift operation stuck before FINISHED). Newly rejected on the Thrift backend: arrays, plain objects, boxed primitives (e.g. `new String('x')`, `new Number(1)`), functions, symbols, and a `DBSQLParameter` that wraps another `DBSQLParameter`. Primitives, `Date`, `Int64`, and `DBSQLParameter` are still accepted. diff --git a/lib/kernel/KernelOperationBackend.ts b/lib/kernel/KernelOperationBackend.ts index 8173a0fe..c593242c 100644 --- a/lib/kernel/KernelOperationBackend.ts +++ b/lib/kernel/KernelOperationBackend.ts @@ -779,9 +779,8 @@ export default class KernelOperationBackend implements IOperationBackend { } const metadata = await this.getResultMetadata(); const handle = await this.getFetchHandle(); - // KernelResultsProvider consumes only `fetchNextBatch`; both the async result - // handle and the blocking statement satisfy that surface. - this.resultsProvider = new KernelResultsProvider(handle as unknown as KernelStatement); + // Only AsyncResultHandle has a fetch lock separate from close/cancel. + this.resultsProvider = new KernelResultsProvider(handle, this.asyncStatement !== undefined); // DECIMAL/BIGINT precision preservation is opt-in via the // `preserveBigNumericPrecision` connection option (default off). The kernel // always delivers native Arrow Decimal128 / Int64, so when enabled the diff --git a/lib/kernel/KernelResultsProvider.ts b/lib/kernel/KernelResultsProvider.ts index 80fb67b5..a5144864 100644 --- a/lib/kernel/KernelResultsProvider.ts +++ b/lib/kernel/KernelResultsProvider.ts @@ -52,6 +52,8 @@ type Settled = { ok: true; value: { ipcBytes: Buffer } | null } | { ok: false; e export default class KernelResultsProvider implements IResultsProvider { private readonly statement: KernelFetchHandle; + private readonly prefetchEnabled: boolean; + // Prefetched next batch so `hasMore()` can be answered without an // extra round-trip. Set by `prime()` (lazy) and by `fetchNext`. private prefetched?: ArrowBatch; @@ -61,8 +63,9 @@ export default class KernelResultsProvider implements IResultsProvider; - constructor(statement: KernelFetchHandle) { + constructor(statement: KernelFetchHandle, prefetchEnabled = false) { this.statement = statement; + this.prefetchEnabled = prefetchEnabled; } public async hasMore(): Promise { @@ -119,8 +122,9 @@ export default class KernelResultsProvider implements IResultsProvider() { + let resolve!: (value: T) => void; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + +function nextTurn(): Promise { + return new Promise((resolve) => { + setImmediate(resolve); + }); +} + // Minimal stub of the napi `Statement` surface that emits a precomputed // Arrow IPC payload per `fetchNextBatch()` call. Used to feed // `KernelOperationBackend` synthetic batches that mirror the kernel's @@ -388,4 +404,107 @@ describe('KernelOperationBackend — M0 datatype round-trip via napi → ArrowRe await backend.close(); expect(stub.closed).to.equal(true); }); + + for (const asyncHandle of [false, true]) { + const handleKind = asyncHandle ? 'AsyncStatement' : 'Statement'; + for (const method of ['close', 'cancel'] as const) { + it(`${handleKind} ${method} completes after an early stop without waiting for an added fetch`, async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const pending = deferred<{ ipcBytes: Buffer } | null>(); + let fetching = false; + const resultHandle = new StatementStub(schemaIpc, []); + const fetchNextBatch = sinon.stub(resultHandle, 'fetchNextBatch').resolves(null); + fetchNextBatch.onFirstCall().resolves({ ipcBytes: ipcFromColumns(schema, { x: [1] }) }); + fetchNextBatch.onSecondCall().resolves({ ipcBytes: ipcFromColumns(schema, { x: [2] }) }); + fetchNextBatch.onThirdCall().callsFake(() => { + fetching = true; + return pending.promise.finally(() => { + fetching = false; + }); + }); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const lifecycleHandle = asyncHandle ? asyncStatement : resultHandle; + const lifecycleCall = sinon.stub(lifecycleHandle, method).callsFake(async () => { + // Model Statement's shared mutex; AsyncStatement owns a separate lock. + if (!asyncHandle && fetching) { + await pending.promise; + } + }); + const context = new ClientContextStub(); + const backend = new KernelOperationBackend( + asyncHandle ? { asyncStatement, context } : { statement: resultHandle, context }, + ); + let completion: Promise | undefined; + try { + expect(await backend.fetchChunk({ limit: 1 })).to.deep.equal([{ x: 1 }]); + expect(fetchNextBatch.callCount).to.equal(asyncHandle ? 3 : 2); + expect(fetching).to.equal(asyncHandle); + + let completed = false; + completion = backend[method]().then(() => { + completed = true; + }); + await nextTurn(); + expect(completed).to.be.true; + expect(lifecycleCall.calledOnce).to.be.true; + expect(fetching).to.equal(asyncHandle); + } finally { + pending.resolve(null); + await completion; + } + }); + } + + it(`${handleKind} decoding errors close the operation without waiting for an added fetch`, async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const pending = deferred<{ ipcBytes: Buffer } | null>(); + let fetching = false; + const resultHandle = new StatementStub(schemaIpc, []); + const fetchNextBatch = sinon.stub(resultHandle, 'fetchNextBatch'); + fetchNextBatch.onFirstCall().resolves({ ipcBytes: ipcFromColumns(schema, { x: [1] }) }); + fetchNextBatch.onSecondCall().callsFake(() => { + fetching = true; + return pending.promise.finally(() => { + fetching = false; + }); + }); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const lifecycleHandle = asyncHandle ? asyncStatement : resultHandle; + const close = sinon.stub(lifecycleHandle, 'close').callsFake(async () => { + if (!asyncHandle && fetching) { + await pending.promise; + } + }); + const context = new ClientContextStub(); + const backend = new KernelOperationBackend( + asyncHandle ? { asyncStatement, context } : { statement: resultHandle, context }, + ); + await backend.getResultMetadata(); + const decode = sinon.stub(RecordBatchReader, 'from').throws(new Error('decode failed')); + let observed: unknown; + const received = backend.fetchChunk({ limit: 1 }).catch((error: unknown) => { + observed = error; + }); + try { + await nextTurn(); + expect(observed).to.be.instanceOf(Error); + expect((observed as Error).message).to.equal('decode failed'); + expect(close.calledOnce).to.be.true; + expect(fetchNextBatch.callCount).to.equal(asyncHandle ? 2 : 1); + expect(fetching).to.equal(asyncHandle); + } finally { + decode.restore(); + pending.resolve(null); + await received; + } + }); + } }); diff --git a/tests/unit/kernel/KernelResultsProvider.test.ts b/tests/unit/kernel/KernelResultsProvider.test.ts index 87ad84a9..79f88254 100644 --- a/tests/unit/kernel/KernelResultsProvider.test.ts +++ b/tests/unit/kernel/KernelResultsProvider.test.ts @@ -48,7 +48,7 @@ describe('KernelResultsProvider', () => { const fetchNextBatch = sinon.stub<[], Promise>(); fetchNextBatch.onFirstCall().returns(first.promise); fetchNextBatch.onSecondCall().returns(second.promise); - const provider = new KernelResultsProvider({ fetchNextBatch }); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); const received = provider.fetchNext(options); expect(fetchNextBatch.callCount).to.equal(1); @@ -76,7 +76,7 @@ describe('KernelResultsProvider', () => { active -= 1; }); }); - const provider = new KernelResultsProvider({ fetchNextBatch }); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); const hasMore = provider.hasMore(); expect(active).to.equal(1); @@ -114,7 +114,7 @@ describe('KernelResultsProvider', () => { if (prefetch) { fetchNextBatch.onFirstCall().resolves(makeBatch([1])); } - const provider = new KernelResultsProvider({ fetchNextBatch }); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); if (prefetch) { await provider.fetchNext(options); } @@ -135,7 +135,7 @@ describe('KernelResultsProvider', () => { if (prefetch) { fetchNextBatch.onFirstCall().resolves(makeBatch([1])); } - const provider = new KernelResultsProvider({ fetchNextBatch }); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); if (prefetch) { expect((await provider.fetchNext(options)).rowCount).to.equal(1); } @@ -153,7 +153,7 @@ describe('KernelResultsProvider', () => { const fetchNextBatch = sinon.stub<[], Promise>(); fetchNextBatch.onFirstCall().resolves(makeBatch([1])); fetchNextBatch.onSecondCall().returns(abandoned.promise); - const provider = new KernelResultsProvider({ fetchNextBatch }); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); const unhandled: unknown[] = []; const onUnhandled = (reason: unknown) => { unhandled.push(reason); @@ -176,7 +176,7 @@ describe('KernelResultsProvider', () => { const batches = [makeBatch([]), makeBatch([]), batch, makeBatch([]), null]; const fetchNextBatch = sinon.stub<[], Promise>(); batches.forEach((value, index) => fetchNextBatch.onCall(index).resolves(value)); - const provider = new KernelResultsProvider({ fetchNextBatch }); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); expect(await provider.hasMore()).to.be.true; expect(await provider.fetchNext(options)).to.deep.equal({ batches: [batch.ipcBytes], rowCount: 3 }); @@ -184,4 +184,23 @@ describe('KernelResultsProvider', () => { expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); expect(fetchNextBatch.callCount).to.equal(batches.length); }); + + it('fetches on demand by default', async () => { + const fetchNextBatch = sinon.stub<[], Promise>(); + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + fetchNextBatch.onSecondCall().resolves(makeBatch([2])); + fetchNextBatch.onThirdCall().resolves(null); + const provider = new KernelResultsProvider({ fetchNextBatch }); + + expect(await provider.hasMore()).to.be.true; + expect(fetchNextBatch.callCount).to.equal(1); + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + await nextTurn(); + expect(fetchNextBatch.callCount).to.equal(1); + + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + expect(fetchNextBatch.callCount).to.equal(2); + expect(await provider.hasMore()).to.be.false; + expect(fetchNextBatch.callCount).to.equal(3); + }); }); From 8b8ffe482f8e60748dc48742bd3486fb7cedc454 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 16:45:09 +0000 Subject: [PATCH 3/6] docs: update kernel prefetch note after live validation Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 87a3f97e..ab9e3c0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Unreleased -- Prefetch one kernel result batch when backed by `AsyncStatement`; terminal `Statement` handles keep fetching on demand so close/cancel do not wait behind an added prefetch. Live throughput validation is pending. +- Prefetch one kernel result batch when backed by `AsyncStatement`; terminal `Statement` handles keep fetching on demand so close/cancel do not wait behind an added prefetch. - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: (last error: )`. `` is `HTTP : ` (or `HTTP ` when the status text is empty) for HTTP failures, and `network error `, `network error ()` (e.g. `aborted`, when there is no `code`) or a bare `network error` for network failures (the raw network-error message, which can contain the presigned URL, is not copied). If no cause can be determined, the ` (last error: …)` suffix is omitted. The `RetryError` instance, `errorCode` and `payload` (including the node-fetch `Request` / `Response`) are unchanged. (#471) - Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470) - Reject non-scalar query parameter values at bind time with a `ParameterError` on both the Thrift and kernel backends, instead of stringifying them (which could leave a Thrift operation stuck before FINISHED). Newly rejected on the Thrift backend: arrays, plain objects, boxed primitives (e.g. `new String('x')`, `new Number(1)`), functions, symbols, and a `DBSQLParameter` that wraps another `DBSQLParameter`. Primitives, `Date`, `Int64`, and `DBSQLParameter` are still accepted. From 9bc404760ee8d4e0fbb7745da4f1967d927d3058 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 18:50:26 +0000 Subject: [PATCH 4/6] fix: serialize kernel priming and dispose abandoned prefetches Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 2 +- lib/kernel/KernelOperationBackend.ts | 16 +- lib/kernel/KernelResultsProvider.ts | 53 ++++- tests/e2e/kernel/fetch-pipelining-e2e.test.ts | 196 ++++++++++++++++++ .../kernel/KernelOperationBackend.test.ts | 25 ++- .../unit/kernel/KernelResultsProvider.test.ts | 94 ++++++++- 6 files changed, 372 insertions(+), 14 deletions(-) create mode 100644 tests/e2e/kernel/fetch-pipelining-e2e.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index ab9e3c0a..57efcda4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Unreleased -- Prefetch one kernel result batch when backed by `AsyncStatement`; terminal `Statement` handles keep fetching on demand so close/cancel do not wait behind an added prefetch. +- Prefetch one kernel result batch for queries submitted asynchronously (`runAsync: true`, or default execution that remains pending after the initial response). Queries that finish in the initial response fetch on demand. Close/cancel discard the prefetch without waiting for it, but an already-started native fetch may finish in the background. (#578) - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: (last error: )`. `` is `HTTP : ` (or `HTTP ` when the status text is empty) for HTTP failures, and `network error `, `network error ()` (e.g. `aborted`, when there is no `code`) or a bare `network error` for network failures (the raw network-error message, which can contain the presigned URL, is not copied). If no cause can be determined, the ` (last error: …)` suffix is omitted. The `RetryError` instance, `errorCode` and `payload` (including the node-fetch `Request` / `Response`) are unchanged. (#471) - Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470) - Reject non-scalar query parameter values at bind time with a `ParameterError` on both the Thrift and kernel backends, instead of stringifying them (which could leave a Thrift operation stuck before FINISHED). Newly rejected on the Thrift backend: arrays, plain objects, boxed primitives (e.g. `new String('x')`, `new Number(1)`), functions, symbols, and a `DBSQLParameter` that wraps another `DBSQLParameter`. Primitives, `Date`, `Int64`, and `DBSQLParameter` are still accepted. diff --git a/lib/kernel/KernelOperationBackend.ts b/lib/kernel/KernelOperationBackend.ts index c593242c..00b2bb10 100644 --- a/lib/kernel/KernelOperationBackend.ts +++ b/lib/kernel/KernelOperationBackend.ts @@ -309,7 +309,7 @@ export default class KernelOperationBackend implements IOperationBackend { // wedged, so nothing downstream forces another close). We still don't // mask the original fetch error, but log the close failure at warn so // the leak is diagnosable rather than completely invisible. - await kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id).catch((closeErr) => { + await this.close().catch((closeErr) => { const cause = closeErr instanceof Error ? closeErr.message : String(closeErr); this.context .getLogger() @@ -443,10 +443,12 @@ export default class KernelOperationBackend implements IOperationBackend { } public async cancel(): Promise { + this.resultsProvider?.dispose(); return kernelCancel(this.lifecycle, this.lifecycleHandle, this.context, this._id); } public async close(): Promise { + this.resultsProvider?.dispose(); return kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id); } @@ -706,7 +708,7 @@ export default class KernelOperationBackend implements IOperationBackend { * original error; warn-logs a close failure so the leak is diagnosable. */ private async bestEffortClose(): Promise { - await kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id).catch((closeErr) => { + await this.close().catch((closeErr) => { const cause = closeErr instanceof Error ? closeErr.message : String(closeErr); this.context .getLogger() @@ -779,8 +781,16 @@ export default class KernelOperationBackend implements IOperationBackend { } const metadata = await this.getResultMetadata(); const handle = await this.getFetchHandle(); + failIfNotActive(this.lifecycle); + if (this.resultSlicer) { + return this.resultSlicer; + } // Only AsyncResultHandle has a fetch lock separate from close/cancel. - this.resultsProvider = new KernelResultsProvider(handle, this.asyncStatement !== undefined); + this.resultsProvider = new KernelResultsProvider( + handle, + this.asyncStatement !== undefined, + this.context.getLogger(), + ); // DECIMAL/BIGINT precision preservation is opt-in via the // `preserveBigNumericPrecision` connection option (default off). The kernel // always delivers native Arrow Decimal128 / Int64, so when enabled the diff --git a/lib/kernel/KernelResultsProvider.ts b/lib/kernel/KernelResultsProvider.ts index a5144864..b85210c6 100644 --- a/lib/kernel/KernelResultsProvider.ts +++ b/lib/kernel/KernelResultsProvider.ts @@ -14,6 +14,7 @@ import IResultsProvider, { ResultsProviderFetchNextOptions } from '../result/IResultsProvider'; import { ArrowBatch } from '../result/utils'; +import IDBSQLLogger, { LogLevel } from '../contracts/IDBSQLLogger'; import { countRowsInIpc, patchIpcBytes } from './KernelArrowIpc'; /** @@ -63,9 +64,33 @@ export default class KernelResultsProvider implements IResultsProvider; - constructor(statement: KernelFetchHandle, prefetchEnabled = false) { + private primePromise?: Promise; + + private disposed = false; + + private readonly logger?: IDBSQLLogger; + + constructor(statement: KernelFetchHandle, prefetchEnabled = false, logger?: IDBSQLLogger) { this.statement = statement; this.prefetchEnabled = prefetchEnabled; + this.logger = logger; + } + + public dispose(): void { + if (this.disposed) { + return; + } + this.disposed = true; + this.exhausted = true; + this.prefetched = undefined; + // The binding cannot abort an AsyncResultHandle fetch already in progress. + const abandoned = this.inflight; + this.inflight = undefined; + abandoned?.then((settled) => { + if (!settled.ok) { + this.logger?.log(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal'); + } + }); } public async hasMore(): Promise { @@ -94,10 +119,14 @@ export default class KernelResultsProvider implements IResultsProvider ({ ok: true, value }), - (error): Settled => ({ ok: false, error }), - ); + try { + this.inflight = this.statement.fetchNextBatch().then( + (value): Settled => ({ ok: true, value }), + (error): Settled => ({ ok: false, error }), + ); + } catch (error) { + this.inflight = Promise.resolve({ ok: false, error }); + } } } @@ -105,7 +134,16 @@ export default class KernelResultsProvider implements IResultsProvider { + private prime(): Promise { + if (!this.primePromise) { + this.primePromise = this.primeNext().finally(() => { + this.primePromise = undefined; + }); + } + return this.primePromise; + } + + private async primeNext(): Promise { // Loop rather than self-recurse: a long run of empty batches drains // iteratively toward the null sentinel without building a deep promise // chain. @@ -114,6 +152,9 @@ export default class KernelResultsProvider implements IResultsProvider(promise: Promise, ms: number): Promise { + let timer: NodeJS.Timeout | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error(`Timed out after ${ms} ms`)), ms); + }), + ]); + } finally { + clearTimeout(timer); + } +} + +// Tunnel TLS unchanged; pause storage responses while warehouse RPCs continue. +async function cloudFetchGate(workspaceHost: string) { + const sockets = new Set(); + const storage = new Set(); + let paused = false; + const server = http.createServer(); + server.on('connect', (request, downstream, head) => { + const target = new URL(`https://${request.url}`); + const upstream = net.connect(Number(target.port || 443), target.hostname); + const isStorage = target.hostname !== workspaceHost; + for (const socket of [downstream, upstream]) { + sockets.add(socket); + socket.on('close', () => sockets.delete(socket)); + socket.on('error', () => { + downstream.destroy(); + upstream.destroy(); + }); + } + if (isStorage) { + storage.add(upstream); + upstream.on('close', () => storage.delete(upstream)); + } + upstream.on('connect', () => { + downstream.write('HTTP/1.1 200 Connection Established\r\n\r\n'); + upstream.write(head); + downstream.pipe(upstream); + upstream.pipe(downstream); + if (isStorage && paused) upstream.pause(); + }); + }); + await new Promise((resolve) => { + server.listen(0, '127.0.0.1', resolve); + }); + return { + port: (server.address() as net.AddressInfo).port, + pause() { + expect(storage.size, 'expected CloudFetch storage tunnels').to.be.greaterThan(0); + paused = true; + storage.forEach((socket) => socket.pause()); + }, + resume() { + paused = false; + storage.forEach((socket) => socket.resume()); + }, + close() { + sockets.forEach((socket) => socket.destroy()); + return new Promise((resolve) => { + server.close(() => resolve()); + }); + }, + }; +} + +describe('kernel fetch pipelining — live CloudFetch lifecycle', function suite() { + this.timeout(180_000); + const host = process.env.DATABRICKS_PECOTESTING_SERVER_HOSTNAME || process.env.E2E_HOST; + const path = process.env.DATABRICKS_PECOTESTING_HTTP_PATH || process.env.E2E_PATH; + const token = process.env.DATABRICKS_PECOTESTING_TOKEN_PERSONAL || process.env.E2E_ACCESS_TOKEN; + + before(function gate() { + if (!host || !path || !token) { + // eslint-disable-next-line no-invalid-this + this.skip(); + } + }); + + for (const method of ['close', 'cancel'] as const) { + it(`${method} returns while a native CloudFetch prefetch remains blocked`, async () => { + const workspaceHost = new URL(`https://${host?.replace(/^https?:\/\//, '')}`).hostname; + const { prototype } = getKernelNative().AsyncResultHandle; + const proxy = await cloudFetchGate(workspaceHost); + const originalFetch = prototype.fetchNextBatch; + const pending = new Set>(); + let fetchCalls = 0; + // Observe real native futures; do not model or replace their locks. + prototype.fetchNextBatch = function trackFetch(this: KernelNativeAsyncResultHandle) { + fetchCalls += 1; + const promise = originalFetch.call(this); + pending.add(promise); + promise.then( + () => pending.delete(promise), + () => pending.delete(promise), + ); + return promise; + }; + let cloudFetch = false; + const client = new DBSQLClient({ + logger: { + getLevel: () => LogLevel.info, + log(_level, message) { + if (message.includes('Using CloudFetch reader:')) cloudFetch = true; + }, + }, + }); + let session: IDBSQLSession | undefined; + let operation: IOperation | undefined; + let consuming: Promise> | undefined; + try { + await client.connect({ + host: workspaceHost, + path: path!, + token: token!, + proxy: { protocol: 'http', host: '127.0.0.1', port: proxy.port }, + useKernel: true, + disableRowMaterialization: true, + telemetryEnabled: false, + } as Parameters[0] & InternalConnectionOptions); + session = await client.openSession({ configuration: { use_cached_result: 'false' } }); + // Exceed the kernel's downloaded-chunk buffer without depending on a table. + operation = await session.executeStatement( + 'SELECT id, sha2(CAST(id AS STRING), 512) AS payload FROM range(0, 5000000)', + { runAsync: true, statementConf: { use_cached_result: 'false' } }, + ); + await operation.finished(); + expect(await operation.fetchChunk({ maxRows: 1 })).to.have.length(1); + proxy.pause(); + + // Drain buffered chunks until the native fetch really waits on storage. + let blocked = false; + while (!blocked) { + consuming = operation.fetchChunk({ maxRows: 65536 }); + // eslint-disable-next-line no-await-in-loop + const rows = await Promise.race([ + consuming, + new Promise((resolve) => { + setTimeout(resolve, 100); + }), + ]); + if (rows === undefined) blocked = pending.size > 0; + else expect(rows.length, 'query exhausted before exercising a blocked prefetch').to.be.greaterThan(0); + if (!blocked) { + // eslint-disable-next-line no-await-in-loop + await consuming; + } + } + expect(fetchCalls).to.be.greaterThan(2); + const blockedFetch = [...pending][0]; + const status = await within(operation[method](), 5_000); + expect(status.isSuccess).to.be.true; + expect(pending.has(blockedFetch), 'lifecycle call waited for the native fetch').to.be.true; + // CloudFetch logging is delivered asynchronously through napi. + expect(cloudFetch).to.be.true; + } finally { + proxy.resume(); + try { + await within(Promise.allSettled([...(consuming ? [consuming] : []), ...pending]), 30_000); + await operation?.close(); + await session?.close(); + await client.close(); + } finally { + prototype.fetchNextBatch = originalFetch; + await proxy.close(); + } + } + }); + } +}); diff --git a/tests/unit/kernel/KernelOperationBackend.test.ts b/tests/unit/kernel/KernelOperationBackend.test.ts index cb79e727..8b913b85 100644 --- a/tests/unit/kernel/KernelOperationBackend.test.ts +++ b/tests/unit/kernel/KernelOperationBackend.test.ts @@ -42,13 +42,16 @@ import { import KernelOperationBackend from '../../../lib/kernel/KernelOperationBackend'; import { TTypeId } from '../../../thrift/TCLIService_types'; import ClientContextStub from '../.stubs/ClientContextStub'; +import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; function deferred() { let resolve!: (value: T) => void; - const promise = new Promise((resolvePromise) => { + let reject!: (error: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { resolve = resolvePromise; + reject = rejectPromise; }); - return { promise, resolve }; + return { promise, resolve, reject }; } function nextTurn(): Promise { @@ -435,6 +438,7 @@ describe('KernelOperationBackend — M0 datatype round-trip via napi → ArrowRe } }); const context = new ClientContextStub(); + const log = sinon.spy(context.logger, 'log'); const backend = new KernelOperationBackend( asyncHandle ? { asyncStatement, context } : { statement: resultHandle, context }, ); @@ -443,7 +447,6 @@ describe('KernelOperationBackend — M0 datatype round-trip via napi → ArrowRe expect(await backend.fetchChunk({ limit: 1 })).to.deep.equal([{ x: 1 }]); expect(fetchNextBatch.callCount).to.equal(asyncHandle ? 3 : 2); expect(fetching).to.equal(asyncHandle); - let completed = false; completion = backend[method]().then(() => { completed = true; @@ -452,6 +455,14 @@ describe('KernelOperationBackend — M0 datatype round-trip via napi → ArrowRe expect(completed).to.be.true; expect(lifecycleCall.calledOnce).to.be.true; expect(fetching).to.equal(asyncHandle); + if (asyncHandle) { + pending.reject(new Error('abandoned prefetch failed')); + } + await nextTurn(); + expect(log.calledWith(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal')).to.equal( + asyncHandle, + ); + expect(fetchNextBatch.callCount).to.equal(asyncHandle ? 3 : 2); } finally { pending.resolve(null); await completion; @@ -484,6 +495,7 @@ describe('KernelOperationBackend — M0 datatype round-trip via napi → ArrowRe } }); const context = new ClientContextStub(); + const log = sinon.spy(context.logger, 'log'); const backend = new KernelOperationBackend( asyncHandle ? { asyncStatement, context } : { statement: resultHandle, context }, ); @@ -500,6 +512,13 @@ describe('KernelOperationBackend — M0 datatype round-trip via napi → ArrowRe expect(close.calledOnce).to.be.true; expect(fetchNextBatch.callCount).to.equal(asyncHandle ? 2 : 1); expect(fetching).to.equal(asyncHandle); + if (asyncHandle) { + pending.reject(new Error('abandoned prefetch failed')); + } + await nextTurn(); + expect(log.calledWith(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal')).to.equal( + asyncHandle, + ); } finally { decode.restore(); pending.resolve(null); diff --git a/tests/unit/kernel/KernelResultsProvider.test.ts b/tests/unit/kernel/KernelResultsProvider.test.ts index 79f88254..2053ae0e 100644 --- a/tests/unit/kernel/KernelResultsProvider.test.ts +++ b/tests/unit/kernel/KernelResultsProvider.test.ts @@ -14,8 +14,9 @@ import { expect } from 'chai'; import sinon from 'sinon'; -import { tableFromArrays, tableToIPC } from 'apache-arrow'; +import { tableFromArrays, tableFromIPC, tableToIPC } from 'apache-arrow'; import KernelResultsProvider, { KernelFetchHandle } from '../../../lib/kernel/KernelResultsProvider'; +import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; type FetchResult = Awaited>; @@ -42,6 +43,97 @@ function nextTurn(): Promise { describe('KernelResultsProvider', () => { const options = { limit: 1000 }; + for (const prefetch of [false, true]) { + for (const fetchFirst of [false, true]) { + it(`preserves every batch with concurrent hasMore/fetchNext (prefetch=${prefetch}, fetchFirst=${fetchFirst})`, async () => { + const first = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>().resolves(null); + fetchNextBatch.onFirstCall().returns(first.promise); + [2, 3, 4].forEach((row, index) => fetchNextBatch.onCall(index + 1).resolves(makeBatch([row]))); + const provider = new KernelResultsProvider({ fetchNextBatch }, prefetch); + const calls = fetchFirst + ? [provider.fetchNext(options), provider.hasMore()] + : [provider.hasMore(), provider.fetchNext(options)]; + first.resolve(makeBatch([1])); + const results = await Promise.all(calls); + const batches = [results[fetchFirst ? 0 : 1] as Awaited>]; + while (await provider.hasMore()) { + batches.push(await provider.fetchNext(options)); + } + const rows = batches.flatMap((batch) => Array.from(tableFromIPC(batch.batches[0]).getChild('value')!)); + expect(rows).to.deep.equal([1, 2, 3, 4]); + expect(fetchNextBatch.callCount).to.equal(5); + }); + } + } + + it('delivers the current batch before surfacing a synchronous lookahead failure', async () => { + const expected = new Error('synchronous fetch failure'); + const fetchNextBatch = sinon.stub<[], Promise>().throws(expected); + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + const error = await provider.fetchNext(options).then( + () => expect.fail('Expected the lookahead error'), + (reason: unknown) => reason, + ); + expect(error).to.equal(expected); + expect(fetchNextBatch.callCount).to.equal(2); + }); + + it('shares a synchronous initial failure between concurrent callers', async () => { + const expected = new Error('synchronous initial failure'); + const fetchNextBatch = sinon.stub<[], Promise>().throws(expected); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + const results = await Promise.allSettled([provider.hasMore(), provider.fetchNext(options)]); + expect(results).to.deep.equal([ + { status: 'rejected', reason: expected }, + { status: 'rejected', reason: expected }, + ]); + expect(fetchNextBatch.calledOnce).to.be.true; + }); + + for (const rejectBeforeDispose of [false, true]) { + it(`discards a failed prefetch and logs it once (already settled=${rejectBeforeDispose})`, async () => { + const pending = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>().returns(pending.promise); + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + const logger = { log: sinon.spy() }; + const provider = new KernelResultsProvider({ fetchNextBatch }, true, logger); + await provider.hasMore(); + if (rejectBeforeDispose) { + pending.reject(new Error('abandoned fetch failed')); + await nextTurn(); + } + + provider.dispose(); + provider.dispose(); + expect(await provider.hasMore()).to.be.false; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); + if (!rejectBeforeDispose) { + pending.reject(new Error('abandoned fetch failed')); + } + await nextTurn(); + expect(logger.log.calledOnceWithExactly(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal')) + .to.be.true; + expect(fetchNextBatch.callCount).to.equal(2); + }); + } + + it('discards an initial fetch that finishes after disposal without starting a lookahead', async () => { + const pending = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>().returns(pending.promise); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + const result = provider.fetchNext(options); + provider.dispose(); + pending.resolve(makeBatch([1])); + + expect(await result).to.deep.equal({ batches: [], rowCount: 0 }); + expect(await provider.hasMore()).to.be.false; + expect(fetchNextBatch.calledOnce).to.be.true; + }); + it('starts the next fetch before the consumer receives the current batch', async () => { const first = deferred(); const second = deferred(); From 14340321a55a28e0422f22dcfb2adf7681e12e3e Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 19:03:39 +0000 Subject: [PATCH 5/6] test: require a sustained CloudFetch wait before lifecycle checks Signed-off-by: Vu Anh Phung --- tests/e2e/kernel/fetch-pipelining-e2e.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/e2e/kernel/fetch-pipelining-e2e.test.ts b/tests/e2e/kernel/fetch-pipelining-e2e.test.ts index eb802989..da7493f0 100644 --- a/tests/e2e/kernel/fetch-pipelining-e2e.test.ts +++ b/tests/e2e/kernel/fetch-pipelining-e2e.test.ts @@ -154,7 +154,7 @@ describe('kernel fetch pipelining — live CloudFetch lifecycle', function suite expect(await operation.fetchChunk({ maxRows: 1 })).to.have.length(1); proxy.pause(); - // Drain buffered chunks until the native fetch really waits on storage. + // Buffered parsing can take >100 ms on CI; require a sustained wait. let blocked = false; while (!blocked) { consuming = operation.fetchChunk({ maxRows: 65536 }); @@ -162,7 +162,7 @@ describe('kernel fetch pipelining — live CloudFetch lifecycle', function suite const rows = await Promise.race([ consuming, new Promise((resolve) => { - setTimeout(resolve, 100); + setTimeout(resolve, 2_000); }), ]); if (rows === undefined) blocked = pending.size > 0; From 062ae68f3d3f5a5aced9fac8612be86cc0a3fc25 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 19:54:04 +0000 Subject: [PATCH 6/6] fix: preserve kernel results when close fails Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 2 +- lib/kernel/KernelOperationBackend.ts | 25 ++++- .../kernel/KernelOperationBackend.test.ts | 104 ++++++++++++++++++ 3 files changed, 125 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 57efcda4..79e1bb14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Unreleased -- Prefetch one kernel result batch for queries submitted asynchronously (`runAsync: true`, or default execution that remains pending after the initial response). Queries that finish in the initial response fetch on demand. Close/cancel discard the prefetch without waiting for it, but an already-started native fetch may finish in the background. (#578) +- Prefetch one kernel result batch for queries submitted asynchronously (`runAsync: true`, or default execution that remains pending after the initial response). Queries that finish in the initial response fetch on demand. Successful close and cancellation discard the prefetch without waiting for it; failed close requests preserve unread batches. An already-started native fetch may finish in the background. (#578) - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: (last error: )`. `` is `HTTP : ` (or `HTTP ` when the status text is empty) for HTTP failures, and `network error `, `network error ()` (e.g. `aborted`, when there is no `code`) or a bare `network error` for network failures (the raw network-error message, which can contain the presigned URL, is not copied). If no cause can be determined, the ` (last error: …)` suffix is omitted. The `RetryError` instance, `errorCode` and `payload` (including the node-fetch `Request` / `Response`) are unchanged. (#471) - Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470) - Reject non-scalar query parameter values at bind time with a `ParameterError` on both the Thrift and kernel backends, instead of stringifying them (which could leave a Thrift operation stuck before FINISHED). Newly rejected on the Thrift backend: arrays, plain objects, boxed primitives (e.g. `new String('x')`, `new Number(1)`), functions, symbols, and a `DBSQLParameter` that wraps another `DBSQLParameter`. Primitives, `Date`, `Int64`, and `DBSQLParameter` are still accepted. diff --git a/lib/kernel/KernelOperationBackend.ts b/lib/kernel/KernelOperationBackend.ts index 00b2bb10..232e8dba 100644 --- a/lib/kernel/KernelOperationBackend.ts +++ b/lib/kernel/KernelOperationBackend.ts @@ -206,6 +206,8 @@ export default class KernelOperationBackend implements IOperationBackend { private resultsProvider?: KernelResultsProvider; + private closePromise?: Promise; + private metadata?: ResultMetadata; private metadataPromise?: Promise; @@ -443,13 +445,26 @@ export default class KernelOperationBackend implements IOperationBackend { } public async cancel(): Promise { - this.resultsProvider?.dispose(); - return kernelCancel(this.lifecycle, this.lifecycleHandle, this.context, this._id); + const result = kernelCancel(this.lifecycle, this.lifecycleHandle, this.context, this._id); + // A pending close makes cancel a no-op without setting the cancelled flag. + if (this.lifecycle.isCancelled) { + this.resultsProvider?.dispose(); + } + return result; } - public async close(): Promise { - this.resultsProvider?.dispose(); - return kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id); + public close(): Promise { + if (!this.closePromise) { + this.closePromise = kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id) + .then((status) => { + this.resultsProvider?.dispose(); + return status; + }) + .finally(() => { + this.closePromise = undefined; + }); + } + return this.closePromise; } // --------------------------------------------------------------------------- diff --git a/tests/unit/kernel/KernelOperationBackend.test.ts b/tests/unit/kernel/KernelOperationBackend.test.ts index 8b913b85..0541c761 100644 --- a/tests/unit/kernel/KernelOperationBackend.test.ts +++ b/tests/unit/kernel/KernelOperationBackend.test.ts @@ -40,6 +40,7 @@ import { } from 'apache-arrow'; import KernelOperationBackend from '../../../lib/kernel/KernelOperationBackend'; +import DBSQLOperation from '../../../lib/DBSQLOperation'; import { TTypeId } from '../../../thrift/TCLIService_types'; import ClientContextStub from '../.stubs/ClientContextStub'; import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; @@ -408,6 +409,109 @@ describe('KernelOperationBackend — M0 datatype round-trip via napi → ArrowRe expect(stub.closed).to.equal(true); }); + for (const concurrent of [false, true]) { + it(`preserves unread rows after ${concurrent ? 'concurrent close calls fail' : 'a close call fails'}`, async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const resultHandle = new StatementStub( + schemaIpc, + [1, 2, 3].map((x) => ipcFromColumns(schema, { x: [x] })), + ); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const pendingClose = deferred(); + const close = sinon.stub(asyncStatement, 'close').resolves(); + close.onFirstCall().returns(pendingClose.promise); + const context = new ClientContextStub(); + const backend = new KernelOperationBackend({ asyncStatement, context }); + const operation = new DBSQLOperation({ backend, context }); + + expect(await operation.fetchChunk({ maxRows: 1 })).to.deep.equal([{ x: 1 }]); + const attempts = Promise.allSettled(concurrent ? [operation.close(), operation.close()] : [operation.close()]); + await nextTurn(); + pendingClose.reject(new Error('close failed')); + const outcomes = await attempts; + expect(outcomes.every((outcome) => outcome.status === 'rejected' && outcome.reason.message === 'close failed')).to + .be.true; + expect(close.calledOnce).to.be.true; + + expect(await operation.hasMoreRows()).to.be.true; + expect(await operation.fetchAll()).to.deep.equal([{ x: 2 }, { x: 3 }]); + expect(await operation.hasMoreRows()).to.be.false; + expect((await operation.close()).isSuccess).to.be.true; + expect(close.callCount).to.equal(2); + }); + } + + it('preserves unread rows when cancel is a no-op during a failed close', async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const resultHandle = new StatementStub( + schemaIpc, + [1, 2, 3].map((x) => ipcFromColumns(schema, { x: [x] })), + ); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const pendingClose = deferred(); + sinon.stub(asyncStatement, 'close').onFirstCall().returns(pendingClose.promise); + const backend = new KernelOperationBackend({ asyncStatement, context: new ClientContextStub() }); + + expect(await backend.fetchChunk({ limit: 1 })).to.deep.equal([{ x: 1 }]); + const closing = backend.close().then( + () => expect.fail('Expected close to fail'), + (reason: unknown) => reason, + ); + expect((await backend.cancel()).isSuccess).to.be.true; + pendingClose.reject(new Error('close failed')); + expect(((await closing) as Error).message).to.equal('close failed'); + expect(await backend.fetchChunk({ limit: 10 })).to.deep.equal([{ x: 2 }, { x: 3 }]); + expect(await backend.hasMore()).to.be.false; + }); + + it('surfaces a pending fetch failure after error-cleanup close fails', async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const pending = deferred<{ ipcBytes: Buffer } | null>(); + const resultHandle = new StatementStub(schemaIpc, []); + const fetchNextBatch = sinon.stub(resultHandle, 'fetchNextBatch').returns(pending.promise); + fetchNextBatch.onFirstCall().resolves({ ipcBytes: ipcFromColumns(schema, { x: [1] }) }); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const close = sinon.stub(asyncStatement, 'close').resolves(); + close.onFirstCall().rejects(new Error('cleanup close failed')); + const context = new ClientContextStub(); + const log = sinon.spy(context.logger, 'log'); + const backend = new KernelOperationBackend({ asyncStatement, context }); + await backend.getResultMetadata(); + const decode = sinon.stub(RecordBatchReader, 'from').throws(new Error('decode failed')); + try { + const error = await backend.fetchChunk({ limit: 1 }).then( + () => expect.fail('Expected a decode error'), + (reason: unknown) => reason, + ); + expect((error as Error).message).to.equal('decode failed'); + expect(close.calledOnce).to.be.true; + expect(log.calledWith(LogLevel.warn, sinon.match('cleanup close failed'))).to.be.true; + + pending.reject(new Error('prefetch failed')); + const nextError = await backend.hasMore().then( + () => expect.fail('Expected the pending fetch error instead of end of results'), + (reason: unknown) => reason, + ); + expect((nextError as Error).message).to.equal('prefetch failed'); + } finally { + decode.restore(); + pending.resolve(null); + await backend.close(); + } + }); + for (const asyncHandle of [false, true]) { const handleKind = asyncHandle ? 'AsyncStatement' : 'Statement'; for (const method of ['close', 'cancel'] as const) {