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,
@@ -206,4 +206,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(() => {}); // times 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

@@ -4,12 +4,25 @@ const WebWorker = (WebWorkerModule as unknown as { default: typeof Worker })
 import { trackBSON, untrackBSON } from './structured-clone-bson.js';
 import type { WorkerResponse } from './worker-types.js';
 
+/** 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 workerPromise: Promise<Worker> | 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>;
+  }
 >();
 
 const isNodeEnv =
@@ -42,59 +55,85 @@ async function getWorkerScriptUrl(): Promise<string> {
   return blobUrl;
 }
 
-async function createWorker(): Promise<Worker> {
+function createWorker(): Promise<Worker> {
   if (worker) {
-    return worker;
+    return Promise.resolve(worker);
   }
+  if (workerPromise) {
+    return workerPromise;
+  }
+
+  workerPromise = (async () => {
+    const scriptUrl = await getWorkerScriptUrl();
+    const newWorker = new WebWorker(scriptUrl, { type: 'module' });
+    const onMessageHandler = (event: MessageEvent<WorkerResponse>) => {
+      const response = event.data;
+      const entry = pending.get(response.id);
+      if (!entry) {
+        return;
+      }
+      clearTimeout(entry.executionTimer);
+      pending.delete(response.id);
+      if (!response.ok) {
+        entry.reject(response.error);
+        return;
+      }
+      try {
+        entry.resolve(untrackBSON(response.result));
+      } catch (err) {
+        entry.reject(err as Error);
+      }
+    };
+
+    const onErrorHandler = (event: ErrorEvent) => {
+      terminateWorker(new Error(event.message || 'Worker error'));
+    };
+
+    const onMessageErrorHandler = () => {
+      terminateWorker(new Error('Worker message could not be deserialized'));
+    };
+
+    newWorker.addEventListener('message', onMessageHandler);
+    newWorker.addEventListener('error', onErrorHandler);
+    newWorker.addEventListener('messageerror', onMessageErrorHandler);
+
+    worker = newWorker;
+    return newWorker;
+  })();
+
+  workerPromise.catch(() => {
+    workerPromise = null;
+  });
 
-  const scriptUrl = await getWorkerScriptUrl();
-  worker = new WebWorker(scriptUrl, { type: 'module' });
-
-  const onMessageHandler = (event: MessageEvent<WorkerResponse>) => {
-    const response = event.data;
-    const entry = pending.get(response.id);
-    if (!entry) {
-      return;
-    }
-    pending.delete(response.id);
-    if (!response.ok) {
-      entry.reject(response.error);
-      return;
-    }
-    try {
-      entry.resolve(untrackBSON(response.result));
-    } catch (err) {
-      entry.reject(err as Error);
-    }
-  };
-
-  const onErrorHandler = (event: ErrorEvent) => {
-    terminateWorker(new Error(event.message || 'Worker error'));
-  };
-
-  const onMessageErrorHandler = () => {
-    terminateWorker(new Error('Worker message could not be deserialized'));
-  };
-
-  worker.addEventListener('message', onMessageHandler);
-  worker.addEventListener('error', onErrorHandler);
-  worker.addEventListener('messageerror', onMessageErrorHandler);
-
-  return worker;
+  return workerPromise;
 }
 
-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 if 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: trackBSON(args),
     });
   } catch (err) {
+    const entry = pending.get(id);
+    if (entry) clearTimeout(entry.executionTimer);
     pending.get(id)?.reject(err as Error);
     pending.delete(id);
   }
@@ -108,9 +147,11 @@ export function terminateWorker(
   if (blobUrl) URL.revokeObjectURL(blobUrl);
 
   worker = null;
+  workerPromise = null;
   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"]
 }