fix(profiling): enable sub-classing wrapped lock types by custom lock implementations

ddtrace/profiling/collector/_lock.py

@@ -313,6 +313,17 @@ def __getattr__(self, name: str) -> Any:
             return getattr(original_class, name)
         raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
 
+    def __mro_entries__(self, bases: Tuple[Any, ...]) -> Tuple[Type[Any], ...]:
+        """Support subclassing the wrapped lock type (PEP 560).
+
+        When custom lock types inherit from a wrapped lock
+        (e.g. neo4j's AsyncRLock that inherits from asyncio.Lock), program error with:
+        > TypeError: _LockAllocatorWrapper.__init__() takes 2 positional arguments but 4 were given
+
+        This method returns the actual object type to be used as the base class.
+        """
+        return (self._original_class,)  # type: ignore[return-value]
+
 
 class LockCollector(collector.CaptureSamplerCollector):
     """Record lock usage."""

releasenotes/notes/profiling-fix-lock-subclassing-00a61ebaa41ef3d0.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    profiling: Fixes a bug where code that sub-classes our wrapped locks crashes with ``TypeError`` during profiling.
+    One example of this is neo4j's ``AsyncRLock``, which inherits from ``asyncio.Lock``:
+    https://github.com/neo4j/neo4j-python-driver/blob/6.x/src/neo4j/_async_compat/concurrency.py#L45

tests/profiling/collector/test_asyncio.py

@@ -19,6 +19,7 @@
 from tests.profiling.collector import test_collector
 from tests.profiling.collector.lock_utils import get_lock_linenos
 from tests.profiling.collector.lock_utils import init_linenos
+from tests.profiling.collector.test_threading import LockAllocatorWrapper
 
 
 init_linenos(__file__)
@@ -208,6 +209,23 @@ async def test_lock_events_tracer(self, tracer: Any) -> None:
             ],
         )
 
+    async def test_subclassing_wrapped_lock(self) -> None:
+        """Test that subclassing of a wrapped lock type when profiling is active."""
+        with self.collector_class(capture_pct=100):
+            assert isinstance(self.lock_class, LockAllocatorWrapper)
+
+            # This should NOT raise TypeError
+            class CustomLock(self.lock_class):  # type: ignore[misc]
+                def __init__(self) -> None:
+                    super().__init__()
+
+            # Verify subclassing and functionality
+            custom_lock: CustomLock = CustomLock()
+            await custom_lock.acquire()
+            assert custom_lock.locked()
+            custom_lock.release()
+            assert not custom_lock.locked()
+
 
 class TestAsyncioLockCollector(BaseAsyncioLockCollectorTest):
     """Test asyncio.Lock profiling."""

tests/profiling/collector/test_threading.py

@@ -21,6 +21,7 @@
 from ddtrace._trace.span import Span
 from ddtrace._trace.tracer import Tracer
 from ddtrace.internal.datadog.profiling import ddup
+from ddtrace.profiling.collector._lock import _LockAllocatorWrapper as LockAllocatorWrapper
 from ddtrace.profiling.collector._lock import _ProfiledLock
 from ddtrace.profiling.collector.threading import ThreadingBoundedSemaphoreCollector
 from ddtrace.profiling.collector.threading import ThreadingLockCollector
@@ -173,16 +174,14 @@ def test_lock_repr(
 
 
 def test_patch():
-    from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
-
     lock = threading.Lock
     collector = ThreadingLockCollector()
     collector.start()
     assert lock == collector._original_lock
     # 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 isinstance(threading.Lock, LockAllocatorWrapper)  # threading.Lock is now wrapped
     assert callable(threading.Lock)  # and it's callable
     collector.stop()
     # After stopping, everything is restored
@@ -1389,6 +1388,26 @@ class BaseSemaphoreTest(BaseThreadingLockCollectorTest):
     particularly those related to internal lock detection and Condition-based implementation.
     """
 
+    def test_subclassing_wrapped_lock(self) -> None:
+        """Test that subclassing of a wrapped lock type works when profiling is active.
+
+        This test is only valid for Semaphore-like types (pure Python classes).
+        threading.Lock and threading.RLock are C types that don't support subclassing
+        through __mro_entries__.
+        """
+        with self.collector_class():
+            assert isinstance(self.lock_class, LockAllocatorWrapper)
+
+            # This should NOT raise TypeError
+            class CustomLock(self.lock_class):  # type: ignore[misc]
+                def __init__(self) -> None:
+                    super().__init__()
+
+            # Verify subclassing and functionality
+            custom_lock: CustomLock = CustomLock()
+            assert custom_lock.acquire()
+            custom_lock.release()
+
     def _verify_no_double_counting(self, marker_name: str, lock_var_name: str) -> None:
         """Helper method to verify no double-counting in profile output.