feat(traces): per-span limits and exception stacktraces

posthog/test/tracing/test_config.py

@@ -1,6 +1,9 @@
 import pytest
 
 from posthog.tracing._config import (
+    DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH,
+    DEFAULT_MAX_ATTRIBUTES_PER_SPAN,
+    DEFAULT_MAX_EVENTS_PER_SPAN,
     DEFAULT_FLUSH_INTERVAL_SECONDS,
     DEFAULT_MAX_EXPORT_BATCH_SIZE,
     DEFAULT_MAX_LIVE_SPANS,
@@ -190,3 +193,39 @@ def __str__(self):
         assert resolved.service_name == "api"
         assert resolved.resource_attributes["team"] == "x"
         assert all(isinstance(key, str) for key in resolved.resource_attributes)
+
+
+class TestSpanLimitKnobs:
+    def test_defaults_to_opentelemetrys_counts_and_a_finite_value_length(self):
+        resolved = resolve_traces_config({})
+        assert (
+            resolved.max_attributes_per_span == DEFAULT_MAX_ATTRIBUTES_PER_SPAN == 128
+        )
+        assert resolved.max_events_per_span == DEFAULT_MAX_EVENTS_PER_SPAN == 128
+        assert resolved.max_attribute_value_length == DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH
+        assert DEFAULT_MAX_ATTRIBUTE_VALUE_LENGTH == 8192
+
+    def test_honours_explicit_values(self):
+        resolved = resolve_traces_config(
+            {
+                "max_attributes_per_span": 10,
+                "max_events_per_span": 5,
+                "max_attribute_value_length": 100,
+            }
+        )
+        assert resolved.max_attributes_per_span == 10
+        assert resolved.max_events_per_span == 5
+        assert resolved.max_attribute_value_length == 100
+
+    @pytest.mark.parametrize("value", [0, -1, 1.5, "128", None, True])
+    def test_an_unusable_value_falls_back_rather_than_dropping_every_span(self, value):
+        resolved = resolve_traces_config(
+            {
+                "max_attributes_per_span": value,
+                "max_events_per_span": value,
+                "max_attribute_value_length": value,
+            }
+        )
+        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

posthog/test/tracing/test_export.py

@@ -950,6 +950,23 @@ def test_a_forked_child_drops_the_inherited_queue_and_timer(self):
         assert [r.name for r in queued(pipeline)] == ["child-span"]
 
 
+class TestResourceAttributes:
+    def test_bounds_resource_attributes_on_every_batch(self):
+        sender = FakeSender(SendOutcome("ok"))
+        pipeline, _, _ = make_traces(
+            sender=sender,
+            max_attribute_value_length=5,
+            resource_attributes={"team": "platform-infrastructure"},
+        )
+        pipeline.start_span("a").end()
+        pipeline.flush()
+        resource = {
+            kv["key"]: kv["value"]
+            for kv in sender.payloads[0]["resourceSpans"][0]["resource"]["attributes"]
+        }
+        assert resource["team"] == {"stringValue": "platf"}
+
+
 def waits_advance(clock, pipeline):
     """Make the exporter's backoff wait move the fake clock instead of sleeping."""
     waited = []

posthog/test/tracing/test_limits.py

@@ -0,0 +1,183 @@
+from unittest import mock
+
+import pytest
+
+from posthog.tracing import _limits as limits_module
+from posthog.tracing._limits import (
+    bound_attributes,
+    truncate_attribute_value,
+    truncate_attributes,
+)
+from posthog.tracing._otlp import (
+    CIRCULAR_VALUE,
+    MAX_VALUE_ITEMS,
+    MAX_VALUE_NODES,
+    TRUNCATED_VALUE,
+    to_any_value,
+)
+from posthog.tracing._sanitize import FUNCTION_VALUE, UNSERIALIZABLE_VALUE
+
+
+class TestTruncateAttributeValue:
+    def test_truncates_a_long_string(self):
+        assert truncate_attribute_value("x" * 40000, 8192) == "x" * 8192
+
+    def test_never_walks_a_scalar(self):
+        with mock.patch.object(limits_module, "_truncate") as walk:
+            assert truncate_attribute_value("short", 8) == "short"
+            assert truncate_attribute_value(7, 8) == 7
+            assert truncate_attribute_value(None, 8) is None
+        assert not walk.called
+
+    def test_returns_a_short_string_unchanged(self):
+        assert truncate_attribute_value("short", 8192) == "short"
+
+    @pytest.mark.parametrize("value", [42, 1.5, True, None, 2**70])
+    def test_leaves_numbers_booleans_and_none_alone(self, value):
+        assert truncate_attribute_value(value, 3) == value
+
+    def test_reaches_strings_nested_in_mappings_and_lists(self):
+        value = {"body": "x" * 40000, "items": ["y" * 20, {"deep": "z" * 20}]}
+        assert truncate_attribute_value(value, 8) == {
+            "body": "x" * 8,
+            "items": ["y" * 8, {"deep": "z" * 8}],
+        }
+
+    def test_does_not_mutate_the_callers_value(self):
+        value = {"body": "x" * 20}
+        truncate_attribute_value(value, 4)
+        assert value == {"body": "x" * 20}
+
+    def test_a_self_referencing_value_terminates_with_the_encoders_marker(self):
+        value: dict = {"name": "n" * 20}
+        value["self"] = value
+        assert truncate_attribute_value(value, 4) == {
+            "name": "nnnn",
+            "self": CIRCULAR_VALUE,
+        }
+
+    def test_siblings_sharing_one_object_are_not_a_cycle(self):
+        shared = {"k": "v" * 10}
+        assert truncate_attribute_value([shared, shared], 2) == [
+            {"k": "vv"},
+            {"k": "vv"},
+        ]
+
+    def test_marks_items_past_the_encoders_item_cap(self):
+        bounded = truncate_attribute_value(["a"] * (MAX_VALUE_ITEMS + 5), 8)
+        assert len(bounded) == MAX_VALUE_ITEMS + 1
+        assert bounded[-1] == TRUNCATED_VALUE
+        # The encoder emits the same shape it would have for the original.
+        assert to_any_value(bounded) == to_any_value(["a"] * (MAX_VALUE_ITEMS + 5))
+
+    def test_stringifies_and_bounds_a_type_the_encoder_would_stringify(self):
+        class Big:
+            def __str__(self):
+                return "b" * 100
+
+        assert truncate_attribute_value(Big(), 10) == "b" * 10
+        assert truncate_attribute_value(b"\x00" * 100, 10) == "b'\\x00\\x00"
+
+    def test_a_key_the_encoder_skips_does_not_spend_the_walks_budget(self):
+        # The encoder drops "" without charging for its value, so the walk must
+        # too, or "x" would ship unbounded once the walk's budget ran out.
+        value = {"": list(range(999)), "a": [list(range(999))] * 9, "x": "A" * 1000}
+        bounded = truncate_attribute_value(value, 100)
+        assert "" not in bounded
+        assert bounded["x"] == "A" * 100
+        encoded = to_any_value(bounded)["kvlistValue"]["values"]
+        x = next(kv for kv in encoded if kv["key"] == "x")
+        assert len(x["value"]["stringValue"]) == 100
+
+    def test_leaves_a_callable_for_the_encoders_marker(self):
+        def handler():
+            pass
+
+        assert truncate_attribute_value(handler, 3) is handler
+        assert truncate_attribute_value({"fn": handler}, 3) == {"fn": handler}
+        assert to_any_value(truncate_attribute_value(handler, 3)) == {
+            "stringValue": FUNCTION_VALUE
+        }
+
+    def test_a_raising_str_costs_only_that_value(self):
+        class Hostile:
+            def __str__(self):
+                raise RuntimeError("no")
+
+        assert truncate_attribute_value({"a": Hostile(), "b": "ok"}, 8) == {
+            "a": UNSERIALIZABLE_VALUE,
+            "b": "ok",
+        }
+
+    def test_a_raising_accessor_costs_only_that_key(self):
+        class Explosive(dict):
+            def __getitem__(self, key):
+                if key == "bad":
+                    raise RuntimeError("no")
+                return super().__getitem__(key)
+
+        assert truncate_attribute_value(Explosive(good="g" * 9, bad=1), 3) == {
+            "good": "ggg",
+            "bad": UNSERIALIZABLE_VALUE,
+        }
+
+
+class TestBoundAttributes:
+    def test_keeps_the_earliest_entries_and_counts_the_rest(self):
+        source = {f"k{i}": i for i in range(130)}
+        attributes, dropped = bound_attributes(source, 128, 8192)
+        assert list(attributes) == [f"k{i}" for i in range(128)]
+        assert dropped == 2
+
+    @pytest.mark.parametrize(
+        "source",
+        [{"a": None, "b": 1, "c": 2}, {"b": 1, "c": 2, "a": None}],
+        ids=["before-the-cap", "past-the-cap"],
+    )
+    def test_a_none_value_spends_no_slot_and_counts_no_drop(self, source):
+        attributes, dropped = bound_attributes(source, 2, 8)
+        assert attributes == {"b": 1, "c": 2}
+        assert dropped == 0
+
+    def test_a_real_value_past_the_cap_counts_a_drop(self):
+        attributes, dropped = bound_attributes({"b": 1, "c": 2, "a": 3}, 2, 8)
+        assert attributes == {"b": 1, "c": 2}
+        assert dropped == 1
+
+    def test_bounds_each_value(self):
+        attributes, _ = bound_attributes({"a": "x" * 20}, 2, 5)
+        assert attributes == {"a": "xxxxx"}
+
+    def test_an_empty_key_spends_no_slot_and_counts_no_drop(self):
+        attributes, dropped = bound_attributes({"": 1, "b": 2, "c": 3}, 2, 8)
+        assert attributes == {"b": 2, "c": 3}
+        assert dropped == 0
+
+
+class TestTruncateAttributes:
+    def test_bounds_every_value_as_a_copy(self):
+        source = {"service.name": "api", "blob": "x" * 20}
+        assert truncate_attributes(source, 4) == {"service.name": "api", "blob": "xxxx"}
+        assert source["blob"] == "x" * 20
+
+
+class TestWalkBounds:
+    def test_walks_no_more_strings_than_the_encoder_would_emit(self):
+        # A thousand paths to one shared list of a thousand strings. Leaves are
+        # charged against the node budget, as in the encoder, so the walk does
+        # not copy every string on every path.
+        inner = ["x" * 50] * 1000
+        bounded = truncate_attribute_value([inner] * 1000, 8)
+        walked = sum(
+            1
+            for items in bounded
+            if items is not inner
+            for item in items
+            if item == "x" * 8
+        )
+        assert 0 < walked <= MAX_VALUE_NODES
+
+    def test_stops_walking_a_mapping_at_the_encoders_item_cap(self):
+        value = {f"k{i}": "v" * 50 for i in range(MAX_VALUE_ITEMS + 50)}
+        bounded = truncate_attribute_value(value, 4)
+        assert len(bounded) == MAX_VALUE_ITEMS

posthog/test/tracing/test_otlp.py

@@ -284,6 +284,39 @@ def test_marks_a_root_span_as_known_not_remote(self):
     def test_marks_a_header_parent_as_remote(self):
         assert build_otlp_span(record(parent_is_remote=True))["flags"] == 0x301
 
+    def test_omits_dropped_counts_when_nothing_was_dropped(self):
+        span = build_otlp_span(record(events=[SpanEventRecord("e", START_NS)]))
+        assert "droppedAttributesCount" not in span
+        assert "droppedEventsCount" not in span
+        assert "droppedAttributesCount" not in span["events"][0]
+
+    def test_emits_dropped_counts_on_the_span_and_its_events(self):
+        span = build_otlp_span(
+            record(
+                dropped_attributes_count=2,
+                dropped_events_count=3,
+                events=[SpanEventRecord("e", START_NS, {"k": 1}, 4)],
+            )
+        )
+        assert span["droppedAttributesCount"] == 2
+        assert span["droppedEventsCount"] == 3
+        assert span["events"][0]["droppedAttributesCount"] == 4
+
+    @pytest.mark.parametrize(
+        "value,expected",
+        [
+            (2**40, 0xFFFFFFFF),
+            (-1, 0),
+            (1.9, 1),
+            ("3", 0),
+            (True, 0),
+            (float("inf"), 0),
+        ],
+    )
+    def test_clamps_a_dropped_count_to_uint32(self, value, expected):
+        span = build_otlp_span(record(dropped_attributes_count=value))
+        assert span.get("droppedAttributesCount", 0) == expected
+
     def test_propagates_an_inbound_sampled_out_flag(self):
         assert build_otlp_span(record(trace_flags="00"))["flags"] == 0x100
 

posthog/test/tracing/test_pipeline.py

@@ -11,14 +11,17 @@
 from posthog.test.tracing.helpers import (
     SPAN_ID,
     TRACE_ID,
+    FakeSender,
     clock,
     fake_timers,
     make,
+    make_traces,
     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._transport import SendOutcome
 from posthog.tracing._span import NOOP_SPAN, PassThroughSpan, RecordingSpan
 
 __all__ = ["clock", "fake_timers"]
@@ -414,6 +417,17 @@ def test_lets_user_attributes_win_on_collision(self):
         pipeline.start_span("a", attributes={"posthogDistinctId": "override"}).end()
         assert queued(pipeline)[0].attributes["posthogDistinctId"] == "override"
 
+    def test_the_join_keys_survive_a_span_at_its_attribute_cap(self):
+        pipeline, _, _ = make(
+            context={"distinct_id": "user-1", "session_id": "sess-1"},
+            max_attributes_per_span=2,
+        )
+        pipeline.start_span("a", attributes={"x": 1, "y": 2, "z": 3}).end()
+        record = queued(pipeline)[0]
+        assert record.attributes["posthogDistinctId"] == "user-1"
+        assert record.attributes["sessionId"] == "sess-1"
+        assert record.dropped_attributes_count == 1
+
     def test_still_records_the_span_when_reading_context_raises(self):
         pipeline, _, _ = make()
         pipeline._get_context = mock.Mock(side_effect=RuntimeError("no context"))
@@ -586,3 +600,29 @@ def test_reinit_after_fork_replaces_locks_without_acquiring_them(self):
         pipeline.reinit_after_fork()
         assert not pipeline._lock.locked()
         pipeline.start_span("a").end()
+
+
+class TestLimitsReachTheExport:
+    def test_bounds_names_and_attributes_with_the_configured_length(self):
+        sender = FakeSender(SendOutcome("ok"))
+        pipeline, _, _ = make_traces(sender=sender, max_attribute_value_length=5)
+        pipeline.start_span("a long name", attributes={"k": "a long value"}).end()
+        pipeline.flush()
+        (span,) = sender.batches()[0]
+        assert span["name"] == "a lon"
+        assert span["attributes"] == [{"key": "k", "value": {"stringValue": "a lon"}}]
+
+    def test_reports_a_spans_limit_drops_once_at_debug(self, caplog):
+        caplog.set_level("DEBUG", logger="posthog")
+        pipeline, _, _ = make(max_attributes_per_span=1, max_events_per_span=1)
+        span = pipeline.start_span("capped", attributes={"a": 1, "b": 2})
+        span.add_event("e1", {"k": 1}).add_event("e2")
+        span.end()
+        messages = [
+            r.getMessage() for r in caplog.records if "Span limits" in r.getMessage()
+        ]
+        assert len(messages) == 1
+        assert messages[0].endswith(
+            'Span limits discarded data from "capped": 1 attributes, 1 events, '
+            "0 event attributes"
+        )