backend/app/core/cache.py

import collections
import sys
import threading
import logging

logger = logging.getLogger(__name__)


class LRUCache:
    """
    基于LRU（最近最少使用）策略的内存缓存。
    支持按条目数量和内存大小进行限制。
    """

    def __init__(self, capacity: int = 100, max_memory_mb: int = 512):
        """
        初始化LRU缓存。
        :param capacity: 最大缓存条目数量。
        :param max_memory_mb: 最大缓存内存限制（MB）。
        """
        self.capacity = capacity
        self.max_memory_bytes = max_memory_mb * 1024 * 1024  # 转换为字节
        self.cache = collections.OrderedDict()
        self.current_memory_bytes = 0
        self.lock = threading.Lock()
        logger.info(f"LRU缓存初始化：容量={capacity}条，最大内存={max_memory_mb}MB")

    def _get_item_size(self, key, value):
        """
        估算单个缓存条目（键值对）的内存大小（字节）。
        注意：sys.getsizeof 不会递归计算对象大小，对于复杂对象可能不准确。
        """
        return sys.getsizeof(key) + sys.getsizeof(value)

    def get(self, key):
        """
        从缓存中获取值。如果键存在，则将其标记为最近使用。
        :param key: 缓存键。
        :return: 对应的值，如果不存在则返回None。
        """
        with self.lock:
            if key not in self.cache:
                return None
            value = self.cache.pop(key)
            self.cache[key] = value
            logger.debug(f"从缓存获取: {key}")
            return value

    def set(self, key, value):
        """
        将键值对存入缓存。如果缓存超出容量或内存限制，则移除最不常用的条目。
        :param key: 缓存键。
        :param value: 要缓存的值。
        """
        with self.lock:
            if key in self.cache:
                old_value = self.cache.pop(key)
                self.current_memory_bytes -= self._get_item_size(key, old_value)

            item_size = self._get_item_size(key, value)

            # 检查是否超出内存限制
            while (
                self.current_memory_bytes + item_size > self.max_memory_bytes
                and len(self.cache) > 0
            ):
                old_key, old_value = self.cache.popitem(last=False)  # 移除最不常用的
                self.current_memory_bytes -= self._get_item_size(old_key, old_value)
                logger.warning(
                    f"缓存内存超出限制，移除最旧条目: {old_key}，"
                    f"当前内存: {self.current_memory_bytes / (1024 * 1024):.2f}MB"
                )

            # 检查是否超出条目容量
            while len(self.cache) >= self.capacity and len(self.cache) > 0:
                old_key, old_value = self.cache.popitem(last=False)  # 移除最不常用的
                self.current_memory_bytes -= self._get_item_size(old_key, old_value)
                logger.warning(f"缓存条目超出容量，移除最旧条目: {old_key}，当前条目数: {len(self.cache)}")

            self.cache[key] = value
            self.current_memory_bytes += item_size
            logger.debug(
                f"设置缓存: {key}，当前条目数: {len(self.cache)}，"
                f"当前内存: {self.current_memory_bytes / (1024 * 1024):.2f}MB"
            )

    def delete(self, key):
        """
        从缓存中删除指定键。
        :param key: 缓存键。
        """
        with self.lock:
            if key in self.cache:
                value = self.cache.pop(key)
                self.current_memory_bytes -= self._get_item_size(key, value)
                logger.debug(f"从缓存删除: {key}")

    def clear(self):
        """
        清空缓存。
        """
        with self.lock:
            self.cache.clear()
            self.current_memory_bytes = 0
            logger.info("缓存已清空。")

    def __len__(self):
        """
        返回缓存中的条目数量。
        """
        with self.lock:
            return len(self.cache)

    def get_current_memory_usage_mb(self) -> float:
        """
        获取当前缓存占用的内存大小（MB）。
        """
        with self.lock:
            return self.current_memory_bytes / (1024 * 1024)


# 全局缓存实例，将在config中配置其容量和内存限制
cache_instance = None


def get_cache_instance(capacity: int = 100, max_memory_mb: int = 512):
    """
    获取LRU缓存的单例实例。
    """
    global cache_instance
    if cache_instance is None:
        cache_instance = LRUCache(capacity, max_memory_mb)
    return cache_instance