fix(Environments): Environment document can sync to DynamoDB before feature states commit

api/core/signals.py

@@ -1,10 +1,7 @@
 import logging
 
-from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
-from django.utils import timezone
 from simple_history.models import HistoricalRecords  # type: ignore[import-untyped]
-from task_processor.task_run_method import TaskRunMethod
 
 from audit import tasks
 from core.models import AbstractBaseAuditableModel
@@ -19,21 +16,10 @@ def create_audit_log_from_historical_record(  # type: ignore[no-untyped-def]
     history_instance,
     **kwargs,
 ):
-    # The environment document in dynamodb is updated based on the post_save signal from the audit log
-    # When creating a new feature, the feature states are created after the feature has been created.
-    # i.e: the below task gets created/scheduled before feature states are created
-    # Usually, there is enough time for the main thread to create the feature states
-    # before the task is executed, but not always.
-    # In those cases, we send the environment document to dynamodb without any feature states for the new feature.
-    # In order to avoid this, either we need to update environment
-    # document when creating feature states
-    # or delay the execution of this task
-    # We prefer to delay the execution of the task because of it's low surface area
-    delay_until = (
-        timezone.now() + timezone.timedelta(seconds=1)  # type: ignore[attr-defined]
-        if settings.TASK_RUN_METHOD == TaskRunMethod.TASK_PROCESSOR
-        else None
-    )
+    # Note: this task can run before the feature states of a newly created feature or
+    # environment exist. `environments.tasks.process_environment_update` guards the
+    # environment document write against that, using `Feature.is_creating` and
+    # `Environment.is_creating`.
     if instance.get_skip_create_audit_log():
         return
 
@@ -62,7 +48,6 @@ def create_audit_log_from_historical_record(  # type: ignore[no-untyped-def]
             "history_user_id": getattr(history_user, "id", None),
             "history_record_class_path": instance.history_record_class_path,
         },
-        delay_until=delay_until,
     )
 
 

api/environments/constants.py

