diff --git a/posthog/test/tracing/test_span.py b/posthog/test/tracing/test_span.py new file mode 100644 index 00000000..1065916c --- /dev/null +++ b/posthog/test/tracing/test_span.py @@ -0,0 +1,646 @@ +import asyncio +import contextvars +import gc +import sys +import threading +import weakref +from contextvars import ContextVar +from datetime import datetime, timezone +from unittest import mock + +import pytest + +from posthog.tracing import _span as span_module +from posthog.tracing._otlp import SpanRecord +from posthog.tracing._sanitize import FALLBACK_SPAN_NAME, UNSERIALIZABLE_VALUE +from posthog.tracing._span import ( + NOOP_SPAN, + ClockAnchor, + NoopSpan, + PassThroughSpan, + RecordingSpan, + Span, + describe_error, + inert_span, +) + +TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736" +SPAN_ID = "00f067aa0ba902b7" +START_NS = 1_700_000_000_000_000_000 + + +def make_span(records=None, **overrides) -> RecordingSpan: + records = records if records is not None else [] + init = dict( + trace_id=TRACE_ID, + span_id=SPAN_ID, + name="checkout", + start_ns=START_NS, + backdated=True, + on_end=records.append, + ) + init.update(overrides) + return RecordingSpan(**init) + + +class TestRecording: + def test_produces_exactly_one_record_on_end(self): + records: list = [] + span = make_span(records) + span.end() + assert len(records) == 1 + assert isinstance(records[0], SpanRecord) + assert records[0].name == "checkout" + assert records[0].trace_id == TRACE_ID + assert records[0].span_id == SPAN_ID + + def test_is_idempotent_on_end(self): + records: list = [] + span = make_span(records) + span.end() + span.end() + assert len(records) == 1 + + def test_ignores_operations_after_end(self): + records: list = [] + span = make_span(records) + span.end() + span.set_attribute("k", "v").add_event("late").set_status("error").update_name( + "other" + ) + assert records[0].attributes == {} + assert records[0].events == [] + assert records[0].status is None + assert records[0].name == "checkout" + + def test_chains_mutators(self): + span = make_span() + assert ( + span.set_attribute("a", 1).set_attributes({"b": 2}).add_event("e") is span + ) + + def test_replaces_the_name_up_until_end(self): + records: list = [] + span = make_span(records, name="HTTP request") + span.update_name("GET /users/:id") + span.end() + assert records[0].name == "GET /users/:id" + + def test_replaces_an_empty_name_rather_than_dropping_the_span(self): + records: list = [] + span = make_span(records) + span.update_name("") + span.end() + assert records[0].name == FALLBACK_SPAN_NAME + + def test_applies_last_write_wins_to_status(self): + records: list = [] + span = make_span(records) + span.set_status("error", "boom").set_status("ok") + span.end() + assert records[0].status is not None + assert records[0].status.code == "ok" + assert records[0].status.message is None + + def test_omits_status_when_never_set(self): + records: list = [] + make_span(records).end() + assert records[0].status is None + + def test_ignores_a_status_code_whose_comparison_raises(self): + class Hostile: + def __eq__(self, other): + raise RuntimeError("no") + + __hash__ = object.__hash__ + + records: list = [] + span = make_span(records) + span.set_status(Hostile()) # type: ignore[arg-type] + span.end() + assert records[0].status is None + + def test_ignores_an_unrecognized_status_along_with_its_message(self): + records: list = [] + span = make_span(records) + span.set_status("weird", "nope") + span.end() + assert records[0].status is None + + def test_never_raises_when_the_pipeline_callback_fails(self): + def explode(record): + raise RuntimeError("queue broken") + + make_span(on_end=explode).end() + + +class TestAttributes: + def test_user_attributes_are_kept_on_the_record(self): + records: list = [] + span = make_span(records) + span.set_attribute("plan", "pro").set_attributes({"n": 1}) + span.end() + assert records[0].attributes == {"plan": "pro", "n": 1} + + def test_marks_only_the_raising_key_on_set_attributes(self): + class Explosive(dict): + def __getitem__(self, key): + if key == "bad": + raise RuntimeError("boom") + return super().__getitem__(key) + + records: list = [] + span = make_span(records) + span.set_attributes(Explosive(good=1, bad=2)) + span.end() + assert records[0].attributes == {"good": 1, "bad": UNSERIALIZABLE_VALUE} + + def test_snapshots_event_attributes_so_a_reused_mapping_cannot_mutate_them(self): + records: list = [] + span = make_span(records) + shared = {"step": 1} + span.add_event("a", shared) + shared["step"] = 2 + span.add_event("b", shared) + span.end() + assert records[0].events[0].attributes == {"step": 1} + assert records[0].events[1].attributes == {"step": 2} + + def test_stringifies_a_non_string_attribute_key(self): + records: list = [] + span = make_span(records) + span.set_attribute(7, "x") # type: ignore[arg-type] + span.end() + assert records[0].attributes == {"7": "x"} + + def test_ignores_a_key_that_cannot_be_stringified(self): + class HostileKey: + def __str__(self): + raise RuntimeError("no") + + records: list = [] + span = make_span(records) + span.set_attribute(HostileKey(), "x") # type: ignore[arg-type] + span.end() + assert records[0].attributes == {} + + +class TestRecordException: + def test_sets_error_status_and_attaches_an_event_without_ending(self): + records: list = [] + span = make_span(records) + span.record_exception(TypeError("boom")) + assert records == [] + span.end() + assert records[0].status is not None + assert records[0].status.code == "error" + assert records[0].status.message == "boom" + assert records[0].events[0].name == "exception" + assert records[0].events[0].attributes == { + "exception.type": "TypeError", + "exception.message": "boom", + } + + def test_overrides_an_explicit_ok(self): + records: list = [] + span = make_span(records) + span.set_status("ok").record_exception(ValueError("x")) + span.end() + assert records[0].status is not None + assert records[0].status.code == "error" + + +class TestMonotonicClock: + def test_measures_duration_against_the_monotonic_reading_not_the_wall_clock(self): + records: list = [] + with mock.patch.object( + span_module.time, "monotonic_ns", side_effect=[10, 10 + 80_000_000] + ): + span = make_span(records, backdated=False) + span.end() + assert records[0].end_ns - records[0].start_ns == 80_000_000 + + def test_never_reports_a_negative_duration_when_the_monotonic_source_goes_backwards( + self, + ): + records: list = [] + with mock.patch.object(span_module.time, "monotonic_ns", side_effect=[100, 50]): + span = make_span(records, backdated=False) + span.end() + assert records[0].end_ns == records[0].start_ns + + def test_places_an_event_inside_the_span_window(self): + records: list = [] + with mock.patch.object( + span_module.time, "monotonic_ns", side_effect=[0, 30_000_000, 80_000_000] + ): + span = make_span(records, backdated=False) + span.add_event("cache miss") + span.end() + event_ns = records[0].events[0].timestamp_ns + assert records[0].start_ns <= event_ns <= records[0].end_ns + assert event_ns == START_NS + 30_000_000 + + def test_uses_the_wall_clock_for_a_backdated_span(self): + records: list = [] + with mock.patch.object(span_module.time, "time_ns", return_value=START_NS + 5): + span = make_span(records, backdated=True) + span.end() + assert records[0].end_ns == START_NS + 5 + + +class TestTimestamps: + def test_records_an_end_at_or_after_the_start(self): + records: list = [] + make_span(records, backdated=False).end() + assert records[0].end_ns >= records[0].start_ns + + def test_honours_an_explicit_end_time(self): + records: list = [] + make_span(records).end(1_700_000_001) + assert records[0].end_ns == START_NS + 10**9 + + def test_accepts_a_datetime_as_an_end_time(self): + records: list = [] + make_span(records).end(datetime(2023, 11, 14, 22, 13, 21, tzinfo=timezone.utc)) + assert records[0].end_ns == START_NS + 10**9 + + def test_corrects_an_end_before_the_start_to_a_zero_duration(self): + records: list = [] + make_span(records).end(1_699_000_000) + assert records[0].end_ns == START_NS + + def test_falls_back_to_the_derived_end_for_an_out_of_range_end_time(self): + records: list = [] + with mock.patch.object(span_module.time, "time_ns", return_value=START_NS + 7): + make_span(records).end(-5) + assert records[0].end_ns == START_NS + 7 + + def test_honours_an_explicit_event_timestamp(self): + records: list = [] + span = make_span(records) + span.add_event("e", timestamp=1_700_000_000.5) + span.end() + assert records[0].events[0].timestamp_ns == START_NS + 500_000_000 + + +class TestContextPropagation: + def test_produces_a_sampled_traceparent(self): + assert make_span().traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-01" + + def test_returns_none_tracestate_when_it_has_none(self): + assert make_span().tracestate() is None + + def test_returns_the_tracestate_it_was_created_with(self): + assert make_span(trace_state="vendor=abc").tracestate() == "vendor=abc" + + def test_propagates_the_trace_flags_it_was_started_with(self): + assert ( + make_span(trace_flags="00").traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-00" + ) + + def test_hands_a_child_the_flags_it_propagates(self): + context = make_span(trace_flags="00", trace_state="v=1")._child_context() + assert (context.trace_id, context.parent_span_id) == (TRACE_ID, SPAN_ID) + assert (context.trace_state, context.trace_flags) == ("v=1", "00") + + +class TestClockAnchor: + def test_a_root_anchors_its_children_to_its_own_start(self): + with mock.patch.object(span_module.time, "monotonic_ns", return_value=500): + span = make_span(backdated=False) + assert span._child_context().clock_anchor == ClockAnchor(START_NS, 500) + + def test_a_backdated_span_hands_its_children_no_anchor(self): + assert make_span(backdated=True)._child_context().clock_anchor is None + + def test_a_child_starts_on_its_anchor_and_passes_the_same_anchor_on(self): + anchor = ClockAnchor(START_NS, 1_000) + with mock.patch.object(span_module.time, "monotonic_ns", return_value=1_250): + child = make_span(backdated=False, clock_anchor=anchor, start_ns=1) + assert child._start_ns == START_NS + 250 + assert child._child_context().clock_anchor == anchor + + def test_a_backdated_child_keeps_its_own_start(self): + child = make_span(backdated=True, clock_anchor=ClockAnchor(START_NS, 1_000)) + assert child._start_ns == START_NS + + def test_records_parent_and_remoteness(self): + records: list = [] + make_span( + records, parent_span_id="b7ad6b7169203331", parent_is_remote=True + ).end() + assert records[0].parent_span_id == "b7ad6b7169203331" + assert records[0].parent_is_remote is True + + +class TestContextManager: + def test_activates_for_the_block_and_ends_on_exit(self): + active: ContextVar = ContextVar("active", default=None) + records: list = [] + span = make_span(records, active_var=active) + assert active.get() is None + with span as entered: + assert entered is span + assert active.get() is span + assert active.get() is None + assert len(records) == 1 + + def test_records_a_raised_exception_and_reraises_it_unchanged(self): + records: list = [] + error = TypeError("boom") + with pytest.raises(TypeError) as raised: + with make_span(records): + raise error + assert raised.value is error + assert records[0].status is not None + assert records[0].status.code == "error" + assert records[0].status.message == "boom" + assert records[0].events[0].attributes == { + "exception.type": "TypeError", + "exception.message": "boom", + } + + @pytest.mark.parametrize( + "error", [GeneratorExit(), asyncio.CancelledError(), KeyboardInterrupt()] + ) + def test_ends_without_recording_a_base_exception_that_is_control_flow(self, error): + records: list = [] + with pytest.raises(type(error)): + with make_span(records): + raise error + assert len(records) == 1 + assert records[0].status is None + assert records[0].events == [] + + def test_a_closed_generator_holding_a_span_is_not_an_error(self): + records: list = [] + + def stream(): + with make_span(records): + yield 1 + yield 2 + + consumer = stream() + next(consumer) + consumer.close() + assert len(records) == 1 + assert records[0].status is None + + def test_treats_an_explicit_ok_status_as_final_when_the_block_raises(self): + records: list = [] + with pytest.raises(ValueError): + with make_span(records) as span: + span.set_status("ok") + raise ValueError("x") + assert records[0].status is not None + assert records[0].status.code == "ok" + assert records[0].events[0].name == "exception" + + def test_deactivates_even_when_the_block_raises(self): + active: ContextVar = ContextVar("active", default=None) + with pytest.raises(RuntimeError): + with make_span(active_var=active): + raise RuntimeError("x") + assert active.get() is None + + def test_is_no_longer_active_when_on_end_runs(self): + active: ContextVar = ContextVar("active", default=None) + seen: list = [] + span = make_span(active_var=active, on_end=lambda _: seen.append(active.get())) + with span: + pass + assert seen == [None] + + def test_exiting_in_a_context_that_never_entered_leaves_it_active(self, caplog): + active: ContextVar = ContextVar("active", default=None) + span = make_span(active_var=active) + span.__enter__() + with caplog.at_level("DEBUG", logger="posthog"): + contextvars.copy_context().run(span.__exit__, None, None, None) + assert active.get() is span + assert "never entered it" in caplog.text + span.__exit__(None, None, None) + assert active.get() is None + + def test_ending_inside_the_block_does_not_double_record(self): + records: list = [] + with make_span(records) as span: + span.end() + assert len(records) == 1 + + def test_nested_blocks_restore_the_outer_span(self): + active: ContextVar = ContextVar("active", default=None) + outer = make_span(active_var=active) + inner = make_span(active_var=active) + with outer: + with inner: + assert active.get() is inner + assert active.get() is outer + assert active.get() is None + + +class TestActivationAcrossContexts: + def test_one_span_entered_in_two_threads_leaves_neither_active(self): + active: ContextVar = ContextVar("active", default=None) + span = make_span(active_var=active) + entered, exited = ( + threading.Barrier(2, timeout=5), + threading.Barrier(2, timeout=5), + ) + seen = {} + + def worker(name, exit_first): + span.__enter__() + entered.wait() + if not exit_first: + exited.wait() + span._deactivate() + if exit_first: + exited.wait() + seen[name] = active.get() + + threads = [ + threading.Thread(target=worker, args=("a", True)), + threading.Thread(target=worker, args=("b", False)), + ] + for t in threads: + t.start() + for t in threads: + t.join(5) + assert seen == {"a": None, "b": None} + + def test_one_span_entered_in_two_tasks_leaves_neither_active(self): + active: ContextVar = ContextVar("active", default=None) + span = make_span(active_var=active) + + async def task(first_in, first_out): + await first_in.wait() + with span: + first_out.set() + await asyncio.sleep(0.01) + return active.get() + + async def main(): + go, a_in = asyncio.Event(), asyncio.Event() + go.set() + return await asyncio.gather(task(go, a_in), task(a_in, asyncio.Event())) + + assert asyncio.run(main()) == [None, None] + + def test_one_span_entered_concurrently_in_many_threads_stays_consistent(self): + active: ContextVar = ContextVar("active", default=None) + span = make_span(active_var=active) + errors: list = [] + seen = {} + + def worker(name): + try: + for _ in range(2000): + with span: + pass + except Exception as e: + errors.append(e) + seen[name] = active.get() + + interval = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + try: + threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(30) + finally: + sys.setswitchinterval(interval) + assert errors == [] + assert span._tokens == [] + assert seen == {i: None for i in range(8)} + + +class TestNoopSpan: + def test_supports_the_full_surface_without_raising(self): + span = NoopSpan() + assert span.set_attribute("k", "v") is span + assert span.set_attributes({"k": "v"}) is span + assert span.add_event("e", {"k": "v"}, 1) is span + assert span.set_status("error", "m") is span + assert span.record_exception(ValueError()) is span + assert span.update_name("n") is span + span.end() + span.end(1) + with span as entered: + assert entered is span + + def test_never_produces_a_traceparent(self): + assert NOOP_SPAN.traceparent() is None + assert NOOP_SPAN.tracestate() is None + + def test_is_never_activated(self): + active: ContextVar = ContextVar("active", default=None) + with inert_span(active_var=active) as span: + assert span is NOOP_SPAN + assert active.get() is None + + def test_all_handles_are_spans(self): + assert isinstance(NOOP_SPAN, Span) + assert isinstance(make_span(), Span) + assert isinstance( + PassThroughSpan("00-" + TRACE_ID + "-" + SPAN_ID + "-01"), Span + ) + + +class TestInertSpan: + def test_returns_the_shared_noop_without_a_usable_parent(self): + assert inert_span() is NOOP_SPAN + assert inert_span("garbage") is NOOP_SPAN + assert inert_span(NOOP_SPAN) is NOOP_SPAN + + def test_echoes_an_inbound_traceparent_flags_and_version_included(self): + span = inert_span(f"01-{TRACE_ID}-{SPAN_ID}-00", "vendor=abc") + assert isinstance(span, PassThroughSpan) + assert span.traceparent() == f"01-{TRACE_ID}-{SPAN_ID}-00" + assert span.tracestate() == "vendor=abc" + + def test_discards_an_invalid_tracestate_without_losing_the_traceparent(self): + span = inert_span(f"00-{TRACE_ID}-{SPAN_ID}-01", "novalue") + assert span.traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-01" + assert span.tracestate() is None + + def test_pass_through_is_activated_by_the_scoped_form(self): + active: ContextVar = ContextVar("active", default=None) + span = inert_span(f"00-{TRACE_ID}-{SPAN_ID}-01", active_var=active) + with span: + assert active.get() is span + assert active.get() is None + + def test_echoes_the_active_pass_through_when_no_parent_is_given(self): + # Tracing off, a span nested inside the one that received the inbound + # trace: it must keep forwarding that trace, not return a no-op. + active: ContextVar = ContextVar("active", default=None) + outer = inert_span(f"00-{TRACE_ID}-{SPAN_ID}-00", "vendor=abc", active) + with outer: + inner = inert_span(active_var=active) + assert isinstance(inner, PassThroughSpan) + assert inner.traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-00" + assert inner.tracestate() == "vendor=abc" + with inner: + assert active.get() is inner + assert active.get() is outer + + def test_echoes_an_explicit_handle_parent(self): + outer = inert_span(f"00-{TRACE_ID}-{SPAN_ID}-01", "vendor=abc") + child = inert_span(outer) + assert child.traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-01" + assert child.tracestate() == "vendor=abc" + + def test_an_explicit_parent_wins_over_the_active_span(self): + active: ContextVar = ContextVar("active", default=None) + with inert_span(f"00-{TRACE_ID}-{SPAN_ID}-01", active_var=active): + assert inert_span(NOOP_SPAN, active_var=active) is NOOP_SPAN + + def test_unwraps_a_one_element_header_list(self): + span = inert_span([f"00-{TRACE_ID}-{SPAN_ID}-01"]) + assert span.traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-01" + + def test_nothing_active_and_no_parent_is_a_noop(self): + active: ContextVar = ContextVar("active", default=None) + assert inert_span(active_var=active) is NOOP_SPAN + + def test_a_parent_whose_traceparent_raises_gives_a_noop(self): + class Hostile(Span): + def traceparent(self): + raise RuntimeError("no") + + assert inert_span(Hostile()) is NOOP_SPAN + + +class TestDescribeError: + @pytest.mark.parametrize( + "error,expected", + [ + (TypeError("boom"), ("TypeError", "boom")), + (ValueError(), ("ValueError", "")), + (KeyboardInterrupt(), ("KeyboardInterrupt", "")), + ("plain string", ("str", "plain string")), + (42, ("int", "42")), + ], + ) + def test_describes_values(self, error, expected): + assert describe_error(error) == expected + + def test_survives_a_value_whose_str_raises(self): + class Hostile(Exception): + def __str__(self): + raise RuntimeError("no") + + assert describe_error(Hostile()) == ("Hostile", "") + + +class TestHandleLifetime: + def test_a_dropped_handle_is_collectable(self): + span = make_span() + ref = weakref.ref(span) + del span + gc.collect() + assert ref() is None diff --git a/posthog/tracing/_span.py b/posthog/tracing/_span.py new file mode 100644 index 00000000..571a10fb --- /dev/null +++ b/posthog/tracing/_span.py @@ -0,0 +1,346 @@ +"""Span handles: the recording span, and the inert handles returned when tracing cannot run.""" + +import logging +import threading +import time +from contextvars import ContextVar, Token +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Mapping, Optional + +from ._otlp import SpanEventRecord, SpanRecord, SpanStatus +from ._sanitize import ( + SpanTimeInput, + attribute_key, + clamp_end_ns, + copy_user_attributes, + resolve_supplied_ns, + safe_str, + sanitize_name, +) +from ._traceparent import ( + TRACE_FLAGS_SAMPLED, + format_traceparent, + normalize_traceparent, + sanitize_tracestate, + traceparent_header, +) +from .span import Span + +log = logging.getLogger("posthog") + + +@dataclass(frozen=True) +class ClockAnchor: + """A root span's start on both clocks, shared by its local descendants.""" + + wall_ns: int + mono_ns: int + + +@dataclass(frozen=True) +class ParentContext: + """What a child span inherits from its parent.""" + + trace_id: str + parent_span_id: str + trace_state: Optional[str] + trace_flags: str + # True when the parent came from a traceparent header. + is_remote: bool = False + # Set only for a local, non-backdated parent. + clock_anchor: Optional[ClockAnchor] = None + + +class NoopSpan(Span): + """Returned when tracing cannot run and there is no inbound trace to forward. + + Never activated; a child started under it is also a no-op. + """ + + +NOOP_SPAN = NoopSpan() + + +class _Activatable: + """Mixin: entering the handle makes it the active span until exit.""" + + _active_var: Optional[ContextVar] + _tokens: List[Token] + _tokens_lock: threading.Lock + + def _activate(self) -> None: + if self._active_var is not None: + with self._tokens_lock: + self._tokens.append(self._active_var.set(self)) + + def _deactivate(self) -> None: + if self._active_var is None: + return + # The same handle can be entered in several threads or tasks at once, + # and a token only resets in the context that created it. + with self._tokens_lock: + for index in range(len(self._tokens) - 1, -1, -1): + try: + self._active_var.reset(self._tokens[index]) + except ValueError: + continue + del self._tokens[index] + return + log.debug( + "Span exited in a context that never entered it; it stays active " + "where it was entered" + ) + + +class PassThroughSpan(_Activatable, NoopSpan): + """An inert handle that echoes an inbound ``traceparent`` and ``tracestate``. + + Records nothing, so a service with tracing off still forwards the trace it + received. Entering it makes it the active span. + """ + + def __init__( + self, + traceparent: str, + tracestate: Optional[str] = None, + active_var: Optional[ContextVar] = None, + ) -> None: + self._traceparent = traceparent + self._tracestate = tracestate + self._active_var = active_var + self._tokens = [] + self._tokens_lock = threading.Lock() + + def traceparent(self) -> Optional[str]: + return self._traceparent + + def tracestate(self) -> Optional[str]: + return self._tracestate + + def __enter__(self) -> "Span": + self._activate() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self._deactivate() + + +def inert_span( + parent: Any = None, + tracestate: Any = None, + active_var: Optional[ContextVar] = None, +) -> Span: + """The handle to return when a span cannot be recorded. + + A pass-through when an inbound context is available, from ``parent`` or + else the active span, so the trace survives nesting; the no-op otherwise. + """ + try: + parent = traceparent_header(parent) + if parent is None and active_var is not None: + parent = active_var.get(None) + if isinstance(parent, Span): + inbound: Any = parent.traceparent() + tracestate = parent.tracestate() + else: + inbound = parent + traceparent = normalize_traceparent(inbound) + if not traceparent: + return NOOP_SPAN + return PassThroughSpan(traceparent, sanitize_tracestate(tracestate), active_var) + except Exception: + log.debug( + "Could not read the span parent; returning a no-op span", exc_info=True + ) + return NOOP_SPAN + + +def describe_error(error: Any) -> "tuple[str, str]": + """The OTel ``exception.type`` / ``exception.message`` pair for a raised value.""" + try: + return type(error).__name__, str(error) + except Exception: + return type(error).__name__, "" + + +class RecordingSpan(_Activatable, Span): + """A span that records and, on ``end()``, hands one record to the pipeline.""" + + def __init__( + self, + *, + trace_id: str, + span_id: str, + name: str, + start_ns: int, + backdated: bool, + on_end: Callable[[SpanRecord], None], + kind: str = "internal", + attributes: Optional[Dict[str, Any]] = None, + parent_span_id: Optional[str] = None, + trace_state: Optional[str] = None, + trace_flags: str = TRACE_FLAGS_SAMPLED, + parent_is_remote: bool = False, + active_var: Optional[ContextVar] = None, + clock_anchor: Optional[ClockAnchor] = None, + ) -> None: + self._trace_id = trace_id + self._span_id = span_id + self._parent_span_id = parent_span_id + self._trace_state = trace_state + self._trace_flags = trace_flags + self._parent_is_remote = parent_is_remote + self._start_mono: Optional[int] = None if backdated else time.monotonic_ns() + # A child of a local parent starts on its root's clock basis, so clock + # rounding or a wall-clock step cannot place it outside its parent. + if clock_anchor is not None and self._start_mono is not None: + start_ns = clock_anchor.wall_ns + (self._start_mono - clock_anchor.mono_ns) + self._start_ns = start_ns + self._clock_anchor: Optional[ClockAnchor] = None + if self._start_mono is not None: + self._clock_anchor = clock_anchor or ClockAnchor(start_ns, self._start_mono) + self._on_end = on_end + self._active_var = active_var + self._tokens = [] + self._tokens_lock = threading.Lock() + + self._name = name + self._kind = kind + self._attributes: Dict[str, Any] = attributes if attributes is not None else {} + self._events: List[SpanEventRecord] = [] + self._status: Optional[SpanStatus] = None + self._ended = False + + def _now_ns(self) -> int: + """Now, on this span's clock basis: start plus monotonic elapsed, else wall clock.""" + if self._start_mono is not None: + return self._start_ns + max(0, time.monotonic_ns() - self._start_mono) + return time.time_ns() + + def _mutable(self, operation: str) -> bool: + if self._ended: + log.debug("Ignoring %s on a span that has already ended", operation) + return False + return True + + def set_attribute(self, key: str, value: Any) -> "Span": + if self._mutable("set_attribute"): + key_str = attribute_key(key) + if key_str is not None: + self._attributes[key_str] = value + return self + + def set_attributes(self, attributes: Mapping[str, Any]) -> "Span": + if self._mutable("set_attributes"): + copy_user_attributes(self._attributes, attributes) + return self + + def add_event( + self, + name: str, + attributes: Optional[Mapping[str, Any]] = None, + timestamp: Optional[SpanTimeInput] = None, + ) -> "Span": + if self._mutable("add_event"): + self._events.append( + SpanEventRecord( + name=sanitize_name(name, "Span event name"), + timestamp_ns=resolve_supplied_ns( + timestamp, self._now_ns(), "event timestamp" + ), + attributes=copy_user_attributes({}, attributes) + if attributes is not None + else None, + ) + ) + return self + + def set_status(self, code: str, message: Optional[str] = None) -> "Span": + if self._mutable("set_status"): + if not isinstance(code, str) or code not in ("ok", "error"): + log.debug('Ignoring an unknown span status; expected "ok" or "error"') + return self + text = None if message is None else safe_str(message) + self._status = SpanStatus(code, text or None) + return self + + @property + def _status_is_explicitly_ok(self) -> bool: + return self._status is not None and self._status.code == "ok" + + def record_exception(self, exception: BaseException) -> "Span": + if self._mutable("record_exception"): + self._record_exception(exception, keep_ok=False) + return self + + def _record_exception(self, exception: BaseException, keep_ok: bool) -> None: + exc_type, message = describe_error(exception) + self.add_event( + "exception", {"exception.type": exc_type, "exception.message": message} + ) + # Only the scoped form treats an explicit `ok` as final. + if not (keep_ok and self._status_is_explicitly_ok): + self.set_status("error", message) + + def update_name(self, name: str) -> "Span": + if self._mutable("update_name"): + self._name = sanitize_name(name, "Span name") + return self + + def traceparent(self) -> Optional[str]: + return format_traceparent(self._trace_id, self._span_id, self._trace_flags) + + def tracestate(self) -> Optional[str]: + return self._trace_state + + def _child_context(self) -> ParentContext: + return ParentContext( + trace_id=self._trace_id, + parent_span_id=self._span_id, + trace_state=self._trace_state, + trace_flags=self._trace_flags, + clock_anchor=self._clock_anchor, + ) + + def end(self, end_time: Optional[SpanTimeInput] = None) -> None: + with self._tokens_lock: + if self._ended: + log.debug("Ignoring end() on a span that has already ended") + return + self._ended = True + + derived = self._now_ns() + resolved = resolve_supplied_ns(end_time, derived, "end time") + record = SpanRecord( + trace_id=self._trace_id, + span_id=self._span_id, + parent_span_id=self._parent_span_id, + trace_state=self._trace_state, + trace_flags=self._trace_flags, + parent_is_remote=self._parent_is_remote, + name=self._name, + kind=self._kind, + status=self._status, + attributes=dict(self._attributes), + events=self._events, + start_ns=self._start_ns, + end_ns=clamp_end_ns(resolved, self._start_ns), + ) + try: + self._on_end(record) + except Exception: + log.debug("Failed to enqueue span", exc_info=True) + + def __enter__(self) -> "Span": + self._activate() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + # Detached before it ends, so a span started from the on_end path + # does not become a child of one that is already over. + self._deactivate() + # GeneratorExit, CancelledError and the like are control flow, not + # failures of the span's work. + if isinstance(exc, Exception) and not self._ended: + self._record_exception(exc, keep_ok=True) + self.end() diff --git a/posthog/tracing/span.py b/posthog/tracing/span.py new file mode 100644 index 00000000..3cddf2fd --- /dev/null +++ b/posthog/tracing/span.py @@ -0,0 +1,143 @@ +"""The span handle type returned by ``start_span``.""" + +from datetime import datetime +from typing import Any, Mapping, Optional, Union + +__all__ = ["Span"] + + +class Span: + """A handle to a span. + + Every method is safe to call on any handle, including after ``end()`` and + on the inert handles returned when tracing is off, so calling code never + branches on whether tracing is running. Entering a handle (``with span:``) + makes it the active span for the block and ends it on exit, recording an + ``Exception`` raised inside it. A ``BaseException`` that is not an + ``Exception`` (``GeneratorExit``, ``CancelledError``, ``KeyboardInterrupt``) + still ends the span but is not recorded as a failure. + """ + + def set_attribute(self, key: str, value: Any) -> "Span": + """Set one attribute on the span; last write wins. Ignored after ``end()``. + + Returns the span, so calls chain. Prefer primitive values: strings, + booleans, integers and floats keep their type. Lists and mappings are + sent as OTLP arrays and maps, but PostHog stores them as serialized + strings. ``None`` removes the key; anything else is stringified. + + Examples: + ```python + span.set_attribute("http.status_code", 200).set_attribute("cache.hit", True) + ``` + """ + return self + + def set_attributes(self, attributes: Mapping[str, Any]) -> "Span": + """Set several attributes at once; last write wins per key. Ignored after ``end()``. + + Returns the span, so calls chain. Same encoding as ``set_attribute()``. + + Examples: + ```python + span.set_attributes({"db.system": "postgresql", "db.rows": 12}) + ``` + """ + return self + + def add_event( + self, + name: str, + attributes: Optional[Mapping[str, Any]] = None, + timestamp: Union[datetime, float, None] = None, + ) -> "Span": + """Add a timestamped event to the span. Ignored after ``end()``. + + ``timestamp`` defaults to now; it accepts a ``datetime``, or seconds + since the epoch as a ``float`` or ``int``. Returns the span, so calls + chain. + + Examples: + ```python + span.add_event("cache.miss", {"key": "user:42"}) + ``` + """ + return self + + def set_status(self, code: str, message: Optional[str] = None) -> "Span": + """Set the span status to ``"ok"`` or ``"error"``. Ignored after ``end()``. + + Unset by default. Any other ``code`` is ignored. ``ok`` is final for the + scoped form: an ``Exception`` raised inside ``with span:`` does not override + it. Returns the span, so calls chain. + + Examples: + ```python + span.set_status("error", "upstream timed out") + ``` + """ + return self + + def record_exception(self, exception: BaseException) -> "Span": + """Record an exception as an ``exception`` event and mark the span ``error``. + + Ignored after ``end()``. The event carries ``exception.type`` and + ``exception.message``. Returns the span, so calls chain. Inside + ``with span:`` a raised ``Exception`` is recorded automatically, so + this is for exceptions that are caught and handled. + + Examples: + ```python + try: + charge(card) + except PaymentError as e: + span.record_exception(e) + ``` + """ + return self + + def update_name(self, name: str) -> "Span": + """Rename the span, for a name only known after it started. Ignored after ``end()``. + + Returns the span, so calls chain. + + Examples: + ```python + span = posthog.start_span("http.request") + span.update_name(f"{request.method} {route.pattern}") + ``` + """ + return self + + def traceparent(self) -> Optional[str]: + """The W3C ``traceparent`` header value to propagate, or ``None``.""" + return None + + def tracestate(self) -> Optional[str]: + """The W3C ``tracestate`` header value to propagate, or ``None``.""" + return None + + def end(self, end_time: Union[datetime, float, None] = None) -> None: + """End the span and queue it for export. Idempotent. + + ``end_time`` defaults to now; it accepts a ``datetime``, or seconds + since the epoch as a ``float`` or ``int``, and is never earlier than + the start. Later calls, and later mutations, no-op. ``with span:`` + ends the span on exit, so call this only for spans started manually. + + Examples: + ```python + span = posthog.start_span("job") + try: + run_job() + finally: + span.end() + ``` + """ + return None + + def __enter__(self) -> "Span": + return self + + def __exit__(self, exc_type, exc, tb) -> None: + return None diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 0650028b..bc64b740 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -1045,6 +1045,7 @@ class posthog.request.DatetimeSerializer class posthog.request.GetResponse(data: Any, etag: Optional[str] = None, not_modified: bool = False) class posthog.request.HTTPAdapterWithSocketOptions(*args, socket_options: Optional[SocketOptions] = None, **kwargs) class posthog.request.QuotaLimitError +class posthog.tracing.span.Span class posthog.types.FeatureFlag(key: str, enabled: bool, variant: Optional[str], reason: Optional[FlagReason], metadata: Union[FlagMetadata, LegacyFlagMetadata]) class posthog.types.FeatureFlagError class posthog.types.FeatureFlagResult(key: str, enabled: bool, variant: Optional[str], payload: Optional[Any], reason: Optional[str]) @@ -1468,6 +1469,15 @@ method posthog.poller.Poller.run() method posthog.poller.Poller.stop() method posthog.request.DatetimeSerializer.default(obj: Any) method posthog.request.HTTPAdapterWithSocketOptions.init_poolmanager(*args, **kwargs) +method posthog.tracing.span.Span.add_event(name: str, attributes: Optional[Mapping[str, Any]] = None, timestamp: Union[datetime, float, None] = None) -> Span +method posthog.tracing.span.Span.end(end_time: Union[datetime, float, None] = None) -> None +method posthog.tracing.span.Span.record_exception(exception: BaseException) -> Span +method posthog.tracing.span.Span.set_attribute(key: str, value: Any) -> Span +method posthog.tracing.span.Span.set_attributes(attributes: Mapping[str, Any]) -> Span +method posthog.tracing.span.Span.set_status(code: str, message: Optional[str] = None) -> Span +method posthog.tracing.span.Span.traceparent() -> Optional[str] +method posthog.tracing.span.Span.tracestate() -> Optional[str] +method posthog.tracing.span.Span.update_name(name: str) -> Span method posthog.types.FeatureFlag.from_json(resp: Any) -> FeatureFlag method posthog.types.FeatureFlag.from_value_and_payload(key: str, value: FlagValue, payload: Any) -> FeatureFlag method posthog.types.FeatureFlag.get_value() -> FlagValue @@ -1559,6 +1569,7 @@ module posthog.metrics_capture module posthog.poller module posthog.request module posthog.tracing +module posthog.tracing.span module posthog.types module posthog.utils module posthog.version