Merge pull request #710 from zzstoatzz/fix/refresh-locks-memory-leak

fix: prevent unbounded memory growth in token refresh locks

backend/src/backend/_internal/atproto/client.py

···
 from typing import Any
 
 from atproto_oauth.models import OAuthSession
+from cachetools import LRUCache
 
 from backend._internal import Session as AuthSession
 from backend._internal import get_oauth_client, get_session, update_session_tokens
 
 logger = logging.getLogger(__name__)
 
-# per-session locks for token refresh to prevent concurrent refresh races
-_refresh_locks: dict[str, asyncio.Lock] = {}
+# per-session locks for token refresh to prevent concurrent refresh races.
+# uses LRUCache (not TTLCache) to bound memory - LRU eviction is safe because:
+# 1. recently-used locks won't be evicted while in use
+# 2. TTL expiration could evict a lock while a coroutine holds it, breaking mutual exclusion
+_refresh_locks: LRUCache[str, asyncio.Lock] = LRUCache(maxsize=10_000)
 
 
 def reconstruct_oauth_session(oauth_data: dict[str, Any]) -> OAuthSession:

backend/tests/test_token_refresh.py

···
 
 import pytest
 from atproto_oauth.models import OAuthSession
+from cachetools import LRUCache
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import ec
 
 from backend._internal import Session as AuthSession
-from backend._internal.atproto.client import _refresh_session_tokens
+from backend._internal.atproto.client import _refresh_locks, _refresh_session_tokens
 
 
 @pytest.fixture
 def mock_auth_session() -> AuthSession:
     """create mock auth session."""
-    # generate a real EC key and serialize it
-    import cryptography.hazmat.backends
-    import cryptography.hazmat.primitives.asymmetric.ec as ec
-    from cryptography.hazmat.primitives import serialization
-
-    private_key = ec.generate_private_key(
-        ec.SECP256R1(), cryptography.hazmat.backends.default_backend()
-    )
+    private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
 
     dpop_key_pem = private_key.private_bytes(
         encoding=serialization.Encoding.PEM,
···
 @pytest.fixture
 def mock_oauth_session() -> OAuthSession:
     """create mock oauth session."""
-    # defer cryptography import to avoid overhead
-    import cryptography.hazmat.backends
-    import cryptography.hazmat.primitives.asymmetric.ec as ec
-
-    # generate a real key for the mock
-    private_key = ec.generate_private_key(
-        ec.SECP256R1(), cryptography.hazmat.backends.default_backend()
-    )
+    private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
 
     return OAuthSession(
         did="did:plc:test123",
···
         new_token = "new-refreshed-token"
 
         async def mock_refresh_session(self, session: OAuthSession) -> OAuthSession:
-            """mock OAuth client refresh with delay to simulate race."""
             nonlocal refresh_call_count
             refresh_call_count += 1
-
-            # simulate network delay
             await asyncio.sleep(0.1)
-
-            # return updated session with new token
             session.access_token = new_token
             return session
 
         async def mock_get_session(session_id: str) -> AuthSession | None:
-            """mock get_session that returns updated tokens after first refresh."""
             if refresh_call_count > 0:
-                # first refresh completed - return new token
                 mock_auth_session.oauth_session["access_token"] = new_token
             return mock_auth_session
 
         async def mock_update_session_tokens(
             session_id: str, oauth_session_data: dict
         ) -> None:
-            """mock session update."""
             mock_auth_session.oauth_session.update(oauth_session_data)
 
-        # create a mock OAuth client with the refresh method
         mock_oauth_client = type(
             "MockOAuthClient", (), {"refresh_session": mock_refresh_session}
         )()
···
                 side_effect=mock_update_session_tokens,
             ),
         ):
-            # launch 5 concurrent refresh attempts
             tasks = [
                 _refresh_session_tokens(mock_auth_session, mock_oauth_session)
                 for _ in range(5)
             ]
             results = await asyncio.gather(*tasks)
 
-            # all should succeed and get the new token
             assert all(result.access_token == new_token for result in results)
