rbox/mywebdav/cache/layer.py

import asyncio
import time
import hashlib
import json
from typing import Dict, Any, Optional, Callable, Set, TypeVar, Generic
from dataclasses import dataclass, field
from collections import OrderedDict
import logging

logger = logging.getLogger(__name__)

T = TypeVar('T')


@dataclass
class CacheEntry:
    value: Any
    created_at: float = field(default_factory=time.time)
    ttl: float = 300.0
    dirty: bool = False

    @property
    def is_expired(self) -> bool:
        return time.time() - self.created_at > self.ttl


class LRUCache:
    def __init__(self, maxsize: int = 10000):
        self.maxsize = maxsize
        self.cache: OrderedDict[str, CacheEntry] = OrderedDict()
        self._lock = asyncio.Lock()

    async def get(self, key: str) -> Optional[CacheEntry]:
        async with self._lock:
            if key in self.cache:
                entry = self.cache[key]
                if entry.is_expired:
                    del self.cache[key]
                    return None
                self.cache.move_to_end(key)
                return entry
            return None

    async def set(self, key: str, value: Any, ttl: float = 300.0, dirty: bool = False):
        async with self._lock:
            if key in self.cache:
                self.cache.move_to_end(key)
            elif len(self.cache) >= self.maxsize:
                oldest_key, oldest_entry = self.cache.popitem(last=False)
                if oldest_entry.dirty:
                    self.cache[oldest_key] = oldest_entry
                    self.cache.move_to_end(oldest_key)
                    if len(self.cache) >= self.maxsize:
                        for k in list(self.cache.keys()):
                            if not self.cache[k].dirty:
                                del self.cache[k]
                                break
            self.cache[key] = CacheEntry(value=value, ttl=ttl, dirty=dirty)

    async def delete(self, key: str) -> bool:
        async with self._lock:
            if key in self.cache:
                del self.cache[key]
                return True
            return False

    async def get_dirty_keys(self) -> Set[str]:
        async with self._lock:
            return {k for k, v in self.cache.items() if v.dirty}

    async def mark_clean(self, key: str):
        async with self._lock:
            if key in self.cache:
                self.cache[key].dirty = False

    async def clear(self):
        async with self._lock:
            self.cache.clear()

    async def invalidate_pattern(self, pattern: str) -> int:
        async with self._lock:
            keys_to_delete = [k for k in self.cache.keys() if pattern in k]
            for k in keys_to_delete:
                del self.cache[k]
            return len(keys_to_delete)


