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/.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..092fb00b2 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 as 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,10 @@ 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 + # 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 59ccb9dca..857bd743e 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,12 @@ 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() + # 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) # 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 +2505,131 @@ 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: + # 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() + ) + 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 = False + 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. + 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. + + 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 +2637,15 @@ 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, 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 @@ -2483,20 +2656,53 @@ 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 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[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 or not traces.has_queued_spans(): + 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: + return lambda _seconds: flush_spans() + return flusher.join + def _is_consumer_thread(self) -> bool: current = threading.current_thread() return any(current in lane.consumers for lane in self._lanes) @@ -2687,6 +2893,19 @@ def _shutdown_once(self, errors: list[Exception]) -> None: self._run_lifecycle_cleanup( "Failed to reset metrics on shutdown", self._metrics.reset, errors ) + 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", 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", @@ -2772,8 +2991,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 +3008,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, waiter: Optional[Callable[[Optional[float]], None]], deadline: float + ) -> None: + 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() + @no_throw() def join(self) -> None: """ @@ -2813,7 +3052,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/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 new file mode 100644 index 000000000..ae7072b88 --- /dev/null +++ b/posthog/test/tracing/test_client_traces.py @@ -0,0 +1,805 @@ +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 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" + + +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 + + +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) +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_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_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() + 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, fake_timers + ): + client = make_client(traces={"max_export_batch_size": 1}) + client.start_span("x").end() + client.start_span("y").end() + requests = [] + client._traces._exporter._send = slow_send(requests) + 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, fake_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, fake_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_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() + 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, 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 = [] + 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, 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, fake_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, 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, fake_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, fake_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, fake_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, fake_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, fake_timers, caplog, sync_mode + ): + caplog.set_level("WARNING", logger="posthog") + hooks = [] + holder = {} + requests = [] + + 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(requests) + 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() + + 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: + 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() + + 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): + 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_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 + 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/test/tracing/test_export.py b/posthog/test/tracing/test_export.py index b4aebf2d1..3d1c532bb 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 @@ -9,10 +10,7 @@ FakeSender, FakeTimer, RealTimer, - clock, - fake_timers, make_traces, - no_jitter, queued, ) from posthog.tracing import _export as export_module @@ -21,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: @@ -940,7 +938,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..1026f3b96 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 @@ -14,8 +15,6 @@ TRACE_ID, FakeSender, FakeTimer, - clock, - fake_timers, make, make_traces, queued, @@ -27,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: @@ -593,14 +592,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/__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/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 3e8cd735d..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,10 +116,14 @@ 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 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. + # with the parent, and the child's active span is the fresh var. self._lock = threading.Lock() + self._active_var = active_var self._live_spans.clear() self._drops.reinit_after_fork() self._exporter.reinit_after_fork() 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()