feat(bson-parsing): terminate worker if execution exceeds 2mins COMPASS-11147

packages/shell-bson-parser/.eslintrc.cjs

@@ -5,5 +5,5 @@ module.exports = {
     tsconfigRootDir: __dirname,
     project: ['./tsconfig-lint.json'],
   },
-  ignorePatterns: ['webpack.worker.config.cjs'],
+  ignorePatterns: ['webpack.worker.config.cjs', 'test/fixtures/**'],
 };

packages/shell-bson-parser/src/index.spec.ts

@@ -7,7 +7,7 @@ import { fileURLToPath } from 'url';
 import * as WebWorkerModule from 'web-worker';
 
 import * as api from './index.js';
-import { terminateWorker } from './worker-client.js';
+import { terminateWorker, callWorker } from './worker-client.js';
 import {
   restrictGlobalScope,
   restrictObjectPrototype,
@@ -27,7 +27,11 @@ describe('shell-bson-parser with webworker processing', function () {
   });
 
   after(function () {
-    process.env.TEST_WORKER_SCRIPT_URL = initialWorkerScriptUrl;
+    if (initialWorkerScriptUrl) {
+      process.env.TEST_WORKER_SCRIPT_URL = initialWorkerScriptUrl;
+    } else {
+      delete process.env.TEST_WORKER_SCRIPT_URL;
+    }
     terminateWorker();
   });
 
@@ -192,4 +196,39 @@ describe('shell-bson-parser with webworker processing', function () {
     // It should not modify the default object proto
     expect(Object.prototype).to.have.property('__proto__');
   });
+
+  describe('execution timeout', function () {
+    const initialWorkerScriptUrl = process.env.TEST_WORKER_SCRIPT_URL;
+
+    beforeEach(function () {
+      terminateWorker();
+      process.env.TEST_WORKER_SCRIPT_URL = '../test/fixtures/slow-worker.mjs';
+    });
+
+    afterEach(function () {
+      terminateWorker();
+      if (initialWorkerScriptUrl) {
+        process.env.TEST_WORKER_SCRIPT_URL = initialWorkerScriptUrl;
+      } else {
+        delete process.env.TEST_WORKER_SCRIPT_URL;
+      }
+    });
+
+    it('rejects a request whose worker thread is wedged past the timeout', async function () {
+      try {
+        await callWorker([1000], { executionTimeoutMs: 500 });
+        expect.fail('Expected callWorker to throw an error due to timeout');
+      } catch (err) {
+        expect((err as Error)?.message).to.equal(
+          'Worker execution timed out after 500ms',
+        );
+      }
+    });
+
+    it('spins up a fresh worker for the next call after a timeout kill', async function () {
+      await callWorker([1000]).catch(() => {}); // timeouts out
+      const result = await callWorker([0]);
+      expect(result).to.equal('done');
+    });
+  });
 });

packages/shell-bson-parser/src/index.ts

@@ -1,12 +1,21 @@
 import type { parse as parseSync } from './parse.js';
 import { toJSString } from './stringify.js';
 import { ParseMode } from './options.js';
+import type { Options } from './options.js';
 import { callWorker, terminateWorker } from './worker-client.js';
+import type { ExecutionOptions } from './worker-client.js';
 
 export const parse = (
-  ...args: Parameters<typeof parseSync>
-): Promise<ReturnType<typeof parseSync>> => callWorker(args);
+  input: string,
+  {
+    executionTimeoutMs,
+    ...parseOptions
+  }: Partial<Options & ExecutionOptions> = {},
+): Promise<ReturnType<typeof parseSync>> => {
+  return callWorker([input, parseOptions], { executionTimeoutMs });
+};
 
 export { ParseMode, toJSString, terminateWorker };
+export type { ExecutionOptions, Options as ParseOptions };
 
 export default parse;

packages/shell-bson-parser/src/worker-client.ts

@@ -6,14 +6,25 @@ import type { WorkerResponse } from './worker-types.js';
 
 /** Close the worker after being idle for 30sec */
 const IDLE_TIMEOUT_MS = 30_000;
+/** Default execution timeout for worker requests */
+const DEFAULT_EXECUTION_TIMEOUT_MS = 120_000;
+
+export type ExecutionOptions = {
+  /** Defaults to `120_000` (2 minutes). */
+  executionTimeoutMs?: number;
+};
 
 let worker: Worker | null = null;
 let idleTimer: ReturnType<typeof setTimeout> | null = null;
 let blobUrl: string | null = null;
 let nextId = 0;
 const pending = new Map<
   number,
-  { resolve: (v: any) => void; reject: (e: Error) => void }
+  {
+    resolve: (v: any) => void;
+    reject: (e: Error) => void;
+    executionTimer: ReturnType<typeof setTimeout>;
+  }
 >();
 
 function scheduleIdleTermination() {
@@ -72,6 +83,7 @@ async function createWorker(): Promise<Worker> {
     if (!entry) {
       return;
     }
+    clearTimeout(entry.executionTimer);
     pending.delete(response.id);
     if (!response.ok) {
       entry.reject(new Error(response.error));
@@ -94,18 +106,32 @@ async function createWorker(): Promise<Worker> {
   return worker;
 }
 
-export async function callWorker<T>(args: unknown[]): Promise<T> {
+export async function callWorker<T>(
+  args: unknown[],
+  executionOptions?: ExecutionOptions,
+): Promise<T> {
   const activeWorker = await createWorker();
   const id = nextId++;
+  const executionTimeoutMs =
+    executionOptions?.executionTimeoutMs ?? DEFAULT_EXECUTION_TIMEOUT_MS;
   const promise = new Promise<T>((resolve, reject) => {
-    pending.set(id, { resolve, reject });
+    const executionTimer = setTimeout(() => {
+      // Terminate the worker is this message is taking too long to execute,
+      // this means all the other pending requests will also be terminated.
+      terminateWorker(
+        new Error(`Worker execution timed out after ${executionTimeoutMs}ms`),
+      );
+    }, executionTimeoutMs);
+    pending.set(id, { resolve, reject, executionTimer });
   });
   try {
     activeWorker.postMessage({
       id,
       args: markBSON(args),
     });
   } catch (err) {
+    const entry = pending.get(id);
+    if (entry) clearTimeout(entry.executionTimer);
     pending.get(id)?.reject(err as Error);
     pending.delete(id);
   } finally {
@@ -126,6 +152,7 @@ export function terminateWorker(
   blobUrl = null;
 
   for (const [id, entry] of pending) {
+    clearTimeout(entry.executionTimer);
     entry.reject(reason);
     pending.delete(id);
   }

packages/shell-bson-parser/test/fixtures/slow-worker.mjs

@@ -0,0 +1,21 @@
+// Avoiding the import of these two functions,
+// adding them here to keep this file self-contained
+function unmarkBSON(value) {
+  return value.data;
+}
+function markBSON(value) {
+  return { data: value, bsonTypes: new Map() };
+}
+self.onmessage = (event) => {
+  const { id, args } = event.data;
+  const [delayMs] = unmarkBSON(args);
+  const start = Date.now();
+  while (Date.now() - start < delayMs) {
+    // noop
+  }
+  self.postMessage({
+    id,
+    ok: true,
+    result: markBSON('done'),
+  });
+};

packages/shell-bson-parser/tsconfig-lint.json

@@ -1,5 +1,5 @@
 {
   "extends": "./tsconfig.json",
   "include": ["**/*"],
-  "exclude": ["node_modules", "dist"]
+  "exclude": ["node_modules", "dist", "test/fixtures"]
 }