Stop file transfers when HTTP clients disconnect

starlette/responses.py

@@ -6,7 +6,8 @@
 import os
 import stat
 import sys
-from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable, Mapping, Sequence
+from contextlib import asynccontextmanager
 from datetime import datetime
 from email.utils import format_datetime, formatdate
 from functools import partial
@@ -365,7 +366,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         http_if_range = headers.get("if-range")
 
         if http_range is None or (http_if_range is not None and not self._should_use_range(http_if_range)):
-            await self._handle_simple(send, send_header_only, send_pathsend)
+            send_file = partial(self._handle_simple, send, send_header_only, send_pathsend)
         else:
             try:
                 ranges = self._parse_range_header(http_range, stat_result.st_size)
@@ -376,24 +377,52 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
                 return await response(scope, receive, send)
 
             if len(ranges) == 0:
-                await self._handle_simple(send, send_header_only, send_pathsend)
+                send_file = partial(self._handle_simple, send, send_header_only, send_pathsend)
             elif len(ranges) == 1:
                 start, end = ranges[0]
-                await self._handle_single_range(send, start, end, stat_result.st_size, send_header_only)
+                send_file = partial(self._handle_single_range, send, start, end, stat_result.st_size, send_header_only)
+                send_pathsend = False
             else:
-                await self._handle_multiple_ranges(send, ranges, stat_result.st_size, send_header_only)
+                send_file = partial(self._handle_multiple_ranges, send, ranges, stat_result.st_size, send_header_only)
+                send_pathsend = False
+
+        spec_version = tuple(map(int, scope.get("asgi", {}).get("spec_version", "2.0").split(".")))
+        if scope_type != "http" or send_header_only or send_pathsend or spec_version >= (2, 4):
+            await send_file()
+        else:
+            async with create_collapsing_task_group() as task_group:
+
+                async def stream_file() -> None:
+                    await send_file()
+                    task_group.cancel_scope.cancel()
+
+                task_group.start_soon(stream_file)
+                while True:
+                    if (await receive())["type"] == "http.disconnect":
+                        task_group.cancel_scope.cancel()
+                        break
 
         if self.background is not None:
             await self.background()
 
+    @asynccontextmanager
+    async def _open_file(self) -> AsyncIterator[anyio.AsyncFile[bytes]]:
+        file = await anyio.open_file(self.path, mode="rb")
+        try:
+            yield file
+        finally:
+            # Closing must finish even when the transfer is cancelled.
+            with anyio.CancelScope(shield=True):
+                await file.aclose()
+
     async def _handle_simple(self, send: Send, send_header_only: bool, send_pathsend: bool) -> None:
         await send({"type": "http.response.start", "status": self.status_code, "headers": self.raw_headers})
         if send_header_only:
             await send({"type": "http.response.body", "body": b"", "more_body": False})
         elif send_pathsend:
             await send({"type": "http.response.pathsend", "path": str(self.path)})
         else:
-            async with await anyio.open_file(self.path, mode="rb") as file:
+            async with self._open_file() as file:
                 more_body = True
                 while more_body:
                     chunk = await file.read(self.chunk_size)
