feat: add permanent environment document cache (#5187)

Author: matthewelwell

api/app/settings/common.py

@@ -28,6 +28,7 @@
 
 from app.routers import ReplicaReadStrategy
 from app.utils import get_numbered_env_vars_with_prefix
+from environments.enums import EnvironmentDocumentCacheMode
 
 django_stubs_ext.monkeypatch()
 
@@ -703,8 +704,30 @@
     "django.core.cache.backends.locmem.LocMemCache",
 )
 
+CACHE_ENVIRONMENT_DOCUMENT_LOCATION = env(
+    "CACHE_ENVIRONMENT_DOCUMENT_LOCATION", default="environment-documents"
+)
+CACHE_ENVIRONMENT_DOCUMENT_BACKEND = env(
+    "CACHE_ENVIRONMENT_DOCUMENT_BACKEND", "django.core.cache.backends.db.DatabaseCache"
+)
+CACHE_ENVIRONMENT_DOCUMENT_MODE = env.enum(
+    "CACHE_ENVIRONMENT_DOCUMENT_MODE",
+    enum=EnvironmentDocumentCacheMode,
+    default=EnvironmentDocumentCacheMode.EXPIRING.value,
+)
 CACHE_ENVIRONMENT_DOCUMENT_SECONDS = env.int("CACHE_ENVIRONMENT_DOCUMENT_SECONDS", 0)
-ENVIRONMENT_DOCUMENT_CACHE_LOCATION = "environment-documents"
+CACHE_ENVIRONMENT_DOCUMENT_OPTIONS = env.json(
+    "CACHE_ENVIRONMENT_DOCUMENT_OPTIONS", default=None
+)
+
+if (
+    CACHE_ENVIRONMENT_DOCUMENT_MODE == EnvironmentDocumentCacheMode.PERSISTENT
+    and CACHE_ENVIRONMENT_DOCUMENT_SECONDS
+):
+    warnings.warn(
+        "Ignoring CACHE_ENVIRONMENT_DOCUMENT_SECONDS variable "
+        'since CACHE_ENVIRONMENT_DOCUMENT_MODE == "PERSISTENT"'
+    )  # pragma: no cover
 
 USER_THROTTLE_CACHE_NAME = "user-throttle"
 USER_THROTTLE_CACHE_BACKEND = env.str(
@@ -761,10 +784,16 @@
         "LOCATION": CHARGEBEE_CACHE_LOCATION,
         "TIMEOUT": 12 * 60 * 60,  # 12 hours
     },
