Lazy load the executor

Author: hectorcast-db

databricks/sdk/oauth.py

@@ -1,3 +1,4 @@
+import sys
 import base64
 import functools
 import hashlib
@@ -205,17 +206,31 @@ class _TokenState(Enum):
 class Refreshable(TokenSource):
     """A token source that supports refreshing expired tokens."""
 
-    _EXECUTOR = ThreadPoolExecutor(max_workers=10)
+    _EXECUTOR = None
+    _EXECUTOR_LOCK = threading.Lock()
     _DEFAULT_STALE_DURATION = timedelta(minutes=3)
 
+    @classmethod
+    def _get_executor(cls):
+        """Lazy initialization of the ThreadPoolExecutor."""
+        if cls._EXECUTOR is None:
+            with cls._EXECUTOR_LOCK:
+                if cls._EXECUTOR is None:
+                    # This thread pool has multiple workers because it is shared by all instances of Refreshable.
+                    cls._EXECUTOR = ThreadPoolExecutor(max_workers=10)
+        return cls._EXECUTOR
+
     def __init__(self,
                  token: Token = None,
                  disable_async: bool = True,
                  stale_duration: timedelta = _DEFAULT_STALE_DURATION):
-        self._lock = threading.Lock()
-        self._token = token
+        # Config properties
         self._stale_duration = stale_duration
         self._disable_async = disable_async
+        # Lock
+        self._lock = threading.Lock()
+        # Non Thread safe properties. Protected by the lock above.
+        self._token = token
         self._is_refreshing = False
         self._refresh_err = False
 
@@ -231,8 +246,9 @@ def _async_token(self) -> Token:
         If the token is stale, triggers an asynchronous refresh.
         If the token is expired, refreshes it synchronously, blocking until the refresh is complete.
         """
-        with self._lock:
-            token_state = self._token_state()
+        token_state = self._token_state()
+
+        with self._lock:    
             token = self._token
 
         if token_state == _TokenState.FRESH:
@@ -244,20 +260,22 @@ def _async_token(self) -> Token:
 
     def _token_state(self) -> _TokenState:
         """Returns the current state of the token."""
-        if not self._token or not self._token.valid:
-            return _TokenState.EXPIRED
-        if not self._token.expiry:
+        with self._lock:
+            if not self._token or not self._token.valid:
+                return _TokenState.EXPIRED
+            if not self._token.expiry:
+                return _TokenState.FRESH
+
+            lifespan = self._token.expiry - datetime.now()
+            if lifespan < timedelta(seconds=0):
+                return _TokenState.EXPIRED
+            if lifespan < self._stale_duration:
+                return _TokenState.STALE
             return _TokenState.FRESH
 
-        lifespan = self._token.expiry - datetime.now()
-        if lifespan < timedelta(seconds=0):
-            return _TokenState.EXPIRED
-        if lifespan < self._stale_duration:
-            return _TokenState.STALE
-        return _TokenState.FRESH
-
     def _blocking_token(self) -> Token:
         """Returns a token, blocking if necessary to refresh it."""
+        state = self._token_state()
         with self._lock:
             # This is important to recover from potential previous failed attempts
             # to refresh the token asynchronously.
@@ -267,7 +285,7 @@ def _blocking_token(self) -> Token:
             # It's possible that the token got refreshed (either by a _blocking_refresh or
             # an _async_refresh call) while this particular call was waiting to acquire
             # the lock. This check avoids refreshing the token again in such cases.
-            if self._token_state() != _TokenState.EXPIRED:
+            if state != _TokenState.EXPIRED:
                 return self._token
 
             self._token = self.refresh()
@@ -276,20 +294,29 @@ def _blocking_token(self) -> Token:
     def _trigger_async_refresh(self):
         """Starts an asynchronous refresh if none is in progress."""
 
-        # Note: _refresh_internal function is not thread safe.
-        # Only call it inside the lock.
         def _refresh_internal():
+            new_token: Token = None
             try:
-                self._token = self.refresh()
-            except Exception:
-                self._refresh_err = True
-            finally:
+                new_token = self.refresh()
+            except Exception as e:
+                # This happens on a thread, so we don't want to propagate the error.
+                # Instead, if there is no new_token for any reason, we will disable async refresh below
+                # But we will do it inside the lock.
+                logger.warning(f'Tried to refresh token asynchronously, but failed: {e}')
+                pass
+            
+            with self._lock:
+                if new_token is not None:
+                    self._token = new_token
+                else:
+                    self._refresh_err = True
                 self._is_refreshing = False
+            
 
         with self._lock:
             if not self._is_refreshing and not self._refresh_err:
                 self._is_refreshing = True
-                self._EXECUTOR.submit(_refresh_internal)
+                Refreshable._get_executor().submit(_refresh_internal)
 
     @abstractmethod
     def refresh(self) -> Token:

tests/test_refreshable.py

@@ -25,7 +25,7 @@ def refresh(self) -> Token:
 
 
 def fail() -> Token:
-    raise Exception("Failed to refresh token")
+    raise Exception("Simulated token refresh failure")
 
 
 def static_token(token: Token, wait: int = 0) -> Callable[[], Token]: