From f63d7eb147fe8a77cf1dbf854649d25dfb5e219f Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:24:07 +0700 Subject: [PATCH 01/27] fix: prevent client lifecycle deadlocks Serialize join and shutdown ownership without holding locks across user callbacks, defer reentrant callback lifecycle work, discard unflushed join queues safely, and rebuild fork-unsafe sync queues. --- .sampo/changesets/lifecycle-deadlocks.md | 5 + posthog/__init__.py | 2 +- posthog/client.py | 226 ++++++++++++--- posthog/test/test_ai_capture_lane.py | 6 +- posthog/test/test_client.py | 344 ++++++++++++++++++++++- posthog/test/test_client_fork.py | 20 +- 6 files changed, 554 insertions(+), 49 deletions(-) create mode 100644 .sampo/changesets/lifecycle-deadlocks.md diff --git a/.sampo/changesets/lifecycle-deadlocks.md b/.sampo/changesets/lifecycle-deadlocks.md new file mode 100644 index 000000000..2170c8e50 --- /dev/null +++ b/.sampo/changesets/lifecycle-deadlocks.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +fix: prevent client lifecycle deadlocks when error callbacks, concurrent `join()`/`shutdown()` calls, or forked sync-mode clients interact with queue and worker teardown. diff --git a/posthog/__init__.py b/posthog/__init__.py index 4e553162e..3eb8ff269 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -1124,7 +1124,7 @@ def flush(timeout_seconds: Optional[float] = 10) -> None: def join() -> None: """ - Block program until the client clears the queue. Used during program shutdown. You should use `shutdown()` directly in most cases. + Stop the client's background workers without flushing queued events. Pending events may be discarded; use `shutdown()` when delivery is required. Examples: ```python diff --git a/posthog/client.py b/posthog/client.py index 74a5d9ebf..4c8be3bef 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -113,7 +113,7 @@ from posthog.version import VERSION -from queue import Queue, Full +from queue import Empty, Full, Queue _configure_posthog_logging() @@ -410,15 +410,29 @@ def flush(self, timeout_seconds: Optional[float]) -> None: self.log.debug("successfully flushed about %s items.", size) def join(self) -> None: - """Pause this lane's consumers and wait for them to exit; a never-started lane is a no-op.""" + """Pause this lane's consumers, wait for them, and discard queued work.""" for consumer in self.consumers: consumer.pause() + for consumer in self.consumers: try: consumer.join() except RuntimeError: # consumer thread has not started pass + dropped = 0 + while True: + try: + self.queue.get_nowait() + except Empty: + break + self.queue.task_done() + dropped += 1 + if dropped: + self.log.warning( + "%s lane discarded %d queued events during join", self.name, dropped + ) + def reset_sync_send_state_after_fork(self) -> None: """Replace sync-send state inherited from threads that did not survive fork.""" self._active_sync_sends = 0 @@ -635,6 +649,17 @@ def __init__( self.debug = debug self.send = send self.sync_mode = sync_mode + self._lifecycle_lock = threading.Lock() + self._lifecycle_condition = threading.Condition(self._lifecycle_lock) + self._lifecycle_in_progress = False + self._lifecycle_owner: Optional[threading.Thread] = None + self._workers_joined = False + self._join_cleanup_complete = False + self._shutdown_requested = False + self._shutdown_complete = False + self._deferred_lifecycle_error: Optional[BaseException] = None + self._deferred_lifecycle_failure = threading.Event() + self._shutdown_complete_event = threading.Event() # Used for session replay URL generation - we don't want the server host here. self.raw_host = normalize_host(host) self.host = determine_server_host(host) @@ -1901,12 +1926,22 @@ def _reinit_after_fork(self): consumer pool are rebuilt (see `_Lane.rebuild_after_fork`). """ for lane in self._lanes: - if self.sync_mode: - lane.reset_sync_send_state_after_fork() - else: - lane.rebuild_after_fork() - - if self.enable_local_evaluation: + lane.rebuild_after_fork() + + shutdown_complete = self._shutdown_complete + self._lifecycle_lock = threading.Lock() + self._lifecycle_condition = threading.Condition(self._lifecycle_lock) + self._lifecycle_in_progress = False + self._lifecycle_owner = None + self._deferred_lifecycle_error = None + self._deferred_lifecycle_failure = threading.Event() + self._shutdown_complete_event = threading.Event() + if shutdown_complete: + self._shutdown_complete_event.set() + + if self._workers_joined: + self.poller = None + elif self.enable_local_evaluation: self.poller = Poller( interval=timedelta(seconds=self.poll_interval), execute=self._load_feature_flags, @@ -2133,6 +2168,8 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: posthog.flush() # Ensures the event is sent immediately ``` """ + if self._defer_from_callback(self.flush, "flush", timeout_seconds): + return try: if timeout_seconds is None: for lane in self._lanes: @@ -2148,24 +2185,148 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: self.log.exception("error flushing queue: %s", e) return + def _is_consumer_thread(self) -> bool: + current = threading.current_thread() + return any(current in lane.consumers for lane in self._lanes) + + def _is_lifecycle_callback_thread(self) -> bool: + current = threading.current_thread() + if self._is_consumer_thread() or current is self.poller: + return True + runner = self._flag_definition_cache_provider_async_runner + return runner is not None and current is runner._thread + + def _start_lifecycle_thread(self, target, name: str, *args) -> None: + def run() -> None: + for attempt in range(2): + try: + target(*args) + return + except BaseException as error: + self.log.exception( + "Deferred %s attempt %d failed", name, attempt + 1 + ) + if attempt == 1: + with self._lifecycle_lock: + self._deferred_lifecycle_error = error + self._deferred_lifecycle_failure.set() + + threading.Thread( + target=run, + name=f"posthog-{name}", + daemon=False, + ).start() + + def _defer_from_callback(self, target, name: str, *args) -> bool: + if not self._is_lifecycle_callback_thread(): + return False + self._start_lifecycle_thread(target, name, *args) + return True + + def _join_once(self) -> None: + if not self._workers_joined: + for lane in self._lanes: + lane.close() + for lane in self._lanes: + lane.wait_for_sync_sends() + for lane in self._lanes: + lane.join() + self._workers_joined = True + + if not self._join_cleanup_complete: + if self.poller: + self.poller.stop() + + # Shutdown the cache provider (release locks, cleanup) + self._shutdown_flag_definition_cache_provider() + self._unregister_duplicate_client() + self._join_cleanup_complete = True + + def _shutdown_once(self) -> None: + if not self._workers_joined: + # Close every lane before draining any of them so no producer can be + # admitted between a completed flush and consumer shutdown. + for lane in self._lanes: + lane.close() + for lane in self._lanes: + lane.wait_for_sync_sends() + self.flush(timeout_seconds=None) + + if self._metrics is not None: + try: + self._metrics.flush() + except Exception: + self.log.exception("Failed to flush metrics on shutdown") + self._metrics.reset() + self._join_once() + self.distinct_ids_feature_flags_reported.clear() + + if self.exception_capture: + self.exception_capture.close() + self._shutdown_complete = True + self._deferred_lifecycle_error = None + self._deferred_lifecycle_failure.clear() + self._shutdown_complete_event.set() + + def _run_lifecycle(self, require_shutdown: bool = False) -> None: + while True: + with self._lifecycle_condition: + if require_shutdown and self._shutdown_complete: + return + if not require_shutdown and ( + self._join_cleanup_complete or self._shutdown_complete + ): + return + if self._lifecycle_in_progress: + if self._is_lifecycle_callback_thread() or ( + threading.current_thread() is self._lifecycle_owner + ): + return + self._lifecycle_condition.wait() + continue + self._lifecycle_in_progress = True + self._lifecycle_owner = threading.current_thread() + + try: + while True: + with self._lifecycle_lock: + run_shutdown = self._shutdown_requested + + if run_shutdown: + self._shutdown_once() + else: + self._join_once() + + with self._lifecycle_condition: + if ( + self._shutdown_requested + and not self._shutdown_complete + and not run_shutdown + ): + continue + self._lifecycle_in_progress = False + self._lifecycle_owner = None + self._lifecycle_condition.notify_all() + return + except BaseException: + with self._lifecycle_condition: + self._lifecycle_in_progress = False + self._lifecycle_owner = None + self._lifecycle_condition.notify_all() + raise + def join(self) -> None: """ - End the consumer thread once the queue is empty. Do not use directly, call `shutdown()` instead. + End the consumer threads without flushing queued events. Do not use directly, call `shutdown()` instead. Examples: ```python posthog.join() ``` """ - for lane in self._lanes: - lane.join() - - if self.poller: - self.poller.stop() - - # Shutdown the cache provider (release locks, cleanup) - self._shutdown_flag_definition_cache_provider() - self._unregister_duplicate_client() + if self._defer_from_callback(self._run_lifecycle, "join"): + return + self._run_lifecycle() def shutdown(self) -> None: """ @@ -2176,25 +2337,16 @@ def shutdown(self) -> None: posthog.shutdown() ``` """ - # Close every lane before draining any of them so no producer can be - # admitted between a completed flush and consumer shutdown. - for lane in self._lanes: - lane.close() - for lane in self._lanes: - lane.wait_for_sync_sends() - - self.flush(timeout_seconds=None) - if self._metrics is not None: - try: - self._metrics.flush() - except Exception: - self.log.exception("Failed to flush metrics on shutdown") - self._metrics.reset() - self.join() - self.distinct_ids_feature_flags_reported.clear() - - if self.exception_capture: - self.exception_capture.close() + with self._lifecycle_lock: + if self._shutdown_complete: + return + self._shutdown_requested = True + if not self._lifecycle_in_progress: + self._deferred_lifecycle_error = None + self._deferred_lifecycle_failure.clear() + if self._defer_from_callback(self._run_lifecycle, "shutdown", True): + return + self._run_lifecycle(require_shutdown=True) def _resolve_flag_definition_cache_provider_result(self, result): if not inspect.isawaitable(result): diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 4cf708a24..61fa203b5 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -314,15 +314,15 @@ def test_fork_rebuild_restarts_analytics_and_resets_ai(self): ) client.join() - def test_fork_rebuild_noop_for_sync_mode(self): + def test_fork_rebuild_replaces_sync_mode_queues(self): client = Client(TEST_API_KEY, sync_mode=True) old_analytics_queue = client._analytics_lane.queue old_ai_queue = client._ai_lane.queue client._reinit_after_fork() - self.assertIs(client._analytics_lane.queue, old_analytics_queue) - self.assertIs(client._ai_lane.queue, old_ai_queue) + self.assertIsNot(client._analytics_lane.queue, old_analytics_queue) + self.assertIsNot(client._ai_lane.queue, old_ai_queue) class TestCaptureAiEventHelper(unittest.TestCase): diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 4d0919cde..3f056a44d 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2164,6 +2164,342 @@ def test_shutdown_flushes_without_timeout(self): mock_flush.assert_called_once_with(timeout_seconds=None) + def test_sync_send_failure_does_not_invoke_async_on_error(self): + on_error = mock.Mock() + client = Client( + FAKE_TEST_API_KEY, + sync_mode=True, + on_error=on_error, + ) + + with mock.patch( + "posthog.client.batch_post", side_effect=Exception("upload failed") + ): + result = client.capture("event", distinct_id="distinct_id") + + self.assertIsNone(result) + on_error.assert_not_called() + self.assertEqual(client._analytics_lane._active_sync_sends, 0) + + def test_on_error_can_request_shutdown_with_pending_work(self): + first_send_started = threading.Event() + release_first_send = threading.Event() + callback_returned = threading.Event() + sent_events = [] + client: Client + + def request(batch): + event = batch[0]["event"] + sent_events.append(event) + if event == "first": + first_send_started.set() + self.assertTrue(release_first_send.wait(2)) + raise Exception("upload failed") + + def on_error(error, batch): + client.shutdown() + client.join() + callback_returned.set() + + client = Client( + FAKE_TEST_API_KEY, + on_error=on_error, + flush_at=1, + flush_interval=0.01, + max_retries=0, + ) + exception_capture = mock.Mock() + exception_capture.close.side_effect = [Exception("cleanup failed"), None] + client.exception_capture = exception_capture + with mock.patch.object(client.consumers[0], "request", side_effect=request): + client.capture("first", distinct_id="distinct_id") + self.assertTrue(first_send_started.wait(1)) + client.capture("second", distinct_id="distinct_id") + release_first_send.set() + + self.assertTrue(callback_returned.wait(1)) + self.assertTrue(client._shutdown_complete_event.wait(3)) + + self.assertEqual(sent_events, ["first", "second"]) + self.assertEqual(client.queue.unfinished_tasks, 0) + self.assertEqual(exception_capture.close.call_count, 2) + self.assertTrue(all(not consumer.is_alive() for consumer in client.consumers)) + + def test_concurrent_join_waits_for_lifecycle_owner(self): + send_started = threading.Event() + release_send = threading.Event() + + def request(batch): + send_started.set() + self.assertTrue(release_send.wait(2)) + + client = Client(FAKE_TEST_API_KEY, flush_at=1) + with mock.patch.object(client.consumers[0], "request", side_effect=request): + client.capture("event", distinct_id="distinct_id") + self.assertTrue(send_started.wait(1)) + + first_join = threading.Thread(target=client.join) + second_join = threading.Thread(target=client.join) + first_join.start() + time.sleep(0.05) + second_join.start() + time.sleep(0.05) + self.assertTrue(second_join.is_alive()) + + release_send.set() + first_join.join(3) + second_join.join(3) + + self.assertFalse(first_join.is_alive()) + self.assertFalse(second_join.is_alive()) + self.assertTrue(client._join_cleanup_complete) + + def test_join_winning_shutdown_race_discards_pending_work_without_deadlock(self): + first_send_started = threading.Event() + release_first_send = threading.Event() + join_started = threading.Event() + sent_events = [] + + def request(batch): + sent_events.extend(event["event"] for event in batch) + first_send_started.set() + self.assertTrue(release_first_send.wait(2)) + + client = Client(FAKE_TEST_API_KEY, flush_at=1) + with mock.patch.object(client.consumers[0], "request", side_effect=request): + client.capture("first", distinct_id="distinct_id") + self.assertTrue(first_send_started.wait(1)) + client.capture("second", distinct_id="distinct_id") + + original_close = client._analytics_lane.close + + def observed_close(): + original_close() + join_started.set() + + with mock.patch.object( + client._analytics_lane, "close", side_effect=observed_close + ): + join_thread = threading.Thread(target=client.join) + shutdown_thread = threading.Thread(target=client.shutdown) + join_thread.start() + self.assertTrue(join_started.wait(1)) + shutdown_thread.start() + release_first_send.set() + join_thread.join(3) + shutdown_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertFalse(shutdown_thread.is_alive()) + self.assertEqual(sent_events, ["first"]) + self.assertEqual(client.queue.unfinished_tasks, 0) + self.assertTrue(client._shutdown_complete) + + def test_shutdown_winning_join_race_drains_pending_work(self): + first_send_started = threading.Event() + release_first_send = threading.Event() + shutdown_started = threading.Event() + sent_events = [] + + def request(batch): + sent_events.extend(event["event"] for event in batch) + if len(sent_events) == 1: + first_send_started.set() + self.assertTrue(release_first_send.wait(2)) + + client = Client(FAKE_TEST_API_KEY, flush_at=1, flush_interval=0.01) + with mock.patch.object(client.consumers[0], "request", side_effect=request): + client.capture("first", distinct_id="distinct_id") + self.assertTrue(first_send_started.wait(1)) + client.capture("second", distinct_id="distinct_id") + + original_close = client._analytics_lane.close + + def observed_close(): + original_close() + shutdown_started.set() + + with mock.patch.object( + client._analytics_lane, "close", side_effect=observed_close + ): + shutdown_thread = threading.Thread(target=client.shutdown) + join_thread = threading.Thread(target=client.join) + shutdown_thread.start() + self.assertTrue(shutdown_started.wait(1)) + join_thread.start() + release_first_send.set() + shutdown_thread.join(3) + join_thread.join(3) + + self.assertFalse(shutdown_thread.is_alive()) + self.assertFalse(join_thread.is_alive()) + self.assertEqual(sent_events, ["first", "second"]) + self.assertEqual(client.queue.unfinished_tasks, 0) + + def test_shutdown_after_join_runs_shutdown_only_cleanup(self): + client = Client(FAKE_TEST_API_KEY) + metrics = mock.Mock() + exception_capture = mock.Mock() + client._metrics = metrics + client.exception_capture = exception_capture + + client.join() + client.shutdown() + + metrics.flush.assert_called_once() + metrics.reset.assert_called_once() + exception_capture.close.assert_called_once() + self.assertTrue(client._shutdown_complete) + + def test_cache_provider_shutdown_can_reenter_join(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=client.join, + ): + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(2) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(client._workers_joined) + + def test_cache_provider_shutdown_can_reenter_join_from_another_thread(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + runner = mock.Mock() + client._flag_definition_cache_provider_async_runner = runner + + def reenter_join(): + reentrant_thread = threading.Thread(target=client.join) + runner._thread = reentrant_thread + reentrant_thread.start() + reentrant_thread.join(1) + self.assertFalse(reentrant_thread.is_alive()) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=reenter_join, + ): + client.join() + + self.assertTrue(client._workers_joined) + self.assertTrue(client._join_cleanup_complete) + + def test_poller_shutdown_request_is_completed_by_join_owner(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + class ReentrantPoller(threading.Thread): + def __init__(self): + super().__init__(daemon=True) + self.stop_requested = threading.Event() + + def run(self): + self.stop_requested.wait(2) + client.shutdown() + + def stop(self): + self.stop_requested.set() + self.join(1) + self.assert_stopped() + + def assert_stopped(self): + if self.is_alive(): + raise AssertionError("poller did not stop") + + poller = ReentrantPoller() + client.poller = poller # type: ignore[assignment] + poller.start() + + client.join() + + self.assertTrue(client._shutdown_complete) + self.assertTrue(client._shutdown_complete_event.is_set()) + + def test_join_retries_auxiliary_cleanup_after_failure(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=[Exception("cleanup failed"), None], + ) as cleanup: + with self.assertRaisesRegex(Exception, "cleanup failed"): + client.join() + self.assertTrue(client._workers_joined) + self.assertFalse(client._join_cleanup_complete) + + client.join() + + self.assertTrue(client._workers_joined) + self.assertTrue(client._join_cleanup_complete) + self.assertEqual(cleanup.call_count, 2) + + def test_pending_shutdown_is_retried_when_join_cleanup_fails(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + cleanup_started = threading.Event() + release_cleanup = threading.Event() + cleanup_calls = 0 + join_errors = [] + + def cleanup(): + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_calls == 1: + cleanup_started.set() + self.assertTrue(release_cleanup.wait(2)) + raise Exception("cleanup failed") + + def run_join(): + try: + client.join() + except Exception as error: + join_errors.append(error) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=cleanup, + ): + join_thread = threading.Thread(target=run_join) + join_thread.start() + self.assertTrue(cleanup_started.wait(1)) + + shutdown_thread = threading.Thread(target=client.shutdown) + shutdown_thread.start() + time.sleep(0.05) + self.assertTrue(shutdown_thread.is_alive()) + + release_cleanup.set() + join_thread.join(2) + shutdown_thread.join(3) + self.assertTrue(client._shutdown_complete_event.is_set()) + + self.assertFalse(join_thread.is_alive()) + self.assertFalse(shutdown_thread.is_alive()) + self.assertEqual(str(join_errors[0]), "cleanup failed") + self.assertEqual(cleanup_calls, 2) + self.assertTrue(client._shutdown_complete) + + def test_shutdown_retries_cleanup_before_publishing_completion(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + exception_capture = mock.Mock() + exception_capture.close.side_effect = [Exception("cleanup failed"), None] + client.exception_capture = exception_capture + + with self.assertRaisesRegex(Exception, "cleanup failed"): + client.shutdown() + self.assertFalse(client._shutdown_complete) + self.assertFalse(client._shutdown_complete_event.is_set()) + + client.shutdown() + + self.assertTrue(client._shutdown_complete) + self.assertTrue(client._shutdown_complete_event.is_set()) + self.assertEqual(exception_capture.close.call_count, 2) + def test_shutdown_waits_for_racing_enqueue_before_draining(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) put_started = threading.Event() @@ -2259,12 +2595,8 @@ def test_synchronous(self): mock_post.assert_called_once() def test_overflow(self): - client = Client(FAKE_TEST_API_KEY, max_queue_size=1) - # Ensure consumer thread is no longer uploading - client.join() - - for i in range(10): - client.capture("test event", distinct_id="distinct_id") + client = Client(FAKE_TEST_API_KEY, max_queue_size=1, thread=0) + client.capture("test event", distinct_id="distinct_id") with self.assertLogs("posthog", level="WARNING") as logs: msg_uuid = client.capture("test event", distinct_id="distinct_id") diff --git a/posthog/test/test_client_fork.py b/posthog/test/test_client_fork.py index ac8f7463d..329f5ccd6 100644 --- a/posthog/test/test_client_fork.py +++ b/posthog/test/test_client_fork.py @@ -164,20 +164,36 @@ def test_reinit_after_fork_replaces_queue_and_consumers( self.assertIs(client.consumers[0].queue, client.queue) self.assertEqual(mock_start.call_count, expected_starts) - def test_reinit_after_fork_resets_sync_send_state_for_sync_mode(self): + def test_reinit_after_fork_replaces_sync_mode_queue_and_locks(self): client = Client(FAKE_TEST_API_KEY, sync_mode=True) lane = client._analytics_lane old_queue = lane.queue old_lock = lane._start_lock old_condition = lane._sync_sends_done + old_lifecycle_lock = client._lifecycle_lock lane._active_sync_sends = 1 client._reinit_after_fork() - self.assertIs(lane.queue, old_queue) + self.assertIsNot(lane.queue, old_queue) self.assertEqual(lane._active_sync_sends, 0) self.assertIsNot(lane._start_lock, old_lock) self.assertIsNot(lane._sync_sends_done, old_condition) + self.assertIsNot(client._lifecycle_lock, old_lifecycle_lock) + + def test_reinit_after_fork_preserves_terminal_client_state(self): + client = Client(FAKE_TEST_API_KEY) + client.join() + + with mock.patch("posthog.client.Poller") as mock_poller: + client._reinit_after_fork() + + self.assertTrue(client._workers_joined) + self.assertTrue(client._analytics_lane._closed) + self.assertEqual(client.consumers, []) + self.assertIsNone(client.poller) + mock_poller.assert_not_called() + self.assertIsNone(client.capture("after join", distinct_id="distinct_id")) @unittest.skipUnless( From 52364e78a40e0c69eb33d0d42f1deef2b3d01ce8 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:47:15 +0700 Subject: [PATCH 02/27] fix: coalesce callback flush helpers Ensure repeated flush requests from lifecycle callbacks share one non-daemon helper instead of accumulating blocked threads. --- posthog/client.py | 24 +++++++++++++++++++++++- posthog/test/test_client.py | 14 ++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/posthog/client.py b/posthog/client.py index 4c8be3bef..324644ec8 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -659,6 +659,8 @@ def __init__( self._shutdown_complete = False self._deferred_lifecycle_error: Optional[BaseException] = None self._deferred_lifecycle_failure = threading.Event() + self._deferred_flush_lock = threading.Lock() + self._deferred_flush_pending = False self._shutdown_complete_event = threading.Event() # Used for session replay URL generation - we don't want the server host here. self.raw_host = normalize_host(host) @@ -1935,6 +1937,8 @@ def _reinit_after_fork(self): self._lifecycle_owner = None self._deferred_lifecycle_error = None self._deferred_lifecycle_failure = threading.Event() + self._deferred_flush_lock = threading.Lock() + self._deferred_flush_pending = False self._shutdown_complete_event = threading.Event() if shutdown_complete: self._shutdown_complete_event.set() @@ -2168,7 +2172,7 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: posthog.flush() # Ensures the event is sent immediately ``` """ - if self._defer_from_callback(self.flush, "flush", timeout_seconds): + if self._defer_flush_from_callback(timeout_seconds): return try: if timeout_seconds is None: @@ -2223,6 +2227,24 @@ def _defer_from_callback(self, target, name: str, *args) -> bool: self._start_lifecycle_thread(target, name, *args) return True + def _defer_flush_from_callback(self, timeout_seconds: Optional[float]) -> bool: + if not self._is_lifecycle_callback_thread(): + return False + with self._deferred_flush_lock: + if self._deferred_flush_pending: + return True + self._deferred_flush_pending = True + + def run() -> None: + try: + self.flush(timeout_seconds) + finally: + with self._deferred_flush_lock: + self._deferred_flush_pending = False + + self._start_lifecycle_thread(run, "flush") + return True + def _join_once(self) -> None: if not self._workers_joined: for lane in self._lanes: diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 3f056a44d..03187a319 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2164,6 +2164,20 @@ def test_shutdown_flushes_without_timeout(self): mock_flush.assert_called_once_with(timeout_seconds=None) + def test_callback_flushes_are_coalesced(self): + client = Client(FAKE_TEST_API_KEY) + + with ( + mock.patch.object( + client, "_is_lifecycle_callback_thread", return_value=True + ), + mock.patch.object(client, "_start_lifecycle_thread") as start_thread, + ): + client.flush() + client.flush() + + start_thread.assert_called_once() + def test_sync_send_failure_does_not_invoke_async_on_error(self): on_error = mock.Mock() client = Client( From b009bf4e8a3d890e6ab10ab75c1299708c1c29ae Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:55:27 +0700 Subject: [PATCH 03/27] fix: preserve deferred flush requests Run one strongest follow-up flush for callback requests that arrive while the coalesced helper is active. --- posthog/client.py | 28 +++++++++++++++++++++++++--- posthog/test/test_client.py | 28 ++++++++++++++++++++++------ 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index 324644ec8..fea7e90e6 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -661,6 +661,8 @@ def __init__( self._deferred_lifecycle_failure = threading.Event() self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False + self._deferred_flush_followup = False + self._deferred_flush_followup_timeout: Optional[float] = None self._shutdown_complete_event = threading.Event() # Used for session replay URL generation - we don't want the server host here. self.raw_host = normalize_host(host) @@ -1939,6 +1941,8 @@ def _reinit_after_fork(self): self._deferred_lifecycle_failure = threading.Event() self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False + self._deferred_flush_followup = False + self._deferred_flush_followup_timeout = None self._shutdown_complete_event = threading.Event() if shutdown_complete: self._shutdown_complete_event.set() @@ -2232,15 +2236,33 @@ def _defer_flush_from_callback(self, timeout_seconds: Optional[float]) -> bool: return False with self._deferred_flush_lock: if self._deferred_flush_pending: + if not self._deferred_flush_followup: + self._deferred_flush_followup = True + self._deferred_flush_followup_timeout = timeout_seconds + elif ( + self._deferred_flush_followup_timeout is not None + and timeout_seconds is not None + ): + self._deferred_flush_followup_timeout = max( + self._deferred_flush_followup_timeout, timeout_seconds + ) + else: + self._deferred_flush_followup_timeout = None return True self._deferred_flush_pending = True def run() -> None: - try: - self.flush(timeout_seconds) - finally: + next_timeout = timeout_seconds + while True: + self.flush(next_timeout) with self._deferred_flush_lock: + if self._deferred_flush_followup: + next_timeout = self._deferred_flush_followup_timeout + self._deferred_flush_followup = False + self._deferred_flush_followup_timeout = None + continue self._deferred_flush_pending = False + return self._start_lifecycle_thread(run, "flush") return True diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 03187a319..37d0e68e4 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2164,19 +2164,35 @@ def test_shutdown_flushes_without_timeout(self): mock_flush.assert_called_once_with(timeout_seconds=None) - def test_callback_flushes_are_coalesced(self): + def test_callback_flushes_are_coalesced_with_strongest_followup(self): client = Client(FAKE_TEST_API_KEY) + first_flush_started = threading.Event() + release_first_flush = threading.Event() + followup_complete = threading.Event() + timeouts = [] + + def flush(timeout_seconds): + timeouts.append(timeout_seconds) + if len(timeouts) == 1: + first_flush_started.set() + self.assertTrue(release_first_flush.wait(2)) + else: + followup_complete.set() with ( mock.patch.object( client, "_is_lifecycle_callback_thread", return_value=True ), - mock.patch.object(client, "_start_lifecycle_thread") as start_thread, + mock.patch.object(client, "flush", side_effect=flush), ): - client.flush() - client.flush() - - start_thread.assert_called_once() + self.assertTrue(client._defer_flush_from_callback(0)) + self.assertTrue(first_flush_started.wait(1)) + self.assertTrue(client._defer_flush_from_callback(1)) + self.assertTrue(client._defer_flush_from_callback(None)) + release_first_flush.set() + self.assertTrue(followup_complete.wait(1)) + + self.assertEqual(timeouts, [0, None]) def test_sync_send_failure_does_not_invoke_async_on_error(self): on_error = mock.Mock() From 429c2a206282f42acc99bb3f9c22bf63f5ec9911 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:00:43 +0700 Subject: [PATCH 04/27] fix: coalesce callback lifecycle helpers Avoid creating one waiting lifecycle thread per callback while preserving deferred cleanup retries. --- posthog/client.py | 44 +++++++++++++++++++++++-------------- posthog/test/test_client.py | 14 ++++++++++++ 2 files changed, 42 insertions(+), 16 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index fea7e90e6..78f9c687b 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -659,6 +659,7 @@ def __init__( self._shutdown_complete = False self._deferred_lifecycle_error: Optional[BaseException] = None self._deferred_lifecycle_failure = threading.Event() + self._deferred_lifecycle_thread_pending = False self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -1939,6 +1940,7 @@ def _reinit_after_fork(self): self._lifecycle_owner = None self._deferred_lifecycle_error = None self._deferred_lifecycle_failure = threading.Event() + self._deferred_lifecycle_thread_pending = False self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -2205,22 +2207,9 @@ def _is_lifecycle_callback_thread(self) -> bool: return runner is not None and current is runner._thread def _start_lifecycle_thread(self, target, name: str, *args) -> None: - def run() -> None: - for attempt in range(2): - try: - target(*args) - return - except BaseException as error: - self.log.exception( - "Deferred %s attempt %d failed", name, attempt + 1 - ) - if attempt == 1: - with self._lifecycle_lock: - self._deferred_lifecycle_error = error - self._deferred_lifecycle_failure.set() - threading.Thread( - target=run, + target=target, + args=args, name=f"posthog-{name}", daemon=False, ).start() @@ -2228,7 +2217,30 @@ def run() -> None: def _defer_from_callback(self, target, name: str, *args) -> bool: if not self._is_lifecycle_callback_thread(): return False - self._start_lifecycle_thread(target, name, *args) + with self._lifecycle_lock: + if self._lifecycle_in_progress or self._deferred_lifecycle_thread_pending: + return True + self._deferred_lifecycle_thread_pending = True + + def run() -> None: + try: + for attempt in range(2): + try: + target(*args) + return + except BaseException as error: + self.log.exception( + "Deferred %s attempt %d failed", name, attempt + 1 + ) + if attempt == 1: + with self._lifecycle_lock: + self._deferred_lifecycle_error = error + self._deferred_lifecycle_failure.set() + finally: + with self._lifecycle_lock: + self._deferred_lifecycle_thread_pending = False + + self._start_lifecycle_thread(run, name) return True def _defer_flush_from_callback(self, timeout_seconds: Optional[float]) -> bool: diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 37d0e68e4..220f9a062 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2164,6 +2164,20 @@ def test_shutdown_flushes_without_timeout(self): mock_flush.assert_called_once_with(timeout_seconds=None) + def test_callback_lifecycle_requests_are_coalesced(self): + client = Client(FAKE_TEST_API_KEY) + + with ( + mock.patch.object( + client, "_is_lifecycle_callback_thread", return_value=True + ), + mock.patch.object(client, "_start_lifecycle_thread") as start_thread, + ): + client.shutdown() + client.shutdown() + + start_thread.assert_called_once() + def test_callback_flushes_are_coalesced_with_strongest_followup(self): client = Client(FAKE_TEST_API_KEY) first_flush_started = threading.Event() From b661138f952383618fa7cbb7f8b899a4ba39c0b2 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:06:49 +0700 Subject: [PATCH 05/27] fix: retain callback lifecycle handoff Keep one bounded deferred lifecycle waiter so callback shutdown requests survive an active owner's cleanup failure. --- posthog/client.py | 2 +- posthog/test/test_client.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index 78f9c687b..418a2937e 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -2218,7 +2218,7 @@ def _defer_from_callback(self, target, name: str, *args) -> bool: if not self._is_lifecycle_callback_thread(): return False with self._lifecycle_lock: - if self._lifecycle_in_progress or self._deferred_lifecycle_thread_pending: + if self._deferred_lifecycle_thread_pending: return True self._deferred_lifecycle_thread_pending = True diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 220f9a062..10abc1979 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2164,8 +2164,9 @@ def test_shutdown_flushes_without_timeout(self): mock_flush.assert_called_once_with(timeout_seconds=None) - def test_callback_lifecycle_requests_are_coalesced(self): + def test_callback_lifecycle_requests_are_coalesced_while_owner_runs(self): client = Client(FAKE_TEST_API_KEY) + client._lifecycle_in_progress = True with ( mock.patch.object( From 88ac9ad309ef68f4103b4c73feb1352002f7f845 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:36:52 +0700 Subject: [PATCH 06/27] fix: integrate drain-aware lifecycle teardown Combine the merged drain signaling behavior from #797 with serialized lifecycle shutdown before merging main. --- posthog/client.py | 99 ++++++++++++++++++++++++------------- posthog/test/test_client.py | 99 +++++++++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+), 33 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index 418a2937e..faed18c95 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -23,7 +23,7 @@ ) from posthog.capture_mode import CaptureMode, _resolve_capture_mode from posthog.capture_v1 import _send_v1_batch -from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE, Consumer +from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE, Consumer, _DrainSignal from posthog.contexts import ( _get_current_context, get_capture_exception_code_variables_context, @@ -308,6 +308,7 @@ def __init__( self._active_sync_sends = 0 self._start_lock = threading.Lock() self._sync_sends_done = threading.Condition(self._start_lock) + self._drain_signal = _DrainSignal(self.queue) if eager_start: self.start() @@ -331,6 +332,7 @@ def _start_locked(self) -> None: capture_mode=self.capture_mode, capture_compression=self.capture_compression, ) + consumer._set_drain_signal(self._drain_signal) self.consumers.append(consumer) if self.send: @@ -386,39 +388,55 @@ def wait_for_sync_sends(self) -> None: self._sync_sends_done.wait() def flush(self, timeout_seconds: Optional[float]) -> None: - """Block until this lane's queue drains, or until `timeout_seconds` elapse.""" + """Block until this lane's queue drains, or until `timeout_seconds` elapse. + + Signals the consumers first so a partial batch is delivered now instead + of waiting out `flush_at` / `flush_interval`. + """ queue = self.queue - size = queue.qsize() - if timeout_seconds is None: - queue.join() - else: - deadline = time.monotonic() + timeout_seconds - with queue.all_tasks_done: - while queue.unfinished_tasks: - remaining = deadline - time.monotonic() - if remaining <= 0: - self.log.warning( - "%s lane flush ran out of budget (%.1fs granted) with %s items pending.", - self.name, - timeout_seconds, - queue.unfinished_tasks, - ) - return - queue.all_tasks_done.wait(remaining) + # Keep the request active only while this flush is waiting. This avoids + # an empty flush changing how events captured after it are batched. + self._drain_signal.request() + try: + size = queue.qsize() + if timeout_seconds is None: + queue.join() + else: + deadline = time.monotonic() + timeout_seconds + with queue.all_tasks_done: + while queue.unfinished_tasks: + remaining = deadline - time.monotonic() + if remaining <= 0: + self.log.warning( + "%s lane flush ran out of budget (%.1fs granted) with %s items pending.", + self.name, + timeout_seconds, + queue.unfinished_tasks, + ) + return + queue.all_tasks_done.wait(remaining) - # Note that this message may not be precise, because of threading. - self.log.debug("successfully flushed about %s items.", size) + # Note that this message may not be precise, because of threading. + self.log.debug("successfully flushed about %s items.", size) + finally: + self._drain_signal.complete() def join(self) -> None: """Pause this lane's consumers, wait for them, and discard queued work.""" - for consumer in self.consumers: - consumer.pause() - for consumer in self.consumers: - try: - consumer.join() - except RuntimeError: - # consumer thread has not started - pass + # Teardown bypasses the batching wait too, so a consumer holding a + # partial batch delivers it instead of exiting `flush_interval` later. + self._drain_signal.request() + try: + for consumer in self.consumers: + consumer.pause() + for consumer in self.consumers: + try: + consumer.join() + except RuntimeError: + # consumer thread has not started + pass + finally: + self._drain_signal.complete() dropped = 0 while True: @@ -450,6 +468,7 @@ def rebuild_after_fork(self) -> None: """ self.queue = Queue(self._max_queue_size) self.reset_sync_send_state_after_fork() + self._drain_signal = _DrainSignal(self.queue) self.consumers = [] self._started = False if self._eager_start: @@ -1698,7 +1717,7 @@ def group_identify( @no_throw() def alias( self, - previous_id: str, + previous_id: ID_TYPES, distinct_id: Optional[str], timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, @@ -1708,8 +1727,11 @@ def alias( Create an alias between two distinct IDs. Args: - previous_id: The previous distinct ID. - distinct_id: The new distinct ID to alias to. + previous_id: The previous distinct ID. Required - the call is dropped + with a warning if it is missing or empty. + distinct_id: The new distinct ID to alias to. Falls back to the + context distinct ID; the call is dropped with a warning if + neither is available. timestamp: The timestamp of the event. uuid: A unique identifier for the event. If provided, it must be a valid UUID string or uuid.UUID instance; invalid values are @@ -1726,10 +1748,21 @@ def alias( Note: This method will not raise exceptions. Errors are logged. """ + previous_id = stringify_id(previous_id) + if not previous_id: + self.log.warning( + "alias() called without a previous_id, dropping the $create_alias event" + ) + return None + (distinct_id, personless) = get_identity_state(distinct_id) if personless: - return None # Personless alias() does nothing - should this throw? + # No alias target was passed and none is available from context. + self.log.warning( + "alias() called without a distinct_id, dropping the $create_alias event" + ) + return None msg: Dict[str, Any] = { "properties": { diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 10abc1979..ebccb141b 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -277,6 +277,18 @@ def test_client_flag_helpers_return_defaults_on_api_error(self, patch_flags): def test_empty_flush(self): self.client.flush() + def test_empty_flush_does_not_drain_a_later_event(self): + with mock.patch("posthog.consumer.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, flush_at=100, flush_interval=0.5) + + client.flush() + client.capture("after empty flush", distinct_id="distinct_id") + time.sleep(0.05) + + mock_post.assert_not_called() + client.flush() + mock_post.assert_called_once() + def test_flush_timeout_returns_when_queue_does_not_drain(self): client = Client(FAKE_TEST_API_KEY, send=False, thread=0) client.queue.put({"event": "stuck"}) @@ -292,6 +304,48 @@ def test_flush_timeout_returns_when_queue_does_not_drain(self): client.queue.get_nowait() client.queue.task_done() + def test_flush_does_not_wait_for_flush_interval(self): + # flush() must attempt delivery now rather than letting the consumer sit + # on a below-flush_at batch until flush_interval elapses. + with mock.patch("posthog.consumer.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, flush_interval=30) + client.capture("event", distinct_id="distinct_id") + + start = time.monotonic() + client.flush() + + self.assertLess(time.monotonic() - start, 5) + self.assertTrue(client.queue.empty()) + mock_post.assert_called_once() + + def test_flush_delivers_when_flush_interval_exceeds_the_flush_timeout(self): + # Waiting out flush_interval meant a flush_interval longer than the + # flush timeout delivered nothing at all. + with mock.patch("posthog.consumer.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, flush_interval=30) + client.capture("event", distinct_id="distinct_id") + + client.flush(timeout_seconds=5) + + mock_post.assert_called_once() + self.assertEqual(client.queue.unfinished_tasks, 0) + + def test_flush_keeps_batches_whole(self): + # Draining early must not turn a full queue into one request per event. + with mock.patch("posthog.consumer.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, flush_at=10, flush_interval=30) + for _ in range(30): + client.capture("event", distinct_id="distinct_id") + + client.flush() + + self.assertTrue(client.queue.empty()) + batch_sizes = [ + len(call.kwargs["batch"]) for call in mock_post.call_args_list + ] + self.assertEqual(sum(batch_sizes), 30) + self.assertLessEqual(len(batch_sizes), 5) + def test_flush_logs_and_returns_on_unexpected_error(self): client = Client(FAKE_TEST_API_KEY, send=False, thread=0) client.queue.put({"event": "stuck"}) @@ -1853,6 +1907,43 @@ def test_basic_alias(self): self.assertEqual(msg["properties"]["distinct_id"], "previousId") self.assertEqual(msg["properties"]["alias"], "distinct_id") + @parameterized.expand( + [ + ("none", None), + ("empty_string", ""), + ] + ) + def test_alias_without_previous_id_is_dropped(self, _name, previous_id): + with mock.patch("posthog.client.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) + with self.assertLogs("posthog", level="WARNING") as logs: + msg_uuid = client.alias(previous_id, "distinct_id") + + self.assertIsNone(msg_uuid) + mock_post.assert_not_called() + self.assertIn("previous_id", logs.output[0]) + + def test_alias_accepts_non_string_previous_id(self): + with mock.patch("posthog.client.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) + msg_uuid = client.alias(0, "distinct_id") + self.assertIsNotNone(msg_uuid) + + mock_post.assert_called_once() + msg = mock_post.call_args[1]["batch"][0] + self.assertEqual(msg["distinct_id"], "0") + self.assertEqual(msg["properties"]["distinct_id"], "0") + + def test_alias_without_distinct_id_is_dropped(self): + with mock.patch("posthog.client.batch_post") as mock_post: + client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) + with self.assertLogs("posthog", level="WARNING") as logs: + msg_uuid = client.alias("previousId", None) + + self.assertIsNone(msg_uuid) + mock_post.assert_not_called() + self.assertIn("distinct_id", logs.output[0]) + @parameterized.expand( [ # test_name, session_id, additional_properties, expected_properties @@ -2545,6 +2636,14 @@ def test_shutdown_retries_cleanup_before_publishing_completion(self): self.assertTrue(client._shutdown_complete_event.is_set()) self.assertEqual(exception_capture.close.call_count, 2) + def test_shutdown_does_not_wait_for_idle_consumers_flush_interval(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=5) + + start = time.monotonic() + client.shutdown() + + self.assertLess(time.monotonic() - start, 1) + def test_shutdown_waits_for_racing_enqueue_before_draining(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) put_started = threading.Event() From 9b31b743d7b0081af7243463cecb90eb8b56cab0 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 5 Aug 2026 17:51:32 +0200 Subject: [PATCH 07/27] fix: close remaining lifecycle handoff races --- posthog/_async_utils.py | 12 +++++++ posthog/client.py | 48 ++++++++++++++++---------- posthog/test/test_client.py | 64 +++++++++++++++++++++++++++++++---- posthog/test/test_consumer.py | 2 +- 4 files changed, 102 insertions(+), 24 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 8bc076a15..2a9300f0e 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -10,6 +10,7 @@ class _BackgroundEventLoopRunner: def __init__(self) -> None: self._loop: asyncio.AbstractEventLoop | None = None self._thread: threading.Thread | None = None + self._closing_threads: set[threading.Thread] = set() self._started = threading.Event() self._lock = threading.Lock() @@ -24,8 +25,13 @@ def close(self) -> None: thread = self._thread self._loop = None self._thread = None + if thread is not None: + self._closing_threads.add(thread) if loop is None or thread is None or loop.is_closed(): + with self._lock: + if thread is not None: + self._closing_threads.discard(thread) return if thread is threading.current_thread(): @@ -35,6 +41,10 @@ def close(self) -> None: loop.call_soon_threadsafe(loop.stop) thread.join() + def owns_thread(self, thread: threading.Thread) -> bool: + with self._lock: + return thread is self._thread or thread in self._closing_threads + @staticmethod async def _await_result(awaitable: Awaitable[Any]) -> Any: return await awaitable @@ -83,3 +93,5 @@ def _run_loop(self) -> None: loop.run_until_complete(loop.shutdown_default_executor()) asyncio.set_event_loop(None) loop.close() + with self._lock: + self._closing_threads.discard(threading.current_thread()) diff --git a/posthog/client.py b/posthog/client.py index faed18c95..1098e51fd 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -679,6 +679,7 @@ def __init__( self._deferred_lifecycle_error: Optional[BaseException] = None self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False + self._deferred_lifecycle_generation = 0 self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -1974,6 +1975,7 @@ def _reinit_after_fork(self): self._deferred_lifecycle_error = None self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False + self._deferred_lifecycle_generation = 0 self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -2237,7 +2239,7 @@ def _is_lifecycle_callback_thread(self) -> bool: if self._is_consumer_thread() or current is self.poller: return True runner = self._flag_definition_cache_provider_async_runner - return runner is not None and current is runner._thread + return runner is not None and runner.owns_thread(current) def _start_lifecycle_thread(self, target, name: str, *args) -> None: threading.Thread( @@ -2247,31 +2249,43 @@ def _start_lifecycle_thread(self, target, name: str, *args) -> None: daemon=False, ).start() - def _defer_from_callback(self, target, name: str, *args) -> bool: + def _defer_lifecycle_from_callback(self, name: str) -> bool: if not self._is_lifecycle_callback_thread(): return False with self._lifecycle_lock: + self._deferred_lifecycle_generation += 1 if self._deferred_lifecycle_thread_pending: return True self._deferred_lifecycle_thread_pending = True def run() -> None: - try: - for attempt in range(2): - try: - target(*args) + attempt = 0 + while True: + with self._lifecycle_lock: + generation = self._deferred_lifecycle_generation + require_shutdown = self._shutdown_requested + try: + self._run_lifecycle(require_shutdown=require_shutdown) + except BaseException as error: + attempt += 1 + self.log.exception("Deferred %s attempt %d failed", name, attempt) + with self._lifecycle_lock: + if generation != self._deferred_lifecycle_generation: + attempt = 0 + continue + if attempt < 2: + continue + self._deferred_lifecycle_error = error + self._deferred_lifecycle_failure.set() + self._deferred_lifecycle_thread_pending = False return - except BaseException as error: - self.log.exception( - "Deferred %s attempt %d failed", name, attempt + 1 - ) - if attempt == 1: - with self._lifecycle_lock: - self._deferred_lifecycle_error = error - self._deferred_lifecycle_failure.set() - finally: + with self._lifecycle_lock: + if generation != self._deferred_lifecycle_generation: + attempt = 0 + continue self._deferred_lifecycle_thread_pending = False + return self._start_lifecycle_thread(run, name) return True @@ -2413,7 +2427,7 @@ def join(self) -> None: posthog.join() ``` """ - if self._defer_from_callback(self._run_lifecycle, "join"): + if self._defer_lifecycle_from_callback("join"): return self._run_lifecycle() @@ -2433,7 +2447,7 @@ def shutdown(self) -> None: if not self._lifecycle_in_progress: self._deferred_lifecycle_error = None self._deferred_lifecycle_failure.clear() - if self._defer_from_callback(self._run_lifecycle, "shutdown", True): + if self._defer_lifecycle_from_callback("shutdown"): return self._run_lifecycle(require_shutdown=True) diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index ebccb141b..16a119dca 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2255,20 +2255,40 @@ def test_shutdown_flushes_without_timeout(self): mock_flush.assert_called_once_with(timeout_seconds=None) - def test_callback_lifecycle_requests_are_coalesced_while_owner_runs(self): + def test_callback_shutdown_escalates_pending_deferred_join(self): client = Client(FAKE_TEST_API_KEY) - client._lifecycle_in_progress = True + first_run_started = threading.Event() + release_first_run = threading.Event() + shutdown_run_complete = threading.Event() + require_shutdown_calls = [] + + def run_lifecycle(require_shutdown=False): + require_shutdown_calls.append(require_shutdown) + if len(require_shutdown_calls) == 1: + first_run_started.set() + self.assertTrue(release_first_run.wait(2)) + else: + shutdown_run_complete.set() with ( mock.patch.object( client, "_is_lifecycle_callback_thread", return_value=True ), - mock.patch.object(client, "_start_lifecycle_thread") as start_thread, + mock.patch.object(client, "_run_lifecycle", side_effect=run_lifecycle), + mock.patch.object( + client, + "_start_lifecycle_thread", + wraps=client._start_lifecycle_thread, + ) as start_thread, ): + client.join() + self.assertTrue(first_run_started.wait(1)) client.shutdown() - client.shutdown() + release_first_run.set() + self.assertTrue(shutdown_run_complete.wait(1)) start_thread.assert_called_once() + self.assertEqual(require_shutdown_calls, [False, True]) def test_callback_flushes_are_coalesced_with_strongest_followup(self): client = Client(FAKE_TEST_API_KEY) @@ -2487,6 +2507,31 @@ def test_shutdown_after_join_runs_shutdown_only_cleanup(self): exception_capture.close.assert_called_once() self.assertTrue(client._shutdown_complete) + def test_async_cache_provider_task_can_reenter_join_during_runner_close(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + finalizer_called = threading.Event() + + class AsyncProvider: + async def shutdown(self): + async def pending_task(): + try: + await asyncio.Event().wait() + finally: + client.join() + finalizer_called.set() + + asyncio.create_task(pending_task()) + await asyncio.sleep(0) + + client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(finalizer_called.is_set()) + self.assertTrue(client._join_cleanup_complete) + def test_cache_provider_shutdown_can_reenter_join(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) @@ -2504,8 +2549,15 @@ def test_cache_provider_shutdown_can_reenter_join(self): def test_cache_provider_shutdown_can_reenter_join_from_another_thread(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) - runner = mock.Mock() - client._flag_definition_cache_provider_async_runner = runner + + class Runner: + _thread = None + + def owns_thread(self, thread): + return thread is self._thread + + runner = Runner() + client._flag_definition_cache_provider_async_runner = runner # type: ignore[assignment] def reenter_join(): reentrant_thread = threading.Thread(target=client.join) diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index 296587bf3..ae78f4788 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -84,7 +84,7 @@ def test_message_only_error_logs_include_posthog_prefix(self) -> None: upload_logs = [ line for line in logs.getvalue().splitlines() if "error uploading" in line ] - self.assertEqual(upload_logs, ["[PostHog] error uploading: boom"]) + self.assertIn("[PostHog] error uploading: boom", upload_logs) def test_flush_interval(self) -> None: # Put _n_ items in the queue, pausing a little bit more than From bfdd442ce44144f020d18d2d90e6ee80ec95e1e1 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 5 Aug 2026 18:02:42 +0200 Subject: [PATCH 08/27] fix: preserve lifecycle delivery and fork ordering --- posthog/__init__.py | 2 +- posthog/client.py | 63 +++++++++++++++++--------------- posthog/test/test_client.py | 25 ++++++++++++- posthog/test/test_client_fork.py | 27 ++++++++++++++ 4 files changed, 84 insertions(+), 33 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index e3a717fc9..a106d512c 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -1124,7 +1124,7 @@ def flush(timeout_seconds: Optional[float] = 10) -> None: def join() -> None: """ - Stop the client's background workers without flushing queued events. Pending events may be discarded; use `shutdown()` when delivery is required. + Block until queued events are delivered and stop the client's background workers. Use `shutdown()` directly in most cases. Examples: ```python diff --git a/posthog/client.py b/posthog/client.py index 1098e51fd..47ac3c3bd 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -8,6 +8,7 @@ import time import warnings import weakref +from contextvars import ContextVar from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Mapping, Optional, Union from uuid import UUID, uuid4 @@ -113,7 +114,7 @@ from posthog.version import VERSION -from queue import Empty, Full, Queue +from queue import Full, Queue _configure_posthog_logging() @@ -422,7 +423,7 @@ def flush(self, timeout_seconds: Optional[float]) -> None: self._drain_signal.complete() def join(self) -> None: - """Pause this lane's consumers, wait for them, and discard queued work.""" + """Pause this lane's consumers and wait for them to exit.""" # Teardown bypasses the batching wait too, so a consumer holding a # partial batch delivers it instead of exiting `flush_interval` later. self._drain_signal.request() @@ -438,19 +439,6 @@ def join(self) -> None: finally: self._drain_signal.complete() - dropped = 0 - while True: - try: - self.queue.get_nowait() - except Empty: - break - self.queue.task_done() - dropped += 1 - if dropped: - self.log.warning( - "%s lane discarded %d queued events during join", self.name, dropped - ) - def reset_sync_send_state_after_fork(self) -> None: """Replace sync-send state inherited from threads that did not survive fork.""" self._active_sync_sends = 0 @@ -680,6 +668,9 @@ def __init__( self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_generation = 0 + self._lifecycle_callback_context: ContextVar[bool] = ContextVar( + "posthog_lifecycle_callback", default=False + ) self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -1976,6 +1967,9 @@ def _reinit_after_fork(self): self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_generation = 0 + self._lifecycle_callback_context = ContextVar( + "posthog_lifecycle_callback", default=False + ) self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -1984,17 +1978,6 @@ def _reinit_after_fork(self): if shutdown_complete: self._shutdown_complete_event.set() - if self._workers_joined: - self.poller = None - elif self.enable_local_evaluation: - self.poller = Poller( - interval=timedelta(seconds=self.poll_interval), - execute=self._load_feature_flags, - ) - self.poller.start() - else: - self.poller = None - # Async runner threads do not survive fork(); recreate lazily on next async cache call. self._flag_definition_cache_provider_async_runner = None self._flag_definition_cache_provider_async_runner_lock = threading.Lock() @@ -2017,6 +2000,18 @@ def _reinit_after_fork(self): reset_sessions() + # Start child threads only after replacing every lock they can touch. + if self._workers_joined: + self.poller = None + elif self.enable_local_evaluation: + self.poller = Poller( + interval=timedelta(seconds=self.poll_interval), + execute=self._load_feature_flags, + ) + self.poller.start() + else: + self.poller = None + def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): # type: (...) -> Optional[str] """Push a new `msg` onto a lane's queue (analytics when unspecified), return the event uuid or None.""" @@ -2235,6 +2230,8 @@ def _is_consumer_thread(self) -> bool: return any(current in lane.consumers for lane in self._lanes) def _is_lifecycle_callback_thread(self) -> bool: + if self._lifecycle_callback_context.get(): + return True current = threading.current_thread() if self._is_consumer_thread() or current is self.poller: return True @@ -2326,12 +2323,14 @@ def run() -> None: self._start_lifecycle_thread(run, "flush") return True - def _join_once(self) -> None: + def _join_once(self, flush_queues: bool = True) -> None: if not self._workers_joined: for lane in self._lanes: lane.close() for lane in self._lanes: lane.wait_for_sync_sends() + if flush_queues: + self.flush(timeout_seconds=None) for lane in self._lanes: lane.join() self._workers_joined = True @@ -2361,7 +2360,7 @@ def _shutdown_once(self) -> None: except Exception: self.log.exception("Failed to flush metrics on shutdown") self._metrics.reset() - self._join_once() + self._join_once(flush_queues=False) self.distinct_ids_feature_flags_reported.clear() if self.exception_capture: @@ -2420,7 +2419,7 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: def join(self) -> None: """ - End the consumer threads without flushing queued events. Do not use directly, call `shutdown()` instead. + Flush queued events and end the consumer threads. Do not use directly, call `shutdown()` instead. Examples: ```python @@ -2460,7 +2459,11 @@ def _resolve_flag_definition_cache_provider_result(self, result): self._flag_definition_cache_provider_async_runner = ( _BackgroundEventLoopRunner() ) - return self._flag_definition_cache_provider_async_runner.run(result) + token = self._lifecycle_callback_context.set(True) + try: + return self._flag_definition_cache_provider_async_runner.run(result) + finally: + self._lifecycle_callback_context.reset(token) def _shutdown_flag_definition_cache_provider(self): if not self._flag_definition_cache_provider: diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 16a119dca..3a66ef789 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2410,7 +2410,7 @@ def request(batch): self.assertFalse(second_join.is_alive()) self.assertTrue(client._join_cleanup_complete) - def test_join_winning_shutdown_race_discards_pending_work_without_deadlock(self): + def test_join_winning_shutdown_race_drains_pending_work_without_deadlock(self): first_send_started = threading.Event() release_first_send = threading.Event() join_started = threading.Event() @@ -2447,7 +2447,7 @@ def observed_close(): self.assertFalse(join_thread.is_alive()) self.assertFalse(shutdown_thread.is_alive()) - self.assertEqual(sent_events, ["first"]) + self.assertEqual(sent_events, ["first", "second"]) self.assertEqual(client.queue.unfinished_tasks, 0) self.assertTrue(client._shutdown_complete) @@ -2507,6 +2507,27 @@ def test_shutdown_after_join_runs_shutdown_only_cleanup(self): exception_capture.close.assert_called_once() self.assertTrue(client._shutdown_complete) + def test_async_cache_provider_executor_can_reenter_join(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + executor_called = threading.Event() + + class AsyncProvider: + async def shutdown(self): + def reenter_join(): + client.join() + executor_called.set() + + await asyncio.to_thread(reenter_join) + + client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(executor_called.is_set()) + self.assertTrue(client._join_cleanup_complete) + def test_async_cache_provider_task_can_reenter_join_during_runner_close(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) finalizer_called = threading.Event() diff --git a/posthog/test/test_client_fork.py b/posthog/test/test_client_fork.py index 329f5ccd6..01d900e93 100644 --- a/posthog/test/test_client_fork.py +++ b/posthog/test/test_client_fork.py @@ -181,6 +181,33 @@ def test_reinit_after_fork_replaces_sync_mode_queue_and_locks(self): self.assertIsNot(lane._sync_sends_done, old_condition) self.assertIsNot(client._lifecycle_lock, old_lifecycle_lock) + def test_reinit_after_fork_replaces_locks_before_starting_poller(self): + client = Client(FAKE_TEST_API_KEY) + client.enable_local_evaluation = True + old_runner_lock = client._flag_definition_cache_provider_async_runner_lock + old_publication_lock = client._flag_definition_publication_lock + old_cache_write_lock = client._flag_definition_cache_write_lock + old_metrics_lock = client._metrics_lock + + def assert_locks_replaced(): + self.assertIsNot( + client._flag_definition_cache_provider_async_runner_lock, + old_runner_lock, + ) + self.assertIsNot( + client._flag_definition_publication_lock, old_publication_lock + ) + self.assertIsNot( + client._flag_definition_cache_write_lock, old_cache_write_lock + ) + self.assertIsNot(client._metrics_lock, old_metrics_lock) + + with mock.patch("posthog.client.Poller") as mock_poller: + mock_poller.return_value.start.side_effect = assert_locks_replaced + client._reinit_after_fork() + + mock_poller.return_value.start.assert_called_once() + def test_reinit_after_fork_preserves_terminal_client_state(self): client = Client(FAKE_TEST_API_KEY) client.join() From 32391d87b884c6088c4073bca926822a47781c0c Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:12:17 +0700 Subject: [PATCH 09/27] fix: handle executor callbacks and zero-worker teardown Propagate lifecycle callback context through the runner default executor and avoid unbounded drains when no consumer can process queued work. --- posthog/_async_utils.py | 9 +++++++++ posthog/client.py | 28 ++++++++++++++++++++++++++-- posthog/test/test_client.py | 14 ++++++++++++-- posthog/test/test_consumer.py | 5 ++++- 4 files changed, 51 insertions(+), 5 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 2a9300f0e..841cded8c 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,9 +1,17 @@ import asyncio import threading from collections.abc import Awaitable +from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from typing import Any +class _ContextThreadPoolExecutor(ThreadPoolExecutor): + def submit(self, fn, /, *args, **kwargs): + context = copy_context() + return super().submit(context.run, fn, *args, **kwargs) + + class _BackgroundEventLoopRunner: """Run awaitables to completion on a reusable background event loop.""" @@ -74,6 +82,7 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: def _run_loop(self) -> None: loop = asyncio.new_event_loop() + loop.set_default_executor(_ContextThreadPoolExecutor()) asyncio.set_event_loop(loop) with self._lock: self._loop = loop diff --git a/posthog/client.py b/posthog/client.py index 47ac3c3bd..fb9e54c2b 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -114,7 +114,7 @@ from posthog.version import VERSION -from queue import Full, Queue +from queue import Empty, Full, Queue _configure_posthog_logging() @@ -422,6 +422,26 @@ def flush(self, timeout_seconds: Optional[float]) -> None: finally: self._drain_signal.complete() + def discard_undrainable_queued_work(self) -> None: + """Balance queued work when this lane has no running sender.""" + if any(consumer.is_alive() for consumer in self.consumers): + return + + dropped = 0 + while True: + try: + self.queue.get_nowait() + except Empty: + break + self.queue.task_done() + dropped += 1 + if dropped: + self.log.warning( + "%s lane discarded %d queued events because no consumer is running", + self.name, + dropped, + ) + def join(self) -> None: """Pause this lane's consumers and wait for them to exit.""" # Teardown bypasses the batching wait too, so a consumer holding a @@ -2330,7 +2350,11 @@ def _join_once(self, flush_queues: bool = True) -> None: for lane in self._lanes: lane.wait_for_sync_sends() if flush_queues: - self.flush(timeout_seconds=None) + for lane in self._lanes: + if any(consumer.is_alive() for consumer in lane.consumers): + lane.flush(timeout_seconds=None) + else: + lane.discard_undrainable_queued_work() for lane in self._lanes: lane.join() self._workers_joined = True diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 3a66ef789..d91fe4c5d 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2507,7 +2507,7 @@ def test_shutdown_after_join_runs_shutdown_only_cleanup(self): exception_capture.close.assert_called_once() self.assertTrue(client._shutdown_complete) - def test_async_cache_provider_executor_can_reenter_join(self): + def test_async_cache_provider_default_executor_can_reenter_join(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) executor_called = threading.Event() @@ -2517,7 +2517,7 @@ def reenter_join(): client.join() executor_called.set() - await asyncio.to_thread(reenter_join) + await asyncio.get_running_loop().run_in_executor(None, reenter_join) client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] join_thread = threading.Thread(target=client.join) @@ -2821,6 +2821,16 @@ def test_overflow(self): self.assertIsNone(msg_uuid) self.assertIn("dropping event", logs.output[0]) + def test_join_discards_queued_work_when_no_consumer_can_drain_it(self): + client = Client(FAKE_TEST_API_KEY, thread=0) + client.capture("test event", distinct_id="distinct_id") + + start = time.monotonic() + client.join() + + self.assertLess(time.monotonic() - start, 1) + self.assertEqual(client.queue.unfinished_tasks, 0) + def test_unicode(self): Client("unicode_key") diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index ae78f4788..f13733863 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -84,7 +84,10 @@ def test_message_only_error_logs_include_posthog_prefix(self) -> None: upload_logs = [ line for line in logs.getvalue().splitlines() if "error uploading" in line ] - self.assertIn("[PostHog] error uploading: boom", upload_logs) + expected_log = "[PostHog] error uploading: boom" + self.assertEqual( + [line for line in upload_logs if line == expected_log], [expected_log] + ) def test_flush_interval(self) -> None: # Put _n_ items in the queue, pausing a little bit more than From b1742862a329b3c0dd10d6da576270378a6ad990 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:21:58 +0700 Subject: [PATCH 10/27] fix: propagate custom executor lifecycle context Wrap all runner executor submissions with callback context and make shutdown discard only queues that have no live consumer. --- posthog/_async_utils.py | 17 ++++++++-------- posthog/client.py | 15 ++++++++------ posthog/test/test_client.py | 40 +++++++++++++++++++++++++++++++++++-- 3 files changed, 56 insertions(+), 16 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 841cded8c..d81e30649 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,17 +1,10 @@ import asyncio import threading from collections.abc import Awaitable -from concurrent.futures import ThreadPoolExecutor from contextvars import copy_context from typing import Any -class _ContextThreadPoolExecutor(ThreadPoolExecutor): - def submit(self, fn, /, *args, **kwargs): - context = copy_context() - return super().submit(context.run, fn, *args, **kwargs) - - class _BackgroundEventLoopRunner: """Run awaitables to completion on a reusable background event loop.""" @@ -82,7 +75,15 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: def _run_loop(self) -> None: loop = asyncio.new_event_loop() - loop.set_default_executor(_ContextThreadPoolExecutor()) + original_run_in_executor = loop.run_in_executor + + def run_in_executor(executor, func, *args): + context = copy_context() + return original_run_in_executor(executor, context.run, func, *args) + + # Providers can use either the default or an explicit executor. Preserve + # callback context across both forms so lifecycle reentry remains safe. + setattr(loop, "run_in_executor", run_in_executor) asyncio.set_event_loop(loop) with self._lock: self._loop = loop diff --git a/posthog/client.py b/posthog/client.py index fb9e54c2b..3c369a5c2 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -2343,6 +2343,13 @@ def run() -> None: self._start_lifecycle_thread(run, "flush") return True + def _flush_or_discard_queues(self) -> None: + for lane in self._lanes: + if any(consumer.is_alive() for consumer in lane.consumers): + lane.flush(timeout_seconds=None) + else: + lane.discard_undrainable_queued_work() + def _join_once(self, flush_queues: bool = True) -> None: if not self._workers_joined: for lane in self._lanes: @@ -2350,11 +2357,7 @@ def _join_once(self, flush_queues: bool = True) -> None: for lane in self._lanes: lane.wait_for_sync_sends() if flush_queues: - for lane in self._lanes: - if any(consumer.is_alive() for consumer in lane.consumers): - lane.flush(timeout_seconds=None) - else: - lane.discard_undrainable_queued_work() + self._flush_or_discard_queues() for lane in self._lanes: lane.join() self._workers_joined = True @@ -2376,7 +2379,7 @@ def _shutdown_once(self) -> None: lane.close() for lane in self._lanes: lane.wait_for_sync_sends() - self.flush(timeout_seconds=None) + self._flush_or_discard_queues() if self._metrics is not None: try: diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index d91fe4c5d..9e37a464b 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -4,6 +4,7 @@ import time import unittest import warnings +from concurrent.futures import ThreadPoolExecutor from datetime import datetime from unittest import mock from uuid import UUID, uuid4 @@ -2248,9 +2249,9 @@ def test_shutdown_clears_feature_flag_called_dedupe_cache(self): self.assertEqual(len(client.distinct_ids_feature_flags_reported), 0) def test_shutdown_flushes_without_timeout(self): - client = Client(FAKE_TEST_API_KEY, send=False, thread=0) + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) - with mock.patch.object(client, "flush") as mock_flush: + with mock.patch.object(client._analytics_lane, "flush") as mock_flush: client.shutdown() mock_flush.assert_called_once_with(timeout_seconds=None) @@ -2528,6 +2529,30 @@ def reenter_join(): self.assertTrue(executor_called.is_set()) self.assertTrue(client._join_cleanup_complete) + def test_async_cache_provider_custom_executor_can_reenter_join(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + executor_called = threading.Event() + + class AsyncProvider: + async def shutdown(self): + def reenter_join(): + client.join() + executor_called.set() + + with ThreadPoolExecutor(max_workers=1) as executor: + await asyncio.get_running_loop().run_in_executor( + executor, reenter_join + ) + + client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(executor_called.is_set()) + self.assertTrue(client._join_cleanup_complete) + def test_async_cache_provider_task_can_reenter_join_during_runner_close(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) finalizer_called = threading.Event() @@ -2831,6 +2856,17 @@ def test_join_discards_queued_work_when_no_consumer_can_drain_it(self): self.assertLess(time.monotonic() - start, 1) self.assertEqual(client.queue.unfinished_tasks, 0) + def test_shutdown_discards_queued_work_when_no_consumer_can_drain_it(self): + client = Client(FAKE_TEST_API_KEY, thread=0) + client.capture("test event", distinct_id="distinct_id") + + start = time.monotonic() + client.shutdown() + + self.assertLess(time.monotonic() - start, 1) + self.assertEqual(client.queue.unfinished_tasks, 0) + self.assertTrue(client._shutdown_complete) + def test_unicode(self): Client("unicode_key") From 222ce5fd7cd299cd5f9fc0ca2938d3088b4a343e Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:30:30 +0700 Subject: [PATCH 11/27] fix: preserve process executor support Propagate lifecycle context through thread executors while leaving process-pool submissions picklable. --- posthog/_async_utils.py | 3 +++ posthog/test/test_client.py | 21 ++++++++++++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index d81e30649..f1a38db54 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,6 +1,7 @@ import asyncio import threading from collections.abc import Awaitable +from concurrent.futures import ProcessPoolExecutor from contextvars import copy_context from typing import Any @@ -78,6 +79,8 @@ def _run_loop(self) -> None: original_run_in_executor = loop.run_in_executor def run_in_executor(executor, func, *args): + if isinstance(executor, ProcessPoolExecutor): + return original_run_in_executor(executor, func, *args) context = copy_context() return original_run_in_executor(executor, context.run, func, *args) diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 9e37a464b..fac363f3f 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -4,7 +4,7 @@ import time import unittest import warnings -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from datetime import datetime from unittest import mock from uuid import UUID, uuid4 @@ -2553,6 +2553,25 @@ def reenter_join(): self.assertTrue(executor_called.is_set()) self.assertTrue(client._join_cleanup_complete) + def test_async_cache_provider_process_executor_remains_supported(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + class AsyncProvider: + result = None + + async def shutdown(self): + with ProcessPoolExecutor(max_workers=1) as executor: + self.result = await asyncio.get_running_loop().run_in_executor( + executor, abs, -6 + ) + + provider = AsyncProvider() + client._flag_definition_cache_provider = provider # type: ignore[assignment] + client.join() + + self.assertEqual(provider.result, 6) + self.assertTrue(client._join_cleanup_complete) + def test_async_cache_provider_task_can_reenter_join_during_runner_close(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) finalizer_called = threading.Event() From b80e7a0e8fa450889e34491f9d8d19ff6e259f8d Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:40:19 +0700 Subject: [PATCH 12/27] fix: make async callback detection executor agnostic Track active provider calls without mutating event loops or wrapping executors, and make unbounded drains recover if all consumers stop. --- posthog/_async_utils.py | 13 ------------- posthog/client.py | 37 +++++++++++++++++++++---------------- posthog/test/test_client.py | 25 +++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 29 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index f1a38db54..2a9300f0e 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,8 +1,6 @@ import asyncio import threading from collections.abc import Awaitable -from concurrent.futures import ProcessPoolExecutor -from contextvars import copy_context from typing import Any @@ -76,17 +74,6 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: def _run_loop(self) -> None: loop = asyncio.new_event_loop() - original_run_in_executor = loop.run_in_executor - - def run_in_executor(executor, func, *args): - if isinstance(executor, ProcessPoolExecutor): - return original_run_in_executor(executor, func, *args) - context = copy_context() - return original_run_in_executor(executor, context.run, func, *args) - - # Providers can use either the default or an explicit executor. Preserve - # callback context across both forms so lifecycle reentry remains safe. - setattr(loop, "run_in_executor", run_in_executor) asyncio.set_event_loop(loop) with self._lock: self._loop = loop diff --git a/posthog/client.py b/posthog/client.py index 3c369a5c2..5f0b1ea6c 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -8,7 +8,6 @@ import time import warnings import weakref -from contextvars import ContextVar from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Mapping, Optional, Union from uuid import UUID, uuid4 @@ -400,12 +399,21 @@ def flush(self, timeout_seconds: Optional[float]) -> None: self._drain_signal.request() try: size = queue.qsize() - if timeout_seconds is None: - queue.join() - else: - deadline = time.monotonic() + timeout_seconds + deadline = ( + None if timeout_seconds is None else time.monotonic() + timeout_seconds + ) + while queue.unfinished_tasks: + if deadline is None and not any( + consumer.is_alive() for consumer in self.consumers + ): + self.discard_undrainable_queued_work() + break with queue.all_tasks_done: - while queue.unfinished_tasks: + if not queue.unfinished_tasks: + break + if deadline is None: + wait_seconds = 0.05 + else: remaining = deadline - time.monotonic() if remaining <= 0: self.log.warning( @@ -415,7 +423,8 @@ def flush(self, timeout_seconds: Optional[float]) -> None: queue.unfinished_tasks, ) return - queue.all_tasks_done.wait(remaining) + wait_seconds = min(0.05, remaining) + queue.all_tasks_done.wait(wait_seconds) # Note that this message may not be precise, because of threading. self.log.debug("successfully flushed about %s items.", size) @@ -688,9 +697,7 @@ def __init__( self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_generation = 0 - self._lifecycle_callback_context: ContextVar[bool] = ContextVar( - "posthog_lifecycle_callback", default=False - ) + self._async_provider_call_active = threading.Event() self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -1987,9 +1994,7 @@ def _reinit_after_fork(self): self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_generation = 0 - self._lifecycle_callback_context = ContextVar( - "posthog_lifecycle_callback", default=False - ) + self._async_provider_call_active = threading.Event() self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -2250,7 +2255,7 @@ def _is_consumer_thread(self) -> bool: return any(current in lane.consumers for lane in self._lanes) def _is_lifecycle_callback_thread(self) -> bool: - if self._lifecycle_callback_context.get(): + if self._async_provider_call_active.is_set(): return True current = threading.current_thread() if self._is_consumer_thread() or current is self.poller: @@ -2486,11 +2491,11 @@ def _resolve_flag_definition_cache_provider_result(self, result): self._flag_definition_cache_provider_async_runner = ( _BackgroundEventLoopRunner() ) - token = self._lifecycle_callback_context.set(True) + self._async_provider_call_active.set() try: return self._flag_definition_cache_provider_async_runner.run(result) finally: - self._lifecycle_callback_context.reset(token) + self._async_provider_call_active.clear() def _shutdown_flag_definition_cache_provider(self): if not self._flag_definition_cache_provider: diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index fac363f3f..21b72124a 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2875,6 +2875,31 @@ def test_join_discards_queued_work_when_no_consumer_can_drain_it(self): self.assertLess(time.monotonic() - start, 1) self.assertEqual(client.queue.unfinished_tasks, 0) + def test_join_discards_remaining_work_if_consumer_stops_during_drain(self): + send_started = threading.Event() + release_send = threading.Event() + + def request(batch): + send_started.set() + self.assertTrue(release_send.wait(2)) + + client = Client(FAKE_TEST_API_KEY, flush_at=1) + consumer = client.consumers[0] + with mock.patch.object(consumer, "request", side_effect=request): + client.capture("first", distinct_id="distinct_id") + self.assertTrue(send_started.wait(1)) + client.capture("second", distinct_id="distinct_id") + + join_thread = threading.Thread(target=client.join) + join_thread.start() + time.sleep(0.05) + consumer.pause() + release_send.set() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertEqual(client.queue.unfinished_tasks, 0) + def test_shutdown_discards_queued_work_when_no_consumer_can_drain_it(self): client = Client(FAKE_TEST_API_KEY, thread=0) client.capture("test event", distinct_id="distinct_id") From dbe018938e7be5aee198fd3cee85e61fa4047b12 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:52:49 +0700 Subject: [PATCH 13/27] fix: scope async lifecycle callback context Propagate callback context through known thread executors without affecting unrelated callers or serialization executors, and balance dequeued work if batching is interrupted. --- posthog/_async_utils.py | 23 +++++++++ posthog/client.py | 15 ++++-- posthog/consumer.py | 93 ++++++++++++++++++++--------------- posthog/test/test_client.py | 4 +- posthog/test/test_consumer.py | 12 +++++ 5 files changed, 99 insertions(+), 48 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 2a9300f0e..bc8244c43 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,9 +1,17 @@ import asyncio import threading from collections.abc import Awaitable +from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from typing import Any +class _ContextThreadPoolExecutor(ThreadPoolExecutor): + def submit(self, fn, /, *args, **kwargs): + context = copy_context() + return super().submit(context.run, fn, *args, **kwargs) + + class _BackgroundEventLoopRunner: """Run awaitables to completion on a reusable background event loop.""" @@ -74,6 +82,21 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: def _run_loop(self) -> None: loop = asyncio.new_event_loop() + loop.set_default_executor(_ContextThreadPoolExecutor()) + original_run_in_executor = loop.run_in_executor + + def run_in_executor(executor, func, *args): + if not isinstance(executor, ThreadPoolExecutor): + return original_run_in_executor(executor, func, *args) + context = copy_context() + return original_run_in_executor(executor, context.run, func, *args) + + try: + setattr(loop, "run_in_executor", run_in_executor) + except (AttributeError, TypeError): + # Some policy-provided loops expose read-only methods. Their default + # executor still propagates context through _ContextThreadPoolExecutor. + pass asyncio.set_event_loop(loop) with self._lock: self._loop = loop diff --git a/posthog/client.py b/posthog/client.py index 5f0b1ea6c..0a2049152 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -8,6 +8,7 @@ import time import warnings import weakref +from contextvars import ContextVar from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Mapping, Optional, Union from uuid import UUID, uuid4 @@ -697,7 +698,9 @@ def __init__( self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_generation = 0 - self._async_provider_call_active = threading.Event() + self._lifecycle_callback_context: ContextVar[bool] = ContextVar( + "posthog_lifecycle_callback", default=False + ) self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -1994,7 +1997,9 @@ def _reinit_after_fork(self): self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_generation = 0 - self._async_provider_call_active = threading.Event() + self._lifecycle_callback_context = ContextVar( + "posthog_lifecycle_callback", default=False + ) self._deferred_flush_lock = threading.Lock() self._deferred_flush_pending = False self._deferred_flush_followup = False @@ -2255,7 +2260,7 @@ def _is_consumer_thread(self) -> bool: return any(current in lane.consumers for lane in self._lanes) def _is_lifecycle_callback_thread(self) -> bool: - if self._async_provider_call_active.is_set(): + if self._lifecycle_callback_context.get(): return True current = threading.current_thread() if self._is_consumer_thread() or current is self.poller: @@ -2491,11 +2496,11 @@ def _resolve_flag_definition_cache_provider_result(self, result): self._flag_definition_cache_provider_async_runner = ( _BackgroundEventLoopRunner() ) - self._async_provider_call_active.set() + token = self._lifecycle_callback_context.set(True) try: return self._flag_definition_cache_provider_async_runner.run(result) finally: - self._async_provider_call_active.clear() + self._lifecycle_callback_context.reset(token) def _shutdown_flag_definition_cache_provider(self): if not self._flag_definition_cache_provider: diff --git a/posthog/consumer.py b/posthog/consumer.py index 046c02e9b..f1afb0208 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -182,51 +182,62 @@ def next(self): start_time = time.monotonic() total_size = 0 + pending_items = 0 - while len(items) < self.flush_at: - # While draining we take only what is already queued, never waiting - # for `flush_interval` to elapse or for `flush_at` to be reached. - draining = self._draining() - remaining = self.flush_interval - (time.monotonic() - start_time) - if not draining and remaining <= 0: - break + try: + while len(items) < self.flush_at: + # While draining we take only what is already queued, never waiting + # for `flush_interval` to elapse or for `flush_at` to be reached. + draining = self._draining() + remaining = self.flush_interval - (time.monotonic() - start_time) + if not draining and remaining <= 0: + break - try: - if draining: - item = queue.get(block=False) - elif self._drain_signal is not None: - item = self._drain_signal.get(timeout=remaining) - else: - item = queue.get(block=True, timeout=remaining) try: - item_size = len(json.dumps(item, cls=DatetimeSerializer).encode()) - except Exception: - # Callback-modified events can still contain invalid mapping - # keys or circular references. Never log the payload here. - self.log.error( - "Unable to serialize queued event for sizing, dropping." - ) - queue.task_done() - continue - if item_size > self.max_msg_size: - # Log only name and size: AI events may carry unredacted - # multimodal payloads that must not leak into logs. - self.log.error( - "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", - item.get("event") if isinstance(item, dict) else type(item), - item_size, - self.max_msg_size // 1024, - self.endpoint, - ) - queue.task_done() - continue - items.append(item) - total_size += item_size - if total_size >= BATCH_SIZE_LIMIT: - self.log.debug("hit batch size limit (size: %d)", total_size) + if draining: + item = queue.get(block=False) + elif self._drain_signal is not None: + item = self._drain_signal.get(timeout=remaining) + else: + item = queue.get(block=True, timeout=remaining) + pending_items += 1 + try: + item_size = len( + json.dumps(item, cls=DatetimeSerializer).encode() + ) + except Exception: + # Callback-modified events can still contain invalid mapping + # keys or circular references. Never log the payload here. + self.log.error( + "Unable to serialize queued event for sizing, dropping." + ) + queue.task_done() + pending_items -= 1 + continue + if item_size > self.max_msg_size: + # Log only name and size: AI events may carry unredacted + # multimodal payloads that must not leak into logs. + self.log.error( + "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", + item.get("event") if isinstance(item, dict) else type(item), + item_size, + self.max_msg_size // 1024, + self.endpoint, + ) + queue.task_done() + pending_items -= 1 + continue + items.append(item) + total_size += item_size + if total_size >= BATCH_SIZE_LIMIT: + self.log.debug("hit batch size limit (size: %d)", total_size) + break + except Empty: break - except Empty: - break + except BaseException: + for _ in range(pending_items): + queue.task_done() + raise return items diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 21b72124a..a5a14f4c0 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2572,7 +2572,7 @@ async def shutdown(self): self.assertEqual(provider.result, 6) self.assertTrue(client._join_cleanup_complete) - def test_async_cache_provider_task_can_reenter_join_during_runner_close(self): + def test_async_cache_provider_executor_can_reenter_join_during_runner_close(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) finalizer_called = threading.Event() @@ -2582,7 +2582,7 @@ async def pending_task(): try: await asyncio.Event().wait() finally: - client.join() + await asyncio.to_thread(client.join) finalizer_called.set() asyncio.create_task(pending_task()) diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index f13733863..4a636fd47 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -53,6 +53,18 @@ def test_dropping_oversize_msg(self) -> None: self.assertTrue(q.empty()) self.assertEqual(q.unfinished_tasks, 0) + def test_next_balances_dequeued_work_if_batching_is_interrupted(self) -> None: + q = Queue() + consumer = Consumer(q, "") + q.put(_track_event()) + + with mock.patch("posthog.consumer.json.dumps", side_effect=SystemExit): + with self.assertRaises(SystemExit): + consumer.next() + + self.assertTrue(q.empty()) + self.assertEqual(q.unfinished_tasks, 0) + def test_max_msg_size_param_raises_per_event_ceiling(self) -> None: q = Queue() consumer = Consumer(q, "", flush_at=1, max_msg_size=4 * MAX_MSG_SIZE) From 72bee65a4ebd9c40ce310f64116ad2a83a76437d Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:04:23 +0700 Subject: [PATCH 14/27] fix: propagate context through generic executors Use a private context-aware event loop with a process-safe registry wrapper so thread-backed executor proxies remain reentrant without breaking serialization executors. --- posthog/_async_utils.py | 68 +++++++++++++++++++++++++------------ posthog/test/test_client.py | 14 ++++++-- 2 files changed, 58 insertions(+), 24 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index bc8244c43..a19b6496b 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,15 +1,54 @@ import asyncio +import os import threading from collections.abc import Awaitable -from concurrent.futures import ThreadPoolExecutor -from contextvars import copy_context +from contextvars import Context, copy_context from typing import Any -class _ContextThreadPoolExecutor(ThreadPoolExecutor): - def submit(self, fn, /, *args, **kwargs): - context = copy_context() - return super().submit(context.run, fn, *args, **kwargs) +_executor_contexts: dict[int, Context] = {} +_executor_contexts_lock = threading.Lock() +_executor_context_token = 0 + + +def _run_with_registered_context(origin_pid, token, func, *args): + if os.getpid() != origin_pid: + return func(*args) + with _executor_contexts_lock: + context = _executor_contexts.get(token) + if context is None: + return func(*args) + return context.run(func, *args) + + +class _ContextEventLoop(asyncio.SelectorEventLoop): + def run_in_executor(self, executor, func, *args): # type: ignore[override] + global _executor_context_token + with _executor_contexts_lock: + _executor_context_token += 1 + token = _executor_context_token + _executor_contexts[token] = copy_context() + + try: + future = super().run_in_executor( + executor, + _run_with_registered_context, + os.getpid(), + token, + func, + *args, + ) + except BaseException: + with _executor_contexts_lock: + _executor_contexts.pop(token, None) + raise + + def remove_context(_): + with _executor_contexts_lock: + _executor_contexts.pop(token, None) + + future.add_done_callback(remove_context) + return future class _BackgroundEventLoopRunner: @@ -81,22 +120,7 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: return self._loop def _run_loop(self) -> None: - loop = asyncio.new_event_loop() - loop.set_default_executor(_ContextThreadPoolExecutor()) - original_run_in_executor = loop.run_in_executor - - def run_in_executor(executor, func, *args): - if not isinstance(executor, ThreadPoolExecutor): - return original_run_in_executor(executor, func, *args) - context = copy_context() - return original_run_in_executor(executor, context.run, func, *args) - - try: - setattr(loop, "run_in_executor", run_in_executor) - except (AttributeError, TypeError): - # Some policy-provided loops expose read-only methods. Their default - # executor still propagates context through _ContextThreadPoolExecutor. - pass + loop = _ContextEventLoop() asyncio.set_event_loop(loop) with self._lock: self._loop = loop diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index a5a14f4c0..0307da8d7 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -4,7 +4,7 @@ import time import unittest import warnings -from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor from datetime import datetime from unittest import mock from uuid import UUID, uuid4 @@ -2533,13 +2533,23 @@ def test_async_cache_provider_custom_executor_can_reenter_join(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) executor_called = threading.Event() + class DelegatingExecutor(Executor): + def __init__(self): + self.executor = ThreadPoolExecutor(max_workers=1) + + def submit(self, fn, /, *args, **kwargs): + return self.executor.submit(fn, *args, **kwargs) + + def shutdown(self, wait=True, *, cancel_futures=False): + self.executor.shutdown(wait=wait, cancel_futures=cancel_futures) + class AsyncProvider: async def shutdown(self): def reenter_join(): client.join() executor_called.set() - with ThreadPoolExecutor(max_workers=1) as executor: + with DelegatingExecutor() as executor: await asyncio.get_running_loop().run_in_executor( executor, reenter_join ) From ebd77d94e812a4c71881d32ad4dc66e33dedcabc Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:14:40 +0700 Subject: [PATCH 15/27] fix: retain executor context through cancellation Carry callback context inside a serialization-safe callable so executor cancellation, process pools, and fork do not depend on shared registry state. --- posthog/_async_utils.py | 79 +++++++++++++++++++++-------------------- 1 file changed, 41 insertions(+), 38 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index a19b6496b..7e7ca350f 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,54 +1,43 @@ import asyncio -import os import threading from collections.abc import Awaitable +from concurrent.futures import ThreadPoolExecutor from contextvars import Context, copy_context from typing import Any -_executor_contexts: dict[int, Context] = {} -_executor_contexts_lock = threading.Lock() -_executor_context_token = 0 +class _PlainExecutorCall: + def __init__(self, func, args, kwargs) -> None: + self._func = func + self._args = args + self._kwargs = kwargs + def __call__(self): + return self._func(*self._args, **self._kwargs) -def _run_with_registered_context(origin_pid, token, func, *args): - if os.getpid() != origin_pid: - return func(*args) - with _executor_contexts_lock: - context = _executor_contexts.get(token) - if context is None: - return func(*args) - return context.run(func, *args) +class _ContextExecutorCall: + """Carry context in-process while remaining safe for serializing executors.""" -class _ContextEventLoop(asyncio.SelectorEventLoop): - def run_in_executor(self, executor, func, *args): # type: ignore[override] - global _executor_context_token - with _executor_contexts_lock: - _executor_context_token += 1 - token = _executor_context_token - _executor_contexts[token] = copy_context() + def __init__(self, context: Context, func, args, kwargs=None) -> None: + self._context = context + self._func = func + self._args = args + self._kwargs = kwargs or {} - try: - future = super().run_in_executor( - executor, - _run_with_registered_context, - os.getpid(), - token, - func, - *args, - ) - except BaseException: - with _executor_contexts_lock: - _executor_contexts.pop(token, None) - raise + def __call__(self): + return self._context.run(self._func, *self._args, **self._kwargs) - def remove_context(_): - with _executor_contexts_lock: - _executor_contexts.pop(token, None) + def __reduce__(self): + # Context objects are not picklable and are process-local. Executors + # that serialize work reconstruct a plain call instead. + return (_PlainExecutorCall, (self._func, self._args, self._kwargs)) - future.add_done_callback(remove_context) - return future + +class _ContextThreadPoolExecutor(ThreadPoolExecutor): + def submit(self, fn, /, *args, **kwargs): + call = _ContextExecutorCall(copy_context(), fn, args, kwargs) + return super().submit(call) class _BackgroundEventLoopRunner: @@ -120,7 +109,21 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: return self._loop def _run_loop(self) -> None: - loop = _ContextEventLoop() + loop = asyncio.new_event_loop() + loop.set_default_executor(_ContextThreadPoolExecutor()) + original_run_in_executor = loop.run_in_executor + + def run_in_executor(executor, func, *args): + call = _ContextExecutorCall(copy_context(), func, args) + return original_run_in_executor(executor, call) + + try: + setattr(loop, "run_in_executor", run_in_executor) + except (AttributeError, TypeError): + # Policy-provided loops can expose read-only methods. Default + # executor calls still propagate context through our executor. + pass + asyncio.set_event_loop(loop) with self._lock: self._loop = loop From 9b444d8fcda0e624f3ab2556b157ab7ac1830cdf Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:26:08 +0700 Subject: [PATCH 16/27] fix: harden async runner startup and platform loops Use a context-aware platform event loop and make startup errors and concurrent close observable without orphaning runner threads. --- posthog/_async_utils.py | 90 ++++++++++++++++++++------------ posthog/test/test_async_utils.py | 49 +++++++++++++++++ 2 files changed, 106 insertions(+), 33 deletions(-) create mode 100644 posthog/test/test_async_utils.py diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 7e7ca350f..fd6a27dc1 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,7 +1,7 @@ import asyncio +import sys import threading from collections.abc import Awaitable -from concurrent.futures import ThreadPoolExecutor from contextvars import Context, copy_context from typing import Any @@ -34,10 +34,16 @@ def __reduce__(self): return (_PlainExecutorCall, (self._func, self._args, self._kwargs)) -class _ContextThreadPoolExecutor(ThreadPoolExecutor): - def submit(self, fn, /, *args, **kwargs): - call = _ContextExecutorCall(copy_context(), fn, args, kwargs) - return super().submit(call) +if sys.platform == "win32": + from asyncio.windows_events import ProactorEventLoop as _PlatformEventLoop +else: + _PlatformEventLoop = asyncio.SelectorEventLoop + + +class _ContextEventLoop(_PlatformEventLoop): + def run_in_executor(self, executor, func, *args): # type: ignore[override] + call = _ContextExecutorCall(copy_context(), func, args) + return super().run_in_executor(executor, call) class _BackgroundEventLoopRunner: @@ -48,6 +54,8 @@ def __init__(self) -> None: self._thread: threading.Thread | None = None self._closing_threads: set[threading.Thread] = set() self._started = threading.Event() + self._startup_error: BaseException | None = None + self._close_requested = False self._lock = threading.Lock() def run(self, awaitable: Awaitable[Any]) -> Any: @@ -59,15 +67,23 @@ def close(self) -> None: with self._lock: loop = self._loop thread = self._thread - self._loop = None - self._thread = None - if thread is not None: + if thread is None: + return + if loop is None: + self._close_requested = True + else: + self._loop = None + self._thread = None self._closing_threads.add(thread) - if loop is None or thread is None or loop.is_closed(): + if loop is None: + if thread is not threading.current_thread(): + thread.join() + return + + if loop.is_closed(): with self._lock: - if thread is not None: - self._closing_threads.discard(thread) + self._closing_threads.discard(thread) return if thread is threading.current_thread(): @@ -95,40 +111,43 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: ): return self._loop - self._started.clear() - self._thread = threading.Thread( - target=self._run_loop, - name="PostHogBackgroundEventLoopRunner", - daemon=True, - ) - self._thread.start() + if self._thread is None or not self._thread.is_alive(): + self._started.clear() + self._startup_error = None + self._close_requested = False + self._thread = threading.Thread( + target=self._run_loop, + name="PostHogBackgroundEventLoopRunner", + daemon=True, + ) + self._thread.start() self._started.wait() with self._lock: + if self._startup_error is not None: + raise self._startup_error assert self._loop is not None return self._loop def _run_loop(self) -> None: - loop = asyncio.new_event_loop() - loop.set_default_executor(_ContextThreadPoolExecutor()) - original_run_in_executor = loop.run_in_executor - - def run_in_executor(executor, func, *args): - call = _ContextExecutorCall(copy_context(), func, args) - return original_run_in_executor(executor, call) - try: - setattr(loop, "run_in_executor", run_in_executor) - except (AttributeError, TypeError): - # Policy-provided loops can expose read-only methods. Default - # executor calls still propagate context through our executor. - pass + loop = _ContextEventLoop() + asyncio.set_event_loop(loop) + except BaseException as error: + with self._lock: + self._startup_error = error + self._thread = None + self._started.set() + return - asyncio.set_event_loop(loop) with self._lock: self._loop = loop + close_requested = self._close_requested self._started.set() + if close_requested: + loop.call_soon(loop.stop) + try: loop.run_forever() finally: @@ -143,5 +162,10 @@ def run_in_executor(executor, func, *args): loop.run_until_complete(loop.shutdown_default_executor()) asyncio.set_event_loop(None) loop.close() + current = threading.current_thread() with self._lock: - self._closing_threads.discard(threading.current_thread()) + if self._thread is current: + self._thread = None + if self._loop is loop: + self._loop = None + self._closing_threads.discard(current) diff --git a/posthog/test/test_async_utils.py b/posthog/test/test_async_utils.py new file mode 100644 index 000000000..6a239e884 --- /dev/null +++ b/posthog/test/test_async_utils.py @@ -0,0 +1,49 @@ +import asyncio +import threading +import unittest +from unittest import mock + +from posthog._async_utils import _BackgroundEventLoopRunner, _ContextEventLoop + + +class TestBackgroundEventLoopRunner(unittest.TestCase): + def test_startup_error_is_reported(self): + runner = _BackgroundEventLoopRunner() + awaitable = asyncio.sleep(0) + + with mock.patch( + "posthog._async_utils._ContextEventLoop", + side_effect=RuntimeError("startup failed"), + ): + with self.assertRaisesRegex(RuntimeError, "startup failed"): + runner.run(awaitable) + + awaitable.close() + + def test_close_during_startup_does_not_orphan_thread(self): + runner = _BackgroundEventLoopRunner() + construction_started = threading.Event() + release_construction = threading.Event() + + def create_loop(): + construction_started.set() + self.assertTrue(release_construction.wait(2)) + return _ContextEventLoop() + + with mock.patch( + "posthog._async_utils._ContextEventLoop", side_effect=create_loop + ): + ensure_thread = threading.Thread(target=runner._ensure_loop) + ensure_thread.start() + self.assertTrue(construction_started.wait(1)) + + close_thread = threading.Thread(target=runner.close) + close_thread.start() + release_construction.set() + ensure_thread.join(2) + close_thread.join(2) + + self.assertFalse(ensure_thread.is_alive()) + self.assertFalse(close_thread.is_alive()) + self.assertIsNone(runner._thread) + self.assertIsNone(runner._loop) From 42596994bbfe138208cdc68baeb9053f15c2be35 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:40:29 +0700 Subject: [PATCH 17/27] fix: serialize async runner run and close Prevent close from publishing or stopping a loop during an active startup/run and close partially initialized loops on startup failure. --- posthog/_async_utils.py | 59 ++++++++++++++++++-------------- posthog/test/test_async_utils.py | 18 +++++++--- 2 files changed, 47 insertions(+), 30 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index fd6a27dc1..eb392cffb 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -57,41 +57,47 @@ def __init__(self) -> None: self._startup_error: BaseException | None = None self._close_requested = False self._lock = threading.Lock() + self._operation_lock = threading.Lock() def run(self, awaitable: Awaitable[Any]) -> Any: - loop = self._ensure_loop() - future = asyncio.run_coroutine_threadsafe(self._await_result(awaitable), loop) - return future.result() + with self._operation_lock: + loop = self._ensure_loop() + future = asyncio.run_coroutine_threadsafe( + self._await_result(awaitable), loop + ) + return future.result() def close(self) -> None: + current = threading.current_thread() with self._lock: - loop = self._loop - thread = self._thread - if thread is None: + if current is self._thread and self._loop is not None: + self._loop.call_soon(self._loop.stop) return - if loop is None: - self._close_requested = True - else: - self._loop = None - self._thread = None - self._closing_threads.add(thread) - - if loop is None: - if thread is not threading.current_thread(): - thread.join() - return - if loop.is_closed(): + with self._operation_lock: with self._lock: - self._closing_threads.discard(thread) - return + loop = self._loop + thread = self._thread + if thread is None: + return + if loop is None: + self._close_requested = True + else: + self._loop = None + self._thread = None + self._closing_threads.add(thread) - if thread is threading.current_thread(): - loop.call_soon(loop.stop) - return + if loop is None: + thread.join() + return + + if loop.is_closed(): + with self._lock: + self._closing_threads.discard(thread) + return - loop.call_soon_threadsafe(loop.stop) - thread.join() + loop.call_soon_threadsafe(loop.stop) + thread.join() def owns_thread(self, thread: threading.Thread) -> bool: with self._lock: @@ -130,10 +136,13 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: return self._loop def _run_loop(self) -> None: + loop = None try: loop = _ContextEventLoop() asyncio.set_event_loop(loop) except BaseException as error: + if loop is not None and not loop.is_closed(): + loop.close() with self._lock: self._startup_error = error self._thread = None diff --git a/posthog/test/test_async_utils.py b/posthog/test/test_async_utils.py index 6a239e884..f606b5aab 100644 --- a/posthog/test/test_async_utils.py +++ b/posthog/test/test_async_utils.py @@ -20,30 +20,38 @@ def test_startup_error_is_reported(self): awaitable.close() - def test_close_during_startup_does_not_orphan_thread(self): + def test_close_waits_for_run_during_startup(self): runner = _BackgroundEventLoopRunner() construction_started = threading.Event() release_construction = threading.Event() + run_errors = [] def create_loop(): construction_started.set() self.assertTrue(release_construction.wait(2)) return _ContextEventLoop() + def run(): + try: + runner.run(asyncio.sleep(0)) + except BaseException as error: + run_errors.append(error) + with mock.patch( "posthog._async_utils._ContextEventLoop", side_effect=create_loop ): - ensure_thread = threading.Thread(target=runner._ensure_loop) - ensure_thread.start() + run_thread = threading.Thread(target=run) + run_thread.start() self.assertTrue(construction_started.wait(1)) close_thread = threading.Thread(target=runner.close) close_thread.start() release_construction.set() - ensure_thread.join(2) + run_thread.join(2) close_thread.join(2) - self.assertFalse(ensure_thread.is_alive()) + self.assertFalse(run_thread.is_alive()) self.assertFalse(close_thread.is_alive()) + self.assertEqual(run_errors, []) self.assertIsNone(runner._thread) self.assertIsNone(runner._loop) From 4de053368d09c8c514076bb2d256e29b54dd7035 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:48:39 +0700 Subject: [PATCH 18/27] fix: coordinate concurrent async runner operations Submit concurrent work without holding completion locks, reject loop-thread synchronous reentry, and define close-during-startup cancellation. --- posthog/_async_utils.py | 65 ++++++++++++++++++-------------- posthog/test/test_async_utils.py | 31 ++++++++++++++- 2 files changed, 65 insertions(+), 31 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index eb392cffb..a0e5e61ad 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -57,47 +57,52 @@ def __init__(self) -> None: self._startup_error: BaseException | None = None self._close_requested = False self._lock = threading.Lock() - self._operation_lock = threading.Lock() def run(self, awaitable: Awaitable[Any]) -> Any: - with self._operation_lock: + if threading.current_thread() is self._thread: + raise RuntimeError("cannot synchronously run from the runner thread") + + while True: loop = self._ensure_loop() - future = asyncio.run_coroutine_threadsafe( - self._await_result(awaitable), loop - ) - return future.result() + with self._lock: + if loop is self._loop and not self._close_requested: + future = asyncio.run_coroutine_threadsafe( + self._await_result(awaitable), loop + ) + break + return future.result() def close(self) -> None: current = threading.current_thread() with self._lock: - if current is self._thread and self._loop is not None: - self._loop.call_soon(self._loop.stop) + loop = self._loop + thread = self._thread + if thread is None: return - - with self._operation_lock: - with self._lock: - loop = self._loop - thread = self._thread - if thread is None: - return - if loop is None: - self._close_requested = True - else: - self._loop = None - self._thread = None - self._closing_threads.add(thread) - + self._close_requested = True if loop is None: + self._startup_error = RuntimeError("runner closed during startup") + else: + self._loop = None + self._thread = None + self._closing_threads.add(thread) + + if loop is None: + if thread is not current: thread.join() - return + return - if loop.is_closed(): - with self._lock: - self._closing_threads.discard(thread) - return + if loop.is_closed(): + with self._lock: + self._closing_threads.discard(thread) + return + + if thread is current: + loop.call_soon(loop.stop) + return - loop.call_soon_threadsafe(loop.stop) - thread.join() + loop.call_soon_threadsafe(loop.stop) + thread.join() def owns_thread(self, thread: threading.Thread) -> bool: with self._lock: @@ -152,6 +157,8 @@ def _run_loop(self) -> None: with self._lock: self._loop = loop close_requested = self._close_requested + if close_requested and self._startup_error is None: + self._startup_error = RuntimeError("runner closed during startup") self._started.set() if close_requested: diff --git a/posthog/test/test_async_utils.py b/posthog/test/test_async_utils.py index f606b5aab..6c82fe1d4 100644 --- a/posthog/test/test_async_utils.py +++ b/posthog/test/test_async_utils.py @@ -32,9 +32,11 @@ def create_loop(): return _ContextEventLoop() def run(): + awaitable = asyncio.sleep(0) try: - runner.run(asyncio.sleep(0)) + runner.run(awaitable) except BaseException as error: + awaitable.close() run_errors.append(error) with mock.patch( @@ -52,6 +54,31 @@ def run(): self.assertFalse(run_thread.is_alive()) self.assertFalse(close_thread.is_alive()) - self.assertEqual(run_errors, []) + self.assertEqual(len(run_errors), 1) + self.assertRegex(str(run_errors[0]), "closed during startup") self.assertIsNone(runner._thread) self.assertIsNone(runner._loop) + + def test_run_from_runner_thread_fails_instead_of_deadlocking(self): + runner = _BackgroundEventLoopRunner() + + async def reenter(): + awaitable = asyncio.sleep(0) + try: + with self.assertRaisesRegex(RuntimeError, "runner thread"): + runner.run(awaitable) + finally: + awaitable.close() + + runner.run(reenter()) + runner.close() + + def test_close_from_runner_thread_allows_fresh_loop(self): + runner = _BackgroundEventLoopRunner() + + async def close_runner(): + runner.close() + + runner.run(close_runner()) + runner.run(asyncio.sleep(0)) + runner.close() From bcc703724cc0c5a9b75bc1ff56d9c0a774a09a09 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 01:00:55 +0700 Subject: [PATCH 19/27] fix: preserve async runner loop policy Use the configured event-loop policy when extensible, fall back only for read-only implementations, and stabilize startup-close coverage. --- posthog/_async_utils.py | 16 ++++++++++++- posthog/test/test_async_utils.py | 39 ++++++++++++++++++++++++++++---- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index a0e5e61ad..91c2099c8 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -143,7 +143,21 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: def _run_loop(self) -> None: loop = None try: - loop = _ContextEventLoop() + loop = asyncio.new_event_loop() + original_run_in_executor = loop.run_in_executor + + def run_in_executor(executor, func, *args): + call = _ContextExecutorCall(copy_context(), func, args) + return original_run_in_executor(executor, call) + + try: + setattr(loop, "run_in_executor", run_in_executor) + except (AttributeError, TypeError): + # Preserve custom policy loops when they are extensible. For a + # read-only implementation, use the platform-equivalent loop + # so explicit executors retain callback context safely. + loop.close() + loop = _ContextEventLoop() asyncio.set_event_loop(loop) except BaseException as error: if loop is not None and not loop.is_closed(): diff --git a/posthog/test/test_async_utils.py b/posthog/test/test_async_utils.py index 6c82fe1d4..d32942d80 100644 --- a/posthog/test/test_async_utils.py +++ b/posthog/test/test_async_utils.py @@ -1,9 +1,10 @@ import asyncio import threading +import time import unittest from unittest import mock -from posthog._async_utils import _BackgroundEventLoopRunner, _ContextEventLoop +from posthog._async_utils import _BackgroundEventLoopRunner class TestBackgroundEventLoopRunner(unittest.TestCase): @@ -12,7 +13,7 @@ def test_startup_error_is_reported(self): awaitable = asyncio.sleep(0) with mock.patch( - "posthog._async_utils._ContextEventLoop", + "posthog._async_utils.asyncio.new_event_loop", side_effect=RuntimeError("startup failed"), ): with self.assertRaisesRegex(RuntimeError, "startup failed"): @@ -26,10 +27,12 @@ def test_close_waits_for_run_during_startup(self): release_construction = threading.Event() run_errors = [] + original_new_event_loop = asyncio.new_event_loop + def create_loop(): construction_started.set() self.assertTrue(release_construction.wait(2)) - return _ContextEventLoop() + return original_new_event_loop() def run(): awaitable = asyncio.sleep(0) @@ -40,7 +43,7 @@ def run(): run_errors.append(error) with mock.patch( - "posthog._async_utils._ContextEventLoop", side_effect=create_loop + "posthog._async_utils.asyncio.new_event_loop", side_effect=create_loop ): run_thread = threading.Thread(target=run) run_thread.start() @@ -48,6 +51,11 @@ def run(): close_thread = threading.Thread(target=runner.close) close_thread.start() + deadline = time.monotonic() + 1 + while not runner._close_requested: + if time.monotonic() >= deadline: + self.fail("close did not reach startup state") + time.sleep(0.001) release_construction.set() run_thread.join(2) close_thread.join(2) @@ -59,6 +67,29 @@ def run(): self.assertIsNone(runner._thread) self.assertIsNone(runner._loop) + def test_runner_preserves_configured_event_loop_policy(self): + runner = _BackgroundEventLoopRunner() + original_policy = asyncio.get_event_loop_policy() + + class Policy(asyncio.DefaultEventLoopPolicy): + loop = None + + def new_event_loop(self): + self.loop = super().new_event_loop() + return self.loop + + policy = Policy() + + async def running_loop(): + return asyncio.get_running_loop() + + try: + asyncio.set_event_loop_policy(policy) + self.assertIs(runner.run(running_loop()), policy.loop) + runner.close() + finally: + asyncio.set_event_loop_policy(original_policy) + def test_run_from_runner_thread_fails_instead_of_deadlocking(self): runner = _BackgroundEventLoopRunner() From 12b0816027c5c680350e36ba2a37564919579eca Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:24:18 +0700 Subject: [PATCH 20/27] docs: clarify lifecycle calls from error callbacks Document direct deferred lifecycle calls and the safe application-thread handoff pattern for callbacks that require blocking shutdown completion. --- posthog/__init__.py | 11 ++++++++++- posthog/client.py | 13 ++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index a106d512c..c61ebd5e0 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -310,7 +310,9 @@ def get_tags() -> Dict[str, Any]: host: PostHog ingestion host. Defaults to the US ingestion endpoint when not set. on_error: Optional callback invoked by background consumers when event upload - fails. + fails. Keep it short and non-blocking. Lifecycle methods can be called + directly and will be deferred, but the callback must not wait for another + thread or task that calls ``flush()``, ``join()``, or ``shutdown()``. debug: Enable verbose SDK logging and re-raise errors from public APIs. send: If False, queueing succeeds but events are not sent to PostHog. sync_mode: If True, send events synchronously instead of using background @@ -1142,6 +1144,13 @@ def shutdown() -> None: """ Flush all messages and cleanly shutdown the client. + This normally blocks until delivery and cleanup finish. Calls made directly + from SDK callbacks such as ``on_error`` are deferred to avoid deadlocking the + worker. If blocking completion is required, signal an application-owned + thread, return from the callback, and call ``shutdown()`` from that thread. + Do not wait inside a callback for another thread or task calling a lifecycle + method. + Examples: ```python from posthog import shutdown diff --git a/posthog/client.py b/posthog/client.py index 0a2049152..25107b3c8 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -583,7 +583,10 @@ def __init__( max_queue_size: Maximum number of events buffered before upload. send: If False, queueing succeeds but events are not sent. on_error: Optional callback invoked by background consumers when an - upload fails. + upload fails. Keep it short and non-blocking. Calling lifecycle + methods directly is safe and deferred, but do not start another + thread or task that calls ``flush()``, ``join()``, or + ``shutdown()`` and then wait for it from the callback. flush_at: Number of queued events that triggers a batch upload. flush_interval: Maximum seconds a background consumer waits before flushing a partial batch. @@ -2471,6 +2474,14 @@ def shutdown(self) -> None: """ Flush all messages and cleanly shutdown the client. Call this before the process ends in serverless environments to avoid data loss. + Normally this method blocks until delivery and cleanup finish. When + called directly from an SDK callback such as ``on_error``, shutdown is + deferred to avoid blocking the worker that invoked the callback. If the + callback must coordinate a blocking shutdown, have it signal an + application-owned thread and return before that thread calls shutdown. + Do not wait inside the callback for another thread or task that calls a + lifecycle method. + Examples: ```python posthog.shutdown() From 6ea9973627504919513bf4bb3fd095ad37a3114b Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Thu, 6 Aug 2026 08:33:04 +0200 Subject: [PATCH 21/27] address lifecycle review feedback --- posthog/client.py | 28 ++++------------------------ posthog/test/test_client.py | 23 +++++++++++++---------- 2 files changed, 17 insertions(+), 34 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index 4eb4a9e4c..c4685795a 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -705,14 +705,11 @@ def __init__( self.sync_mode = sync_mode self._lifecycle_lock = threading.Lock() self._lifecycle_condition = threading.Condition(self._lifecycle_lock) - self._lifecycle_in_progress = False self._lifecycle_owner: Optional[threading.Thread] = None self._workers_joined = False self._join_cleanup_complete = False self._shutdown_requested = False self._shutdown_complete = False - self._deferred_lifecycle_error: Optional[BaseException] = None - self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_generation = 0 self._lifecycle_callback_context: ContextVar[bool] = ContextVar( @@ -722,7 +719,6 @@ def __init__( self._deferred_flush_pending = False self._deferred_flush_followup = False self._deferred_flush_followup_timeout: Optional[float] = None - self._shutdown_complete_event = threading.Event() # Used for session replay URL generation - we don't want the server host here. self.raw_host = normalize_host(host) self.host = determine_server_host(host) @@ -2005,13 +2001,9 @@ def _reinit_after_fork(self): for lane in self._lanes: lane.rebuild_after_fork() - shutdown_complete = self._shutdown_complete self._lifecycle_lock = threading.Lock() self._lifecycle_condition = threading.Condition(self._lifecycle_lock) - self._lifecycle_in_progress = False self._lifecycle_owner = None - self._deferred_lifecycle_error = None - self._deferred_lifecycle_failure = threading.Event() self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_generation = 0 self._lifecycle_callback_context = ContextVar( @@ -2021,9 +2013,6 @@ def _reinit_after_fork(self): self._deferred_flush_pending = False self._deferred_flush_followup = False self._deferred_flush_followup_timeout = None - self._shutdown_complete_event = threading.Event() - if shutdown_complete: - self._shutdown_complete_event.set() # Async runner threads do not survive fork(); recreate lazily on next async cache call. self._flag_definition_cache_provider_async_runner = None @@ -2310,7 +2299,7 @@ def run() -> None: require_shutdown = self._shutdown_requested try: self._run_lifecycle(require_shutdown=require_shutdown) - except BaseException as error: + except BaseException: attempt += 1 self.log.exception("Deferred %s attempt %d failed", name, attempt) with self._lifecycle_lock: @@ -2319,8 +2308,6 @@ def run() -> None: continue if attempt < 2: continue - self._deferred_lifecycle_error = error - self._deferred_lifecycle_failure.set() self._deferred_lifecycle_thread_pending = False return @@ -2420,9 +2407,6 @@ def _shutdown_once(self) -> None: if self.exception_capture: self.exception_capture.close() self._shutdown_complete = True - self._deferred_lifecycle_error = None - self._deferred_lifecycle_failure.clear() - self._shutdown_complete_event.set() def _run_lifecycle(self, require_shutdown: bool = False) -> None: while True: @@ -2433,14 +2417,13 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: self._join_cleanup_complete or self._shutdown_complete ): return - if self._lifecycle_in_progress: + if self._lifecycle_owner is not None: if self._is_lifecycle_callback_thread() or ( threading.current_thread() is self._lifecycle_owner ): return self._lifecycle_condition.wait() continue - self._lifecycle_in_progress = True self._lifecycle_owner = threading.current_thread() try: @@ -2460,17 +2443,16 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: and not run_shutdown ): continue - self._lifecycle_in_progress = False self._lifecycle_owner = None self._lifecycle_condition.notify_all() return except BaseException: with self._lifecycle_condition: - self._lifecycle_in_progress = False self._lifecycle_owner = None self._lifecycle_condition.notify_all() raise + @no_throw() def join(self) -> None: """ Flush queued events and end the consumer threads. Do not use directly, call `shutdown()` instead. @@ -2484,6 +2466,7 @@ def join(self) -> None: return self._run_lifecycle() + @no_throw() def shutdown(self) -> None: """ Flush all messages and cleanly shutdown the client. Call this before the process ends in serverless environments to avoid data loss. @@ -2505,9 +2488,6 @@ def shutdown(self) -> None: if self._shutdown_complete: return self._shutdown_requested = True - if not self._lifecycle_in_progress: - self._deferred_lifecycle_error = None - self._deferred_lifecycle_failure.clear() if self._defer_lifecycle_from_callback("shutdown"): return self._run_lifecycle(require_shutdown=True) diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 0307da8d7..2c40c08df 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -26,6 +26,15 @@ # Legacy single-flag behavior remains covered here; warning emission itself is # asserted in test_evaluate_flags.py. +def _wait_until(predicate, timeout=3): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + pytestmark = [ pytest.mark.filterwarnings( r"ignore:`(feature_enabled|get_feature_flag|get_feature_flag_payload)` is deprecated:DeprecationWarning" @@ -2375,7 +2384,7 @@ def on_error(error, batch): release_first_send.set() self.assertTrue(callback_returned.wait(1)) - self.assertTrue(client._shutdown_complete_event.wait(3)) + self.assertTrue(_wait_until(lambda: client._shutdown_complete)) self.assertEqual(sent_events, ["first", "second"]) self.assertEqual(client.queue.unfinished_tasks, 0) @@ -2679,7 +2688,6 @@ def assert_stopped(self): client.join() self.assertTrue(client._shutdown_complete) - self.assertTrue(client._shutdown_complete_event.is_set()) def test_join_retries_auxiliary_cleanup_after_failure(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) @@ -2689,8 +2697,7 @@ def test_join_retries_auxiliary_cleanup_after_failure(self): "_shutdown_flag_definition_cache_provider", side_effect=[Exception("cleanup failed"), None], ) as cleanup: - with self.assertRaisesRegex(Exception, "cleanup failed"): - client.join() + client.join() self.assertTrue(client._workers_joined) self.assertFalse(client._join_cleanup_complete) @@ -2738,11 +2745,10 @@ def run_join(): release_cleanup.set() join_thread.join(2) shutdown_thread.join(3) - self.assertTrue(client._shutdown_complete_event.is_set()) self.assertFalse(join_thread.is_alive()) self.assertFalse(shutdown_thread.is_alive()) - self.assertEqual(str(join_errors[0]), "cleanup failed") + self.assertEqual(join_errors, []) self.assertEqual(cleanup_calls, 2) self.assertTrue(client._shutdown_complete) @@ -2752,15 +2758,12 @@ def test_shutdown_retries_cleanup_before_publishing_completion(self): exception_capture.close.side_effect = [Exception("cleanup failed"), None] client.exception_capture = exception_capture - with self.assertRaisesRegex(Exception, "cleanup failed"): - client.shutdown() + client.shutdown() self.assertFalse(client._shutdown_complete) - self.assertFalse(client._shutdown_complete_event.is_set()) client.shutdown() self.assertTrue(client._shutdown_complete) - self.assertTrue(client._shutdown_complete_event.is_set()) self.assertEqual(exception_capture.close.call_count, 2) def test_shutdown_does_not_wait_for_idle_consumers_flush_interval(self): From a19299e7a6a9bc14074284e2adde721554e202f2 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:55:26 +0700 Subject: [PATCH 22/27] simplify deferred lifecycle coalescing --- posthog/client.py | 12 +++++------ posthog/test/test_client.py | 43 +++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index c4685795a..1763b0f52 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -711,7 +711,7 @@ def __init__( self._shutdown_requested = False self._shutdown_complete = False self._deferred_lifecycle_thread_pending = False - self._deferred_lifecycle_generation = 0 + self._deferred_lifecycle_dirty = False self._lifecycle_callback_context: ContextVar[bool] = ContextVar( "posthog_lifecycle_callback", default=False ) @@ -2005,7 +2005,7 @@ def _reinit_after_fork(self): self._lifecycle_condition = threading.Condition(self._lifecycle_lock) self._lifecycle_owner = None self._deferred_lifecycle_thread_pending = False - self._deferred_lifecycle_generation = 0 + self._deferred_lifecycle_dirty = False self._lifecycle_callback_context = ContextVar( "posthog_lifecycle_callback", default=False ) @@ -2286,7 +2286,7 @@ def _defer_lifecycle_from_callback(self, name: str) -> bool: if not self._is_lifecycle_callback_thread(): return False with self._lifecycle_lock: - self._deferred_lifecycle_generation += 1 + self._deferred_lifecycle_dirty = True if self._deferred_lifecycle_thread_pending: return True self._deferred_lifecycle_thread_pending = True @@ -2295,7 +2295,7 @@ def run() -> None: attempt = 0 while True: with self._lifecycle_lock: - generation = self._deferred_lifecycle_generation + self._deferred_lifecycle_dirty = False require_shutdown = self._shutdown_requested try: self._run_lifecycle(require_shutdown=require_shutdown) @@ -2303,7 +2303,7 @@ def run() -> None: attempt += 1 self.log.exception("Deferred %s attempt %d failed", name, attempt) with self._lifecycle_lock: - if generation != self._deferred_lifecycle_generation: + if self._deferred_lifecycle_dirty: attempt = 0 continue if attempt < 2: @@ -2312,7 +2312,7 @@ def run() -> None: return with self._lifecycle_lock: - if generation != self._deferred_lifecycle_generation: + if self._deferred_lifecycle_dirty: attempt = 0 continue self._deferred_lifecycle_thread_pending = False diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 2c40c08df..0caf76712 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2299,6 +2299,49 @@ def run_lifecycle(require_shutdown=False): start_thread.assert_called_once() self.assertEqual(require_shutdown_calls, [False, True]) + self.assertTrue( + _wait_until(lambda: not client._deferred_lifecycle_thread_pending) + ) + self.assertFalse(client._deferred_lifecycle_dirty) + + def test_callback_shutdown_escalates_failed_deferred_join(self): + client = Client(FAKE_TEST_API_KEY) + first_run_started = threading.Event() + release_first_run = threading.Event() + shutdown_run_complete = threading.Event() + require_shutdown_calls = [] + + def run_lifecycle(require_shutdown=False): + require_shutdown_calls.append(require_shutdown) + if len(require_shutdown_calls) == 1: + first_run_started.set() + self.assertTrue(release_first_run.wait(2)) + raise Exception("join cleanup failed") + shutdown_run_complete.set() + + with ( + mock.patch.object( + client, "_is_lifecycle_callback_thread", return_value=True + ), + mock.patch.object(client, "_run_lifecycle", side_effect=run_lifecycle), + mock.patch.object( + client, + "_start_lifecycle_thread", + wraps=client._start_lifecycle_thread, + ) as start_thread, + ): + client.join() + self.assertTrue(first_run_started.wait(1)) + client.shutdown() + release_first_run.set() + self.assertTrue(shutdown_run_complete.wait(1)) + + start_thread.assert_called_once() + self.assertEqual(require_shutdown_calls, [False, True]) + self.assertTrue( + _wait_until(lambda: not client._deferred_lifecycle_thread_pending) + ) + self.assertFalse(client._deferred_lifecycle_dirty) def test_callback_flushes_are_coalesced_with_strongest_followup(self): client = Client(FAKE_TEST_API_KEY) From 04b30199e4d0a0945f40ac7914eb4347bd6f2bdd Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:51:28 +0700 Subject: [PATCH 23/27] address automated lifecycle review feedback --- posthog/__init__.py | 13 ++- posthog/_async_utils.py | 95 ++++++++++++------ posthog/client.py | 31 +++--- posthog/test/test_async_utils.py | 164 ++++++++++++++++++++++++++++++- posthog/test/test_client.py | 39 ++++++++ 5 files changed, 296 insertions(+), 46 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index c61ebd5e0..c28802f0e 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -1126,7 +1126,10 @@ def flush(timeout_seconds: Optional[float] = 10) -> None: def join() -> None: """ - Block until queued events are delivered and stop the client's background workers. Use `shutdown()` directly in most cases. + Attempt to process queued events and stop the client's background workers. Use `shutdown()` directly in most cases. + + Failed or undrainable events may be dropped and reported through logging or + ``on_error``; returning does not guarantee server receipt. Examples: ```python @@ -1144,9 +1147,11 @@ def shutdown() -> None: """ Flush all messages and cleanly shutdown the client. - This normally blocks until delivery and cleanup finish. Calls made directly - from SDK callbacks such as ``on_error`` are deferred to avoid deadlocking the - worker. If blocking completion is required, signal an application-owned + This normally blocks until queued events have been attempted and cleanup + finishes. Failed or undrainable events may be dropped and reported through + logging or ``on_error``; returning does not guarantee server receipt. Calls + made directly from SDK callbacks such as ``on_error`` are deferred to avoid + deadlocking the worker. If blocking completion is required, signal an application-owned thread, return from the callback, and call ``shutdown()`` from that thread. Do not wait inside a callback for another thread or task calling a lifecycle method. diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 91c2099c8..0ee6119ce 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,4 +1,5 @@ import asyncio +import inspect import sys import threading from collections.abc import Awaitable @@ -46,6 +47,13 @@ def run_in_executor(self, executor, func, *args): # type: ignore[override] return super().run_in_executor(executor, call) +class _LoopStartup: + def __init__(self) -> None: + self.done = threading.Event() + self.loop: asyncio.AbstractEventLoop | None = None + self.error: BaseException | None = None + + class _BackgroundEventLoopRunner: """Run awaitables to completion on a reusable background event loop.""" @@ -53,24 +61,35 @@ def __init__(self) -> None: self._loop: asyncio.AbstractEventLoop | None = None self._thread: threading.Thread | None = None self._closing_threads: set[threading.Thread] = set() - self._started = threading.Event() - self._startup_error: BaseException | None = None + self._startup: _LoopStartup | None = None self._close_requested = False self._lock = threading.Lock() def run(self, awaitable: Awaitable[Any]) -> Any: if threading.current_thread() is self._thread: + self._close_awaitable(awaitable) raise RuntimeError("cannot synchronously run from the runner thread") - while True: - loop = self._ensure_loop() - with self._lock: - if loop is self._loop and not self._close_requested: - future = asyncio.run_coroutine_threadsafe( - self._await_result(awaitable), loop - ) - break - return future.result() + try: + while True: + loop = self._ensure_loop() + with self._lock: + if loop is self._loop and not self._close_requested: + wrapped = self._await_result(awaitable) + try: + future = asyncio.run_coroutine_threadsafe(wrapped, loop) + except BaseException: + wrapped.close() + raise + break + except BaseException: + self._close_awaitable(awaitable) + raise + try: + return future.result() + finally: + if future.cancelled(): + self._close_awaitable(awaitable) def close(self) -> None: current = threading.current_thread() @@ -81,7 +100,8 @@ def close(self) -> None: return self._close_requested = True if loop is None: - self._startup_error = RuntimeError("runner closed during startup") + if self._startup is not None: + self._startup.error = RuntimeError("runner closed during startup") else: self._loop = None self._thread = None @@ -108,6 +128,11 @@ def owns_thread(self, thread: threading.Thread) -> bool: with self._lock: return thread is self._thread or thread in self._closing_threads + @staticmethod + def _close_awaitable(awaitable: Awaitable[Any]) -> None: + if inspect.iscoroutine(awaitable): + awaitable.close() + @staticmethod async def _await_result(awaitable: Awaitable[Any]) -> Any: return await awaitable @@ -122,25 +147,32 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: ): return self._loop + startup: _LoopStartup if self._thread is None or not self._thread.is_alive(): - self._started.clear() - self._startup_error = None + startup = _LoopStartup() + self._startup = startup self._close_requested = False self._thread = threading.Thread( target=self._run_loop, + args=(startup,), name="PostHogBackgroundEventLoopRunner", daemon=True, ) self._thread.start() - - self._started.wait() - with self._lock: - if self._startup_error is not None: - raise self._startup_error - assert self._loop is not None - return self._loop - - def _run_loop(self) -> None: + else: + existing_startup = self._startup + if existing_startup is None: + raise RuntimeError("event loop startup state is unavailable") + startup = existing_startup + + startup.done.wait() + if startup.error is not None: + raise startup.error + if startup.loop is None: + raise RuntimeError("event loop startup completed without a loop") + return startup.loop + + def _run_loop(self, startup: _LoopStartup) -> None: loop = None try: loop = asyncio.new_event_loop() @@ -163,17 +195,20 @@ def run_in_executor(executor, func, *args): if loop is not None and not loop.is_closed(): loop.close() with self._lock: - self._startup_error = error - self._thread = None - self._started.set() + if startup.error is None: + startup.error = error + if self._thread is threading.current_thread(): + self._thread = None + startup.done.set() return with self._lock: self._loop = loop + startup.loop = loop close_requested = self._close_requested - if close_requested and self._startup_error is None: - self._startup_error = RuntimeError("runner closed during startup") - self._started.set() + if close_requested and startup.error is None: + startup.error = RuntimeError("runner closed during startup") + startup.done.set() if close_requested: loop.call_soon(loop.stop) @@ -198,4 +233,6 @@ def run_in_executor(executor, func, *args): self._thread = None if self._loop is loop: self._loop = None + if self._startup is startup: + self._startup = None self._closing_threads.discard(current) diff --git a/posthog/client.py b/posthog/client.py index 1763b0f52..a71a12607 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -2279,10 +2279,10 @@ def _start_lifecycle_thread(self, target, name: str, *args) -> None: target=target, args=args, name=f"posthog-{name}", - daemon=False, + daemon=True, ).start() - def _defer_lifecycle_from_callback(self, name: str) -> bool: + def _defer_lifecycle_from_callback(self) -> bool: if not self._is_lifecycle_callback_thread(): return False with self._lifecycle_lock: @@ -2297,11 +2297,14 @@ def run() -> None: with self._lifecycle_lock: self._deferred_lifecycle_dirty = False require_shutdown = self._shutdown_requested + operation = "shutdown" if require_shutdown else "join" try: self._run_lifecycle(require_shutdown=require_shutdown) except BaseException: attempt += 1 - self.log.exception("Deferred %s attempt %d failed", name, attempt) + self.log.exception( + "Deferred %s attempt %d failed", operation, attempt + ) with self._lifecycle_lock: if self._deferred_lifecycle_dirty: attempt = 0 @@ -2318,7 +2321,7 @@ def run() -> None: self._deferred_lifecycle_thread_pending = False return - self._start_lifecycle_thread(run, name) + self._start_lifecycle_thread(run, "lifecycle") return True def _defer_flush_from_callback(self, timeout_seconds: Optional[float]) -> bool: @@ -2455,14 +2458,17 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: @no_throw() def join(self) -> None: """ - Flush queued events and end the consumer threads. Do not use directly, call `shutdown()` instead. + Attempt to process queued events and end the consumer threads. Do not use directly, call `shutdown()` instead. + + Failed or undrainable events may be dropped and reported through logging + or ``on_error``; returning does not guarantee server receipt. Examples: ```python posthog.join() ``` """ - if self._defer_lifecycle_from_callback("join"): + if self._defer_lifecycle_from_callback(): return self._run_lifecycle() @@ -2471,10 +2477,13 @@ def shutdown(self) -> None: """ Flush all messages and cleanly shutdown the client. Call this before the process ends in serverless environments to avoid data loss. - Normally this method blocks until delivery and cleanup finish. When - called directly from an SDK callback such as ``on_error``, shutdown is - deferred to avoid blocking the worker that invoked the callback. If the - callback must coordinate a blocking shutdown, have it signal an + Normally this method blocks until queued events have been attempted and + cleanup finishes. Failed or undrainable events may be dropped and + reported through logging or ``on_error``; returning does not guarantee + server receipt. When called directly from an SDK callback such as + ``on_error``, shutdown is deferred to avoid blocking the worker that + invoked the callback. If the callback must coordinate a blocking + shutdown, have it signal an application-owned thread and return before that thread calls shutdown. Do not wait inside the callback for another thread or task that calls a lifecycle method. @@ -2488,7 +2497,7 @@ def shutdown(self) -> None: if self._shutdown_complete: return self._shutdown_requested = True - if self._defer_lifecycle_from_callback("shutdown"): + if self._defer_lifecycle_from_callback(): return self._run_lifecycle(require_shutdown=True) diff --git a/posthog/test/test_async_utils.py b/posthog/test/test_async_utils.py index d32942d80..de7371de4 100644 --- a/posthog/test/test_async_utils.py +++ b/posthog/test/test_async_utils.py @@ -4,7 +4,23 @@ import unittest from unittest import mock -from posthog._async_utils import _BackgroundEventLoopRunner +from posthog._async_utils import _BackgroundEventLoopRunner, _LoopStartup + + +class _PausingEvent: + def __init__(self) -> None: + self._completed = threading.Event() + self.waiter_paused = threading.Event() + self.release_waiter = threading.Event() + + def set(self) -> None: + self._completed.set() + + def wait(self, timeout=None) -> bool: + if not self._completed.wait(timeout): + return False + self.waiter_paused.set() + return self.release_waiter.wait(2) class TestBackgroundEventLoopRunner(unittest.TestCase): @@ -19,7 +35,7 @@ def test_startup_error_is_reported(self): with self.assertRaisesRegex(RuntimeError, "startup failed"): runner.run(awaitable) - awaitable.close() + self.assertIsNone(awaitable.cr_frame) def test_close_waits_for_run_during_startup(self): runner = _BackgroundEventLoopRunner() @@ -67,6 +83,150 @@ def run(): self.assertIsNone(runner._thread) self.assertIsNone(runner._loop) + def test_run_retries_when_close_wins_after_startup_completes(self): + runner = _BackgroundEventLoopRunner() + first_startup = _LoopStartup() + first_startup.done = _PausingEvent() # type: ignore[assignment] + second_startup = _LoopStartup() + run_results = [] + run_errors = [] + + async def result(): + return 42 + + def run(): + try: + run_results.append(runner.run(result())) + except BaseException as error: + run_errors.append(error) + + with mock.patch( + "posthog._async_utils._LoopStartup", + side_effect=[first_startup, second_startup], + ): + run_thread = threading.Thread(target=run) + run_thread.start() + self.assertTrue(first_startup.done.waiter_paused.wait(1)) + + runner.close() + first_startup.done.release_waiter.set() + run_thread.join(2) + + runner.close() + self.assertFalse(run_thread.is_alive()) + self.assertEqual(run_errors, []) + self.assertEqual(run_results, [42]) + + def test_startup_failure_state_is_not_overwritten_by_next_attempt(self): + runner = _BackgroundEventLoopRunner() + first_startup = _LoopStartup() + first_startup.done = _PausingEvent() # type: ignore[assignment] + second_startup = _LoopStartup() + original_new_event_loop = asyncio.new_event_loop + loop_attempt = 0 + first_errors = [] + second_results = [] + + def create_loop(): + nonlocal loop_attempt + loop_attempt += 1 + if loop_attempt == 1: + raise RuntimeError("first startup failed") + return original_new_event_loop() + + async def result(): + return 42 + + first_awaitable = asyncio.sleep(0) + + def first_run(): + try: + runner.run(first_awaitable) + except BaseException as error: + first_errors.append(error) + + def second_run(): + second_results.append(runner.run(result())) + + with ( + mock.patch( + "posthog._async_utils._LoopStartup", + side_effect=[first_startup, second_startup], + ), + mock.patch( + "posthog._async_utils.asyncio.new_event_loop", + side_effect=create_loop, + ), + ): + first_thread = threading.Thread(target=first_run) + first_thread.start() + self.assertTrue(first_startup.done.waiter_paused.wait(1)) + + second_thread = threading.Thread(target=second_run) + second_thread.start() + second_thread.join(2) + + first_startup.done.release_waiter.set() + first_thread.join(2) + + runner.close() + self.assertFalse(first_thread.is_alive()) + self.assertFalse(second_thread.is_alive()) + self.assertEqual(len(first_errors), 1) + self.assertRegex(str(first_errors[0]), "first startup failed") + self.assertEqual(second_results, [42]) + self.assertIsNone(first_awaitable.cr_frame) + + def test_close_closes_awaitable_cancelled_before_first_step(self): + runner = _BackgroundEventLoopRunner() + runner.run(asyncio.sleep(0)) + loop = runner._loop + self.assertIsNotNone(loop) + + loop_blocked = threading.Event() + release_loop = threading.Event() + scheduled = threading.Event() + run_errors = [] + original_run_coroutine_threadsafe = asyncio.run_coroutine_threadsafe + + def block_loop(): + loop_blocked.set() + self.assertTrue(release_loop.wait(2)) + + def schedule(coro, target_loop): + future = original_run_coroutine_threadsafe(coro, target_loop) + scheduled.set() + return future + + loop.call_soon_threadsafe(block_loop) # type: ignore[union-attr] + self.assertTrue(loop_blocked.wait(1)) + awaitable = asyncio.sleep(0) + + def run(): + try: + runner.run(awaitable) + except BaseException as error: + run_errors.append(error) + + with mock.patch( + "posthog._async_utils.asyncio.run_coroutine_threadsafe", + side_effect=schedule, + ): + run_thread = threading.Thread(target=run) + run_thread.start() + self.assertTrue(scheduled.wait(1)) + + close_thread = threading.Thread(target=runner.close) + close_thread.start() + release_loop.set() + run_thread.join(2) + close_thread.join(2) + + self.assertFalse(run_thread.is_alive()) + self.assertFalse(close_thread.is_alive()) + self.assertEqual(len(run_errors), 1) + self.assertIsNone(awaitable.cr_frame) + def test_runner_preserves_configured_event_loop_policy(self): runner = _BackgroundEventLoopRunner() original_policy = asyncio.get_event_loop_policy() diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 0caf76712..6ccd47ada 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2298,6 +2298,7 @@ def run_lifecycle(require_shutdown=False): self.assertTrue(shutdown_run_complete.wait(1)) start_thread.assert_called_once() + self.assertEqual(start_thread.call_args.args[1], "lifecycle") self.assertEqual(require_shutdown_calls, [False, True]) self.assertTrue( _wait_until(lambda: not client._deferred_lifecycle_thread_pending) @@ -2343,6 +2344,44 @@ def run_lifecycle(require_shutdown=False): ) self.assertFalse(client._deferred_lifecycle_dirty) + def test_deferred_lifecycle_worker_is_daemon(self): + client = Client(FAKE_TEST_API_KEY) + target = mock.Mock() + + with mock.patch("posthog.client.threading.Thread") as thread: + client._start_lifecycle_thread(target, "lifecycle") + + thread.assert_called_once_with( + target=target, + args=(), + name="posthog-lifecycle", + daemon=True, + ) + thread.return_value.start.assert_called_once_with() + + def test_deferred_lifecycle_logs_selected_operation(self): + client = Client(FAKE_TEST_API_KEY) + + with ( + mock.patch.object( + client, "_is_lifecycle_callback_thread", return_value=True + ), + mock.patch.object( + client, + "_run_lifecycle", + side_effect=[Exception("cleanup failed"), None], + ), + mock.patch.object(client.log, "exception") as log_exception, + ): + client.shutdown() + self.assertTrue( + _wait_until(lambda: not client._deferred_lifecycle_thread_pending) + ) + + log_exception.assert_called_once_with( + "Deferred %s attempt %d failed", "shutdown", 1 + ) + def test_callback_flushes_are_coalesced_with_strongest_followup(self): client = Client(FAKE_TEST_API_KEY) first_flush_started = threading.Event() From 989f7b04f4d91ce81fbf658cb0d5a10a31fb40c0 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:33:00 +0700 Subject: [PATCH 24/27] bound interpreter exit queue draining --- posthog/client.py | 52 +++++++++++-- posthog/consumer.py | 90 +++++++++++++++++------ posthog/test/test_client.py | 87 ++++++++++++++++++++++ posthog/test/test_consumer.py | 135 ++++++++++++++++++++++++++++++++++ 4 files changed, 334 insertions(+), 30 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index a71a12607..157dcc122 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -120,6 +120,17 @@ _configure_posthog_logging() MAX_DICT_SIZE = 50_000 +_ATEXIT_FLUSH_TIMEOUT_SECONDS = 1.0 +_atexit_deadline: Optional[float] = None +_atexit_deadline_lock = threading.Lock() + + +def _get_atexit_deadline() -> float: + global _atexit_deadline + with _atexit_deadline_lock: + if _atexit_deadline is None: + _atexit_deadline = time.monotonic() + _ATEXIT_FLUSH_TIMEOUT_SECONDS + return _atexit_deadline def get_identity_state(passed) -> tuple[str, bool]: @@ -468,12 +479,11 @@ def discard_undrainable_queued_work(self) -> None: def join(self) -> None: """Pause this lane's consumers and wait for them to exit.""" - # Teardown bypasses the batching wait too, so a consumer holding a - # partial batch delivers it instead of exiting `flush_interval` later. + # Normal teardown bypasses the batching wait so a partial batch is sent. self._drain_signal.request() try: for consumer in self.consumers: - consumer.pause() + consumer._pause(drain=True) for consumer in self.consumers: try: consumer.join() @@ -874,10 +884,9 @@ def __init__( # On program exit, allow the consumer threads to exit cleanly. # This prevents exceptions and a messy shutdown when the # interpreter is destroyed before the daemon threads finish - # execution. However, it is *not* the same as flushing the queue! - # To guarantee all messages have been delivered, you'll still need - # to call flush(). - atexit.register(self.join) + # execution. Exit performs only a short best-effort flush; call + # flush() or shutdown() explicitly when blocking completion matters. + atexit.register(self._atexit) lane_defaults = dict( api_key=self.api_key, @@ -2455,6 +2464,35 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: self._lifecycle_condition.notify_all() raise + @no_throw() + def _atexit(self) -> None: + """Make a bounded delivery attempt, then stop daemon workers.""" + with self._lifecycle_condition: + # A daemon lifecycle worker already owns cleanup. Do not wait for it + # at interpreter exit; the process must remain free to terminate. + if self._lifecycle_owner is not None: + return + self._lifecycle_owner = threading.current_thread() + + try: + try: + for lane in self._lanes: + lane.close() + + deadline = _get_atexit_deadline() + for lane in self._lanes: + lane.flush(max(0.0, deadline - time.monotonic())) + finally: + # Consumers are daemon threads. Publish a non-draining stop to + # every consumer, but do not join in-flight requests at exit. + for lane in self._lanes: + for consumer in lane.consumers: + consumer.pause() + finally: + with self._lifecycle_condition: + self._lifecycle_owner = None + self._lifecycle_condition.notify_all() + @no_throw() def join(self) -> None: """ diff --git a/posthog/consumer.py b/posthog/consumer.py index f1afb0208..bc41f4c15 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -51,8 +51,11 @@ def complete(self) -> None: self._requests -= 1 self._queue.not_empty.notify_all() - def wake(self) -> None: + def stop(self, consumer, drain: bool) -> None: + """Publish a consumer stop under the queue's dequeue lock.""" with self._queue.not_empty: + consumer.running = False + consumer._drain_on_stop = drain self._queue.not_empty.notify_all() def wait_until_inactive_or_work(self, consumer) -> None: @@ -65,12 +68,23 @@ def requested(self) -> bool: with self._queue.mutex: return self._requests > 0 - def get(self, timeout: float): - """Get an item, or wake with ``Empty`` when draining an empty queue.""" + def draining(self, consumer) -> bool: + with self._queue.mutex: + return self._requests > 0 and (consumer.running or consumer._drain_on_stop) + + def get(self, timeout: float, consumer=None): + """Get an item, or wake with ``Empty`` when draining or stopping.""" with self._queue.not_empty: deadline = time.monotonic() + timeout - while not self._queue._qsize(): - if self._requests: + while True: + draining = self._requests > 0 and ( + consumer is None or consumer.running or consumer._drain_on_stop + ) + if consumer is not None and not consumer.running and not draining: + raise Empty + if self._queue._qsize(): + break + if draining: raise Empty remaining = deadline - time.monotonic() if remaining <= 0: @@ -119,6 +133,7 @@ def __init__( self.capture_mode = capture_mode self.capture_compression = capture_compression self._drain_signal: Optional[_DrainSignal] = None + self._drain_on_stop = False # It's important to set running in the constructor: if we are asked to # pause immediately after construction, we might set running to True in # run() *after* we set it to False in pause... and keep running @@ -139,10 +154,15 @@ def run(self): self.log.debug("consumer exited.") def pause(self): - """Pause the consumer.""" - self.running = False + """Pause the consumer without admitting additional queued work.""" + self._pause(drain=False) + + def _pause(self, drain: bool) -> None: if self._drain_signal is not None: - self._drain_signal.wake() + self._drain_signal.stop(self, drain) + else: + self.running = False + self._drain_on_stop = drain def upload(self): """Upload the next batch of items, return whether successful.""" @@ -152,16 +172,19 @@ def upload(self): return False try: - self.request(batch) - success = True - except Exception as e: - self.log.error("error uploading: %s", e) - success = False - if self.on_error: - try: - self.on_error(e, batch) - except Exception as e: - self.log.error("on_error handler failed: %s", e) + if not self._can_upload(): + return False + try: + self.request(batch) + success = True + except Exception as e: + self.log.error("error uploading: %s", e) + success = False + if self.on_error: + try: + self.on_error(e, batch) + except Exception as e: + self.log.error("on_error handler failed: %s", e) finally: # mark items as acknowledged from queue for item in batch: @@ -173,7 +196,20 @@ def _set_drain_signal(self, drain_signal: _DrainSignal) -> None: self._drain_signal = drain_signal def _draining(self) -> bool: - return self._drain_signal.requested if self._drain_signal is not None else False + return ( + self._drain_signal.draining(self) + if self._drain_signal is not None + else False + ) + + def _can_upload(self) -> bool: + if self._drain_signal is None: + return self.running or self._drain_on_stop + # This lock-protected check is the request admission point. A request + # admitted before a stop may finish on the daemon consumer, but a stop + # prevents any later buffered or queued batch from being admitted. + with self.queue.mutex: + return self.running or self._drain_on_stop def next(self): """Return the next batch of items to upload.""" @@ -189,15 +225,18 @@ def next(self): # While draining we take only what is already queued, never waiting # for `flush_interval` to elapse or for `flush_at` to be reached. draining = self._draining() + if not self.running and not draining: + break remaining = self.flush_interval - (time.monotonic() - start_time) if not draining and remaining <= 0: break try: - if draining: - item = queue.get(block=False) - elif self._drain_signal is not None: - item = self._drain_signal.get(timeout=remaining) + if self._drain_signal is not None: + item = self._drain_signal.get( + timeout=0 if draining else remaining, + consumer=self, + ) else: item = queue.get(block=True, timeout=remaining) pending_items += 1 @@ -239,6 +278,11 @@ def next(self): queue.task_done() raise + if not self._can_upload(): + for _ in range(pending_items): + queue.task_done() + return [] + return items def request(self, batch): diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 6ccd47ada..026e5ca66 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -1,5 +1,8 @@ import logging import asyncio +import subprocess +import sys +import textwrap import threading import time import unittest @@ -2249,6 +2252,90 @@ def test_shutdown(self): for consumer in client.consumers: self.assertFalse(consumer.is_alive()) + def test_atexit_registers_bounded_worker_cleanup(self): + with mock.patch("posthog.client.atexit.register") as register: + client = Client(FAKE_TEST_API_KEY) + + register.assert_called_once_with(client._atexit) + self.addCleanup(client.shutdown) + + def test_atexit_bounds_flush_and_stops_consumers_without_joining(self): + client = Client(FAKE_TEST_API_KEY, send=False) + lanes = [mock.Mock(), mock.Mock()] + for lane in lanes: + lane.consumers = [mock.Mock(), mock.Mock()] + client._lanes = lanes + + client._atexit() + + for lane in lanes: + lane.close.assert_called_once_with() + lane.flush.assert_called_once() + timeout = lane.flush.call_args.args[0] + self.assertGreaterEqual(timeout, 0) + self.assertLessEqual(timeout, 1) + lane.wait_for_sync_sends.assert_not_called() + lane.join.assert_not_called() + for consumer in lane.consumers: + consumer.pause.assert_called_once_with() + + def test_atexit_subprocess_does_not_drain_queued_backlog(self): + script = textwrap.dedent( + """ + import threading + import time + from unittest import mock + + from posthog.client import Client + from posthog.consumer import Consumer + + with mock.patch.object(Consumer, "start"): + clients = [ + Client(f"test-key-{index}", flush_at=100, flush_interval=60) + for index in range(3) + ] + + for client in clients: + consumer = client.consumers[0] + consumer.request = lambda batch: time.sleep(10) + consumer.start() + time.sleep(0.1) + + queues = [client.queue for client in clients] + for queue in queues: + for index in range(10): + queue.put({"event": str(index), "distinct_id": "test"}) + + deadline = time.monotonic() + 1 + while any(not queue.empty() for queue in queues): + if time.monotonic() >= deadline: + raise RuntimeError("consumer did not buffer the queued backlog") + time.sleep(0.001) + """ + ) + + subprocess.run( + [sys.executable, "-c", script], + check=True, + capture_output=True, + text=True, + timeout=3, + ) + + def test_atexit_does_not_wait_for_active_lifecycle_owner(self): + client = Client(FAKE_TEST_API_KEY, send=False) + owner = mock.Mock() + client._lifecycle_owner = owner + + lane = mock.Mock() + client._lanes = [lane] + + client._atexit() + + lane.close.assert_not_called() + lane.flush.assert_not_called() + self.assertIs(client._lifecycle_owner, owner) + def test_shutdown_clears_feature_flag_called_dedupe_cache(self): client = Client(FAKE_TEST_API_KEY, send=False, thread=0) client.distinct_ids_feature_flags_reported["user"] = {("flag", True, ())} diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index 4a636fd47..3782ad37f 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -34,6 +34,141 @@ def test_next(self) -> None: next = consumer.next() self.assertEqual(next, [1]) + def test_next_does_not_take_queued_items_after_non_draining_pause(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100) + drain_signal = _DrainSignal(q) + consumer._set_drain_signal(drain_signal) + for item in range(10): + q.put(item) + + consumer.pause() + + self.assertEqual(consumer.next(), []) + self.assertEqual(q.qsize(), 10) + self.assertEqual(q.unfinished_tasks, 10) + + def test_non_draining_pause_overrides_active_flush_signal(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100) + drain_signal = _DrainSignal(q) + consumer._set_drain_signal(drain_signal) + q.put(_track_event()) + + drain_signal.request() + consumer.pause() + try: + self.assertEqual(consumer.next(), []) + self.assertEqual(q.qsize(), 1) + self.assertEqual(q.unfinished_tasks, 1) + finally: + drain_signal.complete() + + def test_non_draining_pause_between_drain_snapshot_and_dequeue(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100) + drain_signal = _DrainSignal(q) + consumer._set_drain_signal(drain_signal) + q.put(_track_event()) + drain_signal.request() + original_get = drain_signal.get + + def pause_then_get(*args, **kwargs): + consumer.pause() + return original_get(*args, **kwargs) + + with mock.patch.object(drain_signal, "get", side_effect=pause_then_get): + self.assertEqual(consumer.next(), []) + + drain_signal.complete() + self.assertEqual(q.qsize(), 1) + self.assertEqual(q.unfinished_tasks, 1) + + def test_pause_publishes_stop_under_queue_dequeue_lock(self) -> None: + q = Queue() + consumer = Consumer(q, "") + drain_signal = _DrainSignal(q) + stop_started = threading.Event() + original_stop = drain_signal.stop + + def observed_stop(target, drain): + stop_started.set() + original_stop(target, drain) + + drain_signal.stop = observed_stop # type: ignore[method-assign] + consumer._set_drain_signal(drain_signal) + + with q.mutex: + pause_thread = threading.Thread(target=consumer.pause) + pause_thread.start() + self.assertTrue(stop_started.wait(1)) + self.assertTrue(consumer.running) + + pause_thread.join(1) + self.assertFalse(pause_thread.is_alive()) + self.assertFalse(consumer.running) + + def test_non_draining_pause_discards_buffered_partial_batch(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100, flush_interval=60) + consumer._set_drain_signal(_DrainSignal(q)) + request_called = threading.Event() + consumer.request = lambda batch: request_called.set() # type: ignore[method-assign] + consumer.start() + q.put(_track_event()) + + deadline = time.monotonic() + 1 + while not q.empty(): + if time.monotonic() >= deadline: + self.fail("consumer did not buffer the queued event") + time.sleep(0.001) + + consumer.pause() + consumer.join(1) + + self.assertFalse(consumer.is_alive()) + self.assertFalse(request_called.is_set()) + self.assertEqual(q.unfinished_tasks, 0) + + def test_pause_does_not_wait_for_active_request(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=1) + consumer._set_drain_signal(_DrainSignal(q)) + request_started = threading.Event() + release_request = threading.Event() + + def request(batch): + request_started.set() + self.assertTrue(release_request.wait(2)) + + consumer.request = request # type: ignore[method-assign] + consumer.start() + q.put(_track_event()) + self.assertTrue(request_started.wait(1)) + + consumer.pause() + + self.assertFalse(consumer.running) + self.assertTrue(consumer.is_alive()) + release_request.set() + consumer.join(1) + self.assertFalse(consumer.is_alive()) + + def test_next_still_takes_queued_items_when_paused_for_drain(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100) + drain_signal = _DrainSignal(q) + consumer._set_drain_signal(drain_signal) + for item in range(10): + q.put(item) + + drain_signal.request() + consumer._pause(drain=True) + try: + self.assertEqual(consumer.next(), list(range(10))) + finally: + drain_signal.complete() + def test_next_limit(self) -> None: q = Queue() flush_at = 50 From 693cb50c6b0d102cc7ee2de9b5b1136588692831 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto <5731772+marandaneto@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:05:11 +0700 Subject: [PATCH 25/27] preserve read-only custom event loops --- posthog/_async_utils.py | 22 ++++------------------ posthog/test/test_async_utils.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 18 deletions(-) diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 0ee6119ce..85a653ce0 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,6 +1,5 @@ import asyncio import inspect -import sys import threading from collections.abc import Awaitable from contextvars import Context, copy_context @@ -35,18 +34,6 @@ def __reduce__(self): return (_PlainExecutorCall, (self._func, self._args, self._kwargs)) -if sys.platform == "win32": - from asyncio.windows_events import ProactorEventLoop as _PlatformEventLoop -else: - _PlatformEventLoop = asyncio.SelectorEventLoop - - -class _ContextEventLoop(_PlatformEventLoop): - def run_in_executor(self, executor, func, *args): # type: ignore[override] - call = _ContextExecutorCall(copy_context(), func, args) - return super().run_in_executor(executor, call) - - class _LoopStartup: def __init__(self) -> None: self.done = threading.Event() @@ -185,11 +172,10 @@ def run_in_executor(executor, func, *args): try: setattr(loop, "run_in_executor", run_in_executor) except (AttributeError, TypeError): - # Preserve custom policy loops when they are extensible. For a - # read-only implementation, use the platform-equivalent loop - # so explicit executors retain callback context safely. - loop.close() - loop = _ContextEventLoop() + # Some policy-provided loops expose a read-only implementation. + # Preserve that loop rather than silently replacing its policy + # semantics; executor context propagation is best-effort there. + pass asyncio.set_event_loop(loop) except BaseException as error: if loop is not None and not loop.is_closed(): diff --git a/posthog/test/test_async_utils.py b/posthog/test/test_async_utils.py index de7371de4..76b3886db 100644 --- a/posthog/test/test_async_utils.py +++ b/posthog/test/test_async_utils.py @@ -250,6 +250,35 @@ async def running_loop(): finally: asyncio.set_event_loop_policy(original_policy) + def test_runner_preserves_read_only_policy_loop(self): + runner = _BackgroundEventLoopRunner() + original_policy = asyncio.get_event_loop_policy() + + class ReadOnlyLoop(asyncio.SelectorEventLoop): + def __setattr__(self, name, value): + if name == "run_in_executor": + raise AttributeError("run_in_executor is read-only") + super().__setattr__(name, value) + + class Policy(asyncio.DefaultEventLoopPolicy): + loop = None + + def new_event_loop(self): + self.loop = ReadOnlyLoop() + return self.loop + + policy = Policy() + + async def running_loop(): + return asyncio.get_running_loop() + + try: + asyncio.set_event_loop_policy(policy) + self.assertIs(runner.run(running_loop()), policy.loop) + runner.close() + finally: + asyncio.set_event_loop_policy(original_policy) + def test_run_from_runner_thread_fails_instead_of_deadlocking(self): runner = _BackgroundEventLoopRunner() From fba9502a77737cfbbb77224f55412d61ff7f5207 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Thu, 6 Aug 2026 15:17:23 +0200 Subject: [PATCH 26/27] fix: harden lifecycle cleanup completion --- posthog/__init__.py | 10 +- posthog/_async_utils.py | 22 +++- posthog/client.py | 159 ++++++++++++++++++++--------- posthog/test/test_async_utils.py | 8 +- posthog/test/test_client.py | 169 ++++++++++++++++++++++++++----- posthog/test/test_client_fork.py | 18 ++++ 6 files changed, 305 insertions(+), 81 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index c28802f0e..6ef2449b5 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -1129,7 +1129,8 @@ def join() -> None: Attempt to process queued events and stop the client's background workers. Use `shutdown()` directly in most cases. Failed or undrainable events may be dropped and reported through logging or - ``on_error``; returning does not guarantee server receipt. + ``on_error``; returning does not guarantee server receipt. Lifecycle cleanup + is attempted once, and cleanup failures are logged without retry. Examples: ```python @@ -1149,9 +1150,10 @@ def shutdown() -> None: This normally blocks until queued events have been attempted and cleanup finishes. Failed or undrainable events may be dropped and reported through - logging or ``on_error``; returning does not guarantee server receipt. Calls - made directly from SDK callbacks such as ``on_error`` are deferred to avoid - deadlocking the worker. If blocking completion is required, signal an application-owned + logging or ``on_error``; returning does not guarantee server receipt. + Lifecycle cleanup is attempted once, and cleanup failures are logged without + retry. Calls made directly from SDK callbacks such as ``on_error`` are deferred + to avoid deadlocking the worker. If blocking completion is required, signal an application-owned thread, return from the callback, and call ``shutdown()`` from that thread. Do not wait inside a callback for another thread or task calling a lifecycle method. diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 85a653ce0..f81d6964b 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,5 +1,6 @@ import asyncio import inspect +import sys import threading from collections.abc import Awaitable from contextvars import Context, copy_context @@ -34,6 +35,18 @@ def __reduce__(self): return (_PlainExecutorCall, (self._func, self._args, self._kwargs)) +if sys.platform == "win32": + from asyncio.windows_events import ProactorEventLoop as _PlatformEventLoop +else: + _PlatformEventLoop = asyncio.SelectorEventLoop + + +class _ContextEventLoop(_PlatformEventLoop): + def run_in_executor(self, executor, func, *args): # type: ignore[override] + call = _ContextExecutorCall(copy_context(), func, args) + return super().run_in_executor(executor, call) + + class _LoopStartup: def __init__(self) -> None: self.done = threading.Event() @@ -172,10 +185,11 @@ def run_in_executor(executor, func, *args): try: setattr(loop, "run_in_executor", run_in_executor) except (AttributeError, TypeError): - # Some policy-provided loops expose a read-only implementation. - # Preserve that loop rather than silently replacing its policy - # semantics; executor context propagation is best-effort there. - pass + # A loop that cannot carry callback context is unsafe for + # lifecycle re-entry from executor threads. Use the equivalent + # context-aware platform loop rather than silently losing it. + loop.close() + loop = _ContextEventLoop() asyncio.set_event_loop(loop) except BaseException as error: if loop is not None and not loop.is_closed(): diff --git a/posthog/client.py b/posthog/client.py index 157dcc122..aa2ea3703 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -10,7 +10,7 @@ import weakref from contextvars import ContextVar from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Mapping, Optional, Union +from typing import Any, Callable, Dict, List, Mapping, Optional, Union from uuid import UUID, uuid4 from typing_extensions import Unpack @@ -483,13 +483,31 @@ def join(self) -> None: self._drain_signal.request() try: for consumer in self.consumers: - consumer._pause(drain=True) + try: + consumer._pause(drain=True) + except Exception: + self.log.exception( + "Failed to pause %s lane consumer during lifecycle cleanup", + self.name, + ) for consumer in self.consumers: try: consumer.join() except RuntimeError: # consumer thread has not started pass + except Exception: + self.log.exception( + "Failed to join %s lane consumer during lifecycle cleanup", + self.name, + ) + try: + self.discard_undrainable_queued_work() + except Exception: + self.log.exception( + "Failed to discard queued %s lane work during lifecycle cleanup", + self.name, + ) finally: self._drain_signal.complete() @@ -499,20 +517,22 @@ def reset_sync_send_state_after_fork(self) -> None: self._start_lock = threading.Lock() self._sync_sends_done = threading.Condition(self._start_lock) - def rebuild_after_fork(self) -> None: + def rebuild_after_fork(self, *, closed: bool) -> None: """Replace fork-unsafe lane state in a forked child. Threads do not survive fork() and queue.Queue internal locks may be in an inconsistent state, so the queue, lock, and consumer pool are replaced. Inherited queue items are not retained as they'll be handled - by the parent process's consumers. An eager lane restarts immediately; - a lazy lane returns to not-started and restarts on next use. + by the parent process's consumers. ``closed`` normalizes every lane to + the client's fork-visible lifecycle state. An eager open lane restarts + immediately; a lazy lane returns to not-started and restarts on next use. """ self.queue = Queue(self._max_queue_size) self.reset_sync_send_state_after_fork() self._drain_signal = _DrainSignal(self.queue) self.consumers = [] self._started = False + self._closed = closed if self._eager_start: self.start() @@ -718,6 +738,7 @@ def __init__( self._lifecycle_owner: Optional[threading.Thread] = None self._workers_joined = False self._join_cleanup_complete = False + self._join_requested = False self._shutdown_requested = False self._shutdown_complete = False self._deferred_lifecycle_thread_pending = False @@ -2007,8 +2028,11 @@ def _reinit_after_fork(self): Python threads do not survive fork(), so each lane's queue and consumer pool are rebuilt (see `_Lane.rebuild_after_fork`). """ + terminal_requested = ( + self._join_requested or self._shutdown_requested or self._workers_joined + ) for lane in self._lanes: - lane.rebuild_after_fork() + lane.rebuild_after_fork(closed=terminal_requested) self._lifecycle_lock = threading.Lock() self._lifecycle_condition = threading.Condition(self._lifecycle_lock) @@ -2046,7 +2070,7 @@ def _reinit_after_fork(self): reset_sessions() # Start child threads only after replacing every lock they can touch. - if self._workers_joined: + if terminal_requested: self.poller = None elif self.enable_local_evaluation: self.poller = Poller( @@ -2301,7 +2325,6 @@ def _defer_lifecycle_from_callback(self) -> bool: self._deferred_lifecycle_thread_pending = True def run() -> None: - attempt = 0 while True: with self._lifecycle_lock: self._deferred_lifecycle_dirty = False @@ -2310,22 +2333,10 @@ def run() -> None: try: self._run_lifecycle(require_shutdown=require_shutdown) except BaseException: - attempt += 1 - self.log.exception( - "Deferred %s attempt %d failed", operation, attempt - ) - with self._lifecycle_lock: - if self._deferred_lifecycle_dirty: - attempt = 0 - continue - if attempt < 2: - continue - self._deferred_lifecycle_thread_pending = False - return + self.log.exception("Deferred %s failed", operation) with self._lifecycle_lock: if self._deferred_lifecycle_dirty: - attempt = 0 continue self._deferred_lifecycle_thread_pending = False return @@ -2369,32 +2380,68 @@ def run() -> None: self._start_lifecycle_thread(run, "flush") return True + def _run_lifecycle_cleanup( + self, log_message: str, cleanup: Callable[[], None] + ) -> None: + """Attempt one cleanup step without preventing later independent steps.""" + try: + cleanup() + except Exception: + self.log.exception(log_message) + def _flush_or_discard_queues(self) -> None: for lane in self._lanes: - if any(consumer.is_alive() for consumer in lane.consumers): - lane.flush(timeout_seconds=None) - else: - lane.discard_undrainable_queued_work() + try: + if any(consumer.is_alive() for consumer in lane.consumers): + lane.flush(timeout_seconds=None) + else: + lane.discard_undrainable_queued_work() + except Exception: + self.log.exception( + "Failed to drain %s lane during lifecycle cleanup", lane.name + ) - def _join_once(self, flush_queues: bool = True) -> None: + def _join_once( + self, flush_queues: bool = True, *, lanes_prepared: bool = False + ) -> None: if not self._workers_joined: - for lane in self._lanes: - lane.close() - for lane in self._lanes: - lane.wait_for_sync_sends() + if not lanes_prepared: + for lane in self._lanes: + self._run_lifecycle_cleanup( + f"Failed to close {lane.name} lane during lifecycle cleanup", + lane.close, + ) + for lane in self._lanes: + self._run_lifecycle_cleanup( + f"Failed waiting for {lane.name} synchronous sends during lifecycle cleanup", + lane.wait_for_sync_sends, + ) if flush_queues: self._flush_or_discard_queues() for lane in self._lanes: - lane.join() + self._run_lifecycle_cleanup( + f"Failed to stop {lane.name} lane during lifecycle cleanup", + lane.join, + ) + # Ordinary cleanup failures are logged by each step. Reaching here + # means every worker cleanup step was attempted once. self._workers_joined = True if not self._join_cleanup_complete: if self.poller: - self.poller.stop() + self._run_lifecycle_cleanup( + "Failed to stop feature flag poller during lifecycle cleanup", + self.poller.stop, + ) - # Shutdown the cache provider (release locks, cleanup) - self._shutdown_flag_definition_cache_provider() - self._unregister_duplicate_client() + self._run_lifecycle_cleanup( + "Failed to shut down feature flag cache provider during lifecycle cleanup", + self._shutdown_flag_definition_cache_provider, + ) + self._run_lifecycle_cleanup( + "Failed to unregister client during lifecycle cleanup", + self._unregister_duplicate_client, + ) self._join_cleanup_complete = True def _shutdown_once(self) -> None: @@ -2402,22 +2449,36 @@ def _shutdown_once(self) -> None: # Close every lane before draining any of them so no producer can be # admitted between a completed flush and consumer shutdown. for lane in self._lanes: - lane.close() + self._run_lifecycle_cleanup( + f"Failed to close {lane.name} lane during shutdown", lane.close + ) for lane in self._lanes: - lane.wait_for_sync_sends() + self._run_lifecycle_cleanup( + f"Failed waiting for {lane.name} synchronous sends during shutdown", + lane.wait_for_sync_sends, + ) self._flush_or_discard_queues() if self._metrics is not None: - try: - self._metrics.flush() - except Exception: - self.log.exception("Failed to flush metrics on shutdown") - self._metrics.reset() - self._join_once(flush_queues=False) - self.distinct_ids_feature_flags_reported.clear() + self._run_lifecycle_cleanup( + "Failed to flush metrics on shutdown", self._metrics.flush + ) + self._run_lifecycle_cleanup( + "Failed to reset metrics on shutdown", self._metrics.reset + ) + self._join_once(flush_queues=False, lanes_prepared=True) + self._run_lifecycle_cleanup( + "Failed to clear feature flag deduplication state on shutdown", + self.distinct_ids_feature_flags_reported.clear, + ) if self.exception_capture: - self.exception_capture.close() + self._run_lifecycle_cleanup( + "Failed to close exception capture on shutdown", + self.exception_capture.close, + ) + # Ordinary cleanup failures are logged by each step. Reaching here + # means every shutdown cleanup step was attempted once. self._shutdown_complete = True def _run_lifecycle(self, require_shutdown: bool = False) -> None: @@ -2499,13 +2560,16 @@ def join(self) -> None: Attempt to process queued events and end the consumer threads. Do not use directly, call `shutdown()` instead. Failed or undrainable events may be dropped and reported through logging - or ``on_error``; returning does not guarantee server receipt. + or ``on_error``; returning does not guarantee server receipt. Lifecycle + cleanup is attempted once, and cleanup failures are logged without retry. Examples: ```python posthog.join() ``` """ + with self._lifecycle_lock: + self._join_requested = True if self._defer_lifecycle_from_callback(): return self._run_lifecycle() @@ -2518,7 +2582,8 @@ def shutdown(self) -> None: Normally this method blocks until queued events have been attempted and cleanup finishes. Failed or undrainable events may be dropped and reported through logging or ``on_error``; returning does not guarantee - server receipt. When called directly from an SDK callback such as + server receipt. Lifecycle cleanup is attempted once, and cleanup failures + are logged without retry. When called directly from an SDK callback such as ``on_error``, shutdown is deferred to avoid blocking the worker that invoked the callback. If the callback must coordinate a blocking shutdown, have it signal an diff --git a/posthog/test/test_async_utils.py b/posthog/test/test_async_utils.py index 76b3886db..1eb7c7837 100644 --- a/posthog/test/test_async_utils.py +++ b/posthog/test/test_async_utils.py @@ -250,7 +250,7 @@ async def running_loop(): finally: asyncio.set_event_loop_policy(original_policy) - def test_runner_preserves_read_only_policy_loop(self): + def test_runner_uses_context_aware_fallback_for_read_only_policy_loop(self): runner = _BackgroundEventLoopRunner() original_policy = asyncio.get_event_loop_policy() @@ -274,9 +274,11 @@ async def running_loop(): try: asyncio.set_event_loop_policy(policy) - self.assertIs(runner.run(running_loop()), policy.loop) - runner.close() + running = runner.run(running_loop()) + self.assertIsNot(running, policy.loop) + self.assertTrue(policy.loop.is_closed()) finally: + runner.close() asyncio.set_event_loop_policy(original_policy) def test_run_from_runner_thread_fails_instead_of_deadlocking(self): diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 026e5ca66..9b7e03437 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -1,3 +1,4 @@ +import contextlib import logging import asyncio import subprocess @@ -2456,8 +2457,8 @@ def test_deferred_lifecycle_logs_selected_operation(self): mock.patch.object( client, "_run_lifecycle", - side_effect=[Exception("cleanup failed"), None], - ), + side_effect=Exception("cleanup failed"), + ) as run_lifecycle, mock.patch.object(client.log, "exception") as log_exception, ): client.shutdown() @@ -2465,9 +2466,8 @@ def test_deferred_lifecycle_logs_selected_operation(self): _wait_until(lambda: not client._deferred_lifecycle_thread_pending) ) - log_exception.assert_called_once_with( - "Deferred %s attempt %d failed", "shutdown", 1 - ) + run_lifecycle.assert_called_once_with(require_shutdown=True) + log_exception.assert_called_once_with("Deferred %s failed", "shutdown") def test_callback_flushes_are_coalesced_with_strongest_followup(self): client = Client(FAKE_TEST_API_KEY) @@ -2544,7 +2544,7 @@ def on_error(error, batch): max_retries=0, ) exception_capture = mock.Mock() - exception_capture.close.side_effect = [Exception("cleanup failed"), None] + exception_capture.close.side_effect = Exception("cleanup failed") client.exception_capture = exception_capture with mock.patch.object(client.consumers[0], "request", side_effect=request): client.capture("first", distinct_id="distinct_id") @@ -2557,7 +2557,7 @@ def on_error(error, batch): self.assertEqual(sent_events, ["first", "second"]) self.assertEqual(client.queue.unfinished_tasks, 0) - self.assertEqual(exception_capture.close.call_count, 2) + exception_capture.close.assert_called_once_with() self.assertTrue(all(not consumer.is_alive() for consumer in client.consumers)) def test_concurrent_join_waits_for_lifecycle_owner(self): @@ -2671,6 +2671,21 @@ def observed_close(): self.assertEqual(sent_events, ["first", "second"]) self.assertEqual(client.queue.unfinished_tasks, 0) + def test_join_publishes_terminal_intent_before_closing_lanes(self): + client = Client(FAKE_TEST_API_KEY, send=False) + original_close = client._analytics_lane.close + + def close_analytics_lane(): + self.assertTrue(client._join_requested) + original_close() + + with mock.patch.object( + client._analytics_lane, "close", side_effect=close_analytics_lane + ): + client.join() + + self.assertTrue(client._join_cleanup_complete) + def test_shutdown_after_join_runs_shutdown_only_cleanup(self): client = Client(FAKE_TEST_API_KEY) metrics = mock.Mock() @@ -2741,6 +2756,38 @@ def reenter_join(): self.assertTrue(executor_called.is_set()) self.assertTrue(client._join_cleanup_complete) + def test_async_cache_provider_read_only_loop_falls_back_without_deadlocking(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + executor_called = threading.Event() + + class ReadOnlyLoop(asyncio.SelectorEventLoop): + def __setattr__(self, name, value): + if name == "run_in_executor": + raise AttributeError("run_in_executor is read-only") + super().__setattr__(name, value) + + class AsyncProvider: + async def shutdown(self): + def reenter_join(): + executor_called.set() + client.join() + + await asyncio.get_running_loop().run_in_executor(None, reenter_join) + + loop = ReadOnlyLoop() + client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] + with mock.patch( + "posthog._async_utils.asyncio.new_event_loop", return_value=loop + ): + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(executor_called.is_set()) + self.assertTrue(loop.is_closed()) + self.assertTrue(client._join_cleanup_complete) + def test_async_cache_provider_process_executor_remains_supported(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) @@ -2858,25 +2905,26 @@ def assert_stopped(self): self.assertTrue(client._shutdown_complete) - def test_join_retries_auxiliary_cleanup_after_failure(self): + def test_join_failure_does_not_retry_or_skip_later_cleanup(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) - with mock.patch.object( - client, - "_shutdown_flag_definition_cache_provider", - side_effect=[Exception("cleanup failed"), None], - ) as cleanup: + with ( + mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=Exception("cleanup failed"), + ) as cleanup, + mock.patch.object(client, "_unregister_duplicate_client") as unregister, + ): client.join() - self.assertTrue(client._workers_joined) - self.assertFalse(client._join_cleanup_complete) - client.join() self.assertTrue(client._workers_joined) self.assertTrue(client._join_cleanup_complete) - self.assertEqual(cleanup.call_count, 2) + cleanup.assert_called_once_with() + unregister.assert_called_once_with() - def test_pending_shutdown_is_retried_when_join_cleanup_fails(self): + def test_pending_shutdown_continues_when_join_cleanup_fails(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) cleanup_started = threading.Event() release_cleanup = threading.Event() @@ -2918,22 +2966,97 @@ def run_join(): self.assertFalse(join_thread.is_alive()) self.assertFalse(shutdown_thread.is_alive()) self.assertEqual(join_errors, []) - self.assertEqual(cleanup_calls, 2) + self.assertEqual(cleanup_calls, 1) self.assertTrue(client._shutdown_complete) - def test_shutdown_retries_cleanup_before_publishing_completion(self): + def test_join_interruption_does_not_publish_completion(self): + client = Client(FAKE_TEST_API_KEY, send=False) + + with ( + mock.patch.object( + client, "_flush_or_discard_queues", side_effect=KeyboardInterrupt + ), + self.assertRaises(KeyboardInterrupt), + ): + client.join() + + self.assertFalse(client._workers_joined) + self.assertFalse(client._join_cleanup_complete) + self.assertIsNone(client._lifecycle_owner) + + def test_shutdown_interruption_does_not_publish_completion(self): + client = Client(FAKE_TEST_API_KEY, send=False) + metrics = mock.Mock() + metrics.flush.side_effect = KeyboardInterrupt + client._metrics = metrics + + with self.assertRaises(KeyboardInterrupt): + client.shutdown() + + self.assertFalse(client._workers_joined) + self.assertFalse(client._join_cleanup_complete) + self.assertFalse(client._shutdown_complete) + self.assertIsNone(client._lifecycle_owner) + + def test_shutdown_failure_is_terminal_and_not_retried(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) exception_capture = mock.Mock() - exception_capture.close.side_effect = [Exception("cleanup failed"), None] + exception_capture.close.side_effect = Exception("cleanup failed") client.exception_capture = exception_capture client.shutdown() - self.assertFalse(client._shutdown_complete) + client.shutdown() + + self.assertTrue(client._shutdown_complete) + exception_capture.close.assert_called_once_with() + + def test_shutdown_failure_does_not_skip_later_cleanup(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01, debug=True) + metrics = mock.Mock() + metrics.reset.side_effect = Exception("reset failed") + dedupe_cache = mock.Mock() + dedupe_cache.clear.side_effect = Exception("clear failed") + exception_capture = mock.Mock() + client._metrics = metrics + client.distinct_ids_feature_flags_reported = dedupe_cache + client.exception_capture = exception_capture client.shutdown() + metrics.flush.assert_called_once_with() + metrics.reset.assert_called_once_with() + dedupe_cache.clear.assert_called_once_with() + exception_capture.close.assert_called_once_with() + self.assertTrue(client._workers_joined) + self.assertTrue(client._join_cleanup_complete) self.assertTrue(client._shutdown_complete) - self.assertEqual(exception_capture.close.call_count, 2) + + def test_shutdown_prepares_each_lane_once(self): + client = Client(FAKE_TEST_API_KEY, send=False) + close_methods = [] + wait_methods = [] + + with contextlib.ExitStack() as stack: + for lane in client._lanes: + close_methods.append( + stack.enter_context( + mock.patch.object(lane, "close", wraps=lane.close) + ) + ) + wait_methods.append( + stack.enter_context( + mock.patch.object( + lane, + "wait_for_sync_sends", + wraps=lane.wait_for_sync_sends, + ) + ) + ) + client.shutdown() + + for close, wait in zip(close_methods, wait_methods): + close.assert_called_once_with() + wait.assert_called_once_with() def test_shutdown_does_not_wait_for_idle_consumers_flush_interval(self): client = Client(FAKE_TEST_API_KEY, flush_interval=5) diff --git a/posthog/test/test_client_fork.py b/posthog/test/test_client_fork.py index 01d900e93..16a64f7b4 100644 --- a/posthog/test/test_client_fork.py +++ b/posthog/test/test_client_fork.py @@ -222,6 +222,24 @@ def test_reinit_after_fork_preserves_terminal_client_state(self): mock_poller.assert_not_called() self.assertIsNone(client.capture("after join", distinct_id="distinct_id")) + def test_reinit_after_fork_normalizes_partially_closed_join_state(self): + client = Client(FAKE_TEST_API_KEY, send=False) + client._join_requested = True + client._analytics_lane._closed = True + client._ai_lane._closed = False + + with mock.patch("posthog.client.Poller") as mock_poller: + client._reinit_after_fork() + + self.assertFalse(client._workers_joined) + self.assertTrue(client._analytics_lane._closed) + self.assertTrue(client._ai_lane._closed) + self.assertEqual(client.consumers, []) + self.assertIsNone(client.poller) + mock_poller.assert_not_called() + self.assertIsNone(client.capture("analytics", distinct_id="distinct_id")) + self.assertIsNone(client._capture_ai("ai", distinct_id="distinct_id")) + @unittest.skipUnless( hasattr(os, "fork") and hasattr(os, "register_at_fork"), From e61a95b8824a6b9b3113d6bf188d420cacbb696d Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Thu, 6 Aug 2026 16:04:48 +0200 Subject: [PATCH 27/27] fix: propagate lifecycle cleanup failures --- posthog/client.py | 131 +++++++++++++++++++++++++----------- posthog/test/test_client.py | 74 +++++++++++++++++++- 2 files changed, 164 insertions(+), 41 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index aa2ea3703..b6cbbf817 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -480,36 +480,58 @@ def discard_undrainable_queued_work(self) -> None: def join(self) -> None: """Pause this lane's consumers and wait for them to exit.""" # Normal teardown bypasses the batching wait so a partial batch is sent. - self._drain_signal.request() + errors: list[Exception] = [] + drain_requested = False try: - for consumer in self.consumers: - try: - consumer._pause(drain=True) - except Exception: - self.log.exception( - "Failed to pause %s lane consumer during lifecycle cleanup", - self.name, - ) - for consumer in self.consumers: - try: - consumer.join() - except RuntimeError: - # consumer thread has not started - pass - except Exception: - self.log.exception( - "Failed to join %s lane consumer during lifecycle cleanup", - self.name, - ) + self._drain_signal.request() + drain_requested = True + except Exception as error: + self.log.exception( + "Failed to request %s lane drain during lifecycle cleanup", self.name + ) + errors.append(error) + + for consumer in self.consumers: try: - self.discard_undrainable_queued_work() - except Exception: + consumer._pause(drain=True) + except Exception as error: self.log.exception( - "Failed to discard queued %s lane work during lifecycle cleanup", + "Failed to pause %s lane consumer during lifecycle cleanup", self.name, ) - finally: - self._drain_signal.complete() + errors.append(error) + for consumer in self.consumers: + try: + consumer.join() + except RuntimeError: + # consumer thread has not started + pass + except Exception as error: + self.log.exception( + "Failed to join %s lane consumer during lifecycle cleanup", + self.name, + ) + errors.append(error) + try: + self.discard_undrainable_queued_work() + except Exception as error: + self.log.exception( + "Failed to discard queued %s lane work during lifecycle cleanup", + self.name, + ) + errors.append(error) + if drain_requested: + try: + self._drain_signal.complete() + except Exception as error: + self.log.exception( + "Failed to complete %s lane drain during lifecycle cleanup", + self.name, + ) + errors.append(error) + + if errors: + raise errors[0] def reset_sync_send_state_after_fork(self) -> None: """Replace sync-send state inherited from threads that did not survive fork.""" @@ -741,6 +763,7 @@ def __init__( self._join_requested = False self._shutdown_requested = False self._shutdown_complete = False + self._lifecycle_cleanup_failed = False self._deferred_lifecycle_thread_pending = False self._deferred_lifecycle_dirty = False self._lifecycle_callback_context: ContextVar[bool] = ContextVar( @@ -2381,28 +2404,37 @@ def run() -> None: return True def _run_lifecycle_cleanup( - self, log_message: str, cleanup: Callable[[], None] + self, + log_message: str, + cleanup: Callable[[], None], + errors: list[Exception], ) -> None: """Attempt one cleanup step without preventing later independent steps.""" try: cleanup() - except Exception: + except Exception as error: self.log.exception(log_message) + errors.append(error) - def _flush_or_discard_queues(self) -> None: + def _flush_or_discard_queues(self, errors: list[Exception]) -> None: for lane in self._lanes: try: if any(consumer.is_alive() for consumer in lane.consumers): lane.flush(timeout_seconds=None) else: lane.discard_undrainable_queued_work() - except Exception: + except Exception as error: self.log.exception( "Failed to drain %s lane during lifecycle cleanup", lane.name ) + errors.append(error) def _join_once( - self, flush_queues: bool = True, *, lanes_prepared: bool = False + self, + errors: list[Exception], + flush_queues: bool = True, + *, + lanes_prepared: bool = False, ) -> None: if not self._workers_joined: if not lanes_prepared: @@ -2410,18 +2442,21 @@ def _join_once( self._run_lifecycle_cleanup( f"Failed to close {lane.name} lane during lifecycle cleanup", lane.close, + errors, ) for lane in self._lanes: self._run_lifecycle_cleanup( f"Failed waiting for {lane.name} synchronous sends during lifecycle cleanup", lane.wait_for_sync_sends, + errors, ) if flush_queues: - self._flush_or_discard_queues() + self._flush_or_discard_queues(errors) for lane in self._lanes: self._run_lifecycle_cleanup( f"Failed to stop {lane.name} lane during lifecycle cleanup", lane.join, + errors, ) # Ordinary cleanup failures are logged by each step. Reaching here # means every worker cleanup step was attempted once. @@ -2432,50 +2467,58 @@ def _join_once( self._run_lifecycle_cleanup( "Failed to stop feature flag poller during lifecycle cleanup", self.poller.stop, + errors, ) self._run_lifecycle_cleanup( "Failed to shut down feature flag cache provider during lifecycle cleanup", self._shutdown_flag_definition_cache_provider, + errors, ) self._run_lifecycle_cleanup( "Failed to unregister client during lifecycle cleanup", self._unregister_duplicate_client, + errors, ) self._join_cleanup_complete = True - def _shutdown_once(self) -> None: + def _shutdown_once(self, errors: list[Exception]) -> None: if not self._workers_joined: # Close every lane before draining any of them so no producer can be # admitted between a completed flush and consumer shutdown. for lane in self._lanes: self._run_lifecycle_cleanup( - f"Failed to close {lane.name} lane during shutdown", lane.close + f"Failed to close {lane.name} lane during shutdown", + lane.close, + errors, ) for lane in self._lanes: self._run_lifecycle_cleanup( f"Failed waiting for {lane.name} synchronous sends during shutdown", lane.wait_for_sync_sends, + errors, ) - self._flush_or_discard_queues() + self._flush_or_discard_queues(errors) if self._metrics is not None: self._run_lifecycle_cleanup( - "Failed to flush metrics on shutdown", self._metrics.flush + "Failed to flush metrics on shutdown", self._metrics.flush, errors ) self._run_lifecycle_cleanup( - "Failed to reset metrics on shutdown", self._metrics.reset + "Failed to reset metrics on shutdown", self._metrics.reset, errors ) - self._join_once(flush_queues=False, lanes_prepared=True) + self._join_once(errors, flush_queues=False, lanes_prepared=True) self._run_lifecycle_cleanup( "Failed to clear feature flag deduplication state on shutdown", self.distinct_ids_feature_flags_reported.clear, + errors, ) if self.exception_capture: self._run_lifecycle_cleanup( "Failed to close exception capture on shutdown", self.exception_capture.close, + errors, ) # Ordinary cleanup failures are logged by each step. Reaching here # means every shutdown cleanup step was attempted once. @@ -2485,10 +2528,14 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: while True: with self._lifecycle_condition: if require_shutdown and self._shutdown_complete: + if self.debug and self._lifecycle_cleanup_failed: + raise RuntimeError("client lifecycle cleanup failed") return if not require_shutdown and ( self._join_cleanup_complete or self._shutdown_complete ): + if self.debug and self._lifecycle_cleanup_failed: + raise RuntimeError("client lifecycle cleanup failed") return if self._lifecycle_owner is not None: if self._is_lifecycle_callback_thread() or ( @@ -2500,14 +2547,15 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: self._lifecycle_owner = threading.current_thread() try: + errors: list[Exception] = [] while True: with self._lifecycle_lock: run_shutdown = self._shutdown_requested if run_shutdown: - self._shutdown_once() + self._shutdown_once(errors) else: - self._join_once() + self._join_once(errors) with self._lifecycle_condition: if ( @@ -2516,6 +2564,9 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: and not run_shutdown ): continue + if errors: + self._lifecycle_cleanup_failed = True + raise errors[0] self._lifecycle_owner = None self._lifecycle_condition.notify_all() return @@ -2598,6 +2649,8 @@ def shutdown(self) -> None: """ with self._lifecycle_lock: if self._shutdown_complete: + if self.debug and self._lifecycle_cleanup_failed: + raise RuntimeError("client lifecycle cleanup failed") return self._shutdown_requested = True if self._defer_lifecycle_from_callback(): diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 9b7e03437..8b335dcd8 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2924,6 +2924,49 @@ def test_join_failure_does_not_retry_or_skip_later_cleanup(self): cleanup.assert_called_once_with() unregister.assert_called_once_with() + def test_concurrent_debug_join_callers_observe_cleanup_failure(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01, debug=True) + cleanup_started = threading.Event() + release_cleanup = threading.Event() + owner_errors = [] + waiter_errors = [] + + def cleanup(): + cleanup_started.set() + self.assertTrue(release_cleanup.wait(2)) + raise ValueError("cleanup failed") + + def run_join(errors): + try: + client.join() + except Exception as error: + errors.append(error) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=cleanup, + ): + owner = threading.Thread(target=run_join, args=(owner_errors,)) + waiter = threading.Thread(target=run_join, args=(waiter_errors,)) + owner.start() + self.assertTrue(cleanup_started.wait(1)) + waiter.start() + time.sleep(0.05) + self.assertTrue(waiter.is_alive()) + + release_cleanup.set() + owner.join(2) + waiter.join(2) + + self.assertFalse(owner.is_alive()) + self.assertFalse(waiter.is_alive()) + self.assertEqual(len(owner_errors), 1) + self.assertRegex(str(owner_errors[0]), "cleanup failed") + self.assertEqual(len(waiter_errors), 1) + self.assertRegex(str(waiter_errors[0]), "client lifecycle cleanup failed") + self.assertTrue(client._join_cleanup_complete) + def test_pending_shutdown_continues_when_join_cleanup_fails(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) cleanup_started = threading.Event() @@ -3010,7 +3053,7 @@ def test_shutdown_failure_is_terminal_and_not_retried(self): self.assertTrue(client._shutdown_complete) exception_capture.close.assert_called_once_with() - def test_shutdown_failure_does_not_skip_later_cleanup(self): + def test_shutdown_failure_is_raised_after_later_cleanup_in_debug_mode(self): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01, debug=True) metrics = mock.Mock() metrics.reset.side_effect = Exception("reset failed") @@ -3021,7 +3064,10 @@ def test_shutdown_failure_does_not_skip_later_cleanup(self): client.distinct_ids_feature_flags_reported = dedupe_cache client.exception_capture = exception_capture - client.shutdown() + with self.assertRaisesRegex(Exception, "reset failed"): + client.shutdown() + with self.assertRaisesRegex(RuntimeError, "client lifecycle cleanup failed"): + client.shutdown() metrics.flush.assert_called_once_with() metrics.reset.assert_called_once_with() @@ -3031,6 +3077,30 @@ def test_shutdown_failure_does_not_skip_later_cleanup(self): self.assertTrue(client._join_cleanup_complete) self.assertTrue(client._shutdown_complete) + def test_lane_join_raises_first_failure_after_attempting_later_cleanup(self): + client = Client(FAKE_TEST_API_KEY, send=False) + lane = client._analytics_lane + consumer = mock.Mock() + first_error = ValueError("pause failed") + consumer._pause.side_effect = first_error + lane.consumers = [consumer] + + with ( + mock.patch.object(lane, "discard_undrainable_queued_work") as discard, + mock.patch.object( + lane._drain_signal, + "complete", + side_effect=RuntimeError("complete failed"), + ) as complete, + self.assertRaises(ValueError) as raised, + ): + lane.join() + + self.assertIs(raised.exception, first_error) + consumer.join.assert_called_once_with() + discard.assert_called_once_with() + complete.assert_called_once_with() + def test_shutdown_prepares_each_lane_once(self): client = Client(FAKE_TEST_API_KEY, send=False) close_methods = []