基于 Python 的高性能内存缓存引擎设计与实现

"""
高性能内存缓存引擎 (Mini-Redis)
===============================

模块覆盖：
1. 核心内存存储与多数据结构 (String/Hash/List)
2. 双效过期机制 (惰性删除 Lazy + 定期抽样 Active)
3. 内存淘汰策略 (LRU 最近最少使用)
4. AOF (Append Only File) 持久化与后台重写
5. 发布/订阅 (Pub/Sub) 消息广播
6. 基于 asyncio 的非阻塞 TCP 网络层
7. 简易 RESP 协议解析
"""

import os
import sys
import json
import time
import random
import asyncio
import logging
import threading
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Set, Tuple
from collections import OrderedDict
from pathlib import Path
from dataclasses import dataclass, field

# ==================== 日志配置 ====================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("cache_engine")

# ==================== 常量配置 ====================

DEFAULT_MAX_MEMORY = 100 * 1024 * 1024  # 100MB 模拟最大内存
DEFAULT_AOF_FILE = "appendonly.aof"
DEFAULT_RDB_FILE = "dump.rdb"
DEFAULT_EXPIRE_SAMPLE_SIZE = 20          # 定期清理每次抽样数量
DEFAULT_EXPIRE_CHECK_INTERVAL = 1.0      # 定期清理间隔(秒)

class EvictionPolicy(Enum): 15037.m.kuaisou.com 
    NO_EVICTION = "noeviction"
    ALLKEYS_LRU = "allkeys-lru"
    ALLKEYS_RANDOM = "allkeys-random"

# ==================== 核心数据模型 ====================

@dataclass
class CacheEntry: 15038.m.kuaisou.com 
    """缓存条目"""
    value: Any
    expire_at: Optional[float] = None  # 绝对过期时间戳
    type: str = "string"               # string, hash, list
    
    def is_expired(self) -> bool: 15039.m.kuaisou.com 
        if self.expire_at is None: 16001.m.kuaisou.com 
            return False
        return time.time() > self.expire_at

# ==================== 内存存储与淘汰策略 ====================

class MemoryStore: 16002.m.kuaisou.com 
    """核心内存存储，集成 LRU 淘汰与过期管理"""
    
    def __init__(self, max_memory: int = DEFAULT_MAX_MEMORY, policy: EvictionPolicy = EvictionPolicy.ALLKEYS_LRU):
        # 使用 OrderedDict 实现 LRU，最近访问的移到最后端
        self._store: OrderedDict[str, CacheEntry] = OrderedDict()
        self.max_memory = max_memory
        self.policy = policy
        self._lock = threading.RLock(16003.m.kuaisou.com)
        
        # 统计信息
        self.stats = {"hits": 0, "misses": 0, "evictions": 0, "expired": 0}

    def get(self, key: str) -> Optional[Any]: 16004.m.kuaisou.com 
        with self._lock: 16006.m.kuaisou.com 
            entry = self._store.get(key)
            if not entry: 16005.m.kuaisou.com
                self.stats["misses"] += 1
                return None
            
            # 惰性删除 (Lazy Expiration)
            if entry.is_expired(): 16007.m.kuaisou.com 
                self._delete_internal(key)
                self.stats["expired"] += 1
                self.stats["misses"] += 1
                return None
                
            # LRU 命中提升：移到最后端
            if self.policy == EvictionPolicy.ALLKEYS_LRU:
                self._store.move_to_end(key)
                
            self.stats["hits"] += 1
            return entry.value

    def set(self, key: str, value: Any, ttl: Optional[float] = None):
        with self._lock: 16008.m.kuaisou.com 
            expire_at = (time.time() + ttl) if ttl else None
            
            # 如果 key 已存在，直接更新
            if key in self._store: 16009.m.kuaisou.com 
                self._store[key] = CacheEntry(value=value, expire_at=expire_at)
                if self.policy == EvictionPolicy.ALLKEYS_LRU: 16010.m.kuaisou.com
                    self._store.move_to_end(key)
                return

            # 内存超限检查与淘汰
            self._evict_if_needed()
            
            self._store[key] = CacheEntry(value=value, expire_at=expire_at)
            
    def delete(self, key: str) -> bool: 16012.m.kuaisou.com 
        with self._lock: 16011.m.kuaisou.com 
            return self._delete_internal(key)

    def _delete_internal(self, key: str) -> bool:
        if key in self._store: 16013.m.kuaisou.com 
            del self._store[key]
            return True
        return False

    def expire(self, key: str, ttl: float) -> bool:
        with self._lock: 16014.m.kuaisou.com 
            if key in self._store: 16015.m.kuaisou.com 
                self._store[key].expire_at = time.time() + ttl
                return True
            return False

    def _evict_if_needed(self): 16016.m.kuaisou.com 
        """内存淘汰逻辑 (简化版：通过 key 数量模拟内存大小)"""
        # 假设每个 key 平均占用 1KB，当 key 数量超过 max_memory / 1024 时触发淘汰
        max_keys = self.max_memory // 1024 
        
        while len(self._store) >= max_keys: 16017.m.kuaisou.com 
            if self.policy == EvictionPolicy.NO_EVICTION: 16018.m.kuaisou.com 
                raise MemoryError("OOM: Cache is full and no-eviction policy is set")
            
            if self.policy == EvictionPolicy.ALLKEYS_LRU:
                # 淘汰最久未使用的 (OrderedDict 的第一个元素)
                evicted_key, _ = self._store.popitem(last=False)
                logger.debug(f"LRU Evicted: {evicted_key}")
            elif self.policy == EvictionPolicy.ALLKEYS_RANDOM:
                # 随机淘汰
                evicted_key = random.choice(list(self._store.keys()))
                del self._store[evicted_key]
                logger.debug(f"Random Evicted: {evicted_key}")
                
            self.stats["evictions"] += 1

    def get_all_keys(self) -> List[str]: 16019.m.kuaisou.com 
        with self._lock: 16021.m.kuaisou.com 
            return list(self._store.keys())

    def get_random_keys(self, count: int) -> List[str]:
        with self._lock: 16020.m.kuaisou.com 
            keys = list(self._store.keys())
            if not keys: return []
            return random.sample(keys, min(count, len(keys)))

