music-assistant-server

1"""
2Implementation of a Stream for the Universal Group Player.
3
4Stream handler for Universal Groups, managing audio distribution to group members.
5Essentially, it multicasts an audio source to multiple client streams, allowing individual
6filter_params for each client.
7"""
8
9from __future__ import annotations
10
11import asyncio
12import logging
13from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence
14from contextlib import suppress
15from typing import TYPE_CHECKING
16
17if TYPE_CHECKING:
18    from music_assistant_models.media_items import AudioFormat
19
20    from music_assistant.helpers.dsp import ComplexFilter
21
22from music_assistant.constants import MASS_LOGGER_NAME
23from music_assistant.helpers.ffmpeg import get_ffmpeg_stream
24from music_assistant.helpers.util import empty_queue
25
26LOGGER = logging.getLogger(f"{MASS_LOGGER_NAME}.providers.ugp_stream")
27
28
29class UGPStream:
30    """
31    Implementation of a Stream for the Universal Group Player.
32
33    Stream handler for Universal Groups, managing audio distribution to group members.
34    Essentially, it multicasts an audio source to multiple client streams, allowing individual
35    filter_params for each client.
36    """
37
38    def __init__(
39        self,
40        audio_source: AsyncGenerator[bytes],
41        audio_format: AudioFormat,
42        base_pcm_format: AudioFormat,
43        queue_id: str | None,
44        session_id: str | None,
45    ) -> None:
46        """Initialize UGP Stream."""
47        self.audio_source = audio_source
48        self.input_format = audio_format
49        self.base_pcm_format = base_pcm_format
50        self.queue_id = queue_id
51        self.session_id = session_id
52        self.subscribers: list[Callable[[bytes], Awaitable[None]]] = []
53        self._task: asyncio.Task[None] | None = None
54        self._done: asyncio.Event = asyncio.Event()
55
56    @property
57    def done(self) -> bool:
58        """Return if this stream is already done."""
59        return self._done.is_set() and self._task is not None and self._task.done()
60
61    async def stop(self) -> None:
62        """Stop/cancel the stream."""
63        if self._done.is_set():
64            return
65        if self._task and not self._task.done():
66            self._task.cancel()
67            with suppress(asyncio.CancelledError):
68                await self._task
69        self._done.set()
70
71    async def subscribe_raw(self) -> AsyncGenerator[bytes]:
72        """
73        Subscribe to the raw/unaltered audio stream.
74
75        The returned stream has the format `self.base_pcm_format`.
76        """
77        # start the runner as soon as the (first) client connects
78        if not self._task:
79            self._task = asyncio.create_task(self._runner())
80        queue: asyncio.Queue[bytes] = asyncio.Queue(10)
81        try:
82            self.subscribers.append(queue.put)
83            while True:
84                chunk = await queue.get()
85                if not chunk:
86                    break
87                yield chunk
88        finally:
89            self.subscribers.remove(queue.put)
90            empty_queue(queue)
91            del queue
92
93    async def get_stream(
94        self, output_format: AudioFormat, filter_params: Sequence[str | ComplexFilter] | None = None
95    ) -> AsyncGenerator[bytes]:
96        """Subscribe to the client specific audio stream."""
97        # start the runner as soon as the (first) client connects
98        async for chunk in get_ffmpeg_stream(
99            audio_input=self.subscribe_raw(),
100            input_format=self.base_pcm_format,
101            output_format=output_format,
102            filter_params=filter_params,
103        ):
104            yield chunk
105
106    async def _runner(self) -> None:
107        """Run the stream for the given audio source."""
108        await asyncio.sleep(0.25)  # small delay to allow subscribers to connect
109        try:
110            async for chunk in get_ffmpeg_stream(
111                audio_input=self.audio_source,
112                input_format=self.input_format,
113                output_format=self.base_pcm_format,
114                # we don't allow the player to buffer too much ahead so we use readrate limiting
115                extra_input_args=["-readrate", "1.1", "-readrate_initial_burst", "5"],
116            ):
117                await asyncio.gather(
118                    *[sub(chunk) for sub in self.subscribers],
119                    return_exceptions=True,
120                )
121        except asyncio.CancelledError:
122            LOGGER.debug("UGP stream runner cancelled")
123            raise
124        except Exception as err:
125            LOGGER.error("UGP stream runner error: %s", err, exc_info=err)
126        finally:
127            # empty chunk when done
128            await asyncio.gather(*[sub(b"") for sub in self.subscribers], return_exceptions=True)
129            self._done.set()
130