fix: flush() delivers pending events without waiting out flush_interval

.sampo/changesets/flush-bypasses-flush-interval.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+fix: `flush()` no longer waits out `flush_interval` before delivering a partial batch. A consumer holding fewer than `flush_at` events now sends them as soon as `flush()` (or `shutdown()`) asks it to, instead of blocking the caller for the rest of the batching window — which previously made `flush()` deliver nothing at all when `flush_interval` was longer than the flush timeout. Timer-based batching without an explicit flush is unchanged.

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,38 +388,54 @@ 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 and wait for them to exit; a never-started lane is a no-op."""
-        for consumer in self.consumers:
-            consumer.pause()
-            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()
+                try:
+                    consumer.join()
+                except RuntimeError:
+                    # consumer thread has not started
+                    pass
+        finally:
+            self._drain_signal.complete()
 
     def reset_sync_send_state_after_fork(self) -> None:
         """Replace sync-send state inherited from threads that did not survive fork."""
@@ -436,6 +454,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:

posthog/consumer.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Optional
 import json
 import logging
 import time
@@ -31,10 +31,56 @@
 # in case we want to lower it in the future.
 BATCH_SIZE_LIMIT = 5 * 1024 * 1024
 
-
 _configure_posthog_logging()
 
 
+class _DrainSignal:
+    """Wake queue consumers while one or more drain requests are active."""
+
+    def __init__(self, queue) -> None:
+        self._queue = queue
+        self._requests = 0
+
+    def request(self) -> None:
+        with self._queue.not_empty:
+            self._requests += 1
+            self._queue.not_empty.notify_all()
+
+    def complete(self) -> None:
+        with self._queue.not_empty:
+            self._requests -= 1
+            self._queue.not_empty.notify_all()
+
+    def wake(self) -> None:
+        with self._queue.not_empty:
+            self._queue.not_empty.notify_all()
+
+    def wait_until_inactive_or_work(self, consumer) -> None:
+        with self._queue.not_empty:
+            while self._requests and not self._queue._qsize() and consumer.running:
+                self._queue.not_empty.wait()
+
+    @property
+    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."""
+        with self._queue.not_empty:
+            deadline = time.monotonic() + timeout
+            while not self._queue._qsize():
+                if self._requests:
+                    raise Empty
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise Empty
+                self._queue.not_empty.wait(remaining)
+            item = self._queue._get()
+            self._queue.not_full.notify()
+            return item
+
+
 class Consumer(Thread):
     """Consumes the messages from the client's queue."""
 
@@ -72,6 +118,7 @@ def __init__(
         self.max_msg_size = max_msg_size
         self.capture_mode = capture_mode
         self.capture_compression = capture_compression
+        self._drain_signal: Optional[_DrainSignal] = None
         # 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
@@ -86,12 +133,16 @@ def run(self):
         self.log.debug("consumer is running...")
         while self.running:
             self.upload()
+            if self._drain_signal is not None:
+                self._drain_signal.wait_until_inactive_or_work(self)
 
         self.log.debug("consumer exited.")
 
     def pause(self):
         """Pause the consumer."""
         self.running = False
+        if self._drain_signal is not None:
+            self._drain_signal.wake()
 
     def upload(self):
         """Upload the next batch of items, return whether successful."""
@@ -118,6 +169,12 @@ def upload(self):
 
         return success
 
+    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
+
     def next(self):
         """Return the next batch of items to upload."""
         queue = self.queue
@@ -127,11 +184,20 @@ def next(self):
         total_size = 0
 
         while len(items) < self.flush_at:
-            elapsed = time.monotonic() - start_time
-            if elapsed >= self.flush_interval:
+            # 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:
-                item = queue.get(block=True, timeout=self.flush_interval - elapsed)
+                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:

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"})
@@ -2164,6 +2218,14 @@ def test_shutdown_flushes_without_timeout(self):
 
         mock_flush.assert_called_once_with(timeout_seconds=None)
 
+    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()