diff --git a/src/native/span/span_data.rs b/src/native/span/span_data.rs index bba1b219481..01a76d6fa21 100644 --- a/src/native/span/span_data.rs +++ b/src/native/span/span_data.rs @@ -74,17 +74,100 @@ impl SpanData { self.trace_id = id; self._trace_id_py = None; } +} + +const HTTP_STATUS_CODE_KEY: &str = "http.status_code"; + +/// Convert one Python key/value pair to native attribute storage. +/// +/// DEV: Keep Python coercion outside a mutable SpanData borrow. Arbitrary +/// `__str__` and `__index__` implementations can start nested spans and re-enter +/// the native context provider, which needs to borrow the active SpanData. +fn extract_attribute( + key: &Bound<'_, PyAny>, + value: &Bound<'_, PyAny>, +) -> Option<(AttrKey, AttributeValue)> { + let key_str = key.cast::().ok()?; + let is_http_status_code = key_str.to_str().unwrap_or("") == HTTP_STATUS_CODE_KEY; + let attr_key = AttrKey::new(key_str.clone().unbind()); + + // http.status_code must always be a string in meta. + // Fast path: typed contract is `str`, so most callers already pass a PyString. + // Only fall back to str() for non-string inputs (e.g. an int 200). + if is_http_status_code { + let s = if let Ok(s) = value.cast::() { + s.clone() + } else { + value.str().ok()? + }; + return Some((attr_key, AttributeValue::Str(s.unbind()))); + } - /// Setdefault helper for `_set_default_attributes`: insert one key/value pair only if - /// the key is not already present in either meta or metrics. - fn set_default_attribute_entry(&mut self, k: &Bound<'_, PyAny>, v: &Bound<'_, PyAny>) { - if !self.has_attribute(k) { - let _ = self.set_attribute(k, v); + // str → Str + if let Ok(s) = value.cast::() { + return Some((attr_key, AttributeValue::Str(s.clone().unbind()))); + } + + // float → Float (drop NaN/Inf) + // Check before int because some types (e.g. numpy.float64) implement __float__ + // but not __index__, so PyFloat succeeds and PyInt would fail. + if let Ok(f) = value.cast::() { + let n = f.value(); + if n.is_nan() || n.is_infinite() { + return None; } + return Some((attr_key, AttributeValue::Float(n))); } + + // int (catches bool and numpy.int* via __index__) → Int. + // extract::() succeeds for bool (True → 1, False → 0) and for any + // type implementing __index__. Python ints that overflow i64 fall through + // to the str() fallback below. + if let Ok(n) = value.extract::() { + return Some((attr_key, AttributeValue::Int(n))); + } + + // bytes → UTF-8 decoded Str (with U+FFFD replacements for invalid sequences) + if let Ok(b) = value.cast::() { + let decoded = String::from_utf8_lossy(b.as_bytes()); + let py_str = PyString::new(key.py(), &decoded); + return Some((attr_key, AttributeValue::Str(py_str.unbind()))); + } + + // Fallback: str(value) — covers Python ints that overflow i64, arbitrary objects, etc. + let s = value.str().ok()?; + Some((attr_key, AttributeValue::Str(s.unbind()))) } -const HTTP_STATUS_CODE_KEY: &str = "http.status_code"; +fn set_default_attribute( + slf: &Bound<'_, SpanData>, + key: &Bound<'_, PyAny>, + value: &Bound<'_, PyAny>, +) { + let Ok(key_str) = key.cast::() else { + return; + }; + let Ok(key_text) = key_str.to_str() else { + return; + }; + if slf.borrow().attributes.contains_key(key_text) { + return; + } + if let Some((attr_key, attr_value)) = extract_attribute(key, value) { + // Re-entrant coercion may have inserted the key after the check above; if so, + // `attr_value` is discarded rather than stored. A discarded str subclass can run + // Python finalization when its last reference is dropped, so release the SpanData + // borrow before dropping it (same hazard as the `replaced` value in `set_attribute`). + let discarded = match slf.borrow_mut().attributes.entry(attr_key) { + std::collections::hash_map::Entry::Occupied(_) => Some(attr_value), + std::collections::hash_map::Entry::Vacant(e) => { + e.insert(attr_value); + None + } + }; + drop(discarded); + } +} #[pyo3::pymethods] impl SpanData { @@ -509,75 +592,19 @@ impl SpanData { /// basis (bytes → UTF-8 decoded str, oversized ints → str, arbitrary objects → str). #[pyo3(name = "_set_attribute")] fn set_attribute( - &mut self, + slf: &Bound<'_, Self>, key: &Bound<'_, PyAny>, value: &Bound<'_, PyAny>, ) -> pyo3::PyResult<()> { - let Ok(key_str) = key.cast::() else { - return Ok(()); - }; - let attr_key = AttrKey::new(key_str.clone().unbind()); - - // http.status_code must always be a string in meta. - // Fast path: typed contract is `str`, so most callers already pass a PyString. - // Only fall back to str() for non-string inputs (e.g. an int 200). - if key_str.to_str().unwrap_or("") == HTTP_STATUS_CODE_KEY { - let s = if let Ok(s) = value.cast::() { - s.clone() - } else { - let Ok(s) = value.str() else { - return Ok(()); - }; - s + if let Some((attr_key, attr_value)) = extract_attribute(key, value) { + let replaced = { + let mut span = slf.borrow_mut(); + span.attributes.insert(attr_key, attr_value) }; - self.attributes - .insert(attr_key, AttributeValue::Str(s.unbind())); - return Ok(()); + // A replaced str subclass can run Python finalization when its last + // reference is dropped, so release the SpanData borrow first. + drop(replaced); } - - // str → Str - if let Ok(s) = value.cast::() { - self.attributes - .insert(attr_key, AttributeValue::Str(s.clone().unbind())); - return Ok(()); - } - - // float → Float (drop NaN/Inf) - // Check before int because some types (e.g. numpy.float64) implement __float__ - // but not __index__, so PyFloat succeeds and PyInt would fail. - if let Ok(f) = value.cast::() { - let n = f.value(); - if n.is_nan() || n.is_infinite() { - return Ok(()); - } - self.attributes.insert(attr_key, AttributeValue::Float(n)); - return Ok(()); - } - - // int (catches bool and numpy.int* via __index__) → Int. - // extract::() succeeds for bool (True → 1, False → 0) and for any - // type implementing __index__. Python ints that overflow i64 fall through - // to the str() fallback below. - if let Ok(n) = value.extract::() { - self.attributes.insert(attr_key, AttributeValue::Int(n)); - return Ok(()); - } - - // bytes → UTF-8 decoded Str (with U+FFFD replacements for invalid sequences) - if let Ok(b) = value.cast::() { - let decoded = String::from_utf8_lossy(b.as_bytes()); - let py_str = PyString::new(key.py(), &decoded); - self.attributes - .insert(attr_key, AttributeValue::Str(py_str.unbind())); - return Ok(()); - } - - // Fallback: str(value) — covers Python ints that overflow i64, arbitrary objects, etc. - let Ok(s) = value.str() else { - return Ok(()); - }; - self.attributes - .insert(attr_key, AttributeValue::Str(s.unbind())); Ok(()) } @@ -588,10 +615,10 @@ impl SpanData { /// neither, the call is a no-op. Invalid value types follow the same coercion rules as /// `_set_attribute`. #[pyo3(name = "_set_attributes")] - fn set_attributes(&mut self, attrs: &Bound<'_, PyAny>) -> pyo3::PyResult<()> { + fn set_attributes(slf: &Bound<'_, Self>, attrs: &Bound<'_, PyAny>) -> pyo3::PyResult<()> { if let Ok(d) = attrs.cast_exact::() { for (k, v) in d.iter() { - let _ = self.set_attribute(&k, &v); + let _ = Self::set_attribute(slf, &k, &v); } } else if let Ok(m) = attrs.cast::() { if let Ok(items) = m.items() { @@ -605,7 +632,7 @@ impl SpanData { let Ok(v) = pair.get_item(1) else { continue; }; - let _ = self.set_attribute(&k, &v); + let _ = Self::set_attribute(slf, &k, &v); } } } @@ -743,10 +770,13 @@ impl SpanData { /// Used by callers that previously called `_update_tags_from_context`. /// Callers handle any locking on the source dict themselves. #[pyo3(name = "_set_default_attributes")] - fn set_default_attributes(&mut self, values: &Bound<'_, PyAny>) -> pyo3::PyResult<()> { + fn set_default_attributes( + slf: &Bound<'_, Self>, + values: &Bound<'_, PyAny>, + ) -> pyo3::PyResult<()> { if let Ok(d) = values.cast_exact::() { for (k, v) in d.iter() { - self.set_default_attribute_entry(&k, &v); + set_default_attribute(slf, &k, &v); } } else if let Ok(m) = values.cast::() { if let Ok(items) = m.items() { @@ -760,7 +790,7 @@ impl SpanData { let Ok(v) = pair.get_item(1) else { continue; }; - self.set_default_attribute_entry(&k, &v); + set_default_attribute(slf, &k, &v); } } } diff --git a/tests/tracer/test_span_tags.py b/tests/tracer/test_span_tags.py index b0f29143fe8..dbc72e1aa2e 100644 --- a/tests/tracer/test_span_tags.py +++ b/tests/tracer/test_span_tags.py @@ -10,6 +10,7 @@ import mock import pytest +from ddtrace._trace.provider import DefaultContextProvider from ddtrace.constants import _SPAN_MEASURED_KEY from ddtrace.constants import ENV_KEY from ddtrace.constants import MANUAL_DROP_KEY @@ -20,6 +21,7 @@ from ddtrace.constants import USER_REJECT from ddtrace.constants import VERSION_KEY from ddtrace.ext import http +from ddtrace.internal.ci_visibility.context import CIContextProvider from ddtrace.trace import Span from tests.utils import assert_is_measured from tests.utils import assert_is_not_measured @@ -189,6 +191,41 @@ def __repr__(self): s.set_tag("a", Foo()) +@pytest.mark.parametrize("context_provider_class", [DefaultContextProvider, CIContextProvider]) +@pytest.mark.parametrize("setter", ["set_tag", "_set_attributes", "_set_default_attributes"]) +def test_attribute_string_coercion_can_read_active_span(tracer, context_provider_class, setter): + tracer.context_provider = context_provider_class() + + class ReentrantTag: + def __str__(self): + active = tracer.current_span() + assert active is span + with tracer.trace("nested") as nested: + assert nested._parent is span + assert nested.trace_id == span.trace_id + return str(active.trace_id) + + with tracer.trace("test") as span: + if setter == "set_tag": + span.set_tag("reentrant", ReentrantTag()) + else: + getattr(span, setter)({"reentrant": ReentrantTag()}) + assert span.get_tag("reentrant") == str(span.trace_id) + + +def test_set_default_attributes_preserves_value_set_during_coercion(): + span = Span(name="test.span") + + class ReentrantTag: + def __str__(self): + span.set_tag("reentrant", "set-during-coercion") + return "default" + + span._set_default_attributes({"reentrant": ReentrantTag()}) + + assert span.get_tag("reentrant") == "set-during-coercion" + + @mock.patch("ddtrace._trace.span.log") def test_numeric_tags_none(span_log): s = Span(name="test.span")