fix(mongodb-runner): make close operation more resilient (#571)

Store identifying information about a server in a database, then only actually
run `process.kill()` on servers where this identifying information matches
what we expect it to match.

Also, do not kill processes to which we cannot connect, and instead assume
they have already been closed before.

Fixes: #483

Author: addaleax

packages/mongodb-runner/src/mongocluster.spec.ts

@@ -258,10 +258,10 @@ describe('MongoCluster', function () {
       secondaries: 0,
     });
     cluster = await MongoCluster.deserialize(cluster.serialize());
-    const { ok } = await cluster.withClient(async (client) => {
-      return await client.db('admin').command({ ping: 1 });
+    const doc = await cluster.withClient(async (client) => {
+      return await client.db('config').collection('mongodbrunner').findOne();
     });
-    expect(ok).to.equal(1);
+    expect(doc?._id).to.be.a('string');
     await cluster.close();
   });
 });

packages/mongodb-runner/src/mongoserver.ts

@@ -9,11 +9,11 @@ import {
   isFailureToSetupListener,
 } from './mongologreader';
 import { Readable } from 'stream';
-import type { Document } from 'mongodb';
+import type { Document, MongoClientOptions } from 'mongodb';
 import { MongoClient } from 'mongodb';
 import path from 'path';
 import { once } from 'events';
