wehub-resource-sync
2382 lines
90 KiB
Python
2382 lines
90 KiB
Python
import os
|
||
import sys
|
||
import asyncio
|
||
import multiprocessing as mp
|
||
import uuid
|
||
from multiprocessing.synchronize import Lock as ProcessLock
|
||
from multiprocessing import Manager
|
||
import time
|
||
import logging
|
||
from contextvars import ContextVar
|
||
from typing import Any, Dict, List, Mapping, Optional, Union, TypeVar, Generic
|
||
|
||
from lightrag.constants import (
|
||
DEFAULT_GLOBAL_SLOT_HEARTBEAT_TTL,
|
||
DEFAULT_GLOBAL_SLOT_SUSPECT_GRACE,
|
||
DEFAULT_GLOBAL_SLOT_WAITER_STALE_TTL,
|
||
DEFAULT_QUEUE_STATS_STALE_TTL,
|
||
)
|
||
from lightrag.exceptions import PipelineNotInitializedError
|
||
|
||
DEBUG_LOCKS = False
|
||
|
||
|
||
# Define a direct print function for critical logs that must be visible in all processes
|
||
def direct_log(message, enable_output: bool = True, level: str = "DEBUG"):
|
||
"""
|
||
Log a message directly to stderr to ensure visibility in all processes,
|
||
including the Gunicorn master process.
|
||
|
||
Args:
|
||
message: The message to log
|
||
level: Log level for message (control the visibility of the message by comparing with the current logger level)
|
||
enable_output: Enable or disable log message (Force to turn off the message,)
|
||
"""
|
||
if not enable_output:
|
||
return
|
||
|
||
# Get the current logger level from the lightrag logger
|
||
try:
|
||
from lightrag.utils import logger
|
||
|
||
current_level = logger.getEffectiveLevel()
|
||
except ImportError:
|
||
# Fallback if lightrag.utils is not available
|
||
current_level = 20 # INFO
|
||
|
||
# Convert string level to numeric level for comparison
|
||
level_mapping = {
|
||
"DEBUG": 10, # DEBUG
|
||
"INFO": 20, # INFO
|
||
"WARNING": 30, # WARNING
|
||
"ERROR": 40, # ERROR
|
||
"CRITICAL": 50, # CRITICAL
|
||
}
|
||
message_level = level_mapping.get(level.upper(), logging.DEBUG)
|
||
|
||
if message_level >= current_level:
|
||
print(f"{level}: {message}", file=sys.stderr, flush=True)
|
||
|
||
|
||
T = TypeVar("T")
|
||
LockType = Union[ProcessLock, asyncio.Lock]
|
||
|
||
_is_multiprocess = None
|
||
_workers = None
|
||
_manager = None
|
||
|
||
# Global singleton data for multi-process keyed locks
|
||
_lock_registry: Optional[Dict[str, mp.synchronize.Lock]] = None
|
||
_lock_registry_count: Optional[Dict[str, int]] = None
|
||
_lock_cleanup_data: Optional[Dict[str, time.time]] = None
|
||
_registry_guard = None
|
||
# Timeout for keyed locks in seconds (Default 300)
|
||
CLEANUP_KEYED_LOCKS_AFTER_SECONDS = 300
|
||
# Cleanup pending list threshold for triggering cleanup (Default 500)
|
||
CLEANUP_THRESHOLD = 500
|
||
# Minimum interval between cleanup operations in seconds (Default 30)
|
||
MIN_CLEANUP_INTERVAL_SECONDS = 30
|
||
# Track the earliest cleanup time for efficient cleanup triggering (multiprocess locks only)
|
||
_earliest_mp_cleanup_time: Optional[float] = None
|
||
# Track the last cleanup time to enforce minimum interval (multiprocess locks only)
|
||
_last_mp_cleanup_time: Optional[float] = None
|
||
|
||
_initialized = None
|
||
|
||
# Default workspace for backward compatibility
|
||
_default_workspace: Optional[str] = None
|
||
|
||
# shared data for storage across processes
|
||
_shared_dicts: Optional[Dict[str, Any]] = None
|
||
_init_flags: Optional[Dict[str, bool]] = None # namespace -> initialized
|
||
_update_flags: Optional[Dict[str, bool]] = None # namespace -> updated
|
||
|
||
# locks for mutex access
|
||
_internal_lock: Optional[LockType] = None
|
||
_data_init_lock: Optional[LockType] = None
|
||
# Manager for all keyed locks
|
||
_storage_keyed_lock: Optional["KeyedUnifiedLock"] = None
|
||
|
||
# async locks for coroutine synchronization in multiprocess mode
|
||
_async_locks: Optional[Dict[str, asyncio.Lock]] = None
|
||
|
||
_debug_n_locks_acquired: int = 0
|
||
|
||
# --- Cross-worker global concurrency gate + queue stats aggregation ---
|
||
#
|
||
# Read-only configuration set once by the FIRST initialize_share_data() call
|
||
# (the gunicorn master, before fork — workers inherit it as a module global).
|
||
# Later no-arg calls (e.g. LightRAG.__post_init__) hit the `_initialized`
|
||
# guard and never overwrite it.
|
||
_global_concurrency_limits: Optional[Dict[str, int]] = None
|
||
|
||
# Separator between the queue name and the per-pid suffix in queue-stats
|
||
# namespace keys. \x1f (ASCII unit separator) cannot appear in queue names.
|
||
# (Concurrency gate state needs no separator: one key per group.)
|
||
KEY_SEP = "\x1f"
|
||
|
||
_CONCURRENCY_LEASE_NAMESPACE = "concurrency_leases"
|
||
_QUEUE_STATS_NAMESPACE = "queue_stats"
|
||
|
||
# Heartbeat / staleness parameters (module-level so tests can monkeypatch).
|
||
_heartbeat_ttl: float = DEFAULT_GLOBAL_SLOT_HEARTBEAT_TTL
|
||
_suspect_grace: float = DEFAULT_GLOBAL_SLOT_SUSPECT_GRACE
|
||
_queue_stats_stale_ttl: float = DEFAULT_QUEUE_STATS_STALE_TTL
|
||
_waiter_stale_ttl: float = DEFAULT_GLOBAL_SLOT_WAITER_STALE_TTL
|
||
|
||
# Per-process cached namespace references (avoid the internal lock on every
|
||
# publish). Reset by initialize_share_data()/finalize_share_data().
|
||
_lease_ns_cache: Optional[Dict[str, Any]] = None
|
||
_queue_stats_ns_cache: Optional[Dict[str, Any]] = None
|
||
|
||
# Rate limiting for acquire-failure warnings (fail-closed path).
|
||
_ACQUIRE_FAILURE_LOG_INTERVAL = 30.0
|
||
_last_acquire_failure_log: float = 0.0
|
||
|
||
|
||
def get_final_namespace(namespace: str, workspace: str | None = None):
|
||
global _default_workspace
|
||
if workspace is None:
|
||
workspace = _default_workspace
|
||
|
||
if workspace is None:
|
||
direct_log(
|
||
f"Error: Invoke namespace operation without workspace, pid={os.getpid()}",
|
||
level="ERROR",
|
||
)
|
||
raise ValueError("Invoke namespace operation without workspace")
|
||
|
||
final_namespace = f"{workspace}:{namespace}" if workspace else f"{namespace}"
|
||
return final_namespace
|
||
|
||
|
||
def inc_debug_n_locks_acquired():
|
||
global _debug_n_locks_acquired
|
||
if DEBUG_LOCKS:
|
||
_debug_n_locks_acquired += 1
|
||
print(f"DEBUG: Keyed Lock acquired, total: {_debug_n_locks_acquired:>5}")
|
||
|
||
|
||
def dec_debug_n_locks_acquired():
|
||
global _debug_n_locks_acquired
|
||
if DEBUG_LOCKS:
|
||
if _debug_n_locks_acquired > 0:
|
||
_debug_n_locks_acquired -= 1
|
||
print(f"DEBUG: Keyed Lock released, total: {_debug_n_locks_acquired:>5}")
|
||
else:
|
||
raise RuntimeError("Attempting to release lock when no locks are acquired")
|
||
|
||
|
||
def get_debug_n_locks_acquired():
|
||
global _debug_n_locks_acquired
|
||
return _debug_n_locks_acquired
|
||
|
||
|
||
class UnifiedLock(Generic[T]):
|
||
"""Provide a unified lock interface type for asyncio.Lock and multiprocessing.Lock"""
|
||
|
||
def __init__(
|
||
self,
|
||
lock: Union[ProcessLock, asyncio.Lock],
|
||
is_async: bool,
|
||
name: str = "unnamed",
|
||
enable_logging: bool = True,
|
||
async_lock: Optional[asyncio.Lock] = None,
|
||
):
|
||
self._lock = lock
|
||
self._is_async = is_async
|
||
self._pid = os.getpid() # for debug only
|
||
self._name = name # for debug only
|
||
self._enable_logging = enable_logging # for debug only
|
||
self._async_lock = async_lock # auxiliary lock for coroutine synchronization
|
||
|
||
async def _acquire_mp_lock_in_executor(self) -> None:
|
||
"""Acquire the multiprocess lock without blocking the event loop.
|
||
|
||
Cancellation safety: if this coroutine is cancelled while the
|
||
executor thread is still blocked inside ``acquire()``, the thread
|
||
cannot be interrupted and WILL take the lock eventually — with no
|
||
owner left to release it, every process would deadlock. The shield +
|
||
done-callback below returns such an orphaned acquisition immediately.
|
||
"""
|
||
loop = asyncio.get_running_loop()
|
||
acquire_future = loop.run_in_executor(None, self._lock.acquire)
|
||
try:
|
||
await asyncio.shield(acquire_future)
|
||
except asyncio.CancelledError:
|
||
|
||
def _release_orphaned_acquire(f) -> None:
|
||
if f.cancelled() or f.exception() is not None:
|
||
return
|
||
try:
|
||
self._lock.release()
|
||
except Exception:
|
||
pass
|
||
|
||
acquire_future.add_done_callback(_release_orphaned_acquire)
|
||
raise
|
||
|
||
async def __aenter__(self) -> "UnifiedLock[T]":
|
||
async_gate_acquired = False
|
||
try:
|
||
# If in multiprocess mode and async lock exists, acquire it first
|
||
if not self._is_async and self._async_lock is not None:
|
||
await self._async_lock.acquire()
|
||
async_gate_acquired = True
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Acquired async lock '{self._name}",
|
||
level="DEBUG",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
|
||
# Acquire the main lock
|
||
# Note: self._lock should never be None here as the check has been moved
|
||
# to get_internal_lock() and get_data_init_lock() functions
|
||
if self._is_async:
|
||
await self._lock.acquire()
|
||
else:
|
||
# A Manager lock proxy blocks the calling thread until every
|
||
# other PROCESS ahead of us releases — unbounded. Offload to
|
||
# the default executor so this process's event loop keeps
|
||
# serving while we wait (the async gate above already
|
||
# serializes this process's coroutines, so at most one
|
||
# executor thread per lock key is ever parked here).
|
||
await self._acquire_mp_lock_in_executor()
|
||
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Acquired lock {self._name} (async={self._is_async})",
|
||
level="INFO",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
return self
|
||
except asyncio.CancelledError:
|
||
# Cancellation can arrive while awaiting the executor-offloaded
|
||
# mp acquire (any orphaned acquisition is returned inside
|
||
# _acquire_mp_lock_in_executor). Roll back the per-process gate
|
||
# we already hold so this process's other coroutines never
|
||
# deadlock on it.
|
||
if async_gate_acquired:
|
||
self._async_lock.release()
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Lock acquisition cancelled '{self._name}'",
|
||
level="WARNING",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
raise
|
||
except Exception as e:
|
||
# If main lock acquisition fails, release the async lock if it was acquired
|
||
if async_gate_acquired:
|
||
self._async_lock.release()
|
||
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Failed to acquire lock '{self._name}': {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
raise
|
||
|
||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||
main_lock_released = False
|
||
async_lock_released = False
|
||
try:
|
||
# Release main lock first
|
||
if self._lock is not None:
|
||
if self._is_async:
|
||
self._lock.release()
|
||
else:
|
||
self._lock.release()
|
||
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Released lock {self._name} (async={self._is_async})",
|
||
level="INFO",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
main_lock_released = True
|
||
|
||
# Then release async lock if in multiprocess mode
|
||
if not self._is_async and self._async_lock is not None:
|
||
self._async_lock.release()
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Released async lock {self._name}",
|
||
level="DEBUG",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
async_lock_released = True
|
||
|
||
except Exception as e:
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Failed to release lock '{self._name}': {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
# If main lock release failed but async lock hasn't been attempted yet, try to release it
|
||
if (
|
||
not main_lock_released
|
||
and not async_lock_released
|
||
and not self._is_async
|
||
and self._async_lock is not None
|
||
):
|
||
try:
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Attempting to release async lock after main lock failure",
|
||
level="DEBUG",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
self._async_lock.release()
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Successfully released async lock after main lock failure",
|
||
level="INFO",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
except Exception as inner_e:
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Failed to release async lock after main lock failure: {inner_e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
raise
|
||
|
||
def __enter__(self) -> "UnifiedLock[T]":
|
||
"""For backward compatibility"""
|
||
try:
|
||
if self._is_async:
|
||
raise RuntimeError("Use 'async with' for shared_storage lock")
|
||
|
||
# Acquire the main lock
|
||
# Note: self._lock should never be None here as the check has been moved
|
||
# to get_internal_lock() and get_data_init_lock() functions
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Acquiring lock {self._name} (sync)",
|
||
level="DEBUG",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
self._lock.acquire()
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Acquired lock {self._name} (sync)",
|
||
level="INFO",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
return self
|
||
except Exception as e:
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Failed to acquire lock '{self._name}' (sync): {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
raise
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
"""For backward compatibility"""
|
||
try:
|
||
if self._is_async:
|
||
raise RuntimeError("Use 'async with' for shared_storage lock")
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Releasing lock '{self._name}' (sync)",
|
||
level="DEBUG",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
self._lock.release()
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Released lock {self._name} (sync)",
|
||
level="INFO",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
except Exception as e:
|
||
direct_log(
|
||
f"== Lock == Process {self._pid}: Failed to release lock '{self._name}' (sync): {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
raise
|
||
|
||
def locked(self) -> bool:
|
||
if self._is_async:
|
||
return self._lock.locked()
|
||
else:
|
||
return self._lock.locked()
|
||
|
||
|
||
def _get_combined_key(factory_name: str, key: str) -> str:
|
||
"""Return the combined key for the factory and key."""
|
||
return f"{factory_name}:{key}"
|
||
|
||
|
||
def _perform_lock_cleanup(
|
||
lock_type: str,
|
||
cleanup_data: Dict[str, float],
|
||
lock_registry: Optional[Dict[str, Any]],
|
||
lock_count: Optional[Dict[str, int]],
|
||
earliest_cleanup_time: Optional[float],
|
||
last_cleanup_time: Optional[float],
|
||
current_time: float,
|
||
threshold_check: bool = True,
|
||
) -> tuple[int, Optional[float], Optional[float]]:
|
||
"""
|
||
Generic lock cleanup function to unify cleanup logic for both multiprocess and async locks.
|
||
|
||
Args:
|
||
lock_type: Lock type identifier ("mp" or "async")
|
||
cleanup_data: Cleanup data dictionary
|
||
lock_registry: Lock registry dictionary (can be None for async locks)
|
||
lock_count: Lock count dictionary (can be None for async locks)
|
||
earliest_cleanup_time: Earliest cleanup time
|
||
last_cleanup_time: Last cleanup time
|
||
current_time: Current time
|
||
threshold_check: Whether to check threshold condition (default True, set to False in cleanup_expired_locks)
|
||
|
||
Returns:
|
||
tuple: (cleaned_count, new_earliest_time, new_last_cleanup_time)
|
||
"""
|
||
if len(cleanup_data) == 0:
|
||
return 0, earliest_cleanup_time, last_cleanup_time
|
||
|
||
# If threshold check is needed and threshold not reached, return directly
|
||
if threshold_check and len(cleanup_data) < CLEANUP_THRESHOLD:
|
||
return 0, earliest_cleanup_time, last_cleanup_time
|
||
|
||
# Time rollback detection
|
||
if last_cleanup_time is not None and current_time < last_cleanup_time:
|
||
direct_log(
|
||
f"== {lock_type} Lock == Time rollback detected, resetting cleanup time",
|
||
level="WARNING",
|
||
enable_output=False,
|
||
)
|
||
last_cleanup_time = None
|
||
|
||
# Check cleanup conditions
|
||
has_expired_locks = (
|
||
earliest_cleanup_time is not None
|
||
and current_time - earliest_cleanup_time > CLEANUP_KEYED_LOCKS_AFTER_SECONDS
|
||
)
|
||
|
||
interval_satisfied = (
|
||
last_cleanup_time is None
|
||
or current_time - last_cleanup_time > MIN_CLEANUP_INTERVAL_SECONDS
|
||
)
|
||
|
||
if not (has_expired_locks and interval_satisfied):
|
||
return 0, earliest_cleanup_time, last_cleanup_time
|
||
|
||
try:
|
||
cleaned_count = 0
|
||
new_earliest_time = None
|
||
|
||
# Calculate total count before cleanup
|
||
total_cleanup_len = len(cleanup_data)
|
||
|
||
# Perform cleanup operation
|
||
for cleanup_key, cleanup_time in list(cleanup_data.items()):
|
||
if current_time - cleanup_time > CLEANUP_KEYED_LOCKS_AFTER_SECONDS:
|
||
# Remove from cleanup data
|
||
cleanup_data.pop(cleanup_key, None)
|
||
|
||
# Remove from lock registry if exists
|
||
if lock_registry is not None:
|
||
lock_registry.pop(cleanup_key, None)
|
||
if lock_count is not None:
|
||
lock_count.pop(cleanup_key, None)
|
||
|
||
cleaned_count += 1
|
||
else:
|
||
# Track the earliest time among remaining locks
|
||
if new_earliest_time is None or cleanup_time < new_earliest_time:
|
||
new_earliest_time = cleanup_time
|
||
|
||
# Update state only after successful cleanup
|
||
if cleaned_count > 0:
|
||
new_last_cleanup_time = current_time
|
||
|
||
# Log cleanup results
|
||
next_cleanup_in = max(
|
||
(new_earliest_time + CLEANUP_KEYED_LOCKS_AFTER_SECONDS - current_time)
|
||
if new_earliest_time
|
||
else float("inf"),
|
||
MIN_CLEANUP_INTERVAL_SECONDS,
|
||
)
|
||
|
||
if lock_type == "async":
|
||
direct_log(
|
||
f"== {lock_type} Lock == Cleaned up {cleaned_count}/{total_cleanup_len} expired {lock_type} locks, "
|
||
f"next cleanup in {next_cleanup_in:.1f}s",
|
||
enable_output=False,
|
||
level="INFO",
|
||
)
|
||
else:
|
||
direct_log(
|
||
f"== {lock_type} Lock == Cleaned up {cleaned_count}/{total_cleanup_len} expired locks, "
|
||
f"next cleanup in {next_cleanup_in:.1f}s",
|
||
enable_output=False,
|
||
level="INFO",
|
||
)
|
||
|
||
return cleaned_count, new_earliest_time, new_last_cleanup_time
|
||
else:
|
||
return 0, earliest_cleanup_time, last_cleanup_time
|
||
|
||
except Exception as e:
|
||
direct_log(
|
||
f"== {lock_type} Lock == Cleanup failed: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
return 0, earliest_cleanup_time, last_cleanup_time
|
||
|
||
|
||
def _get_or_create_shared_raw_mp_lock(
|
||
factory_name: str, key: str
|
||
) -> Optional[mp.synchronize.Lock]:
|
||
"""Return the *singleton* manager.Lock() proxy for keyed lock, creating if needed."""
|
||
if not _is_multiprocess:
|
||
return None
|
||
|
||
with _registry_guard:
|
||
combined_key = _get_combined_key(factory_name, key)
|
||
raw = _lock_registry.get(combined_key)
|
||
count = _lock_registry_count.get(combined_key)
|
||
if raw is None:
|
||
raw = _manager.Lock()
|
||
_lock_registry[combined_key] = raw
|
||
count = 0
|
||
else:
|
||
if count is None:
|
||
raise RuntimeError(
|
||
f"Shared-Data lock registry for {factory_name} is corrupted for key {key}"
|
||
)
|
||
if (
|
||
count == 0 and combined_key in _lock_cleanup_data
|
||
): # Reusing an key waiting for cleanup, remove it from cleanup list
|
||
_lock_cleanup_data.pop(combined_key)
|
||
count += 1
|
||
_lock_registry_count[combined_key] = count
|
||
return raw
|
||
|
||
|
||
def _release_shared_raw_mp_lock(factory_name: str, key: str):
|
||
"""Release the *singleton* manager.Lock() proxy for *key*."""
|
||
if not _is_multiprocess:
|
||
return
|
||
|
||
global _earliest_mp_cleanup_time, _last_mp_cleanup_time
|
||
|
||
with _registry_guard:
|
||
combined_key = _get_combined_key(factory_name, key)
|
||
raw = _lock_registry.get(combined_key)
|
||
count = _lock_registry_count.get(combined_key)
|
||
if raw is None and count is None:
|
||
return
|
||
elif raw is None or count is None:
|
||
raise RuntimeError(
|
||
f"Shared-Data lock registry for {factory_name} is corrupted for key {key}"
|
||
)
|
||
|
||
count -= 1
|
||
if count < 0:
|
||
raise RuntimeError(
|
||
f"Attempting to release lock for {key} more times than it was acquired"
|
||
)
|
||
|
||
_lock_registry_count[combined_key] = count
|
||
|
||
current_time = time.time()
|
||
if count == 0:
|
||
_lock_cleanup_data[combined_key] = current_time
|
||
|
||
# Update earliest multiprocess cleanup time (only when earlier)
|
||
if (
|
||
_earliest_mp_cleanup_time is None
|
||
or current_time < _earliest_mp_cleanup_time
|
||
):
|
||
_earliest_mp_cleanup_time = current_time
|
||
|
||
# Use generic cleanup function
|
||
cleaned_count, new_earliest_time, new_last_cleanup_time = _perform_lock_cleanup(
|
||
lock_type="mp",
|
||
cleanup_data=_lock_cleanup_data,
|
||
lock_registry=_lock_registry,
|
||
lock_count=_lock_registry_count,
|
||
earliest_cleanup_time=_earliest_mp_cleanup_time,
|
||
last_cleanup_time=_last_mp_cleanup_time,
|
||
current_time=current_time,
|
||
threshold_check=True,
|
||
)
|
||
|
||
# Update global state if cleanup was performed
|
||
if cleaned_count > 0:
|
||
_earliest_mp_cleanup_time = new_earliest_time
|
||
_last_mp_cleanup_time = new_last_cleanup_time
|
||
|
||
|
||
class KeyedUnifiedLock:
|
||
"""
|
||
Manager for unified keyed locks, supporting both single and multi-process
|
||
|
||
• Keeps only a table of async keyed locks locally
|
||
• Fetches the multi-process keyed lock on every acquire
|
||
• Builds a fresh `UnifiedLock` each time, so `enable_logging`
|
||
(or future options) can vary per call.
|
||
• Supports dynamic namespaces specified at lock usage time
|
||
"""
|
||
|
||
def __init__(self, *, default_enable_logging: bool = True) -> None:
|
||
self._default_enable_logging = default_enable_logging
|
||
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_cleanup_data: Dict[
|
||
str, time.time
|
||
] = {} # local keyed locks timeout
|
||
self._mp_locks: Dict[
|
||
str, mp.synchronize.Lock
|
||
] = {} # multi-process lock proxies
|
||
self._earliest_async_cleanup_time: Optional[float] = (
|
||
None # track earliest async cleanup time
|
||
)
|
||
self._last_async_cleanup_time: Optional[float] = (
|
||
None # track last async cleanup time for minimum interval
|
||
)
|
||
|
||
def __call__(
|
||
self, namespace: str, keys: list[str], *, enable_logging: Optional[bool] = None
|
||
):
|
||
"""
|
||
Ergonomic helper so you can write:
|
||
|
||
async with storage_keyed_lock("namespace", ["key1", "key2"]):
|
||
...
|
||
"""
|
||
if enable_logging is None:
|
||
enable_logging = self._default_enable_logging
|
||
return _KeyedLockContext(
|
||
self,
|
||
namespace=namespace,
|
||
keys=keys,
|
||
enable_logging=enable_logging,
|
||
)
|
||
|
||
def _get_or_create_async_lock(self, combined_key: str) -> asyncio.Lock:
|
||
async_lock = self._async_lock.get(combined_key)
|
||
count = self._async_lock_count.get(combined_key, 0)
|
||
if async_lock is None:
|
||
async_lock = asyncio.Lock()
|
||
self._async_lock[combined_key] = async_lock
|
||
elif count == 0 and combined_key in self._async_lock_cleanup_data:
|
||
self._async_lock_cleanup_data.pop(combined_key)
|
||
count += 1
|
||
self._async_lock_count[combined_key] = count
|
||
return async_lock
|
||
|
||
def _release_async_lock(self, combined_key: str):
|
||
count = self._async_lock_count.get(combined_key, 0)
|
||
count -= 1
|
||
|
||
current_time = time.time()
|
||
if count == 0:
|
||
self._async_lock_cleanup_data[combined_key] = current_time
|
||
|
||
# Update earliest async cleanup time (only when earlier)
|
||
if (
|
||
self._earliest_async_cleanup_time is None
|
||
or current_time < self._earliest_async_cleanup_time
|
||
):
|
||
self._earliest_async_cleanup_time = current_time
|
||
self._async_lock_count[combined_key] = count
|
||
|
||
# Use generic cleanup function
|
||
cleaned_count, new_earliest_time, new_last_cleanup_time = _perform_lock_cleanup(
|
||
lock_type="async",
|
||
cleanup_data=self._async_lock_cleanup_data,
|
||
lock_registry=self._async_lock,
|
||
lock_count=self._async_lock_count,
|
||
earliest_cleanup_time=self._earliest_async_cleanup_time,
|
||
last_cleanup_time=self._last_async_cleanup_time,
|
||
current_time=current_time,
|
||
threshold_check=True,
|
||
)
|
||
|
||
# Update instance state if cleanup was performed
|
||
if cleaned_count > 0:
|
||
self._earliest_async_cleanup_time = new_earliest_time
|
||
self._last_async_cleanup_time = new_last_cleanup_time
|
||
|
||
def _get_lock_for_key(
|
||
self, namespace: str, key: str, enable_logging: bool = False
|
||
) -> UnifiedLock:
|
||
# 1. Create combined key for this namespace:key combination
|
||
combined_key = _get_combined_key(namespace, key)
|
||
|
||
# 2. get (or create) the per‑process async gate for this combined key
|
||
# Is synchronous, so no need to acquire a lock
|
||
async_lock = self._get_or_create_async_lock(combined_key)
|
||
|
||
# 3. fetch the shared raw lock
|
||
raw_lock = _get_or_create_shared_raw_mp_lock(namespace, key)
|
||
is_multiprocess = raw_lock is not None
|
||
if not is_multiprocess:
|
||
raw_lock = async_lock
|
||
|
||
# 4. build a *fresh* UnifiedLock with the chosen logging flag
|
||
if is_multiprocess:
|
||
return UnifiedLock(
|
||
lock=raw_lock,
|
||
is_async=False, # manager.Lock is synchronous
|
||
name=combined_key,
|
||
enable_logging=enable_logging,
|
||
async_lock=async_lock, # prevents event‑loop blocking
|
||
)
|
||
else:
|
||
return UnifiedLock(
|
||
lock=raw_lock,
|
||
is_async=True,
|
||
name=combined_key,
|
||
enable_logging=enable_logging,
|
||
async_lock=None, # No need for async lock in single process mode
|
||
)
|
||
|
||
def _release_lock_for_key(self, namespace: str, key: str):
|
||
combined_key = _get_combined_key(namespace, key)
|
||
self._release_async_lock(combined_key)
|
||
_release_shared_raw_mp_lock(namespace, key)
|
||
|
||
def cleanup_expired_locks(self) -> Dict[str, Any]:
|
||
"""
|
||
Cleanup expired locks for both async and multiprocess locks following the same
|
||
conditions as _release_shared_raw_mp_lock and _release_async_lock functions.
|
||
|
||
Only performs cleanup when both has_expired_locks and interval_satisfied conditions are met
|
||
to avoid too frequent cleanup operations.
|
||
|
||
Since async and multiprocess locks work together, this method cleans up
|
||
both types of expired locks and returns comprehensive statistics.
|
||
|
||
Returns:
|
||
Dict containing cleanup statistics and current status:
|
||
{
|
||
"process_id": 12345,
|
||
"cleanup_performed": {
|
||
"mp_cleaned": 5,
|
||
"async_cleaned": 3
|
||
},
|
||
"current_status": {
|
||
"total_mp_locks": 10,
|
||
"pending_mp_cleanup": 2,
|
||
"total_async_locks": 8,
|
||
"pending_async_cleanup": 1
|
||
}
|
||
}
|
||
"""
|
||
global _lock_registry, _lock_registry_count, _lock_cleanup_data
|
||
global _registry_guard, _earliest_mp_cleanup_time, _last_mp_cleanup_time
|
||
|
||
cleanup_stats = {"mp_cleaned": 0, "async_cleaned": 0}
|
||
|
||
current_time = time.time()
|
||
|
||
# 1. Cleanup multiprocess locks using generic function
|
||
if (
|
||
_is_multiprocess
|
||
and _lock_registry is not None
|
||
and _registry_guard is not None
|
||
):
|
||
try:
|
||
with _registry_guard:
|
||
if _lock_cleanup_data is not None:
|
||
# Use generic cleanup function without threshold check
|
||
cleaned_count, new_earliest_time, new_last_cleanup_time = (
|
||
_perform_lock_cleanup(
|
||
lock_type="mp",
|
||
cleanup_data=_lock_cleanup_data,
|
||
lock_registry=_lock_registry,
|
||
lock_count=_lock_registry_count,
|
||
earliest_cleanup_time=_earliest_mp_cleanup_time,
|
||
last_cleanup_time=_last_mp_cleanup_time,
|
||
current_time=current_time,
|
||
threshold_check=False, # Force cleanup in cleanup_expired_locks
|
||
)
|
||
)
|
||
|
||
# Update global state if cleanup was performed
|
||
if cleaned_count > 0:
|
||
_earliest_mp_cleanup_time = new_earliest_time
|
||
_last_mp_cleanup_time = new_last_cleanup_time
|
||
cleanup_stats["mp_cleaned"] = cleaned_count
|
||
|
||
except Exception as e:
|
||
direct_log(
|
||
f"Error during multiprocess lock cleanup: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
# 2. Cleanup async locks using generic function
|
||
try:
|
||
# Use generic cleanup function without threshold check
|
||
cleaned_count, new_earliest_time, new_last_cleanup_time = (
|
||
_perform_lock_cleanup(
|
||
lock_type="async",
|
||
cleanup_data=self._async_lock_cleanup_data,
|
||
lock_registry=self._async_lock,
|
||
lock_count=self._async_lock_count,
|
||
earliest_cleanup_time=self._earliest_async_cleanup_time,
|
||
last_cleanup_time=self._last_async_cleanup_time,
|
||
current_time=current_time,
|
||
threshold_check=False, # Force cleanup in cleanup_expired_locks
|
||
)
|
||
)
|
||
|
||
# Update instance state if cleanup was performed
|
||
if cleaned_count > 0:
|
||
self._earliest_async_cleanup_time = new_earliest_time
|
||
self._last_async_cleanup_time = new_last_cleanup_time
|
||
cleanup_stats["async_cleaned"] = cleaned_count
|
||
|
||
except Exception as e:
|
||
direct_log(
|
||
f"Error during async lock cleanup: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
# 3. Get current status after cleanup
|
||
current_status = self.get_lock_status()
|
||
|
||
return {
|
||
"process_id": os.getpid(),
|
||
"cleanup_performed": cleanup_stats,
|
||
"current_status": current_status,
|
||
}
|
||
|
||
def get_lock_status(self) -> Dict[str, int]:
|
||
"""
|
||
Get current status of both async and multiprocess locks.
|
||
|
||
Returns comprehensive lock counts for both types of locks since
|
||
they work together in the keyed lock system.
|
||
|
||
Returns:
|
||
Dict containing lock counts:
|
||
{
|
||
"total_mp_locks": 10,
|
||
"pending_mp_cleanup": 2,
|
||
"total_async_locks": 8,
|
||
"pending_async_cleanup": 1
|
||
}
|
||
"""
|
||
global _lock_registry_count, _lock_cleanup_data, _registry_guard
|
||
|
||
status = {
|
||
"total_mp_locks": 0,
|
||
"pending_mp_cleanup": 0,
|
||
"total_async_locks": 0,
|
||
"pending_async_cleanup": 0,
|
||
}
|
||
|
||
try:
|
||
# Count multiprocess locks
|
||
if _is_multiprocess and _lock_registry_count is not None:
|
||
if _registry_guard is not None:
|
||
with _registry_guard:
|
||
status["total_mp_locks"] = len(_lock_registry_count)
|
||
if _lock_cleanup_data is not None:
|
||
status["pending_mp_cleanup"] = len(_lock_cleanup_data)
|
||
|
||
# Count async locks
|
||
status["total_async_locks"] = len(self._async_lock_count)
|
||
status["pending_async_cleanup"] = len(self._async_lock_cleanup_data)
|
||
|
||
except Exception as e:
|
||
direct_log(
|
||
f"Error getting keyed lock status: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
return status
|
||
|
||
|
||
class _KeyedLockContext:
|
||
def __init__(
|
||
self,
|
||
parent: KeyedUnifiedLock,
|
||
namespace: str,
|
||
keys: list[str],
|
||
enable_logging: bool,
|
||
) -> None:
|
||
self._parent = parent
|
||
self._namespace = namespace
|
||
|
||
# The sorting is critical to ensure proper lock and release order
|
||
# to avoid deadlocks
|
||
self._keys = sorted(keys)
|
||
self._enable_logging = (
|
||
enable_logging
|
||
if enable_logging is not None
|
||
else parent._default_enable_logging
|
||
)
|
||
self._ul: Optional[List[Dict[str, Any]]] = None # set in __aenter__
|
||
|
||
# ----- enter -----
|
||
async def __aenter__(self):
|
||
if self._ul is not None:
|
||
raise RuntimeError("KeyedUnifiedLock already acquired in current context")
|
||
|
||
self._ul = []
|
||
|
||
try:
|
||
# Acquire locks for all keys in the namespace
|
||
for key in self._keys:
|
||
lock = None
|
||
entry = None
|
||
|
||
try:
|
||
# 1. Get lock object (reference count is incremented here)
|
||
lock = self._parent._get_lock_for_key(
|
||
self._namespace, key, enable_logging=self._enable_logging
|
||
)
|
||
|
||
# 2. Immediately create and add entry to list (critical for rollback to work)
|
||
entry = {
|
||
"key": key,
|
||
"lock": lock,
|
||
"entered": False,
|
||
"debug_inc": False,
|
||
"ref_incremented": True, # Mark that reference count has been incremented
|
||
}
|
||
self._ul.append(
|
||
entry
|
||
) # Add immediately after _get_lock_for_key for rollback to work
|
||
|
||
# 3. Try to acquire the lock
|
||
# Use try-finally to ensure state is updated atomically
|
||
lock_acquired = False
|
||
try:
|
||
await lock.__aenter__()
|
||
lock_acquired = True # Lock successfully acquired
|
||
finally:
|
||
if lock_acquired:
|
||
entry["entered"] = True
|
||
inc_debug_n_locks_acquired()
|
||
entry["debug_inc"] = True
|
||
|
||
except asyncio.CancelledError:
|
||
# Lock acquisition was cancelled
|
||
# The finally block above ensures entry["entered"] is correct
|
||
direct_log(
|
||
f"Lock acquisition cancelled for key {key}",
|
||
level="WARNING",
|
||
enable_output=self._enable_logging,
|
||
)
|
||
raise
|
||
except Exception as e:
|
||
# Other exceptions, log and re-raise
|
||
direct_log(
|
||
f"Lock acquisition failed for key {key}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
raise
|
||
|
||
return self
|
||
|
||
except BaseException:
|
||
# Critical: if any exception occurs (including CancelledError) during lock acquisition,
|
||
# we must rollback all already acquired locks to prevent lock leaks
|
||
# Use shield to ensure rollback completes
|
||
await asyncio.shield(self._rollback_acquired_locks())
|
||
raise
|
||
|
||
async def _rollback_acquired_locks(self):
|
||
"""Rollback all acquired locks in case of exception during __aenter__"""
|
||
if not self._ul:
|
||
return
|
||
|
||
async def rollback_single_entry(entry):
|
||
"""Rollback a single lock acquisition"""
|
||
key = entry["key"]
|
||
lock = entry["lock"]
|
||
debug_inc = entry["debug_inc"]
|
||
entered = entry["entered"]
|
||
ref_incremented = entry.get(
|
||
"ref_incremented", True
|
||
) # Default to True for safety
|
||
|
||
errors = []
|
||
|
||
# 1. If lock was acquired, release it
|
||
if entered:
|
||
try:
|
||
await lock.__aexit__(None, None, None)
|
||
except Exception as e:
|
||
errors.append(("lock_exit", e))
|
||
direct_log(
|
||
f"Lock rollback error for key {key}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
# 2. Release reference count (if it was incremented)
|
||
if ref_incremented:
|
||
try:
|
||
self._parent._release_lock_for_key(self._namespace, key)
|
||
except Exception as e:
|
||
errors.append(("ref_release", e))
|
||
direct_log(
|
||
f"Lock rollback reference release error for key {key}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
# 3. Decrement debug counter
|
||
if debug_inc:
|
||
try:
|
||
dec_debug_n_locks_acquired()
|
||
except Exception as e:
|
||
errors.append(("debug_dec", e))
|
||
direct_log(
|
||
f"Lock rollback counter decrementing error for key {key}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
return errors
|
||
|
||
# Release already acquired locks in reverse order
|
||
for entry in reversed(self._ul):
|
||
# Use shield to protect each lock's rollback
|
||
try:
|
||
await asyncio.shield(rollback_single_entry(entry))
|
||
except Exception as e:
|
||
# Log but continue rolling back other locks
|
||
direct_log(
|
||
f"Lock rollback unexpected error for {entry['key']}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
self._ul = None
|
||
|
||
# ----- exit -----
|
||
async def __aexit__(self, exc_type, exc, tb):
|
||
if self._ul is None:
|
||
return
|
||
|
||
async def release_all_locks():
|
||
"""Release all locks with comprehensive error handling, protected from cancellation"""
|
||
|
||
async def release_single_entry(entry, exc_type, exc, tb):
|
||
"""Release a single lock with full protection"""
|
||
key = entry["key"]
|
||
lock = entry["lock"]
|
||
debug_inc = entry["debug_inc"]
|
||
entered = entry["entered"]
|
||
|
||
errors = []
|
||
|
||
# 1. Release the lock
|
||
if entered:
|
||
try:
|
||
await lock.__aexit__(exc_type, exc, tb)
|
||
except Exception as e:
|
||
errors.append(("lock_exit", e))
|
||
direct_log(
|
||
f"Lock release error for key {key}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
# 2. Release reference count
|
||
try:
|
||
self._parent._release_lock_for_key(self._namespace, key)
|
||
except Exception as e:
|
||
errors.append(("ref_release", e))
|
||
direct_log(
|
||
f"Lock release reference error for key {key}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
# 3. Decrement debug counter
|
||
if debug_inc:
|
||
try:
|
||
dec_debug_n_locks_acquired()
|
||
except Exception as e:
|
||
errors.append(("debug_dec", e))
|
||
direct_log(
|
||
f"Lock release counter decrementing error for key {key}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
return errors
|
||
|
||
all_errors = []
|
||
|
||
# Release locks in reverse order
|
||
# This entire loop is protected by the outer shield
|
||
for entry in reversed(self._ul):
|
||
try:
|
||
errors = await release_single_entry(entry, exc_type, exc, tb)
|
||
for error_type, error in errors:
|
||
all_errors.append((entry["key"], error_type, error))
|
||
except Exception as e:
|
||
all_errors.append((entry["key"], "unexpected", e))
|
||
direct_log(
|
||
f"Lock release unexpected error for {entry['key']}: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
|
||
return all_errors
|
||
|
||
# CRITICAL: Protect the entire release process with shield
|
||
# This ensures that even if cancellation occurs, all locks are released
|
||
try:
|
||
all_errors = await asyncio.shield(release_all_locks())
|
||
except Exception as e:
|
||
direct_log(
|
||
f"Critical error during __aexit__ cleanup: {e}",
|
||
level="ERROR",
|
||
enable_output=True,
|
||
)
|
||
all_errors = []
|
||
finally:
|
||
# Always clear the lock list, even if shield was cancelled
|
||
self._ul = None
|
||
|
||
# If there were release errors and no other exception, raise the first release error
|
||
if all_errors and exc_type is None:
|
||
raise all_errors[0][2] # (key, error_type, error)
|
||
|
||
|
||
def get_internal_lock(enable_logging: bool = False) -> UnifiedLock:
|
||
"""return unified storage lock for data consistency"""
|
||
if _internal_lock is None:
|
||
raise RuntimeError(
|
||
"Shared data not initialized. Call initialize_share_data() before using locks!"
|
||
)
|
||
async_lock = _async_locks.get("internal_lock") if _is_multiprocess else None
|
||
return UnifiedLock(
|
||
lock=_internal_lock,
|
||
is_async=not _is_multiprocess,
|
||
name="internal_lock",
|
||
enable_logging=enable_logging,
|
||
async_lock=async_lock,
|
||
)
|
||
|
||
|
||
# Workspace based storage_lock is implemented by get_storage_keyed_lock instead.
|
||
# Workspace based pipeline_status_lock is implemented by get_storage_keyed_lock instead.
|
||
# No need to implement graph_db_lock:
|
||
# data integrity is ensured by entity level keyed-lock and allowing only one process to hold pipeline at a time.
|
||
|
||
|
||
def get_storage_keyed_lock(
|
||
keys: str | list[str], namespace: str = "default", enable_logging: bool = False
|
||
) -> _KeyedLockContext:
|
||
"""Return unified storage keyed lock for ensuring atomic operations across different namespaces"""
|
||
global _storage_keyed_lock
|
||
if _storage_keyed_lock is None:
|
||
raise RuntimeError("Shared-Data is not initialized")
|
||
if isinstance(keys, str):
|
||
keys = [keys]
|
||
return _storage_keyed_lock(namespace, keys, enable_logging=enable_logging)
|
||
|
||
|
||
def get_data_init_lock(enable_logging: bool = False) -> UnifiedLock:
|
||
"""return unified data initialization lock for ensuring atomic data initialization"""
|
||
if _data_init_lock is None:
|
||
raise RuntimeError(
|
||
"Shared data not initialized. Call initialize_share_data() before using locks!"
|
||
)
|
||
async_lock = _async_locks.get("data_init_lock") if _is_multiprocess else None
|
||
return UnifiedLock(
|
||
lock=_data_init_lock,
|
||
is_async=not _is_multiprocess,
|
||
name="data_init_lock",
|
||
enable_logging=enable_logging,
|
||
async_lock=async_lock,
|
||
)
|
||
|
||
|
||
def cleanup_keyed_lock() -> Dict[str, Any]:
|
||
"""
|
||
Force cleanup of expired keyed locks and return comprehensive status information.
|
||
|
||
This function actively cleans up expired locks for both async and multiprocess locks,
|
||
then returns detailed statistics about the cleanup operation and current lock status.
|
||
|
||
Returns:
|
||
Same as cleanup_expired_locks in KeyedUnifiedLock
|
||
"""
|
||
global _storage_keyed_lock
|
||
|
||
# Check if shared storage is initialized
|
||
if not _initialized or _storage_keyed_lock is None:
|
||
return {
|
||
"process_id": os.getpid(),
|
||
"cleanup_performed": {"mp_cleaned": 0, "async_cleaned": 0},
|
||
"current_status": {
|
||
"total_mp_locks": 0,
|
||
"pending_mp_cleanup": 0,
|
||
"total_async_locks": 0,
|
||
"pending_async_cleanup": 0,
|
||
},
|
||
}
|
||
|
||
return _storage_keyed_lock.cleanup_expired_locks()
|
||
|
||
|
||
def get_keyed_lock_status() -> Dict[str, Any]:
|
||
"""
|
||
Get current status of keyed locks without performing cleanup.
|
||
|
||
This function provides a read-only view of the current lock counts
|
||
for both multiprocess and async locks, including pending cleanup counts.
|
||
|
||
Returns:
|
||
Same as get_lock_status in KeyedUnifiedLock
|
||
"""
|
||
global _storage_keyed_lock
|
||
|
||
# Check if shared storage is initialized
|
||
if not _initialized or _storage_keyed_lock is None:
|
||
return {
|
||
"process_id": os.getpid(),
|
||
"total_mp_locks": 0,
|
||
"pending_mp_cleanup": 0,
|
||
"total_async_locks": 0,
|
||
"pending_async_cleanup": 0,
|
||
}
|
||
|
||
status = _storage_keyed_lock.get_lock_status()
|
||
status["process_id"] = os.getpid()
|
||
return status
|
||
|
||
|
||
def initialize_share_data(
|
||
workers: int = 1,
|
||
global_concurrency_limits: Optional[Mapping[str, int]] = None,
|
||
):
|
||
"""
|
||
Initialize shared storage data for single or multi-process mode.
|
||
|
||
When used with Gunicorn's preload feature, this function is called once in the
|
||
master process before forking worker processes, allowing all workers to share
|
||
the same initialized data.
|
||
|
||
In single-process mode, this function is called in FASTAPI lifespan function.
|
||
|
||
The function determines whether to use cross-process shared variables for data storage
|
||
based on the number of workers. If workers=1, it uses thread locks and local dictionaries.
|
||
If workers>1, it uses process locks and shared dictionaries managed by multiprocessing.Manager.
|
||
|
||
Args:
|
||
workers (int): Number of worker processes. If 1, single-process mode is used.
|
||
If > 1, multi-process mode with shared memory is used.
|
||
global_concurrency_limits: Optional mapping of concurrency group name
|
||
(e.g. "llm:extract", "embedding", "rerank") to the
|
||
cross-worker global max concurrency for that group.
|
||
Read-only after initialization; later calls (which hit
|
||
the already-initialized guard) never overwrite it.
|
||
"""
|
||
global \
|
||
_manager, \
|
||
_workers, \
|
||
_is_multiprocess, \
|
||
_lock_registry, \
|
||
_lock_registry_count, \
|
||
_lock_cleanup_data, \
|
||
_registry_guard, \
|
||
_internal_lock, \
|
||
_data_init_lock, \
|
||
_shared_dicts, \
|
||
_init_flags, \
|
||
_initialized, \
|
||
_update_flags, \
|
||
_async_locks, \
|
||
_storage_keyed_lock, \
|
||
_earliest_mp_cleanup_time, \
|
||
_last_mp_cleanup_time, \
|
||
_global_concurrency_limits, \
|
||
_lease_ns_cache, \
|
||
_queue_stats_ns_cache
|
||
|
||
# Check if already initialized
|
||
if _initialized:
|
||
direct_log(
|
||
f"Process {os.getpid()} Shared-Data already initialized (multiprocess={_is_multiprocess})"
|
||
)
|
||
return
|
||
|
||
_workers = workers
|
||
_global_concurrency_limits = (
|
||
{
|
||
str(group): int(limit)
|
||
for group, limit in global_concurrency_limits.items()
|
||
if limit is not None and int(limit) > 0
|
||
}
|
||
if global_concurrency_limits
|
||
else {}
|
||
)
|
||
_lease_ns_cache = None
|
||
_queue_stats_ns_cache = None
|
||
if _global_concurrency_limits:
|
||
direct_log(
|
||
f"Process {os.getpid()} Global concurrency limits: {_global_concurrency_limits}",
|
||
level="INFO",
|
||
)
|
||
|
||
if workers > 1:
|
||
_is_multiprocess = True
|
||
_manager = Manager()
|
||
_lock_registry = _manager.dict()
|
||
_lock_registry_count = _manager.dict()
|
||
_lock_cleanup_data = _manager.dict()
|
||
_registry_guard = _manager.RLock()
|
||
_internal_lock = _manager.Lock()
|
||
_data_init_lock = _manager.Lock()
|
||
_shared_dicts = _manager.dict()
|
||
_init_flags = _manager.dict()
|
||
_update_flags = _manager.dict()
|
||
|
||
_storage_keyed_lock = KeyedUnifiedLock()
|
||
|
||
# Initialize async locks for multiprocess mode
|
||
_async_locks = {
|
||
"internal_lock": asyncio.Lock(),
|
||
"graph_db_lock": asyncio.Lock(),
|
||
"data_init_lock": asyncio.Lock(),
|
||
}
|
||
|
||
direct_log(
|
||
f"Process {os.getpid()} Shared-Data created for Multiple Process (workers={workers})"
|
||
)
|
||
else:
|
||
_is_multiprocess = False
|
||
_internal_lock = asyncio.Lock()
|
||
_data_init_lock = asyncio.Lock()
|
||
_shared_dicts = {}
|
||
_init_flags = {}
|
||
_update_flags = {}
|
||
_async_locks = None # No need for async locks in single process mode
|
||
|
||
_storage_keyed_lock = KeyedUnifiedLock()
|
||
direct_log(f"Process {os.getpid()} Shared-Data created for Single Process")
|
||
|
||
# Initialize multiprocess cleanup times
|
||
_earliest_mp_cleanup_time = None
|
||
_last_mp_cleanup_time = None
|
||
|
||
# Mark as initialized
|
||
_initialized = True
|
||
|
||
|
||
async def initialize_pipeline_status(workspace: str | None = None):
|
||
"""
|
||
Initialize pipeline_status share data with default values.
|
||
This function could be called before during FASTAPI lifespan for each worker.
|
||
|
||
Args:
|
||
workspace: Optional workspace identifier for pipeline_status of specific workspace.
|
||
If None or empty string, uses the default workspace set by
|
||
set_default_workspace().
|
||
"""
|
||
pipeline_namespace = await get_namespace_data(
|
||
"pipeline_status", first_init=True, workspace=workspace
|
||
)
|
||
|
||
async with get_internal_lock():
|
||
# Check if already initialized by checking for required fields
|
||
if "busy" in pipeline_namespace:
|
||
return
|
||
|
||
# Create a shared list object for history_messages
|
||
history_messages = _manager.list() if _is_multiprocess else []
|
||
pipeline_namespace.update(
|
||
{
|
||
"autoscanned": False, # Auto-scan started
|
||
"busy": False, # Control concurrent processes
|
||
# Destructive subset of ``busy``: clear / delete jobs that
|
||
# DROP storages or remove input files. Concurrent enqueue
|
||
# would race against the drop and silently lose the
|
||
# accepted document, so reservation and the enqueue
|
||
# last-line guard reject when this is True. ``busy`` on
|
||
# its own (the processing loop) remains compatible with
|
||
# concurrent enqueue via request_pending.
|
||
"destructive_busy": False,
|
||
"scanning": False, # /documents/scan task running (whole lifecycle)
|
||
# Exclusive subset of ``scanning``: only True during the
|
||
# scan's *classification* phase, when run_scanning_process
|
||
# is reading doc_status to classify files (PROCESSED →
|
||
# archive, FAILED-without-full_docs → retry-as-new, etc.)
|
||
# and possibly deleting stale stubs. After classification
|
||
# the scan transitions to its processing phase (which
|
||
# behaves like any other busy processing run) and clears
|
||
# this flag, allowing concurrent uploads to land in
|
||
# doc_status while the scan-driven processing finishes.
|
||
"scanning_exclusive": False,
|
||
# Counter of upload/insert endpoints that have passed the
|
||
# idle preflight but whose background enqueue has not yet
|
||
# run. Closes the preflight-to-background race: scan
|
||
# refuses to start while this is > 0 so the bg task is
|
||
# guaranteed to see scanning=False at enqueue time.
|
||
"pending_enqueues": 0,
|
||
"job_name": "-", # Current job name (indexing files/indexing texts)
|
||
"job_start": None, # Job start time
|
||
"docs": 0, # Total number of documents to be indexed
|
||
"batchs": 0, # Number of batches for processing documents
|
||
"cur_batch": 0, # Current processing batch
|
||
"request_pending": False, # Flag for pending request for processing
|
||
"latest_message": "", # Latest message from pipeline processing
|
||
"history_messages": history_messages, # 使用共享列表对象
|
||
}
|
||
)
|
||
|
||
final_namespace = get_final_namespace("pipeline_status", workspace)
|
||
direct_log(
|
||
f"Process {os.getpid()} Pipeline namespace '{final_namespace}' initialized"
|
||
)
|
||
|
||
|
||
async def get_update_flag(namespace: str, workspace: str | None = None):
|
||
"""
|
||
Create a namespace's update flag for a workers.
|
||
Returen the update flag to caller for referencing or reset.
|
||
"""
|
||
global _update_flags
|
||
if _update_flags is None:
|
||
raise ValueError("Try to create namespace before Shared-Data is initialized")
|
||
|
||
final_namespace = get_final_namespace(namespace, workspace)
|
||
|
||
async with get_internal_lock():
|
||
if final_namespace not in _update_flags:
|
||
if _is_multiprocess and _manager is not None:
|
||
_update_flags[final_namespace] = _manager.list()
|
||
else:
|
||
_update_flags[final_namespace] = []
|
||
direct_log(
|
||
f"Process {os.getpid()} initialized updated flags for namespace: [{final_namespace}]"
|
||
)
|
||
|
||
if _is_multiprocess and _manager is not None:
|
||
new_update_flag = _manager.Value("b", False)
|
||
else:
|
||
# Create a simple mutable object to store boolean value for compatibility with mutiprocess
|
||
class MutableBoolean:
|
||
def __init__(self, initial_value=False):
|
||
self.value = initial_value
|
||
|
||
new_update_flag = MutableBoolean(False)
|
||
|
||
_update_flags[final_namespace].append(new_update_flag)
|
||
return new_update_flag
|
||
|
||
|
||
async def set_all_update_flags(namespace: str, workspace: str | None = None):
|
||
"""Set all update flag of namespace indicating all workers need to reload data from files"""
|
||
global _update_flags
|
||
if _update_flags is None:
|
||
raise ValueError("Try to create namespace before Shared-Data is initialized")
|
||
|
||
final_namespace = get_final_namespace(namespace, workspace)
|
||
|
||
async with get_internal_lock():
|
||
if final_namespace not in _update_flags:
|
||
raise ValueError(f"Namespace {final_namespace} not found in update flags")
|
||
# Update flags for both modes
|
||
for i in range(len(_update_flags[final_namespace])):
|
||
_update_flags[final_namespace][i].value = True
|
||
|
||
|
||
async def clear_all_update_flags(namespace: str, workspace: str | None = None):
|
||
"""Clear all update flag of namespace indicating all workers need to reload data from files"""
|
||
global _update_flags
|
||
if _update_flags is None:
|
||
raise ValueError("Try to create namespace before Shared-Data is initialized")
|
||
|
||
final_namespace = get_final_namespace(namespace, workspace)
|
||
|
||
async with get_internal_lock():
|
||
if final_namespace not in _update_flags:
|
||
raise ValueError(f"Namespace {final_namespace} not found in update flags")
|
||
# Update flags for both modes
|
||
for i in range(len(_update_flags[final_namespace])):
|
||
_update_flags[final_namespace][i].value = False
|
||
|
||
|
||
async def get_all_update_flags_status(workspace: str | None = None) -> Dict[str, list]:
|
||
"""
|
||
Get update flags status for all namespaces.
|
||
|
||
Returns:
|
||
Dict[str, list]: A dictionary mapping namespace names to lists of update flag statuses
|
||
"""
|
||
if _update_flags is None:
|
||
return {}
|
||
|
||
if workspace is None:
|
||
workspace = get_default_workspace()
|
||
|
||
result = {}
|
||
async with get_internal_lock():
|
||
for namespace, flags in _update_flags.items():
|
||
# Check if namespace has a workspace prefix (contains ':')
|
||
if ":" in namespace:
|
||
# Namespace has workspace prefix like "space1:pipeline_status"
|
||
# Only include if workspace matches the prefix
|
||
# Use rsplit to split from the right since workspace can contain colons
|
||
namespace_split = namespace.rsplit(":", 1)
|
||
if not workspace or namespace_split[0] != workspace:
|
||
continue
|
||
else:
|
||
# Namespace has no workspace prefix like "pipeline_status"
|
||
# Only include if we're querying the default (empty) workspace
|
||
if workspace:
|
||
continue
|
||
|
||
worker_statuses = []
|
||
for flag in flags:
|
||
if _is_multiprocess:
|
||
worker_statuses.append(flag.value)
|
||
else:
|
||
worker_statuses.append(flag)
|
||
result[namespace] = worker_statuses
|
||
|
||
return result
|
||
|
||
|
||
async def try_initialize_namespace(
|
||
namespace: str, workspace: str | None = None
|
||
) -> bool:
|
||
"""
|
||
Returns True if the current worker(process) gets initialization permission for loading data later.
|
||
The worker does not get the permission is prohibited to load data from files.
|
||
"""
|
||
global _init_flags, _manager
|
||
|
||
if _init_flags is None:
|
||
raise ValueError("Try to create nanmespace before Shared-Data is initialized")
|
||
|
||
final_namespace = get_final_namespace(namespace, workspace)
|
||
|
||
async with get_internal_lock():
|
||
if final_namespace not in _init_flags:
|
||
_init_flags[final_namespace] = True
|
||
direct_log(
|
||
f"Process {os.getpid()} ready to initialize storage namespace: [{final_namespace}]"
|
||
)
|
||
return True
|
||
direct_log(
|
||
f"Process {os.getpid()} storage namespace already initialized: [{final_namespace}]"
|
||
)
|
||
|
||
return False
|
||
|
||
|
||
async def get_namespace_data(
|
||
namespace: str, first_init: bool = False, workspace: str | None = None
|
||
) -> Dict[str, Any]:
|
||
"""get the shared data reference for specific namespace
|
||
|
||
Args:
|
||
namespace: The namespace to retrieve
|
||
first_init: If True, allows pipeline_status namespace to create namespace if it doesn't exist.
|
||
Prevent getting pipeline_status namespace without initialize_pipeline_status().
|
||
This parameter is used internally by initialize_pipeline_status().
|
||
workspace: Workspace identifier (may be empty string for global namespace)
|
||
"""
|
||
if _shared_dicts is None:
|
||
direct_log(
|
||
f"Error: Try to getnanmespace before it is initialized, pid={os.getpid()}",
|
||
level="ERROR",
|
||
)
|
||
raise ValueError("Shared dictionaries not initialized")
|
||
|
||
final_namespace = get_final_namespace(namespace, workspace)
|
||
|
||
async with get_internal_lock():
|
||
if final_namespace not in _shared_dicts:
|
||
# Special handling for pipeline_status namespace
|
||
if (
|
||
final_namespace.endswith(":pipeline_status")
|
||
or final_namespace == "pipeline_status"
|
||
) and not first_init:
|
||
# Check if pipeline_status should have been initialized but wasn't
|
||
# This helps users to call initialize_pipeline_status() before get_namespace_data()
|
||
raise PipelineNotInitializedError(final_namespace)
|
||
|
||
# For other namespaces or when allow_create=True, create them dynamically
|
||
if _is_multiprocess and _manager is not None:
|
||
_shared_dicts[final_namespace] = _manager.dict()
|
||
else:
|
||
_shared_dicts[final_namespace] = {}
|
||
|
||
return _shared_dicts[final_namespace]
|
||
|
||
|
||
class NamespaceLock:
|
||
"""
|
||
Reusable namespace lock wrapper that creates a fresh context on each use.
|
||
|
||
This class solves the lock re-entrance and concurrent coroutine issues by using
|
||
contextvars.ContextVar to provide per-coroutine storage. Each coroutine gets its
|
||
own independent lock context, preventing state interference between concurrent
|
||
coroutines using the same NamespaceLock instance.
|
||
|
||
Example:
|
||
lock = NamespaceLock("my_namespace", "workspace1")
|
||
|
||
# Can be used multiple times safely
|
||
async with lock:
|
||
await do_something()
|
||
|
||
# Can even be used concurrently without deadlock
|
||
await asyncio.gather(
|
||
coroutine_1(lock), # Each gets its own context
|
||
coroutine_2(lock) # No state interference
|
||
)
|
||
"""
|
||
|
||
def __init__(
|
||
self, namespace: str, workspace: str | None = None, enable_logging: bool = False
|
||
):
|
||
self._namespace = namespace
|
||
self._workspace = workspace
|
||
self._enable_logging = enable_logging
|
||
# Use ContextVar to provide per-coroutine storage for lock context
|
||
# This ensures each coroutine has its own independent context
|
||
self._ctx_var: ContextVar[Optional[_KeyedLockContext]] = ContextVar(
|
||
"lock_ctx", default=None
|
||
)
|
||
|
||
async def __aenter__(self):
|
||
"""Create a fresh context each time we enter"""
|
||
# Check if this coroutine already has an active lock context
|
||
if self._ctx_var.get() is not None:
|
||
raise RuntimeError(
|
||
"NamespaceLock already acquired in current coroutine context"
|
||
)
|
||
|
||
final_namespace = get_final_namespace(self._namespace, self._workspace)
|
||
ctx = get_storage_keyed_lock(
|
||
["default_key"],
|
||
namespace=final_namespace,
|
||
enable_logging=self._enable_logging,
|
||
)
|
||
|
||
# Acquire the lock first, then store context only after successful acquisition
|
||
# This prevents the ContextVar from being set if acquisition fails (e.g., due to cancellation),
|
||
# which would permanently brick the lock
|
||
result = await ctx.__aenter__()
|
||
self._ctx_var.set(ctx)
|
||
return result
|
||
|
||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||
"""Exit the current context and clean up"""
|
||
# Retrieve this coroutine's context
|
||
ctx = self._ctx_var.get()
|
||
if ctx is None:
|
||
raise RuntimeError("NamespaceLock exited without being entered")
|
||
|
||
result = await ctx.__aexit__(exc_type, exc_val, exc_tb)
|
||
# Clear this coroutine's context
|
||
self._ctx_var.set(None)
|
||
return result
|
||
|
||
|
||
def get_namespace_lock(
|
||
namespace: str, workspace: str | None = None, enable_logging: bool = False
|
||
) -> NamespaceLock:
|
||
"""Get a reusable namespace lock wrapper.
|
||
|
||
This function returns a NamespaceLock instance that can be used multiple times
|
||
safely, even in concurrent scenarios. Each use creates a fresh lock context
|
||
internally, preventing lock re-entrance errors.
|
||
|
||
Args:
|
||
namespace: The namespace to get the lock for.
|
||
workspace: Workspace identifier (may be empty string for global namespace)
|
||
enable_logging: Whether to enable lock operation logging
|
||
|
||
Returns:
|
||
NamespaceLock: A reusable lock wrapper that can be used with 'async with'
|
||
|
||
Example:
|
||
lock = get_namespace_lock("pipeline_status", workspace="space1")
|
||
|
||
# Can be used multiple times
|
||
async with lock:
|
||
await do_something()
|
||
|
||
async with lock:
|
||
await do_something_else()
|
||
"""
|
||
return NamespaceLock(namespace, workspace, enable_logging)
|
||
|
||
|
||
def finalize_share_data():
|
||
"""
|
||
Release shared resources and clean up.
|
||
|
||
This function should be called when the application is shutting down
|
||
to properly release shared resources and avoid memory leaks.
|
||
|
||
In multi-process mode, it shuts down the Manager and releases all shared objects.
|
||
In single-process mode, it simply resets the global variables.
|
||
"""
|
||
global \
|
||
_manager, \
|
||
_is_multiprocess, \
|
||
_internal_lock, \
|
||
_data_init_lock, \
|
||
_shared_dicts, \
|
||
_init_flags, \
|
||
_initialized, \
|
||
_update_flags, \
|
||
_async_locks, \
|
||
_default_workspace, \
|
||
_global_concurrency_limits, \
|
||
_lease_ns_cache, \
|
||
_queue_stats_ns_cache
|
||
|
||
# Check if already initialized
|
||
if not _initialized:
|
||
direct_log(
|
||
f"Process {os.getpid()} storage data not initialized, nothing to finalize"
|
||
)
|
||
return
|
||
|
||
direct_log(
|
||
f"Process {os.getpid()} finalizing storage data (multiprocess={_is_multiprocess})"
|
||
)
|
||
|
||
# In multi-process mode, shut down the Manager
|
||
if _is_multiprocess and _manager is not None:
|
||
try:
|
||
# Clear shared resources before shutting down Manager
|
||
if _shared_dicts is not None:
|
||
# Clear pipeline status history messages first if exists
|
||
try:
|
||
pipeline_status = _shared_dicts.get("pipeline_status", {})
|
||
if "history_messages" in pipeline_status:
|
||
pipeline_status["history_messages"].clear()
|
||
except Exception:
|
||
pass # Ignore any errors during history messages cleanup
|
||
_shared_dicts.clear()
|
||
if _init_flags is not None:
|
||
_init_flags.clear()
|
||
if _update_flags is not None:
|
||
# Clear each namespace's update flags list and Value objects
|
||
try:
|
||
for namespace in _update_flags:
|
||
flags_list = _update_flags[namespace]
|
||
if isinstance(flags_list, list):
|
||
# Clear Value objects in the list
|
||
for flag in flags_list:
|
||
if hasattr(
|
||
flag, "value"
|
||
): # Check if it's a Value object
|
||
flag.value = False
|
||
flags_list.clear()
|
||
except Exception:
|
||
pass # Ignore any errors during update flags cleanup
|
||
_update_flags.clear()
|
||
|
||
# Shut down the Manager - this will automatically clean up all shared resources
|
||
_manager.shutdown()
|
||
direct_log(f"Process {os.getpid()} Manager shutdown complete")
|
||
except Exception as e:
|
||
direct_log(
|
||
f"Process {os.getpid()} Error shutting down Manager: {e}", level="ERROR"
|
||
)
|
||
|
||
# Reset global variables
|
||
_manager = None
|
||
_initialized = None
|
||
_is_multiprocess = None
|
||
_shared_dicts = None
|
||
_init_flags = None
|
||
_internal_lock = None
|
||
_data_init_lock = None
|
||
_update_flags = None
|
||
_async_locks = None
|
||
_default_workspace = None
|
||
_global_concurrency_limits = None
|
||
_lease_ns_cache = None
|
||
_queue_stats_ns_cache = None
|
||
|
||
direct_log(f"Process {os.getpid()} storage data finalization complete")
|
||
|
||
|
||
def set_default_workspace(workspace: str | None = None):
|
||
"""
|
||
Set default workspace for namespace operations for backward compatibility.
|
||
|
||
This allows get_namespace_data(),get_namespace_lock() or initialize_pipeline_status() to
|
||
automatically use the correct workspace when called without workspace parameters,
|
||
maintaining compatibility with legacy code that doesn't pass workspace explicitly.
|
||
|
||
Args:
|
||
workspace: Workspace identifier (may be empty string for global namespace)
|
||
"""
|
||
global _default_workspace
|
||
if workspace is None:
|
||
workspace = ""
|
||
_default_workspace = workspace
|
||
direct_log(
|
||
f"Default workspace set to: '{_default_workspace}' (empty means global)",
|
||
level="DEBUG",
|
||
)
|
||
|
||
|
||
def get_default_workspace() -> str:
|
||
"""
|
||
Get default workspace for backward compatibility.
|
||
|
||
Returns:
|
||
The default workspace string. Empty string means global namespace. None means not set.
|
||
"""
|
||
global _default_workspace
|
||
return _default_workspace
|
||
|
||
|
||
def get_pipeline_status_lock(
|
||
enable_logging: bool = False, workspace: str = None
|
||
) -> NamespaceLock:
|
||
"""Return unified storage lock for pipeline status data consistency.
|
||
|
||
This function is for compatibility with legacy code only.
|
||
"""
|
||
global _default_workspace
|
||
actual_workspace = workspace if workspace else _default_workspace
|
||
return get_namespace_lock(
|
||
"pipeline_status", workspace=actual_workspace, enable_logging=enable_logging
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Cross-worker global concurrency gate (lease + heartbeat semantics)
|
||
# ---------------------------------------------------------------------------
|
||
#
|
||
# Each group's whole gate state lives under a SINGLE key of the
|
||
# workspace-less "concurrency_leases" namespace::
|
||
#
|
||
# ns[group] = {
|
||
# "leases": {lease_id: {"pid": int, "updated_at": float,
|
||
# ("suspect_since": float)}},
|
||
# "waiters": {str(pid): {"pid": int, "wait_start": float,
|
||
# "last_poll": float}},
|
||
# }
|
||
#
|
||
# Whole-value replacement only — in multiprocess mode the namespace is a
|
||
# Manager dict, so in-place mutation of a retrieved value would NOT persist
|
||
# across processes. The single-key layout is deliberate: every proxy access
|
||
# is one IPC round trip to the manager process, so an acquire attempt under
|
||
# the group's keyed lock costs exactly one read (plus at most one write)
|
||
# instead of scanning per-lease keys. Reaping, capacity counting and waiter
|
||
# ranking all run on the local copy.
|
||
#
|
||
# Self-healing: holders refresh ``updated_at`` from their 5s health-check
|
||
# heartbeat. A lease is reclaimed when its owner PID is dead (immediately)
|
||
# or when its heartbeat expired AND the suspect grace elapsed (protects
|
||
# live-but-momentarily-stalled owners from false reclamation). Long-running
|
||
# tasks are never reclaimed as long as their owner keeps renewing.
|
||
#
|
||
# Best-effort cap, not a strict provider-side invariant: the lease table is
|
||
# the admission source of truth, so the cap is exact while holders keep
|
||
# renewing. A slot is reclaimed only when its owner PID is gone or its
|
||
# heartbeat has expired beyond the suspect grace. That prevents permanent
|
||
# capacity leaks after kill -9 / OOM and similar external termination, but
|
||
# the provider may still be finishing the abandoned HTTP request until its
|
||
# own timeout/connection close. During that window, a newly admitted caller
|
||
# can overlap with the abandoned provider-side call. Long event-loop stalls
|
||
# have the same shape after heartbeat TTL + suspect grace, though normal long
|
||
# calls are protected by regular renewal. A reclaimed lease is never
|
||
# resurrected (renew_global_slots refuses to re-insert a popped lease), which
|
||
# keeps the internal gate self-healing.
|
||
|
||
|
||
def is_share_data_initialized() -> bool:
|
||
"""Return True once initialize_share_data() has run in this process."""
|
||
return bool(_initialized)
|
||
|
||
|
||
def is_global_concurrency_limited(group: Optional[str]) -> bool:
|
||
"""Synchronous, cacheable check: does ``group`` have a global limit?
|
||
|
||
Reads only the module-level read-only configuration — no IPC. Returns
|
||
False when shared data is not initialized, no limits were configured,
|
||
or the group has no positive limit.
|
||
"""
|
||
if not group or not _global_concurrency_limits:
|
||
return False
|
||
limit = _global_concurrency_limits.get(group)
|
||
return limit is not None and limit > 0
|
||
|
||
|
||
def get_global_concurrency_limit(group: Optional[str]) -> Optional[int]:
|
||
"""Return the configured global limit for ``group`` (None if unlimited)."""
|
||
if not group or not _global_concurrency_limits:
|
||
return None
|
||
return _global_concurrency_limits.get(group)
|
||
|
||
|
||
def _pid_alive(pid: int) -> bool:
|
||
"""Best-effort liveness probe; errs on the side of 'alive'."""
|
||
if pid == os.getpid():
|
||
return True
|
||
try:
|
||
os.kill(pid, 0)
|
||
except ProcessLookupError:
|
||
return False
|
||
except OSError:
|
||
# PermissionError and friends: the process exists (or we cannot
|
||
# tell) — treat as alive so we never reclaim a live owner's lease.
|
||
return True
|
||
return True
|
||
|
||
|
||
async def _get_lease_namespace() -> Dict[str, Any]:
|
||
global _lease_ns_cache
|
||
if _lease_ns_cache is None:
|
||
_lease_ns_cache = await get_namespace_data(
|
||
_CONCURRENCY_LEASE_NAMESPACE, workspace=""
|
||
)
|
||
return _lease_ns_cache
|
||
|
||
|
||
async def _get_queue_stats_namespace() -> Dict[str, Any]:
|
||
global _queue_stats_ns_cache
|
||
if _queue_stats_ns_cache is None:
|
||
_queue_stats_ns_cache = await get_namespace_data(
|
||
_QUEUE_STATS_NAMESPACE, workspace=""
|
||
)
|
||
return _queue_stats_ns_cache
|
||
|
||
|
||
def _empty_gate_state() -> Dict[str, Any]:
|
||
return {"leases": {}, "waiters": {}}
|
||
|
||
|
||
def _load_gate_state(ns: Dict[str, Any], group: str) -> Dict[str, Any]:
|
||
"""Return a local, independently mutable copy of a group's gate state.
|
||
|
||
Exactly one proxy read. The nested dicts are copied so that mutating the
|
||
result never aliases the stored value (a plain dict in single-process
|
||
mode) — callers mutate the copy and write it back whole.
|
||
"""
|
||
raw = ns.get(group)
|
||
if raw is None:
|
||
return _empty_gate_state()
|
||
state = dict(raw)
|
||
state["leases"] = {
|
||
lease_id: dict(lease)
|
||
for lease_id, lease in dict(state.get("leases") or {}).items()
|
||
}
|
||
state["waiters"] = {
|
||
pid_key: dict(waiter)
|
||
for pid_key, waiter in dict(state.get("waiters") or {}).items()
|
||
}
|
||
return state
|
||
|
||
|
||
def _reap_gate_state(state: Dict[str, Any], now: float) -> tuple[int, bool]:
|
||
"""Reclaim dead/expired leases on a local state copy (no IPC).
|
||
|
||
Returns ``(live_lease_count, changed)``. Suspect handling: a lease whose
|
||
heartbeat expired while its PID is still alive is first marked with
|
||
``suspect_since`` and reclaimed only after ``_suspect_grace`` elapses
|
||
without a renewal; a renewal (fresh ``updated_at``) clears the suspect
|
||
mark. Dead PIDs are reclaimed immediately. Suspect leases still count
|
||
toward capacity so the global limit is never exceeded.
|
||
|
||
Waiter records are reaped in the same pass: a process whose lease was
|
||
just reclaimed (timed out / died), whose PID is dead, or whose record
|
||
has not been refreshed within ``_waiter_stale_ttl`` must not keep
|
||
occupying the longest-waiter seat — a ghost favored waiter would push
|
||
every live waiter onto the deferred backoff and waste freed slots.
|
||
"""
|
||
leases: Dict[str, Any] = state["leases"]
|
||
waiters: Dict[str, Any] = state["waiters"]
|
||
live = 0
|
||
changed = False
|
||
reclaimed_pids = set()
|
||
# Liveness is constant within this synchronous pass (no await, fixed
|
||
# ``now``), so memoize the probe: a PID holding many leases of this group
|
||
# is checked once instead of once per lease. NEVER cache across calls —
|
||
# PID reuse and staleness could reclaim a live owner's lease.
|
||
alive_cache: Dict[int, bool] = {}
|
||
|
||
def _alive(pid: int) -> bool:
|
||
cached = alive_cache.get(pid)
|
||
if cached is None:
|
||
cached = _pid_alive(pid)
|
||
alive_cache[pid] = cached
|
||
return cached
|
||
|
||
for lease_id in list(leases.keys()):
|
||
lease = leases[lease_id]
|
||
pid = lease.get("pid")
|
||
updated_at = lease.get("updated_at", 0.0)
|
||
if pid is None or not _alive(pid):
|
||
leases.pop(lease_id)
|
||
changed = True
|
||
if pid is not None:
|
||
reclaimed_pids.add(pid)
|
||
continue
|
||
if now - updated_at > _heartbeat_ttl:
|
||
suspect_since = lease.get("suspect_since")
|
||
if suspect_since is None:
|
||
lease["suspect_since"] = now
|
||
changed = True
|
||
live += 1
|
||
elif now - suspect_since > _suspect_grace:
|
||
leases.pop(lease_id)
|
||
reclaimed_pids.add(pid)
|
||
changed = True
|
||
else:
|
||
live += 1
|
||
else:
|
||
if "suspect_since" in lease:
|
||
lease.pop("suspect_since", None)
|
||
changed = True
|
||
live += 1
|
||
|
||
for pid_key in list(waiters.keys()):
|
||
waiter = waiters[pid_key]
|
||
pid = waiter.get("pid")
|
||
last_poll = waiter.get("last_poll", 0.0)
|
||
if (
|
||
pid is None
|
||
or pid in reclaimed_pids
|
||
or not _alive(pid)
|
||
or now - last_poll > _waiter_stale_ttl
|
||
):
|
||
waiters.pop(pid_key)
|
||
changed = True
|
||
return live, changed
|
||
|
||
|
||
def _log_acquire_failure(group: str, error: Exception) -> None:
|
||
global _last_acquire_failure_log
|
||
now = time.time()
|
||
if now - _last_acquire_failure_log >= _ACQUIRE_FAILURE_LOG_INTERVAL:
|
||
_last_acquire_failure_log = now
|
||
direct_log(
|
||
f"Process {os.getpid()} failed to acquire global slot for group "
|
||
f"'{group}' (fail-closed, task stays queued): {error}",
|
||
level="WARNING",
|
||
)
|
||
|
||
|
||
def _is_longest_live_waiter(state: Dict[str, Any], pid: int, now: float) -> bool:
|
||
"""Is ``pid`` the longest-waiting live poller in this gate state?
|
||
|
||
Operates on a local state copy after the reap pass has dropped
|
||
dead/stale waiter records — every remaining record belongs to a live,
|
||
actively polling process.
|
||
"""
|
||
my_start = None
|
||
others_min = None
|
||
for waiter in state["waiters"].values():
|
||
wait_start = waiter.get("wait_start", now)
|
||
if waiter.get("pid") == pid:
|
||
my_start = wait_start
|
||
elif others_min is None or wait_start < others_min:
|
||
others_min = wait_start
|
||
if my_start is None:
|
||
return False
|
||
return others_min is None or my_start <= others_min
|
||
|
||
|
||
async def _acquire_global_slot(
|
||
group: str, track_wait: bool
|
||
) -> tuple[Optional[str], bool]:
|
||
"""Shared implementation for the two acquire entry points.
|
||
|
||
Returns ``(lease_id, is_priority_waiter)``. When ``track_wait`` is set,
|
||
a failed attempt registers/refreshes this process's waiter record
|
||
(``wait_start`` set once per waiting episode, ``last_poll`` refreshed on
|
||
every attempt) and reports whether this process is the longest-waiting
|
||
live poller; a successful attempt always clears the record.
|
||
|
||
IPC budget under the keyed lock: one state read, plus one write only
|
||
when something changed (reap effects, waiter registration, or a new
|
||
lease) — a plain failed attempt on an unchanged gate writes nothing.
|
||
"""
|
||
limit = get_global_concurrency_limit(group)
|
||
if limit is None or limit <= 0:
|
||
return None, False
|
||
try:
|
||
ns = await _get_lease_namespace()
|
||
async with get_storage_keyed_lock(
|
||
group, namespace=_CONCURRENCY_LEASE_NAMESPACE, enable_logging=False
|
||
):
|
||
now = time.time()
|
||
state = _load_gate_state(ns, group)
|
||
in_use, changed = _reap_gate_state(state, now)
|
||
pid = os.getpid()
|
||
pid_key = str(pid)
|
||
if in_use >= limit:
|
||
if not track_wait:
|
||
if changed:
|
||
ns[group] = state
|
||
return None, False
|
||
waiter = state["waiters"].get(pid_key) or {
|
||
"pid": pid,
|
||
"wait_start": now,
|
||
}
|
||
waiter["last_poll"] = now
|
||
state["waiters"][pid_key] = waiter
|
||
ns[group] = state
|
||
return None, _is_longest_live_waiter(state, pid, now)
|
||
lease_id = uuid.uuid4().hex
|
||
state["leases"][lease_id] = {"pid": pid, "updated_at": now}
|
||
# Got a slot: this process is no longer waiting. Resetting here
|
||
# (rather than keeping seniority) is what de-prioritizes a
|
||
# backlog-heavy process after each win, yielding approximate
|
||
# round-robin across processes under sustained contention.
|
||
state["waiters"].pop(pid_key, None)
|
||
ns[group] = state
|
||
return lease_id, True
|
||
except Exception as e:
|
||
_log_acquire_failure(group, e)
|
||
return None, False
|
||
|
||
|
||
async def try_acquire_global_slot(group: str) -> Optional[str]:
|
||
"""Try to claim one global concurrency slot for ``group`` (non-blocking).
|
||
|
||
Returns a lease id on success, or None when the group is at capacity.
|
||
Any shared-storage error is fail-closed: returns None (with a
|
||
rate-limited warning) so the caller keeps the task queued and retries —
|
||
capacity is never exceeded due to infrastructure errors.
|
||
|
||
This plain variant never registers waiter records — use
|
||
:func:`try_acquire_global_slot_tracked` from polling loops that want
|
||
longest-waiter fairness.
|
||
"""
|
||
lease_id, _ = await _acquire_global_slot(group, track_wait=False)
|
||
return lease_id
|
||
|
||
|
||
async def try_acquire_global_slot_tracked(group: str) -> tuple[Optional[str], bool]:
|
||
"""Acquire variant for polling loops: ``(lease_id, is_priority_waiter)``.
|
||
|
||
On failure the caller's waiter record is registered/refreshed and the
|
||
second element reports whether this process is currently the
|
||
longest-waiting live poller of the group. Pollers should keep the
|
||
fastest poll interval when favored and back off (bounded) otherwise —
|
||
a soft FIFO across worker processes with no hard gate: any poller that
|
||
finds a free slot still takes it, so a sleeping favored waiter can
|
||
never leave capacity idle indefinitely. Fail-closed errors report
|
||
``(None, False)``.
|
||
"""
|
||
return await _acquire_global_slot(group, track_wait=True)
|
||
|
||
|
||
async def clear_slot_waiter(group: str) -> None:
|
||
"""Drop this process's waiter record for ``group`` (idempotent).
|
||
|
||
Called when a wrapper shuts down so a no-longer-polling process never
|
||
lingers in the longest-waiter seat; the stale TTL and the reap pass
|
||
cover crashes where this cleanup never runs.
|
||
"""
|
||
if not _initialized:
|
||
return
|
||
ns = await _get_lease_namespace()
|
||
async with get_storage_keyed_lock(
|
||
group, namespace=_CONCURRENCY_LEASE_NAMESPACE, enable_logging=False
|
||
):
|
||
state = _load_gate_state(ns, group)
|
||
if state["waiters"].pop(str(os.getpid()), None) is not None:
|
||
ns[group] = state
|
||
|
||
|
||
async def global_slot_waiters(group: str) -> List[Dict[str, Any]]:
|
||
"""Snapshot of processes actively polling for a slot of ``group``.
|
||
|
||
Returns ``[{"pid": ..., "waited": seconds}, ...]`` sorted by descending
|
||
wait time; stale records (not refreshed within the waiter TTL) are
|
||
skipped. Read-only and lock-free — intended for observability.
|
||
"""
|
||
if not _initialized:
|
||
return []
|
||
ns = await _get_lease_namespace()
|
||
now = time.time()
|
||
state = _load_gate_state(ns, group)
|
||
waiters = []
|
||
for waiter in state["waiters"].values():
|
||
if now - waiter.get("last_poll", 0.0) > _waiter_stale_ttl:
|
||
continue
|
||
waiters.append(
|
||
{
|
||
"pid": waiter.get("pid"),
|
||
"waited": max(0.0, now - waiter.get("wait_start", now)),
|
||
}
|
||
)
|
||
return sorted(waiters, key=lambda w: -w["waited"])
|
||
|
||
|
||
async def release_global_slot(group: str, lease_id: str) -> None:
|
||
"""Release a previously acquired global slot (idempotent).
|
||
|
||
Raises on shared-storage errors — callers that must never propagate
|
||
(e.g. worker ``finally`` blocks) wrap this and queue the lease for a
|
||
later retry; the heartbeat TTL guarantees eventual reclamation anyway.
|
||
"""
|
||
ns = await _get_lease_namespace()
|
||
async with get_storage_keyed_lock(
|
||
group, namespace=_CONCURRENCY_LEASE_NAMESPACE, enable_logging=False
|
||
):
|
||
state = _load_gate_state(ns, group)
|
||
if state["leases"].pop(lease_id, None) is not None:
|
||
ns[group] = state
|
||
|
||
|
||
async def renew_global_slots(group: str, lease_ids) -> None:
|
||
"""Refresh the heartbeat of this process's held leases for ``group``.
|
||
|
||
A renewal rewrites the lease whole (clearing any ``suspect_since``
|
||
mark). Leases that have already been reclaimed are NOT resurrected —
|
||
re-inserting could exceed the configured limit; the suspect grace
|
||
exists precisely to make false reclamation unlikely.
|
||
"""
|
||
lease_ids = list(lease_ids)
|
||
if not lease_ids:
|
||
return
|
||
ns = await _get_lease_namespace()
|
||
async with get_storage_keyed_lock(
|
||
group, namespace=_CONCURRENCY_LEASE_NAMESPACE, enable_logging=False
|
||
):
|
||
now = time.time()
|
||
state = _load_gate_state(ns, group)
|
||
changed = False
|
||
for lease_id in lease_ids:
|
||
if lease_id in state["leases"]:
|
||
state["leases"][lease_id] = {"pid": os.getpid(), "updated_at": now}
|
||
changed = True
|
||
if changed:
|
||
ns[group] = state
|
||
|
||
|
||
async def reconcile_global_slots(group: str) -> int:
|
||
"""Run the lease reaper for ``group``; return surviving lease count."""
|
||
ns = await _get_lease_namespace()
|
||
async with get_storage_keyed_lock(
|
||
group, namespace=_CONCURRENCY_LEASE_NAMESPACE, enable_logging=False
|
||
):
|
||
state = _load_gate_state(ns, group)
|
||
live, changed = _reap_gate_state(state, time.time())
|
||
if changed:
|
||
ns[group] = state
|
||
return live
|
||
|
||
|
||
async def global_concurrency_in_use(group: str) -> int:
|
||
"""Approximate count of currently held global slots for ``group``.
|
||
|
||
Lock-free single read — intended for observability.
|
||
"""
|
||
ns = await _get_lease_namespace()
|
||
return len(_load_gate_state(ns, group)["leases"])
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Cross-worker queue stats (best-effort, debounced snapshots)
|
||
# ---------------------------------------------------------------------------
|
||
#
|
||
# Each worker process publishes per-queue snapshots under
|
||
# ``f"{queue_name}{KEY_SEP}{pid}"`` in the workspace-less "queue_stats"
|
||
# namespace (whole-value replacement). The local closure counters remain
|
||
# the source of truth; the shared area only needs to be "fresh enough"
|
||
# (event-triggered publishes debounced by the caller + 5s heartbeat flush).
|
||
|
||
# Flat counter fields summed across workers during aggregation. These are
|
||
# the existing get_queue_stats() fields (schema compatibility for /health
|
||
# and the webui) plus the new global_slot_waits / physical_queued counters.
|
||
QUEUE_STATS_SUM_FIELDS = (
|
||
"queued",
|
||
"running",
|
||
"in_flight",
|
||
"worker_count",
|
||
"submitted_total",
|
||
"completed_total",
|
||
"failed_total",
|
||
"cancelled_total",
|
||
"rejected_total",
|
||
"global_slot_waits",
|
||
"physical_queued",
|
||
)
|
||
|
||
|
||
async def publish_queue_stats(queue_name: str, snapshot: Dict[str, Any]) -> None:
|
||
"""Publish this process's snapshot for ``queue_name`` (whole replacement).
|
||
|
||
The snapshot must carry ``pid`` and ``updated_at`` (wall-clock time) so
|
||
aggregation can reap stale entries. Best-effort by contract: callers
|
||
must tolerate exceptions.
|
||
"""
|
||
if not _initialized:
|
||
return
|
||
ns = await _get_queue_stats_namespace()
|
||
ns[f"{queue_name}{KEY_SEP}{os.getpid()}"] = dict(snapshot)
|
||
|
||
|
||
async def unpublish_queue_stats(queue_name: str) -> None:
|
||
"""Remove this process's snapshot for ``queue_name`` (idempotent)."""
|
||
if not _initialized:
|
||
return
|
||
ns = await _get_queue_stats_namespace()
|
||
ns.pop(f"{queue_name}{KEY_SEP}{os.getpid()}", None)
|
||
|
||
|
||
async def aggregate_queue_stats(queue_name: str) -> Dict[str, Any]:
|
||
"""Aggregate all workers' published snapshots for ``queue_name``.
|
||
|
||
Sums the flat counter fields across live snapshots and returns them
|
||
together with ``reporting_workers`` and the raw ``per_worker`` map.
|
||
Entries owned by dead PIDs or older than the stale TTL are reaped —
|
||
re-checked under the internal lock against the previously observed
|
||
``updated_at`` so a snapshot republished concurrently is never deleted.
|
||
"""
|
||
ns = await _get_queue_stats_namespace()
|
||
now = time.time()
|
||
prefix = f"{queue_name}{KEY_SEP}"
|
||
per_worker: Dict[str, Dict[str, Any]] = {}
|
||
stale: List[tuple] = []
|
||
for key in [k for k in ns.keys() if k.startswith(prefix)]:
|
||
raw = ns.get(key)
|
||
if raw is None:
|
||
continue
|
||
snap = dict(raw)
|
||
pid = snap.get("pid")
|
||
updated_at = snap.get("updated_at", 0.0)
|
||
if (pid is not None and not _pid_alive(pid)) or (
|
||
now - updated_at > _queue_stats_stale_ttl
|
||
):
|
||
stale.append((key, updated_at))
|
||
continue
|
||
per_worker[str(pid)] = snap
|
||
|
||
if stale:
|
||
async with get_internal_lock():
|
||
for key, seen_updated_at in stale:
|
||
current = ns.get(key)
|
||
if current is None:
|
||
continue
|
||
if dict(current).get("updated_at", 0.0) != seen_updated_at:
|
||
continue # republished since we looked — keep it
|
||
ns.pop(key, None)
|
||
|
||
aggregated: Dict[str, Any] = {
|
||
field: sum(int(snap.get(field, 0) or 0) for snap in per_worker.values())
|
||
for field in QUEUE_STATS_SUM_FIELDS
|
||
}
|
||
aggregated["reporting_workers"] = len(per_worker)
|
||
aggregated["per_worker"] = per_worker
|
||
return aggregated
|