9e3c64df03
Bug 1a - RuntimeError when _registry_guard is None: - Added explicit check for _registry_guard initialization - Now raises clear RuntimeError instead of cryptic TypeError - Helps users understand they need to call initialize_share_data() first Bug 1b - Workspace async_locks not visible across processes: - Created new _workspace_async_locks dict for per-process storage - Fixed issue where async_locks modifications in one process were invisible to others - This is correct design since asyncio.Lock objects cannot be pickled/shared Why per-process async_locks: - asyncio.Lock objects cannot be shared across processes - Each process needs its own asyncio.Lock instances for coroutine sync - Cross-process sync is handled by Manager.RLock() in _sync_locks - Within-process async sync is handled by per-process asyncio.Lock Testing: - All 17 existing workspace lock tests pass - Added 3 new tests specifically for bug verification - Total 20 tests passing Impact: - Fixes potential race conditions in multiprocess scenarios - Ensures proper synchronization both across and within processes - Maintains backward compatibility
84 lines
3.2 KiB
Python
84 lines
3.2 KiB
Python
"""
|
|
Test bug fixes for workspace lock implementation.
|
|
|
|
This test module specifically verifies the bug fixes for:
|
|
1. RuntimeError when _registry_guard is None
|
|
2. Per-process workspace async locks in multiprocess mode
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import patch
|
|
from lightrag.kg.shared_storage import (
|
|
_get_workspace_lock,
|
|
initialize_share_data,
|
|
)
|
|
|
|
|
|
def test_error_when_not_initialized():
|
|
"""Test that _get_workspace_lock raises RuntimeError when called before initialize_share_data().
|
|
|
|
Bug Fix: Previously, calling _get_workspace_lock() before initialize_share_data()
|
|
would cause TypeError: 'NoneType' object does not support the context manager protocol
|
|
because _registry_guard was None.
|
|
|
|
Now it should raise a clear RuntimeError with helpful message.
|
|
"""
|
|
# Mock _is_multiprocess to be True to trigger the bug path
|
|
with patch("lightrag.kg.shared_storage._is_multiprocess", True):
|
|
with patch("lightrag.kg.shared_storage._registry_guard", None):
|
|
with pytest.raises(RuntimeError, match="Shared data not initialized"):
|
|
# This should raise RuntimeError, not TypeError
|
|
_get_workspace_lock(
|
|
"test_lock", None, workspace="tenant1", enable_logging=False
|
|
)
|
|
|
|
|
|
def test_workspace_async_locks_per_process():
|
|
"""Test that workspace async locks are stored per-process.
|
|
|
|
Bug Fix: Previously, workspace async locks were added to _async_locks dict,
|
|
which is a regular Python dict in multiprocess mode. This meant modifications
|
|
in one process were not visible to other processes.
|
|
|
|
Now workspace async locks are stored in _workspace_async_locks, which is
|
|
per-process by design (since asyncio.Lock cannot be shared across processes).
|
|
"""
|
|
from lightrag.kg import shared_storage
|
|
|
|
# Initialize in multiprocess mode
|
|
initialize_share_data(workers=2)
|
|
|
|
# Get a workspace lock
|
|
lock1 = shared_storage.get_storage_lock(workspace="tenant1")
|
|
|
|
# Verify the workspace async lock exists in _workspace_async_locks
|
|
assert "tenant1:storage_lock" in shared_storage._workspace_async_locks
|
|
|
|
# Get the same workspace lock again
|
|
lock2 = shared_storage.get_storage_lock(workspace="tenant1")
|
|
|
|
# Both locks should reference the same async_lock instance
|
|
# (within the same process)
|
|
assert lock1._async_lock is lock2._async_lock
|
|
|
|
# Verify it's the same as the one stored in _workspace_async_locks
|
|
assert lock1._async_lock is shared_storage._workspace_async_locks["tenant1:storage_lock"]
|
|
|
|
|
|
def test_multiple_workspace_locks_different_async_locks():
|
|
"""Test that different workspaces have different async locks."""
|
|
from lightrag.kg import shared_storage
|
|
|
|
# Initialize in multiprocess mode
|
|
initialize_share_data(workers=2)
|
|
|
|
# Get locks for different workspaces
|
|
lock1 = shared_storage.get_storage_lock(workspace="tenant1")
|
|
lock2 = shared_storage.get_storage_lock(workspace="tenant2")
|
|
|
|
# They should have different async_lock instances
|
|
assert lock1._async_lock is not lock2._async_lock
|
|
|
|
# Verify both are in _workspace_async_locks
|
|
assert "tenant1:storage_lock" in shared_storage._workspace_async_locks
|
|
assert "tenant2:storage_lock" in shared_storage._workspace_async_locks
|