fix(flags): prevent stale definition publication

.sampo/changesets/ardent-baroness-tapio.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Prevent stale feature flag definition publication

posthog/client.py

@@ -627,6 +627,11 @@ 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_fetch_generation = 0
+        self._flag_definition_published_generation = 0
+        self._flag_definition_cache_generation = 0
+        self._flag_definition_publication_lock = threading.Lock()
+        self._flag_definition_cache_write_lock = threading.RLock()
         self._flag_definition_cache_provider = flag_definition_cache_provider
         self._flag_definition_cache_provider_async_runner: Optional[
             _BackgroundEventLoopRunner
@@ -1864,6 +1869,11 @@ def _reinit_after_fork(self):
         self._flag_definition_cache_provider_async_runner = None
         self._flag_definition_cache_provider_async_runner_lock = threading.Lock()
 
+        # A parent thread may have been publishing or caching flag definitions at
+        # fork time.
+        self._flag_definition_publication_lock = threading.Lock()
+        self._flag_definition_cache_write_lock = threading.RLock()
+
         # Metrics locks may have been held by a parent thread at fork time; replace
         # them (never acquire them) so the child can't deadlock on a vanished holder.
         self._metrics_lock = threading.Lock()
@@ -2216,91 +2226,140 @@ def _fetch_feature_flags_from_api(self):
             )
             return
 
-        try:
-            # Store old flags to detect changes
-            old_flags_by_key: dict[str, dict] = self.feature_flags_by_key or {}
+        with self._flag_definition_publication_lock:
+            self._flag_definition_fetch_generation += 1
+            fetch_generation = self._flag_definition_fetch_generation
+            request_etag = self._flags_etag
 
+        cache_data_to_store: Optional[FlagDefinitionCacheData] = None
+        try:
             response = get(
                 personal_api_key,
                 f"/flags/definitions?token={self.api_key}&send_cohorts",
                 self.host,
                 timeout=10,
-                etag=self._flags_etag,
+                etag=request_etag,
             )
 
-            # Update stored ETag (clear if server stops sending one)
-            self._flags_etag = response.etag
-
-            # If 304 Not Modified, flags haven't changed - skip processing
-            if response.not_modified:
-                self.log.debug(
-                    "[FEATURE FLAGS] Flags not modified (304), using cached data"
-                )
-                self._last_feature_flag_poll = datetime.now(tz=timezone.utc)
-                return
-
-            if response.data is None:
-                self.log.error(
-                    "[FEATURE FLAGS] Unexpected empty response data in non-304 response"
-                )
-                return
-
-            self._update_flag_state(response.data, old_flags_by_key=old_flags_by_key)
+            with self._flag_definition_publication_lock:
+                if fetch_generation <= self._flag_definition_published_generation:
+                    self.log.debug(
+                        "[FEATURE FLAGS] Ignoring stale flag definition response"
+                    )
+                    self._last_feature_flag_poll = datetime.now(tz=timezone.utc)
+                    return
 
-            # Store in external cache if provider is configured
-            if self._flag_definition_cache_provider:
-                try:
-                    self._resolve_flag_definition_cache_provider_result(
-                        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 {},
-                                "minimal_flag_called_events": self._minimal_flag_called_events,
-                            }
+                # A 304 is valid only for the ETag used by this request. Another
+                # overlapping response may already have installed newer definitions.
+                if response.not_modified:
+                    if self._flags_etag != request_etag:
+                        self.log.debug(
+                            "[FEATURE FLAGS] Ignoring stale 304 flag definition response"
                         )
-                    )
-                except Exception as e:
-                    self.log.error(f"[FEATURE FLAGS] Cache provider store error: {e}")
-                    # Flags are already in memory, so continue normally
+                        self._last_feature_flag_poll = datetime.now(tz=timezone.utc)
+                        return
 
-        except APIError as e:
-            if e.status == 401:
-                detail = (
-                    f"Error loading feature flags: {e.message}. "
-                    "Please verify both your project_api_key and secret_key. "
-                    "More information: https://posthog.com/docs/api/overview"
-                )
-                self.log.error("[FEATURE FLAGS] %s", detail)
-                self.feature_flags = []
-                self.group_type_mapping = {}
-                self.cohorts = {}
+                    self._flags_etag = response.etag
+                    self._flag_definition_published_generation = fetch_generation
+                    self._flag_definition_cache_generation = fetch_generation
+                    self.log.debug(
+                        "[FEATURE FLAGS] Flags not modified (304), using cached data"
+                    )
+                    self._last_feature_flag_poll = datetime.now(tz=timezone.utc)
+                    return
 
-                if self.flag_cache:
-                    self.flag_cache.clear()
+                if response.data is None:
+                    self.log.error(
+                        "[FEATURE FLAGS] Unexpected empty response data in non-304 response"
+                    )
+                    return
 