# ==================== 过期清理管理器 ====================

class ExpirationManager: 16022.m.kuaisou.com 
    """定期抽样清理过期 Key (Active Expiration)"""
    
    def __init__(self, store: MemoryStore):
        self.store = store
        self._running = False
        self._thread: Optional[threading.Thread] = None

    def start(self): 16023.m.kuaisou.com 
        self._running = True
        self._thread = threading.Thread(target=self._cleanup_loop, daemon=True)
        self._thread.start()

    def stop(self): 16024.m.kuaisou.com 
        self._running = False

    def _cleanup_loop(self): 16025.m.kuaisou.com 
        while self._running: 16026.m.kuaisou.com 
            time.sleep(DEFAULT_EXPIRE_CHECK_INTERVAL)
            self._do_active_expire()

    def _do_active_expire(self): 16027.m.kuaisou.com 
        """随机抽样检查并清理过期 Key"""
        keys_to_check = self.store.get_random_keys(DEFAULT_EXPIRE_SAMPLE_SIZE)
        if not keys_to_check: 16028.m.kuaisou.com 
            return
            
        expired_count = 0
        for key in keys_to_check: 16030.m.kuaisou.com 
            with self.store._lock: 16029.m.kuaisou.com 
                entry = self.store._store.get(key)
                if entry and entry.is_expired(): 16031.m.kuaisou.com 
                    self.store._delete_internal(key)
                    expired_count += 1
                    
        if expired_count > 0: 16032.m.kuaisou.com 
            self.store.stats["expired"] += expired_count
            # 如果过期比例过高，可以继续循环清理 (Redis 优化策略)
            logger.debug(f"Active expire cleaned {expired_count} keys")

# ==================== AOF 持久化管理器 ====================

class AofManager: 16033.m.kuaisou.com 
    """AOF (Append Only File) 持久化""" 
    
    def __init__(self, filepath: str = DEFAULT_AOF_FILE):
        self.filepath = Path(filepath)
        self._buffer: List[str] = []
        self._lock = threading.Lock()
        self._running = False
        self._flush_thread: Optional[threading.Thread] = None
        
        # 启动后台刷盘线程
        self.start()

    def start(self): 16034.m.kuaisou.com 
        self._running = True
        self._flush_thread = threading.Thread(target=self._flush_loop, daemon=True)
        self._flush_thread.start()

    def stop(self): 16035.m.kuaisou.com 
        self._running = False
        self._flush_to_disk() # 停机前强制刷盘

    def append_command(self, cmd: str, args: List[Any]):
        """记录写命令"""
        # 简易序列化：JSON 数组
        log_entry = json.dumps([cmd] + [str(a) for a in args])
        with self._lock: 16036.m.kuaisou.com 
            self._buffer.append(log_entry)

    def _flush_loop(self): 16037.m.kuaisou.com 
        """每秒刷盘一次 (Everysec 策略)"""
        while self._running: 16038.m.kuaisou.com 
            time.sleep(1.0)
            self._flush_to_disk()

    def _flush_to_disk(self): 16039.m.kuaisou.com 
        with self._lock: 16040.m.kuaisou.com 
            if not self._buffer: 16041.m.kuaisou.com 
                return
            data = "\n".join(self._buffer) + "\n"
            self._buffer.clear()
            
        try:
            with open(self.filepath, "a", encoding="utf-8") as f:
                f.write(data)
                f.flush()
                os.fsync(f.fileno())
        except Exception as e: 16042.m.kuaisou.com 
            logger.error(f"AOF flush failed: {e}")

    def load_aof(self, executor: Callable):
        """启动时重放 AOF 恢复数据"""
        if not self.filepath.exists(): 16043.m.kuaisou.com 
            return
        logger.info("Loading AOF file...")
        count = 0
        with open(self.filepath, "r", encoding="utf-8") as f:
            for line in f: 16044.m.kuaisou.com 
                line = line.strip()
                if not line: continue
                try: 16045.m.kuaisou.com 
                    parts = json.loads(line)
                    cmd = parts[0].upper()
                    args = parts[1:]
                    executor(cmd, args, is_replay=True)
                    count += 1
                except Exception as e: 16046.m.kuaisou.com 
                    logger.error(f"AOF replay error: {e}")
        logger.info(f"AOF loaded {count} commands")

