fix(cohorts): guard cohort recalculation Redis calls

products/cohorts/backend/models/dependencies.py

@@ -204,6 +204,22 @@ def _invalidate_team_behavioral_cohort_cache(team_id: int) -> None:
 invalidate_team_behavioral_cohort_cache = _invalidate_team_behavioral_cohort_cache
 
 
+def _on_cohort_behavioral_cache_invalidated(team_id: int) -> None:
+    # Runs synchronously inside Cohort.save()/delete() under autocommit, so an unguarded Redis error
+    # here would escape the save, the same failure the _on_cohort_changed guard prevents. The
+    # behavioral-cohort cache is best-effort with a TTL and rebuilds on the next read, so a failed
+    # invalidation is safe to log and continue past. The public alias stays unguarded on purpose: the
+    # backfill finalizer wraps its own call and accounts for the failure.
+    try:
+        _invalidate_team_behavioral_cohort_cache(team_id)
+    except Exception as error:
+        logger.exception(
+            "failed_to_invalidate_team_behavioral_cohort_cache",
+            team_id=team_id,
+            error=str(error),
+        )
+
+
 def extract_cohort_dependencies(cohort: Cohort) -> set[int]:
     """
     Extract cohort dependencies from the given cohort.
@@ -315,22 +331,33 @@ def warm_team_cohort_dependency_cache(team_id: int, batch_size: int = 1000):
 
 
 def _on_cohort_changed(cohort: Cohort, always_invalidate: bool = False):
-    new_dependencies = extract_cohort_dependencies(cohort)
-    existing_dependencies = cache.get(_cohort_dependencies_key(cohort.id))
-    dependencies_changed = existing_dependencies is None or set(existing_dependencies) != new_dependencies
+    # Under autocommit this runs synchronously inside Cohort.save(), so a Redis error here escapes the
+    # save and can leave the cohort stuck with is_calculating=True. The dependency caches rebuild on the
+    # next change or warm pass, so a failed refresh is safe to log and continue past.
+    try:
+        new_dependencies = extract_cohort_dependencies(cohort)
+        existing_dependencies = cache.get(_cohort_dependencies_key(cohort.id))
+        dependencies_changed = existing_dependencies is None or set(existing_dependencies) != new_dependencies
 
-    # If the dependencies haven't changed, no need to refresh the cache
-    if not always_invalidate and not cohort.deleted and not dependencies_changed:
-        return
+        # If the dependencies haven't changed, no need to refresh the cache
+        if not always_invalidate and not cohort.deleted and not dependencies_changed:
+            return
 
-    cache.delete(_cohort_dependencies_key(cohort.id))
-    cache.delete(_cohort_dependents_key(cohort.id))
+        cache.delete(_cohort_dependencies_key(cohort.id))
+        cache.delete(_cohort_dependents_key(cohort.id))
 
-    if existing_dependencies:
-        for dep_id in existing_dependencies:
-            cache.delete(_cohort_dependents_key(dep_id))
+        if existing_dependencies:
+            for dep_id in existing_dependencies:
+                cache.delete(_cohort_dependents_key(dep_id))
 
-    warm_team_cohort_dependency_cache(cohort.team_id)
+        warm_team_cohort_dependency_cache(cohort.team_id)
+    except Exception as error:
+        logger.exception(
+            "failed_to_refresh_cohort_dependency_cache",
+            cohort_id=cohort.pk,
+            team_id=cohort.team_id,
+            error=str(error),
+        )
 
 
 def _has_backfillable_filters(cohort: Cohort, kind: CohortBackfillKind) -> bool:
@@ -443,7 +470,7 @@ def cohort_changed(sender, instance, **kwargs):
         return
 
     transaction.on_commit(lambda: _on_cohort_changed(instance))
-    transaction.on_commit(lambda: _invalidate_team_behavioral_cohort_cache(instance.team_id))
+    transaction.on_commit(lambda: _on_cohort_behavioral_cache_invalidated(instance.team_id))
 
 
 @receiver(post_save, sender=Cohort)
@@ -561,7 +588,7 @@ def cohort_deleted(sender, instance, **kwargs):
     Clear and rebuild dependency caches when cohort is deleted.
     """
     transaction.on_commit(lambda: _on_cohort_changed(instance, always_invalidate=True))
-    transaction.on_commit(lambda: _invalidate_team_behavioral_cohort_cache(instance.team_id))
+    transaction.on_commit(lambda: _on_cohort_behavioral_cache_invalidated(instance.team_id))
 
 
 @receiver(post_delete, sender=Team)

products/cohorts/backend/models/test/test_dependencies.py

@@ -258,6 +258,26 @@ def test_warm_team_cohort_dependency_cache_refreshes_ttl(self) -> None:
             args, kwargs = mock_set_many.call_args
             self.assertEqual(kwargs.get("timeout"), DEPENDENCY_CACHE_TIMEOUT)
 
