diff --git a/products/cohorts/backend/models/dependencies.py b/products/cohorts/backend/models/dependencies.py index faf34f9e3f67..7dec859bdd6b 100644 --- a/products/cohorts/backend/models/dependencies.py +++ b/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) diff --git a/products/cohorts/backend/models/test/test_dependencies.py b/products/cohorts/backend/models/test/test_dependencies.py index 72bc880e26f0..00a38826c820 100644 --- a/products/cohorts/backend/models/test/test_dependencies.py +++ b/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( diff --git a/products/cohorts/backend/models/test/test_util.py b/products/cohorts/backend/models/test/test_util.py index 6f99e403a56e..e6d7bd9e3978 100644 --- a/products/cohorts/backend/models/test/test_util.py +++ b/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() diff --git a/products/cohorts/backend/models/util.py b/products/cohorts/backend/models/util.py index ba751136c7c3..2b5088b621d1 100644 --- a/products/cohorts/backend/models/util.py +++ b/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