Make file_upload handle cancellations better (#181163)

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)
             )

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"