feat: added new polling class with CockatielPolicy

Author: lwin-kyaw

packages/shield-controller/package.json

@@ -48,7 +48,9 @@
   },
   "dependencies": {
     "@metamask/base-controller": "^8.4.2",
-    "@metamask/utils": "^11.8.1"
+    "@metamask/controller-utils": "^11.14.1",
+    "@metamask/utils": "^11.8.1",
+    "cockatiel": "^3.1.2"
   },
   "devDependencies": {
     "@babel/runtime": "^7.23.9",

packages/shield-controller/src/polling-with-policy.test.ts

@@ -0,0 +1,141 @@
+import { PollingWithCockatielPolicy } from './polling-with-policy';
+import { delay } from '../tests/utils';
+import { HttpError } from '@metamask/controller-utils';
+
+describe('PollingWithCockatielPolicy', () => {
+  it('should return the success result', async () => {
+    const policy = new PollingWithCockatielPolicy();
+    const result = await policy.start('test', async () => {
+      return 'test';
+    });
+    expect(result).toBe('test');
+  });
+
+  it('should retry the request and complete successfully', async () => {
+    const policy = new PollingWithCockatielPolicy();
+    let invocationCount = 0;
+    const mockRequestFn = jest
+      .fn()
+      .mockImplementation(async (_abortSignal: AbortSignal) => {
+        invocationCount += 1;
+        return new Promise((resolve, reject) => {
+          setTimeout(() => {
+            // eslint-disable-next-line jest/no-conditional-in-test
+            if (invocationCount < 3) {
+              reject(new HttpError(412));
+            }
+            resolve('test');
+          }, 100);
+        });
+      });
+    const result = await policy.start('test', mockRequestFn);
+    expect(result).toBe('test');
+    expect(mockRequestFn).toHaveBeenCalledTimes(3);
+  });
+
+  it('should not retry when the error is not retryable', async () => {
+    const policy = new PollingWithCockatielPolicy();
+    const mockRequestFn = jest
+      .fn()
+      .mockImplementation(async (_abortSignal: AbortSignal) => {
+        return new Promise((_resolve, reject) => {
+          reject(new HttpError(500, 'Internal server error'));
+        });
+      });
+    await expect(policy.start('test', mockRequestFn)).rejects.toThrow(
+      'Internal server error',
+    );
+    expect(mockRequestFn).toHaveBeenCalledTimes(1);
+  });
+
+  it('should throw an error when the retry exceeds the max retries', async () => {
+    const policy = new PollingWithCockatielPolicy({
+      maxRetries: 3,
+    });
+
+    const requestFn = jest
+      .fn()
+      .mockImplementation(async (_abortSignal: AbortSignal) => {
+        return new Promise((_resolve, reject) => {
+          setTimeout(() => {
+            reject(new HttpError(412, 'Results are not available yet'));
+          }, 100);
+        });
+      });
+
+    const result = policy.start('test', requestFn);
+    await expect(result).rejects.toThrow('Results are not available yet');
+    expect(requestFn).toHaveBeenCalledTimes(4);
+  });
+
+  it('should throw a `Request Cancelled` error when the request is aborted', async () => {
+    const policy = new PollingWithCockatielPolicy({
+      maxRetries: 3,
+    });
+
+    const requestFn = jest
+      .fn()
+      .mockImplementation(async (abortSignal: AbortSignal) => {
+        return new Promise((_resolve, reject) => {
+          setTimeout(() => {
+            // eslint-disable-next-line jest/no-conditional-in-test
+            if (abortSignal.aborted) {
+              reject(new Error('test error'));
+            }
+            reject(new Error('test error'));
+          }, 100);
+        });
+      });
+
+    const result = policy.start('test', requestFn);
+    await delay(10);
+    policy.abortPendingRequest('test');
+    await expect(result).rejects.toThrow('Request cancelled');
+  });
+
+  it('should throw a `Request Cancelled` error when a new request is started with the same request id', async () => {
+    const policy = new PollingWithCockatielPolicy();
+
+    const requestFn = jest
+      .fn()
+      .mockImplementation(async (abortSignal: AbortSignal) => {
+        return new Promise((resolve, reject) => {
+          setTimeout(() => {
+            // eslint-disable-next-line jest/no-conditional-in-test
+            if (abortSignal.aborted) {
+              reject(new Error('test error'));
+            }
+            resolve('test');
+          }, 100);
+        });
+      });
+
+    const result = policy.start('test', requestFn);
+    await delay(10);
+    const secondResult = policy.start('test', requestFn);
+    await expect(result).rejects.toThrow('Request cancelled');
+    expect(await secondResult).toBe('test');
+  });
+
+  it('should resolve the result when two requests are started with the different request ids', async () => {
+    const policy = new PollingWithCockatielPolicy();
+
+    const requestFn = (result: string) =>
+      jest.fn().mockImplementation(async (abortSignal: AbortSignal) => {
+        return new Promise((resolve, reject) => {
+          // eslint-disable-next-line jest/no-conditional-in-test
+          if (abortSignal.aborted) {
+            reject(new Error('test error'));
+          }
+          setTimeout(() => {
+            resolve(result);
+          }, 100);
+        });
+      });
+
+    const result = policy.start('test', requestFn('test'));
+    const secondResult = policy.start('test2', requestFn('test2'));
+    expect(await result).toBe('test');
+    expect(await secondResult).toBe('test2');
+  });
+});

packages/shield-controller/src/polling-with-policy.ts

@@ -0,0 +1,65 @@
+import {
+  createServicePolicy,
+  HttpError,
+  type CreateServicePolicyOptions,
+  type ServicePolicy,
+} from '@metamask/controller-utils';
+import { handleWhen } from 'cockatiel';
+
+export type RequestFn<ReturnType> = (
+  signal: AbortSignal,
+) => Promise<ReturnType>;
+
+export class PollingWithCockatielPolicy {
+  readonly #policy: ServicePolicy;
+
+  readonly #requestEntry = new Map<string, AbortController>();
+
+  constructor(policyOptions: CreateServicePolicyOptions = {}) {
+    const retryFilterPolicy = handleWhen(this.#shouldRetry);
+    this.#policy = createServicePolicy({
+      ...policyOptions,
+      retryFilterPolicy,
+    });
+  }
+
+  async start<ReturnType>(requestId: string, requestFn: RequestFn<ReturnType>) {
+    this.abortPendingRequest(requestId);
+    const abortController = this.addNewRequestEntry(requestId);
+
+    try {
+      const result = await this.#policy.execute(
+        async ({ signal: abortSignal }) => {
+          return requestFn(abortSignal);
+        },
+        abortController.signal,
+      );
+      return result;
+    } catch (error) {
+      if (abortController.signal.aborted) {
+        throw new Error('Request cancelled');
+      }
+      throw error;
+    }
+  }
+
+  addNewRequestEntry(requestId: string) {
+    const abortController = new AbortController();
+    this.#requestEntry.set(requestId, abortController);
+    return abortController;
+  }
+
+  abortPendingRequest(requestId: string) {
+    const abortController = this.#requestEntry.get(requestId);
+    abortController?.abort();
+  }
+
+  #shouldRetry(error: Error): boolean {
+    if (error instanceof HttpError) {
+      // Note: we don't retry on 4xx errors, only on 5xx errors.
+      // but we won't retry on 400 coz it means that the request body is invalid.
+      return error.httpStatus > 400 && error.httpStatus < 500;
+    }
+    return false;
+  }
+}

