feat: Add optional non blocking refresh for sync auth code

google/auth/_refresh_worker.py

@@ -13,39 +13,29 @@
 # limitations under the License.
 
 import logging
-import queue
 import threading
 
 import google.auth.exceptions as e
 
-WORKER_TIMEOUT_SECONDS = 5
-
 _LOGGER = logging.getLogger(__name__)
 
 
-class RefreshWorker:
+class RefreshThreadManager:
     """
-    A worker that will perform a non-blocking refresh of credentials.
+    Organizes exactly one background job that refresh a token.
     """
 
-    MAX_REFRESH_QUEUE_SIZE = 1
-    MAX_ERROR_QUEUE_SIZE = 2
-
     def __init__(self):
-        """Initializes the worker."""
+        """Initializes the manager."""
 
-        self._refresh_queue = queue.Queue(self.MAX_REFRESH_QUEUE_SIZE)
-        # Bound the error queue to avoid infinitely growing the heap.
-        self._error_queue = queue.Queue(self.MAX_ERROR_QUEUE_SIZE)
         self._worker = None
+        self._lock = threading.Lock()
 
     def _need_worker(self):
         return self._worker is None or not self._worker.is_alive()
 
-    def _spawn_worker(self):
-        self._worker = RefreshThread(
-            work_queue=self._refresh_queue, error_queue=self._error_queue
-        )
+    def _spawn_worker(self, cred, request):
+        self._worker = RefreshThread(cred=cred, request=request)
         self._worker.start()
 
     def start_refresh(self, cred, request):
@@ -62,101 +52,53 @@ def start_refresh(self, cred, request):
                 "Unable to start refresh. cred and request must be valid and instantiated objects."
             )
 
-        # This test case is covered by the unit tests but sometimes the cover
-        # check can flake due to the schdule.
-        #
-        # Specifially this test is covered by test_refresh_dead_worker
-        if not self._refresh_queue.empty():  # pragma: NO COVER
-            if self._need_worker():
-                self._spawn_worker()
-            return
+        with self._lock:
+            if self._worker is not None and self._worker._error_info is not None:
+                raise e.RefreshError(
+                    f"Could not start a background refresh. The background refresh previously failed with {self._worker._error_info}."
+                ) from self._worker._error_info
 
-        try:
-            self._refresh_queue.put_nowait((cred, request))
-        except queue.Full:
-            return
-
-        # This test case is covered by the unit tests but sometimes the cover
-        # check can flake due to the schdule.
-        if self._need_worker():  # pragma: NO COVER
-            self._spawn_worker()
-
-    def error_queue_full(self):
-        """
-      True if the refresh worker error queue is full. False if it is not yet full.
-
-      Returns:
-        bool
-      """
-        return self._error_queue.full()
-
-    def flush_error_queue(self):
-        """
-      Drop all errors in the error queue.
-      """
-        try:
-            while not self._error_queue.empty():
-                _ = self._error_queue.get_nowait()
-        # This condition is unlikely but there is a possibility that an
-        # error gets queued between the empty and get calls
-        except queue.Empty:  # pragma: NO COVER
-            pass
+            if self._need_worker():  # pragma: NO COVER
+                self._spawn_worker(cred, request)
 
     def get_error(self):
         """
-        Returns the first error in the error queue. It is recommended to flush the full error queue to root cause refresh failures.
+        Returns the error that occurred in the refresh thread. Clears the error once called.
+
         Returns:
           Optional[exceptions.Exception]
         """
-        try:
-            return self._error_queue.get_nowait()
-        except queue.Empty:
+        if not self._worker:
             return None
+        err, self._worker._error_info = self._worker._error_info, None
+        return err
 
 
 class RefreshThread(threading.Thread):
     """
     Thread that refreshes credentials.
     """
 
-    def __init__(self, work_queue, error_queue, **kwargs):
+    def __init__(self, cred, request, **kwargs):
         """Initializes the thread.
 
         Args:
-            work_queue: A queue of credentials and request tuples.
-            error_queue: A queue containing errors that prevented a credential refresh.
+            cred: A Credential object to refresh.
+            request: A Request object used to perform a credential refresh.
             **kwargs: Additional keyword arguments.
         """
 
         super().__init__(**kwargs)
-        self._work_queue = work_queue
-        self._error_queue = error_queue
+        self._cred = cred
+        self._request = request
+        self._error_info = None
 
     def run(self):
         """
-        Gets credentials and request objects from a queue.
-
-        The thread will block until a work item appears from the queue.
-
-        Once the refresh has completed, the thread will mark the queue task
-        as complete, and exit.
+        Perform the credential refresh.
         """
         try:
-            cred, request = self._work_queue.get(timeout=WORKER_TIMEOUT_SECONDS)
-        except queue.Empty:
-            _LOGGER.error(
-                f"Timed out waiting for refresh work after {WORKER_TIMEOUT_SECONDS} seconds. This could mean there is a race condition, work starvation, or other logic error in the refresh code."
-            )
-            return
-        try:
-            cred.refresh(request)
+            self._cred.refresh(self._request)
         except Exception as err:  # pragma: NO COVER
-            # This condition is covered by the unit test test_refresh_error but
-            # it can be flaky due to the scheduler.
             _LOGGER.error(f"Background refresh failed due to: {err}")
