Improve timeout handling

packages/snaps-controllers/coverage.json

@@ -1,5 +1,5 @@
 {
-  "branches": 91.01,
+  "branches": 91.05,
   "functions": 96.56,
   "lines": 97.8,
   "statements": 97.46

packages/snaps-controllers/src/services/AbstractExecutionService.test.ts

@@ -1,3 +1,4 @@
+import { BasePostMessageStream } from '@metamask/post-message-stream';
 import { HandlerType } from '@metamask/snaps-utils';
 import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils';
 
@@ -126,5 +127,31 @@ describe('AbstractExecutionService', () => {
     ).rejects.toThrow(
       'Invalid JSON-RPC request: At path: params -- Expected the value to satisfy a union of `record | array`, but received: [object Object].',
     );
+
+    await service.terminateAllSnaps();
+  });
+
+  it('throws an error if execution environment fails to respond to ping', async () => {
+    const { service } = createService(MockExecutionService);
+
+    class MockStream extends BasePostMessageStream {
+      protected _postMessage(_data?: unknown): void {
+        // no-op
+      }
+    }
+
+    // @ts-expect-error Accessing private property and returning unusable worker.
+    service.initEnvStream = async () =>
+      Promise.resolve({ worker: null, stream: new MockStream() });
+
+    await expect(
+      service.executeSnap({
+        snapId: MOCK_SNAP_ID,
+        sourceCode: `
+        console.log('foo');
+      `,
+        endowments: ['console'],
+      }),
+    ).rejects.toThrow('The Snaps execution environment failed to start.');
   });
 });

packages/snaps-controllers/src/services/AbstractExecutionService.ts

@@ -13,6 +13,7 @@ import type {
 import {
   Duration,
   assertIsJsonRpcRequest,
+  inMilliseconds,
   isJsonRpcNotification,
   isObject,
 } from '@metamask/utils';
@@ -37,6 +38,7 @@ export type SetupSnapProvider = (snapId: string, stream: Duplex) => void;
 export type ExecutionServiceArgs = {
   setupSnapProvider: SetupSnapProvider;
   messenger: ExecutionServiceMessenger;
+  pingTimeout?: number;
   terminationTimeout?: number;
 };
 
@@ -70,19 +72,23 @@ export abstract class AbstractExecutionService<WorkerType>
 
   #messenger: ExecutionServiceMessenger;
 
+  #pingTimeout: number;
+
   #terminationTimeout: number;
 
   constructor({
     setupSnapProvider,
     messenger,
-    terminationTimeout = Duration.Second,
+    pingTimeout = inMilliseconds(2, Duration.Second),
+    terminationTimeout = inMilliseconds(1, Duration.Second),
   }: ExecutionServiceArgs) {
     this.#snapRpcHooks = new Map();
     this.jobs = new Map();
     this.setupSnapProvider = setupSnapProvider;
     this.#snapToJobMap = new Map();
     this.#jobToSnapMap = new Map();
     this.#messenger = messenger;
+    this.#pingTimeout = pingTimeout;
     this.#terminationTimeout = terminationTimeout;
 
     this.registerMessageHandlers();
@@ -338,15 +344,23 @@ export abstract class AbstractExecutionService<WorkerType>
       throw new Error(`Snap "${snapData.snapId}" is already being executed.`);
     }
 
+    // This may resolve even if the environment has failed to start up fully
     const job = await this.initJob();
     this.#mapSnapAndJob(snapData.snapId, job.id);
 
     // Ping the worker to ensure that it started up
-    await this.command(job.id, {
-      jsonrpc: '2.0',
-      method: 'ping',
-      id: nanoid(),
-    });
+    const pingResult = await withTimeout(
+      this.command(job.id, {
+        jsonrpc: '2.0',
+        method: 'ping',
+        id: nanoid(),
+      }),
+      this.#pingTimeout,
+    );
+
+    if (pingResult === hasTimedOut) {
+      throw new Error('The Snaps execution environment failed to start.');
+    }
 
     const rpcStream = job.streams.rpc;
 

packages/snaps-controllers/src/snaps/SnapController.test.ts

@@ -472,7 +472,7 @@ describe('SnapController', () => {
           id: 1,
         },
       }),
-    ).rejects.toThrow(/request timed out/u);
+    ).rejects.toThrow(`${snap.id} failed to respond to the request in time.`);
 
     expect(snapController.state.snaps[snap.id].status).toBe('crashed');
     snapController.destroy();
@@ -1141,7 +1141,7 @@ describe('SnapController', () => {
           id: 1,
         },
       }),
-    ).rejects.toThrow(/request timed out/u);
+    ).rejects.toThrow(`${snap.id} failed to respond to the request in time.`);
     expect(snapController.state.snaps[snap.id].status).toBe('crashed');
 
     snapController.destroy();
@@ -1226,7 +1226,7 @@ describe('SnapController', () => {
           id: 1,
         },
       }),
-    ).rejects.toThrow(/request timed out/u);
+    ).rejects.toThrow(`${snap.id} failed to respond to the request in time.`);
     expect(snapController.state.snaps[snap.id].status).toBe('crashed');
 
     snapController.destroy();
