feat: Optimize segment lookup for large segments. (#235)

Author: kinyoklion

packages/shared/sdk-server/__tests__/store/serialization.test.ts

@@ -6,6 +6,7 @@ import {
   deserializePatch,
   replacer,
   reviver,
+  serializeSegment,
 } from '../../src/store/serialization';
 
 const flagWithAttributeNameInClause = {
@@ -339,3 +340,109 @@ it('given bad json', () => {
   expect(deserializePatch(data)).toBeUndefined();
   expect(deserializeDelete(data)).toBeUndefined();
 });
+
+it('deserialization creates a set for a large number of includes/excludes', () => {
+  const included = [...Array(500).keys()].map((i) => (i + 1).toString());
+  const excluded = [...Array(500).keys()].map((i) => (i + 10).toString());
+
+  const jsonString = makeSerializedPatchData(undefined, {
+    key: 'test-segment-1',
+    included,
+    excluded,
+    includedContexts: [],
+    excludedContexts: [],
+    salt: 'saltyA',
+    rules: [],
+    version: 0,
+    deleted: false,
+  });
+
+  const res = deserializePatch(jsonString);
+  const segment = res?.data as Segment;
+  expect(segment.included).toBeUndefined();
+  expect(segment.excluded).toBeUndefined();
+
+  expect([...segment.generated_includedSet!]).toEqual(included);
+  expect([...segment.generated_excludedSet!]).toEqual(excluded);
+});
+
+it('deserialization creates a set for a large number of included/excluded context values', () => {
+  const included = [...Array(500).keys()].map((i) => (i + 10).toString());
+  const excluded = [...Array(500).keys()].map((i) => (i + 1).toString());
+
+  const jsonString = makeSerializedPatchData(undefined, {
+    key: 'test-segment-1',
+    included: [],
+    excluded: [],
+    includedContexts: [{ contextKind: 'org', values: included }],
+    excludedContexts: [{ contextKind: 'user', values: excluded }],
+    salt: 'saltyA',
+    rules: [],
+    version: 0,
+    deleted: false,
+  });
+
+  const res = deserializePatch(jsonString);
+  const segment = res?.data as Segment;
+
+  expect([...segment.includedContexts![0].generated_valuesSet!]).toEqual(included);
+  expect([...segment.excludedContexts![0].generated_valuesSet!]).toEqual(excluded);
+});
+
+it('serialization converts sets back to arrays for included/excluded', () => {
+  const included = [...Array(500).keys()].map((i) => (i + 1).toString());
+  const excluded = [...Array(500).keys()].map((i) => (i + 10).toString());
+
+  const jsonString = makeSerializedPatchData(undefined, {
+    key: 'test-segment-1',
+    included,
+    excluded,
+    includedContexts: [],
+    excludedContexts: [],
+    salt: 'saltyA',
+    rules: [],
+    version: 0,
+    deleted: false,
+  });
+
+  const res = deserializePatch(jsonString);
+  const segment = res?.data as Segment;
+
+  const serializedSegment = serializeSegment(segment);
+  // Just json parse. We don't want it to automatically re-populate the sets.
+  const jsonDeserialized = JSON.parse(serializedSegment);
+
+  expect(jsonDeserialized.included).toEqual(included);
+  expect(jsonDeserialized.excluded).toEqual(excluded);
+  expect(jsonDeserialized.generated_includedSet).toBeUndefined();
+  expect(jsonDeserialized.generated_excludedSet).toBeUndefined();
+});
+
+it('serialization converts sets back to arrays for includedContexts/excludedContexts', () => {
+  const included = [...Array(500).keys()].map((i) => (i + 1).toString());
+  const excluded = [...Array(500).keys()].map((i) => (i + 10).toString());
+
+  const jsonString = makeSerializedPatchData(undefined, {
+    key: 'test-segment-1',
+    included: [],
+    excluded: [],
+    includedContexts: [{ contextKind: 'org', values: included }],
+    excludedContexts: [{ contextKind: 'user', values: excluded }],
+    salt: 'saltyA',
+    rules: [],
+    version: 0,
+    deleted: false,
+  });
+
+  const res = deserializePatch(jsonString);
+  const segment = res?.data as Segment;
+
+  const serializedSegment = serializeSegment(segment);
+  // Just json parse. We don't want it to automatically re-populate the sets.
+  const jsonDeserialized = JSON.parse(serializedSegment);
+
+  expect(jsonDeserialized.includedContexts[0].values).toEqual(included);
+  expect(jsonDeserialized.excludedContexts[0].values).toEqual(excluded);
+  expect(jsonDeserialized.includedContexts[0].generated_valuesSet).toBeUndefined();
+  expect(jsonDeserialized.excludedContexts[0].generated_valuesSet).toBeUndefined();
+});

packages/shared/sdk-server/src/evaluation/data/Segment.ts

@@ -7,6 +7,7 @@ import { Versioned } from './Versioned';
 export interface Segment extends Versioned {
   included?: string[];
   excluded?: string[];
+
   includedContexts?: SegmentTarget[];
   excludedContexts?: SegmentTarget[];
   rules?: SegmentRule[];
@@ -17,4 +18,9 @@ export interface Segment extends Versioned {
 
   // This field is not part of the schema, but it is populated during parsing.
   bucketByAttributeReference?: AttributeReference;
+
+  // When there are a large number targets for a segment then
+  // we put them into sets during de-serialization.
+  generated_includedSet?: Set<string>;
+  generated_excludedSet?: Set<string>;
 }

packages/shared/sdk-server/src/evaluation/data/SegmentTarget.ts

@@ -1,4 +1,8 @@
 export interface SegmentTarget {
   contextKind: string;
   values: string[];
+  // When there are a large number of values we put them into a set.
+  // This set is generated during deserialization, and changed back to a list
+  // during serialization.
+  generated_valuesSet?: Set<string>;
 }

packages/shared/sdk-server/src/evaluation/matchSegmentTargets.ts

@@ -7,18 +7,33 @@ function segmentSearch(
   context: Context,
   contextTargets?: SegmentTarget[],
   userTargets?: string[],
+  userTargetSet?: Set<string>,
 ): boolean {
   if (contextTargets) {
     for (let targetIndex = 0; targetIndex < contextTargets.length; targetIndex += 1) {
       const target = contextTargets[targetIndex];
       const key = context.key(target.contextKind);
-      if (key && target.values.includes(key)) {
-        return true;
+      if (key) {
+        if (target.generated_valuesSet) {
+          // Only check generated_valuesSet if present.
+          if (target.generated_valuesSet.has(key)) {
+            return true;
+          }
+        } else if (target.values.includes(key)) {
+          return true;
+        }
       }
     }
   }
 
-  if (userTargets) {
+  if (userTargetSet) {
+    const userKey = context.key('user');
+    if (userKey) {
+      if (userTargetSet.has(userKey)) {
+        return true;
+      }
+    }
+  } else if (userTargets) {
     const userKey = context.key('user');
     if (userKey) {
       if (userTargets.includes(userKey)) {
@@ -33,11 +48,21 @@ export default function matchSegmentTargets(
   segment: Segment,
   context: Context,
 ): boolean | undefined {
-  const included = segmentSearch(context, segment.includedContexts, segment.included);
+  const included = segmentSearch(
+    context,
+    segment.includedContexts,
+    segment.included,
+    segment.generated_includedSet,
+  );
   if (included) {
     return true;
   }
-  const excluded = segmentSearch(context, segment.excludedContexts, segment.excluded);
+  const excluded = segmentSearch(
+    context,
+    segment.excludedContexts,
+    segment.excluded,
+    segment.generated_excludedSet,
+  );
   if (excluded) {
     // The match was an exclusion, so it should be negated.
     return !excluded;

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

@@ -10,6 +10,9 @@ import { Rollout } from '../evaluation/data/Rollout';
 import { Segment } from '../evaluation/data/Segment';
 import VersionedDataKinds, { VersionedDataKind } from './VersionedDataKinds';
 
+// The max size where we use an array instead of a set.
+const TARGET_LIST_ARRAY_CUTOFF = 100;
+
 /**
  * @internal
  */
@@ -51,6 +54,30 @@ export function replacer(this: any, key: string, value: any): any {
       return undefined;
     }
   }
+  if (value.generated_includedSet) {
+    value.included = [...value.generated_includedSet];
+    delete value.generated_includedSet;
+  }
+  if (value.generated_excludedSet) {
+    value.excluded = [...value.generated_excludedSet];
+    delete value.generated_excludedSet;
+  }
+  if (value.includedContexts) {
+    value.includedContexts.forEach((target: any) => {
+      if (target.generated_valuesSet) {
+        target.values = [...target.generated_valuesSet];
+      }
+      delete target.generated_valuesSet;
+    });
+  }
+  if (value.excludedContexts) {
+    value.excludedContexts.forEach((target: any) => {
+      if (target.generated_valuesSet) {
+        target.values = [...target.generated_valuesSet];
+      }
+      delete target.generated_valuesSet;
+    });
+  }
   return value;
 }
 
@@ -104,6 +131,35 @@ export function processFlag(flag: Flag) {
  * @internal
  */
 export function processSegment(segment: Segment) {
+  if (segment?.included?.length && segment.included.length > TARGET_LIST_ARRAY_CUTOFF) {
+    segment.generated_includedSet = new Set(segment.included);
+    delete segment.included;
+  }
+  if (segment?.excluded?.length && segment.excluded.length > TARGET_LIST_ARRAY_CUTOFF) {
+    segment.generated_excludedSet = new Set(segment.excluded);
+    delete segment.excluded;
+  }
+
+  if (segment?.includedContexts?.length) {
+    segment.includedContexts.forEach((target) => {
+      if (target?.values?.length && target.values.length > TARGET_LIST_ARRAY_CUTOFF) {
+        target.generated_valuesSet = new Set(target.values);
+        // Currently typing is non-optional, so we don't delete it.
+        target.values = [];
+      }
+    });
+  }
+
+  if (segment?.excludedContexts?.length) {
+    segment.excludedContexts.forEach((target) => {
+      if (target?.values?.length && target.values.length > TARGET_LIST_ARRAY_CUTOFF) {
+        target.generated_valuesSet = new Set(target.values);
+        // Currently typing is non-optional, so we don't delete it.
+        target.values = [];
+      }
+    });
+  }
+
   segment?.rules?.forEach((rule) => {
     if (rule.bucketBy) {
       // Rules before U2C would have had literals for attributes.