diff --git a/ddtrace/profiling/collector/asyncio.py b/ddtrace/profiling/collector/asyncio.py index 3c8e610da0f..115b4db36f6 100644 --- a/ddtrace/profiling/collector/asyncio.py +++ b/ddtrace/profiling/collector/asyncio.py @@ -1,7 +1,5 @@ -from asyncio.locks import Lock -import typing +import asyncio -from .. import collector from . import _lock @@ -9,25 +7,33 @@ class _ProfiledAsyncioLock(_lock._ProfiledLock): pass +class _ProfiledAsyncioSemaphore(_lock._ProfiledLock): + pass + + +class _ProfiledAsyncioBoundedSemaphore(_lock._ProfiledLock): + pass + + class AsyncioLockCollector(_lock.LockCollector): """Record asyncio.Lock usage.""" PROFILED_LOCK_CLASS = _ProfiledAsyncioLock + MODULE = asyncio + PATCHED_LOCK_NAME = "Lock" + + +class AsyncioSemaphoreCollector(_lock.LockCollector): + """Record asyncio.Semaphore usage.""" + + PROFILED_LOCK_CLASS = _ProfiledAsyncioSemaphore + MODULE = asyncio + PATCHED_LOCK_NAME = "Semaphore" + + +class AsyncioBoundedSemaphoreCollector(_lock.LockCollector): + """Record asyncio.BoundedSemaphore usage.""" - def _start_service(self) -> None: - """Start collecting lock usage.""" - try: - import asyncio - except ImportError as e: - raise collector.CollectorUnavailable(e) - self._asyncio_module = asyncio - return super(AsyncioLockCollector, self)._start_service() - - def _get_patch_target(self) -> typing.Type[Lock]: - return self._asyncio_module.Lock - - def _set_patch_target( - self, - value: typing.Any, - ) -> None: - self._asyncio_module.Lock = value # type: ignore[misc] + PROFILED_LOCK_CLASS = _ProfiledAsyncioBoundedSemaphore + MODULE = asyncio + PATCHED_LOCK_NAME = "BoundedSemaphore" diff --git a/ddtrace/profiling/profiler.py b/ddtrace/profiling/profiler.py index 7d5163421d7..e1cfbfc0acd 100644 --- a/ddtrace/profiling/profiler.py +++ b/ddtrace/profiling/profiler.py @@ -218,6 +218,8 @@ def start_collector(collector_class: Type[collector.Collector]) -> None: ("threading", lambda _: start_collector(threading.ThreadingSemaphoreCollector)), ("threading", lambda _: start_collector(threading.ThreadingBoundedSemaphoreCollector)), ("asyncio", lambda _: start_collector(asyncio.AsyncioLockCollector)), + ("asyncio", lambda _: start_collector(asyncio.AsyncioSemaphoreCollector)), + ("asyncio", lambda _: start_collector(asyncio.AsyncioBoundedSemaphoreCollector)), ] for module, hook in self._collectors_on_import: diff --git a/releasenotes/notes/Added-support-for-profiling-of-asyncio.BoundedSemaphore-objects-to-the-Python-Lock-profiler-5bf06391b0b047f1.yaml b/releasenotes/notes/Added-support-for-profiling-of-asyncio.BoundedSemaphore-objects-to-the-Python-Lock-profiler-5bf06391b0b047f1.yaml new file mode 100644 index 00000000000..29a4c99c873 --- /dev/null +++ b/releasenotes/notes/Added-support-for-profiling-of-asyncio.BoundedSemaphore-objects-to-the-Python-Lock-profiler-5bf06391b0b047f1.yaml @@ -0,0 +1,7 @@ +--- +features: + - | + profiling: Add support for ``asyncio.BoundedSemaphore`` locking type profiling in Python. + Unlike ``threading.BoundedSemaphore``, the asyncio version does not use internal locks, + so no internal lock detection is needed. + diff --git a/releasenotes/notes/Added-support-for-profiling-of-asyncio.Semaphore-objects-to-the-Python-Lock-profiler-193fd8150d084ea6.yaml b/releasenotes/notes/Added-support-for-profiling-of-asyncio.Semaphore-objects-to-the-Python-Lock-profiler-193fd8150d084ea6.yaml new file mode 100644 index 00000000000..9e02aba16d0 --- /dev/null +++ b/releasenotes/notes/Added-support-for-profiling-of-asyncio.Semaphore-objects-to-the-Python-Lock-profiler-193fd8150d084ea6.yaml @@ -0,0 +1,7 @@ +--- +features: + - | + profiling: Add support for ``asyncio.Semaphore`` locking type profiling in Python. + Also refactored asyncio collectors to use base class attributes (``MODULE`` and ``PATCHED_LOCK_NAME``) + for consistency with the threading module collectors. + diff --git a/tests/profiling/collector/test_asyncio.py b/tests/profiling/collector/test_asyncio.py index f72455678cf..e7386e73aa5 100644 --- a/tests/profiling/collector/test_asyncio.py +++ b/tests/profiling/collector/test_asyncio.py @@ -3,13 +3,17 @@ import glob import os import sys +from typing import Type +from typing import Union import uuid import pytest from ddtrace import ext from ddtrace.internal.datadog.profiling import ddup -from ddtrace.profiling.collector import asyncio as collector_asyncio +from ddtrace.profiling.collector.asyncio import AsyncioBoundedSemaphoreCollector +from ddtrace.profiling.collector.asyncio import AsyncioLockCollector +from ddtrace.profiling.collector.asyncio import AsyncioSemaphoreCollector from tests.profiling.collector import pprof_utils from tests.profiling.collector import test_collector from tests.profiling.collector.lock_utils import get_lock_linenos @@ -20,16 +24,56 @@ PY_311_OR_ABOVE = sys.version_info[:2] >= (3, 11) +# Type aliases for supported classes +LockTypeClass = Union[Type[asyncio.Lock], Type[asyncio.Semaphore], Type[asyncio.BoundedSemaphore]] +LockTypeInst = Union[asyncio.Lock, asyncio.Semaphore, asyncio.BoundedSemaphore] -def test_repr(): - test_collector._test_repr( - collector_asyncio.AsyncioLockCollector, - "AsyncioLockCollector(status=, capture_pct=1.0, nframes=64, tracer=None)", # noqa: E501 - ) +CollectorTypeClass = Union[ + Type[AsyncioLockCollector], + Type[AsyncioSemaphoreCollector], + Type[AsyncioBoundedSemaphoreCollector], +] +CollectorTypeInst = Union[AsyncioLockCollector, AsyncioSemaphoreCollector, AsyncioBoundedSemaphoreCollector] + + +@pytest.mark.parametrize( + "collector_class,expected_repr", + [ + ( + AsyncioLockCollector, + "AsyncioLockCollector(status=, capture_pct=1.0, nframes=64, tracer=None)", + ), + ( + AsyncioSemaphoreCollector, + "AsyncioSemaphoreCollector(status=, capture_pct=1.0, nframes=64, tracer=None)", # noqa: E501 + ), + ( + AsyncioBoundedSemaphoreCollector, + "AsyncioBoundedSemaphoreCollector(status=, capture_pct=1.0, nframes=64, tracer=None)", # noqa: E501 + ), + ], +) +def test_collector_repr(collector_class: CollectorTypeClass, expected_repr: str) -> None: + test_collector._test_repr(collector_class, expected_repr) @pytest.mark.asyncio -class TestAsyncioLockCollector: +class BaseAsyncioLockCollectorTest: + """Base test class for asyncio lock collectors. + + Child classes must implement: + - collector_class: The collector class to test + - lock_class: The asyncio lock class to test + """ + + @property + def collector_class(self) -> CollectorTypeClass: + raise NotImplementedError("Child classes must implement collector_class") + + @property + def lock_class(self) -> LockTypeClass: + raise NotImplementedError("Child classes must implement lock_class") + def setup_method(self, method): self.test_name = method.__qualname__ if PY_311_OR_ABOVE else method.__name__ self.output_prefix = "/tmp" + os.sep + self.test_name @@ -51,16 +95,17 @@ def teardown_method(self): except Exception as e: print("Error while deleting file: ", e) - async def test_asyncio_lock_events(self): - with collector_asyncio.AsyncioLockCollector(capture_pct=100): - lock = asyncio.Lock() # !CREATE! test_asyncio_lock_events - await lock.acquire() # !ACQUIRE! test_asyncio_lock_events + async def test_lock_events(self): + """Test basic acquire/release event profiling.""" + with self.collector_class(capture_pct=100): + lock = self.lock_class() # !CREATE! asyncio_test_lock_events + await lock.acquire() # !ACQUIRE! asyncio_test_lock_events assert lock.locked() - lock.release() # !RELEASE! test_asyncio_lock_events + lock.release() # !RELEASE! asyncio_test_lock_events ddup.upload() - linenos = get_lock_linenos("test_asyncio_lock_events") + linenos = get_lock_linenos("asyncio_test_lock_events") profile = pprof_utils.parse_newest_profile(self.output_filename) expected_thread_id = _thread.get_ident() pprof_utils.assert_lock_events( @@ -85,29 +130,30 @@ async def test_asyncio_lock_events(self): ], ) - async def test_asyncio_lock_events_tracer(self, tracer): + async def test_lock_events_tracer(self, tracer): + """Test event profiling with tracer integration.""" tracer._endpoint_call_counter_span_processor.enable() resource = str(uuid.uuid4()) span_type = ext.SpanTypes.WEB - with collector_asyncio.AsyncioLockCollector(capture_pct=100, tracer=tracer): - lock = asyncio.Lock() # !CREATE! test_asyncio_lock_events_tracer_1 - await lock.acquire() # !ACQUIRE! test_asyncio_lock_events_tracer_1 + with self.collector_class(capture_pct=100, tracer=tracer): + lock = self.lock_class() # !CREATE! asyncio_test_lock_events_tracer_1 + await lock.acquire() # !ACQUIRE! asyncio_test_lock_events_tracer_1 with tracer.trace("test", resource=resource, span_type=span_type) as t: - lock2 = asyncio.Lock() # !CREATE! test_asyncio_lock_events_tracer_2 - await lock2.acquire() # !ACQUIRE! test_asyncio_lock_events_tracer_2 - lock.release() # !RELEASE! test_asyncio_lock_events_tracer_1 + lock2 = self.lock_class() # !CREATE! asyncio_test_lock_events_tracer_2 + await lock2.acquire() # !ACQUIRE! asyncio_test_lock_events_tracer_2 + lock.release() # !RELEASE! asyncio_test_lock_events_tracer_1 span_id = t.span_id - lock2.release() # !RELEASE! test_asyncio_lock_events_tracer_2 + lock2.release() # !RELEASE! asyncio_test_lock_events_tracer_2 - lock_ctx = asyncio.Lock() # !CREATE! test_asyncio_lock_events_tracer_3 - async with lock_ctx: # !ACQUIRE! !RELEASE! test_asyncio_lock_events_tracer_3 + lock_ctx = self.lock_class() # !CREATE! asyncio_test_lock_events_tracer_3 + async with lock_ctx: # !ACQUIRE! !RELEASE! asyncio_test_lock_events_tracer_3 pass ddup.upload(tracer=tracer) - linenos_1 = get_lock_linenos("test_asyncio_lock_events_tracer_1") - linenos_2 = get_lock_linenos("test_asyncio_lock_events_tracer_2") - linenos_3 = get_lock_linenos("test_asyncio_lock_events_tracer_3", with_stmt=True) + linenos_1 = get_lock_linenos("asyncio_test_lock_events_tracer_1") + linenos_2 = get_lock_linenos("asyncio_test_lock_events_tracer_2") + linenos_3 = get_lock_linenos("asyncio_test_lock_events_tracer_3", with_stmt=True) profile = pprof_utils.parse_newest_profile(self.output_filename) expected_thread_id = _thread.get_ident() @@ -167,3 +213,53 @@ async def test_asyncio_lock_events_tracer(self, tracer): ), ], ) + + +class TestAsyncioLockCollector(BaseAsyncioLockCollectorTest): + """Test asyncio.Lock profiling.""" + + @property + def collector_class(self): + return AsyncioLockCollector + + @property + def lock_class(self): + return asyncio.Lock + + +class TestAsyncioSemaphoreCollector(BaseAsyncioLockCollectorTest): + """Test asyncio.Semaphore profiling.""" + + @property + def collector_class(self): + return AsyncioSemaphoreCollector + + @property + def lock_class(self): + return asyncio.Semaphore + + +class TestAsyncioBoundedSemaphoreCollector(BaseAsyncioLockCollectorTest): + """Test asyncio.BoundedSemaphore profiling.""" + + @property + def collector_class(self): + return AsyncioBoundedSemaphoreCollector + + @property + def lock_class(self): + return asyncio.BoundedSemaphore + + async def test_bounded_behavior_preserved(self): + """Test that profiling wrapper preserves BoundedSemaphore's bounded behavior. + + This verifies the wrapper doesn't interfere with BoundedSemaphore's unique characteristic: + raising ValueError when releasing beyond the initial value. + """ + with self.collector_class(capture_pct=100): + bs = asyncio.BoundedSemaphore(1) + await bs.acquire() + bs.release() + # BoundedSemaphore should raise ValueError when releasing more than initial value + with pytest.raises(ValueError, match="BoundedSemaphore released too many times"): + bs.release() diff --git a/tests/profiling/test_profiler.py b/tests/profiling/test_profiler.py index 09db302abe8..8224d8c0bd6 100644 --- a/tests/profiling/test_profiler.py +++ b/tests/profiling/test_profiler.py @@ -138,6 +138,8 @@ def test_default_collectors(): pass else: assert any(isinstance(c, asyncio.AsyncioLockCollector) for c in p._profiler._collectors) + assert any(isinstance(c, asyncio.AsyncioSemaphoreCollector) for c in p._profiler._collectors) + assert any(isinstance(c, asyncio.AsyncioBoundedSemaphoreCollector) for c in p._profiler._collectors) p.stop(flush=False)