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 9469407b..52fbf282 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, @@ -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: diff --git a/posthog/consumer.py b/posthog/consumer.py index b60e156f..ff9d9aae 100644 --- a/posthog/consumer.py +++ b/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 diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 137fbeff..46de3b4a 100644 --- a/posthog/test/test_client.py +++ b/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"}) diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index ab582193..f2911d9b 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 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) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 90f4f940..9379287d 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -235,6 +235,7 @@ alias posthog.client.DEFAULT_CODE_VARIABLES_DETECT_SECRETS -> posthog.exception_ alias posthog.client.DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS alias posthog.client.DEFAULT_CODE_VARIABLES_MASK_PATTERNS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_MASK_PATTERNS alias posthog.client.DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS +alias posthog.client.DrainSignal -> posthog.consumer.DrainSignal alias posthog.client.EVENTS_ENDPOINT -> posthog.request.EVENTS_ENDPOINT alias posthog.client.ExceptionArg -> posthog.args.ExceptionArg alias posthog.client.ExceptionCapture -> posthog.exception_capture.ExceptionCapture @@ -586,6 +587,7 @@ attribute posthog.consumer.Consumer.api_key = api_key attribute posthog.consumer.Consumer.capture_compression = capture_compression attribute posthog.consumer.Consumer.capture_mode = capture_mode attribute posthog.consumer.Consumer.daemon = True +attribute posthog.consumer.Consumer.drain_signal = drain_signal attribute posthog.consumer.Consumer.endpoint = endpoint attribute posthog.consumer.Consumer.flush_at = flush_at attribute posthog.consumer.Consumer.flush_interval = flush_interval @@ -599,6 +601,8 @@ attribute posthog.consumer.Consumer.queue = queue attribute posthog.consumer.Consumer.retries = retries attribute posthog.consumer.Consumer.running = True attribute posthog.consumer.Consumer.timeout = timeout +attribute posthog.consumer.DRAIN_POLL_INTERVAL = 0.05 +attribute posthog.consumer.DrainSignal.generation: int attribute posthog.consumer.MAX_MSG_SIZE = 900 * 1024 attribute posthog.contexts.ContextScope.capture_exception_code_variables: Optional[bool] = None attribute posthog.contexts.ContextScope.capture_exceptions = capture_exceptions @@ -907,7 +911,8 @@ class posthog.capture_compression.CaptureCompression class posthog.capture_mode.CaptureMode class posthog.capture_v1.CaptureV1Error(status: int | str, message: str, *, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None) class posthog.client.Client(project_api_key: str, host=None, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, gzip=False, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_mode: Optional[Union[CaptureMode, str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, _use_ai_lane=False, _enable_multimodal_capture=False) -class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, gzip=False, retries=10, timeout=15, historical_migration=False, endpoint=EVENTS_ENDPOINT, max_msg_size=MAX_MSG_SIZE, capture_mode=CaptureMode.V0, capture_compression=CaptureCompression.NONE) +class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, gzip=False, retries=10, timeout=15, historical_migration=False, endpoint=EVENTS_ENDPOINT, max_msg_size=MAX_MSG_SIZE, capture_mode=CaptureMode.V0, capture_compression=CaptureCompression.NONE, drain_signal: Optional[DrainSignal] = None) +class posthog.consumer.DrainSignal() class posthog.contexts.ContextScope(parent=None, fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) class posthog.exception_capture.ExceptionCapture(client: Client, rate_limiting_enabled=False, bucket_size=DEFAULT_BUCKET_SIZE, refill_rate=DEFAULT_REFILL_RATE, refill_interval_seconds=DEFAULT_REFILL_INTERVAL_SECONDS) class posthog.exception_utils.AnnotatedValue(value, metadata) @@ -1263,6 +1268,7 @@ method posthog.consumer.Consumer.pause() method posthog.consumer.Consumer.request(batch) method posthog.consumer.Consumer.run() method posthog.consumer.Consumer.upload() +method posthog.consumer.DrainSignal.request() -> None method posthog.contexts.ContextScope.add_tag(key: str, value: Any) method posthog.contexts.ContextScope.collect_tags() -> Dict[str, Any] method posthog.contexts.ContextScope.get_capture_exception_code_variables() -> Optional[bool]