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

ddtrace/profiling/collector/_lock.py

@@ -314,6 +314,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).
+
+        Without this, when custom lock types inherit from a wrapped lock
+        (e.g. neo4j's AsyncRLock that inherits from asyncio.Lock), the program would 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

@@ -51,6 +51,35 @@ def teardown_method(self):
             except Exception as e:
                 print("Error while deleting file: ", e)
 
+    async def test_subclassing_wrapped_lock(self) -> None:
+        """Test that subclassing of a wrapped lock type when profiling is active."""
+        from typing import Optional
+
+        from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
+
+        with collector_asyncio.AsyncioLockCollector(capture_pct=100):
+            assert isinstance(asyncio.Lock, _LockAllocatorWrapper)
+
+            # This should NOT raise TypeError
+            class CustomLock(asyncio.Lock):  # type: ignore[misc]
+                def __init__(self) -> None:
+                    super().__init__()
+                    self._owner: Optional[int] = None
+                    self._count: int = 0
+
+            # Verify subclassing and functionality
+            custom_lock: CustomLock = CustomLock()
+            assert hasattr(custom_lock, "_owner")
+            assert hasattr(custom_lock, "_count")
+            assert custom_lock._owner is None
+            assert custom_lock._count == 0
+
+            # Test async acquire/release
+            await custom_lock.acquire()
+            assert custom_lock.locked()
+            custom_lock.release()
+            assert not custom_lock.locked()
+
     async def test_asyncio_lock_events(self):
         with collector_asyncio.AsyncioLockCollector(capture_pct=100):
             lock = asyncio.Lock()  # !CREATE! test_asyncio_lock_events

tests/profiling/collector/test_threading.py

@@ -1378,6 +1378,32 @@ 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__.
+        """
+        from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
+
+        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__()
+                    self.custom_attr: str = "test"
+
+            # Verify subclassing and functionality
+            custom_lock: CustomLock = CustomLock()
+            assert hasattr(custom_lock, "custom_attr")
+            assert custom_lock.custom_attr == "test"
+
+            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.