From 42b55d6563bc21d40f21625cdfa92eab3aec84ec Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Fri, 7 Feb 2025 09:43:16 -0800 Subject: [PATCH 1/3] experimenting with wrappingcontext for dd-trace-api integration --- .../contrib/internal/dd_trace_api/patch.py | 45 ++++++++++++++++++- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/ddtrace/contrib/internal/dd_trace_api/patch.py b/ddtrace/contrib/internal/dd_trace_api/patch.py index c4b82bfa6b1..45a480682d7 100644 --- a/ddtrace/contrib/internal/dd_trace_api/patch.py +++ b/ddtrace/contrib/internal/dd_trace_api/patch.py @@ -4,17 +4,53 @@ from typing import List from typing import Optional from typing import Tuple +from typing import TypeVar import weakref import dd_trace_api +from wrapt.importer import when_imported import ddtrace +from ddtrace.internal.logger import get_logger +from ddtrace.internal.wrapping.context import WrappingContext _DD_HOOK_NAME = "dd.hook" _TRACER_KEY = "Tracer" _STUB_TO_REAL = weakref.WeakKeyDictionary() _STUB_TO_REAL[dd_trace_api.tracer] = ddtrace.tracer +log = get_logger(__name__) +T = TypeVar("T") + + +class DDTraceAPIWrappingContextBase(WrappingContext): + def _handle_return(self) -> None: + method_name = self.__frame__.f_code.co_name + stub_self = self.get_local("self") + api_return_value = self.get_local("retval") + _call_on_real_instance(stub_self, method_name, api_return_value, self.get_local("name")) + + def _handle_enter(self) -> None: + pass + + def __enter__(self) -> "DDTraceAPIWrappingContextBase": + super().__enter__() + + try: + self._handle_enter() + except Exception: # noqa: E722 + log.debug("Error handling dd_trace_api instrumentation enter", exc_info=True) + + return self + + def __return__(self, value: T) -> T: + """Always return the original value no matter what our instrumentation does""" + try: + self._handle_return() + except Exception: # noqa: E722 + log.debug("Error handling instrumentation return", exc_info=True) + + return value def _proxy_span_arguments(args: List, kwargs: Dict) -> Tuple[List, Dict]: @@ -57,12 +93,17 @@ def get_version() -> str: def patch(tracer=None): if getattr(dd_trace_api, "__datadog_patch", False): return - dd_trace_api.__datadog_patch = True _STUB_TO_REAL[dd_trace_api.tracer] = tracer - if not getattr(dd_trace_api, "__dd_has_audit_hook", False): + if False and not getattr(dd_trace_api, "__dd_has_audit_hook", False): addaudithook(_hook) dd_trace_api.__dd_has_audit_hook = True + @when_imported("dd_trace_api") + def _(m): + DDTraceAPIWrappingContextBase(m.Tracer.start_span).wrap() + + dd_trace_api.__datadog_patch = True + def unpatch(): if not getattr(dd_trace_api, "__datadog_patch", False): From f9bae759a603167e9d05a6cfd8f62087ffdeba48 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Fri, 7 Feb 2025 11:59:34 -0800 Subject: [PATCH 2/3] use wrappingcontext instead of sysaudit --- .../contrib/internal/dd_trace_api/patch.py | 39 +++++++++++++++++-- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/ddtrace/contrib/internal/dd_trace_api/patch.py b/ddtrace/contrib/internal/dd_trace_api/patch.py index 45a480682d7..f54a7d270b4 100644 --- a/ddtrace/contrib/internal/dd_trace_api/patch.py +++ b/ddtrace/contrib/internal/dd_trace_api/patch.py @@ -1,3 +1,4 @@ +import inspect from sys import addaudithook from typing import Any from typing import Dict @@ -25,10 +26,16 @@ class DDTraceAPIWrappingContextBase(WrappingContext): def _handle_return(self) -> None: - method_name = self.__frame__.f_code.co_name - stub_self = self.get_local("self") - api_return_value = self.get_local("retval") - _call_on_real_instance(stub_self, method_name, api_return_value, self.get_local("name")) + _call_on_real_instance( + self.get_local("self"), + self.__frame__.f_code.co_name, + self.get_local("retval"), + **{ + param: self.get_local(param) + for param in inspect.signature(self.__wrapped__).parameters.keys() + if param != "self" + } + ) def _handle_enter(self) -> None: pass @@ -101,6 +108,16 @@ def patch(tracer=None): @when_imported("dd_trace_api") def _(m): DDTraceAPIWrappingContextBase(m.Tracer.start_span).wrap() + DDTraceAPIWrappingContextBase(m.Tracer.trace).wrap() + DDTraceAPIWrappingContextBase(m.Tracer.current_span).wrap() + DDTraceAPIWrappingContextBase(m.Tracer.current_root_span).wrap() + DDTraceAPIWrappingContextBase(m.Span.finish).wrap() + DDTraceAPIWrappingContextBase(m.Span.set_exc_info).wrap() + DDTraceAPIWrappingContextBase(m.Span.finish_with_ancestors).wrap() + DDTraceAPIWrappingContextBase(m.Span.set_tags).wrap() + DDTraceAPIWrappingContextBase(m.Span.set_traceback).wrap() + DDTraceAPIWrappingContextBase(m.Span.__enter__).wrap() + DDTraceAPIWrappingContextBase(m.Span.__exit__).wrap() dd_trace_api.__datadog_patch = True @@ -110,3 +127,17 @@ def unpatch(): return dd_trace_api.__datadog_patch = False # NB sys.addaudithook's cannot be removed + + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.start_span).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.trace).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.current_span).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.current_root_span).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.finish).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.set_exc_info).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.finish_with_ancestors).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.set_tags).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.set_traceback).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.__enter__).unwrap() + DDTraceAPIWrappingContextBase.extract(dd_trace_api.Span.__exit__).unwrap() + + dd_trace_api.__datadog_patch = False From 975d9e9588b508b1b1a771166b0c385180ee23ec Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Mon, 10 Feb 2025 14:55:14 -0800 Subject: [PATCH 3/3] avoid using sysaudit --- ddtrace/contrib/internal/dd_trace_api/patch.py | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/ddtrace/contrib/internal/dd_trace_api/patch.py b/ddtrace/contrib/internal/dd_trace_api/patch.py index f54a7d270b4..5779bd2f72e 100644 --- a/ddtrace/contrib/internal/dd_trace_api/patch.py +++ b/ddtrace/contrib/internal/dd_trace_api/patch.py @@ -1,5 +1,4 @@ import inspect -from sys import addaudithook from typing import Any from typing import Dict from typing import List @@ -34,7 +33,7 @@ def _handle_return(self) -> None: param: self.get_local(param) for param in inspect.signature(self.__wrapped__).parameters.keys() if param != "self" - } + }, ) def _handle_enter(self) -> None: @@ -84,15 +83,6 @@ def _call_on_real_instance( _STUB_TO_REAL[retval_from_api] = retval_from_impl -def _hook(name, hook_args): - """Called in response to `sys.audit` events""" - if name != _DD_HOOK_NAME or not dd_trace_api.__datadog_patch: - return - args = hook_args[0][0] - api_return_value, stub_self, method_name = args[0:3] - _call_on_real_instance(stub_self, method_name, api_return_value, *args[3:], **hook_args[0][1]) - - def get_version() -> str: return getattr(dd_trace_api, "__version__", "") @@ -101,9 +91,6 @@ def patch(tracer=None): if getattr(dd_trace_api, "__datadog_patch", False): return _STUB_TO_REAL[dd_trace_api.tracer] = tracer - if False and not getattr(dd_trace_api, "__dd_has_audit_hook", False): - addaudithook(_hook) - dd_trace_api.__dd_has_audit_hook = True @when_imported("dd_trace_api") def _(m): @@ -126,7 +113,6 @@ def unpatch(): if not getattr(dd_trace_api, "__datadog_patch", False): return dd_trace_api.__datadog_patch = False - # NB sys.addaudithook's cannot be removed DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.start_span).unwrap() DDTraceAPIWrappingContextBase.extract(dd_trace_api.Tracer.trace).unwrap()