diff --git a/ddtrace/profiling/collector/_lock.py b/ddtrace/profiling/collector/_lock.py index e7a2f32753f..d8158b88aab 100644 --- a/ddtrace/profiling/collector/_lock.py +++ b/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.""" diff --git a/releasenotes/notes/profiling-fix-lock-subclassing-00a61ebaa41ef3d0.yaml b/releasenotes/notes/profiling-fix-lock-subclassing-00a61ebaa41ef3d0.yaml new file mode 100644 index 00000000000..039338f19a8 --- /dev/null +++ b/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 diff --git a/tests/profiling/collector/test_asyncio.py b/tests/profiling/collector/test_asyncio.py index f72455678cf..8a65675854f 100644 --- a/tests/profiling/collector/test_asyncio.py +++ b/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 diff --git a/tests/profiling/collector/test_threading.py b/tests/profiling/collector/test_threading.py index 1360f21b68f..5a6f8b88ec9 100644 --- a/tests/profiling/collector/test_threading.py +++ b/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.