diff --git a/posthog/test/tracing/test_config.py b/posthog/test/tracing/test_config.py index ae3c5230..4cfd7041 100644 --- a/posthog/test/tracing/test_config.py +++ b/posthog/test/tracing/test_config.py @@ -221,3 +221,52 @@ def test_an_unusable_value_falls_back_rather_than_dropping_every_span(self, valu assert resolved.max_attributes_per_span == DEFAULT_MAX_ATTRIBUTES_PER_SPAN assert resolved.max_events_per_span == DEFAULT_MAX_EVENTS_PER_SPAN assert resolved.max_attribute_value_length == DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH + + +class TestBeforeSpanSendConfig: + def test_accepts_one_hook_or_a_list(self): + def hook(span): + return span + + assert resolve_traces_config({"before_span_send": hook}).before_span_send == ( + hook, + ) + assert resolve_traces_config( + {"before_span_send": [hook, hook]} + ).before_span_send == (hook, hook) + + def test_defaults_to_no_hooks(self): + assert resolve_traces_config({}).before_span_send == () + + def test_skips_falsy_entries_silently(self, caplog): + caplog.set_level("WARNING", logger="posthog") + + def hook(span): + return span + + resolved = resolve_traces_config({"before_span_send": [None, False, hook]}) + assert resolved.before_span_send == (hook,) + assert not caplog.records + + def test_ignores_and_warns_about_entries_that_are_not_callable(self, caplog): + caplog.set_level("WARNING", logger="posthog") + + def hook(span): + return span + + resolved = resolve_traces_config({"before_span_send": ["scrub", hook]}) + assert resolved.before_span_send == (hook,) + assert any("1 of 2" in r.getMessage() for r in caplog.records) + + def test_a_hook_whose_truthiness_raises_is_still_resolved(self): + class Hook: + def __bool__(self): + raise RuntimeError("no") + + def __call__(self, span): + return span + + hook = Hook() + assert resolve_traces_config({"before_span_send": hook}).before_span_send == ( + hook, + ) diff --git a/posthog/test/tracing/test_pipeline.py b/posthog/test/tracing/test_pipeline.py index 45dfc560..dda3ac2d 100644 --- a/posthog/test/tracing/test_pipeline.py +++ b/posthog/test/tracing/test_pipeline.py @@ -2,6 +2,7 @@ import logging import threading import time +import warnings import weakref from types import SimpleNamespace from unittest import mock @@ -12,6 +13,7 @@ SPAN_ID, TRACE_ID, FakeSender, + FakeTimer, clock, fake_timers, make, @@ -20,9 +22,10 @@ ) from posthog.tracing import _pipeline as pipeline_module from posthog.tracing import _span as span_module +from posthog.tracing._config import resolve_traces_config from posthog.tracing._drops import DropLog -from posthog.tracing._transport import SendOutcome from posthog.tracing._span import NOOP_SPAN, PassThroughSpan, RecordingSpan +from posthog.tracing._transport import SendOutcome __all__ = ["clock", "fake_timers"] @@ -581,3 +584,366 @@ def test_reports_a_spans_limit_drops_once_at_debug(self, caplog): 'Span limits discarded data from "capped": 1 attributes, 1 events, ' "0 event attributes" ) + + +class TestBeforeSpanSend: + def test_a_hook_returning_none_drops_the_span_quietly(self, caplog): + caplog.set_level("WARNING", logger="posthog") + pipeline, _, _ = make(before_span_send=lambda span: None) + pipeline.start_span("a").end() + assert queued(pipeline) == [] + assert any( + "Dropping 1 span(s): before_span_send dropped it" in r.getMessage() + for r in caplog.records + ) + + def test_a_raising_hook_drops_the_span_rather_than_exporting_it(self, caplog): + caplog.set_level("WARNING", logger="posthog") + + def broken(span): + raise RuntimeError("scrubber bug") + + pipeline, _, _ = make(before_span_send=broken) + pipeline.start_span("a", attributes={"password": "hunter2"}).end() + assert queued(pipeline) == [] + assert any("before_span_send failed" in r.getMessage() for r in caplog.records) + + def test_the_hook_sees_plain_values_not_the_wire_encoding(self): + seen = {} + + def hook(span): + seen.update(span) + return span + + pipeline, _, _ = make(before_span_send=hook, context={"distinct_id": "u1"}) + pipeline.start_span("checkout", kind="server", attributes={"userId": 42}).end() + assert seen["attributes"] == {"posthogDistinctId": "u1", "userId": 42} + assert seen["name"] == "checkout" + assert seen["kind"] == "server" + assert seen["status"] is None + assert len(seen["trace_id"]) == 32 and len(seen["span_id"]) == 16 + assert seen["parent_span_id"] is None + assert isinstance(seen["start_time_ns"], int) + assert seen["end_time_ns"] >= seen["start_time_ns"] + + def test_redacts_in_place(self): + def scrub(span): + span["attributes"].pop("http.request.header.authorization", None) + return span + + pipeline, _, _ = make(before_span_send=scrub) + pipeline.start_span( + "a", attributes={"http.request.header.authorization": "Bearer x", "ok": 1} + ).end() + assert queued(pipeline)[0].attributes == {"ok": 1} + + def test_identity_fields_are_readable_but_not_writable(self, caplog): + caplog.set_level("DEBUG", logger="posthog") + + def forge(span): + span["trace_id"] = "f" * 32 + span["span_id"] = "e" * 16 + return span + + pipeline, _, _ = make(before_span_send=forge) + span = pipeline.start_span("a", parent=f"00-{TRACE_ID}-{SPAN_ID}-00") + span_id = span._span_id + span.end() + record = queued(pipeline)[0] + assert (record.trace_id, record.span_id) == (TRACE_ID, span_id) + assert record.parent_span_id == SPAN_ID + assert any("identity field" in r.getMessage() for r in caplog.records) + + def test_propagation_state_survives_a_hook_that_rebuilds_the_dict(self): + def rebuild(span): + return {k: v for k, v in span.items() if k != "trace_id"} + + pipeline, _, _ = make(before_span_send=rebuild) + pipeline.start_span( + "a", parent=f"00-{TRACE_ID}-{SPAN_ID}-00", tracestate="vendor=abc" + ).end() + record = queued(pipeline)[0] + assert record.trace_id == TRACE_ID + assert record.trace_flags == "00" + assert record.parent_is_remote is True + assert record.trace_state == "vendor=abc" + + def test_runs_a_list_left_to_right_and_the_first_none_stops_it(self): + calls = [] + + def first(span): + calls.append("first") + span["attributes"]["step"] = 1 + return span + + def second(span): + calls.append("second") + return None + + def third(span): + calls.append("third") + return span + + pipeline, _, _ = make(before_span_send=[first, second, third]) + pipeline.start_span("a").end() + assert calls == ["first", "second"] + assert queued(pipeline) == [] + + def test_ignores_non_callable_entries_and_exports_unhooked(self, caplog): + caplog.set_level("WARNING", logger="posthog") + pipeline, _, _ = make(before_span_send=["not a hook", None]) + pipeline.start_span("a").end() + assert len(queued(pipeline)) == 1 + assert any("not callable" in r.getMessage() for r in caplog.records) + + def test_a_hook_can_remove_the_stacktrace(self): + def strip_stacks(span): + for event in span["events"]: + event["attributes"].pop("exception.stacktrace", None) + return span + + pipeline, _, _ = make(before_span_send=strip_stacks) + with pytest.raises(ValueError): + with pipeline.start_span("job"): + raise ValueError("boom") + (event,) = queued(pipeline)[0].events + assert event.attributes == { + "exception.type": "ValueError", + "exception.message": "boom", + } + + def test_reapplies_the_caps_to_what_the_hook_added(self): + def enrich(span): + for i in range(10): + span["attributes"][f"extra{i}"] = "x" * 50 + span["events"].extend({"name": f"e{i}"} for i in range(5)) + return span + + pipeline, _, _ = make( + before_span_send=enrich, + context={"distinct_id": "u1"}, + max_attributes_per_span=3, + max_events_per_span=2, + max_attribute_value_length=10, + ) + pipeline.start_span("a", attributes={"first": 1}).end() + record = queued(pipeline)[0] + assert list(record.attributes) == [ + "posthogDistinctId", + "first", + "extra0", + "extra1", + ] + assert record.attributes["extra0"] == "x" * 10 + assert record.dropped_attributes_count == 8 + assert [e.name for e in record.events] == ["e0", "e1"] + assert record.dropped_events_count == 3 + + def test_keeps_the_counts_the_span_already_dropped(self): + pipeline, _, _ = make( + before_span_send=lambda span: span, max_attributes_per_span=1 + ) + span = pipeline.start_span("a", attributes={"a": 1, "b": 2}) + span.add_event("e", {"k1": 1}) + span.end() + assert queued(pipeline)[0].dropped_attributes_count == 1 + + @pytest.mark.parametrize( + "result", + [ + "not a dict", + 42, + {"name": "a"}, + {"attributes": {}, "events": "no"}, + ], + ) + def test_drops_a_value_that_is_not_a_span_dict(self, result): + pipeline, _, _ = make(before_span_send=lambda span: result) + pipeline.start_span("a").end() + assert queued(pipeline) == [] + + def test_drops_the_span_from_an_async_hook_without_a_never_awaited_warning(self): + async def hook(span): + return span + + pipeline, _, _ = make(before_span_send=hook) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + pipeline.start_span("a").end() + gc.collect() + assert queued(pipeline) == [] + assert not [w for w in caught if "never awaited" in str(w.message)] + + def test_resanitizes_names_times_status_and_events(self, caplog): + caplog.set_level("DEBUG", logger="posthog") + + def mangle(span): + span["name"] = "" + span["start_time_ns"] = -5 + span["end_time_ns"] = "later" + span["status"] = {"code": "maybe"} + span["events"].append({"name": None, "timestamp_ns": 2**70}) + span["events"].append(None) + return span + + pipeline, _, _ = make(before_span_send=mangle) + span = pipeline.start_span("a", start_time=1_700_000_000) + span.set_status("error", "real failure") + span.end(end_time=1_700_000_005) + record = queued(pipeline)[0] + assert record.name == "unknown" + assert record.start_ns == 1_700_000_000 * 10**9 + assert record.end_ns == 1_700_000_005 * 10**9 + assert record.status.code == "error" + assert record.status.message == "real failure" + (event,) = record.events + assert event.name == "unknown" + assert event.timestamp_ns == record.start_ns + assert any( + "not an epoch-nanosecond int" in r.getMessage() for r in caplog.records + ) + + def test_clears_the_status_when_the_hook_sets_none(self): + def clear(span): + span["status"] = None + return span + + pipeline, _, _ = make(before_span_send=clear) + span = pipeline.start_span("a") + span.set_status("error", "x") + span.end() + assert queued(pipeline)[0].status is None + + def test_a_hook_that_ends_a_span_of_its_own_does_not_deadlock(self): + pipeline, _, _ = make() + seen = [] + + def nested(span): + if span["name"] == "outer": + pipeline.start_span("from-hook").end() + seen.append(span["name"]) + return span + + pipeline._config = resolve_traces_config({"before_span_send": nested}) + pipeline.start_span("outer").end() + assert sorted(seen) == ["from-hook", "outer"] + assert sorted(r.name for r in queued(pipeline)) == ["from-hook", "outer"] + + def test_is_not_called_for_a_span_whose_client_was_disabled(self): + client = SimpleNamespace(disabled=False, send=True) + hook = mock.Mock(side_effect=lambda span: span) + pipeline, _, _ = make(client=client, before_span_send=hook) + span = pipeline.start_span("a") + client.disabled = True + span.end() + assert not hook.called + + def test_a_span_whose_hook_outlives_close_is_not_queued(self): + # shutdown() closes tracing while another thread's hook is running: + # that span must not reach a queue nothing will flush. + pipeline, _, _ = make_traces() + + def close_mid_hook(span): + pipeline.close() + return span + + pipeline._config = resolve_traces_config({"before_span_send": close_mid_hook}) + timers_before = len(FakeTimer.instances) + pipeline.start_span("a").end() + assert queued(pipeline) == [] + assert len(FakeTimer.instances) == timers_before + + def test_the_limit_report_is_logged_with_no_tracing_lock_held(self): + pipeline, _, _ = make_traces(max_attributes_per_span=1) + held = [] + original = pipeline_module._report_limit_drops + + def spy(record): + held.append((pipeline._lock.locked(), pipeline._exporter._lock.locked())) + original(record) + + with mock.patch.object(pipeline_module, "_report_limit_drops", spy): + pipeline.start_span("a", attributes={"x": 1, "y": 2}).end() + assert held == [(False, False)] + + +class TestBeforeSpanSendBounds: + def test_bounds_a_status_message_the_hook_sets(self): + def long_message(span): + span["status"] = {"code": "error", "message": "m" * 100} + return span + + pipeline, _, _ = make( + before_span_send=long_message, max_attribute_value_length=5 + ) + pipeline.start_span("a").end() + assert queued(pipeline)[0].status.message == "mmmmm" + + def test_keeps_a_falsy_status_message_the_hook_sets(self): + def zero_message(span): + span["status"] = {"code": "error", "message": 0} + return span + + pipeline, _, _ = make(before_span_send=zero_message) + pipeline.start_span("a").end() + assert queued(pipeline)[0].status.message == "0" + + def test_a_status_message_whose_str_raises_keeps_the_span(self): + class Hostile: + def __str__(self): + raise RuntimeError("no") + + def hostile_message(span): + span["status"] = {"code": "error", "message": Hostile()} + return span + + pipeline, _, _ = make(before_span_send=hostile_message) + pipeline.start_span("a").end() + assert queued(pipeline)[0].status.message == "[Unserializable]" + + def test_a_hook_can_scrub_the_auto_context_keys(self): + def scrub(span): + span["attributes"].pop("posthogDistinctId") + span["attributes"].pop("sessionId") + return span + + pipeline, _, _ = make( + before_span_send=scrub, context={"distinct_id": "u", "session_id": "s"} + ) + pipeline.start_span("a").end() + assert queued(pipeline)[0].attributes == {} + + def test_recaps_event_attributes_the_hook_widens(self): + def widen(span): + span["events"][0]["attributes"].update({f"k{i}": i for i in range(200)}) + return span + + pipeline, _, _ = make(before_span_send=widen) + span = pipeline.start_span("a") + span.add_event("e", {"first": 1}) + span.end() + (event,) = queued(pipeline)[0].events + assert len(event.attributes) == 128 + assert event.dropped_attributes_count == 73 + + def test_keeps_an_events_dropped_count_when_the_hook_rebuilds_the_events(self): + def rebuild(span): + span["events"] = [dict(event) for event in span["events"]] + return span + + pipeline, _, _ = make(before_span_send=rebuild) + span = pipeline.start_span("a") + span.add_event("wide", {f"k{i}": i for i in range(130)}) + span.end() + assert queued(pipeline)[0].events[0].dropped_attributes_count == 2 + + def test_drops_a_dict_missing_a_required_key(self, caplog): + caplog.set_level("WARNING", logger="posthog") + + def incomplete(span): + return {"attributes": {}, "events": []} + + pipeline, _, _ = make(before_span_send=incomplete) + pipeline.start_span("a").end() + assert queued(pipeline) == [] + assert any("unusable record" in r.getMessage() for r in caplog.records) diff --git a/posthog/tracing/_before_span_send.py b/posthog/tracing/_before_span_send.py new file mode 100644 index 00000000..500cfdf0 --- /dev/null +++ b/posthog/tracing/_before_span_send.py @@ -0,0 +1,210 @@ +"""The ``before_span_send`` hook chain, where a finished span is scrubbed or dropped. + +The hook receives a plain dict, like the events ``before_send`` hook. It is the +scrubbing point, so a hook that raises drops the span rather than letting the +unscrubbed record through. +""" + +import inspect +import logging +from typing import Any, Dict, Mapping, Optional + +from ._config import MAX_ATTRIBUTES_PER_EVENT, ResolvedTracesConfig +from ._drops import DropLog +from ._limits import apply_span_limits +from ._otlp import SpanEventRecord, SpanRecord, SpanStatus, non_negative_count +from ._sanitize import ( + MAX_TIMESTAMP_NS, + MIN_TIMESTAMP_NS, + clamp_end_ns, + copy_user_attributes, + safe_str, + sanitize_name, +) + +log = logging.getLogger("posthog") + +_READ_ONLY_ID_KEYS = ("trace_id", "span_id", "parent_span_id") +_REQUIRED_KEYS = ("name", "kind", "start_time_ns", "end_time_ns") + + +def run_before_span_send( + record: SpanRecord, config: ResolvedTracesConfig, drops: DropLog +) -> Optional[SpanRecord]: + """Run the hook chain, returning the span to queue or ``None`` to drop it. + + The ids are read-only: they are restored after every hook, since a + rewritten id would orphan children already sent. + """ + if not config.before_span_send: + return record + + identity = (record.trace_id, record.span_id, record.parent_span_id) + # The span's own write order, so the caps keep the earliest-set entries. + keys_before_hook = list(record.attributes) + data = _hook_view(record) + try: + for hook in config.before_span_send: + result: Any = hook(data) + if result is None: + drops.record(1, "before_span_send dropped it") + return None + if not isinstance(result, Mapping): + if inspect.iscoroutine(result): + # An async hook is not awaited; closed so it does not warn. + result.close() + log.debug( + "before_span_send did not return a span dict; dropping the span" + ) + drops.record(1, "before_span_send returned an unusable record") + return None + data = _keep_identity(result, identity) + rebuilt = _rebuild(record, data, config) + if rebuilt is None: + log.debug("before_span_send did not return a span dict; dropping the span") + drops.record(1, "before_span_send returned an unusable record") + return None + apply_span_limits( + rebuilt, + record.auto_attribute_keys, + config.max_attributes_per_span, + config.max_events_per_span, + MAX_ATTRIBUTES_PER_EVENT, + config.max_attribute_value_length, + keys_before_hook, + ) + return rebuilt + except Exception: + log.debug( + "before_span_send failed; dropping the span rather than exporting it " + "unscrubbed", + exc_info=True, + ) + drops.record(1, "before_span_send failed") + return None + + +def _hook_view(record: SpanRecord) -> Dict[str, Any]: + return { + "trace_id": record.trace_id, + "span_id": record.span_id, + "parent_span_id": record.parent_span_id, + "name": record.name, + "kind": record.kind, + "status": ( + {"code": record.status.code, "message": record.status.message} + if record.status is not None + else None + ), + "attributes": dict(record.attributes), + "events": [ + { + "name": event.name, + "timestamp_ns": event.timestamp_ns, + "attributes": dict(event.attributes or {}), + "dropped_attributes_count": event.dropped_attributes_count, + } + for event in record.events + ], + "start_time_ns": record.start_ns, + "end_time_ns": record.end_ns, + } + + +def _keep_identity(result: Mapping, identity: tuple) -> Dict[str, Any]: + data = result if isinstance(result, dict) else dict(result) + if any(data.get(key) != value for key, value in zip(_READ_ONLY_ID_KEYS, identity)): + log.debug( + "before_span_send changed a span identity field; keeping the original ids" + ) + data.update(zip(_READ_ONLY_ID_KEYS, identity)) + return data + + +def _rebuild( + record: SpanRecord, data: Mapping, config: ResolvedTracesConfig +) -> Optional[SpanRecord]: + """A span record from what the chain returned, sanitized as ``end()`` would. + + ``None`` when it is not a span dict, rather than exporting a span of + fallbacks joinable to nothing. + """ + attributes = data.get("attributes") + events = data.get("events") + if ( + not isinstance(attributes, Mapping) + or not isinstance(events, (list, tuple)) + or any(key not in data for key in _REQUIRED_KEYS) + ): + return None + + max_length = config.max_attribute_value_length + start_ns = _valid_ns(data["start_time_ns"], record.start_ns) + end_ns = clamp_end_ns(_valid_ns(data["end_time_ns"], record.end_ns), start_ns) + kind = data["kind"] + + rebuilt_events = [] + for event in events: + try: + event_attributes = event.get("attributes") + rebuilt_events.append( + SpanEventRecord( + name=sanitize_name( + event.get("name"), "Span event name", max_length + ), + timestamp_ns=_valid_ns(event.get("timestamp_ns"), start_ns), + attributes=copy_user_attributes({}, event_attributes) + if event_attributes is not None + else None, + dropped_attributes_count=non_negative_count( + event.get("dropped_attributes_count") + ), + ) + ) + except Exception: + log.debug("before_span_send left an unreadable span event; dropping it") + + return SpanRecord( + trace_id=record.trace_id, + span_id=record.span_id, + parent_span_id=record.parent_span_id, + # SDK bookkeeping the hook is not shown, so it cannot erase it. + trace_state=record.trace_state, + trace_flags=record.trace_flags, + parent_is_remote=record.parent_is_remote, + dropped_attributes_count=record.dropped_attributes_count, + dropped_events_count=record.dropped_events_count, + auto_attribute_keys=record.auto_attribute_keys, + name=sanitize_name(data["name"], "Span name", max_length), + kind=kind if isinstance(kind, str) else record.kind, + status=_hook_status(data.get("status"), record.status), + attributes=copy_user_attributes({}, attributes), + events=rebuilt_events, + start_ns=start_ns, + end_ns=end_ns, + ) + + +def _hook_status(value: Any, original: Optional[SpanStatus]) -> Optional[SpanStatus]: + if value is None: + return None + if isinstance(value, Mapping) and value.get("code") in ("ok", "error"): + message = value.get("message") + text = None if message is None else safe_str(message) + return SpanStatus(value["code"], text or None) + # An unknown code would lose an error the span really had. + log.debug("before_span_send set an unknown span status; keeping the original") + return original + + +def _valid_ns(value: Any, fallback: int) -> int: + if ( + isinstance(value, int) + and not isinstance(value, bool) + and MIN_TIMESTAMP_NS <= value <= MAX_TIMESTAMP_NS + ): + return value + log.debug( + "before_span_send set a time that is not an epoch-nanosecond int; keeping the original" + ) + return fallback diff --git a/posthog/tracing/_config.py b/posthog/tracing/_config.py index e8abc9c2..11bdf9ef 100644 --- a/posthog/tracing/_config.py +++ b/posthog/tracing/_config.py @@ -6,7 +6,7 @@ import logging import math from dataclasses import dataclass, field -from typing import Any, Dict, Mapping, Optional +from typing import Any, Callable, Dict, Mapping, Optional, Tuple from ._sanitize import attribute_key @@ -47,6 +47,7 @@ class ResolvedTracesConfig: max_attributes_per_span: int = DEFAULT_MAX_ATTRIBUTES_PER_SPAN max_events_per_span: int = DEFAULT_MAX_EVENTS_PER_SPAN max_attribute_value_length: int = DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH + before_span_send: Tuple[Callable[[dict], Optional[dict]], ...] = () def _positive_number(config: Mapping, key: str, default: float) -> float: @@ -117,6 +118,29 @@ def _usable_resource_attributes(value: Any) -> Dict[str, Any]: return attributes +def _resolve_before_span_send(value: Any) -> Tuple[Callable, ...]: + """Keep only the callable hooks, in order. + + Anything else is dropped rather than called: a hook that raises drops every + span, which would leave tracing silently off. + """ + if value is None: + return () + supplied = list(value) if isinstance(value, (list, tuple)) else [value] + # `[enabled and scrub]` yields None or False: no hook, rather than a broken one. + supplied = [hook for hook in supplied if hook is not None and hook is not False] + hooks = tuple(hook for hook in supplied if callable(hook)) + if len(hooks) != len(supplied): + log.warning( + "Ignoring %s of %s traces before_span_send entries that are not callable. " + "Spans export without them, so whatever they were redacting is not " + "redacted.", + len(supplied) - len(hooks), + len(supplied), + ) + return hooks + + def resolve_traces_config( config: Any, host_resource_attributes: Optional[Mapping[str, str]] = None ) -> ResolvedTracesConfig: @@ -173,4 +197,5 @@ def resolve_traces_config( max_attribute_value_length=_positive_int( config, "max_attribute_value_length", DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH ), + before_span_send=_resolve_before_span_send(config.get("before_span_send")), ) diff --git a/posthog/tracing/_limits.py b/posthog/tracing/_limits.py index 35ad7d00..a7c35a79 100644 --- a/posthog/tracing/_limits.py +++ b/posthog/tracing/_limits.py @@ -7,7 +7,7 @@ """ from datetime import date -from typing import Any, Dict, Mapping, Tuple +from typing import AbstractSet, Any, Dict, List, Mapping, Sequence, Tuple from ._otlp import ( CIRCULAR_VALUE, @@ -15,6 +15,9 @@ MAX_VALUE_ITEMS, MAX_VALUE_NODES, TRUNCATED_VALUE, + SpanRecord, + SpanStatus, + non_negative_count, ) from ._sanitize import UNSERIALIZABLE_VALUE, attribute_key @@ -148,3 +151,69 @@ def truncate_attributes(attributes: Mapping, max_length: int) -> Dict[str, Any]: key: truncate_attribute_value(value, max_length) for key, value in attributes.items() } + + +def _ordered_keys(attributes: Mapping, keys_before_hook: Sequence[str]) -> List[Any]: + """The keys the span set first, in its order, so the earliest-set entries win.""" + keys = list(attributes.keys()) + present = set(keys) + before = [key for key in keys_before_hook if key in present] + seen = set(before) + return before + [key for key in keys if key not in seen] + + +def apply_span_limits( + record: SpanRecord, + auto_keys: AbstractSet[str], + max_attributes: int, + max_events: int, + max_attributes_per_event: int, + max_length: int, + keys_before_hook: Sequence[str] = (), +) -> None: + """Re-apply the per-span caps after a ``before_span_send`` hook, which + bypasses the span's own writer. Counts add to what the span already dropped.""" + attributes: Dict[str, Any] = {} + kept = 0 + dropped_attributes = 0 + for key in _ordered_keys(record.attributes, keys_before_hook): + value = record.attributes[key] + if value is None: + continue + if key not in auto_keys: + if kept >= max_attributes: + dropped_attributes += 1 + continue + kept += 1 + attributes[key] = truncate_attribute_value(value, max_length) + record.attributes = attributes + if dropped_attributes: + record.dropped_attributes_count = ( + non_negative_count(record.dropped_attributes_count) + dropped_attributes + ) + + kept_events: list = [] + dropped_events = 0 + for event in record.events: + if len(kept_events) >= max_events: + dropped_events += 1 + continue + if event.attributes: + event.attributes, dropped = bound_attributes( + event.attributes, max_attributes_per_event, max_length + ) + if dropped: + event.dropped_attributes_count = ( + non_negative_count(event.dropped_attributes_count) + dropped + ) + kept_events.append(event) + record.events = kept_events + if dropped_events: + record.dropped_events_count = ( + non_negative_count(record.dropped_events_count) + dropped_events + ) + + if record.status is not None and record.status.message: + record.status = SpanStatus( + record.status.code, truncate_string(record.status.message, max_length) + ) diff --git a/posthog/tracing/_otlp.py b/posthog/tracing/_otlp.py index 6fb15070..de8dbb74 100644 --- a/posthog/tracing/_otlp.py +++ b/posthog/tracing/_otlp.py @@ -9,7 +9,7 @@ import platform from dataclasses import dataclass, field from datetime import date, datetime -from typing import Any, Dict, List, Mapping, Optional +from typing import Any, Dict, FrozenSet, List, Mapping, Optional from ..version import VERSION from ._sanitize import FUNCTION_VALUE, UNSERIALIZABLE_VALUE, attribute_key, safe_str @@ -78,6 +78,9 @@ class SpanRecord: # User attributes and events the per-span caps refused. dropped_attributes_count: int = 0 dropped_events_count: int = 0 + # Keys the SDK attached itself, exempt from the attribute cap when it is + # re-applied after before_span_send. + auto_attribute_keys: FrozenSet[str] = frozenset() MAX_UINT32 = 0xFFFFFFFF diff --git a/posthog/tracing/_pipeline.py b/posthog/tracing/_pipeline.py index 10fd8c8b..22f6f77c 100644 --- a/posthog/tracing/_pipeline.py +++ b/posthog/tracing/_pipeline.py @@ -10,6 +10,7 @@ from contextvars import ContextVar from typing import Any, Callable, Dict, List, Mapping, Optional, Protocol +from ._before_span_send import run_before_span_send from ._config import ResolvedTracesConfig from ._drops import DropLog from ._ids import new_span_id, new_trace_id @@ -291,8 +292,11 @@ def _on_span_end(self, record: SpanRecord) -> None: if disabled: self._drops.record(1, "the client is disabled") else: - _report_limit_drops(record) - self._exporter.enqueue(record) + # The hook is application code, so it runs with no lock held. + hooked = run_before_span_send(record, self._config, self._drops) + if hooked is not None: + _report_limit_drops(hooked) + self._exporter.enqueue(hooked) self._drops.warn_if_due() diff --git a/posthog/tracing/_span.py b/posthog/tracing/_span.py index 07e17d1a..794a6cb2 100644 --- a/posthog/tracing/_span.py +++ b/posthog/tracing/_span.py @@ -412,6 +412,7 @@ def end(self, end_time: Optional[SpanTimeInput] = None) -> None: end_ns=clamp_end_ns(resolved, self._start_ns), dropped_attributes_count=self._dropped_attributes, dropped_events_count=self._dropped_events, + auto_attribute_keys=self._auto_keys, ) try: self._on_end(record)