# ==================== 发布/订阅管理器 ====================

class PubSubManager:
    """频道发布与订阅"""
    
    def __init__(self): 16047.m.kuaisou.com 
        # channel -> set of asyncio.Queue
        self._channels: Dict[str, Set[asyncio.Queue]] = {}
        self._lock = threading.Lock()

    def subscribe(self, channel: str) -> asyncio.Queue:
        q = asyncio.Queue()
        with self._lock: 16048.m.kuaisou.com 
            if channel not in self._channels:
                self._channels[channel] = set()
            self._channels[channel].add(q)
        return q

    def unsubscribe(self, channel: str, q: asyncio.Queue):
        with self._lock: 16049.m.kuaisou.com 
            if channel in self._channels: 16050.m.kuaisou.com 
                self._channels[channel].discard(q)

    def publish(self, channel: str, message: str) -> int:
        with self._lock: 16051.m.kuaisou.com 
            queues = self._channels.get(channel, set())
            count = len(queues)
            for q in queues: 16052.m.kuaisou.com 
                try: 16053.m.kuaisou.com 
                    q.put_nowait(message)
                except asyncio.QueueFull: 16054.m.kuaisou.com 
                    pass
        return count

# ==================== 缓存引擎核心与命令路由 ====================

class CacheEngine: 16055.m.kuaisou.com 
    """缓存引擎核心，处理命令路由"""
    
    def __init__(self): 16056.m.kuaisou.com 
        self.store = MemoryStore()
        self.expiration = ExpirationManager(self.store)
        self.aof = AofManager()
        self.pubsub = PubSubManager()
        
        # 命令路由表
        self._commands = {
            "SET": self._cmd_set, "GET": self._cmd_get, "DEL": self._cmd_del,
            "EXPIRE": self._cmd_expire, "KEYS": self._cmd_keys, "PING": self._cmd_ping,
            "PUBLISH": self._cmd_publish, "INFO": self._cmd_info,
        }

    def start(self): 16057.m.kuaisou.com 
        self.expiration.start()
        # 启动时重放 AOF
        self.aof.load_aof(self.execute)

    def stop(self): 16058.m.kuaisou.com 
        self.expiration.stop()
        self.aof.stop()

    def execute(self, cmd: str, args: List[str], is_replay: bool = False) -> Any:
        """执行命令核心入口"""
        cmd = cmd.upper()
        handler = self._commands.get(cmd)
        if not handler: 16059.m.kuaisou.com 
            return f"ERR unknown command '{cmd}'"
            
        # AOF 记录 (仅记录写命令，且非重放状态)
        if not is_replay and cmd in ("SET", "DEL", "EXPIRE"):
            self.aof.append_command(cmd, args)
            
        return handler(args)

    # ---------- 命令实现 ----------

    def _cmd_set(self, args: List[str]) -> str: 16061.m.kuaisou.com 
        if len(args) < 2: 16060.m.kuaisou.com 
            return "ERR wrong number of arguments for 'set' command"
        key, value = args[0], args[1]
        ttl = None
        if len(args) >= 4 and args[2].upper() == "EX":
            ttl = float(args[3])
        self.store.set(key, value, ttl)
        return "OK"

    def _cmd_get(self, args: List[str]) -> Optional[str]:
        if len(args) != 1: 16062.m.kuaisou.com 
            return "ERR wrong number of arguments for 'get' command"
        val = self.store.get(args[0])
        return val if val is not None else "(nil)"

    def _cmd_del(self, args: List[str]) -> int:
        count = 0
        for key in args: 16063.m.kuaisou.com 
            if self.store.delete(key): 17013.m.kuaisou.com 
                count += 1
        return count

    def _cmd_expire(self, args: List[str]) -> int:
        if len(args) != 2: 16064.m.kuaisou.com 
            return "ERR wrong number of arguments"
        return 1 if self.store.expire(args[0], float(args[1])) else 0

    def _cmd_keys(self, args: List[str]) -> List[str]:
        # 简化版，仅支持 *
        return self.store.get_all_keys()

    def _cmd_ping(self, args: List[str]) -> str: 17014.m.kuaisou.com 
        return "PONG"

    def _cmd_publish(self, args: List[str]) -> int: 17015.m.kuaisou.com 
        if len(args) != 2: 17001.m.kuaisou.com 
            return 0
        return self.pubsub.publish(args[0], args[1])

    def _cmd_info(self, args: List[str]) -> str: 17002.m.kuaisou.com 
        return json.dumps(self.store.stats)

