feature(profiling): add support for direct profiling of threading.Semaphore objects

Author: vlad-scherbich

ddtrace/profiling/collector/threading.py

@@ -18,6 +18,10 @@ class _ProfiledThreadingRLock(_lock._ProfiledLock):
     pass
 
 
+class _ProfiledThreadingSemaphore(_lock._ProfiledLock):
+    pass
+
+
 class ThreadingLockCollector(_lock.LockCollector):
     """Record threading.Lock usage."""
 
@@ -48,6 +52,21 @@ def _set_patch_target(
         threading.RLock = value
 
 
+class ThreadingSemaphoreCollector(_lock.LockCollector):
+    """Record threading.Semaphore usage."""
+
+    PROFILED_LOCK_CLASS = _ProfiledThreadingSemaphore
+
+    def _get_patch_target(self) -> typing.Type[threading.Semaphore]:
+        return threading.Semaphore
+
+    def _set_patch_target(
+        self,
+        value: typing.Any,
+    ) -> None:
+        threading.Semaphore = value
+
+
 # Also patch threading.Thread so echion can track thread lifetimes
 def init_stack_v2() -> None:
     if config.stack.enabled and stack_v2.is_available:

tests/profiling_v2/collector/test_threading.py

@@ -21,6 +21,7 @@
 from ddtrace.internal.datadog.profiling import ddup
 from ddtrace.profiling.collector.threading import ThreadingLockCollector
 from ddtrace.profiling.collector.threading import ThreadingRLockCollector
+from ddtrace.profiling.collector.threading import ThreadingSemaphoreCollector
 from tests.profiling.collector import pprof_utils
 from tests.profiling.collector import test_collector
 from tests.profiling.collector.lock_utils import LineNo
@@ -30,11 +31,12 @@
 
 
 # Type aliases for supported classes
-LockClassType = Union[Type[threading.Lock], Type[threading.RLock]]
-CollectorClassType = Union[Type[ThreadingLockCollector], Type[ThreadingRLockCollector]]
+LockClassType = Union[Type[threading.Lock], Type[threading.RLock], Type[threading.Semaphore]]
+CollectorClassType = Union[Type[ThreadingLockCollector], Type[ThreadingRLockCollector], Type[ThreadingSemaphoreCollector]]
 # threading.Lock and threading.RLock are factory functions that return _thread types.
 # We reference the underlying _thread types directly to avoid creating instances at import time.
-LockClassInst = Union[_thread.LockType, _thread.RLock]
+# threading.Semaphore is a Python class, not a factory function.
+LockClassInst = Union[_thread.LockType, _thread.RLock, threading.Semaphore]
 
 # Module-level globals for testing global lock profiling
 _test_global_lock: LockClassInst
@@ -1328,3 +1330,222 @@ def test_lock_getattr(self) -> None:
             # After releasing, it should not be owned
             lock.release()
             assert not lock._is_owned()
+
+
+class TestThreadingSemaphoreCollector(BaseThreadingLockCollectorTest):
+    """Test Semaphore profiling"""
+
+    @property
+    def collector_class(self) -> Type[ThreadingSemaphoreCollector]:
+        return ThreadingSemaphoreCollector
+
+    @property
+    def lock_class(self) -> Type[threading.Semaphore]:
+        return threading.Semaphore
+
+    def test_semaphore_with_value(self) -> None:
+        """Test that Semaphore works with different initial values."""
+        with self.collector_class(capture_pct=100):
+            from ddtrace.profiling.collector._lock import _ProfiledLock
+
+            # Test with value=1
+            sem1 = self.lock_class(1)
+            assert isinstance(sem1, _ProfiledLock)
+            assert sem1.acquire(timeout=1)
+            assert not sem1.acquire(timeout=0.01)  # Should block
+            sem1.release()
+
+            # Test with value=3
+            sem3 = self.lock_class(3)
+            assert isinstance(sem3, _ProfiledLock)
+            for i in range(3):
+                assert sem3.acquire(timeout=1), f"Acquire {i+1} failed"
+            assert not sem3.acquire(timeout=0.01)  # Should block at 4th
+            for i in range(3):
+                sem3.release()
+
+            # Test with default value (1)
+            sem_default = self.lock_class()
+            assert isinstance(sem_default, _ProfiledLock)
+            assert sem_default.acquire(timeout=1)
+            assert not sem_default.acquire(timeout=0.01)  # Should block
+            sem_default.release()
+
+    def test_semaphore_multiple_acquires(self) -> None:
+        """Test that Semaphore correctly handles multiple acquires."""
+        with self.collector_class(capture_pct=100):
+            from ddtrace.profiling.collector._lock import _ProfiledLock
+
+            sem = self.lock_class(2)
+            assert isinstance(sem, _ProfiledLock)
+
+            # Should be able to acquire twice
+            assert sem.acquire(timeout=1)
+            assert sem.acquire(timeout=1)
+
+            # Third acquire should fail (timeout)
+            assert not sem.acquire(timeout=0.01)
+
+            # Release one and try again
+            sem.release()
+            assert sem.acquire(timeout=1)
+
+            # Clean up
+            sem.release()
+            sem.release()
+
+    def test_semaphore_non_blocking_acquire(self) -> None:
+        """Test non-blocking acquire behavior."""
+        with self.collector_class(capture_pct=100):
+            from ddtrace.profiling.collector._lock import _ProfiledLock
+
+            sem = self.lock_class(2)
+            assert isinstance(sem, _ProfiledLock)
+
+            # Non-blocking acquires should succeed immediately
+            assert sem.acquire(blocking=False)
+            assert sem.acquire(blocking=False)
+
+            # Third should fail immediately (not timeout)
+            assert not sem.acquire(blocking=False)
+
+            # Clean up
+            sem.release()
+            sem.release()
+
+    def test_semaphore_concurrent_threads(self) -> None:
+        """Test that multiple threads can hold semaphore simultaneously."""
+        import time
+
+        with self.collector_class(capture_pct=100):
+            from ddtrace.profiling.collector._lock import _ProfiledLock
+
+            sem = self.lock_class(3)  # Allow 3 threads
+            assert isinstance(sem, _ProfiledLock)
+
+            results = []
+            threads_holding = []
+
+            def worker(worker_id):
+                # Acquire the semaphore
+                if sem.acquire(timeout=2):
+                    threads_holding.append(worker_id)
+                    results.append(f"worker-{worker_id}-acquired")
+                    time.sleep(0.05)  # Hold it briefly
+                    results.append(f"worker-{worker_id}-releasing")
+                    sem.release()
+                else:
+                    results.append(f"worker-{worker_id}-timeout")
+
+            # Start 5 threads, but only 3 can hold semaphore at once
+            threads = []
+            for i in range(5):
+                t = threading.Thread(target=worker, args=(i,))
+                threads.append(t)
+                t.start()
+
+            # Wait for all threads
+            for t in threads:
+                t.join(timeout=5)
+
+            # All workers should have acquired and released
+            assert len([r for r in results if "acquired" in r]) == 5
+            assert len([r for r in results if "releasing" in r]) == 5
+
+            # At most 3 threads should have been holding at once
+            # (This is approximate since we're just checking they all got through)
+
+    def test_semaphore_blocking_contention(self) -> None:
+        """Test that threads block when semaphore is at capacity."""
+        import time
+
+        with self.collector_class(capture_pct=100):
+            from ddtrace.profiling.collector._lock import _ProfiledLock
+
+            sem = self.lock_class(1)  # Only 1 thread allowed
+            assert isinstance(sem, _ProfiledLock)
+
+            acquired_times = []
+            released_times = []
+
+            def holder():
+                """Hold the semaphore for a while"""
+                sem.acquire()
+                acquired_times.append(time.time())
+                time.sleep(0.1)  # Hold for 100ms
+                released_times.append(time.time())
+                sem.release()
+
+            def waiter():
+                """Try to acquire - should wait"""
+                start = time.time()
+                sem.acquire(timeout=1)
+                acquired_times.append(time.time())
+                wait_time = time.time() - start
+                sem.release()
+                return wait_time
+
+            # Start holder thread
+            holder_thread = threading.Thread(target=holder)
+            holder_thread.start()
+            time.sleep(0.01)  # Ensure holder gets it first
+
+            # Start waiter thread - should block
+            waiter_result = []
+            waiter_thread = threading.Thread(target=lambda: waiter_result.append(waiter()))
+            waiter_thread.start()
+
+            # Wait for both
+            holder_thread.join(timeout=2)
+            waiter_thread.join(timeout=2)
+
+            # Waiter should have waited for holder to release
+            assert len(waiter_result) == 1
+            assert waiter_result[0] >= 0.08  # Should have waited ~100ms
+
+    def test_semaphore_zero_value(self) -> None:
+        """Test semaphore with value=0 (initially blocking)."""
+        import time
+
+        with self.collector_class(capture_pct=100):
+            from ddtrace.profiling.collector._lock import _ProfiledLock
+
+            sem = self.lock_class(0)  # No permits available
+            assert isinstance(sem, _ProfiledLock)
+
+            # Immediate acquire should fail
+            assert not sem.acquire(blocking=False)
+            assert not sem.acquire(timeout=0.01)
+
+            # Release to make one available
+            sem.release()
+
+            # Now acquire should work
+            assert sem.acquire(timeout=1)
+            sem.release()
+
+    def test_semaphore_bounded_semaphore(self) -> None:
+        """Test that BoundedSemaphore gets wrapped and prevents over-release."""
+        with self.collector_class(capture_pct=100):
+            from ddtrace.profiling.collector._lock import _ProfiledLock
+
+            # BoundedSemaphore might be wrapped if it inherits from Semaphore patching
+            bounded_sem = threading.BoundedSemaphore(2)
+
+            # Check if it's wrapped (may not be if BoundedSemaphore needs separate collector)
+            if isinstance(bounded_sem, _ProfiledLock):
+                # Test basic functionality
+                assert bounded_sem.acquire(timeout=1)
+                bounded_sem.release()
+
+                # BoundedSemaphore should raise ValueError on over-release
+                bounded_sem.acquire()
+                bounded_sem.release()
+                try:
+                    bounded_sem.release()  # Over-release
+                    assert False, "Should have raised ValueError"
+                except ValueError:
+                    pass  # Expected
+            else:
+                # Document that BoundedSemaphore is not currently wrapped
+                print("Note: BoundedSemaphore not wrapped, needs separate collector")