Tracing: wait on an event instead of polling in the batch worker

src/agents/tracing/processors.py

@@ -568,6 +568,9 @@ def __init__(
         self._max_batch_size = max_batch_size
         self._schedule_delay = schedule_delay
         self._shutdown_event = threading.Event()
+        # Set when the worker has something to do before its next scheduled export: the queue
+        # reached the trigger size, or we are shutting down.
+        self._wake_event = threading.Event()
 
         # The queue size threshold at which we export immediately.
         self._export_trigger_size = max(1, int(max_queue_size * export_trigger_ratio))
@@ -602,6 +605,8 @@ def on_trace_start(self, trace: Trace) -> None:
             self._queue.put_nowait(trace)
         except queue.Full:
             logger.warning("Queue is full, dropping trace.")
+            return
+        self._wake_worker_if_queue_is_full_enough()
 
     def on_trace_end(self, trace: Trace) -> None:
         # We send traces via on_trace_start, so we don't need to do anything here.
@@ -619,12 +624,20 @@ def on_span_end(self, span: Span[Any]) -> None:
             self._queue.put_nowait(span)
         except queue.Full:
             logger.warning("Queue is full, dropping span.")
+            return
+        self._wake_worker_if_queue_is_full_enough()
+
+    def _wake_worker_if_queue_is_full_enough(self) -> None:
+        """Wake the worker early when the queue reaches the export trigger size."""
+        if self._queue.qsize() >= self._export_trigger_size:
+            self._wake_event.set()
 
     def shutdown(self, timeout: float | None = None):
         """
         Called when the application stops. We signal our thread to stop, then join it.
         """
         self._shutdown_event.set()
+        self._wake_event.set()
         if timeout is not None:
             request_exporter_shutdown = getattr(self._exporter, "_request_shutdown", None)
             if callable(request_exporter_shutdown):
@@ -660,9 +673,14 @@ def _run(self):
                 self._export_batches()
                 # Reset the next scheduled flush time
                 self._next_export_time = time.monotonic() + self._schedule_delay
-            else:
-                # Sleep a short interval so we don't busy-wait.
-                time.sleep(0.2)
+                continue
+
+            # Block until the next scheduled flush rather than polling. Producers set the
+            # event once the queue reaches the trigger size and shutdown() sets it to break
+            # out immediately, so an idle process wakes once per schedule_delay instead of
+            # five times a second.
+            self._wake_event.wait(max(0.0, self._next_export_time - time.monotonic()))
+            self._wake_event.clear()
 
         # Final drain after shutdown
         self._export_batches(deadline=self._shutdown_deadline)

tests/test_trace_processor.py

@@ -282,6 +282,52 @@ def export(self, items: list[Trace | Span[Any]]) -> None:
     assert exporter.call_count >= 3
 
 
+def test_batch_trace_processor_queue_trigger_wakes_the_idle_worker() -> None:
+    """The worker parks until its next scheduled export, so producers have to wake it."""
+    exported = threading.Event()
+
+    class SignalingExporter(TracingExporter):
+        def export(self, items: list[Trace | Span[Any]]) -> None:
+            exported.set()
+
+    processor = BatchTraceProcessor(
+        exporter=SignalingExporter(),
+        max_queue_size=4,
+        schedule_delay=60.0,
+        export_trigger_ratio=0.5,
+    )
+
+    processor.on_span_end(get_span(processor))
+    # One span is below the two-item trigger, so the worker has no reason to wake yet.
+    assert not exported.wait(timeout=0.2)
+
+    processor.on_span_end(get_span(processor))
+    assert exported.wait(timeout=2.0), "reaching the trigger size must wake the worker"
+
+    processor.shutdown(timeout=2.0)
+
+
+def test_batch_trace_processor_shutdown_wakes_the_idle_worker() -> None:
+    """shutdown() must not wait out the schedule delay before the final drain."""
+    exported: list[Trace | Span[Any]] = []
+
+    class RecordingExporter(TracingExporter):
+        def export(self, items: list[Trace | Span[Any]]) -> None:
+            exported.extend(items)
+
+    processor = BatchTraceProcessor(exporter=RecordingExporter(), schedule_delay=60.0)
+    processor.on_span_end(get_span(processor))
+
+    start = time.monotonic()
+    processor.shutdown(timeout=2.0)
+    elapsed = time.monotonic() - start
+
+    assert elapsed < 1.0
+    assert processor._worker_thread is not None
+    assert not processor._worker_thread.is_alive()
+    assert len(exported) == 1
+
+
 @pytest.mark.parametrize(
     ("adjusted_wall_time", "adjusted_monotonic_time", "expected_scheduled_exports"),
     [
@@ -300,7 +346,6 @@ class ControlledTime:
         def __init__(self) -> None:
             self.wall_time = 1000.0
             self.monotonic_time = 100.0
-            self.sleep_calls = 0
 
         def time(self) -> float:
             return self.wall_time
@@ -309,12 +354,25 @@ def monotonic(self) -> float:
             return self.monotonic_time
 
         def sleep(self, _seconds: float) -> None:
-            self.sleep_calls += 1
-            if self.sleep_calls == 1:
-                self.wall_time = adjusted_wall_time
-                self.monotonic_time = adjusted_monotonic_time
+            raise AssertionError("the worker must wait on its event, not poll with sleep")
+
+    class ControlledWakeEvent:
+        """Stands in for the worker's wake event and drives the loop from its waits."""
+
+        def __init__(self) -> None:
+            self.wait_calls = 0
+
+        def wait(self, _timeout: float | None = None) -> bool:
+            self.wait_calls += 1
+            if self.wait_calls == 1:
+                controlled_time.wall_time = adjusted_wall_time
+                controlled_time.monotonic_time = adjusted_monotonic_time
             else:
                 processor._shutdown_event.set()
+            return False
+
+        def clear(self) -> None:
+            pass
 
     controlled_time = ControlledTime()
     monkeypatch.setattr("agents.tracing.processors.time", controlled_time)
@@ -324,6 +382,7 @@ def sleep(self, _seconds: float) -> None:
         schedule_delay=1.0,
         export_trigger_ratio=1.0,
     )
+    monkeypatch.setattr(processor, "_wake_event", ControlledWakeEvent())
     processor._queue.put_nowait(get_span(processor))
     scheduled_export = object()
     export_deadlines: list[float | None | object] = []