fix: Improve serialization/deserialization overhead with EdgeFeatureStore (#914)

**Requirements**

- [X] I have added test coverage for new or changed functionality
- [X] I have followed the repository's [pull request submission
guidelines](../blob/main/CONTRIBUTING.md#submitting-pull-requests)
- [X] I have validated my changes against all supported platform
versions

**Related issues**

#907

**Describe the solution you've provided**

Vercel's edge config already returns the JSON.parsed payload but it was
being JSON.stringified and JSON.parsed an additional time which was
adding overhead. EdgeFeatureStore has been improved to be able to handle
either stringified or parsed data so that this extra
serialization/deserialization step doesn't need to occur.

Author: abarker-launchdarkly

packages/sdk/vercel/src/index.ts

@@ -49,11 +49,10 @@ export type { LDClient };
 export const init = (sdkKey: string, edgeConfig: EdgeConfigClient, options: LDOptions = {}) => {
   const logger = options.logger ?? BasicLogger.get();
 
-  // vercel does not support string gets so we have to stringify it
   const edgeProvider: EdgeProvider = {
     get: async (rootKey: string) => {
-      const json = await edgeConfig.get(rootKey);
-      return json ? JSON.stringify(json) : null;
+      const json = await edgeConfig.get<Record<string, any>>(rootKey);
+      return json || null;
     },
   };
 

packages/shared/sdk-server-edge/__tests__/api/EdgeFeatureStore.test.ts

@@ -18,7 +18,6 @@ describe('EdgeFeatureStore', () => {
   let asyncFeatureStore: AsyncStoreFacade;
 
   beforeEach(() => {
-    mockGet.mockImplementation(() => Promise.resolve(JSON.stringify(testData)));
     featureStore = new EdgeFeatureStore(mockEdgeProvider, sdkKey, 'MockEdgeProvider', mockLogger);
     asyncFeatureStore = new AsyncStoreFacade(featureStore);
   });
@@ -27,7 +26,50 @@ describe('EdgeFeatureStore', () => {
     jest.resetAllMocks();
   });
 
-  describe('get', () => {
+  describe('get (string payload)', () => {
+    beforeEach(() => {
+      mockGet.mockImplementation(() => Promise.resolve(JSON.stringify(testData)));
+    });
+
+    test('get flag', async () => {
+      const flag = await asyncFeatureStore.get({ namespace: 'features' }, 'testFlag1');
+
+      expect(mockGet).toHaveBeenCalledWith(kvKey);
+      expect(flag).toMatchObject(testData.flags.testFlag1);
+    });
+
+    test('invalid flag key', async () => {
+      const flag = await asyncFeatureStore.get({ namespace: 'features' }, 'invalid');
+
+      expect(flag).toBeUndefined();
+    });
+
+    test('get segment', async () => {
+      const segment = await asyncFeatureStore.get({ namespace: 'segments' }, 'testSegment1');
+
+      expect(mockGet).toHaveBeenCalledWith(kvKey);
+      expect(segment).toMatchObject(testData.segments.testSegment1);
+    });
+
+    test('invalid segment key', async () => {
+      const segment = await asyncFeatureStore.get({ namespace: 'segments' }, 'invalid');
+
+      expect(segment).toBeUndefined();
+    });
+
+    test('invalid kv key', async () => {
+      mockGet.mockImplementation(() => Promise.resolve(null));
+      const flag = await asyncFeatureStore.get({ namespace: 'features' }, 'testFlag1');
+
+      expect(flag).toBeNull();
+    });
+  });
+
+  describe('get (object payload)', () => {
+    beforeEach(() => {
+      mockGet.mockImplementation(() => Promise.resolve(testData));
+    });
+
     test('get flag', async () => {
       const flag = await asyncFeatureStore.get({ namespace: 'features' }, 'testFlag1');
 
@@ -62,7 +104,44 @@ describe('EdgeFeatureStore', () => {
     });
   });
 
-  describe('all', () => {
+  describe('all (string payload)', () => {
+    beforeEach(() => {
+      mockGet.mockImplementation(() => Promise.resolve(JSON.stringify(testData)));
+    });
+
+    test('all flags', async () => {
+      const flags = await asyncFeatureStore.all({ namespace: 'features' });
+
+      expect(mockGet).toHaveBeenCalledWith(kvKey);
+      expect(flags).toMatchObject(testData.flags);
+    });
+
+    test('all segments', async () => {
+      const segment = await asyncFeatureStore.all({ namespace: 'segments' });
+
+      expect(mockGet).toHaveBeenCalledWith(kvKey);
+      expect(segment).toMatchObject(testData.segments);
+    });
+
+    test('invalid DataKind', async () => {
+      const flag = await asyncFeatureStore.all({ namespace: 'InvalidDataKind' });
+
+      expect(flag).toEqual({});
+    });
+
+    test('invalid kv key', async () => {
+      mockGet.mockImplementation(() => Promise.resolve(null));
+      const segment = await asyncFeatureStore.all({ namespace: 'segments' });
+
+      expect(segment).toEqual({});
+    });
+  });
+
+  describe('all (object payload)', () => {
+    beforeEach(() => {
+      mockGet.mockImplementation(() => Promise.resolve(testData));
+    });
+
     test('all flags', async () => {
       const flags = await asyncFeatureStore.all({ namespace: 'features' });
 

packages/shared/sdk-server-edge/src/api/EdgeFeatureStore.ts

@@ -6,12 +6,12 @@ import type {
   LDFeatureStoreKindData,
   LDLogger,
 } from '@launchdarkly/js-server-sdk-common';
-import { deserializePoll, noop } from '@launchdarkly/js-server-sdk-common';
+import { deserializePoll, noop, reviveFullPayload } from '@launchdarkly/js-server-sdk-common';
 
 import Cache from './cache';
 
 export interface EdgeProvider {
-  get: (rootKey: string) => Promise<string | null | undefined>;
+  get: (rootKey: string) => Promise<string | Record<string, any> | null | undefined>;
 }
 
 export class EdgeFeatureStore implements LDFeatureStore {
@@ -96,7 +96,11 @@ export class EdgeFeatureStore implements LDFeatureStore {
       throw new Error(`${this._rootKey} is not found in KV.`);
     }
 
-    payload = deserializePoll(providerData);
+    payload =
+      typeof providerData === 'string'
+        ? deserializePoll(providerData)
+        : reviveFullPayload(providerData);
+
     if (!payload) {
       throw new Error(`Error deserializing ${this._rootKey}`);
     }

packages/shared/sdk-server/src/store/index.ts

@@ -1,7 +1,7 @@
 import AsyncStoreFacade from './AsyncStoreFacade';
 import AsyncTransactionalStoreFacade from './AsyncTransactionalStoreFacade';
 import PersistentDataStoreWrapper from './PersistentDataStoreWrapper';
-import { deserializePoll } from './serialization';
+import { deserializePoll, reviveFullPayload } from './serialization';
 import TransactionalFeatureStore from './TransactionalFeatureStore';
 
 export {
@@ -10,4 +10,5 @@ export {
   PersistentDataStoreWrapper,
   TransactionalFeatureStore,
   deserializePoll,
+  reviveFullPayload,
 };

packages/shared/sdk-server/src/store/serialization.ts

@@ -237,6 +237,27 @@ function tryParse(data: string): any {
   }
 }
 
+/**
+ * This function is intended for usage inside LaunchDarkly SDKs.
+ * This function should NOT be used by customer applications.
+ * This function may be changed or removed without a major version.
+ *
+ * @param payload Payload data from launchdarkly.
+ * @returns The revived and processed data.
+ */
+export function reviveFullPayload(payload: Record<string, any>): FlagsAndSegments {
+  const flagsAndSegments = payload as FlagsAndSegments;
+  Object.values(flagsAndSegments?.flags || []).forEach((flag) => {
+    processFlag(flag);
+  });
+
+  Object.values(flagsAndSegments?.segments || []).forEach((segment) => {
+    processSegment(segment);
+  });
+
+  return flagsAndSegments;
+}
+
 /**
  * @internal
  */
@@ -253,13 +274,7 @@ export function deserializeAll(data: string): AllData | undefined {
     return undefined;
   }
 
-  Object.values(parsed?.data?.flags || []).forEach((flag) => {
-    processFlag(flag);
-  });
-
-  Object.values(parsed?.data?.segments || []).forEach((segment) => {
-    processSegment(segment);
-  });
+  reviveFullPayload(parsed?.data);
   return parsed;
 }
 
@@ -278,13 +293,7 @@ export function deserializePoll(data: string): FlagsAndSegments | undefined {
     return undefined;
   }
 
-  Object.values(parsed?.flags || []).forEach((flag) => {
-    processFlag(flag);
-  });
-
-  Object.values(parsed?.segments || []).forEach((segment) => {
-    processSegment(segment);
-  });
+  reviveFullPayload(parsed);
   return parsed;
 }