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

optimize: improve lock cleanup performance with threshold-based strategy

Browse source 
- Add CLEANUP_THRESHOLD constant (100) to control cleanup frequency
- Modify _release_shared_raw_mp_lock to only scan when cleanup list exceeds threshold
- Modify _release_async_lock to only scan when cleanup list exceeds threshold
This commit is contained in:
yangdx 2025-07-11 23:43:40 +08:00
parent ad99d9ba5a
commit a64c767298
1 changed files with 38 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
lightrag/kg/shared_storage.py
View file

@ -37,6 +37,8 @@ _lock_cleanup_data: Optional[Dict[str, time.time]] = None
_registry_guard = None _registry_guard = None
# Timeout for keyed locks in seconds # Timeout for keyed locks in seconds
CLEANUP_KEYED_LOCKS_AFTER_SECONDS = 300 CLEANUP_KEYED_LOCKS_AFTER_SECONDS = 300
# Threshold for triggering cleanup - only clean when pending list exceeds this size
CLEANUP_THRESHOLD = 100
_initialized = None _initialized = None
@ -60,6 +62,7 @@ _async_locks: Optional[Dict[str, asyncio.Lock]] = None
DEBUG_LOCKS = False DEBUG_LOCKS = False
_debug_n_locks_acquired: int = 0 _debug_n_locks_acquired: int = 0
def inc_debug_n_locks_acquired(): def inc_debug_n_locks_acquired():
global _debug_n_locks_acquired global _debug_n_locks_acquired
if DEBUG_LOCKS: if DEBUG_LOCKS:
@ -324,6 +327,8 @@ def _release_shared_raw_mp_lock(factory_name: str, key: str):
if count == 0: if count == 0:
_lock_cleanup_data[combined_key] = current_time _lock_cleanup_data[combined_key] = current_time
# Only perform cleanup when the pending cleanup list exceeds threshold
if len(_lock_cleanup_data) > CLEANUP_THRESHOLD:
for cleanup_key, cleanup_time in list(_lock_cleanup_data.items()): for cleanup_key, cleanup_time in list(_lock_cleanup_data.items()):
if current_time - cleanup_time > CLEANUP_KEYED_LOCKS_AFTER_SECONDS: if current_time - cleanup_time > CLEANUP_KEYED_LOCKS_AFTER_SECONDS:
_lock_registry.pop(cleanup_key, None) _lock_registry.pop(cleanup_key, None)
@ -347,9 +352,15 @@ class KeyedUnifiedLock:
self._factory_name = factory_name self._factory_name = factory_name
self._default_enable_logging = default_enable_logging self._default_enable_logging = default_enable_logging
self._async_lock: Dict[str, asyncio.Lock] = {} # local keyed locks self._async_lock: Dict[str, asyncio.Lock] = {} # local keyed locks
self._async_lock_count: Dict[str, int] = {} # local keyed locks referenced count self._async_lock_count: Dict[
self._async_lock_cleanup_data: Dict[str, time.time] = {} # local keyed locks timeout str, int
self._mp_locks: Dict[str, mp.synchronize.Lock] = {} # multi-process lock proxies ] = {} # local keyed locks referenced count
self._async_lock_cleanup_data: Dict[
str, time.time
] = {} # local keyed locks timeout
self._mp_locks: Dict[
str, mp.synchronize.Lock
] = {} # multi-process lock proxies
def __call__(self, keys: list[str], *, enable_logging: Optional[bool] = None): def __call__(self, keys: list[str], *, enable_logging: Optional[bool] = None):
""" """
@ -388,7 +399,11 @@ class KeyedUnifiedLock:
self._async_lock_cleanup_data[key] = current_time self._async_lock_cleanup_data[key] = current_time
self._async_lock_count[key] = count self._async_lock_count[key] = count
for cleanup_key, cleanup_time in list(self._async_lock_cleanup_data.items()): # Only perform cleanup when the pending cleanup list exceeds threshold
if len(self._async_lock_cleanup_data) > CLEANUP_THRESHOLD:
for cleanup_key, cleanup_time in list(
self._async_lock_cleanup_data.items()
):
if current_time - cleanup_time > CLEANUP_KEYED_LOCKS_AFTER_SECONDS: if current_time - cleanup_time > CLEANUP_KEYED_LOCKS_AFTER_SECONDS:
self._async_lock.pop(cleanup_key) self._async_lock.pop(cleanup_key)
self._async_lock_count.pop(cleanup_key) self._async_lock_count.pop(cleanup_key)