fix: add asyncio.Lock to prevent session/connector leak on concurrent recreation

When multiple requests detect a closed shared session simultaneously,
they would each create a new aiohttp.ClientSession, leaking intermediate
sessions and their TCP connectors. Added double-checked locking pattern
with asyncio.Lock to ensure only one coroutine recreates the session.

Added concurrent recreation test case.
This commit is contained in:
voidborne-d
2026-03-17 08:08:44 +00:00
parent 7b66c970e9
commit ab4fda2eeb
2 changed files with 85 additions and 11 deletions
+34 -11
View File
@@ -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
@@ -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