diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 265a311d39..376bf07f2b 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,3 +1,4 @@ +import asyncio from typing import TYPE_CHECKING, Any, Literal, Optional from fastapi import HTTPException, status @@ -123,11 +124,24 @@ def get_team_id_from_data(data: dict) -> Optional[str]: return None +_shared_session_lock: Optional[asyncio.Lock] = None + + +def _get_shared_session_lock() -> asyncio.Lock: + """Lazily create the shared session lock (must be called within a running event loop).""" + global _shared_session_lock + if _shared_session_lock is None: + _shared_session_lock = asyncio.Lock() + return _shared_session_lock + + async def add_shared_session_to_data(data: dict) -> None: """ Add shared aiohttp session for connection reuse (prevents cold starts). If the session was closed (e.g. due to network interruption or idle timeout), automatically recreates it so connection pooling is restored. + Uses an asyncio.Lock to prevent race conditions where multiple concurrent + requests could each create a new session, leaking intermediate ones. Silently continues without session reuse if import fails or session is unavailable. Args: @@ -146,23 +160,32 @@ async def add_shared_session_to_data(data: dict) -> None: ) elif session is not None and session.closed: # Session was created at startup but has since closed — recreate it - verbose_proxy_logger.warning( - f"SESSION REUSE: Shared aiohttp session is closed (ID: {id(session)}), recreating..." - ) - new_session = await proxy_server._initialize_shared_aiohttp_session() - if new_session is not None: - proxy_server.shared_aiohttp_session = new_session - data["shared_session"] = new_session - else: - verbose_proxy_logger.info( - "SESSION REUSE: Failed to recreate shared session, continuing without session reuse" + # Use lock to prevent concurrent recreation (avoids session/connector leak) + lock = _get_shared_session_lock() + async with lock: + # Double-check under lock — another coroutine may have already recreated it + session = proxy_server.shared_aiohttp_session + if session is not None and not session.closed: + data["shared_session"] = session + return + + verbose_proxy_logger.warning( + f"SESSION REUSE: Shared aiohttp session is closed (ID: {id(session)}), recreating..." ) + new_session = await proxy_server._initialize_shared_aiohttp_session() + if new_session is not None: + proxy_server.shared_aiohttp_session = new_session + data["shared_session"] = new_session + else: + verbose_proxy_logger.info( + "SESSION REUSE: Failed to recreate shared session, continuing without session reuse" + ) else: verbose_proxy_logger.info( "SESSION REUSE: No shared session available for this request" ) except Exception: - # Silently continue without session reuse if import fails or session unavailable + # Silently continue without session reuse if import fails or session is unavailable pass diff --git a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py index e87ca5a163..b71a61d2c2 100644 --- a/tests/test_litellm/proxy/test_aiohttp_session_recovery.py +++ b/tests/test_litellm/proxy/test_aiohttp_session_recovery.py @@ -97,3 +97,54 @@ async def test_add_shared_session_no_session_available(): data = {} await add_shared_session_to_data(data) assert "shared_session" not in data + + +@pytest.mark.asyncio +async def test_add_shared_session_concurrent_recreation_uses_lock(): + """When multiple coroutines detect a closed session concurrently, + only one should recreate it (double-checked locking via asyncio.Lock).""" + import litellm.proxy.route_llm_request as route_module + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.route_llm_request import add_shared_session_to_data + + # Reset the module-level lock so each test is isolated + route_module._shared_session_lock = None + + closed_session = MagicMock() + closed_session.closed = True + + new_session = MagicMock() + new_session.closed = False + + call_count = 0 + + async def mock_init(): + nonlocal call_count + call_count += 1 + # Simulate some async work + await asyncio.sleep(0.01) + proxy_server_module.shared_aiohttp_session = new_session + return new_session + + with patch.object( + proxy_server_module, + "shared_aiohttp_session", + closed_session, + ): + with patch.object( + proxy_server_module, + "_initialize_shared_aiohttp_session", + side_effect=mock_init, + ): + # Launch 5 concurrent calls + results = [{} for _ in range(5)] + await asyncio.gather( + *(add_shared_session_to_data(d) for d in results) + ) + + # Only 1 coroutine should have called _initialize (the rest see the + # re-checked session as open under the lock) + assert call_count == 1, f"Expected 1 init call, got {call_count}" + # All should have the new session + for d in results: + assert d.get("shared_session") is new_session