Improve shutdown of esphome ffmpeg proxy (#129326)

* Improve shutdown of esphome ffmpeg proxy

* Add test
This commit is contained in:
Erik Montnemery 2024-10-30 13:46:05 +01:00 committed by GitHub
parent db81edfb2b
commit b4e69bab71
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 59 additions and 3 deletions

View file

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

View file

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