Merge branch 'perf/local-eval-regression' into feat/flags-context-values

# Conflicts:
#	flagsmith/local_evaluation.py

Author: khvn26

flagsmith/flagsmith.py

@@ -18,7 +18,6 @@
 from flagsmith.local_evaluation import LocalEvaluation
 from flagsmith.mappers import (
     map_context_and_identity_data_to_context,
-    map_environment_document_to_context,
     map_environment_document_to_environment_updated_at,
     map_segment_results_to_identity_segments,
     resolve_trait_values,
@@ -150,10 +149,8 @@ def __init__(
             )
 
         if self.offline_handler:
-            self._local_evaluation = LocalEvaluation(
-                map_environment_document_to_context(
-                    self.offline_handler.get_environment()
-                )
+            self._local_evaluation = LocalEvaluation.from_environment_document(
+                self.offline_handler.get_environment()
             )
 
         if not self.offline_mode:
@@ -467,8 +464,8 @@ def update_environment(self) -> None:
             logger.exception("Error retrieving environment document from API")
         else:
             try:
-                self._local_evaluation = LocalEvaluation(
-                    map_environment_document_to_context(environment_data)
+                self._local_evaluation = LocalEvaluation.from_environment_document(
+                    environment_data
                 )
                 self._environment_updated_at = (
                     map_environment_document_to_environment_updated_at(

flagsmith/local_evaluation.py

@@ -7,10 +7,11 @@
 from flag_engine import engine
 from flag_engine.context.types import SegmentContext, SegmentRule
 
+from flagsmith.api.types import EnvironmentModel
 from flagsmith.mappers import (
-    get_identity_overrides_identifiers,
     map_context_and_identity_data_to_context,
-    map_identifiers_to_identity_overrides_rule,
+    map_environment_document_to_context,
+    map_environment_document_to_identity_overrides_segments,
 )
 from flagsmith.models import (
     Flag,
@@ -61,40 +62,49 @@ class LocalEvaluation:
     Built once each time the environment document is refreshed, so the
     work here is amortised across every evaluation until the next refresh.
 
-    Identity overrides are mapped to segments matching the identifiers
-    they apply to. Left in the context, every one of them would be
-    matched on every identity evaluation. Instead, they're indexed by
-    identifier, and each identity is evaluated with only its own.
+    Identity overrides are kept out of the context, by identifier. Each
+    identity is evaluated with its own identity overrides segment only.
     """
 
-    def __init__(self, context: SDKEvaluationContext) -> None:
-        segments: SegmentContexts = {}
+    def __init__(
+        self,
+        context: SDKEvaluationContext,
+        identity_overrides_segments: typing.Optional[
+            typing.Mapping[str, SegmentContext[SegmentMetadata, FeatureMetadata]]
+        ] = None,
+    ) -> None:
+        self.context = context
         self._identity_free_segments: SegmentContexts = {}
-        self._identity_overrides_segments: SegmentContexts = {}
         reads_flags = reads_identity = False
-        self._identity_overrides_segment_keys: typing.Dict[str, typing.List[str]] = {}
         for segment_key, segment_context in (context.get("segments") or {}).items():
-            identifiers = get_identity_overrides_identifiers(segment_context)
-            if identifiers is None:
-                segments[segment_key] = segment_context
-                reads = _get_context_reads(segment_context["rules"])
-                reads_flags |= reads.flags
-                reads_identity |= reads.identity
-                if not reads.identity:
-                    self._identity_free_segments[segment_key] = segment_context
-                continue
-            self._identity_overrides_segments[segment_key] = segment_context
-            for identifier in identifiers:
-                self._identity_overrides_segment_keys.setdefault(identifier, []).append(
-                    segment_key
-                )
-
+            reads = _get_context_reads(segment_context["rules"])
+            reads_flags |= reads.flags
+            reads_identity |= reads.identity
+            if not reads.identity:
+                self._identity_free_segments[segment_key] = segment_context
         self.reads = _ContextReads(flags=reads_flags, identity=reads_identity)
-
-        # The context without identity overrides segments.
-        self.context: SDKEvaluationContext = {**context, "segments": segments}
         self.overrides_index: SegmentOverridesIndex = build_segment_overrides_index(
-            self.context
+            context
+        )
+        self._identity_overrides = {
+            identifier: (segment_context, self._index_overrides(segment_context))
+            for identifier, segment_context in (
+                identity_overrides_segments or {}
+            ).items()
+        }
+
+    @classmethod
+    def from_environment_document(
+        cls,
+        environment_document: EnvironmentModel,
+    ) -> LocalEvaluation:
+        return cls(
+            context=map_environment_document_to_context(environment_document),
+            identity_overrides_segments=(
+                map_environment_document_to_identity_overrides_segments(
+                    environment_document
+                )
+            ),
         )
 
     @cached_property
@@ -113,35 +123,33 @@ def get_identity_context(
         identifier: str,
         traits: typing.Optional[TraitMapping] = None,
     ) -> typing.Tuple[SDKEvaluationContext, SegmentOverridesIndex]:
-        """Evaluation context and overrides index for an identity.
-
-        The context carries the identity's own identity overrides segments,
-        narrowed to its identifier, so that their size doesn't depend on how
-        many identities share the same overrides.
-        """
+        """Evaluation context and overrides index for an identity, including
+        its own identity overrides segment, if any."""
         context = map_context_and_identity_data_to_context(
             context=self.context,
             identifier=identifier,
             traits=traits,
         )
-        if not (segment_keys := self._identity_overrides_segment_keys.get(identifier)):
+        if (identity_overrides := self._identity_overrides.get(identifier)) is None:
             return context, self.overrides_index
 
-        identity_segments: SegmentContexts = {}
-        identity_overrides_index: SegmentOverridesIndex = {}
-        for segment_key in segment_keys:
-            segment_context: SegmentContext[SegmentMetadata, FeatureMetadata] = {
-                **self._identity_overrides_segments[segment_key],
-                "rules": [map_identifiers_to_identity_overrides_rule([identifier])],
-            }
-            identity_segments[segment_key] = segment_context
-            for override in segment_context.get("overrides") or ():
-                feature_name = override["name"]
-                if feature_name not in identity_overrides_index:
-                    identity_overrides_index[feature_name] = dict(
-                        self.overrides_index.get(feature_name) or {}
-                    )
-                identity_overrides_index[feature_name][segment_key] = segment_context
-
-        context["segments"] = {**(context.get("segments") or {}), **identity_segments}
+        segment_context, identity_overrides_index = identity_overrides
+        context["segments"] = {
+            **(context.get("segments") or {}),
+            segment_context["key"]: segment_context,
+        }
         return context, ChainMap(identity_overrides_index, self.overrides_index)
+
+    def _index_overrides(
+        self,
+        segment_context: SegmentContext[SegmentMetadata, FeatureMetadata],
+    ) -> SegmentOverridesIndex:
+        """The overrides index, for the features `segment_context` overrides."""
+        return {
+            feature_name: {
+                **(self.overrides_index.get(feature_name) or {}),
+                segment_context["key"]: segment_context,
+            }
+            for override in segment_context.get("overrides") or ()
+            if (feature_name := override["name"])
+        }

flagsmith/mappers.py

@@ -1,7 +1,6 @@
 import json
 import typing
 import uuid
-from collections import defaultdict
 from datetime import datetime, timezone
 
 import sseclient
@@ -18,7 +17,6 @@
 from flagsmith.api.types import (
     EnvironmentModel,
     FeatureStateModel,
-    IdentityModel,
     SegmentRuleModel,
 )
 from flagsmith.models import Segment
@@ -31,13 +29,8 @@
 )
 from flagsmith.utils.datetime import fromisoformat
 
-OverrideKey = typing.Tuple[
-    int,
-    str,
-    bool,
-    typing.Any,
-]
-OverridesKey = typing.Tuple[OverrideKey, ...]
+# Each identity is evaluated with its own identity overrides segment only.
+IDENTITY_OVERRIDES_SEGMENT_KEY = "identity_overrides"
 
 
 def map_segment_results_to_identity_segments(
@@ -116,85 +109,57 @@ def map_environment_document_to_context(
             )
         },
         "segments": {
-            **{
-                (segment_key := str(segment_id := segment["id"])): {
-                    "key": segment_key,
-                    "name": segment["name"],
-                    "rules": _map_environment_document_rules_to_context_rules(
-                        segment["rules"]
-                    ),
-                    "overrides": list(
-                        _map_environment_document_feature_states_to_feature_contexts(
-                            segment.get("feature_states") or []
-                        )
-                    ),
-                    "metadata": SegmentMetadata(
-                        id=segment_id,
-                        source="api",
-                    ),
-                }
-                for segment in environment_document["project"]["segments"]
-            },
-            **_map_identity_overrides_to_segments(
-                environment_document.get("identity_overrides") or []
-            ),
+            (segment_key := str(segment_id := segment["id"])): {
+                "key": segment_key,
+                "name": segment["name"],
+                "rules": _map_environment_document_rules_to_context_rules(
+                    segment["rules"]
+                ),
+                "overrides": list(
+                    _map_environment_document_feature_states_to_feature_contexts(
+                        segment.get("feature_states") or []
+                    )
+                ),
+                "metadata": SegmentMetadata(
+                    id=segment_id,
+                    source="api",
+                ),
+            }
+            for segment in environment_document["project"]["segments"]
         },
     }
 
 
-def _map_identity_overrides_to_segments(
-    identity_overrides: list[IdentityModel],
-) -> dict[str, SegmentContext[SegmentMetadata, FeatureMetadata]]:
-    features_to_identifiers: typing.Dict[
-        OverridesKey,
-        typing.List[str],
-    ] = defaultdict(list)
-    for identity_override in identity_overrides:
-        identity_features = identity_override["identity_features"]
-        if not identity_features:
-            continue
-        overrides_key = tuple(
-            (
-                feature_state["feature"]["id"],
-                feature_state["feature"]["name"],
-                feature_state["enabled"],
-                feature_state["feature_state_value"],
-            )
-            for feature_state in sorted(
-                identity_features,
-                key=lambda feature_state: feature_state["feature"]["name"],
-            )
-        )
-        features_to_identifiers[overrides_key].append(identity_override["identifier"])
-    segment_contexts: typing.Dict[
-        str,
-        SegmentContext[
-            SegmentMetadata,
-            FeatureMetadata,
-        ],
-    ] = {}
-    for overrides_key, identifiers in features_to_identifiers.items():
-        # Create a segment context for each unique set of overrides
-        # Generate a unique key to avoid collisions
-        segment_key = str(hash(overrides_key))
-        segment_contexts[segment_key] = SegmentContext(
-            key=segment_key,
-            name="identity_overrides",
-            rules=[map_identifiers_to_identity_overrides_rule(identifiers)],
-            overrides=[
+def map_environment_document_to_identity_overrides_segments(
+    environment_document: EnvironmentModel,
+) -> typing.Dict[str, SegmentContext[SegmentMetadata, FeatureMetadata]]:
+    """Map identity overrides to segment contexts by identifier, each matching
+    its identity alone and carrying its overrides."""
+    return {
+        identity_override["identifier"]: {
+            "key": IDENTITY_OVERRIDES_SEGMENT_KEY,
+            "name": "identity_overrides",
+            "rules": [
+                map_identifiers_to_identity_overrides_rule(
+                    [identity_override["identifier"]]
+                )
+            ],
+            "overrides": [
                 {
                     "key": "",  # Identity overrides never carry multivariate options
-                    "name": feature_name,
-                    "enabled": feature_enabled,
-                    "value": feature_value,
+                    "name": feature_state["feature"]["name"],
+                    "enabled": feature_state["enabled"],
+                    "value": feature_state["feature_state_value"],
                     "priority": float("-inf"),  # Highest possible priority
-                    "metadata": {"id": feature_id},
+                    "metadata": {"id": feature_state["feature"]["id"]},
                 }
-                for feature_id, feature_name, feature_enabled, feature_value in overrides_key
+                for feature_state in identity_features
             ],
-            metadata=SegmentMetadata(source="identity_overrides"),
-        )
-    return segment_contexts
+            "metadata": SegmentMetadata(source="identity_overrides"),
+        }
+        for identity_override in environment_document.get("identity_overrides") or []
+        if (identity_features := identity_override["identity_features"])
+    }
 
 
 def map_identifiers_to_identity_overrides_rule(
@@ -208,18 +173,6 @@ def map_identifiers_to_identity_overrides_rule(
     return {"type": "ALL", "conditions": [condition]}
 
 
-def get_identity_overrides_identifiers(
-    segment_context: SegmentContext[SegmentMetadata, FeatureMetadata],
-) -> typing.Optional[typing.List[str]]:
-    """Identifiers an identity-overrides segment applies to, or `None` for
-    any other segment."""
-    if (segment_context.get("metadata") or {}).get("source") != "identity_overrides":
-        return None
-    [rule] = segment_context["rules"]
-    [condition] = rule["conditions"]
-    return typing.cast(InOperatorSegmentCondition, condition)["value"]
-
-
 def _map_environment_document_rules_to_context_rules(
     rules: list[SegmentRuleModel],
 ) -> list[SegmentRule]: