Skip to content

Commit 8b01f85

Browse files
feat(profiling): profile asyncio.BoundedSemaphore primitives with Python Lock profiler
1 parent 15c46fb commit 8b01f85

File tree

5 files changed

+67
-7
lines changed

5 files changed

+67
-7
lines changed

ddtrace/profiling/collector/asyncio.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,10 @@ class _ProfiledAsyncioSemaphore(_lock._ProfiledLock):
1111
pass
1212

1313

14+
class _ProfiledAsyncioBoundedSemaphore(_lock._ProfiledLock):
15+
pass
16+
17+
1418
class AsyncioLockCollector(_lock.LockCollector):
1519
"""Record asyncio.Lock usage."""
1620

@@ -25,3 +29,11 @@ class AsyncioSemaphoreCollector(_lock.LockCollector):
2529
PROFILED_LOCK_CLASS = _ProfiledAsyncioSemaphore
2630
MODULE = asyncio
2731
PATCHED_LOCK_NAME = "Semaphore"
32+
33+
34+
class AsyncioBoundedSemaphoreCollector(_lock.LockCollector):
35+
"""Record asyncio.BoundedSemaphore usage."""
36+
37+
PROFILED_LOCK_CLASS = _ProfiledAsyncioBoundedSemaphore
38+
MODULE = asyncio
39+
PATCHED_LOCK_NAME = "BoundedSemaphore"

ddtrace/profiling/profiler.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,7 @@ def start_collector(collector_class: Type[collector.Collector]) -> None:
219219
("threading", lambda _: start_collector(threading.ThreadingBoundedSemaphoreCollector)),
220220
("asyncio", lambda _: start_collector(asyncio.AsyncioLockCollector)),
221221
("asyncio", lambda _: start_collector(asyncio.AsyncioSemaphoreCollector)),
222+
("asyncio", lambda _: start_collector(asyncio.AsyncioBoundedSemaphoreCollector)),
222223
]
223224

224225
for module, hook in self._collectors_on_import:
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
---
2+
features:
3+
- |
4+
profiling: Add support for ``asyncio.BoundedSemaphore`` locking type profiling in Python.
5+
Unlike ``threading.BoundedSemaphore``, the asyncio version does not use internal locks,
6+
so no internal lock detection is needed.
7+

tests/profiling/collector/test_asyncio.py

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
from ddtrace import ext
1313
from ddtrace.internal.datadog.profiling import ddup
14+
from ddtrace.profiling.collector.asyncio import AsyncioBoundedSemaphoreCollector
1415
from ddtrace.profiling.collector.asyncio import AsyncioLockCollector
1516
from ddtrace.profiling.collector.asyncio import AsyncioSemaphoreCollector
1617
from tests.profiling.collector import pprof_utils
@@ -23,12 +24,16 @@
2324

2425
PY_311_OR_ABOVE = sys.version_info[:2] >= (3, 11)
2526

26-
# Type aliases for collector and lock types
27-
CollectorType = Union[
27+
# Type aliases for supported classes
28+
LockTypeClass = Union[Type[asyncio.Lock], Type[asyncio.Semaphore], Type[asyncio.BoundedSemaphore]]
29+
LockTypeInst = Union[asyncio.Lock, asyncio.Semaphore, asyncio.BoundedSemaphore]
30+
31+
CollectorTypeClass = Union[
2832
Type[AsyncioLockCollector],
2933
Type[AsyncioSemaphoreCollector],
34+
Type[AsyncioBoundedSemaphoreCollector],
3035
]
31-
LockType = Union[Type[asyncio.Lock], Type[asyncio.Semaphore]]
36+
CollectorTypeInst = Union[AsyncioLockCollector, AsyncioSemaphoreCollector, AsyncioBoundedSemaphoreCollector]
3237

3338

3439
@pytest.mark.parametrize(
@@ -42,9 +47,13 @@
4247
AsyncioSemaphoreCollector,
4348
"AsyncioSemaphoreCollector(status=<ServiceStatus.STOPPED: 'stopped'>, capture_pct=1.0, nframes=64, tracer=None)", # noqa: E501
4449
),
50+
(
51+
AsyncioBoundedSemaphoreCollector,
52+
"AsyncioBoundedSemaphoreCollector(status=<ServiceStatus.STOPPED: 'stopped'>, capture_pct=1.0, nframes=64, tracer=None)",
53+
),
4554
],
4655
)
47-
def test_collector_repr(collector_class: CollectorType, expected_repr: str) -> None:
56+
def test_collector_repr(collector_class: CollectorTypeClass, expected_repr: str) -> None:
4857
test_collector._test_repr(collector_class, expected_repr)
4958

5059

@@ -59,19 +68,19 @@ class BaseAsyncioLockCollectorTest:
5968
"""
6069

6170
@property
62-
def collector_class(self) -> CollectorType:
71+
def collector_class(self) -> CollectorTypeClass:
6372
raise NotImplementedError("Child classes must implement collector_class")
6473

6574
@property
66-
def lock_class(self) -> LockType:
75+
def lock_class(self) -> LockTypeClass:
6776
raise NotImplementedError("Child classes must implement lock_class")
6877

6978
@property
7079
def lock_init_args(self) -> tuple:
7180
"""Arguments to pass to lock constructor. Override for Semaphore-like locks."""
7281
return ()
7382

74-
def create_lock(self) -> Union[asyncio.Lock, asyncio.Semaphore]:
83+
def create_lock(self) -> LockTypeInst:
7584
"""Create a lock instance with the appropriate arguments."""
7685
return self.lock_class(*self.lock_init_args)
7786

@@ -242,3 +251,33 @@ def lock_class(self):
242251
@property
243252
def lock_init_args(self):
244253
return (2,) # Initial semaphore value
254+
255+
256+
class TestAsyncioBoundedSemaphoreCollector(BaseAsyncioLockCollectorTest):
257+
"""Test asyncio.BoundedSemaphore profiling."""
258+
259+
@property
260+
def collector_class(self):
261+
return AsyncioBoundedSemaphoreCollector
262+
263+
@property
264+
def lock_class(self):
265+
return asyncio.BoundedSemaphore
266+
267+
@property
268+
def lock_init_args(self):
269+
return (2,) # Initial semaphore value
270+
271+
async def test_bounded_behavior_preserved(self):
272+
"""Test that profiling wrapper preserves BoundedSemaphore's bounded behavior.
273+
274+
This verifies the wrapper doesn't interfere with BoundedSemaphore's unique characteristic:
275+
raising ValueError when releasing beyond the initial value.
276+
"""
277+
with self.collector_class(capture_pct=100):
278+
bs = asyncio.BoundedSemaphore(1)
279+
await bs.acquire()
280+
bs.release()
281+
# BoundedSemaphore should raise ValueError when releasing more than initial value
282+
with pytest.raises(ValueError, match="BoundedSemaphore released too many times"):
283+
bs.release()

tests/profiling/test_profiler.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,7 @@ def test_default_collectors():
139139
else:
140140
assert any(isinstance(c, asyncio.AsyncioLockCollector) for c in p._profiler._collectors)
141141
assert any(isinstance(c, asyncio.AsyncioSemaphoreCollector) for c in p._profiler._collectors)
142+
assert any(isinstance(c, asyncio.AsyncioBoundedSemaphoreCollector) for c in p._profiler._collectors)
142143
p.stop(flush=False)
143144

144145

0 commit comments

Comments
 (0)