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,
@@ -306,6 +306,7 @@ def __init__(
         self._started = False
         self._closed = False
         self._start_lock = threading.Lock()
+        self._drain_signal = DrainSignal()
         if eager_start:
             self.start()
 
@@ -334,6 +335,7 @@ def start(self):
                     max_msg_size=self.max_msg_size,
                     capture_mode=self.capture_mode,
                     capture_compression=self.capture_compression,
+                    drain_signal=self._drain_signal,
                 )
                 self.consumers.append(consumer)
 
@@ -359,8 +361,15 @@ def close(self) -> None:
             self._closed = True
 
     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
+        # Must happen before we start waiting: a consumer parked on a partial
+        # batch only learns to send it from this signal.
+        self._drain_signal.request()
         size = queue.qsize()
         if timeout_seconds is None:
             queue.join()
@@ -384,6 +393,9 @@ def flush(self, timeout_seconds: Optional[float]) -> None:
 
     def join(self) -> None:
         """Pause this lane's consumers and wait for them to exit; a never-started lane is a no-op."""
+        # 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()
         for consumer in self.consumers:
             consumer.pause()
             try:
@@ -403,6 +415,7 @@ def rebuild_after_fork(self) -> None:
         """
         self.queue = Queue(self._max_queue_size)
         self._start_lock = threading.Lock()
+        self._drain_signal = DrainSignal()
         self.consumers = []
         self._started = False
         if self._eager_start:

posthog/consumer.py

@@ -1,6 +1,7 @@
-from typing import Any
+from typing import Any, Optional
 import json
 import logging
+import threading
 import time
 from threading import Thread
 
@@ -31,10 +32,42 @@
 # in case we want to lower it in the future.
 BATCH_SIZE_LIMIT = 5 * 1024 * 1024
 
+# How long a consumer that is already accumulating a batch may block on the
+# queue before re-checking its drain signal. An idle consumer (nothing
+# accumulated) still parks for the whole `flush_interval`, because anything a
+# caller enqueued before calling `flush()` is already in the queue and wakes the
+# blocking `get` on its own.
+DRAIN_POLL_INTERVAL = 0.05
+
 
 _configure_posthog_logging()
 
 
+class DrainSignal:
+    """Cross-thread "stop batching and send what is pending" signal.
+
+    Explicit flushes must not wait for `flush_at` or `flush_interval`, but a
+    consumer accumulating a partial batch is parked on its queue and cannot see
+    a plain flag flip. `flush()` bumps a generation counter here; each consumer
+    remembers the generation it last saw its queue empty at, so a request stays
+    pending until that consumer has actually handed off everything it holds.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._generation = 0
+
+    def request(self) -> None:
+        """Ask every consumer sharing this signal to deliver what it has now."""
+        with self._lock:
+            self._generation += 1
+
+    @property
+    def generation(self) -> int:
+        with self._lock:
+            return self._generation
+
+
 class Consumer(Thread):
     """Consumes the messages from the client's queue."""
 
@@ -56,6 +89,7 @@ def __init__(
         max_msg_size=MAX_MSG_SIZE,
         capture_mode=CaptureMode.V0,
         capture_compression=CaptureCompression.NONE,
+        drain_signal: Optional[DrainSignal] = None,
     ):
         """Create a consumer thread."""
         Thread.__init__(self)