-                if self.debug:
-                    raise APIError(status=401, message=detail)
-            elif e.status == 402:
-                self.log.warning(
-                    "[FEATURE FLAGS] PostHog feature flags quota limited, resetting feature flag data.  Learn more about billing limits at https://posthog.com/docs/billing/limits-alerts"
+                old_flags_by_key: dict[str, dict] = self.feature_flags_by_key or {}
+                self._update_flag_state(
+                    response.data, old_flags_by_key=old_flags_by_key
                 )
-                # Reset all feature flag data when quota limited
-                self.feature_flags = []
-                self.group_type_mapping = {}
-                self.cohorts = {}
 
-                # Clear flag cache when quota limited
-                if self.flag_cache:
-                    self.flag_cache.clear()
+                if self._flag_definition_cache_provider:
+                    cache_data_to_store = {
+                        "flags": self.feature_flags or [],
+                        "group_type_mapping": self.group_type_mapping or {},
+                        "cohorts": self.cohorts or {},
+                        "minimal_flag_called_events": self._minimal_flag_called_events,
+                    }
+
+                # Publish the ETag only after its matching flag state is installed.
+                self._flags_etag = response.etag
+                self._flag_definition_published_generation = fetch_generation
+                self._flag_definition_cache_generation = fetch_generation
+
+            if cache_data_to_store and self._flag_definition_cache_provider:
+                # Keep provider I/O out of the publication lock. The separate lock
+                # preserves cache write order without delaying newer API fetches or
+                # in-memory publication.
+                with self._flag_definition_cache_write_lock:
+                    with self._flag_definition_publication_lock:
+                        should_store = (
+                            fetch_generation == self._flag_definition_cache_generation
+                        )
+                    if should_store:
+                        try:
+                            self._resolve_flag_definition_cache_provider_result(
+                                self._flag_definition_cache_provider.on_flag_definitions_received(
+                                    cache_data_to_store
+                                )
+                            )
+                        except Exception as e:
+                            self.log.error(
+                                f"[FEATURE FLAGS] Cache provider store error: {e}"
+                            )
+                            # Flags are already in memory, so continue normally
 
-                if self.debug:
-                    raise APIError(
-                        status=402,
-                        message="PostHog feature flags quota limited",
+        except APIError as e:
+            with self._flag_definition_publication_lock:
+                if fetch_generation <= self._flag_definition_published_generation:
+                    self.log.debug("[FEATURE FLAGS] Ignoring stale API error response")
+                elif e.status == 401:
+                    detail = (
+                        f"Error loading feature flags: {e.message}. "
+                        "Please verify both your project_api_key and secret_key. "
+                        "More information: https://posthog.com/docs/api/overview"
                     )
-            else:
-                self.log.error(f"[FEATURE FLAGS] Error loading feature flags: {e}")
+                    self.log.error("[FEATURE FLAGS] %s", detail)
+                    self.feature_flags = []
+                    self.group_type_mapping = {}
+                    self.cohorts = {}
+                    self._flags_etag = None
+                    self._flag_definition_published_generation = fetch_generation
+                    self._flag_definition_cache_generation = fetch_generation
+
+                    if self.flag_cache:
+                        self.flag_cache.clear()
+
+                    if self.debug:
+                        raise APIError(status=401, message=detail)
+                elif e.status == 402:
+                    self.log.warning(
+                        "[FEATURE FLAGS] PostHog feature flags quota limited, resetting feature flag data.  Learn more about billing limits at https://posthog.com/docs/billing/limits-alerts"
+                    )
+                    # Reset all feature flag data when quota limited
+                    self.feature_flags = []
+                    self.group_type_mapping = {}
+                    self.cohorts = {}
+                    self._flags_etag = None
+                    self._flag_definition_published_generation = fetch_generation
+                    self._flag_definition_cache_generation = fetch_generation
+
+                    # Clear flag cache when quota limited
+                    if self.flag_cache:
+                        self.flag_cache.clear()
+
+                    if self.debug:
+                        raise APIError(
+                            status=402,
+                            message="PostHog feature flags quota limited",
+                        )
+                else:
+                    self.log.error(f"[FEATURE FLAGS] Error loading feature flags: {e}")
         except Exception as e:
             self.log.warning(
                 "[FEATURE FLAGS] Fetching feature flags failed with following error. We will retry in %s seconds."

posthog/test/test_client_fork.py

@@ -92,6 +92,26 @@ def test_reinit_after_fork_clears_poller_when_local_evaluation_disabled(self):
 
         self.assertIsNone(client.poller)
 
+    def test_reinit_after_fork_replaces_flag_definition_locks(self):
+        client = Client(FAKE_TEST_API_KEY, send=False)
+        old_publication_lock = client._flag_definition_publication_lock
+        old_cache_write_lock = client._flag_definition_cache_write_lock
+        old_publication_lock.acquire()
+        old_cache_write_lock.acquire()
+
+        try:
+            client._reinit_after_fork()
+            for new_lock, old_lock in (
+                (client._flag_definition_publication_lock, old_publication_lock),
+                (client._flag_definition_cache_write_lock, old_cache_write_lock),
+            ):
+                self.assertIsNot(new_lock, old_lock)
+                self.assertTrue(new_lock.acquire(blocking=False))
+                new_lock.release()
+        finally:
+            old_cache_write_lock.release()
+            old_publication_lock.release()
+
     @mock.patch("posthog.client.reset_sessions")
     def test_reinit_after_fork_resets_sessions(self, mock_reset_sessions):
         client = Client(FAKE_TEST_API_KEY, send=False)