fix: add token expiry buffer to prevent expired token usage and lock concurrent refreshes

openfga_sdk/oauth2.py

@@ -4,18 +4,30 @@
 import random
 import sys
 
+from dataclasses import dataclass
 from datetime import datetime, timedelta
 
 import urllib3
 
 from openfga_sdk.configuration import Configuration
-from openfga_sdk.constants import USER_AGENT
+from openfga_sdk.constants import (
+    TOKEN_EXPIRY_JITTER_IN_SEC,
+    TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC,
+    USER_AGENT,
+)
 from openfga_sdk.credentials import Credentials
 from openfga_sdk.exceptions import AuthenticationError
 from openfga_sdk.telemetry.attributes import TelemetryAttributes
 from openfga_sdk.telemetry.telemetry import Telemetry
 
 
+@dataclass(frozen=True)
+class _TokenState:
+    access_token: str
+    expiry_time: datetime
+    expiry_buffer: float
+
+
 def jitter(loop_count, min_wait_in_ms):
     """
     Generate a random jitter value for exponential backoff
@@ -34,8 +46,8 @@ def jitter(loop_count, min_wait_in_ms):
 class OAuth2Client:
     def __init__(self, credentials: Credentials, configuration=None):
         self._credentials = credentials
-        self._access_token = None
-        self._access_expiry_time = None
+        self._token_state: _TokenState | None = None
+        self._lock = asyncio.Lock()
         self._telemetry = Telemetry()
 
         if configuration is None:
@@ -45,13 +57,13 @@ def __init__(self, credentials: Credentials, configuration=None):
 
     def _token_valid(self):
         """
-        Return whether token is valid
+        Return whether token is valid (with proactive expiry buffer to avoid using near-expired tokens)
         """
-        if self._access_token is None or self._access_expiry_time is None:
+        state = self._token_state  # atomic snapshot — either old or new, never torn
+        if state is None:
             return False
-        if self._access_expiry_time < datetime.now():
-            return False
-        return True
+        remaining = (state.expiry_time - datetime.now()).total_seconds()
+        return remaining > state.expiry_buffer
 
     async def _obtain_token(self, client):
         """
@@ -76,7 +88,9 @@ async def _obtain_token(self, client):
         # Add scope parameter if scopes are configured
         if configuration.scopes is not None:
             if isinstance(configuration.scopes, list):
