Stop teardown landing on a stream fetch that is still running

index.js

@@ -345,22 +345,34 @@ 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; }
 
   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 = () => {

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', () => {

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<T>(p: Promise<T>, 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)
+})