review comments

Author: parinporecha

posthog/client.py

@@ -3,6 +3,7 @@
 import os
 import sys
 import warnings
+import weakref
 from datetime import datetime, timedelta
 from typing import Any, Dict, Optional, Union
 from uuid import uuid4
@@ -334,7 +335,8 @@ def __init__(
                     consumer.start()
 
         if hasattr(os, "register_at_fork"):
-            os.register_at_fork(after_in_child=self._reinit_after_fork)
+            weak_self = weakref.ref(self)
+            os.register_at_fork(after_in_child=lambda: Client._reinit_after_fork_weak(weak_self))
 
     def new_context(self, fresh=False, capture_exceptions=True):
         """
@@ -1084,6 +1086,17 @@ def capture_exception(
         except Exception as e:
             self.log.exception(f"Failed to capture exception: {e}")
 
+    @staticmethod
+    def _reinit_after_fork_weak(weak_self):
+        """
+        Reinitialize the client after a fork.
+        Garbage collected if the client is deleted.
+        """
+        self = weak_self()
+        if self is None:
+            return
+        self._reinit_after_fork()
+
     def _reinit_after_fork(self):
         """Reinitialize queue and consumer threads in a forked child process.
 

posthog/integrations/celery.py

@@ -207,6 +207,7 @@ def _on_after_task_publish(self, *args, **kwargs):
             logger.exception("Failed to capture Celery after_task_publish lifecycle event")
 
     def _on_task_prerun(self, *args, **kwargs):
+        context_manager = None
         try:
             task_id = kwargs.get("task_id")
             if not task_id:
@@ -222,17 +223,16 @@ def _on_task_prerun(self, *args, **kwargs):
             )
             task_name = task_properties.get("celery_task_name")
 
-            context_manager = contexts.new_context(
-                fresh=True,  # to prevent context bleed across tasks
-                capture_exceptions=False,  # Celery catches task exceptions internally and
-                                           # delivers them via task_failure signal, so they
-                                           # never propagate through the context manager.
-                                           # We capture them in _on_task_failure.
-                client=self.client,
-            )
-            context_manager.__enter__()
-
             if request is not None:
+                context_manager = contexts.new_context(
+                    fresh=True,  # to prevent context bleed across tasks
+                    capture_exceptions=False,  # Celery catches task exceptions internally and
+                                               # delivers them via task_failure signal, so they
+                                               # never propagate through the context manager.
+                                               # We capture them in _on_task_failure.
+                    client=self.client,
+                )
+                context_manager.__enter__()
                 request._posthog_ctx = context_manager
                 request._posthog_start = time.monotonic()
 
@@ -246,6 +246,11 @@ def _on_task_prerun(self, *args, **kwargs):
                 self._capture_event("celery task started", properties=task_properties)
         except Exception:
             logger.exception("Failed to process Celery task_prerun")
+            if context_manager is not None:
+                try:
+                    context_manager.__exit__(None, None, None)
+                except Exception:
+                    pass
 
     def _on_task_success(self, *args, **kwargs):
         self._handle_task_end("success", **kwargs)

posthog/test/integrations/test_celery_integration.py

@@ -463,6 +463,26 @@ def test_extract_headers_supports_request_dict_shape(self):
 
         self.assertEqual(headers, {CONTEXT_DISTINCT_ID_HEADER: "user-1"})
 
+    def test_prerun_exits_context_on_failure_after_entry(self):
+        mock_client = Mock()
+        integration = PosthogCeleryIntegration(client=mock_client)
+
+        request = SimpleNamespace(
+            headers={},
+            delivery_info={},
+            hostname="worker-1",
+            retries=0,
+        )
+        task = SimpleNamespace(name="app.tasks.boom", request=request)
+
+        ctx_before = contexts._get_current_context()
+
+        with patch.object(integration, "_apply_propagated_identity", side_effect=RuntimeError("boom")):
+            integration._on_task_prerun(sender=task, task_id="task-leak")
+
+        ctx_after = contexts._get_current_context()
+        self.assertIs(ctx_after, ctx_before)
+
     def test_extract_propagated_tags_invalid_json_returns_empty_dict(self):
         integration = PosthogCeleryIntegration()
         request = SimpleNamespace(headers={CONTEXT_TAGS_HEADER: "{bad json"})

posthog/test/test_client.py

@@ -2735,8 +2735,17 @@ def test_registers_at_fork_hook(self, mock_register_at_fork):
 
         mock_register_at_fork.assert_called_once()
         after_in_child = mock_register_at_fork.call_args.kwargs["after_in_child"]
-        self.assertEqual(after_in_child.__self__, client)
-        self.assertEqual(after_in_child.__name__, "_reinit_after_fork")
+
+        with mock.patch.object(client, "_reinit_after_fork") as mock_reinit:
+            after_in_child()
+            mock_reinit.assert_called_once()
+
+    @mock.patch("posthog.client.os.register_at_fork")
+    def test_register_at_fork_noop_after_client_gc(self, mock_register_at_fork):
+        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail)
+        after_in_child = mock_register_at_fork.call_args.kwargs["after_in_child"]
+        del client
+        after_in_child()
 
     @parameterized.expand([(True, 1), (False, 0)])
     def test_reinit_after_fork_replaces_queue_and_consumers(self, send, expected_starts):