308fbc4fb4
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012vJkmBGKKqMoEMzvsXG3y2
54 lines
2.9 KiB
Python
54 lines
2.9 KiB
Python
"""Потоковая отдача файла с детерминированным освобождением источника (T9.1, Review Focus 4 плана Среза 3b).
|
||
|
||
Starlette `StreamingResponse.stream_response` не вызывает `aclose()` у `body_iterator`: при обрыве клиента или если тело
|
||
не начали итерировать поток S3 освобождался бы только сборщиком мусора. `ClosingStreamingResponse` закрывает итератор
|
||
в `finally`, а `ChainedStream` гарантирует, что `aclose` закрывает и исходный поток, даже если обёртку не начинали
|
||
читать (у неначатого асинхронного генератора `finally` не выполняется, поэтому обёртка — объект, а не генератор).
|
||
"""
|
||
|
||
from collections.abc import AsyncGenerator
|
||
|
||
import anyio
|
||
from fastapi.responses import StreamingResponse
|
||
from starlette.types import Send
|
||
|
||
|
||
class ChainedStream:
|
||
"""Уже прочитанная первая порция и остальной поток `rest`; `aclose` всегда закрывает `rest` (там `close` + `release`
|
||
ответа S3), в том числе при завершении, ошибке и обрыве потребителя, начатом и не начатом чтении."""
|
||
|
||
def __init__(self, first: bytes, rest: AsyncGenerator[bytes, None]) -> None:
|
||
self._first = first
|
||
self._rest = rest
|
||
|
||
def __aiter__(self) -> "ChainedStream":
|
||
return self
|
||
|
||
async def __anext__(self) -> bytes:
|
||
if self._first:
|
||
chunk, self._first = self._first, b""
|
||
return chunk
|
||
return await anext(self._rest)
|
||
|
||
async def aclose(self) -> None:
|
||
"""Закрыть исходный поток и сбросить префетченную порцию: после закрытия `__anext__` её не отдаёт
|
||
(идемпотентно: у завершённого генератора `aclose` — no-op)."""
|
||
self._first = b""
|
||
await self._rest.aclose()
|
||
|
||
|
||
class ClosingStreamingResponse(StreamingResponse):
|
||
"""`StreamingResponse`, закрывающий `body_iterator` (`aclose`) после отдачи, ошибки или обрыва клиента.
|
||
|
||
Закрытие идёт под `CancelScope(shield=True)`: при обрыве клиента задача отдачи отменяется, а освобождение
|
||
соединения S3 требует `await`."""
|
||
|
||
async def stream_response(self, send: Send) -> None:
|
||
try:
|
||
await super().stream_response(send)
|
||
finally:
|
||
aclose = getattr(self.body_iterator, "aclose", None)
|
||
if aclose is not None:
|
||
with anyio.CancelScope(shield=True):
|
||
await aclose()
|