From 79335aa54c2b5192cc83b1af6948f385f9557dbd Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Fri, 17 Jul 2026 16:31:45 +0000 Subject: [PATCH 01/14] refactor: add `MessageIO.writeMessage`/`readMessage` and use them for the `PRELOGIN` exchange This adds two static methods to `MessageIO` that read and write TDS messages directly on a stream, without going through the `Message`, `IncomingMessageStream` and `OutgoingMessageStream` machinery: * `MessageIO.writeMessage` writes a payload (a sync or async iterable of `Buffer`s) to a stream, wrapped in TDS packets, respecting backpressure. If iterating the payload throws, the message is terminated with the `IGNORE` flag set and the error is re-thrown. * `MessageIO.readMessage` is an async generator that yields the packet contents of the next message on the stream. Bytes that arrive after the message's last packet are pushed back onto the stream for the next read. The `PRELOGIN` exchange during connection establishment is the first consumer: it now runs directly on the raw socket, and the `MessageIO` instance is only constructed once the prelogin response was fully read. Setting up per-message stream objects is a significant part of the message handling cost, especially before v8 optimizations kick in. Comparing the new functions against the existing stream classes (messages/second, Node.js v24.18.0, median of 3 runs): | n | IncomingMessageStream | readMessage | | |---|---|---|---| | 100 | 9,889 | 22,441 | +127% | | 1,000 | 26,244 | 58,069 | +121% | | 10,000 | 76,942 | 150,728 | +96% | | 100,000 | 173,808 | 293,468 | +69% | | n | OutgoingMessageStream | writeMessage | | |---|---|---|---| | 100 | 9,353 | 16,518 | +77% | | 1,000 | 29,759 | 48,655 | +64% | | 10,000 | 91,387 | 127,079 | +39% | | 100,000 | 156,812 | 191,138 | +22% | Co-Authored-By: Claude Fable 5 --- .../message-io/incoming-message-stream.js | 41 ++ .../message-io/outgoing-message-stream.js | 70 +++ benchmarks/message-io/read-message.js | 38 ++ benchmarks/message-io/write-message.js | 41 ++ src/connection.ts | 43 +- src/message-io.ts | 256 +++++++++- test/unit/message-io-test.ts | 451 +++++++++++++++++- 7 files changed, 919 insertions(+), 21 deletions(-) create mode 100644 benchmarks/message-io/incoming-message-stream.js create mode 100644 benchmarks/message-io/outgoing-message-stream.js create mode 100644 benchmarks/message-io/read-message.js create mode 100644 benchmarks/message-io/write-message.js diff --git a/benchmarks/message-io/incoming-message-stream.js b/benchmarks/message-io/incoming-message-stream.js new file mode 100644 index 000000000..2d3aba868 --- /dev/null +++ b/benchmarks/message-io/incoming-message-stream.js @@ -0,0 +1,41 @@ +const { createBenchmark } = require('../common'); +const { Readable } = require('stream'); + +const Debug = require('tedious/lib/debug'); +const IncomingMessageStream = require('tedious/lib/incoming-message-stream'); +const { Packet } = require('tedious/lib/packet'); + +const bench = createBenchmark(main, { + n: [100, 1000, 10000, 100000] +}); + +function main({ n }) { + const debug = new Debug(); + + const stream = Readable.from((async function*() { + for (let i = 0; i < n; i++) { + const packet = new Packet(2); + packet.last(true); + packet.addData(Buffer.from([1, 2, 3, 4, 5, 6, 7, 8, 9])); + + yield packet.buffer; + } + })()); + + const incoming = new IncomingMessageStream(debug); + stream.pipe(incoming); + + bench.start(); + + (async function() { + let total = 0; + + for await (const message of incoming) { + for await (const chunk of message) { + total += chunk.length; + } + } + + bench.end(n); + })(); +} diff --git a/benchmarks/message-io/outgoing-message-stream.js b/benchmarks/message-io/outgoing-message-stream.js new file mode 100644 index 000000000..c8c5aecbf --- /dev/null +++ b/benchmarks/message-io/outgoing-message-stream.js @@ -0,0 +1,70 @@ +const { createBenchmark } = require('../common'); +const { Duplex } = require('stream'); + +const Debug = require('tedious/lib/debug'); +const OutgoingMessageStream = require('tedious/lib/outgoing-message-stream'); +const Message = require('tedious/lib/message'); + +const bench = createBenchmark(main, { + n: [100, 1000, 10000, 100000] +}); + +function main({ n }) { + const debug = new Debug(); + + const stream = new Duplex({ + read() {}, + write(chunk, encoding, callback) { + // Just consume the data + callback(); + } + }); + + const payload = [ + Buffer.alloc(1024), + Buffer.alloc(1024), + Buffer.alloc(1024), + Buffer.alloc(256), + Buffer.alloc(256), + Buffer.alloc(256), + Buffer.alloc(256), + ]; + + const out = new OutgoingMessageStream(debug, { + packetSize: 8 + 1024 + }); + out.pipe(stream); + + bench.start(); + + function writeNextMessage(i) { + if (i === n) { + out.end(); + out.once('finish', () => { + bench.end(n); + }); + return; + } + + const message = new Message({ type: 2, resetConnection: false }); + out.write(message); + + for (const chunk of payload) { + message.write(chunk); + } + + message.end(); + + if (out.needsDrain) { + out.once('drain', () => { + writeNextMessage(i + 1); + }); + } else { + process.nextTick(() => { + writeNextMessage(i + 1); + }); + } + } + + writeNextMessage(0); +} diff --git a/benchmarks/message-io/read-message.js b/benchmarks/message-io/read-message.js new file mode 100644 index 000000000..c3594ed76 --- /dev/null +++ b/benchmarks/message-io/read-message.js @@ -0,0 +1,38 @@ +const { createBenchmark } = require('../common'); +const { Readable } = require('stream'); + +const Debug = require('tedious/lib/debug'); +const MessageIO = require('tedious/lib/message-io'); +const { Packet } = require('tedious/lib/packet'); + +const bench = createBenchmark(main, { + n: [100, 1000, 10000, 100000] +}); + +function main({ n }) { + const debug = new Debug(); + + const stream = Readable.from((async function*() { + for (let i = 0; i < n; i++) { + const packet = new Packet(2); + packet.last(true); + packet.addData(Buffer.from([1, 2, 3, 4, 5, 6, 7, 8, 9])); + + yield packet.buffer; + } + })()); + + (async function() { + bench.start(); + + let total = 0; + + for (let i = 0; i < n; i++) { + for await (const chunk of MessageIO.readMessage(stream, debug)) { + total += chunk.length; + } + } + + bench.end(n); + })(); +} diff --git a/benchmarks/message-io/write-message.js b/benchmarks/message-io/write-message.js new file mode 100644 index 000000000..545e78ddf --- /dev/null +++ b/benchmarks/message-io/write-message.js @@ -0,0 +1,41 @@ +const { createBenchmark } = require('../common'); +const { Duplex } = require('stream'); + +const Debug = require('tedious/lib/debug'); +const MessageIO = require('tedious/lib/message-io'); + +const bench = createBenchmark(main, { + n: [100, 1000, 10000, 100000] +}); + +function main({ n }) { + const debug = new Debug(); + + const stream = new Duplex({ + read() {}, + write(chunk, encoding, callback) { + // Just consume the data + callback(); + } + }); + + const payload = [ + Buffer.alloc(1024), + Buffer.alloc(1024), + Buffer.alloc(1024), + Buffer.alloc(256), + Buffer.alloc(256), + Buffer.alloc(256), + Buffer.alloc(256), + ]; + + (async function() { + bench.start(); + + for (let i = 0; i < n; i++) { + await MessageIO.writeMessage(stream, debug, 8 + 1024, 2, payload); + } + + bench.end(n); + })(); +} diff --git a/src/connection.ts b/src/connection.ts index 0ff2f90e3..baec7435d 100644 --- a/src/connection.ts +++ b/src/connection.ts @@ -2092,17 +2092,32 @@ class Connection extends EventEmitter { socket.setKeepAlive(true, KEEP_ALIVE_INITIAL_DELAY); - this.messageIo = new MessageIO(socket, this.config.options.packetSize, this.debug); - this.messageIo.on('secure', (cleartext) => { this.emit('secure', cleartext); }); - this.socket = socket; this.debug.log('connected to ' + this.config.server + ':' + this.config.options.port); - this.sendPreLogin(); + try { + await this.sendPreLogin(socket); + } catch (err: any) { + signal.throwIfAborted(); + + throw this.wrapSocketError(err); + } this.transitionTo(this.STATE.SENT_PRELOGIN); - const preloginResponse = await this.readPreloginResponse(signal); + + let preloginResponse; + try { + preloginResponse = await this.readPreloginResponse(socket, signal); + } catch (err: any) { + signal.throwIfAborted(); + + throw this.wrapSocketError(err); + } + + this.messageIo = new MessageIO(socket, this.config.options.packetSize, this.debug); + this.messageIo.on('secure', (cleartext) => { this.emit('secure', cleartext); }); + await this.performTlsNegotiation(preloginResponse, signal); this.sendLogin7Packet(); @@ -2456,7 +2471,7 @@ class Connection extends EventEmitter { /** * @private */ - sendPreLogin() { + async sendPreLogin(socket: net.Socket) { const [, major, minor, build] = /^(\d+)\.(\d+)\.(\d+)/.exec(version) ?? ['0.0.0', '0', '0', '0']; const payload = new PreloginPayload({ // If encrypt setting is set to 'strict', then we should have already done the encryption before calling @@ -2466,7 +2481,7 @@ class Connection extends EventEmitter { version: { major: Number(major), minor: Number(minor), build: Number(build), subbuild: 0 } }); - this.messageIo.sendMessage(TYPE.PRELOGIN, payload.data); + await MessageIO.writeMessage(socket, this.debug, this.config.options.packetSize, TYPE.PRELOGIN, [payload.data]); this.debug.payload(function() { return payload.toString(' '); }); @@ -3365,18 +3380,12 @@ class Connection extends EventEmitter { }); } - async readPreloginResponse(signal: AbortSignal): Promise { + async readPreloginResponse(socket: net.Socket, signal: AbortSignal): Promise { let messageBuffer = Buffer.alloc(0); await withAbortRace(signal, async (signalAborted) => { - const message = await Promise.race([ - this.messageIo.readMessage().catch((err) => { - throw this.wrapSocketError(err); - }), - signalAborted - ]); + const iterator = MessageIO.readMessage(socket, this.debug); - const iterator = message[Symbol.asyncIterator](); try { while (true) { const { done, value } = await Promise.race([ @@ -3391,9 +3400,7 @@ class Connection extends EventEmitter { messageBuffer = Buffer.concat([messageBuffer, value]); } } finally { - if (iterator.return) { - await iterator.return(); - } + await iterator.return(); } }); diff --git a/src/message-io.ts b/src/message-io.ts index d068d73a0..e22f54513 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -1,6 +1,7 @@ import DuplexPair from 'native-duplexpair'; -import { Duplex } from 'stream'; +import BufferList from 'bl'; +import { Duplex, type Readable, type Writable } from 'stream'; import * as tls from 'tls'; import { isIP, Socket } from 'net'; import { EventEmitter } from 'events'; @@ -8,7 +9,8 @@ import { EventEmitter } from 'events'; import Debug from './debug'; import Message from './message'; -import { TYPE } from './packet'; +import { HEADER_LENGTH, Packet, TYPE } from './packet'; +import { ConnectionError } from './errors'; import IncomingMessageStream from './incoming-message-stream'; import OutgoingMessageStream from './outgoing-message-stream'; @@ -188,6 +190,256 @@ class MessageIO extends EventEmitter { return result.value; } + + /** + * Write a message with the given `type` and `payload` to the given `stream`, + * wrapping it into TDS packets of the given `packetSize`. + * + * Respects backpressure from the stream, waiting for it to drain before + * writing more data. + * + * If iterating the `payload` throws, the message is terminated with a final + * packet that has the `IGNORE` flag set (telling the server to disregard the + * message) and the error is re-thrown. Errors from the stream itself are + * thrown as-is. + * + * @param stream The stream to write the message to. + * @param debug The debug instance to use for logging. + * @param packetSize The maximum packet size to use. + * @param type The type of the message to write. + * @param payload The payload to write. + * @param resetConnection Whether the server should reset the connection when processing the message. + */ + static async writeMessage(stream: Writable, debug: Debug, packetSize: number, type: number, payload: AsyncIterable | Iterable, resetConnection = false): Promise { + if (!stream.writable) { + throw new Error('Premature close'); + } + + let drain: PromiseWithResolvers | null = null; + + const onDrain = () => { + if (drain) { + const { resolve } = drain; + drain = null; + resolve(); + } + }; + + const onError = (err: Error) => { + if (drain) { + const { reject } = drain; + drain = null; + reject(err); + } + }; + + const waitForDrain = () => { + drain = Promise.withResolvers(); + return drain.promise; + }; + + stream.on('drain', onDrain); + stream.on('close', onDrain); + stream.on('error', onError); + + try { + const bl = new BufferList(); + const length = packetSize - HEADER_LENGTH; + let packetNumber = 0; + + const writePacket = async (packet: Packet) => { + debug.packet('Sent', packet); + debug.data(packet); + + if (stream.write(packet.buffer) === false) { + await waitForDrain(); + } + }; + + let iterator; + if ((payload as AsyncIterable)[Symbol.asyncIterator]) { + iterator = (payload as AsyncIterable)[Symbol.asyncIterator](); + } else { + iterator = (payload as Iterable)[Symbol.iterator](); + } + + while (true) { + let value, done; + try { + ({ value, done } = await iterator.next()); + } catch (err) { + // The payload errored while being iterated. If the stream is still + // writable, terminate the message with the `IGNORE` flag set so the + // server disregards everything sent so far. + if (stream.writable) { + const packet = new Packet(type); + packet.packetId(packetNumber += 1); + packet.resetConnection(resetConnection); + packet.last(true); + packet.ignore(true); + + await writePacket(packet); + } + + throw err; + } + + if (done) { + break; + } + + bl.append(value); + + while (bl.length > length) { + const data = bl.slice(0, length); + bl.consume(length); + + const packet = new Packet(type); + packet.packetId(packetNumber += 1); + packet.resetConnection(resetConnection); + packet.addData(data); + + await writePacket(packet); + } + } + + const data = bl.slice(); + bl.consume(data.length); + + const packet = new Packet(type); + packet.packetId(packetNumber += 1); + packet.resetConnection(resetConnection); + packet.last(true); + packet.addData(data); + + await writePacket(packet); + } finally { + stream.removeListener('drain', onDrain); + stream.removeListener('close', onDrain); + stream.removeListener('error', onError); + } + } + + /** + * Read the next TDS message from the given `stream`. + * + * Returns an async generator that yields the data of the message's packets + * as they arrive. The generator throws if the stream emits an error or is + * closed before the message was fully read. + * + * Any bytes following the message's last packet (e.g. the start of the next + * message) are pushed back onto the stream, to be consumed by the next read. + * + * @param stream The stream to read the message from. + * @param debug The debug instance to use for logging. + */ + static async *readMessage(stream: Readable, debug: Debug): AsyncGenerator { + if (!stream.readable) { + throw new Error('Premature close'); + } + + const bl = new BufferList(); + + let error: Error | null = null; + let closed = false; + let waiting: PromiseWithResolvers | null = null; + + const onReadable = () => { + if (waiting) { + const { resolve } = waiting; + waiting = null; + resolve(); + } + }; + + const onError = (err: Error) => { + error = err; + + if (waiting) { + const { reject } = waiting; + waiting = null; + reject(err); + } + }; + + const onClose = () => { + closed = true; + + if (waiting) { + const { reject } = waiting; + waiting = null; + reject(new Error('Premature close')); + } + }; + + stream.on('readable', onReadable); + stream.on('error', onError); + stream.on('close', onClose); + + try { + while (true) { + if (error) { + throw error; + } + + if (closed) { + throw new Error('Premature close'); + } + + let chunk: Buffer; + while ((chunk = stream.read()) !== null) { + bl.append(chunk); + + // The packet header is always 8 bytes of length. + while (bl.length >= HEADER_LENGTH) { + // Get the full packet length + const length = bl.readUInt16BE(2); + if (length < HEADER_LENGTH) { + throw new ConnectionError('Unable to process incoming packet'); + } + + if (bl.length < length) { + break; + } + + const data = bl.slice(0, length); + bl.consume(length); + + const packet = new Packet(data); + debug.packet('Received', packet); + debug.data(packet); + + yield packet.data(); + + // Did the stream error while we yielded? + if (error) { + throw error; + } + + if (packet.isLast()) { + // This was the last packet of the message. Any data left in the + // buffer belongs to the next message (e.g. the response to an + // `ATTENTION` message sent by the client while reading an + // incoming response), so push it back onto the stream. + if (bl.length) { + stream.unshift(bl.slice()); + } + + return; + } + } + } + + // Wait for the stream to become readable again (or error out or close). + waiting = Promise.withResolvers(); + await waiting.promise; + } + } finally { + stream.removeListener('readable', onReadable); + stream.removeListener('error', onError); + stream.removeListener('close', onClose); + } + } } export default MessageIO; diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index ab1376dfe..70be87574 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -5,18 +5,467 @@ import { promisify } from 'util'; import DuplexPair from 'native-duplexpair'; import { checkServerIdentity, type PeerCertificate, TLSSocket } from 'tls'; import { readFileSync } from 'fs'; -import { Duplex } from 'stream'; +import { Duplex, Readable } from 'stream'; +import BufferListStream from 'bl'; import Debug from '../../src/debug'; import MessageIO from '../../src/message-io'; import Message from '../../src/message'; import { Packet, TYPE } from '../../src/packet'; +import { ConnectionError } from '../../src/errors'; const packetType = 2; const packetSize = 8 + 4; const delay = promisify(setTimeout); +function assertNoDanglingEventListeners(stream: Duplex) { + assert.strictEqual(stream.listenerCount('error'), 0); + assert.strictEqual(stream.listenerCount('drain'), 0); + assert.strictEqual(stream.listenerCount('readable'), 0); +} + +function splitPackets(data: Buffer): Packet[] { + const packets = []; + + let offset = 0; + while (offset < data.length) { + const length = data.readUInt16BE(offset + 2); + packets.push(new Packet(data.subarray(offset, offset + length))); + offset += length; + } + + return packets; +} + +describe('MessageIO.writeMessage', function() { + let debug: Debug; + + beforeEach(function() { + debug = new Debug(); + }); + + it('wraps the given payload into a TDS packet and writes it to the given stream', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new BufferListStream(); + + await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload]); + + const buf = stream.read(); + assert.instanceOf(buf, Buffer); + + const packet = new Packet(buf); + assert.strictEqual(packet.type(), packetType); + assert.strictEqual(packet.length(), payload.length + 8); + assert.isTrue(packet.isLast()); + assert.deepEqual(packet.data(), payload); + + assert.isNull(stream.read()); + assertNoDanglingEventListeners(stream); + }); + + it('splits payloads that are larger than the packet size across multiple packets', async function() { + const payload = Buffer.from([1, 2, 3, 4, 5, 6]); + const stream = new BufferListStream(); + + // `packetSize` allows for 4 bytes of data per packet + await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload]); + + const [firstPacket, secondPacket, ...rest] = splitPackets(stream.read()); + assert.lengthOf(rest, 0); + + assert.strictEqual(firstPacket.packetId(), 1); + assert.isFalse(firstPacket.isLast()); + assert.deepEqual(firstPacket.data(), payload.subarray(0, 4)); + + assert.strictEqual(secondPacket.packetId(), 2); + assert.isTrue(secondPacket.isLast()); + assert.deepEqual(secondPacket.data(), payload.subarray(4)); + }); + + it('writes an empty final packet for an empty payload', async function() { + const stream = new BufferListStream(); + + await MessageIO.writeMessage(stream, debug, packetSize, packetType, []); + + const packet = new Packet(stream.read()); + assert.isTrue(packet.isLast()); + assert.deepEqual(packet.data(), Buffer.alloc(0)); + + assert.isNull(stream.read()); + }); + + it('terminates the message with the ignore flag set when iterating the payload errors', async function() { + const payload = Buffer.from([1, 2, 3, 4, 5, 6]); + const stream = new BufferListStream(); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, (async function*() { + yield payload; + throw new Error('iteration error'); + })()); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'iteration error'); + } + + assert(hadError); + + const [firstPacket, lastPacket, ...rest] = splitPackets(stream.read()); + assert.lengthOf(rest, 0); + + // The part of the payload that filled a whole packet was sent before the + // error occurred, the rest is discarded. + assert.isFalse(firstPacket.isLast()); + assert.deepEqual(firstPacket.data(), payload.subarray(0, 4)); + + assert.isTrue(lastPacket.isLast()); + assert.include(lastPacket.statusAsString(), 'IGNORE'); + assert.deepEqual(lastPacket.data(), Buffer.alloc(0)); + + assertNoDanglingEventListeners(stream); + }); + + it('handles errors from the payload while the stream is waiting for drain', async function() { + const payload = Buffer.from([1, 2, 3, 4]); + + const callbacks: Array<() => void> = []; + const stream = new Duplex({ + write(chunk, encoding, callback) { + // Collect all callbacks so that we can simulate draining the stream later + callbacks.push(callback); + }, + read() {}, + + // instantly return false on write requests to indicate that the stream needs to drain + highWaterMark: 1 + }); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, (async function*() { + yield payload; + + // Simulate draining the stream after the exception was thrown + setTimeout(() => { + let cb; + while (cb = callbacks.shift()) { + cb(); + } + }, 20); + + throw new Error('iteration error'); + })()); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'iteration error'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); + + it('handles errors on the stream during writing', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new Duplex({ + write(chunk, encoding, callback) { + callback(new Error('write error')); + }, + read() {} + }); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload]); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'write error'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); + + it('handles errors on the stream while waiting for the stream to drain', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new Duplex({ + write(chunk, encoding, callback) { + // never call callback so that the stream never drains + }, + read() {}, + + // instantly return false on write requests to indicate that the stream needs to drain + highWaterMark: 1 + }); + + setTimeout(() => { + assert(stream.writableNeedDrain); + stream.destroy(new Error('write error')); + }, 20); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload, payload, payload]); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'write error'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); + + it('handles errors on the stream while waiting for more payload data', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new Duplex({ + write(chunk, encoding, callback) { + // never call callback so that the stream never drains + }, + read() {}, + + // instantly return false on write requests to indicate that the stream needs to drain + highWaterMark: 1 + }); + + setTimeout(() => { + assert(stream.writableNeedDrain); + stream.destroy(new Error('write error')); + }, 20); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, (async function*() { + yield payload; + yield payload; + yield payload; + })()); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'write error'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); + + it('throws when the given stream is not writable', async function() { + const stream = new BufferListStream(); + stream.end(); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, []); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'Premature close'); + } + + assert(hadError); + }); +}); + +describe('MessageIO.readMessage', function() { + let debug: Debug; + + beforeEach(function() { + debug = new Debug(); + }); + + it('reads a message consisting of a single TDS packet from the given stream', async function() { + const payload = Buffer.from([1, 2, 3]); + const packet = new Packet(packetType); + packet.last(true); + packet.addData(payload); + + const stream = new BufferListStream(); + stream.write(packet.buffer); + + const chunks = []; + for await (const chunk of MessageIO.readMessage(stream, debug)) { + chunks.push(chunk); + } + + assert.deepEqual(chunks, [payload]); + assertNoDanglingEventListeners(stream); + }); + + it('reads a message that spans multiple TDS packets', async function() { + const payload = Buffer.from([1, 2, 3]); + + const firstPacket = new Packet(packetType); + firstPacket.addData(payload.subarray(0, 2)); + + const lastPacket = new Packet(packetType); + lastPacket.last(true); + lastPacket.addData(payload.subarray(2)); + + const stream = new BufferListStream(); + stream.write(firstPacket.buffer); + stream.write(lastPacket.buffer); + + const chunks = []; + for await (const chunk of MessageIO.readMessage(stream, debug)) { + chunks.push(chunk); + } + + assert.deepEqual(Buffer.concat(chunks), payload); + }); + + it('reads packets that arrive in chunks that do not align with packet boundaries', async function() { + const payload = Buffer.from([1, 2, 3]); + + const firstPacket = new Packet(packetType); + firstPacket.addData(payload.subarray(0, 2)); + + const lastPacket = new Packet(packetType); + lastPacket.last(true); + lastPacket.addData(payload.subarray(2)); + + const data = Buffer.concat([firstPacket.buffer, lastPacket.buffer]); + + const stream = new Readable({ read() {} }); + + // Deliver the data byte by byte. + for (let i = 0; i < data.length; i++) { + setTimeout(() => { + stream.push(data.subarray(i, i + 1)); + }, i); + } + + const chunks = []; + for await (const chunk of MessageIO.readMessage(stream, debug)) { + chunks.push(chunk); + } + + assert.deepEqual(Buffer.concat(chunks), payload); + }); + + it('pushes bytes that follow the last packet of a message back onto the stream', async function() { + const payload = Buffer.from([1, 2, 3]); + const packet = new Packet(packetType); + packet.last(true); + packet.addData(payload); + + const trailingBytes = Buffer.from([9, 9, 9, 9]); + + const stream = new BufferListStream(); + stream.write(Buffer.concat([packet.buffer, trailingBytes])); + + const chunks = []; + for await (const chunk of MessageIO.readMessage(stream, debug)) { + chunks.push(chunk); + } + + assert.deepEqual(chunks, [payload]); + assert.deepEqual(stream.read(), trailingBytes); + }); + + it('handles errors while reading from the stream', async function() { + const stream = Readable.from((async function*(): AsyncGenerator { + throw new Error('read error'); + })()); + + let hadError = false; + try { + const message = MessageIO.readMessage(stream, debug); + while (!(await message.next()).done) { + // Discard the message contents. + } + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'read error'); + } + + assert(hadError); + }); + + it('throws when the stream is closed before the message was fully read', async function() { + const packet = new Packet(packetType); + packet.addData(Buffer.from([1, 2, 3])); + + const stream = new Readable({ read() {} }); + stream.push(packet.buffer); + + setTimeout(() => { + stream.destroy(); + }, 20); + + let hadError = false; + try { + const message = MessageIO.readMessage(stream, debug); + while (!(await message.next()).done) { + // Discard the message contents. + } + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'Premature close'); + } + + assert(hadError); + }); + + it('throws when the given stream is not readable', async function() { + const stream = new Readable({ read() {} }); + stream.destroy(); + + let hadError = false; + try { + const message = MessageIO.readMessage(stream, debug); + while (!(await message.next()).done) { + // Discard the message contents. + } + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'Premature close'); + } + + assert(hadError); + }); + + it('throws when receiving a packet with an invalid length', async function() { + const invalidPacket = Buffer.alloc(8); + invalidPacket.writeUInt16BE(4, 2); + + const stream = new BufferListStream(); + stream.write(invalidPacket); + + let hadError = false; + try { + const message = MessageIO.readMessage(stream, debug); + while (!(await message.next()).done) { + // Discard the message contents. + } + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, ConnectionError); + assert.strictEqual(err.message, 'Unable to process incoming packet'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); +}); + describe('MessageIO', function() { let server: Server; let serverConnection: Socket; From f3112fbe737f224838c14df3d0eb667b9b73b29e Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Fri, 17 Jul 2026 16:58:05 +0000 Subject: [PATCH 02/14] refactor: support `AbortSignal` in `MessageIO.writeMessage`/`readMessage` Both methods now accept an optional abort signal via an options object. For `readMessage` this is a correctness requirement, not a convenience: calling `.return()` on the generator while it is suspended waiting for stream data is queued behind the pending `.next()` call, so a caller racing reads against an abort promise and cleaning up via `.return()` deadlocks when the abort fires while the stream is open but quiet - exactly the connect-timeout scenario. Aborting the signal settles the pending read from the inside instead. For `writeMessage`, an abort interrupts waiting for the stream to drain or for more payload data. This can leave a partially written message on the stream, so it is only meant for teardown paths; cooperative cancellation that keeps the TDS stream aligned still goes through the payload-error path that terminates the message with the `IGNORE` flag. The `PRELOGIN` exchange now passes the connect signal down, and `readPreloginResponse` collapses to a plain `for await` loop. Co-Authored-By: Claude Fable 5 --- src/connection.ts | 29 ++--- src/message-io.ts | 75 +++++++++++-- test/unit/message-io-test.ts | 202 ++++++++++++++++++++++++++++++++++- 3 files changed, 275 insertions(+), 31 deletions(-) diff --git a/src/connection.ts b/src/connection.ts index baec7435d..737048070 100644 --- a/src/connection.ts +++ b/src/connection.ts @@ -2097,7 +2097,7 @@ class Connection extends EventEmitter { this.debug.log('connected to ' + this.config.server + ':' + this.config.options.port); try { - await this.sendPreLogin(socket); + await this.sendPreLogin(socket, signal); } catch (err: any) { signal.throwIfAborted(); @@ -2471,7 +2471,7 @@ class Connection extends EventEmitter { /** * @private */ - async sendPreLogin(socket: net.Socket) { + async sendPreLogin(socket: net.Socket, signal: AbortSignal) { const [, major, minor, build] = /^(\d+)\.(\d+)\.(\d+)/.exec(version) ?? ['0.0.0', '0', '0', '0']; const payload = new PreloginPayload({ // If encrypt setting is set to 'strict', then we should have already done the encryption before calling @@ -2481,7 +2481,7 @@ class Connection extends EventEmitter { version: { major: Number(major), minor: Number(minor), build: Number(build), subbuild: 0 } }); - await MessageIO.writeMessage(socket, this.debug, this.config.options.packetSize, TYPE.PRELOGIN, [payload.data]); + await MessageIO.writeMessage(socket, this.debug, this.config.options.packetSize, TYPE.PRELOGIN, [payload.data], { signal }); this.debug.payload(function() { return payload.toString(' '); }); @@ -3383,26 +3383,9 @@ class Connection extends EventEmitter { async readPreloginResponse(socket: net.Socket, signal: AbortSignal): Promise { let messageBuffer = Buffer.alloc(0); - await withAbortRace(signal, async (signalAborted) => { - const iterator = MessageIO.readMessage(socket, this.debug); - - try { - while (true) { - const { done, value } = await Promise.race([ - iterator.next(), - signalAborted - ]); - - if (done) { - break; - } - - messageBuffer = Buffer.concat([messageBuffer, value]); - } - } finally { - await iterator.return(); - } - }); + for await (const data of MessageIO.readMessage(socket, this.debug, { signal })) { + messageBuffer = Buffer.concat([messageBuffer, data]); + } const preloginPayload = new PreloginPayload(messageBuffer); this.debug.payload(function() { diff --git a/src/message-io.ts b/src/message-io.ts index e22f54513..37bdb94d6 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -203,14 +203,25 @@ class MessageIO extends EventEmitter { * message) and the error is re-thrown. Errors from the stream itself are * thrown as-is. * + * If the given `signal` is aborted, writing stops and the signal's abort + * reason is thrown. Note that this can leave a partially written message on + * the stream, so it must only be used when the connection is being torn down + * anyway. To cancel a message in a way that keeps the TDS stream aligned, + * throw from the `payload` iterable instead. + * * @param stream The stream to write the message to. * @param debug The debug instance to use for logging. * @param packetSize The maximum packet size to use. * @param type The type of the message to write. * @param payload The payload to write. - * @param resetConnection Whether the server should reset the connection when processing the message. + * @param options.resetConnection Whether the server should reset the connection when processing the message. + * @param options.signal An abort signal to stop writing the message. */ - static async writeMessage(stream: Writable, debug: Debug, packetSize: number, type: number, payload: AsyncIterable | Iterable, resetConnection = false): Promise { + static async writeMessage(stream: Writable, debug: Debug, packetSize: number, type: number, payload: AsyncIterable | Iterable, options: { resetConnection?: boolean, signal?: AbortSignal } = {}): Promise { + const { resetConnection = false, signal } = options; + + signal?.throwIfAborted(); + if (!stream.writable) { throw new Error('Premature close'); } @@ -233,9 +244,24 @@ class MessageIO extends EventEmitter { } }; + let abortPromise: Promise | null = null; + let onAbort: (() => void) | null = null; + + if (signal) { + const { promise, reject } = Promise.withResolvers(); + + // Prevent unhandled rejections if the signal is aborted while + // nothing is currently racing against `abortPromise`. + promise.catch(() => {}); + + abortPromise = promise; + onAbort = () => { reject(signal.reason); }; + signal.addEventListener('abort', onAbort, { once: true }); + } + const waitForDrain = () => { drain = Promise.withResolvers(); - return drain.promise; + return abortPromise ? Promise.race([drain.promise, abortPromise]) : drain.promise; }; stream.on('drain', onDrain); @@ -266,12 +292,15 @@ class MessageIO extends EventEmitter { while (true) { let value, done; try { - ({ value, done } = await iterator.next()); + const result = iterator.next(); + ({ value, done } = abortPromise ? await Promise.race([Promise.resolve(result), abortPromise]) : await result); } catch (err) { // The payload errored while being iterated. If the stream is still // writable, terminate the message with the `IGNORE` flag set so the - // server disregards everything sent so far. - if (stream.writable) { + // server disregards everything sent so far. If the signal was + // aborted instead, the connection is being torn down and the + // message is left unterminated. + if (stream.writable && !signal?.aborted) { const packet = new Packet(type); packet.packetId(packetNumber += 1); packet.resetConnection(resetConnection); @@ -317,6 +346,10 @@ class MessageIO extends EventEmitter { stream.removeListener('drain', onDrain); stream.removeListener('close', onDrain); stream.removeListener('error', onError); + + if (signal && onAbort) { + signal.removeEventListener('abort', onAbort); + } } } @@ -330,10 +363,20 @@ class MessageIO extends EventEmitter { * Any bytes following the message's last packet (e.g. the start of the next * message) are pushed back onto the stream, to be consumed by the next read. * + * If the given `signal` is aborted, reading stops and the signal's abort + * reason is thrown. This also interrupts waiting for more data on a quiet + * stream, which an external `.return()` or `.throw()` call can not do (it + * would be queued behind the pending read). + * * @param stream The stream to read the message from. * @param debug The debug instance to use for logging. + * @param options.signal An abort signal to stop reading the message. */ - static async *readMessage(stream: Readable, debug: Debug): AsyncGenerator { + static async *readMessage(stream: Readable, debug: Debug, options: { signal?: AbortSignal } = {}): AsyncGenerator { + const { signal } = options; + + signal?.throwIfAborted(); + if (!stream.readable) { throw new Error('Premature close'); } @@ -372,6 +415,20 @@ class MessageIO extends EventEmitter { } }; + let onAbort: (() => void) | null = null; + if (signal) { + onAbort = () => { + error ??= signal.reason; + + if (waiting) { + const { reject } = waiting; + waiting = null; + reject(signal.reason); + } + }; + signal.addEventListener('abort', onAbort, { once: true }); + } + stream.on('readable', onReadable); stream.on('error', onError); stream.on('close', onClose); @@ -438,6 +495,10 @@ class MessageIO extends EventEmitter { stream.removeListener('readable', onReadable); stream.removeListener('error', onError); stream.removeListener('close', onClose); + + if (signal && onAbort) { + signal.removeEventListener('abort', onAbort); + } } } } diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index 70be87574..b76976c8a 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -1,5 +1,5 @@ import { type AddressInfo, createConnection, createServer, Server, Socket } from 'net'; -import { once } from 'events'; +import { getEventListeners, once } from 'events'; import { assert } from 'chai'; import { promisify } from 'util'; import DuplexPair from 'native-duplexpair'; @@ -275,6 +275,97 @@ describe('MessageIO.writeMessage', function() { assert(hadError); }); + + it('throws when the given signal is already aborted', async function() { + const stream = new BufferListStream(); + const signal = AbortSignal.abort(new Error('canceled')); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, [Buffer.from([1, 2, 3])], { signal }); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'canceled'); + } + + assert(hadError); + + // Nothing was written. + assert.isNull(stream.read()); + assertNoDanglingEventListeners(stream); + }); + + it('stops writing when the given signal is aborted while waiting for the stream to drain', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new Duplex({ + write(chunk, encoding, callback) { + // never call callback so that the stream never drains + }, + read() {}, + + // instantly return false on write requests to indicate that the stream needs to drain + highWaterMark: 1 + }); + + const controller = new AbortController(); + setTimeout(() => { + assert(stream.writableNeedDrain); + controller.abort(new Error('canceled')); + }, 20); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload, payload, payload], { signal: controller.signal }); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'canceled'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); + }); + + it('stops writing when the given signal is aborted while waiting for more payload data', async function() { + const payload = Buffer.from([1, 2, 3, 4, 5, 6]); + const stream = new BufferListStream(); + + const controller = new AbortController(); + setTimeout(() => { + controller.abort(new Error('canceled')); + }, 20); + + let hadError = false; + try { + await MessageIO.writeMessage(stream, debug, packetSize, packetType, (async function*() { + yield payload; + + // Stall forever, waiting for more data that never arrives. + await new Promise(() => {}); + })(), { signal: controller.signal }); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'canceled'); + } + + assert(hadError); + + // The part of the payload that filled a whole packet was written, but + // no `IGNORE` terminator: the message is intentionally left unterminated + // as the connection is being torn down. + const packets = splitPackets(stream.read()); + assert.lengthOf(packets, 1); + assert.isFalse(packets[0].isLast()); + + assertNoDanglingEventListeners(stream); + assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); + }); }); describe('MessageIO.readMessage', function() { @@ -441,6 +532,115 @@ describe('MessageIO.readMessage', function() { assert(hadError); }); + it('throws when the given signal is already aborted', async function() { + const stream = new BufferListStream(); + const signal = AbortSignal.abort(new Error('canceled')); + + let hadError = false; + try { + const message = MessageIO.readMessage(stream, debug, { signal }); + while (!(await message.next()).done) { + // Discard the message contents. + } + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'canceled'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); + + it('stops reading when the given signal is aborted while waiting for data on a quiet stream', async function() { + // A stream that stays open but never produces data, like a connection + // to a server that accepted the connection but never responds. + const stream = new Readable({ read() {} }); + + const controller = new AbortController(); + setTimeout(() => { + controller.abort(new Error('canceled')); + }, 20); + + let hadError = false; + try { + const message = MessageIO.readMessage(stream, debug, { signal: controller.signal }); + while (!(await message.next()).done) { + // Discard the message contents. + } + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'canceled'); + } + + assert(hadError); + assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); + }); + + it('stops reading when the given signal is aborted after a partial message was read', async function() { + const firstPacket = new Packet(packetType); + firstPacket.addData(Buffer.from([1, 2, 3])); + + const stream = new Readable({ read() {} }); + stream.push(firstPacket.buffer); + + const controller = new AbortController(); + setTimeout(() => { + controller.abort(new Error('canceled')); + }, 20); + + const chunks = []; + let hadError = false; + try { + for await (const chunk of MessageIO.readMessage(stream, debug, { signal: controller.signal })) { + chunks.push(chunk); + } + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'canceled'); + } + + assert(hadError); + assert.deepEqual(chunks, [Buffer.from([1, 2, 3])]); + assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); + }); + + it('settles pending reads when the given signal is aborted, instead of deadlocking on `.return()`', async function() { + // Regression test: calling `.return()` on the generator while it is + // suspended waiting for stream data is queued behind the pending + // `.next()` call and would deadlock. Aborting the signal settles the + // pending read, after which `.return()` resolves normally. + const stream = new Readable({ read() {} }); + + const controller = new AbortController(); + const message = MessageIO.readMessage(stream, debug, { signal: controller.signal }); + + const pendingRead = message.next(); + + controller.abort(new Error('canceled')); + + let hadError = false; + try { + await pendingRead; + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'canceled'); + } + + assert(hadError); + + // The generator has completed; `.return()` settles immediately. + assert.deepEqual(await message.return(), { value: undefined, done: true }); + assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); + }); + it('throws when receiving a packet with an invalid length', async function() { const invalidPacket = Buffer.alloc(8); invalidPacket.writeUInt16BE(4, 2); From d84e7ad5bdb4e8f5f52ee2d5673d6b28f5e07fa7 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Fri, 17 Jul 2026 17:08:36 +0000 Subject: [PATCH 03/14] refactor: make `writeMessage`/`readMessage` module-level functions Instead of static methods on `MessageIO`, `writeMessage` and `readMessage` are now plain module-level functions, matching the precedent set by `connector.ts`. The `debug` parameter moves into the options object and becomes optional, shrinking the positional parameter lists: writeMessage(stream, packetSize, type, payload, { debug, resetConnection, signal }) readMessage(stream, { debug, signal }) Since `module.exports = MessageIO` replaces the named exports on the CommonJS side, they are explicitly re-attached for `require` users. Co-Authored-By: Claude Fable 5 --- benchmarks/message-io/read-message.js | 4 +- benchmarks/message-io/write-message.js | 4 +- src/connection.ts | 6 +- src/message-io.ts | 496 +++++++++++++------------ test/unit/message-io-test.ts | 66 ++-- 5 files changed, 284 insertions(+), 292 deletions(-) diff --git a/benchmarks/message-io/read-message.js b/benchmarks/message-io/read-message.js index c3594ed76..0bd70320b 100644 --- a/benchmarks/message-io/read-message.js +++ b/benchmarks/message-io/read-message.js @@ -2,7 +2,7 @@ const { createBenchmark } = require('../common'); const { Readable } = require('stream'); const Debug = require('tedious/lib/debug'); -const MessageIO = require('tedious/lib/message-io'); +const { readMessage } = require('tedious/lib/message-io'); const { Packet } = require('tedious/lib/packet'); const bench = createBenchmark(main, { @@ -28,7 +28,7 @@ function main({ n }) { let total = 0; for (let i = 0; i < n; i++) { - for await (const chunk of MessageIO.readMessage(stream, debug)) { + for await (const chunk of readMessage(stream, { debug })) { total += chunk.length; } } diff --git a/benchmarks/message-io/write-message.js b/benchmarks/message-io/write-message.js index 545e78ddf..8580980f0 100644 --- a/benchmarks/message-io/write-message.js +++ b/benchmarks/message-io/write-message.js @@ -2,7 +2,7 @@ const { createBenchmark } = require('../common'); const { Duplex } = require('stream'); const Debug = require('tedious/lib/debug'); -const MessageIO = require('tedious/lib/message-io'); +const { writeMessage } = require('tedious/lib/message-io'); const bench = createBenchmark(main, { n: [100, 1000, 10000, 100000] @@ -33,7 +33,7 @@ function main({ n }) { bench.start(); for (let i = 0; i < n; i++) { - await MessageIO.writeMessage(stream, debug, 8 + 1024, 2, payload); + await writeMessage(stream, 8 + 1024, 2, payload, { debug }); } bench.end(n); diff --git a/src/connection.ts b/src/connection.ts index 737048070..5154bcb0d 100644 --- a/src/connection.ts +++ b/src/connection.ts @@ -29,7 +29,7 @@ import NTLMResponsePayload from './ntlm-payload'; import Request from './request'; import RpcRequestPayload from './rpcrequest-payload'; import SqlBatchPayload from './sqlbatch-payload'; -import MessageIO from './message-io'; +import MessageIO, { readMessage, writeMessage } from './message-io'; import { Parser as TokenStreamParser } from './token/token-stream-parser'; import { Transaction, ISOLATION_LEVEL, assertValidIsolationLevel } from './transaction'; import { ConnectionError, RequestError } from './errors'; @@ -2481,7 +2481,7 @@ class Connection extends EventEmitter { version: { major: Number(major), minor: Number(minor), build: Number(build), subbuild: 0 } }); - await MessageIO.writeMessage(socket, this.debug, this.config.options.packetSize, TYPE.PRELOGIN, [payload.data], { signal }); + await writeMessage(socket, this.config.options.packetSize, TYPE.PRELOGIN, [payload.data], { debug: this.debug, signal }); this.debug.payload(function() { return payload.toString(' '); }); @@ -3383,7 +3383,7 @@ class Connection extends EventEmitter { async readPreloginResponse(socket: net.Socket, signal: AbortSignal): Promise { let messageBuffer = Buffer.alloc(0); - for await (const data of MessageIO.readMessage(socket, this.debug, { signal })) { + for await (const data of readMessage(socket, { debug: this.debug, signal })) { messageBuffer = Buffer.concat([messageBuffer, data]); } diff --git a/src/message-io.ts b/src/message-io.ts index 37bdb94d6..10a0aeae6 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -190,318 +190,322 @@ class MessageIO extends EventEmitter { return result.value; } +} - /** - * Write a message with the given `type` and `payload` to the given `stream`, - * wrapping it into TDS packets of the given `packetSize`. - * - * Respects backpressure from the stream, waiting for it to drain before - * writing more data. - * - * If iterating the `payload` throws, the message is terminated with a final - * packet that has the `IGNORE` flag set (telling the server to disregard the - * message) and the error is re-thrown. Errors from the stream itself are - * thrown as-is. - * - * If the given `signal` is aborted, writing stops and the signal's abort - * reason is thrown. Note that this can leave a partially written message on - * the stream, so it must only be used when the connection is being torn down - * anyway. To cancel a message in a way that keeps the TDS stream aligned, - * throw from the `payload` iterable instead. - * - * @param stream The stream to write the message to. - * @param debug The debug instance to use for logging. - * @param packetSize The maximum packet size to use. - * @param type The type of the message to write. - * @param payload The payload to write. - * @param options.resetConnection Whether the server should reset the connection when processing the message. - * @param options.signal An abort signal to stop writing the message. - */ - static async writeMessage(stream: Writable, debug: Debug, packetSize: number, type: number, payload: AsyncIterable | Iterable, options: { resetConnection?: boolean, signal?: AbortSignal } = {}): Promise { - const { resetConnection = false, signal } = options; +/** + * Write a message with the given `type` and `payload` to the given `stream`, + * wrapping it into TDS packets of the given `packetSize`. + * + * Respects backpressure from the stream, waiting for it to drain before + * writing more data. + * + * If iterating the `payload` throws, the message is terminated with a final + * packet that has the `IGNORE` flag set (telling the server to disregard the + * message) and the error is re-thrown. Errors from the stream itself are + * thrown as-is. + * + * If the given `signal` is aborted, writing stops and the signal's abort + * reason is thrown. Note that this can leave a partially written message on + * the stream, so it must only be used when the connection is being torn down + * anyway. To cancel a message in a way that keeps the TDS stream aligned, + * throw from the `payload` iterable instead. + * + * @param stream The stream to write the message to. + * @param packetSize The maximum packet size to use. + * @param type The type of the message to write. + * @param payload The payload to write. + * @param options.debug A debug instance to log packets to. + * @param options.resetConnection Whether the server should reset the connection when processing the message. + * @param options.signal An abort signal to stop writing the message. + */ +export async function writeMessage(stream: Writable, packetSize: number, type: number, payload: AsyncIterable | Iterable, options: { debug?: Debug, resetConnection?: boolean, signal?: AbortSignal } = {}): Promise { + const { debug, resetConnection = false, signal } = options; + + signal?.throwIfAborted(); + + if (!stream.writable) { + throw new Error('Premature close'); + } - signal?.throwIfAborted(); + let drain: PromiseWithResolvers | null = null; - if (!stream.writable) { - throw new Error('Premature close'); + const onDrain = () => { + if (drain) { + const { resolve } = drain; + drain = null; + resolve(); } + }; - let drain: PromiseWithResolvers | null = null; - - const onDrain = () => { - if (drain) { - const { resolve } = drain; - drain = null; - resolve(); - } - }; - - const onError = (err: Error) => { - if (drain) { - const { reject } = drain; - drain = null; - reject(err); - } - }; - - let abortPromise: Promise | null = null; - let onAbort: (() => void) | null = null; - - if (signal) { - const { promise, reject } = Promise.withResolvers(); - - // Prevent unhandled rejections if the signal is aborted while - // nothing is currently racing against `abortPromise`. - promise.catch(() => {}); - - abortPromise = promise; - onAbort = () => { reject(signal.reason); }; - signal.addEventListener('abort', onAbort, { once: true }); + const onError = (err: Error) => { + if (drain) { + const { reject } = drain; + drain = null; + reject(err); } + }; - const waitForDrain = () => { - drain = Promise.withResolvers(); - return abortPromise ? Promise.race([drain.promise, abortPromise]) : drain.promise; - }; - - stream.on('drain', onDrain); - stream.on('close', onDrain); - stream.on('error', onError); + let abortPromise: Promise | null = null; + let onAbort: (() => void) | null = null; - try { - const bl = new BufferList(); - const length = packetSize - HEADER_LENGTH; - let packetNumber = 0; + if (signal) { + const { promise, reject } = Promise.withResolvers(); - const writePacket = async (packet: Packet) => { - debug.packet('Sent', packet); - debug.data(packet); + // Prevent unhandled rejections if the signal is aborted while + // nothing is currently racing against `abortPromise`. + promise.catch(() => {}); - if (stream.write(packet.buffer) === false) { - await waitForDrain(); - } - }; + abortPromise = promise; + onAbort = () => { reject(signal.reason); }; + signal.addEventListener('abort', onAbort, { once: true }); + } - let iterator; - if ((payload as AsyncIterable)[Symbol.asyncIterator]) { - iterator = (payload as AsyncIterable)[Symbol.asyncIterator](); - } else { - iterator = (payload as Iterable)[Symbol.iterator](); - } + const waitForDrain = () => { + drain = Promise.withResolvers(); + return abortPromise ? Promise.race([drain.promise, abortPromise]) : drain.promise; + }; - while (true) { - let value, done; - try { - const result = iterator.next(); - ({ value, done } = abortPromise ? await Promise.race([Promise.resolve(result), abortPromise]) : await result); - } catch (err) { - // The payload errored while being iterated. If the stream is still - // writable, terminate the message with the `IGNORE` flag set so the - // server disregards everything sent so far. If the signal was - // aborted instead, the connection is being torn down and the - // message is left unterminated. - if (stream.writable && !signal?.aborted) { - const packet = new Packet(type); - packet.packetId(packetNumber += 1); - packet.resetConnection(resetConnection); - packet.last(true); - packet.ignore(true); - - await writePacket(packet); - } + stream.on('drain', onDrain); + stream.on('close', onDrain); + stream.on('error', onError); - throw err; - } + try { + const bl = new BufferList(); + const length = packetSize - HEADER_LENGTH; + let packetNumber = 0; - if (done) { - break; - } + const writePacket = async (packet: Packet) => { + debug?.packet('Sent', packet); + debug?.data(packet); - bl.append(value); + if (stream.write(packet.buffer) === false) { + await waitForDrain(); + } + }; - while (bl.length > length) { - const data = bl.slice(0, length); - bl.consume(length); + let iterator; + if ((payload as AsyncIterable)[Symbol.asyncIterator]) { + iterator = (payload as AsyncIterable)[Symbol.asyncIterator](); + } else { + iterator = (payload as Iterable)[Symbol.iterator](); + } + while (true) { + let value, done; + try { + const result = iterator.next(); + ({ value, done } = abortPromise ? await Promise.race([Promise.resolve(result), abortPromise]) : await result); + } catch (err) { + // The payload errored while being iterated. If the stream is still + // writable, terminate the message with the `IGNORE` flag set so the + // server disregards everything sent so far. If the signal was + // aborted instead, the connection is being torn down and the + // message is left unterminated. + if (stream.writable && !signal?.aborted) { const packet = new Packet(type); packet.packetId(packetNumber += 1); packet.resetConnection(resetConnection); - packet.addData(data); + packet.last(true); + packet.ignore(true); await writePacket(packet); } + + throw err; + } + + if (done) { + break; } - const data = bl.slice(); - bl.consume(data.length); + bl.append(value); - const packet = new Packet(type); - packet.packetId(packetNumber += 1); - packet.resetConnection(resetConnection); - packet.last(true); - packet.addData(data); + while (bl.length > length) { + const data = bl.slice(0, length); + bl.consume(length); - await writePacket(packet); - } finally { - stream.removeListener('drain', onDrain); - stream.removeListener('close', onDrain); - stream.removeListener('error', onError); + const packet = new Packet(type); + packet.packetId(packetNumber += 1); + packet.resetConnection(resetConnection); + packet.addData(data); - if (signal && onAbort) { - signal.removeEventListener('abort', onAbort); + await writePacket(packet); } } - } - /** - * Read the next TDS message from the given `stream`. - * - * Returns an async generator that yields the data of the message's packets - * as they arrive. The generator throws if the stream emits an error or is - * closed before the message was fully read. - * - * Any bytes following the message's last packet (e.g. the start of the next - * message) are pushed back onto the stream, to be consumed by the next read. - * - * If the given `signal` is aborted, reading stops and the signal's abort - * reason is thrown. This also interrupts waiting for more data on a quiet - * stream, which an external `.return()` or `.throw()` call can not do (it - * would be queued behind the pending read). - * - * @param stream The stream to read the message from. - * @param debug The debug instance to use for logging. - * @param options.signal An abort signal to stop reading the message. - */ - static async *readMessage(stream: Readable, debug: Debug, options: { signal?: AbortSignal } = {}): AsyncGenerator { - const { signal } = options; + const data = bl.slice(); + bl.consume(data.length); + + const packet = new Packet(type); + packet.packetId(packetNumber += 1); + packet.resetConnection(resetConnection); + packet.last(true); + packet.addData(data); - signal?.throwIfAborted(); + await writePacket(packet); + } finally { + stream.removeListener('drain', onDrain); + stream.removeListener('close', onDrain); + stream.removeListener('error', onError); - if (!stream.readable) { - throw new Error('Premature close'); + if (signal && onAbort) { + signal.removeEventListener('abort', onAbort); } + } +} - const bl = new BufferList(); +/** + * Read the next TDS message from the given `stream`. + * + * Returns an async generator that yields the data of the message's packets + * as they arrive. The generator throws if the stream emits an error or is + * closed before the message was fully read. + * + * Any bytes following the message's last packet (e.g. the start of the next + * message) are pushed back onto the stream, to be consumed by the next read. + * + * If the given `signal` is aborted, reading stops and the signal's abort + * reason is thrown. This also interrupts waiting for more data on a quiet + * stream, which an external `.return()` or `.throw()` call can not do (it + * would be queued behind the pending read). + * + * @param stream The stream to read the message from. + * @param options.debug A debug instance to log packets to. + * @param options.signal An abort signal to stop reading the message. + */ +export async function* readMessage(stream: Readable, options: { debug?: Debug, signal?: AbortSignal } = {}): AsyncGenerator { + const { debug, signal } = options; + + signal?.throwIfAborted(); + + if (!stream.readable) { + throw new Error('Premature close'); + } - let error: Error | null = null; - let closed = false; - let waiting: PromiseWithResolvers | null = null; + const bl = new BufferList(); - const onReadable = () => { - if (waiting) { - const { resolve } = waiting; - waiting = null; - resolve(); - } - }; + let error: Error | null = null; + let closed = false; + let waiting: PromiseWithResolvers | null = null; - const onError = (err: Error) => { - error = err; + const onReadable = () => { + if (waiting) { + const { resolve } = waiting; + waiting = null; + resolve(); + } + }; - if (waiting) { - const { reject } = waiting; - waiting = null; - reject(err); - } - }; + const onError = (err: Error) => { + error = err; + + if (waiting) { + const { reject } = waiting; + waiting = null; + reject(err); + } + }; - const onClose = () => { - closed = true; + const onClose = () => { + closed = true; + + if (waiting) { + const { reject } = waiting; + waiting = null; + reject(new Error('Premature close')); + } + }; + + let onAbort: (() => void) | null = null; + if (signal) { + onAbort = () => { + error ??= signal.reason; if (waiting) { const { reject } = waiting; waiting = null; - reject(new Error('Premature close')); + reject(signal.reason); } }; + signal.addEventListener('abort', onAbort, { once: true }); + } - let onAbort: (() => void) | null = null; - if (signal) { - onAbort = () => { - error ??= signal.reason; - - if (waiting) { - const { reject } = waiting; - waiting = null; - reject(signal.reason); - } - }; - signal.addEventListener('abort', onAbort, { once: true }); - } + stream.on('readable', onReadable); + stream.on('error', onError); + stream.on('close', onClose); - stream.on('readable', onReadable); - stream.on('error', onError); - stream.on('close', onClose); + try { + while (true) { + if (error) { + throw error; + } - try { - while (true) { - if (error) { - throw error; - } + if (closed) { + throw new Error('Premature close'); + } - if (closed) { - throw new Error('Premature close'); - } + let chunk: Buffer; + while ((chunk = stream.read()) !== null) { + bl.append(chunk); - let chunk: Buffer; - while ((chunk = stream.read()) !== null) { - bl.append(chunk); + // The packet header is always 8 bytes of length. + while (bl.length >= HEADER_LENGTH) { + // Get the full packet length + const length = bl.readUInt16BE(2); + if (length < HEADER_LENGTH) { + throw new ConnectionError('Unable to process incoming packet'); + } - // The packet header is always 8 bytes of length. - while (bl.length >= HEADER_LENGTH) { - // Get the full packet length - const length = bl.readUInt16BE(2); - if (length < HEADER_LENGTH) { - throw new ConnectionError('Unable to process incoming packet'); - } + if (bl.length < length) { + break; + } - if (bl.length < length) { - break; - } + const data = bl.slice(0, length); + bl.consume(length); - const data = bl.slice(0, length); - bl.consume(length); + const packet = new Packet(data); + debug?.packet('Received', packet); + debug?.data(packet); - const packet = new Packet(data); - debug.packet('Received', packet); - debug.data(packet); + yield packet.data(); - yield packet.data(); + // Did the stream error while we yielded? + if (error) { + throw error; + } - // Did the stream error while we yielded? - if (error) { - throw error; + if (packet.isLast()) { + // This was the last packet of the message. Any data left in the + // buffer belongs to the next message (e.g. the response to an + // `ATTENTION` message sent by the client while reading an + // incoming response), so push it back onto the stream. + if (bl.length) { + stream.unshift(bl.slice()); } - if (packet.isLast()) { - // This was the last packet of the message. Any data left in the - // buffer belongs to the next message (e.g. the response to an - // `ATTENTION` message sent by the client while reading an - // incoming response), so push it back onto the stream. - if (bl.length) { - stream.unshift(bl.slice()); - } - - return; - } + return; } } - - // Wait for the stream to become readable again (or error out or close). - waiting = Promise.withResolvers(); - await waiting.promise; } - } finally { - stream.removeListener('readable', onReadable); - stream.removeListener('error', onError); - stream.removeListener('close', onClose); - if (signal && onAbort) { - signal.removeEventListener('abort', onAbort); - } + // Wait for the stream to become readable again (or error out or close). + waiting = Promise.withResolvers(); + await waiting.promise; + } + } finally { + stream.removeListener('readable', onReadable); + stream.removeListener('error', onError); + stream.removeListener('close', onClose); + + if (signal && onAbort) { + signal.removeEventListener('abort', onAbort); } } } export default MessageIO; module.exports = MessageIO; +// The `module.exports` assignment above replaces the named exports on the +// CommonJS side, so re-attach them. +module.exports.writeMessage = writeMessage; +module.exports.readMessage = readMessage; diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index b76976c8a..0d69bc57b 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -9,7 +9,7 @@ import { Duplex, Readable } from 'stream'; import BufferListStream from 'bl'; import Debug from '../../src/debug'; -import MessageIO from '../../src/message-io'; +import MessageIO, { readMessage, writeMessage } from '../../src/message-io'; import Message from '../../src/message'; import { Packet, TYPE } from '../../src/packet'; import { ConnectionError } from '../../src/errors'; @@ -38,18 +38,12 @@ function splitPackets(data: Buffer): Packet[] { return packets; } -describe('MessageIO.writeMessage', function() { - let debug: Debug; - - beforeEach(function() { - debug = new Debug(); - }); - +describe('writeMessage', function() { it('wraps the given payload into a TDS packet and writes it to the given stream', async function() { const payload = Buffer.from([1, 2, 3]); const stream = new BufferListStream(); - await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload]); + await writeMessage(stream, packetSize, packetType, [payload]); const buf = stream.read(); assert.instanceOf(buf, Buffer); @@ -69,7 +63,7 @@ describe('MessageIO.writeMessage', function() { const stream = new BufferListStream(); // `packetSize` allows for 4 bytes of data per packet - await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload]); + await writeMessage(stream, packetSize, packetType, [payload]); const [firstPacket, secondPacket, ...rest] = splitPackets(stream.read()); assert.lengthOf(rest, 0); @@ -86,7 +80,7 @@ describe('MessageIO.writeMessage', function() { it('writes an empty final packet for an empty payload', async function() { const stream = new BufferListStream(); - await MessageIO.writeMessage(stream, debug, packetSize, packetType, []); + await writeMessage(stream, packetSize, packetType, []); const packet = new Packet(stream.read()); assert.isTrue(packet.isLast()); @@ -101,7 +95,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, (async function*() { + await writeMessage(stream, packetSize, packetType, (async function*() { yield payload; throw new Error('iteration error'); })()); @@ -146,7 +140,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, (async function*() { + await writeMessage(stream, packetSize, packetType, (async function*() { yield payload; // Simulate draining the stream after the exception was thrown @@ -181,7 +175,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload]); + await writeMessage(stream, packetSize, packetType, [payload]); } catch (err: any) { hadError = true; @@ -212,7 +206,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload, payload, payload]); + await writeMessage(stream, packetSize, packetType, [payload, payload, payload]); } catch (err: any) { hadError = true; @@ -243,7 +237,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, (async function*() { + await writeMessage(stream, packetSize, packetType, (async function*() { yield payload; yield payload; yield payload; @@ -265,7 +259,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, []); + await writeMessage(stream, packetSize, packetType, []); } catch (err: any) { hadError = true; @@ -282,7 +276,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, [Buffer.from([1, 2, 3])], { signal }); + await writeMessage(stream, packetSize, packetType, [Buffer.from([1, 2, 3])], { signal }); } catch (err: any) { hadError = true; @@ -317,7 +311,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, [payload, payload, payload], { signal: controller.signal }); + await writeMessage(stream, packetSize, packetType, [payload, payload, payload], { signal: controller.signal }); } catch (err: any) { hadError = true; @@ -341,7 +335,7 @@ describe('MessageIO.writeMessage', function() { let hadError = false; try { - await MessageIO.writeMessage(stream, debug, packetSize, packetType, (async function*() { + await writeMessage(stream, packetSize, packetType, (async function*() { yield payload; // Stall forever, waiting for more data that never arrives. @@ -368,13 +362,7 @@ describe('MessageIO.writeMessage', function() { }); }); -describe('MessageIO.readMessage', function() { - let debug: Debug; - - beforeEach(function() { - debug = new Debug(); - }); - +describe('readMessage', function() { it('reads a message consisting of a single TDS packet from the given stream', async function() { const payload = Buffer.from([1, 2, 3]); const packet = new Packet(packetType); @@ -385,7 +373,7 @@ describe('MessageIO.readMessage', function() { stream.write(packet.buffer); const chunks = []; - for await (const chunk of MessageIO.readMessage(stream, debug)) { + for await (const chunk of readMessage(stream)) { chunks.push(chunk); } @@ -408,7 +396,7 @@ describe('MessageIO.readMessage', function() { stream.write(lastPacket.buffer); const chunks = []; - for await (const chunk of MessageIO.readMessage(stream, debug)) { + for await (const chunk of readMessage(stream)) { chunks.push(chunk); } @@ -437,7 +425,7 @@ describe('MessageIO.readMessage', function() { } const chunks = []; - for await (const chunk of MessageIO.readMessage(stream, debug)) { + for await (const chunk of readMessage(stream)) { chunks.push(chunk); } @@ -456,7 +444,7 @@ describe('MessageIO.readMessage', function() { stream.write(Buffer.concat([packet.buffer, trailingBytes])); const chunks = []; - for await (const chunk of MessageIO.readMessage(stream, debug)) { + for await (const chunk of readMessage(stream)) { chunks.push(chunk); } @@ -471,7 +459,7 @@ describe('MessageIO.readMessage', function() { let hadError = false; try { - const message = MessageIO.readMessage(stream, debug); + const message = readMessage(stream); while (!(await message.next()).done) { // Discard the message contents. } @@ -498,7 +486,7 @@ describe('MessageIO.readMessage', function() { let hadError = false; try { - const message = MessageIO.readMessage(stream, debug); + const message = readMessage(stream); while (!(await message.next()).done) { // Discard the message contents. } @@ -518,7 +506,7 @@ describe('MessageIO.readMessage', function() { let hadError = false; try { - const message = MessageIO.readMessage(stream, debug); + const message = readMessage(stream); while (!(await message.next()).done) { // Discard the message contents. } @@ -538,7 +526,7 @@ describe('MessageIO.readMessage', function() { let hadError = false; try { - const message = MessageIO.readMessage(stream, debug, { signal }); + const message = readMessage(stream, { signal }); while (!(await message.next()).done) { // Discard the message contents. } @@ -565,7 +553,7 @@ describe('MessageIO.readMessage', function() { let hadError = false; try { - const message = MessageIO.readMessage(stream, debug, { signal: controller.signal }); + const message = readMessage(stream, { signal: controller.signal }); while (!(await message.next()).done) { // Discard the message contents. } @@ -595,7 +583,7 @@ describe('MessageIO.readMessage', function() { const chunks = []; let hadError = false; try { - for await (const chunk of MessageIO.readMessage(stream, debug, { signal: controller.signal })) { + for await (const chunk of readMessage(stream, { signal: controller.signal })) { chunks.push(chunk); } } catch (err: any) { @@ -618,7 +606,7 @@ describe('MessageIO.readMessage', function() { const stream = new Readable({ read() {} }); const controller = new AbortController(); - const message = MessageIO.readMessage(stream, debug, { signal: controller.signal }); + const message = readMessage(stream, { signal: controller.signal }); const pendingRead = message.next(); @@ -650,7 +638,7 @@ describe('MessageIO.readMessage', function() { let hadError = false; try { - const message = MessageIO.readMessage(stream, debug); + const message = readMessage(stream); while (!(await message.next()).done) { // Discard the message contents. } From 68315681cd6b6e0a1128c84d503410691e326d29 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Fri, 17 Jul 2026 21:53:42 +0000 Subject: [PATCH 04/14] refactor: add a `cancelSignal` option to `writeMessage` While the `signal` option is for connection teardown (stop writing, throw, stream may be left misaligned), the new `cancelSignal` option expresses graceful request cancellation as defined by the protocol: stop consuming the payload, discard buffered data, and terminate the message with the `IGNORE` flag set. This is a normal protocol outcome, so the returned promise resolves normally, the TDS stream stays aligned, and the server sends a short response to the ignored message. Cancellation also skips any pending drain waits: nothing further will be written for this message, so there is no reason to keep the caller waiting on a stalled peer. This replaces the need for callers to express cancellation by throwing a sentinel error from the payload iterable and swallowing it again. When both signals are aborted, teardown wins. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 84 ++++++++++++++++++++++++---- test/unit/message-io-test.ts | 104 +++++++++++++++++++++++++++++++++++ 2 files changed, 177 insertions(+), 11 deletions(-) diff --git a/src/message-io.ts b/src/message-io.ts index 10a0aeae6..7cf7561bc 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -204,11 +204,17 @@ class MessageIO extends EventEmitter { * message) and the error is re-thrown. Errors from the stream itself are * thrown as-is. * + * If the given `cancelSignal` is aborted, the remaining payload is discarded + * and the message is terminated with the `IGNORE` flag set, telling the + * server to disregard it. This is a normal protocol outcome: the returned + * promise resolves, the TDS stream stays aligned, and the server will send a + * (short) response to the ignored message. Callers can check + * `cancelSignal.aborted` to distinguish this from a fully sent message. + * * If the given `signal` is aborted, writing stops and the signal's abort * reason is thrown. Note that this can leave a partially written message on * the stream, so it must only be used when the connection is being torn down - * anyway. To cancel a message in a way that keeps the TDS stream aligned, - * throw from the `payload` iterable instead. + * anyway. When both signals are aborted, `signal` wins. * * @param stream The stream to write the message to. * @param packetSize The maximum packet size to use. @@ -216,10 +222,11 @@ class MessageIO extends EventEmitter { * @param payload The payload to write. * @param options.debug A debug instance to log packets to. * @param options.resetConnection Whether the server should reset the connection when processing the message. - * @param options.signal An abort signal to stop writing the message. + * @param options.cancelSignal An abort signal to cancel the message while keeping the connection usable. + * @param options.signal An abort signal to stop writing the message during connection teardown. */ -export async function writeMessage(stream: Writable, packetSize: number, type: number, payload: AsyncIterable | Iterable, options: { debug?: Debug, resetConnection?: boolean, signal?: AbortSignal } = {}): Promise { - const { debug, resetConnection = false, signal } = options; +export async function writeMessage(stream: Writable, packetSize: number, type: number, payload: AsyncIterable | Iterable, options: { debug?: Debug, resetConnection?: boolean, cancelSignal?: AbortSignal, signal?: AbortSignal } = {}): Promise { + const { debug, resetConnection = false, cancelSignal, signal } = options; signal?.throwIfAborted(); @@ -260,9 +267,38 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n signal.addEventListener('abort', onAbort, { once: true }); } + const CANCEL = Symbol('cancel'); + + let canceled = cancelSignal?.aborted ?? false; + let cancelPromise: Promise | null = null; + let onCancel: (() => void) | null = null; + + if (cancelSignal && !canceled) { + const { promise, resolve } = Promise.withResolvers(); + + cancelPromise = promise; + onCancel = () => { + canceled = true; + resolve(CANCEL); + }; + cancelSignal.addEventListener('abort', onCancel, { once: true }); + } + const waitForDrain = () => { drain = Promise.withResolvers(); - return abortPromise ? Promise.race([drain.promise, abortPromise]) : drain.promise; + + if (abortPromise || cancelPromise) { + const contenders: Promise[] = [drain.promise]; + if (abortPromise) { + contenders.push(abortPromise); + } + if (cancelPromise) { + contenders.push(cancelPromise); + } + return Promise.race(contenders); + } + + return drain.promise; }; stream.on('drain', onDrain); @@ -290,11 +326,30 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n iterator = (payload as Iterable)[Symbol.iterator](); } - while (true) { + while (!canceled) { let value, done; try { const result = iterator.next(); - ({ value, done } = abortPromise ? await Promise.race([Promise.resolve(result), abortPromise]) : await result); + + let raceResult: IteratorResult | typeof CANCEL; + if (abortPromise || cancelPromise) { + const contenders: Promise | typeof CANCEL>[] = [Promise.resolve(result)]; + if (abortPromise) { + contenders.push(abortPromise); + } + if (cancelPromise) { + contenders.push(cancelPromise); + } + raceResult = await Promise.race(contenders); + } else { + raceResult = await result; + } + + if (raceResult === CANCEL) { + break; + } + + ({ value, done } = raceResult); } catch (err) { // The payload errored while being iterated. If the stream is still // writable, terminate the message with the `IGNORE` flag set so the @@ -320,7 +375,7 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n bl.append(value); - while (bl.length > length) { + while (!canceled && bl.length > length) { const data = bl.slice(0, length); bl.consume(length); @@ -333,13 +388,16 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } } - const data = bl.slice(); - bl.consume(data.length); + // On cancellation, any buffered payload data is discarded and the final + // packet is flagged so the server ignores the whole message. + const data = canceled ? Buffer.alloc(0) : bl.slice(); + bl.consume(bl.length); const packet = new Packet(type); packet.packetId(packetNumber += 1); packet.resetConnection(resetConnection); packet.last(true); + packet.ignore(canceled); packet.addData(data); await writePacket(packet); @@ -351,6 +409,10 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n if (signal && onAbort) { signal.removeEventListener('abort', onAbort); } + + if (cancelSignal && onCancel) { + cancelSignal.removeEventListener('abort', onCancel); + } } } diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index 0d69bc57b..380a04b67 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -270,6 +270,110 @@ describe('writeMessage', function() { assert(hadError); }); + it('terminates the message with the ignore flag set when the given cancel signal is aborted while writing', async function() { + const payload = Buffer.from([1, 2, 3, 4, 5, 6]); + const stream = new BufferListStream(); + + const controller = new AbortController(); + setTimeout(() => { + controller.abort(); + }, 20); + + await writeMessage(stream, packetSize, packetType, (async function*() { + yield payload; + + // Stall forever, waiting for more data that never arrives. + await new Promise(() => {}); + })(), { cancelSignal: controller.signal }); + + const [firstPacket, lastPacket, ...rest] = splitPackets(stream.read()); + assert.lengthOf(rest, 0); + + // The part of the payload that filled a whole packet was sent before the + // cancellation, the rest is discarded and the message is terminated with + // the `IGNORE` flag set. + assert.isFalse(firstPacket.isLast()); + assert.deepEqual(firstPacket.data(), payload.subarray(0, 4)); + + assert.isTrue(lastPacket.isLast()); + assert.include(lastPacket.statusAsString(), 'IGNORE'); + assert.deepEqual(lastPacket.data(), Buffer.alloc(0)); + + assertNoDanglingEventListeners(stream); + assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); + }); + + it('writes an ignored message without consuming the payload when the given cancel signal is already aborted', async function() { + const stream = new BufferListStream(); + const cancelSignal = AbortSignal.abort(); + + let payloadConsumed = false; + + await writeMessage(stream, packetSize, packetType, (async function*() { + payloadConsumed = true; + yield Buffer.from([1, 2, 3]); + })(), { cancelSignal }); + + assert.isFalse(payloadConsumed); + + const [packet, ...rest] = splitPackets(stream.read()); + assert.lengthOf(rest, 0); + + assert.isTrue(packet.isLast()); + assert.include(packet.statusAsString(), 'IGNORE'); + assert.deepEqual(packet.data(), Buffer.alloc(0)); + + assertNoDanglingEventListeners(stream); + }); + + it('stops waiting for the stream to drain when the given cancel signal is aborted', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new Duplex({ + write(chunk, encoding, callback) { + // never call callback so that the stream never drains + }, + read() {}, + + // instantly return false on write requests to indicate that the stream needs to drain + highWaterMark: 1 + }); + + const controller = new AbortController(); + setTimeout(() => { + assert(stream.writableNeedDrain); + controller.abort(); + }, 20); + + // Resolves normally: the remaining payload is discarded and the + // `IGNORE`-flagged terminator is written without further drain waits. + await writeMessage(stream, packetSize, packetType, [payload, payload, payload], { cancelSignal: controller.signal }); + + assertNoDanglingEventListeners(stream); + assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); + }); + + it('prefers teardown over cancellation when both signals are aborted', async function() { + const stream = new BufferListStream(); + const signal = AbortSignal.abort(new Error('teardown')); + const cancelSignal = AbortSignal.abort(); + + let hadError = false; + try { + await writeMessage(stream, packetSize, packetType, [Buffer.from([1, 2, 3])], { cancelSignal, signal }); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'teardown'); + } + + assert(hadError); + + // Nothing was written. + assert.isNull(stream.read()); + assertNoDanglingEventListeners(stream); + }); + it('throws when the given signal is already aborted', async function() { const stream = new BufferListStream(); const signal = AbortSignal.abort(new Error('canceled')); From 6f836608caf06b2e76d3aba54bdb68b9e429c3a9 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 07:30:42 +0000 Subject: [PATCH 05/14] fix: prevent `readMessage` hang when the stream closes during chunk processing After yielding a chunk, `readMessage` only re-checked the `error` flag on re-entry. A close without an error event (e.g. `socket.destroy()` or a plain FIN) arriving while the generator was suspended at `yield` was recorded in the `closed` flag, but nothing re-checked it: the generator would fall through to waiting for stream events that already fired and hang forever. Re-check `closed` after the yield, mirroring the `error` check. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 7 ++++++- test/unit/message-io-test.ts | 31 +++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/message-io.ts b/src/message-io.ts index 7cf7561bc..98f28759e 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -531,11 +531,16 @@ export async function* readMessage(stream: Readable, options: { debug?: Debug, s yield packet.data(); - // Did the stream error while we yielded? + // Did the stream error out or close while we yielded? The events + // have already fired, so the wait below would never settle. if (error) { throw error; } + if (closed) { + throw new Error('Premature close'); + } + if (packet.isLast()) { // This was the last packet of the message. Any data left in the // buffer belongs to the next message (e.g. the response to an diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index 380a04b67..cbbe97eac 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -604,6 +604,37 @@ describe('readMessage', function() { assert(hadError); }); + it('throws when the stream is closed while a yielded chunk is being processed', async function() { + const payload = Buffer.from([1, 2, 3]); + const packet = new Packet(packetType); + packet.addData(payload); + + const stream = new Readable({ read() {} }); + stream.push(packet.buffer); + + const message = readMessage(stream); + + const { value } = await message.next(); + assert.deepEqual(value, payload); + + // The generator is now suspended at `yield`. Close the stream without + // an error while the consumer is processing the chunk. + stream.destroy(); + await once(stream, 'close'); + + let hadError = false; + try { + await message.next(); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'Premature close'); + } + + assert(hadError); + }); + it('throws when the given stream is not readable', async function() { const stream = new Readable({ read() {} }); stream.destroy(); From 3ab3eea3ca792fd27498cb27d56192b3a8ec9758 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 21:10:28 +0000 Subject: [PATCH 06/14] perf: build outgoing packets in a single allocation Instead of building a header-only `Packet`, slicing the packet data out of the buffer list, and appending it via `Buffer.concat` (three allocations, with every payload byte copied twice), `writeMessage` now allocates the packet buffer once, writes the header fields in place, and copies the payload data directly out of the buffer list. A `Packet` object is only constructed (as a copy-free wrapper) when a debug instance is attached. Also restores the synchronous-iterator fast path from #1656: sync payloads can neither stall nor interleave with signal events between chunks, so they skip the per-chunk promise machinery entirely. `writeMessage` throughput (messages/second, Node.js v24.18.0, median of 3 runs): | n | before | after | | |---|---|---|---| | 100 | 17,913 | 23,155 | +29% | | 1,000 | 48,569 | 64,873 | +34% | | 10,000 | 128,604 | 175,395 | +36% | | 100,000 | 187,042 | 255,030 | +36% | `readMessage` gets the same debug-gating (constructing the `Packet` wrapper only when needed) and reads the status flags directly from the packet header; its throughput is unchanged. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 125 ++++++++++++++++++++++++++-------------------- src/packet.ts | 2 +- 2 files changed, 72 insertions(+), 55 deletions(-) diff --git a/src/message-io.ts b/src/message-io.ts index 98f28759e..c6f496ef0 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -9,7 +9,7 @@ import { EventEmitter } from 'events'; import Debug from './debug'; import Message from './message'; -import { HEADER_LENGTH, Packet, TYPE } from './packet'; +import { HEADER_LENGTH, OFFSET, Packet, STATUS, TYPE } from './packet'; import { ConnectionError } from './errors'; import IncomingMessageStream from './incoming-message-stream'; @@ -310,46 +310,81 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n const length = packetSize - HEADER_LENGTH; let packetNumber = 0; - const writePacket = async (packet: Packet) => { - debug?.packet('Sent', packet); - debug?.data(packet); + const baseStatus = resetConnection ? STATUS.RESETCONNECTION : STATUS.NORMAL; + + // Build a packet buffer in a single allocation, copying the packet data + // directly out of the buffered payload. This avoids the double copy + // (buffer list -> data buffer -> packet buffer) and the extra allocations + // that building a header-only `Packet` and appending data to it incur. + const buildPacket = (dataLength: number, status: number) => { + const buffer = Buffer.allocUnsafe(HEADER_LENGTH + dataLength); + buffer[OFFSET.Type] = type; + buffer[OFFSET.Status] = status; + buffer.writeUInt16BE(HEADER_LENGTH + dataLength, OFFSET.Length); + buffer.writeUInt16BE(0, OFFSET.SPID); + buffer[OFFSET.PacketID] = (packetNumber += 1) % 256; + buffer[OFFSET.Window] = 0; + + if (dataLength) { + bl.copy(buffer, HEADER_LENGTH, 0, dataLength); + bl.consume(dataLength); + } + + return buffer; + }; - if (stream.write(packet.buffer) === false) { + const writePacket = async (buffer: Buffer) => { + if (debug) { + const packet = new Packet(buffer); + debug.packet('Sent', packet); + debug.data(packet); + } + + if (stream.write(buffer) === false) { await waitForDrain(); } }; - let iterator; + let iterator: Iterator | AsyncIterator; + let isAsync; if ((payload as AsyncIterable)[Symbol.asyncIterator]) { + isAsync = true; iterator = (payload as AsyncIterable)[Symbol.asyncIterator](); } else { + isAsync = false; iterator = (payload as Iterable)[Symbol.iterator](); } while (!canceled) { let value, done; try { - const result = iterator.next(); + if (!isAsync) { + // Synchronous payloads can neither stall nor interleave with + // signal events between chunks, so skip the promise machinery. + ({ value, done } = (iterator as Iterator).next()); + } else { + const result = (iterator as AsyncIterator).next(); - let raceResult: IteratorResult | typeof CANCEL; - if (abortPromise || cancelPromise) { - const contenders: Promise | typeof CANCEL>[] = [Promise.resolve(result)]; - if (abortPromise) { - contenders.push(abortPromise); + let raceResult: IteratorResult | typeof CANCEL; + if (abortPromise || cancelPromise) { + const contenders: Promise | typeof CANCEL>[] = [Promise.resolve(result)]; + if (abortPromise) { + contenders.push(abortPromise); + } + if (cancelPromise) { + contenders.push(cancelPromise); + } + raceResult = await Promise.race(contenders); + } else { + raceResult = await result; } - if (cancelPromise) { - contenders.push(cancelPromise); + + if (raceResult === CANCEL) { + break; } - raceResult = await Promise.race(contenders); - } else { - raceResult = await result; - } - if (raceResult === CANCEL) { - break; + ({ value, done } = raceResult); } - - ({ value, done } = raceResult); } catch (err) { // The payload errored while being iterated. If the stream is still // writable, terminate the message with the `IGNORE` flag set so the @@ -357,13 +392,7 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n // aborted instead, the connection is being torn down and the // message is left unterminated. if (stream.writable && !signal?.aborted) { - const packet = new Packet(type); - packet.packetId(packetNumber += 1); - packet.resetConnection(resetConnection); - packet.last(true); - packet.ignore(true); - - await writePacket(packet); + await writePacket(buildPacket(0, baseStatus | STATUS.EOM | STATUS.IGNORE)); } throw err; @@ -376,31 +405,17 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n bl.append(value); while (!canceled && bl.length > length) { - const data = bl.slice(0, length); - bl.consume(length); - - const packet = new Packet(type); - packet.packetId(packetNumber += 1); - packet.resetConnection(resetConnection); - packet.addData(data); - - await writePacket(packet); + await writePacket(buildPacket(length, baseStatus)); } } // On cancellation, any buffered payload data is discarded and the final // packet is flagged so the server ignores the whole message. - const data = canceled ? Buffer.alloc(0) : bl.slice(); - bl.consume(bl.length); - - const packet = new Packet(type); - packet.packetId(packetNumber += 1); - packet.resetConnection(resetConnection); - packet.last(true); - packet.ignore(canceled); - packet.addData(data); - - await writePacket(packet); + if (canceled) { + await writePacket(buildPacket(0, baseStatus | STATUS.EOM | STATUS.IGNORE)); + } else { + await writePacket(buildPacket(bl.length, baseStatus | STATUS.EOM)); + } } finally { stream.removeListener('drain', onDrain); stream.removeListener('close', onDrain); @@ -525,11 +540,13 @@ export async function* readMessage(stream: Readable, options: { debug?: Debug, s const data = bl.slice(0, length); bl.consume(length); - const packet = new Packet(data); - debug?.packet('Received', packet); - debug?.data(packet); + if (debug) { + const packet = new Packet(data); + debug.packet('Received', packet); + debug.data(packet); + } - yield packet.data(); + yield data.subarray(HEADER_LENGTH); // Did the stream error out or close while we yielded? The events // have already fired, so the wait below would never settle. @@ -541,7 +558,7 @@ export async function* readMessage(stream: Readable, options: { debug?: Debug, s throw new Error('Premature close'); } - if (packet.isLast()) { + if (data[OFFSET.Status] & STATUS.EOM) { // This was the last packet of the message. Any data left in the // buffer belongs to the next message (e.g. the response to an // `ATTENTION` message sent by the client while reading an diff --git a/src/packet.ts b/src/packet.ts index c31500ea5..3a06cca74 100644 --- a/src/packet.ts +++ b/src/packet.ts @@ -21,7 +21,7 @@ for (const name in TYPE) { typeByValue[TYPE[name]] = name; } -const STATUS: { [key: string]: number } = { +export const STATUS: { [key: string]: number } = { NORMAL: 0x00, EOM: 0x01, IGNORE: 0x02, From 8efee798fa7e1ccf2149e9b16bbc5e86cd6067cd Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 21:38:03 +0000 Subject: [PATCH 07/14] fix: fail `writeMessage` cleanly when the stream closes mid-write An errorless stream close (e.g. a clean `destroy()` or a remote FIN) was wired to the same handler as `'drain'`, resolving a pending backpressure wait as if the stream had drained. Depending on where the close landed, `writeMessage` would either resolve successfully with an incompletely written message (close during the final packet's drain wait), or hang forever (close during an intermediate drain wait - writes to a destroyed stream report their error only to the absent write callback, so no `'error'` event ever fires). A close while waiting for more payload data was not observed at all. Stream errors and closes now reject a single persistent failure promise that is raced against every wait - drain waits and payload waits alike - so a stream failure surfaces as `Premature close` (or the stream error) no matter where the writer currently is, mirroring `readMessage`'s behavior. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 61 ++++++++++++----------- test/unit/message-io-test.ts | 96 ++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 30 deletions(-) diff --git a/src/message-io.ts b/src/message-io.ts index c6f496ef0..515f88ad5 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -244,12 +244,22 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } }; + // A promise that rejects when the stream errors out or is closed before + // the message was fully written. It is raced against every wait below, so + // that a stream failure can neither leave the writer stuck nor let it + // resolve successfully with an incompletely written message. + const { promise: failurePromise, reject: rejectWithFailure } = Promise.withResolvers(); + + // Prevent unhandled rejections if the stream fails while nothing is + // currently racing against `failurePromise`. + failurePromise.catch(() => {}); + const onError = (err: Error) => { - if (drain) { - const { reject } = drain; - drain = null; - reject(err); - } + rejectWithFailure(err); + }; + + const onClose = () => { + rejectWithFailure(new Error('Premature close')); }; let abortPromise: Promise | null = null; @@ -287,22 +297,18 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n const waitForDrain = () => { drain = Promise.withResolvers(); - if (abortPromise || cancelPromise) { - const contenders: Promise[] = [drain.promise]; - if (abortPromise) { - contenders.push(abortPromise); - } - if (cancelPromise) { - contenders.push(cancelPromise); - } - return Promise.race(contenders); + const contenders: Promise[] = [drain.promise, failurePromise]; + if (abortPromise) { + contenders.push(abortPromise); } - - return drain.promise; + if (cancelPromise) { + contenders.push(cancelPromise); + } + return Promise.race(contenders); }; stream.on('drain', onDrain); - stream.on('close', onDrain); + stream.on('close', onClose); stream.on('error', onError); try { @@ -365,19 +371,14 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } else { const result = (iterator as AsyncIterator).next(); - let raceResult: IteratorResult | typeof CANCEL; - if (abortPromise || cancelPromise) { - const contenders: Promise | typeof CANCEL>[] = [Promise.resolve(result)]; - if (abortPromise) { - contenders.push(abortPromise); - } - if (cancelPromise) { - contenders.push(cancelPromise); - } - raceResult = await Promise.race(contenders); - } else { - raceResult = await result; + const contenders: Promise | typeof CANCEL>[] = [Promise.resolve(result), failurePromise]; + if (abortPromise) { + contenders.push(abortPromise); } + if (cancelPromise) { + contenders.push(cancelPromise); + } + const raceResult = await Promise.race(contenders); if (raceResult === CANCEL) { break; @@ -418,7 +419,7 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } } finally { stream.removeListener('drain', onDrain); - stream.removeListener('close', onDrain); + stream.removeListener('close', onClose); stream.removeListener('error', onError); if (signal && onAbort) { diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index cbbe97eac..d4b3b9e7a 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -270,6 +270,102 @@ describe('writeMessage', function() { assert(hadError); }); + it('throws when the stream is closed while waiting for the final packet to drain', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new Duplex({ + write(chunk, encoding, callback) { + // never call callback so that the stream never drains + }, + read() {}, + + // instantly return false on write requests to indicate that the stream needs to drain + highWaterMark: 1 + }); + + setTimeout(() => { + assert(stream.writableNeedDrain); + + // Close the stream without an error. + stream.destroy(); + }, 20); + + let hadError = false; + try { + // A single small payload: the only drain wait is the one for the + // final packet, after which the message would be considered sent. + await writeMessage(stream, packetSize, packetType, [payload]); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'Premature close'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); + + it('throws when the stream is closed while waiting for an intermediate packet to drain', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new Duplex({ + write(chunk, encoding, callback) { + // never call callback so that the stream never drains + }, + read() {}, + + // instantly return false on write requests to indicate that the stream needs to drain + highWaterMark: 1 + }); + + setTimeout(() => { + assert(stream.writableNeedDrain); + + // Close the stream without an error. + stream.destroy(); + }, 20); + + let hadError = false; + try { + await writeMessage(stream, packetSize, packetType, [payload, payload, payload]); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'Premature close'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); + + it('throws when the stream is closed while waiting for more payload data', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new BufferListStream(); + + setTimeout(() => { + // Close the stream without an error. + stream.destroy(); + }, 20); + + let hadError = false; + try { + await writeMessage(stream, packetSize, packetType, (async function*() { + yield payload; + + // Stall forever, waiting for more data that never arrives. + await new Promise(() => {}); + })()); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'Premature close'); + } + + assert(hadError); + assertNoDanglingEventListeners(stream); + }); + it('terminates the message with the ignore flag set when the given cancel signal is aborted while writing', async function() { const payload = Buffer.from([1, 2, 3, 4, 5, 6]); const stream = new BufferListStream(); From 84f9b16f254d35c7b7895643fa18832375b4fbf4 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 21:40:28 +0000 Subject: [PATCH 08/14] fix: close the payload iterator when `writeMessage` exits early When a message write ended early - cancellation, teardown, or a stream failure - the payload iterator was abandoned without calling `.return()`, so `finally` blocks in generator payloads never ran. Harmless for plain array payloads, but streaming payloads (bulk load, large parameter payloads) rely on those blocks to release resources. The `.return()` call is deliberately not awaited: if the payload is suspended on a pending `next()`, the call is queued behind it and only settles once that read settles - awaiting it could block forever. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 19 +++++++++- test/unit/message-io-test.ts | 73 ++++++++++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/src/message-io.ts b/src/message-io.ts index 515f88ad5..fe86b2922 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -311,6 +311,9 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n stream.on('close', onClose); stream.on('error', onError); + let iterator: Iterator | AsyncIterator | null = null; + let payloadConsumed = false; + try { const bl = new BufferList(); const length = packetSize - HEADER_LENGTH; @@ -351,7 +354,6 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } }; - let iterator: Iterator | AsyncIterator; let isAsync; if ((payload as AsyncIterable)[Symbol.asyncIterator]) { isAsync = true; @@ -400,6 +402,7 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } if (done) { + payloadConsumed = true; break; } @@ -418,6 +421,20 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n await writePacket(buildPacket(bl.length, baseStatus | STATUS.EOM)); } } finally { + // If the payload was not fully consumed (cancellation, teardown, or a + // stream failure), close its iterator so `finally` blocks in generator + // payloads can release their resources. This is deliberately not + // awaited: if the payload is currently suspended on a pending `next()`, + // the `return()` call is queued behind it and only settles once that + // read settles - awaiting it here could block forever. + if (iterator && !payloadConsumed && iterator.return) { + try { + Promise.resolve(iterator.return()).catch(() => {}); + } catch { + // Ignore errors from closing the payload iterator. + } + } + stream.removeListener('drain', onDrain); stream.removeListener('close', onClose); stream.removeListener('error', onError); diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index d4b3b9e7a..e71887a2b 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -448,6 +448,79 @@ describe('writeMessage', function() { assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); }); + it('closes the payload iterator when the given cancel signal is aborted', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new BufferListStream(); + + const controller = new AbortController(); + setTimeout(() => { + controller.abort(); + }, 20); + + let cleanedUp = false; + + await writeMessage(stream, packetSize, packetType, (async function*() { + try { + yield payload; + + // Stall on something that eventually settles, so the deferred + // `.return()` call can be processed by the generator. + await delay(40); + + yield payload; + } finally { + cleanedUp = true; + } + })(), { cancelSignal: controller.signal }); + + // The message write finished at cancellation time, ... + assert.isFalse(cleanedUp); + + // ... while the payload iterator is closed once it becomes settleable. + await delay(40); + assert.isTrue(cleanedUp); + }); + + it('closes the payload iterator when the given signal is aborted', async function() { + const payload = Buffer.from([1, 2, 3]); + const stream = new BufferListStream(); + + const controller = new AbortController(); + setTimeout(() => { + controller.abort(new Error('teardown')); + }, 20); + + let cleanedUp = false; + + let hadError = false; + try { + await writeMessage(stream, packetSize, packetType, (async function*() { + try { + yield payload; + + // Stall on something that eventually settles, so the deferred + // `.return()` call can be processed by the generator. + await delay(40); + + yield payload; + } finally { + cleanedUp = true; + } + })(), { signal: controller.signal }); + } catch (err: any) { + hadError = true; + + assert.instanceOf(err, Error); + assert.strictEqual(err.message, 'teardown'); + } + + assert(hadError); + assert.isFalse(cleanedUp); + + await delay(40); + assert.isTrue(cleanedUp); + }); + it('prefers teardown over cancellation when both signals are aborted', async function() { const stream = new BufferListStream(); const signal = AbortSignal.abort(new Error('teardown')); From 9114842df7de51621d2211beb9892fd90eb5dd3e Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 21:40:43 +0000 Subject: [PATCH 09/14] test: cover the `resetConnection` option of `writeMessage` Asserts that the `RESETCONNECTION` status bit is set on every packet of the message, not just the first or last one, matching the behavior of `OutgoingMessageStream`. Co-Authored-By: Claude Fable 5 --- test/unit/message-io-test.ts | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index e71887a2b..557894b88 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -77,6 +77,21 @@ describe('writeMessage', function() { assert.deepEqual(secondPacket.data(), payload.subarray(4)); }); + it('sets the reset connection status bit on every packet', async function() { + const payload = Buffer.from([1, 2, 3, 4, 5, 6]); + const stream = new BufferListStream(); + + // `packetSize` allows for 4 bytes of data per packet + await writeMessage(stream, packetSize, packetType, [payload], { resetConnection: true }); + + const packets = splitPackets(stream.read()); + assert.lengthOf(packets, 2); + + for (const packet of packets) { + assert.include(packet.statusAsString(), 'RESETCONNECTION'); + } + }); + it('writes an empty final packet for an empty payload', async function() { const stream = new BufferListStream(); From 7a20de3427d5339d56b1ee6bb11c50ad21cce440 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 21:40:58 +0000 Subject: [PATCH 10/14] test: also check for dangling `close` listeners Both `writeMessage` and `readMessage` register a `close` listener, so the leak-check helper should assert its removal too. Co-Authored-By: Claude Fable 5 --- test/unit/message-io-test.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index 557894b88..557c3b9cb 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -23,6 +23,7 @@ function assertNoDanglingEventListeners(stream: Duplex) { assert.strictEqual(stream.listenerCount('error'), 0); assert.strictEqual(stream.listenerCount('drain'), 0); assert.strictEqual(stream.listenerCount('readable'), 0); + assert.strictEqual(stream.listenerCount('close'), 0); } function splitPackets(data: Buffer): Packet[] { From 98b9ea00ff47d1213750cc3f3b8551bddcc25127 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 21:41:16 +0000 Subject: [PATCH 11/14] docs: document `readMessage`'s consume-to-completion contract Stopping iteration before the message's last packet discards buffered bytes and leaves the TDS stream misaligned - and this is also the reason `readMessage` has no `cancelSignal` option: a canceled request still has to read the server's response to its end. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/message-io.ts b/src/message-io.ts index fe86b2922..3515126a4 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -459,6 +459,12 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n * Any bytes following the message's last packet (e.g. the start of the next * message) are pushed back onto the stream, to be consumed by the next read. * + * The generator must be consumed until the message's last packet: stopping + * iteration early discards any bytes buffered beyond the last yielded chunk + * and leaves the TDS stream misaligned. This is also why there is no + * `cancelSignal` option here - a canceled request still has to read the + * server's response to its end to keep the stream aligned. + * * If the given `signal` is aborted, reading stops and the signal's abort * reason is thrown. This also interrupts waiting for more data on a quiet * stream, which an external `.return()` or `.throw()` call can not do (it From 29d00cb69e9807f3b7450af5d6c7bd7acc328e95 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 21:46:36 +0000 Subject: [PATCH 12/14] refactor: decide cancellation from the `canceled` flag instead of a sentinel value The cancel promise resolved with a `CANCEL` symbol that the payload wait's race result was compared against. Cancellation now follows the same pattern as the other signals: the promise is only a wakeup for pending waits, while the `canceled` flag carries the state. `onCancel` sets the flag before resolving, so whenever the wakeup wins a race the flag is already observable - the same rule the drain wait and the loop conditions already used. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/message-io.ts b/src/message-io.ts index 3515126a4..098bc300b 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -277,19 +277,21 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n signal.addEventListener('abort', onAbort, { once: true }); } - const CANCEL = Symbol('cancel'); - + // Cancellation follows the same pattern as the other signals: the promise + // is only a wakeup for pending waits, while the `canceled` flag carries the + // state. `onCancel` sets the flag before resolving, so whenever the wakeup + // wins a race, the flag is already observable. let canceled = cancelSignal?.aborted ?? false; - let cancelPromise: Promise | null = null; + let cancelPromise: Promise | null = null; let onCancel: (() => void) | null = null; if (cancelSignal && !canceled) { - const { promise, resolve } = Promise.withResolvers(); + const { promise, resolve } = Promise.withResolvers(); cancelPromise = promise; onCancel = () => { canceled = true; - resolve(CANCEL); + resolve(); }; cancelSignal.addEventListener('abort', onCancel, { once: true }); } @@ -373,7 +375,7 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } else { const result = (iterator as AsyncIterator).next(); - const contenders: Promise | typeof CANCEL>[] = [Promise.resolve(result), failurePromise]; + const contenders: Promise[] = [Promise.resolve(result), failurePromise]; if (abortPromise) { contenders.push(abortPromise); } @@ -382,11 +384,17 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } const raceResult = await Promise.race(contenders); - if (raceResult === CANCEL) { + // Cancellation may have fired while we were waiting - either its + // wakeup won the race, or a chunk arrived in the same tick and + // would be discarded anyway. + if (canceled) { break; } - ({ value, done } = raceResult); + // The only resolving contenders are the iterator result and the + // cancel wakeup, and the latter implies `canceled` - so at this + // point, the race result is always an iterator result. + ({ value, done } = raceResult as IteratorResult); } } catch (err) { // The payload errored while being iterated. If the stream is still From 5fb7db70168aea9269d55276fb9223d7982a5559 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 21:54:55 +0000 Subject: [PATCH 13/14] refactor: tighten the types of the race contender arrays The drain wait only awaits settlement, so its contenders are `Promise`. The payload wait's race carries the iterator result, so typing its contenders as `Promise | void>` turns the subsequent assertion into a plain union narrowing. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/message-io.ts b/src/message-io.ts index 098bc300b..dd5a07df0 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -299,7 +299,7 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n const waitForDrain = () => { drain = Promise.withResolvers(); - const contenders: Promise[] = [drain.promise, failurePromise]; + const contenders: Promise[] = [drain.promise, failurePromise]; if (abortPromise) { contenders.push(abortPromise); } @@ -375,7 +375,7 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n } else { const result = (iterator as AsyncIterator).next(); - const contenders: Promise[] = [Promise.resolve(result), failurePromise]; + const contenders: Promise | void>[] = [Promise.resolve(result), failurePromise]; if (abortPromise) { contenders.push(abortPromise); } From a1e08bb70d63f1addb968fd888ed90d2f0838c86 Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Sun, 19 Jul 2026 22:01:07 +0000 Subject: [PATCH 14/14] fix: skip drain waits when the cancel signal was already aborted at call time A cancellation that happened mid-write skips pending drain waits via its settled wakeup promise, but a cancel signal that was already aborted when `writeMessage` was called never created that promise - so the `IGNORE` terminator's drain wait could block on a stalled peer until teardown. Use a pre-settled wakeup in that case, making both cancellation paths behave identically. Co-Authored-By: Claude Fable 5 --- src/message-io.ts | 22 ++++++++++++++-------- test/unit/message-io-test.ts | 20 ++++++++++++++++++++ 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/src/message-io.ts b/src/message-io.ts index dd5a07df0..af16e36c0 100644 --- a/src/message-io.ts +++ b/src/message-io.ts @@ -285,15 +285,21 @@ export async function writeMessage(stream: Writable, packetSize: number, type: n let cancelPromise: Promise | null = null; let onCancel: (() => void) | null = null; - if (cancelSignal && !canceled) { - const { promise, resolve } = Promise.withResolvers(); + if (cancelSignal) { + if (canceled) { + // Already canceled when called: use a settled wakeup so that waits + // behave the same as when the cancellation happens mid-write. + cancelPromise = Promise.resolve(); + } else { + const { promise, resolve } = Promise.withResolvers(); - cancelPromise = promise; - onCancel = () => { - canceled = true; - resolve(); - }; - cancelSignal.addEventListener('abort', onCancel, { once: true }); + cancelPromise = promise; + onCancel = () => { + canceled = true; + resolve(); + }; + cancelSignal.addEventListener('abort', onCancel, { once: true }); + } } const waitForDrain = () => { diff --git a/test/unit/message-io-test.ts b/test/unit/message-io-test.ts index 557c3b9cb..094ef0373 100644 --- a/test/unit/message-io-test.ts +++ b/test/unit/message-io-test.ts @@ -464,6 +464,26 @@ describe('writeMessage', function() { assert.lengthOf(getEventListeners(controller.signal, 'abort'), 0); }); + it('does not wait for the stream to drain when the given cancel signal is already aborted', async function() { + const stream = new Duplex({ + write(chunk, encoding, callback) { + // never call callback so that the stream never drains + }, + read() {}, + + // instantly return false on write requests to indicate that the stream needs to drain + highWaterMark: 1 + }); + + const cancelSignal = AbortSignal.abort(); + + // Resolves normally even though the terminator write is backpressured: + // nothing further will be written for this message. + await writeMessage(stream, packetSize, packetType, [Buffer.from([1, 2, 3])], { cancelSignal }); + + assertNoDanglingEventListeners(stream); + }); + it('closes the payload iterator when the given cancel signal is aborted', async function() { const payload = Buffer.from([1, 2, 3]); const stream = new BufferListStream();