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;
+  /** True when after of onClose, onError, or onTimeout was called */
   closed: boolean;
+  /** True only after destroy has been called */
   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 resouce 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,6 +489,8 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
     this.removeAllListeners(Connection.PINNED);
     this.removeAllListeners(Connection.UNPINNED);
 
+    this[kMessageStream].destroy();
+
     options = Object.assign({ force: false }, options);
     if (this[kStream] == null || this.destroyed) {
       this.destroyed = true;
@@ -488,12 +511,19 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
       return;
     }
 
-    this[kStream].end(() => {
+    if (this[kStream].destroyed) {
       this.destroyed = true;
       if (typeof callback === 'function') {
         callback();
       }
-    });
+    } else {
+      this[kStream].end(() => {
+        this.destroyed = true;
+        if (typeof callback === 'function') {
+          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 {
+  destroyed: boolean;
   address() {
     // is never called
   }
@@ -34,6 +35,10 @@ class FakeSocket extends EventEmitter {
   }
   destroy() {
     // is called, has no side effects
+    this.destroyed = true;
+  }
+  end(cb) {
+    cb();
   }
   get remoteAddress() {
     return 'iLoveJavaScript';
@@ -389,7 +394,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 +438,85 @@ describe('new Connection()', function () {
       expect(connection).to.have.property('closed', false);
       expect(connection).to.have.property(kDelayedTimeoutId, null);
     });
+
+    it('should destroy 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('should destroy 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('should destroy 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 +570,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('should end tcp socket and destroy messageStream', () => {
+      connection.destroy();
+      clock.tick(1);
+      expect(messageStream.destroy).to.have.been.calledOnce;
+      expect(driverSocket.end).to.have.been.calledOnce;
+    });
+
+    it('should 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;
+    });
+  });
 });