Refactor auth token refresh logic

Author: vgvoleg

tests/aio/test_credentials.py

@@ -5,6 +5,8 @@
 import tempfile
 import os
 import json
+import asyncio
+from unittest.mock import patch, AsyncMock
 
 import tests.auth.test_credentials
 import tests.oauth2_token_exchange
@@ -112,3 +114,101 @@ def serve(s):
     except Exception:
         os.remove(cfg_file_name)
         raise
+
+
+@pytest.mark.asyncio
+async def test_token_lazy_refresh():
+    credentials = ServiceAccountCredentialsForTest(
+        tests.auth.test_credentials.SERVICE_ACCOUNT_ID,
+        tests.auth.test_credentials.ACCESS_KEY_ID,
+        tests.auth.test_credentials.PRIVATE_KEY,
+        "localhost:0",
+    )
+
+    mock_response = {"access_token": "token_v1", "expires_in": 3600}
+    credentials._make_token_request = AsyncMock(return_value=mock_response)
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000
+
+        token1 = await credentials.token()
+        assert token1 == "token_v1"
+        assert credentials._make_token_request.call_count == 1
+
+        token2 = await credentials.token()
+        assert token2 == "token_v1"
+        assert credentials._make_token_request.call_count == 1
+
+        mock_time.return_value = 2000
+        credentials._make_token_request.return_value = {"access_token": "token_v2", "expires_in": 3600}
+
+        token3 = await credentials.token()
+        assert token3 == "token_v2"
+        assert credentials._make_token_request.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_token_double_check_locking():
+    credentials = ServiceAccountCredentialsForTest(
+        tests.auth.test_credentials.SERVICE_ACCOUNT_ID,
+        tests.auth.test_credentials.ACCESS_KEY_ID,
+        tests.auth.test_credentials.PRIVATE_KEY,
+        "localhost:0",
+    )
+
+    call_count = 0
+
+    async def mock_make_request():
+        nonlocal call_count
+        call_count += 1
+        await asyncio.sleep(0.01)
+        return {"access_token": f"token_v{call_count}", "expires_in": 3600}
+
+    credentials._make_token_request = mock_make_request
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000
+
+        tasks = [credentials.token() for _ in range(10)]
+        results = await asyncio.gather(*tasks)
+
+        assert len(set(results)) == 1
+        assert call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_token_expiration_calculation():
+    credentials = ServiceAccountCredentialsForTest(
+        tests.auth.test_credentials.SERVICE_ACCOUNT_ID,
+        tests.auth.test_credentials.ACCESS_KEY_ID,
+        tests.auth.test_credentials.PRIVATE_KEY,
+        "localhost:0",
+    )
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000
+
+        credentials._make_token_request = AsyncMock(return_value={"access_token": "token", "expires_in": 3600})
+
+        await credentials.token()
+
+        expected_expires = 1000 + min(1800, 3600 / 4)
+        assert credentials._expires_in == expected_expires
+
+
+@pytest.mark.asyncio
+async def test_token_refresh_error_handling():
+    credentials = ServiceAccountCredentialsForTest(
+        tests.auth.test_credentials.SERVICE_ACCOUNT_ID,
+        tests.auth.test_credentials.ACCESS_KEY_ID,
+        tests.auth.test_credentials.PRIVATE_KEY,
+        "localhost:0",
+    )
+
+    credentials._make_token_request = AsyncMock(side_effect=Exception("Network error"))
+
+    with pytest.raises(Exception) as exc_info:
+        await credentials.token()
+
+    assert "Network error" in str(exc_info.value)
+    assert credentials.last_error == "Network error"

tests/auth/test_static_credentials.py

@@ -1,5 +1,6 @@
 import pytest
 import ydb
+from unittest.mock import patch, MagicMock
 
 
 USERNAME = "root"
@@ -45,3 +46,86 @@ def test_static_credentials_wrong_creds(endpoint, database):
     with pytest.raises(ydb.ConnectionFailure):
         with ydb.Driver(driver_config=driver_config) as driver:
             driver.wait(5, fail_fast=True)
