fix(experimentation): keep variant assignment stable across rollout updates (#7912)

Author: gagantrivedi

api/experimentation/services.py

@@ -17,6 +17,7 @@
 from audit.models import AuditLog
 from audit.related_object_type import RelatedObjectType
 from core.dataclasses import AuthorData
+from environments.tasks import rebuild_environment_document
 from experimentation.constants import (
     CONTROL_VARIANT_KEY,
     EXPERIMENT_FLAG,
@@ -57,7 +58,10 @@
 from features.models import FeatureState
 from features.value_types import BOOLEAN, INTEGER, STRING
 from features.versioning.dataclasses import FlagChangeSet
-from features.versioning.versioning_service import update_flag
+from features.versioning.versioning_service import (
+    update_flag,
+    update_multivariate_values,
+)
 from integrations.flagsmith.client import get_openfeature_client
 from segments.models import Condition, Segment, SegmentRule
 
@@ -574,6 +578,58 @@ def _sync_rollout_segment(experiment: Experiment, rollout_percentage: float) ->
     return segment
 
 
+def _get_live_rollout_override(experiment: Experiment) -> FeatureState | None:
+    return (
+        FeatureState.objects.get_live_feature_states(
+            environment=experiment.environment,
+            additional_filters=Q(
+                feature_segment__segment_id=experiment.rollout_segment_id,
+                identity__isnull=True,
+            ),
+            feature_id=experiment.feature_id,
+        )
+        .order_by("-id")
+        .first()
+    )
+
+
+def _update_live_feature_state(
+    feature_state: FeatureState, change_set: FlagChangeSet
+) -> None:
+    feature_state.enabled = change_set.enabled
+    feature_state.save()
+    feature_state.feature_state_value.set_value(
+        change_set.feature_state_value, change_set.type_
+    )
+    feature_state.feature_state_value.save()
+    update_multivariate_values(feature_state, change_set.multivariate_values)
+
+
+def _update_rollout_in_place(experiment: Experiment, change_set: FlagChangeSet) -> None:
+    """Write the rollout-segment override, keeping variant assignment stable.
+
+    Under v2 versioning, ``update_flag`` clones the override into a fresh feature
+    state on every call. Since the multivariate split is salted on the feature
+    state id, that would re-randomise control/variant for already-enrolled
+    identities on each rollout update. Once the override exists, mutate it in
+    place instead and rebuild the environment document by hand (no version is
+    published). Creating the override, and v1 versioning, still go through
+    ``update_flag``, which already reuses the feature state.
+
+    This is a temporary solution until we find a permanent fix for the
+    underlying salting issue: https://github.com/Flagsmith/flagsmith/issues/7913
+    """
+    if experiment.environment.use_v2_feature_versioning and (
+        override := _get_live_rollout_override(experiment)
+    ):
+        _update_live_feature_state(override, change_set)
+        rebuild_environment_document.delay(
+            kwargs={"environment_id": experiment.environment_id}
+        )
+        return
+    update_flag(experiment.environment, experiment.feature, change_set)
+
+
 def apply_experiment_rollout(experiment: Experiment, spec: RolloutSpec) -> None:
     if experiment.status == ExperimentStatus.COMPLETED:
         raise ValidationError(
@@ -582,9 +638,8 @@ def apply_experiment_rollout(experiment: Experiment, spec: RolloutSpec) -> None:
     validate_rollout_spec(experiment, spec)
     with transaction.atomic():
         segment = _sync_rollout_segment(experiment, spec.rollout_percentage)
-        update_flag(
-            experiment.environment,
-            experiment.feature,
+        _update_rollout_in_place(
+            experiment,
             FlagChangeSet(
                 author=spec.author,
                 enabled=spec.enabled,
@@ -638,9 +693,8 @@ def enable_experiment_rollout(experiment: Experiment, author: AuthorData) -> Non
         return
 
     value = rollout["feature_state_value"]
-    update_flag(
-        experiment.environment,
-        experiment.feature,
+    _update_rollout_in_place(
+        experiment,
         FlagChangeSet(
             author=author,
             enabled=True,

api/features/versioning/versioning_service.py

@@ -188,7 +188,7 @@ def _update_flag_for_versioning_v2(
         change_set.feature_state_value,
         change_set.type_,
     )
-    _update_multivariate_values(target_feature_state, change_set.multivariate_values)
+    update_multivariate_values(target_feature_state, change_set.multivariate_values)
 
     if change_set.segment_id is not None and change_set.segment_priority is not None:
         _update_segment_priority(target_feature_state, change_set.segment_priority)
@@ -241,7 +241,7 @@ def _update_flag_for_versioning_v1(
         change_set.feature_state_value,
         change_set.type_,
     )
-    _update_multivariate_values(target_feature_state, change_set.multivariate_values)
+    update_multivariate_values(target_feature_state, change_set.multivariate_values)
 
     if change_set.segment_id is not None and change_set.segment_priority is not None:
         _update_segment_priority(target_feature_state, change_set.segment_priority)
@@ -256,7 +256,7 @@ def _update_feature_state_value(
     fsv.save()
 
 
-def _update_multivariate_values(
+def update_multivariate_values(
     feature_state: FeatureState,
     values: list[MultivariateValueChangeSet] | None,
 ) -> None:
@@ -374,7 +374,7 @@ def _update_flag_v2_for_versioning_v2(
                 override.feature_state_value,
                 override.type_,
             )
-            _update_multivariate_values(segment_state, override.multivariate_values)
+            update_multivariate_values(segment_state, override.multivariate_values)
 
             if override.priority is not None:
                 _update_segment_priority(segment_state, override.priority)
@@ -393,7 +393,7 @@ def _update_flag_v2_for_versioning_v2(
                 override.feature_state_value,
                 override.type_,
             )
-            _update_multivariate_values(segment_state, override.multivariate_values)
+            update_multivariate_values(segment_state, override.multivariate_values)
 
     new_version.publish(
         published_by=change_set.author.user,
@@ -445,7 +445,7 @@ def _update_flag_v2_for_versioning_v1(
                 override.feature_state_value,
                 override.type_,
             )
-            _update_multivariate_values(segment_state, override.multivariate_values)
+            update_multivariate_values(segment_state, override.multivariate_values)
         else:
             assert len(segment_states) == 1
             segment_state = list(segment_states.values())[0]
@@ -457,7 +457,7 @@ def _update_flag_v2_for_versioning_v1(
                 override.feature_state_value,
                 override.type_,
             )
-            _update_multivariate_values(segment_state, override.multivariate_values)
+            update_multivariate_values(segment_state, override.multivariate_values)
 
             if override.priority is not None:
                 _update_segment_priority(segment_state, override.priority)

api/tests/unit/experimentation/test_services.py

@@ -2,6 +2,7 @@
 from datetime import datetime, timezone
 
 import pytest
+from django.db.models import Q
 from flag_engine.segments.constants import PERCENTAGE_SPLIT
 from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
@@ -1668,3 +1669,68 @@ def test_enable_experiment_rollout__no_rollout__no_op(
 
     # Then nothing is written
     update_flag.assert_not_called()
+
+
+def test_apply_experiment_rollout__reapplied_under_v2__keeps_variant_assignment(
+    environment_v2_versioning: Environment,
+    multivariate_feature: Feature,
+    multivariate_options: list[MultivariateFeatureOption],
+    admin_user: FFAdminUser,
+) -> None:
+    # Given a running experiment whose rollout splits two variants 50/50
+    option_a, option_b, _ = multivariate_options
+    experiment = Experiment.objects.create(
+        environment=environment_v2_versioning,
+        feature=multivariate_feature,
+        name="exp",
+        hypothesis="h",
+        status=ExperimentStatus.RUNNING,
+    )
+    spec = RolloutSpec(
+        enabled=True,
+        rollout_percentage=50.0,
+        feature_state_value="control",
+        value_type="string",
+        multivariate_values=[
+            MultivariateValueChangeSet(option_a.id, 50.0),
+            MultivariateValueChangeSet(option_b.id, 50.0),
+        ],
+        author=AuthorData(user=admin_user),
+    )
+    identity_hash_keys = [f"identity-{i}" for i in range(50)]
+
+    def variant_assignment() -> dict[str, int]:
+        override = (
+            FeatureState.objects.get_live_feature_states(
+                environment=experiment.environment,
+                additional_filters=Q(
+                    feature_segment__segment=experiment.rollout_segment,
+                    identity__isnull=True,
+                ),
+                feature_id=experiment.feature_id,
+            )
+            .prefetch_related(
+                "multivariate_feature_state_values__multivariate_feature_option"
+            )
+            .latest("id")
+        )
+        assignment: dict[str, int] = {}
+        for key in identity_hash_keys:
+            option = override.get_multivariate_feature_state_value(key)
+            # The 50/50 split allocates 100%, so every identity lands on an option.
+            assert isinstance(option, MultivariateFeatureOption)
+            assignment[key] = option.id
+        return assignment
+
+    # When the rollout is applied, then re-applied unchanged (e.g. tuned while
+    # the experiment is running)
+    services.apply_experiment_rollout(experiment, spec)
+    experiment.refresh_from_db()
+    before = variant_assignment()
+
+    services.apply_experiment_rollout(experiment, spec)
+    after = variant_assignment()
+
+    # Then every already-enrolled identity keeps the variant it was first
+    # assigned; tuning the rollout must not re-randomise the split.
+    assert before == after

docs/docs/deployment-self-hosting/observability/_events-catalogue.md

@@ -494,7 +494,7 @@ Attributes:
 ### `warehouse.connection.connected`
 
 Logged at `info` from:
- - `api/experimentation/services.py:684`
+ - `api/experimentation/services.py:738`
 
 Attributes:
  - `environment.id`
@@ -503,7 +503,7 @@ Attributes:
 ### `warehouse.connection.test_event_sent`
 
 Logged at `info` from:
- - `api/experimentation/services.py:664`
+ - `api/experimentation/services.py:718`
 
 Attributes:
  - `environment.id`
@@ -512,7 +512,7 @@ Attributes:
 ### `warehouse.srm.overallocated`
 
 Logged at `error` from:
- - `api/experimentation/services.py:392`
+ - `api/experimentation/services.py:396`
 
 Attributes:
  - `environment.id`
@@ -522,7 +522,7 @@ Attributes:
 ### `warehouse.srm.unkeyed_variant`
 
 Logged at `error` from:
- - `api/experimentation/services.py:378`
+ - `api/experimentation/services.py:382`
 
 Attributes:
  - `environment.id`