@@ -1401,7 +1401,7 @@ describe('SnapController', () => {
     const snapController = getSnapController(
       getSnapControllerOptions({
         messenger,
-        maxRequestTime: 50,
+        maxInitTime: 50,
         state: {
           snaps: getPersistedSnapsState(),
         },
@@ -1421,7 +1421,7 @@ describe('SnapController', () => {
     );
 
     await expect(snapController.startSnap(snap.id)).rejects.toThrow(
-      /request timed out/u,
+      `${snap.id} failed to start.`,
     );
 
     snapController.destroy();
@@ -1583,7 +1583,7 @@ describe('SnapController', () => {
           id: 1,
         },
       }),
-    ).rejects.toThrow(/request timed out/u);
+    ).rejects.toThrow(`${snap.id} failed to respond to the request in time.`);
     expect(snapController.state.snaps[snap.id].status).toBe('crashed');
 
     await snapController.removeSnap(snap.id);

packages/snaps-controllers/src/snaps/SnapController.ts

@@ -598,6 +598,12 @@ type SnapControllerArgs = {
    */
   maxRequestTime?: number;
 
+  /**
+   * The maximum amount of time a snap may take to initialize, including
+   * the time it takes for the execution environment to start.
+   */
+  maxInitTime?: number;
+
   /**
    * The npm registry URL that will be used to fetch published snaps.
    */
@@ -699,6 +705,8 @@ export class SnapController extends BaseController<
   // This property cannot be hash private yet because of tests.
   private readonly maxRequestTime: number;
 
+  #maxInitTime: number;
+
   #detectSnapLocation: typeof detectSnapLocation;
 
   #snapsRuntimeData: Map<SnapId, SnapRuntimeData>;
@@ -723,6 +731,7 @@ export class SnapController extends BaseController<
     idleTimeCheckInterval = inMilliseconds(5, Duration.Second),
     maxIdleTime = inMilliseconds(30, Duration.Second),
     maxRequestTime = inMilliseconds(60, Duration.Second),
+    maxInitTime = inMilliseconds(60, Duration.Second),
     fetchFunction = globalThis.fetch.bind(globalThis),
     featureFlags = {},
     detectSnapLocation: detectSnapLocationFunction = detectSnapLocation,
@@ -778,6 +787,7 @@ export class SnapController extends BaseController<
     this.#idleTimeCheckInterval = idleTimeCheckInterval;
     this.#maxIdleTime = maxIdleTime;
     this.maxRequestTime = maxRequestTime;
+    this.#maxInitTime = maxInitTime;
     this.#detectSnapLocation = detectSnapLocationFunction;
     this._onUnhandledSnapError = this._onUnhandledSnapError.bind(this);
     this._onOutboundRequest = this._onOutboundRequest.bind(this);
@@ -2406,12 +2416,18 @@ export class SnapController extends BaseController<
 
     try {
       const runtime = this.#getRuntimeExpect(snapId);
-      const result = await this.#executeWithTimeout(
+      const result = await withTimeout(
         this.messagingSystem.call('ExecutionService:executeSnap', {
           ...snapData,
           endowments: await this.#getEndowments(snapId),
         }),
+        this.#maxInitTime,
       );
+
+      if (result === hasTimedOut) {
+        throw new Error(`${snapId} failed to start.`);
+      }
+
       this.#transition(snapId, SnapStatusEvents.Start);
       // We treat the initialization of the snap as the first request, for idle timing purposes.
       runtime.lastRequest = Date.now();
@@ -2899,10 +2915,13 @@ export class SnapController extends BaseController<
 
       // This will either get the result or reject due to the timeout.
       try {
-        const result = await this.#executeWithTimeout(
-          handleRpcRequestPromise,
-          timer,
-        );
+        const result = await withTimeout(handleRpcRequestPromise, timer);
+
+        if (result === hasTimedOut) {
+          throw new Error(
+            `${snapId} failed to respond to the request in time.`,
+          );
+        }
 
         await this.#assertSnapRpcRequestResult(snapId, handlerType, result);
 
@@ -3033,26 +3052,6 @@ export class SnapController extends BaseController<
     }
   }
 
-  /**
-   * Awaits the specified promise and rejects if the promise doesn't resolve
-   * before the timeout.
-   *
-   * @param promise - The promise to await.
-   * @param timer - An optional timer object to control the timeout.
-   * @returns The result of the promise or rejects if the promise times out.
-   * @template PromiseValue - The value of the Promise.
-   */
-  async #executeWithTimeout<PromiseValue>(
-    promise: Promise<PromiseValue>,
-    timer?: Timer,
-  ): Promise<PromiseValue> {
-    const result = await withTimeout(promise, timer ?? this.maxRequestTime);
-    if (result === hasTimedOut) {
-      throw new Error('The request timed out.');
-    }
-    return result;
-  }
-
   #recordSnapRpcRequestStart(snapId: SnapId, requestId: unknown, timer: Timer) {
     const runtime = this.#getRuntimeExpect(snapId);
     runtime.pendingInboundRequests.push({ requestId, timer });