Fix session refreshing (#33)

MarshalX · May 20, 2023 · dd54de9 · dd54de9
1 parent 8048aa4
commit dd54de9
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 6 deletions.
diff --git a/atproto/codegen/clients/generate_async_client.py b/atproto/codegen/clients/generate_async_client.py
@@ -18,10 +18,13 @@ def gen_client(input_filename: str, output_filename: str) -> None:
         '_invoke',
     ]
 
+    code = code.replace('from threading', 'from asyncio')
     code = code.replace('client.raw', 'client.async_raw')
     code = code.replace('class Client', 'class AsyncClient')
     code = code.replace('ClientRaw', 'AsyncClientRaw')
 
+    code = code.replace('with self', 'async with self')
+
     code = code.replace('def', 'async def')
     code = code.replace('async def __', 'def __')
 

diff --git a/atproto/xrpc_client/client/async_client.py b/atproto/xrpc_client/client/async_client.py
@@ -5,6 +5,7 @@
 ##################################################################
 
 import typing as t
+from asyncio import Lock
 from datetime import datetime
 
 from atproto.xrpc_client import models
@@ -29,11 +30,18 @@ def __init__(self, base_url: str = None):
         self._refresh_jwt: t.Optional[str] = None
         self._refresh_jwt_payload: t.Optional['JwtPayload'] = None
 
+        self._refresh_lock = Lock()
+
         self.me: t.Optional[models.AppBskyActorDefs.ProfileViewDetailed] = None
 
     async def _invoke(self, invoke_type: 'InvokeType', **kwargs) -> 'Response':
-        if self.me and self._should_refresh_session():
-            await self._refresh_and_set_session()
+        session_refreshing = kwargs.pop('session_refreshing', False)
+        if session_refreshing:
+            return await super()._invoke(invoke_type, **kwargs)
+
+        async with self._refresh_lock:
+            if self._access_jwt and self._should_refresh_session():
+                await self._refresh_and_set_session()
 
         return await super()._invoke(invoke_type, **kwargs)
 
@@ -47,7 +55,7 @@ async def _get_and_set_session(self, login: str, password: str) -> models.ComAtp
 
     async def _refresh_and_set_session(self) -> models.ComAtprotoServerRefreshSession.Response:
         refresh_session = await self.com.atproto.server.refresh_session(
-            headers=self._get_auth_headers(self._refresh_jwt)
+            headers=self._get_auth_headers(self._refresh_jwt), session_refreshing=True
         )
         self._set_session(refresh_session)
 

diff --git a/atproto/xrpc_client/client/client.py b/atproto/xrpc_client/client/client.py
@@ -1,5 +1,6 @@
 import typing as t
 from datetime import datetime
+from threading import Lock
 
 from atproto.xrpc_client import models
 from atproto.xrpc_client.client.methods_mixin import SessionMethodsMixin
@@ -23,11 +24,18 @@ def __init__(self, base_url: str = None):
         self._refresh_jwt: t.Optional[str] = None
         self._refresh_jwt_payload: t.Optional['JwtPayload'] = None
 
+        self._refresh_lock = Lock()
+
         self.me: t.Optional[models.AppBskyActorDefs.ProfileViewDetailed] = None
 
     def _invoke(self, invoke_type: 'InvokeType', **kwargs) -> 'Response':
-        if self.me and self._should_refresh_session():
-            self._refresh_and_set_session()
+        session_refreshing = kwargs.pop('session_refreshing', False)
+        if session_refreshing:
+            return super()._invoke(invoke_type, **kwargs)
+
+        with self._refresh_lock:
+            if self._access_jwt and self._should_refresh_session():
+                self._refresh_and_set_session()
 
         return super()._invoke(invoke_type, **kwargs)
 
@@ -38,7 +46,9 @@ def _get_and_set_session(self, login: str, password: str) -> models.ComAtprotoSe
         return session
 
     def _refresh_and_set_session(self) -> models.ComAtprotoServerRefreshSession.Response:
-        refresh_session = self.com.atproto.server.refresh_session(headers=self._get_auth_headers(self._refresh_jwt))
+        refresh_session = self.com.atproto.server.refresh_session(
+            headers=self._get_auth_headers(self._refresh_jwt), session_refreshing=True
+        )
         self._set_session(refresh_session)
 
         return refresh_session