fix(client): lock async session setup against concurrent callers

src/simple_github/client.py

@@ -112,10 +112,14 @@ def _get_gql_session(self) -> SyncClientSession:
             assert isinstance(self._gql_session, SyncClientSession)
             return self._gql_session
 
-        # Create a new session with updated token.
-        self._prev_token = token
-        if self._gql_client:
-            self._gql_client.close_sync()
+        # Drop the old session before building the new one, so a failure below
+        # leaves nothing cached and the next caller retries.
+        prev_client = self._gql_client
+        self._prev_token = None
+        self._gql_client = None
+        self._gql_session = None
+        if prev_client:
+            prev_client.close_sync()
 
         headers = {
             "Accept": "application/vnd.github+json",
@@ -124,12 +128,14 @@ def _get_gql_session(self) -> SyncClientSession:
             headers["Authorization"] = f"Bearer {token}"
 
         transport = RequestsHTTPTransport(url=GITHUB_GRAPHQL_ENDPOINT, headers=headers)
-        self._gql_client = GqlClient(
-            transport=transport, fetch_schema_from_transport=False
-        )
-        self._gql_session = self._gql_client.connect_sync()
-        assert isinstance(self._gql_session, SyncClientSession)
-        return self._gql_session
+        client = GqlClient(transport=transport, fetch_schema_from_transport=False)
+        session = client.connect_sync()
+        assert isinstance(session, SyncClientSession)
+
+        self._gql_client = client
+        self._gql_session = session
+        self._prev_token = token
+        return session
 
     def _get_requests_session(self) -> Session:
         session = self._get_gql_session()
@@ -242,6 +248,10 @@ def execute(self, query: str, variables: RequestData = None) -> dict[str, Any]:
 
 
 class AsyncClient(Client):
+    def __init__(self, auth: "Auth"):
+        super().__init__(auth)
+        self._session_lock = asyncio.Lock()
+
     async def __aenter__(self):
         return self
 
@@ -265,31 +275,41 @@ async def _get_gql_session(self) -> ReconnectingAsyncClientSession:
         Returns:
             aiohttp.ClientSession: An AIOHTTP session object.
         """
-        token = await self.get_token()
-        if token == self._prev_token:
-            assert isinstance(self._gql_session, ReconnectingAsyncClientSession)
-            return self._gql_session
-
-        # Create a new session with updated token.
-        self._prev_token = token
-        if self._gql_client:
-            await self._gql_client.close_async()
-
-        headers = {
-            "Accept": "application/vnd.github+json",
-        }
-        if token:
-            headers["Authorization"] = f"Bearer {token}"
-
-        transport = AIOHTTPTransport(
-            url=GITHUB_GRAPHQL_ENDPOINT, headers=headers, ssl=True
-        )
-        self._gql_client = GqlClient(
-            transport=transport, fetch_schema_from_transport=False
-        )
-        self._gql_session = await self._gql_client.connect_async(reconnecting=True)
-        assert isinstance(self._gql_session, ReconnectingAsyncClientSession)
-        return self._gql_session
+        # The lock spans `get_token()` because `AppInstallationAuth` drives an
+        # async generator, which two callers cannot advance at once.
+        async with self._session_lock:
+            token = await self.get_token()
+
+            if token == self._prev_token:
+                assert isinstance(self._gql_session, ReconnectingAsyncClientSession)
+                return self._gql_session
+
+            # Drop the old session before building the new one, so a failure
+            # below leaves nothing cached and the next caller retries.
+            prev_client = self._gql_client
+            self._prev_token = None
+            self._gql_client = None
+            self._gql_session = None
+            if prev_client:
+                await prev_client.close_async()
+
+            headers = {
+                "Accept": "application/vnd.github+json",
+            }
+            if token:
+                headers["Authorization"] = f"Bearer {token}"
+
+            transport = AIOHTTPTransport(
+                url=GITHUB_GRAPHQL_ENDPOINT, headers=headers, ssl=True
+            )
+            client = GqlClient(transport=transport, fetch_schema_from_transport=False)
+            session = await client.connect_async(reconnecting=True)
+            assert isinstance(session, ReconnectingAsyncClientSession)
+
+            self._gql_client = client
+            self._gql_session = session
+            self._prev_token = token
+            return session
 
     async def _get_aiohttp_session(self) -> ClientSession:
         session = await self._get_gql_session()

test/test_auth.py

@@ -1,11 +1,12 @@
+import asyncio
 import time
 from unittest import mock
 
 import jwt
 import pytest
 
 from simple_github.auth import AppAuth, AppInstallationAuth, PublicAuth, TokenAuth
-from simple_github.client import GITHUB_API_ENDPOINT
+from simple_github.client import GITHUB_API_ENDPOINT, AsyncClient
 
 
 @pytest.mark.asyncio
@@ -83,3 +84,26 @@ async def test_app_installation_auth_get_token(
     with mock.patch.object(time, "time", return_value=cur + 3600 - 59):
         new_token = await auth.get_token()
         assert new_token != token
+
+
+@pytest.mark.asyncio
+async def test_concurrent_app_installation_auth(aioresponses, privkey: str):
+    """Concurrent callers must not advance the token generator at once."""
+    inst_id = 100
+    owner = "mozilla"
+
+    aioresponses.get(
+        f"{GITHUB_API_ENDPOINT}/app/installations",
+        status=200,
+        payload=[{"id": inst_id, "account": {"login": owner}}],
+    )
+    aioresponses.post(
+        f"{GITHUB_API_ENDPOINT}/app/installations/{inst_id}/access_tokens",
+        status=200,
+        payload={"token": "111"},
+    )
+    aioresponses.get(f"{GITHUB_API_ENDPOINT}/octocat", payload={}, repeat=True)
+
+    client = AsyncClient(auth=AppInstallationAuth(AppAuth(42, privkey), owner))
+    await asyncio.gather(*(client.get("/octocat") for _ in range(2)))
+    await client.close()

test/test_client.py

@@ -1,3 +1,4 @@
+import asyncio
 from unittest import mock
 
 import pytest
@@ -60,6 +61,17 @@ async def test_async_client_get_session(async_client):
     }
 
 
+@pytest.mark.asyncio
+async def test_async_client_get_session_concurrent(async_client):
+    """Concurrent callers all get a session, none observe a partial setup."""
+    client = async_client
+
+    sessions = await asyncio.gather(*(client._get_aiohttp_session() for _ in range(25)))
+
+    assert isinstance(client._gql_session, ReconnectingAsyncClientSession)
+    assert all(s == sessions[0] for s in sessions)
+
+
 @pytest.mark.asyncio
 async def test_async_client_get_session_no_token(async_client):
     client = async_client
@@ -70,6 +82,26 @@ async def test_async_client_get_session_no_token(async_client):
     }
 
 
+@pytest.mark.asyncio
+async def test_async_client_recovers_from_connection_failure(async_client):
+    """A failed connect leaves nothing cached, so the next call retries."""
+    client = async_client
+
+    with mock.patch.object(
+        GqlClient, "connect_async", side_effect=OSError("no route to host")
+    ):
+        with pytest.raises(OSError):
+            await client._get_gql_session()
+
+    assert client._prev_token is None
+
+    session = await client._get_aiohttp_session()
+    assert dict(session._default_headers) == {
+        "Accept": "application/vnd.github+json",
+        "Authorization": f"Bearer {client.auth._token}",
+    }
+
+
 def test_sync_client_get_session(sync_client):
     client = sync_client
     assert client._gql_client is None
@@ -108,6 +140,25 @@ def test_sync_client_get_session_no_token(sync_client):
     }
 
 
+def test_sync_client_recovers_from_connection_failure(sync_client):
+    """A failed connect leaves nothing cached, so the next call retries."""
+    client = sync_client
+
+    with mock.patch.object(
+        GqlClient, "connect_sync", side_effect=OSError("no route to host")
+    ):
+        with pytest.raises(OSError):
+            client._get_gql_session()
+
+    assert client._prev_token is None
+
+    client._get_requests_session()
+    assert dict(client._gql_session.transport.headers) == {
+        "Accept": "application/vnd.github+json",
+        "Authorization": f"Bearer {client.auth._token}",
+    }
+
+
 @pytest.mark.asyncio
 async def test_async_client_rest(aioresponses, async_client):
     client = async_client