280 lines
11 KiB
Python
280 lines
11 KiB
Python
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())
|