test(NODE-4262): simplify leak checker for startSession fixes

test/integration/client-side-encryption/client_side_encryption.prose.test.js

@@ -1,12 +1,13 @@
 'use strict';
 const BSON = require('bson');
-const chai = require('chai');
+const { expect } = require('chai');
 const fs = require('fs');
 const path = require('path');
+
 const { deadlockTests } = require('./client_side_encryption.prose.deadlock');
+const { dropCollection, APMEventCollector } = require('../shared');
 
-const expect = chai.expect;
-chai.use(require('chai-subset'));
+const { EJSON } = BSON;
 const { LEGACY_HELLO_COMMAND } = require('../../../src/constants');
 
 const getKmsProviders = (localKey, kmipEndpoint, azureEndpoint, gcpEndpoint) => {
@@ -53,10 +54,6 @@ describe('Client Side Encryption Prose Tests', metadata, function () {
   const keyVaultCollName = 'datakeys';
   const keyVaultNamespace = `${keyVaultDbName}.${keyVaultCollName}`;
 
-  const shared = require('../shared');
-  const dropCollection = shared.dropCollection;
-  const APMEventCollector = shared.APMEventCollector;
-
   const LOCAL_KEY = Buffer.from(
     'Mng0NCt4ZHVUYUJCa1kxNkVyNUR1QURhZ2h2UzR2d2RrZzh0cFBwM3R6NmdWMDFBMUN3YkQ5aXRRMkhGRGdQV09wOGVNYUMxT2k3NjZKelhaQmRCZGJkTXVyZG9uSjFk',
     'base64'
@@ -337,9 +334,6 @@ describe('Client Side Encryption Prose Tests', metadata, function () {
   // and confirming that the externalClient is firing off keyVault requests during
   // encrypted operations
   describe('External Key Vault Test', function () {
-    const fs = require('fs');
-    const path = require('path');
-    const { EJSON } = BSON;
     function loadExternal(file) {
       return EJSON.parse(
         fs.readFileSync(path.resolve(__dirname, '../../spec/client-side-encryption/external', file))
@@ -541,9 +535,6 @@ describe('Client Side Encryption Prose Tests', metadata, function () {
   });
 
   describe('BSON size limits and batch splitting', function () {
-    const fs = require('fs');
-    const path = require('path');
-    const { EJSON } = BSON;
     function loadLimits(file) {
       return EJSON.parse(
         fs.readFileSync(path.resolve(__dirname, '../../spec/client-side-encryption/limits', file))
@@ -621,7 +612,7 @@ describe('Client Side Encryption Prose Tests', metadata, function () {
       }
     });
 
-    after(function () {
+    afterEach(function () {
       return this.client && this.client.close();
     });
 

test/integration/client-side-encryption/client_side_encryption.spec.test.js

@@ -57,9 +57,17 @@ const skippedAuthTests = [
 // TODO(NODE-4006): Investigate csfle test "operation fails with maxWireVersion < 8"
 const skippedMaxWireVersionTest = 'operation fails with maxWireVersion < 8';
 
-const SKIPPED_TESTS = new Set(
-  isAuthEnabled ? skippedAuthTests.concat(skippedMaxWireVersionTest) : [skippedMaxWireVersionTest]
-);
+const SKIPPED_TESTS = new Set([
+  ...(isAuthEnabled
+    ? skippedAuthTests.concat(skippedMaxWireVersionTest)
+    : [skippedMaxWireVersionTest]),
+  // TODO(NODE-4288): Fix FLE 2 tests
+  'default state collection names are applied',
+  'drop removes all state collections',
+  'CreateCollection from encryptedFields.',
+  'DropCollection from encryptedFields',
+  'DropCollection from remote encryptedFields'
+]);
 
 describe('Client Side Encryption', function () {
   const testContext = new TestRunnerContext();

test/mocha_mongodb.json

@@ -5,8 +5,7 @@
     "ts-node/register",
     "test/tools/runner/chai-addons.js",
     "test/tools/runner/hooks/configuration.js",
-    "test/tools/runner/hooks/client_leak_checker.js",
-    "test/tools/runner/hooks/session_leak_checker.js"
+    "test/tools/runner/hooks/leak_checker.ts"
   ],
   "extension": ["js", "ts"],
   "ui": "test/tools/runner/metadata_ui.js",

test/tools/runner/hooks/leak_checker.ts

@@ -0,0 +1,182 @@
+/* eslint-disable @typescript-eslint/no-this-alias */
+import { expect } from 'chai';
+import * as chalk from 'chalk';
+import * as net from 'net';
+
+import { MongoClient } from '../../../../src/mongo_client';
+import { ServerSessionPool } from '../../../../src/sessions';
+
+class LeakChecker {
+  static originalAcquire: typeof ServerSessionPool.prototype.acquire;
+  static originalRelease: typeof ServerSessionPool.prototype.release;
+  static kAcquiredCount: symbol;
+
+  static originalConnect: typeof MongoClient.prototype.connect;
+  static originalClose: typeof MongoClient.prototype.close;
+  static kConnectCount: symbol;
+
+  static {
+    this.originalAcquire = ServerSessionPool.prototype.acquire;
+    this.originalRelease = ServerSessionPool.prototype.release;
+    this.kAcquiredCount = Symbol('acquiredCount');
+    this.originalConnect = MongoClient.prototype.connect;
+    this.originalClose = MongoClient.prototype.close;
+    this.kConnectCount = Symbol('connectedCount');
+  }
+
+  clients: Set<MongoClient>;
+  sessionPools: Set<ServerSessionPool>;
+
+  constructor(public titlePath: string) {
+    this.clients = new Set<MongoClient>();
+    this.sessionPools = new Set<ServerSessionPool>();
+  }
+
+  setupSessionLeakChecker() {
+    const leakChecker = this;
+    ServerSessionPool.prototype.acquire = function (...args) {
+      leakChecker.sessionPools.add(this);
+
+      this[LeakChecker.kAcquiredCount] ??= 0;
+      this[LeakChecker.kAcquiredCount] += 1;
+
+      return LeakChecker.originalAcquire.call(this, ...args);
+    };
+
+    ServerSessionPool.prototype.release = function (...args) {
+      if (!(LeakChecker.kAcquiredCount in this)) {
+        throw new Error('releasing before acquiring even once??');
+      } else {
+        this[LeakChecker.kAcquiredCount] -= 1;
+      }
+
+      return LeakChecker.originalRelease.call(this, ...args);
+    };
+  }
+
+  setupClientLeakChecker() {
+    const leakChecker = this;
+    MongoClient.prototype.connect = function (...args) {
+      leakChecker.clients.add(this);
+      this[LeakChecker.kConnectCount] ??= 0;
+
+      const lastArg = args[args.length - 1];
+      const lastArgIsCallback = typeof lastArg === 'function';
+      if (lastArgIsCallback) {
+        const argsWithoutCallback = args.slice(0, args.length - 1);
+        return LeakChecker.originalConnect.call(this, ...argsWithoutCallback, (error, client) => {
+          if (error == null) {
+            this[LeakChecker.kConnectCount] += 1; // only increment on successful connects
+          }
+          return lastArg(error, client);
+        });
+      } else {
+        return LeakChecker.originalConnect.call(this, ...args).then(client => {
+          this[LeakChecker.kConnectCount] += 1; // only increment on successful connects
+          return client;
+        });
+      }
+    };
+
+    MongoClient.prototype.close = function (...args) {
+      this[LeakChecker.kConnectCount] ??= 0; // prevents NaN, its fine to call close on a client that never called connect
+      this[LeakChecker.kConnectCount] -= 1;
+      return LeakChecker.originalClose.call(this, ...args);
+    };
+  }
+
+  setup() {
+    this.setupSessionLeakChecker();
+    this.setupClientLeakChecker();
+  }
+
+  reset() {
+    for (const sessionPool of this.sessionPools) {
+      delete sessionPool[LeakChecker.kAcquiredCount];
+    }
+    ServerSessionPool.prototype.acquire = LeakChecker.originalAcquire;
+    ServerSessionPool.prototype.release = LeakChecker.originalRelease;
+    this.sessionPools.clear();
+
+    for (const client of this.clients) {
+      delete client[LeakChecker.kConnectCount];
+    }
+    MongoClient.prototype.connect = LeakChecker.originalConnect;
+    MongoClient.prototype.close = LeakChecker.originalClose;
+    this.clients.clear();
+  }
+
+  assert() {
+    for (const pool of this.sessionPools) {
+      expect(pool[LeakChecker.kAcquiredCount], 'ServerSessionPool acquired count').to.equal(0);
+    }
+    for (const client of this.clients) {
+      expect(client[LeakChecker.kConnectCount], 'MongoClient connect count').to.be.lessThanOrEqual(
+        0
+      );
+    }
+  }
+}
+
+let currentLeakChecker: LeakChecker | null;
+
+const leakCheckerBeforeEach = async function () {
+  currentLeakChecker = new LeakChecker(this.currentTest.fullTitle());
+  currentLeakChecker.setup();
+};
+
+const leakCheckerAfterEach = async function () {
+  let thrownError: Error | undefined;
+  try {
+    currentLeakChecker.assert();
+  } catch (error) {
+    thrownError = error;
+  }
+
+  currentLeakChecker?.reset();
+  currentLeakChecker = null;
+
+  if (thrownError instanceof Error) {
+    this.test.error(thrownError);
+  }
+};
+
+const TRACE_SOCKETS = process.env.TRACE_SOCKETS === 'true' ? true : false;
+const kSocketId = Symbol('socketId');
+const originalCreateConnection = net.createConnection;
+let socketCounter = 0n;
+
+const socketLeakCheckBeforeAll = function socketLeakCheckBeforeAll() {
+  // @ts-expect-error: Typescript says this is readonly, but it is not at runtime
+  net.createConnection = options => {
+    const socket = originalCreateConnection(options);
+    socket[kSocketId] = socketCounter.toString().padStart(5, '0');
+    socketCounter++;
+    return socket;
+  };
+};
+
+const filterHandlesForSockets = function (handle: any): handle is net.Socket {
+  // Stdio are instanceof Socket so look for fd to be null
+  return handle?.fd == null && handle instanceof net.Socket && handle?.destroyed !== true;
+};
+
+const socketLeakCheckAfterEach: Mocha.AsyncFunc = async function socketLeakCheckAfterEach() {
+  const indent = '  '.repeat(this.currentTest.titlePath().length + 1);
+
+  const handles = (process as any)._getActiveHandles();
+  const sockets: net.Socket[] = handles.filter(handle => filterHandlesForSockets(handle));
+
+  for (const socket of sockets) {
+    console.log(
+      chalk.yellow(
+        `${indent}⚡︎ socket ${socket[kSocketId]} not destroyed [${socket.localAddress}:${socket.localPort} → ${socket.remoteAddress}:${socket.remotePort}]`
+      )
+    );
+  }
+};
+
+const beforeAll = TRACE_SOCKETS ? [socketLeakCheckBeforeAll] : [];
+const beforeEach = [leakCheckerBeforeEach];
+const afterEach = [leakCheckerAfterEach, ...(TRACE_SOCKETS ? [socketLeakCheckAfterEach] : [])];
+module.exports = { mochaHooks: { beforeAll, beforeEach, afterEach } };