0bd162a416
Why this change is needed:
The finalize_share_data() function was not properly cleaning up workspace
lock-related global variables (_sync_locks, _workspace_async_locks, and
lock registry variables). This caused stale references to remain after
finalization, leading to EOFError or BrokenPipeError when trying to
re-initialize or when processes tried to use locks after the Manager
was shut down.
How it solves it:
1. Added comprehensive cleanup of all Manager.dict proxies before Manager
shutdown (_sync_locks, _lock_registry, _lock_registry_count,
_lock_cleanup_data)
2. Added cleanup of per-process _workspace_async_locks dictionary
3. Reset all lock-related globals to None at end of finalization:
- _workers, _lock_registry, _lock_registry_count, _lock_cleanup_data
- _registry_guard, _storage_keyed_lock, _sync_locks
- _workspace_async_locks, _earliest_mp_cleanup_time,
_last_mp_cleanup_time
Impact:
- Prevents EOFError/BrokenPipeError in production deployments
- Enables safe re-initialization after finalization
- Critical for proper resource cleanup in multi-process deployments
- Fixes memory leaks from stale lock references
Testing:
- Added 3 comprehensive tests in test_finalization_cleanup.py
- All 23 workspace lock tests pass (17 original + 3 bug fixes + 3 finalization)
- Tests verify clean re-initialization after finalization in both
single-process and multiprocess modes
92 lines
3.4 KiB
Python
92 lines
3.4 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,
|
|
finalize_share_data,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def cleanup_shared_storage():
|
|
"""Ensure shared storage is cleaned up after each test."""
|
|
yield
|
|
finalize_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
|