feat(cohorts): apply membership deltas to edge identities in batches

api/cohorts/constants.py

@@ -0,0 +1,10 @@
+COHORT_SYSTEM_TRAIT_KEY_PREFIX = "flagsmith_cohort_"
+COHORT_MEMBERSHIP_APPLY_BATCH_SIZE = 100
+COHORT_MEMBERSHIP_APPLY_MAX_BATCHES_PER_RUN = 10
+DYNAMODB_THROTTLING_ERROR_CODES = frozenset(
+    {
+        "ProvisionedThroughputExceededException",
+        "RequestLimitExceeded",
+        "ThrottlingException",
+    }
+)

api/cohorts/metrics.py

@@ -0,0 +1,9 @@
+import prometheus_client
+
+flagsmith_cohorts_membership_deltas_applied_total = prometheus_client.Counter(
+    "flagsmith_cohorts_membership_deltas_applied_total",
+    "Total number of cohort membership ledger rows transitioned to their "
+    "applied state after the corresponding identity document write. "
+    "The `operation` label is either `add` or `remove`.",
+    ["operation"],
+)

api/cohorts/models.py

@@ -1,5 +1,6 @@
 from django.db import models
 
+from cohorts.constants import COHORT_SYSTEM_TRAIT_KEY_PREFIX
 from core.models import SoftDeleteExportableModel
 
 
@@ -26,6 +27,10 @@ class Cohort(SoftDeleteExportableModel):
     version = models.PositiveIntegerField(default=0)
     created_at = models.DateTimeField(auto_now_add=True)
 
