fix(NODE-4834): ensure that MessageStream is destroyed when connections are destroyed

src/cmap/connection.ts

@@ -170,7 +170,9 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
   address: string;
   socketTimeoutMS: number;
   monitorCommands: boolean;
+  /** Indicates that the connection (including underlying TCP socket) has been closed. */
   closed: boolean;
+  /** Indicates that the connection has been explicitly destroyed. In the destroyed state, the connection is also closed */
   destroyed: boolean;
   lastHelloMS?: number;
   serverApi?: ServerApi;
@@ -309,11 +311,18 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     this[kLastUseTime] = now();
   }
 
+  /**
+   * @remarks
+   * onError is called when an error is propagated up from the socket behind
+   * kStream to kMessageStream. This occurs prior to the closing of the socket,
+   * so the resource must be return to the operating system here.
+   */
   onError(error: Error) {
-    if (this.closed) {
+    if (this.closed || this.destroyed) {
       return;
     }
 
+    this[kMessageStream].destroy(error);
     this[kStream].destroy(error);
 
     this.closed = true;
@@ -326,10 +335,16 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     this.emit(Connection.CLOSE);
   }
 
+  /**
+   * @remarks
+   * onClose is called when the socket underlying kStream has been closed and
+   * the resource is returned to the operating system
+   */
   onClose() {
-    if (this.closed) {
+    if (this.closed || this.destroyed) {
       return;
     }
+    this[kMessageStream].destroy();
 
     this.closed = true;
 
@@ -342,12 +357,18 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     this.emit(Connection.CLOSE);
   }
 
+  /**
+   * @remarks
+   * onTimeout is called when the tcp socket underlying kStream times out. This
+   * occurs prior to the closing of the socket, so the resource must be returned
+   * to the operating system here */
   onTimeout() {
-    if (this.closed) {
+    if (this.closed || this.destroyed) {
       return;
     }
 
     this[kDelayedTimeoutId] = setTimeout(() => {
+      this[kMessageStream].destroy();
       this[kStream].destroy();
 
       this.closed = true;
@@ -468,32 +489,28 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     this.removeAllListeners(Connection.PINNED);
     this.removeAllListeners(Connection.UNPINNED);
 
-    options = Object.assign({ force: false }, options);
-    if (this[kStream] == null || this.destroyed) {
-      this.destroyed = true;
-      if (typeof callback === 'function') {
-        callback();
-      }
+    this[kMessageStream].destroy();
 
-      return;
+    options = Object.assign({ force: false }, options);
+    if (this.destroyed) {
+      return callback?.();
     }
 
     if (options.force) {
       this[kStream].destroy();
       this.destroyed = true;
-      if (typeof callback === 'function') {
-        callback();
-      }
-
-      return;
+      return callback?.();
     }
 
-    this[kStream].end(() => {
+    if (this[kStream].writableEnded) {
       this.destroyed = true;
-      if (typeof callback === 'function') {
-        callback();
-      }
-    });
+      return callback?.();
+    } else {
+      this[kStream].end(() => {
+        this.destroyed = true;
+        callback?.();
+      });
+    }
   }
 
   command(

test/unit/cmap/connection.test.ts

@@ -26,6 +26,7 @@ const connectionOptionsDefaults = {
 
 /** The absolute minimum socket API needed by Connection as of writing this test */
 class FakeSocket extends EventEmitter {
+  writableEnded: boolean;
   address() {
     // is never called
   }
@@ -34,6 +35,11 @@ class FakeSocket extends EventEmitter {
   }
   destroy() {
     // is called, has no side effects
+    this.writableEnded = true;
+  }
+  end(cb) {
+    // nextTick to simulate I/O delay
+    process.nextTick(cb);
   }
   get remoteAddress() {
     return 'iLoveJavaScript';
@@ -389,7 +395,7 @@ describe('new Connection()', function () {
       connection = sinon.spy(new Connection(driverSocket, connectionOptionsDefaults));
       const messageStreamSymbol = getSymbolFrom(connection, 'messageStream');
       kDelayedTimeoutId = getSymbolFrom(connection, 'delayedTimeoutId');
-      messageStream = connection[messageStreamSymbol];
+      messageStream = sinon.spy(connection[messageStreamSymbol]);
     });
 
     afterEach(() => {
@@ -433,6 +439,85 @@ describe('new Connection()', function () {
       expect(connection).to.have.property('closed', false);
       expect(connection).to.have.property(kDelayedTimeoutId, null);
     });
+
+    it('destroys the message stream and socket', () => {
+      expect(connection).to.have.property(kDelayedTimeoutId, null);
+
+      driverSocket.emit('timeout');
+
+      clock.tick(1);
+
+      expect(connection.onTimeout).to.have.been.calledOnce;
+      expect(connection).to.have.property(kDelayedTimeoutId).that.is.instanceOf(NodeJSTimeoutClass);
+
+      expect(messageStream.destroy).to.have.been.calledOnce;
+      expect(driverSocket.destroy).to.have.been.calledOnce;
+    });
+  });
+
+  describe('onError()', () => {
+    let connection: sinon.SinonSpiedInstance<Connection>;
+    let clock: sinon.SinonFakeTimers;
+    let timerSandbox: sinon.SinonFakeTimers;
+    let driverSocket: sinon.SinonSpiedInstance<FakeSocket>;
+    let messageStream: MessageStream;
+    beforeEach(() => {
+      timerSandbox = createTimerSandbox();
+      clock = sinon.useFakeTimers();
+      driverSocket = sinon.spy(new FakeSocket());
+      // @ts-expect-error: driverSocket does not fully satisfy the stream type, but that's okay
+      connection = sinon.spy(new Connection(driverSocket, connectionOptionsDefaults));
+      const messageStreamSymbol = getSymbolFrom(connection, 'messageStream');
+      messageStream = sinon.spy(connection[messageStreamSymbol]);
+    });
+
+    afterEach(() => {
+      timerSandbox.restore();
+      clock.restore();
+    });
+
+    it('destroys the message stream and socket', () => {
+      messageStream.emit('error');
+
+      clock.tick(1);
+
+      expect(connection.onError).to.have.been.calledOnce;
+
+      expect(messageStream.destroy).to.have.been.calledOnce;
+      expect(driverSocket.destroy).to.have.been.calledOnce;
+    });
+  });
+
+  describe('onClose()', () => {
+    let connection: sinon.SinonSpiedInstance<Connection>;
+    let clock: sinon.SinonFakeTimers;
+    let timerSandbox: sinon.SinonFakeTimers;
+    let driverSocket: sinon.SinonSpiedInstance<FakeSocket>;
+    let messageStream: MessageStream;
+    beforeEach(() => {
+      timerSandbox = createTimerSandbox();
+      clock = sinon.useFakeTimers();
+
+      driverSocket = sinon.spy(new FakeSocket());
+      // @ts-expect-error: driverSocket does not fully satisfy the stream type, but that's okay
+      connection = sinon.spy(new Connection(driverSocket, connectionOptionsDefaults));
+      const messageStreamSymbol = getSymbolFrom(connection, 'messageStream');
+      messageStream = sinon.spy(connection[messageStreamSymbol]);
+    });
+
+    afterEach(() => {
+      timerSandbox.restore();
+      clock.restore();
+    });
+
+    it('destroys the message stream', () => {
+      driverSocket.emit('close');
+
+      clock.tick(1);
+
+      expect(connection.onClose).to.have.been.calledOnce;
+      expect(messageStream.destroy).to.have.been.calledOnce;
+    });
   });
 
   describe('.hasSessionSupport', function () {
@@ -486,4 +571,44 @@ describe('new Connection()', function () {
       });
     });
   });
+
+  describe('destroy()', () => {
+    let connection: sinon.SinonSpiedInstance<Connection>;
+    let clock: sinon.SinonFakeTimers;
+    let timerSandbox: sinon.SinonFakeTimers;
+    let driverSocket: sinon.SinonSpiedInstance<FakeSocket>;
+    let messageStream: MessageStream;
+    beforeEach(() => {
+      timerSandbox = createTimerSandbox();
+      clock = sinon.useFakeTimers();
+
+      driverSocket = sinon.spy(new FakeSocket());
+      // @ts-expect-error: driverSocket does not fully satisfy the stream type, but that's okay
+      connection = sinon.spy(new Connection(driverSocket, connectionOptionsDefaults));
+      const messageStreamSymbol = getSymbolFrom(connection, 'messageStream');
+      messageStream = sinon.spy(connection[messageStreamSymbol]);
+    });
+
+    afterEach(() => {
+      timerSandbox.restore();
+      clock.restore();
+    });
+
+    it('ends the tcp socket and destroys the messageStream', () => {
+      connection.destroy();
+      clock.tick(1);
+      expect(messageStream.destroy).to.have.been.calledOnce;
+      expect(driverSocket.end).to.have.been.calledOnce;
+    });
+
+    it('does not call stream.end after onClose, onTimeout, or onError', () => {
+      messageStream.emit('error');
+      clock.tick(1);
+      expect(connection.onError).to.have.been.calledOnce;
+      expect(driverSocket.destroy).to.have.been.calledOnce;
+      connection.destroy();
+      clock.tick(1);
+      expect(driverSocket.end).to.not.have.been.called;
+    });
+  });
 });