refactor: add MessageIO.writeMessage/readMessage and use them for the PRELOGIN exchange

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);
+  })();
+}

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);
+}

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 { readMessage } = 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 readMessage(stream, { debug })) {
+        total += chunk.length;
+      }
+    }
+
+    bench.end(n);
+  })();
+}

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 { writeMessage } = 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 writeMessage(stream, 8 + 1024, 2, payload, { debug });
+    }
+
+    bench.end(n);
+  })();
+}

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';
@@ -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, signal);
+          } 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, 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
@@ -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 writeMessage(socket, this.config.options.packetSize, TYPE.PRELOGIN, [payload.data], { debug: this.debug, signal });
     this.debug.payload(function() {
       return payload.toString('  ');
     });
@@ -3365,37 +3380,12 @@ class Connection extends EventEmitter {
     });
   }
 
-  async readPreloginResponse(signal: AbortSignal): Promise<PreloginPayload> {
+  async readPreloginResponse(socket: net.Socket, signal: AbortSignal): Promise<PreloginPayload> {
     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 = message[Symbol.asyncIterator]();
-      try {
-        while (true) {
-          const { done, value } = await Promise.race([
-            iterator.next(),
-            signalAborted
-          ]);
-
-          if (done) {
-            break;
-          }
-
-          messageBuffer = Buffer.concat([messageBuffer, value]);
-        }
-      } finally {
-        if (iterator.return) {
-          await iterator.return();
-        }
-      }
-    });
+    for await (const data of readMessage(socket, { debug: this.debug, signal })) {
+      messageBuffer = Buffer.concat([messageBuffer, data]);
+    }
 
     const preloginPayload = new PreloginPayload(messageBuffer);
     this.debug.payload(function() {