diff --git a/posthog/test/tracing/helpers.py b/posthog/test/tracing/helpers.py new file mode 100644 index 00000000..e9786bac --- /dev/null +++ b/posthog/test/tracing/helpers.py @@ -0,0 +1,94 @@ +"""Shared fakes for the tracing pipeline tests.""" + +import threading +import time +from contextvars import ContextVar +from types import SimpleNamespace +from unittest import mock + +import pytest + +from posthog.tracing._config import resolve_traces_config +from posthog.tracing._drops import DropLog +from posthog.tracing._pipeline import PostHogTraces + +TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736" +SPAN_ID = "00f067aa0ba902b7" + + +class FakeTimer: + """Records the delay it was armed with; fires only when a test says so.""" + + instances: list = [] + + def __init__(self, delay, fn): + self.delay = delay + self.fn = fn + self.daemon = False + self.started = False + self.cancelled = False + FakeTimer.instances.append(self) + + def start(self): + self.started = True + + def cancel(self): + self.cancelled = True + + def fire(self): + self.fn() + + +class RecordingExporter: + """Stands in for the export queue: keeps every record it is handed.""" + + def __init__(self): + self.records: list = [] + self.closed = False + self.reinitialized = False + + def enqueue(self, record): + self.records.append(record) + + def flush(self, timeout=None): + pass + + def close(self): + self.closed = True + + def warn_if_queued(self): + pass + + 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 +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) + exporter = RecordingExporter() + active: ContextVar = ContextVar("active", default=None) + resolved = resolve_traces_config(config) + drops = DropLog(resolved.flush_interval) + pipeline = PostHogTraces( + client, resolved, lambda: context or {}, active, exporter, drops + ) + return pipeline, exporter, active + + +def queued(pipeline): + return pipeline._exporter.records diff --git a/posthog/test/tracing/test_config.py b/posthog/test/tracing/test_config.py new file mode 100644 index 00000000..a9648562 --- /dev/null +++ b/posthog/test/tracing/test_config.py @@ -0,0 +1,184 @@ +import pytest + +from posthog.tracing._config import ( + DEFAULT_FLUSH_INTERVAL_SECONDS, + DEFAULT_MAX_EXPORT_BATCH_SIZE, + DEFAULT_MAX_LIVE_SPANS, + DEFAULT_MAX_QUEUE_SIZE, + DEFAULT_MAX_SPAN_AGE_SECONDS, + ResolvedTracesConfig, + resolve_traces_config, +) + + +class TestDefaults: + def test_applies_the_documented_defaults(self): + assert resolve_traces_config({}) == ResolvedTracesConfig( + flush_interval=DEFAULT_FLUSH_INTERVAL_SECONDS, + max_export_batch_size=DEFAULT_MAX_EXPORT_BATCH_SIZE, + max_queue_size=DEFAULT_MAX_QUEUE_SIZE, + max_live_spans=DEFAULT_MAX_LIVE_SPANS, + max_span_age=DEFAULT_MAX_SPAN_AGE_SECONDS, + ) + + def test_leaves_service_name_unset_so_the_encoder_supplies_unknown_service(self): + assert resolve_traces_config({}).service_name is None + + @pytest.mark.parametrize("config", [None, "nope", 42, ["a"]]) + def test_a_non_dict_config_falls_back_to_defaults(self, config): + assert resolve_traces_config(config) == resolve_traces_config({}) + + +class TestExplicitValues: + def test_honours_explicit_values(self): + resolved = resolve_traces_config( + { + "service_name": "api", + "service_version": "1.2.3", + "environment": "prod", + "flush_interval": 2, + "max_export_batch_size": 100, + "max_queue_size": 400, + "max_live_spans": 50, + "max_span_age": 60, + } + ) + assert resolved == ResolvedTracesConfig( + service_name="api", + service_version="1.2.3", + environment="prod", + flush_interval=2.0, + max_export_batch_size=100, + max_queue_size=400, + max_live_spans=50, + max_span_age=60.0, + ) + + @pytest.mark.parametrize( + "value", + [0, -1, 0.5, 512.7, float("nan"), float("inf"), "512", True, None], + ) + def test_falls_back_for_an_unusable_batch_size(self, value): + assert ( + resolve_traces_config( + {"max_export_batch_size": value} + ).max_export_batch_size + == DEFAULT_MAX_EXPORT_BATCH_SIZE + ) + + def test_accepts_a_whole_number_float_batch_size(self): + assert ( + resolve_traces_config( + {"max_export_batch_size": 100.0} + ).max_export_batch_size + == 100 + ) + + def test_an_unusable_knob_keeps_the_rest_of_the_config(self): + resolved = resolve_traces_config( + {"service_name": "api", "max_live_spans": float("inf")} + ) + assert resolved.service_name == "api" + assert resolved.max_live_spans == DEFAULT_MAX_LIVE_SPANS + + @pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf"), "5", False]) + def test_falls_back_for_an_unusable_flush_interval(self, value): + assert ( + resolve_traces_config({"flush_interval": value}).flush_interval + == DEFAULT_FLUSH_INTERVAL_SECONDS + ) + + @pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf")]) + def test_falls_back_for_unusable_live_span_bounds(self, value): + resolved = resolve_traces_config( + {"max_live_spans": value, "max_span_age": value} + ) + assert resolved.max_live_spans == DEFAULT_MAX_LIVE_SPANS + assert resolved.max_span_age == DEFAULT_MAX_SPAN_AGE_SECONDS + + def test_keeps_the_queue_at_least_as_large_as_the_export_batch(self): + resolved = resolve_traces_config({"max_export_batch_size": 4096}) + assert resolved.max_queue_size == 4096 + + def test_floors_an_explicit_queue_size_at_the_batch_size(self): + resolved = resolve_traces_config( + {"max_export_batch_size": 10, "max_queue_size": 3} + ) + assert resolved.max_queue_size == 10 + + def test_ignores_a_non_string_named_field(self): + assert resolve_traces_config({"service_name": 42}).service_name is None + + +class TestResourceAttributes: + def test_lets_otlp_resource_attributes_override_the_named_fields(self): + resolved = resolve_traces_config( + { + "service_name": "named", + "resource_attributes": {"service.name": "from-attrs", "region": "eu"}, + } + ) + assert resolved.service_name == "from-attrs" + assert resolved.resource_attributes == { + "service.name": "from-attrs", + "region": "eu", + } + + def test_attaches_host_attributes_and_lets_user_attributes_override_them(self): + resolved = resolve_traces_config( + {"resource_attributes": {"os.name": "Custom"}}, + {"os.name": "Linux", "os.version": "6.1"}, + ) + assert resolved.resource_attributes == { + "os.name": "Custom", + "os.version": "6.1", + } + + def test_ignores_a_non_dict_value(self): + assert ( + resolve_traces_config({"resource_attributes": ["a"]}).resource_attributes + == {} + ) + + def test_drops_an_identity_key_that_is_not_a_string(self): + resolved = resolve_traces_config( + { + "service_name": "named", + "resource_attributes": { + "service.name": 42, + "deployment.environment": 1, + }, + } + ) + assert resolved.service_name == "named" + assert resolved.environment is None + assert "service.name" not in resolved.resource_attributes + + def test_keeps_the_readable_attributes_when_one_accessor_raises(self): + class Explosive(dict): + def __getitem__(self, key): + if key == "bad": + raise RuntimeError("boom") + return super().__getitem__(key) + + resolved = resolve_traces_config( + {"resource_attributes": Explosive(good=1, bad=2)} + ) + assert resolved.resource_attributes == {"good": 1} + + +class TestHostileResourceAttributeKeys: + def test_drops_only_a_key_that_cannot_be_stringified(self): + class HostileKey: + def __str__(self): + raise RuntimeError("no") + + resolved = resolve_traces_config( + { + "service_name": "api", + "resource_attributes": {HostileKey(): 1, "team": "x"}, + } + ) + assert resolved.service_name == "api" + assert resolved.resource_attributes["team"] == "x" + assert all(isinstance(key, str) for key in resolved.resource_attributes) diff --git a/posthog/test/tracing/test_pipeline.py b/posthog/test/tracing/test_pipeline.py new file mode 100644 index 00000000..eeb1180d --- /dev/null +++ b/posthog/test/tracing/test_pipeline.py @@ -0,0 +1,543 @@ +import gc +import logging +import threading +import time +import weakref +from types import SimpleNamespace +from unittest import mock + +import pytest + +from posthog.test.tracing.helpers import ( + SPAN_ID, + TRACE_ID, + clock, + fake_timers, + make, + queued, +) +from posthog.tracing import _pipeline as pipeline_module +from posthog.tracing import _span as span_module +from posthog.tracing._drops import DropLog +from posthog.tracing._span import NOOP_SPAN, PassThroughSpan, RecordingSpan + +__all__ = ["clock", "fake_timers"] + + +class TestStartSpan: + def test_enqueues_exactly_one_record_per_span(self): + pipeline, _, _ = make() + pipeline.start_span("a").end() + assert len(queued(pipeline)) == 1 + + def test_gives_a_root_span_a_fresh_trace_id_and_no_parent(self): + pipeline, _, _ = make() + pipeline.start_span("a").end() + pipeline.start_span("b").end() + a, b = queued(pipeline) + assert len(a.trace_id) == 32 and len(a.span_id) == 16 + assert a.parent_span_id is None + assert a.trace_id != b.trace_id + + def test_does_not_activate_the_span_it_returns(self): + pipeline, _, active = make() + span = pipeline.start_span("a") + assert active.get() is None + assert active.get() is None + span.end() + + def test_parents_a_child_to_an_explicit_span_handle(self): + pipeline, _, _ = make() + parent = pipeline.start_span("parent") + child = pipeline.start_span("child", parent=parent) + child.end() + parent.end() + child_record, parent_record = queued(pipeline) + assert child_record.trace_id == parent_record.trace_id + assert child_record.parent_span_id == parent_record.span_id + assert child_record.parent_is_remote is False + + def test_defaults_kind_to_internal_and_honours_an_explicit_kind(self): + pipeline, _, _ = make() + pipeline.start_span("a").end() + pipeline.start_span("b", kind="server").end() + assert [r.kind for r in queued(pipeline)] == ["internal", "server"] + + def test_returns_an_inert_handle_when_the_client_is_disabled(self): + pipeline, _, _ = make(client=SimpleNamespace(disabled=True, send=True)) + span = pipeline.start_span("a") + assert span is NOOP_SPAN + span.end() + assert queued(pipeline) == [] + + def test_makes_a_child_of_an_inert_handle_inert_rather_than_an_orphan(self): + pipeline, _, _ = make() + child = pipeline.start_span("child", parent=NOOP_SPAN) + assert child is NOOP_SPAN + child = pipeline.start_span( + "child", parent=PassThroughSpan(f"00-{TRACE_ID}-{SPAN_ID}-01", "v=1") + ) + # Inert like its parent, but still carrying the inbound context onward. + assert isinstance(child, PassThroughSpan) + assert child.traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-01" + assert child.tracestate() == "v=1" + child.end() + assert queued(pipeline) == [] + + def test_never_raises_out_of_start_span(self): + pipeline, _, _ = make() + with mock.patch.object( + pipeline, "_start_span", side_effect=RuntimeError("boom") + ): + span = pipeline.start_span("a") + assert span is NOOP_SPAN + + +class TestStartTime: + def test_backdates_the_span_to_a_supplied_start(self): + pipeline, _, _ = make() + pipeline.start_span("a", start_time=1_700_000_000).end() + assert queued(pipeline)[0].start_ns == 1_700_000_000 * 10**9 + + def test_falls_back_to_now_for_an_unusable_start(self): + pipeline, _, _ = make() + with mock.patch.object(time, "time_ns", return_value=42 * 10**9): + pipeline.start_span("a", start_time="yesterday").end() + assert queued(pipeline)[0].start_ns == 42 * 10**9 + + +class TestClockBasis: + """Children of a local parent share its root's clock, so they stay inside it.""" + + @pytest.fixture + def clocks(self): + state = {"wall": 0, "mono": 0} + with ( + mock.patch.object(time, "time_ns", lambda: state["wall"]), + mock.patch.object(span_module.time, "time_ns", lambda: state["wall"]), + mock.patch.object(span_module.time, "monotonic_ns", lambda: state["mono"]), + ): + yield lambda wall_ms, mono_ms: state.update( + wall=int(wall_ms * 1_000_000), mono=int(mono_ms * 1_000_000) + ) + + @staticmethod + def assert_within(inner, outer): + assert inner.start_ns >= outer.start_ns + assert inner.end_ns <= outer.end_ns + + def test_keeps_a_child_inside_its_parent_across_sub_millisecond_starts( + self, clocks + ): + pipeline, _, _ = make() + clocks(1_700_000_001_000, 0.9) + parent = pipeline.start_span("POST /checkout") + clocks(1_700_000_001_001, 1) + child = pipeline.start_span("http.post payments", parent=parent) + clocks(1_700_000_001_005, 5) + child.end() + clocks(1_700_000_001_005, 5.05) + parent.end() + child_record, parent_record = queued(pipeline) + self.assert_within(child_record, parent_record) + + def test_keeps_nested_active_spans_inside_their_parents_across_a_clock_step( + self, clocks + ): + pipeline, _, _ = make() + clocks(1_700_000_001_000, 100) + with pipeline.start_span("root"): + clocks(1_700_000_000_510, 110) + with pipeline.start_span("child"): + clocks(1_700_000_000_515, 115) + with pipeline.start_span("grandchild"): + clocks(1_700_000_000_520, 120) + clocks(1_700_000_000_525, 125) + clocks(1_700_000_000_530, 130) + grandchild, child, root = queued(pipeline) + self.assert_within(child, root) + self.assert_within(grandchild, child) + + def test_keeps_its_own_clock_under_a_remote_or_backdated_parent_or_backdated( + self, clocks + ): + pipeline, _, _ = make() + clocks(1_700_000_001_000, 100) + parent = pipeline.start_span("parent") + backdated_parent = pipeline.start_span( + "backdated parent", start_time=1_700_000_000 + ) + clocks(1_700_000_000_500, 110) + pipeline.start_span("remote child", parent=f"00-{TRACE_ID}-{SPAN_ID}-01").end() + pipeline.start_span("child of backdated", parent=backdated_parent).end() + pipeline.start_span( + "backdated child", parent=parent, start_time=1_700_000_000.2 + ).end() + assert [r.start_ns for r in queued(pipeline)] == [ + 1_700_000_000_500_000_000, + 1_700_000_000_500_000_000, + 1_700_000_000_200_000_000, + ] + + def test_keeps_an_explicit_start_time_that_equals_the_current_time(self, clocks): + pipeline, _, _ = make() + clocks(1_700_000_001_000, 100) + parent = pipeline.start_span("parent") + clocks(1_700_000_000_500, 110) + pipeline.start_span("child", parent=parent, start_time=1_700_000_000.5).end() + assert queued(pipeline)[0].start_ns == 1_700_000_000_500_000_000 + + +class TestTraceContinuation: + def test_continues_a_remote_trace_from_a_traceparent_string(self): + pipeline, _, _ = make() + pipeline.start_span("a", parent=f"00-{TRACE_ID}-{SPAN_ID}-01").end() + record = queued(pipeline)[0] + assert record.trace_id == TRACE_ID + assert record.parent_span_id == SPAN_ID + assert record.parent_is_remote is True + assert record.trace_flags == "01" + + def test_continues_a_trace_the_caller_sampled_out_and_propagates_the_flag(self): + pipeline, _, _ = make() + span = pipeline.start_span("a", parent=f"00-{TRACE_ID}-{SPAN_ID}-00") + assert span.traceparent().endswith("-00") + span.end() + assert queued(pipeline)[0].trace_flags == "00" + + def test_preserves_tracestate_opaquely_and_passes_it_to_children(self): + pipeline, _, _ = make() + parent = pipeline.start_span( + "a", parent=f"00-{TRACE_ID}-{SPAN_ID}-01", tracestate="vendor=abc" + ) + child = pipeline.start_span("b", parent=parent, tracestate="ignored=1") + assert child.tracestate() == "vendor=abc" + child.end() + parent.end() + assert [r.trace_state for r in queued(pipeline)] == ["vendor=abc", "vendor=abc"] + + def test_starts_a_fresh_root_on_a_malformed_traceparent_without_raising(self): + pipeline, _, _ = make() + pipeline.start_span("a", parent="garbage").end() + record = queued(pipeline)[0] + assert record.parent_span_id is None + assert record.trace_id != TRACE_ID + + def test_ignores_an_unusable_parent_and_falls_back_to_the_active_span(self): + pipeline, _, _ = make() + with pipeline.start_span("outer") as outer: + pipeline.start_span("inner", parent=["dup", "header"]).end() + assert queued(pipeline)[0].parent_span_id == outer._span_id + + def test_a_local_child_of_a_sampled_out_trace_keeps_the_00_flag(self): + pipeline, _, _ = make() + parent = pipeline.start_span("a", parent=f"00-{TRACE_ID}-{SPAN_ID}-00") + child = pipeline.start_span("b", parent=parent) + assert child.traceparent().endswith("-00") + child.end() + parent.end() + assert [r.trace_flags for r in queued(pipeline)] == ["00", "00"] + assert [r.parent_is_remote for r in queued(pipeline)] == [False, True] + + def test_two_header_values_with_nothing_active_start_a_fresh_root(self): + pipeline, _, _ = make() + headers = [f"00-{TRACE_ID}-{SPAN_ID}-01", f"00-{TRACE_ID}-{SPAN_ID}-00"] + pipeline.start_span("a", parent=headers).end() + record = queued(pipeline)[0] + assert record.parent_span_id is None and record.trace_id != TRACE_ID + + def test_a_parent_whose_traceparent_raises_gives_an_inert_span(self): + class Hostile(NOOP_SPAN.__class__): + def traceparent(self): + raise RuntimeError("no") + + pipeline, _, _ = make() + assert pipeline.start_span("a", parent=Hostile()) is NOOP_SPAN + + def test_continues_a_trace_from_a_one_element_header_list(self): + pipeline, _, _ = make() + pipeline.start_span("a", parent=[f"00-{TRACE_ID}-{SPAN_ID}-01"]).end() + assert queued(pipeline)[0].trace_id == TRACE_ID + assert queued(pipeline)[0].parent_span_id == SPAN_ID + + def test_an_active_pass_through_parents_a_recorded_span_as_remote(self): + # A pass-through is active when an earlier span in this trace could not + # be recorded; the inbound trace must survive it, not restart. + pipeline, _, active = make() + pass_through = PassThroughSpan( + f"00-{TRACE_ID}-{SPAN_ID}-00", "vendor=abc", active + ) + with pass_through: + pipeline.start_span("child").end() + record = queued(pipeline)[0] + assert record.trace_id == TRACE_ID + assert record.parent_span_id == SPAN_ID + assert record.parent_is_remote is True + assert record.trace_flags == "00" + assert record.trace_state == "vendor=abc" + + +class TestActiveSpan: + def test_nests_spans_started_inside_a_with_block(self): + pipeline, _, active = make() + with pipeline.start_span("parent") as parent: + assert active.get() is parent + pipeline.start_span("child").end() + assert active.get() is None + child, parent_record = queued(pipeline) + assert child.parent_span_id == parent_record.span_id + assert child.trace_id == parent_record.trace_id + + def test_lets_an_explicit_parent_override_the_active_span(self): + pipeline, _, _ = make() + other = pipeline.start_span("other") + with pipeline.start_span("active"): + pipeline.start_span("child", parent=other).end() + assert queued(pipeline)[0].parent_span_id == other._span_id + other.end() + + def test_records_a_raised_error_and_reraises_it_unmodified(self): + pipeline, _, _ = make() + error = ValueError("boom") + with pytest.raises(ValueError) as raised: + with pipeline.start_span("job"): + raise error + assert raised.value is error + record = queued(pipeline)[0] + assert record.status.code == "error" + assert record.events[0].attributes["exception.type"] == "ValueError" + + def test_isolates_concurrent_threads_from_each_other(self): + pipeline, _, active = make() + seen = {} + barrier = threading.Barrier(2, timeout=5) + + def work(name): + with pipeline.start_span(name) as span: + barrier.wait() + seen[name] = active.get() 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} + + +class TestPassThroughWhenTracingCannotRun: + def test_echoes_an_inbound_traceparent_flags_included_when_disabled(self): + pipeline, _, _ = make(client=SimpleNamespace(disabled=True, send=True)) + span = pipeline.start_span( + "a", parent=f"01-{TRACE_ID}-{SPAN_ID}-00", tracestate="v=1" + ) + assert span.traceparent() == f"01-{TRACE_ID}-{SPAN_ID}-00" + assert span.tracestate() == "v=1" + span.end() + assert queued(pipeline) == [] + + def test_activates_the_pass_through_handle_so_it_can_propagate(self): + pipeline, _, active = make(client=SimpleNamespace(disabled=True, send=True)) + with pipeline.start_span("a", parent=f"00-{TRACE_ID}-{SPAN_ID}-01") as span: + assert active.get() is span + assert active.get() is None + + def test_passes_the_inbound_context_through_when_the_live_span_limit_refuses(self): + pipeline, _, _ = make(max_live_spans=1) + held = pipeline.start_span("held") + span = pipeline.start_span("refused", parent=f"00-{TRACE_ID}-{SPAN_ID}-01") + assert span.traceparent() == f"00-{TRACE_ID}-{SPAN_ID}-01" + held.end() + + +class TestAutoContext: + def test_attaches_the_distinct_id_and_session_id_as_join_keys(self): + pipeline, _, _ = make(context={"distinct_id": "user-1", "session_id": "sess-1"}) + pipeline.start_span("a").end() + assert queued(pipeline)[0].attributes == { + "posthogDistinctId": "user-1", + "sessionId": "sess-1", + } + + def test_omits_keys_with_no_value(self): + pipeline, _, _ = make(context={"distinct_id": "", "session_id": None}) + pipeline.start_span("a").end() + assert queued(pipeline)[0].attributes == {} + + def test_freezes_the_snapshot_at_span_start(self): + context = {"distinct_id": "a"} + pipeline, _, _ = make(context=context) + span = pipeline.start_span("a") + context["distinct_id"] = "b" + span.end() + assert queued(pipeline)[0].attributes["posthogDistinctId"] == "a" + + def test_lets_user_attributes_win_on_collision(self): + pipeline, _, _ = make(context={"distinct_id": "a"}) + pipeline.start_span("a", attributes={"posthogDistinctId": "override"}).end() + assert queued(pipeline)[0].attributes["posthogDistinctId"] == "override" + + def test_still_records_the_span_when_reading_context_raises(self): + pipeline, _, _ = make() + pipeline._get_context = mock.Mock(side_effect=RuntimeError("no context")) + pipeline.start_span("a").end() + assert queued(pipeline)[0].attributes == {} + + +class TestGating: + def test_drops_a_span_whose_client_was_disabled_mid_trace_without_raising(self): + client = SimpleNamespace(disabled=False, send=True) + pipeline, _, _ = make(client=client) + span = pipeline.start_span("a") + client.disabled = True + span.end() + assert queued(pipeline) == [] + assert pipeline._live_spans == {} + + def test_a_close_that_lands_mid_start_makes_the_span_inert(self): + pipeline, _, _ = make() + resolve = pipeline._resolve_parent + + def close_then_resolve(parent, tracestate): + pipeline.close() + return resolve(parent, tracestate) + + with mock.patch.object(pipeline, "_resolve_parent", close_then_resolve): + span = pipeline.start_span("late") + assert span is NOOP_SPAN + assert pipeline._live_spans == {} + + +class TestLiveSpanBounds: + def test_returns_the_slot_when_building_the_span_fails(self): + pipeline, _, _ = make(max_live_spans=1) + with mock.patch.object( + pipeline_module, "copy_user_attributes", side_effect=RuntimeError("no") + ): + assert pipeline.start_span("a") is NOOP_SPAN + assert pipeline._live_spans == {} + assert isinstance(pipeline.start_span("b"), RecordingSpan) + + def test_returns_an_inert_handle_once_max_live_spans_are_live(self): + pipeline, _, _ = make(max_live_spans=2) + a, b = pipeline.start_span("a"), pipeline.start_span("b") + assert pipeline.start_span("c") is NOOP_SPAN + a.end() + b.end() + + def test_frees_the_slot_when_a_span_ends(self): + pipeline, _, _ = make(max_live_spans=1) + pipeline.start_span("a").end() + assert isinstance(pipeline.start_span("b"), RecordingSpan) + + def test_never_exports_a_span_evicted_for_exceeding_max_span_age(self, clock): + pipeline, _, _ = make(max_span_age=10) + leaked = pipeline.start_span("leaked") + clock["now"] += 11 + pipeline.start_span("fresh").end() + leaked.end() + assert [r.name for r in queued(pipeline)] == ["fresh"] + + def test_returns_the_slot_on_age_eviction_so_a_leak_cannot_disable_tracing( + self, clock + ): + pipeline, _, _ = make(max_live_spans=1, max_span_age=10) + pipeline.start_span("leaked") + assert pipeline.start_span("refused") is NOOP_SPAN + clock["now"] += 11 + assert isinstance(pipeline.start_span("recovered"), RecordingSpan) + + def test_ages_from_start_span_not_from_a_caller_supplied_start_time(self, clock): + pipeline, _, _ = make(max_span_age=10) + backdated = pipeline.start_span("old", start_time=1) + pipeline.start_span("probe").end() + backdated.end() + assert [r.name for r in queued(pipeline)] == ["probe", "old"] + + def test_holds_only_ids_and_floats_so_a_dropped_handle_is_collectable(self): + pipeline, _, _ = make() + span = pipeline.start_span("leaked") + ref = weakref.ref(span) + del span + gc.collect() + assert ref() is None + assert all( + isinstance(k, str) and isinstance(v, float) + for k, v in pipeline._live_spans.items() + ) + + +class TestDropLog: + def test_names_reasons_in_the_order_they_happened(self, caplog): + caplog.set_level("WARNING", logger="posthog") + drops = DropLog(5) + drops.record(1, "the queue is full") + drops.record(2, "before_span_send dropped it") + drops.record(1, "the queue is full") + drops.warn_if_due(force=True) + (record,) = caplog.records + assert record.getMessage().endswith( + "Dropping 4 span(s): the queue is full; before_span_send dropped it" + ) + + def test_a_raising_log_handler_does_not_escape(self): + class Raising(logging.Handler): + def emit(self, record): + raise RuntimeError("handler broke") + + handler = Raising() + logging.getLogger("posthog").addHandler(handler) + try: + drops = DropLog(5) + drops.record(1, "the queue is full") + drops.warn_if_due(force=True) + finally: + logging.getLogger("posthog").removeHandler(handler) + + def test_a_forked_child_does_not_wait_out_the_parents_warning_interval( + self, caplog + ): + caplog.set_level("WARNING", logger="posthog") + drops = DropLog(5) + drops.record(1, "the queue is full") + drops.warn_if_due() + drops.reinit_after_fork() + drops.record(1, "the queue is full") + drops.warn_if_due() + assert len(caplog.records) == 2 + + +class TestCloseAndFork: + def test_close_counts_spans_still_open_and_drops_them_when_they_end(self, caplog): + caplog.set_level("WARNING", logger="posthog") + pipeline, _, _ = make() + open_span = pipeline.start_span("open") + pipeline.close() + open_span.end() + assert queued(pipeline) == [] + assert any( + "Dropping 1 span(s): they were still open at shutdown" in r.getMessage() + for r in caplog.records + ) + + def test_close_makes_later_spans_inert_and_closes_the_exporter(self): + pipeline, exporter, _ = make() + pipeline.start_span("live") + pipeline.close() + assert exporter.closed + assert pipeline._live_spans == {} + assert pipeline.start_span("late") is NOOP_SPAN + + def test_a_forked_child_drops_the_parents_live_spans(self): + pipeline, exporter, _ = make() + pipeline.start_span("live") + pipeline.reinit_after_fork() + 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() + assert not pipeline._lock.locked() + pipeline.start_span("a").end() diff --git a/posthog/tracing/_config.py b/posthog/tracing/_config.py new file mode 100644 index 00000000..8ba8b755 --- /dev/null +++ b/posthog/tracing/_config.py @@ -0,0 +1,157 @@ +"""Resolution of the ``traces={...}`` client option. + +An unusable value falls back to its documented default with a warning. +""" + +import logging +import math +from dataclasses import dataclass, field +from typing import Any, Dict, Mapping, Optional + +from ._sanitize import attribute_key + +log = logging.getLogger("posthog") + +# OpenTelemetry's BatchSpanProcessor defaults; a full batch stays well under the +# ingestion service's default 2 MiB body limit. +DEFAULT_FLUSH_INTERVAL_SECONDS = 5.0 +DEFAULT_MAX_EXPORT_BATCH_SIZE = 512 +DEFAULT_MAX_QUEUE_SIZE = 2048 + +# Well above realistic concurrency; production traces routinely exceed ten minutes. +DEFAULT_MAX_LIVE_SPANS = 10_000 +DEFAULT_MAX_SPAN_AGE_SECONDS = 3600.0 + +# Resource attributes the server attributes spans by; they must be strings. +_RESOURCE_IDENTITY_KEYS = ("service.name", "service.version", "deployment.environment") + + +@dataclass(frozen=True) +class ResolvedTracesConfig: + service_name: Optional[str] = None + service_version: Optional[str] = None + environment: Optional[str] = None + resource_attributes: Dict[str, Any] = field(default_factory=dict) + flush_interval: float = DEFAULT_FLUSH_INTERVAL_SECONDS + max_export_batch_size: int = DEFAULT_MAX_EXPORT_BATCH_SIZE + max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE + max_live_spans: int = DEFAULT_MAX_LIVE_SPANS + max_span_age: float = DEFAULT_MAX_SPAN_AGE_SECONDS + + +def _positive_number(config: Mapping, key: str, default: float) -> float: + value = config.get(key, default) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not value > 0 + # An infinite interval cannot be armed as a timer. + or not math.isfinite(value) + ): + log.warning( + "Ignoring traces %s %r: expected a positive number of seconds", key, value + ) + return default + return float(value) + + +def _positive_int(config: Mapping, key: str, default: int) -> int: + value = config.get(key, default) + is_integer = isinstance(value, int) or ( + isinstance(value, float) and value.is_integer() + ) + if isinstance(value, bool) or not is_integer or not value >= 1: + log.warning("Ignoring traces %s %r: expected a positive integer", key, value) + return default + return int(value) + + +def _optional_string(config: Mapping, key: str) -> Optional[str]: + value = config.get(key) + if value is None or value == "": + return None + if not isinstance(value, str): + log.warning("Ignoring traces %s %r: expected a string", key, value) + return None + return value + + +def _usable_resource_attributes(value: Any) -> Dict[str, Any]: + if value is None: + return {} + if not isinstance(value, Mapping): + log.warning( + "Ignoring traces resource_attributes: expected a dict, got %s", + type(value).__name__, + ) + return {} + attributes: Dict[str, Any] = {} + try: + keys = list(value.keys()) + except Exception: + return {} + for key in keys: + try: + item = value[key] + except Exception: + continue + key_str = attribute_key(key) + if key_str is None: + continue + if key_str in _RESOURCE_IDENTITY_KEYS and not isinstance(item, str): + log.warning( + "Ignoring traces resource attribute %s: expected a string", key_str + ) + continue + attributes[key_str] = item + return attributes + + +def resolve_traces_config( + config: Any, host_resource_attributes: Optional[Mapping[str, str]] = None +) -> ResolvedTracesConfig: + """Validate the ``traces`` option, falling back to documented defaults per key. + + Host resource attributes (``os.name``, ``os.version``) merge first so user + ``resource_attributes`` win; a string ``service.name`` / ``service.version`` + / ``deployment.environment`` there wins over the named fields. + """ + if not isinstance(config, Mapping): + if config is not None: + log.warning( + "Ignoring traces config: expected a dict, got %s", type(config).__name__ + ) + config = {} + + resource_attributes: Dict[str, Any] = dict(host_resource_attributes or {}) + resource_attributes.update( + _usable_resource_attributes(config.get("resource_attributes")) + ) + + max_export_batch_size = _positive_int( + config, "max_export_batch_size", DEFAULT_MAX_EXPORT_BATCH_SIZE + ) + # The queue must hold at least one full batch, or the depth trigger never fires. + max_queue_size = max( + _positive_int(config, "max_queue_size", DEFAULT_MAX_QUEUE_SIZE), + max_export_batch_size, + ) + + return ResolvedTracesConfig( + service_name=resource_attributes.get("service.name") + or _optional_string(config, "service_name"), + service_version=resource_attributes.get("service.version") + or _optional_string(config, "service_version"), + environment=resource_attributes.get("deployment.environment") + or _optional_string(config, "environment"), + resource_attributes=resource_attributes, + flush_interval=_positive_number( + config, "flush_interval", DEFAULT_FLUSH_INTERVAL_SECONDS + ), + max_export_batch_size=max_export_batch_size, + max_queue_size=max_queue_size, + max_live_spans=_positive_int(config, "max_live_spans", DEFAULT_MAX_LIVE_SPANS), + max_span_age=_positive_number( + config, "max_span_age", DEFAULT_MAX_SPAN_AGE_SECONDS + ), + ) diff --git a/posthog/tracing/_drops.py b/posthog/tracing/_drops.py new file mode 100644 index 00000000..947055e1 --- /dev/null +++ b/posthog/tracing/_drops.py @@ -0,0 +1,54 @@ +"""Dropped-span accounting shared by span creation and export.""" + +import logging +import threading +import time +from typing import Dict + +log = logging.getLogger("posthog") + + +class DropLog: + """Counts dropped spans and warns at most once per interval, naming every reason. + + ``record()`` never logs: its callers hold their own locks, and a logging + handler is application code. ``warn_if_due()`` does, with no lock held. + """ + + def __init__(self, interval: float) -> None: + self._interval = interval + self._lock = threading.Lock() + self._count = 0 + # A dict for its order: reasons are named in the order they happened. + self._reasons: Dict[str, None] = {} + self._last_warning_at = 0.0 + + def record(self, count: int, reason: str) -> None: + with self._lock: + self._count += count + self._reasons[reason] = None + + def warn_if_due(self, force: bool = False) -> None: + with self._lock: + now = time.monotonic() + if not self._count or ( + not force and now - self._last_warning_at < self._interval + ): + return + message = "Dropping {} span(s): {}".format( + self._count, "; ".join(self._reasons) + ) + self._count = 0 + self._reasons.clear() + self._last_warning_at = now + try: + log.warning(message) + except Exception: + # A raising logging handler must not surface through span creation. + pass + + def reinit_after_fork(self) -> None: + self._lock = threading.Lock() + self._count = 0 + self._reasons.clear() + self._last_warning_at = 0.0 diff --git a/posthog/tracing/_pipeline.py b/posthog/tracing/_pipeline.py new file mode 100644 index 00000000..49225f55 --- /dev/null +++ b/posthog/tracing/_pipeline.py @@ -0,0 +1,290 @@ +"""Span creation, parenting and the end-of-span gates. + +Ended spans go to the exporter. Every public method is safe from any thread and +never raises into the host. +""" + +import logging +import threading +import time +from contextvars import ContextVar +from typing import Any, Callable, Dict, List, Mapping, Optional, Protocol + +from ._config import ResolvedTracesConfig +from ._drops import DropLog +from ._ids import new_span_id, new_trace_id +from ._otlp import SpanRecord +from ._sanitize import ( + copy_user_attributes, + resolve_start_ns, + sanitize_name, + to_epoch_ns, +) +from ._span import ( + ParentContext, + PassThroughSpan, + RecordingSpan, + Span, + inert_span, +) +from ._traceparent import parse_traceparent, sanitize_tracestate, traceparent_header + +log = logging.getLogger("posthog") + +_CONTEXT_ATTRIBUTE_KEYS = ( + ("distinct_id", "posthogDistinctId"), + ("session_id", "sessionId"), +) + +GetContextFn = Callable[[], Mapping[str, Any]] + + +class Exporter(Protocol): + def enqueue(self, record: SpanRecord) -> None: ... + + def flush(self, timeout: Optional[float] = None) -> None: ... + + def close(self) -> None: ... + + def warn_if_queued(self) -> None: ... + + def reinit_after_fork(self) -> None: ... + + +class PostHogTraces: + def __init__( + self, + client: Any, + config: ResolvedTracesConfig, + get_context: GetContextFn, + active_var: ContextVar, + exporter: Exporter, + drops: DropLog, + ) -> None: + self._client = client + self._config = config + self._get_context = get_context + self._active_var = active_var + self._exporter = exporter + self._drops = drops + self._lock = threading.Lock() + self._closed = False + # span id -> monotonic start. Never the span itself, so a dropped handle + # stays collectable. Insertion order is start order. + self._live_spans: Dict[str, float] = {} + + def start_span( + self, + name: str, + *, + kind: Optional[str] = None, + attributes: Optional[Mapping[str, Any]] = None, + parent: Any = None, + tracestate: Any = None, + start_time: Any = None, + ) -> Span: + """Start a span without making it active; always returns a handle.""" + try: + return self._start_span( + name, kind, attributes, parent, tracestate, start_time + ) + except Exception: + log.debug("start_span failed; returning an inert span", exc_info=True) + return inert_span(parent, tracestate, self._active_var) + finally: + self._drops.warn_if_due() + + def flush(self, timeout: Optional[float] = None) -> None: + self._exporter.flush(timeout) + + def close(self) -> None: + """Stop tracing: later spans are inert, and open ones are dropped when they end.""" + with self._lock: + self._closed = True + open_spans = len(self._live_spans) + self._live_spans.clear() + if open_spans: + self._drops.record(open_spans, "they were still open at shutdown") + self._exporter.close() + self._drops.warn_if_due(force=True) + + def warn_if_queued(self) -> None: + """Warn about spans still queued at exit, without discarding them.""" + self._exporter.warn_if_queued() + self._drops.warn_if_due(force=True) + + def reinit_after_fork(self) -> None: + # Runs in the forked child before user code; the parent's spans stay + # with the parent. + self._lock = threading.Lock() + self._live_spans.clear() + self._drops.reinit_after_fork() + self._exporter.reinit_after_fork() + + def _start_span( + self, + name: str, + kind: Optional[str], + attributes: Optional[Mapping[str, Any]], + parent: Any, + tracestate: Any, + start_time: Any, + ) -> Span: + if self._closed or getattr(self._client, "disabled", False): + return inert_span(parent, tracestate, self._active_var) + + parent = traceparent_header(parent) + if parent is not None and not isinstance(parent, (str, RecordingSpan)): + if isinstance(parent, Span): + # A child of an inert handle is inert too, still forwarding any + # inbound context. + return inert_span(parent, tracestate, self._active_var) + # Two header values, or another tracer's span. + log.debug("Ignoring an unusable span parent") + parent = None + + parent_context = self._resolve_parent(parent, tracestate) + + with self._lock: + if self._is_closed(): + # close() may have run since the check above. + return inert_span(parent, tracestate, self._active_var) + # Swept first, so a process that leaked its way to the bound + # recovers once the leaks age out. + aged = self._evict_aged_spans_locked() + at_limit = len(self._live_spans) >= self._config.max_live_spans + if not at_limit: + span_id = new_span_id() + # Aged from this call, not from a caller-supplied start_time. + self._live_spans[span_id] = time.monotonic() + if aged: + self._drops.record( + aged, + "they were still live after {:g}s".format(self._config.max_span_age), + ) + if at_limit: + self._drops.record( + 1, + "the live-span limit ({}) was reached; spans are being started " + "and never ended".format(self._config.max_live_spans), + ) + return inert_span(parent, tracestate, self._active_var) + + try: + return self._build_span( + span_id, name, kind, attributes, parent_context, start_time + ) + except Exception: + # The reservation would otherwise hold its slot until age eviction. + with self._lock: + self._live_spans.pop(span_id, None) + raise + + def _build_span( + self, + span_id: str, + name: str, + kind: Optional[str], + attributes: Optional[Mapping[str, Any]], + parent_context: Optional[ParentContext], + start_time: Any, + ) -> RecordingSpan: + now_ns = time.time_ns() + start_ns = resolve_start_ns(start_time, now_ns) + auto_attributes = self._auto_context_attributes() + span_attributes = copy_user_attributes(dict(auto_attributes), attributes) + + return RecordingSpan( + trace_id=parent_context.trace_id if parent_context else new_trace_id(), + span_id=span_id, + name=sanitize_name(name, "Span name"), + start_ns=start_ns, + backdated=start_ns != now_ns, + on_end=self._on_span_end, + kind=kind if isinstance(kind, str) and kind else "internal", + attributes=span_attributes, + parent_span_id=parent_context.parent_span_id if parent_context else None, + trace_state=parent_context.trace_state if parent_context else None, + trace_flags=parent_context.trace_flags if parent_context else "01", + parent_is_remote=parent_context.is_remote if parent_context else False, + active_var=self._active_var, + # The caller's own start_time wins over the parent's clock basis. + clock_anchor=parent_context.clock_anchor + if parent_context and to_epoch_ns(start_time) is None + else None, + ) + + def _is_closed(self) -> bool: + # A method, so the re-check under the lock is not narrowed away. + return self._closed + + def _resolve_parent(self, parent: Any, tracestate: Any) -> Optional[ParentContext]: + """An explicit parent, else the active span, else a fresh root.""" + if isinstance(parent, str): + remote = self._remote_context(parent, sanitize_tracestate(tracestate)) + if remote is None: + log.debug("Ignoring malformed traceparent; starting a new trace") + return remote + if isinstance(parent, RecordingSpan): + # The child inherits the parent's tracestate; the argument is ignored. + return parent._child_context() + active = self._active_var.get(None) + if isinstance(active, RecordingSpan): + return active._child_context() + if isinstance(active, PassThroughSpan): + # An earlier span in this trace was not recorded; the trace goes on. + return self._remote_context(active.traceparent(), active.tracestate()) + return None + + @staticmethod + def _remote_context( + traceparent: Any, tracestate: Optional[str] + ) -> Optional[ParentContext]: + remote = parse_traceparent(traceparent) + if remote is None: + return None + return ParentContext( + trace_id=remote.trace_id, + parent_span_id=remote.span_id, + trace_state=tracestate, + trace_flags=remote.flags, + is_remote=True, + ) + + def _auto_context_attributes(self) -> Dict[str, Any]: + try: + context = self._get_context() or {} + except Exception: + log.debug("Failed to read the request context for a span", exc_info=True) + return {} + attributes: Dict[str, Any] = {} + for source_key, wire_key in _CONTEXT_ATTRIBUTE_KEYS: + value = context.get(source_key) + if value: + attributes[wire_key] = value + return attributes + + def _evict_aged_spans_locked(self) -> int: + # An evicted span is never exported (its end() finds no entry), so a + # leak returns its slot rather than disabling tracing. + cutoff = time.monotonic() - self._config.max_span_age + aged: List[str] = [] + for span_id, started_at in self._live_spans.items(): + if started_at > cutoff: + break + aged.append(span_id) + for span_id in aged: + del self._live_spans[span_id] + return len(aged) + + def _on_span_end(self, record: SpanRecord) -> None: + with self._lock: + # A miss means it was evicted for age, or tracing was shut down. + if self._live_spans.pop(record.span_id, None) is None: + return + disabled = getattr(self._client, "disabled", False) + if disabled: + self._drops.record(1, "the client is disabled") + else: + self._exporter.enqueue(record) + self._drops.warn_if_due()