+
+
+def test_token_lazy_refresh():
+    credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
+
+    mock_response = {"access_token": "token_v1", "expires_in": 3600}
+    credentials._make_token_request = MagicMock(return_value=mock_response)
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000
+
+        token1 = credentials.token
+        assert token1 == "token_v1"
+        assert credentials._make_token_request.call_count == 1
+
+        token2 = credentials.token
+        assert token2 == "token_v1"
+        assert credentials._make_token_request.call_count == 1
+
+        mock_time.return_value = 2000
+        credentials._make_token_request.return_value = {"access_token": "token_v2", "expires_in": 3600}
+
+        token3 = credentials.token
+        assert token3 == "token_v2"
+        assert credentials._make_token_request.call_count == 2
+
+
+def test_token_double_check_locking():
+    credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
+
+    call_count = 0
+
+    def mock_make_request():
+        nonlocal call_count
+        call_count += 1
+        return {"access_token": f"token_v{call_count}", "expires_in": 3600}
+
+    credentials._make_token_request = mock_make_request
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000
+
+        import threading
+
+        results = []
+
+        def get_token():
+            results.append(credentials.token)
+
+        threads = [threading.Thread(target=get_token) for _ in range(10)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        assert len(set(results)) == 1
+        assert call_count == 1
+
+
+def test_token_expiration_calculation():
+    credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000
+
+        credentials._make_token_request = MagicMock(return_value={"access_token": "token", "expires_in": 3600})
+
+        credentials.token
+
+        expected_expires = 1000 + min(1800, 3600 / 4)
+        assert credentials._expires_in == expected_expires
+
+
+def test_token_refresh_error_handling():
+    credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
+
+    credentials._make_token_request = MagicMock(side_effect=Exception("Network error"))
+
+    with pytest.raises(ydb.ConnectionError) as exc_info:
+        credentials.token
+
+    assert "Network error" in str(exc_info.value)
+    assert credentials.last_error == "Network error"

ydb/aio/credentials.py

@@ -10,57 +10,10 @@
 YDB_AUTH_TICKET_HEADER = "x-ydb-auth-ticket"
 
 
-class _OneToManyValue(object):
-    def __init__(self):
-        self._value = None
-        self._condition = asyncio.Condition()
-
-    async def consume(self, timeout=3):
-        async with self._condition:
-            if self._value is None:
-                try:
-                    await asyncio.wait_for(self._condition.wait(), timeout=timeout)
-                except Exception:
-                    return self._value
-            return self._value
-
-    async def update(self, n_value):
-        async with self._condition:
-            prev_value = self._value
-            self._value = n_value
-            if prev_value is None:
-                self._condition.notify_all()
-
-
-class _AtMostOneExecution(object):
-    def __init__(self):
-        self._can_schedule = True
-        self._lock = asyncio.Lock()  # Lock to guarantee only one execution
-
-    async def _wrapped_execution(self, callback):
-        await self._lock.acquire()
-        try:
-            res = callback()
-            if asyncio.iscoroutine(res):
-                await res
-        except Exception:
-            pass
-
-        finally:
-            self._lock.release()
-            self._can_schedule = True
-
-    def submit(self, callback):
-        if self._can_schedule:
-            self._can_schedule = False
-            asyncio.ensure_future(self._wrapped_execution(callback))
-
-
 class AbstractExpiringTokenCredentials(credentials.AbstractExpiringTokenCredentials):
     def __init__(self):
         super(AbstractExpiringTokenCredentials, self).__init__()
-        self._tp = _AtMostOneExecution()
-        self._cached_token = _OneToManyValue()
+        self._token_lock = asyncio.Lock()
 
     @abc.abstractmethod
     async def _make_token_request(self):
@@ -72,51 +25,45 @@ async def get_auth_token(self) -> str:
                 return token
         return ""
 
-    async def _refresh(self):
+    async def _refresh_token(self):
         current_time = time.time()
-        self._log_refresh_start(current_time)
 
         try:
-            auth_metadata = await self._make_token_request()
-            await self._cached_token.update(auth_metadata["access_token"])
-            self._update_expiration_info(auth_metadata)
-            self.logger.info(
-                "Token refresh successful. current_time %s, refresh_in %s",
+            self.logger.debug(
+                "Refreshing token async, current_time: %s, expires_in: %s",
                 current_time,
-                self._refresh_in,
+                self._expires_in,
             )
 
-        except (KeyboardInterrupt, SystemExit):
-            return
+            token_response = await self._make_token_request()
+            self._update_token_info(token_response, current_time)
 
-        except Exception as e:
-            self.last_error = str(e)
-            await asyncio.sleep(1)
-            self._tp.submit(self._refresh)
+            self.logger.info("Token refreshed successfully async, expires_in: %s", self._expires_in)
+            self.last_error = None
 
-        except BaseException as e:
+        except Exception as e:
             self.last_error = str(e)
-            raise
+            self.logger.error("Failed to refresh token async: %s", e)
+            raise issues.ConnectionError(
+                "%s: %s.\n%s" % (self.__class__.__name__, self.last_error, self.extra_error_message)
+            )
 
     async def token(self):
-        current_time = time.time()
-        if current_time > self._refresh_in:
-            self._tp.submit(self._refresh)
-
-        cached_token = await self._cached_token.consume(timeout=3)
-        if cached_token is None:
-            if self.last_error is None:
-                raise issues.ConnectionError(
-                    "%s: timeout occurred while waiting for token.\n%s"
-                    % (
-                        self.__class__.__name__,
-                        self.extra_error_message,
-                    )
-                )
+        if self._is_token_valid():
+            return self._cached_token
+
+        async with self._token_lock:
+            if self._is_token_valid():
+                return self._cached_token
+
+            await self._refresh_token()
+
+        if self._cached_token is None:
             raise issues.ConnectionError(
-                "%s: %s.\n%s" % (self.__class__.__name__, self.last_error, self.extra_error_message)
+                "%s: No token available.\n%s" % (self.__class__.__name__, self.extra_error_message)
             )
-        return cached_token
+
+        return self._cached_token
 
     async def auth_metadata(self):
         return [(credentials.YDB_AUTH_TICKET_HEADER, await self.token())]

ydb/aio/iam.py

@@ -102,7 +102,6 @@ def __init__(self, metadata_url=None):
         super(MetadataUrlCredentials, self).__init__()
         assert aiohttp is not None, "Install aiohttp library to use metadata credentials provider"
         self._metadata_url = auth.DEFAULT_METADATA_URL if metadata_url is None else metadata_url
-        self._tp.submit(self._refresh)
         self.extra_error_message = "Check that metadata service configured properly and application deployed in VM or function at Yandex.Cloud."
 
     async def _make_token_request(self):