Repository: DataDog/dd-trace-py

Issue:
"Tracing Across Processes" example crashes
<!--
Thanks for taking the time for reporting an issue!

Before reporting an issue on dd-trace-py, please be sure to provide all
necessary information.

If you're hitting a bug, make sure that you're using the latest version of this
library.
-->

### Summary of problem

Hello, I cannot get "tracing across processes" to work. Python crashes while trying to pickle the tracer context RLock. The minimal reproduction example is [the example from the ddtrace advanced usage docs](https://ddtrace.readthedocs.io/en/stable/advanced_usage.html#tracing-across-processes). I also cannot get this to work with any prior ddtrace version.

### Which version of dd-trace-py are you using?

ddtrace: 1.5.3
python: 3.9.13

### How can we reproduce your problem?

Run the `Tracing Across Processes` example [from the docs](https://ddtrace.readthedocs.io/en/stable/advanced_usage.html#tracing-across-processes)

```python
from multiprocessing import Process
import time
from ddtrace import tracer

def _target(ctx):
    tracer.context_provider.activate(ctx)
    with tracer.trace("proc"):
        time.sleep(1)
    tracer.shutdown()

with tracer.trace("work"):
    proc = Process(target=_target, args=(tracer.current_trace_context(),))
    proc.start()
    time.sleep(1)
    proc.join()
```

### What is the result that you get?

```
Traceback (most recent call last):
  File "/Users/jboulanger/Library/Application Support/JetBrains/PyCharm2021.2/scratches/scratch.py", line 15, in <module>
    proc.start()
  File "/usr/local/Cellar/python@3.9/3.9.13_3/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/usr/local/Cellar/python@3.9/3.9.13_3/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/usr/local/Cellar/python@3.9/3.9.13_3/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/usr/local/Cellar/python@3.9/3.9.13_3/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/usr/local/Cellar/python@3.9/3.9.13_3/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/usr/local/Cellar/python@3.9/3.9.13_3/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/usr/local/Cellar/python@3.9/3.9.13_3/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
TypeError: cannot pickle '_thread.RLock' object
Process finished with exit code 1
```

### What is the result that you expected?

I expected this to complete without crashing


Discussion:

Resolving patch:
diff --git a/ddtrace/context.py b/ddtrace/context.py
--- a/ddtrace/context.py
+++ b/ddtrace/context.py
@@ -14,10 +14,20 @@
 
 
 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."""


Test patch:
diff --git a/tests/integration/test_context_snapshots.py b/tests/integration/test_context_snapshots.py
new file mode 100644
--- /dev/null
+++ b/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
diff --git a/tests/snapshots/tests.integration.test_context_snapshots.test_context_multiprocess.json b/tests/snapshots/tests.integration.test_context_snapshots.test_context_multiprocess.json
new file mode 100644
--- /dev/null
+++ b/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
+     }]]
diff --git a/tests/tracer/test_context.py b/tests/tracer/test_context.py
--- a/tests/tracer/test_context.py
+++ b/tests/tracer/test_context.py
@@ -1,3 +1,5 @@
+import pickle
+
 import pytest
 
 from ddtrace.context import Context
@@ -75,3 +77,19 @@ def validate_traceparent(context, sampled_expected):
     span = Span("span_c")
     span.context.sampling_priority = 1
     validate_traceparent(span.context, "01")
+
+
+@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