fix(tracing): ensure Context is serializable (backport #4432) (#4438)

* fix(tracing): ensure Context is serializable (#4432)

ddtrace.context.Context object is not serializable, meaning it cannot be pickled/shared between processes.

This breaks the example usage we have in our documentation for passing context through to other threads.

e.g. Process(target=_target, args=(ctx, ))

This fix added __getstate__ and __setstate__ methods to Context class to have pickle ignore the RLock which cannot be serialized.

Co-authored-by: Munir Abdinur <[email protected]>
Co-authored-by: Kyle Verhoog <[email protected]>
(cherry picked from commit 96e6bca)

# Conflicts:
#	ddtrace/context.py
#	tests/tracer/test_context.py

* fix conflict

* Update tests/tracer/test_context.py

* merge conflict: remove traceparent from context

* Update ddtrace/context.py

Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: Munir Abdinur <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Julien Danjou <[email protected]>

Author: 4 people

ddtrace/context.py

@@ -13,11 +13,21 @@
 from .internal.logger import get_logger
 
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
+    from typing import Tuple
+
     from .span import Span
     from .span import _MetaDictType
     from .span import _MetricDictType
 
+    _ContextState = Tuple[
+        Optional[int],  # trace_id
+        Optional[int],  # span_id
+        _MetaDictType,  # _meta
+        _MetricDictType,  # _metrics
+    ]
+
+
 log = get_logger(__name__)
 
 
@@ -63,6 +73,22 @@ def __init__(
             # https://github.com/DataDog/dd-trace-py/blob/a1932e8ddb704d259ea8a3188d30bf542f59fd8d/ddtrace/tracer.py#L489-L508
             self._lock = threading.RLock()
 
+    def __getstate__(self):
+        # type: () -> _ContextState
+        return (
+            self.trace_id,
+            self.span_id,
+            self._meta,
+            self._metrics,
+            # Note: self._lock is not serializable
+        )
+
+    def __setstate__(self, state):
+        # type: (_ContextState) -> None
+        self.trace_id, self.span_id, self._meta, self._metrics = state
+        # We cannot serialize and lock, so we must recreate it unless we already have one
+        self._lock = threading.RLock()
+
     def _with_span(self, span):
         # type: (Span) -> Context
         """Return a shallow copy of the context with the given span."""

releasenotes/notes/fix-context-serializable-2c9768cbe4738b73.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    tracing: make ``ddtrace.context.Context`` serializable which fixes distributed tracing across processes.

tests/integration/test_context_snapshots.py

@@ -0,0 +1,44 @@
+import pytest
+
+from tests.utils import snapshot
+
+from .test_integration import AGENT_VERSION
+
+
+pytestmark = pytest.mark.skipif(AGENT_VERSION != "testagent", reason="Tests only compatible with a testagent")
+
+
+@snapshot()
+def test_context_multiprocess(run_python_code_in_subprocess):
+    # Testing example from our docs:
+    # https://ddtrace.readthedocs.io/en/stable/advanced_usage.html#tracing-across-processes
+    code = """
+from multiprocessing import Process
+import time
+
+from ddtrace import tracer
+
+
+def _target(ctx):
+    tracer.context_provider.activate(ctx)
+    with tracer.trace("proc"):
+        time.sleep(0.1)
+    tracer.shutdown()
+
+
+def main():
+    with tracer.trace("work"):
+        proc = Process(target=_target, args=(tracer.current_trace_context(), ))
+        proc.start()
+        time.sleep(0.25)
+        proc.join()
+
+
+if __name__ == "__main__":
+    main()
+    """
+
+    stdout, stderr, status, _ = run_python_code_in_subprocess(code=code)
+    assert status == 0, (stdout, stderr)
+    assert stdout == b"", stderr
+    assert stderr == b"", stdout

tests/snapshots/tests.integration.test_context_snapshots.test_context_multiprocess.json

@@ -0,0 +1,42 @@
+[[
+  {
+    "name": "work",
+    "service": null,
+    "resource": "work",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "runtime-id": "f706b29e0e8049178c7f1e0ac8d01ab7"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "system.pid": 25193
+    },
+    "duration": 259538000,
+    "start": 1667237294717521000
+  },
+     {
+       "name": "proc",
+       "service": null,
+       "resource": "proc",
+       "trace_id": 0,
+       "span_id": 2,
+       "parent_id": 1,
+       "meta": {
+         "_dd.p.dm": "-0",
+         "runtime-id": "38ca0ac0547f4097b2e030ebff1064c7"
+       },
+       "metrics": {
+         "_dd.top_level": 1,
+         "_dd.tracer_kr": 1.0,
+         "_sampling_priority_v1": 1,
+         "system.pid": 25194
+       },
+       "duration": 100317000,
+       "start": 1667237294727339000
+     }]]

tests/tracer/test_context.py

@@ -1,3 +1,5 @@
+import pickle
+
 import pytest
 
 from ddtrace.context import Context
@@ -48,3 +50,19 @@ def test_eq(ctx1, ctx2):
 )
 def test_not_eq(ctx1, ctx2):
     assert ctx1 != ctx2
+
+
+@pytest.mark.parametrize(
+    "context",
+    [
+        Context(),
+        Context(trace_id=123, span_id=321),
+        Context(trace_id=123, span_id=321, dd_origin="synthetics", sampling_priority=2),
+        Context(trace_id=123, span_id=321, meta={"meta": "value"}, metrics={"metric": 4.556}),
+    ],
+)
+def test_context_serializable(context):
+    # type: (Context) -> None
+    state = pickle.dumps(context)
+    restored = pickle.loads(state)
+    assert context == restored