@@ -1,3 +1,9 @@
+# How many times an environment document write is deferred while the environments or
+# features it covers are still being created, and the base of the exponential backoff
+# (in seconds) between those attempts, i.e. 1s, 2s, 4s.
+ENVIRONMENT_DOCUMENT_WRITE_MAX_DEFERRALS = 3
+ENVIRONMENT_DOCUMENT_WRITE_DEFERRAL_SECONDS = 1
+
 IDENTITY_INTEGRATIONS_RELATION_NAMES = [
     "amplitude_config",
     "heap_config",

api/environments/dynamodb/wrappers/environment_wrapper.py

@@ -21,6 +21,7 @@
 from environments.metrics import (
     flagsmith_dynamo_environment_document_compression_ratio,
     flagsmith_dynamo_environment_document_size_bytes,
+    flagsmith_environment_document_writes_total,
 )
 from integrations.flagsmith.client import get_openfeature_client
 from util.mappers import (
@@ -71,37 +72,70 @@ def _write_environments(self, environments: Iterable["Environment"]) -> None:
         )
 
         assert self.table
-        with self.table.batch_writer() as writer:
-            for environment in environments:
-                organisation = environment.project.organisation
-                if openfeature_client.get_boolean_value(
-                    "compress_dynamo_documents",
-                    default_value=False,
-                    evaluation_context=organisation.openfeature_evaluation_context,
-                ):
-                    result = self._map_compressed_environment_document(environment)
-                    writer.put_item(Item=result.document)
-
-                    flagsmith_dynamo_environment_document_size_bytes.labels(
-                        table=self.get_table_name(),
-                        compressed="true",
-                    ).observe(result.compressed_size_bytes)
-                    flagsmith_dynamo_environment_document_compression_ratio.labels(
-                        table=self.get_table_name(),
-                    ).observe(result.compression_ratio)
-                    logger.info(
-                        "environment-document-compressed",
-                        environment_id=environment.id,
-                        environment_api_key=environment.api_key,
+        # `batch_writer` only buffers documents, flushing them when it fills up and
+        # when the context exits, so nothing is known to be written until then.
+        attempted = 0
+        # (environment id, feature states count, document size in bytes)
+        pending_events: list[tuple[int, int, int]] = []
+        try:
+            with self.table.batch_writer() as writer:
+                for environment in environments:
+                    attempted += 1
+                    organisation = environment.project.organisation
+                    if openfeature_client.get_boolean_value(
+                        "compress_dynamo_documents",
+                        default_value=False,
+                        evaluation_context=organisation.openfeature_evaluation_context,
+                    ):
+                        result = self._map_compressed_environment_document(environment)
+                        writer.put_item(Item=result.document)
+
+                        flagsmith_dynamo_environment_document_size_bytes.labels(
+                            table=self.get_table_name(),
+                            compressed="true",
+                        ).observe(result.compressed_size_bytes)
+                        flagsmith_dynamo_environment_document_compression_ratio.labels(
+                            table=self.get_table_name(),
+                        ).observe(result.compression_ratio)
+                        logger.info(
+                            "environment-document-compressed",
+                            environment_id=environment.id,
+                            environment_api_key=environment.api_key,
+                        )
+                        document_bytes = result.compressed_size_bytes
+                        feature_states_count = result.feature_states_count
+                    else:
+                        item = self._map_environment_document(environment)
+                        writer.put_item(Item=item)
+
+                        document_bytes = estimate_document_size(item)
+                        feature_states_count = len(item["feature_states"])
+                        flagsmith_dynamo_environment_document_size_bytes.labels(
+                            table=self.get_table_name(),
+                            compressed="false",
+                        ).observe(document_bytes)
+
+                    pending_events.append(
+                        (environment.id, feature_states_count, document_bytes)
                     )
-                else:
-                    item = self._map_environment_document(environment)
-                    writer.put_item(Item=item)
-
-                    flagsmith_dynamo_environment_document_size_bytes.labels(
-                        table=self.get_table_name(),
-                        compressed="false",
-                    ).observe(estimate_document_size(item))
+        except Exception:
+            # A failed batch does not report which of its documents were persisted,
+            # so every document it carried is counted.
+            flagsmith_environment_document_writes_total.labels(result="failure").inc(
+                attempted
+            )
+            raise
+
+        for environment_id, feature_states_count, document_bytes in pending_events:
+            logger.info(
+                "environment_document.written",
+                environment__id=environment_id,
+                feature_states__count=feature_states_count,
+                document__bytes=document_bytes,
+            )
+        flagsmith_environment_document_writes_total.labels(result="success").inc(
+            attempted
+        )
 
 
 class DynamoEnvironmentWrapper(BaseDynamoEnvironmentWrapper):

api/environments/metrics.py

@@ -14,6 +14,14 @@
     ["result"],
 )
 
+flagsmith_environment_document_writes_total = prometheus_client.Counter(
+    "flagsmith_environment_document_writes_total",
+    "Environment documents written to DynamoDB. `result` label is either `success` or "
+    "`failure`. A batch write that fails counts every document it carried, since "
+    "DynamoDB does not report which of them were persisted.",
+    ["result"],
+)
+
 flagsmith_dynamo_environment_document_size_bytes = prometheus_client.Histogram(
     "flagsmith_dynamo_environment_document_size_bytes",
     "Size of environment documents written to DynamoDB.",

api/environments/models.py

@@ -17,6 +17,7 @@
     AFTER_DELETE,
     AFTER_SAVE,
     AFTER_UPDATE,
+    BEFORE_CREATE,
     BEFORE_SAVE,
     LifecycleModel,
     hook,
@@ -82,7 +83,9 @@
 class Environment(
     LifecycleModel,  # type: ignore[misc]
     abstract_base_auditable_model_factory(  # type: ignore[misc]
-        change_details_excluded_fields=["updated_at"],
+        # `is_creating` is an internal lifecycle flag, and is cleared without a
+        # historical record, so its recorded value is always stale.
+        change_details_excluded_fields=["updated_at", "is_creating"],
         historical_records_excluded_fields=["uuid"],
     ),
     SoftDeleteObject,  # type: ignore[misc]
@@ -180,9 +183,21 @@ class Environment(
     class Meta:
         ordering = ["id"]
 
+    @hook(BEFORE_CREATE)  # type: ignore[misc]
+    def mark_as_creating(self) -> None:
+        # The environment row is committed before `create_feature_states` seeds its
+        # initial feature states, so anything reading the environment in between sees
+        # it without any flags. `is_creating` marks that window for the environment
+        # document writer. Clones set the flag in `clone()` instead, since create
+        # hooks don't fire for them.
+        self.is_creating = True
+
     @hook(AFTER_CREATE)  # type: ignore[misc]
     def create_feature_states(self) -> None:
         FeatureState.create_initial_feature_states_for_environment(environment=self)
+        self.is_creating = False
+        # Update the row directly to avoid re-triggering hooks for a lifecycle flag.
+        Environment.objects.filter(id=self.id).update(is_creating=False)
 
     @hook(AFTER_UPDATE)  # type: ignore[misc]
     def clear_environment_cache(self) -> None:

api/environments/tasks.py

@@ -1,17 +1,27 @@
+from datetime import timedelta
+
+import structlog
+from django.conf import settings
 from django.db.models import Prefetch, Q
 from django.utils import timezone
 from task_processor.decorators import (
     register_task_handler,
 )
 from task_processor.models import TaskPriority
+from task_processor.task_run_method import TaskRunMethod
 
 from audit.models import AuditLog
+from environments.constants import (
+    ENVIRONMENT_DOCUMENT_WRITE_DEFERRAL_SECONDS,
+    ENVIRONMENT_DOCUMENT_WRITE_MAX_DEFERRALS,
+)
 from environments.dynamodb import DynamoIdentityWrapper
 from environments.models import (
     Environment,
     environment_v2_wrapper,
     environment_wrapper,
 )
+from features.models import Feature
 from features.multivariate.models import MultivariateFeatureStateValue
 from features.versioning.models import EnvironmentFeatureVersion
 from features.versioning.versioning_service import (
@@ -22,16 +32,61 @@
     send_environment_update_message_for_project,
 )
 
+logger = structlog.get_logger("environments")
+
 
 @register_task_handler(priority=TaskPriority.HIGH)
 def rebuild_environment_document(environment_id: int) -> None:
     Environment.write_environment_documents(environment_id=environment_id)
 
 
+def _is_seeding_feature_states(audit_log: AuditLog) -> bool:
+    """
+    Whether any environment or feature the audit log covers is still seeding its
+    initial feature states, and would therefore be written to DynamoDB incomplete.
+    """
+    environments_filter = Q(project_id=audit_log.project_id)
+    if audit_log.environment_id:
+        environments_filter &= Q(id=audit_log.environment_id)
+
+    return bool(
+        Environment.objects.filter(environments_filter, is_creating=True).exists()
+        or Feature.objects.filter(
+            project_id=audit_log.project_id, is_creating=True
+        ).exists()
+    )
+
+
 @register_task_handler(priority=TaskPriority.HIGHEST)
-def process_environment_update(audit_log_id: int):  # type: ignore[no-untyped-def]
+def process_environment_update(audit_log_id: int, deferrals: int = 0):  # type: ignore[no-untyped-def]
     audit_log = AuditLog.objects.get(id=audit_log_id)
 
+    # Deferring relies on `delay_until`, which is a no-op without the task processor —
+    # re-enqueueing there would drop the write and leave the document stale.
+    if (
+        settings.TASK_RUN_METHOD == TaskRunMethod.TASK_PROCESSOR
+        and _is_seeding_feature_states(audit_log)
+    ):
+        if deferrals < ENVIRONMENT_DOCUMENT_WRITE_MAX_DEFERRALS:
+            process_environment_update.delay(
+                kwargs={"audit_log_id": audit_log_id, "deferrals": deferrals + 1},
+                delay_until=timezone.now()
+                + timedelta(
+                    seconds=ENVIRONMENT_DOCUMENT_WRITE_DEFERRAL_SECONDS * 2**deferrals
+                ),
+            )
+            return
+
+        # Writing a document that may be missing feature states is preferable to not
+        # writing one at all, since a stuck `is_creating` flag would otherwise stop
+        # every future update to this environment.
+        logger.warning(
+            "environment_document.written_while_seeding",
+            audit_log__id=audit_log_id,
+            environment__id=audit_log.environment_id,
+            project__id=audit_log.project_id,
+        )
+
     # Send environment document to dynamodb
     Environment.write_environment_documents(
         environment_id=audit_log.environment_id, project_id=audit_log.project_id

api/features/migrations/0068_add_feature_is_creating.py

@@ -0,0 +1,21 @@
+# Generated by Django 5.2.16 on 2026-08-07 23:55
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("features", "0067_add_feature_state_mv_hashing_salt"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="feature",
+            name="is_creating",
+            field=models.BooleanField(
+                default=False,
+                help_text="Attribute used to indicate when a feature is still being created, and its initial feature states are not yet guaranteed to exist",
+            ),
+        ),
+    ]

api/features/models.py

@@ -92,7 +92,9 @@
 class Feature(  # type: ignore[django-manager-missing]
     SoftDeleteExportableModel,
     CustomLifecycleModelMixin,
-    abstract_base_auditable_model_factory(["uuid"]),  # type: ignore[misc]
+    # `is_creating` is an internal lifecycle flag, and is cleared without a historical
+    # record, so it is kept out of the feature's history entirely.
+    abstract_base_auditable_model_factory(["uuid", "is_creating"]),  # type: ignore[misc]
 ):
     name = models.CharField(max_length=2000)
     created_date = models.DateTimeField("DateCreated", auto_now_add=True)
@@ -128,6 +130,14 @@ class Feature(  # type: ignore[django-manager-missing]
 
     is_server_key_only = models.BooleanField(default=False)
 
+    is_creating = models.BooleanField(
+        default=False,
+        help_text=(
+            "Attribute used to indicate when a feature is still being created, and its"
+            " initial feature states are not yet guaranteed to exist"
+        ),
+    )
+
     history_record_class_path = "features.models.HistoricalFeature"
     related_object_type = RelatedObjectType.FEATURE
 
@@ -178,9 +188,20 @@ def create_gitlab_comment(self) -> None:
                 args=(self.name, self.id, self.project_id),
             )
 
+    @hook(BEFORE_CREATE)  # type: ignore[misc]
+    def mark_as_creating(self) -> None:
+        # The feature row is committed before `create_feature_states` seeds its feature
+        # states, so anything reading the project in between sees a feature with no
+        # states in any environment. `is_creating` marks that window for the
+        # environment document writer.
+        self.is_creating = True
+
     @hook(AFTER_CREATE)
     def create_feature_states(self):  # type: ignore[no-untyped-def]
         FeatureState.create_initial_feature_states_for_feature(feature=self)
+        self.is_creating = False
+        # Update the row directly to avoid re-triggering hooks for a lifecycle flag.
+        Feature.objects.filter(id=self.id).update(is_creating=False)
 
     @hook(AFTER_SAVE)  # type: ignore[misc]
     def delete_identity_overrides(self) -> None: