Remove 'wrapt' dependency from the Lock Profiler

vlad-scherbich · vlad-scherbich · commit c240aa2d0bc3 · 2025-11-03T10:50:46.000-05:00
diff --git a/ddtrace/profiling/collector/_lock.py b/ddtrace/profiling/collector/_lock.py
@@ -17,8 +17,6 @@
 from typing import Tuple
 from typing import Type
 
-import wrapt
-
 from ddtrace.internal.datadog.profiling import ddup
 from ddtrace.profiling import _threading
 from ddtrace.profiling import collector
@@ -34,22 +32,24 @@ def _current_thread() -> Tuple[int, str]:
     return thread_id, _threading.get_thread_name(thread_id)
 
 
-# We need to know if wrapt is compiled in C or not. If it's not using the C module, then the wrappers function will
-# appear in the stack trace and we need to hide it.
-WRAPT_C_EXT: bool
-if os.environ.get("WRAPT_DISABLE_EXTENSIONS"):
-    WRAPT_C_EXT = False
-else:
-    try:
-        import wrapt._wrappers as _w  # noqa: F401
-    except ImportError:
-        WRAPT_C_EXT = False
-    else:
-        WRAPT_C_EXT = True
-        del _w
+class _ProfiledLock:
+    """Lightweight lock wrapper that profiles lock acquire/release operations.
+
+    This is a simple delegating wrapper that intercepts lock methods without
+    the overhead of a full proxy object.
+    """
 
+    __slots__ = (
+        "__wrapped__",
+        "_self_tracer",
+        "_self_max_nframes",
+        "_self_capture_sampler",
+        "_self_endpoint_collection_enabled",
+        "_self_init_loc",
+        "_self_acquired_at",
+        "_self_name",
+    )
 
-class _ProfiledLock(wrapt.ObjectProxy):
     def __init__(
         self,
         wrapped: Any,
@@ -58,12 +58,13 @@ def __init__(
         capture_sampler: collector.CaptureSampler,
         endpoint_collection_enabled: bool,
     ) -> None:
-        wrapt.ObjectProxy.__init__(self, wrapped)
+        self.__wrapped__: Any = wrapped
         self._self_tracer: Optional[Tracer] = tracer
         self._self_max_nframes: int = max_nframes
         self._self_capture_sampler: collector.CaptureSampler = capture_sampler
         self._self_endpoint_collection_enabled: bool = endpoint_collection_enabled
-        frame: FrameType = sys._getframe(2 if WRAPT_C_EXT else 3)
+        # Frame depth: 0=__init__, 1=_profiled_allocate_lock, 2=_LockAllocatorWrapper.__call__, 3=caller
+        frame: FrameType = sys._getframe(3)
         code: CodeType = frame.f_code
         self._self_init_loc: str = "%s:%d" % (os.path.basename(code.co_filename), frame.f_lineno)
         self._self_acquired_at: int = 0
@@ -134,11 +135,6 @@ def acquire(self, *args: Any, **kwargs: Any) -> Any:
         return self._acquire(self.__wrapped__.acquire, *args, **kwargs)
 
     def _release(self, inner_func: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
-        # The underlying threading.Lock class is implemented using C code, and
-        # it doesn't have the __dict__ attribute. So we can't do
-        # self.__dict__.pop("_self_acquired_at", None) to remove the attribute.
-        # Instead, we need to use the following workaround to retrieve and
-        # remove the attribute.
         start: Optional[int] = getattr(self, "_self_acquired_at", None)
         try:
             # Though it should generally be avoided to call release() from
@@ -213,9 +209,14 @@ def _find_self_name(self, var_dict: Dict[str, Any]) -> Optional[str]:
                 return name
             if config.lock.name_inspect_dir:
                 for attribute in dir(value):
-                    if not attribute.startswith("__") and getattr(value, attribute) is self:
-                        self._self_name = attribute
-                        return attribute
+                    try:
+                        if not attribute.startswith("__") and getattr(value, attribute) is self:
+                            self._self_name = attribute
+                            return attribute
+                    except AttributeError:
+                        # With __slots__, accessing unset attributes raises AttributeError
+                        # (e.g., _self_acquired_at after it's been deleted in _release)
+                        continue
         return None
 
     # Get lock acquire/release call location and variable name the lock is assigned to
@@ -251,12 +252,38 @@ def _maybe_update_self_name(self) -> None:
         if not self._self_name:
             self._self_name = ""
 
+    # Delegate remaining lock methods to the wrapped lock
+    def locked(self) -> bool:
+        """Return True if lock is currently held."""
+        return self.__wrapped__.locked()
+
+    def __repr__(self) -> str:
+        return f"<_ProfiledLock({self.__wrapped__!r}) at {self._self_init_loc}>"
+
+    # Support for being used in with statements
+    def __bool__(self) -> bool:
+        return True
+
+
+class _LockAllocatorWrapper:
+    """Wrapper for lock allocator functions that prevents method binding.
+
+    When a function is stored as a class attribute and accessed via an instance,
+    Python's descriptor protocol normally binds it as a method. This wrapper
+    prevents that behavior by implementing __get__ to always return self,
+    similar to how staticmethod works, but as a callable object.
+    """
+
+    __slots__ = ("_func",)
+
+    def __init__(self, func: Callable[..., Any]) -> None:
+        self._func: Callable[..., Any] = func
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._func(*args, **kwargs)
 
-class FunctionWrapper(wrapt.FunctionWrapper):
-    # Override the __get__ method: whatever happens, _allocate_lock is always considered by Python like a "static"
-    # method, even when used as a class attribute. Python never tried to "bind" it to a method, because it sees it is a
-    # builtin function. Override default wrapt behavior here that tries to detect bound method.
-    def __get__(self, instance: Any, owner: Optional[Type] = None) -> FunctionWrapper:  # type: ignore
+    def __get__(self, instance: Any, owner: Optional[Type] = None) -> _LockAllocatorWrapper:
+        # Always return self, never bind as a method
         return self
 
 
@@ -303,9 +330,9 @@ def patch(self) -> None:
         # Nobody should use locks from `_thread`; if they do so, then it's deliberate and we don't profile.
         self._original = self._get_patch_target()
 
-        # TODO: `instance` is unused
-        def _allocate_lock(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> _ProfiledLock:
-            lock: Any = wrapped(*args, **kwargs)
+        # Create a simple wrapper function that returns profiled locks
+        def _profiled_allocate_lock(*args: Any, **kwargs: Any) -> _ProfiledLock:
+            lock: Any = self._original(*args, **kwargs)
             return self.PROFILED_LOCK_CLASS(
                 lock,
                 self.tracer,
@@ -314,7 +341,9 @@ def _allocate_lock(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> _Prof
                 self.endpoint_collection_enabled,
             )
 
-        self._set_patch_target(FunctionWrapper(self._original, _allocate_lock))
+        # Wrap the function to prevent it from being bound as a method when
+        # accessed as a class attribute (e.g., Foo.lock_class = threading.Lock)
+        self._set_patch_target(_LockAllocatorWrapper(_profiled_allocate_lock))
 
     def unpatch(self) -> None:
         """Unpatch the threading module for tracking lock allocation."""
diff --git a/releasenotes/notes/remove-wrapt-lock-profiler-c5f2e83097e0ff28.yaml b/releasenotes/notes/remove-wrapt-lock-profiler-c5f2e83097e0ff28.yaml
@@ -0,0 +1,4 @@
+---
+other:
+  - |
+    profiling: This removes the ``wrapt`` library dependency from the Lock Profiler implementation, improving performance and reducing overhead during lock instrumentation.
diff --git a/tests/profiling_v2/collector/test_threading.py b/tests/profiling_v2/collector/test_threading.py
@@ -1,6 +1,9 @@
+from __future__ import absolute_import
+
 import _thread
 import glob
 import os
+import sys
 import threading
 from typing import Callable
 from typing import List
@@ -88,94 +91,67 @@ def test_repr(
     test_collector._test_repr(collector_class, expected_repr)
 
 
-@pytest.mark.parametrize(
-    "lock_class,collector_class",
-    [
-        (threading.Lock, ThreadingLockCollector),
-        (threading.RLock, ThreadingRLockCollector),
-    ],
-)
-def test_patch(
-    lock_class: LockClassType,
-    collector_class: CollectorClassType,
-) -> None:
-    lock: LockClassType = lock_class
-    collector: ThreadingLockCollector | ThreadingRLockCollector = collector_class()
+def test_patch():
+    from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
+
+    lock = threading.Lock
+    collector = ThreadingLockCollector()
     collector.start()
     assert lock == collector._original
-    # wrapt makes this true
-    assert lock == lock_class
+    # After patching, threading.Lock is replaced with our wrapper
+    # The old reference (lock) points to the original builtin Lock class
+    assert lock != threading.Lock  # They're different after patching
+    assert isinstance(threading.Lock, _LockAllocatorWrapper)  # threading.Lock is now wrapped
+    assert callable(threading.Lock)  # and it's callable
     collector.stop()
-    assert lock == lock_class
-    assert collector._original == lock_class
-
-
-@pytest.mark.subprocess(
-    env=dict(WRAPT_DISABLE_EXTENSIONS="True", DD_PROFILING_FILE_PATH=__file__),
-)
-def test_wrapt_disable_extensions() -> None:
-    import os
-    import threading
-
-    from ddtrace.internal.datadog.profiling import ddup
-    from ddtrace.profiling.collector import _lock
-    from ddtrace.profiling.collector.threading import ThreadingLockCollector
-    from tests.profiling.collector import pprof_utils
-    from tests.profiling.collector.lock_utils import LineNo
-    from tests.profiling.collector.lock_utils import get_lock_linenos
-    from tests.profiling.collector.lock_utils import init_linenos
-    from tests.profiling.collector.pprof_utils import pprof_pb2
-
-    assert ddup.is_available, "ddup is not available"
-
-    # Set up the ddup exporter
-    test_name: str = "test_wrapt_disable_extensions"
-    pprof_prefix: str = "/tmp" + os.sep + test_name
-    output_filename: str = pprof_prefix + "." + str(os.getpid())
-    ddup.config(
-        env="test", service=test_name, version="my_version", output_filename=pprof_prefix
-    )  # pyright: ignore[reportCallIssue]
-    ddup.start()
-
-    init_linenos(os.environ["DD_PROFILING_FILE_PATH"])
-
-    # WRAPT_DISABLE_EXTENSIONS is a flag that can be set to disable the C extension
-    # for wrapt. It's not set by default in dd-trace-py, but it can be set by
-    # users. This test checks that the collector works even if the flag is set.
-    assert os.environ.get("WRAPT_DISABLE_EXTENSIONS")
-    assert _lock.WRAPT_C_EXT is False
-
-    with ThreadingLockCollector(capture_pct=100):
-        th_lock: threading.Lock = threading.Lock()  # !CREATE! test_wrapt_disable_extensions
-        with th_lock:  # !ACQUIRE! !RELEASE! test_wrapt_disable_extensions
-            pass
-
-    ddup.upload()  # pyright: ignore[reportCallIssue]
-
-    expected_filename: str = "test_threading.py"
-
-    linenos: LineNo = get_lock_linenos("test_wrapt_disable_extensions", with_stmt=True)
-
-    profile: pprof_pb2.Profile = pprof_utils.parse_newest_profile(output_filename)
-    pprof_utils.assert_lock_events(
-        profile,
-        expected_acquire_events=[
-            pprof_utils.LockAcquireEvent(
-                caller_name="<module>",
-                filename=expected_filename,
-                linenos=linenos,
-                lock_name="th_lock",
-            )
-        ],
-        expected_release_events=[
-            pprof_utils.LockReleaseEvent(
-                caller_name="<module>",
-                filename=expected_filename,
-                linenos=linenos,
-                lock_name="th_lock",
-            )
-        ],
-    )
+    # After stopping, everything is restored
+    assert lock == threading.Lock
+    assert collector._original == threading.Lock
+
+
+@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="only works on linux")
+@pytest.mark.subprocess(err=None)
+# For macOS: Could print 'Error uploading' but okay to ignore since we are checking if native_id is set
+def test_user_threads_have_native_id():
+    from os import getpid
+    from threading import Thread
+    from threading import _MainThread
+    from threading import current_thread
+    from time import sleep
+
+    from ddtrace.profiling import profiler
+
+    # DEV: We used to run this test with ddtrace_run=True passed into the
+    # subprocess decorator, but that caused this to be flaky for Python 3.8.x
+    # with gevent. When it failed for that specific venv, current_thread()
+    # returned a DummyThread instead of a _MainThread.
+    p = profiler.Profiler()
+    p.start()
+
+    main = current_thread()
+    assert isinstance(main, _MainThread)
+    # We expect the current thread to have the same ID as the PID
+    assert main.native_id == getpid(), (main.native_id, getpid())
+
+    t = Thread(target=lambda: None)
+    t.start()
+
+    for _ in range(10):
+        try:
+            # The TID should be higher than the PID, but not too high
+            assert 0 < t.native_id - getpid() < 100, (t.native_id, getpid())
+        except AttributeError:
+            # The native_id attribute is set by the thread so we might have to
+            # wait a bit for it to be set.
+            sleep(0.1)
+        else:
+            break
+    else:
+        raise AssertionError("Thread.native_id not set")
+
+    t.join()
+
+    p.stop()
 
 
 # This test has to be run in a subprocess because it calls gevent.monkey.patch_all()

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +---
 +other:
 +  - |
 +    profiling: This removes the ``wrapt`` library dependency from the Lock Profiler implementation, improving performance and reducing overhead during lock instrumentation.