-
-            # but OAuth client should only be called once (the lock worked!)
             assert refresh_call_count == 1
 
     async def test_refresh_failure_uses_fallback(
···
         async def mock_refresh_session_fails(
             self, session: OAuthSession
         ) -> OAuthSession:
-            """mock refresh that always fails."""
             nonlocal refresh_called
             refresh_called = True
             await asyncio.sleep(0.05)
···
         get_session_calls = 0
 
         async def mock_get_session(session_id: str) -> AuthSession | None:
-            """mock get_session that returns updated tokens on retry."""
             nonlocal get_session_calls
             get_session_calls += 1
-
-            # on retry (after failure), return new tokens as if another request succeeded
             if get_session_calls >= 2:
                 mock_auth_session.oauth_session["access_token"] = new_token
-
             return mock_auth_session
 
         async def mock_update_session_tokens(
             session_id: str, oauth_session_data: dict
         ) -> None:
-            """mock session update."""
             mock_auth_session.oauth_session.update(oauth_session_data)
 
-        # create a mock OAuth client with the failing refresh method
         mock_oauth_client = type(
             "MockOAuthClient", (), {"refresh_session": mock_refresh_session_fails}
         )()
···
                 side_effect=mock_update_session_tokens,
             ),
         ):
-            # this should fail to refresh but succeed via fallback
             result = await _refresh_session_tokens(
                 mock_auth_session, mock_oauth_session
             )
 
-            # verify it tried to refresh
             assert refresh_called
-
-            # verify it fell back to reloaded tokens
             assert result.access_token == new_token
 
     async def test_second_request_skips_refresh_if_already_done(
···
         new_token = "already-refreshed-token"
 
         async def mock_refresh_session(self, session: OAuthSession) -> OAuthSession:
-            """mock OAuth client refresh."""
             nonlocal refresh_call_count
             refresh_call_count += 1
             await asyncio.sleep(0.1)
···
         get_session_calls = 0
 
         async def mock_get_session(session_id: str) -> AuthSession | None:
-            """mock get_session that simulates first refresh completing quickly."""
             nonlocal get_session_calls
             get_session_calls += 1
-
-            # on second+ call, act like refresh already happened
             if get_session_calls > 1:
                 mock_auth_session.oauth_session["access_token"] = new_token
-
             return mock_auth_session
 
         async def mock_update_session_tokens(
             session_id: str, oauth_session_data: dict
         ) -> None:
-            """mock session update."""
             mock_auth_session.oauth_session.update(oauth_session_data)
 
-        # create a mock OAuth client with the refresh method
         mock_oauth_client = type(
             "MockOAuthClient", (), {"refresh_session": mock_refresh_session}
         )()
···
                 side_effect=mock_update_session_tokens,
             ),
         ):
-            # first refresh
             result1 = await _refresh_session_tokens(
                 mock_auth_session, mock_oauth_session
             )
             assert result1.access_token == new_token
 
-            # second refresh attempt should skip network call
             result2 = await _refresh_session_tokens(
                 mock_auth_session, mock_oauth_session
             )
             assert result2.access_token == new_token
 
-            # OAuth client should have been called exactly once
             assert refresh_call_count == 1
+
+
+class TestRefreshLocksCache:
+    """test _refresh_locks cache behavior (memory leak prevention)."""
+
+    def test_same_session_returns_same_lock(self):
+        """same session_id should return the same lock instance."""
+        _refresh_locks.clear()
+
+        _refresh_locks["session-a"] = asyncio.Lock()
+        lock1 = _refresh_locks["session-a"]
+        lock2 = _refresh_locks["session-a"]
+
+        assert lock1 is lock2
+
+    def test_different_sessions_have_different_locks(self):
+        """different session_ids should have different lock instances."""
+        _refresh_locks.clear()
+
+        _refresh_locks["session-a"] = asyncio.Lock()
+        _refresh_locks["session-b"] = asyncio.Lock()
+
+        assert _refresh_locks["session-a"] is not _refresh_locks["session-b"]
+
+    def test_cache_is_bounded_by_maxsize(self):
+        """cache should evict entries when full (LRU behavior)."""
+        _refresh_locks.clear()
+
+        assert _refresh_locks.maxsize == 10_000
+
+        for i in range(100):
+            _refresh_locks[f"session-{i}"] = asyncio.Lock()
+
+        assert len(_refresh_locks) == 100
+
+    def test_lru_eviction_order(self):
+        """LRU cache should evict least recently used entries first."""
+        small_cache: LRUCache[str, asyncio.Lock] = LRUCache(maxsize=3)
+
+        small_cache["a"] = asyncio.Lock()
+        small_cache["b"] = asyncio.Lock()
+        small_cache["c"] = asyncio.Lock()
+
+        _ = small_cache["a"]
+        small_cache["d"] = asyncio.Lock()
+
+        assert "a" in small_cache
+        assert "b" not in small_cache
+        assert "c" in small_cache
+        assert "d" in small_cache