diff --git a/.sampo/changesets/regal-guardian-tuulikki.md b/.sampo/changesets/regal-guardian-tuulikki.md new file mode 100644 index 000000000..eb305618d --- /dev/null +++ b/.sampo/changesets/regal-guardian-tuulikki.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Make client shutdown an atomic terminal boundary diff --git a/posthog/client.py b/posthog/client.py index 9469407bd..b2122e504 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -305,10 +305,38 @@ def __init__( self.consumers: List[Consumer] = [] self._started = False self._closed = False + self._active_sync_sends = 0 self._start_lock = threading.Lock() + self._sync_sends_done = threading.Condition(self._start_lock) if eager_start: self.start() + def _start_locked(self) -> None: + if self._started or self._closed: + return + for _ in range(self._thread_count): + consumer = Consumer( + self.queue, + self.api_key, + host=self.host, + on_error=self.on_error, + flush_at=self.flush_at, + flush_interval=self.flush_interval, + gzip=self.gzip, + retries=self.max_retries, + timeout=self.timeout, + historical_migration=self.historical_migration, + endpoint=self.endpoint, + max_msg_size=self.max_msg_size, + capture_mode=self.capture_mode, + capture_compression=self.capture_compression, + ) + self.consumers.append(consumer) + + if self.send: + consumer.start() + self._started = True + def start(self): """Construct this lane's consumer pool, starting its threads when sending is enabled. @@ -316,48 +344,47 @@ def start(self): one pool. """ with self._start_lock: - if self._started or self._closed: - return - for _ in range(self._thread_count): - consumer = Consumer( - self.queue, - self.api_key, - host=self.host, - on_error=self.on_error, - flush_at=self.flush_at, - flush_interval=self.flush_interval, - gzip=self.gzip, - retries=self.max_retries, - timeout=self.timeout, - historical_migration=self.historical_migration, - endpoint=self.endpoint, - max_msg_size=self.max_msg_size, - capture_mode=self.capture_mode, - capture_compression=self.capture_compression, - ) - self.consumers.append(consumer) - - if self.send: - consumer.start() - self._started = True + self._start_locked() def enqueue(self, msg) -> bool: - """Queue `msg` for upload, starting the lane on its first event.""" - if self._closed: - return False - if not self._started: - self.start() + """Atomically admit and queue `msg`, starting the lane on its first event.""" + with self._start_lock: + if self._closed: + return False + self._start_locked() + try: + self.queue.put(msg, block=False) + return True + except Full: + return False + + def run_sync_if_open(self, send) -> bool: + """Run a synchronous send admitted before closure, and report whether it ran.""" + with self._sync_sends_done: + if self._closed: + return False + self._active_sync_sends += 1 + try: - self.queue.put(msg, block=False) - return True - except Full: - return False + send() + finally: + with self._sync_sends_done: + self._active_sync_sends -= 1 + if not self._active_sync_sends: + self._sync_sends_done.notify_all() + return True def close(self) -> None: - """Terminal: refuse all future enqueues and consumer starts.""" + """Terminal: atomically refuse all future queue and sync admissions.""" with self._start_lock: self._closed = True + def wait_for_sync_sends(self) -> None: + """Wait for synchronous sends admitted before close to finish.""" + with self._sync_sends_done: + while self._active_sync_sends: + 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.""" queue = self.queue @@ -392,6 +419,12 @@ def join(self) -> None: # consumer thread has not started pass + 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 + self._start_lock = threading.Lock() + self._sync_sends_done = threading.Condition(self._start_lock) + def rebuild_after_fork(self) -> None: """Replace fork-unsafe lane state in a forked child. @@ -402,7 +435,7 @@ def rebuild_after_fork(self) -> None: a lazy lane returns to not-started and restarts on next use. """ self.queue = Queue(self._max_queue_size) - self._start_lock = threading.Lock() + self.reset_sync_send_state_after_fork() self.consumers = [] self._started = False if self._eager_start: @@ -1847,8 +1880,10 @@ 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`). """ - if not self.sync_mode: - for lane in self._lanes: + 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: @@ -1953,38 +1988,55 @@ def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): self.log.debug("queueing: %s", msg) - # if send is False, return msg as if it was successfully queued + # if send is False, return msg as if it was successfully queued, unless + # shutdown has already closed this lane's admission. if not self.send: - return sent_uuid + if lane.run_sync_if_open(lambda: None): + return sent_uuid + self.log.warning( + "%s lane received event %s after shutdown, dropping it", + lane.name, + msg["event"], + ) + return None if self.sync_mode: self.log.debug("enqueued with blocking %s.", msg["event"]) - # Sync mode bypasses the lane's queue but keeps its wire config: - # the AI lane is pinned to v0, so its events post to the AI - # endpoint regardless of `capture_mode`. - if lane.capture_mode == CaptureMode.V1: - _send_v1_batch( + + def send_sync() -> None: + # Sync mode bypasses the lane's queue but keeps its wire config: + # the AI lane is pinned to v0, so its events post to the AI + # endpoint regardless of `capture_mode`. + if lane.capture_mode == CaptureMode.V1: + _send_v1_batch( + self.api_key, + self.host, + [msg], + compression=self.capture_compression, + timeout=self.timeout, + max_retries=self.max_retries, + historical_migration=self.historical_migration, + ) + return + + batch_post( self.api_key, self.host, - [msg], - compression=self.capture_compression, + gzip=self.gzip, timeout=self.timeout, - max_retries=self.max_retries, + batch=[msg], historical_migration=self.historical_migration, + path=lane.endpoint, ) - return sent_uuid - batch_post( - self.api_key, - self.host, - gzip=self.gzip, - timeout=self.timeout, - batch=[msg], - historical_migration=self.historical_migration, - path=lane.endpoint, + if lane.run_sync_if_open(send_sync): + return sent_uuid + self.log.warning( + "%s lane received event %s after shutdown, dropping it", + lane.name, + msg["event"], ) - - return sent_uuid + return None if lane.enqueue(msg): self.log.debug("enqueued %s.", msg["event"]) @@ -2097,6 +2149,13 @@ 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: @@ -2105,8 +2164,6 @@ def shutdown(self) -> None: self.log.exception("Failed to flush metrics on shutdown") self._metrics.reset() self.join() - for lane in self._lanes: - lane.close() self.distinct_ids_feature_flags_reported.clear() if self.exception_capture: diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index fd798f7bc..a61d947d1 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -1,5 +1,6 @@ import logging import asyncio +import threading import time import unittest import warnings @@ -2161,6 +2162,88 @@ def test_shutdown_flushes_without_timeout(self): mock_flush.assert_called_once_with(timeout_seconds=None) + def test_shutdown_waits_for_racing_enqueue_before_draining(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + put_started = threading.Event() + release_put = threading.Event() + shutdown_done = threading.Event() + original_put = client.queue.put + capture_result = [] + + def blocking_put(*args, **kwargs): + put_started.set() + self.assertTrue(release_put.wait(2)) + return original_put(*args, **kwargs) + + capture_thread = threading.Thread( + target=lambda: capture_result.append( + client.capture("racing event", distinct_id="distinct_id") + ) + ) + shutdown_thread = threading.Thread( + target=lambda: (client.shutdown(), shutdown_done.set()) + ) + + with mock.patch.object(client.queue, "put", side_effect=blocking_put): + capture_thread.start() + self.assertTrue(put_started.wait(2)) + shutdown_thread.start() + try: + self.assertFalse(shutdown_done.wait(0.1)) + finally: + release_put.set() + + capture_thread.join(2) + shutdown_thread.join(2) + + self.assertFalse(capture_thread.is_alive()) + self.assertFalse(shutdown_thread.is_alive()) + self.assertTrue(shutdown_done.is_set()) + self.assertIsNotNone(capture_result[0]) + self.assertTrue(client.queue.empty()) + + def test_shutdown_waits_for_sync_send_and_rejects_later_sends(self): + client = Client(FAKE_TEST_API_KEY, sync_mode=True) + send_started = threading.Event() + release_send = threading.Event() + shutdown_done = threading.Event() + capture_result = [] + + def blocking_post(*args, **kwargs): + send_started.set() + self.assertTrue(release_send.wait(2)) + + capture_thread = threading.Thread( + target=lambda: capture_result.append( + client.capture("in-flight event", distinct_id="distinct_id") + ) + ) + shutdown_thread = threading.Thread( + target=lambda: (client.shutdown(), shutdown_done.set()) + ) + + with mock.patch("posthog.client.batch_post", side_effect=blocking_post) as post: + capture_thread.start() + self.assertTrue(send_started.wait(2)) + shutdown_thread.start() + try: + self.assertFalse(shutdown_done.wait(0.1)) + finally: + release_send.set() + + capture_thread.join(2) + shutdown_thread.join(2) + later_result = client.capture( + "post-shutdown event", distinct_id="distinct_id" + ) + + self.assertFalse(capture_thread.is_alive()) + self.assertFalse(shutdown_thread.is_alive()) + self.assertTrue(shutdown_done.is_set()) + self.assertIsNotNone(capture_result[0]) + self.assertIsNone(later_result) + post.assert_called_once() + def test_synchronous(self): with mock.patch("posthog.client.batch_post") as mock_post: client = Client(FAKE_TEST_API_KEY, sync_mode=True) diff --git a/posthog/test/test_client_fork.py b/posthog/test/test_client_fork.py index 294dcffb0..37b0a3eaa 100644 --- a/posthog/test/test_client_fork.py +++ b/posthog/test/test_client_fork.py @@ -144,13 +144,20 @@ 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_noop_for_sync_mode(self): + def test_reinit_after_fork_resets_sync_send_state_for_sync_mode(self): client = Client(FAKE_TEST_API_KEY, sync_mode=True) - old_queue = client.queue + lane = client._analytics_lane + old_queue = lane.queue + old_lock = lane._start_lock + old_condition = lane._sync_sends_done + lane._active_sync_sends = 1 client._reinit_after_fork() - self.assertIs(client.queue, old_queue) + self.assertIs(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) @unittest.skipUnless(