Add bandit assignment cache (#120)

* test in place

* use an assignment cache to deduplicate bandit action assignments

* feedback from self-review of PR

* fix type and lint errors

* bump minor version

* just hash whole cache value

* feedback from PR

* dedupe when queuing assignments as well

* Bump package.json version correctly

Co-authored-by: Leo Romanovsky <[email protected]>

* feedback from PR

---------

Co-authored-by: Leo Romanovsky <[email protected]>

Author: aarsilv

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@eppo/js-client-sdk-common",
-  "version": "4.0.2",
+  "version": "4.1.0",
   "description": "Eppo SDK for client-side JavaScript applications (base for both web and react native)",
   "main": "dist/index.js",
   "files": [

src/cache/abstract-assignment-cache.spec.ts

@@ -5,7 +5,7 @@ import {
 } from './abstract-assignment-cache';
 
 describe('NonExpiringInMemoryAssignmentCache', () => {
-  it('read and write entries', () => {
+  it('read and write variation entries', () => {
     const cache = new NonExpiringInMemoryAssignmentCache();
     const key1 = { subjectKey: 'a', flagKey: 'b', allocationKey: 'c', variationKey: 'd' };
     const key2 = { subjectKey: '1', flagKey: '2', allocationKey: '3', variationKey: '4' };
@@ -14,7 +14,7 @@ describe('NonExpiringInMemoryAssignmentCache', () => {
     expect(cache.has(key2)).toBeFalsy();
     cache.set(key2);
     expect(cache.has(key2)).toBeTruthy();
-    // this makes an assumption about the internal implementation of the cache, which is not ideal
+    // this makes an assumption about the internal implementation of the cache, which is not ideal,
     // but it's the only way to test the cache without exposing the internal state
     expect(Array.from(cache.entries())).toEqual([
       [assignmentCacheKeyToString(key1), assignmentCacheValueToString(key1)],
@@ -24,4 +24,24 @@ describe('NonExpiringInMemoryAssignmentCache', () => {
     expect(cache.has({ ...key1, allocationKey: 'c1' })).toBeFalsy();
     expect(cache.has({ ...key2, variationKey: 'd1' })).toBeFalsy();
   });
+
+  it('read and write bandit entries', () => {
+    const cache = new NonExpiringInMemoryAssignmentCache();
+    const key1 = { subjectKey: 'a', flagKey: 'b', banditKey: 'c', actionKey: 'd' };
+    const key2 = { subjectKey: '1', flagKey: '2', banditKey: '3', actionKey: '4' };
+    cache.set(key1);
+    expect(cache.has(key1)).toBeTruthy();
+    expect(cache.has(key2)).toBeFalsy();
+    cache.set(key2);
+    expect(cache.has(key2)).toBeTruthy();
+    // this makes an assumption about the internal implementation of the cache, which is not ideal,
+    // but it's the only way to test the cache without exposing the internal state
+    expect(Array.from(cache.entries())).toEqual([
+      [assignmentCacheKeyToString(key1), assignmentCacheValueToString(key1)],
+      [assignmentCacheKeyToString(key2), assignmentCacheValueToString(key2)],
+    ]);
+
+    expect(cache.has({ ...key1, banditKey: 'c1' })).toBeFalsy();
+    expect(cache.has({ ...key2, actionKey: 'd1' })).toBeFalsy();
+  });
 });

src/cache/abstract-assignment-cache.ts

@@ -2,28 +2,37 @@ import { getMD5Hash } from '../obfuscation';
 
 import { LRUCache } from './lru-cache';
 
-export type AssignmentCacheValue = {
-  allocationKey: string;
-  variationKey: string;
-};
-
+/**
+ * Assignment cache keys are only on the subject and flag level, while the entire value is used
+ * for uniqueness checking. This way that if an assigned variation or bandit action changes for a
+ * flag, it evicts the old one. Then, if an older assignment is later reassigned, it will be treated
+ * as new.
+ */
 export type AssignmentCacheKey = {
   subjectKey: string;
   flagKey: string;
 };
 
+export type CacheKeyPair<T extends string, U extends string> = {
+  [K in T]: string;
+} & {
+  [K in U]: string;
+};
+
+type VariationCacheValue = CacheKeyPair<'allocationKey', 'variationKey'>;
+type BanditCacheValue = CacheKeyPair<'banditKey', 'actionKey'>;
+export type AssignmentCacheValue = VariationCacheValue | BanditCacheValue;
+
 export type AssignmentCacheEntry = AssignmentCacheKey & AssignmentCacheValue;
 
 /** Converts an {@link AssignmentCacheKey} to a string. */
 export function assignmentCacheKeyToString({ subjectKey, flagKey }: AssignmentCacheKey): string {
   return getMD5Hash([subjectKey, flagKey].join(';'));
 }
 
-export function assignmentCacheValueToString({
-  allocationKey,
-  variationKey,
-}: AssignmentCacheValue): string {
-  return getMD5Hash([allocationKey, variationKey].join(';'));
+/** Converts an {@link AssignmentCacheValue} to a string. */
+export function assignmentCacheValueToString(cacheValue: AssignmentCacheValue): string {
+  return getMD5Hash(Object.values(cacheValue).join(';'));
 }
 
 export interface AsyncMap<K, V> {

src/client/eppo-client-with-bandits.spec.ts

@@ -8,10 +8,11 @@ import {
 } from '../../test/testHelpers';
 import ApiEndpoints from '../api-endpoints';
 import { IAssignmentEvent, IAssignmentLogger } from '../assignment-logger';
-import { BanditEvaluator } from '../bandit-evaluator';
+import { BanditEvaluation, BanditEvaluator } from '../bandit-evaluator';
 import { IBanditEvent, IBanditLogger } from '../bandit-logger';
 import ConfigurationRequestor from '../configuration-requestor';
 import { MemoryOnlyConfigurationStore } from '../configuration-store/memory.store';
+import { Evaluator, FlagEvaluation } from '../evaluator';
 import {
   AllocationEvaluationCode,
   IFlagEvaluationDetails,
@@ -20,7 +21,7 @@ import FetchHttpClient from '../http-client';
 import { BanditVariation, BanditParameters, Flag } from '../interfaces';
 import { Attributes, ContextAttributes } from '../types';
 
-import EppoClient from './eppo-client';
+import EppoClient, { IAssignmentDetails } from './eppo-client';
 
 describe('EppoClient Bandits E2E test', () => {
   const flagStore = new MemoryOnlyConfigurationStore<Flag>();
@@ -204,6 +205,8 @@ describe('EppoClient Bandits E2E test', () => {
     });
 
     it('Flushed queued logging events when a logger is set', () => {
+      client.useLRUInMemoryAssignmentCache(5);
+      client.useLRUInMemoryBanditAssignmentCache(5);
       client.setAssignmentLogger(null as unknown as IAssignmentLogger);
       client.setBanditLogger(null as unknown as IBanditLogger);
       const banditAssignment = client.getBanditAction(
@@ -220,6 +223,20 @@ describe('EppoClient Bandits E2E test', () => {
       expect(mockLogAssignment).not.toHaveBeenCalled();
       expect(mockLogBanditAction).not.toHaveBeenCalled();
 
+      const repeatAssignment = client.getBanditAction(
+        flagKey,
+        subjectKey,
+        subjectAttributes,
+        actions,
+        'control',
+      );
+
+      expect(repeatAssignment.variation).toBe('banner_bandit');
+      expect(repeatAssignment.action).toBe('adidas');
+
+      expect(mockLogAssignment).not.toHaveBeenCalled();
+      expect(mockLogBanditAction).not.toHaveBeenCalled();
+
       client.setAssignmentLogger({ logAssignment: mockLogAssignment });
       client.setBanditLogger({ logBanditAction: mockLogBanditAction });
 
@@ -429,5 +446,137 @@ describe('EppoClient Bandits E2E test', () => {
         expect(mockLogBanditAction.mock.calls[1][0].actionProbability).toBeCloseTo(0.256);
       });
     });
+
+    describe('Assignment logging deduplication', () => {
+      let mockEvaluateFlag: jest.SpyInstance;
+      let mockEvaluateBandit: jest.SpyInstance;
+      // The below two variables allow easily changing what the mock evaluation functions return throughout the test
+      let variationToReturn: string;
+      let actionToReturn: string | null;
+
+      // Convenience method for repeatedly making the exact same assignment call
+      function requestClientBanditAction(): Omit<IAssignmentDetails<string>, 'evaluationDetails'> {
+        return client.getBanditAction(
+          flagKey,
+          subjectKey,
+          subjectAttributes,
+          ['toyota', 'honda'],
+          'control',
+        );
+      }
+
+      beforeAll(() => {
+        mockEvaluateFlag = jest
+          .spyOn(Evaluator.prototype, 'evaluateFlag')
+          .mockImplementation(() => {
+            return {
+              flagKey,
+              subjectKey,
+              subjectAttributes,
+              allocationKey: 'mock-allocation',
+              variation: { key: variationToReturn, value: variationToReturn },
+              extraLogging: {},
+              doLog: true,
+              flagEvaluationDetails: {
+                flagEvaluationCode: 'MATCH',
+                flagEvaluationDescription: 'Mocked evaluation',
+              },
+            } as FlagEvaluation;
+          });
+
+        mockEvaluateBandit = jest
+          .spyOn(BanditEvaluator.prototype, 'evaluateBandit')
+          .mockImplementation(() => {
+            return {
+              flagKey,
+              subjectKey,
+              subjectAttributes: { numericAttributes: {}, categoricalAttributes: {} },
+              actionKey: actionToReturn,
+              actionAttributes: { numericAttributes: {}, categoricalAttributes: {} },
+              actionScore: 10,
+              actionWeight: 0.5,
+              gamma: 1.0,
+              optimalityGap: 5,
+            } as BanditEvaluation;
+          });
+      });
+
+      beforeEach(() => {
+        client.useNonExpiringInMemoryAssignmentCache();
+        client.useNonExpiringInMemoryBanditAssignmentCache();
+      });
+
+      afterEach(() => {
+        client.disableAssignmentCache();
+        client.disableBanditAssignmentCache();
+      });
+
+      afterAll(() => {
+        mockEvaluateFlag.mockClear();
+        mockEvaluateBandit.mockClear();
+      });
+
+      it('handles bandit actions appropriately', async () => {
+        // First assign to non-bandit variation
+        variationToReturn = 'non-bandit';
+        actionToReturn = null;
+        const firstNonBanditAssignment = requestClientBanditAction();
+
+        expect(firstNonBanditAssignment.variation).toBe('non-bandit');
+        expect(firstNonBanditAssignment.action).toBeNull();
+        expect(mockLogAssignment).toHaveBeenCalledTimes(1); // new variation assignment
+        expect(mockLogBanditAction).not.toHaveBeenCalled(); // no bandit assignment
+
+        // Assign bandit action
+        variationToReturn = 'banner_bandit';
+        actionToReturn = 'toyota';
+        const firstBanditAssignment = requestClientBanditAction();
+
+        expect(firstBanditAssignment.variation).toBe('banner_bandit');
+        expect(firstBanditAssignment.action).toBe('toyota');
+        expect(mockLogAssignment).toHaveBeenCalledTimes(2); // new variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(1); // new bandit assignment
+
+        // Repeat bandit action assignment
+        variationToReturn = 'banner_bandit';
+        actionToReturn = 'toyota';
+        const secondBanditAssignment = requestClientBanditAction();
+
+        expect(secondBanditAssignment.variation).toBe('banner_bandit');
+        expect(secondBanditAssignment.action).toBe('toyota');
+        expect(mockLogAssignment).toHaveBeenCalledTimes(2); // repeat variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(1); // repeat bandit assignment
+
+        // New bandit action assignment
+        variationToReturn = 'banner_bandit';
+        actionToReturn = 'honda';
+        const thirdBanditAssignment = requestClientBanditAction();
+
+        expect(thirdBanditAssignment.variation).toBe('banner_bandit');
+        expect(thirdBanditAssignment.action).toBe('honda');
+        expect(mockLogAssignment).toHaveBeenCalledTimes(2); // repeat variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(2); // new bandit assignment
+
+        // Flip-flop to an earlier action assignment
+        variationToReturn = 'banner_bandit';
+        actionToReturn = 'toyota';
+        const fourthBanditAssignment = requestClientBanditAction();
+
+        expect(fourthBanditAssignment.variation).toBe('banner_bandit');
+        expect(fourthBanditAssignment.action).toBe('toyota');
+        expect(mockLogAssignment).toHaveBeenCalledTimes(2); // repeat variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(3); // "new" bandit assignment
+
+        // Flip-flop back to non-bandit assignment
+        variationToReturn = 'non-bandit';
+        actionToReturn = null;
+        const secondNonBanditAssignment = requestClientBanditAction();
+
+        expect(secondNonBanditAssignment.variation).toBe('non-bandit');
+        expect(secondNonBanditAssignment.action).toBeNull();
+        expect(mockLogAssignment).toHaveBeenCalledTimes(3); // "new" variation assignment
+        expect(mockLogBanditAction).toHaveBeenCalledTimes(3); // no bandit assignment
+      });
+    });
   });
 });

src/client/eppo-client.ts

@@ -75,6 +75,7 @@ export default class EppoClient {
   private banditLogger?: IBanditLogger;
   private isGracefulFailureMode = true;
   private assignmentCache?: AssignmentCache;
+  private banditAssignmentCache?: AssignmentCache;
   private requestPoller?: IPoller;
   private readonly evaluator = new Evaluator();
   private readonly banditEvaluator = new BanditEvaluator();
@@ -651,16 +652,34 @@ export default class EppoClient {
   }
 
   private logBanditAction(banditEvent: IBanditEvent): void {
-    if (!this.banditLogger) {
-      // No bandit logger set; enqueue the event in case a logger is later set
-      if (this.queuedBanditEvents.length < MAX_EVENT_QUEUE_SIZE) {
-        this.queuedBanditEvents.push(banditEvent);
-      }
+    // First we check if this bandit action has been logged before
+    const subjectKey = banditEvent.subject;
+    const flagKey = banditEvent.featureFlag;
+    const banditKey = banditEvent.bandit;
+    const actionKey = banditEvent.action ?? '__eppo_no_action';
+
+    const banditAssignmentCacheProperties = {
+      flagKey,
+      subjectKey,
+      banditKey,
+      actionKey,
+    };
+
+    if (this.banditAssignmentCache?.has(banditAssignmentCacheProperties)) {
+      // Ignore repeat assignment
       return;
     }
-    // If here, we have a logger
+
+    // If here, we have a logger and a new assignment to be logged
     try {
-      this.banditLogger.logBanditAction(banditEvent);
+      if (this.banditLogger) {
+        this.banditLogger.logBanditAction(banditEvent);
+      } else if (this.queuedBanditEvents.length < MAX_EVENT_QUEUE_SIZE) {
+        // If no logger defined, queue up the events (up to a max) to flush if a logger is later defined
+        this.queuedBanditEvents.push(banditEvent);
+      }
+      // Record in the assignment cache, if active, to deduplicate subsequent repeat assignments
+      this.banditAssignmentCache?.set(banditAssignmentCacheProperties);
     } catch (err) {
       logger.warn('Error encountered logging bandit action', err);
     }
@@ -904,6 +923,22 @@ export default class EppoClient {
     this.assignmentCache = cache;
   }
 
+  public disableBanditAssignmentCache() {
+    this.banditAssignmentCache = undefined;
+  }
+
+  public useNonExpiringInMemoryBanditAssignmentCache() {
+    this.banditAssignmentCache = new NonExpiringInMemoryAssignmentCache();
+  }
+
+  public useLRUInMemoryBanditAssignmentCache(maxSize: number) {
+    this.banditAssignmentCache = new LRUInMemoryAssignmentCache(maxSize);
+  }
+
+  public useCustomBanditAssignmentCache(cache: AssignmentCache) {
+    this.banditAssignmentCache = cache;
+  }
+
   public setIsGracefulFailureMode(gracefulFailureMode: boolean) {
     this.isGracefulFailureMode = gracefulFailureMode;
   }
@@ -956,14 +991,14 @@ export default class EppoClient {
       }
     }
 
-    // assignment logger may be null while waiting for initialization
-    if (!this.assignmentLogger) {
-      this.queuedAssignmentEvents.length < MAX_EVENT_QUEUE_SIZE &&
-        this.queuedAssignmentEvents.push(event);
-      return;
-    }
     try {
-      this.assignmentLogger.logAssignment(event);
+      if (this.assignmentLogger) {
+        this.assignmentLogger.logAssignment(event);
+      } else if (this.queuedAssignmentEvents.length < MAX_EVENT_QUEUE_SIZE) {
+        // assignment logger may be null while waiting for initialization, queue up events (up to a max)
+        // to be flushed when set
+        this.queuedAssignmentEvents.push(event);
+      }
       this.assignmentCache?.set({
         flagKey,
         subjectKey,