""" 數據庫緩存層 - 減少 Firestore 調用頻率 實現多級緩存策略，大幅降低數據庫讀取次數 """ import logging import asyncio from datetime import datetime, timedelta from typing import Dict, Any, Optional, List, Tuple from collections import OrderedDict import hashlib import json logger = logging.getLogger("DatabaseCache") class LRUCache: """LRU (Least Recently Used) 緩存實現""" def __init__(self, max_size: int = 1000, ttl_seconds: int = 300): """ 初始化 LRU 緩存 Args: max_size: 最大緩存條目數 ttl_seconds: 緩存過期時間（秒） """ self.cache: OrderedDict = OrderedDict() self.max_size = max_size self.ttl = timedelta(seconds=ttl_seconds) self.hits = 0 self.misses = 0 self._lock = asyncio.Lock() async def get(self, key: str) -> Optional[Any]: """獲取緩存值""" async with self._lock: if key not in self.cache: self.misses += 1 return None value, expire_time = self.cache[key] # 檢查是否過期 if datetime.now() > expire_time: del self.cache[key] self.misses += 1 return None # 移到最後（表示最近使用） self.cache.move_to_end(key) self.hits += 1 return value async def set(self, key: str, value: Any): """設置緩存值""" async with self._lock: expire_time = datetime.now() + self.ttl if key in self.cache: # 更新現有值 self.cache[key] = (value, expire_time) self.cache.move_to_end(key) else: # 新增值 self.cache[key] = (value, expire_time) # 如果超過最大容量，移除最舊的 if len(self.cache) > self.max_size: oldest_key = next(iter(self.cache)) del self.cache[oldest_key] logger.debug(f"LRU 緩存已滿，移除最舊條目: {oldest_key}") async def delete(self, key: str): """刪除緩存值""" async with self._lock: if key in self.cache: del self.cache[key] async def clear(self): """清空緩存""" async with self._lock: self.cache.clear() self.hits = 0 self.misses = 0 def get_stats(self) -> Dict[str, Any]: """獲取緩存統計""" total = self.hits + self.misses hit_rate = (self.hits / total * 100) if total > 0 else 0 return { "size": len(self.cache), "max_size": self.max_size, "hits": self.hits, "misses": self.misses, "hit_rate": f"{hit_rate:.2f}%" } class DatabaseCache: """數據庫緩存管理器""" def __init__(self): # 不同數據類型使用不同的緩存策略 self.user_cache = LRUCache(max_size=500, ttl_seconds=600) # 用戶資料：10分鐘 self.chat_cache = LRUCache(max_size=300, ttl_seconds=300) # 對話資料：5分鐘 self.message_cache = LRUCache(max_size=1000, ttl_seconds=180) # 消息歷史：3分鐘 self.memory_cache = LRUCache(max_size=200, ttl_seconds=900) # 記憶：15分鐘 # 寫入緩衝區（批量寫入優化） self.write_buffer: Dict[str, List[Dict[str, Any]]] = { "messages": [], "memories": [], "chats": [] } self.write_buffer_lock = asyncio.Lock() self.write_buffer_max_size = 50 # 達到50條時觸發批量寫入 self.write_buffer_timeout = 10 # 10秒後強制寫入 # 請求合併（同一查詢只執行一次） self.pending_requests: Dict[str, asyncio.Future] = {} self.pending_lock = asyncio.Lock() # 其他快取：環境、反地理、路徑 self.env_ctx_cache = LRUCache(max_size=1000, ttl_seconds=600) # 使用者環境快取：10 分鐘 self.geo_cache = LRUCache(max_size=5000, ttl_seconds=604800) # 反地理快取：7 天 self.route_cache = LRUCache(max_size=5000, ttl_seconds=86400) # 路線快取：1 天 logger.info("數據庫緩存管理器初始化完成") def _generate_cache_key(self, operation: str, **kwargs) -> str: """生成緩存鍵""" # 將參數排序後生成哈希，確保相同參數生成相同鍵 params_str = json.dumps(kwargs, sort_keys=True, default=str) key_hash = hashlib.md5(f"{operation}:{params_str}".encode()).hexdigest() return f"{operation}:{key_hash}" async def get_user_cached(self, user_id: str) -> Optional[Dict[str, Any]]: """獲取緩存的用戶資料""" cache_key = self._generate_cache_key("user", user_id=user_id) return await self.user_cache.get(cache_key) async def set_user_cache(self, user_id: str, user_data: Dict[str, Any]): """設置用戶緩存""" cache_key = self._generate_cache_key("user", user_id=user_id) await self.user_cache.set(cache_key, user_data) async def get_chat_cached(self, chat_id: str) -> Optional[Dict[str, Any]]: """獲取緩存的對話資料""" cache_key = self._generate_cache_key("chat", chat_id=chat_id) return await self.chat_cache.get(cache_key) async def set_chat_cache(self, chat_id: str, chat_data: Dict[str, Any]): """設置對話緩存""" cache_key = self._generate_cache_key("chat", chat_id=chat_id) await self.chat_cache.set(cache_key, chat_data) async def invalidate_chat_cache(self, chat_id: str): """使對話緩存失效（當對話更新時調用）""" cache_key = self._generate_cache_key("chat", chat_id=chat_id) await self.chat_cache.delete(cache_key) async def get_user_chats_cached(self, user_id: str) -> Optional[List[Dict[str, Any]]]: """獲取緩存的用戶對話列表""" cache_key = self._generate_cache_key("user_chats", user_id=user_id) return await self.chat_cache.get(cache_key) async def set_user_chats_cache(self, user_id: str, chats: List[Dict[str, Any]]): """設置用戶對話列表緩存""" cache_key = self._generate_cache_key("user_chats", user_id=user_id) await self.chat_cache.set(cache_key, chats) async def invalidate_user_chats_cache(self, user_id: str): """使用戶對話列表緩存失效""" cache_key = self._generate_cache_key("user_chats", user_id=user_id) await self.chat_cache.delete(cache_key) async def get_memories_cached(self, user_id: str, memory_type: Optional[str] = None) -> Optional[List[Dict[str, Any]]]: """獲取緩存的記憶""" cache_key = self._generate_cache_key("memories", user_id=user_id, memory_type=memory_type) return await self.memory_cache.get(cache_key) async def set_memories_cache(self, user_id: str, memories: List[Dict[str, Any]], memory_type: Optional[str] = None): """設置記憶緩存""" cache_key = self._generate_cache_key("memories", user_id=user_id, memory_type=memory_type) await self.memory_cache.set(cache_key, memories) async def coalesce_request(self, cache_key: str, fetch_func): """ 請求合併：如果同一個請求正在執行，等待其結果而不是重複執行 Args: cache_key: 請求的唯一標識 fetch_func: 實際的數據獲取函數（async） Returns: 查詢結果 """ async with self.pending_lock: # 檢查是否有相同請求正在執行 if cache_key in self.pending_requests: logger.debug(f"請求合併：等待現有請求 {cache_key}") # 等待現有請求完成 future = self.pending_requests[cache_key] else: # 創建新的請求 future = asyncio.create_task(fetch_func()) self.pending_requests[cache_key] = future try: # 等待結果 result = await future return result finally: # 清理已完成的請求 async with self.pending_lock: if cache_key in self.pending_requests: del self.pending_requests[cache_key] async def buffer_write(self, collection: str, data: Dict[str, Any]) -> bool: """ 緩衝寫入：先存入緩衝區，達到一定數量或時間後批量寫入 Args: collection: 集合名稱 (messages/memories/chats) data: 要寫入的數據 Returns: True 如果已觸發批量寫入 """ async with self.write_buffer_lock: if collection not in self.write_buffer: self.write_buffer[collection] = [] # 添加時間戳 data["_buffered_at"] = datetime.now() self.write_buffer[collection].append(data) # 檢查是否需要立即寫入 if len(self.write_buffer[collection]) >= self.write_buffer_max_size: logger.info(f"寫入緩衝區已滿 ({collection})，觸發批量寫入") return True return False async def flush_write_buffer(self, collection: Optional[str] = None) -> Dict[str, int]: """ 清空寫入緩衝區，執行批量寫入 Args: collection: 指定集合名稱，None 表示清空所有 Returns: 每個集合的寫入數量 """ async with self.write_buffer_lock: collections_to_flush = [collection] if collection else list(self.write_buffer.keys()) result = {} for coll in collections_to_flush: if coll not in self.write_buffer or not self.write_buffer[coll]: result[coll] = 0 continue items = self.write_buffer[coll] self.write_buffer[coll] = [] result[coll] = len(items) logger.info(f"批量寫入 {coll}: {len(items)} 條記錄") # 這裡需要實際的批量寫入邏輯 # 將在 database.py 中實現 return result async def get_buffer_size(self) -> Dict[str, int]: """獲取各緩衝區的大小""" async with self.write_buffer_lock: return {k: len(v) for k, v in self.write_buffer.items()} def get_all_stats(self) -> Dict[str, Any]: """獲取所有緩存統計""" return { "user_cache": self.user_cache.get_stats(), "chat_cache": self.chat_cache.get_stats(), "message_cache": self.message_cache.get_stats(), "memory_cache": self.memory_cache.get_stats(), "env_ctx_cache": self.env_ctx_cache.get_stats(), "geo_cache": self.geo_cache.get_stats(), "route_cache": self.route_cache.get_stats(), "write_buffer": {k: len(v) for k, v in self.write_buffer.items()} } async def clear_all(self): """清空所有緩存""" await self.user_cache.clear() await self.chat_cache.clear() await self.message_cache.clear() await self.memory_cache.clear() logger.info("所有緩存已清空") # ===== 環境/地理/路線 快取 API ===== async def get_env_ctx_cached(self, user_id: str) -> Optional[Dict[str, Any]]: key = self._generate_cache_key("env_ctx", user_id=user_id) return await self.env_ctx_cache.get(key) async def set_env_ctx_cache(self, user_id: str, ctx: Dict[str, Any]): key = self._generate_cache_key("env_ctx", user_id=user_id) await self.env_ctx_cache.set(key, ctx) async def get_geo_cached(self, geohash7: str) -> Optional[Dict[str, Any]]: key = self._generate_cache_key("geo", geohash=geohash7) return await self.geo_cache.get(key) async def set_geo_cache(self, geohash7: str, payload: Dict[str, Any]): key = self._generate_cache_key("geo", geohash=geohash7) await self.geo_cache.set(key, payload) async def get_route_cached(self, cache_key: str) -> Optional[Dict[str, Any]]: key = self._generate_cache_key("route", key=cache_key) return await self.route_cache.get(key) async def set_route_cache(self, cache_key: str, payload: Dict[str, Any]): key = self._generate_cache_key("route", key=cache_key) await self.route_cache.set(key, payload) async def get_tdx_cached(self, cache_key: str) -> Optional[Any]: """獲取 TDX API 快取資料""" return await self.route_cache.get(cache_key) async def set_tdx_cache(self, cache_key: str, data: Any, ttl: int = 60): """設置 TDX API 快取資料（使用 route_cache，因為 TDX 也是路線相關）""" await self.route_cache.set(cache_key, data) # 全局緩存實例 db_cache = DatabaseCache() async def periodic_cache_maintenance(): """定期緩存維護任務""" while True: try: await asyncio.sleep(300) # 每5分鐘執行一次 # 輸出緩存統計 stats = db_cache.get_all_stats() logger.info(f"緩存統計: {stats}") # 檢查寫入緩衝區 buffer_size = await db_cache.get_buffer_size() if any(size > 0 for size in buffer_size.values()): logger.info(f"清空寫入緩衝區: {buffer_size}") await db_cache.flush_write_buffer() except Exception as e: logger.error(f"緩存維護任務出錯: {e}")