266 lines
8.7 KiB
Python
Raw Normal View History

import asyncio
import time
import hashlib
import uuid
from typing import Dict, Optional, Set
from dataclasses import dataclass, field
from contextlib import asynccontextmanager
import logging
logger = logging.getLogger(__name__)
@dataclass
class LockInfo:
token: str
owner: str
user_id: int
acquired_at: float = field(default_factory=time.time)
timeout: float = 30.0
extend_count: int = 0
@property
def is_expired(self) -> bool:
return time.time() - self.acquired_at > self.timeout
@dataclass
class LockEntry:
lock: asyncio.Lock
info: Optional[LockInfo] = None
waiters: int = 0
class LockManager:
LOCK_PATTERNS = {
"file_upload": "lock:user:{user_id}:upload:{path_hash}",
"quota_update": "lock:user:{user_id}:quota",
"folder_create": "lock:user:{user_id}:folder:{parent_id}:create",
"file_create": "lock:user:{user_id}:file:{parent_id}:create:{name_hash}",
"webdav_lock": "lock:webdav:{path_hash}",
"invoice_gen": "lock:billing:invoice:{user_id}",
"user_update": "lock:user:{user_id}:update",
}
def __init__(self, default_timeout: float = 30.0, cleanup_interval: float = 60.0):
self.default_timeout = default_timeout
self.cleanup_interval = cleanup_interval
self.locks: Dict[str, LockEntry] = {}
self._global_lock = asyncio.Lock()
self._cleanup_task: Optional[asyncio.Task] = None
self._running = False
async def start(self):
self._running = True
self._cleanup_task = asyncio.create_task(self._background_cleanup())
logger.info("LockManager started")
async def stop(self):
self._running = False
if self._cleanup_task:
self._cleanup_task.cancel()
try:
await self._cleanup_task
except asyncio.CancelledError:
pass
logger.info("LockManager stopped")
def build_lock_key(self, lock_type: str, **kwargs) -> str:
template = self.LOCK_PATTERNS.get(lock_type)
if not template:
return f"lock:custom:{lock_type}:{':'.join(str(v) for v in kwargs.values())}"
for key, value in kwargs.items():
if "hash" in key and not isinstance(value, str):
kwargs[key] = hashlib.md5(str(value).encode()).hexdigest()[:16]
return template.format(**kwargs)
async def _get_or_create_lock(self, resource: str) -> LockEntry:
async with self._global_lock:
if resource not in self.locks:
self.locks[resource] = LockEntry(lock=asyncio.Lock())
return self.locks[resource]
@asynccontextmanager
async def acquire(self, resource: str, timeout: Optional[float] = None,
owner: str = "", user_id: int = 0):
if timeout is None:
timeout = self.default_timeout
entry = await self._get_or_create_lock(resource)
async with self._global_lock:
entry.waiters += 1
try:
try:
await asyncio.wait_for(entry.lock.acquire(), timeout=timeout)
except asyncio.TimeoutError:
async with self._global_lock:
entry.waiters -= 1
raise TimeoutError(f"Failed to acquire lock for {resource} within {timeout}s")
token = str(uuid.uuid4())
entry.info = LockInfo(
token=token,
owner=owner,
user_id=user_id,
timeout=timeout
)
logger.debug(f"Lock acquired: {resource} by {owner}")
try:
yield token
finally:
entry.lock.release()
entry.info = None
async with self._global_lock:
entry.waiters -= 1
logger.debug(f"Lock released: {resource}")
except Exception:
async with self._global_lock:
if entry.waiters > 0:
entry.waiters -= 1
raise
async def try_acquire(self, resource: str, owner: str = "",
user_id: int = 0, timeout: float = 0) -> Optional[str]:
entry = await self._get_or_create_lock(resource)
if entry.lock.locked():
if entry.info and entry.info.is_expired:
pass
elif timeout > 0:
try:
await asyncio.wait_for(entry.lock.acquire(), timeout=timeout)
except asyncio.TimeoutError:
return None
else:
return None
else:
await entry.lock.acquire()
token = str(uuid.uuid4())
entry.info = LockInfo(
token=token,
owner=owner,
user_id=user_id,
timeout=self.default_timeout
)
return token
async def release(self, resource: str, token: str) -> bool:
async with self._global_lock:
if resource not in self.locks:
return False
entry = self.locks[resource]
if not entry.info or entry.info.token != token:
return False
entry.lock.release()
entry.info = None
logger.debug(f"Lock released via token: {resource}")
return True
async def extend(self, resource: str, token: str, extension: float = 30.0) -> bool:
async with self._global_lock:
if resource not in self.locks:
return False
entry = self.locks[resource]
if not entry.info or entry.info.token != token:
return False
entry.info.acquired_at = time.time()
entry.info.timeout = extension
entry.info.extend_count += 1
return True
async def get_lock_info(self, resource: str) -> Optional[LockInfo]:
async with self._global_lock:
if resource in self.locks:
return self.locks[resource].info
return None
async def is_locked(self, resource: str) -> bool:
async with self._global_lock:
if resource in self.locks:
entry = self.locks[resource]
if entry.lock.locked():
if entry.info and not entry.info.is_expired:
return True
return False
async def force_release(self, resource: str, user_id: int) -> bool:
async with self._global_lock:
if resource not in self.locks:
return False
entry = self.locks[resource]
if not entry.info:
return False
if entry.info.user_id != user_id:
return False
if entry.lock.locked():
entry.lock.release()
entry.info = None
return True
async def _background_cleanup(self):
while self._running:
try:
await asyncio.sleep(self.cleanup_interval)
await self._cleanup_expired()
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"Error in lock cleanup: {e}")
async def _cleanup_expired(self):
async with self._global_lock:
expired = []
for resource, entry in self.locks.items():
if entry.info and entry.info.is_expired and entry.waiters == 0:
expired.append(resource)
for resource in expired:
async with self._global_lock:
if resource in self.locks:
entry = self.locks[resource]
if entry.lock.locked() and entry.info and entry.info.is_expired:
entry.lock.release()
entry.info = None
logger.debug(f"Cleaned up expired lock: {resource}")
async def get_stats(self) -> Dict:
async with self._global_lock:
total_locks = len(self.locks)
active_locks = sum(1 for e in self.locks.values() if e.lock.locked())
waiting = sum(e.waiters for e in self.locks.values())
return {
"total_locks": total_locks,
"active_locks": active_locks,
"waiting_requests": waiting,
}
_lock_manager: Optional[LockManager] = None
async def init_lock_manager(default_timeout: float = 30.0) -> LockManager:
global _lock_manager
_lock_manager = LockManager(default_timeout=default_timeout)
await _lock_manager.start()
return _lock_manager
async def shutdown_lock_manager():
global _lock_manager
if _lock_manager:
await _lock_manager.stop()
_lock_manager = None
def get_lock_manager() -> LockManager:
if not _lock_manager:
raise RuntimeError("Lock manager not initialized")
return _lock_manager