From 579d099b3bfb1c46a3c09a859c28b52546618863 Mon Sep 17 00:00:00 2001 From: Andrew Chen Date: Mon, 3 Aug 2026 16:10:59 -0700 Subject: [PATCH] fix: recover from broken S3 audio streams instead of wedging A single mid-stream failure while proxying audio left a dead connection in the shared boto3 client's pool. Every later get_object then returned valid headers with an empty body, so uvicorn aborted the response and Cloudflare served 520 for all playback until the process was restarted. Track bytes streamed against Content-Length and drop the cached client when a stream errors or ends short, so the damage is limited to the one request that hit the dead connection. Connection-level failures on get_object reset it too, and now surface as 502 rather than an unhandled error. Client disconnects (seeking) raise GeneratorExit and are deliberately left alone -- they are routine and leave the connection usable. Co-Authored-By: Claude Opus 5 --- app/audio.py | 50 +++++++++++++++++++++++---- app/s3_client.py | 10 ++++++ tests/test_shows_router.py | 69 +++++++++++++++++++++++++++++++++++++- 3 files changed, 122 insertions(+), 7 deletions(-) diff --git a/app/audio.py b/app/audio.py index 1f6a80b..945a6e7 100644 --- a/app/audio.py +++ b/app/audio.py @@ -1,12 +1,15 @@ import asyncio +import logging from collections.abc import Iterator from dataclasses import dataclass from typing import Any -from botocore.exceptions import ClientError +from botocore.exceptions import BotoCoreError, ClientError from app.config import get_settings -from app.s3_client import get_s3_client +from app.s3_client import get_s3_client, reset_s3_client + +logger = logging.getLogger(__name__) AUDIO_CHUNK_SIZE = 64 * 1024 AUDIO_URL_EXPIRES_IN = 3600 @@ -49,13 +52,41 @@ class AudioStream: media_type: str -def _stream_body(body: Any) -> Iterator[bytes]: +def _stream_body(body: Any, s3_key: str, expected_length: int | None) -> Iterator[bytes]: + """Yield the S3 body in chunks, healing the client pool if upstream fails. + + A stream that raises or ends short of `Content-Length` means the connection + broke mid-response. uvicorn aborts such a response without finishing it, + which cloudflared reports as `unexpected EOF` and Cloudflare turns into a + 520 -- and the dead connection stays in the shared client's pool, so every + later request is truncated the same way until the process restarts. Dropping + the cached client here keeps the damage to the one request that hit it. + + A client that disconnects mid-playback (seeking) closes this generator and + raises `GeneratorExit`, which is deliberately not treated as upstream + breakage -- it is routine and the connection is still fine. + """ + sent = 0 try: while True: chunk = body.read(AUDIO_CHUNK_SIZE) if not chunk: break + sent += len(chunk) yield chunk + except Exception: + logger.warning("audio stream for %s failed after %d bytes", s3_key, sent, exc_info=True) + reset_s3_client() + raise + else: + if expected_length is not None and sent < expected_length: + logger.warning( + "audio stream for %s ended early: sent %d of %d bytes", + s3_key, + sent, + expected_length, + ) + reset_s3_client() finally: body.close() @@ -91,15 +122,22 @@ async def open_audio_stream(s3_key: str, range_header: str | None) -> AudioStrea if code in {"NoSuchKey", "404"}: raise AudioNotFound from e raise AudioUpstreamError(str(e)) from e + except BotoCoreError as e: + # Connection-level failures (closed/refused/timed out) mean the pooled + # connection is unusable; rebuild the client so the retry can succeed. + logger.warning("audio fetch for %s failed to connect", s3_key, exc_info=True) + reset_s3_client() + raise AudioUpstreamError(str(e)) from e + content_length: int | None = s3_response.get("ContentLength") headers: dict[str, str] = {"Accept-Ranges": "bytes"} - if "ContentLength" in s3_response: - headers["Content-Length"] = str(s3_response["ContentLength"]) + if content_length is not None: + headers["Content-Length"] = str(content_length) if "ContentRange" in s3_response: headers["Content-Range"] = s3_response["ContentRange"] return AudioStream( - body=_stream_body(s3_response["Body"]), + body=_stream_body(s3_response["Body"], s3_key, content_length), status_code=206 if "ContentRange" in s3_response else 200, headers=headers, media_type=_media_type_for_key(s3_key), diff --git a/app/s3_client.py b/app/s3_client.py index b751df1..0e6002f 100644 --- a/app/s3_client.py +++ b/app/s3_client.py @@ -24,3 +24,13 @@ def get_s3_client() -> Any: response_checksum_validation="when_required", ), ) + + +def reset_s3_client() -> None: + """Drop the cached client so the next caller builds a fresh connection pool. + + A connection that dies mid-response can be left behind in the shared client's + pool; every later `get_object` then returns valid headers with an empty body. + Rebuilding the client is the only way back out of that state. + """ + get_s3_client.cache_clear() diff --git a/tests/test_shows_router.py b/tests/test_shows_router.py index 229865a..a7dec1e 100644 --- a/tests/test_shows_router.py +++ b/tests/test_shows_router.py @@ -1,8 +1,9 @@ +import contextlib import io from datetime import UTC, date, datetime from unittest.mock import AsyncMock, MagicMock, patch -from botocore.exceptions import ClientError +from botocore.exceptions import ClientError, EndpointConnectionError, ResponseStreamingError from fastapi.testclient import TestClient @@ -394,6 +395,72 @@ def test_audio_other_client_error_502(client: TestClient, mock_conn: AsyncMock) assert "AccessDenied" in response.json()["detail"] +def test_audio_full_stream_keeps_cached_client(client: TestClient, mock_conn: AsyncMock) -> None: + mock_conn.fetchval.return_value = "k.m4a" + s3_mock = MagicMock() + s3_mock.get_object.return_value = {"Body": io.BytesIO(b"COMPLETE"), "ContentLength": 8} + + with ( + patch("app.audio.get_s3_client", return_value=s3_mock), + patch("app.audio.reset_s3_client") as reset_mock, + ): + response = client.get("/api/shows/episodes/11/audio") + + assert response.status_code == 200 + assert response.content == b"COMPLETE" + reset_mock.assert_not_called() + + +def test_audio_short_stream_resets_cached_client(client: TestClient, mock_conn: AsyncMock) -> None: + mock_conn.fetchval.return_value = "k.m4a" + s3_mock = MagicMock() + # Upstream advertises 900 bytes but delivers 5: the pooled connection is dead. + s3_mock.get_object.return_value = {"Body": io.BytesIO(b"SHORT"), "ContentLength": 900} + + with ( + patch("app.audio.get_s3_client", return_value=s3_mock), + patch("app.audio.reset_s3_client") as reset_mock, + ): + client.get("/api/shows/episodes/11/audio") + + reset_mock.assert_called_once() + + +def test_audio_stream_error_resets_cached_client(client: TestClient, mock_conn: AsyncMock) -> None: + mock_conn.fetchval.return_value = "k.m4a" + body = MagicMock() + body.read.side_effect = ResponseStreamingError(error="connection broken") + s3_mock = MagicMock() + s3_mock.get_object.return_value = {"Body": body, "ContentLength": 900} + + with ( + patch("app.audio.get_s3_client", return_value=s3_mock), + patch("app.audio.reset_s3_client") as reset_mock, + contextlib.suppress(ResponseStreamingError), + ): + client.get("/api/shows/episodes/11/audio") + + reset_mock.assert_called_once() + body.close.assert_called_once() + + +def test_audio_connection_error_resets_cached_client( + client: TestClient, mock_conn: AsyncMock +) -> None: + mock_conn.fetchval.return_value = "k.m4a" + s3_mock = MagicMock() + s3_mock.get_object.side_effect = EndpointConnectionError(endpoint_url="https://s3.invalid") + + with ( + patch("app.audio.get_s3_client", return_value=s3_mock), + patch("app.audio.reset_s3_client") as reset_mock, + ): + response = client.get("/api/shows/episodes/11/audio") + + assert response.status_code == 502 + reset_mock.assert_called_once() + + def test_audio_url_returns_presigned_url(client: TestClient, mock_conn: AsyncMock) -> None: mock_conn.fetchval.return_value = "shows/rthk-radio1/x.m4a" s3_mock = MagicMock()