From e8a30fb86773f82b17d582e9aacd32553741b864 Mon Sep 17 00:00:00 2001 From: "Gabriele N. Tornetta" Date: Mon, 17 Nov 2025 13:30:11 +0000 Subject: [PATCH 1/2] chore: lazy context wrapping We implement a lazy context wrapping mechanism whereby the bytecode context wrapping is performed on first invocation of the function. Note that this still requires the original function to be instrumented via bytecode manipulations. --- ddtrace/internal/wrapping/__init__.py | 7 +++++ ddtrace/internal/wrapping/context.py | 10 ++++++++ tests/internal/test_wrapping.py | 37 +++++++++++++++++++++++++++ 3 files changed, 54 insertions(+) diff --git a/ddtrace/internal/wrapping/__init__.py b/ddtrace/internal/wrapping/__init__.py index 852c99dc151..1aa0235abce 100644 --- a/ddtrace/internal/wrapping/__init__.py +++ b/ddtrace/internal/wrapping/__init__.py @@ -336,3 +336,10 @@ def unwrap(wf, wrapper): except AttributeError: # The function is not wrapped so we return it as is. return cast(FunctionType, wf) + + +def wrap_once(f, wrapper): + def trampoline(_, args, kwargs): + return wrapper(unwrap(f, trampoline), args, kwargs) + + return wrap(f, trampoline) diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index 2e5be4b1013..1477fb572ac 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -14,6 +14,7 @@ from ddtrace.internal.assembly import Assembly from ddtrace.internal.utils.inspection import link_function_to_code +from ddtrace.internal.wrapping import wrap_once T = t.TypeVar("T") @@ -366,6 +367,15 @@ def wrap(self) -> None: def unwrap(self) -> None: raise NotImplementedError + def wrap_lazy(self) -> None: + """Perform the bytecode wrapping on first invocation.""" + + def wrapper(f, args, kwargs): + self.wrap() + return f(*args, **kwargs) + + wrap_once(self.__wrapped__, wrapper) + # This is the public interface exported by this module class WrappingContext(BaseWrappingContext): diff --git a/tests/internal/test_wrapping.py b/tests/internal/test_wrapping.py index 3610f0d452a..3360b53fb74 100644 --- a/tests/internal/test_wrapping.py +++ b/tests/internal/test_wrapping.py @@ -10,6 +10,7 @@ from ddtrace.internal.wrapping import is_wrapped_with from ddtrace.internal.wrapping import unwrap from ddtrace.internal.wrapping import wrap +from ddtrace.internal.wrapping import wrap_once from ddtrace.internal.wrapping.context import WrappingContext from ddtrace.internal.wrapping.context import _UniversalWrappingContext @@ -926,3 +927,39 @@ def foo(): new_method_count = len([_ for _ in gc.get_objects() if type(_).__name__ == "method"]) assert new_method_count <= method_count + 1 + + +def test_wrap_once(): + c = 0 + + def wrapper(f, args, kwargs): + nonlocal c + c += 1 + return f(*args, **kwargs) + + def foo(): + pass + + wrap_once(foo, wrapper) + + for _ in range(10): + foo() + + assert c == 1 + + +def test_wrapping_context_lazy(): + free = 42 + + def foo(): + return free + + (wc := DummyWrappingContext(foo)).wrap_lazy() + + assert not DummyWrappingContext.is_wrapped(foo) + + assert foo() == free + + assert DummyWrappingContext.is_wrapped(foo) + + assert wc.entered From 809ce1090657450e76714bb1fdaec58e202ed4ac Mon Sep 17 00:00:00 2001 From: "Gabriele N. Tornetta" Date: Mon, 17 Nov 2025 18:33:30 +0000 Subject: [PATCH 2/2] make it thread-safe --- ddtrace/internal/wrapping/__init__.py | 7 ---- ddtrace/internal/wrapping/context.py | 54 ++++++++++++++++++++++----- tests/internal/test_wrapping.py | 47 ++++++++++++----------- 3 files changed, 69 insertions(+), 39 deletions(-) diff --git a/ddtrace/internal/wrapping/__init__.py b/ddtrace/internal/wrapping/__init__.py index 1aa0235abce..852c99dc151 100644 --- a/ddtrace/internal/wrapping/__init__.py +++ b/ddtrace/internal/wrapping/__init__.py @@ -336,10 +336,3 @@ def unwrap(wf, wrapper): except AttributeError: # The function is not wrapped so we return it as is. return cast(FunctionType, wf) - - -def wrap_once(f, wrapper): - def trampoline(_, args, kwargs): - return wrapper(unwrap(f, trampoline), args, kwargs) - - return wrap(f, trampoline) diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index 1477fb572ac..2389631106a 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -13,8 +13,13 @@ from bytecode import Bytecode from ddtrace.internal.assembly import Assembly +from ddtrace.internal.forksafe import Lock from ddtrace.internal.utils.inspection import link_function_to_code -from ddtrace.internal.wrapping import wrap_once +from ddtrace.internal.wrapping import WrappedFunction +from ddtrace.internal.wrapping import Wrapper +from ddtrace.internal.wrapping import is_wrapped_with +from ddtrace.internal.wrapping import unwrap +from ddtrace.internal.wrapping import wrap T = t.TypeVar("T") @@ -367,15 +372,6 @@ def wrap(self) -> None: def unwrap(self) -> None: raise NotImplementedError - def wrap_lazy(self) -> None: - """Perform the bytecode wrapping on first invocation.""" - - def wrapper(f, args, kwargs): - self.wrap() - return f(*args, **kwargs) - - wrap_once(self.__wrapped__, wrapper) - # This is the public interface exported by this module class WrappingContext(BaseWrappingContext): @@ -416,6 +412,44 @@ def unwrap(self) -> None: _UniversalWrappingContext.extract(f).unregister(self) +class LazyWrappingContext(WrappingContext): + def __init__(self, f: FunctionType): + super().__init__(f) + + self._trampoline: t.Optional[Wrapper] = None + self._trampoline_lock = Lock() + + def wrap(self) -> None: + """Perform the bytecode wrapping on first invocation.""" + with (tl := self._trampoline_lock): + if self._trampoline is not None: + return + + def trampoline(_, args, kwargs): + with tl: + f = t.cast(WrappedFunction, self.__wrapped__) + if is_wrapped_with(self.__wrapped__, trampoline): + f = unwrap(f, trampoline) + super(LazyWrappingContext, self).wrap() + return f(*args, **kwargs) + + wrap(self.__wrapped__, trampoline) + + self._trampoline = trampoline + + def unwrap(self) -> None: + with self._trampoline_lock: + if self._trampoline is None: + return + + if self.is_wrapped(self.__wrapped__): + super().unwrap() + else: + unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline) + + self._trampoline = None + + class ContextWrappedFunction(Protocol): """A wrapped function.""" diff --git a/tests/internal/test_wrapping.py b/tests/internal/test_wrapping.py index 3360b53fb74..40e048f5bf3 100644 --- a/tests/internal/test_wrapping.py +++ b/tests/internal/test_wrapping.py @@ -10,7 +10,7 @@ from ddtrace.internal.wrapping import is_wrapped_with from ddtrace.internal.wrapping import unwrap from ddtrace.internal.wrapping import wrap -from ddtrace.internal.wrapping import wrap_once +from ddtrace.internal.wrapping.context import LazyWrappingContext from ddtrace.internal.wrapping.context import WrappingContext from ddtrace.internal.wrapping.context import _UniversalWrappingContext @@ -929,37 +929,40 @@ def foo(): assert new_method_count <= method_count + 1 -def test_wrap_once(): - c = 0 - - def wrapper(f, args, kwargs): - nonlocal c - c += 1 - return f(*args, **kwargs) +def test_wrapping_context_lazy(): + free = 42 def foo(): - pass + return free - wrap_once(foo, wrapper) + class DummyLazyWrappingContext(LazyWrappingContext): + def __init__(self, f): + super().__init__(f) - for _ in range(10): - foo() + self.count = 0 - assert c == 1 + def __enter__(self): + self.count += 1 + return super().__enter__() + (wc := DummyLazyWrappingContext(foo)).wrap() -def test_wrapping_context_lazy(): - free = 42 + assert not DummyLazyWrappingContext.is_wrapped(foo) - def foo(): - return free + for _ in range(n := 10): + assert foo() == free - (wc := DummyWrappingContext(foo)).wrap_lazy() + assert DummyLazyWrappingContext.is_wrapped(foo) - assert not DummyWrappingContext.is_wrapped(foo) + assert wc.count == n - assert foo() == free + wc.count = 0 - assert DummyWrappingContext.is_wrapped(foo) + wc.unwrap() - assert wc.entered + for _ in range(10): + assert not DummyLazyWrappingContext.is_wrapped(foo) + + assert foo() == free + + assert wc.count == 0