@@ -410,7 +439,7 @@ async def _handle_single_range(
         if send_header_only:
             await send({"type": "http.response.body", "body": b"", "more_body": False})
         else:
-            async with await anyio.open_file(self.path, mode="rb") as file:
+            async with self._open_file() as file:
                 await file.seek(start)
                 more_body = True
                 while more_body:
@@ -438,7 +467,7 @@ async def _handle_multiple_ranges(
         if send_header_only:
             await send({"type": "http.response.body", "body": b"", "more_body": False})
         else:
-            async with await anyio.open_file(self.path, mode="rb") as file:
+            async with self._open_file() as file:
                 for start, end in ranges:
                     await send({"type": "http.response.body", "body": header_generator(start, end), "more_body": True})
                     await file.seek(start)

tests/test_responses.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import datetime as dt
+import os
 import sys
 import time
 from collections.abc import AsyncGenerator, AsyncIterator, Iterator
 from dataclasses import dataclass
 from http.cookies import SimpleCookie
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import anyio
 import pytest
@@ -395,6 +396,130 @@ async def send(message: Message) -> None:
     )
 
 
+@pytest.fixture
+def file_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, list[anyio.AsyncFile[bytes]]]:
+    path = tmp_path / "file.bin"
+    path.write_bytes(b"x" * (8 * FileResponse.chunk_size))
+    files: list[anyio.AsyncFile[bytes]] = []
+    open_file = anyio.open_file
+
+    async def track_open_file(path: str | os.PathLike[str], mode: Literal["rb"]) -> anyio.AsyncFile[bytes]:
+        file = await open_file(path, mode=mode)
+        files.append(file)
+        return file
+
+    monkeypatch.setattr(anyio, "open_file", track_open_file)
+    return path, files
+
+
+@pytest.fixture(params=[None, b"bytes=0-", b"bytes=0-131071,196608-524287"])
+def scope(request: pytest.FixtureRequest) -> Scope:
+    return {
+        "type": "http",
+        "method": "GET",
+        "headers": [] if request.param is None else [(b"range", request.param)],
+        "extensions": {"http.response.pathsend": {}} if request.param is not None else {},
+    }
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("spec_version", [None, "2.0", "2.3"])
+async def test_file_response_stops_on_disconnect(
+    file_path: tuple[Path, list[anyio.AsyncFile[bytes]]], scope: Scope, spec_version: str | None
+) -> None:
+    path, files = file_path
+    if spec_version is not None:
+        scope["asgi"] = {"spec_version": spec_version}
+    disconnected = anyio.Event()
+    received_request = False
+    submitted = 0
+    background_ran = False
+
+    async def receive() -> Message:
+        nonlocal received_request
+        if not received_request:
+            received_request = True
+            return {"type": "http.request", "body": b"", "more_body": False}
+        await disconnected.wait()
+        return {"type": "http.disconnect"}
+
+    async def send(message: Message) -> None:
+        nonlocal submitted
+        if message["type"] == "http.response.body":
+            submitted += len(message["body"])
+            if submitted >= FileResponse.chunk_size:
+                disconnected.set()
+
+    async def cleanup() -> None:
+        nonlocal background_ran
+        assert len(files) == 1
+        assert files[0].closed
+        await anyio.sleep(0)
+        background_ran = True
+
+    with anyio.fail_after(5):
+        await FileResponse(path, background=BackgroundTask(cleanup))(scope, receive, send)
+
+    assert FileResponse.chunk_size <= submitted < 3 * FileResponse.chunk_size
+    assert background_ran
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("spec_version", ["2.0", "2.3", "2.4"])
+async def test_file_response_closes_on_cancellation(
+    file_path: tuple[Path, list[anyio.AsyncFile[bytes]]], scope: Scope, spec_version: str
+) -> None:
+    path, files = file_path
+    scope["asgi"] = {"spec_version": spec_version}
+
+    async def receive() -> Message:
+        await anyio.sleep_forever()
+        pytest.fail("The disconnect listener should be cancelled")  # pragma: no cover - sleep never returns
+
+    async def send(message: Message) -> None:
+        if message["type"] == "http.response.body":
+            cancel_scope.cancel()
+            await anyio.sleep_forever()
+
+    async def cleanup() -> None:
+        pytest.fail(
+            "Background tasks should not run after external cancellation"
+        )  # pragma: no cover - failure sentinel
+
+    with anyio.fail_after(5), anyio.CancelScope() as cancel_scope:
+        await FileResponse(path, background=BackgroundTask(cleanup))(scope, receive, send)
+
+    assert cancel_scope.cancelled_caught
+    assert len(files) == 1
+    assert files[0].closed
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("spec_version", ["2.0", "2.3", "2.4"])
+async def test_file_response_closes_on_send_error(
+    file_path: tuple[Path, list[anyio.AsyncFile[bytes]]], scope: Scope, spec_version: str
+) -> None:
+    path, files = file_path
+    scope["asgi"] = {"spec_version": spec_version}
+    error = OSError("Disconnected")
+
+    async def receive() -> Message:
+        assert spec_version != "2.4"
+        await anyio.sleep_forever()
+        pytest.fail("The disconnect listener should be cancelled")  # pragma: no cover - sleep never returns
+
+    async def send(message: Message) -> None:
+        if message["type"] == "http.response.body":
+            raise error
+
+    with anyio.fail_after(5), pytest.raises(OSError) as exc:
+        await FileResponse(path)(scope, receive, send)
+
+    assert exc.value is error
+    assert len(files) == 1
+    assert files[0].closed
+
+
 def test_set_cookie(test_client_factory: TestClientFactory, monkeypatch: pytest.MonkeyPatch) -> None:
     # Mock time used as a reference for `Expires` by stdlib `SimpleCookie`.
     mocked_now = dt.datetime(2037, 1, 22, 12, 0, 0, tzinfo=dt.timezone.utc)
@@ -1018,7 +1143,8 @@ class SmallChunkSizeFileResponse(FileResponse):
     start_message: dict[str, Any] = {}
 
     async def receive() -> Message:
-        raise NotImplementedError("Should not be called!")
+        await anyio.sleep_forever()
+        pytest.fail("The disconnect listener should be cancelled")  # pragma: no cover - sleep never returns
 
     async def send(message: Message) -> None:
         if message["type"] == "http.response.start":