From daf319cc57979161d30cb580c0914a807b334ccc Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Tue, 3 Nov 2026 16:36:04 +0000 Subject: [PATCH] fix(runtime): bound the tracer flush wait on streaming responses Since 1.25.2, sync-generator streams (including async-generator entrypoints, which are bridged to a sync generator) end with an uncapped _flush_tracer_provider(). OTel SDK 1.45's BatchProcessor.force_flush ignores timeout_millis, and AgentCore Runtime's OTEL_EXPORTER_OTLP_TIMEOUT=5000 is read as seconds, so with no route to the OTLP endpoint the response streams every event and then never closes. Run force_flush on a daemon thread and stop waiting after the cap, lowered from 30 s to 5 s. An overrunning flush keeps _flush_lock until it returns, so later calls skip instead of stacking threads. This also frees the async path's executor thread at the cap, so a stuck flush no longer blocks event loop shutdown. Streaming responses no longer flush in _handle_invocation before they are returned: the stream wrapper flushes when the stream ends, and the extra flush could only hold back the response headers. --- src/bedrock_agentcore/runtime/app.py | 7 +- src/bedrock_agentcore/runtime/tracing.py | 35 ++++- tests/bedrock_agentcore/runtime/test_app.py | 141 +++++++++++++++++- .../bedrock_agentcore/runtime/test_tracing.py | 47 ++++++ 4 files changed, 221 insertions(+), 9 deletions(-) diff --git a/src/bedrock_agentcore/runtime/app.py b/src/bedrock_agentcore/runtime/app.py index efae2faa..10afdc2c 100644 --- a/src/bedrock_agentcore/runtime/app.py +++ b/src/bedrock_agentcore/runtime/app.py @@ -542,6 +542,8 @@ async def _handle_invocation(self, request): request_context = self._build_request_context(request) start_time = time.time() + # Streaming responses flush when the stream ends; flushing here as well would only delay the headers. + streaming = False try: payload = await request.json() @@ -578,9 +580,11 @@ async def _handle_invocation(self, request): duration = time.time() - start_time if inspect.isgenerator(result): self.logger.info("Returning streaming response (generator) (%.3fs)", duration) + streaming = True return StreamingResponse(self._sync_stream_with_error_handling(result), media_type="text/event-stream") elif inspect.isasyncgen(result): self.logger.info("Returning streaming response (async generator) (%.3fs)", duration) + streaming = True return StreamingResponse(self._stream_with_error_handling(result), media_type="text/event-stream") # If handler returned a Starlette Response directly, pass it through. @@ -614,7 +618,8 @@ async def _handle_invocation(self, request): self.logger.exception("Invocation failed (%.3fs)", duration) return JSONResponse({"error": str(e)}, status_code=500) finally: - await _flush_tracer_provider_async() + if not streaming: + await _flush_tracer_provider_async() def _handle_ping(self, request): try: diff --git a/src/bedrock_agentcore/runtime/tracing.py b/src/bedrock_agentcore/runtime/tracing.py index 449c13ee..5c41e75f 100644 --- a/src/bedrock_agentcore/runtime/tracing.py +++ b/src/bedrock_agentcore/runtime/tracing.py @@ -35,6 +35,11 @@ _registration_lock = threading.Lock() _registered_on: Optional[object] = None # the TracerProvider instance we last registered on _flush_lock = threading.Lock() +# How long a response waits for buffered spans to export. A healthy OTLP export takes tens of milliseconds; +# 5 s leaves wide headroom, matches the export timeout AgentCore Runtime intends to set +# (OTEL_EXPORTER_OTLP_TIMEOUT=5000 ms), and keeps a stuck exporter from delaying the end of a response +# by more than a few seconds. Spans the flush misses stay queued for the background exporter. +_FLUSH_TIMEOUT_MILLIS = 5000 def _ensure_baggage_processor_registered() -> None: @@ -65,12 +70,36 @@ def _ensure_baggage_processor_registered() -> None: logger.debug("Could not register BaggageSpanProcessor", exc_info=True) -def _flush_tracer_provider(timeout_millis: int = 30000) -> None: - """Flush buffered spans before AgentCore Runtime freezes the microVM.""" +def _flush_tracer_provider(timeout_millis: int = _FLUSH_TIMEOUT_MILLIS) -> None: + """Flush buffered spans before AgentCore Runtime freezes the microVM, waiting at most ``timeout_millis``. + + ``force_flush`` cannot be trusted to honor its timeout (OTel SDK 1.45's ``BatchProcessor`` ignores it and + exports until the queue drains, which never ends if the OTLP endpoint is unreachable), so it runs on a + daemon thread that the caller stops waiting for. A flush that overruns keeps ``_flush_lock`` until it + returns, so later calls skip instead of piling up threads behind it. + """ if not _flush_lock.acquire(blocking=False): logger.debug("Tracer provider flush already in progress; skipping duplicate flush") return + try: + flush_thread = threading.Thread( + target=_force_flush_and_release, args=(timeout_millis,), name="agentcore-trace-flush", daemon=True + ) + flush_thread.start() + except Exception: + _flush_lock.release() + logger.debug("Could not start tracer provider flush", exc_info=True) + return + + flush_thread.join(timeout_millis / 1000) + if flush_thread.is_alive(): + logger.warning( + "Tracer provider flush did not finish within %d ms; continuing without waiting for it", timeout_millis + ) + + +def _force_flush_and_release(timeout_millis: int) -> None: try: from opentelemetry import trace @@ -83,7 +112,7 @@ def _flush_tracer_provider(timeout_millis: int = 30000) -> None: _flush_lock.release() -async def _flush_tracer_provider_async(timeout_millis: int = 30000) -> None: +async def _flush_tracer_provider_async(timeout_millis: int = _FLUSH_TIMEOUT_MILLIS) -> None: """Flush buffered spans without blocking the event loop or response indefinitely.""" try: await asyncio.wait_for( diff --git a/tests/bedrock_agentcore/runtime/test_app.py b/tests/bedrock_agentcore/runtime/test_app.py index 7ce404c6..e6315915 100644 --- a/tests/bedrock_agentcore/runtime/test_app.py +++ b/tests/bedrock_agentcore/runtime/test_app.py @@ -2,6 +2,7 @@ import concurrent.futures import contextlib import contextvars +import functools import json import os import threading @@ -14,6 +15,7 @@ from starlette.testclient import TestClient import bedrock_agentcore.runtime.app as _app_module +import bedrock_agentcore.runtime.tracing as _tracing_module from bedrock_agentcore.config_bundle.bundle import ConfigBundleRef from bedrock_agentcore.runtime import BedrockAgentCoreApp from bedrock_agentcore.runtime.app import _parse_runtime_arn @@ -21,6 +23,25 @@ from bedrock_agentcore.runtime.models import is_forwardable_header +class _TimeoutIgnoringTracerProvider: + """Tracer provider whose ``force_flush`` ignores ``timeout_millis``, like OTel SDK 1.45's BatchProcessor. + + With nothing queued, a flush returns at once. Once the handler has produced spans, a flush blocks the way + an export to an unreachable OTLP endpoint does, until ``release`` is set. + """ + + def __init__(self): + self.spans_pending = False + self.blocked_flushes = 0 + self.release = threading.Event() + + def force_flush(self, timeout_millis=30000): + if self.spans_pending: + self.blocked_flushes += 1 + self.release.wait(timeout=10) + return True + + class TestBedrockAgentCoreApp: def test_bedrock_agentcore_initialization(self): """Test BedrockAgentCoreApp initializes with correct name and routes.""" @@ -166,7 +187,7 @@ def handler(payload): mock_flush.assert_called_once_with() def test_sync_stream_flushes_after_generator_finishes(self): - """Sync streams flush once when prepared and again after their final chunk.""" + """Sync streams flush only after their final chunk, not before the response starts.""" bedrock_agentcore = BedrockAgentCoreApp() @bedrock_agentcore.entrypoint @@ -182,11 +203,11 @@ def handler(payload): response = TestClient(bedrock_agentcore).post("/invocations", json={}) assert response.status_code == 200 - mock_async_flush.assert_awaited_once_with() + mock_async_flush.assert_not_awaited() mock_sync_flush.assert_called_once_with() - def test_async_stream_flushes_after_generator_finishes(self): - """Async streams flush once when prepared and again after their final chunk.""" + def test_async_generator_entrypoint_flushes_after_stream_finishes(self): + """Async-generator entrypoints are bridged to a sync stream, which flushes once at the end.""" bedrock_agentcore = BedrockAgentCoreApp() @bedrock_agentcore.entrypoint @@ -202,9 +223,119 @@ async def handler(payload): response = TestClient(bedrock_agentcore).post("/invocations", json={}) assert response.status_code == 200 - mock_async_flush.assert_awaited_once_with() + mock_async_flush.assert_not_awaited() mock_sync_flush.assert_called_once_with() + def test_returned_async_generator_flushes_after_stream_finishes(self): + """A handler that returns an async generator streams on the async path, which flushes once at the end.""" + bedrock_agentcore = BedrockAgentCoreApp() + + async def stream(): + yield {"ok": True} + + @bedrock_agentcore.entrypoint + def handler(payload): + return stream() + + with ( + patch( + "bedrock_agentcore.runtime.app._flush_tracer_provider_async", new_callable=AsyncMock + ) as mock_async_flush, + patch("bedrock_agentcore.runtime.app._flush_tracer_provider") as mock_sync_flush, + ): + response = TestClient(bedrock_agentcore).post("/invocations", json={}) + + assert response.status_code == 200 + assert 'data: {"ok": true}' in response.text + mock_async_flush.assert_awaited_once_with() + mock_sync_flush.assert_not_called() + + def test_stream_headers_not_delayed_by_flush_of_earlier_spans(self): + """Spans queued before a streaming request must not hold back the response headers.""" + provider = _TimeoutIgnoringTracerProvider() + provider.spans_pending = True + bedrock_agentcore = BedrockAgentCoreApp() + + @bedrock_agentcore.entrypoint + def handler(payload): + yield {"ok": True} + + try: + with ( + patch("opentelemetry.trace.get_tracer_provider", return_value=provider), + patch( + "bedrock_agentcore.runtime.app._flush_tracer_provider_async", + functools.partial(_tracing_module._flush_tracer_provider_async, timeout_millis=3000), + ), + patch( + "bedrock_agentcore.runtime.app._flush_tracer_provider", + functools.partial(_tracing_module._flush_tracer_provider, timeout_millis=100), + ), + ): + begin = time.monotonic() + response = TestClient(bedrock_agentcore).post("/invocations", json={}) + elapsed = time.monotonic() - begin + finally: + provider.release.set() + assert _tracing_module._flush_lock.acquire(timeout=5) + _tracing_module._flush_lock.release() + + assert response.status_code == 200 + assert elapsed < 2, f"response took {elapsed:.2f}s; a pre-stream flush held it back" + assert provider.blocked_flushes == 1 + + @pytest.mark.parametrize("handler_kind", ["sync_generator", "async_generator"]) + def test_stream_closes_when_force_flush_ignores_timeout(self, handler_kind): + """A stuck exporter must not hold the stream open (OTel SDK 1.45 ignores ``timeout_millis``). + + Async-generator entrypoints are bridged to a sync generator, so both kinds end in + ``_sync_stream_with_error_handling``. + """ + provider = _TimeoutIgnoringTracerProvider() + bedrock_agentcore = BedrockAgentCoreApp() + + if handler_kind == "sync_generator": + + @bedrock_agentcore.entrypoint + def handler(payload): + provider.spans_pending = True + yield {"ok": True} + + else: + + @bedrock_agentcore.entrypoint + async def handler(payload): + provider.spans_pending = True + yield {"ok": True} + + result = {} + + def post(): + result["response"] = TestClient(bedrock_agentcore).post("/invocations", json={}) + + client_thread = threading.Thread(target=post, daemon=True) + try: + with ( + patch("opentelemetry.trace.get_tracer_provider", return_value=provider), + patch( + "bedrock_agentcore.runtime.app._flush_tracer_provider", + functools.partial(_tracing_module._flush_tracer_provider, timeout_millis=100), + ), + ): + client_thread.start() + client_thread.join(timeout=5) + assert not client_thread.is_alive(), "stream stayed open while force_flush was blocked" + finally: + provider.release.set() + client_thread.join(timeout=5) + # Let the abandoned flush finish so it does not hold the module-level lock for later tests. + assert _tracing_module._flush_lock.acquire(timeout=5) + _tracing_module._flush_lock.release() + + assert provider.blocked_flushes == 1 + assert result["response"].status_code == 200 + assert 'data: {"ok": true}' in result["response"].text + def test_async_handler_without_context(self): """Test async handler without context parameter.""" bedrock_agentcore = BedrockAgentCoreApp() diff --git a/tests/bedrock_agentcore/runtime/test_tracing.py b/tests/bedrock_agentcore/runtime/test_tracing.py index 21f79867..2586bcf8 100644 --- a/tests/bedrock_agentcore/runtime/test_tracing.py +++ b/tests/bedrock_agentcore/runtime/test_tracing.py @@ -1,6 +1,7 @@ """Tests for BaggageSpanProcessor and auto-registration helper.""" import threading +import time from unittest.mock import MagicMock, patch import pytest @@ -8,6 +9,7 @@ from bedrock_agentcore.runtime.tracing import ( BaggageSpanProcessor, _ensure_baggage_processor_registered, + _flush_lock, _flush_tracer_provider, _flush_tracer_provider_async, ) @@ -303,6 +305,51 @@ def blocking_flush(**_): mock_provider.force_flush.assert_called_once() + def test_flush_wait_is_bounded_when_provider_ignores_timeout(self): + """OTel SDK 1.45 ``force_flush`` ignores ``timeout_millis``; the caller must still get control back.""" + started = threading.Event() + release = threading.Event() + mock_provider = MagicMock() + + def blocking_flush(**_): + started.set() + release.wait(timeout=5) + + mock_provider.force_flush.side_effect = blocking_flush + + try: + with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider): + begin = time.monotonic() + _flush_tracer_provider(timeout_millis=50) + assert time.monotonic() - begin < 2 + assert started.is_set() + + # The abandoned flush still holds the lock, so later flushes skip instead of stacking threads. + _flush_tracer_provider(timeout_millis=50) + assert mock_provider.force_flush.call_count == 1 + finally: + release.set() + + # Once the abandoned flush returns it releases the lock, and the next flush runs normally. + assert _flush_lock.acquire(timeout=5) + _flush_lock.release() + mock_provider.force_flush.side_effect = None + with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider): + _flush_tracer_provider(timeout_millis=50) + assert mock_provider.force_flush.call_count == 2 + + def test_lock_is_released_when_flush_thread_cannot_start(self): + mock_provider = MagicMock() + + with ( + patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider), + patch("threading.Thread.start", side_effect=RuntimeError("can't start new thread")), + ): + _flush_tracer_provider() + + mock_provider.force_flush.assert_not_called() + assert not _flush_lock.locked() + async def test_async_flush_runs_on_worker_thread(self): event_loop_thread = threading.get_ident() worker_threads = []