Improve error handling in file_upload (#181078)

homeassistant/components/file_upload/__init__.py

@@ -176,27 +176,37 @@ class FileUploadView(HomeAssistantView):
 
         fut: asyncio.Future[None] | None = None
         try:
-            fut = hass.async_add_executor_job(_sync_queue_consumer)
-            megabytes_sending = 0
-            while chunk := await file_field_reader.read_chunk(ONE_MEGABYTE):
-                megabytes_sending += 1
-                if megabytes_sending % 5 != 0:
-                    queue.put_nowait((chunk, None))
-                    continue
+            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
-
-            queue.put_nowait(None)  # terminate queue consumer
-        finally:
-            if fut is not None:
-                await fut
+                    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
+        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.
+            await hass.async_add_executor_job(
+                lambda: shutil.rmtree(file_dir, ignore_errors=True)
+            )
+            raise
 
         file_upload_data.files[file_id] = filename
 

tests/components/file_upload/test_init.py

@@ -1,15 +1,19 @@
 """Test the File Upload integration."""
 
+import asyncio
 from contextlib import contextmanager
+from io import StringIO
 from pathlib import Path
 from random import getrandbits
 from typing import Any
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
+from aiohttp import BodyPartReader
 import pytest
 
 from homeassistant.components import file_upload
-from homeassistant.components.file_upload import DOMAIN
+from homeassistant.components.file_upload import DOMAIN, FileUploadView
+from homeassistant.components.http import KEY_HASS
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
@@ -170,3 +174,111 @@ async def test_upload_large_file_fails(
     response = await res.content.read()
 
     assert b"Boom" in response
+
+
+async def test_upload_stream_error_releases_lock(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    large_file_io: StringIO,
+) -> None:
+    """Test a mid-upload stream error propagates and releases the upload lock.
+
+    Models a client disconnect: aiohttp raises from read_chunk when the
+    connection is lost mid-transfer. If the queue consumer's terminating sentinel
+    is skipped on the error path, awaiting the consumer deadlocks while holding
+    the upload lock, wedging every later upload; this guards that regression.
+    """
+    assert await async_setup_component(hass, DOMAIN, {})
+    client = await hass_client()
+
+    with (
+        patch(
+            # Patch temp dir name to avoid tests fail running in parallel
+            "homeassistant.components.file_upload.TEMP_DIR_NAME",
+            file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}",
+        ),
+        patch.object(
+            BodyPartReader,
+            "read_chunk",
+            AsyncMock(
+                side_effect=[b"partial", ConnectionResetError("Connection lost")]
+            ),
+        ),
+    ):
+        # Bound the request so a reintroduced deadlock fails fast instead of hanging
+        async with asyncio.timeout(10):
+            res = await client.post("/api/file_upload", data={"file": large_file_io})
+
+    assert res.status == 500
+
+    # The failed upload must not leave a partially written file orphaned on disk
+    file_upload_data = hass.data[file_upload.DOMAIN]
+    assert list(file_upload_data.temp_dir.iterdir()) == []
+
+    # The upload lock must have been released: a subsequent normal upload succeeds
+    large_file_io.seek(0)
+    with patch(
+        "homeassistant.components.file_upload.TEMP_DIR_NAME",
+        file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}",
+    ):
+        async with asyncio.timeout(10):
+            res = await client.post("/api/file_upload", data={"file": large_file_io})
+
+    assert res.status == 200
+
+
+async def test_upload_cancelled_releases_consumer(hass: HomeAssistant) -> None:
+    """Test cancelling an upload mid-transfer does not deadlock the consumer.
+
+    Driven at the view level because the test HTTP client cannot produce a true
+    task cancellation mid-request. Without delivering the queue sentinel on the
+    cancellation path, awaiting the consumer future would hang forever.
+    """
+    assert await async_setup_component(hass, DOMAIN, {})
+    view = FileUploadView()
+
+    first_chunk_sent = asyncio.Event()
+    blocked = asyncio.Event()  # never set, so the stream blocks until cancelled
+
+    class _BlockingPart:
+        """Fake BodyPartReader that blocks after yielding one chunk."""
+
+        name = "file"
+        filename = "blocking.bin"
+
+        async def read_chunk(self, size: int) -> bytes:
+            if first_chunk_sent.is_set():
+                await blocked.wait()
+                return b""
+            first_chunk_sent.set()
+            return b"chunk"
+
+    part = _BlockingPart()
+
+    class _Reader:
+        async def next(self) -> _BlockingPart:
+            return part
+
+    class _Request:
+        app = {KEY_HASS: hass}
+
+        async def multipart(self) -> _Reader:
+            return _Reader()
+
+    with (
+        patch(
+            "homeassistant.components.file_upload.TEMP_DIR_NAME",
+            file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}",
+        ),
+        patch("homeassistant.components.file_upload.BodyPartReader", _BlockingPart),
+    ):
+        task = asyncio.create_task(view._upload_file(_Request()))
+        await first_chunk_sent.wait()
+        task.cancel()
+        # asyncio.wait does not cancel on timeout, so a deadlocked task stays
+        # pending and the assertion fails fast instead of the run hanging.
+        _done, pending = await asyncio.wait({task}, timeout=10)
+
+    assert not pending
+    with pytest.raises(asyncio.CancelledError):
+        task.result()