gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Restore thread safety to MongoDB client manager

Browse source 
- Protected client creation with lock
- Protected client release with lock
This commit is contained in:
yangdx 2025-08-12 16:42:53 +08:00
parent 0b2c3d06c7
commit 41f8ef05b9
1 changed files with 27 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
lightrag/kg/mongo_impl.py
View file

@ -40,9 +40,11 @@ GRAPH_BFS_MODE = os.getenv("MONGO_GRAPH_BFS_MODE", "bidirectional")
class ClientManager: class ClientManager:
_instances = {"db": None, "ref_count": 0} _instances = {"db": None, "ref_count": 0}
_lock = asyncio.Lock()
@classmethod @classmethod
async def get_client(cls) -> AsyncMongoClient: async def get_client(cls) -> AsyncMongoClient:
async with cls._lock:
if cls._instances["db"] is None: if cls._instances["db"] is None:
uri = os.environ.get( uri = os.environ.get(
"MONGO_URI", "MONGO_URI",
@ -65,6 +67,7 @@ class ClientManager:
@classmethod @classmethod
async def release_client(cls, db: AsyncDatabase): async def release_client(cls, db: AsyncDatabase):
async with cls._lock:
if db is not None: if db is not None:
if db is cls._instances["db"]: if db is cls._instances["db"]:
cls._instances["ref_count"] -= 1 cls._instances["ref_count"] -= 1