packages/shield-controller/tests/utils.ts

@@ -106,3 +106,15 @@ export function setupCoverageResultReceived(
     baseMessenger.subscribe('ShieldController:coverageResultReceived', handler);
   });
 }
+
+/**
+ * Delay for a specified amount of time.
+ *
+ * @param ms - The number of milliseconds to delay.
+ * @returns A promise that resolves after the specified amount of time.
+ */
+export function delay(ms: number): Promise<void> {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms);
+  });
+}

yarn.lock

@@ -4727,11 +4727,13 @@ __metadata:
     "@lavamoat/preinstall-always-fail": "npm:^2.1.0"
     "@metamask/auto-changelog": "npm:^3.4.4"
     "@metamask/base-controller": "npm:^8.4.2"
+    "@metamask/controller-utils": "npm:^11.14.1"
     "@metamask/signature-controller": "npm:^34.0.2"
     "@metamask/transaction-controller": "npm:^60.10.0"
     "@metamask/utils": "npm:^11.8.1"
     "@ts-bridge/cli": "npm:^0.6.1"
     "@types/jest": "npm:^27.4.1"
+    cockatiel: "npm:^3.1.2"
     deepmerge: "npm:^4.2.2"
     jest: "npm:^27.5.1"
     lodash: "npm:^4.17.21"