+    @parameterized.expand(
+        [
+            # Dependency-refresh callback: cache.get raises inside _on_cohort_changed.
+            ("dependency_refresh", "products.cohorts.backend.models.dependencies.cache.get"),
+            # Behavioral-cache callback: cache.delete_many raises inside the invalidation callback.
+            ("behavioral_cache", "products.cohorts.backend.models.dependencies.cache.delete_many"),
+        ]
+    )
+    def test_cohort_save_survives_redis_error_in_commit_callback(self, _name: str, redis_target: str) -> None:
+        # A non-recalculation save fires both cohort_changed commit callbacks synchronously (see
+        # mock_transaction). A Redis error in either must not escape Cohort.save() and leave the cohort
+        # stuck calculating.
+        cohort = self._create_cohort(name="Test Cohort")
+        with mock.patch(redis_target, side_effect=Exception("redis unavailable")):
+            cohort.name = "renamed"
+            cohort.save()
+
+        cohort.refresh_from_db()
+        self.assertEqual(cohort.name, "renamed")
+
     def test_cache_miss_get_cohort_dependencies(self) -> None:
         cohort_a = self._create_cohort(name="Test Cohort A")
         cohort_b = self._create_cohort(

products/cohorts/backend/models/test/test_util.py

@@ -4,9 +4,10 @@
 from unittest.mock import MagicMock, patch
 
 from django.db import DEFAULT_DB_ALIAS, OperationalError
-from django.test import SimpleTestCase
+from django.test import SimpleTestCase, override_settings
 
 from clickhouse_driver.errors import SocketTimeoutError
+from kombu.exceptions import OperationalError as BrokerOperationalError
 from parameterized import parameterized
 from pydantic import (
     BaseModel,
@@ -38,6 +39,7 @@
     insert_cohort_query_actors_into_ch,
     parse_error_code,
     print_cohort_hogql_query,
+    run_cohort_query,
     simplified_cohort_filter_properties,
     sort_cohorts_topologically,
     validate_actors_query_for_cohort,
@@ -1360,10 +1362,17 @@ def test_original_error_surfaces_when_recovery_history_save_fails(self):
         assert mock_history.save.call_count == 2
         mock_connections[DEFAULT_DB_ALIAS].close.assert_called_once()
 
-    def test_original_error_surfaces_when_reset_calculating_save_fails(self):
-        # calculate_people_ch's finally block resets is_calculating on the same connection the recovery
-        # bookkeeping save uses. If that reset write hits the dropped connection it must not raise out of
-        # finally and mask the real calculation error - it runs through the same reconnect-and-retry.
+    @parameterized.expand(
+        [
+            # A dropped Postgres connection, and a Redis error escaping the save's post_save signals.
+            ("connection_error", OperationalError("the connection is closed")),
+            ("redis_error", Exception("redis unavailable")),
+        ]
+    )
+    def test_original_error_surfaces_when_reset_calculating_save_fails(self, _name, bookkeeping_error):
+        # calculate_people_ch's finally block resets is_calculating through the recovery bookkeeping.
+        # If that reset write fails - a dropped connection, or a Redis error from the save's signals -
+        # it must not raise out of finally and mask the real calculation error.
         cohort = _create_cohort(
             team=self.team,
             name="c",
@@ -1376,11 +1385,11 @@ def test_original_error_surfaces_when_reset_calculating_save_fails(self):
                 "products.cohorts.backend.models.util.recalculate_cohortpeople",
                 side_effect=real_error,
             ),
-            # The is_calculating reset write fails on the dropped connection, initial attempt and retry.
+            # The is_calculating reset write fails, initial attempt and retry.
             patch.object(
                 Cohort,
                 "_safe_reset_calculating_state",
-                side_effect=OperationalError("the connection is closed"),
+                side_effect=bookkeeping_error,
             ) as mock_reset,
             # Patch connections so the reconnect doesn't close the real test transaction's connection.
             patch("products.cohorts.backend.models.util.connections") as mock_connections,
@@ -1391,3 +1400,25 @@ def test_original_error_surfaces_when_reset_calculating_save_fails(self):
         self.assertIs(ctx.exception, real_error)
         assert mock_reset.call_count == 2
         mock_connections[DEFAULT_DB_ALIAS].close.assert_called_once()
+
+
+class TestRunCohortQueryScheduling(SimpleTestCase):
+    @override_settings(TEST=False, IN_EVAL_TESTING=False)
+    def test_survives_broker_error_scheduling_initial_stats_task(self):
+        # The initial stats task is scheduled before the query runs. A broker blip there must not
+        # abort the recalculation before any ClickHouse work happens - run_cohort_query has to swallow
+        # it and still return fn's result, the same way the reschedule call after fn is guarded.
+        history = MagicMock()
+        with patch(
+            "posthog.tasks.calculate_cohort.collect_cohort_query_stats.apply_async",
+            side_effect=BrokerOperationalError("broker unavailable"),
+        ) as mock_apply:
+            result, _end_time = run_cohort_query(
+                lambda: "calc-result",
+                cohort_id=123,
+                history=history,
+                query="SELECT 1",
+            )
+
+        assert result == "calc-result"
+        mock_apply.assert_called_once()

products/cohorts/backend/models/util.py

@@ -8,7 +8,7 @@
 from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
 from django.conf import settings
-from django.db import DEFAULT_DB_ALIAS, InterfaceError, OperationalError, connections
+from django.db import DEFAULT_DB_ALIAS, connections
 from django.utils import timezone
 
 import structlog
@@ -156,15 +156,16 @@ def save_recovery_bookkeeping(save_fn: Callable[[], None], *, cohort_id: int, te
     """Persist post-calculation bookkeeping, surviving a Postgres connection dropped mid-recalculation.
 
     A long recalculation can outlive its connection (the server closes it unexpectedly); the first
-    write afterwards then raises a connection error. Left unguarded on the error/finally recovery
-    path, that cascades into "the connection is closed" - burying the real root-cause error and
-    leaving the cohort stuck with is_calculating=True. Reconnect and retry once so the bookkeeping
-    still lands and the original error is what propagates; if the retry fails too, swallow it (a
-    recovery write must never mask the failure it is recording).
+    write afterwards then raises a connection error. The save also fires post_save signal handlers
+    that touch Redis, so a Redis blip can escape here too. Left unguarded on the error/finally
+    recovery path, either failure buries the real root-cause error and leaves the cohort stuck with
+    is_calculating=True. Reconnect and retry once so the bookkeeping still lands and the original
+    error is what propagates; if the retry fails too, swallow it (a recovery write must never mask
+    the failure it is recording).
     """
     try:
         save_fn()
-    except (InterfaceError, OperationalError):
+    except Exception:
         connections[DEFAULT_DB_ALIAS].close()  # next query opens a fresh connection
         try:
             save_fn()
@@ -212,10 +213,16 @@ def run_cohort_query(
         # Schedule delayed task to collect stats after query_log_archive is synced
         # Only if we have a history record to update and not in test mode
         if history and query and not (settings.TEST or settings.IN_EVAL_TESTING):
-            delayed_task = collect_cohort_query_stats.apply_async(
-                args=[cohort_tag, cohort_id, start_time.isoformat(), history.id, query],
-                countdown=COHORT_QUERY_TIMEOUT_SECONDS + COHORT_STATS_COLLECTION_DELAY_SECONDS,
-            )
+            # Best-effort stats telemetry scheduled before the query runs, so a broker blip here must
+            # not abort the recalculation before any ClickHouse work happens. Log and continue with no
+            # delayed task, matching the guard on the reschedule below.
+            try:
+                delayed_task = collect_cohort_query_stats.apply_async(
+                    args=[cohort_tag, cohort_id, start_time.isoformat(), history.id, query],
+                    countdown=COHORT_QUERY_TIMEOUT_SECONDS + COHORT_STATS_COLLECTION_DELAY_SECONDS,
+                )
+            except Exception as error:
+                logger.warning("cohort_stats_collection_scheduling_failed", cohort_id=cohort_id, error=str(error))
 
     try:
         result = fn(*args, **kwargs)
@@ -224,14 +231,20 @@ def run_cohort_query(
         # If calculation succeeded and we scheduled a delayed task, cancel it and run immediately
         # This avoids waiting the full timeout when the query completed quickly
         if delayed_task and history and query and not settings.TEST:
-            if delayed_task.state in ["PENDING", "RECEIVED"]:
-                delayed_task.revoke()  # Cancel the delayed task
+            # The recalculation result is already computed. This rescheduling is best-effort stats
+            # bookkeeping that talks to the Celery result backend and broker, so a Redis blip must
+            # not discard the result. The delayed task still runs after its original countdown.
+            try:
+                if delayed_task.state in ["PENDING", "RECEIVED"]:
+                    delayed_task.revoke()  # Cancel the delayed task
 
-            # Run immediately since the query already completed
-            collect_cohort_query_stats.apply_async(
-                args=[cohort_tag, cohort_id, start_time.isoformat(), history.id, query],
-                countdown=COHORT_STATS_COLLECTION_DELAY_SECONDS,
-            )
+                # Run immediately since the query already completed
+                collect_cohort_query_stats.apply_async(
+                    args=[cohort_tag, cohort_id, start_time.isoformat(), history.id, query],
+                    countdown=COHORT_STATS_COLLECTION_DELAY_SECONDS,
+                )
+            except Exception as error:
+                logger.warning("cohort_stats_collection_scheduling_failed", cohort_id=cohort_id, error=str(error))
 
         return result, end_time