From 749367b8011c4f272582e9a23b25920eef49e44a Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Sat, 5 Sep 2026 10:22:03 +0200 Subject: [PATCH] Make file_upload handle cancellations better (#181163) --- .../components/file_upload/__init__.py | 101 +++++++++++------- tests/components/file_upload/test_init.py | 81 +++++++++++++- 2 files changed, 140 insertions(+), 42 deletions(-) diff --git a/homeassistant/components/file_upload/__init__.py b/homeassistant/components/file_upload/__init__.py index b535b080767d..9562b77e337a 100644 --- a/homeassistant/components/file_upload/__init__.py +++ b/homeassistant/components/file_upload/__init__.py @@ -109,6 +109,61 @@ class FileUploadData: return self.file_dir(file_id) / self.files[file_id] +async def _receive_file_field( + hass: HomeAssistant, + file_field_reader: BodyPartReader, + file_path: Path, +) -> None: + """Stream a multipart file field to file_path using an executor writer.""" + queue: SimpleQueue[tuple[bytes, asyncio.Future[None] | None] | None] = SimpleQueue() + + def _sync_queue_consumer() -> None: + with file_path.open("wb") as file_handle: + while True: + if (_chunk_future := queue.get()) is None: + break + _chunk, _future = _chunk_future + if _future is not None: + hass.loop.call_soon_threadsafe(_future.set_result, None) + file_handle.write(_chunk) + + fut: asyncio.Future[None] | None = None + try: + fut = hass.async_add_executor_job(_sync_queue_consumer) + chunks_sent = 0 + while chunk := await file_field_reader.read_chunk(ONE_MEGABYTE): + chunks_sent += 1 + if chunks_sent % 5 != 0: + queue.put_nowait((chunk, None)) + continue + + chunk_future = hass.loop.create_future() + queue.put_nowait((chunk, chunk_future)) + await asyncio.wait((fut, chunk_future), return_when=asyncio.FIRST_COMPLETED) + if fut.done(): + # The executor job failed + break + finally: + # Always terminate the queue consumer, also if the stream raised or the task + # was cancelled, otherwise awaiting the consumer future deadlocks. + queue.put_nowait(None) + if fut is not None: + # The executor thread can't be cancelled and is guaranteed to finish once + # it reads the sentinel queued above. Await it even if this task is + # cancelled: awaiting only through a shield keeps the executor future + # itself uncancelled, so the thread is fully done (file written and closed) + # before the caller cleans up. Re-raise any cancellation after. + cancelled: asyncio.CancelledError | None = None + while not fut.done(): + try: + await asyncio.shield(fut) + except asyncio.CancelledError as err: + cancelled = err + if cancelled is not None: + raise cancelled + fut.result() + + class FileUploadView(HomeAssistantView): """HTTP View to upload files.""" @@ -159,50 +214,14 @@ class FileUploadView(HomeAssistantView): file_upload_data = hass.data[_DATA] file_dir = file_upload_data.file_dir(file_id) - queue: SimpleQueue[tuple[bytes, asyncio.Future[None] | None] | None] = ( - SimpleQueue() - ) - def _sync_queue_consumer() -> None: - file_dir.mkdir() - with (file_dir / filename).open("wb") as file_handle: - while True: - if (_chunk_future := queue.get()) is None: - break - _chunk, _future = _chunk_future - if _future is not None: - hass.loop.call_soon_threadsafe(_future.set_result, None) - file_handle.write(_chunk) - - fut: asyncio.Future[None] | None = None try: - try: - fut = hass.async_add_executor_job(_sync_queue_consumer) - chunks_sent = 0 - while chunk := await file_field_reader.read_chunk(ONE_MEGABYTE): - chunks_sent += 1 - if chunks_sent % 5 != 0: - queue.put_nowait((chunk, None)) - continue - - chunk_future = hass.loop.create_future() - queue.put_nowait((chunk, chunk_future)) - await asyncio.wait( - (fut, chunk_future), return_when=asyncio.FIRST_COMPLETED - ) - if fut.done(): - # The executor job failed - break - finally: - # Always terminate the queue consumer, also if the stream raised or - # the task was cancelled. - queue.put_nowait(None) - if fut is not None: - await fut + await hass.async_add_executor_job(file_dir.mkdir) + await _receive_file_field(hass, file_field_reader, file_dir / filename) except Exception, asyncio.CancelledError: - # Upload failed: the consumer has finished and closed the file (inner - # finally above), so removing the directory now cannot race the writer. - # ignore_errors covers a failure that happened before the dir was created. + # Upload failed: _receive_file_field has joined the writer and closed the + # file, so removing the directory now cannot race the writer. ignore_errors + # covers a failure that happened before the directory was created. await hass.async_add_executor_job( lambda: shutil.rmtree(file_dir, ignore_errors=True) ) diff --git a/tests/components/file_upload/test_init.py b/tests/components/file_upload/test_init.py index 6f940e63fc07..6062a25cb35e 100644 --- a/tests/components/file_upload/test_init.py +++ b/tests/components/file_upload/test_init.py @@ -5,7 +5,8 @@ from contextlib import contextmanager from io import StringIO from pathlib import Path from random import getrandbits -from typing import Any +import threading +from typing import Any, Self from unittest.mock import AsyncMock, patch from aiohttp import BodyPartReader @@ -282,3 +283,81 @@ async def test_upload_cancelled_releases_consumer(hass: HomeAssistant) -> None: assert not pending with pytest.raises(asyncio.CancelledError): task.result() + + +async def test_receive_file_field_cancelled_while_joining_writer( + hass: HomeAssistant, tmp_path: Path +) -> None: + """Test a cancel while joining the writer waits for the writer to finish. + + A cancellation delivered while _receive_file_field is awaiting the executor + writer (the whole field already streamed, sentinel queued) must not return + until the writer thread has finished, so the caller's cleanup cannot race it. + """ + file_path = tmp_path / "uploaded.bin" + writing_started = asyncio.Event() + release_writer = threading.Event() # blocks the writer thread mid-write + writes: list[bytes] = [] + blocked_once = False + + class _BlockingHandle: + """A file handle whose first write parks the writer thread until released.""" + + def __enter__(self) -> Self: + return self + + def __exit__(self, *exc: object) -> bool: + return False + + def write(self, data: bytes) -> int: + nonlocal blocked_once + if not blocked_once: + blocked_once = True + hass.loop.call_soon_threadsafe(writing_started.set) + release_writer.wait() + writes.append(bytes(data)) + return len(data) + + real_open = Path.open + + def _blocking_open(self: Path, *args: object, **kwargs: object) -> object: + if self != file_path: + return real_open(self, *args, **kwargs) + return _BlockingHandle() + + chunks = iter([b"chunk1", b"chunk2"]) + + class _Part: + """Fake BodyPartReader yielding two chunks then EOF.""" + + async def read_chunk(self, size: int) -> bytes: + return next(chunks, b"") + + with patch.object(Path, "open", _blocking_open): + task = asyncio.create_task( + file_upload._receive_file_field(hass, _Part(), file_path) + ) + try: + # The writer thread is now blocked on its first write, so the task has + # queued every chunk plus the sentinel and is parked at the join; let it + # settle there so the cancel lands on the join. + await writing_started.wait() + for _ in range(3): + await asyncio.sleep(0) + task.cancel() + for _ in range(10): + await asyncio.sleep(0) + # Without the cancellation-safe join the task would finish here (returning + # while the writer thread runs on); the fix keeps it waiting for the writer. + assert not task.done() + finally: + # Always release the writer so a failed assertion can't leak the blocked + # thread and hang teardown. + release_writer.set() + _done, pending = await asyncio.wait({task}, timeout=10) + + assert not pending + with pytest.raises(asyncio.CancelledError): + task.result() + # The writer finished writing both chunks before the cancellation propagated. + assert b"".join(writes) == b"chunk1chunk2"