diff --git a/index.js b/index.js index ec00ed5..71eb794 100644 --- a/index.js +++ b/index.js @@ -345,6 +345,7 @@ class ChdbQueryStream { this._format = format; this._signal = signal; this._closed = false; + this._inflight = null; // the fetch currently on a libuv thread, if any } get closed() { return this._closed; } @@ -352,15 +353,26 @@ class ChdbQueryStream { async *[Symbol.asyncIterator]() { try { while (true) { + // cancel() may have landed while the previous fetch was in flight. Issuing + // another one would race the destroy it deferred onto that fetch. + if (this._closed) break; if (this._signal && this._signal.aborted) { throw new ChdbAbortError('Stream aborted'); } let raw; try { - raw = await chdbNode.StreamFetch(this._handle); + // Tracked like every other native op (see withAbortTimeout): a fetch + // runs on a libuv thread against this stream's connection, so a + // Session.close() or a teardown that does not wait for it destroys the + // connection mid-op and aborts the engine for the whole process. + // Streaming was the one path that skipped tracking. + this._inflight = trackNative(chdbNode.StreamFetch(this._handle)); + raw = await this._inflight; } catch (e) { this._closed = true; throw new ChdbStreamError(asQueryError(e).message, { cause: e }); + } finally { + this._inflight = null; } if (raw.done) { this._closed = true; break; } yield new StreamChunk(raw.bytes, raw.numRows, this._format); @@ -386,7 +398,20 @@ class ChdbQueryStream { cancel() { if (this._closed) return; this._closed = true; - try { chdbNode.StreamCancel(this._handle); } catch (_) { /* best effort */ } + const destroy = () => { + try { chdbNode.StreamCancel(this._handle); } catch (_) { /* best effort */ } + }; + // StreamCancel destroys the stream handle synchronously, and the fetch worker + // reads that handle without holding a lock. Destroying it while a fetch is + // executing pulls it out from under the worker. Wait for the fetch, and + // register the wait so a Session.close() defers behind the destroy too. + if (this._inflight) { + const deferred = this._inflight.then(destroy, destroy); + pendingNativeOps.add(deferred); + deferred.finally(() => pendingNativeOps.delete(deferred)); + } else { + destroy(); + } } } @@ -656,6 +681,17 @@ class Session { for (const sig of SESSION_SIGNALS) process.removeListener(sig, this.#signalHandler); this.#signalHandler = null; } + // Release a stream still open on this connection before the connection goes. + // Its handle points into the connection, so the stream's own finally would + // later call StreamCancel against a connection that no longer exists. Doing + // it here also orders the two teardowns correctly: a cancel that has to wait + // for an in-flight fetch registers that wait in pendingNativeOps, which the + // deferral below then waits for, so the stream handle is always destroyed + // before the connection it lives in. + if (this._activeStream && !this._activeStream.closed) { + try { this._activeStream.cancel(); } catch (_) { /* best effort */ } + } + this._activeStream = null; const conn = this.connection; this.connection = null; const teardown = () => { diff --git a/test/v3/async-stress.test.ts b/test/v3/async-stress.test.ts index c4fd8fa..bdd5af5 100644 --- a/test/v3/async-stress.test.ts +++ b/test/v3/async-stress.test.ts @@ -11,10 +11,12 @@ const HEAVY = (n: number) => `SELECT count() FROM numbers(${n})` // The storm/race cases below run dozens of heavy queries in sequence and take // ~15-30s on a fast runner. Give them a generous timeout so a slow/loaded CI -// runner never hits the 30s default and kills a test mid-flight — an abandoned -// in-flight native op (plain queryAsync is not tracked for drain) would collide -// with the next test on the single in-process engine and abort it (code 236), -// cascading failures into every later file. +// runner never hits the 30s default and kills a test mid-flight — a test killed +// at its timeout is torn down without the global afterEach draining its in-flight +// native op, which then collides with the next test on the single in-process +// engine and aborts it (code 236), cascading failures into every later file. +// (Plain queryAsync is tracked for drain; withAbortTimeout tracks every async +// query. This comment used to claim it was not.) const STRESS_TIMEOUT_MS = 120_000 describe('async concurrency — correctness & no deadlock', () => { diff --git a/test/v3/stream-cancel-race.test.ts b/test/v3/stream-cancel-race.test.ts new file mode 100644 index 0000000..ce1d9e9 --- /dev/null +++ b/test/v3/stream-cancel-race.test.ts @@ -0,0 +1,89 @@ +import { describe, it, expect } from 'vitest' +import { Session, queryAsync } from '../../index.js' +// @ts-expect-error — internal test helper, not in the type surface +import { _drainPendingOps } from '../../index.js' + +// Teardown landing on a stream fetch that is still running. +// +// A fetch runs on a libuv thread against the session's connection. StreamCancel +// destroys the stream handle synchronously and CloseConnection destroys the +// connection, and the fetch worker reads both without holding a lock, so either +// one landing mid-fetch pulls the memory out from under the worker. libchdb +// answers that by aborting the shared in-process engine, which fails every later +// query in the process, and on macOS can leave the worker blocked inside +// chdb_stream_fetch_result so its promise never settles. +// +// All of it was reachable because streaming was the one native async path that +// never registered itself in pendingNativeOps, so neither close() nor the +// suite-wide teardown drain knew a fetch was in flight. +// +// One row that needs a real per-row computation, so the first fetch takes about a +// second and teardown lands well inside it. count() over numbers() returns in +// ~15ms — fast enough that these would pass whether the races are handled or not. +const SLOW_ONE_ROW = 'SELECT max(sipHash64(number)) FROM numbers(200000000)' + +// The macOS symptom was a fetch whose promise never settled. Bound the wait so +// that failure reads as this assertion rather than as a suite timeout. +async function settlesWithin(p: Promise, ms: number): Promise<'settled' | 'hung'> { + let timer: NodeJS.Timeout + const hung = new Promise<'hung'>((r) => { timer = setTimeout(() => r('hung'), ms) }) + const settled = p.then(() => 'settled' as const, () => 'settled' as const) + try { + return await Promise.race([settled, hung]) + } finally { + clearTimeout(timer!) + } +} + +describe('teardown that lands on an in-flight stream fetch', () => { + it('cancel() waits for the fetch instead of destroying the handle under it', async () => { + const s = new Session() + try { + const stream = s.queryStream(SLOW_ONE_ROW, { format: 'JSONEachRow' }) + const it = stream[Symbol.asyncIterator]() + + // next() runs the generator body up to the await, so the fetch is dispatched + // before it returns and is still running on the next line. Asserted rather + // than assumed: if a change ever makes the fetch settle first, this test + // would quietly stop covering the race it exists for. + const first = it.next() + expect((stream as unknown as { _inflight: unknown })._inflight).toBeTruthy() + + stream.cancel() + expect(await settlesWithin(first, 20_000)).toBe('settled') + await _drainPendingOps() // let the deferred destroy land + + // The engine survived, and the cancelled stream was released cleanly enough + // that the session takes another one. + const rows: number[] = [] + for await (const row of s + .queryStream('SELECT number AS n FROM numbers(3)') + .rows<{ n: number }>()) { + rows.push(row.n) + } + expect(rows).toEqual([0, 1, 2]) + } finally { + s.close() + } + }, 60_000) + + it('close() releases the stream and defers the connection until the fetch drains', async () => { + const s = new Session() + const stream = s.queryStream(SLOW_ONE_ROW, { format: 'JSONEachRow' }) + const it = stream[Symbol.asyncIterator]() + + const first = it.next() + expect((stream as unknown as { _inflight: unknown })._inflight).toBeTruthy() + + s.close() // must cancel the stream and hold the connection until the fetch ends + expect(await settlesWithin(first, 20_000)).toBe('settled') + await _drainPendingOps() + + // Standalone query() refuses to run while any session connection is still + // registered ("a session (path=…) is active"), so this answering proves both + // that the engine is alive and that the deferred teardown released the + // connection instead of leaking it into the rest of the process. + const r = await queryAsync('SELECT 1', { format: 'CSV' }) + expect(r.text().trim()).toBe('1') + }, 60_000) +})