chore: add token revocation, caching, concurrency, and WebDAV lock modules
This commit is contained in:
@@ -0,0 +1,9 @@
|
||||
from .locks import LockManager, get_lock_manager
|
||||
from .atomic import AtomicOperations, get_atomic_ops
|
||||
from .webdav_locks import PersistentWebDAVLocks, get_webdav_locks
|
||||
|
||||
__all__ = [
|
||||
"LockManager", "get_lock_manager",
|
||||
"AtomicOperations", "get_atomic_ops",
|
||||
"PersistentWebDAVLocks", "get_webdav_locks",
|
||||
]
|
||||
@@ -0,0 +1,151 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
from typing import Optional, Tuple, Any
|
||||
from dataclasses import dataclass
|
||||
import logging
|
||||
|
||||
from .locks import get_lock_manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class QuotaCheckResult:
|
||||
allowed: bool
|
||||
current_usage: int
|
||||
quota: int
|
||||
requested: int
|
||||
remaining: int
|
||||
|
||||
|
||||
class AtomicOperations:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
async def atomic_quota_check_and_update(
|
||||
self,
|
||||
user,
|
||||
delta: int,
|
||||
save_callback=None
|
||||
) -> QuotaCheckResult:
|
||||
lock_manager = get_lock_manager()
|
||||
lock_key = lock_manager.build_lock_key("quota_update", user_id=user.id)
|
||||
|
||||
async with lock_manager.acquire(lock_key, timeout=10.0, owner="quota_update", user_id=user.id):
|
||||
current = user.used_storage_bytes
|
||||
quota = user.storage_quota_bytes
|
||||
new_usage = current + delta
|
||||
|
||||
result = QuotaCheckResult(
|
||||
allowed=new_usage <= quota,
|
||||
current_usage=current,
|
||||
quota=quota,
|
||||
requested=delta,
|
||||
remaining=max(0, quota - current)
|
||||
)
|
||||
|
||||
if result.allowed and save_callback:
|
||||
user.used_storage_bytes = new_usage
|
||||
await save_callback(user)
|
||||
logger.debug(f"Quota updated for user {user.id}: {current} -> {new_usage}")
|
||||
|
||||
return result
|
||||
|
||||
async def atomic_file_create(
|
||||
self,
|
||||
user,
|
||||
parent_id: Optional[int],
|
||||
name: str,
|
||||
check_exists_callback,
|
||||
create_callback,
|
||||
):
|
||||
lock_manager = get_lock_manager()
|
||||
name_hash = hashlib.md5(name.encode()).hexdigest()[:16]
|
||||
lock_key = lock_manager.build_lock_key(
|
||||
"file_create",
|
||||
user_id=user.id,
|
||||
parent_id=parent_id or 0,
|
||||
name_hash=name_hash
|
||||
)
|
||||
|
||||
async with lock_manager.acquire(lock_key, timeout=10.0, owner="file_create", user_id=user.id):
|
||||
existing = await check_exists_callback()
|
||||
if existing:
|
||||
raise FileExistsError(f"File '{name}' already exists in this location")
|
||||
|
||||
result = await create_callback()
|
||||
logger.debug(f"File created atomically: {name} for user {user.id}")
|
||||
return result
|
||||
|
||||
async def atomic_folder_create(
|
||||
self,
|
||||
user,
|
||||
parent_id: Optional[int],
|
||||
name: str,
|
||||
check_exists_callback,
|
||||
create_callback,
|
||||
):
|
||||
lock_manager = get_lock_manager()
|
||||
lock_key = lock_manager.build_lock_key(
|
||||
"folder_create",
|
||||
user_id=user.id,
|
||||
parent_id=parent_id or 0
|
||||
)
|
||||
|
||||
async with lock_manager.acquire(lock_key, timeout=10.0, owner="folder_create", user_id=user.id):
|
||||
existing = await check_exists_callback()
|
||||
if existing:
|
||||
raise FileExistsError(f"Folder '{name}' already exists in this location")
|
||||
|
||||
result = await create_callback()
|
||||
logger.debug(f"Folder created atomically: {name} for user {user.id}")
|
||||
return result
|
||||
|
||||
async def atomic_file_update(
|
||||
self,
|
||||
user,
|
||||
file_id: int,
|
||||
update_callback,
|
||||
):
|
||||
lock_manager = get_lock_manager()
|
||||
lock_key = f"lock:file:{file_id}:update"
|
||||
|
||||
async with lock_manager.acquire(lock_key, timeout=30.0, owner="file_update", user_id=user.id):
|
||||
result = await update_callback()
|
||||
return result
|
||||
|
||||
async def atomic_batch_operation(
|
||||
self,
|
||||
user,
|
||||
operation_id: str,
|
||||
items: list,
|
||||
operation_callback,
|
||||
):
|
||||
lock_manager = get_lock_manager()
|
||||
lock_key = f"lock:batch:{user.id}:{operation_id}"
|
||||
|
||||
async with lock_manager.acquire(lock_key, timeout=60.0, owner="batch_op", user_id=user.id):
|
||||
results = []
|
||||
errors = []
|
||||
for item in items:
|
||||
try:
|
||||
result = await operation_callback(item)
|
||||
results.append(result)
|
||||
except Exception as e:
|
||||
errors.append({"item": item, "error": str(e)})
|
||||
return {"results": results, "errors": errors}
|
||||
|
||||
|
||||
_atomic_ops: Optional[AtomicOperations] = None
|
||||
|
||||
|
||||
def init_atomic_ops() -> AtomicOperations:
|
||||
global _atomic_ops
|
||||
_atomic_ops = AtomicOperations()
|
||||
return _atomic_ops
|
||||
|
||||
|
||||
def get_atomic_ops() -> AtomicOperations:
|
||||
if not _atomic_ops:
|
||||
return AtomicOperations()
|
||||
return _atomic_ops
|
||||
@@ -0,0 +1,265 @@
|
||||
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
|
||||
@@ -0,0 +1,319 @@
|
||||
import asyncio
|
||||
import time
|
||||
import hashlib
|
||||
import uuid
|
||||
from typing import Dict, Optional
|
||||
from dataclasses import dataclass, field, asdict
|
||||
import json
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WebDAVLockInfo:
|
||||
token: str
|
||||
path: str
|
||||
path_hash: str
|
||||
owner: str
|
||||
user_id: int
|
||||
scope: str = "exclusive"
|
||||
depth: str = "0"
|
||||
timeout: int = 3600
|
||||
created_at: float = field(default_factory=time.time)
|
||||
|
||||
@property
|
||||
def is_expired(self) -> bool:
|
||||
return time.time() - self.created_at > self.timeout
|
||||
|
||||
@property
|
||||
def remaining_seconds(self) -> int:
|
||||
remaining = self.timeout - (time.time() - self.created_at)
|
||||
return max(0, int(remaining))
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return asdict(self)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> "WebDAVLockInfo":
|
||||
return cls(**data)
|
||||
|
||||
|
||||
class PersistentWebDAVLocks:
|
||||
def __init__(self, db_manager=None):
|
||||
self.db_manager = db_manager
|
||||
self.locks: Dict[str, WebDAVLockInfo] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
self._cleanup_task: Optional[asyncio.Task] = None
|
||||
self._running = False
|
||||
self._persistence_enabled = False
|
||||
|
||||
async def start(self, db_manager=None):
|
||||
if db_manager:
|
||||
self.db_manager = db_manager
|
||||
self._persistence_enabled = True
|
||||
await self._load_locks_from_db()
|
||||
self._running = True
|
||||
self._cleanup_task = asyncio.create_task(self._background_cleanup())
|
||||
logger.info("PersistentWebDAVLocks 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
|
||||
if self._persistence_enabled:
|
||||
await self._save_all_locks()
|
||||
logger.info("PersistentWebDAVLocks stopped")
|
||||
|
||||
def _hash_path(self, path: str) -> str:
|
||||
return hashlib.sha256(path.encode()).hexdigest()[:32]
|
||||
|
||||
async def acquire_lock(
|
||||
self,
|
||||
path: str,
|
||||
owner: str,
|
||||
user_id: int,
|
||||
timeout: int = 3600,
|
||||
scope: str = "exclusive",
|
||||
depth: str = "0"
|
||||
) -> Optional[str]:
|
||||
path = path.strip("/")
|
||||
path_hash = self._hash_path(path)
|
||||
|
||||
async with self._lock:
|
||||
if path_hash in self.locks:
|
||||
existing = self.locks[path_hash]
|
||||
if not existing.is_expired:
|
||||
if existing.user_id == user_id:
|
||||
existing.created_at = time.time()
|
||||
existing.timeout = timeout
|
||||
await self._persist_lock(existing)
|
||||
return existing.token
|
||||
return None
|
||||
else:
|
||||
del self.locks[path_hash]
|
||||
await self._delete_lock_from_db(path_hash)
|
||||
|
||||
token = f"opaquelocktoken:{uuid.uuid4()}"
|
||||
lock_info = WebDAVLockInfo(
|
||||
token=token,
|
||||
path=path,
|
||||
path_hash=path_hash,
|
||||
owner=owner,
|
||||
user_id=user_id,
|
||||
scope=scope,
|
||||
depth=depth,
|
||||
timeout=timeout
|
||||
)
|
||||
self.locks[path_hash] = lock_info
|
||||
await self._persist_lock(lock_info)
|
||||
logger.debug(f"WebDAV lock acquired: {path} by {owner}")
|
||||
return token
|
||||
|
||||
async def refresh_lock(self, path: str, token: str, timeout: int = 3600) -> bool:
|
||||
path = path.strip("/")
|
||||
path_hash = self._hash_path(path)
|
||||
|
||||
async with self._lock:
|
||||
if path_hash not in self.locks:
|
||||
return False
|
||||
lock_info = self.locks[path_hash]
|
||||
if lock_info.token != token:
|
||||
return False
|
||||
lock_info.created_at = time.time()
|
||||
lock_info.timeout = timeout
|
||||
await self._persist_lock(lock_info)
|
||||
return True
|
||||
|
||||
async def release_lock(self, path: str, token: str) -> bool:
|
||||
path = path.strip("/")
|
||||
path_hash = self._hash_path(path)
|
||||
|
||||
async with self._lock:
|
||||
if path_hash not in self.locks:
|
||||
return False
|
||||
lock_info = self.locks[path_hash]
|
||||
if lock_info.token != token:
|
||||
return False
|
||||
del self.locks[path_hash]
|
||||
await self._delete_lock_from_db(path_hash)
|
||||
logger.debug(f"WebDAV lock released: {path}")
|
||||
return True
|
||||
|
||||
async def check_lock(self, path: str) -> Optional[WebDAVLockInfo]:
|
||||
path = path.strip("/")
|
||||
path_hash = self._hash_path(path)
|
||||
|
||||
async with self._lock:
|
||||
if path_hash in self.locks:
|
||||
lock_info = self.locks[path_hash]
|
||||
if lock_info.is_expired:
|
||||
del self.locks[path_hash]
|
||||
await self._delete_lock_from_db(path_hash)
|
||||
return None
|
||||
return lock_info
|
||||
return None
|
||||
|
||||
async def get_lock_by_token(self, token: str) -> Optional[WebDAVLockInfo]:
|
||||
async with self._lock:
|
||||
for lock_info in self.locks.values():
|
||||
if lock_info.token == token and not lock_info.is_expired:
|
||||
return lock_info
|
||||
return None
|
||||
|
||||
async def is_locked(self, path: str, user_id: Optional[int] = None) -> bool:
|
||||
lock_info = await self.check_lock(path)
|
||||
if not lock_info:
|
||||
return False
|
||||
if user_id is not None and lock_info.user_id == user_id:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def force_unlock(self, path: str, user_id: int) -> bool:
|
||||
path = path.strip("/")
|
||||
path_hash = self._hash_path(path)
|
||||
|
||||
async with self._lock:
|
||||
if path_hash not in self.locks:
|
||||
return False
|
||||
lock_info = self.locks[path_hash]
|
||||
if lock_info.user_id != user_id:
|
||||
return False
|
||||
del self.locks[path_hash]
|
||||
await self._delete_lock_from_db(path_hash)
|
||||
return True
|
||||
|
||||
async def get_user_locks(self, user_id: int) -> list:
|
||||
async with self._lock:
|
||||
return [
|
||||
lock_info for lock_info in self.locks.values()
|
||||
if lock_info.user_id == user_id and not lock_info.is_expired
|
||||
]
|
||||
|
||||
async def _persist_lock(self, lock_info: WebDAVLockInfo):
|
||||
if not self._persistence_enabled or not self.db_manager:
|
||||
return
|
||||
try:
|
||||
async with self.db_manager.get_master_connection() as conn:
|
||||
await conn.execute("""
|
||||
INSERT OR REPLACE INTO webdav_locks
|
||||
(path_hash, path, token, owner, user_id, scope, timeout, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
lock_info.path_hash,
|
||||
lock_info.path,
|
||||
lock_info.token,
|
||||
lock_info.owner,
|
||||
lock_info.user_id,
|
||||
lock_info.scope,
|
||||
lock_info.timeout,
|
||||
lock_info.created_at
|
||||
))
|
||||
await conn.commit()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to persist WebDAV lock: {e}")
|
||||
|
||||
async def _delete_lock_from_db(self, path_hash: str):
|
||||
if not self._persistence_enabled or not self.db_manager:
|
||||
return
|
||||
try:
|
||||
async with self.db_manager.get_master_connection() as conn:
|
||||
await conn.execute(
|
||||
"DELETE FROM webdav_locks WHERE path_hash = ?",
|
||||
(path_hash,)
|
||||
)
|
||||
await conn.commit()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete WebDAV lock from db: {e}")
|
||||
|
||||
async def _load_locks_from_db(self):
|
||||
if not self.db_manager:
|
||||
return
|
||||
try:
|
||||
async with self.db_manager.get_master_connection() as conn:
|
||||
cursor = await conn.execute("SELECT * FROM webdav_locks")
|
||||
rows = await cursor.fetchall()
|
||||
for row in rows:
|
||||
lock_info = WebDAVLockInfo(
|
||||
token=row[3],
|
||||
path=row[2],
|
||||
path_hash=row[1],
|
||||
owner=row[4],
|
||||
user_id=row[5],
|
||||
scope=row[6],
|
||||
timeout=row[7],
|
||||
created_at=row[8]
|
||||
)
|
||||
if not lock_info.is_expired:
|
||||
self.locks[lock_info.path_hash] = lock_info
|
||||
else:
|
||||
await self._delete_lock_from_db(lock_info.path_hash)
|
||||
logger.info(f"Loaded {len(self.locks)} WebDAV locks from database")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load WebDAV locks: {e}")
|
||||
|
||||
async def _save_all_locks(self):
|
||||
if not self._persistence_enabled:
|
||||
return
|
||||
for lock_info in list(self.locks.values()):
|
||||
if not lock_info.is_expired:
|
||||
await self._persist_lock(lock_info)
|
||||
|
||||
async def _background_cleanup(self):
|
||||
while self._running:
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
await self._cleanup_expired()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error in WebDAV lock cleanup: {e}")
|
||||
|
||||
async def _cleanup_expired(self):
|
||||
async with self._lock:
|
||||
expired = [
|
||||
path_hash for path_hash, lock_info in self.locks.items()
|
||||
if lock_info.is_expired
|
||||
]
|
||||
for path_hash in expired:
|
||||
del self.locks[path_hash]
|
||||
await self._delete_lock_from_db(path_hash)
|
||||
if expired:
|
||||
logger.debug(f"Cleaned up {len(expired)} expired WebDAV locks")
|
||||
|
||||
async def get_stats(self) -> dict:
|
||||
async with self._lock:
|
||||
total = len(self.locks)
|
||||
active = sum(1 for l in self.locks.values() if not l.is_expired)
|
||||
return {
|
||||
"total_locks": total,
|
||||
"active_locks": active,
|
||||
"expired_locks": total - active,
|
||||
}
|
||||
|
||||
|
||||
_webdav_locks: Optional[PersistentWebDAVLocks] = None
|
||||
|
||||
|
||||
async def init_webdav_locks(db_manager=None) -> PersistentWebDAVLocks:
|
||||
global _webdav_locks
|
||||
_webdav_locks = PersistentWebDAVLocks()
|
||||
await _webdav_locks.start(db_manager)
|
||||
return _webdav_locks
|
||||
|
||||
|
||||
async def shutdown_webdav_locks():
|
||||
global _webdav_locks
|
||||
if _webdav_locks:
|
||||
await _webdav_locks.stop()
|
||||
_webdav_locks = None
|
||||
|
||||
|
||||
def get_webdav_locks() -> PersistentWebDAVLocks:
|
||||
if not _webdav_locks:
|
||||
raise RuntimeError("WebDAV locks not initialized")
|
||||
return _webdav_locks
|
||||
Reference in New Issue
Block a user