From 9bc400dbaf9989ffa9b21feacda8a6646e078597 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Cloux?= <clement.cloux@ynov.com>
Date: Mon, 13 Mar 2023 21:41:17 +0100
Subject: [PATCH] fix(NODE-5106): refactor connect locking and add test

---
 src/mongo_client.ts                   | 97 +++++++++++++++------------
 test/integration/mongo_client.test.ts | 50 ++++++++++++++
 2 files changed, 103 insertions(+), 44 deletions(-)
 create mode 100644 test/integration/mongo_client.test.ts

diff --git a/src/mongo_client.ts b/src/mongo_client.ts
index fa286b2efa..97d6e9232e 100644
--- a/src/mongo_client.ts
+++ b/src/mongo_client.ts
@@ -407,63 +407,72 @@ export class MongoClient extends TypedEventEmitter<MongoClientEvents> {
    * @see docs.mongodb.org/manual/reference/connection-string/
    */
   async connect(): Promise<this> {
-    if (this.topology && this.topology.isConnected()) {
-      return this;
-    }
-
     if (this.connectionLock) {
       return this.connectionLock;
     }
 
-    this.connectionLock = (async () => {
-      const options = this[kOptions];
+    try {
+      this.connectionLock = this._connect();
+      await this.connectionLock;
+    } finally {
+      // release
+      this.connectionLock = undefined;
+    }
 
-      if (typeof options.srvHost === 'string') {
-        const hosts = await resolveSRVRecord(options);
+    return this;
+  }
 
-        for (const [index, host] of hosts.entries()) {
-          options.hosts[index] = host;
-        }
-      }
+  /**
+   * Create a topology to open the connection, must be locked to avoid topology leaks in concurrency scenario.
+   * Locking is enforced by the connect method.
+   *
+   * When decorators available put implementation back to original connect methods
+   * and enforce locking via a dedicated decorator.
+   * @see https://github.com/mongodb/node-mongodb-native/pull/3596#discussion_r1134211500
+   */
+  private async _connect(): Promise<this> {
+    if (this.topology && this.topology.isConnected()) {
+      return this;
+    }
 
-      const topology = new Topology(options.hosts, options);
-      // Events can be emitted before initialization is complete so we have to
-      // save the reference to the topology on the client ASAP if the event handlers need to access it
-      this.topology = topology;
-      topology.client = this;
+    const options = this[kOptions];
 
-      topology.once(Topology.OPEN, () => this.emit('open', this));
+    if (typeof options.srvHost === 'string') {
+      const hosts = await resolveSRVRecord(options);
 
-      for (const event of MONGO_CLIENT_EVENTS) {
-        topology.on(event, (...args: any[]) => this.emit(event, ...(args as any)));
+      for (const [index, host] of hosts.entries()) {
+        options.hosts[index] = host;
       }
+    }
 
-      const topologyConnect = async () => {
-        try {
-          await promisify(callback => topology.connect(options, callback))();
-        } catch (error) {
-          topology.close({ force: true });
-          throw error;
-        }
-      };
-
-      if (this.autoEncrypter) {
-        const initAutoEncrypter = promisify(callback => this.autoEncrypter?.init(callback));
-        await initAutoEncrypter();
-        await topologyConnect();
-        await options.encrypter.connectInternalClient();
-      } else {
-        await topologyConnect();
-      }
+    const topology = new Topology(options.hosts, options);
+    // Events can be emitted before initialization is complete so we have to
+    // save the reference to the topology on the client ASAP if the event handlers need to access it
+    this.topology = topology;
+    topology.client = this;
 
-      return this;
-    })();
+    topology.once(Topology.OPEN, () => this.emit('open', this));
 
-    try {
-      await this.connectionLock;
-    } finally {
-      // release
-      this.connectionLock = undefined;
+    for (const event of MONGO_CLIENT_EVENTS) {
+      topology.on(event, (...args: any[]) => this.emit(event, ...(args as any)));
+    }
+
+    const topologyConnect = async () => {
+      try {
+        await promisify(callback => topology.connect(options, callback))();
+      } catch (error) {
+        topology.close({ force: true });
+        throw error;
+      }
+    };
+
+    if (this.autoEncrypter) {
+      const initAutoEncrypter = promisify(callback => this.autoEncrypter?.init(callback));
+      await initAutoEncrypter();
+      await topologyConnect();
+      await options.encrypter.connectInternalClient();
+    } else {
+      await topologyConnect();
     }
 
     return this;
diff --git a/test/integration/mongo_client.test.ts b/test/integration/mongo_client.test.ts
new file mode 100644
index 0000000000..7de30e161f
--- /dev/null
+++ b/test/integration/mongo_client.test.ts
@@ -0,0 +1,50 @@
+import { expect } from 'chai';
+import * as sinon from 'sinon';
+
+import { MongoClient } from '../../src';
+
+describe('MongoClient', () => {
+  let client: MongoClient;
+  let topologyOpenEvents;
+
+  beforeEach(async function () {
+    client = this.configuration.newClient();
+    topologyOpenEvents = [];
+    client.on('open', event => topologyOpenEvents.push(event));
+  });
+
+  afterEach(async function () {
+    await client.close();
+  });
+
+  it('Concurrents client connect correctly locked (only one topology created)', async function () {
+    await Promise.all([client.connect(), client.connect(), client.connect()]);
+
+    expect(topologyOpenEvents).to.have.lengthOf(1);
+    expect(client.topology?.isConnected()).to.be.true;
+  });
+
+  it('Failed client connect must properly release lock', async function () {
+    const internalConnectStub = sinon.stub(client, '_connect' as keyof MongoClient);
+    internalConnectStub.onFirstCall().rejects();
+
+    // first call rejected to simulate a connection failure
+    try {
+      await client.connect();
+    } catch (err) {
+      expect(err).to.exist;
+    }
+
+    internalConnectStub.restore();
+
+    // second call should connect
+    try {
+      await client.connect();
+    } catch (err) {
+      expect.fail(`client connect throwed unexpected error`);
+    }
+
+    expect(topologyOpenEvents).to.have.lengthOf(1);
+    expect(client.topology?.isConnected()).to.be.true;
+  });
+});