+    @property
+    def system_trait_key(self) -> str:
+        return f"{COHORT_SYSTEM_TRAIT_KEY_PREFIX}{self.uuid}"
+
     class Meta:
         constraints = [
             # Exactly one active cohort feeds a segment: two cohorts on one

api/cohorts/services.py

@@ -0,0 +1,68 @@
+import structlog
+from django.db.models import QuerySet
+from django.utils import timezone
+
+from cohorts.constants import COHORT_MEMBERSHIP_APPLY_BATCH_SIZE
+from cohorts.metrics import flagsmith_cohorts_membership_deltas_applied_total
+from cohorts.models import Cohort, CohortMembership, CohortMembershipState
+from environments.dynamodb import DynamoIdentityWrapper
+
+logger = structlog.get_logger("cohorts")
+
+_PENDING_STATES = [
+    CohortMembershipState.PENDING_ADD,
+    CohortMembershipState.PENDING_REMOVE,
+]
+
+
+def pending_memberships(cohort: Cohort) -> "QuerySet[CohortMembership]":
+    return CohortMembership.objects.filter(cohort=cohort, state__in=_PENDING_STATES)
+
+
+def apply_pending_memberships(cohort: Cohort) -> bool:
+    identity_wrapper = DynamoIdentityWrapper()
+    environment_api_key: str = cohort.environment.api_key
+    trait_key = cohort.system_trait_key
+    batch = list(
+        pending_memberships(cohort).order_by("id")[:COHORT_MEMBERSHIP_APPLY_BATCH_SIZE]
+    )
+    if not batch:
+        return False
+    added_ids: list[int] = []
+    removed_ids: list[int] = []
+    for row in batch:
+        if row.state == CohortMembershipState.PENDING_ADD:
+            identity_wrapper.set_system_trait(
+                environment_api_key=environment_api_key,
+                identifier=row.identifier,
+                trait_key=trait_key,
+            )
+            added_ids.append(row.id)
+        else:
+            identity_wrapper.unset_system_trait(
+                environment_api_key=environment_api_key,
+                identifier=row.identifier,
+                trait_key=trait_key,
+            )
+            removed_ids.append(row.id)
+    added_count = CohortMembership.objects.filter(
+        id__in=added_ids, state=CohortMembershipState.PENDING_ADD
+    ).update(state=CohortMembershipState.APPLIED, updated_at=timezone.now())
+    removed_count, _ = CohortMembership.objects.filter(
+        id__in=removed_ids, state=CohortMembershipState.PENDING_REMOVE
+    ).delete()
+    flagsmith_cohorts_membership_deltas_applied_total.labels(operation="add").inc(
+        added_count
+    )
+    flagsmith_cohorts_membership_deltas_applied_total.labels(operation="remove").inc(
+        removed_count
+    )
+    if added_count or removed_count:
+        logger.info(
+            "membership.applied",
+            cohort__id=cohort.id,
+            environment__id=cohort.environment_id,
+            adds__count=added_count,
+            removes__count=removed_count,
+        )
+    return pending_memberships(cohort).exists()

api/cohorts/tasks.py

@@ -0,0 +1,41 @@
+from datetime import timedelta
+
+import structlog
+from botocore.exceptions import ClientError
+from task_processor.decorators import register_task_handler
+from task_processor.exceptions import TaskBackoffError
+
+from cohorts import services
+from cohorts.constants import (
+    COHORT_MEMBERSHIP_APPLY_MAX_BATCHES_PER_RUN,
+    DYNAMODB_THROTTLING_ERROR_CODES,
+)
+from cohorts.models import Cohort
+from environments.dynamodb import DynamoIdentityWrapper
+
+logger = structlog.get_logger("cohorts")
+
+
+@register_task_handler(timeout=timedelta(minutes=5))
+def apply_cohort_membership_deltas(cohort_id: int) -> None:
+    log = logger.bind(cohort__id=cohort_id)
+    if (cohort := Cohort.objects.filter(id=cohort_id).first()) is None:
+        log.info("membership.apply.skipped", reason="cohort_missing")
+        return
+    if not (
+        cohort.environment.project.enable_dynamo_db
+        and DynamoIdentityWrapper().is_enabled
+    ):
+        log.info("membership.apply.skipped", reason="not_edge")
+        return
+    try:
+        for _ in range(COHORT_MEMBERSHIP_APPLY_MAX_BATCHES_PER_RUN):
+            if not services.apply_pending_memberships(cohort):
+                return
+    except ClientError as exc:
+        if exc.response["Error"]["Code"] in DYNAMODB_THROTTLING_ERROR_CODES:
+            log.warning("membership.apply.throttled")
+            raise TaskBackoffError() from exc
+        raise
+    # Still pending after this run's batch cap; continue in a fresh task run.
+    apply_cohort_membership_deltas.delay(kwargs={"cohort_id": cohort_id})

api/environments/dynamodb/constants.py

@@ -9,6 +9,8 @@
 DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT = 25
 IDENTITIES_PAGINATION_LIMIT = 1000
 
+SYSTEM_TRAIT_WRITE_MAX_ATTEMPTS = 3
+
 # DynamoDB max item size is 400 KB (409,600 bytes).
 DOCUMENT_SIZE_HISTOGRAM_BUCKETS = (
     1_000,

api/environments/dynamodb/wrappers/exceptions.py

@@ -1,6 +1,16 @@
 from decimal import Decimal
 
 
+class SystemTraitWriteRaceError(Exception):
+    def __init__(self, composite_key: str) -> None:
+        super().__init__(
+            f"Gave up writing a system trait for identity {composite_key!r}: "
+            "concurrent writers kept changing the document between read and "
+            "conditional write."
+        )
+        self.composite_key = composite_key
+
+
 class CapacityBudgetExceeded(Exception):
     def __init__(
         self,

api/environments/dynamodb/wrappers/identity_wrapper.py

@@ -5,19 +5,29 @@
 from typing import Iterable
 
 from boto3.dynamodb.conditions import Attr, Key
+from botocore.exceptions import ClientError
 from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
 from rest_framework.exceptions import NotFound
 
 from edge_api.identities.search import EdgeIdentitySearchData
-from environments.dynamodb.constants import IDENTITIES_PAGINATION_LIMIT
-from environments.dynamodb.wrappers.exceptions import CapacityBudgetExceeded
+from environments.dynamodb.constants import (
+    IDENTITIES_PAGINATION_LIMIT,
+    SYSTEM_TRAIT_WRITE_MAX_ATTEMPTS,
+)
+from environments.dynamodb.wrappers.exceptions import (
+    CapacityBudgetExceeded,
+    SystemTraitWriteRaceError,
+)
 from util.engine_models.context.mappers import (
     is_context_in_segment,
     map_environment_identity_to_context,
 )
 from util.engine_models.identities.models import IdentityModel
-from util.mappers import map_identity_to_identity_document
+from util.mappers import (
+    map_engine_identity_to_identity_document,
+    map_identity_to_identity_document,
+)
 
 from .base import BaseDynamoWrapper
 
@@ -34,6 +44,17 @@
 logger = logging.getLogger(__name__)
 
 
+def _system_trait_value_matches(
+    stored_value: object,
+    document_value: bool | int | Decimal | str,
+) -> bool:
+    # The bool check stops `Decimal(1) == True` false positives.
+    return (
+        isinstance(stored_value, bool) == isinstance(document_value, bool)
+        and stored_value == document_value
+    )
+
+
 class DynamoIdentityWrapper(BaseDynamoWrapper):
     def __init__(self) -> None:
         super().__init__()
@@ -63,6 +84,117 @@ def write_identities(self, identities: Iterable["Identity"]):  # type: ignore[no
     def get_item(self, composite_key: str) -> typing.Optional[dict]:  # type: ignore[type-arg]
         return self.table.get_item(Key={"composite_key": composite_key}).get("Item")  # type: ignore[union-attr]
 
+    def set_system_trait(
+        self,
+        *,
+        environment_api_key: str,
+        identifier: str,
+        trait_key: str,
+        trait_value: bool | int | float | str = True,
+    ) -> None:
+        """Idempotently set a system trait on an identity document.
+
+        Writes only touch the `system_traits.<trait_key>` attribute, so
+        concurrent writes to other attributes are never overwritten; the
+        document is created if missing. Each write is conditional on the
+        document shape just read — a lost race re-reads and retries, and
+        `SystemTraitWriteRaceError` is raised once attempts are exhausted.
+
+        Assumes stored documents never carry `system_traits` as NULL — the
+        document mapper omits the attribute when unset.
+        """
+        composite_key = IdentityModel.generate_composite_key(
+            environment_api_key, identifier
+        )
+        # DynamoDB rejects floats and returns all numbers as Decimal.
+        document_value: bool | int | Decimal | str = (
+            Decimal(str(trait_value)) if isinstance(trait_value, float) else trait_value
+        )
+        for _ in range(SYSTEM_TRAIT_WRITE_MAX_ATTEMPTS):
+            # Strongly consistent read: a replication-lagged hint would burn
+            # retry attempts on conditional writes that can never succeed.
+            document = self.table.get_item(  # type: ignore[union-attr]
+                Key={"composite_key": composite_key}, ConsistentRead=True
+            ).get("Item")
+            system_traits = document.get("system_traits") if document else None
+            if isinstance(system_traits, dict) and _system_trait_value_matches(
+                system_traits.get(trait_key), document_value
+            ):
+                return
+            try:
+                if document is None:
+                    self.table.put_item(  # type: ignore[union-attr]
+                        Item=map_engine_identity_to_identity_document(
+                            IdentityModel(
+                                identifier=identifier,
+                                environment_api_key=environment_api_key,
+                                system_traits={trait_key: trait_value},
+                            )
+                        ),
+                        ConditionExpression="attribute_not_exists(composite_key)",
+                    )
+                elif isinstance(system_traits, dict):
+                    self.table.update_item(  # type: ignore[union-attr]
+                        Key={"composite_key": composite_key},
+                        UpdateExpression="SET system_traits.#tk = :value",
+                        ConditionExpression="attribute_exists(system_traits)",
+                        ExpressionAttributeNames={"#tk": trait_key},
+                        ExpressionAttributeValues={":value": document_value},
+                    )
+                else:
+                    # If another writer created system_traits after we read
+                    # the document, this write does nothing and their traits
+                    # survive; the returned attributes tell us which happened.
+                    response = self.table.update_item(  # type: ignore[union-attr]
+                        Key={"composite_key": composite_key},
+                        UpdateExpression=(
+                            "SET system_traits = if_not_exists(system_traits, :init)"
+                        ),
+                        # Without this condition, update_item would re-create
+                        # a just-deleted identity as an empty document
+                        # containing nothing but this trait.
+                        ConditionExpression="attribute_exists(composite_key)",
+                        ExpressionAttributeValues={
+                            ":init": {trait_key: document_value}
+                        },
+                        ReturnValues="ALL_NEW",
+                    )
+                    written_traits = response["Attributes"].get("system_traits")
+                    if isinstance(written_traits, dict) and _system_trait_value_matches(
+                        written_traits.get(trait_key), document_value
+                    ):
+                        return
+                    continue
+                return
+            except ClientError as exc:
+                if exc.response["Error"]["Code"] != "ConditionalCheckFailedException":
+                    raise
+        raise SystemTraitWriteRaceError(composite_key)
+
+    def unset_system_trait(
+        self,
+        *,
+        environment_api_key: str,
+        identifier: str,
+        trait_key: str,
+    ) -> None:
+        """Idempotently remove a system trait from an identity document."""
+        composite_key = IdentityModel.generate_composite_key(
+            environment_api_key, identifier
+        )
+        try:
+            self.table.update_item(  # type: ignore[union-attr]
+                Key={"composite_key": composite_key},
+                UpdateExpression="REMOVE system_traits.#tk",
+                # Failing this condition covers every no-op case at once:
+                # missing document, missing system_traits, or trait already absent.
+                ConditionExpression="attribute_exists(system_traits.#tk)",
+                ExpressionAttributeNames={"#tk": trait_key},
+            )
+        except ClientError as exc:
+            if exc.response["Error"]["Code"] != "ConditionalCheckFailedException":
+                raise
+
     def delete_item(self, composite_key: str):  # type: ignore[no-untyped-def]
         self.table.delete_item(Key={"composite_key": composite_key})  # type: ignore[union-attr]
 

api/tests/unit/cohorts/conftest.py

@@ -0,0 +1,26 @@
+import pytest
+
+from cohorts.models import Cohort
+from environments.models import Environment
+from projects.models import Project
+from segments.models import Segment
+
+
+@pytest.fixture()
+def cohort(environment: Environment, segment: Segment) -> Cohort:
+    cohort: Cohort = Cohort.objects.create(environment=environment, segment=segment)
+    return cohort
+
+
+@pytest.fixture()
+def edge_cohort(
+    dynamo_enabled_project: Project,
+    dynamo_enabled_project_environment_one: Environment,
+) -> Cohort:
+    segment = Segment.objects.create(
+        name="edge segment", project=dynamo_enabled_project
+    )
+    cohort: Cohort = Cohort.objects.create(
+        environment=dynamo_enabled_project_environment_one, segment=segment
+    )
+    return cohort