/
/
/
1"""Implementation of a simple multi-client stream task/job."""
2
3from __future__ import annotations
4
5import asyncio
6import logging
7from collections.abc import AsyncGenerator, Sequence
8from contextlib import aclosing, suppress
9from typing import TYPE_CHECKING
10
11from music_assistant.helpers.ffmpeg import get_ffmpeg_stream
12from music_assistant.helpers.util import empty_queue
13
14if TYPE_CHECKING:
15 from music_assistant_models.media_items import AudioFormat
16
17 from music_assistant.helpers.dsp import ComplexFilter
18
19LOGGER = logging.getLogger(__name__)
20
21
22class MultiClientStream:
23 """Implementation of a simple multi-client (audio) stream task/job."""
24
25 def __init__(
26 self,
27 audio_source: AsyncGenerator[bytes],
28 audio_format: AudioFormat,
29 queue_id: str | None,
30 session_id: str | None,
31 expected_clients: int = 0,
32 ) -> None:
33 """Initialize MultiClientStream."""
34 self.audio_source = audio_source
35 self.audio_format = audio_format
36 self.queue_id = queue_id
37 self.session_id = session_id
38 self.subscribers: list[asyncio.Queue[bytes]] = []
39 self.expected_clients = expected_clients
40 self.task = asyncio.create_task(self._runner())
41
42 @property
43 def done(self) -> bool:
44 """Return if this stream is already done."""
45 return self.task.done()
46
47 async def stop(self) -> None:
48 """Stop/cancel the stream."""
49 if self.done:
50 return
51 self.task.cancel()
52 with suppress(asyncio.CancelledError):
53 await self.task
54 for sub_queue in list(self.subscribers):
55 empty_queue(sub_queue)
56
57 async def get_stream(
58 self,
59 output_format: AudioFormat,
60 filter_params: Sequence[str | ComplexFilter] | None = None,
61 ) -> AsyncGenerator[bytes]:
62 """Get (client specific encoded) ffmpeg stream."""
63 async for chunk in get_ffmpeg_stream(
64 audio_input=self.subscribe_raw(),
65 input_format=self.audio_format,
66 output_format=output_format,
67 filter_params=filter_params,
68 ):
69 yield chunk
70
71 async def subscribe_raw(self) -> AsyncGenerator[bytes]:
72 """Subscribe to the raw/unaltered audio stream."""
73 queue: asyncio.Queue[bytes] = asyncio.Queue(2)
74 try:
75 self.subscribers.append(queue)
76 while True:
77 chunk = await queue.get()
78 if chunk == b"":
79 break
80 yield chunk
81 finally:
82 with suppress(ValueError):
83 self.subscribers.remove(queue)
84
85 async def _runner(self) -> None:
86 """Run the stream for the given audio source."""
87 expected_clients = self.expected_clients or 1
88 # wait for first/all subscriber
89 count = 0
90 while count < 50:
91 await asyncio.sleep(0.1)
92 count += 1
93 if len(self.subscribers) >= expected_clients:
94 break
95 LOGGER.debug(
96 "Starting multi-client stream with %s/%s clients",
97 len(self.subscribers),
98 self.expected_clients,
99 )
100 # aclosing so giving up below tears the source down; returning out of an async for
101 # leaves it suspended, and this object keeps it referenced so nothing collects it
102 async with aclosing(self.audio_source):
103 async for chunk in self.audio_source:
104 fail_count = 0
105 while len(self.subscribers) == 0:
106 await asyncio.sleep(0.1)
107 fail_count += 1
108 if fail_count > 50:
109 LOGGER.warning("No clients connected, stopping stream")
110 return
111 await asyncio.gather(
112 *[sub.put(chunk) for sub in self.subscribers], return_exceptions=True
113 )
114 # EOF: send empty chunk
115 await asyncio.gather(*[sub.put(b"") for sub in self.subscribers], return_exceptions=True)
116