feat: add enterprise write buffer and data access layer with batched activity and usage record flushing
This commit is contained in:
@@ -0,0 +1,13 @@
|
||||
from .dal import DataAccessLayer, get_dal, init_dal, shutdown_dal
|
||||
from .write_buffer import WriteBuffer, get_write_buffer, init_write_buffer, shutdown_write_buffer
|
||||
|
||||
__all__ = [
|
||||
"DataAccessLayer",
|
||||
"get_dal",
|
||||
"init_dal",
|
||||
"shutdown_dal",
|
||||
"WriteBuffer",
|
||||
"get_write_buffer",
|
||||
"init_write_buffer",
|
||||
"shutdown_write_buffer",
|
||||
]
|
||||
@@ -0,0 +1,335 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import time
|
||||
from typing import Dict, List, Any, Optional, Tuple, Union
|
||||
from dataclasses import dataclass, field
|
||||
from collections import OrderedDict
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheEntry:
|
||||
value: Any
|
||||
created_at: float = field(default_factory=time.time)
|
||||
ttl: float = 300.0
|
||||
access_count: int = 0
|
||||
|
||||
@property
|
||||
def is_expired(self) -> bool:
|
||||
return time.time() - self.created_at > self.ttl
|
||||
|
||||
|
||||
class DALCache:
|
||||
def __init__(self, maxsize: int = 50000):
|
||||
self.maxsize = maxsize
|
||||
self._data: OrderedDict[str, CacheEntry] = OrderedDict()
|
||||
self._lock = asyncio.Lock()
|
||||
self._stats = {"hits": 0, "misses": 0, "evictions": 0}
|
||||
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
async with self._lock:
|
||||
if key in self._data:
|
||||
entry = self._data[key]
|
||||
if entry.is_expired:
|
||||
del self._data[key]
|
||||
self._stats["misses"] += 1
|
||||
return None
|
||||
entry.access_count += 1
|
||||
self._data.move_to_end(key)
|
||||
self._stats["hits"] += 1
|
||||
return entry.value
|
||||
self._stats["misses"] += 1
|
||||
return None
|
||||
|
||||
async def set(self, key: str, value: Any, ttl: float = 300.0):
|
||||
async with self._lock:
|
||||
if len(self._data) >= self.maxsize:
|
||||
oldest_key = next(iter(self._data))
|
||||
del self._data[oldest_key]
|
||||
self._stats["evictions"] += 1
|
||||
self._data[key] = CacheEntry(value=value, ttl=ttl)
|
||||
|
||||
async def delete(self, key: str):
|
||||
async with self._lock:
|
||||
self._data.pop(key, None)
|
||||
|
||||
async def invalidate_prefix(self, prefix: str) -> int:
|
||||
async with self._lock:
|
||||
keys_to_delete = [k for k in self._data.keys() if k.startswith(prefix)]
|
||||
for k in keys_to_delete:
|
||||
del self._data[k]
|
||||
return len(keys_to_delete)
|
||||
|
||||
async def clear(self):
|
||||
async with self._lock:
|
||||
self._data.clear()
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
total = self._stats["hits"] + self._stats["misses"]
|
||||
hit_rate = (self._stats["hits"] / total * 100) if total > 0 else 0
|
||||
return {
|
||||
**self._stats,
|
||||
"size": len(self._data),
|
||||
"hit_rate": round(hit_rate, 2),
|
||||
}
|
||||
|
||||
|
||||
class DataAccessLayer:
|
||||
TTL_USER = 600.0
|
||||
TTL_FOLDER = 120.0
|
||||
TTL_FILE = 120.0
|
||||
TTL_PATH = 60.0
|
||||
TTL_CONTENTS = 30.0
|
||||
|
||||
def __init__(self, cache_size: int = 50000):
|
||||
self._cache = DALCache(maxsize=cache_size)
|
||||
self._path_cache: Dict[str, Tuple[Any, Any, bool]] = {}
|
||||
self._path_lock = asyncio.Lock()
|
||||
|
||||
async def start(self):
|
||||
logger.info("DataAccessLayer started")
|
||||
|
||||
async def stop(self):
|
||||
await self._cache.clear()
|
||||
async with self._path_lock:
|
||||
self._path_cache.clear()
|
||||
logger.info("DataAccessLayer stopped")
|
||||
|
||||
def _user_key(self, user_id: int) -> str:
|
||||
return f"user:{user_id}"
|
||||
|
||||
def _user_by_name_key(self, username: str) -> str:
|
||||
return f"user:name:{username}"
|
||||
|
||||
def _folder_key(self, user_id: int, folder_id: Optional[int]) -> str:
|
||||
return f"u:{user_id}:f:{folder_id or 'root'}"
|
||||
|
||||
def _file_key(self, user_id: int, file_id: int) -> str:
|
||||
return f"u:{user_id}:file:{file_id}"
|
||||
|
||||
def _folder_contents_key(self, user_id: int, folder_id: Optional[int]) -> str:
|
||||
return f"u:{user_id}:contents:{folder_id or 'root'}"
|
||||
|
||||
def _path_key(self, user_id: int, path: str) -> str:
|
||||
path_hash = hashlib.md5(path.encode()).hexdigest()[:12]
|
||||
return f"u:{user_id}:path:{path_hash}"
|
||||
|
||||
async def get_user_by_id(self, user_id: int) -> Optional[Any]:
|
||||
key = self._user_key(user_id)
|
||||
cached = await self._cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from ..models import User
|
||||
user = await User.get_or_none(id=user_id)
|
||||
if user:
|
||||
await self._cache.set(key, user, ttl=self.TTL_USER)
|
||||
await self._cache.set(self._user_by_name_key(user.username), user, ttl=self.TTL_USER)
|
||||
return user
|
||||
|
||||
async def get_user_by_username(self, username: str) -> Optional[Any]:
|
||||
key = self._user_by_name_key(username)
|
||||
cached = await self._cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from ..models import User
|
||||
user = await User.get_or_none(username=username)
|
||||
if user:
|
||||
await self._cache.set(key, user, ttl=self.TTL_USER)
|
||||
await self._cache.set(self._user_key(user.id), user, ttl=self.TTL_USER)
|
||||
return user
|
||||
|
||||
async def invalidate_user(self, user_id: int, username: Optional[str] = None):
|
||||
await self._cache.delete(self._user_key(user_id))
|
||||
if username:
|
||||
await self._cache.delete(self._user_by_name_key(username))
|
||||
await self._cache.invalidate_prefix(f"u:{user_id}:")
|
||||
|
||||
async def get_folder(
|
||||
self,
|
||||
user_id: int,
|
||||
folder_id: Optional[int] = None,
|
||||
name: Optional[str] = None,
|
||||
parent_id: Optional[int] = None,
|
||||
) -> Optional[Any]:
|
||||
if folder_id:
|
||||
key = self._folder_key(user_id, folder_id)
|
||||
cached = await self._cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from ..models import Folder
|
||||
|
||||
if folder_id:
|
||||
folder = await Folder.get_or_none(id=folder_id, owner_id=user_id, is_deleted=False)
|
||||
elif name is not None:
|
||||
folder = await Folder.get_or_none(
|
||||
name=name, parent_id=parent_id, owner_id=user_id, is_deleted=False
|
||||
)
|
||||
else:
|
||||
return None
|
||||
|
||||
if folder:
|
||||
await self._cache.set(self._folder_key(user_id, folder.id), folder, ttl=self.TTL_FOLDER)
|
||||
return folder
|
||||
|
||||
async def get_file(
|
||||
self,
|
||||
user_id: int,
|
||||
file_id: Optional[int] = None,
|
||||
name: Optional[str] = None,
|
||||
parent_id: Optional[int] = None,
|
||||
) -> Optional[Any]:
|
||||
if file_id:
|
||||
key = self._file_key(user_id, file_id)
|
||||
cached = await self._cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from ..models import File
|
||||
|
||||
if file_id:
|
||||
file = await File.get_or_none(id=file_id, owner_id=user_id, is_deleted=False)
|
||||
elif name is not None:
|
||||
file = await File.get_or_none(
|
||||
name=name, parent_id=parent_id, owner_id=user_id, is_deleted=False
|
||||
)
|
||||
else:
|
||||
return None
|
||||
|
||||
if file:
|
||||
await self._cache.set(self._file_key(user_id, file.id), file, ttl=self.TTL_FILE)
|
||||
return file
|
||||
|
||||
async def get_folder_contents(
|
||||
self,
|
||||
user_id: int,
|
||||
folder_id: Optional[int] = None,
|
||||
) -> Tuple[List[Any], List[Any]]:
|
||||
key = self._folder_contents_key(user_id, folder_id)
|
||||
cached = await self._cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from ..models import Folder, File
|
||||
|
||||
folders = await Folder.filter(
|
||||
owner_id=user_id, parent_id=folder_id, is_deleted=False
|
||||
).all()
|
||||
files = await File.filter(
|
||||
owner_id=user_id, parent_id=folder_id, is_deleted=False
|
||||
).all()
|
||||
|
||||
result = (folders, files)
|
||||
await self._cache.set(key, result, ttl=self.TTL_CONTENTS)
|
||||
return result
|
||||
|
||||
async def resolve_path(
|
||||
self,
|
||||
user_id: int,
|
||||
path_str: str,
|
||||
) -> Tuple[Optional[Any], Optional[Any], bool]:
|
||||
path_str = path_str.strip("/")
|
||||
if not path_str:
|
||||
return None, None, True
|
||||
|
||||
key = self._path_key(user_id, path_str)
|
||||
cached = await self._cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from ..models import Folder, File
|
||||
|
||||
parts = [p for p in path_str.split("/") if p]
|
||||
current_folder = None
|
||||
current_folder_id = None
|
||||
|
||||
for i, part in enumerate(parts[:-1]):
|
||||
folder = await Folder.get_or_none(
|
||||
name=part, parent_id=current_folder_id, owner_id=user_id, is_deleted=False
|
||||
)
|
||||
if not folder:
|
||||
result = (None, None, False)
|
||||
await self._cache.set(key, result, ttl=self.TTL_PATH)
|
||||
return result
|
||||
current_folder = folder
|
||||
current_folder_id = folder.id
|
||||
await self._cache.set(
|
||||
self._folder_key(user_id, folder.id), folder, ttl=self.TTL_FOLDER
|
||||
)
|
||||
|
||||
last_part = parts[-1]
|
||||
|
||||
folder = await Folder.get_or_none(
|
||||
name=last_part, parent_id=current_folder_id, owner_id=user_id, is_deleted=False
|
||||
)
|
||||
if folder:
|
||||
await self._cache.set(
|
||||
self._folder_key(user_id, folder.id), folder, ttl=self.TTL_FOLDER
|
||||
)
|
||||
result = (folder, current_folder, True)
|
||||
await self._cache.set(key, result, ttl=self.TTL_PATH)
|
||||
return result
|
||||
|
||||
file = await File.get_or_none(
|
||||
name=last_part, parent_id=current_folder_id, owner_id=user_id, is_deleted=False
|
||||
)
|
||||
if file:
|
||||
await self._cache.set(
|
||||
self._file_key(user_id, file.id), file, ttl=self.TTL_FILE
|
||||
)
|
||||
result = (file, current_folder, True)
|
||||
await self._cache.set(key, result, ttl=self.TTL_PATH)
|
||||
return result
|
||||
|
||||
result = (None, current_folder, False)
|
||||
await self._cache.set(key, result, ttl=self.TTL_PATH)
|
||||
return result
|
||||
|
||||
async def invalidate_folder(self, user_id: int, folder_id: Optional[int] = None):
|
||||
await self._cache.delete(self._folder_key(user_id, folder_id))
|
||||
await self._cache.delete(self._folder_contents_key(user_id, folder_id))
|
||||
await self._cache.invalidate_prefix(f"u:{user_id}:path:")
|
||||
|
||||
async def invalidate_file(self, user_id: int, file_id: int, parent_id: Optional[int] = None):
|
||||
await self._cache.delete(self._file_key(user_id, file_id))
|
||||
if parent_id is not None:
|
||||
await self._cache.delete(self._folder_contents_key(user_id, parent_id))
|
||||
await self._cache.invalidate_prefix(f"u:{user_id}:path:")
|
||||
|
||||
async def invalidate_path(self, user_id: int, path: str):
|
||||
key = self._path_key(user_id, path)
|
||||
await self._cache.delete(key)
|
||||
|
||||
async def invalidate_user_paths(self, user_id: int):
|
||||
await self._cache.invalidate_prefix(f"u:{user_id}:path:")
|
||||
await self._cache.invalidate_prefix(f"u:{user_id}:contents:")
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
return self._cache.get_stats()
|
||||
|
||||
|
||||
_dal: Optional[DataAccessLayer] = None
|
||||
|
||||
|
||||
async def init_dal(cache_size: int = 50000) -> DataAccessLayer:
|
||||
global _dal
|
||||
_dal = DataAccessLayer(cache_size=cache_size)
|
||||
await _dal.start()
|
||||
return _dal
|
||||
|
||||
|
||||
async def shutdown_dal():
|
||||
global _dal
|
||||
if _dal:
|
||||
await _dal.stop()
|
||||
_dal = None
|
||||
|
||||
|
||||
def get_dal() -> DataAccessLayer:
|
||||
if not _dal:
|
||||
raise RuntimeError("DataAccessLayer not initialized")
|
||||
return _dal
|
||||
@@ -0,0 +1,178 @@
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
from typing import Dict, List, Any, Optional, Callable, Awaitable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from collections import defaultdict
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WriteType(Enum):
|
||||
ACTIVITY = "activity"
|
||||
USAGE_RECORD = "usage_record"
|
||||
FILE_ACCESS = "file_access"
|
||||
WEBDAV_PROPERTY = "webdav_property"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BufferedWrite:
|
||||
write_type: WriteType
|
||||
data: Dict[str, Any]
|
||||
created_at: float = field(default_factory=time.time)
|
||||
priority: int = 0
|
||||
|
||||
|
||||
class WriteBuffer:
|
||||
def __init__(
|
||||
self,
|
||||
flush_interval: float = 60.0,
|
||||
max_buffer_size: int = 1000,
|
||||
immediate_types: Optional[set] = None,
|
||||
):
|
||||
self.flush_interval = flush_interval
|
||||
self.max_buffer_size = max_buffer_size
|
||||
self.immediate_types = immediate_types or set()
|
||||
self._buffers: Dict[WriteType, List[BufferedWrite]] = defaultdict(list)
|
||||
self._lock = asyncio.Lock()
|
||||
self._flush_task: Optional[asyncio.Task] = None
|
||||
self._running = False
|
||||
self._handlers: Dict[WriteType, Callable[[List[Dict]], Awaitable[None]]] = {}
|
||||
self._stats = {
|
||||
"buffered": 0,
|
||||
"flushed": 0,
|
||||
"immediate": 0,
|
||||
"errors": 0,
|
||||
}
|
||||
|
||||
async def start(self):
|
||||
self._running = True
|
||||
self._flush_task = asyncio.create_task(self._background_flusher())
|
||||
logger.info(f"WriteBuffer started (interval={self.flush_interval}s)")
|
||||
|
||||
async def stop(self):
|
||||
self._running = False
|
||||
if self._flush_task:
|
||||
self._flush_task.cancel()
|
||||
try:
|
||||
await self._flush_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
await self.flush_all()
|
||||
logger.info("WriteBuffer stopped")
|
||||
|
||||
def register_handler(
|
||||
self,
|
||||
write_type: WriteType,
|
||||
handler: Callable[[List[Dict]], Awaitable[None]],
|
||||
):
|
||||
self._handlers[write_type] = handler
|
||||
logger.debug(f"Registered write handler for {write_type.value}")
|
||||
|
||||
async def buffer(
|
||||
self,
|
||||
write_type: WriteType,
|
||||
data: Dict[str, Any],
|
||||
priority: int = 0,
|
||||
):
|
||||
if write_type in self.immediate_types:
|
||||
await self._execute_immediate(write_type, data)
|
||||
return
|
||||
|
||||
async with self._lock:
|
||||
self._buffers[write_type].append(
|
||||
BufferedWrite(write_type=write_type, data=data, priority=priority)
|
||||
)
|
||||
self._stats["buffered"] += 1
|
||||
|
||||
if len(self._buffers[write_type]) >= self.max_buffer_size:
|
||||
await self._flush_type(write_type)
|
||||
|
||||
async def _execute_immediate(self, write_type: WriteType, data: Dict[str, Any]):
|
||||
handler = self._handlers.get(write_type)
|
||||
if handler:
|
||||
try:
|
||||
await handler([data])
|
||||
self._stats["immediate"] += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Immediate write failed for {write_type.value}: {e}")
|
||||
self._stats["errors"] += 1
|
||||
else:
|
||||
logger.warning(f"No handler for immediate write type: {write_type.value}")
|
||||
|
||||
async def _flush_type(self, write_type: WriteType):
|
||||
if not self._buffers[write_type]:
|
||||
return
|
||||
|
||||
writes = self._buffers[write_type]
|
||||
self._buffers[write_type] = []
|
||||
|
||||
handler = self._handlers.get(write_type)
|
||||
if handler:
|
||||
try:
|
||||
data_list = [w.data for w in writes]
|
||||
await handler(data_list)
|
||||
self._stats["flushed"] += len(writes)
|
||||
logger.debug(f"Flushed {len(writes)} {write_type.value} writes")
|
||||
except Exception as e:
|
||||
logger.error(f"Flush failed for {write_type.value}: {e}")
|
||||
self._stats["errors"] += len(writes)
|
||||
async with self._lock:
|
||||
self._buffers[write_type].extend(writes)
|
||||
|
||||
async def flush_all(self):
|
||||
async with self._lock:
|
||||
types_to_flush = list(self._buffers.keys())
|
||||
|
||||
for write_type in types_to_flush:
|
||||
async with self._lock:
|
||||
await self._flush_type(write_type)
|
||||
|
||||
async def _background_flusher(self):
|
||||
while self._running:
|
||||
try:
|
||||
await asyncio.sleep(self.flush_interval)
|
||||
await self.flush_all()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error in write buffer flusher: {e}")
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
buffer_sizes = {t.value: len(b) for t, b in self._buffers.items()}
|
||||
return {
|
||||
**self._stats,
|
||||
"buffer_sizes": buffer_sizes,
|
||||
"total_buffered": sum(buffer_sizes.values()),
|
||||
}
|
||||
|
||||
|
||||
_write_buffer: Optional[WriteBuffer] = None
|
||||
|
||||
|
||||
async def init_write_buffer(
|
||||
flush_interval: float = 60.0,
|
||||
max_buffer_size: int = 1000,
|
||||
) -> WriteBuffer:
|
||||
global _write_buffer
|
||||
_write_buffer = WriteBuffer(
|
||||
flush_interval=flush_interval,
|
||||
max_buffer_size=max_buffer_size,
|
||||
)
|
||||
await _write_buffer.start()
|
||||
return _write_buffer
|
||||
|
||||
|
||||
async def shutdown_write_buffer():
|
||||
global _write_buffer
|
||||
if _write_buffer:
|
||||
await _write_buffer.stop()
|
||||
_write_buffer = None
|
||||
|
||||
|
||||
def get_write_buffer() -> WriteBuffer:
|
||||
if not _write_buffer:
|
||||
raise RuntimeError("WriteBuffer not initialized")
|
||||
return _write_buffer
|
||||
Reference in New Issue
Block a user