fix(traces): guard span writes and end() with the handle's lock

A same-key write from two threads could reserve two slots, and a write
that lost the race to end() could land after the record was taken. The
cap check, the event count and the end() snapshot now run under the
handle's lock; the value walk still runs outside it.

Author: turnipdabeets

posthog/test/tracing/test_span.py

@@ -3,6 +3,7 @@
 import gc
 import sys
 import threading
+import time
 import weakref
 from contextvars import ContextVar
 from datetime import datetime, timezone
@@ -823,6 +824,68 @@ def test_add_event_with_a_non_mapping_logs_and_keeps_the_event(self, caplog):
         assert "expected a mapping" in caplog.text
 
 
+class TestConcurrentWrites:
+    def _run(self, worker, count=8):
+        threads = [threading.Thread(target=worker, args=(i,)) for i in range(count)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join(30)
+
+    def test_writes_to_one_key_from_many_threads_spend_one_slot(self):
+        records: list = []
+        span = make_span(records, max_attributes=1)
+
+        def worker(i):
+            for n in range(300):
+                span.set_attribute("k", n)
+
+        self._run(worker)
+        span.end()
+        assert list(records[0].attributes) == ["k"]
+        assert records[0].dropped_attributes_count == 0
+
+    def test_events_from_many_threads_never_exceed_the_cap(self):
+        records: list = []
+        span = make_span(records, max_events=100)
+
+        def worker(i):
+            for n in range(50):
+                span.add_event(f"{i}-{n}")
+
+        self._run(worker)
+        span.end()
+        assert len(records[0].events) == 100
+        assert records[0].dropped_events_count == 300
+
+    def test_a_write_racing_end_never_lands_after_the_record(self):
+        records: list = []
+        span = make_span(records)
+        stop = threading.Event()
+
+        def writer(i):
+            n = 0
+            while not stop.is_set():
+                span.set_attribute("k", n)
+                span.add_event("e")
+                n += 1
+
+        threads = [threading.Thread(target=writer, args=(i,)) for i in range(4)]
+        for thread in threads:
+            thread.start()
+        time.sleep(0.02)
+        span.end()
+        stop.set()
+        for thread in threads:
+            thread.join(30)
+        record = records[0]
+        assert len(records) == 1
+        assert span._events == record.events
+        assert {k: v for k, v in span._attributes.items() if v is not None} == (
+            record.attributes
+        )
+
+
 class TestExceptionStacktrace:
     def test_record_exception_attaches_the_stack_of_a_raised_exception(self):
         records: list = []

posthog/tracing/_otlp.py

@@ -290,7 +290,9 @@ def _to_otlp_event(event: SpanEventRecord) -> dict:
         attributes, cut = encode_attributes(event.attributes)
         if attributes:
             encoded["attributes"] = attributes
-    dropped = non_negative_count(non_negative_count(event.dropped_attributes_count) + cut)
+    dropped = non_negative_count(
+        non_negative_count(event.dropped_attributes_count) + cut
+    )
     if dropped:
         encoded["droppedAttributesCount"] = dropped
     return encoded

posthog/tracing/_span.py

@@ -74,19 +74,20 @@ class _Activatable:
 
     _active_var: Optional[ContextVar]
     _tokens: List[Token]
-    _tokens_lock: threading.Lock
+    # Guards the tokens and, on a recording span, every write and end().
+    _lock: threading.Lock
 
     def _activate(self) -> None:
         if self._active_var is not None:
-            with self._tokens_lock:
+            with self._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:
+        with self._lock:
             for index in range(len(self._tokens) - 1, -1, -1):
                 try:
                     self._active_var.reset(self._tokens[index])
@@ -117,7 +118,7 @@ def __init__(
         self._tracestate = tracestate
         self._active_var = active_var
         self._tokens = []
-        self._tokens_lock = threading.Lock()
+        self._lock = threading.Lock()
 
     def traceparent(self) -> Optional[str]:
         return self._traceparent
@@ -244,7 +245,7 @@ def __init__(
         self._on_end = on_end
         self._active_var = active_var
         self._tokens = []
-        self._tokens_lock = threading.Lock()
+        self._lock = threading.Lock()
 
         self._name = name
         self._kind = kind
@@ -258,11 +259,11 @@ def __init__(
         self._dropped_attributes = 0
         self._dropped_events = 0
         self._attributes: Dict[str, Any] = {}
-        for key, value in (attributes or {}).items():
-            self._write_attribute(key, value)
         self._events: List[SpanEventRecord] = []
         self._status: Optional[SpanStatus] = None
         self._ended = False
+        for key, value in (attributes or {}).items():
+            self._write_attribute(key, value)
 
     def _now_ns(self) -> int:
         """Now, on this span's clock basis: start plus monotonic elapsed, else wall clock."""
@@ -282,21 +283,27 @@ def _write_attribute(self, key: str, value: Any) -> None:
             # The encoder drops it, so it must not spend a slot.
             log.debug("Dropping an attribute with an empty key")
             return
-        if value is None:
-            # None removes the key, freeing its slot.
-            if key in self._attributes and key not in self._auto_keys:
-                self._user_attribute_count -= 1
-            self._attributes.pop(key, None)
-            return
-        # Checked before the value is walked, which is the costly part.
-        if key not in self._auto_keys and key not in self._attributes:
-            if self._user_attribute_count >= self._max_attributes:
-                self._dropped_attributes += 1
+        with self._lock:
+            if self._ended:
                 return
-            self._user_attribute_count += 1
-        self._attributes[key] = truncate_attribute_value(
-            value, self._max_attribute_value_length
-        )
+            if value is None:
+                # None removes the key, freeing its slot.
+                if key in self._attributes and key not in self._auto_keys:
+                    self._user_attribute_count -= 1
+                self._attributes.pop(key, None)
+                return
+            # Checked before the value is walked, which is the costly part.
+            if key not in self._auto_keys and key not in self._attributes:
+                if self._user_attribute_count >= self._max_attributes:
+                    self._dropped_attributes += 1
+                    return
+                self._user_attribute_count += 1
+                # Reserved, so a concurrent write to the same key sees it taken.
+                self._attributes[key] = None
+        bounded = truncate_attribute_value(value, self._max_attribute_value_length)
+        with self._lock:
+            if not self._ended:
+                self._attributes[key] = bounded
 
     def set_attribute(self, key: str, value: Any) -> "Span":
         if self._mutable("set_attribute"):
@@ -318,11 +325,14 @@ def add_event(
         timestamp: Optional[SpanTimeInput] = None,
     ) -> "Span":
         if self._mutable("add_event"):
-            # A recorded exception spends a slot like any other event.
-            if self._user_event_count >= self._max_events:
-                self._dropped_events += 1
-                return self
-            self._user_event_count += 1
+            with self._lock:
+                if self._ended:
+                    return self
+                # A recorded exception spends a slot like any other event.
+                if self._user_event_count >= self._max_events:
+                    self._dropped_events += 1
+                    return self
+                self._user_event_count += 1
             bounded: Optional[Dict[str, Any]] = None
             dropped = 0
             if attributes is not None:
@@ -331,18 +341,19 @@ def add_event(
                     MAX_ATTRIBUTES_PER_EVENT,
                     self._max_attribute_value_length,
                 )
-            self._events.append(
-                SpanEventRecord(
-                    name=sanitize_name(
-                        name, "Span event name", self._max_attribute_value_length
-                    ),
-                    timestamp_ns=resolve_supplied_ns(
-                        timestamp, self._now_ns(), "event timestamp"
-                    ),
-                    attributes=bounded,
-                    dropped_attributes_count=dropped,
-                )
+            event = SpanEventRecord(
+                name=sanitize_name(
+                    name, "Span event name", self._max_attribute_value_length
+                ),
+                timestamp_ns=resolve_supplied_ns(
+                    timestamp, self._now_ns(), "event timestamp"
+                ),
+                attributes=bounded,
+                dropped_attributes_count=dropped,
             )
+            with self._lock:
+                if not self._ended:
+                    self._events.append(event)
         return self
 
     def set_status(self, code: str, message: Optional[str] = None) -> "Span":
@@ -351,12 +362,15 @@ def set_status(self, code: str, message: Optional[str] = None) -> "Span":
                 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(
+            status = SpanStatus(
                 code,
                 truncate_string(text, self._max_attribute_value_length)
                 if text
                 else None,
             )
+            with self._lock:
+                if not self._ended:
+                    self._status = status
         return self
 
     @property
@@ -379,9 +393,12 @@ def _record_exception(self, exception: BaseException, keep_ok: bool) -> None:
 
     def update_name(self, name: str) -> "Span":
         if self._mutable("update_name"):
-            self._name = sanitize_name(
+            sanitized = sanitize_name(
                 name, "Span name", self._max_attribute_value_length
             )
+            with self._lock:
+                if not self._ended:
+                    self._name = sanitized
         return self
 
     def traceparent(self) -> Optional[str]:
@@ -400,11 +417,22 @@ def _child_context(self) -> ParentContext:
         )
 
     def end(self, end_time: Optional[SpanTimeInput] = None) -> None:
-        with self._tokens_lock:
+        with self._lock:
             if self._ended:
                 log.debug("Ignoring end() on a span that has already ended")
                 return
             self._ended = True
+            # Snapshotted under the lock: a write that lost the race to end()
+            # must not land in the record or after it.
+            attributes = {
+                key: value
+                for key, value in self._attributes.items()
+                if value is not None
+            }
+            events = list(self._events)
+            name, status = self._name, self._status
+            dropped_attributes = self._dropped_attributes
+            dropped_events = self._dropped_events
 
         derived = self._now_ns()
         resolved = resolve_supplied_ns(end_time, derived, "end time")
@@ -415,15 +443,15 @@ def end(self, end_time: Optional[SpanTimeInput] = None) -> None:
             trace_state=self._trace_state,
             trace_flags=self._trace_flags,
             parent_is_remote=self._parent_is_remote,
-            name=self._name,
+            name=name,
             kind=self._kind,
-            status=self._status,
-            attributes=dict(self._attributes),
-            events=self._events,
+            status=status,
+            attributes=attributes,
+            events=events,
             start_ns=self._start_ns,
             end_ns=clamp_end_ns(resolved, self._start_ns),
-            dropped_attributes_count=self._dropped_attributes,
-            dropped_events_count=self._dropped_events,
+            dropped_attributes_count=dropped_attributes,
+            dropped_events_count=dropped_events,
         )
         try:
             self._on_end(record)