1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
|
import threading
from functools import partial
from typing import Any
from typing import Callable
from py_zipkin import storage
_orig_Thread_start = threading.Thread.start
_orig_Thread_run = threading.Thread.run
def _Thread_pre_start(self: Any) -> None:
self._orig_tracer = None
if storage.has_default_tracer():
self._orig_tracer = storage.get_default_tracer().copy()
def _Thread_wrap_run(self: Any, actual_run_fn: Callable[[], None]) -> None:
# This executes in the new OS thread
if self._orig_tracer:
# Inject our copied Tracer into our thread-local-storage
storage.set_default_tracer(self._orig_tracer)
try:
actual_run_fn()
finally:
# I think this is probably a good idea for the same reasons the
# parent class deletes __target, __args, and __kwargs
if self._orig_tracer:
del self._orig_tracer
def patch_threading() -> None: # pragma: no cover
"""Monkey-patch threading module to work better with tracing."""
def _new_start(self: Any) -> None:
_Thread_pre_start(self)
_orig_Thread_start(self)
def _new_run(self: Any) -> None:
_Thread_wrap_run(self, partial(_orig_Thread_run, self))
threading.Thread.start = _new_start # type: ignore[method-assign]
threading.Thread.run = _new_run # type: ignore[method-assign]
def unpatch_threading() -> None: # pragma: no cover
threading.Thread.start = _orig_Thread_start # type: ignore[method-assign]
threading.Thread.run = _orig_Thread_run # type: ignore[method-assign]
|
