Make file_upload handle cancellations better (#181711)

homeassistant/components/file_upload/__init__.py

@@ -118,6 +118,7 @@ async def _receive_file_field(
     queue: SimpleQueue[tuple[bytes, asyncio.Future[None] | None] | None] = SimpleQueue()
 
     def _sync_queue_consumer() -> None:
+        file_path.parent.mkdir()
         with file_path.open("wb") as file_handle:
             while True:
                 if (_chunk_future := queue.get()) is None:
@@ -128,6 +129,7 @@ async def _receive_file_field(
                 file_handle.write(_chunk)
 
     fut: asyncio.Future[None] | None = None
+    cancelled: asyncio.CancelledError | None = None
     try:
         fut = hass.async_add_executor_job(_sync_queue_consumer)
         chunks_sent = 0
@@ -143,23 +145,31 @@ async def _receive_file_field(
             if fut.done():
                 # The executor job failed
                 break
+    except asyncio.CancelledError as err:
+        # Remember a cancellation from the streaming loop so the join below re-raises
+        # it instead of a later writer error.
+        cancelled = err
+        raise
     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
+            # it reads the sentinel queued above. Wait for it even if this task is
+            # cancelled: asyncio.wait neither cancels the future nor raises its
+            # exception, so the thread is fully done (file written and closed) before
+            # the caller cleans up. The loop re-waits through repeated cancellations.
             while not fut.done():
                 try:
-                    await asyncio.shield(fut)
+                    await asyncio.wait({fut})
                 except asyncio.CancelledError as err:
                     cancelled = err
             if cancelled is not None:
+                # A cancellation takes precedence over a writer error; retrieve the
+                # writer result so its exception isn't flagged as never-retrieved.
+                if not fut.cancelled():
+                    fut.exception()
                 raise cancelled
             fut.result()
 
@@ -216,12 +226,12 @@ class FileUploadView(HomeAssistantView):
         file_dir = file_upload_data.file_dir(file_id)
 
         try:
-            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: _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.
+            # Upload failed: _receive_file_field has joined the writer, which created
+            # the directory 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

@@ -284,6 +284,10 @@ async def test_upload_cancelled_releases_consumer(hass: HomeAssistant) -> None:
     with pytest.raises(asyncio.CancelledError):
         task.result()
 
+    # The cancelled upload must not orphan its file directory on disk.
+    file_upload_data = hass.data[file_upload.DOMAIN]
+    assert list(file_upload_data.temp_dir.iterdir()) == []
+
 
 async def test_receive_file_field_cancelled_while_joining_writer(
     hass: HomeAssistant, tmp_path: Path
@@ -294,7 +298,8 @@ async def test_receive_file_field_cancelled_while_joining_writer(
     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"
+    # Nested under a directory that does not exist yet, so the writer's mkdir runs.
+    file_path = tmp_path / "upload_dir" / "uploaded.bin"
     writing_started = asyncio.Event()
     release_writer = threading.Event()  # blocks the writer thread mid-write
     writes: list[bytes] = []
@@ -347,8 +352,8 @@ async def test_receive_file_field_cancelled_while_joining_writer(
             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.
+            # The task must still be waiting for the writer thread to finish before it
+            # returns, so the caller's cleanup cannot race the writer.
             assert not task.done()
         finally:
             # Always release the writer so a failed assertion can't leak the blocked
@@ -361,3 +366,149 @@ async def test_receive_file_field_cancelled_while_joining_writer(
         task.result()
     # The writer finished writing both chunks before the cancellation propagated.
     assert b"".join(writes) == b"chunk1chunk2"
+    # The writer created the parent directory as part of the joined job, so cleanup
+    # cannot race an in-flight mkdir.
+    assert file_path.parent.is_dir()
+
+
+async def test_receive_file_field_cancel_wins_over_writer_error(
+    hass: HomeAssistant, tmp_path: Path
+) -> None:
+    """Test a writer error during the join does not mask the cancellation.
+
+    When the task is cancelled while joining the writer and the writer then fails,
+    the caller must still observe CancelledError, not the writer's exception.
+    """
+    file_path = tmp_path / "upload_dir" / "uploaded.bin"
+    writing_started = asyncio.Event()
+    release_writer = threading.Event()  # blocks the writer thread mid-write
+
+    class _FailingHandle:
+        """A file handle whose first write parks the writer, then fails."""
+
+        def __enter__(self) -> Self:
+            return self
+
+        def __exit__(self, *exc: object) -> bool:
+            return False
+
+        def write(self, data: bytes) -> int:
+            hass.loop.call_soon_threadsafe(writing_started.set)
+            release_writer.wait()
+            raise OSError("write failed")
+
+    real_open = Path.open
+
+    def _failing_open(self: Path, *args: object, **kwargs: object) -> object:
+        if self != file_path:
+            return real_open(self, *args, **kwargs)
+        return _FailingHandle()
+
+    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", _failing_open):
+        task = asyncio.create_task(
+            file_upload._receive_file_field(hass, _Part(), file_path)
+        )
+        try:
+            # Let the task settle at the join with the writer blocked mid-write, so
+            # the cancel lands on the join and the writer fails afterwards.
+            await writing_started.wait()
+            for _ in range(3):
+                await asyncio.sleep(0)
+            task.cancel()
+            for _ in range(10):
+                await asyncio.sleep(0)
+            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
+    # The cancellation wins over the writer's OSError.
+    assert task.cancelled()
+    with pytest.raises(asyncio.CancelledError):
+        task.result()
+
+
+async def test_receive_file_field_cancel_in_stream_wins_over_writer_error(
+    hass: HomeAssistant, tmp_path: Path
+) -> None:
+    """Test a cancel in the streaming loop is not masked by a writer error.
+
+    When the cancellation lands in read_chunk (not the join) and the writer has
+    already failed, the caller must still observe CancelledError, not the writer's
+    exception.
+    """
+    file_path = tmp_path / "upload_dir" / "uploaded.bin"
+    writer_failed = asyncio.Event()  # set once the writer thread has raised
+    reading_blocked = asyncio.Event()  # set when the stream parks on its 2nd read
+    blocked = asyncio.Event()  # never set, so the 2nd read blocks until cancelled
+
+    class _FailingHandle:
+        """A file handle whose first write fails the writer thread."""
+
+        def __enter__(self) -> Self:
+            return self
+
+        def __exit__(self, *exc: object) -> bool:
+            return False
+
+        def write(self, data: bytes) -> int:
+            hass.loop.call_soon_threadsafe(writer_failed.set)
+            raise OSError("write failed")
+
+    real_open = Path.open
+
+    def _failing_open(self: Path, *args: object, **kwargs: object) -> object:
+        if self != file_path:
+            return real_open(self, *args, **kwargs)
+        return _FailingHandle()
+
+    reads = 0
+
+    class _Part:
+        """Fake BodyPartReader yielding one chunk, then blocking on the next read."""
+
+        async def read_chunk(self, size: int) -> bytes:
+            nonlocal reads
+            reads += 1
+            if reads == 1:
+                return b"chunk1"
+            reading_blocked.set()
+            await blocked.wait()
+            return b""
+
+    with patch.object(Path, "open", _failing_open):
+        task = asyncio.create_task(
+            file_upload._receive_file_field(hass, _Part(), file_path)
+        )
+        try:
+            # The stream is parked on its second read and the writer has failed, so
+            # the cancel lands in the streaming loop with the writer future already
+            # done with an error.
+            await reading_blocked.wait()
+            await writer_failed.wait()
+            for _ in range(5):
+                await asyncio.sleep(0)
+            task.cancel()
+            for _ in range(10):
+                await asyncio.sleep(0)
+        finally:
+            # Always unblock the stream so a failed assertion can't hang teardown.
+            blocked.set()
+        _done, pending = await asyncio.wait({task}, timeout=10)
+
+    assert not pending
+    # The cancellation wins over the writer's OSError.
+    assert task.cancelled()
+    with pytest.raises(asyncio.CancelledError):
+        task.result()