From 79335aa54c2b5192cc83b1af6948f385f9557dbd Mon Sep 17 00:00:00 2001 From: Arthur Schreiber Date: Fri, 17 Jul 2026 16:31:45 +0000 Subject: [PATCH 1/4] 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 2/4] 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 3/4] 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 4/4] 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'));