From f776410384cf106fb9eccd4dd6cb20e2b230858c Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Sun, 2 Aug 2026 18:58:40 -0700 Subject: [PATCH 1/2] perf(tracking): make MixpanelProvider dispatch non-blocking (bounded queue + daemon worker + real flush) (BE-5868) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MixpanelProvider.track() posted inline on the calling thread, so every consented invocation paid a synchronous HTTP round-trip — worst case ~10s against a blackholed endpoint — before the wrapped command body ran (@track_command fires its event first; `run` fires execution_start before submitting the workflow). Dispatch is now queue-and-drain: a bounded (256) queue drained by a single daemon worker. track() is put_nowait + drop-on-overflow with a debug line; flush() is a real queue.join() bounded by _flush_all_providers' existing shared 5s daemon deadline, exactly as PostHog's client.flush() already is. No atexit hook of our own, and the lazy mixpanel import is unchanged. --- comfy_cli/tracking.py | 72 ++++++-- tests/comfy_cli/test_tracking_providers.py | 184 +++++++++++++++++++-- 2 files changed, 236 insertions(+), 20 deletions(-) diff --git a/comfy_cli/tracking.py b/comfy_cli/tracking.py index bddb5549..bd4aaf5a 100644 --- a/comfy_cli/tracking.py +++ b/comfy_cli/tracking.py @@ -5,6 +5,7 @@ import json import logging as logginglib import os +import queue import sys import threading import time @@ -204,6 +205,12 @@ def flush(self) -> None: ... class MixpanelProvider: + # A CLI invocation emits ~1-3 events, so this cap is effectively unreachable + # outside pathological cases. It exists so a wedged worker (blackholed + # endpoint) can't let an unbounded backlog accumulate — dropping is the + # correct trade for best-effort telemetry. + _QUEUE_MAX = 256 + def __init__(self, token: str): self.client = None if token: @@ -212,24 +219,69 @@ def __init__(self, token: str): from mixpanel import Mixpanel # mixpanel-python's default Consumer uses request_timeout=None → an - # unbounded, synchronous requests.post on the main thread, so a - # blackholed telemetry endpoint (accepts TCP, never responds) hangs the - # CLI indefinitely (BE-3354/BE-3403). track() sends inline on the calling - # thread and flush() is a no-op, so this bound is the ONLY thing guarding - # the hot event path — it isn't covered by the atexit daemon deadline. - # retry_limit=1 (default is 4 with backoff) keeps a blackholed send to a - # single ~10s attempt instead of ~40s+ across retries. + # unbounded, synchronous requests.post, so a blackholed telemetry + # endpoint (accepts TCP, never responds) hangs whichever thread sends + # (BE-3354/BE-3403). Sends now happen on the worker below rather than + # the caller's thread, so this bound caps how long ONE send can occupy + # the queue — and with it, how much of the atexit drain's shared 5s + # deadline a single in-flight event can consume. retry_limit=1 + # (default is 4 with backoff) keeps a blackholed send to a single ~10s + # attempt instead of ~40s+ across retries. self.client = Mixpanel(token, consumer=MixpanelConsumer(request_timeout=10, retry_limit=1)) + # Dispatch is queue-and-drain so track() never blocks the caller + # (BE-5868): @track_command fires its event *before* running the + # wrapped command body, and `run` fires execution_start before + # submitting the workflow, so an inline send put a synchronous HTTP + # round-trip on the hot path of every consented invocation. + self._queue: queue.Queue = queue.Queue(maxsize=self._QUEUE_MAX) + # daemon=True with NO atexit hook of our own, and no shutdown + # sentinel: `_flush_all_providers` is the single bounded shutdown + # drain path (BE-3403), and the worker just dies with the process. + # Constructed lazily on the first dispatched event (`_get_providers` + # is only reached from `_dispatch`), so a run that sends nothing — + # `comfy --help`, shell completion, no consent — never starts it. + self._worker = threading.Thread(target=self._run, daemon=True, name="mixpanel-telemetry") + self._worker.start() self.enabled = self.client is not None + def _run(self) -> None: + """Drain the queue forever, one send at a time. + + A single worker over a single FIFO preserves per-process event ordering. + Accepted semantic shift: mixpanel-python stamps an event's `time` when + `client.track()` runs, which moves from call time to dequeue time — + sub-second skew in practice, no dashboard impact. + """ + while True: + event_name, distinct_id, properties = self._queue.get() + try: + self.client.track(distinct_id=distinct_id, event_name=event_name, properties=properties) + except Exception as e: # noqa: BLE001 + # debug, not warning: a telemetry failure must never surface on + # the user's stderr (cf. the urllib3/posthog silencing above). + logging.debug(f"Failed to send mixpanel event {event_name}: {e}") + finally: + # In `finally` so a raising send can't leave flush()'s + # queue.join() waiting on a task that is never marked done. + self._queue.task_done() + def track(self, event_name: str, distinct_id: str | None, properties: dict[str, Any]) -> None: if self.client is None or distinct_id is None: return - self.client.track(distinct_id=distinct_id, event_name=event_name, properties=properties) + try: + self._queue.put_nowait((event_name, distinct_id, dict(properties))) + except queue.Full: + logging.debug(f"mixpanel queue full; dropping event {event_name}") def flush(self) -> None: - # mixpanel-python ships per-call over sync HTTP; nothing to drain. - return + if self.client is None: + return + # Waits for the queue to drain completely, including the send already in + # flight. Unbounded here by design, exactly like PostHog's client.flush(): + # boundedness comes from the caller — `_flush_all_providers` runs this in + # a daemon thread and joins it against the shared `_FLUSH_DEADLINE_SECONDS` + # deadline, logging and abandoning it on timeout. + self._queue.join() class PostHogProvider: diff --git a/tests/comfy_cli/test_tracking_providers.py b/tests/comfy_cli/test_tracking_providers.py index ae599d2d..57ac736d 100644 --- a/tests/comfy_cli/test_tracking_providers.py +++ b/tests/comfy_cli/test_tracking_providers.py @@ -9,6 +9,7 @@ import importlib import logging import threading +import time from unittest.mock import MagicMock, patch import pytest @@ -60,11 +61,25 @@ def _posthog_capture_kwargs(client_mock): return kwargs +def _mixpanel_track_kwargs(mp_provider): + """Drain the Mixpanel worker, then return the last ``track(...)`` kwargs. + + ``MixpanelProvider`` dispatches through a bounded queue drained by a daemon + worker (BE-5868), so the send has not necessarily happened by the time + ``track_event()`` returns. Every assertion on the mocked client goes through + a ``flush()`` first. + """ + mp_provider.flush() + _, kwargs = mp_provider.client.track.call_args + return kwargs + + class TestDualFanOut: def test_track_event_fans_out_to_both_providers(self, tracking_with_two_providers): tracking_mod, mp_provider, ph_provider = tracking_with_two_providers tracking_mod.track_event("some_event", {"k": "v"}) + mp_provider.flush() mp_provider.client.track.assert_called_once() ph_provider.client.capture.assert_called_once() @@ -73,6 +88,7 @@ def test_opt_out_short_circuits_both_providers(self, tracking_with_two_providers tracking_mod.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "False") tracking_mod.track_event("some_event") + mp_provider.flush() mp_provider.client.track.assert_not_called() ph_provider.client.capture.assert_not_called() @@ -82,6 +98,9 @@ def test_one_provider_raising_does_not_block_the_other(self, tracking_with_two_p tracking_mod.track_event("some_event") + # The Mixpanel send now raises on its worker thread, not the caller's; + # flush() still returns (the worker's finally: task_done()). + mp_provider.flush() # Mixpanel raised but PostHog still got the call. ph_provider.client.capture.assert_called_once() @@ -91,7 +110,8 @@ def test_provider_order_does_not_matter_for_failure_isolation(self, tracking_wit tracking_mod.track_event("some_event") - # PostHog raised but Mixpanel still got the call (it ran first). + # PostHog raised but Mixpanel still got the call (it was enqueued first). + mp_provider.flush() mp_provider.client.track.assert_called_once() @@ -130,8 +150,7 @@ def test_mixpanel_does_not_receive_posthog_standard_props(self, tracking_with_tw tracking_mod, mp_provider, _ = tracking_with_two_providers tracking_mod.track_event("any_event") - _, kwargs = mp_provider.client.track.call_args - props = kwargs["properties"] + props = _mixpanel_track_kwargs(mp_provider)["properties"] assert "environment" not in props assert "surface" not in props assert "source" not in props @@ -168,7 +187,7 @@ def test_mixpanel_name_kwarg_routes_to_mixpanel_only(self, tracking_with_two_pro tracking_mod, mp_provider, ph_provider = tracking_with_two_providers tracking_mod.track_event("execution_start", mixpanel_name="run") - _, mp_kwargs = mp_provider.client.track.call_args + mp_kwargs = _mixpanel_track_kwargs(mp_provider) assert mp_kwargs["event_name"] == "run" # PostHog receives the canonical name, prefixed; Mixpanel keeps "run". @@ -179,7 +198,7 @@ def test_posthog_prefixes_event_while_mixpanel_stays_bare(self, tracking_with_tw tracking_mod, mp_provider, ph_provider = tracking_with_two_providers tracking_mod.track_event("execution_success") - _, mp_kwargs = mp_provider.client.track.call_args + mp_kwargs = _mixpanel_track_kwargs(mp_provider) ph_kwargs = _posthog_capture_kwargs(ph_provider.client) # Mixpanel keeps the bare name for stream continuity; PostHog is namespaced. assert mp_kwargs["event_name"] == "execution_success" @@ -192,7 +211,7 @@ def test_top_level_event_is_prefixed(self, tracking_with_two_providers): tracking_mod.track_event("install") # Mixpanel bare, PostHog namespaced. - _, mp_kwargs = mp_provider.client.track.call_args + mp_kwargs = _mixpanel_track_kwargs(mp_provider) assert mp_kwargs["event_name"] == "install" assert _posthog_capture_kwargs(ph_provider.client)["event"] == "cli:install" @@ -266,6 +285,153 @@ def test_posthog_track_skips_when_distinct_id_is_none(self, tracking_with_two_pr ph_provider.client.capture.assert_not_called() +def _mixpanel_provider_with_mock_client(): + """A real ``MixpanelProvider`` — bounded queue and daemon worker included — + whose SDK client is a MagicMock, so the worker's send never hits the network.""" + provider = MixpanelProvider("token-mp") + provider.client = MagicMock() + return provider + + +def _sent_event_names(provider): + return [call.kwargs["event_name"] for call in provider.client.track.call_args_list] + + +class TestMixpanelNonBlockingDispatch: + """BE-5868: ``track()`` used to post inline on the calling thread, so every + consented invocation paid a synchronous HTTP round-trip (worst case ~10s + against a blackholed endpoint) *before* the wrapped command body ran. + Dispatch is now a bounded queue drained by a daemon worker.""" + + def test_track_returns_while_a_send_is_still_in_flight(self): + provider = _mixpanel_provider_with_mock_client() + started, release = threading.Event(), threading.Event() + + def _blocking_send(**_kwargs): + started.set() + release.wait(timeout=30) + + provider.client.track.side_effect = _blocking_send + try: + # Occupy the worker so the next track() provably overlaps a live send. + provider.track("blocker", "test-distinct-id", {}) + assert started.wait(timeout=5), "worker never picked up the queued event" + + start = time.monotonic() + provider.track("payload", "test-distinct-id", {"k": "v"}) + elapsed = time.monotonic() - start + assert elapsed < 1.0, f"track() blocked on the in-flight send ({elapsed:.1f}s)" + finally: + release.set() + + provider.flush() + assert _sent_event_names(provider) == ["blocker", "payload"] + payload_kwargs = provider.client.track.call_args_list[1].kwargs + assert payload_kwargs["distinct_id"] == "test-distinct-id" + assert payload_kwargs["event_name"] == "payload" + assert payload_kwargs["properties"] == {"k": "v"} + + def test_flush_drains_every_queued_event_in_submission_order(self): + # A single FIFO drained by a single worker preserves per-process ordering. + provider = _mixpanel_provider_with_mock_client() + for i in range(20): + provider.track(f"event-{i}", "test-distinct-id", {"i": i}) + + provider.flush() + assert _sent_event_names(provider) == [f"event-{i}" for i in range(20)] + + def test_overflow_drops_events_without_blocking_or_raising(self, monkeypatch, caplog): + """A wedged worker must never turn into back-pressure on the CLI. The + queue is bounded; overflow drops with a debug line (never a warning — + telemetry failures must not surface on the user's stderr).""" + monkeypatch.setattr(MixpanelProvider, "_QUEUE_MAX", 2) + provider = _mixpanel_provider_with_mock_client() + started, release = threading.Event(), threading.Event() + + def _blocking_send(**_kwargs): + started.set() + release.wait(timeout=30) + + provider.client.track.side_effect = _blocking_send + try: + provider.track("blocker", "test-distinct-id", {}) + assert started.wait(timeout=5), "worker never picked up the queued event" + + with caplog.at_level(logging.DEBUG): + start = time.monotonic() + for i in range(50): + provider.track(f"event-{i}", "test-distinct-id", {}) + elapsed = time.monotonic() - start + assert elapsed < 1.0, f"track() blocked once the queue filled ({elapsed:.1f}s)" + finally: + release.set() + + provider.flush() + # Two slots behind the blocked send; everything past that is dropped. + assert _sent_event_names(provider) == ["blocker", "event-0", "event-1"] + drops = [r for r in caplog.records if "queue full" in r.getMessage()] + assert drops, "an overflow drop must leave a debug breadcrumb" + assert all(r.levelno == logging.DEBUG for r in drops) + + def test_flush_returns_even_when_the_send_raises(self): + """The worker marks each item done in a ``finally``, so a raising send + can't leave ``flush()``' ``queue.join()`` waiting forever.""" + provider = _mixpanel_provider_with_mock_client() + provider.client.track.side_effect = RuntimeError("mixpanel down") + provider.track("boom", "test-distinct-id", {}) + + finished = threading.Event() + # Run flush() off-thread so a wedge fails the test instead of hanging it. + threading.Thread(target=lambda: (provider.flush(), finished.set()), daemon=True).start() + assert finished.wait(timeout=10), "flush() wedged after a raising send" + provider.client.track.assert_called_once() + + def test_exit_stays_bounded_when_a_send_hangs(self): + """End to end against a blackholed endpoint: ``flush()`` is deliberately + unbounded (it joins the queue), and ``_flush_all_providers`` is what + bounds it — a daemon thread joined against the shared 5s deadline, same + as PostHog's internally-unbounded ``client.flush()`` (BE-3403).""" + import comfy_cli.tracking as tracking_mod + + provider = _mixpanel_provider_with_mock_client() + release = threading.Event() + provider.client.track.side_effect = lambda **_kwargs: release.wait(timeout=60) + try: + provider.track("hangs", "test-distinct-id", {}) + + with patch.object(tracking_mod, "PROVIDERS", [provider]): + start = time.monotonic() + tracking_mod._flush_all_providers() + elapsed = time.monotonic() - start + + budget = tracking_mod._FLUSH_DEADLINE_SECONDS + 3.0 + assert elapsed < budget, f"exit was not bounded by the flush deadline (took {elapsed:.1f}s)" + finally: + release.set() + + def test_worker_is_a_daemon_and_the_provider_registers_no_atexit_hook(self): + """Design constraint (BE-3403): ``_flush_all_providers`` is the ONLY + shutdown drain path. The worker must die with the process rather than + join it, and the provider must not add a second, unbounded atexit hook — + the exact mistake ``PostHogProvider`` has to actively unregister.""" + import comfy_cli.tracking as tracking_mod + + with patch.object(tracking_mod, "atexit") as fake_atexit: + provider = MixpanelProvider("token-mp") + + assert provider._worker.daemon is True + assert provider._worker.is_alive() + fake_atexit.register.assert_not_called() + + def test_disabled_provider_track_and_flush_are_inert(self): + # No token → no client, no queue, no worker: neither call may raise. + provider = MixpanelProvider("") + assert provider.enabled is False + assert not hasattr(provider, "_queue") + provider.track("any_event", "test-distinct-id", {}) + provider.flush() + + class TestRedactionThroughFanOut: def test_api_key_redaction_reaches_both_providers(self, tracking_with_two_providers): tracking_mod, mp_provider, ph_provider = tracking_with_two_providers @@ -276,7 +442,7 @@ def fake_cmd(workflow, api_key=None): fake_cmd(workflow="wf.json", api_key="sk-supersecret") - _, mp_kwargs = mp_provider.client.track.call_args + mp_kwargs = _mixpanel_track_kwargs(mp_provider) ph_kwargs = _posthog_capture_kwargs(ph_provider.client) assert mp_kwargs["properties"]["api_key"] == "" assert ph_kwargs["properties"]["api_key"] == "" @@ -303,7 +469,7 @@ def download(_ctx=None, url=None, set_civitai_api_token=None, set_hf_api_token=N set_hf_api_token="hf-secret", ) - _, mp_kwargs = mp_provider.client.track.call_args + mp_kwargs = _mixpanel_track_kwargs(mp_provider) ph_kwargs = _posthog_capture_kwargs(ph_provider.client) for properties in (mp_kwargs["properties"], ph_kwargs["properties"]): assert "_ctx" not in properties @@ -345,8 +511,6 @@ def test_flush_returns_before_deadline_when_a_provider_hangs(self): runs each flush in a daemon thread and joins with a ~5s timeout, so a 60s-hanging provider is abandoned rather than allowed to hang the CLI (BE-3354/BE-3403). Bounds the total at well under the 60s hang.""" - import time - import comfy_cli.tracking as tracking_mod release = threading.Event() From 7b31a377e2b8cc1c214a5b739b7bb81e3fd697da Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Sun, 2 Aug 2026 21:22:36 -0700 Subject: [PATCH 2/2] fix(tracking): harden the mixpanel worker and drain telemetry before os._exit (BE-5868) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cursor review follow-ups on the non-blocking dispatch change. - launch's os._exit paths now drain. @track_command fires its event *before* the command body runs, so inline sends were already delivered by the time `comfy launch --background` called os._exit; with queue-and-drain dispatch the event was still queued and dropped on every backgrounded launch. Adds tracking.flush_for_hard_exit() and routes launch.py's exits through _hard_exit(), which drains (bounded, best-effort) then os._exit()s. __main__.py's broken-pipe exit keeps skipping the drain, as documented there. - The worker loop is guarded as a whole, on BaseException. Nothing detects or restarts this thread, so any escape — an unpack failure, a MemoryError out of the SDK, a raise from inside the except handler — wedged telemetry for the rest of the process. - flush() is bounded and gives up on a dead worker. queue.join() is unconditional and unbounded, so a worker that never came back blocked every later flush() forever; flush() is public, and direct callers have no deadline of their own. Replaces join()/task_done() with a counter + condition. - track() deepcopies the payload. Serialization moved to the worker, so the shallow dict() copy left nested values aliased to objects the command body could still mutate — shipping post-mutation contents, or racing mixpanel's json.dumps into "dictionary changed size during iteration". - A Thread.start() RuntimeError ("can't start new thread") degrades to an inert provider instead of escaping into _get_providers' logging.warning, which would put a telemetry-side resource problem on the user's stderr. - The atexit flush-timeout log drops to debug. It was unreachable for Mixpanel while flush() was a no-op; now that it drains a queue, a slow endpoint would print to stderr after the terminal envelope. Co-Authored-By: Claude Opus 5 --- comfy_cli/command/launch.py | 38 ++++- comfy_cli/tracking.py | 136 +++++++++++++++-- tests/comfy_cli/test_tracking_providers.py | 167 ++++++++++++++++++++- 3 files changed, 314 insertions(+), 27 deletions(-) diff --git a/comfy_cli/command/launch.py b/comfy_cli/command/launch.py index 48edcf07..6ffca2d6 100644 --- a/comfy_cli/command/launch.py +++ b/comfy_cli/command/launch.py @@ -31,6 +31,30 @@ console = Console() +def _hard_exit(code: int) -> None: + """`os._exit(code)`, but drain telemetry on the way out. + + Every exit path in this module uses `os._exit` (plain `sys.exit` doesn't + work once the redirector threads are running), which skips atexit handlers — + so `comfy_cli.tracking`'s shutdown drain never runs here. That cost nothing + while Mixpanel sent inline from `track()`: `@track_command` fires the + `launch` event *before* the command body runs, so it was already delivered. + Now that dispatch is queue-and-drain (BE-5868) the event is still queued at + this point, and without an explicit drain every `comfy launch` — background + success included — silently drops its own telemetry. + + The drain is bounded (~5s worst case, same budget as the atexit hook) and + best-effort; the process exits with `code` regardless. + """ + try: + from comfy_cli import tracking + + tracking.flush_for_hard_exit() + except BaseException: # noqa: BLE001 # pragma: no cover - defensive + pass + os._exit(code) + + def _get_manager_flags() -> list[str]: """Get manager flags based on config mode.""" mode = resolve_manager_gui_mode(not_installed_value=None) @@ -198,15 +222,15 @@ def redirector_stdout(): if reboot_path is None: print("[bold red]ComfyUI is not installed.[/bold red]\n") - os._exit(1) + _hard_exit(1) if not os.path.exists(reboot_path): - os._exit(process.returncode) + _hard_exit(process.returncode) os.remove(reboot_path) except KeyboardInterrupt: if process is not None: - os._exit(1) + _hard_exit(1) def launch( @@ -343,7 +367,7 @@ def background_launch(extra, frontend_pr=None): log = asyncio.run(launch_and_monitor(cmd, listen, port)) # Reaching here means the monitor returned without seeing the success line - # (the success path emits its envelope and os._exit(0)s inside the monitor). + # (the success path emits its envelope and _hard_exit(0)s inside the monitor). if log is not None: print( Panel( @@ -362,7 +386,7 @@ def background_launch(extra, frontend_pr=None): details={"log": _bounded_log(log)} if log else None, ) # NOTE: os.exit(0) doesn't work - os._exit(1) + _hard_exit(1) def background_log_path(port, workspace: str | None = None) -> str: @@ -439,7 +463,7 @@ async def launch_and_monitor(cmd, listen, port): command="launch", details={"log_path": log_path}, ) - os._exit(1) + _hard_exit(1) # Record the log path up front so `comfy logs` can surface a crash log even # when startup fails before the success marker below (where the running @@ -489,7 +513,7 @@ def _handle(line): _emit_launch_success(listen, port, process.pid) # NOTE: os.exit(0) doesn't work. - os._exit(0) + _hard_exit(0) if logging_flag: log.append(line) diff --git a/comfy_cli/tracking.py b/comfy_cli/tracking.py index bd4aaf5a..ab778df7 100644 --- a/comfy_cli/tracking.py +++ b/comfy_cli/tracking.py @@ -1,6 +1,7 @@ from __future__ import annotations import atexit +import copy import functools import json import logging as logginglib @@ -204,6 +205,21 @@ def track(self, event_name: str, distinct_id: str | None, properties: dict[str, def flush(self) -> None: ... +def _log_telemetry_debug(message: str) -> None: + """Best-effort debug log that can never propagate out of telemetry code. + + Used on the Mixpanel worker's failure paths. The worker is a daemon thread + that can outlive `_flush_all_providers`' deadline and so log *after* the + stdlib's own `logging.shutdown` atexit hook has closed the handlers; on top + of that, a raise from inside an `except` handler would kill the worker for + the rest of the process. Neither is worth a user-visible traceback. + """ + try: + logging.debug(message) + except BaseException: # noqa: BLE001 # pragma: no cover - defensive + pass + + class MixpanelProvider: # A CLI invocation emits ~1-3 events, so this cap is effectively unreachable # outside pathological cases. It exists so a wedged worker (blackholed @@ -234,6 +250,14 @@ def __init__(self, token: str): # submitting the workflow, so an inline send put a synchronous HTTP # round-trip on the hot path of every consented invocation. self._queue: queue.Queue = queue.Queue(maxsize=self._QUEUE_MAX) + # flush() waits on this counter rather than `queue.join()`: + # `join()` is unconditional and unbounded, so a worker that never + # comes back would block every later flush() forever — including + # direct callers of the public method, who have no deadline of their + # own. A counter plus a condition lets flush() wake the moment the + # queue drains, and give up on a deadline or a dead worker. + self._pending = 0 + self._drained = threading.Condition() # daemon=True with NO atexit hook of our own, and no shutdown # sentinel: `_flush_all_providers` is the single bounded shutdown # drain path (BE-3403), and the worker just dies with the process. @@ -241,9 +265,28 @@ def __init__(self, token: str): # is only reached from `_dispatch`), so a run that sends nothing — # `comfy --help`, shell completion, no consent — never starts it. self._worker = threading.Thread(target=self._run, daemon=True, name="mixpanel-telemetry") - self._worker.start() + try: + self._worker.start() + except RuntimeError as e: + # "can't start new thread" under thread/FD/memory pressure. A + # telemetry-side resource problem must not become a user-visible + # provider-construction failure (`_get_providers` reports those + # with logging.warning, i.e. on the user's stderr), so degrade to + # an inert provider instead of letting it escape. + _log_telemetry_debug(f"could not start the mixpanel worker; disabling mixpanel telemetry: {e}") + self.client = None self.enabled = self.client is not None + def _mark_done(self) -> None: + """Retire one dequeued event and wake `flush()` once nothing is left.""" + try: + with self._drained: + self._pending -= 1 + if self._pending <= 0: + self._drained.notify_all() + except BaseException: # noqa: BLE001 # pragma: no cover - defensive + pass + def _run(self) -> None: """Drain the queue forever, one send at a time. @@ -253,35 +296,71 @@ def _run(self) -> None: sub-second skew in practice, no dashboard impact. """ while True: - event_name, distinct_id, properties = self._queue.get() + # The WHOLE body is guarded, not just the send: nothing detects or + # restarts this thread, so anything that escapes wedges telemetry for + # the rest of the process — every later event silently dropped and + # every flush() burning its full deadline. BaseException rather than + # Exception because a MemoryError or SystemExit out of the SDK would + # do exactly that. + item = None + event_name = "" try: + item = self._queue.get() + event_name, distinct_id, properties = item self.client.track(distinct_id=distinct_id, event_name=event_name, properties=properties) - except Exception as e: # noqa: BLE001 + except BaseException as e: # noqa: BLE001 # debug, not warning: a telemetry failure must never surface on # the user's stderr (cf. the urllib3/posthog silencing above). - logging.debug(f"Failed to send mixpanel event {event_name}: {e}") + _log_telemetry_debug(f"Failed to send mixpanel event {event_name}: {e}") finally: - # In `finally` so a raising send can't leave flush()'s - # queue.join() waiting on a task that is never marked done. - self._queue.task_done() + # In `finally`, and keyed on having actually dequeued something, + # so a raising send can't leave flush() waiting on an event that + # is never retired — nor retire one that was never taken. + if item is not None: + self._mark_done() def track(self, event_name: str, distinct_id: str | None, properties: dict[str, Any]) -> None: if self.client is None or distinct_id is None: return + # deepcopy, not dict(): `@track_command` fires its event *before* running + # the wrapped body, and serialization now happens on the worker instead of + # inline, so a shallow copy leaves nested values (Typer multi-value + # options, feedback score dicts) aliased to objects the body can still + # mutate — shipping post-mutation contents, or racing mixpanel's + # json.dumps into "dictionary changed size during iteration" and dropping + # the event. Copying here restores the old snapshot-at-call-time + # semantics; a value deepcopy can't handle falls back to the shallow copy. try: - self._queue.put_nowait((event_name, distinct_id, dict(properties))) - except queue.Full: - logging.debug(f"mixpanel queue full; dropping event {event_name}") + payload = copy.deepcopy(properties) + except Exception: # noqa: BLE001 + payload = dict(properties) + with self._drained: + try: + self._queue.put_nowait((event_name, distinct_id, payload)) + except queue.Full: + _log_telemetry_debug(f"mixpanel queue full; dropping event {event_name}") + return + # Counted under the same lock the worker takes to retire an event, so + # a send that completes before we return here can't decrement first. + self._pending += 1 def flush(self) -> None: if self.client is None: return # Waits for the queue to drain completely, including the send already in - # flight. Unbounded here by design, exactly like PostHog's client.flush(): - # boundedness comes from the caller — `_flush_all_providers` runs this in - # a daemon thread and joins it against the shared `_FLUSH_DEADLINE_SECONDS` - # deadline, logging and abandoning it on timeout. - self._queue.join() + # flight — but bounded, and abandoned outright if the worker is gone. + # `_flush_all_providers` supplies its own deadline at exit; this one is + # for every other caller of what is, after all, a public method. + deadline = time.monotonic() + _FLUSH_DEADLINE_SECONDS + with self._drained: + while self._pending > 0: + remaining = deadline - time.monotonic() + if remaining <= 0 or not self._worker.is_alive(): + _log_telemetry_debug(f"mixpanel flush gave up with {self._pending} event(s) still queued") + return + # Short waits so a worker that dies mid-drain is noticed promptly + # rather than at the deadline. + self._drained.wait(timeout=min(remaining, 0.1)) class PostHogProvider: @@ -648,7 +727,32 @@ def _flush_all_providers() -> None: logging.warning(f"telemetry flush join failed for {type(provider).__name__}: {e}") continue if t.is_alive(): - logging.warning(f"telemetry flush timed out for {type(provider).__name__}; dropping in-flight events") + # debug, not warning: this fires purely because a telemetry endpoint + # is slow or blackholed, and it fires at exit — i.e. it would print + # to the user's stderr *after* the terminal envelope. That is the one + # thing this module refuses to do for a telemetry failure. It was + # unreachable for Mixpanel while flush() was a no-op; it isn't now + # that flush() actually drains a queue (BE-5868). + logging.debug(f"telemetry flush timed out for {type(provider).__name__}; dropping in-flight events") + + +def flush_for_hard_exit() -> None: + """Drain telemetry before an `os._exit`, which skips atexit handlers. + + `comfy launch` terminates through `os._exit` on both its background-success + and failure paths, so `_flush_all_providers` never runs there. That was + harmless while MixpanelProvider sent inline from `track()` — the `launch` + event was already delivered before the command body ran. Now that dispatch is + queue-and-drain (BE-5868) the event is still sitting in the queue at that + point, so those paths have to drain explicitly or drop it every time. + + Bounded by the same `_FLUSH_DEADLINE_SECONDS` budget as the atexit hook, and + best-effort: nothing it does may keep the caller from exiting. + """ + try: + _flush_all_providers() + except BaseException: # noqa: BLE001 # pragma: no cover - defensive + pass atexit.register(_flush_all_providers) diff --git a/tests/comfy_cli/test_tracking_providers.py b/tests/comfy_cli/test_tracking_providers.py index 57ac736d..98765d7b 100644 --- a/tests/comfy_cli/test_tracking_providers.py +++ b/tests/comfy_cli/test_tracking_providers.py @@ -387,10 +387,12 @@ def test_flush_returns_even_when_the_send_raises(self): provider.client.track.assert_called_once() def test_exit_stays_bounded_when_a_send_hangs(self): - """End to end against a blackholed endpoint: ``flush()`` is deliberately - unbounded (it joins the queue), and ``_flush_all_providers`` is what - bounds it — a daemon thread joined against the shared 5s deadline, same - as PostHog's internally-unbounded ``client.flush()`` (BE-3403).""" + """End to end against a blackholed endpoint: ``_flush_all_providers`` runs + each provider's ``flush()`` in a daemon thread joined against the shared + 5s deadline, so a hung send costs the exit path that much and no more — + same treatment as PostHog's internally-unbounded ``client.flush()`` + (BE-3403). (Mixpanel's ``flush()`` also self-bounds; this covers the + caller-side guarantee, which is what holds for every provider.)""" import comfy_cli.tracking as tracking_mod provider = _mixpanel_provider_with_mock_client() @@ -432,6 +434,163 @@ def test_disabled_provider_track_and_flush_are_inert(self): provider.flush() +class TestMixpanelWorkerSurvivability: + """A single worker drains the whole queue and nothing restarts it, so any + escape from the loop wedges telemetry for the rest of the process: every + later event is dropped and every ``flush()`` burns its full deadline.""" + + def test_a_baseexception_from_the_sdk_does_not_kill_the_worker(self): + class _Boom(BaseException): + """Not an ``Exception`` — e.g. a ``MemoryError``/``SystemExit`` escape.""" + + provider = _mixpanel_provider_with_mock_client() + provider.client.track.side_effect = [_Boom("kaboom"), None] + + provider.track("first", "test-distinct-id", {}) + provider.track("second", "test-distinct-id", {}) + provider.flush() + + assert provider._worker.is_alive(), "worker died on a non-Exception failure" + assert _sent_event_names(provider) == ["first", "second"] + + def test_flush_gives_up_promptly_when_the_worker_is_gone(self): + """``queue.join()`` would block forever here, and ``flush()`` is public — + a direct caller has no deadline of its own to fall back on.""" + provider = _mixpanel_provider_with_mock_client() + provider.track("stranded", "test-distinct-id", {}) + # Impersonate a dead worker with the event still outstanding. + dead = threading.Thread(target=lambda: None) + dead.start() + dead.join() + provider._worker = dead + provider._pending = 1 + + start = time.monotonic() + provider.flush() + elapsed = time.monotonic() - start + assert elapsed < 2.0, f"flush() waited on a dead worker ({elapsed:.1f}s)" + + def test_flush_is_bounded_when_the_worker_never_finishes(self, monkeypatch): + import comfy_cli.tracking as tracking_mod + + monkeypatch.setattr(tracking_mod, "_FLUSH_DEADLINE_SECONDS", 0.5) + provider = _mixpanel_provider_with_mock_client() + release = threading.Event() + provider.client.track.side_effect = lambda **_kwargs: release.wait(timeout=60) + try: + provider.track("hangs", "test-distinct-id", {}) + + start = time.monotonic() + provider.flush() + elapsed = time.monotonic() - start + assert elapsed < 5.0, f"flush() ignored its deadline (took {elapsed:.1f}s)" + finally: + release.set() + + def test_worker_start_failure_degrades_to_an_inert_provider(self): + """``Thread.start`` raises ``RuntimeError: can't start new thread`` under + thread/FD/memory pressure. ``_get_providers`` reports a construction + failure with ``logging.warning`` — i.e. on the user's stderr — so a + telemetry-side resource problem must not escape as one.""" + with patch.object(threading.Thread, "start", side_effect=RuntimeError("can't start new thread")): + provider = MixpanelProvider("token-mp") + + assert provider.enabled is False + assert provider.client is None + # Still inert rather than raising, even though the queue exists. + provider.track("any_event", "test-distinct-id", {}) + provider.flush() + + +class TestMixpanelPropertySnapshot: + """``@track_command`` fires its event *before* running the wrapped body, and + serialization now happens on the worker instead of inline, so the payload has + to be snapshotted at enqueue time or the body can mutate it out from under + the send.""" + + def test_nested_values_are_snapshotted_at_track_time(self): + provider = _mixpanel_provider_with_mock_client() + started, release = threading.Event(), threading.Event() + + def _blocking_send(**kwargs): + if kwargs["event_name"] == "blocker": + started.set() + release.wait(timeout=30) + + provider.client.track.side_effect = _blocking_send + try: + # Hold the worker so the mutation below provably lands before the send. + provider.track("blocker", "test-distinct-id", {}) + assert started.wait(timeout=5), "worker never picked up the queued event" + + nested = {"flags": ["--fast-deps"]} + provider.track("payload", "test-distinct-id", {"nested": nested}) + nested["flags"].append("--no-deps") + finally: + release.set() + + provider.flush() + sent = provider.client.track.call_args_list[1].kwargs["properties"] + assert sent == {"nested": {"flags": ["--fast-deps"]}} + + def test_an_uncopyable_value_falls_back_to_a_shallow_copy(self): + provider = _mixpanel_provider_with_mock_client() + + class _NoDeepCopy: + def __deepcopy__(self, memo): + raise TypeError("cannot deepcopy this") + + sentinel = _NoDeepCopy() + provider.track("payload", "test-distinct-id", {"obj": sentinel}) + provider.flush() + + assert provider.client.track.call_args.kwargs["properties"]["obj"] is sentinel + + +class TestHardExitDrain: + """``comfy launch`` leaves through ``os._exit``, which skips atexit handlers. + That was free while Mixpanel sent inline from ``track()`` (the ``launch`` + event was delivered before the command body ran); with queue-and-drain + dispatch those paths have to drain explicitly or drop the event every time.""" + + def test_flush_for_hard_exit_drains_and_swallows(self): + import comfy_cli.tracking as tracking_mod + + provider = MagicMock() + with patch.object(tracking_mod, "PROVIDERS", [provider]): + tracking_mod.flush_for_hard_exit() + provider.flush.assert_called_once() + + with patch.object(tracking_mod, "_flush_all_providers", side_effect=RuntimeError("drain blew up")): + tracking_mod.flush_for_hard_exit() # must not propagate into the exit path + + def test_launch_hard_exit_drains_before_os_exit(self): + import comfy_cli.tracking as tracking_mod + from comfy_cli.command import launch as launch_mod + + calls = [] + with ( + patch.object(tracking_mod, "flush_for_hard_exit", side_effect=lambda: calls.append("drain")), + patch.object(launch_mod.os, "_exit", side_effect=lambda code: calls.append(("exit", code))), + ): + launch_mod._hard_exit(3) + + assert calls == ["drain", ("exit", 3)], "telemetry must be drained before the process is torn down" + + def test_launch_exits_even_if_the_drain_raises(self): + import comfy_cli.tracking as tracking_mod + from comfy_cli.command import launch as launch_mod + + exits = [] + with ( + patch.object(tracking_mod, "flush_for_hard_exit", side_effect=RuntimeError("drain blew up")), + patch.object(launch_mod.os, "_exit", side_effect=lambda code: exits.append(code)), + ): + launch_mod._hard_exit(1) + + assert exits == [1] + + class TestRedactionThroughFanOut: def test_api_key_redaction_reaches_both_providers(self, tracking_with_two_providers): tracking_mod, mp_provider, ph_provider = tracking_with_two_providers