fix: Async auth lock

src/firebolt/client/auth/base.py

@@ -2,7 +2,7 @@
 from time import time
 from typing import AsyncGenerator, Generator, Optional
 
-from anyio import Lock
+from anyio import Lock, get_current_task
 from httpx import Auth as HttpxAuth
 from httpx import Request, Response, codes
 
@@ -149,16 +149,27 @@ async def async_auth_flow(
         if self.requires_request_body:
             await request.aread()
 
-        async with self._lock:
-            flow = self.auth_flow(request)
-            request = next(flow)
+        if not self.token or self.expired:
+            await self._lock.acquire()
+            # If another task has already updated the token,
+            # we don't need to hold the lock
+            if self.token and not self.expired:
+                self._lock.release()
 
-            while True:
-                response = yield request
-                if self.requires_response_body:
-                    await response.aread()
+        flow = self.auth_flow(request)
+        request = next(flow)
 
-                try:
-                    request = flow.send(response)
-                except StopIteration:
-                    break
+        while True:
+            response = yield request
+            if self.requires_response_body:
+                await response.aread()
+
+            try:
+                request = flow.send(response)
+            except StopIteration:
+                break
+            finally:
+                # token gets updated only after flow.send is called
+                # so unlock only after that
+                if self._lock.locked() and self._lock._owner_task == get_current_task():
+                    self._lock.release()

tests/unit/client/V1/test_client_async.py

@@ -1,3 +1,5 @@
+import random
+from queue import Queue
 from re import Pattern, compile
 from types import MethodType
 from typing import Any, Callable
@@ -12,7 +14,7 @@
 from firebolt.client.auth import Token, UsernamePassword
 from firebolt.utils.urls import AUTH_URL
 from firebolt.utils.util import fix_url_schema
-from tests.unit.conftest import Response
+from tests.unit.conftest import Response, retry_if_failed
 
 
 async def test_client_retry(
@@ -157,7 +159,7 @@ def check_credentials(
     httpx_mock.add_callback(check_credentials, url=auth_url)
 
     async with AsyncClient(
-        auth=UsernamePassword(test_username, test_password),
+        auth=UsernamePassword(test_username, test_password, False),
         api_endpoint=server,
     ) as c:
         c._send_handling_redirects = MethodType(mock_send_handling_redirects, c)
@@ -167,3 +169,54 @@ def check_credentials(
                 nursery.start_soon(c.get, url)
 
     assert checked_creds_times == 1
+
+
+# test that client requests are truly concurrent
+# and are executed not in order that they were started
+# but in order of completion
+@retry_if_failed(3)
+async def test_true_concurent_requests(
+    httpx_mock: HTTPXMock,
+    test_username: str,
+    test_password: str,
+    auth_url: str,
+    auth_callback: Callable,
+    server: str,
+):
+    url = "https://url"
+    CONCURENT_COUNT = 10
+
+    queue = Queue(CONCURENT_COUNT)
+
+    # create callback that uses check_token_callback but also pushes URl to a queue
+    async def check_token_callback_with_queue(request: Request, **kwargs) -> Response:
+        nonlocal queue
+        queue.put(str(request.url))
+        return Response(status_code=codes.OK, headers={"content-length": "0"})
+
+    async def mock_send_handling_redirects(self, *args: Any, **kwargs: Any) -> Response:
+        # simulate network delay so the context switches
+        # random delay to make sure that requests are not executed in order
+        await sleep(0.1 * random.random())
+        return await AsyncClient._send_handling_redirects(self, *args, **kwargs)
+
+    httpx_mock.add_callback(auth_callback, url=auth_url)
+
+    httpx_mock.add_callback(check_token_callback_with_queue, url=compile(f"{url}/."))
+
+    urls = [f"{url}/{i}" for i in range(CONCURENT_COUNT)]
+    async with AsyncClient(
+        auth=UsernamePassword(test_username, test_password),
+        api_endpoint=server,
+    ) as c:
+        c._send_handling_redirects = MethodType(mock_send_handling_redirects, c)
+        async with open_nursery() as nursery:
+            for url in urls:
+                nursery.start_soon(c.get, url)
+
+    assert queue.qsize() == CONCURENT_COUNT
+    # Make sure the order is random and not sequential
+    assert list(queue.queue) != urls
+    # Cover the case when requests might be queued in reverse order
+    urls.reverse()
+    assert list(queue.queue) != urls

tests/unit/client/V2/test_client_async.py

@@ -1,3 +1,5 @@
+import random
+from queue import Queue
 from re import Pattern, compile
 from types import MethodType
 from typing import Any, Callable
@@ -11,7 +13,7 @@
 from firebolt.client.auth import Auth, ClientCredentials
 from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL
 from firebolt.utils.util import fix_url_schema
-from tests.unit.conftest import Response
+from tests.unit.conftest import Response, retry_if_failed
 
 
 async def test_client_retry(
@@ -186,3 +188,56 @@ def check_credentials(
                 nursery.start_soon(c.get, url)
 
     assert checked_creds_times == 1
+
+
+# test that client requests are truly concurrent
+# and are executed not in order that they were started
+# but in order of completion
+@retry_if_failed(3)
+async def test_true_concurent_requests(
+    httpx_mock: HTTPXMock,
+    account_name: str,
+    client_id: str,
+    client_secret: str,
+    auth_url: str,
+    auth_callback: Callable,
+    server: str,
+):
+    url = "https://url"
+    CONCURENT_COUNT = 10
+
+    queue = Queue(CONCURENT_COUNT)
+
+    # create callback that uses check_token_callback but also pushes URl to a queue
+    async def check_token_callback_with_queue(request: Request, **kwargs) -> Response:
+        nonlocal queue
+        queue.put(str(request.url))
+        return Response(status_code=codes.OK, headers={"content-length": "0"})
+
+    async def mock_send_handling_redirects(self, *args: Any, **kwargs: Any) -> Response:
+        # simulate network delay so the context switches
+        # random delay to make sure that requests are not executed in order
+        await sleep(0.1 * random.random())
+        return await AsyncClient._send_handling_redirects(self, *args, **kwargs)
+
+    httpx_mock.add_callback(auth_callback, url=auth_url)
+
+    httpx_mock.add_callback(check_token_callback_with_queue, url=compile(f"{url}/."))
+
+    urls = [f"{url}/{i}" for i in range(CONCURENT_COUNT)]
+    async with AsyncClient(
+        auth=ClientCredentials(client_id, client_secret),
+        api_endpoint=server,
+        account_name=account_name,
+    ) as c:
+        c._send_handling_redirects = MethodType(mock_send_handling_redirects, c)
+        async with open_nursery() as nursery:
+            for url in urls:
+                nursery.start_soon(c.get, url)
+
+    assert queue.qsize() == CONCURENT_COUNT
+    # Make sure the order is random and not sequential
+    assert list(queue.queue) != urls
+    # Cover the case when requests might be queued in reverse order
+    urls.reverse()
+    assert list(queue.queue) != urls

tests/unit/conftest.py

@@ -1,3 +1,4 @@
+import functools
 from re import Pattern, compile
 from typing import Callable
 
@@ -373,3 +374,22 @@ def settings(
         account_name=account_name,
     )
     return seett
+
+
+# Retry decorator that allows to retry test N number of times in case one
+# ot the asserts fail
+def retry_if_failed(num_retries):
+    def decorator(func):
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):
+            for i in range(num_retries):
+                try:
+                    await func(*args, **kwargs)
+                    break
+                except AssertionError as e:
+                    if i == num_retries - 1:
+                        raise e
+
+        return wrapper
+
+    return decorator