-            if not self._error_queue.full():
-                self._error_queue.put_nowait(err)
-
-        # The coverage tool is not able to capturre this line, but it is covered
-        # by test_start_refresh in the unit tests.
-        self._work_queue.task_done()
+            self._error_info = err

google/auth/credentials.py

@@ -22,7 +22,7 @@
 from google.auth import _helpers, environment_vars
 from google.auth import exceptions
 from google.auth import metrics
-from google.auth._refresh_worker import RefreshWorker
+from google.auth._refresh_worker import RefreshThreadManager
 
 
 class Credentials(metaclass=abc.ABCMeta):
@@ -62,7 +62,7 @@ def __init__(self):
         """
 
         self._use_non_blocking_refresh = False
-        self._refresh_worker = RefreshWorker()
+        self._refresh_worker = RefreshThreadManager()
 
     @property
     def expired(self):
@@ -71,30 +71,45 @@ def expired(self):
         Note that credentials can be invalid but not expired because
         Credentials with :attr:`expiry` set to None is considered to never
         expire.
+
+        .. deprecated:: v2.24.0
+          Prefer checking :attr:`token_state` instead.
         """
         if not self.expiry:
             return False
-
-        return _helpers.utcnow() >= self.expiry
+        # Remove some threshold from expiry to err on the side of reporting
+        # expiration early so that we avoid the 401-refresh-retry loop.
+        skewed_expiry = self.expiry - _helpers.REFRESH_THRESHOLD
+        return _helpers.utcnow() >= skewed_expiry
 
     @property
     def valid(self):
         """Checks the validity of the credentials.
 
         This is True if the credentials have a :attr:`token` and the token
         is not :attr:`expired`.
+
+        .. deprecated:: v2.24.0
+          Prefer checking :attr:`token_state` instead.
         """
         return self.token is not None and not self.expired
 
     @property
     def token_state(self):
-        if not self.valid or self.expired:
+        """
+        See `:obj:`TokenState`
+        """
+        if self.token is None:
             return TokenState.INVALID
 
         # Credentials that can't expire are always treated as fresh.
-        if not self.expiry:
+        if self.expiry is None:
             return TokenState.FRESH
 
+        expired = _helpers.utcnow() >= self.expiry
+        if expired:
+            return TokenState.INVALID
+
         is_stale = _helpers.utcnow() >= (self.expiry - _helpers.REFRESH_THRESHOLD)
         if is_stale:
             return TokenState.STALE
@@ -171,6 +186,17 @@ def apply(self, headers, token=None):
         if self.quota_project_id:
             headers["x-goog-user-project"] = self.quota_project_id
 
+    def _blocking_refresh(self, request):
+        if not self.valid:
+            self.refresh(request)
+
+    def _non_blocking_refresh(self, request):
+        if self.token_state == TokenState.STALE:
+            self._refresh_worker.start_refresh(self, request)
+
+        if self.token_state == TokenState.INVALID:
+            self.refresh(request)
+
     def before_request(self, request, method, url, headers):
         """Performs credential-specific before request logic.
 
@@ -188,21 +214,10 @@ def before_request(self, request, method, url, headers):
         # pylint: disable=unused-argument
         # (Subclasses may use these arguments to ascertain information about
         # the http request.)
-
-        if self.token_state == TokenState.FRESH:
-            self._refresh_worker.flush_error_queue()
-
-        if self.token_state == TokenState.STALE:
-            if (
-                self._use_non_blocking_refresh
-                and not self._refresh_worker.error_queue_full()
-            ):
-                self._refresh_worker.start_refresh(self, request)
-            else:
-                self.refresh(request)
-
-        if self.token_state == TokenState.INVALID:
-            self.refresh(request)
+        if self._use_non_blocking_refresh:
+            self._non_blocking_refresh(request)
+        else:
+            self._blocking_refresh(request)
 
         metrics.add_metric_header(headers, self._metric_header_for_usage())
         self.apply(headers)
@@ -212,9 +227,8 @@ def with_non_blocking_refresh(self):
 
     def get_background_refresh_error(self):
         """
-        Returns the first error in the background error queue. It is recommended to flush the full error queue to root cause refresh failures.
+        Returns the error in from a failed background refresh. Once called, the error will be flushed.
 
-        This error queue is populated by the token refreshes performed in a background thread.
         Returns:
           Optional[exceptions.Exception]
         """
@@ -469,6 +483,13 @@ def signer(self):
 
 
 class TokenState(Enum):
+    """
+    Tracks the state of a token.
+    FRESH: The token is not expired and can be used normally.
+    STALE: The token is close to expired, and should be refreshed. The token can be used normally.
+    INVALID: The token is expired or invalid. The token cannot be used for a normal operation.
+    """
+
     FRESH = 1
     STALE = 2
     INVALID = 3

google/auth/impersonated_credentials.py

@@ -260,8 +260,7 @@ def _update_token(self, request):
 
         # Refresh our source credentials if it is not valid.
         if (
-            not self._source_credentials.valid
-            or self._source_credentials.token_state == credentials.TokenState.STALE
+            self._source_credentials.token_state == credentials.TokenState.STALE
             or self._source_credentials.token_state == credentials.TokenState.INVALID
         ):
             self._source_credentials.refresh(request)