diff --git a/posthog/apps.py b/posthog/apps.py index a080decbd310..1696bc61f83c 100644 --- a/posthog/apps.py +++ b/posthog/apps.py @@ -73,6 +73,17 @@ def ready(self): "service": settings.OTEL_SERVICE_NAME, "environment": os.getenv("OTEL_SERVICE_ENVIRONMENT"), } + # Config for the SDK's `client.metrics` API. The pinned SDK version predates + # the metrics API and ignores this attr; once posthoganalytics is bumped to + # >=7.23 it's picked up by setup(), so metrics get a real service.name + # instead of 'unknown_service'. + posthoganalytics.metrics = { # type: ignore[attr-defined] + # Same fallback as the OTel trace resource (otel_instrumentation.py) — + # metrics and traces from one process must share a service identity. + "service_name": settings.OTEL_SERVICE_NAME or "posthog-django-default", + "service_version": os.getenv("COMMIT_SHA"), + "environment": os.getenv("OTEL_SERVICE_ENVIRONMENT"), + } if str_to_bool(os.environ.get("TEMPORAL_DISABLE_EXCEPTION_VARIABLE_CAPTURE", "false")): posthoganalytics.capture_exception_code_variables = False diff --git a/posthog/celery.py b/posthog/celery.py index 8cdc55719f55..9579934b4520 100644 --- a/posthog/celery.py +++ b/posthog/celery.py @@ -1,6 +1,7 @@ import os import time import errno +import threading from django.dispatch import receiver @@ -219,6 +220,9 @@ def on_worker_start(**kwargs) -> None: _initialize_worker_metrics() +_ANALYTICS_METRICS_FLUSH_TIMEOUT_SECONDS = 5.0 + + @worker_process_shutdown.connect def on_worker_process_shutdown(**kwargs) -> None: """Remove metric files for this child so recycled workers don't leak stale data.""" @@ -227,6 +231,33 @@ def on_worker_process_shutdown(**kwargs) -> None: multiprocess.mark_process_dead(os.getpid()) + # Flush the posthoganalytics SDK's final metrics window: `client.metrics` + # aggregates in memory and flushes on an interval, so a recycled child + # (--max-tasks-per-child) would otherwise drop up to one interval of samples. + # Inert on SDK versions without the metrics API and on untouched/disabled clients. + import posthoganalytics # noqa: PLC0415 — keep the SDK off the celery import path + + try: + client = posthoganalytics.default_client + metrics = getattr(client, "metrics", None) if client is not None else None + if metrics is not None: + # Bound the wait: an unreachable metrics endpoint must not hold a + # recycling child hostage (same hazard otel_instrumentation.py caps + # with a 5s force_flush). The daemon thread is abandoned on timeout. + def _flush() -> None: + try: + metrics.flush() + except Exception: + logger.warning("posthoganalytics_metrics_flush_failed", exc_info=True) + + flush_thread = threading.Thread(target=_flush, name="posthoganalytics-metrics-flush", daemon=True) + flush_thread.start() + flush_thread.join(timeout=_ANALYTICS_METRICS_FLUSH_TIMEOUT_SECONDS) + if flush_thread.is_alive(): + logger.warning("posthoganalytics_metrics_flush_timed_out") + except Exception: + logger.warning("posthoganalytics_metrics_flush_failed", exc_info=True) + # Set up clickhouse query instrumentation @task_prerun.connect diff --git a/posthog/test/test_celery.py b/posthog/test/test_celery.py index 5b315b3b7176..3d09fa1d4fb5 100644 --- a/posthog/test/test_celery.py +++ b/posthog/test/test_celery.py @@ -1,9 +1,70 @@ +import threading + from unittest import TestCase -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +import posthoganalytics +from parameterized import parameterized +import posthog.celery +from posthog.celery import on_worker_process_shutdown from posthog.tasks.tasks import clickhouse_errors_count +class TestWorkerShutdownFlushesAnalyticsMetrics(TestCase): + def test_flushes_sdk_metrics_tail_window(self) -> None: + client = MagicMock() + with patch.object(posthoganalytics, "default_client", client): + on_worker_process_shutdown() + client.metrics.flush.assert_called_once() + + def test_hung_flush_does_not_stall_worker_recycling(self) -> None: + release = threading.Event() + flush_completed = threading.Event() + + def hung_flush() -> None: + release.wait(timeout=10) + flush_completed.set() + + client = MagicMock(**{"metrics.flush.side_effect": hung_flush}) + try: + with ( + patch.object(posthog.celery, "_ANALYTICS_METRICS_FLUSH_TIMEOUT_SECONDS", 0.05), + patch.object(posthoganalytics, "default_client", client), + ): + on_worker_process_shutdown() + # The handler must abandon the hung flush, not wait it out. + assert not flush_completed.is_set() + finally: + release.set() + + @parameterized.expand( + [ + ("no_default_client", lambda: None), + # The pinned SDK version has no `metrics` API — the hook must stay + # inert (a bare `client.metrics.flush()` would raise on every + # worker recycle until the dependency is bumped). + ( + "real_client_on_pinned_sdk_version", + lambda: posthoganalytics.Client("phc_test", sync_mode=True, disabled=True), + ), + ("flush_raises", lambda: MagicMock(**{"metrics.flush.side_effect": RuntimeError("network down")})), + ] + ) + def test_handler_never_breaks_worker_shutdown(self, _name: str, client_factory) -> None: + with patch.object(posthoganalytics, "default_client", client_factory()): + on_worker_process_shutdown() + + +class TestAnalyticsMetricsConfig(TestCase): + def test_apps_ready_configures_module_level_metrics(self) -> None: + # Deleting the "unused" attr assignment in apps.py before the SDK bump + # would make the bump silently ship service_name='unknown_service'. + config = getattr(posthoganalytics, "metrics", None) + assert isinstance(config, dict) + assert config["service_name"] + + class TestCeleryMetrics(TestCase): @patch("posthog.clickhouse.client.sync_execute") @patch("posthog.metrics.push_to_gateway")