From b4e69bab71ecda4b742e9420c70086900c45fc73 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Wed, 30 Oct 2024 13:46:05 +0100
Subject: [PATCH] Improve shutdown of esphome ffmpeg proxy (#129326)

* Improve shutdown of esphome ffmpeg proxy

* Add test
---
 .../components/esphome/ffmpeg_proxy.py        | 16 +++++--
 tests/components/esphome/test_ffmpeg_proxy.py | 46 +++++++++++++++++++
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/esphome/ffmpeg_proxy.py b/homeassistant/components/esphome/ffmpeg_proxy.py
index d750fcca572..cefe87f49ba 100644
--- a/homeassistant/components/esphome/ffmpeg_proxy.py
+++ b/homeassistant/components/esphome/ffmpeg_proxy.py
@@ -194,7 +194,11 @@ class FFmpegConvertResponse(web.StreamResponse):
         # Only one conversion process per device is allowed
         self.convert_info.proc = proc
 
-        await self._write_ffmpeg_data(request, writer, proc)
+        # Create background task which will be cancelled when home assistant shuts down
+        write_task = self.hass.async_create_background_task(
+            self._write_ffmpeg_data(request, writer, proc), "ESPHome media proxy"
+        )
+        await write_task
 
     async def _write_ffmpeg_data(
         self,
@@ -215,6 +219,11 @@ class FFmpegConvertResponse(web.StreamResponse):
             ):
                 await self.write(chunk)
         except asyncio.CancelledError:
+            _LOGGER.debug("ffmpeg transcoding cancelled")
+            # Abort the transport, we don't wait for ESPHome to drain the write buffer;
+            # it may need a very long time or never finish if the player is paused.
+            if request.transport:
+                request.transport.abort()
             raise  # don't log error
         except:
             _LOGGER.exception("Unexpected error during ffmpeg conversion")
@@ -234,8 +243,9 @@ class FFmpegConvertResponse(web.StreamResponse):
             if proc.returncode is None:
                 proc.kill()
 
-            # Close connection
-            await writer.write_eof()
+            # Close connection by writing EOF unless already closing
+            if request.transport and not request.transport.is_closing():
+                await writer.write_eof()
 
 
 class FFmpegProxyView(HomeAssistantView):
diff --git a/tests/components/esphome/test_ffmpeg_proxy.py b/tests/components/esphome/test_ffmpeg_proxy.py
index 403da008498..295d8d2fda9 100644
--- a/tests/components/esphome/test_ffmpeg_proxy.py
+++ b/tests/components/esphome/test_ffmpeg_proxy.py
@@ -9,6 +9,7 @@ from unittest.mock import patch
 from urllib.request import pathname2url
 import wave
 
+from aiohttp import client_exceptions
 import mutagen
 import pytest
 
@@ -286,3 +287,48 @@ async def test_max_conversions_per_device(
             for url in urls[1:]:
                 req = await client.get(url)
                 assert req.status == HTTPStatus.OK
+
+
+async def test_abort_on_shutdown(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+) -> None:
+    """Test we abort on Home Assistant shutdown."""
+    device_id = "1234"
+
+    await async_setup_component(hass, esphome.DOMAIN, {esphome.DOMAIN: {}})
+    client = await hass_client()
+
+    with tempfile.NamedTemporaryFile(mode="wb+", suffix=".wav") as temp_file:
+        with wave.open(temp_file.name, "wb") as wav_file:
+            wav_file.setframerate(16000)
+            wav_file.setsampwidth(2)
+            wav_file.setnchannels(1)
+            wav_file.writeframes(bytes(16000 * 2))  # 1s
+
+        wav_url = pathname2url(temp_file.name)
+        convert_id = "test-id"
+        url = f"/api/esphome/ffmpeg_proxy/{device_id}/{convert_id}.mp3"
+
+        wav_url = pathname2url(temp_file.name)
+        url = async_create_proxy_url(
+            hass,
+            device_id,
+            wav_url,
+            media_format="wav",
+            rate=22050,
+            channels=2,
+            width=2,
+        )
+
+        # Get URL and start reading
+        req = await client.get(url)
+        assert req.status == HTTPStatus.OK
+        initial_mp3_data = await req.content.read(4)
+        assert initial_mp3_data == b"RIFF"
+
+        # Shut down Home Assistant
+        await hass.async_stop()
+
+        with pytest.raises(client_exceptions.ClientPayloadError):
+            await req.content.read()