use hybrid approach

Author: vgvoleg

tests/aio/test_credentials.py

@@ -6,7 +6,7 @@
 import os
 import json
 import asyncio
-from unittest.mock import patch, AsyncMock
+from unittest.mock import patch, AsyncMock, MagicMock
 
 import tests.auth.test_credentials
 import tests.oauth2_token_exchange
@@ -125,6 +125,8 @@ async def test_token_lazy_refresh():
         "localhost:0",
     )
 
+    credentials._tp.submit = MagicMock()
+
     mock_response = {"access_token": "token_v1", "expires_in": 3600}
     credentials._make_token_request = AsyncMock(return_value=mock_response)
 
@@ -139,7 +141,7 @@ async def test_token_lazy_refresh():
         assert token2 == "token_v1"
         assert credentials._make_token_request.call_count == 1
 
-        mock_time.return_value = 2000
+        mock_time.return_value = 1000 + 3600 - 30 + 1
         credentials._make_token_request.return_value = {"access_token": "token_v2", "expires_in": 3600}
 
         token3 = await credentials.token()
@@ -156,6 +158,8 @@ async def test_token_double_check_locking():
         "localhost:0",
     )
 
+    credentials._tp.submit = MagicMock()
+
     call_count = 0
 
     async def mock_make_request():
@@ -185,14 +189,16 @@ async def test_token_expiration_calculation():
         "localhost:0",
     )
 
+    credentials._tp.submit = MagicMock()
+
     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)
+        expected_expires = 1000 + 3600 - 30
         assert credentials._expires_in == expected_expires
 
 
@@ -205,10 +211,55 @@ async def test_token_refresh_error_handling():
         "localhost:0",
     )
 
+    credentials._tp.submit = MagicMock()
+
     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"
+
+
+@pytest.mark.asyncio
+async def test_hybrid_background_and_sync_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",
+    )
+
+    call_count = 0
+    background_calls = []
+
+    async def mock_make_request():
+        nonlocal call_count
+        call_count += 1
+        return {"access_token": f"token_v{call_count}", "expires_in": 3600}
+
+    def mock_submit(callback):
+        background_calls.append(callback)
+
+    credentials._make_token_request = mock_make_request
+    credentials._tp.submit = mock_submit
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000
+
+        token1 = await credentials.token()
+        assert token1 == "token_v1"
+        assert call_count == 1
+        assert len(background_calls) == 0
+
+        mock_time.return_value = 1000 + min(1800, 3600 / 10) + 1
+        token2 = await credentials.token()
+        assert token2 == "token_v1"
+        assert call_count == 1
+        assert len(background_calls) == 1
+
+        mock_time.return_value = 1000 + 3600 - 30 + 1
+        token3 = await credentials.token()
+        assert token3 == "token_v2"
+        assert call_count == 2

tests/auth/test_static_credentials.py

@@ -51,6 +51,8 @@ def test_static_credentials_wrong_creds(endpoint, database):
 def test_token_lazy_refresh():
     credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
 
+    credentials._tp.submit = MagicMock()
+
     mock_response = {"access_token": "token_v1", "expires_in": 3600}
     credentials._make_token_request = MagicMock(return_value=mock_response)
 
@@ -65,7 +67,7 @@ def test_token_lazy_refresh():
         assert token2 == "token_v1"
         assert credentials._make_token_request.call_count == 1
 
-        mock_time.return_value = 2000
+        mock_time.return_value = 1000 + 3600 - 30 + 1
         credentials._make_token_request.return_value = {"access_token": "token_v2", "expires_in": 3600}
 
         token3 = credentials.token
@@ -75,6 +77,7 @@ def test_token_lazy_refresh():
 
 def test_token_double_check_locking():
     credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
+    credentials._tp.submit = MagicMock()
 
     call_count = 0
 
@@ -108,24 +111,66 @@ def get_token():
 def test_token_expiration_calculation():
     credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
 
+    credentials._tp.submit = MagicMock()
+
     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)
+        expected_expires = 1000 + 3600 - 30
         assert credentials._expires_in == expected_expires
 
 
 def test_token_refresh_error_handling():
     credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
-
+    credentials._tp.submit = MagicMock()
     credentials._make_token_request = MagicMock(side_effect=Exception("Network error"))
 
