mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-09 14:22:12 +00:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user