diff --git a/comfy_cli/tracking.py b/comfy_cli/tracking.py index 580bceb1..6a4e3e37 100644 --- a/comfy_cli/tracking.py +++ b/comfy_cli/tracking.py @@ -6,10 +6,13 @@ import logging as logginglib import os import sys +import threading +import time import uuid from typing import Any, Protocol import typer +from mixpanel import Consumer as MixpanelConsumer from mixpanel import Mixpanel from posthog import Posthog @@ -139,7 +142,15 @@ def flush(self) -> None: ... class MixpanelProvider: def __init__(self, token: str): - self.client = Mixpanel(token) if token else None + # 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. + self.client = Mixpanel(token, consumer=MixpanelConsumer(request_timeout=10, retry_limit=1)) if token else None self.enabled = self.client is not None def track(self, event_name: str, distinct_id: str | None, properties: dict[str, Any]) -> None: @@ -166,7 +177,18 @@ def __init__(self, token: str, host: str): if not token: return # disable_geoip=False lets PostHog enrich events with IP-derived location. - self.client = Posthog(project_api_key=token, host=host, disable_geoip=False) + # max_retries/timeout tighten the consumer drain budget from the posthog + # 7.x defaults (3 × 15s ≈ 50s worst case) to ~21s, so the atexit flush + # can't linger on a blackholed endpoint after the terminal envelope is + # already on stdout (BE-3354/BE-3403). + self.client = Posthog(project_api_key=token, host=host, disable_geoip=False, max_retries=1, timeout=10) + # Posthog's constructor registers its own atexit.register(self.join), + # which runs self.join() synchronously on the main thread at shutdown — + # independently of _flush_all_providers and NOT bounded by its daemon + # deadline. Against a blackholed endpoint that join can still block ~21s + # after the terminal envelope is on stdout, defeating this change. Drop it + # so our bounded flush is the only shutdown drain path (BE-3403). + atexit.unregister(self.client.join) self.enabled = True def track(self, event_name: str, distinct_id: str | None, properties: dict[str, Any]) -> None: @@ -407,12 +429,44 @@ def init_tracking(enable_tracking: bool): track_event("install") +def _flush_one(provider: TelemetryProvider) -> None: + try: + provider.flush() + except Exception as e: # noqa: BLE001 + logging.warning(f"Failed to flush telemetry provider {type(provider).__name__}: {e}") + + +_FLUSH_DEADLINE_SECONDS = 5.0 + + def _flush_all_providers() -> None: + # Telemetry is best-effort by contract: a blackholed endpoint (accepts TCP, + # never responds) must never let this atexit hook wedge every consumer of the + # CLI's stdout after the terminal envelope is already emitted (BE-3329/BE-3403). + # Start every provider's flush in a daemon thread, then join them all against + # a SINGLE shared deadline so total exit delay stays ~5s regardless of how + # many providers there are (a per-provider join would make it 5s × N). + # Dropping in-flight events beats hanging the process. t.start()/t.join() are + # wrapped defensively so nothing this hook does can raise and print a + # traceback to stderr after the terminal envelope. + deadline = time.monotonic() + _FLUSH_DEADLINE_SECONDS + threads: list[tuple[threading.Thread, TelemetryProvider]] = [] for provider in PROVIDERS: + t = threading.Thread(target=_flush_one, args=(provider,), daemon=True) + try: + t.start() + except RuntimeError as e: # e.g. a thread-creation race during shutdown + logging.warning(f"could not start telemetry flush for {type(provider).__name__}: {e}") + continue + threads.append((t, provider)) + for t, provider in threads: try: - provider.flush() - except Exception as e: # noqa: BLE001 - logging.warning(f"Failed to flush telemetry provider {type(provider).__name__}: {e}") + t.join(timeout=max(0.0, deadline - time.monotonic())) + except RuntimeError as e: # pragma: no cover - defensive + 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") atexit.register(_flush_all_providers) diff --git a/tests/comfy_cli/test_tracking_providers.py b/tests/comfy_cli/test_tracking_providers.py index 591cc754..51b3b316 100644 --- a/tests/comfy_cli/test_tracking_providers.py +++ b/tests/comfy_cli/test_tracking_providers.py @@ -6,6 +6,7 @@ ``workflow_run_id`` on the canonical execution lifecycle events. """ +import threading from unittest.mock import MagicMock, patch import pytest @@ -216,6 +217,42 @@ def test_mixpanel_with_empty_token_is_disabled(self): provider = MixpanelProvider("") assert provider.enabled is False + def test_mixpanel_client_has_bounded_request_timeout(self): + """Regression guard for BE-3354/BE-3403: mixpanel-python's default + Consumer uses request_timeout=None → an unbounded requests.post that + hangs the CLI forever on a blackholed endpoint. The provider must build + its Mixpanel client with an explicit 10s consumer timeout so this can't + silently regress to None. retry_limit must also be 1: the Consumer + default is 4 with backoff, which would let a blackholed send run ~40s+ + across retries even with a 10s per-attempt timeout — and this send is + synchronous on the main thread, not covered by the atexit deadline.""" + provider = MixpanelProvider("token-mp") + assert provider.client is not None + consumer = provider.client._consumer + assert consumer._request_timeout == 10 + # retry_limit is threaded into the session's urllib3 Retry.total. + adapter = consumer._session.get_adapter("https://api.mixpanel.com") + assert adapter.max_retries.total == 1 + + def test_posthog_unregisters_its_own_atexit_join(self): + """Regression guard for BE-3403: Posthog's constructor registers its own + ``atexit.register(self.join)``, which flushes synchronously on the main + thread at shutdown, unbounded by ``_flush_all_providers``' 5s daemon + deadline. Against a blackholed endpoint that join can block ~21s after + the terminal envelope. The provider must unregister it so the bounded + flush is the only shutdown drain path.""" + import comfy_cli.tracking as tracking_mod + + fake_client = MagicMock() + with ( + patch.object(tracking_mod, "Posthog", return_value=fake_client), + patch.object(tracking_mod, "atexit") as fake_atexit, + ): + provider = PostHogProvider("phc_test", "https://t.comfy.org") + + assert provider.enabled is True + fake_atexit.unregister.assert_called_once_with(fake_client.join) + def test_posthog_track_skips_when_distinct_id_is_none(self, tracking_with_two_providers): tracking_mod, _, ph_provider = tracking_with_two_providers with patch.object(tracking_mod, "user_id", None): @@ -296,3 +333,35 @@ def test_flush_swallows_provider_errors(self): tracking_mod._flush_all_providers() p2.flush.assert_called_once() + + def test_flush_returns_before_deadline_when_a_provider_hangs(self): + """A provider whose flush() blocks (e.g. a blackholed telemetry endpoint) + must not wedge the atexit hook past its per-provider deadline. The hook + 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() + + class _HangingProvider: + def flush(self): + # Blocks until the test tears down; the deadline must fire first. + release.wait(timeout=60) + + fast = MagicMock() + try: + with patch.object(tracking_mod, "PROVIDERS", [_HangingProvider(), fast]): + start = time.monotonic() + tracking_mod._flush_all_providers() + elapsed = time.monotonic() - start + + # Deadline is 5s per provider; allow generous slack but stay far + # under the 60s the hanging provider would otherwise cost. + assert elapsed < 8.0, f"flush did not honor the deadline (took {elapsed:.1f}s)" + # The healthy provider after the hanging one is still drained. + fast.flush.assert_called_once() + finally: + release.set()