-                scope_str = " ".join(s.strip() for s in configuration.scopes if s and s.strip())
+                scope_str = " ".join(
+                    s.strip() for s in configuration.scopes if s and s.strip()
+                )
             else:
                 scope_str = (
                     configuration.scopes.strip()
@@ -136,10 +150,15 @@ async def _obtain_token(self, client):
                     raise AuthenticationError(http_resp=raw_response)
 
                 if api_response.get("expires_in") and api_response.get("access_token"):
-                    self._access_expiry_time = datetime.now() + timedelta(
-                        seconds=int(api_response.get("expires_in"))
+                    self._token_state = _TokenState(
+                        access_token=api_response.get("access_token"),
+                        expiry_time=datetime.now()
+                        + timedelta(seconds=int(api_response.get("expires_in"))),
+                        expiry_buffer=(
+                            TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC
+                            + random.random() * TOKEN_EXPIRY_JITTER_IN_SEC
+                        ),
                     )
-                    self._access_token = api_response.get("access_token")
                     self._telemetry.metrics.credentialsRequest(
                         attributes={
                             TelemetryAttributes.fga_client_request_client_id: configuration.client_id
@@ -154,8 +173,8 @@ async def get_authentication_header(self, client):
         """
         If configured, return the header for authentication
         """
-        # check to see token is valid
         if not self._token_valid():
-            # In this case, the token is not valid, we need to get the refresh the token
-            await self._obtain_token(client)
-        return {"Authorization": f"Bearer {self._access_token}"}
+            async with self._lock:
+                if not self._token_valid():
+                    await self._obtain_token(client)
+        return {"Authorization": f"Bearer {self._token_state.access_token}"}

openfga_sdk/sync/oauth2.py

@@ -2,20 +2,33 @@
 import math
 import random
 import sys
+import threading
 import time
 
+from dataclasses import dataclass
 from datetime import datetime, timedelta
 
 import urllib3
 
 from openfga_sdk.configuration import Configuration
-from openfga_sdk.constants import USER_AGENT
+from openfga_sdk.constants import (
+    TOKEN_EXPIRY_JITTER_IN_SEC,
+    TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC,
+    USER_AGENT,
+)
 from openfga_sdk.credentials import Credentials
 from openfga_sdk.exceptions import AuthenticationError
 from openfga_sdk.telemetry.attributes import TelemetryAttributes
 from openfga_sdk.telemetry.telemetry import Telemetry
 
 
+@dataclass(frozen=True)
+class _TokenState:
+    access_token: str
+    expiry_time: datetime
+    expiry_buffer: float
+
+
 def jitter(loop_count, min_wait_in_ms):
     """
     Generate a random jitter value for exponential backoff
@@ -34,8 +47,8 @@ def jitter(loop_count, min_wait_in_ms):
 class OAuth2Client:
     def __init__(self, credentials: Credentials, configuration=None):
         self._credentials = credentials
-        self._access_token = None
-        self._access_expiry_time = None
+        self._token_state: _TokenState | None = None
+        self._lock = threading.Lock()
         self._telemetry = Telemetry()
 
         if configuration is None:
@@ -45,13 +58,13 @@ def __init__(self, credentials: Credentials, configuration=None):
 
     def _token_valid(self):
         """
-        Return whether token is valid
+        Return whether token is valid (with proactive expiry buffer to avoid using near-expired tokens)
         """
-        if self._access_token is None or self._access_expiry_time is None:
+        state = self._token_state  # atomic snapshot — either old or new, never torn
+        if state is None:
             return False
-        if self._access_expiry_time < datetime.now():
-            return False
-        return True
+        remaining = (state.expiry_time - datetime.now()).total_seconds()
+        return remaining > state.expiry_buffer
 
     def _obtain_token(self, client):
         """
@@ -76,7 +89,9 @@ def _obtain_token(self, client):
         # Add scope parameter if scopes are configured
         if configuration.scopes is not None:
             if isinstance(configuration.scopes, list):
-                scope_str = " ".join(s.strip() for s in configuration.scopes if s and s.strip())
+                scope_str = " ".join(
+                    s.strip() for s in configuration.scopes if s and s.strip()
+                )
             else:
                 scope_str = (
                     configuration.scopes.strip()
@@ -136,10 +151,15 @@ def _obtain_token(self, client):
                     raise AuthenticationError(http_resp=raw_response)
 
                 if api_response.get("expires_in") and api_response.get("access_token"):
-                    self._access_expiry_time = datetime.now() + timedelta(
-                        seconds=int(api_response.get("expires_in"))
+                    self._token_state = _TokenState(
+                        access_token=api_response.get("access_token"),
+                        expiry_time=datetime.now()
+                        + timedelta(seconds=int(api_response.get("expires_in"))),
+                        expiry_buffer=(
+                            TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC
+                            + random.random() * TOKEN_EXPIRY_JITTER_IN_SEC
+                        ),
                     )
-                    self._access_token = api_response.get("access_token")
                     self._telemetry.metrics.credentialsRequest(
                         attributes={
                             TelemetryAttributes.fga_client_request_client_id: configuration.client_id
@@ -154,8 +174,8 @@ def get_authentication_header(self, client):
         """
         If configured, return the header for authentication
         """
-        # check to see token is valid
         if not self._token_valid():
-            # In this case, the token is not valid, we need to get the refresh the token
-            self._obtain_token(client)
-        return {"Authorization": f"Bearer {self._access_token}"}
+            with self._lock:
+                if not self._token_valid():
+                    self._obtain_token(client)
+        return {"Authorization": f"Bearer {self._token_state.access_token}"}