class CacheLayer:
    CACHE_KEYS = {
        "user_profile": "user:{user_id}:profile",
        "folder_contents": "user:{user_id}:folder:{folder_id}:contents",
        "file_metadata": "user:{user_id}:file:{file_id}:meta",
        "storage_quota": "user:{user_id}:quota",
        "folder_tree": "user:{user_id}:tree",
        "share_info": "share:{share_id}:info",
        "webdav_lock": "webdav:lock:{path_hash}",
        "rate_limit": "ratelimit:{limit_type}:{key}",
    }

    TTL_CONFIG = {
        "user_profile": 300.0,
        "folder_contents": 30.0,
        "file_metadata": 120.0,
        "storage_quota": 60.0,
        "folder_tree": 60.0,
        "share_info": 300.0,
        "webdav_lock": 3600.0,
        "rate_limit": 60.0,
        "default": 300.0,
    }

    def __init__(self, maxsize: int = 10000, flush_interval: int = 30):
        self.l1_cache = LRUCache(maxsize=maxsize)
        self.flush_interval = flush_interval
        self.dirty_keys: Set[str] = set()
        self._flush_callbacks: Dict[str, Callable] = {}
        self._flush_task: Optional[asyncio.Task] = None
        self._running = False
        self._stats = {
            "hits": 0,
            "misses": 0,
            "sets": 0,
            "invalidations": 0,
        }

    async def start(self):
        self._running = True
        self._flush_task = asyncio.create_task(self._background_flusher())
        logger.info("CacheLayer started")

    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_dirty()
        logger.info("CacheLayer stopped")

    def register_flush_callback(self, key_pattern: str, callback: Callable):
        self._flush_callbacks[key_pattern] = callback

    async def get(self, key: str, loader: Optional[Callable] = None) -> Optional[Any]:
        entry = await self.l1_cache.get(key)
        if entry:
            self._stats["hits"] += 1
            return entry.value

        self._stats["misses"] += 1

        if loader:
            try:
                value = await loader() if asyncio.iscoroutinefunction(loader) else loader()
                if value is not None:
                    ttl = self._get_ttl_for_key(key)
                    await self.set(key, value, ttl=ttl)
                return value
            except Exception as e:
                logger.error(f"Cache loader failed for key {key}: {e}")
                return None
        return None

    async def set(self, key: str, value: Any, ttl: Optional[float] = None, persist: bool = False):
        if ttl is None:
            ttl = self._get_ttl_for_key(key)
        await self.l1_cache.set(key, value, ttl=ttl, dirty=persist)
        if persist:
            self.dirty_keys.add(key)
        self._stats["sets"] += 1

    async def delete(self, key: str):
        await self.l1_cache.delete(key)
        self.dirty_keys.discard(key)

    async def invalidate(self, key: str, cascade: bool = True):
        await self.l1_cache.delete(key)
        self.dirty_keys.discard(key)
        self._stats["invalidations"] += 1

        if cascade:
            parts = key.split(":")
            if len(parts) >= 2 and parts[0] == "user":
                user_id = parts[1]
                await self.invalidate_user_cache(int(user_id))

    async def invalidate_user_cache(self, user_id: int):
        pattern = f"user:{user_id}:"
        count = await self.l1_cache.invalidate_pattern(pattern)
        logger.debug(f"Invalidated {count} cache entries for user {user_id}")

    async def invalidate_pattern(self, pattern: str) -> int:
        count = await self.l1_cache.invalidate_pattern(pattern)
        self._stats["invalidations"] += count
        return count

    async def flush_dirty(self):
        dirty_keys = await self.l1_cache.get_dirty_keys()
        for key in dirty_keys:
            entry = await self.l1_cache.get(key)
            if entry and entry.dirty:
                for pattern, callback in self._flush_callbacks.items():
                    if pattern in key:
                        try:
                            await callback(key, entry.value)
                            await self.l1_cache.mark_clean(key)
                        except Exception as e:
                            logger.error(f"Flush callback failed for {key}: {e}")
                        break
        self.dirty_keys.clear()

    def _get_ttl_for_key(self, key: str) -> float:
        for key_type, ttl in self.TTL_CONFIG.items():
            if key_type in key:
                return ttl
        return self.TTL_CONFIG["default"]

    async def _background_flusher(self):
        while self._running:
            try:
                await asyncio.sleep(self.flush_interval)
                await self.flush_dirty()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in cache background flusher: {e}")

    def get_stats(self) -> Dict[str, int]:
        total = self._stats["hits"] + self._stats["misses"]
        hit_rate = (self._stats["hits"] / total * 100) if total > 0 else 0
        return {
            **self._stats,
            "hit_rate_percent": round(hit_rate, 2),
        }

    def build_key(self, key_type: str, **kwargs) -> str:
        template = self.CACHE_KEYS.get(key_type)
        if not template:
            raise ValueError(f"Unknown cache key type: {key_type}")
        return template.format(**kwargs)


_cache: Optional[CacheLayer] = None


async def init_cache(maxsize: int = 10000, flush_interval: int = 30) -> CacheLayer:
    global _cache
    _cache = CacheLayer(maxsize=maxsize, flush_interval=flush_interval)
    await _cache.start()
    return _cache


async def shutdown_cache():
    global _cache
    if _cache:
        await _cache.stop()
        _cache = None


def get_cache() -> CacheLayer:
    if not _cache:
        raise RuntimeError("Cache not initialized")
    return _cache