yujunjun/LightRAG
1
0
Fork 0
mirror of https://github.com/HKUDS/LightRAG.git synced 2025-11-22 13:06:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

51
lightrag/kg/mongo_impl.py
View File

@ -40,36 +40,39 @@ 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:
if cls._instances["db"] is None: async with cls._lock:
uri = os.environ.get( if cls._instances["db"] is None:
"MONGO_URI", uri = os.environ.get(
config.get( "MONGO_URI",
"mongodb", config.get(
"uri", "mongodb",
fallback="mongodb://root:root@localhost:27017/", "uri",
), fallback="mongodb://root:root@localhost:27017/",
) ),
database_name = os.environ.get( )
"MONGO_DATABASE", database_name = os.environ.get(
config.get("mongodb", "database", fallback="LightRAG"), "MONGO_DATABASE",
) config.get("mongodb", "database", fallback="LightRAG"),
client = AsyncMongoClient(uri) )
db = client.get_database(database_name) client = AsyncMongoClient(uri)
cls._instances["db"] = db db = client.get_database(database_name)
cls._instances["ref_count"] = 0 cls._instances["db"] = db
cls._instances["ref_count"] += 1 cls._instances["ref_count"] = 0
return cls._instances["db"] cls._instances["ref_count"] += 1
return cls._instances["db"]
@classmethod @classmethod
async def release_client(cls, db: AsyncDatabase): async def release_client(cls, db: AsyncDatabase):
if db is not None: async with cls._lock:
if db is cls._instances["db"]: if db is not None:
cls._instances["ref_count"] -= 1 if db is cls._instances["db"]:
if cls._instances["ref_count"] == 0: cls._instances["ref_count"] -= 1
cls._instances["db"] = None if cls._instances["ref_count"] == 0:
cls._instances["db"] = None
@final @final