class AsyncIteratorPool(collabc.AsyncGenerator, AsyncPool[T]):
"""An asynchronous pool that wraps another async iterator."""
def __init__(self, base: AsyncIterator[T]):
self._base = base
self._basegen = isinstance(base, collabc.AsyncGenerator)
self._yl = deque()
self._rl = deque()
self._ac = None
self._stopped = False
async def asend(self, value: OpT) -> OpT:
if self._stopped: raise StopAsyncIteration
if self._ac is None: self._ac = ACondition()
if value is None:
async with self._ac:
if len(self._rl) == 0:
try:
if self._basegen:
yv = await self._base.asend(None)
else:
yv = await self._base.__anext__()
self._yl.append(yv)
return yv
except (Exception, GeneratorExit) as exc:
self._stopped = True
self._ac.notify_all()
raise StopAsyncIteration from exc
else:
yv = self._rl.popleft()
self._yl.append(yv)
return yv
else:
async with self._ac:
if value in self._yl:
self._yl.remove(value)
self._rl.append(value)
async def athrow(self, typ, val=None, tb=None):
try:
if self._basegen:
return await self._base.athrow(typ, val, tb)
else:
return await super().athrow(typ, val, tb)
except (Exception, GeneratorExit) as exc:
self._stopped = True
if self._ac is not None:
await self._ac.acquire()
self._ac.notify_all()
self._ac.release()
raise StopAsyncIteration from exc
class SetBasedAsyncPopPool(collabc.AsyncGenerator, AsyncPopulationPool[H]):
"""An asynchronous population pool backed by a set."""
def __init__(self, ivals: Optional[AbstractSet[H]] = None):
self._stopped = False
self._set = Set[H]() if ivals is None else {ivals}
self._ac = None
self._stopped = False
async def apopulate(self, val: H, *args: H) -> None:
#Can't populate a closed pool:
if self._stopped: raise StopAsyncIteration
if self._ac is None: self._ac = ACondition()
if args is None or len(args) == 0:
self._set.add(val)
count = 1
else:
argset = {args}
argset.add(val)
count = len(argset)
self._set |= argset
async with self._ac:
self._ac.notify(count)
async def asend(self, value: Optional[H]) -> Optional[H]:
if self._stopped: raise StopAsyncIteration
if self._ac is None: self._ac = ACondition()
if value is None:
async with self._ac:
while len(self._set) == 0:
self._ac.wait()
if self._stopped: raise StopAsyncIteration
return self._set.pop()
else:
if value not in self._set:
self._set.add(value)
await self._ac.acquire()
self._ac.notify()
self._ac.release()
async def athrow(self, typ, val=None, tb=None) -> None:
try:
return await super().athrow(typ, val, tb)
except (Exception, GeneratorExit) as exc:
self._stopped = True
if self._ac is not None:
await self._ac.acquire()
self._ac.notify_all()
self._ac.release()
raise StopAsyncIteration from exc
class CommunicationTokenCredential(object):
"""Credential type used for authenticating to an Azure Communication service.
:param str token: The token used to authenticate to an Azure Communication service
:keyword token_refresher: The token refresher to provide capacity to fetch fresh token
:raises: TypeError
"""
_ON_DEMAND_REFRESHING_INTERVAL_MINUTES = 2
def __init__(self,
token, # type: str
**kwargs
):
token_refresher = kwargs.pop('token_refresher', None)
communication_token_refresh_options = CommunicationTokenRefreshOptions(token=token,
token_refresher=token_refresher)
self._token = communication_token_refresh_options.get_token()
self._token_refresher = communication_token_refresh_options.get_token_refresher()
self._lock = Condition(Lock())
self._some_thread_refreshing = False
def get_token(self):
# type () -> ~azure.core.credentials.AccessToken
"""The value of the configured token.
:rtype: ~azure.core.credentials.AccessToken
"""
if not self._token_refresher or not self._token_expiring():
return self._token
should_this_thread_refresh = False
with self._lock:
while self._token_expiring():
if self._some_thread_refreshing:
if self._is_currenttoken_valid():
return self._token
self._wait_till_inprogress_thread_finish_refreshing()
else:
should_this_thread_refresh = True
self._some_thread_refreshing = True
break
if should_this_thread_refresh:
try:
newtoken = self._token_refresher() # pylint:disable=not-callable
with self._lock:
self._token = newtoken
self._some_thread_refreshing = False
self._lock.notify_all()
except:
with self._lock:
self._some_thread_refreshing = False
self._lock.notify_all()
raise
return self._token
def _wait_till_inprogress_thread_finish_refreshing(self):
self._lock.release()
self._lock.acquire()
def _token_expiring(self):
return self._token.expires_on - self._get_utc_now_as_int() <\
timedelta(minutes=self._ON_DEMAND_REFRESHING_INTERVAL_MINUTES).total_seconds()
def _is_currenttoken_valid(self):
return self._get_utc_now_as_int() < self._token.expires_on
@classmethod
def _get_utc_now_as_int(cls):
current_utc_datetime = datetime.utcnow().replace(tzinfo=TZ_UTC)
current_utc_datetime_as_int = _convert_datetime_to_utc_int(current_utc_datetime)
return current_utc_datetime_as_int
class CommunicationTokenCredential(object):
"""Credential type used for authenticating to an Azure Communication service.
:param str token: The token used to authenticate to an Azure Communication service
:raises: TypeError
"""
ON_DEMAND_REFRESHING_INTERVAL_MINUTES = 2
def __init__(
self,
token, # type: str
token_refresher=None):
# type: (str) -> None
if not isinstance(token, six.string_types):
raise TypeError("token must be a string.")
self._token = create_access_token(token)
self._token_refresher = token_refresher
self._lock = Condition(Lock())
self._some_thread_refreshing = False
def get_token(self):
# type () -> ~azure.core.credentials.AccessToken
"""The value of the configured token.
:rtype: ~azure.core.credentials.AccessToken
"""
if not self._token_refresher or not self._token_expiring():
return self._token
should_this_thread_refresh = False
with self._lock:
while self._token_expiring():
if self._some_thread_refreshing:
if self._is_currenttoken_valid():
return self._token
self._wait_till_inprogress_thread_finish_refreshing()
else:
should_this_thread_refresh = True
self._some_thread_refreshing = True
break
if should_this_thread_refresh:
try:
newtoken = self._token_refresher()
with self._lock:
self._token = newtoken
self._some_thread_refreshing = False
self._lock.notify_all()
except:
with self._lock:
self._some_thread_refreshing = False
self._lock.notify_all()
raise
return self._token
def _wait_till_inprogress_thread_finish_refreshing(self):
self._lock.release()
self._lock.acquire()
def _token_expiring(self):
return self._token.expires_on - self._get_utc_now() <\
timedelta(minutes=self.ON_DEMAND_REFRESHING_INTERVAL_MINUTES)
def _is_currenttoken_valid(self):
return self._get_utc_now() < self._token.expires_on
@classmethod
def _get_utc_now(cls):
return datetime.now().replace(tzinfo=TZ_UTC)
class CommunicationTokenCredential(object):
"""Credential type used for authenticating to an Azure Communication service.
:param str token: The token used to authenticate to an Azure Communication service.
:keyword token_refresher: The async token refresher to provide capacity to fetch a fresh token.
The returned token must be valid (expiration date must be in the future).
:paramtype token_refresher: Callable[[], Awaitable[AccessToken]]
:keyword bool proactive_refresh: Whether to refresh the token proactively or not.
If the proactive refreshing is enabled ('proactive_refresh' is true), the credential will use
a background thread to attempt to refresh the token within 10 minutes before the cached token expires,
the proactive refresh will request a new token by calling the 'token_refresher' callback.
When 'proactive_refresh is enabled', the Credential object must be either run within a context manager
or the 'close' method must be called once the object usage has been finished.
:raises: TypeError if paramater 'token' is not a string
:raises: ValueError if the 'proactive_refresh' is enabled without providing the 'token_refresher' function.
"""
_ON_DEMAND_REFRESHING_INTERVAL_MINUTES = 2
_DEFAULT_AUTOREFRESH_INTERVAL_MINUTES = 10
def __init__(self, token: str, **kwargs: Any):
if not isinstance(token, six.string_types):
raise TypeError("Token must be a string.")
self._token = create_access_token(token)
self._token_refresher = kwargs.pop('token_refresher', None)
self._proactive_refresh = kwargs.pop('proactive_refresh', False)
if (self._proactive_refresh and self._token_refresher is None):
raise ValueError(
"When 'proactive_refresh' is True, 'token_refresher' must not be None."
)
self._timer = None
self._async_mutex = Lock()
if sys.version_info[:3] == (3, 10, 0):
# Workaround for Python 3.10 bug(https://bugs.python.org/issue45416):
getattr(self._async_mutex, '_get_loop', lambda: None)()
self._lock = Condition(self._async_mutex)
self._some_thread_refreshing = False
self._is_closed = Event()
async def get_token(self, *scopes, **kwargs): # pylint: disable=unused-argument
# type (*str, **Any) -> AccessToken
"""The value of the configured token.
:rtype: ~azure.core.credentials.AccessToken
"""
if self._proactive_refresh and self._is_closed.is_set():
raise RuntimeError(
"An instance of CommunicationTokenCredential cannot be reused once it has been closed."
)
if not self._token_refresher or not self._is_token_expiring_soon(
self._token):
return self._token
await self._update_token_and_reschedule()
return self._token
async def _update_token_and_reschedule(self):
should_this_thread_refresh = False
async with self._lock:
while self._is_token_expiring_soon(self._token):
if self._some_thread_refreshing:
if self._is_token_valid(self._token):
return self._token
await self._wait_till_lock_owner_finishes_refreshing()
else:
should_this_thread_refresh = True
self._some_thread_refreshing = True
break
if should_this_thread_refresh:
try:
new_token = await self._token_refresher()
if not self._is_token_valid(new_token):
raise ValueError(
"The token returned from the token_refresher is expired."
)
async with self._lock:
self._token = new_token
self._some_thread_refreshing = False
self._lock.notify_all()
except:
async with self._lock:
self._some_thread_refreshing = False
self._lock.notify_all()
raise
if self._proactive_refresh:
self._schedule_refresh()
return self._token
def _schedule_refresh(self):
if self._is_closed.is_set():
return
if self._timer is not None:
self._timer.cancel()
token_ttl = self._token.expires_on - get_current_utc_as_int()
if self._is_token_expiring_soon(self._token):
# Schedule the next refresh for when it reaches a certain percentage of the remaining lifetime.
timespan = token_ttl // 2
else:
# Schedule the next refresh for when it gets in to the soon-to-expire window.
timespan = token_ttl - timedelta(
minutes=self._DEFAULT_AUTOREFRESH_INTERVAL_MINUTES
).total_seconds()
self._timer = AsyncTimer(timespan, self._update_token_and_reschedule)
self._timer.start()
async def _wait_till_lock_owner_finishes_refreshing(self):
self._lock.release()
await self._lock.acquire()
def _is_token_expiring_soon(self, token):
if self._proactive_refresh:
interval = timedelta(
minutes=self._DEFAULT_AUTOREFRESH_INTERVAL_MINUTES)
else:
interval = timedelta(
minutes=self._ON_DEMAND_REFRESHING_INTERVAL_MINUTES)
return ((token.expires_on - get_current_utc_as_int()) <
interval.total_seconds())
@classmethod
def _is_token_valid(cls, token):
return get_current_utc_as_int() < token.expires_on
async def __aenter__(self):
if self._proactive_refresh:
if self._is_closed.is_set():
raise RuntimeError(
"An instance of CommunicationTokenCredential cannot be reused once it has been closed."
)
self._schedule_refresh()
return self
async def __aexit__(self, *args):
await self.close()
async def close(self) -> None:
if self._timer is not None:
self._timer.cancel()
self._timer = None
self._is_closed.set()
