Ensure thread safety in storage update callbacks

- Added storage lock in index_done_callback
- Fixed potential race conditions
This commit is contained in:
yangdx
2025-03-24 02:11:59 +08:00
parent 7e8a2c0e9b
commit ff9cb2138d
4 changed files with 42 additions and 39 deletions

View File

@@ -343,6 +343,7 @@ class FaissVectorDBStorage(BaseVectorStorage):
self._id_to_meta = {} self._id_to_meta = {}
async def index_done_callback(self) -> None: async def index_done_callback(self) -> None:
async with self._storage_lock:
# Check if storage was updated by another process # Check if storage was updated by another process
if is_multiprocess and self.storage_updated.value: if is_multiprocess and self.storage_updated.value:
# Storage was updated by another process, reload data instead of saving # Storage was updated by another process, reload data instead of saving

View File

@@ -206,6 +206,7 @@ class NanoVectorDBStorage(BaseVectorStorage):
async def index_done_callback(self) -> bool: async def index_done_callback(self) -> bool:
"""Save data to disk""" """Save data to disk"""
async with self._storage_lock:
# Check if storage was updated by another process # Check if storage was updated by another process
if is_multiprocess and self.storage_updated.value: if is_multiprocess and self.storage_updated.value:
# Storage was updated by another process, reload data instead of saving # Storage was updated by another process, reload data instead of saving

View File

@@ -401,6 +401,7 @@ class NetworkXStorage(BaseGraphStorage):
async def index_done_callback(self) -> bool: async def index_done_callback(self) -> bool:
"""Save data to disk""" """Save data to disk"""
async with self._storage_lock:
# Check if storage was updated by another process # Check if storage was updated by another process
if is_multiprocess and self.storage_updated.value: if is_multiprocess and self.storage_updated.value:
# Storage was updated by another process, reload data instead of saving # Storage was updated by another process, reload data instead of saving