From 9b8a559aed2ea1f594e0f1c94f14d64131ed7eb8 Mon Sep 17 00:00:00 2001
From: Tim Yiu <137842098+tyiuhc@users.noreply.github.com>
Date: Mon, 29 Jan 2024 13:52:00 -0800
Subject: [PATCH] fix: Improve remote evaluation fetch retry logic (#96)

---
 .../src/experimentClient.ts                   | 13 +++--
 .../experiment-browser/test/client.test.ts    | 48 +++++++++++++++++++
 .../experiment-core/src/api/evaluation-api.ts |  6 ++-
 .../experiment-core/src/evaluation/error.ts   |  9 ++++
 packages/experiment-core/src/index.ts         |  1 +
 5 files changed, 73 insertions(+), 4 deletions(-)
 create mode 100644 packages/experiment-core/src/evaluation/error.ts

diff --git a/packages/experiment-browser/src/experimentClient.ts b/packages/experiment-browser/src/experimentClient.ts
index 42f3ff9c..0abbb3fc 100644
--- a/packages/experiment-browser/src/experimentClient.ts
+++ b/packages/experiment-browser/src/experimentClient.ts
@@ -7,6 +7,7 @@ import {
   EvaluationApi,
   EvaluationEngine,
   EvaluationFlag,
+  FetchError,
   FlagApi,
   Poller,
   SdkEvaluationApi,
@@ -37,7 +38,6 @@ import {
   isLocalEvaluationMode,
   isNullOrUndefined,
   isNullUndefinedOrEmpty,
-  isRemoteEvaluationMode,
 } from './util';
 import { Backoff } from './util/backoff';
 import {
@@ -627,7 +627,7 @@ export class ExperimentClient implements Client {
 
     this.debug(`[Experiment] Fetch all: retry=${retry}`);
 
-    // Proactively cancel retries if active in order to avoid unecessary API
+    // Proactively cancel retries if active in order to avoid unnecessary API
     // requests. A new failure will restart the retries.
     if (retry) {
       this.stopRetries();
@@ -638,7 +638,7 @@ export class ExperimentClient implements Client {
       await this.storeVariants(variants, options);
       return variants;
     } catch (e) {
-      if (retry) {
+      if (retry && this.shouldRetryFetch(e)) {
         void this.startRetries(user, options);
       }
       throw e;
@@ -819,6 +819,13 @@ export class ExperimentClient implements Client {
       console.debug(message, ...optionalParams);
     }
   }
+
+  private shouldRetryFetch(e: Error): boolean {
+    if (e instanceof FetchError) {
+      return e.statusCode < 400 || e.statusCode >= 500 || e.statusCode === 429;
+    }
+    return true;
+  }
 }
 
 type SourceVariant = {
diff --git a/packages/experiment-browser/test/client.test.ts b/packages/experiment-browser/test/client.test.ts
index 6f9a63c5..2ad5e7fe 100644
--- a/packages/experiment-browser/test/client.test.ts
+++ b/packages/experiment-browser/test/client.test.ts
@@ -1,4 +1,5 @@
 import { AnalyticsConnector } from '@amplitude/analytics-connector';
+import { FetchError } from '@amplitude/experiment-core';
 
 import {
   ExperimentAnalyticsProvider,
@@ -259,9 +260,11 @@ class TestAnalyticsProvider
   track(): void {
     return;
   }
+
   setUserProperty(): void {
     return;
   }
+
   unsetUserProperty(): void {
     return;
   }
@@ -1048,3 +1051,48 @@ describe('start', () => {
     expect(fetchSpy).toBeCalledTimes(0);
   });
 });
+
+describe('fetch retry with different response codes', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+  test.each([
+    [300, 'Fetch Exception 300', 1],
+    [400, 'Fetch Exception 400', 0],
+    [429, 'Fetch Exception 429', 1],
+    [500, 'Fetch Exception 500', 1],
+    [0, 'Other Exception', 1],
+  ])(
+    'responseCode=%p, errorMessage=%p, retryCalled=%p',
+    async (responseCode, errorMessage, retryCalled) => {
+      const client = new ExperimentClient(API_KEY, {
+        retryFetchOnFailure: true,
+      });
+      mockClientStorage(client);
+
+      jest
+        .spyOn(ExperimentClient.prototype as any, 'doFetch')
+        .mockImplementation(
+          async (user?: ExperimentUser, options?: FetchOptions) => {
+            return new Promise<ExperimentClient>((resolve, reject) => {
+              if (responseCode === 0) {
+                reject(new Error(errorMessage));
+              } else {
+                reject(new FetchError(responseCode, errorMessage));
+              }
+            });
+          },
+        );
+      const retryMock = jest.spyOn(
+        ExperimentClient.prototype as any,
+        'startRetries',
+      );
+      try {
+        await client.fetch({ user_id: 'test_user' });
+      } catch (e) {
+        // catch error
+      }
+      expect(retryMock).toHaveBeenCalledTimes(retryCalled);
+    },
+  );
+});
diff --git a/packages/experiment-core/src/api/evaluation-api.ts b/packages/experiment-core/src/api/evaluation-api.ts
index 21316638..2db43323 100644
--- a/packages/experiment-core/src/api/evaluation-api.ts
+++ b/packages/experiment-core/src/api/evaluation-api.ts
@@ -1,5 +1,6 @@
 import { Base64 } from 'js-base64';
 
+import { FetchError } from '../evaluation/error';
 import { EvaluationVariant } from '../evaluation/flag';
 import { HttpClient } from '../transport/http';
 
@@ -55,7 +56,10 @@ export class SdkEvaluationApi implements EvaluationApi {
       timeoutMillis: options?.timeoutMillis,
     });
     if (response.status != 200) {
-      throw Error(`Fetch error response: status=${response.status}`);
+      throw new FetchError(
+        response.status,
+        `Fetch error response: status=${response.status}`,
+      );
     }
     return JSON.parse(response.body);
   }
diff --git a/packages/experiment-core/src/evaluation/error.ts b/packages/experiment-core/src/evaluation/error.ts
new file mode 100644
index 00000000..b24e516d
--- /dev/null
+++ b/packages/experiment-core/src/evaluation/error.ts
@@ -0,0 +1,9 @@
+export class FetchError extends Error {
+  statusCode: number;
+
+  constructor(statusCode: number, message: string) {
+    super(message);
+    this.statusCode = statusCode;
+    Object.setPrototypeOf(this, FetchError.prototype);
+  }
+}
diff --git a/packages/experiment-core/src/index.ts b/packages/experiment-core/src/index.ts
index 2932f5ae..0cddac25 100644
--- a/packages/experiment-core/src/index.ts
+++ b/packages/experiment-core/src/index.ts
@@ -15,3 +15,4 @@ export { FlagApi, SdkFlagApi } from './api/flag-api';
 export { HttpClient, HttpRequest, HttpResponse } from './transport/http';
 export { Poller } from './util/poller';
 export { safeGlobal } from './util/global';
+export { FetchError } from './evaluation/error';