fix: Improve serialization/deserialization overhead with EdgeFeatureStore

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;
 }