Skip to content

Commit

Permalink
test(NODE-4262): simplify leak checker for startSession fixes (#3281)
Browse files Browse the repository at this point in the history
  • Loading branch information
nbbeeken committed Jun 3, 2022
1 parent 0936b58 commit ed50ef5
Show file tree
Hide file tree
Showing 6 changed files with 199 additions and 222 deletions.
Original file line number Diff line number Diff line change
@@ -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) => {
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -621,7 +612,7 @@ describe('Client Side Encryption Prose Tests', metadata, function () {
}
});

after(function () {
afterEach(function () {
return this.client && this.client.close();
});

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
3 changes: 1 addition & 2 deletions test/mocha_mongodb.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
55 changes: 0 additions & 55 deletions test/tools/runner/hooks/client_leak_checker.js

This file was deleted.

182 changes: 182 additions & 0 deletions test/tools/runner/hooks/leak_checker.ts
Original file line number Diff line number Diff line change
@@ -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 } };

0 comments on commit ed50ef5

Please sign in to comment.