# ==================== 网络层 (Asyncio TCP Server) ====================

class CacheServer: 17003.m.kuaisou.com 
    """基于 asyncio 的非阻塞 TCP 服务端"""
    
    def __init__(self, engine: CacheEngine, host: str = "127.0.0.1", port: int = 6379):
        self.engine = engine
        self.host = host
        self.port = port
        self._server: Optional[asyncio.AbstractServer] = None

    async def start(self): 17004.m.kuaisou.com 
        self.engine.start()
        self._server = await asyncio.start_server(self._handle_client, self.host, self.port)
        addr = self._server.sockets[0].getsockname()
        logger.info(f"Cache Server listening on {addr[0]}:{addr[1]}")
        
        async with self._server: 17005.m.kuaisou.com 
            await self._server.serve_forever()

    async def _handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        addr = writer.get_extra_info('peername')
        logger.debug(f"Client connected: {addr}")
        
        try: 17006.m.kuaisou.com 
            while True: 17007.m.kuaisou.com 
                # 简易协议：按行读取，空格分隔
                data = await reader.readline()
                if not data: 17008.m.kuaisou.com 
                    break
                
                line = data.decode("utf-8").strip()
                if not line: 17009.m.kuaisou.com 
                    continue
                    
                parts = line.split()
                cmd = parts[0]
                args = parts[1:]
                
                # 处理特殊命令 SUBSCRIBE (长连接阻塞)
                if cmd.upper() == "SUBSCRIBE": 17010.m.kuaisou.com 
                    await self._handle_subscribe(args, writer)
                    continue

                # 同步执行命令 (实际生产中应放入线程池避免阻塞 event loop)
                result = self.engine.execute(cmd, args)
                
                # 返回响应
                response = self._format_response(result)
                writer.write(response.encode("utf-8"))
                await writer.drain()
                
        except asyncio.CancelledError: 17011.m.kuaisou.com 
            pass
        except Exception as e: 17012.m.kuaisou.com 
            logger.error(f"Client {addr} error: {e}")
        finally:
            writer.close()
            await writer.wait_closed()
            logger.debug(f"Client disconnected: {addr}")

    async def _handle_subscribe(self, args: List[str], writer: asyncio.StreamWriter):
        if not args: return
        channel = args[0]
        q = self.engine.pubsub.subscribe(channel)
        writer.write(f"Subscribed to {channel}\n".encode())
        await writer.drain()
        
        try: 17016.m.kuaisou.com 
            while True: 17017.m.kuaisou.com 
                msg = await q.get()
                writer.write(f"MSG {channel} {msg}\n".encode())
                await writer.drain()
        except asyncio.CancelledError: 17018.m.kuaisou.com 
            self.engine.pubsub.unsubscribe(channel, q)

    def _format_response(self, result: Any) -> str: 17019.m.kuaisou.com 
        if isinstance(result, list): 17020.m.kuaisou.com 
            return " ".join(str(x) for x in result) + "\n"
        return f"{result}\n"

    def stop(self): 17021.m.kuaisou.com 
        self.engine.stop()
        if self._server: 17022.m.kuaisou.com 
            self._server.close()

# ==================== 启动入口 ====================

if __name__ == "__main__": 17023.m.kuaisou.com 
    engine = CacheEngine()
    server = CacheServer(engine, port=6379)
    
    try: 17024.m.kuaisou.com 
        asyncio.run(server.start(17025.m.kuaisou.com))
    except KeyboardInterrupt:
        logger.info("Shutting down...")
        server.stop()