From 4449c664098dc87d3eb21735e56eb622f3bbb2fc Mon Sep 17 00:00:00 2001 From: Anna Garcia Date: Thu, 10 Sep 2026 21:43:27 -0400 Subject: [PATCH 1/4] feat(traces): wire tracing into the client Makes tracing usable. Adds the `traces` client option (tracing stays off until it is set), Client.start_span / get_active_span and their posthog module-level counterparts, with the active span scoped per client so two clients never parent to each other's spans. flush() drains spans alongside events within the same budget; shutdown() gives queued spans a final flush of up to 30 s and warns about any it discards; an exit flush bounded by the existing exit deadline covers scripts that never call shutdown(), and a forked child drops the parent's spans. Export failures, limits and the hook are documented on the option. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01TkZAsCciW4PV8ZdcCHmAbA --- .sampo/changesets/traces-spans.md | 5 + posthog/__init__.py | 96 +++ posthog/client.py | 240 ++++++- posthog/test/tracing/test_client_traces.py | 716 +++++++++++++++++++++ posthog/tracing/__init__.py | 2 - references/public_api_snapshot.txt | 9 +- 6 files changed, 1057 insertions(+), 11 deletions(-) create mode 100644 .sampo/changesets/traces-spans.md create mode 100644 posthog/test/tracing/test_client_traces.py diff --git a/.sampo/changesets/traces-spans.md b/.sampo/changesets/traces-spans.md new file mode 100644 index 000000000..c23cb9163 --- /dev/null +++ b/.sampo/changesets/traces-spans.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +Add distributed tracing (alpha): `start_span()` and `get_active_span()` record spans and export them to PostHog as OTLP, with no OpenTelemetry dependency, when the new `traces` client option is set. diff --git a/posthog/__init__.py b/posthog/__init__.py index 0431bc883..8e86afe1c 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -12,6 +12,7 @@ from posthog.capture_compression import CaptureCompression as CaptureCompression from posthog.capture_mode import CaptureMode as CaptureMode from posthog.client import Client +from posthog.tracing.span import Span from posthog.async_client import AsyncClient as AsyncClient from posthog.async_client import AsyncPosthog as AsyncPosthog from posthog.exception_capture import ExceptionCapture @@ -340,6 +341,20 @@ def get_tags() -> Dict[str, Any]: ``service_version``, ``environment``, ``flush_interval``, ...). Applied when ``setup()`` builds the global client, or on a later ``setup()`` call if the metrics API hasn't been used yet. + traces: Config dict for distributed tracing: ``service_name``, + ``service_version``, ``environment``, ``resource_attributes``, + ``flush_interval`` (5 s), ``max_queue_size`` (2048), + ``max_export_batch_size`` (512), ``max_live_spans`` (10000), + ``max_span_age`` (3600 s), ``max_attributes_per_span`` (128), + ``max_events_per_span`` (128), ``max_attribute_value_length`` (8192). + ``before_span_send`` is a callable, or a list run in order, + that receives each finished span as a dict (``trace_id``, ``span_id`` + and ``parent_span_id`` are read-only) and returns it, edited, or + ``None`` to drop it; a hook that raises drops the span. Tracing is off + until set. Spans export on a background timer even with ``sync_mode``; + serverless handlers should call ``flush()`` before returning. Applied + when ``setup()`` builds the global client, or on a later ``setup()`` + call if no span has been started yet. enable_exception_autocapture: Automatically capture uncaught exceptions. log_captured_exceptions: Also log exceptions captured by error tracking. project_root: Root path used to determine in-app exception stack frames. @@ -396,6 +411,7 @@ def get_tags() -> Dict[str, Any]: feature_flags_request_max_retries = 1 # type: int super_properties = None # type: Optional[Dict] metrics = None # type: Optional[Dict] +traces = None # type: Optional[Dict] enable_exception_autocapture = False # type: bool log_captured_exceptions = False # type: bool # Used to determine in app paths for exception autocapture. Defaults to the current working directory @@ -1197,6 +1213,83 @@ def join() -> None: _proxy("join") +def start_span( + name: str, + *, + kind: Optional[str] = None, + attributes: Optional[Mapping[str, Any]] = None, + parent: Union[Span, str, None] = None, + tracestate: Optional[str] = None, + start_time: Union[datetime.datetime, float, None] = None, +) -> Span: + """ + Start a span for distributed tracing. Alpha. + + Returns a span handle. Use it as a context manager to make it the active + span for the block and end it on exit (recording a raised exception on the + way out); or call ``end()`` yourself for a span that cannot wrap a block. + Spans started inside the block nest under it automatically. Always returns + a usable handle, even when tracing is off, so calling code never branches. + + Args: + name: A low-cardinality operation name, e.g. ``GET /users/:id``. + Variable values belong in attributes, not the name. + kind: ``internal`` (default), ``server``, ``client``, ``producer`` or + ``consumer``. + attributes: Initial attributes. + parent: A span handle, or an inbound W3C ``traceparent`` header value + to continue a remote trace. Defaults to the active span. + tracestate: The inbound ``tracestate`` header accompanying a + ``traceparent`` string ``parent``; preserved and propagated. + start_time: A ``datetime`` or epoch seconds, to backdate the span. + + Examples: + ```python + import posthog + posthog.traces = {"service_name": "checkout-api"} + + with posthog.start_span("POST /checkout", parent=request.headers.get("traceparent")) as span: + span.set_attribute("plan", user.plan) + with posthog.start_span("db.query", kind="client"): + ... + outgoing_headers = {"traceparent": span.traceparent()} + ``` + + Category: + Tracing + """ + return _proxy( + "start_span", + name, + kind=kind, + attributes=attributes, + parent=parent, + tracestate=tracestate, + start_time=start_time, + ) + + +def get_active_span() -> Optional[Span]: + """ + The span that is active in the current context, or ``None``. Alpha. + + Only entering a span (``with posthog.start_span(...) as span:``) makes it + active; a span started manually is not. Use it to propagate the trace to + the next service: ``span.traceparent()`` is the header value. + + Examples: + ```python + span = posthog.get_active_span() + if span is not None: + headers["traceparent"] = span.traceparent() + ``` + + Category: + Tracing + """ + return _proxy("get_active_span") + + def shutdown() -> None: """ Flush all messages and cleanly shutdown the client. @@ -1259,6 +1352,7 @@ def setup() -> Client: feature_flags_request_max_retries=feature_flags_request_max_retries, super_properties=super_properties, metrics=metrics, + traces=traces, # TODO: Currently this monitoring begins only when the Client is initialised (which happens when you do something with the SDK) # This kind of initialisation is very annoying for exception capture. We need to figure out a way around this, # or deprecate this proxy option fully (it's already in the process of deprecation, no new clients should be using this method since like 5-6 months) @@ -1299,6 +1393,8 @@ def setup() -> Client: # already forced setup()) still applies until the metrics API is first used. if default_client._metrics is None: default_client._metrics_config = metrics + if default_client._traces is None: + default_client._traces_config = traces return default_client diff --git a/posthog/client.py b/posthog/client.py index 59ccb9dca..06236f659 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -19,6 +19,13 @@ from posthog._disabled_lane_queue import _DisabledLaneQueue from posthog.args import ID_TYPES, ExceptionArg, OptionalCaptureArgs, OptionalSetArgs from posthog.metrics_capture import PostHogMetrics +from posthog.tracing._config import resolve_traces_config +from posthog.tracing._drops import DropLog +from posthog.tracing._export import SpanExporter +from posthog.tracing._otlp import host_resource_attributes +from posthog.tracing._pipeline import PostHogTraces +from posthog.tracing.span import Span +from posthog.tracing._span import inert_span as _inert_span from posthog.capture_compression import ( CaptureCompression, _resolve_capture_compression, @@ -126,6 +133,9 @@ MAX_DICT_SIZE = 50_000 _ATEXIT_FLUSH_TIMEOUT_SECONDS = 1.0 +# The final span flush shutdown() allows; a request already in flight may +# overrun it by up to the client's request timeout. +_TRACES_SHUTDOWN_FLUSH_SECONDS = 30.0 _atexit_deadline: Optional[float] = None _atexit_deadline_lock = threading.Lock() @@ -714,6 +724,7 @@ def __init__( capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, + traces: Optional[dict] = None, ): """ Initialize a new PostHog client instance. @@ -796,6 +807,20 @@ def __init__( ``$trace_id``/``$span_id`` values passed in ``properties`` win. Exception events (``capture_exception``) always attach these IDs regardless of this setting. Defaults to False. + traces: Config dict for distributed tracing: ``service_name``, + ``service_version``, ``environment``, ``resource_attributes``, + ``flush_interval`` (5 s), ``max_queue_size`` (2048), + ``max_export_batch_size`` (512), ``max_live_spans`` (10000), + ``max_span_age`` (3600 s), ``max_attributes_per_span`` (128), + ``max_events_per_span`` (128), ``max_attribute_value_length`` + (8192). ``before_span_send`` is a callable, or a + list run in order, that receives each finished span as a dict + (``trace_id``, ``span_id`` and ``parent_span_id`` are read-only) + and returns it, edited, or ``None`` to drop it; a hook that + raises drops the span. Tracing is off until this is provided. + Spans export on a background timer even with ``sync_mode``; + serverless handlers should call ``flush()`` before returning. + Defaults to None. code_variables_mask_patterns: Variable-name patterns to mask when capturing code variables. code_variables_ignore_patterns: Variable-name patterns to omit when @@ -907,6 +932,14 @@ def __init__( self._metrics_config = metrics self._metrics: Optional[PostHogMetrics] = None self._metrics_lock = threading.Lock() + self._traces_config: Any = traces + self._traces: Optional[PostHogTraces] = None + self._traces_lock = threading.Lock() + # The active span, scoped to this client so two instances in one + # process never parent to each other's spans. + self._active_span_var: ContextVar[Optional[Span]] = ContextVar( + "posthog_active_span", default=None + ) # `_use_ai_lane` / `_enable_multimodal_capture` are deprecated aliases. self.enable_full_ai_capture = ( enable_full_ai_capture is True @@ -2242,6 +2275,10 @@ def _reinit_after_fork(self): self._metrics_lock = threading.Lock() if self._metrics is not None: self._metrics._reinit_after_fork() + self._traces_lock = threading.Lock() + if self._traces is not None: + self._traces.reinit_after_fork() + self._active_span_var.set(None) # If using Redis cache, we must reinitialize to get a fresh connection (fork-safe). # If using Memory cache, we keep it as-is to benefit from the inherited warm cache. @@ -2466,6 +2503,127 @@ def metrics(self) -> PostHogMetrics: self._metrics = PostHogMetrics(self, None) return self._metrics + @property + def _traces_pipeline(self) -> Optional[PostHogTraces]: + if ( + self._traces_config is None + or self._traces_config is False + or self._shutdown_requested + ): + return self._traces + if self._traces is None: + with self._traces_lock: + if self._traces is None: + try: + config = resolve_traces_config( + self._traces_config, host_resource_attributes() + ) + drops = DropLog(config.flush_interval) + self._traces = PostHogTraces( + self, + config, + self._tracing_context, + self._active_span_var, + SpanExporter(self, config, drops), + drops, + ) + # Sync mode has no exit hook for events; spans still + # queue for the timer, so they need their own. + if self.sync_mode and self.send: + atexit.register(self._atexit_spans) + except Exception: + # Off rather than defaults: defaults would drop a + # before_span_send hook and export unscrubbed spans. + self.log.exception("Error initializing traces; tracing is off") + self._traces_config = None + return self._traces + + def _tracing_context(self) -> Dict[str, Optional[str]]: + # The request context's identity, which the Django middleware fills + # from the X-POSTHOG-* headers: the span's person/session join keys. + return { + "distinct_id": get_context_distinct_id(), + "session_id": get_context_session_id(), + } + + def start_span( + self, + name: str, + *, + kind: Optional[str] = None, + attributes: Optional[Mapping[str, Any]] = None, + parent: Union[Span, str, None] = None, + tracestate: Optional[str] = None, + start_time: Union[datetime, float, None] = None, + ) -> Span: + """ + Start a span for distributed tracing. Alpha. + + Returns a span handle. Use it as a context manager to make it the active + span for the block and end it on exit (recording a raised exception on + the way out); or call ``end()`` yourself for a span that cannot wrap a + block. Spans started inside the block nest under it automatically. + Always returns a usable handle, even when tracing is off, so calling + code never branches. + + Args: + name: A low-cardinality operation name, e.g. ``GET /users/:id``. + Variable values belong in attributes, not the name. + kind: ``internal`` (default), ``server``, ``client``, ``producer`` + or ``consumer``. + attributes: Initial attributes. + parent: A span handle, or an inbound W3C ``traceparent`` header + value to continue a remote trace. Defaults to the active span. + tracestate: The inbound ``tracestate`` header accompanying a + ``traceparent`` string ``parent``; preserved and propagated. + start_time: A ``datetime`` or epoch seconds, to backdate the span. + + Examples: + ```python + posthog = Posthog("", traces={"service_name": "checkout-api"}) + + with posthog.start_span("POST /checkout", parent=request.headers.get("traceparent")) as span: + span.set_attribute("plan", user.plan) + with posthog.start_span("db.query", kind="client"): + ... + outgoing_headers = {"traceparent": span.traceparent()} + ``` + + Category: + Tracing + """ + pipeline = self._traces_pipeline + if pipeline is None: + return _inert_span(parent, tracestate, self._active_span_var) + return pipeline.start_span( + name, + kind=kind, + attributes=attributes, + parent=parent, + tracestate=tracestate, + start_time=start_time, + ) + + def get_active_span(self) -> Optional[Span]: + """ + The span that is active in the current context, or ``None``. Alpha. + + Only entering a span (``with posthog.start_span(...) as span:``) makes + it active; a span started manually is not. Use it to propagate the + trace to the next service: ``span.traceparent()`` is the header value. + + Examples: + ```python + span = posthog.get_active_span() + if span is not None: + headers["traceparent"] = span.traceparent() + ``` + + Category: + Tracing + """ + return self._active_span_var.get() + def flush(self, timeout_seconds: Optional[float] = 10) -> None: """ Force a flush from the internal queue to the server. Do not use directly, call `shutdown()` instead. @@ -2473,6 +2631,10 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: Args: timeout_seconds: Maximum seconds to wait for the queue to flush. Defaults to 10 seconds. Pass ``None`` to wait indefinitely. + Queued spans are sent at the same time, within the same + budget: at least one span request is attempted even when the + budget is already spent, no further one starts once it is, and + each request is bounded by ``timeout``. Examples: ```python @@ -2483,20 +2645,49 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: if self._defer_flush_from_callback(timeout_seconds): return try: + # Spans drain with events: serverless handlers call flush(), not + # shutdown(), and leaving spans on their own timer would lose them. + span_flush = self._start_span_flush(timeout_seconds) if timeout_seconds is None: for lane in self._lanes: lane.flush(None) - return - - # The timeout is a total budget shared by the lanes, so flush() - # returns within roughly `timeout_seconds` overall. - deadline = time.monotonic() + timeout_seconds - for lane in self._lanes: - lane.flush(max(0.0, deadline - time.monotonic())) + else: + deadline = time.monotonic() + timeout_seconds + for lane in self._lanes: + lane.flush(max(0.0, deadline - time.monotonic())) + if span_flush is not None: + # The first span request is exempt from the budget and bounded + # only by the request timeout, so the join is not. + span_flush.join() except Exception as e: self.log.exception("error flushing queue: %s", e) return + def _start_span_flush( + self, timeout_seconds: Optional[float] + ) -> Optional[threading.Thread]: + """Flush spans alongside the events, so a handler waits one round trip, not two.""" + traces = self._traces + if traces is None: + return None + + def flush_spans() -> None: + try: + traces.flush(timeout_seconds) + except Exception as e: + self.log.exception("error flushing spans: %s", e) + + flusher = threading.Thread( + target=flush_spans, name="posthog-span-flush", daemon=True + ) + try: + flusher.start() + except RuntimeError: + # No new threads at interpreter shutdown; flush on this one. + flush_spans() + return None + return flusher + def _is_consumer_thread(self) -> bool: current = threading.current_thread() return any(current in lane.consumers for lane in self._lanes) @@ -2687,6 +2878,16 @@ def _shutdown_once(self, errors: list[Exception]) -> None: self._run_lifecycle_cleanup( "Failed to reset metrics on shutdown", self._metrics.reset, errors ) + if self._traces is not None: + traces = self._traces + self._run_lifecycle_cleanup( + "Failed to flush spans on shutdown", + lambda: traces.flush(_TRACES_SHUTDOWN_FLUSH_SECONDS), + errors, + ) + self._run_lifecycle_cleanup( + "Failed to close traces on shutdown", self._traces.close, errors + ) self._join_once(errors, flush_queues=False, lanes_prepared=True) self._run_lifecycle_cleanup( "Failed to clear feature flag deduplication state on shutdown", @@ -2772,8 +2973,12 @@ def _atexit(self) -> None: lane.close() deadline = _get_atexit_deadline() + span_flush = self._start_span_flush( + max(0.0, deadline - time.monotonic()) + ) for lane in self._lanes: lane.flush(max(0.0, deadline - time.monotonic())) + self._join_span_flush(span_flush, deadline) finally: # Consumers are daemon threads. Publish a non-draining stop to # every consumer, but do not join in-flight requests at exit. @@ -2785,6 +2990,22 @@ def _atexit(self) -> None: self._lifecycle_owner = None self._lifecycle_condition.notify_all() + @no_throw() + def _atexit_spans(self) -> None: + # The span timer is a daemon thread that dies at exit. + deadline = _get_atexit_deadline() + span_flush = self._start_span_flush(max(0.0, deadline - time.monotonic())) + self._join_span_flush(span_flush, deadline) + + def _join_span_flush( + self, flusher: Optional[threading.Thread], deadline: float + ) -> None: + if flusher is not None: + flusher.join(max(0.0, deadline - time.monotonic())) + if self._traces is not None: + # Not close(): an app's own shutdown() hook may still run and send them. + self._traces.warn_if_queued() + @no_throw() def join(self) -> None: """ @@ -2813,7 +3034,10 @@ def shutdown(self) -> None: 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 + server receipt. Queued spans get one final flush of up to 30 s (plus a + request already in flight); any it cannot send are discarded with a + warning, as are spans still open. + 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 diff --git a/posthog/test/tracing/test_client_traces.py b/posthog/test/tracing/test_client_traces.py new file mode 100644 index 000000000..76d24eaf1 --- /dev/null +++ b/posthog/test/tracing/test_client_traces.py @@ -0,0 +1,716 @@ +import asyncio +import gzip +import json +import threading +import time +from unittest import mock + +import pytest + +import posthog +from posthog import Posthog +from posthog.client import Client +from posthog.contexts import identify_context, new_context, set_context_session +from posthog.test.tracing.helpers import FakeTimer +from posthog.tracing._transport import OK +from posthog.tracing._span import NOOP_SPAN, RecordingSpan, Span +from posthog.version import VERSION + +FAKE_API_KEY = "phc_test_key" +TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736" +SPAN_ID = "00f067aa0ba902b7" + + +def make_client(**kwargs): + kwargs.setdefault("host", "https://us.example.com") + kwargs.setdefault("sync_mode", True) + return Client(FAKE_API_KEY, **kwargs) + + +def mock_session(status_code=200): + session = mock.Mock() + session.post.return_value = mock.Mock(status_code=status_code, headers={}) + return session + + +@pytest.fixture +def no_timers(): + # No background drain racing the test. + with mock.patch.object(threading, "Timer", FakeTimer): + yield + + +@pytest.fixture(autouse=True) +def no_real_network(): + with mock.patch( + "posthog.tracing._transport._get_session", return_value=mock_session() + ): + yield + + +def flush_and_capture(client): + session = mock_session() + with mock.patch("posthog.tracing._transport._get_session", return_value=session): + client.flush() + if not session.post.called: + return None, None, None + args, kwargs = session.post.call_args + return json.loads(gzip.decompress(kwargs["data"])), args[0], kwargs + + +def spans_from(payload): + return payload["resourceSpans"][0]["scopeSpans"][0]["spans"] + + +def resource_from(payload): + return { + kv["key"]: kv["value"] + for kv in payload["resourceSpans"][0]["resource"]["attributes"] + } + + +class TestConfiguration: + def test_is_off_until_the_traces_option_is_supplied(self): + client = make_client() + span = client.start_span("x") + assert span is NOOP_SPAN + span.end() + assert client._traces is None + client.shutdown() + + def test_still_runs_a_with_block_when_tracing_is_off(self): + client = make_client() + ran = False + with client.start_span("job") as span: + ran = True + assert isinstance(span, Span) + assert client.get_active_span() is None + assert ran + client.shutdown() + + def test_passes_an_inbound_trace_through_a_service_that_has_tracing_off(self): + client = make_client() + with client.start_span("x", parent=f"00-{TRACE_ID}-{SPAN_ID}-00") as span: + assert client.get_active_span() is span + assert span.traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-00" + client.shutdown() + + def test_a_span_nested_in_a_pass_through_keeps_forwarding_the_trace(self): + client = make_client() + inbound = f"00-{TRACE_ID}-{SPAN_ID}-00" + with client.start_span("request", parent=inbound, tracestate="vendor=abc"): + with client.start_span("db.query") as nested: + assert nested.traceparent() == inbound + assert nested.tracestate() == "vendor=abc" + assert client.get_active_span().traceparent() == inbound + client.shutdown() + + def test_traces_false_leaves_tracing_off(self): + client = make_client(traces=False) + assert client.start_span("x") is NOOP_SPAN + assert client._traces is None + client.shutdown() + + def test_a_parent_whose_traceparent_raises_gives_a_noop_with_tracing_off(self): + class Hostile(Span): + def traceparent(self): + raise RuntimeError("no") + + client = make_client() + assert client.start_span("x", parent=Hostile()) is NOOP_SPAN + client.shutdown() + + def test_a_bad_traces_config_degrades_to_defaults(self): + client = make_client(traces="nope") + assert isinstance(client.start_span("x"), RecordingSpan) + client.shutdown() + + def test_never_starts_a_pipeline_on_a_client_without_traces(self): + client = make_client() + client.flush() + client.shutdown() + assert client._traces is None + + +class TestTransport: + def test_posts_to_the_traces_endpoint_with_bearer_auth(self): + client = make_client(traces={"service_name": "api"}) + client.start_span("x").end() + payload, url, kwargs = flush_and_capture(client) + assert url == "https://us.example.com/i/v1/traces" + assert kwargs["headers"]["Authorization"] == "Bearer phc_test_key" + assert "token=" not in url + assert len(spans_from(payload)) == 1 + client.shutdown() + + def test_sends_the_service_name_sdk_identity_and_host_os(self): + client = make_client(traces={"service_name": "api"}) + client.start_span("x").end() + payload, _, _ = flush_and_capture(client) + resource = resource_from(payload) + assert resource["service.name"] == {"stringValue": "api"} + assert resource["telemetry.sdk.name"] == {"stringValue": "posthog-python"} + assert resource["telemetry.sdk.version"] == {"stringValue": VERSION} + assert "os.name" in resource + assert payload["resourceSpans"][0]["scopeSpans"][0]["scope"] == { + "name": "posthog-python", + "version": VERSION, + } + client.shutdown() + + def test_lets_configured_resource_attributes_override_the_host_os(self): + client = make_client(traces={"resource_attributes": {"os.name": "Custom"}}) + client.start_span("x").end() + payload, _, _ = flush_and_capture(client) + assert resource_from(payload)["os.name"] == {"stringValue": "Custom"} + client.shutdown() + + def test_exports_well_formed_ids_and_string_nanosecond_timestamps(self): + client = make_client(traces={}) + client.start_span("x", attributes={"n": 7}).end() + payload, _, _ = flush_and_capture(client) + (span,) = spans_from(payload) + assert len(span["traceId"]) == 32 and len(span["spanId"]) == 16 + assert span["startTimeUnixNano"].isdigit() and span["endTimeUnixNano"].isdigit() + assert int(span["endTimeUnixNano"]) >= int(span["startTimeUnixNano"]) + assert {"key": "n", "value": {"intValue": "7"}} in span["attributes"] + client.shutdown() + + def test_send_false_records_but_never_posts(self): + client = make_client(send=False, traces={}) + client.start_span("x").end() + _, url, _ = flush_and_capture(client) + assert url is None + assert client._traces._exporter._queue == [] + client.shutdown() + + +class TestActiveSpanContext: + def test_nests_spans_started_inside_a_with_block(self): + client = make_client(traces={}) + with client.start_span("parent"): + client.start_span("child").end() + payload, _, _ = flush_and_capture(client) + child, parent_record = spans_from(payload) + assert child["parentSpanId"] == parent_record["spanId"] + assert child["traceId"] == parent_record["traceId"] + client.shutdown() + + def test_keeps_the_span_active_across_an_await(self): + client = make_client(traces={}) + + async def handler(): + with client.start_span("request") as span: + await asyncio.sleep(0) + assert client.get_active_span() is span + with client.start_span("inner"): + await asyncio.sleep(0) + await asyncio.sleep(0) + assert client.get_active_span() is span + + asyncio.run(handler()) + payload, _, _ = flush_and_capture(client) + inner, request = spans_from(payload) + assert inner["parentSpanId"] == request["spanId"] + client.shutdown() + + def test_isolates_concurrent_tasks_from_each_other(self): + client = make_client(traces={}) + + async def task(name): + with client.start_span(name) as span: + await asyncio.sleep(0.01) + assert client.get_active_span() is span + + async def run(): + await asyncio.gather(task("a"), task("b")) + + asyncio.run(run()) + payload, _, _ = flush_and_capture(client) + assert all("parentSpanId" not in s for s in spans_from(payload)) + client.shutdown() + + def test_parents_each_tasks_children_to_that_tasks_root(self): + client = make_client(traces={}) + + async def task(name): + with client.start_span(name): + await asyncio.sleep(0.01) + client.start_span(name + ".child").end() + + async def run(): + await asyncio.gather(task("a"), task("b")) + + asyncio.run(run()) + payload, _, _ = flush_and_capture(client) + spans = {s["name"]: s for s in spans_from(payload)} + for name in ("a", "b"): + assert spans[name + ".child"]["parentSpanId"] == spans[name]["spanId"] + assert spans[name + ".child"]["traceId"] == spans[name]["traceId"] + client.shutdown() + + def test_isolates_concurrent_threads_from_each_other(self): + client = make_client(traces={}) + barrier = threading.Barrier(2, timeout=5) + seen = {} + + def work(name): + with client.start_span(name) as span: + barrier.wait() + seen[name] = client.get_active_span() is span + + threads = [threading.Thread(target=work, args=(n,)) for n in ("a", "b")] + for t in threads: + t.start() + for t in threads: + t.join() + assert seen == {"a": True, "b": True} + client.shutdown() + + def test_two_clients_do_not_share_an_active_span(self): + first = make_client(traces={}) + second = make_client(traces={}) + with first.start_span("a"): + assert second.get_active_span() is None + second.start_span("b").end() + (record,) = second._traces._exporter._queue + assert record.parent_span_id is None + first.shutdown() + second.shutdown() + + def test_reads_none_outside_any_block(self): + client = make_client(traces={}) + assert client.get_active_span() is None + client.shutdown() + + +class TestAutoContext: + def test_attaches_the_request_distinct_id_and_session_id(self): + client = make_client(traces={}) + with new_context(fresh=True, capture_exceptions=False): + identify_context("user-1") + set_context_session("sess-1") + client.start_span("x").end() + payload, _, _ = flush_and_capture(client) + (span,) = spans_from(payload) + assert {"key": "posthogDistinctId", "value": {"stringValue": "user-1"}} in span[ + "attributes" + ] + assert {"key": "sessionId", "value": {"stringValue": "sess-1"}} in span[ + "attributes" + ] + client.shutdown() + + def test_omits_the_keys_outside_a_request_context(self): + client = make_client(traces={}) + with new_context(fresh=True, capture_exceptions=False): + client.start_span("x").end() + payload, _, _ = flush_and_capture(client) + (span,) = spans_from(payload) + assert "attributes" not in span + client.shutdown() + + +class TestDistributedTracing: + def test_continues_a_trace_from_an_inbound_traceparent_and_propagates_the_flag( + self, + ): + client = make_client(traces={}) + with client.start_span("x", parent=f"00-{TRACE_ID}-{SPAN_ID}-00") as span: + outgoing = span.traceparent() + assert outgoing.startswith(f"00-{TRACE_ID}-") and outgoing.endswith("-00") + payload, _, _ = flush_and_capture(client) + (record,) = spans_from(payload) + assert record["traceId"] == TRACE_ID + assert record["parentSpanId"] == SPAN_ID + assert record["flags"] == 0x300 + client.shutdown() + + def test_records_a_raised_error_and_rethrows_it_unchanged(self): + client = make_client(traces={}) + error = RuntimeError("boom") + try: + with client.start_span("x"): + raise error + except RuntimeError as raised: + assert raised is error + payload, _, _ = flush_and_capture(client) + (record,) = spans_from(payload) + assert record["status"] == {"code": 2, "message": "boom"} + client.shutdown() + + +class TestLifecycle: + def test_flush_drains_queued_spans(self): + client = make_client(traces={}) + client.start_span("x").end() + payload, _, _ = flush_and_capture(client) + assert len(spans_from(payload)) == 1 + assert client._traces._exporter._queue == [] + client.shutdown() + + def test_flush_resolves_when_the_span_export_fails(self): + client = make_client(traces={}) + client.start_span("x").end() + with mock.patch( + "posthog.tracing._transport._get_session", return_value=mock_session(500) + ): + client.flush() + assert len(client._traces._exporter._queue) == 1 + client.shutdown() + + def test_flush_stops_starting_span_requests_once_its_budget_is_spent( + self, no_timers + ): + client = make_client(traces={"max_export_batch_size": 1}) + client.start_span("x").end() + client.start_span("y").end() + requests = [] + + def slow_send(pipeline_client, payload): + requests.append(payload) + time.sleep(0.2) + return OK + + client._traces._exporter._send = slow_send + client.flush(timeout_seconds=0.05) + assert len(requests) == 1 + assert len(client._traces._exporter._queue) == 1 + client._traces._exporter._send = lambda pipeline_client, payload: OK + client.shutdown() + + def test_flush_without_a_timeout_drains_every_queued_span(self): + client = make_client(traces={"max_export_batch_size": 1}) + session = mock_session() + # Spans end inside the patch: at a batch size of 1 each end arms the + # depth trigger, whose background drain must post to this session too. + with mock.patch( + "posthog.tracing._transport._get_session", return_value=session + ): + client.start_span("x").end() + client.start_span("y").end() + client.flush(timeout_seconds=None) + assert session.post.call_count == 2 + assert client._traces._exporter._queue == [] + client.shutdown() + + def test_flush_sends_spans_while_events_are_still_draining(self, no_timers): + client = make_client(traces={}, sync_mode=False) + client.start_span("x").end() + span_sent = threading.Event() + + def send(pipeline_client, payload): + span_sent.set() + return OK + + client._traces._exporter._send = send + overlapped = [] + for lane in client._lanes: + lane.flush = lambda timeout: overlapped.append(span_sent.wait(2)) + client.flush() + assert overlapped and all(overlapped) + client.shutdown() + + def test_flush_sends_spans_inline_when_no_thread_can_start(self, no_timers): + client = make_client(traces={}) + client.start_span("x").end() + with mock.patch("posthog.client.threading.Thread") as thread: + thread.return_value.start.side_effect = RuntimeError( + "can't create new thread at interpreter shutdown" + ) + payload, _, _ = flush_and_capture(client) + assert len(spans_from(payload)) == 1 + client.shutdown() + + def test_shutdown_flushes_pending_spans(self): + client = make_client(traces={}) + client.start_span("x").end() + session = mock_session() + with mock.patch( + "posthog.tracing._transport._get_session", return_value=session + ): + client.shutdown() + assert session.post.called + assert client._traces._exporter._queue == [] + + def test_shutdown_bounds_the_final_span_flush_and_warns_about_the_rest( + self, no_timers, caplog + ): + caplog.set_level("WARNING", logger="posthog") + client = make_client(traces={"max_export_batch_size": 1}) + for name in ("a", "b", "c"): + client.start_span(name).end() + requests = [] + + def slow_send(pipeline_client, payload): + requests.append(payload) + time.sleep(0.2) + return OK + + client._traces._exporter._send = slow_send + with mock.patch("posthog.client._TRACES_SHUTDOWN_FLUSH_SECONDS", 0.05): + client.shutdown() + assert len(requests) == 1 + assert any("Discarding 2 span(s)" in r.getMessage() for r in caplog.records) + + def test_tracing_is_inert_after_shutdown(self, no_timers): + client = make_client(traces={}) + client.start_span("before").end() + client.shutdown() + assert client.start_span("late") is NOOP_SPAN + assert client._traces._exporter._flush_timer is None + + def test_tracing_never_starts_after_shutdown(self, no_timers): + client = make_client(traces={}) + client.shutdown() + assert client.start_span("late") is NOOP_SPAN + assert client._traces is None + + def test_exit_drains_spans_the_timer_would_have_sent(self, no_timers): + client = make_client(traces={}, sync_mode=False) + client.start_span("x").end() + session = mock_session() + # The exit deadline is process-wide and set once; start a fresh one. + with ( + mock.patch("posthog.tracing._transport._get_session", return_value=session), + mock.patch("posthog.client._atexit_deadline", None), + ): + client._atexit() + assert session.post.called + # The exit hook leaves tracing open for a later shutdown(). + client.start_span("y").end() + assert len(client._traces._exporter._queue) == 1 + session = mock_session() + with mock.patch( + "posthog.tracing._transport._get_session", return_value=session + ): + client.shutdown() + assert session.post.called + + def test_exit_flushes_spans_alongside_events_that_use_up_the_budget( + self, no_timers + ): + client = make_client(traces={}, sync_mode=False) + client.start_span("x").end() + session = mock_session() + + def slow_lane_flush(timeout_seconds): + time.sleep(0.3) + + with ( + mock.patch("posthog.tracing._transport._get_session", return_value=session), + mock.patch("posthog.client._ATEXIT_FLUSH_TIMEOUT_SECONDS", 0.2), + mock.patch("posthog.client._atexit_deadline", None), + mock.patch.object(client._lanes[0], "flush", slow_lane_flush), + ): + client._atexit() + assert session.post.called + client.shutdown() + + def test_exit_does_not_wait_on_a_hung_span_request(self, no_timers, caplog): + caplog.set_level("WARNING", logger="posthog") + client = make_client(traces={}, sync_mode=False) + client.start_span("x").end() + release = threading.Event() + session = mock.Mock() + + def hung_post(*args, **kwargs): + release.wait(5) + return mock.Mock(status_code=200, headers={}) + + session.post.side_effect = hung_post + with ( + mock.patch("posthog.tracing._transport._get_session", return_value=session), + mock.patch("posthog.client._ATEXIT_FLUSH_TIMEOUT_SECONDS", 0.2), + mock.patch("posthog.client._atexit_deadline", None), + ): + try: + started = time.monotonic() + client._atexit() + elapsed = time.monotonic() - started + finally: + release.set() + assert session.post.called + assert elapsed < 1 + assert any( + "1 span(s) were still queued at exit" in r.getMessage() + for r in caplog.records + ) + client.shutdown() + + @pytest.mark.parametrize( + "sync_mode, hook", [(False, "_atexit"), (True, "_atexit_spans")] + ) + def test_exit_warns_about_spans_it_could_not_send( + self, no_timers, caplog, sync_mode, hook + ): + caplog.set_level("WARNING", logger="posthog") + client = make_client(traces={}, sync_mode=sync_mode) + client.start_span("x").end() + with ( + mock.patch( + "posthog.tracing._transport._get_session", + return_value=mock_session(503), + ), + mock.patch("posthog.client._atexit_deadline", None), + ): + getattr(client, hook)() + assert any( + "1 span(s) were still queued at exit" in r.getMessage() + for r in caplog.records + ) + client.shutdown() + + @pytest.mark.parametrize("sync_mode", [True, False]) + def test_an_app_exit_hook_registered_earlier_still_gets_to_flush_spans( + self, no_timers, caplog, sync_mode + ): + caplog.set_level("WARNING", logger="posthog") + hooks = [] + holder = {} + requests = [] + + def slow_send(pipeline_client, payload): + requests.append(payload) + time.sleep(0.2) + return OK + + with ( + mock.patch("posthog.client.atexit.register", side_effect=hooks.append), + mock.patch("posthog.client._ATEXIT_FLUSH_TIMEOUT_SECONDS", 0.1), + mock.patch("posthog.client._atexit_deadline", None), + ): + # The app's own hook, registered before the client's. + hooks.append(lambda: holder["client"].shutdown()) + client = holder["client"] = make_client( + traces={"max_export_batch_size": 1}, sync_mode=sync_mode + ) + client.start_span("a").end() + client._traces._exporter._send = slow_send + client.start_span("b").end() + client.start_span("c").end() + assert len(hooks) == 2 + # atexit runs hooks last-registered first. + for hook in reversed(hooks): + hook() + assert len(requests) == 3 + assert not any("Discarding" in r.getMessage() for r in caplog.records) + + def test_sync_mode_registers_the_span_exit_drain_when_tracing_starts(self): + with mock.patch("posthog.client.atexit.register") as register: + client = make_client(traces={}, sync_mode=True) + register.assert_not_called() + client.start_span("x").end() + client.start_span("y").end() + register.assert_called_once_with(client._atexit_spans) + client.shutdown() + + @pytest.mark.parametrize("traces", [False, None]) + def test_sync_mode_without_tracing_registers_no_exit_hook(self, traces): + with mock.patch("posthog.client.atexit.register") as register: + client = make_client(traces=traces, sync_mode=True) + client.start_span("x").end() + register.assert_not_called() + client.shutdown() + + def test_the_span_exit_drain_leaves_sync_mode_events_alone(self): + client = make_client(traces={}, sync_mode=True) + client.start_span("x").end() + with mock.patch("posthog.client._atexit_deadline", None): + client._atexit_spans() + with mock.patch("posthog.client.batch_post") as batch_post: + client.capture("after-exit", distinct_id="d") + batch_post.assert_called_once() + assert batch_post.call_args[1]["batch"][0]["event"] == "after-exit" + client.shutdown() + + @pytest.mark.parametrize("traces", [{}, None]) + def test_background_mode_registers_the_exit_hook_once(self, traces): + with mock.patch("posthog.client.atexit.register") as register: + client = make_client(traces=traces, sync_mode=False) + register.assert_called_once_with(client._atexit) + client.start_span("x").end() + register.assert_called_once_with(client._atexit) + client.shutdown() + + def test_a_forked_child_drops_the_inherited_queue(self): + client = make_client(traces={}) + client.start_span("parent-span").end() + client._reinit_after_fork() + assert client._traces._exporter._queue == [] + client.start_span("child-span").end() + assert [r.name for r in client._traces._exporter._queue] == ["child-span"] + client.shutdown() + + def test_a_forked_child_does_not_inherit_the_active_span(self): + client = make_client(traces={}) + with client.start_span("parent-span"): + client._reinit_after_fork() + assert client.get_active_span() is None + client.shutdown() + + +class TestModuleLevelApi: + def _with_module_client(self, traces, body): + saved = ( + posthog.default_client, + posthog.api_key, + posthog.host, + posthog.sync_mode, + posthog.traces, + ) + posthog.default_client = None + posthog.api_key = FAKE_API_KEY + posthog.host = "https://us.example.com" + posthog.sync_mode = True + posthog.traces = traces + try: + body() + finally: + if posthog.default_client is not None: + posthog.default_client.shutdown() + ( + posthog.default_client, + posthog.api_key, + posthog.host, + posthog.sync_mode, + posthog.traces, + ) = saved + + def test_module_config_flows_to_the_default_client(self): + def body(): + with posthog.start_span("x") as span: + assert posthog.get_active_span() is span + payload, _, _ = flush_and_capture(posthog.default_client) + assert resource_from(payload)["service.name"] == { + "stringValue": "module-configured" + } + + self._with_module_client({"service_name": "module-configured"}, body) + + def test_module_config_set_after_setup_applies_until_first_use(self): + def body(): + posthog.setup() + posthog.traces = {"service_name": "late-configured"} + posthog.setup() + posthog.start_span("x").end() + payload, _, _ = flush_and_capture(posthog.default_client) + assert resource_from(payload)["service.name"] == { + "stringValue": "late-configured" + } + + self._with_module_client(None, body) + + def test_module_start_span_is_inert_without_config(self): + def body(): + assert posthog.start_span("x") is NOOP_SPAN + assert posthog.get_active_span() is None + + self._with_module_client(None, body) + + def test_posthog_alias_accepts_the_traces_option(self): + client = Posthog( + FAKE_API_KEY, host="https://us.example.com", sync_mode=True, traces={} + ) + assert isinstance(client.start_span("x"), RecordingSpan) + client.shutdown() diff --git a/posthog/tracing/__init__.py b/posthog/tracing/__init__.py index 7d17992ef..05dca68cc 100644 --- a/posthog/tracing/__init__.py +++ b/posthog/tracing/__init__.py @@ -1,3 +1 @@ """Distributed tracing: spans exported as OTLP JSON to PostHog.""" - -__all__: list[str] = [] diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 1b6089d35..a0270e0a0 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -28,6 +28,7 @@ alias posthog.OptionalCaptureArgs -> posthog.args.OptionalCaptureArgs alias posthog.OptionalSetArgs -> posthog.args.OptionalSetArgs alias posthog.RequiresServerEvaluation -> posthog.feature_flags.RequiresServerEvaluation alias posthog.SocketOptions -> posthog.request.SocketOptions +alias posthog.Span -> posthog.tracing.span.Span alias posthog.VERSION -> posthog.version.VERSION alias posthog.ai.PromptResult -> posthog.ai.prompts.PromptResult alias posthog.ai.PromptSource -> posthog.ai.prompts.PromptSource @@ -268,6 +269,7 @@ alias posthog.client.RequestsTimeout -> posthog.request.RequestsTimeout alias posthog.client.RequiresServerEvaluation -> posthog.feature_flags.RequiresServerEvaluation alias posthog.client.SendFeatureFlagsOptions -> posthog.types.SendFeatureFlagsOptions alias posthog.client.SizeLimitedDict -> posthog.utils.SizeLimitedDict +alias posthog.client.Span -> posthog.tracing.span.Span alias posthog.client.VERSION -> posthog.version.VERSION alias posthog.client.batch_post -> posthog.request.batch_post alias posthog.client.clean -> posthog.utils.clean @@ -887,6 +889,7 @@ attribute posthog.secret_key = None attribute posthog.send = True attribute posthog.super_properties = None attribute posthog.sync_mode = False +attribute posthog.traces = None attribute posthog.types.BeforeSendCallback = Callable[[dict[str, Any]], Optional[dict[str, Any]]] attribute posthog.types.BeforeSpanSendCallback = Callable[[dict[str, Any]], Optional[dict[str, Any]]] attribute posthog.types.FeatureFlag.enabled: bool @@ -1011,7 +1014,7 @@ class posthog.bucketed_rate_limiter.BucketedRateLimiter(bucket_size: Number, ref class posthog.capture_compression.CaptureCompression class posthog.capture_mode.CaptureMode class posthog.capture_v1.CaptureV1Error(status: int | str, message: str, *, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None) -class posthog.client.Client(project_api_key: str, host=None, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, gzip=False, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_mode: Optional[Union[CaptureMode, str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False) +class posthog.client.Client(project_api_key: str, host=None, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, gzip=False, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_mode: Optional[Union[CaptureMode, str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, gzip=False, retries=10, timeout=15, historical_migration=False, endpoint=EVENTS_ENDPOINT, max_msg_size=MAX_MSG_SIZE, capture_mode=CaptureMode.V0, capture_compression=CaptureCompression.NONE) class posthog.contexts.ContextScope(parent=None, fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) class posthog.exception_capture.ExceptionCapture(client: Client, rate_limiting_enabled=False, bucket_size=DEFAULT_BUCKET_SIZE, refill_rate=DEFAULT_REFILL_RATE, refill_interval_seconds=DEFAULT_REFILL_INTERVAL_SECONDS) @@ -1214,6 +1217,7 @@ function posthog.feature_flags.relative_date_parse_for_feature_flag_matching(val function posthog.feature_flags.resolve_bucketing_value(flag, distinct_id, device_id=None) function posthog.feature_flags.variant_lookup_table(feature_flag) function posthog.flush(timeout_seconds: Optional[float] = 10) -> None +function posthog.get_active_span() -> Optional[Span] function posthog.get_all_flags(distinct_id: ID_TYPES, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None, flag_keys_to_evaluate: Optional[list[str]] = None) -> Optional[dict[str, FlagValue]] function posthog.get_all_flags_and_payloads(distinct_id: ID_TYPES, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None, flag_keys_to_evaluate: Optional[list[str]] = None) -> FlagsAndPayloads function posthog.get_feature_flag(key: str, distinct_id: ID_TYPES, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, send_feature_flag_events: bool = True, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None) -> Optional[FlagValue] @@ -1262,6 +1266,7 @@ function posthog.set_context_session(session_id: str) function posthog.set_once(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] function posthog.setup() -> Client function posthog.shutdown() -> None +function posthog.start_span(name: str, *, kind: Optional[str] = None, attributes: Optional[Mapping[str, Any]] = None, parent: Union[Span, str, None] = None, tracestate: Optional[str] = None, start_time: Union[datetime.datetime, float, None] = None) -> Span function posthog.tag(name: str, value: Any) function posthog.types.normalize_flags_response(resp: Any) -> FlagsResponse function posthog.types.to_flags_and_payloads(resp: FlagsResponse) -> FlagsAndPayloads @@ -1375,6 +1380,7 @@ method posthog.client.Client.evaluate_flags(distinct_id: Optional[ID_TYPES] = No method posthog.client.Client.feature_enabled(key: str, distinct_id: ID_TYPES, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, send_feature_flag_events: bool = True, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None) -> Optional[bool] method posthog.client.Client.feature_flag_definitions() method posthog.client.Client.flush(timeout_seconds: Optional[float] = 10) -> None +method posthog.client.Client.get_active_span() -> Optional[Span] method posthog.client.Client.get_all_flags(distinct_id: ID_TYPES, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, flag_keys_to_evaluate: Optional[list[str]] = None, device_id: Optional[str] = None) -> Optional[dict[str, Union[bool, str]]] method posthog.client.Client.get_all_flags_and_payloads(distinct_id: ID_TYPES, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, flag_keys_to_evaluate: Optional[list[str]] = None, device_id: Optional[str] = None) -> FlagsAndPayloads method posthog.client.Client.get_feature_flag(key: str, distinct_id: ID_TYPES, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, send_feature_flag_events: bool = True, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None) -> Optional[FlagValue] @@ -1397,6 +1403,7 @@ method posthog.client.Client.set_context_device_id(device_id: str) -> None method posthog.client.Client.set_context_session(session_id: str) -> None method posthog.client.Client.set_once(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] method posthog.client.Client.shutdown() -> None +method posthog.client.Client.start_span(name: str, *, kind: Optional[str] = None, attributes: Optional[Mapping[str, Any]] = None, parent: Union[Span, str, None] = None, tracestate: Optional[str] = None, start_time: Union[datetime, float, None] = None) -> Span method posthog.client.Client.tag(name: str, value: Any) -> None method posthog.consumer.Consumer.next() method posthog.consumer.Consumer.pause() From 38c72ba78efe862ee416c1c472d65b6e45691cf8 Mon Sep 17 00:00:00 2001 From: Anna Garcia Date: Mon, 14 Sep 2026 18:58:14 -0400 Subject: [PATCH 2/4] fix(traces): keep shutdown and fork isolation airtight for the client Shutdown reads the pipeline under the same lock initialization publishes it with, so an initialization in flight is either closed by shutdown or sees the request and stays off. A forked child gets a fresh active-span variable: an inherited handle exiting in the child resets the old one, which would have restored the parent process's outer span. --- posthog/client.py | 15 ++++++---- posthog/test/tracing/test_client_traces.py | 34 ++++++++++++++++++++++ posthog/tracing/_pipeline.py | 4 ++- 3 files changed, 47 insertions(+), 6 deletions(-) diff --git a/posthog/client.py b/posthog/client.py index 06236f659..261b055b6 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -2276,9 +2276,11 @@ def _reinit_after_fork(self): if self._metrics is not None: self._metrics._reinit_after_fork() self._traces_lock = threading.Lock() + # A fresh variable: an inherited handle resets the old one on exit, + # which would restore the parent process's outer span. + self._active_span_var = ContextVar("posthog_active_span", default=None) if self._traces is not None: - self._traces.reinit_after_fork() - self._active_span_var.set(None) + self._traces.reinit_after_fork(self._active_span_var) # If using Redis cache, we must reinitialize to get a fresh connection (fork-safe). # If using Memory cache, we keep it as-is to benefit from the inherited warm cache. @@ -2513,7 +2515,9 @@ def _traces_pipeline(self) -> Optional[PostHogTraces]: return self._traces if self._traces is None: with self._traces_lock: - if self._traces is None: + # Re-checked under the lock, which shutdown takes before it + # reads the pipeline, so neither can miss the other. + if self._traces is None and not self._shutdown_requested: try: config = resolve_traces_config( self._traces_config, host_resource_attributes() @@ -2878,15 +2882,16 @@ def _shutdown_once(self, errors: list[Exception]) -> None: self._run_lifecycle_cleanup( "Failed to reset metrics on shutdown", self._metrics.reset, errors ) - if self._traces is not None: + with self._traces_lock: traces = self._traces + if traces is not None: self._run_lifecycle_cleanup( "Failed to flush spans on shutdown", lambda: traces.flush(_TRACES_SHUTDOWN_FLUSH_SECONDS), errors, ) self._run_lifecycle_cleanup( - "Failed to close traces on shutdown", self._traces.close, errors + "Failed to close traces on shutdown", traces.close, errors ) self._join_once(errors, flush_queues=False, lanes_prepared=True) self._run_lifecycle_cleanup( diff --git a/posthog/test/tracing/test_client_traces.py b/posthog/test/tracing/test_client_traces.py index 76d24eaf1..c09524704 100644 --- a/posthog/test/tracing/test_client_traces.py +++ b/posthog/test/tracing/test_client_traces.py @@ -460,6 +460,24 @@ def test_tracing_is_inert_after_shutdown(self, no_timers): assert client.start_span("late") is NOOP_SPAN assert client._traces._exporter._flush_timer is None + def test_shutdown_closes_a_pipeline_still_initializing(self, no_timers): + client = make_client(traces={}) + shutdown = threading.Thread(target=client.shutdown) + resolve = posthog.client.resolve_traces_config + + def resolve_while_shutting_down(*args): + shutdown.start() + time.sleep(0.05) + return resolve(*args) + + with mock.patch( + "posthog.client.resolve_traces_config", resolve_while_shutting_down + ): + client.start_span("racing").end() + shutdown.join() + assert client._traces._closed + assert client._traces._exporter._flush_timer is None + def test_tracing_never_starts_after_shutdown(self, no_timers): client = make_client(traces={}) client.shutdown() @@ -649,6 +667,22 @@ def test_a_forked_child_does_not_inherit_the_active_span(self): assert client.get_active_span() is None client.shutdown() + def test_an_inherited_span_exiting_in_the_child_does_not_restore_its_parent( + self, + ): + client = make_client(traces={}) + with client.start_span("outer") as outer: + with client.start_span("inner"): + client._reinit_after_fork() + assert client.get_active_span() is None + child = client.start_span("child") + child.end() + (record,) = client._traces._exporter._queue + assert record.name == "child" + assert record.parent_span_id is None + assert record.trace_id != outer.traceparent().split("-")[1] + client.shutdown() + class TestModuleLevelApi: def _with_module_client(self, traces, body): diff --git a/posthog/tracing/_pipeline.py b/posthog/tracing/_pipeline.py index 3e8cd735d..d558800eb 100644 --- a/posthog/tracing/_pipeline.py +++ b/posthog/tracing/_pipeline.py @@ -114,10 +114,12 @@ def warn_if_queued(self) -> None: self._exporter.warn_if_queued() self._drops.warn_if_due(force=True) - def reinit_after_fork(self) -> None: + def reinit_after_fork(self, active_var: Optional[ContextVar] = None) -> None: # Runs in the forked child before user code; the parent's spans stay # with the parent. self._lock = threading.Lock() + if active_var is not None: + self._active_var = active_var self._live_spans.clear() self._drops.reinit_after_fork() self._exporter.reinit_after_fork() From 810c868e7c7e9871ad1e9a7ca16fab50ed3cb1e8 Mon Sep 17 00:00:00 2001 From: Anna Garcia Date: Wed, 16 Sep 2026 21:12:17 -0400 Subject: [PATCH 3/4] fix(traces): keep the module-level setup from turning an explicit traces config off setup() runs on every module-level call and re-applied posthog.traces whenever no pipeline existed yet, so a default client built with its own traces config lost it on the first capture(), and a failed init was retried and re-logged on every call. The module option now applies only where the client has none, and a failed init latches as False. Shutdown unregisters the sync-mode exit drain so the client is collectable, the flush docstring states its worst case, Span is re-exported for pyright strict and checked in CI, the fork ContextVar is required, and the start_span docstring says what a forked child inherits. --- .github/scripts/check_strict_types.sh | 5 ++- posthog/__init__.py | 6 ++-- posthog/client.py | 10 ++++-- posthog/test/tracing/test_client_traces.py | 36 ++++++++++++++++++++++ posthog/test/tracing/test_export.py | 3 +- posthog/test/tracing/test_pipeline.py | 5 +-- posthog/tracing/_pipeline.py | 7 ++--- 7 files changed, 60 insertions(+), 12 deletions(-) diff --git a/.github/scripts/check_strict_types.sh b/.github/scripts/check_strict_types.sh index d25938fa2..bc95aa8d7 100755 --- a/.github/scripts/check_strict_types.sh +++ b/.github/scripts/check_strict_types.sh @@ -26,8 +26,11 @@ all_flags: dict[str, FlagValue] | None = posthog.get_all_flags("user", groups=gr enabled: bool | None = posthog.feature_enabled("flag", "user", groups=groups) payload: object | None = client.get_feature_flag_payload("flag", "user", groups=groups) evaluations: FeatureFlagEvaluations = posthog.evaluate_flags(123, groups=groups) +span: posthog.Span = client.start_span("job") +active: posthog.Span | None = posthog.get_active_span() +span.end() -_ = (flag_value, all_flags, enabled, payload, evaluations) +_ = (flag_value, all_flags, enabled, payload, evaluations, active) PY "$tmp/.venv/bin/python" - <<'PY' > "$tmp/public_api_access.py" diff --git a/posthog/__init__.py b/posthog/__init__.py index 8e86afe1c..092fb00b2 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -12,7 +12,7 @@ from posthog.capture_compression import CaptureCompression as CaptureCompression from posthog.capture_mode import CaptureMode as CaptureMode from posthog.client import Client -from posthog.tracing.span import Span +from posthog.tracing.span import Span as Span from posthog.async_client import AsyncClient as AsyncClient from posthog.async_client import AsyncPosthog as AsyncPosthog from posthog.exception_capture import ExceptionCapture @@ -1393,7 +1393,9 @@ def setup() -> Client: # already forced setup()) still applies until the metrics API is first used. if default_client._metrics is None: default_client._metrics_config = metrics - if default_client._traces is None: + # traces=None means off, so the module option applies only where the client + # has none; False is the latch of an init that failed and stays off. + if traces is not None and default_client._traces_config is None: default_client._traces_config = traces return default_client diff --git a/posthog/client.py b/posthog/client.py index 261b055b6..0deefd584 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -2539,7 +2539,7 @@ def _traces_pipeline(self) -> Optional[PostHogTraces]: # Off rather than defaults: defaults would drop a # before_span_send hook and export unscrubbed spans. self.log.exception("Error initializing traces; tracing is off") - self._traces_config = None + self._traces_config = False return self._traces def _tracing_context(self) -> Dict[str, Optional[str]]: @@ -2578,6 +2578,8 @@ def start_span( attributes: Initial attributes. parent: A span handle, or an inbound W3C ``traceparent`` header value to continue a remote trace. Defaults to the active span. + A forked child starts with no active span; pass the parent + span to continue a trace across a fork. tracestate: The inbound ``tracestate`` header accompanying a ``traceparent`` string ``parent``; preserved and propagated. start_time: A ``datetime`` or epoch seconds, to backdate the span. @@ -2638,7 +2640,9 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: Queued spans are sent at the same time, within the same budget: at least one span request is attempted even when the budget is already spent, no further one starts once it is, and - each request is bounded by ``timeout``. + each request is bounded by ``timeout``. The wait for that first + request is not cut short, so a flush can take up to + ``timeout_seconds`` plus ``timeout`` in the worst case. Examples: ```python @@ -2893,6 +2897,8 @@ def _shutdown_once(self, errors: list[Exception]) -> None: self._run_lifecycle_cleanup( "Failed to close traces on shutdown", traces.close, errors ) + # The sync-mode exit drain, so a shut-down client is collectable. + atexit.unregister(self._atexit_spans) self._join_once(errors, flush_queues=False, lanes_prepared=True) self._run_lifecycle_cleanup( "Failed to clear feature flag deduplication state on shutdown", diff --git a/posthog/test/tracing/test_client_traces.py b/posthog/test/tracing/test_client_traces.py index c09524704..88e0b06c2 100644 --- a/posthog/test/tracing/test_client_traces.py +++ b/posthog/test/tracing/test_client_traces.py @@ -125,6 +125,16 @@ def test_a_bad_traces_config_degrades_to_defaults(self): assert isinstance(client.start_span("x"), RecordingSpan) client.shutdown() + def test_a_failed_init_is_not_retried_on_the_next_call(self): + client = make_client(traces={}) + with mock.patch( + "posthog.client.resolve_traces_config", side_effect=RuntimeError("no") + ) as resolve: + assert client.start_span("x") is NOOP_SPAN + assert client.start_span("y") is NOOP_SPAN + assert resolve.call_count == 1 + client.shutdown() + def test_never_starts_a_pipeline_on_a_client_without_traces(self): client = make_client() client.flush() @@ -623,6 +633,13 @@ def test_sync_mode_registers_the_span_exit_drain_when_tracing_starts(self): register.assert_called_once_with(client._atexit_spans) client.shutdown() + def test_shutdown_unregisters_the_sync_mode_exit_drain(self): + client = make_client(traces={}, sync_mode=True) + client.start_span("x").end() + with mock.patch("posthog.client.atexit.unregister") as unregister: + client.shutdown() + unregister.assert_called_once_with(client._atexit_spans) + @pytest.mark.parametrize("traces", [False, None]) def test_sync_mode_without_tracing_registers_no_exit_hook(self, traces): with mock.patch("posthog.client.atexit.register") as register: @@ -735,6 +752,25 @@ def body(): self._with_module_client(None, body) + def test_setup_leaves_an_explicitly_configured_default_client_alone(self): + def body(): + posthog.default_client = make_client(traces={"service_name": "explicit"}) + posthog.setup() + assert isinstance(posthog.start_span("x"), RecordingSpan) + + self._with_module_client(None, body) + + def test_setup_does_not_retry_a_traces_init_that_failed(self): + def body(): + with mock.patch( + "posthog.client.resolve_traces_config", side_effect=RuntimeError("no") + ) as resolve: + assert posthog.start_span("x") is NOOP_SPAN + assert posthog.start_span("y") is NOOP_SPAN + assert resolve.call_count == 1 + + self._with_module_client({"service_name": "broken"}, body) + def test_module_start_span_is_inert_without_config(self): def body(): assert posthog.start_span("x") is NOOP_SPAN diff --git a/posthog/test/tracing/test_export.py b/posthog/test/tracing/test_export.py index b4aebf2d1..9de73af8f 100644 --- a/posthog/test/tracing/test_export.py +++ b/posthog/test/tracing/test_export.py @@ -1,5 +1,6 @@ import threading import time +from contextvars import ContextVar from types import SimpleNamespace from unittest import mock @@ -940,7 +941,7 @@ def test_a_forked_child_drops_the_inherited_queue_and_timer(self): pipeline.start_span("parent-span").end() assert queued(pipeline) and pipeline._exporter._flush_timer is not None pipeline._exporter._max_export_batch_size = 1 - pipeline.reinit_after_fork() + pipeline.reinit_after_fork(ContextVar("child-active", default=None)) assert queued(pipeline) == [] assert pipeline._exporter._flush_timer is None assert ( diff --git a/posthog/test/tracing/test_pipeline.py b/posthog/test/tracing/test_pipeline.py index da6dbab47..f03529c41 100644 --- a/posthog/test/tracing/test_pipeline.py +++ b/posthog/test/tracing/test_pipeline.py @@ -1,6 +1,7 @@ import gc import logging import threading +from contextvars import ContextVar import time import warnings import weakref @@ -593,14 +594,14 @@ def test_close_makes_later_spans_inert_and_closes_the_exporter(self): def test_a_forked_child_drops_the_parents_live_spans(self): pipeline, exporter, _ = make() pipeline.start_span("live") - pipeline.reinit_after_fork() + pipeline.reinit_after_fork(ContextVar("child-active", default=None)) assert pipeline._live_spans == {} assert exporter.reinitialized def test_reinit_after_fork_replaces_locks_without_acquiring_them(self): pipeline, _, _ = make() pipeline._lock.acquire() - pipeline.reinit_after_fork() + pipeline.reinit_after_fork(ContextVar("child-active", default=None)) assert not pipeline._lock.locked() pipeline.start_span("a").end() diff --git a/posthog/tracing/_pipeline.py b/posthog/tracing/_pipeline.py index d558800eb..896bd1a66 100644 --- a/posthog/tracing/_pipeline.py +++ b/posthog/tracing/_pipeline.py @@ -114,12 +114,11 @@ def warn_if_queued(self) -> None: self._exporter.warn_if_queued() self._drops.warn_if_due(force=True) - def reinit_after_fork(self, active_var: Optional[ContextVar] = None) -> None: + def reinit_after_fork(self, active_var: ContextVar) -> None: # Runs in the forked child before user code; the parent's spans stay - # with the parent. + # with the parent, and the child's active span is the fresh var. self._lock = threading.Lock() - if active_var is not None: - self._active_var = active_var + self._active_var = active_var self._live_spans.clear() self._drops.reinit_after_fork() self._exporter.reinit_after_fork() From 71b482d537fddd48e0dbf7d85ea3fb8c4c810ba1 Mon Sep 17 00:00:00 2001 From: Anna Garcia Date: Wed, 16 Sep 2026 22:08:19 -0400 Subject: [PATCH 4/4] fix(traces): flush the lanes before an inline span flush, and skip the span thread when nothing is queued When no thread could start at exit, the span flush ran before the event lanes and could spend the whole exit budget. _start_span_flush now hands back a waiter the caller runs after the lanes, and starts nothing when the span queue is empty. The flush docstring describes the retry-within- budget contract. Tracing test fixtures move to a conftest, and the client tests reuse the shared ids and one slow sender. --- posthog/client.py | 41 +++++---- posthog/test/tracing/conftest.py | 32 +++++++ posthog/test/tracing/helpers.py | 29 +------ posthog/test/tracing/test_client_traces.py | 99 +++++++++++++--------- posthog/test/tracing/test_export.py | 5 +- posthog/test/tracing/test_pipeline.py | 4 +- posthog/tracing/_export.py | 4 + posthog/tracing/_pipeline.py | 5 ++ 8 files changed, 129 insertions(+), 90 deletions(-) create mode 100644 posthog/test/tracing/conftest.py diff --git a/posthog/client.py b/posthog/client.py index 0deefd584..857bd743e 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -2639,10 +2639,13 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: Defaults to 10 seconds. Pass ``None`` to wait indefinitely. Queued spans are sent at the same time, within the same budget: at least one span request is attempted even when the - budget is already spent, no further one starts once it is, and - each request is bounded by ``timeout``. The wait for that first - request is not cut short, so a flush can take up to - ``timeout_seconds`` plus ``timeout`` in the worst case. + budget is already spent, a retriable failure is retried after + its backoff while budget remains, no other request starts once + it is spent, and each request is bounded by ``timeout``. The + wait for the last request is not cut short, so a flush can + take up to ``timeout_seconds`` plus ``timeout`` in the worst + case. A span flush already in flight for the whole wait is + left to finish instead. Examples: ```python @@ -2664,19 +2667,25 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: for lane in self._lanes: lane.flush(max(0.0, deadline - time.monotonic())) if span_flush is not None: - # The first span request is exempt from the budget and bounded - # only by the request timeout, so the join is not. - span_flush.join() + # The last span request is bounded only by the request + # timeout, so the wait is not. + span_flush(None) except Exception as e: self.log.exception("error flushing queue: %s", e) return def _start_span_flush( self, timeout_seconds: Optional[float] - ) -> Optional[threading.Thread]: - """Flush spans alongside the events, so a handler waits one round trip, not two.""" + ) -> Optional[Callable[[Optional[float]], None]]: + """Flush spans alongside the events, so a handler waits one round trip, not two. + + Returns a waiter taking the seconds to wait, or ``None`` when nothing + is queued. When no thread can start (interpreter shutdown), the waiter + runs the flush on the calling thread, after the caller has flushed the + event lanes. + """ traces = self._traces - if traces is None: + if traces is None or not traces.has_queued_spans(): return None def flush_spans() -> None: @@ -2691,10 +2700,8 @@ def flush_spans() -> None: try: flusher.start() except RuntimeError: - # No new threads at interpreter shutdown; flush on this one. - flush_spans() - return None - return flusher + return lambda _seconds: flush_spans() + return flusher.join def _is_consumer_thread(self) -> bool: current = threading.current_thread() @@ -3009,10 +3016,10 @@ def _atexit_spans(self) -> None: self._join_span_flush(span_flush, deadline) def _join_span_flush( - self, flusher: Optional[threading.Thread], deadline: float + self, waiter: Optional[Callable[[Optional[float]], None]], deadline: float ) -> None: - if flusher is not None: - flusher.join(max(0.0, deadline - time.monotonic())) + if waiter is not None: + waiter(max(0.0, deadline - time.monotonic())) if self._traces is not None: # Not close(): an app's own shutdown() hook may still run and send them. self._traces.warn_if_queued() diff --git a/posthog/test/tracing/conftest.py b/posthog/test/tracing/conftest.py new file mode 100644 index 000000000..6baf9bc02 --- /dev/null +++ b/posthog/test/tracing/conftest.py @@ -0,0 +1,32 @@ +"""Fixtures shared by the tracing tests.""" + +import threading +import time +from unittest import mock + +import pytest + +from posthog.test.tracing.helpers import FakeTimer +from posthog.tracing import _export as export_module + + +@pytest.fixture +def fake_timers(): + """Timers that fire only when a test says so.""" + FakeTimer.instances = [] + with mock.patch.object(threading, "Timer", FakeTimer): + yield FakeTimer + + +@pytest.fixture +def no_jitter(): + # Backoff delays are asserted exactly; TestJitter covers the spread. + with mock.patch.object(export_module, "_draw_jitter", return_value=1.0): + yield + + +@pytest.fixture +def clock(): + state = {"now": 1000.0} + with mock.patch.object(time, "monotonic", lambda: state["now"]): + yield state diff --git a/posthog/test/tracing/helpers.py b/posthog/test/tracing/helpers.py index 9c0790bc0..4ad041572 100644 --- a/posthog/test/tracing/helpers.py +++ b/posthog/test/tracing/helpers.py @@ -1,14 +1,9 @@ """Shared fakes for the tracing pipeline and export tests.""" import threading -import time from contextvars import ContextVar from types import SimpleNamespace -from unittest import mock -import pytest - -from posthog.tracing import _export as export_module from posthog.tracing._config import resolve_traces_config from posthog.tracing._drops import DropLog from posthog.tracing._export import SpanExporter @@ -80,31 +75,13 @@ def close(self): def warn_if_queued(self): pass + def has_queued(self): + return bool(self.records) + def reinit_after_fork(self): self.reinitialized = True -@pytest.fixture(autouse=True) -def fake_timers(): - FakeTimer.instances = [] - with mock.patch.object(threading, "Timer", FakeTimer): - yield FakeTimer - - -@pytest.fixture(autouse=True) -def no_jitter(): - # Backoff delays are asserted exactly; TestJitter covers the spread. - with mock.patch.object(export_module, "_draw_jitter", return_value=1.0): - yield - - -@pytest.fixture -def clock(): - state = {"now": 1000.0} - with mock.patch.object(time, "monotonic", lambda: state["now"]): - yield state - - def make(client=None, context=None, **config): """A pipeline whose ended spans collect on a ``RecordingExporter``.""" client = client or SimpleNamespace(disabled=False, send=True) diff --git a/posthog/test/tracing/test_client_traces.py b/posthog/test/tracing/test_client_traces.py index 88e0b06c2..ae7072b88 100644 --- a/posthog/test/tracing/test_client_traces.py +++ b/posthog/test/tracing/test_client_traces.py @@ -11,14 +11,12 @@ from posthog import Posthog from posthog.client import Client from posthog.contexts import identify_context, new_context, set_context_session -from posthog.test.tracing.helpers import FakeTimer +from posthog.test.tracing.helpers import SPAN_ID, TRACE_ID from posthog.tracing._transport import OK from posthog.tracing._span import NOOP_SPAN, RecordingSpan, Span from posthog.version import VERSION FAKE_API_KEY = "phc_test_key" -TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736" -SPAN_ID = "00f067aa0ba902b7" def make_client(**kwargs): @@ -33,11 +31,15 @@ def mock_session(status_code=200): return session -@pytest.fixture -def no_timers(): - # No background drain racing the test. - with mock.patch.object(threading, "Timer", FakeTimer): - yield +def slow_send(requests, delay=0.2): + """A sender that records each payload and takes ``delay`` seconds to answer.""" + + def send(pipeline_client, payload): + requests.append(payload) + time.sleep(delay) + return OK + + return send @pytest.fixture(autouse=True) @@ -135,6 +137,14 @@ def test_a_failed_init_is_not_retried_on_the_next_call(self): assert resolve.call_count == 1 client.shutdown() + def test_a_non_callable_hook_turns_tracing_off(self, caplog): + caplog.set_level("ERROR", logger="posthog") + client = make_client(traces={"before_span_send": "scrub"}) + assert client.start_span("x") is NOOP_SPAN + assert "Error initializing traces" in caplog.text + assert "not callable" in caplog.text + client.shutdown() + def test_never_starts_a_pipeline_on_a_client_without_traces(self): client = make_client() client.flush() @@ -370,19 +380,13 @@ def test_flush_resolves_when_the_span_export_fails(self): client.shutdown() def test_flush_stops_starting_span_requests_once_its_budget_is_spent( - self, no_timers + self, fake_timers ): client = make_client(traces={"max_export_batch_size": 1}) client.start_span("x").end() client.start_span("y").end() requests = [] - - def slow_send(pipeline_client, payload): - requests.append(payload) - time.sleep(0.2) - return OK - - client._traces._exporter._send = slow_send + client._traces._exporter._send = slow_send(requests) client.flush(timeout_seconds=0.05) assert len(requests) == 1 assert len(client._traces._exporter._queue) == 1 @@ -404,7 +408,7 @@ def test_flush_without_a_timeout_drains_every_queued_span(self): assert client._traces._exporter._queue == [] client.shutdown() - def test_flush_sends_spans_while_events_are_still_draining(self, no_timers): + def test_flush_sends_spans_while_events_are_still_draining(self, fake_timers): client = make_client(traces={}, sync_mode=False) client.start_span("x").end() span_sent = threading.Event() @@ -421,7 +425,7 @@ def send(pipeline_client, payload): assert overlapped and all(overlapped) client.shutdown() - def test_flush_sends_spans_inline_when_no_thread_can_start(self, no_timers): + def test_flush_sends_spans_inline_when_no_thread_can_start(self, fake_timers): client = make_client(traces={}) client.start_span("x").end() with mock.patch("posthog.client.threading.Thread") as thread: @@ -432,6 +436,32 @@ def test_flush_sends_spans_inline_when_no_thread_can_start(self, no_timers): assert len(spans_from(payload)) == 1 client.shutdown() + def test_flush_starts_no_span_thread_when_nothing_is_queued(self, fake_timers): + client = make_client(traces={}) + client.start_span("x").end() + client.flush() + with mock.patch("posthog.client.threading.Thread") as thread: + client.flush() + thread.assert_not_called() + client.shutdown() + + def test_exit_flushes_the_lanes_before_an_inline_span_flush(self, fake_timers): + client = make_client(traces={}, sync_mode=False) + client.start_span("x").end() + order = [] + client._traces.flush = lambda timeout: order.append("spans") + for lane in client._lanes: + lane.flush = lambda timeout, _lane=lane: order.append("lanes") + with ( + mock.patch("posthog.client.threading.Thread") as thread, + mock.patch("posthog.client._atexit_deadline", None), + ): + thread.return_value.start.side_effect = RuntimeError("no threads") + client._atexit() + assert order[0] == "lanes" + assert order[-1] == "spans" + client.shutdown() + def test_shutdown_flushes_pending_spans(self): client = make_client(traces={}) client.start_span("x").end() @@ -444,33 +474,27 @@ def test_shutdown_flushes_pending_spans(self): assert client._traces._exporter._queue == [] def test_shutdown_bounds_the_final_span_flush_and_warns_about_the_rest( - self, no_timers, caplog + self, fake_timers, caplog ): caplog.set_level("WARNING", logger="posthog") client = make_client(traces={"max_export_batch_size": 1}) for name in ("a", "b", "c"): client.start_span(name).end() requests = [] - - def slow_send(pipeline_client, payload): - requests.append(payload) - time.sleep(0.2) - return OK - - client._traces._exporter._send = slow_send + client._traces._exporter._send = slow_send(requests) with mock.patch("posthog.client._TRACES_SHUTDOWN_FLUSH_SECONDS", 0.05): client.shutdown() assert len(requests) == 1 assert any("Discarding 2 span(s)" in r.getMessage() for r in caplog.records) - def test_tracing_is_inert_after_shutdown(self, no_timers): + def test_tracing_is_inert_after_shutdown(self, fake_timers): client = make_client(traces={}) client.start_span("before").end() client.shutdown() assert client.start_span("late") is NOOP_SPAN assert client._traces._exporter._flush_timer is None - def test_shutdown_closes_a_pipeline_still_initializing(self, no_timers): + def test_shutdown_closes_a_pipeline_still_initializing(self, fake_timers): client = make_client(traces={}) shutdown = threading.Thread(target=client.shutdown) resolve = posthog.client.resolve_traces_config @@ -488,13 +512,13 @@ def resolve_while_shutting_down(*args): assert client._traces._closed assert client._traces._exporter._flush_timer is None - def test_tracing_never_starts_after_shutdown(self, no_timers): + def test_tracing_never_starts_after_shutdown(self, fake_timers): client = make_client(traces={}) client.shutdown() assert client.start_span("late") is NOOP_SPAN assert client._traces is None - def test_exit_drains_spans_the_timer_would_have_sent(self, no_timers): + def test_exit_drains_spans_the_timer_would_have_sent(self, fake_timers): client = make_client(traces={}, sync_mode=False) client.start_span("x").end() session = mock_session() @@ -516,7 +540,7 @@ def test_exit_drains_spans_the_timer_would_have_sent(self, no_timers): assert session.post.called def test_exit_flushes_spans_alongside_events_that_use_up_the_budget( - self, no_timers + self, fake_timers ): client = make_client(traces={}, sync_mode=False) client.start_span("x").end() @@ -535,7 +559,7 @@ def slow_lane_flush(timeout_seconds): assert session.post.called client.shutdown() - def test_exit_does_not_wait_on_a_hung_span_request(self, no_timers, caplog): + def test_exit_does_not_wait_on_a_hung_span_request(self, fake_timers, caplog): caplog.set_level("WARNING", logger="posthog") client = make_client(traces={}, sync_mode=False) client.start_span("x").end() @@ -570,7 +594,7 @@ def hung_post(*args, **kwargs): "sync_mode, hook", [(False, "_atexit"), (True, "_atexit_spans")] ) def test_exit_warns_about_spans_it_could_not_send( - self, no_timers, caplog, sync_mode, hook + self, fake_timers, caplog, sync_mode, hook ): caplog.set_level("WARNING", logger="posthog") client = make_client(traces={}, sync_mode=sync_mode) @@ -591,18 +615,13 @@ def test_exit_warns_about_spans_it_could_not_send( @pytest.mark.parametrize("sync_mode", [True, False]) def test_an_app_exit_hook_registered_earlier_still_gets_to_flush_spans( - self, no_timers, caplog, sync_mode + self, fake_timers, caplog, sync_mode ): caplog.set_level("WARNING", logger="posthog") hooks = [] holder = {} requests = [] - def slow_send(pipeline_client, payload): - requests.append(payload) - time.sleep(0.2) - return OK - with ( mock.patch("posthog.client.atexit.register", side_effect=hooks.append), mock.patch("posthog.client._ATEXIT_FLUSH_TIMEOUT_SECONDS", 0.1), @@ -614,7 +633,7 @@ def slow_send(pipeline_client, payload): traces={"max_export_batch_size": 1}, sync_mode=sync_mode ) client.start_span("a").end() - client._traces._exporter._send = slow_send + client._traces._exporter._send = slow_send(requests) client.start_span("b").end() client.start_span("c").end() assert len(hooks) == 2 diff --git a/posthog/test/tracing/test_export.py b/posthog/test/tracing/test_export.py index 9de73af8f..3d1c532bb 100644 --- a/posthog/test/tracing/test_export.py +++ b/posthog/test/tracing/test_export.py @@ -10,10 +10,7 @@ FakeSender, FakeTimer, RealTimer, - clock, - fake_timers, make_traces, - no_jitter, queued, ) from posthog.tracing import _export as export_module @@ -22,7 +19,7 @@ from posthog.tracing._span import NOOP_SPAN from posthog.tracing._transport import TOO_LARGE_LOCALLY, SendOutcome -__all__ = ["clock", "fake_timers", "no_jitter"] +pytestmark = pytest.mark.usefixtures("fake_timers", "no_jitter") class TestExport: diff --git a/posthog/test/tracing/test_pipeline.py b/posthog/test/tracing/test_pipeline.py index f03529c41..1026f3b96 100644 --- a/posthog/test/tracing/test_pipeline.py +++ b/posthog/test/tracing/test_pipeline.py @@ -15,8 +15,6 @@ TRACE_ID, FakeSender, FakeTimer, - clock, - fake_timers, make, make_traces, queued, @@ -28,7 +26,7 @@ from posthog.tracing._span import NOOP_SPAN, PassThroughSpan, RecordingSpan from posthog.tracing._transport import SendOutcome -__all__ = ["clock", "fake_timers"] +pytestmark = pytest.mark.usefixtures("fake_timers", "no_jitter") class TestStartSpan: diff --git a/posthog/tracing/_export.py b/posthog/tracing/_export.py index f7dbdacfc..e68385f78 100644 --- a/posthog/tracing/_export.py +++ b/posthog/tracing/_export.py @@ -244,6 +244,10 @@ def close(self) -> None: discarded, ) + def has_queued(self) -> bool: + with self._lock: + return bool(self._queue) + def warn_if_queued(self) -> None: with self._lock: queued = len(self._queue) diff --git a/posthog/tracing/_pipeline.py b/posthog/tracing/_pipeline.py index 896bd1a66..b09d25dff 100644 --- a/posthog/tracing/_pipeline.py +++ b/posthog/tracing/_pipeline.py @@ -49,6 +49,8 @@ def close(self) -> None: ... def warn_if_queued(self) -> None: ... + def has_queued(self) -> bool: ... + def reinit_after_fork(self) -> None: ... @@ -114,6 +116,9 @@ def warn_if_queued(self) -> None: self._exporter.warn_if_queued() self._drops.warn_if_due(force=True) + def has_queued_spans(self) -> bool: + return self._exporter.has_queued() + def reinit_after_fork(self, active_var: ContextVar) -> None: # Runs in the forked child before user code; the parent's spans stay # with the parent, and the child's active span is the fresh var.