-    ENVIRONMENT_DOCUMENT_CACHE_LOCATION: {
-        "BACKEND": "django.core.cache.backends.db.DatabaseCache",
-        "LOCATION": ENVIRONMENT_DOCUMENT_CACHE_LOCATION,
-        "timeout": CACHE_ENVIRONMENT_DOCUMENT_SECONDS,
+    CACHE_ENVIRONMENT_DOCUMENT_LOCATION: {
+        "BACKEND": CACHE_ENVIRONMENT_DOCUMENT_BACKEND,
+        "LOCATION": CACHE_ENVIRONMENT_DOCUMENT_LOCATION,
+        "TIMEOUT": (
+            None
+            if CACHE_ENVIRONMENT_DOCUMENT_MODE
+            == EnvironmentDocumentCacheMode.PERSISTENT
+            else CACHE_ENVIRONMENT_DOCUMENT_SECONDS
+        ),
+        "OPTIONS": CACHE_ENVIRONMENT_DOCUMENT_OPTIONS or {},
     },
     GET_FLAGS_ENDPOINT_CACHE_NAME: {
         "BACKEND": GET_FLAGS_ENDPOINT_CACHE_BACKEND,

api/environments/admin.py

@@ -2,6 +2,8 @@
 from __future__ import unicode_literals
 
 from django.contrib import admin
+from django.db.models import QuerySet
+from django.http import HttpRequest
 
 from .models import Environment, Webhook
 from .tasks import rebuild_environment_document
@@ -30,6 +32,8 @@ class EnvironmentAdmin(admin.ModelAdmin):  # type: ignore[type-arg]
     inlines = (WebhookInline,)
 
     @admin.action(description="Rebuild selected environment documents")
-    def rebuild_environments(self, request, queryset):  # type: ignore[no-untyped-def]
+    def rebuild_environments(
+        self, request: HttpRequest, queryset: QuerySet[Environment]
+    ) -> None:
         for environment in queryset:
             rebuild_environment_document.delay(args=(environment.id,))

api/environments/enums.py

@@ -0,0 +1,6 @@
+import enum
+
+
+class EnvironmentDocumentCacheMode(enum.Enum):
+    PERSISTENT = "PERSISTENT"
+    EXPIRING = "EXPIRING"

api/environments/metrics.py

@@ -0,0 +1,10 @@
+import prometheus_client
+
+CACHE_HIT = "CACHE_HIT"
+CACHE_MISS = "CACHE_MISS"
+
+flagsmith_environment_document_cache_queries_total = prometheus_client.Counter(
+    "flagsmith_environment_document_cache_queries_total",
+    "Results of cache retrieval for environment document (hit or miss)",
+    ["result"],
+)

api/environments/models.py

@@ -38,20 +38,29 @@
     DynamoEnvironmentV2Wrapper,
     DynamoEnvironmentWrapper,
 )
+from environments.enums import EnvironmentDocumentCacheMode
 from environments.exceptions import EnvironmentHeaderNotPresentError
 from environments.managers import EnvironmentManager
+from environments.metrics import (
+    CACHE_HIT,
+    CACHE_MISS,
+    flagsmith_environment_document_cache_queries_total,
+)
 from features.models import Feature, FeatureSegment, FeatureState
 from features.multivariate.models import MultivariateFeatureStateValue
 from metadata.models import Metadata
 from projects.models import Project
 from segments.models import Segment
-from util.mappers import map_environment_to_sdk_document
+from util.mappers import (
+    map_environment_to_environment_document,
+    map_environment_to_sdk_document,
+)
 from webhooks.models import AbstractBaseExportableWebhookModel
 
 logger = logging.getLogger(__name__)
 
 environment_cache = caches[settings.ENVIRONMENT_CACHE_NAME]
-environment_document_cache = caches[settings.ENVIRONMENT_DOCUMENT_CACHE_LOCATION]
+environment_document_cache = caches[settings.CACHE_ENVIRONMENT_DOCUMENT_LOCATION]
 environment_segments_cache = caches[settings.ENVIRONMENT_SEGMENTS_CACHE_NAME]
 bad_environments_cache = caches[settings.BAD_ENVIRONMENTS_CACHE_LOCATION]
 
@@ -149,22 +158,36 @@ class Environment(
     class Meta:
         ordering = ["id"]
 
-    @hook(AFTER_CREATE)
-    def create_feature_states(self):  # type: ignore[no-untyped-def]
+    @hook(AFTER_CREATE)  # type: ignore[misc]
+    def create_feature_states(self) -> None:
         FeatureState.create_initial_feature_states_for_environment(environment=self)
 
-    @hook(AFTER_UPDATE)
-    def clear_environment_cache(self):  # type: ignore[no-untyped-def]
+    @hook(AFTER_UPDATE)  # type: ignore[misc]
+    def clear_environment_cache(self) -> None:
         # TODO: this could rebuild the cache itself (using an async task)
         environment_cache.delete(self.initial_value("api_key"))
 
-    @hook(AFTER_DELETE)
-    def delete_from_dynamo(self):  # type: ignore[no-untyped-def]
+    @hook(AFTER_UPDATE, when="api_key", has_changed=True)  # type: ignore[misc]
+    def update_environment_document_cache(self) -> None:
+        environment_document_cache.delete(self.initial_value("api_key"))
+        self.write_environment_documents(self.id)
+
+    @hook(AFTER_DELETE)  # type: ignore[misc]
+    def delete_from_dynamo(self) -> None:
         if self.project.enable_dynamo_db and environment_wrapper.is_enabled:
             from environments.tasks import delete_environment_from_dynamo
 
             delete_environment_from_dynamo.delay(args=(self.api_key, self.id))
 
+    @hook(AFTER_DELETE)  # type: ignore[misc]
+    def delete_environment_document_from_cache(self) -> None:
+        if (
+            settings.CACHE_ENVIRONMENT_DOCUMENT_MODE
+            == EnvironmentDocumentCacheMode.PERSISTENT
+            or settings.CACHE_ENVIRONMENT_DOCUMENT_SECONDS > 0
+        ):
+            environment_document_cache.delete(self.api_key)
+
     def __str__(self):  # type: ignore[no-untyped-def]
         return "Project %s - Environment %s" % (self.project.name, self.name)
 
@@ -245,7 +268,7 @@ def get_from_cache(cls, api_key):  # type: ignore[no-untyped-def]
             logger.info("Environment with api_key %s does not exist" % api_key)
 
     @classmethod
-    def write_environments_to_dynamodb(
+    def write_environment_documents(
         cls,
         environment_id: int = None,  # type: ignore[assignment]
         project_id: int = None,  # type: ignore[assignment]
@@ -281,17 +304,31 @@ def write_environments_to_dynamodb(
         # project (which should always be the case). Since we're working with fairly
         # small querysets here, this shouldn't have a noticeable impact on performance.
         project: Project | None = getattr(environments[0], "project", None)
-        for environment in environments[1:]:
-            if not environment.project == project:
-                raise RuntimeError("Environments must all belong to the same project.")
-
-        if not all([project, project.enable_dynamo_db, environment_wrapper.is_enabled]):  # type: ignore[union-attr]
+        if project is None:  # pragma: no cover
             return
 
-        environment_wrapper.write_environments(environments)
+        for environment in environments[1:]:
+            if not environment.project == project:  # pragma: no cover
+                raise RuntimeError("Environments must all belong to the same project.")
 
-        if project.edge_v2_environments_migrated and environment_v2_wrapper.is_enabled:  # type: ignore[union-attr]
-            environment_v2_wrapper.write_environments(environments)
+        if project.enable_dynamo_db and environment_wrapper.is_enabled:
+            environment_wrapper.write_environments(environments)
+
+            if (
+                project.edge_v2_environments_migrated
+                and environment_v2_wrapper.is_enabled
+            ):
+                environment_v2_wrapper.write_environments(environments)
+        elif (
+            settings.CACHE_ENVIRONMENT_DOCUMENT_MODE
+            == EnvironmentDocumentCacheMode.PERSISTENT
+        ):
+            environment_document_cache.set_many(
+                {
+                    e.api_key: map_environment_to_environment_document(e)
+                    for e in environments
+                }
+            )
 
     def get_feature_state(
         self,
@@ -364,7 +401,11 @@ def get_environment_document(
         cls,
         api_key: str,
     ) -> dict[str, typing.Any]:
-        if settings.CACHE_ENVIRONMENT_DOCUMENT_SECONDS > 0:
+        if (
+            settings.CACHE_ENVIRONMENT_DOCUMENT_SECONDS > 0
+            or settings.CACHE_ENVIRONMENT_DOCUMENT_MODE
+            == EnvironmentDocumentCacheMode.PERSISTENT
+        ):
             return cls._get_environment_document_from_cache(api_key)
         return cls._get_environment_document_from_db(api_key)
 
@@ -386,9 +427,14 @@ def _get_environment_document_from_cache(
         api_key: str,
     ) -> dict[str, typing.Any]:
         environment_document = environment_document_cache.get(api_key)
-        if not environment_document:
+        if not (cache_hit := environment_document is not None):
             environment_document = cls._get_environment_document_from_db(api_key)
             environment_document_cache.set(api_key, environment_document)
+
+        flagsmith_environment_document_cache_queries_total.labels(
+            result=CACHE_HIT if cache_hit else CACHE_MISS,
+        ).inc()
+
         return environment_document  # type: ignore[no-any-return]
 
     @classmethod

api/environments/tasks.py

@@ -19,15 +19,15 @@
 
 @register_task_handler(priority=TaskPriority.HIGH)
 def rebuild_environment_document(environment_id: int) -> None:
-    Environment.write_environments_to_dynamodb(environment_id=environment_id)
+    Environment.write_environment_documents(environment_id=environment_id)
 
 
 @register_task_handler(priority=TaskPriority.HIGHEST)
 def process_environment_update(audit_log_id: int):  # type: ignore[no-untyped-def]
     audit_log = AuditLog.objects.get(id=audit_log_id)
 
     # Send environment document to dynamodb
-    Environment.write_environments_to_dynamodb(
+    Environment.write_environment_documents(
         environment_id=audit_log.environment_id, project_id=audit_log.project_id
     )
 

api/integrations/common/models.py

@@ -35,7 +35,7 @@ def write_environment_to_dynamodb(self):  # type: ignore[no-untyped-def]
                 self.__class__.__name__,
             )
             return
-        Environment.write_environments_to_dynamodb(environment_id=self.environment_id)
+        Environment.write_environment_documents(environment_id=self.environment_id)
 
     @hook(AFTER_UPDATE)
     def clear_environment_cache(self):  # type: ignore[no-untyped-def]

api/integrations/webhook/models.py

@@ -21,4 +21,4 @@ class WebhookConfiguration(
     @hook(AFTER_SAVE)
     @hook(AFTER_DELETE)
     def write_environment_to_dynamodb(self):  # type: ignore[no-untyped-def]
-        Environment.write_environments_to_dynamodb(environment_id=self.environment_id)
+        Environment.write_environment_documents(environment_id=self.environment_id)

api/projects/tasks.py

@@ -11,7 +11,7 @@
 def write_environments_to_dynamodb(project_id: int) -> None:
     from environments.models import Environment
 
-    Environment.write_environments_to_dynamodb(project_id=project_id)
+    Environment.write_environment_documents(project_id=project_id)
 
 
 @register_task_handler()

api/tests/integration/conftest.py

@@ -2,13 +2,17 @@
 import uuid
 
 import pytest
+from django.core.cache import BaseCache
+from django.core.cache.backends.locmem import LocMemCache
 from django.test import Client as DjangoClient
 from django.urls import reverse
 from pytest_django.fixtures import SettingsWrapper
+from pytest_mock import MockerFixture
 from rest_framework import status
 from rest_framework.test import APIClient
 
 from app.utils import create_hash
+from environments.enums import EnvironmentDocumentCacheMode
 from organisations.models import Organisation
 from tests.integration.helpers import create_mv_option_with_api
 
@@ -524,3 +528,15 @@ def identity_featurestate(admin_client, environment, feature, identity):  # type
         url, data=json.dumps(data), content_type="application/json"
     )
     return response.json()["id"]
+
+
+@pytest.fixture()
+def persistent_environment_document_cache(
+    settings: SettingsWrapper,
+    mocker: MockerFixture,
+    environment: int,
+) -> BaseCache:
+    settings.CACHE_ENVIRONMENT_DOCUMENT_MODE = EnvironmentDocumentCacheMode.PERSISTENT
+    cache = LocMemCache(name="environment_document", params={})
+    mocker.patch("environments.models.environment_document_cache", cache)
+    return cache