Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/cache_coordinator.py: 21%

1""" 

2Event-driven cache coordinator to prevent cache stampede. 

3 

4Use this when many requests can miss the same cache key at once (e.g. after 

5expiry or restart). Without coordination, they would all run the expensive 

6load (DB query, API call) in parallel and overload the backend. 

7 

8This module ensures only one request performs the load; the rest wait for a 

9signal and then read the freshly cached value. Reuse it for any cache-aside 

10pattern: global spend, feature flags, config, or other shared read-through data. 

11""" 

12 

13import asyncio 

14import time 

15from collections.abc import Awaitable, Callable 

16from typing import Final, Protocol, TypeVar 

17 

18from litellm._logging import verbose_proxy_logger 

19 

20T = TypeVar("T") 

21 

22 

23class AsyncCacheProtocol(Protocol[T]): 

24 """Protocol for cache backends used by EventDrivenCacheCoordinator. 

25 

26 Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params 

27 before ``**kwargs``), not only ``(key, **kwargs)``, so overloads validate. 

28 """ 

29 

30 async def async_get_cache( 30 ↛ exitline 30 didn't return from function 'async_get_cache' because

31 self, 

32 key: str, 

33 parent_otel_span: object = None, 

34 local_only: bool = False, 

35 **kwargs: object, 

36 ) -> T | None: ... 

37 

38 async def async_set_cache( 38 ↛ exitline 38 didn't return from function 'async_set_cache' because

39 self, 

40 key: str, 

41 value: T, 

42 local_only: bool = False, 

43 **kwargs: object, 

44 ) -> object: ... 

45 

46 

47class EventDrivenCacheCoordinator: 

48 """ 

49 Coordinates a single in-flight load per logical resource to prevent cache stampede. 

50 

51 Pattern: 

52 - First request: loads data (e.g. DB query), caches it, then signals waiters. 

53 - Other requests: wait for the signal, then read from cache. 

54 

55 Create one instance per resource (e.g. one for global spend, one for feature flags). 

56 

57 Args: 

58 log_prefix: Prefix for debug log messages. 

59 """ 

60 

61 def __init__(self, log_prefix: str = "[CACHE]"): 

62 self._lock = asyncio.Lock() 

63 self._event: asyncio.Event | None = None 

64 self._query_in_progress = False 

65 self._log_prefix = log_prefix 

66 

67 async def _get_cached(self, cache_key: str, cache: AsyncCacheProtocol[T]) -> T | None: 

68 """Return value from cache if present, else None.""" 

69 return await cache.async_get_cache(key=cache_key) 

70 

71 def _log_cache_hit(self, value: object) -> None: 

72 if self._log_prefix: 

73 verbose_proxy_logger.debug("%s Cache hit, value: %s", self._log_prefix, value) 

74 

75 def _log_cache_miss(self) -> None: 

76 if self._log_prefix: 

77 verbose_proxy_logger.debug("%s Cache miss", self._log_prefix) 

78 

79 async def _claim_role(self) -> asyncio.Event | None: 

80 """ 

81 Under lock: return event to wait on if load is in progress, else set us as loader and return None. 

82 """ 

83 async with self._lock: 

84 if self._query_in_progress and self._event is not None: 

85 if self._log_prefix: 

86 verbose_proxy_logger.debug("%s Load in flight, waiting for signal", self._log_prefix) 

87 return self._event 

88 self._query_in_progress = True 

89 self._event = asyncio.Event() 

90 if self._log_prefix: 

91 verbose_proxy_logger.debug( 

92 "%s Starting load (will signal others when done)", 

93 self._log_prefix, 

94 ) 

95 return None 

96 

97 async def _wait_for_signal_and_get( 

98 self, 

99 event: asyncio.Event, 

100 cache_key: str, 

101 cache: AsyncCacheProtocol[T], 

102 ) -> T | None: 

103 """Wait for loader to finish, then read from cache.""" 

104 await event.wait() 

105 if self._log_prefix: 

106 verbose_proxy_logger.debug("%s Signal received, reading from cache", self._log_prefix) 

107 value: Final[T | None] = await cache.async_get_cache(key=cache_key) 

108 if value is not None and self._log_prefix: 

109 verbose_proxy_logger.debug( 

110 "%s Cache filled by other request, value: %s", 

111 self._log_prefix, 

112 value, 

113 ) 

114 elif value is None and self._log_prefix: 

115 verbose_proxy_logger.debug("%s Signal received but cache still empty", self._log_prefix) 

116 return value 

117 

118 async def _load_and_cache( 

119 self, 

120 cache_key: str, 

121 cache: AsyncCacheProtocol[T], 

122 load_fn: Callable[[], Awaitable[T]], 

123 ) -> T | None: 

124 """Double-check cache, run load_fn, set cache, return value. Caller must call _signal_done in finally.""" 

125 value = await cache.async_get_cache(key=cache_key) 

126 if value is not None: 

127 if self._log_prefix: 

128 verbose_proxy_logger.debug( 

129 "%s Cache filled while acquiring lock, value: %s", 

130 self._log_prefix, 

131 value, 

132 ) 

133 return value 

134 

135 if self._log_prefix: 

136 verbose_proxy_logger.debug("%s Running load", self._log_prefix) 

137 start: Final = time.perf_counter() 

138 value = await load_fn() 

139 elapsed_ms: Final = (time.perf_counter() - start) * 1000 

140 if self._log_prefix: 

141 verbose_proxy_logger.debug( 

142 "%s Load completed in %.2fms, result: %s", 

143 self._log_prefix, 

144 elapsed_ms, 

145 value, 

146 ) 

147 

148 await cache.async_set_cache(key=cache_key, value=value) 

149 if self._log_prefix: 

150 verbose_proxy_logger.debug("%s Result cached", self._log_prefix) 

151 return value 

152 

153 async def _signal_done(self) -> None: 

154 """Reset loader state and signal all waiters.""" 

155 async with self._lock: 

156 self._query_in_progress = False 

157 if self._event is not None: 

158 if self._log_prefix: 

159 verbose_proxy_logger.debug("%s Signaling all waiting requests", self._log_prefix) 

160 self._event.set() 

161 self._event = None 

162 

163 async def get_or_load( 

164 self, 

165 cache_key: str, 

166 cache: AsyncCacheProtocol[T], 

167 load_fn: Callable[[], Awaitable[T]], 

168 ) -> T | None: 

169 """ 

170 Return cached value or load it once and signal waiters. 

171 

172 - cache_key: Key to read/write in the cache. 

173 - cache: Object with async_get_cache(key) and async_set_cache(key, value). 

174 - load_fn: Async callable that performs the load (e.g. DB query). No args. 

175 Return value is cached and returned. If it raises, waiters are 

176 still signaled so they can retry or handle empty cache. 

177 

178 Returns the value from cache or from load_fn, or None if load failed or 

179 cache was still empty after waiting. 

180 """ 

181 value: Final = await self._get_cached(cache_key, cache) 

182 if value is not None: 

183 self._log_cache_hit(value) 

184 return value 

185 

186 self._log_cache_miss() 

187 event_to_wait: Final = await self._claim_role() 

188 

189 if event_to_wait is not None: 

190 return await self._wait_for_signal_and_get(event_to_wait, cache_key, cache) 

191 

192 try: 

193 result: Final = await self._load_and_cache(cache_key, cache, load_fn) 

194 return result 

195 finally: 

196 await self._signal_done()