diff --git a/.sampo/changesets/lifecycle-deadlocks.md b/.sampo/changesets/lifecycle-deadlocks.md new file mode 100644 index 00000000..2170c8e5 --- /dev/null +++ b/.sampo/changesets/lifecycle-deadlocks.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +fix: prevent client lifecycle deadlocks when error callbacks, concurrent `join()`/`shutdown()` calls, or forked sync-mode clients interact with queue and worker teardown. diff --git a/posthog/__init__.py b/posthog/__init__.py index e45fae7e..6ef2449b 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -310,7 +310,9 @@ def get_tags() -> Dict[str, Any]: host: PostHog ingestion host. Defaults to the US ingestion endpoint when not set. on_error: Optional callback invoked by background consumers when event upload - fails. + fails. Keep it short and non-blocking. Lifecycle methods can be called + directly and will be deferred, but the callback must not wait for another + thread or task that calls ``flush()``, ``join()``, or ``shutdown()``. debug: Enable verbose SDK logging and re-raise errors from public APIs. send: If False, queueing succeeds but events are not sent to PostHog. sync_mode: If True, send events synchronously instead of using background @@ -1124,7 +1126,11 @@ def flush(timeout_seconds: Optional[float] = 10) -> None: def join() -> None: """ - Block program until the client clears the queue. Used during program shutdown. You should use `shutdown()` directly in most cases. + Attempt to process queued events and stop the client's background workers. Use `shutdown()` directly in most cases. + + Failed or undrainable events may be dropped and reported through logging or + ``on_error``; returning does not guarantee server receipt. Lifecycle cleanup + is attempted once, and cleanup failures are logged without retry. Examples: ```python @@ -1142,6 +1148,16 @@ def shutdown() -> None: """ Flush all messages and cleanly shutdown the client. + This normally blocks until queued events have been attempted and cleanup + finishes. Failed or undrainable events may be dropped and reported through + logging or ``on_error``; returning does not guarantee server receipt. + Lifecycle cleanup is attempted once, and cleanup failures are logged without + retry. Calls made directly from SDK callbacks such as ``on_error`` are deferred + to avoid deadlocking the worker. If blocking completion is required, signal an application-owned + thread, return from the callback, and call ``shutdown()`` from that thread. + Do not wait inside a callback for another thread or task calling a lifecycle + method. + Examples: ```python from posthog import shutdown diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index 8bc076a1..f81d6964 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -1,40 +1,138 @@ import asyncio +import inspect +import sys import threading from collections.abc import Awaitable +from contextvars import Context, copy_context from typing import Any +class _PlainExecutorCall: + def __init__(self, func, args, kwargs) -> None: + self._func = func + self._args = args + self._kwargs = kwargs + + def __call__(self): + return self._func(*self._args, **self._kwargs) + + +class _ContextExecutorCall: + """Carry context in-process while remaining safe for serializing executors.""" + + def __init__(self, context: Context, func, args, kwargs=None) -> None: + self._context = context + self._func = func + self._args = args + self._kwargs = kwargs or {} + + def __call__(self): + return self._context.run(self._func, *self._args, **self._kwargs) + + def __reduce__(self): + # Context objects are not picklable and are process-local. Executors + # that serialize work reconstruct a plain call instead. + return (_PlainExecutorCall, (self._func, self._args, self._kwargs)) + + +if sys.platform == "win32": + from asyncio.windows_events import ProactorEventLoop as _PlatformEventLoop +else: + _PlatformEventLoop = asyncio.SelectorEventLoop + + +class _ContextEventLoop(_PlatformEventLoop): + def run_in_executor(self, executor, func, *args): # type: ignore[override] + call = _ContextExecutorCall(copy_context(), func, args) + return super().run_in_executor(executor, call) + + +class _LoopStartup: + def __init__(self) -> None: + self.done = threading.Event() + self.loop: asyncio.AbstractEventLoop | None = None + self.error: BaseException | None = None + + class _BackgroundEventLoopRunner: """Run awaitables to completion on a reusable background event loop.""" def __init__(self) -> None: self._loop: asyncio.AbstractEventLoop | None = None self._thread: threading.Thread | None = None - self._started = threading.Event() + self._closing_threads: set[threading.Thread] = set() + self._startup: _LoopStartup | None = None + self._close_requested = False self._lock = threading.Lock() def run(self, awaitable: Awaitable[Any]) -> Any: - loop = self._ensure_loop() - future = asyncio.run_coroutine_threadsafe(self._await_result(awaitable), loop) - return future.result() + if threading.current_thread() is self._thread: + self._close_awaitable(awaitable) + raise RuntimeError("cannot synchronously run from the runner thread") + + try: + while True: + loop = self._ensure_loop() + with self._lock: + if loop is self._loop and not self._close_requested: + wrapped = self._await_result(awaitable) + try: + future = asyncio.run_coroutine_threadsafe(wrapped, loop) + except BaseException: + wrapped.close() + raise + break + except BaseException: + self._close_awaitable(awaitable) + raise + try: + return future.result() + finally: + if future.cancelled(): + self._close_awaitable(awaitable) def close(self) -> None: + current = threading.current_thread() with self._lock: loop = self._loop thread = self._thread - self._loop = None - self._thread = None + if thread is None: + return + self._close_requested = True + if loop is None: + if self._startup is not None: + self._startup.error = RuntimeError("runner closed during startup") + else: + self._loop = None + self._thread = None + self._closing_threads.add(thread) - if loop is None or thread is None or loop.is_closed(): + if loop is None: + if thread is not current: + thread.join() return - if thread is threading.current_thread(): + if loop.is_closed(): + with self._lock: + self._closing_threads.discard(thread) + return + + if thread is current: loop.call_soon(loop.stop) return loop.call_soon_threadsafe(loop.stop) thread.join() + def owns_thread(self, thread: threading.Thread) -> bool: + with self._lock: + return thread is self._thread or thread in self._closing_threads + + @staticmethod + def _close_awaitable(awaitable: Awaitable[Any]) -> None: + if inspect.iscoroutine(awaitable): + awaitable.close() + @staticmethod async def _await_result(awaitable: Awaitable[Any]) -> Any: return await awaitable @@ -49,25 +147,71 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop: ): return self._loop - self._started.clear() - self._thread = threading.Thread( - target=self._run_loop, - name="PostHogBackgroundEventLoopRunner", - daemon=True, - ) - self._thread.start() + startup: _LoopStartup + if self._thread is None or not self._thread.is_alive(): + startup = _LoopStartup() + self._startup = startup + self._close_requested = False + self._thread = threading.Thread( + target=self._run_loop, + args=(startup,), + name="PostHogBackgroundEventLoopRunner", + daemon=True, + ) + self._thread.start() + else: + existing_startup = self._startup + if existing_startup is None: + raise RuntimeError("event loop startup state is unavailable") + startup = existing_startup - self._started.wait() - with self._lock: - assert self._loop is not None - return self._loop + startup.done.wait() + if startup.error is not None: + raise startup.error + if startup.loop is None: + raise RuntimeError("event loop startup completed without a loop") + return startup.loop + + def _run_loop(self, startup: _LoopStartup) -> None: + loop = None + try: + loop = asyncio.new_event_loop() + original_run_in_executor = loop.run_in_executor + + def run_in_executor(executor, func, *args): + call = _ContextExecutorCall(copy_context(), func, args) + return original_run_in_executor(executor, call) + + try: + setattr(loop, "run_in_executor", run_in_executor) + except (AttributeError, TypeError): + # A loop that cannot carry callback context is unsafe for + # lifecycle re-entry from executor threads. Use the equivalent + # context-aware platform loop rather than silently losing it. + loop.close() + loop = _ContextEventLoop() + asyncio.set_event_loop(loop) + except BaseException as error: + if loop is not None and not loop.is_closed(): + loop.close() + with self._lock: + if startup.error is None: + startup.error = error + if self._thread is threading.current_thread(): + self._thread = None + startup.done.set() + return - def _run_loop(self) -> None: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) with self._lock: self._loop = loop - self._started.set() + startup.loop = loop + close_requested = self._close_requested + if close_requested and startup.error is None: + startup.error = RuntimeError("runner closed during startup") + startup.done.set() + + if close_requested: + loop.call_soon(loop.stop) try: loop.run_forever() @@ -83,3 +227,12 @@ def _run_loop(self) -> None: loop.run_until_complete(loop.shutdown_default_executor()) asyncio.set_event_loop(None) loop.close() + current = threading.current_thread() + with self._lock: + if self._thread is current: + self._thread = None + if self._loop is loop: + self._loop = None + if self._startup is startup: + self._startup = None + self._closing_threads.discard(current) diff --git a/posthog/client.py b/posthog/client.py index a10ef6bc..b6cbbf81 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -8,8 +8,9 @@ import time import warnings import weakref +from contextvars import ContextVar from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Mapping, Optional, Union +from typing import Any, Callable, Dict, List, Mapping, Optional, Union from uuid import UUID, uuid4 from typing_extensions import Unpack @@ -113,12 +114,23 @@ from posthog.version import VERSION -from queue import Queue, Full +from queue import Empty, Full, Queue _configure_posthog_logging() MAX_DICT_SIZE = 50_000 +_ATEXIT_FLUSH_TIMEOUT_SECONDS = 1.0 +_atexit_deadline: Optional[float] = None +_atexit_deadline_lock = threading.Lock() + + +def _get_atexit_deadline() -> float: + global _atexit_deadline + with _atexit_deadline_lock: + if _atexit_deadline is None: + _atexit_deadline = time.monotonic() + _ATEXIT_FLUSH_TIMEOUT_SECONDS + return _atexit_deadline def get_identity_state(passed) -> tuple[str, bool]: @@ -413,12 +425,21 @@ def flush(self, timeout_seconds: Optional[float]) -> None: self._drain_signal.request() try: size = queue.qsize() - if timeout_seconds is None: - queue.join() - else: - deadline = time.monotonic() + timeout_seconds + deadline = ( + None if timeout_seconds is None else time.monotonic() + timeout_seconds + ) + while queue.unfinished_tasks: + if deadline is None and not any( + consumer.is_alive() for consumer in self.consumers + ): + self.discard_undrainable_queued_work() + break with queue.all_tasks_done: - while queue.unfinished_tasks: + if not queue.unfinished_tasks: + break + if deadline is None: + wait_seconds = 0.05 + else: remaining = deadline - time.monotonic() if remaining <= 0: self.log.warning( @@ -428,28 +449,89 @@ def flush(self, timeout_seconds: Optional[float]) -> None: queue.unfinished_tasks, ) return - queue.all_tasks_done.wait(remaining) + wait_seconds = min(0.05, remaining) + queue.all_tasks_done.wait(wait_seconds) # 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 discard_undrainable_queued_work(self) -> None: + """Balance queued work when this lane has no running sender.""" + if any(consumer.is_alive() for consumer in self.consumers): + return + + dropped = 0 + while True: + try: + self.queue.get_nowait() + except Empty: + break + self.queue.task_done() + dropped += 1 + if dropped: + self.log.warning( + "%s lane discarded %d queued events because no consumer is running", + self.name, + dropped, + ) + 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() + """Pause this lane's consumers and wait for them to exit.""" + # Normal teardown bypasses the batching wait so a partial batch is sent. + errors: list[Exception] = [] + drain_requested = False try: - for consumer in self.consumers: - consumer.pause() - try: - consumer.join() - except RuntimeError: - # consumer thread has not started - pass - finally: - self._drain_signal.complete() + self._drain_signal.request() + drain_requested = True + except Exception as error: + self.log.exception( + "Failed to request %s lane drain during lifecycle cleanup", self.name + ) + errors.append(error) + + for consumer in self.consumers: + try: + consumer._pause(drain=True) + except Exception as error: + self.log.exception( + "Failed to pause %s lane consumer during lifecycle cleanup", + self.name, + ) + errors.append(error) + for consumer in self.consumers: + try: + consumer.join() + except RuntimeError: + # consumer thread has not started + pass + except Exception as error: + self.log.exception( + "Failed to join %s lane consumer during lifecycle cleanup", + self.name, + ) + errors.append(error) + try: + self.discard_undrainable_queued_work() + except Exception as error: + self.log.exception( + "Failed to discard queued %s lane work during lifecycle cleanup", + self.name, + ) + errors.append(error) + if drain_requested: + try: + self._drain_signal.complete() + except Exception as error: + self.log.exception( + "Failed to complete %s lane drain during lifecycle cleanup", + self.name, + ) + errors.append(error) + + if errors: + raise errors[0] def reset_sync_send_state_after_fork(self) -> None: """Replace sync-send state inherited from threads that did not survive fork.""" @@ -457,20 +539,22 @@ def reset_sync_send_state_after_fork(self) -> None: self._start_lock = threading.Lock() self._sync_sends_done = threading.Condition(self._start_lock) - def rebuild_after_fork(self) -> None: + def rebuild_after_fork(self, *, closed: bool) -> None: """Replace fork-unsafe lane state in a forked child. Threads do not survive fork() and queue.Queue internal locks may be in an inconsistent state, so the queue, lock, and consumer pool are replaced. Inherited queue items are not retained as they'll be handled - by the parent process's consumers. An eager lane restarts immediately; - a lazy lane returns to not-started and restarts on next use. + by the parent process's consumers. ``closed`` normalizes every lane to + the client's fork-visible lifecycle state. An eager open lane restarts + immediately; a lazy lane returns to not-started and restarts on next use. """ 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 + self._closed = closed if self._eager_start: self.start() @@ -565,7 +649,10 @@ def __init__( max_queue_size: Maximum number of events buffered before upload. send: If False, queueing succeeds but events are not sent. on_error: Optional callback invoked by background consumers when an - upload fails. + upload fails. Keep it short and non-blocking. Calling lifecycle + methods directly is safe and deferred, but do not start another + thread or task that calls ``flush()``, ``join()``, or + ``shutdown()`` and then wait for it from the callback. flush_at: Number of queued events that triggers a batch upload. flush_interval: Maximum seconds a background consumer waits before flushing a partial batch. @@ -668,6 +755,24 @@ def __init__( self.debug = debug self.send = send self.sync_mode = sync_mode + self._lifecycle_lock = threading.Lock() + self._lifecycle_condition = threading.Condition(self._lifecycle_lock) + self._lifecycle_owner: Optional[threading.Thread] = None + self._workers_joined = False + self._join_cleanup_complete = False + self._join_requested = False + self._shutdown_requested = False + self._shutdown_complete = False + self._lifecycle_cleanup_failed = False + self._deferred_lifecycle_thread_pending = False + self._deferred_lifecycle_dirty = False + self._lifecycle_callback_context: ContextVar[bool] = ContextVar( + "posthog_lifecycle_callback", default=False + ) + self._deferred_flush_lock = threading.Lock() + self._deferred_flush_pending = False + self._deferred_flush_followup = False + self._deferred_flush_followup_timeout: Optional[float] = None # Used for session replay URL generation - we don't want the server host here. self.raw_host = normalize_host(host) self.host = determine_server_host(host) @@ -823,10 +928,9 @@ def __init__( # On program exit, allow the consumer threads to exit cleanly. # This prevents exceptions and a messy shutdown when the # interpreter is destroyed before the daemon threads finish - # execution. However, it is *not* the same as flushing the queue! - # To guarantee all messages have been delivered, you'll still need - # to call flush(). - atexit.register(self.join) + # execution. Exit performs only a short best-effort flush; call + # flush() or shutdown() explicitly when blocking completion matters. + atexit.register(self._atexit) lane_defaults = dict( api_key=self.api_key, @@ -1947,20 +2051,24 @@ 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`). """ + terminal_requested = ( + self._join_requested or self._shutdown_requested or self._workers_joined + ) 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: - self.poller = Poller( - interval=timedelta(seconds=self.poll_interval), - execute=self._load_feature_flags, - ) - self.poller.start() - else: - self.poller = None + lane.rebuild_after_fork(closed=terminal_requested) + + self._lifecycle_lock = threading.Lock() + self._lifecycle_condition = threading.Condition(self._lifecycle_lock) + self._lifecycle_owner = None + self._deferred_lifecycle_thread_pending = False + self._deferred_lifecycle_dirty = False + self._lifecycle_callback_context = ContextVar( + "posthog_lifecycle_callback", default=False + ) + self._deferred_flush_lock = threading.Lock() + self._deferred_flush_pending = False + self._deferred_flush_followup = False + self._deferred_flush_followup_timeout = None # Async runner threads do not survive fork(); recreate lazily on next async cache call. self._flag_definition_cache_provider_async_runner = None @@ -1984,6 +2092,18 @@ def _reinit_after_fork(self): reset_sessions() + # Start child threads only after replacing every lock they can touch. + if terminal_requested: + self.poller = None + elif self.enable_local_evaluation: + self.poller = Poller( + interval=timedelta(seconds=self.poll_interval), + execute=self._load_feature_flags, + ) + self.poller.start() + else: + self.poller = None + def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): # type: (...) -> Optional[str] """Push a new `msg` onto a lane's queue (analytics when unspecified), return the event uuid or None.""" @@ -2180,6 +2300,8 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: posthog.flush() # Ensures the event is sent immediately ``` """ + if self._defer_flush_from_callback(timeout_seconds): + return try: if timeout_seconds is None: for lane in self._lanes: @@ -2195,53 +2317,345 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: self.log.exception("error flushing queue: %s", e) return + def _is_consumer_thread(self) -> bool: + current = threading.current_thread() + return any(current in lane.consumers for lane in self._lanes) + + def _is_lifecycle_callback_thread(self) -> bool: + if self._lifecycle_callback_context.get(): + return True + current = threading.current_thread() + if self._is_consumer_thread() or current is self.poller: + return True + runner = self._flag_definition_cache_provider_async_runner + return runner is not None and runner.owns_thread(current) + + def _start_lifecycle_thread(self, target, name: str, *args) -> None: + threading.Thread( + target=target, + args=args, + name=f"posthog-{name}", + daemon=True, + ).start() + + def _defer_lifecycle_from_callback(self) -> bool: + if not self._is_lifecycle_callback_thread(): + return False + with self._lifecycle_lock: + self._deferred_lifecycle_dirty = True + if self._deferred_lifecycle_thread_pending: + return True + self._deferred_lifecycle_thread_pending = True + + def run() -> None: + while True: + with self._lifecycle_lock: + self._deferred_lifecycle_dirty = False + require_shutdown = self._shutdown_requested + operation = "shutdown" if require_shutdown else "join" + try: + self._run_lifecycle(require_shutdown=require_shutdown) + except BaseException: + self.log.exception("Deferred %s failed", operation) + + with self._lifecycle_lock: + if self._deferred_lifecycle_dirty: + continue + self._deferred_lifecycle_thread_pending = False + return + + self._start_lifecycle_thread(run, "lifecycle") + return True + + def _defer_flush_from_callback(self, timeout_seconds: Optional[float]) -> bool: + if not self._is_lifecycle_callback_thread(): + return False + with self._deferred_flush_lock: + if self._deferred_flush_pending: + if not self._deferred_flush_followup: + self._deferred_flush_followup = True + self._deferred_flush_followup_timeout = timeout_seconds + elif ( + self._deferred_flush_followup_timeout is not None + and timeout_seconds is not None + ): + self._deferred_flush_followup_timeout = max( + self._deferred_flush_followup_timeout, timeout_seconds + ) + else: + self._deferred_flush_followup_timeout = None + return True + self._deferred_flush_pending = True + + def run() -> None: + next_timeout = timeout_seconds + while True: + self.flush(next_timeout) + with self._deferred_flush_lock: + if self._deferred_flush_followup: + next_timeout = self._deferred_flush_followup_timeout + self._deferred_flush_followup = False + self._deferred_flush_followup_timeout = None + continue + self._deferred_flush_pending = False + return + + self._start_lifecycle_thread(run, "flush") + return True + + def _run_lifecycle_cleanup( + self, + log_message: str, + cleanup: Callable[[], None], + errors: list[Exception], + ) -> None: + """Attempt one cleanup step without preventing later independent steps.""" + try: + cleanup() + except Exception as error: + self.log.exception(log_message) + errors.append(error) + + def _flush_or_discard_queues(self, errors: list[Exception]) -> None: + for lane in self._lanes: + try: + if any(consumer.is_alive() for consumer in lane.consumers): + lane.flush(timeout_seconds=None) + else: + lane.discard_undrainable_queued_work() + except Exception as error: + self.log.exception( + "Failed to drain %s lane during lifecycle cleanup", lane.name + ) + errors.append(error) + + def _join_once( + self, + errors: list[Exception], + flush_queues: bool = True, + *, + lanes_prepared: bool = False, + ) -> None: + if not self._workers_joined: + if not lanes_prepared: + for lane in self._lanes: + self._run_lifecycle_cleanup( + f"Failed to close {lane.name} lane during lifecycle cleanup", + lane.close, + errors, + ) + for lane in self._lanes: + self._run_lifecycle_cleanup( + f"Failed waiting for {lane.name} synchronous sends during lifecycle cleanup", + lane.wait_for_sync_sends, + errors, + ) + if flush_queues: + self._flush_or_discard_queues(errors) + for lane in self._lanes: + self._run_lifecycle_cleanup( + f"Failed to stop {lane.name} lane during lifecycle cleanup", + lane.join, + errors, + ) + # Ordinary cleanup failures are logged by each step. Reaching here + # means every worker cleanup step was attempted once. + self._workers_joined = True + + if not self._join_cleanup_complete: + if self.poller: + self._run_lifecycle_cleanup( + "Failed to stop feature flag poller during lifecycle cleanup", + self.poller.stop, + errors, + ) + + self._run_lifecycle_cleanup( + "Failed to shut down feature flag cache provider during lifecycle cleanup", + self._shutdown_flag_definition_cache_provider, + errors, + ) + self._run_lifecycle_cleanup( + "Failed to unregister client during lifecycle cleanup", + self._unregister_duplicate_client, + errors, + ) + self._join_cleanup_complete = True + + def _shutdown_once(self, errors: list[Exception]) -> None: + if not self._workers_joined: + # 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: + self._run_lifecycle_cleanup( + f"Failed to close {lane.name} lane during shutdown", + lane.close, + errors, + ) + for lane in self._lanes: + self._run_lifecycle_cleanup( + f"Failed waiting for {lane.name} synchronous sends during shutdown", + lane.wait_for_sync_sends, + errors, + ) + self._flush_or_discard_queues(errors) + + if self._metrics is not None: + self._run_lifecycle_cleanup( + "Failed to flush metrics on shutdown", self._metrics.flush, errors + ) + self._run_lifecycle_cleanup( + "Failed to reset metrics on shutdown", self._metrics.reset, errors + ) + self._join_once(errors, flush_queues=False, lanes_prepared=True) + self._run_lifecycle_cleanup( + "Failed to clear feature flag deduplication state on shutdown", + self.distinct_ids_feature_flags_reported.clear, + errors, + ) + + if self.exception_capture: + self._run_lifecycle_cleanup( + "Failed to close exception capture on shutdown", + self.exception_capture.close, + errors, + ) + # Ordinary cleanup failures are logged by each step. Reaching here + # means every shutdown cleanup step was attempted once. + self._shutdown_complete = True + + def _run_lifecycle(self, require_shutdown: bool = False) -> None: + while True: + with self._lifecycle_condition: + if require_shutdown and self._shutdown_complete: + if self.debug and self._lifecycle_cleanup_failed: + raise RuntimeError("client lifecycle cleanup failed") + return + if not require_shutdown and ( + self._join_cleanup_complete or self._shutdown_complete + ): + if self.debug and self._lifecycle_cleanup_failed: + raise RuntimeError("client lifecycle cleanup failed") + return + if self._lifecycle_owner is not None: + if self._is_lifecycle_callback_thread() or ( + threading.current_thread() is self._lifecycle_owner + ): + return + self._lifecycle_condition.wait() + continue + self._lifecycle_owner = threading.current_thread() + + try: + errors: list[Exception] = [] + while True: + with self._lifecycle_lock: + run_shutdown = self._shutdown_requested + + if run_shutdown: + self._shutdown_once(errors) + else: + self._join_once(errors) + + with self._lifecycle_condition: + if ( + self._shutdown_requested + and not self._shutdown_complete + and not run_shutdown + ): + continue + if errors: + self._lifecycle_cleanup_failed = True + raise errors[0] + self._lifecycle_owner = None + self._lifecycle_condition.notify_all() + return + except BaseException: + with self._lifecycle_condition: + self._lifecycle_owner = None + self._lifecycle_condition.notify_all() + raise + + @no_throw() + def _atexit(self) -> None: + """Make a bounded delivery attempt, then stop daemon workers.""" + with self._lifecycle_condition: + # A daemon lifecycle worker already owns cleanup. Do not wait for it + # at interpreter exit; the process must remain free to terminate. + if self._lifecycle_owner is not None: + return + self._lifecycle_owner = threading.current_thread() + + try: + try: + for lane in self._lanes: + lane.close() + + deadline = _get_atexit_deadline() + for lane in self._lanes: + lane.flush(max(0.0, deadline - time.monotonic())) + finally: + # Consumers are daemon threads. Publish a non-draining stop to + # every consumer, but do not join in-flight requests at exit. + for lane in self._lanes: + for consumer in lane.consumers: + consumer.pause() + finally: + with self._lifecycle_condition: + self._lifecycle_owner = None + self._lifecycle_condition.notify_all() + + @no_throw() def join(self) -> None: """ - End the consumer thread once the queue is empty. Do not use directly, call `shutdown()` instead. + Attempt to process queued events and end the consumer threads. Do not use directly, call `shutdown()` instead. + + Failed or undrainable events may be dropped and reported through logging + or ``on_error``; returning does not guarantee server receipt. Lifecycle + cleanup is attempted once, and cleanup failures are logged without retry. Examples: ```python posthog.join() ``` """ - for lane in self._lanes: - lane.join() - - if self.poller: - self.poller.stop() - - # Shutdown the cache provider (release locks, cleanup) - self._shutdown_flag_definition_cache_provider() - self._unregister_duplicate_client() + with self._lifecycle_lock: + self._join_requested = True + if self._defer_lifecycle_from_callback(): + return + self._run_lifecycle() + @no_throw() def shutdown(self) -> None: """ Flush all messages and cleanly shutdown the client. Call this before the process ends in serverless environments to avoid data loss. + Normally this method blocks until queued events have been attempted and + cleanup finishes. Failed or undrainable events may be dropped and + reported through logging or ``on_error``; returning does not guarantee + server receipt. Lifecycle cleanup is attempted once, and cleanup failures + are logged without retry. When called directly from an SDK callback such as + ``on_error``, shutdown is deferred to avoid blocking the worker that + invoked the callback. If the callback must coordinate a blocking + shutdown, have it signal an + application-owned thread and return before that thread calls shutdown. + Do not wait inside the callback for another thread or task that calls a + lifecycle method. + Examples: ```python 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: - self._metrics.flush() - except Exception: - self.log.exception("Failed to flush metrics on shutdown") - self._metrics.reset() - self.join() - self.distinct_ids_feature_flags_reported.clear() - - if self.exception_capture: - self.exception_capture.close() + with self._lifecycle_lock: + if self._shutdown_complete: + if self.debug and self._lifecycle_cleanup_failed: + raise RuntimeError("client lifecycle cleanup failed") + return + self._shutdown_requested = True + if self._defer_lifecycle_from_callback(): + return + self._run_lifecycle(require_shutdown=True) def _resolve_flag_definition_cache_provider_result(self, result): if not inspect.isawaitable(result): @@ -2252,7 +2666,11 @@ def _resolve_flag_definition_cache_provider_result(self, result): self._flag_definition_cache_provider_async_runner = ( _BackgroundEventLoopRunner() ) - return self._flag_definition_cache_provider_async_runner.run(result) + token = self._lifecycle_callback_context.set(True) + try: + return self._flag_definition_cache_provider_async_runner.run(result) + finally: + self._lifecycle_callback_context.reset(token) def _shutdown_flag_definition_cache_provider(self): if not self._flag_definition_cache_provider: diff --git a/posthog/consumer.py b/posthog/consumer.py index 046c02e9..bc41f4c1 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -51,8 +51,11 @@ def complete(self) -> None: self._requests -= 1 self._queue.not_empty.notify_all() - def wake(self) -> None: + def stop(self, consumer, drain: bool) -> None: + """Publish a consumer stop under the queue's dequeue lock.""" with self._queue.not_empty: + consumer.running = False + consumer._drain_on_stop = drain self._queue.not_empty.notify_all() def wait_until_inactive_or_work(self, consumer) -> None: @@ -65,12 +68,23 @@ 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.""" + def draining(self, consumer) -> bool: + with self._queue.mutex: + return self._requests > 0 and (consumer.running or consumer._drain_on_stop) + + def get(self, timeout: float, consumer=None): + """Get an item, or wake with ``Empty`` when draining or stopping.""" with self._queue.not_empty: deadline = time.monotonic() + timeout - while not self._queue._qsize(): - if self._requests: + while True: + draining = self._requests > 0 and ( + consumer is None or consumer.running or consumer._drain_on_stop + ) + if consumer is not None and not consumer.running and not draining: + raise Empty + if self._queue._qsize(): + break + if draining: raise Empty remaining = deadline - time.monotonic() if remaining <= 0: @@ -119,6 +133,7 @@ def __init__( self.capture_mode = capture_mode self.capture_compression = capture_compression self._drain_signal: Optional[_DrainSignal] = None + self._drain_on_stop = False # 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 @@ -139,10 +154,15 @@ def run(self): self.log.debug("consumer exited.") def pause(self): - """Pause the consumer.""" - self.running = False + """Pause the consumer without admitting additional queued work.""" + self._pause(drain=False) + + def _pause(self, drain: bool) -> None: if self._drain_signal is not None: - self._drain_signal.wake() + self._drain_signal.stop(self, drain) + else: + self.running = False + self._drain_on_stop = drain def upload(self): """Upload the next batch of items, return whether successful.""" @@ -152,16 +172,19 @@ def upload(self): return False try: - self.request(batch) - success = True - except Exception as e: - self.log.error("error uploading: %s", e) - success = False - if self.on_error: - try: - self.on_error(e, batch) - except Exception as e: - self.log.error("on_error handler failed: %s", e) + if not self._can_upload(): + return False + try: + self.request(batch) + success = True + except Exception as e: + self.log.error("error uploading: %s", e) + success = False + if self.on_error: + try: + self.on_error(e, batch) + except Exception as e: + self.log.error("on_error handler failed: %s", e) finally: # mark items as acknowledged from queue for item in batch: @@ -173,7 +196,20 @@ 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 + return ( + self._drain_signal.draining(self) + if self._drain_signal is not None + else False + ) + + def _can_upload(self) -> bool: + if self._drain_signal is None: + return self.running or self._drain_on_stop + # This lock-protected check is the request admission point. A request + # admitted before a stop may finish on the daemon consumer, but a stop + # prevents any later buffered or queued batch from being admitted. + with self.queue.mutex: + return self.running or self._drain_on_stop def next(self): """Return the next batch of items to upload.""" @@ -182,51 +218,70 @@ def next(self): start_time = time.monotonic() total_size = 0 + pending_items = 0 - while len(items) < self.flush_at: - # 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: + while len(items) < self.flush_at: + # 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() + if not self.running and not draining: + break + remaining = self.flush_interval - (time.monotonic() - start_time) + if not draining and remaining <= 0: + break - try: - 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: - # Callback-modified events can still contain invalid mapping - # keys or circular references. Never log the payload here. - self.log.error( - "Unable to serialize queued event for sizing, dropping." - ) - queue.task_done() - continue - if item_size > self.max_msg_size: - # Log only name and size: AI events may carry unredacted - # multimodal payloads that must not leak into logs. - self.log.error( - "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", - item.get("event") if isinstance(item, dict) else type(item), - item_size, - self.max_msg_size // 1024, - self.endpoint, - ) - queue.task_done() - continue - items.append(item) - total_size += item_size - if total_size >= BATCH_SIZE_LIMIT: - self.log.debug("hit batch size limit (size: %d)", total_size) + if self._drain_signal is not None: + item = self._drain_signal.get( + timeout=0 if draining else remaining, + consumer=self, + ) + else: + item = queue.get(block=True, timeout=remaining) + pending_items += 1 + try: + item_size = len( + json.dumps(item, cls=DatetimeSerializer).encode() + ) + except Exception: + # Callback-modified events can still contain invalid mapping + # keys or circular references. Never log the payload here. + self.log.error( + "Unable to serialize queued event for sizing, dropping." + ) + queue.task_done() + pending_items -= 1 + continue + if item_size > self.max_msg_size: + # Log only name and size: AI events may carry unredacted + # multimodal payloads that must not leak into logs. + self.log.error( + "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", + item.get("event") if isinstance(item, dict) else type(item), + item_size, + self.max_msg_size // 1024, + self.endpoint, + ) + queue.task_done() + pending_items -= 1 + continue + items.append(item) + total_size += item_size + if total_size >= BATCH_SIZE_LIMIT: + self.log.debug("hit batch size limit (size: %d)", total_size) + break + except Empty: break - except Empty: - break + except BaseException: + for _ in range(pending_items): + queue.task_done() + raise + + if not self._can_upload(): + for _ in range(pending_items): + queue.task_done() + return [] return items diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 4cf708a2..61fa203b 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -314,15 +314,15 @@ def test_fork_rebuild_restarts_analytics_and_resets_ai(self): ) client.join() - def test_fork_rebuild_noop_for_sync_mode(self): + def test_fork_rebuild_replaces_sync_mode_queues(self): client = Client(TEST_API_KEY, sync_mode=True) old_analytics_queue = client._analytics_lane.queue old_ai_queue = client._ai_lane.queue client._reinit_after_fork() - self.assertIs(client._analytics_lane.queue, old_analytics_queue) - self.assertIs(client._ai_lane.queue, old_ai_queue) + self.assertIsNot(client._analytics_lane.queue, old_analytics_queue) + self.assertIsNot(client._ai_lane.queue, old_ai_queue) class TestCaptureAiEventHelper(unittest.TestCase): diff --git a/posthog/test/test_async_utils.py b/posthog/test/test_async_utils.py new file mode 100644 index 00000000..1eb7c783 --- /dev/null +++ b/posthog/test/test_async_utils.py @@ -0,0 +1,306 @@ +import asyncio +import threading +import time +import unittest +from unittest import mock + +from posthog._async_utils import _BackgroundEventLoopRunner, _LoopStartup + + +class _PausingEvent: + def __init__(self) -> None: + self._completed = threading.Event() + self.waiter_paused = threading.Event() + self.release_waiter = threading.Event() + + def set(self) -> None: + self._completed.set() + + def wait(self, timeout=None) -> bool: + if not self._completed.wait(timeout): + return False + self.waiter_paused.set() + return self.release_waiter.wait(2) + + +class TestBackgroundEventLoopRunner(unittest.TestCase): + def test_startup_error_is_reported(self): + runner = _BackgroundEventLoopRunner() + awaitable = asyncio.sleep(0) + + with mock.patch( + "posthog._async_utils.asyncio.new_event_loop", + side_effect=RuntimeError("startup failed"), + ): + with self.assertRaisesRegex(RuntimeError, "startup failed"): + runner.run(awaitable) + + self.assertIsNone(awaitable.cr_frame) + + def test_close_waits_for_run_during_startup(self): + runner = _BackgroundEventLoopRunner() + construction_started = threading.Event() + release_construction = threading.Event() + run_errors = [] + + original_new_event_loop = asyncio.new_event_loop + + def create_loop(): + construction_started.set() + self.assertTrue(release_construction.wait(2)) + return original_new_event_loop() + + def run(): + awaitable = asyncio.sleep(0) + try: + runner.run(awaitable) + except BaseException as error: + awaitable.close() + run_errors.append(error) + + with mock.patch( + "posthog._async_utils.asyncio.new_event_loop", side_effect=create_loop + ): + run_thread = threading.Thread(target=run) + run_thread.start() + self.assertTrue(construction_started.wait(1)) + + close_thread = threading.Thread(target=runner.close) + close_thread.start() + deadline = time.monotonic() + 1 + while not runner._close_requested: + if time.monotonic() >= deadline: + self.fail("close did not reach startup state") + time.sleep(0.001) + release_construction.set() + run_thread.join(2) + close_thread.join(2) + + self.assertFalse(run_thread.is_alive()) + self.assertFalse(close_thread.is_alive()) + self.assertEqual(len(run_errors), 1) + self.assertRegex(str(run_errors[0]), "closed during startup") + self.assertIsNone(runner._thread) + self.assertIsNone(runner._loop) + + def test_run_retries_when_close_wins_after_startup_completes(self): + runner = _BackgroundEventLoopRunner() + first_startup = _LoopStartup() + first_startup.done = _PausingEvent() # type: ignore[assignment] + second_startup = _LoopStartup() + run_results = [] + run_errors = [] + + async def result(): + return 42 + + def run(): + try: + run_results.append(runner.run(result())) + except BaseException as error: + run_errors.append(error) + + with mock.patch( + "posthog._async_utils._LoopStartup", + side_effect=[first_startup, second_startup], + ): + run_thread = threading.Thread(target=run) + run_thread.start() + self.assertTrue(first_startup.done.waiter_paused.wait(1)) + + runner.close() + first_startup.done.release_waiter.set() + run_thread.join(2) + + runner.close() + self.assertFalse(run_thread.is_alive()) + self.assertEqual(run_errors, []) + self.assertEqual(run_results, [42]) + + def test_startup_failure_state_is_not_overwritten_by_next_attempt(self): + runner = _BackgroundEventLoopRunner() + first_startup = _LoopStartup() + first_startup.done = _PausingEvent() # type: ignore[assignment] + second_startup = _LoopStartup() + original_new_event_loop = asyncio.new_event_loop + loop_attempt = 0 + first_errors = [] + second_results = [] + + def create_loop(): + nonlocal loop_attempt + loop_attempt += 1 + if loop_attempt == 1: + raise RuntimeError("first startup failed") + return original_new_event_loop() + + async def result(): + return 42 + + first_awaitable = asyncio.sleep(0) + + def first_run(): + try: + runner.run(first_awaitable) + except BaseException as error: + first_errors.append(error) + + def second_run(): + second_results.append(runner.run(result())) + + with ( + mock.patch( + "posthog._async_utils._LoopStartup", + side_effect=[first_startup, second_startup], + ), + mock.patch( + "posthog._async_utils.asyncio.new_event_loop", + side_effect=create_loop, + ), + ): + first_thread = threading.Thread(target=first_run) + first_thread.start() + self.assertTrue(first_startup.done.waiter_paused.wait(1)) + + second_thread = threading.Thread(target=second_run) + second_thread.start() + second_thread.join(2) + + first_startup.done.release_waiter.set() + first_thread.join(2) + + runner.close() + self.assertFalse(first_thread.is_alive()) + self.assertFalse(second_thread.is_alive()) + self.assertEqual(len(first_errors), 1) + self.assertRegex(str(first_errors[0]), "first startup failed") + self.assertEqual(second_results, [42]) + self.assertIsNone(first_awaitable.cr_frame) + + def test_close_closes_awaitable_cancelled_before_first_step(self): + runner = _BackgroundEventLoopRunner() + runner.run(asyncio.sleep(0)) + loop = runner._loop + self.assertIsNotNone(loop) + + loop_blocked = threading.Event() + release_loop = threading.Event() + scheduled = threading.Event() + run_errors = [] + original_run_coroutine_threadsafe = asyncio.run_coroutine_threadsafe + + def block_loop(): + loop_blocked.set() + self.assertTrue(release_loop.wait(2)) + + def schedule(coro, target_loop): + future = original_run_coroutine_threadsafe(coro, target_loop) + scheduled.set() + return future + + loop.call_soon_threadsafe(block_loop) # type: ignore[union-attr] + self.assertTrue(loop_blocked.wait(1)) + awaitable = asyncio.sleep(0) + + def run(): + try: + runner.run(awaitable) + except BaseException as error: + run_errors.append(error) + + with mock.patch( + "posthog._async_utils.asyncio.run_coroutine_threadsafe", + side_effect=schedule, + ): + run_thread = threading.Thread(target=run) + run_thread.start() + self.assertTrue(scheduled.wait(1)) + + close_thread = threading.Thread(target=runner.close) + close_thread.start() + release_loop.set() + run_thread.join(2) + close_thread.join(2) + + self.assertFalse(run_thread.is_alive()) + self.assertFalse(close_thread.is_alive()) + self.assertEqual(len(run_errors), 1) + self.assertIsNone(awaitable.cr_frame) + + def test_runner_preserves_configured_event_loop_policy(self): + runner = _BackgroundEventLoopRunner() + original_policy = asyncio.get_event_loop_policy() + + class Policy(asyncio.DefaultEventLoopPolicy): + loop = None + + def new_event_loop(self): + self.loop = super().new_event_loop() + return self.loop + + policy = Policy() + + async def running_loop(): + return asyncio.get_running_loop() + + try: + asyncio.set_event_loop_policy(policy) + self.assertIs(runner.run(running_loop()), policy.loop) + runner.close() + finally: + asyncio.set_event_loop_policy(original_policy) + + def test_runner_uses_context_aware_fallback_for_read_only_policy_loop(self): + runner = _BackgroundEventLoopRunner() + original_policy = asyncio.get_event_loop_policy() + + class ReadOnlyLoop(asyncio.SelectorEventLoop): + def __setattr__(self, name, value): + if name == "run_in_executor": + raise AttributeError("run_in_executor is read-only") + super().__setattr__(name, value) + + class Policy(asyncio.DefaultEventLoopPolicy): + loop = None + + def new_event_loop(self): + self.loop = ReadOnlyLoop() + return self.loop + + policy = Policy() + + async def running_loop(): + return asyncio.get_running_loop() + + try: + asyncio.set_event_loop_policy(policy) + running = runner.run(running_loop()) + self.assertIsNot(running, policy.loop) + self.assertTrue(policy.loop.is_closed()) + finally: + runner.close() + asyncio.set_event_loop_policy(original_policy) + + def test_run_from_runner_thread_fails_instead_of_deadlocking(self): + runner = _BackgroundEventLoopRunner() + + async def reenter(): + awaitable = asyncio.sleep(0) + try: + with self.assertRaisesRegex(RuntimeError, "runner thread"): + runner.run(awaitable) + finally: + awaitable.close() + + runner.run(reenter()) + runner.close() + + def test_close_from_runner_thread_allows_fresh_loop(self): + runner = _BackgroundEventLoopRunner() + + async def close_runner(): + runner.close() + + runner.run(close_runner()) + runner.run(asyncio.sleep(0)) + runner.close() diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index f030ff94..8b335dcd 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -1,9 +1,14 @@ +import contextlib import logging import asyncio +import subprocess +import sys +import textwrap import threading import time import unittest import warnings +from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor from datetime import datetime from unittest import mock from uuid import UUID, uuid4 @@ -25,6 +30,15 @@ # Legacy single-flag behavior remains covered here; warning emission itself is # asserted in test_evaluate_flags.py. +def _wait_until(predicate, timeout=3): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + pytestmark = [ pytest.mark.filterwarnings( r"ignore:`(feature_enabled|get_feature_flag|get_feature_flag_payload)` is deprecated:DeprecationWarning" @@ -2239,6 +2253,90 @@ def test_shutdown(self): for consumer in client.consumers: self.assertFalse(consumer.is_alive()) + def test_atexit_registers_bounded_worker_cleanup(self): + with mock.patch("posthog.client.atexit.register") as register: + client = Client(FAKE_TEST_API_KEY) + + register.assert_called_once_with(client._atexit) + self.addCleanup(client.shutdown) + + def test_atexit_bounds_flush_and_stops_consumers_without_joining(self): + client = Client(FAKE_TEST_API_KEY, send=False) + lanes = [mock.Mock(), mock.Mock()] + for lane in lanes: + lane.consumers = [mock.Mock(), mock.Mock()] + client._lanes = lanes + + client._atexit() + + for lane in lanes: + lane.close.assert_called_once_with() + lane.flush.assert_called_once() + timeout = lane.flush.call_args.args[0] + self.assertGreaterEqual(timeout, 0) + self.assertLessEqual(timeout, 1) + lane.wait_for_sync_sends.assert_not_called() + lane.join.assert_not_called() + for consumer in lane.consumers: + consumer.pause.assert_called_once_with() + + def test_atexit_subprocess_does_not_drain_queued_backlog(self): + script = textwrap.dedent( + """ + import threading + import time + from unittest import mock + + from posthog.client import Client + from posthog.consumer import Consumer + + with mock.patch.object(Consumer, "start"): + clients = [ + Client(f"test-key-{index}", flush_at=100, flush_interval=60) + for index in range(3) + ] + + for client in clients: + consumer = client.consumers[0] + consumer.request = lambda batch: time.sleep(10) + consumer.start() + time.sleep(0.1) + + queues = [client.queue for client in clients] + for queue in queues: + for index in range(10): + queue.put({"event": str(index), "distinct_id": "test"}) + + deadline = time.monotonic() + 1 + while any(not queue.empty() for queue in queues): + if time.monotonic() >= deadline: + raise RuntimeError("consumer did not buffer the queued backlog") + time.sleep(0.001) + """ + ) + + subprocess.run( + [sys.executable, "-c", script], + check=True, + capture_output=True, + text=True, + timeout=3, + ) + + def test_atexit_does_not_wait_for_active_lifecycle_owner(self): + client = Client(FAKE_TEST_API_KEY, send=False) + owner = mock.Mock() + client._lifecycle_owner = owner + + lane = mock.Mock() + client._lanes = [lane] + + client._atexit() + + lane.close.assert_not_called() + lane.flush.assert_not_called() + self.assertIs(client._lifecycle_owner, owner) + def test_shutdown_clears_feature_flag_called_dedupe_cache(self): client = Client(FAKE_TEST_API_KEY, send=False, thread=0) client.distinct_ids_feature_flags_reported["user"] = {("flag", True, ())} @@ -2248,13 +2346,788 @@ def test_shutdown_clears_feature_flag_called_dedupe_cache(self): self.assertEqual(len(client.distinct_ids_feature_flags_reported), 0) def test_shutdown_flushes_without_timeout(self): - client = Client(FAKE_TEST_API_KEY, send=False, thread=0) + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) - with mock.patch.object(client, "flush") as mock_flush: + with mock.patch.object(client._analytics_lane, "flush") as mock_flush: client.shutdown() mock_flush.assert_called_once_with(timeout_seconds=None) + def test_callback_shutdown_escalates_pending_deferred_join(self): + client = Client(FAKE_TEST_API_KEY) + first_run_started = threading.Event() + release_first_run = threading.Event() + shutdown_run_complete = threading.Event() + require_shutdown_calls = [] + + def run_lifecycle(require_shutdown=False): + require_shutdown_calls.append(require_shutdown) + if len(require_shutdown_calls) == 1: + first_run_started.set() + self.assertTrue(release_first_run.wait(2)) + else: + shutdown_run_complete.set() + + with ( + mock.patch.object( + client, "_is_lifecycle_callback_thread", return_value=True + ), + mock.patch.object(client, "_run_lifecycle", side_effect=run_lifecycle), + mock.patch.object( + client, + "_start_lifecycle_thread", + wraps=client._start_lifecycle_thread, + ) as start_thread, + ): + client.join() + self.assertTrue(first_run_started.wait(1)) + client.shutdown() + release_first_run.set() + self.assertTrue(shutdown_run_complete.wait(1)) + + start_thread.assert_called_once() + self.assertEqual(start_thread.call_args.args[1], "lifecycle") + self.assertEqual(require_shutdown_calls, [False, True]) + self.assertTrue( + _wait_until(lambda: not client._deferred_lifecycle_thread_pending) + ) + self.assertFalse(client._deferred_lifecycle_dirty) + + def test_callback_shutdown_escalates_failed_deferred_join(self): + client = Client(FAKE_TEST_API_KEY) + first_run_started = threading.Event() + release_first_run = threading.Event() + shutdown_run_complete = threading.Event() + require_shutdown_calls = [] + + def run_lifecycle(require_shutdown=False): + require_shutdown_calls.append(require_shutdown) + if len(require_shutdown_calls) == 1: + first_run_started.set() + self.assertTrue(release_first_run.wait(2)) + raise Exception("join cleanup failed") + shutdown_run_complete.set() + + with ( + mock.patch.object( + client, "_is_lifecycle_callback_thread", return_value=True + ), + mock.patch.object(client, "_run_lifecycle", side_effect=run_lifecycle), + mock.patch.object( + client, + "_start_lifecycle_thread", + wraps=client._start_lifecycle_thread, + ) as start_thread, + ): + client.join() + self.assertTrue(first_run_started.wait(1)) + client.shutdown() + release_first_run.set() + self.assertTrue(shutdown_run_complete.wait(1)) + + start_thread.assert_called_once() + self.assertEqual(require_shutdown_calls, [False, True]) + self.assertTrue( + _wait_until(lambda: not client._deferred_lifecycle_thread_pending) + ) + self.assertFalse(client._deferred_lifecycle_dirty) + + def test_deferred_lifecycle_worker_is_daemon(self): + client = Client(FAKE_TEST_API_KEY) + target = mock.Mock() + + with mock.patch("posthog.client.threading.Thread") as thread: + client._start_lifecycle_thread(target, "lifecycle") + + thread.assert_called_once_with( + target=target, + args=(), + name="posthog-lifecycle", + daemon=True, + ) + thread.return_value.start.assert_called_once_with() + + def test_deferred_lifecycle_logs_selected_operation(self): + client = Client(FAKE_TEST_API_KEY) + + with ( + mock.patch.object( + client, "_is_lifecycle_callback_thread", return_value=True + ), + mock.patch.object( + client, + "_run_lifecycle", + side_effect=Exception("cleanup failed"), + ) as run_lifecycle, + mock.patch.object(client.log, "exception") as log_exception, + ): + client.shutdown() + self.assertTrue( + _wait_until(lambda: not client._deferred_lifecycle_thread_pending) + ) + + run_lifecycle.assert_called_once_with(require_shutdown=True) + log_exception.assert_called_once_with("Deferred %s failed", "shutdown") + + def test_callback_flushes_are_coalesced_with_strongest_followup(self): + client = Client(FAKE_TEST_API_KEY) + first_flush_started = threading.Event() + release_first_flush = threading.Event() + followup_complete = threading.Event() + timeouts = [] + + def flush(timeout_seconds): + timeouts.append(timeout_seconds) + if len(timeouts) == 1: + first_flush_started.set() + self.assertTrue(release_first_flush.wait(2)) + else: + followup_complete.set() + + with ( + mock.patch.object( + client, "_is_lifecycle_callback_thread", return_value=True + ), + mock.patch.object(client, "flush", side_effect=flush), + ): + self.assertTrue(client._defer_flush_from_callback(0)) + self.assertTrue(first_flush_started.wait(1)) + self.assertTrue(client._defer_flush_from_callback(1)) + self.assertTrue(client._defer_flush_from_callback(None)) + release_first_flush.set() + self.assertTrue(followup_complete.wait(1)) + + self.assertEqual(timeouts, [0, None]) + + def test_sync_send_failure_does_not_invoke_async_on_error(self): + on_error = mock.Mock() + client = Client( + FAKE_TEST_API_KEY, + sync_mode=True, + on_error=on_error, + ) + + with mock.patch( + "posthog.client.batch_post", side_effect=Exception("upload failed") + ): + result = client.capture("event", distinct_id="distinct_id") + + self.assertIsNone(result) + on_error.assert_not_called() + self.assertEqual(client._analytics_lane._active_sync_sends, 0) + + def test_on_error_can_request_shutdown_with_pending_work(self): + first_send_started = threading.Event() + release_first_send = threading.Event() + callback_returned = threading.Event() + sent_events = [] + client: Client + + def request(batch): + event = batch[0]["event"] + sent_events.append(event) + if event == "first": + first_send_started.set() + self.assertTrue(release_first_send.wait(2)) + raise Exception("upload failed") + + def on_error(error, batch): + client.shutdown() + client.join() + callback_returned.set() + + client = Client( + FAKE_TEST_API_KEY, + on_error=on_error, + flush_at=1, + flush_interval=0.01, + max_retries=0, + ) + exception_capture = mock.Mock() + exception_capture.close.side_effect = Exception("cleanup failed") + client.exception_capture = exception_capture + with mock.patch.object(client.consumers[0], "request", side_effect=request): + client.capture("first", distinct_id="distinct_id") + self.assertTrue(first_send_started.wait(1)) + client.capture("second", distinct_id="distinct_id") + release_first_send.set() + + self.assertTrue(callback_returned.wait(1)) + self.assertTrue(_wait_until(lambda: client._shutdown_complete)) + + self.assertEqual(sent_events, ["first", "second"]) + self.assertEqual(client.queue.unfinished_tasks, 0) + exception_capture.close.assert_called_once_with() + self.assertTrue(all(not consumer.is_alive() for consumer in client.consumers)) + + def test_concurrent_join_waits_for_lifecycle_owner(self): + send_started = threading.Event() + release_send = threading.Event() + + def request(batch): + send_started.set() + self.assertTrue(release_send.wait(2)) + + client = Client(FAKE_TEST_API_KEY, flush_at=1) + with mock.patch.object(client.consumers[0], "request", side_effect=request): + client.capture("event", distinct_id="distinct_id") + self.assertTrue(send_started.wait(1)) + + first_join = threading.Thread(target=client.join) + second_join = threading.Thread(target=client.join) + first_join.start() + time.sleep(0.05) + second_join.start() + time.sleep(0.05) + self.assertTrue(second_join.is_alive()) + + release_send.set() + first_join.join(3) + second_join.join(3) + + self.assertFalse(first_join.is_alive()) + self.assertFalse(second_join.is_alive()) + self.assertTrue(client._join_cleanup_complete) + + def test_join_winning_shutdown_race_drains_pending_work_without_deadlock(self): + first_send_started = threading.Event() + release_first_send = threading.Event() + join_started = threading.Event() + sent_events = [] + + def request(batch): + sent_events.extend(event["event"] for event in batch) + first_send_started.set() + self.assertTrue(release_first_send.wait(2)) + + client = Client(FAKE_TEST_API_KEY, flush_at=1) + with mock.patch.object(client.consumers[0], "request", side_effect=request): + client.capture("first", distinct_id="distinct_id") + self.assertTrue(first_send_started.wait(1)) + client.capture("second", distinct_id="distinct_id") + + original_close = client._analytics_lane.close + + def observed_close(): + original_close() + join_started.set() + + with mock.patch.object( + client._analytics_lane, "close", side_effect=observed_close + ): + join_thread = threading.Thread(target=client.join) + shutdown_thread = threading.Thread(target=client.shutdown) + join_thread.start() + self.assertTrue(join_started.wait(1)) + shutdown_thread.start() + release_first_send.set() + join_thread.join(3) + shutdown_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertFalse(shutdown_thread.is_alive()) + self.assertEqual(sent_events, ["first", "second"]) + self.assertEqual(client.queue.unfinished_tasks, 0) + self.assertTrue(client._shutdown_complete) + + def test_shutdown_winning_join_race_drains_pending_work(self): + first_send_started = threading.Event() + release_first_send = threading.Event() + shutdown_started = threading.Event() + sent_events = [] + + def request(batch): + sent_events.extend(event["event"] for event in batch) + if len(sent_events) == 1: + first_send_started.set() + self.assertTrue(release_first_send.wait(2)) + + client = Client(FAKE_TEST_API_KEY, flush_at=1, flush_interval=0.01) + with mock.patch.object(client.consumers[0], "request", side_effect=request): + client.capture("first", distinct_id="distinct_id") + self.assertTrue(first_send_started.wait(1)) + client.capture("second", distinct_id="distinct_id") + + original_close = client._analytics_lane.close + + def observed_close(): + original_close() + shutdown_started.set() + + with mock.patch.object( + client._analytics_lane, "close", side_effect=observed_close + ): + shutdown_thread = threading.Thread(target=client.shutdown) + join_thread = threading.Thread(target=client.join) + shutdown_thread.start() + self.assertTrue(shutdown_started.wait(1)) + join_thread.start() + release_first_send.set() + shutdown_thread.join(3) + join_thread.join(3) + + self.assertFalse(shutdown_thread.is_alive()) + self.assertFalse(join_thread.is_alive()) + self.assertEqual(sent_events, ["first", "second"]) + self.assertEqual(client.queue.unfinished_tasks, 0) + + def test_join_publishes_terminal_intent_before_closing_lanes(self): + client = Client(FAKE_TEST_API_KEY, send=False) + original_close = client._analytics_lane.close + + def close_analytics_lane(): + self.assertTrue(client._join_requested) + original_close() + + with mock.patch.object( + client._analytics_lane, "close", side_effect=close_analytics_lane + ): + client.join() + + self.assertTrue(client._join_cleanup_complete) + + def test_shutdown_after_join_runs_shutdown_only_cleanup(self): + client = Client(FAKE_TEST_API_KEY) + metrics = mock.Mock() + exception_capture = mock.Mock() + client._metrics = metrics + client.exception_capture = exception_capture + + client.join() + client.shutdown() + + metrics.flush.assert_called_once() + metrics.reset.assert_called_once() + exception_capture.close.assert_called_once() + self.assertTrue(client._shutdown_complete) + + def test_async_cache_provider_default_executor_can_reenter_join(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + executor_called = threading.Event() + + class AsyncProvider: + async def shutdown(self): + def reenter_join(): + client.join() + executor_called.set() + + await asyncio.get_running_loop().run_in_executor(None, reenter_join) + + client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(executor_called.is_set()) + self.assertTrue(client._join_cleanup_complete) + + def test_async_cache_provider_custom_executor_can_reenter_join(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + executor_called = threading.Event() + + class DelegatingExecutor(Executor): + def __init__(self): + self.executor = ThreadPoolExecutor(max_workers=1) + + def submit(self, fn, /, *args, **kwargs): + return self.executor.submit(fn, *args, **kwargs) + + def shutdown(self, wait=True, *, cancel_futures=False): + self.executor.shutdown(wait=wait, cancel_futures=cancel_futures) + + class AsyncProvider: + async def shutdown(self): + def reenter_join(): + client.join() + executor_called.set() + + with DelegatingExecutor() as executor: + await asyncio.get_running_loop().run_in_executor( + executor, reenter_join + ) + + client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(executor_called.is_set()) + self.assertTrue(client._join_cleanup_complete) + + def test_async_cache_provider_read_only_loop_falls_back_without_deadlocking(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + executor_called = threading.Event() + + class ReadOnlyLoop(asyncio.SelectorEventLoop): + def __setattr__(self, name, value): + if name == "run_in_executor": + raise AttributeError("run_in_executor is read-only") + super().__setattr__(name, value) + + class AsyncProvider: + async def shutdown(self): + def reenter_join(): + executor_called.set() + client.join() + + await asyncio.get_running_loop().run_in_executor(None, reenter_join) + + loop = ReadOnlyLoop() + client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] + with mock.patch( + "posthog._async_utils.asyncio.new_event_loop", return_value=loop + ): + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(executor_called.is_set()) + self.assertTrue(loop.is_closed()) + self.assertTrue(client._join_cleanup_complete) + + def test_async_cache_provider_process_executor_remains_supported(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + class AsyncProvider: + result = None + + async def shutdown(self): + with ProcessPoolExecutor(max_workers=1) as executor: + self.result = await asyncio.get_running_loop().run_in_executor( + executor, abs, -6 + ) + + provider = AsyncProvider() + client._flag_definition_cache_provider = provider # type: ignore[assignment] + client.join() + + self.assertEqual(provider.result, 6) + self.assertTrue(client._join_cleanup_complete) + + def test_async_cache_provider_executor_can_reenter_join_during_runner_close(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + finalizer_called = threading.Event() + + class AsyncProvider: + async def shutdown(self): + async def pending_task(): + try: + await asyncio.Event().wait() + finally: + await asyncio.to_thread(client.join) + finalizer_called.set() + + asyncio.create_task(pending_task()) + await asyncio.sleep(0) + + client._flag_definition_cache_provider = AsyncProvider() # type: ignore[assignment] + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(finalizer_called.is_set()) + self.assertTrue(client._join_cleanup_complete) + + def test_cache_provider_shutdown_can_reenter_join(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=client.join, + ): + join_thread = threading.Thread(target=client.join) + join_thread.start() + join_thread.join(2) + + self.assertFalse(join_thread.is_alive()) + self.assertTrue(client._workers_joined) + + def test_cache_provider_shutdown_can_reenter_join_from_another_thread(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + class Runner: + _thread = None + + def owns_thread(self, thread): + return thread is self._thread + + runner = Runner() + client._flag_definition_cache_provider_async_runner = runner # type: ignore[assignment] + + def reenter_join(): + reentrant_thread = threading.Thread(target=client.join) + runner._thread = reentrant_thread + reentrant_thread.start() + reentrant_thread.join(1) + self.assertFalse(reentrant_thread.is_alive()) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=reenter_join, + ): + client.join() + + self.assertTrue(client._workers_joined) + self.assertTrue(client._join_cleanup_complete) + + def test_poller_shutdown_request_is_completed_by_join_owner(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + class ReentrantPoller(threading.Thread): + def __init__(self): + super().__init__(daemon=True) + self.stop_requested = threading.Event() + + def run(self): + self.stop_requested.wait(2) + client.shutdown() + + def stop(self): + self.stop_requested.set() + self.join(1) + self.assert_stopped() + + def assert_stopped(self): + if self.is_alive(): + raise AssertionError("poller did not stop") + + poller = ReentrantPoller() + client.poller = poller # type: ignore[assignment] + poller.start() + + client.join() + + self.assertTrue(client._shutdown_complete) + + def test_join_failure_does_not_retry_or_skip_later_cleanup(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + + with ( + mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=Exception("cleanup failed"), + ) as cleanup, + mock.patch.object(client, "_unregister_duplicate_client") as unregister, + ): + client.join() + client.join() + + self.assertTrue(client._workers_joined) + self.assertTrue(client._join_cleanup_complete) + cleanup.assert_called_once_with() + unregister.assert_called_once_with() + + def test_concurrent_debug_join_callers_observe_cleanup_failure(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01, debug=True) + cleanup_started = threading.Event() + release_cleanup = threading.Event() + owner_errors = [] + waiter_errors = [] + + def cleanup(): + cleanup_started.set() + self.assertTrue(release_cleanup.wait(2)) + raise ValueError("cleanup failed") + + def run_join(errors): + try: + client.join() + except Exception as error: + errors.append(error) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=cleanup, + ): + owner = threading.Thread(target=run_join, args=(owner_errors,)) + waiter = threading.Thread(target=run_join, args=(waiter_errors,)) + owner.start() + self.assertTrue(cleanup_started.wait(1)) + waiter.start() + time.sleep(0.05) + self.assertTrue(waiter.is_alive()) + + release_cleanup.set() + owner.join(2) + waiter.join(2) + + self.assertFalse(owner.is_alive()) + self.assertFalse(waiter.is_alive()) + self.assertEqual(len(owner_errors), 1) + self.assertRegex(str(owner_errors[0]), "cleanup failed") + self.assertEqual(len(waiter_errors), 1) + self.assertRegex(str(waiter_errors[0]), "client lifecycle cleanup failed") + self.assertTrue(client._join_cleanup_complete) + + def test_pending_shutdown_continues_when_join_cleanup_fails(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + cleanup_started = threading.Event() + release_cleanup = threading.Event() + cleanup_calls = 0 + join_errors = [] + + def cleanup(): + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_calls == 1: + cleanup_started.set() + self.assertTrue(release_cleanup.wait(2)) + raise Exception("cleanup failed") + + def run_join(): + try: + client.join() + except Exception as error: + join_errors.append(error) + + with mock.patch.object( + client, + "_shutdown_flag_definition_cache_provider", + side_effect=cleanup, + ): + join_thread = threading.Thread(target=run_join) + join_thread.start() + self.assertTrue(cleanup_started.wait(1)) + + shutdown_thread = threading.Thread(target=client.shutdown) + shutdown_thread.start() + time.sleep(0.05) + self.assertTrue(shutdown_thread.is_alive()) + + release_cleanup.set() + join_thread.join(2) + shutdown_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertFalse(shutdown_thread.is_alive()) + self.assertEqual(join_errors, []) + self.assertEqual(cleanup_calls, 1) + self.assertTrue(client._shutdown_complete) + + def test_join_interruption_does_not_publish_completion(self): + client = Client(FAKE_TEST_API_KEY, send=False) + + with ( + mock.patch.object( + client, "_flush_or_discard_queues", side_effect=KeyboardInterrupt + ), + self.assertRaises(KeyboardInterrupt), + ): + client.join() + + self.assertFalse(client._workers_joined) + self.assertFalse(client._join_cleanup_complete) + self.assertIsNone(client._lifecycle_owner) + + def test_shutdown_interruption_does_not_publish_completion(self): + client = Client(FAKE_TEST_API_KEY, send=False) + metrics = mock.Mock() + metrics.flush.side_effect = KeyboardInterrupt + client._metrics = metrics + + with self.assertRaises(KeyboardInterrupt): + client.shutdown() + + self.assertFalse(client._workers_joined) + self.assertFalse(client._join_cleanup_complete) + self.assertFalse(client._shutdown_complete) + self.assertIsNone(client._lifecycle_owner) + + def test_shutdown_failure_is_terminal_and_not_retried(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + exception_capture = mock.Mock() + exception_capture.close.side_effect = Exception("cleanup failed") + client.exception_capture = exception_capture + + client.shutdown() + client.shutdown() + + self.assertTrue(client._shutdown_complete) + exception_capture.close.assert_called_once_with() + + def test_shutdown_failure_is_raised_after_later_cleanup_in_debug_mode(self): + client = Client(FAKE_TEST_API_KEY, flush_interval=0.01, debug=True) + metrics = mock.Mock() + metrics.reset.side_effect = Exception("reset failed") + dedupe_cache = mock.Mock() + dedupe_cache.clear.side_effect = Exception("clear failed") + exception_capture = mock.Mock() + client._metrics = metrics + client.distinct_ids_feature_flags_reported = dedupe_cache + client.exception_capture = exception_capture + + with self.assertRaisesRegex(Exception, "reset failed"): + client.shutdown() + with self.assertRaisesRegex(RuntimeError, "client lifecycle cleanup failed"): + client.shutdown() + + metrics.flush.assert_called_once_with() + metrics.reset.assert_called_once_with() + dedupe_cache.clear.assert_called_once_with() + exception_capture.close.assert_called_once_with() + self.assertTrue(client._workers_joined) + self.assertTrue(client._join_cleanup_complete) + self.assertTrue(client._shutdown_complete) + + def test_lane_join_raises_first_failure_after_attempting_later_cleanup(self): + client = Client(FAKE_TEST_API_KEY, send=False) + lane = client._analytics_lane + consumer = mock.Mock() + first_error = ValueError("pause failed") + consumer._pause.side_effect = first_error + lane.consumers = [consumer] + + with ( + mock.patch.object(lane, "discard_undrainable_queued_work") as discard, + mock.patch.object( + lane._drain_signal, + "complete", + side_effect=RuntimeError("complete failed"), + ) as complete, + self.assertRaises(ValueError) as raised, + ): + lane.join() + + self.assertIs(raised.exception, first_error) + consumer.join.assert_called_once_with() + discard.assert_called_once_with() + complete.assert_called_once_with() + + def test_shutdown_prepares_each_lane_once(self): + client = Client(FAKE_TEST_API_KEY, send=False) + close_methods = [] + wait_methods = [] + + with contextlib.ExitStack() as stack: + for lane in client._lanes: + close_methods.append( + stack.enter_context( + mock.patch.object(lane, "close", wraps=lane.close) + ) + ) + wait_methods.append( + stack.enter_context( + mock.patch.object( + lane, + "wait_for_sync_sends", + wraps=lane.wait_for_sync_sends, + ) + ) + ) + client.shutdown() + + for close, wait in zip(close_methods, wait_methods): + close.assert_called_once_with() + wait.assert_called_once_with() + def test_shutdown_does_not_wait_for_idle_consumers_flush_interval(self): client = Client(FAKE_TEST_API_KEY, flush_interval=5) @@ -2358,12 +3231,8 @@ def test_synchronous(self): mock_post.assert_called_once() def test_overflow(self): - client = Client(FAKE_TEST_API_KEY, max_queue_size=1) - # Ensure consumer thread is no longer uploading - client.join() - - for i in range(10): - client.capture("test event", distinct_id="distinct_id") + client = Client(FAKE_TEST_API_KEY, max_queue_size=1, thread=0) + client.capture("test event", distinct_id="distinct_id") with self.assertLogs("posthog", level="WARNING") as logs: msg_uuid = client.capture("test event", distinct_id="distinct_id") @@ -2371,6 +3240,52 @@ def test_overflow(self): self.assertIsNone(msg_uuid) self.assertIn("dropping event", logs.output[0]) + def test_join_discards_queued_work_when_no_consumer_can_drain_it(self): + client = Client(FAKE_TEST_API_KEY, thread=0) + client.capture("test event", distinct_id="distinct_id") + + start = time.monotonic() + client.join() + + self.assertLess(time.monotonic() - start, 1) + self.assertEqual(client.queue.unfinished_tasks, 0) + + def test_join_discards_remaining_work_if_consumer_stops_during_drain(self): + send_started = threading.Event() + release_send = threading.Event() + + def request(batch): + send_started.set() + self.assertTrue(release_send.wait(2)) + + client = Client(FAKE_TEST_API_KEY, flush_at=1) + consumer = client.consumers[0] + with mock.patch.object(consumer, "request", side_effect=request): + client.capture("first", distinct_id="distinct_id") + self.assertTrue(send_started.wait(1)) + client.capture("second", distinct_id="distinct_id") + + join_thread = threading.Thread(target=client.join) + join_thread.start() + time.sleep(0.05) + consumer.pause() + release_send.set() + join_thread.join(3) + + self.assertFalse(join_thread.is_alive()) + self.assertEqual(client.queue.unfinished_tasks, 0) + + def test_shutdown_discards_queued_work_when_no_consumer_can_drain_it(self): + client = Client(FAKE_TEST_API_KEY, thread=0) + client.capture("test event", distinct_id="distinct_id") + + start = time.monotonic() + client.shutdown() + + self.assertLess(time.monotonic() - start, 1) + self.assertEqual(client.queue.unfinished_tasks, 0) + self.assertTrue(client._shutdown_complete) + def test_unicode(self): Client("unicode_key") diff --git a/posthog/test/test_client_fork.py b/posthog/test/test_client_fork.py index ac8f7463..16a64f7b 100644 --- a/posthog/test/test_client_fork.py +++ b/posthog/test/test_client_fork.py @@ -164,20 +164,81 @@ 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_resets_sync_send_state_for_sync_mode(self): + def test_reinit_after_fork_replaces_sync_mode_queue_and_locks(self): client = Client(FAKE_TEST_API_KEY, sync_mode=True) lane = client._analytics_lane old_queue = lane.queue old_lock = lane._start_lock old_condition = lane._sync_sends_done + old_lifecycle_lock = client._lifecycle_lock lane._active_sync_sends = 1 client._reinit_after_fork() - self.assertIs(lane.queue, old_queue) + self.assertIsNot(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) + self.assertIsNot(client._lifecycle_lock, old_lifecycle_lock) + + def test_reinit_after_fork_replaces_locks_before_starting_poller(self): + client = Client(FAKE_TEST_API_KEY) + client.enable_local_evaluation = True + old_runner_lock = client._flag_definition_cache_provider_async_runner_lock + old_publication_lock = client._flag_definition_publication_lock + old_cache_write_lock = client._flag_definition_cache_write_lock + old_metrics_lock = client._metrics_lock + + def assert_locks_replaced(): + self.assertIsNot( + client._flag_definition_cache_provider_async_runner_lock, + old_runner_lock, + ) + self.assertIsNot( + client._flag_definition_publication_lock, old_publication_lock + ) + self.assertIsNot( + client._flag_definition_cache_write_lock, old_cache_write_lock + ) + self.assertIsNot(client._metrics_lock, old_metrics_lock) + + with mock.patch("posthog.client.Poller") as mock_poller: + mock_poller.return_value.start.side_effect = assert_locks_replaced + client._reinit_after_fork() + + mock_poller.return_value.start.assert_called_once() + + def test_reinit_after_fork_preserves_terminal_client_state(self): + client = Client(FAKE_TEST_API_KEY) + client.join() + + with mock.patch("posthog.client.Poller") as mock_poller: + client._reinit_after_fork() + + self.assertTrue(client._workers_joined) + self.assertTrue(client._analytics_lane._closed) + self.assertEqual(client.consumers, []) + self.assertIsNone(client.poller) + mock_poller.assert_not_called() + self.assertIsNone(client.capture("after join", distinct_id="distinct_id")) + + def test_reinit_after_fork_normalizes_partially_closed_join_state(self): + client = Client(FAKE_TEST_API_KEY, send=False) + client._join_requested = True + client._analytics_lane._closed = True + client._ai_lane._closed = False + + with mock.patch("posthog.client.Poller") as mock_poller: + client._reinit_after_fork() + + self.assertFalse(client._workers_joined) + self.assertTrue(client._analytics_lane._closed) + self.assertTrue(client._ai_lane._closed) + self.assertEqual(client.consumers, []) + self.assertIsNone(client.poller) + mock_poller.assert_not_called() + self.assertIsNone(client.capture("analytics", distinct_id="distinct_id")) + self.assertIsNone(client._capture_ai("ai", distinct_id="distinct_id")) @unittest.skipUnless( diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index 296587bf..3782ad37 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -34,6 +34,141 @@ def test_next(self) -> None: next = consumer.next() self.assertEqual(next, [1]) + def test_next_does_not_take_queued_items_after_non_draining_pause(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100) + drain_signal = _DrainSignal(q) + consumer._set_drain_signal(drain_signal) + for item in range(10): + q.put(item) + + consumer.pause() + + self.assertEqual(consumer.next(), []) + self.assertEqual(q.qsize(), 10) + self.assertEqual(q.unfinished_tasks, 10) + + def test_non_draining_pause_overrides_active_flush_signal(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100) + drain_signal = _DrainSignal(q) + consumer._set_drain_signal(drain_signal) + q.put(_track_event()) + + drain_signal.request() + consumer.pause() + try: + self.assertEqual(consumer.next(), []) + self.assertEqual(q.qsize(), 1) + self.assertEqual(q.unfinished_tasks, 1) + finally: + drain_signal.complete() + + def test_non_draining_pause_between_drain_snapshot_and_dequeue(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100) + drain_signal = _DrainSignal(q) + consumer._set_drain_signal(drain_signal) + q.put(_track_event()) + drain_signal.request() + original_get = drain_signal.get + + def pause_then_get(*args, **kwargs): + consumer.pause() + return original_get(*args, **kwargs) + + with mock.patch.object(drain_signal, "get", side_effect=pause_then_get): + self.assertEqual(consumer.next(), []) + + drain_signal.complete() + self.assertEqual(q.qsize(), 1) + self.assertEqual(q.unfinished_tasks, 1) + + def test_pause_publishes_stop_under_queue_dequeue_lock(self) -> None: + q = Queue() + consumer = Consumer(q, "") + drain_signal = _DrainSignal(q) + stop_started = threading.Event() + original_stop = drain_signal.stop + + def observed_stop(target, drain): + stop_started.set() + original_stop(target, drain) + + drain_signal.stop = observed_stop # type: ignore[method-assign] + consumer._set_drain_signal(drain_signal) + + with q.mutex: + pause_thread = threading.Thread(target=consumer.pause) + pause_thread.start() + self.assertTrue(stop_started.wait(1)) + self.assertTrue(consumer.running) + + pause_thread.join(1) + self.assertFalse(pause_thread.is_alive()) + self.assertFalse(consumer.running) + + def test_non_draining_pause_discards_buffered_partial_batch(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100, flush_interval=60) + consumer._set_drain_signal(_DrainSignal(q)) + request_called = threading.Event() + consumer.request = lambda batch: request_called.set() # type: ignore[method-assign] + consumer.start() + q.put(_track_event()) + + deadline = time.monotonic() + 1 + while not q.empty(): + if time.monotonic() >= deadline: + self.fail("consumer did not buffer the queued event") + time.sleep(0.001) + + consumer.pause() + consumer.join(1) + + self.assertFalse(consumer.is_alive()) + self.assertFalse(request_called.is_set()) + self.assertEqual(q.unfinished_tasks, 0) + + def test_pause_does_not_wait_for_active_request(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=1) + consumer._set_drain_signal(_DrainSignal(q)) + request_started = threading.Event() + release_request = threading.Event() + + def request(batch): + request_started.set() + self.assertTrue(release_request.wait(2)) + + consumer.request = request # type: ignore[method-assign] + consumer.start() + q.put(_track_event()) + self.assertTrue(request_started.wait(1)) + + consumer.pause() + + self.assertFalse(consumer.running) + self.assertTrue(consumer.is_alive()) + release_request.set() + consumer.join(1) + self.assertFalse(consumer.is_alive()) + + def test_next_still_takes_queued_items_when_paused_for_drain(self) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=100) + drain_signal = _DrainSignal(q) + consumer._set_drain_signal(drain_signal) + for item in range(10): + q.put(item) + + drain_signal.request() + consumer._pause(drain=True) + try: + self.assertEqual(consumer.next(), list(range(10))) + finally: + drain_signal.complete() + def test_next_limit(self) -> None: q = Queue() flush_at = 50 @@ -53,6 +188,18 @@ def test_dropping_oversize_msg(self) -> None: self.assertTrue(q.empty()) self.assertEqual(q.unfinished_tasks, 0) + def test_next_balances_dequeued_work_if_batching_is_interrupted(self) -> None: + q = Queue() + consumer = Consumer(q, "") + q.put(_track_event()) + + with mock.patch("posthog.consumer.json.dumps", side_effect=SystemExit): + with self.assertRaises(SystemExit): + consumer.next() + + self.assertTrue(q.empty()) + self.assertEqual(q.unfinished_tasks, 0) + def test_max_msg_size_param_raises_per_event_ceiling(self) -> None: q = Queue() consumer = Consumer(q, "", flush_at=1, max_msg_size=4 * MAX_MSG_SIZE) @@ -84,7 +231,10 @@ def test_message_only_error_logs_include_posthog_prefix(self) -> None: upload_logs = [ line for line in logs.getvalue().splitlines() if "error uploading" in line ] - self.assertEqual(upload_logs, ["[PostHog] error uploading: boom"]) + expected_log = "[PostHog] error uploading: boom" + self.assertEqual( + [line for line in upload_logs if line == expected_log], [expected_log] + ) def test_flush_interval(self) -> None: # Put _n_ items in the queue, pausing a little bit more than