diff --git a/ddtrace/profiling/collector/_lock.py b/ddtrace/profiling/collector/_lock.py index f0e75f7a161..228d8624042 100644 --- a/ddtrace/profiling/collector/_lock.py +++ b/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.""" 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 aae9e845968..f7397be53ff 100644 --- a/tests/profiling/collector/test_asyncio.py +++ b/tests/profiling/collector/test_asyncio.py @@ -15,6 +15,7 @@ from ddtrace import ext from ddtrace.internal.datadog.profiling import ddup +from ddtrace.profiling.collector._lock import _LockAllocatorWrapper as LockAllocatorWrapper from ddtrace.profiling.collector.asyncio import AsyncioBoundedSemaphoreCollector from ddtrace.profiling.collector.asyncio import AsyncioLockCollector from ddtrace.profiling.collector.asyncio import AsyncioSemaphoreCollector @@ -214,6 +215,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.""" diff --git a/tests/profiling/collector/test_threading.py b/tests/profiling/collector/test_threading.py index ff42782c5ea..80f3fbb2935 100644 --- a/tests/profiling/collector/test_threading.py +++ b/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 @@ -163,8 +164,6 @@ def test_lock_repr( def test_patch(): - from ddtrace.profiling.collector._lock import _LockAllocatorWrapper - lock = threading.Lock collector = ThreadingLockCollector() collector.start() @@ -172,7 +171,7 @@ def test_patch(): # 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 @@ -1379,6 +1378,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.