diff --git a/.sampo/changesets/flush-bypasses-flush-interval.md b/.sampo/changesets/flush-bypasses-flush-interval.md new file mode 100644 index 00000000..f2556709 --- /dev/null +++ b/.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. diff --git a/posthog/client.py b/posthog/client.py index 74a5d9eb..b350daff 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,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: diff --git a/posthog/consumer.py b/posthog/consumer.py index a80f536d..046c02e9 100644 --- a/posthog/consumer.py +++ b/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: diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 4d0919cd..982711fa 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"}) @@ -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() diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index e9e6bf08..296587bf 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -1,4 +1,5 @@ import json +import threading import time import unittest from datetime import datetime, timedelta, timezone @@ -15,7 +16,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 @@ -76,7 +77,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 @@ -85,12 +93,12 @@ def test_flush_interval(self) -> None: q = Queue() flush_interval = 0.3 consumer = Consumer(q, TEST_API_KEY, flush_at=10, flush_interval=flush_interval) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with mock.patch.object(consumer, "request") as mock_request: consumer.start() for i in range(3): q.put(_track_event("python event %d" % i)) time.sleep(flush_interval * 1.1) - self.assertEqual(mock_post.call_count, 3) + self.assertEqual(mock_request.call_count, 3) def test_multiple_uploads_per_interval(self) -> None: # Put _flush_at*2_ items in the queue at once, then pause for @@ -165,6 +173,179 @@ 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(q) + consumer = Consumer(q, TEST_API_KEY, flush_at=100, flush_interval=30) + consumer._set_drain_signal(signal) + q.put(_track_event("first")) + q.put(_track_event("second")) + signal.request() + + start = time.monotonic() + batch = consumer.next() + signal.complete() + + 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(q) + flush_at = 10 + consumer = Consumer(q, TEST_API_KEY, flush_at=flush_at, flush_interval=30) + consumer._set_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) + signal.complete() + + def test_completed_drain_request_restores_normal_batching(self) -> None: + # Once the caller completes its request, later batches must go back to + # normal timer-based batching instead of inheriting a stale drain. + q = Queue() + signal = _DrainSignal(q) + flush_interval = 0.2 + consumer = Consumer( + q, + TEST_API_KEY, + flush_at=100, + flush_interval=flush_interval, + ) + consumer._set_drain_signal(signal) + q.put(_track_event()) + signal.request() + self.assertEqual(len(consumer.next()), 1) + signal.complete() + + start = time.monotonic() + self.assertEqual(consumer.next(), []) + self.assertGreaterEqual(time.monotonic() - start, flush_interval * 0.5) + + def test_overlapping_drain_requests_remain_active_until_all_complete(self) -> None: + q = Queue() + signal = _DrainSignal(q) + + signal.request() + signal.request() + signal.complete() + self.assertTrue(signal.requested) + + signal.complete() + self.assertFalse(signal.requested) + + 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(q) + consumer = Consumer(q, TEST_API_KEY, flush_at=100, flush_interval=30) + consumer._set_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) + signal.complete() + 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(q) + consumer = Consumer(q, TEST_API_KEY, flush_at=100, flush_interval=30) + consumer._set_drain_signal(signal) + q.put(_track_event()) + threading.Timer(0.1, signal.request).start() + + start = time.monotonic() + batch = consumer.next() + signal.complete() + + self.assertEqual(len(batch), 1) + self.assertLess(time.monotonic() - start, 5) + + def test_drain_signal_wakes_an_idle_consumer(self) -> None: + q = Queue() + signal = _DrainSignal(q) + consumer = Consumer(q, TEST_API_KEY, flush_at=100, flush_interval=2) + consumer._set_drain_signal(signal) + threading.Timer(0.1, signal.request).start() + + start = time.monotonic() + batch = consumer.next() + signal.complete() + + self.assertEqual(batch, []) + self.assertLess(time.monotonic() - start, 1) + + def test_idle_consumer_parks_while_drain_waits_for_an_upload(self) -> None: + q = Queue() + signal = _DrainSignal(q) + upload_started = threading.Event() + release_upload = threading.Event() + idle_returned = threading.Event() + idle_next_calls = 0 + + uploading = Consumer(q, TEST_API_KEY, flush_at=1, flush_interval=30) + idle = Consumer(q, TEST_API_KEY, flush_at=100, flush_interval=30) + uploading._set_drain_signal(signal) + idle._set_drain_signal(signal) + + def blocking_request(batch) -> None: + upload_started.set() + self.assertTrue(release_upload.wait(2)) + + original_idle_next = idle.next + + def counted_idle_next(): + nonlocal idle_next_calls + batch = original_idle_next() + idle_next_calls += 1 + idle_returned.set() + return batch + + uploading.request = blocking_request + idle.next = counted_idle_next + q.put(_track_event()) + uploading.start() + self.assertTrue(upload_started.wait(1)) + idle.start() + signal.request() + + try: + self.assertTrue(idle_returned.wait(1)) + time.sleep(0.1) + self.assertEqual(idle_next_calls, 1) + finally: + uploading.pause() + idle.pause() + release_upload.set() + signal.complete() + uploading.join(2) + idle.join(2) + + self.assertFalse(uploading.is_alive()) + self.assertFalse(idle.is_alive()) + + 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) @@ -181,25 +362,22 @@ def test_max_batch_size(self) -> None: # Let's capture 8MB of data to trigger two batches n_msgs = int(8_000_000 / msg_size) - def mock_post_fn(_: str, data: str, **kwargs: Any) -> mock.Mock: - res = mock.Mock() - res.status_code = 200 - request_size = len(data.encode()) + def mock_send_fn(batch: list[dict[str, Any]], _path: str) -> None: + request_size = len(json.dumps({"batch": batch}).encode()) # Batches close after the first message bringing it bigger than BATCH_SIZE_LIMIT, let's add 10% of margin self.assertTrue( request_size < (5 * 1024 * 1024) * 1.1, "batch size (%d) higher than limit" % request_size, ) - return res - with mock.patch( - "posthog.request._session.post", side_effect=mock_post_fn - ) as mock_post: + with mock.patch.object( + consumer, "_send", side_effect=mock_send_fn + ) as mock_send: consumer.start() for _ in range(0, n_msgs + 2): q.put(track) q.join() - self.assertEqual(mock_post.call_count, 2) + self.assertEqual(mock_send.call_count, 2) def test_request_sleeps_with_retry_after(self) -> None: error = APIError(429, "Too Many Requests", retry_after=5.0)