refreshable.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279
  1. import random
  2. import asyncio
  3. import threading
  4. import logging
  5. import time
  6. import atexit
  7. from datetime import datetime
  8. from enum import Enum
  9. from typing import Callable, Generic, TypeVar, Coroutine, Any
  10. from threading import Semaphore
  11. from concurrent.futures.thread import ThreadPoolExecutor
  12. from alibabacloud_credentials.exceptions import CredentialException
  13. from alibabacloud_credentials_api import ICredentials
  14. log = logging.getLogger('credentials')
  15. log.setLevel(logging.INFO)
  16. ch = logging.StreamHandler()
  17. log.addHandler(ch)
  18. T = TypeVar('T')
  19. INT64_MAX = 2 ** 63 - 1
  20. MAX_CONCURRENT_REFRESHES = 100
  21. CONCURRENT_REFRESH_LEASES = Semaphore(MAX_CONCURRENT_REFRESHES)
  22. EXECUTOR = ThreadPoolExecutor(max_workers=INT64_MAX, thread_name_prefix='non-blocking-refresh')
  23. def _shutdown_handler():
  24. EXECUTOR.shutdown(wait=False)
  25. atexit.register(_shutdown_handler)
  26. def _jitter_time(now: int, jitter_start: int, jitter_end: int) -> int:
  27. jitter_amount = random.randint(jitter_start, jitter_end)
  28. return now + jitter_amount
  29. def _max_stale_failure_jitter(num_failures: int) -> int:
  30. backoff_millis = max(10 * 1000, (1 << num_failures - 1) * 100)
  31. return backoff_millis
  32. class Credentials(ICredentials):
  33. def __init__(self, *,
  34. access_key_id: str = None,
  35. access_key_secret: str = None,
  36. security_token: str = None,
  37. expiration: int = None,
  38. provider_name: str = None):
  39. self._access_key_id = access_key_id
  40. self._access_key_secret = access_key_secret
  41. self._security_token = security_token
  42. self._expiration = expiration
  43. self._provider_name = provider_name
  44. def get_access_key_id(self) -> str:
  45. return self._access_key_id
  46. def get_access_key_secret(self) -> str:
  47. return self._access_key_secret
  48. def get_security_token(self) -> str:
  49. return self._security_token
  50. def get_expiration(self) -> int:
  51. return self._expiration
  52. def get_provider_name(self) -> str:
  53. return self._provider_name
  54. class StaleValueBehavior(Enum):
  55. """
  56. Strictly treat the stale time. Never return a stale cached value (except when the supplier returns an expired
  57. value, in which case the supplier will return the value but only for a very short period of time to prevent
  58. overloading the underlying supplier).
  59. """
  60. STRICT = 0
  61. """
  62. Allow stale values to be returned from the cache. Value retrieval will never fail, as long as the cache has
  63. succeeded when calling the underlying supplier at least once.
  64. """
  65. ALLOW = 1
  66. class RefreshResult(Generic[T]):
  67. def __init__(self, *,
  68. value: T,
  69. stale_time: int = INT64_MAX,
  70. prefetch_time: int = INT64_MAX):
  71. self._value = value
  72. self._stale_time = stale_time
  73. self._prefetch_time = prefetch_time
  74. def value(self) -> T:
  75. return self._value
  76. def stale_time(self) -> int:
  77. return self._stale_time
  78. def prefetch_time(self) -> int:
  79. return self._prefetch_time
  80. class PrefetchStrategy:
  81. def prefetch(self, action: Callable):
  82. raise NotImplementedError
  83. async def prefetch_async(self, action: Callable):
  84. raise NotImplementedError
  85. class NonBlocking(PrefetchStrategy):
  86. def prefetch(self, action: Callable):
  87. if not CONCURRENT_REFRESH_LEASES.acquire(False):
  88. log.warning('Skipping a background refresh task because there are too many other tasks running.')
  89. return
  90. try:
  91. EXECUTOR.submit(action)
  92. except KeyboardInterrupt:
  93. _shutdown_handler()
  94. except Exception as t:
  95. log.warning(f'Exception occurred when submitting background task.', exc_info=True)
  96. finally:
  97. CONCURRENT_REFRESH_LEASES.release()
  98. async def prefetch_async(self, action: Callable):
  99. def run_asyncio_loop():
  100. loop = asyncio.new_event_loop()
  101. asyncio.set_event_loop(loop)
  102. loop.run_until_complete(action())
  103. loop.close()
  104. self.prefetch(run_asyncio_loop)
  105. class OneCallerBlocks(PrefetchStrategy):
  106. def prefetch(self, action: Callable):
  107. action()
  108. async def prefetch_async(self, action: Callable):
  109. await action()
  110. class RefreshCachedSupplier(Generic[T]):
  111. STALE_TIME = 15 * 60 # seconds
  112. REFRESH_BLOCKING_MAX_WAIT = 5 # seconds
  113. def __init__(self, refresh_callable: Callable[[], RefreshResult[T]],
  114. refresh_callable_async: Callable[[], Coroutine[Any, Any, RefreshResult[T]]],
  115. stale_value_behavior: StaleValueBehavior = StaleValueBehavior.STRICT,
  116. prefetch_strategy: PrefetchStrategy = OneCallerBlocks()):
  117. self._refresh_callable = refresh_callable
  118. self._refresh_callable_async = refresh_callable_async
  119. self._stale_value_behavior = stale_value_behavior
  120. self._prefetch_strategy = prefetch_strategy
  121. self._consecutive_refresh_failures = 0
  122. self._cached_value = None
  123. self._refresh_lock = threading.Lock()
  124. def _sync_call(self) -> T:
  125. if self._cache_is_stale():
  126. log.debug('Refreshing synchronously')
  127. self._refresh_cache()
  128. elif self._should_initiate_cache_prefetch():
  129. log.debug(f'Prefetching using strategy: {self._prefetch_strategy.__class__.__name__}')
  130. self._prefetch_cache()
  131. return self._cached_value.value()
  132. async def _async_call(self) -> T:
  133. if self._cache_is_stale():
  134. log.debug('Refreshing synchronously')
  135. await self._refresh_cache_async()
  136. elif self._should_initiate_cache_prefetch():
  137. log.debug(f'Prefetching using strategy: {self._prefetch_strategy.__class__.__name__}')
  138. await self._prefetch_cache_async()
  139. return self._cached_value.value()
  140. def _cache_is_stale(self) -> bool:
  141. if self._cached_value is None:
  142. return True
  143. return int(time.mktime(time.localtime())) >= self._cached_value.stale_time()
  144. def _should_initiate_cache_prefetch(self) -> bool:
  145. if self._cached_value is None:
  146. return True
  147. return int(time.mktime(time.localtime())) >= self._cached_value.prefetch_time()
  148. def _prefetch_cache(self):
  149. self._prefetch_strategy.prefetch(self._refresh_cache)
  150. def _refresh_cache(self):
  151. acquired = self._refresh_lock.acquire(timeout=RefreshCachedSupplier.REFRESH_BLOCKING_MAX_WAIT)
  152. try:
  153. if self._cache_is_stale() or self._should_initiate_cache_prefetch():
  154. try:
  155. self._cached_value = self._handle_fetched_success(self._refresh_callable())
  156. except Exception as ex:
  157. self._cached_value = self._handle_fetched_failure(ex)
  158. finally:
  159. if acquired:
  160. self._refresh_lock.release()
  161. async def _prefetch_cache_async(self):
  162. await self._prefetch_strategy.prefetch_async(self._refresh_cache_async)
  163. async def _refresh_cache_async(self):
  164. acquired = self._refresh_lock.acquire(timeout=RefreshCachedSupplier.REFRESH_BLOCKING_MAX_WAIT)
  165. try:
  166. if self._cache_is_stale() or self._should_initiate_cache_prefetch():
  167. try:
  168. self._cached_value = self._handle_fetched_success(await self._refresh_callable_async())
  169. except Exception as ex:
  170. self._cached_value = self._handle_fetched_failure(ex)
  171. finally:
  172. if acquired:
  173. self._refresh_lock.release()
  174. def _handle_fetched_success(self, value: RefreshResult[T]) -> RefreshResult[T]:
  175. log.debug(f'Refresh credentials successfully, retrieved value is {value}, cached value is {self._cached_value}')
  176. self._consecutive_refresh_failures = 0
  177. now = int(time.mktime(time.localtime()))
  178. # 过期时间大于15分钟,不用管
  179. if now < value.stale_time():
  180. log.debug(
  181. f'Retrieved value stale time is {datetime.fromtimestamp(value.stale_time())}. Using staleTime of {datetime.fromtimestamp(value.stale_time())}')
  182. return value
  183. # 不足或等于15分钟,但未过期,下次会再次刷新
  184. if now < value.stale_time() + RefreshCachedSupplier.STALE_TIME:
  185. log.warning(
  186. f'Retrieved value stale time is in the past ({datetime.fromtimestamp(value.stale_time())}). Using staleTime of {datetime.fromtimestamp(now)}')
  187. return RefreshResult(value=value.value(), stale_time=now, prefetch_time=value.prefetch_time())
  188. log.warning(
  189. f'Retrieved value expiration time of the credential is in the past ({datetime.fromtimestamp(value.stale_time() + RefreshCachedSupplier.STALE_TIME)}). Trying use the cached value.')
  190. # 已过期,看缓存,缓存若大于15分钟,返回缓存,若小于15分钟,则根据策略判断是立刻重试还是稍后重试
  191. if self._cached_value is None:
  192. raise CredentialException('No cached value was found.')
  193. elif now < self._cached_value.stale_time():
  194. log.warning(
  195. f'Cached value staleTime is {datetime.fromtimestamp(self._cached_value.stale_time())}. Using staleTime of {datetime.fromtimestamp(self._cached_value.stale_time())}')
  196. return self._cached_value
  197. elif self._stale_value_behavior == StaleValueBehavior.STRICT:
  198. log.warning(
  199. f'Cached value expiration is in the past ({datetime.fromtimestamp(self._cached_value.stale_time())}). Using expiration of {datetime.fromtimestamp(now + 1)}')
  200. return RefreshResult(value=self._cached_value.value(), stale_time=now + 1,
  201. prefetch_time=self._cached_value.prefetch_time())
  202. else: # ALLOW
  203. extended_stale_time = now + int((50 * 1000 + random.randint(0, 20 * 1000 + 1)) / 1000)
  204. log.warning(
  205. f'Cached value expiration has been extended to {datetime.fromtimestamp(extended_stale_time)} because the downstream service returned a time in the past: {datetime.fromtimestamp(self._cached_value.stale_time())}')
  206. return RefreshResult(value=self._cached_value.value(), stale_time=extended_stale_time,
  207. prefetch_time=self._cached_value.prefetch_time())
  208. def _handle_fetched_failure(self, exception: Exception) -> RefreshResult[T]:
  209. log.warning(f'Refresh credentials failed, cached value is {self._cached_value}, error: {exception}')
  210. if not self._cached_value:
  211. log.exception(exception)
  212. raise exception
  213. now = int(time.mktime(time.localtime()))
  214. if now < self._cached_value.stale_time():
  215. return self._cached_value
  216. self._consecutive_refresh_failures += 1
  217. if self._stale_value_behavior == StaleValueBehavior.STRICT:
  218. log.exception(exception)
  219. raise exception
  220. else: # ALLOW
  221. new_stale_time = int(
  222. _jitter_time(now * 1000, 1000, _max_stale_failure_jitter(self._consecutive_refresh_failures)) / 1000)
  223. log.warning(
  224. f'Cached value expiration has been extended to {datetime.fromtimestamp(new_stale_time)} because calling the downstream service failed (consecutive failures: {self._consecutive_refresh_failures}).')
  225. return RefreshResult(value=self._cached_value.value(), stale_time=new_stale_time,
  226. prefetch_time=self._cached_value.prefetch_time())