/
/
/
1"""Helper for adding buffering to async (audio) generators."""
2
3from __future__ import annotations
4
5import asyncio
6import contextlib
7from collections.abc import AsyncGenerator, Callable
8from functools import wraps
9from typing import Any, Final, ParamSpec
10
11from music_assistant.helpers.util import close_async_generator, empty_queue
12
13# Type variables for the buffered decorator
14_P = ParamSpec("_P")
15
16DEFAULT_BUFFER_SIZE: Final = 30
17DEFAULT_MIN_BUFFER_BEFORE_YIELD: Final = 5
18
19# Keep strong references to producer tasks to prevent garbage collection
20# The event loop only keeps weak references to tasks
21_ACTIVE_PRODUCER_TASKS: set[asyncio.Task[Any]] = set()
22
23
24async def _finalize_producer(
25 generator: AsyncGenerator[bytes],
26 completed_naturally: bool,
27 buffer: asyncio.Queue[bytes | None],
28 threshold_reached: asyncio.Event,
29 cancelled: asyncio.Event,
30) -> None:
31 """Release any waiting consumer and signal the end of the stream."""
32 threshold_reached.set()
33 # Close the upstream generator on any early-exit path, even if it already
34 # produced some chunks before the consumer stopped.
35 if not completed_naturally:
36 close_task = asyncio.create_task(close_async_generator(generator))
37 try:
38 await asyncio.shield(close_task)
39 except asyncio.CancelledError:
40 await close_task
41 raise
42 # Signal end of stream by putting None
43 # We must wait for space in the queue if needed, otherwise the consumer may
44 # hang waiting for data that will never come
45 if not cancelled.is_set():
46 await buffer.put(None)
47
48
49async def _shutdown_producer(
50 producer_task: asyncio.Task[None],
51 buffer: asyncio.Queue[bytes | None],
52 cancelled: asyncio.Event,
53) -> None:
54 """Stop the producer task without blocking indefinitely on a slow source."""
55 # Signal the producer to stop
56 cancelled.set()
57 # Drain the queue to unblock the producer if it's waiting on put()
58 empty_queue(buffer)
59 # Wait for the producer to finish cleanly with a timeout to prevent blocking
60 with contextlib.suppress(asyncio.CancelledError, RuntimeError, asyncio.TimeoutError):
61 await asyncio.wait_for(asyncio.shield(producer_task), timeout=1.0)
62 # Force-cancel producer if still stuck on a slow read to prevent resource leaks
63 if not producer_task.done():
64 producer_task.cancel()
65 with contextlib.suppress(asyncio.CancelledError, RuntimeError, asyncio.TimeoutError):
66 await asyncio.wait_for(producer_task, timeout=1.0)
67
68
69async def buffered(
70 generator: AsyncGenerator[bytes],
71 buffer_size: int = DEFAULT_BUFFER_SIZE,
72 min_buffer_before_yield: int = DEFAULT_MIN_BUFFER_BEFORE_YIELD,
73) -> AsyncGenerator[bytes]:
74 """
75 Add buffering to an async generator that yields (chunks of) bytes.
76
77 This function uses an asyncio.Queue to decouple the producer (reading from the stream)
78 from the consumer (yielding to the client). The producer runs in a separate task and
79 fills the buffer, while the consumer yields from the buffer.
80
81 Args:
82 generator: The async generator to buffer
83 buffer_size: Maximum number of chunks to buffer (default: 30)
84 min_buffer_before_yield: Minimum chunks to buffer before starting to yield (default: 5)
85
86 Example:
87 async for chunk in buffered(my_generator(), buffer_size=100):
88 process(chunk)
89 """
90 buffer: asyncio.Queue[bytes | None] = asyncio.Queue(maxsize=buffer_size)
91 producer_error: Exception | None = None
92 threshold_reached = asyncio.Event()
93 cancelled = asyncio.Event()
94 min_buffer_before_yield = max(1, min(min_buffer_before_yield, buffer_size))
95
96 if buffer_size <= 1:
97 # No buffering needed, yield directly
98 async for chunk in generator:
99 yield chunk
100 return
101
102 async def producer() -> None:
103 """
104 Read from the original generator and fill the buffer.
105
106 Note: When the buffer is full, buffer.put() will naturally wait for the consumer
107 to drain items. This is the intended buffering behavior and may trigger asyncio
108 "slow callback" warnings (typically 0.1-0.2s) which are harmless and expected.
109 These warnings are filtered out in the main logging configuration.
110 """
111 nonlocal producer_error
112 completed_naturally = False
113 stopped_early = False
114 try:
115 async for chunk in generator:
116 if cancelled.is_set():
117 # Consumer has stopped, exit cleanly
118 stopped_early = True
119 break
120 await buffer.put(chunk)
121 if not threshold_reached.is_set() and buffer.qsize() >= min_buffer_before_yield:
122 threshold_reached.set()
123 # Yield to event loop every chunk to prevent blocking
124 await asyncio.sleep(0)
125 completed_naturally = not stopped_early
126 except Exception as err:
127 producer_error = err
128 if isinstance(err, asyncio.CancelledError):
129 raise
130 finally:
131 await _finalize_producer(
132 generator=generator,
133 completed_naturally=completed_naturally,
134 buffer=buffer,
135 threshold_reached=threshold_reached,
136 cancelled=cancelled,
137 )
138
139 # Start the producer task
140 loop = asyncio.get_running_loop()
141 producer_task = loop.create_task(producer())
142
143 # Keep a strong reference to prevent garbage collection issues
144 # The event loop only keeps weak references to tasks
145 _ACTIVE_PRODUCER_TASKS.add(producer_task)
146
147 # Remove from set when done
148 producer_task.add_done_callback(_ACTIVE_PRODUCER_TASKS.discard)
149
150 try:
151 # Wait for initial buffer to fill
152 await threshold_reached.wait()
153
154 # Consume from buffer and yield
155 while True:
156 data = await buffer.get()
157 if data is None:
158 # End of stream
159 if producer_error:
160 raise producer_error
161 break
162 yield data
163
164 finally:
165 await asyncio.shield(_shutdown_producer(producer_task, buffer, cancelled))
166
167
168def use_buffer(
169 buffer_size: int = DEFAULT_BUFFER_SIZE,
170 min_buffer_before_yield: int = DEFAULT_MIN_BUFFER_BEFORE_YIELD,
171) -> Callable[
172 [Callable[_P, AsyncGenerator[bytes]]],
173 Callable[_P, AsyncGenerator[bytes]],
174]:
175 """
176 Add buffering to async generator functions that yield bytes (decorator).
177
178 This decorator uses an asyncio.Queue to decouple the producer (reading from the stream)
179 from the consumer (yielding to the client). The producer runs in a separate task and
180 fills the buffer, while the consumer yields from the buffer.
181
182 Args:
183 buffer_size: Maximum number of chunks to buffer (default: 30)
184 min_buffer_before_yield: Minimum chunks to buffer before starting to yield (default: 5)
185
186 Example:
187 @use_buffer(buffer_size=100)
188 async def my_stream() -> AsyncGenerator[bytes, None]:
189 ...
190 """
191
192 def decorator(
193 func: Callable[_P, AsyncGenerator[bytes]],
194 ) -> Callable[_P, AsyncGenerator[bytes]]:
195 @wraps(func)
196 async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> AsyncGenerator[bytes]:
197 async for chunk in buffered(
198 func(*args, **kwargs),
199 buffer_size=buffer_size,
200 min_buffer_before_yield=min_buffer_before_yield,
201 ):
202 yield chunk
203
204 return wrapper
205
206 return decorator
207