feat(NODE-4139): streaming protocol message changes

src/cmap/connection.ts

@@ -293,6 +293,16 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     this[kHello] = response;
   }
 
+  // Set the whether the message stream is for a monitoring connection using the
+  // streaming protocol.
+  set isStreamingProtocol(value: boolean) {
+    this[kMessageStream].isStreamingProtocol = value;
+  }
+
+  get isStreamingProtocol(): boolean {
+    return this[kMessageStream].isStreamingProtocol;
+  }
+
   get serviceId(): ObjectId | undefined {
     return this.hello?.serviceId;
   }

src/cmap/message_stream.ts

@@ -53,13 +53,19 @@ export class MessageStream extends Duplex {
   maxBsonMessageSize: number;
   /** @internal */
   [kBuffer]: BufferPool;
+  /** @internal */
+  isStreamingProtocol = false;
 
   constructor(options: MessageStreamOptions = {}) {
     super(options);
     this.maxBsonMessageSize = options.maxBsonMessageSize || kDefaultMaxBsonMessageSize;
     this[kBuffer] = new BufferPool();
   }
 
+  get buffer(): BufferPool {
+    return this[kBuffer];
+  }
+
   override _write(chunk: Buffer, _: unknown, callback: Callback<Buffer>): void {
     this[kBuffer].append(chunk);
     processIncomingData(this, callback);
@@ -165,12 +171,20 @@ function processIncomingData(stream: MessageStream, callback: Callback<Buffer>)
   let ResponseType = messageHeader.opCode === OP_MSG ? BinMsg : Response;
   if (messageHeader.opCode !== OP_COMPRESSED) {
     const messageBody = message.slice(MESSAGE_HEADER_SIZE);
-    stream.emit('message', new ResponseType(message, messageHeader, messageBody));
 
-    if (buffer.length >= 4) {
+    // If we are a monitoring message stream using the streaming protocol and
+    // there is more in the buffer that can be read, skip processing since we
+    // want the last hello command response that is in the buffer.
+    if (stream.isStreamingProtocol && buffer.length >= 4) {
       processIncomingData(stream, callback);
     } else {
-      callback();
+      stream.emit('message', new ResponseType(message, messageHeader, messageBody));
+
+      if (buffer.length >= 4) {
+        processIncomingData(stream, callback);
+      } else {
+        callback();
+      }
     }
 
     return;
@@ -198,12 +212,19 @@ function processIncomingData(stream: MessageStream, callback: Callback<Buffer>)
       return;
     }
 
-    stream.emit('message', new ResponseType(message, messageHeader, messageBody));
-
-    if (buffer.length >= 4) {
+    // If we are a monitoring message stream using the streaming protocol and
+    // there is more in the buffer that can be read, skip processing since we
+    // want the last hello command response that is in the buffer.
+    if (stream.isStreamingProtocol && buffer.length >= 4) {
       processIncomingData(stream, callback);
     } else {
-      callback();
+      stream.emit('message', new ResponseType(message, messageHeader, messageBody));
+
+      if (buffer.length >= 4) {
+        processIncomingData(stream, callback);
+      } else {
+        callback();
+      }
     }
   });
 }

src/sdam/monitor.ts

@@ -279,6 +279,10 @@ function checkServer(monitor: Monitor, callback: Callback<Document | null>) {
       // if we are using the streaming protocol then we immediately issue another `started`
       // event, otherwise the "check" is complete and return to the main monitor loop
       if (isAwaitable && hello.topologyVersion) {
+        // Tell the connection that we are using the streaming protocol so that the
+        // connection's message stream will only read the last hello on the buffer.
+        connection.isStreamingProtocol = true;
+
         monitor.emit(
           Server.SERVER_HEARTBEAT_STARTED,
           new ServerHeartbeatStartedEvent(monitor.address)

test/tools/utils.ts

@@ -1,7 +1,10 @@
 import { EJSON } from 'bson';
+import * as BSON from 'bson';
 import { expect } from 'chai';
 import { inspect, promisify } from 'util';
 
+import { OP_MSG } from '../../src/cmap/wire_protocol/constants';
+import { Document } from '../../src/index';
 import { Logger } from '../../src/logger';
 import { deprecateOptions, DeprecateOptionsConfig } from '../../src/utils';
 import { runUnifiedSuite } from './unified-spec-runner/runner';
@@ -343,6 +346,24 @@ export class TestBuilder {
   }
 }
 
+export function generateOpMsgBuffer(document: Document): Buffer {
+  const header = Buffer.alloc(4 * 4 + 4);
+
+  const typeBuffer = Buffer.alloc(1);
+  typeBuffer[0] = 0;
+
+  const docBuffer = BSON.serialize(document);
+
+  const totalLength = header.length + typeBuffer.length + docBuffer.length;
+
+  header.writeInt32LE(totalLength, 0);
+  header.writeInt32LE(0, 4);
+  header.writeInt32LE(0, 8);
+  header.writeInt32LE(OP_MSG, 12);
+  header.writeUInt32LE(0, 16);
+  return Buffer.concat([header, typeBuffer, docBuffer]);
+}
+
 export class UnifiedTestSuiteBuilder {
   private _description = 'Default Description';
   private _schemaVersion = '1.0';

test/unit/cmap/message_stream.test.js

@@ -5,6 +5,7 @@ const { MessageStream } = require('../../../src/cmap/message_stream');
 const { Msg } = require('../../../src/cmap/commands');
 const expect = require('chai').expect;
 const { LEGACY_HELLO_COMMAND } = require('../../../src/constants');
+const { generateOpMsgBuffer } = require('../../tools/utils');
 
 function bufferToStream(buffer) {
   const stream = new Readable();
@@ -18,117 +19,103 @@ function bufferToStream(buffer) {
   return stream;
 }
 
-describe('Message Stream', function () {
-  describe('reading', function () {
-    [
-      {
-        description: 'valid OP_REPLY',
-        data: Buffer.from(
-          '370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000',
-          'hex'
-        ),
-        documents: [{ [LEGACY_HELLO_COMMAND]: 1 }]
-      },
-      {
-        description: 'valid multiple OP_REPLY',
-        expectedMessageCount: 4,
-        data: Buffer.from(
-          '370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000' +
-            '370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000' +
-            '370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000' +
-            '370000000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000',
-          'hex'
-        ),
-        documents: [{ [LEGACY_HELLO_COMMAND]: 1 }]
-      },
-      {
-        description: 'valid OP_REPLY (partial)',
-        data: [
-          Buffer.from('37', 'hex'),
-          Buffer.from('0000', 'hex'),
-          Buffer.from(
-            '000100000001000000010000000000000000000000000000000000000001000000130000001069736d6173746572000100000000',
-            'hex'
-          )
-        ],
-        documents: [{ [LEGACY_HELLO_COMMAND]: 1 }]
-      },
-
-      {
-        description: 'valid OP_MSG',
-        data: Buffer.from(
-          '370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000',
-          'hex'
-        ),
-        documents: [{ $db: 'admin', [LEGACY_HELLO_COMMAND]: 1 }]
-      },
-      {
-        description: 'valid multiple OP_MSG',
-        expectedMessageCount: 4,
-        data: Buffer.from(
-          '370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000' +
-            '370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000' +
-            '370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000' +
-            '370000000100000000000000dd0700000000000000220000001069736d6173746572000100000002246462000600000061646d696e0000',
-          'hex'
-        ),
-        documents: [{ $db: 'admin', [LEGACY_HELLO_COMMAND]: 1 }]
-      },
-
-      {
-        description: 'Invalid message size (negative)',
-        data: Buffer.from('ffffffff', 'hex'),
-        error: 'Invalid message size: -1'
-      },
-      {
-        description: 'Invalid message size (exceeds maximum)',
-        data: Buffer.from('01000004', 'hex'),
-        error: 'Invalid message size: 67108865, max allowed: 67108864'
-      }
-    ].forEach(test => {
-      it(test.description, function (done) {
-        const error = test.error;
-        const expectedMessageCount = test.expectedMessageCount || 1;
-        const inputStream = bufferToStream(test.data);
+describe('MessageStream', function () {
+  context('when the stream uses the streaming protocol', function () {
+    const response = { isWritablePrimary: true };
+    let firstHello;
+    let secondHello;
+    let thirdHello;
+
+    beforeEach(function () {
+      firstHello = generateOpMsgBuffer(response);
+      secondHello = generateOpMsgBuffer(response);
+      thirdHello = generateOpMsgBuffer(response);
+    });
+
+    it('only reads the last message in the buffer', function (done) {
+      const inputStream = bufferToStream(Buffer.concat([firstHello, secondHello, thirdHello]));
+      const messageStream = new MessageStream();
+      messageStream.isStreamingProtocol = true;
+
+      messageStream.once('message', msg => {
+        msg.parse();
+        expect(msg).to.have.property('documents').that.deep.equals([response]);
+        // Make sure there is nothing left in the buffer.
+        expect(messageStream.buffer.length).to.equal(0);
+        done();
+      });
+
+      inputStream.pipe(messageStream);
+    });
+  });
+
+  context('when the stream is not using the streaming protocol', function () {
+    context('when the messages are valid', function () {
+      const response = { isWritablePrimary: true };
+      let firstHello;
+      let secondHello;
+      let thirdHello;
+      let messageCount = 0;
+
+      beforeEach(function () {
+        firstHello = generateOpMsgBuffer(response);
+        secondHello = generateOpMsgBuffer(response);
+        thirdHello = generateOpMsgBuffer(response);
+      });
+
+      it('reads all messages in the buffer', function (done) {
+        const inputStream = bufferToStream(Buffer.concat([firstHello, secondHello, thirdHello]));
         const messageStream = new MessageStream();
 
-        let messageCount = 0;
         messageStream.on('message', msg => {
           messageCount++;
-          if (error) {
-            done(new Error(`expected error: ${error}`));
-            return;
+          msg.parse();
+          expect(msg).to.have.property('documents').that.deep.equals([response]);
+          // Test will not complete until 3 messages processed.
+          if (messageCount === 3) {
+            done();
           }
+        });
 
-          msg.parse();
+        inputStream.pipe(messageStream);
+      });
+    });
 
-          if (test.documents) {
-            expect(msg).to.have.property('documents').that.deep.equals(test.documents);
-          }
+    context('when the messages are invalid', function () {
+      context('when the message size is negative', function () {
+        it('emits an error', function (done) {
+          const inputStream = bufferToStream(Buffer.from('ffffffff', 'hex'));
+          const messageStream = new MessageStream();
 
-          if (messageCount === expectedMessageCount) {
+          messageStream.on('error', err => {
+            expect(err).to.have.property('message').that.equals('Invalid message size: -1');
             done();
-          }
+          });
+
+          inputStream.pipe(messageStream);
         });
+      });
 
-        messageStream.on('error', err => {
-          if (error == null) {
-            done(err);
-            return;
-          }
+      context('when the message size exceeds the bson maximum', function () {
+        it('emits an error', function (done) {
+          const inputStream = bufferToStream(Buffer.from('01000004', 'hex'));
+          const messageStream = new MessageStream();
 
-          expect(err).to.have.property('message').that.equals(error);
+          messageStream.on('error', err => {
+            expect(err)
+              .to.have.property('message')
+              .that.equals('Invalid message size: 67108865, max allowed: 67108864');
+            done();
+          });
 
-          done();
+          inputStream.pipe(messageStream);
         });
-
-        inputStream.pipe(messageStream);
       });
     });
   });
 
-  describe('writing', function () {
-    it('should write a message to the stream', function (done) {
+  context('when writing to the message stream', function () {
+    it('pushes the message', function (done) {
       const readableStream = new Readable({ read() {} });
       const writeableStream = new Writable({
         write: (chunk, _, callback) => {