fix(client): make shutdown atomic

.sampo/changesets/regal-guardian-tuulikki.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Make client shutdown an atomic terminal boundary

posthog/client.py

@@ -305,59 +305,86 @@ 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.
 
         Idempotent and thread-safe, so concurrent first captures start exactly
         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:

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)

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(