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

Author: vlad-scherbich

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).
+
+        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."""

tests/profiling/collector/test_asyncio.py

@@ -51,6 +51,48 @@ def teardown_method(self):
             except Exception as e:
                 print("Error while deleting file: ", e)
 
+    def test_subclassing_wrapped_lock(self):
+        """Test that subclassing asyncio.Lock works when profiling is active.
+
+        Regression test for a bug reported by neo4j users where their library defines:
+            class AsyncRLock(asyncio.Lock): ...
+
+        With profiling enabled, this would fail with:
+            TypeError: _LockAllocatorWrapper.__init__() takes 2 positional arguments but 4 were given
+
+        The fix implements __mro_entries__ on _LockAllocatorWrapper (PEP 560).
+        """
+        from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
+
+        with collector_asyncio.AsyncioLockCollector(capture_pct=100):
+            # Verify the wrapper is in place
+            assert isinstance(asyncio.Lock, _LockAllocatorWrapper)
+
+            # This should NOT raise TypeError
+            class AsyncRLock(asyncio.Lock):
+                """A custom lock that extends asyncio.Lock (like neo4j does)."""
+
+                def __init__(self) -> None:
+                    super().__init__()
+                    self._owner = None
+                    self._count = 0
+
+            # Verify the subclass can be instantiated
+            custom_lock = AsyncRLock()
+            assert hasattr(custom_lock, "_owner")
+            assert hasattr(custom_lock, "_count")
+            assert custom_lock._owner is None
+            assert custom_lock._count == 0
+
+            # Verify async lock functionality still works
+            async def test_async_lock():
+                await custom_lock.acquire()
+                assert custom_lock.locked()
+                custom_lock.release()
+                assert not custom_lock.locked()
+
+            asyncio.get_event_loop().run_until_complete(test_async_lock())
+
     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

@@ -564,6 +564,38 @@ def __init__(self, lock_class: LockTypeClass) -> None:
             # Try this way too
             Foobar(self.lock_class)
 
+    def test_subclassing_wrapped_lock(self) -> None:
+        """Test that subclassing a wrapped lock type works correctly.
+
+        Regression test for a bug where third-party libraries (e.g. neo4j) that subclass
+        asyncio.Lock would fail with:
+            TypeError: _LockAllocatorWrapper.__init__() takes 2 positional arguments but 4 were given
+
+        The fix implements __mro_entries__ on _LockAllocatorWrapper (PEP 560).
+        """
+        from ddtrace.profiling.collector._lock import _LockAllocatorWrapper
+
+        with self.collector_class():
+            # Verify the wrapper is in place
+            assert isinstance(self.lock_class, _LockAllocatorWrapper)
+
+            # This should NOT raise TypeError
+            class CustomLock(self.lock_class):  # type: ignore[misc]
+                """A custom lock that extends the wrapped lock type."""
+
+                def __init__(self) -> None:
+                    super().__init__()
+                    self.custom_attr = "test"
+
+            # Verify the subclass works correctly
+            custom_lock = CustomLock()
+            assert hasattr(custom_lock, "custom_attr")
+            assert custom_lock.custom_attr == "test"
+
+            # Verify lock functionality still works
+            assert custom_lock.acquire()
+            custom_lock.release()
+
     def test_lock_events(self) -> None:
         # The first argument is the recorder.Recorder which is used for the
         # v1 exporter. We don't need it for the v2 exporter.