| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279 |
- import random
- import asyncio
- import threading
- import logging
- import time
- import atexit
- from datetime import datetime
- from enum import Enum
- from typing import Callable, Generic, TypeVar, Coroutine, Any
- from threading import Semaphore
- from concurrent.futures.thread import ThreadPoolExecutor
- from alibabacloud_credentials.exceptions import CredentialException
- from alibabacloud_credentials_api import ICredentials
- log = logging.getLogger('credentials')
- log.setLevel(logging.INFO)
- ch = logging.StreamHandler()
- log.addHandler(ch)
- T = TypeVar('T')
- INT64_MAX = 2 ** 63 - 1
- MAX_CONCURRENT_REFRESHES = 100
- CONCURRENT_REFRESH_LEASES = Semaphore(MAX_CONCURRENT_REFRESHES)
- EXECUTOR = ThreadPoolExecutor(max_workers=INT64_MAX, thread_name_prefix='non-blocking-refresh')
- def _shutdown_handler():
- EXECUTOR.shutdown(wait=False)
- atexit.register(_shutdown_handler)
- def _jitter_time(now: int, jitter_start: int, jitter_end: int) -> int:
- jitter_amount = random.randint(jitter_start, jitter_end)
- return now + jitter_amount
- def _max_stale_failure_jitter(num_failures: int) -> int:
- backoff_millis = max(10 * 1000, (1 << num_failures - 1) * 100)
- return backoff_millis
- class Credentials(ICredentials):
- def __init__(self, *,
- access_key_id: str = None,
- access_key_secret: str = None,
- security_token: str = None,
- expiration: int = None,
- provider_name: str = None):
- self._access_key_id = access_key_id
- self._access_key_secret = access_key_secret
- self._security_token = security_token
- self._expiration = expiration
- self._provider_name = provider_name
- def get_access_key_id(self) -> str:
- return self._access_key_id
- def get_access_key_secret(self) -> str:
- return self._access_key_secret
- def get_security_token(self) -> str:
- return self._security_token
- def get_expiration(self) -> int:
- return self._expiration
- def get_provider_name(self) -> str:
- return self._provider_name
- class StaleValueBehavior(Enum):
- """
- Strictly treat the stale time. Never return a stale cached value (except when the supplier returns an expired
- value, in which case the supplier will return the value but only for a very short period of time to prevent
- overloading the underlying supplier).
- """
- STRICT = 0
- """
- Allow stale values to be returned from the cache. Value retrieval will never fail, as long as the cache has
- succeeded when calling the underlying supplier at least once.
- """
- ALLOW = 1
- class RefreshResult(Generic[T]):
- def __init__(self, *,
- value: T,
- stale_time: int = INT64_MAX,
- prefetch_time: int = INT64_MAX):
- self._value = value
- self._stale_time = stale_time
- self._prefetch_time = prefetch_time
- def value(self) -> T:
- return self._value
- def stale_time(self) -> int:
- return self._stale_time
- def prefetch_time(self) -> int:
- return self._prefetch_time
- class PrefetchStrategy:
- def prefetch(self, action: Callable):
- raise NotImplementedError
- async def prefetch_async(self, action: Callable):
- raise NotImplementedError
- class NonBlocking(PrefetchStrategy):
- def prefetch(self, action: Callable):
- if not CONCURRENT_REFRESH_LEASES.acquire(False):
- log.warning('Skipping a background refresh task because there are too many other tasks running.')
- return
- try:
- EXECUTOR.submit(action)
- except KeyboardInterrupt:
- _shutdown_handler()
- except Exception as t:
- log.warning(f'Exception occurred when submitting background task.', exc_info=True)
- finally:
- CONCURRENT_REFRESH_LEASES.release()
- async def prefetch_async(self, action: Callable):
- def run_asyncio_loop():
- loop = asyncio.new_event_loop()
- asyncio.set_event_loop(loop)
- loop.run_until_complete(action())
- loop.close()
- self.prefetch(run_asyncio_loop)
- class OneCallerBlocks(PrefetchStrategy):
- def prefetch(self, action: Callable):
- action()
- async def prefetch_async(self, action: Callable):
- await action()
- class RefreshCachedSupplier(Generic[T]):
- STALE_TIME = 15 * 60 # seconds
- REFRESH_BLOCKING_MAX_WAIT = 5 # seconds
- def __init__(self, refresh_callable: Callable[[], RefreshResult[T]],
- refresh_callable_async: Callable[[], Coroutine[Any, Any, RefreshResult[T]]],
- stale_value_behavior: StaleValueBehavior = StaleValueBehavior.STRICT,
- prefetch_strategy: PrefetchStrategy = OneCallerBlocks()):
- self._refresh_callable = refresh_callable
- self._refresh_callable_async = refresh_callable_async
- self._stale_value_behavior = stale_value_behavior
- self._prefetch_strategy = prefetch_strategy
- self._consecutive_refresh_failures = 0
- self._cached_value = None
- self._refresh_lock = threading.Lock()
- def _sync_call(self) -> T:
- if self._cache_is_stale():
- log.debug('Refreshing synchronously')
- self._refresh_cache()
- elif self._should_initiate_cache_prefetch():
- log.debug(f'Prefetching using strategy: {self._prefetch_strategy.__class__.__name__}')
- self._prefetch_cache()
- return self._cached_value.value()
- async def _async_call(self) -> T:
- if self._cache_is_stale():
- log.debug('Refreshing synchronously')
- await self._refresh_cache_async()
- elif self._should_initiate_cache_prefetch():
- log.debug(f'Prefetching using strategy: {self._prefetch_strategy.__class__.__name__}')
- await self._prefetch_cache_async()
- return self._cached_value.value()
- def _cache_is_stale(self) -> bool:
- if self._cached_value is None:
- return True
- return int(time.mktime(time.localtime())) >= self._cached_value.stale_time()
- def _should_initiate_cache_prefetch(self) -> bool:
- if self._cached_value is None:
- return True
- return int(time.mktime(time.localtime())) >= self._cached_value.prefetch_time()
- def _prefetch_cache(self):
- self._prefetch_strategy.prefetch(self._refresh_cache)
- def _refresh_cache(self):
- acquired = self._refresh_lock.acquire(timeout=RefreshCachedSupplier.REFRESH_BLOCKING_MAX_WAIT)
- try:
- if self._cache_is_stale() or self._should_initiate_cache_prefetch():
- try:
- self._cached_value = self._handle_fetched_success(self._refresh_callable())
- except Exception as ex:
- self._cached_value = self._handle_fetched_failure(ex)
- finally:
- if acquired:
- self._refresh_lock.release()
- async def _prefetch_cache_async(self):
- await self._prefetch_strategy.prefetch_async(self._refresh_cache_async)
- async def _refresh_cache_async(self):
- acquired = self._refresh_lock.acquire(timeout=RefreshCachedSupplier.REFRESH_BLOCKING_MAX_WAIT)
- try:
- if self._cache_is_stale() or self._should_initiate_cache_prefetch():
- try:
- self._cached_value = self._handle_fetched_success(await self._refresh_callable_async())
- except Exception as ex:
- self._cached_value = self._handle_fetched_failure(ex)
- finally:
- if acquired:
- self._refresh_lock.release()
- def _handle_fetched_success(self, value: RefreshResult[T]) -> RefreshResult[T]:
- log.debug(f'Refresh credentials successfully, retrieved value is {value}, cached value is {self._cached_value}')
- self._consecutive_refresh_failures = 0
- now = int(time.mktime(time.localtime()))
- # 过期时间大于15分钟,不用管
- if now < value.stale_time():
- log.debug(
- f'Retrieved value stale time is {datetime.fromtimestamp(value.stale_time())}. Using staleTime of {datetime.fromtimestamp(value.stale_time())}')
- return value
- # 不足或等于15分钟,但未过期,下次会再次刷新
- if now < value.stale_time() + RefreshCachedSupplier.STALE_TIME:
- log.warning(
- f'Retrieved value stale time is in the past ({datetime.fromtimestamp(value.stale_time())}). Using staleTime of {datetime.fromtimestamp(now)}')
- return RefreshResult(value=value.value(), stale_time=now, prefetch_time=value.prefetch_time())
- log.warning(
- 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.')
- # 已过期,看缓存,缓存若大于15分钟,返回缓存,若小于15分钟,则根据策略判断是立刻重试还是稍后重试
- if self._cached_value is None:
- raise CredentialException('No cached value was found.')
- elif now < self._cached_value.stale_time():
- log.warning(
- f'Cached value staleTime is {datetime.fromtimestamp(self._cached_value.stale_time())}. Using staleTime of {datetime.fromtimestamp(self._cached_value.stale_time())}')
- return self._cached_value
- elif self._stale_value_behavior == StaleValueBehavior.STRICT:
- log.warning(
- f'Cached value expiration is in the past ({datetime.fromtimestamp(self._cached_value.stale_time())}). Using expiration of {datetime.fromtimestamp(now + 1)}')
- return RefreshResult(value=self._cached_value.value(), stale_time=now + 1,
- prefetch_time=self._cached_value.prefetch_time())
- else: # ALLOW
- extended_stale_time = now + int((50 * 1000 + random.randint(0, 20 * 1000 + 1)) / 1000)
- log.warning(
- 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())}')
- return RefreshResult(value=self._cached_value.value(), stale_time=extended_stale_time,
- prefetch_time=self._cached_value.prefetch_time())
- def _handle_fetched_failure(self, exception: Exception) -> RefreshResult[T]:
- log.warning(f'Refresh credentials failed, cached value is {self._cached_value}, error: {exception}')
- if not self._cached_value:
- log.exception(exception)
- raise exception
- now = int(time.mktime(time.localtime()))
- if now < self._cached_value.stale_time():
- return self._cached_value
- self._consecutive_refresh_failures += 1
- if self._stale_value_behavior == StaleValueBehavior.STRICT:
- log.exception(exception)
- raise exception
- else: # ALLOW
- new_stale_time = int(
- _jitter_time(now * 1000, 1000, _max_stale_failure_jitter(self._consecutive_refresh_failures)) / 1000)
- log.warning(
- 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}).')
- return RefreshResult(value=self._cached_value.value(), stale_time=new_stale_time,
- prefetch_time=self._cached_value.prefetch_time())
|