feat: Add FlagDefinitionCacheProvider interface

dustinbyrne · dustinbyrne · commit 6b39b6232ac2 · 2025-12-10T13:36:52.000-05:00
diff --git a/posthog/__init__.py b/posthog/__init__.py
@@ -22,6 +22,10 @@
     InconclusiveMatchError as InconclusiveMatchError,
     RequiresServerEvaluation as RequiresServerEvaluation,
 )
+from posthog.flag_definition_cache import (
+    FlagDefinitionCacheData as FlagDefinitionCacheData,
+    FlagDefinitionCacheProvider as FlagDefinitionCacheProvider,
+)
 from posthog.request import (
     disable_connection_reuse as disable_connection_reuse,
     enable_keep_alive as enable_keep_alive,
diff --git a/posthog/client.py b/posthog/client.py
@@ -28,6 +28,10 @@
     RequiresServerEvaluation,
     match_feature_flag_properties,
 )
+from posthog.flag_definition_cache import (
+    FlagDefinitionCacheData,
+    FlagDefinitionCacheProvider,
+)
 from posthog.poller import Poller
 from posthog.request import (
     DEFAULT_HOST,
@@ -184,6 +188,7 @@ def __init__(
         before_send=None,
         flag_fallback_cache_url=None,
         enable_local_evaluation=True,
+        flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None,
         capture_exception_code_variables=False,
         code_variables_mask_patterns=None,
         code_variables_ignore_patterns=None,
@@ -233,6 +238,7 @@ def __init__(
         self.flag_cache = self._initialize_flag_cache(flag_fallback_cache_url)
         self.flag_definition_version = 0
         self._flags_etag: Optional[str] = None
+        self._flag_definition_cache_provider = flag_definition_cache_provider
         self.disabled = disabled
         self.disable_geoip = disable_geoip
         self.historical_migration = historical_migration
@@ -357,9 +363,9 @@ def feature_flags(self, flags):
             for flag in self._feature_flags
             if flag.get("key") is not None
         }
-        assert self.feature_flags_by_key is not None, (
-            "feature_flags_by_key should be initialized when feature_flags is set"
-        )
+        assert (
+            self.feature_flags_by_key is not None
+        ), "feature_flags_by_key should be initialized when feature_flags is set"
 
     def get_feature_variants(
         self,
@@ -1169,17 +1175,25 @@ def join(self):
             posthog.join()
             ```
         """
-        for consumer in self.consumers:
-            consumer.pause()
-            try:
-                consumer.join()
-            except RuntimeError:
-                # consumer thread has not started
-                pass
+        if self.consumers:
+            for consumer in self.consumers:
+                consumer.pause()
+                try:
+                    consumer.join()
+                except RuntimeError:
+                    # consumer thread has not started
+                    pass
 
         if self.poller:
             self.poller.stop()
 
+        # Shutdown the cache provider (release locks, cleanup)
+        if self._flag_definition_cache_provider:
+            try:
+                self._flag_definition_cache_provider.shutdown()
+            except Exception as e:
+                self.log.error(f"[FEATURE FLAGS] Cache provider shutdown error: {e}")
+
     def shutdown(self):
         """
         Flush all messages and cleanly shutdown the client. Call this before the process ends in serverless environments to avoid data loss.
@@ -1195,7 +1209,71 @@ def shutdown(self):
         if self.exception_capture:
             self.exception_capture.close()
 
+    def _update_flag_state(
+        self, data: FlagDefinitionCacheData, old_flags_by_key: Optional[dict] = None
+    ) -> None:
+        """Update internal flag state from cache data and invalidate evaluation cache if changed."""
+        self.feature_flags = data.get("flags") or []
+        self.group_type_mapping = data.get("group_type_mapping") or {}
+        self.cohorts = data.get("cohorts") or {}
+
+        # Invalidate evaluation cache if flag definitions changed
+        if (
+            self.flag_cache
+            and old_flags_by_key is not None
+            and old_flags_by_key != (self.feature_flags_by_key or {})
+        ):
+            old_version = self.flag_definition_version
+            self.flag_definition_version += 1
+            self.flag_cache.invalidate_version(old_version)
+
     def _load_feature_flags(self):
+        should_fetch = True
+        if self._flag_definition_cache_provider:
+            try:
+                should_fetch = (
+                    self._flag_definition_cache_provider.should_fetch_flag_definitions()
+                )
+            except Exception as e:
+                self.log.error(
+                    f"[FEATURE FLAGS] Cache provider should_fetch error: {e}"
+                )
+                # Fail-safe: fetch from API if cache provider errors
+                should_fetch = True
+
+        # If not fetching, try to get from cache
+        if not should_fetch and self._flag_definition_cache_provider:
+            try:
+                cached_data = (
+                    self._flag_definition_cache_provider.get_flag_definitions()
+                )
+                if cached_data:
+                    self.log.debug(
+                        "[FEATURE FLAGS] Using cached flag definitions from external cache"
+                    )
+                    self._update_flag_state(
+                        cached_data, old_flags_by_key=self.feature_flags_by_key or {}
+                    )
+                    self._last_feature_flag_poll = datetime.now(tz=tzutc())
+                    return
+                else:
+                    # Emergency fallback: if cache is empty and we have no flags, fetch anyway.
+                    # There's really no other way of recovering in this case.
+                    if not self.feature_flags:
+                        self.log.debug(
+                            "[FEATURE FLAGS] Cache empty and no flags loaded, falling back to API fetch"
+                        )
+                        should_fetch = True
+            except Exception as e:
+                self.log.error(f"[FEATURE FLAGS] Cache provider get error: {e}")
+                # Fail-safe: fetch from API if cache provider errors
+                should_fetch = True
+
+        if should_fetch:
+            self._fetch_feature_flags_from_api()
+
+    def _fetch_feature_flags_from_api(self):
+        """Fetch feature flags from the PostHog API."""
         try:
             # Store old flags to detect changes
             old_flags_by_key: dict[str, dict] = self.feature_flags_by_key or {}
@@ -1225,17 +1303,21 @@ def _load_feature_flags(self):
                 )
                 return
 
-            self.feature_flags = response.data["flags"] or []
-            self.group_type_mapping = response.data["group_type_mapping"] or {}
-            self.cohorts = response.data["cohorts"] or {}
+            self._update_flag_state(response.data, old_flags_by_key=old_flags_by_key)
 
-            # Check if flag definitions changed and update version
-            if self.flag_cache and old_flags_by_key != (
-                self.feature_flags_by_key or {}
-            ):
-                old_version = self.flag_definition_version
-                self.flag_definition_version += 1
-                self.flag_cache.invalidate_version(old_version)
+            # Store in external cache if provider is configured
+            if self._flag_definition_cache_provider:
+                try:
+                    self._flag_definition_cache_provider.on_flag_definitions_received(
+                        {
+                            "flags": self.feature_flags or [],
+                            "group_type_mapping": self.group_type_mapping or {},
+                            "cohorts": self.cohorts or {},
+                        }
+                    )
+                except Exception as e:
+                    self.log.error(f"[FEATURE FLAGS] Cache provider store error: {e}")
+                    # Flags are already in memory, so continue normally
 
         except APIError as e:
             if e.status == 401:
@@ -1649,9 +1731,9 @@ def _locally_evaluate_flag(
         response = None
 
         if self.feature_flags:
-            assert self.feature_flags_by_key is not None, (
-                "feature_flags_by_key should be initialized when feature_flags is set"
-            )
+            assert (
+                self.feature_flags_by_key is not None
+            ), "feature_flags_by_key should be initialized when feature_flags is set"
             # Local evaluation
             flag = self.feature_flags_by_key.get(key)
             if flag:
diff --git a/posthog/flag_definition_cache.py b/posthog/flag_definition_cache.py
@@ -0,0 +1,127 @@
+"""
+Flag Definition Cache Provider interface for multi-worker environments.
+
+EXPERIMENTAL: This API may change in future minor version bumps.
+
+This module provides an interface for external caching of feature flag definitions,
+enabling multi-worker environments (Kubernetes, load-balanced servers, serverless
+functions) to share flag definitions and reduce API calls.
+
+Usage:
+
+    from posthog import Posthog
+    from posthog.flag_definition_cache import FlagDefinitionCacheProvider
+
+    cache = RedisFlagDefinitionCache(redis_client, "my-team")
+    posthog = Posthog(
+        "<project_api_key>",
+        personal_api_key="<personal_api_key>",
+        flag_definition_cache_provider=cache,
+    )
+"""
+
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+
+from typing_extensions import Required, TypedDict
+
+
+class FlagDefinitionCacheData(TypedDict):
+    """
+    Data structure for cached flag definitions.
+
+    Attributes:
+        flags: List of feature flag definition dictionaries from the API.
+        group_type_mapping: Mapping of group type indices to group names.
+        cohorts: Dictionary of cohort definitions for local evaluation.
+    """
+
+    flags: Required[List[Dict[str, Any]]]
+    group_type_mapping: Required[Dict[str, str]]
+    cohorts: Required[Dict[str, Any]]
+
+
+@runtime_checkable
+class FlagDefinitionCacheProvider(Protocol):
+    """
+    Interface for external caching of feature flag definitions.
+
+    Enables multi-worker environments to share flag definitions, reducing API
+    calls while ensuring all workers have consistent data.
+
+    EXPERIMENTAL: This API may change in future minor version bumps.
+
+    The four methods handle the complete lifecycle of flag definition caching:
+
+    1. `should_fetch_flag_definitions()` - Called before each poll to determine
+       if this worker should fetch new definitions. Use for distributed lock
+       coordination to ensure only one worker fetches at a time.
+
+    2. `get_flag_definitions()` - Called when `should_fetch_flag_definitions()`
+       returns False. Returns cached definitions if available.
+
+    3. `on_flag_definitions_received()` - Called after successfully fetching
+       new definitions from the API. Store the data in your external cache
+       and release any locks.
+
+    4. `shutdown()` - Called when the PostHog client shuts down. Release any
+       distributed locks and clean up resources.
+
+    Error Handling:
+        All methods are wrapped in try/except. Errors will be logged but will
+        never break flag evaluation. On error:
+        - `should_fetch_flag_definitions()` errors default to fetching (fail-safe)
+        - `get_flag_definitions()` errors fall back to API fetch
+        - `on_flag_definitions_received()` errors are logged but flags remain in memory
+        - `shutdown()` errors are logged but shutdown continues
+    """
+
+    def get_flag_definitions(self) -> Optional[FlagDefinitionCacheData]:
+        """
+        Retrieve cached flag definitions.
+
+        Returns:
+            Cached flag definitions if available and valid, None otherwise.
+            Returning None will trigger a fetch from the API if this worker
+            has no flags loaded yet.
+        """
+        ...
+
+    def should_fetch_flag_definitions(self) -> bool:
+        """
+        Determine whether this instance should fetch new flag definitions.
+
+        Use this for distributed lock coordination. Only one worker should
+        return True to avoid thundering herd problems. A typical implementation
+        uses a distributed lock (e.g., Redis SETNX) that expires after the
+        poll interval.
+
+        Returns:
+            True if this instance should fetch from the API, False otherwise.
+            When False, the client will call `get_flag_definitions()` to
+            retrieve cached data instead.
+        """
+        ...
+
+    def on_flag_definitions_received(self, data: FlagDefinitionCacheData) -> None:
+        """
+        Called after successfully receiving new flag definitions from PostHog.
+
+        Use this to store the data in your external cache and release any
+        distributed locks acquired in `should_fetch_flag_definitions()`.
+
+        Args:
+            data: The flag definitions to cache, containing flags,
+                  group_type_mapping, and cohorts.
+        """
+        ...
+
+    def shutdown(self) -> None:
+        """
+        Called when the PostHog client shuts down.
+
+        Use this to release any distributed locks and clean up resources.
+        This method is called even if `should_fetch_flag_definitions()`
+        returned False, so implementations should handle the case where
+        no lock was acquired.
+        """
+        ...
diff --git a/posthog/test/test_flag_definition_cache.py b/posthog/test/test_flag_definition_cache.py
