fix: Fix issue that caused the feature store get function impl to be called twice (#178)

Story details: https://app.shortcut.com/launchdarkly/story/207639

This PR adds a new `CacheableStoreProvider` wrapper to Akamai Edge SDK
to prevent redundant requests from being made to the origin store. This
cacheable store provider ensure we're making only a single call per
request thus removing duplicate sub-requests call which could be costly
to our customers. With this fix, there will be only a single call made
per method call `variation`, `variationDetail` or `allFlagsState` is
called.

Also added some unit test to validate duplicate calls are not made.

Author: ctawiah

packages/sdk/akamai-base/src/index.ts

@@ -14,9 +14,7 @@ import {
   LDClient,
   LDOptions,
   EdgeProvider,
-  EdgeFeatureStore,
 } from '@launchdarkly/akamai-edgeworker-sdk-common';
-import { BasicLogger } from '@launchdarkly/js-server-sdk-common';
 
 export * from '@launchdarkly/akamai-edgeworker-sdk-common';
 
@@ -35,15 +33,12 @@ export const init = ({
   options = {},
   sdkKey,
   featureStoreProvider,
-}: AkamaiLDClientParams): LDClient => {
-  const logger = options.logger ?? BasicLogger.get();
-
-  return initEdge({
+}: AkamaiLDClientParams): LDClient =>
+  initEdge({
     sdkKey,
     options,
-    edgeFeatureStore: new EdgeFeatureStore(featureStoreProvider, sdkKey, 'Akamai', logger),
+    featureStoreProvider,
     platformName: 'Akamai EdgeWorker',
     sdkName: '@launchdarkly/akamai-server-base-sdk',
     sdkVersion: '__LD_VERSION__',
   });
-};

packages/sdk/akamai-edgekv/src/index.ts

@@ -9,13 +9,7 @@
  * @packageDocumentation
  */
 
-import {
-  init as initEdge,
-  LDClient,
-  LDOptions,
-  EdgeFeatureStore,
-} from '@launchdarkly/akamai-edgeworker-sdk-common';
-import { BasicLogger } from '@launchdarkly/js-server-sdk-common';
+import { init as initEdge, LDClient, LDOptions } from '@launchdarkly/akamai-edgeworker-sdk-common';
 import EdgeKVProvider from './edgekv/edgeKVProvider';
 
 export * from '@launchdarkly/akamai-edgeworker-sdk-common';
@@ -38,14 +32,12 @@ export const init = ({
   options = {},
   sdkKey,
 }: AkamaiLDClientParams): LDClient => {
-  const logger = options.logger ?? BasicLogger.get();
   const edgekvProvider = new EdgeKVProvider({ namespace, group });
-  const featureStore = new EdgeFeatureStore(edgekvProvider, sdkKey, 'Akamai', logger);
 
   return initEdge({
     sdkKey,
     options,
-    edgeFeatureStore: featureStore,
+    featureStoreProvider: edgekvProvider,
     platformName: 'Akamai EdgeWorker',
     sdkName: '@launchdarkly/akamai-server-edgekv-sdk',
     sdkVersion: '__LD_VERSION__',

packages/shared/akamai-edgeworker-sdk/src/__tests__/index.test.ts

@@ -0,0 +1,112 @@
+import { EdgeProvider, LDLogger, LDMultiKindContext, LDSingleKindContext, init } from '../..';
+
+import * as testData from './testData.json';
+
+const createClient = (sdkKey: string, mockLogger: LDLogger, mockEdgeProvider: EdgeProvider) =>
+  init({
+    sdkKey,
+    options: {
+      logger: mockLogger,
+    },
+    featureStoreProvider: mockEdgeProvider,
+    platformName: 'platform-name',
+    sdkName: 'Akamai',
+    sdkVersion: '0.0.1',
+  });
+
+describe('EdgeWorker', () => {
+  const sdkKey = 'sdkKey';
+
+  const mockLogger: LDLogger = {
+    error: jest.fn(),
+    warn: jest.fn(),
+    info: jest.fn(),
+    debug: jest.fn(),
+  };
+
+  const mockEdgeProvider: EdgeProvider = {
+    get: jest.fn(),
+  };
+
+  const mockGet = mockEdgeProvider.get as jest.Mock;
+
+  beforeEach(() => {
+    mockGet.mockImplementation(() => Promise.resolve(JSON.stringify(testData)));
+  });
+
+  afterEach(() => {
+    jest.resetAllMocks();
+  });
+
+  it('should call edge providers get method only once', async () => {
+    const client = createClient(sdkKey, mockLogger, mockEdgeProvider);
+    await client.waitForInitialization();
+    await client.allFlagsState({ kind: 'multi', l: { key: 'key' } });
+
+    expect(mockGet).toHaveBeenCalledTimes(1);
+  });
+
+  it('should call edge providers get method only 3 times', async () => {
+    const client = createClient(sdkKey, mockLogger, mockEdgeProvider);
+    await client.waitForInitialization();
+
+    const context: LDMultiKindContext = { kind: 'multi', l: { key: 'key' } };
+
+    await client.allFlagsState(context, { clientSideOnly: true });
+    await client.variation('testFlag1', context, false);
+    await client.variationDetail('testFlag1', context, false);
+
+    expect(mockGet).toHaveBeenCalledTimes(3);
+  });
+
+  it('should successfully return data for allFlagsState', async () => {
+    const client = createClient(sdkKey, mockLogger, mockEdgeProvider);
+    await client.waitForInitialization();
+
+    const context: LDMultiKindContext = { kind: 'multi', l: { key: 'key' } };
+
+    const allFlags = await client.allFlagsState(context, { clientSideOnly: true });
+    expect(allFlags.toJSON()).toEqual({
+      $flagsState: {
+        testFlag1: { debugEventsUntilDate: 2000, variation: 0, version: 2 },
+        testFlag2: { debugEventsUntilDate: 2000, variation: 0, version: 2 },
+      },
+      $valid: true,
+      testFlag1: true,
+      testFlag2: true,
+    });
+  });
+
+  it('should should successfully evaluate flags using a flag key', async () => {
+    const client = createClient(sdkKey, mockLogger, mockEdgeProvider);
+    await client.waitForInitialization();
+
+    const context: LDMultiKindContext = { kind: 'multi', l: { key: 'key' } };
+
+    const flagValue = await client.variation('testFlag1', context, false);
+    expect(flagValue).toEqual(true);
+  });
+
+  it('should should successfully return flag evaluation details', async () => {
+    const client = createClient(sdkKey, mockLogger, mockEdgeProvider);
+    await client.waitForInitialization();
+
+    const context: LDMultiKindContext = { kind: 'multi', l: { key: 'key' } };
+
+    const detail = await client.variationDetail('testFlag1', context, false);
+    expect(detail).toEqual({ reason: { kind: 'FALLTHROUGH' }, value: true, variationIndex: 0 });
+  });
+
+  it('should should successfully evaluate flags with segment data', async () => {
+    const client = createClient(sdkKey, mockLogger, mockEdgeProvider);
+    await client.waitForInitialization();
+
+    const context: LDSingleKindContext = {
+      kind: 'user',
+      key: 'return-false-for-segment-target',
+    };
+
+    const flagValue = await client.variation('testFlag3', context, false);
+    expect(flagValue).toEqual(false);
+  });
+});

packages/shared/akamai-edgeworker-sdk/src/__tests__/testData.json

@@ -99,7 +99,7 @@
         "usingMobileKey": true,
         "usingEnvironmentId": true
       },
-      "clientSide": true,
+      "clientSide": false,
       "salt": "aef830243d6640d0a973be89988e008d",
       "trackEvents": false,
       "trackEventsFallthrough": false,
@@ -114,7 +114,7 @@
       "tags": [],
       "creationDate": 1676063792158,
       "key": "testSegment1",
-      "included": [],
+      "included": ["return-false-for-segment-target"],
       "excluded": [],
       "includedContexts": [],
       "excludedContexts": [],

packages/shared/akamai-edgeworker-sdk/src/api/LDClient.ts

@@ -1,14 +1,71 @@
 // eslint-disable-next-line max-classes-per-file
-import { LDClientImpl, LDOptions } from '@launchdarkly/js-server-sdk-common';
+import {
+  LDClient as LDClientType,
+  LDClientImpl,
+  LDOptions,
+  LDContext,
+  LDFlagValue,
+  LDEvaluationDetail,
+  LDFlagsStateOptions,
+  LDFlagsState,
+} from '@launchdarkly/js-server-sdk-common';
 import EdgePlatform from '../platform';
 import { createCallbacks, createOptions } from '../utils';
+import CacheableStoreProvider from '../featureStore/cacheableStoreProvider';
+
+export interface CustomLDOptions extends LDOptions {}
 
 /**
  * The LaunchDarkly Akamai SDK edge client object.
  */
-export default class LDClient extends LDClientImpl {
+class LDClient extends LDClientImpl {
+  private cacheableStoreProvider!: CacheableStoreProvider;
+
   // sdkKey is only used to query featureStore, not to initialize with LD servers
-  constructor(sdkKey: string, platform: EdgePlatform, options: LDOptions) {
+  constructor(
+    sdkKey: string,
+    platform: EdgePlatform,
+    options: LDOptions,
+    storeProvider: CacheableStoreProvider
+  ) {
     super(sdkKey, platform, createOptions(options), createCallbacks());
+    this.cacheableStoreProvider = storeProvider;
+  }
+
+  override waitForInitialization(): Promise<LDClientType> {
+    // we need to resolve the promise immediately because Akamai's runtime doesnt
+    // have a setimeout so everything executes synchronously.
+    return Promise.resolve(this);
+  }
+
+  override async variation(
+    key: string,
+    context: LDContext,
+    defaultValue: LDFlagValue,
+    callback?: (err: any, res: LDFlagValue) => void
+  ): Promise<LDFlagValue> {
+    await this.cacheableStoreProvider.prefetchPayloadFromOriginStore();
+    return super.variation(key, context, defaultValue, callback);
+  }
+
+  override async variationDetail(
+    key: string,
+    context: LDContext,
+    defaultValue: LDFlagValue,
+    callback?: (err: any, res: LDEvaluationDetail) => void
+  ): Promise<LDEvaluationDetail> {
+    await this.cacheableStoreProvider.prefetchPayloadFromOriginStore();
+    return super.variationDetail(key, context, defaultValue, callback);
+  }
+
+  override async allFlagsState(
+    context: LDContext,
+    options?: LDFlagsStateOptions,
+    callback?: (err: Error | null, res: LDFlagsState) => void
+  ): Promise<LDFlagsState> {
+    await this.cacheableStoreProvider.prefetchPayloadFromOriginStore();
+    return super.allFlagsState(context, options, callback);
   }
 }
+
+export default LDClient;

packages/shared/akamai-edgeworker-sdk/src/featureStore/cacheableStoreProvider.ts

@@ -0,0 +1,36 @@
+import { EdgeProvider } from '.';
+
+/**
+ * Wraps around an edge provider to cache a copy of the sdk payload locally an explicit request is made to refetch data from the origin.
+ * The wrapper is neccessary to ensure that we dont make redundant sub-requests from Akamai to fetch an entire environment payload.
+ */
+export default class CacheableStoreProvider implements EdgeProvider {
+  cache: string | null | undefined;
+
+  constructor(private readonly edgeProvider: EdgeProvider, private readonly rootKey: string) {}
+
+  /**
+   * Get data from the edge provider feature store.
+   * @param rootKey
+   * @returns
+   */
+  async get(rootKey: string): Promise<string | null | undefined> {
+    if (!this.cache) {
+      this.cache = await this.edgeProvider.get(rootKey);
+    }
+
+    return this.cache;
+  }
+
+  /**
+   * Invalidates cache and fetch environment payload data from origin. The result of this data is cached in memory.
+   * You should only call this function within a feature store to pre-fetch and cache payload data in environments
+   * where its expensive to make multiple outbound requests to the origin
+   * @param rootKey
+   * @returns
+   */
+  async prefetchPayloadFromOriginStore(rootKey?: string): Promise<string | null | undefined> {
+    this.cache = undefined; // clear the cache so that new data can be fetched from the origin
+    return this.get(rootKey || this.rootKey);
+  }
+}

packages/shared/akamai-edgeworker-sdk/src/featureStore/index.ts

@@ -12,6 +12,13 @@ export interface EdgeProvider {
   get: (rootKey: string) => Promise<string | null | undefined>;
 }
 
+/**
+ * Builds the root key needed to retrieve environment payload from the feature store
+ * @param sdkKey string
+ * @returns
+ */
+export const buildRootKey = (sdkKey: string) => `LD-Env-${sdkKey}`;
+
 export class EdgeFeatureStore implements LDFeatureStore {
   private readonly rootKey: string;
 
@@ -21,7 +28,7 @@ export class EdgeFeatureStore implements LDFeatureStore {
     private readonly description: string,
     private logger: LDLogger
   ) {
-    this.rootKey = `LD-Env-${sdkKey}`;
+    this.rootKey = buildRootKey(this.sdkKey);
   }
 
   async get(

packages/shared/akamai-edgeworker-sdk/src/index.ts

@@ -1,14 +1,10 @@
-import {
-  BasicLogger,
-  LDFeatureStore,
-  LDOptions as LDOptionsCommon,
-} from '@launchdarkly/js-server-sdk-common';
+import { BasicLogger, LDOptions as LDOptionsCommon } from '@launchdarkly/js-server-sdk-common';
 import { validateOptions } from './utils';
 import LDClient from './api/LDClient';
 import EdgePlatform from './platform';
 import createPlatformInfo from './platform/info';
-import type { EdgeProvider } from './featureStore';
-import { EdgeFeatureStore } from './featureStore';
+import { EdgeProvider, buildRootKey, EdgeFeatureStore } from './featureStore';
+import CacheableStoreProvider from './featureStore/cacheableStoreProvider';
 
 /**
  * The Launchdarkly Edge SDKs configuration options. Only logger is officially
@@ -29,18 +25,25 @@ export { EdgeFeatureStore, EdgeProvider, LDOptions, LDOptionsInternal };
 type BaseSDKParams = {
   sdkKey: string;
   options?: LDOptions;
-  edgeFeatureStore: LDFeatureStore;
+  featureStoreProvider: EdgeProvider;
   platformName: string;
   sdkName: string;
   sdkVersion: string;
 };
 
 export const init = (params: BaseSDKParams): LDClient => {
-  const { sdkKey, options = {}, edgeFeatureStore, platformName, sdkName, sdkVersion } = params;
+  const { sdkKey, options = {}, featureStoreProvider, platformName, sdkName, sdkVersion } = params;
 
   const logger = options.logger ?? BasicLogger.get();
-  const ldOptions = {
-    featureStore: edgeFeatureStore,
+
+  const cachableStoreProvider = new CacheableStoreProvider(
+    featureStoreProvider,
+    buildRootKey(sdkKey)
+  );
+  const featureStore = new EdgeFeatureStore(cachableStoreProvider, sdkKey, 'Akamai', logger);
+
+  const ldOptions: LDOptionsCommon = {
+    featureStore,
     logger,
     ...options,
   };
@@ -49,5 +52,5 @@ export const init = (params: BaseSDKParams): LDClient => {
   validateOptions(params.sdkKey, ldOptions);
   const platform = createPlatformInfo(platformName, sdkName, sdkVersion);
 
-  return new LDClient(sdkKey, new EdgePlatform(platform), ldOptions);
+  return new LDClient(sdkKey, new EdgePlatform(platform), ldOptions, cachableStoreProvider);
 };