Skip to content

Commit 83480a0

Browse files
mabdinurmabdinur
authored
fix(context): ensures all fields in Context are pickable (#7949)
Resolves: #7849 We will need to make baggage pickable as well. ## Checklist - [x] Change(s) are motivated and described in the PR description. - [x] Testing strategy is described if automated tests are not included in the PR. - [x] Risk is outlined (performance impact, potential for breakage, maintainability, etc). - [x] Change is maintainable (easy to change, telemetry, documentation). - [x] [Library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) are followed. If no release note is required, add label `changelog/no-changelog`. - [x] Documentation is included (in-code, generated user docs, [public corp docs](https://github.com/DataDog/documentation/)). - [x] Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [ ] Title is accurate. - [ ] No unnecessary changes are introduced. - [ ] Description motivates each change. - [ ] Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes unless absolutely necessary. - [ ] Testing strategy adequately addresses listed risk(s). - [ ] Change is maintainable (easy to change, telemetry, documentation). - [ ] Release note makes sense to a user of the library. - [ ] Reviewer has explicitly acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment. - [ ] Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting) - [ ] If this PR touches code that signs or publishes builds or packages, or handles credentials of any kind, I've requested a review from `@DataDog/security-design-and-guidance`. - [ ] This PR doesn't touch any of that.
1 parent 4c63636 commit 83480a0

File tree

4 files changed

+122
-37
lines changed

4 files changed

+122
-37
lines changed

‎ddtrace/context.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
Optional[int], # span_id
3232
_MetaDictType, # _meta
3333
_MetricDictType, # _metrics
34+
list[SpanLink],
35+
dict[str, Any],
3436
]
3537

3638

@@ -89,12 +91,14 @@ def __getstate__(self):
8991
self.span_id,
9092
self._meta,
9193
self._metrics,
94+
self._span_links,
95+
self._baggage
9296
# Note: self._lock is not serializable
9397
)
9498

9599
def __setstate__(self, state):
96100
# type: (_ContextState) -> None
97-
self.trace_id, self.span_id, self._meta, self._metrics = state
101+
self.trace_id, self.span_id, self._meta, self._metrics, self._span_links, self._baggage = state
98102
# We cannot serialize and lock, so we must recreate it unless we already have one
99103
self._lock = threading.RLock()
100104

@@ -246,17 +250,20 @@ def __eq__(self, other):
246250
and self.span_id == other.span_id
247251
and self._meta == other._meta
248252
and self._metrics == other._metrics
253+
and self._span_links == other._span_links
254+
and self._baggage == other._baggage
249255
)
250256
return False
251257

252258
def __repr__(self):
253259
# type: () -> str
254-
return "Context(trace_id=%s, span_id=%s, _meta=%s, _metrics=%s, _span_links=%s)" % (
260+
return "Context(trace_id=%s, span_id=%s, _meta=%s, _metrics=%s, _span_links=%s, _baggage=%s)" % (
255261
self.trace_id,
256262
self.span_id,
257263
self._meta,
258264
self._metrics,
259265
self._span_links,
266+
self._baggage,
260267
)
261268

262269
__str__ = __repr__

‎releasenotes/notes/make-span-links-pickable-d57b92508f5d3a74.yaml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
---
2+
fixes:
3+
- |
4+
tracing: Ensures all fields in ``ddtrace.context.Context`` are picklable.

‎tests/tracer/test_context.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from ddtrace.context import Context
88
from ddtrace.span import Span
9+
from ddtrace.tracing._span_link import SpanLink
910

1011

1112
@pytest.mark.parametrize(
@@ -88,6 +89,24 @@ def validate_traceparent(context, sampled_expected):
8889
Context(trace_id=123, span_id=321),
8990
Context(trace_id=123, span_id=321, dd_origin="synthetics", sampling_priority=2),
9091
Context(trace_id=123, span_id=321, meta={"meta": "value"}, metrics={"metric": 4.556}),
92+
Context(
93+
trace_id=123,
94+
span_id=321,
95+
meta={"meta": "value"},
96+
metrics={"metric": 4.556},
97+
span_links=[
98+
SpanLink(
99+
trace_id=123 << 64,
100+
span_id=456,
101+
tracestate="congo=t61rcWkgMzE",
102+
flags=1,
103+
attributes={"link.name": "some_name"},
104+
)
105+
],
106+
),
107+
Context(
108+
trace_id=123, span_id=321, meta={"meta": "value"}, metrics={"metric": 4.556}, baggage={"some_value": 1}
109+
),
91110
],
92111
)
93112
def test_context_serializable(context):

‎tests/tracer/test_propagation.py‎

Lines changed: 90 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import json
33
import logging
44
import os
5+
import pickle
56

67
import pytest
78

