Improvements for credentials refresh under high load (#658)

Author: kujenga

Diff for: gcsfs/credentials.py

@@ -5,6 +5,7 @@
 import textwrap
 import threading
 import warnings
+from datetime import datetime
 
 import google.auth as gauth
 import google.auth.compute_engine
@@ -16,6 +17,8 @@
 from google.oauth2.credentials import Credentials
 from google_auth_oauthlib.flow import InstalledAppFlow
 
+from gcsfs.retry import HttpError
+
 logger = logging.getLogger("gcsfs.credentials")
 
 tfile = os.path.join(os.path.expanduser("~"), ".gcs_tokens")
@@ -99,7 +102,6 @@ def _connect_cloud(self):
             raise ValueError("Invalid gcloud credentials") from error
 
     def _connect_cache(self):
-
         if len(self.tokens) == 0:
             raise ValueError("No cached tokens")
 
@@ -167,19 +169,41 @@ def _connect_token(self, token):
         if self.credentials.valid:
             self.credentials.apply(self.heads)
 
-    def maybe_refresh(self):
-        # this uses requests and is blocking
+    def _credentials_valid(self, refresh_buffer):
+        return (
+            self.credentials.valid
+            # In addition to checking current validity, we ensure that there is
+            # not a near-future expiry to avoid errors when expiration hits.
+            and self.credentials.expiry
+            and (self.credentials.expiry - datetime.utcnow()).total_seconds()
+            > refresh_buffer
+        )
+
+    def maybe_refresh(self, refresh_buffer=300):
+        """
+        Check and refresh credentials if needed
+        """
         if self.credentials is None:
             return  # anon
-        if self.credentials.valid:
-            return  # still good
+
+        if self._credentials_valid(refresh_buffer):
+            return  # still good, with buffer
+
         with requests.Session() as session:
             req = Request(session)
             with self.lock:
-                if self.credentials.valid:
-                    return  # repeat to avoid race (but don't want lock in common case)
+                if self._credentials_valid(refresh_buffer):
+                    return  # repeat check to avoid race conditions
+
                 logger.debug("GCS refresh")
-                self.credentials.refresh(req)
+                try:
+                    self.credentials.refresh(req)
+                except gauth.exceptions.RefreshError as error:
+                    # Re-raise as HttpError with a 401 code and the expected message
+                    raise HttpError(
+                        {"code": 401, "message": "Invalid Credentials"}
+                    ) from error
+
                 # https://github.com/fsspec/filesystem_spec/issues/565
                 self.credentials.apply(self.heads)
 

Diff for: gcsfs/retry.py

@@ -71,6 +71,9 @@ def is_retriable(exception):
     """Returns True if this exception is retriable."""
 
     if isinstance(exception, HttpError):
+        # Add 401 to retriable errors when it's an auth expiration issue
+        if exception.code == 401 and "Invalid Credentials" in str(exception.message):
+            return True
         return exception.code in errs
 
     return isinstance(exception, RETRIABLE_EXCEPTIONS)