-    with pytest.raises(ydb.ConnectionError) as exc_info:
-        credentials.token
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000 + 3600
+
+        with pytest.raises(ydb.ConnectionError) as exc_info:
+            credentials.token
+
+        assert "Network error" in str(exc_info.value)
+        assert credentials.last_error == "Network error"
+
+
+def test_hybrid_background_and_sync_refresh():
+    credentials = ydb.StaticCredentials.from_user_password(USERNAME, PASSWORD)
+
+    call_count = 0
+    background_calls = []
+
+    def mock_make_request():
+        nonlocal call_count
+        call_count += 1
+        return {"access_token": f"token_v{call_count}", "expires_in": 3600}
 
-    assert "Network error" in str(exc_info.value)
-    assert credentials.last_error == "Network error"
+    def mock_submit(callback):
+        background_calls.append(callback)
+
+    credentials._make_token_request = mock_make_request
+    credentials._tp.submit = mock_submit
+
+    with patch("time.time") as mock_time:
+        mock_time.return_value = 1000
+
+        token1 = credentials.token
+        assert token1 == "token_v1"
+        assert call_count == 1
+        assert len(background_calls) == 0
+
+        mock_time.return_value = 1000 + min(1800, 3600 / 10) + 1
+        token2 = credentials.token
+        assert token2 == "token_v1"
+        assert call_count == 1
+        assert len(background_calls) == 1
+
+        mock_time.return_value = 1000 + 3600 - 30 + 1
+        token3 = credentials.token
+        assert token3 == "token_v2"
+        assert call_count == 2

ydb/aio/credentials.py

@@ -2,6 +2,7 @@
 import asyncio
 import logging
 import time
+from concurrent import futures
 
 from ydb import credentials
 from ydb import issues
@@ -10,10 +11,31 @@
 YDB_AUTH_TICKET_HEADER = "x-ydb-auth-ticket"
 
 
+class AtMostOneExecution(object):
+    def __init__(self):
+        self._can_schedule = True
+        self._lock = asyncio.Lock()
+
+    async def wrapped_execution(self, callback):
+        async with self._lock:
+            try:
+                res = await callback()
+            except Exception:
+                pass
+            finally:
+                self._can_schedule = True
+
+    def submit(self, callback):
+        if self._can_schedule:
+            self._can_schedule = False
+            asyncio.create_task(self.wrapped_execution(callback))
+
+
 class AbstractExpiringTokenCredentials(credentials.AbstractExpiringTokenCredentials):
     def __init__(self):
         super(AbstractExpiringTokenCredentials, self).__init__()
         self._token_lock = asyncio.Lock()
+        self._tp = AtMostOneExecution()
 
     @abc.abstractmethod
     async def _make_token_request(self):
@@ -25,14 +47,12 @@ async def get_auth_token(self) -> str:
                 return token
         return ""
 
-    async def _refresh_token(self):
+    async def _refresh_token(self, should_raise=False):
         current_time = time.time()
 
         try:
             self.logger.debug(
-                "Refreshing token async, current_time: %s, expires_in: %s",
-                current_time,
-                self._expires_in,
+                "Refreshing token async, current_time: %s, expires_in: %s", current_time, self._expires_in
             )
 
             token_response = await self._make_token_request()
@@ -44,19 +64,23 @@ async def _refresh_token(self):
         except Exception as e:
             self.last_error = str(e)
             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)
-            )
+            if should_raise:
+                raise issues.ConnectionError(
+                    "%s: %s.\n%s" % (self.__class__.__name__, self.last_error, self.extra_error_message)
+                )
 
     async def token(self):
         if self._is_token_valid():
+            if self._should_refresh():
+                self._tp.submit(self._refresh_token)
+
             return self._cached_token
 
         async with self._token_lock:
             if self._is_token_valid():
                 return self._cached_token
 
-            await self._refresh_token()
+            await self._refresh_token(should_raise=True)
 
         return self._cached_token
 

ydb/aio/iam.py

@@ -103,6 +103,7 @@ def __init__(self, metadata_url=None):
         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.extra_error_message = "Check that metadata service configured properly and application deployed in VM or function at Yandex.Cloud."
+        self._tp.submit(self._refresh_token)
 
     async def _make_token_request(self):
         timeout = aiohttp.ClientTimeout(total=2)