@@ -977,7 +978,7 @@ def test_extract_tracestate(caplog, ts_string, expected_tuple, expected_logging,
977978
"trace_id": TRACE_ID,
978979
"span_id": 67667974448284343,
979980
"meta": {
980-
"tracestate": "dd=s:2;o:rum;t.dm:-4;t.usr.id:baz64,congo=t61rcWkgMzE",
981+
"tracestate": TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACESTATE],
981982
"_dd.p.dm": "-4",
982983
"_dd.p.usr.id": "baz64",
983984
"_dd.origin": "rum",
@@ -1416,7 +1417,21 @@ def test_extract_tracecontext(headers, expected_context):
14161417
"span_id": 67667974448284343,
14171418
"sampling_priority": 2,
14181419
"dd_origin": "rum",
1419-
"tracestate": "dd=s:2;o:rum;t.dm:-4;t.usr.id:baz64,congo=t61rcWkgMzE",
1420+
"meta": {
1421+
"traceparent": TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACEPARENT],
1422+
"tracestate": TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACESTATE],
1423+
"_dd.p.dm": "-4",
1424+
"_dd.p.usr.id": "baz64",
1425+
},
1426+
"span_links": [
1427+
SpanLink(
1428+
trace_id=13088165645273925489,
1429+
span_id=5678,
1430+
tracestate=None,
1431+
flags=1,
1432+
attributes={"reason": "terminated_context", "context_headers": "datadog"},
1433+
)
1434+
],
14201435
},
14211436
),
14221437
(
@@ -1435,6 +1450,29 @@ def test_extract_tracecontext(headers, expected_context):
14351450
"span_id": 5678,
14361451
"sampling_priority": 1,
14371452
"dd_origin": "synthetics",
1453+
"span_links": [
1454+
SpanLink(
1455+
trace_id=TRACE_ID,
1456+
span_id=11744061942159299346,
1457+
tracestate=None,
1458+
flags=1,
1459+
attributes={"reason": "terminated_context", "context_headers": "b3multi"},
1460+
),
1461+
SpanLink(
1462+
trace_id=TRACE_ID,
1463+
span_id=16453819474850114513,
1464+
tracestate=None,
1465+
flags=1,
1466+
attributes={"reason": "terminated_context", "context_headers": "b3"},
1467+
),
1468+
SpanLink(
1469+
trace_id=TRACE_ID,
1470+
span_id=67667974448284343,
1471+
tracestate=TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACESTATE],
1472+
flags=1,
1473+
attributes={"reason": "terminated_context", "context_headers": "tracecontext"},
1474+
),
1475+
],
14381476
},
14391477
),
14401478
(
@@ -1451,6 +1489,29 @@ def test_extract_tracecontext(headers, expected_context):
14511489
"span_id": 5678,
14521490
"sampling_priority": 1,
14531491
"dd_origin": "synthetics",
1492+
"span_links": [
1493+
SpanLink(
1494+
trace_id=TRACE_ID,
1495+
span_id=11744061942159299346,
1496+
tracestate=None,
1497+
flags=1,
1498+
attributes={"reason": "terminated_context", "context_headers": "b3multi"},
1499+
),
1500+
SpanLink(
1501+
trace_id=TRACE_ID,
1502+
span_id=16453819474850114513,
1503+
tracestate=None,
1504+
flags=1,
1505+
attributes={"reason": "terminated_context", "context_headers": "b3"},
1506+
),
1507+
SpanLink(
1508+
trace_id=TRACE_ID,
1509+
span_id=67667974448284343,
1510+
tracestate=TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACESTATE],
1511+
flags=1,
1512+
attributes={"reason": "terminated_context", "context_headers": "tracecontext"},
1513+
),
1514+
],
14541515
},
14551516
),
14561517
(
@@ -1609,7 +1670,7 @@ def test_extract_tracecontext(headers, expected_context):
16091670
"span_id": 5678,
16101671
"sampling_priority": 1,
16111672
"dd_origin": "synthetics",
1612-
"tracestate": "dd=s:2;o:rum;t.dm:-4;t.usr.id:baz64,congo=t61rcWkgMzE",
1673+
"meta": {"tracestate": TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACESTATE]},
16131674
},
16141675
),
16151676
# testing that tracestate is not added when tracecontext style comes later and does not match first style's trace-id
@@ -1622,6 +1683,15 @@ def test_extract_tracecontext(headers, expected_context):
16221683
"span_id": 5678,
16231684
"sampling_priority": 1,
16241685
"dd_origin": "synthetics",
1686+
"span_links": [
1687+
SpanLink(
1688+
trace_id=TRACE_ID,
1689+
span_id=67667974448284343,
1690+
tracestate=TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACESTATE],
1691+
flags=1,
1692+
attributes={"reason": "terminated_context", "context_headers": "tracecontext"},
1693+
)
1694+
],
16251695
},
16261696
),
16271697
(
@@ -1652,7 +1722,10 @@ def test_extract_tracecontext(headers, expected_context):
16521722
"span_id": 67667974448284343,
16531723
"sampling_priority": 2,
16541724
"dd_origin": "rum",
1655-
"tracestate": "dd=s:2;o:rum",
1725+
"meta": {
1726+
"tracestate": "dd=s:2;o:rum",
1727+
"traceparent": TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACEPARENT],
1728+
},
16561729
},
16571730
),
16581731
]
@@ -1663,41 +1736,23 @@ def test_propagation_extract_env(name, styles, headers, expected_context, run_py
16631736
# Execute the test code in isolation to ensure env variables work as expected
16641737
code = """
16651738
import json
1666-
1739+
import pickle
1740+
from ddtrace.context import Context
16671741
from ddtrace.propagation.http import HTTPPropagator
16681742
1669-
16701743
context = HTTPPropagator.extract({!r})
1671-
if context is None:
1672-
print("null")
1673-
else:
1674-
if context._meta.get("tracestate"):
1675-
print(json.dumps({{
1676-
"trace_id": context.trace_id,
1677-
"span_id": context.span_id,
1678-
"sampling_priority": context.sampling_priority,
1679-
"dd_origin": context.dd_origin,
1680-
"tracestate": context._meta.get("tracestate")
1681-
}}))
1682-
else:
1683-
print(json.dumps({{
1684-
"trace_id": context.trace_id,
1685-
"span_id": context.span_id,
1686-
"sampling_priority": context.sampling_priority,
1687-
"dd_origin": context.dd_origin,
1688-
}}))
1744+
expected_context = Context(**pickle.loads({!r}))
1745+
assert context == expected_context, f"Expected {{expected_context}} but got {{context}}"
16891746
""".format(
1690-
headers
1747+
headers,
1748+
pickle.dumps(expected_context),
16911749
)
16921750
env = os.environ.copy()
16931751
if styles is not None:
16941752
env["DD_TRACE_PROPAGATION_STYLE"] = ",".join(styles)
16951753
stdout, stderr, status, _ = run_python_code_in_subprocess(code=code, env=env)
16961754
assert status == 0, (stdout, stderr)
16971755

1698-
result = json.loads(stdout.decode())
1699-
assert result == expected_context
1700-
17011756

17021757
@pytest.mark.parametrize("name,styles,headers,expected_context", EXTRACT_FIXTURES)
17031758
def test_propagation_extract_w_config(name, styles, headers, expected_context, run_python_code_in_subprocess):
@@ -1814,7 +1869,7 @@ def test_DD_TRACE_PROPAGATION_STYLE_EXTRACT_overrides_DD_TRACE_PROPAGATION_STYLE
18141869
],
18151870
ALL_HEADERS,
18161871
Context(
1817-
trace_id=171395628812617415352188477958425669623,
1872+
trace_id=TRACE_ID,
18181873
span_id=67667974448284343,
18191874
meta={
18201875
"traceparent": TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACEPARENT],
@@ -1844,7 +1899,7 @@ def test_DD_TRACE_PROPAGATION_STYLE_EXTRACT_overrides_DD_TRACE_PROPAGATION_STYLE
18441899
span_id=67667974448284343,
18451900
meta={
18461901
"traceparent": "00-000000000000000064fe8b2a57d3eff7-00f067aa0ba902b7-01",
1847-
"tracestate": "dd=s:2;o:rum;t.dm:-4;t.usr.id:baz64,congo=t61rcWkgMzE",
1902+
"tracestate": TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACESTATE],
18481903
"_dd.p.dm": "-4",
18491904
"_dd.p.usr.id": "baz64",
18501905
"_dd.origin": "rum",
@@ -1873,7 +1928,7 @@ def test_DD_TRACE_PROPAGATION_STYLE_EXTRACT_overrides_DD_TRACE_PROPAGATION_STYLE
18731928
],
18741929
ALL_HEADERS,
18751930
Context(
1876-
trace_id=171395628812617415352188477958425669623,
1931+
trace_id=TRACE_ID,
18771932
span_id=67667974448284343,
18781933
meta={
18791934
"traceparent": TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACEPARENT],
@@ -1912,21 +1967,21 @@ def test_DD_TRACE_PROPAGATION_STYLE_EXTRACT_overrides_DD_TRACE_PROPAGATION_STYLE
19121967
metrics={"_sampling_priority_v1": 1},
19131968
span_links=[
19141969
SpanLink(
1915-
trace_id=171395628812617415352188477958425669623,
1970+
trace_id=TRACE_ID,
19161971
span_id=67667974448284343,
1917-
tracestate="dd=s:2;o:rum;t.dm:-4;t.usr.id:baz64,congo=t61rcWkgMzE",
1972+
tracestate=TRACECONTEXT_HEADERS_VALID[_HTTP_HEADER_TRACESTATE],
19181973
flags=1,
19191974
attributes={"reason": "terminated_context", "context_headers": "tracecontext"},
19201975
),
19211976
SpanLink(
1922-
trace_id=171395628812617415352188477958425669623,
1977+
trace_id=TRACE_ID,
19231978
span_id=16453819474850114513,
19241979
tracestate=None,
19251980
flags=1,
19261981
attributes={"reason": "terminated_context", "context_headers": "b3"},
19271982
),
19281983
SpanLink(
1929-
trace_id=171395628812617415352188477958425669623,
1984+
trace_id=TRACE_ID,
19301985
span_id=11744061942159299346,
19311986
tracestate=None,
19321987
flags=1,

0 commit comments

Comments
 (0)