-import { uuid, debug } from './util';
+import { uuid, debug, pick } from './util';
 
 export interface MongoServerOptions {
   binDir?: string;
@@ -24,32 +24,52 @@ export interface MongoServerOptions {
   docker?: string | string[]; // Image or docker options
 }
 
+interface SerializedServerProperties {
+  _id: string;
+  pid?: number;
+  port?: number;
+  dbPath?: string;
+  startTime: string;
+  hasInsertedMetadataCollEntry: boolean;
+}
+
 export class MongoServer {
+  private uuid: string = uuid();
   private buildInfo?: Document;
   private childProcess?: ChildProcess;
   private pid?: number;
   private port?: number;
   private dbPath?: string;
   private closing = false;
+  private startTime = new Date().toISOString();
+  private hasInsertedMetadataCollEntry = false;
 
   private constructor() {
     /* see .start() */
   }
 
-  serialize(): unknown /* JSON-serializable */ {
+  serialize(): SerializedServerProperties {
     return {
+      _id: this.uuid,
       pid: this.pid,
       port: this.port,
       dbPath: this.dbPath,
+      startTime: this.startTime,
+      hasInsertedMetadataCollEntry: this.hasInsertedMetadataCollEntry,
     };
   }
 
-  static async deserialize(serialized: any): Promise<MongoServer> {
+  static async deserialize(
+    serialized: SerializedServerProperties,
+  ): Promise<MongoServer> {
     const srv = new MongoServer();
-    srv.pid = serialized.pid;
+    srv.uuid = serialized._id;
     srv.port = serialized.port;
-    srv.dbPath = serialized.dbPath;
-    await srv._populateBuildInfo();
+    srv.closing = !!(await srv._populateBuildInfo('restore-check'));
+    if (!srv.closing) {
+      srv.pid = serialized.pid;
+      srv.dbPath = serialized.dbPath;
+    }
     return srv;
   }
 
@@ -116,7 +136,7 @@ export class MongoServer {
     const srv = new MongoServer();
 
     if (!options.docker) {
-      const dbPath = path.join(options.tmpDir, `db-${uuid()}`);
+      const dbPath = path.join(options.tmpDir, `db-${srv.uuid}`);
       await fs.mkdir(dbPath, { recursive: true });
       srv.dbPath = dbPath;
     }
@@ -234,7 +254,10 @@ export class MongoServer {
       logEntryStream.resume();
 
       srv.port = port;
-      await srv._populateBuildInfo();
+      const buildInfoError = await srv._populateBuildInfo('insert-new');
+      if (buildInfoError) {
+        debug('failed to get buildInfo', buildInfoError);
+      }
     } catch (err) {
       await srv.close();
       throw err;
@@ -270,24 +293,83 @@ export class MongoServer {
     this.dbPath = undefined;
   }
 
-  private async _populateBuildInfo(): Promise<void> {
-    if (this.buildInfo?.version) return;
-    this.buildInfo = await this.withClient(
-      async (client) => await client.db('admin').command({ buildInfo: 1 }),
-    );
+  private async _ensureMatchingMetadataColl(
+    client: MongoClient,
+    mode: 'insert-new' | 'restore-check',
+  ): Promise<void> {
+    const hello = await client.db('admin').command({ hello: 1 });
+    const isMongoS = hello.msg === 'isdbgrid';
+    const insertedInfo = pick(this.serialize(), [
+      '_id',
+      'pid',
+      'port',
+      'dbPath',
+      'startTime',
+    ]);
+    const runnerColl = client
+      .db(isMongoS ? 'config' : 'local')
+      .collection<
+        Omit<SerializedServerProperties, 'hasInsertedMetadataCollEntry'>
+      >('mongodbrunner');
+    debug('ensuring metadata collection entry', insertedInfo, { isMongoS });
+    if (mode === 'insert-new') {
+      await runnerColl.insertOne(insertedInfo);
+      debug('inserted metadata collection entry', insertedInfo);
+      this.hasInsertedMetadataCollEntry = true;
+    } else {
+      if (!this.hasInsertedMetadataCollEntry) {
+        debug(
+          'skipping metadata collection match check as we never inserted metadata',
+        );
+        return;
+      }
+      const match = await runnerColl.findOne();
+      debug('read metadata collection entry', insertedInfo, match);
+      if (!match) {
+        throw new Error(
+          'Cannot find mongodbrunner entry, assuming that this instance was not started by mongodb-runner',
+        );
+      }
+      if (match._id !== insertedInfo._id) {
+        throw new Error(
+          `Mismatched mongodbrunner entry: ${JSON.stringify(match)} !== ${JSON.stringify(insertedInfo)}`,
+        );
+      }
+    }
+  }
+
+  private async _populateBuildInfo(
+    mode: 'insert-new' | 'restore-check',
+  ): Promise<Error | null> {
+    try {
+      // directConnection + retryWrites let us write to `local` db on secondaries
+      const clientOpts = { retryWrites: false };
+      this.buildInfo = await this.withClient(async (client) => {
+        // Insert the metadata entry, except if we're a freshly started mongos
+        // (which does not have its own storage to persist)
+        await this._ensureMatchingMetadataColl(client, mode);
+        return await client.db('admin').command({ buildInfo: 1 });
+      }, clientOpts);
+    } catch (err) {
+      debug('failed to get buildInfo, treating as closed server', err);
+      return err as Error;
+    }
     debug(
       'got server build info through client',
       this.serverVersion,
       this.serverVariant,
     );
+    return null;
   }
 
   async withClient<Fn extends (client: MongoClient) => any>(
     fn: Fn,
+    clientOptions: MongoClientOptions = {},
   ): Promise<ReturnType<Fn>> {
-    const client = await MongoClient.connect(
-      `mongodb://${this.hostport}/?directConnection=true`,
-    );
+    const client = await MongoClient.connect(`mongodb://${this.hostport}/`, {
+      directConnection: true,
+      ...clientOptions,
+    });
     try {
       return await fn(client);
     } finally {

packages/mongodb-runner/src/runner-helpers.ts

@@ -84,14 +84,9 @@ export async function stop(argv: {
   id?: string;
   all?: boolean;
 }) {
-  const toStop: Array<StoredInstance> = [];
-  for await (const instance of instances(argv)) {
-    if (instance.id === argv.id || argv.all) toStop.push(instance);
-  }
-  await Promise.all(
-    toStop.map(async ({ filepath, serialized }) => {
-      await (await MongoCluster.deserialize(serialized)).close();
-      await fs.rm(filepath);
-    }),
-  );
+  await parallelForEach(instances(argv), async (instance: StoredInstance) => {
+    if (instance.id !== argv.id && !argv.all) return;
+    await (await MongoCluster.deserialize(instance.serialized)).close();
+    await fs.rm(instance.filepath);
+  });
 }

packages/mongodb-runner/src/util.ts

@@ -28,3 +28,16 @@ export async function parallelForEach<T>(
 
   return await Promise.allSettled(result);
 }
+
+export function pick<T extends object, K extends keyof T>(
+  obj: T,
+  keys: K[],
+): Pick<T, K> {
+  const ret: Partial<Pick<T, K>> = {};
+  for (const key of Object.keys(obj) as K[]) {
+    if (keys.includes(key)) {
+      ret[key] = obj[key];
+    }
+  }
+  return ret as Pick<T, K>;
+}