@@ -72,6 +106,10 @@ def __init__(
         self.max_msg_size = max_msg_size
         self.capture_mode = capture_mode
         self.capture_compression = capture_compression
+        self.drain_signal = drain_signal
+        # Start level with the signal: a consumer built after a flush must not
+        # inherit that flush's pending request.
+        self._drain_seen = drain_signal.generation if drain_signal else 0
         # 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
@@ -118,6 +156,10 @@ def upload(self):
 
         return success
 
+    def _drain_generation(self) -> int:
+        """The drain request generation currently visible to this consumer."""
+        return self.drain_signal.generation if self.drain_signal is not None else 0
+
     def next(self):
         """Return the next batch of items to upload."""
         queue = self.queue
@@ -127,11 +169,27 @@ 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.
+            drain_generation = self._drain_generation()
+            draining = drain_generation != self._drain_seen
+            remaining = self.flush_interval - (time.monotonic() - start_time)
+            if not draining and remaining <= 0:
                 break
+
+            # A partial batch is pending, so break the wait into slices to
+            # notice a drain request that arrives while we are parked here. An
+            # idle consumer still parks for the whole interval: anything a
+            # caller enqueued before flush() is already in the queue and wakes
+            # the blocking get() by itself.
+            sliced = bool(items) and self.drain_signal is not None
+            timeout = min(remaining, DRAIN_POLL_INTERVAL) if sliced else remaining
+
             try:
-                item = queue.get(block=True, timeout=self.flush_interval - elapsed)
+                if draining:
+                    item = queue.get(block=False)
+                else:
+                    item = queue.get(block=True, timeout=timeout)
                 item_size = len(json.dumps(item, cls=DatetimeSerializer).encode())
                 if item_size > self.max_msg_size:
                     # Log only name and size: AI events may carry unredacted
@@ -151,6 +209,17 @@ def next(self):
                     self.log.debug("hit batch size limit (size: %d)", total_size)
                     break
             except Empty:
+                if draining:
+                    # Everything that flush was waiting for is in `items` now.
+                    # Recording the generation read at the top of this iteration
+                    # (rather than the current one) leaves a request that landed
+                    # while we were draining pending for the next batch.
+                    self._drain_seen = drain_generation
+                    break
+                if timeout < remaining:
+                    # Only a poll slice expired, not the batch window: keep
+                    # accumulating so batching is unchanged without a flush.
+                    continue
                 break
 
         return items

posthog/test/test_client.py

@@ -275,6 +275,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"})

posthog/test/test_consumer.py

@@ -1,4 +1,5 @@
 import json
+import threading
 import time
 import unittest
 from typing import Any
@@ -13,7 +14,7 @@
 
 from posthog.capture_compression import CaptureCompression
 from posthog.capture_mode import CaptureMode
-from posthog.consumer import MAX_MSG_SIZE, Consumer
+from posthog.consumer import MAX_MSG_SIZE, Consumer, DrainSignal
 from posthog.request import AI_EVENTS_ENDPOINT, EVENTS_ENDPOINT, APIError
 from posthog.test.logging_helpers import capture_message_only_logs
 from posthog.test.test_utils import TEST_API_KEY
@@ -74,7 +75,14 @@ def test_message_only_error_logs_include_posthog_prefix(self) -> None:
                 success = consumer.upload()
 
         self.assertFalse(success)
-        self.assertEqual(logs.getvalue().strip(), "[PostHog] error uploading: boom")
+        # `capture_message_only_logs` taps the process-wide "posthog" logger and
+        # `upload()` spans a whole flush_interval, so background threads left by
+        # other tests can log into the same stream. Assert on the line under
+        # test rather than on the entire capture.
+        upload_logs = [
+            line for line in logs.getvalue().splitlines() if "error uploading" in line
+        ]
+        self.assertEqual(upload_logs, ["[PostHog] error uploading: boom"])
 
     def test_flush_interval(self) -> None:
         # Put _n_ items in the queue, pausing a little bit more than
@@ -154,6 +162,103 @@ def test_pause(self) -> None:
         consumer.pause()
         self.assertFalse(consumer.running)
 
+    def test_drain_signal_returns_partial_batch_without_waiting(self) -> None:
+        # A drain request means "send what is queued now", so `next()` must not
+        # hold a below-flush_at batch back for the rest of flush_interval.
+        q = Queue()
+        signal = DrainSignal()
+        consumer = Consumer(
+            q, TEST_API_KEY, flush_at=100, flush_interval=30, drain_signal=signal
+        )
+        q.put(_track_event("first"))
+        q.put(_track_event("second"))
+        signal.request()
+
+        start = time.monotonic()
+        batch = consumer.next()
+
+        self.assertEqual(len(batch), 2)
+        self.assertLess(time.monotonic() - start, 5)
+
+    def test_drain_signal_still_respects_flush_at(self) -> None:
+        # Draining must not degrade batching into one request per event.
+        q = Queue()
+        signal = DrainSignal()
+        flush_at = 10
+        consumer = Consumer(
+            q, TEST_API_KEY, flush_at=flush_at, flush_interval=30, drain_signal=signal
+        )
+        for i in range(flush_at * 3):
+            q.put(_track_event("python event %d" % i))
+        signal.request()
+
+        self.assertEqual(len(consumer.next()), flush_at)
+
+    def test_drain_signal_is_satisfied_once_the_queue_empties(self) -> None:
+        # Once the queue has been observed empty the request is served, so the
+        # consumer goes back to normal timer-based batching instead of spinning.
+        q = Queue()
+        signal = DrainSignal()
+        flush_interval = 0.2
+        consumer = Consumer(
+            q,
+            TEST_API_KEY,
+            flush_at=100,
+            flush_interval=flush_interval,
+            drain_signal=signal,
+        )
+        q.put(_track_event())
+        signal.request()
+        self.assertEqual(len(consumer.next()), 1)
+
+        start = time.monotonic()
+        self.assertEqual(consumer.next(), [])
+        self.assertGreaterEqual(time.monotonic() - start, flush_interval * 0.5)
+
+    def test_consecutive_drain_requests_each_drain_immediately(self) -> None:
+        # A later flush must not be served by an earlier flush's bookkeeping.
+        q = Queue()
+        signal = DrainSignal()
+        consumer = Consumer(
+            q, TEST_API_KEY, flush_at=100, flush_interval=30, drain_signal=signal
+        )
+
+        for i in range(3):
+            q.put(_track_event("python event %d" % i))
+            signal.request()
+            start = time.monotonic()
+            self.assertEqual(len(consumer.next()), 1)
+            self.assertLess(time.monotonic() - start, 5)
+
+    def test_drain_signal_wakes_a_consumer_mid_batch(self) -> None:
+        # The realistic ordering: the consumer is already parked on a partial
+        # batch when flush() signals it.
+        q = Queue()
+        signal = DrainSignal()
+        consumer = Consumer(
+            q, TEST_API_KEY, flush_at=100, flush_interval=30, drain_signal=signal
+        )
+        q.put(_track_event())
+        threading.Timer(0.1, signal.request).start()
+
+        start = time.monotonic()
+        batch = consumer.next()
+
+        self.assertEqual(len(batch), 1)
+        self.assertLess(time.monotonic() - start, 5)
+
+    def test_without_drain_signal_batching_is_unchanged(self) -> None:
+        q = Queue()
+        flush_interval = 0.3
+        consumer = Consumer(
+            q, TEST_API_KEY, flush_at=100, flush_interval=flush_interval
+        )
+        q.put(_track_event())
+
+        start = time.monotonic()
+        self.assertEqual(len(consumer.next()), 1)
+        self.assertGreaterEqual(time.monotonic() - start, flush_interval * 0.5)
+
     def test_max_batch_size(self) -> None:
         q = Queue()
         consumer = Consumer(q, TEST_API_KEY, flush_at=100000, flush_interval=3)