From 16515df8526a0c1b884987294eaca848c38f7388 Mon Sep 17 00:00:00 2001
From: Ryan Lamb <4955475+kinyoklion@users.noreply.github.com>
Date: Tue, 14 Nov 2023 13:41:37 -0800
Subject: [PATCH] fix: Better handle waiting for initialization for failure
 cases. (#314)

Affects #312
---
 .../sdk-server/__tests__/LDClientImpl.test.ts | 42 +++++++++++++++++++
 .../shared/sdk-server/src/LDClientImpl.ts     | 26 +++++++++++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/packages/shared/sdk-server/__tests__/LDClientImpl.test.ts b/packages/shared/sdk-server/__tests__/LDClientImpl.test.ts
index ad58d5e2f..b36d966cb 100644
--- a/packages/shared/sdk-server/__tests__/LDClientImpl.test.ts
+++ b/packages/shared/sdk-server/__tests__/LDClientImpl.test.ts
@@ -51,6 +51,25 @@ describe('LDClientImpl', () => {
     expect(callbacks.onError).not.toBeCalled();
   });
 
+  it('wait for initialization completes even if initialization completes before it is called', (done) => {
+    setupMockStreamingProcessor();
+    client = createClient();
+
+    setTimeout(async () => {
+      const initializedClient = await client.waitForInitialization();
+      expect(initializedClient).toEqual(client);
+      done();
+    }, 10);
+  });
+
+  it('waiting for initialization the second time produces the same result', async () => {
+    client = createClient();
+    await client.waitForInitialization();
+
+    const initializedClient = await client.waitForInitialization();
+    expect(initializedClient).toEqual(client);
+  });
+
   it('fires ready event in offline mode', async () => {
     client = createClient({ offline: true });
     const initializedClient = await client.waitForInitialization();
@@ -74,6 +93,29 @@ describe('LDClientImpl', () => {
     expect(callbacks.onError).toBeCalled();
   });
 
+  it('initialization promise is rejected even if the failure happens before wait is called', (done) => {
+    setupMockStreamingProcessor(true);
+    client = createClient();
+
+    setTimeout(async () => {
+      await expect(client.waitForInitialization()).rejects.toThrow('failed');
+
+      expect(client.initialized()).toBeFalsy();
+      expect(callbacks.onReady).not.toBeCalled();
+      expect(callbacks.onFailed).toBeCalled();
+      expect(callbacks.onError).toBeCalled();
+      done();
+    }, 10);
+  });
+
+  it('waiting a second time results in the same rejection', async () => {
+    setupMockStreamingProcessor(true);
+    client = createClient();
+
+    await expect(client.waitForInitialization()).rejects.toThrow('failed');
+    await expect(client.waitForInitialization()).rejects.toThrow('failed');
+  });
+
   it('isOffline returns true in offline mode', () => {
     client = createClient({ offline: true });
     expect(client.isOffline()).toEqual(true);
diff --git a/packages/shared/sdk-server/src/LDClientImpl.ts b/packages/shared/sdk-server/src/LDClientImpl.ts
index ef0818bc0..c0dbc6b07 100644
--- a/packages/shared/sdk-server/src/LDClientImpl.ts
+++ b/packages/shared/sdk-server/src/LDClientImpl.ts
@@ -90,6 +90,8 @@ export default class LDClientImpl implements LDClient {
 
   private initReject?: (err: Error) => void;
 
+  private rejectionReason: Error | undefined;
+
   private initializedPromise?: Promise<LDClient>;
 
   private logger?: LDLogger;
@@ -231,9 +233,30 @@ export default class LDClientImpl implements LDClient {
   }
 
   waitForInitialization(): Promise<LDClient> {
+    // An initialization promise is only created if someone is going to use that promise.
+    // If we always created an initialization promise, and there was no call waitForInitialization
+    // by the time the promise was rejected, then that would result in an unhandled promise
+    // rejection.
+
+    // Initialization promise was created by a previous call to waitForInitialization.
+    if (this.initializedPromise) {
+      return this.initializedPromise;
+    }
+
+    // Initialization completed before waitForInitialization was called, so we have completed
+    // and there was no promise. So we make a resolved promise and return it.
     if (this.initState === InitState.Initialized) {
-      return Promise.resolve(this);
+      this.initializedPromise = Promise.resolve(this);
+      return this.initializedPromise;
     }
+
+    // Initialization failed before waitForInitialization was called, so we have completed
+    // and there was no promise. So we make a rejected promise and return it.
+    if (this.initState === InitState.Failed) {
+      this.initializedPromise = Promise.reject(this.rejectionReason);
+      return this.initializedPromise;
+    }
+
     if (!this.initializedPromise) {
       this.initializedPromise = new Promise((resolve, reject) => {
         this.initResolve = resolve;
@@ -703,6 +726,7 @@ export default class LDClientImpl implements LDClient {
 
     if (!this.initialized()) {
       this.initState = InitState.Failed;
+      this.rejectionReason = error;
       this.initReject?.(error);
     }
   }