/
/
/
1"""Context manager using asyncio_throttle that catches and re-raises RetriesExhausted."""
2
3import asyncio
4import functools
5import logging
6import random
7import time
8from collections import deque
9from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
10from contextlib import asynccontextmanager
11from contextvars import ContextVar
12from email.utils import parsedate_to_datetime
13from types import TracebackType
14from typing import Any, Concatenate, Protocol
15
16from music_assistant_models.errors import (
17 RateLimited,
18 ResourceTemporarilyUnavailable,
19 RetriesExhausted,
20)
21
22from music_assistant.constants import MASS_LOGGER_NAME
23from music_assistant.helpers.datetime import utc
24
25LOGGER = logging.getLogger(f"{MASS_LOGGER_NAME}.throttle_retry")
26
27BYPASS_THROTTLER: ContextVar[bool] = ContextVar("BYPASS_THROTTLER", default=False)
28
29# Cap exponential backoff to prevent absurd wait times
30MAX_BACKOFF = 120
31
32# Cap a server-provided Retry-After, in case it is absurd or hostile
33MAX_RETRY_AFTER = 3600
34
35
36def parse_retry_after(value: str | None) -> int:
37 """
38 Parse a Retry-After header value per RFC 9110 Section 10.2.3.
39
40 Supports both valid formats: delay-seconds (integer) and HTTP-date.
41
42 :param value: The raw Retry-After header value, or None if absent.
43 :returns: Non-negative integer seconds to wait, or 0 if unparsable/absent.
44 """
45 if value is None:
46 return 0
47 # Try delay-seconds (non-negative integer) first â the common case
48 try:
49 return max(0, int(value))
50 except ValueError, TypeError:
51 pass
52 # Try HTTP-date format (e.g., "Fri, 31 Dec 1999 23:59:59 GMT")
53 try:
54 target = parsedate_to_datetime(value)
55 delta = (target - utc()).total_seconds()
56 return max(0, int(delta))
57 except ValueError, TypeError:
58 return 0
59
60
61class Throttler:
62 """
63 asyncio_throttle (https://github.com/hallazzang/asyncio-throttle).
64
65 With improvements:
66 - Accurate sleep without "busy waiting" (PR #4)
67 - Return the delay caused by acquire()
68 """
69
70 def __init__(self, rate_limit: int, period: float = 1.0) -> None:
71 """Initialize the Throttler."""
72 self.rate_limit = rate_limit
73 self.period = period
74 self._task_logs: deque[float] = deque()
75
76 async def acquire(self) -> float:
77 """Acquire a free slot from the Throttler, returns the throttled time."""
78 cur_time = time.monotonic()
79 start_time = cur_time
80 while True:
81 self._flush()
82 if len(self._task_logs) < self.rate_limit:
83 break
84 # sleep the exact amount of time until the oldest task can be flushed
85 time_to_release = self._task_logs[0] + self.period - cur_time
86 await asyncio.sleep(time_to_release)
87 cur_time = time.monotonic()
88
89 self._task_logs.append(cur_time)
90 return cur_time - start_time # exactly 0 if not throttled
91
92 async def __aenter__(self) -> float:
93 """Wait until the lock is acquired, return the time delay."""
94 return await self.acquire()
95
96 async def __aexit__(
97 self,
98 exc_type: type[BaseException] | None,
99 exc_val: BaseException | None,
100 exc_tb: TracebackType | None,
101 ) -> bool | None:
102 """Nothing to do on exit."""
103
104 def _flush(self) -> None:
105 now = time.monotonic()
106 while self._task_logs:
107 if now - self._task_logs[0] > self.period:
108 self._task_logs.popleft()
109 else:
110 break
111
112
113class ThrottlerManager:
114 """Throttler manager that extends asyncio Throttle by retrying."""
115
116 def __init__(
117 self, rate_limit: int, period: float = 1, retry_attempts: int = 5, initial_backoff: int = 5
118 ):
119 """Initialize the AsyncThrottledContextManager."""
120 self.retry_attempts = retry_attempts
121 self.initial_backoff = initial_backoff
122 self.throttler = Throttler(rate_limit, period)
123
124 @asynccontextmanager
125 async def acquire(self) -> AsyncGenerator[float]:
126 """Acquire a free slot from the Throttler, returns the throttled time."""
127 if BYPASS_THROTTLER.get():
128 yield 0
129 else:
130 yield await self.throttler.acquire()
131
132 @asynccontextmanager
133 async def bypass(self) -> AsyncGenerator[None]:
134 """Bypass the throttler."""
135 try:
136 token = BYPASS_THROTTLER.set(True)
137 yield None
138 finally:
139 BYPASS_THROTTLER.reset(token)
140
141
142class _Throttleable(Protocol):
143 """Protocol for objects that can use the @throttle_with_retries decorator."""
144
145 @property
146 def logger(self) -> logging.Logger: ...
147
148 @property
149 def throttler(self) -> ThrottlerManager: ...
150
151
152def throttle_with_retries[ProviderT: _Throttleable, **P, R](
153 func: Callable[Concatenate[ProviderT, P], Awaitable[R]],
154) -> Callable[Concatenate[ProviderT, P], Coroutine[Any, Any, R]]:
155 """Call async function using the throttler with retries."""
156
157 @functools.wraps(func)
158 async def wrapper(self: ProviderT, *args: P.args, **kwargs: P.kwargs) -> R:
159 """Call async function using the throttler with retries."""
160 throttler = self.throttler
161 exp_backoff = throttler.initial_backoff
162 async with throttler.acquire() as delay:
163 if delay != 0:
164 self.logger.debug(
165 "%s was delayed for %.3f secs due to throttling", func.__name__, delay
166 )
167 for attempt in range(throttler.retry_attempts):
168 try:
169 return await func(self, *args, **kwargs)
170 except ResourceTemporarilyUnavailable as e:
171 self.logger.info(
172 f"Attempt {attempt + 1}/{throttler.retry_attempts} failed: {e}"
173 )
174 if attempt < throttler.retry_attempts - 1:
175 server_wait = min(max(float(e.backoff_time), 0.0), MAX_RETRY_AFTER)
176 if isinstance(e, RateLimited):
177 # Retry-After is a floor, not a target: escalate above it,
178 # jittering up only so we never retry sooner than asked
179 base = max(server_wait, min(exp_backoff, MAX_BACKOFF))
180 sleep_time = base * random.uniform(1.0, 1.1)
181 exp_backoff = min(exp_backoff * 2, MAX_BACKOFF)
182 elif server_wait:
183 # Server named a recovery time â respect it, with citizen jitter
184 sleep_time = server_wait * random.uniform(1.0, 1.1)
185 else:
186 # No server guidance â exponential backoff with jitter
187 sleep_time = min(exp_backoff * random.uniform(0.75, 1.25), MAX_BACKOFF)
188 exp_backoff = min(exp_backoff * 2, MAX_BACKOFF)
189 self.logger.info(f"Retrying in {sleep_time:.1f} seconds...")
190 await asyncio.sleep(sleep_time)
191 else: # noqa: PLW0120
192 msg = f"Retries exhausted, failed after {throttler.retry_attempts} attempts"
193 raise RetriesExhausted(msg)
194
195 return wrapper
196