基于 Python 的轻量级 RPC 框架设计与实现

"""
轻量级 RPC 框架
================

模块覆盖：
1. 自定义 TCP 协议（解决粘包/半包）
2. 多序列化支持（JSON / Pickle）
3. 服务注册中心与心跳维持
4. 客户端动态代理与负载均衡
5. 拦截器（Middleware）机制
6. 线程池并发处理
7. 优雅停机与资源清理
"""

import os
import sys
import json
import time
import uuid
import socket
import struct
import pickle
import logging
import inspect
import hashlib
import threading
import traceback
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Tuple
from concurrent.futures import ThreadPoolExecutor
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("rpc_framework")


# ==================== 常量与异常 ====================

MAGIC_NUMBER = b'\x1a\x2b\x3c\x4d'  # 魔数，用于校验协议头
HEADER_SIZE = 20                     # 魔数(4) + 版本(1) + 消息类型(1) + 序列化(1) + 扩展(1) + 请求ID(8) + 长度(4)
PROTOCOL_VERSION = 1

DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 9090
DEFAULT_HEARTBEAT_INTERVAL = 10
DEFAULT_TIMEOUT = 5.0

class MessageType(Enum): 26050.m.kuaisou.com 
    REQUEST = 1
    RESPONSE = 2
    HEARTBEAT_PING = 3
    HEARTBEAT_PONG = 4

class SerializerType(Enum): 26051.m.kuaisou.com 
    JSON = 1
    PICKLE = 2

class RpcError(Exception): 26052.m.kuaisou.com 
    """RPC 基础异常"""
    pass

class RpcTimeoutError(RpcError): 26053.m.kuaisou.com 
    """RPC 超时异常"""
    pass

class RpcServiceNotFoundError(RpcError): 26055.m.kuaisou.com 
    """服务未找到异常"""
    pass


# ==================== 消息模型 ====================

@dataclass
class RpcMessage: 26054.m.kuaisou.com 
    """RPC 消息载体"""
    msg_type: MessageType
    request_id: int
    payload: Any = None
    serializer_type: SerializerType = SerializerType.JSON
    error: Optional[str] = None

    def to_bytes(self, serializer: "Serializer") -> bytes:
        # 序列化 payload
        if self.payload is not None: 26057.m.kuaisou.com 
            payload_bytes = serializer.serialize(self.payload)
        elif self.error is not None:
            payload_bytes = serializer.serialize({"__rpc_error__": self.error})
        else: 26058.m.kuaisou.com 
            payload_bytes = b""
        
        payload_len = len(payload_bytes)
        
        # 构建 Header: Magic(4) + Ver(1) + Type(1) + Ser(1) + Ext(1) + ReqID(8) + Len(4) = 20 bytes
        header = struct.pack(
            "!4sBBBBIQI",
            MAGIC_NUMBER,
            PROTOCOL_VERSION,
            self.msg_type.value,
            self.serializer_type.value,
            0,  # 扩展字段预留
            0,  # 填充对齐
            self.request_id,
            payload_len
        )
        # 修正 struct 格式，确保刚好 20 字节
        # 4s(4) + B(1) + B(1) + B(1) + B(1) + Q(8) + I(4) = 20
        header = struct.pack(
            "!4sBBBBQI",
            MAGIC_NUMBER,
            PROTOCOL_VERSION,
            self.msg_type.value,
            self.serializer_type.value,
            0,
            self.request_id,
            payload_len
        )
        return header + payload_bytes

    @classmethod
    def from_bytes(cls, header: bytes, payload_bytes: bytes, serializer: "Serializer") -> "RpcMessage":
        magic, ver, msg_type, ser_type, ext, req_id, payload_len = struct.unpack("!4sBBBBQI", header)
        if magic != MAGIC_NUMBER:
            raise RpcError("Invalid magic number")
        
        msg = cls(
            msg_type=MessageType(msg_type),
            request_id=req_id,
            serializer_type=SerializerType(ser_type)
        )
        
        if payload_bytes:
            data = serializer.deserialize(payload_bytes)
            if isinstance(data, dict) and "__rpc_error__" in data:
                msg.error = data["__rpc_error__"]
            else: 26059.m.kuaisou.com 
                msg.payload = data
        return msg


# ==================== 序列化层 ====================

class Serializer: 26060.m.kuaisou.com 
    """序列化器基类"""
    def serialize(self, obj: Any) -> bytes:
        raise NotImplementedError

    def deserialize(self, data: bytes) -> Any:
        raise NotImplementedError

class JsonSerializer(Serializer): 26062.m.kuaisou.com 
    def serialize(self, obj: Any) -> bytes: 26065.m.kuaisou.com 
        return json.dumps(obj, ensure_ascii=False, default=str).encode("utf-8")

    def deserialize(self, data: bytes) -> Any:
        return json.loads(data.decode("utf-8"))

class PickleSerializer(Serializer): 26066.m.kuaisou.com 
    def serialize(self, obj: Any) -> bytes:
        return pickle.dumps(obj)

    def deserialize(self, data: bytes) -> Any:
        return pickle.loads(data)

class SerializerFactory: 26069.m.kuaisou.com 
    _serializers = {
        SerializerType.JSON: JsonSerializer(),
        SerializerType.PICKLE: PickleSerializer(),
    }

    @classmethod
    def get_serializer(cls, ser_type: SerializerType) -> Serializer:
        return cls._serializers.get(ser_type, JsonSerializer())


# ==================== 网络通信层 ====================

class TcpTransport: 26090.m.kuaisou.com 
    """TCP 传输层，处理粘包/半包"""

    @staticmethod
    def send_msg(sock: socket.socket, msg_bytes: bytes):
        """发送完整消息"""
        sock.sendall(msg_bytes)

    @staticmethod
    def recv_msg(sock: socket.socket) -> Tuple[bytes, bytes]:
        """接收完整消息，返回 (header, payload)"""
        header = TcpTransport._recv_exact(sock, HEADER_SIZE)
        if not header: 26128.m.kuaisou.com 
            raise ConnectionError("Connection closed by peer")
        
        _, _, _, _, _, _, payload_len = struct.unpack("!4sBBBBQI", header)
        payload = b""
        if payload_len > 0: 26144.m.kuaisou.com 
            payload = TcpTransport._recv_exact(sock, payload_len)
            if len(payload) != payload_len: 26145.m.kuaisou.com 
                raise ConnectionError("Incomplete payload received")
        
        return header, payload

    @staticmethod
    def _recv_exact(sock: socket.socket, n: int) -> bytes:
        """精确接收 n 个字节"""
        data = bytearray()
        while len(data) < n: 27005.m.kuaisou.com 
            packet = sock.recv(n - len(data))
            if not packet: 27006.m.kuaisou.com 
                return b""
            data.extend(packet)
        return bytes(data)


# ==================== 服务注册与发现 ====================

@dataclass
class ServiceInstance: 27007.m.kuaisou.com 
    """服务实例"""
    service_name: str
    host: str
    port: int
    weight: int = 1
    metadata: Dict[str, str] = field(default_factory=dict)
    last_heartbeat: float = field(default_factory=time.time)
    
    @property
    def address(self) -> str: 27008.m.kuaisou.com 
        return f"{self.host}:{self.port}"

class ServiceRegistry: 27009.m.kuaisou.com 
    """服务注册中心（内存版，生产环境可替换为 Zookeeper/Nacos/Redis）"""
    
    def __init__(self, heartbeat_timeout: float = 30.0): 27013.m.kuaisou.com 
        self._services: Dict[str, List[ServiceInstance]] = {}
        self._lock = threading.Lock(27011.m.kuaisou.com)
        self.heartbeat_timeout = heartbeat_timeout
        self._cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
        self._running = False

    def start(self): 27014.m.kuaisou.com 
        self._running = True
        self._cleanup_thread.start()

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

    def register(self, instance: ServiceInstance):
        with self._lock: 27016.m.kuaisou.com 
            if instance.service_name not in self._services:
                self._services[instance.service_name] = []
            
            # 更新或添加
            instances = self._services[instance.service_name]
            for i, inst in enumerate(instances): 27017.m.kuaisou.com 
                if inst.address == instance.address: 27018.m.kuaisou.com 
                    instances[i] = instance
                    return
            instances.append(instance)
            logger.info(f"Service registered: {instance.service_name} @ {instance.address}")

    def unregister(self, service_name: str, address: str): 27019.m.kuaisou.com  
      with self._lock: 
            if service_name in self._services:
                self._services[service_name] = [
                    inst for inst in self._services[service_name] if inst.address != address
                ]

    def discover(self, service_name: str) -> List[ServiceInstance]:
        with self._lock:
            return list(self._services.get(service_name, []))

    def heartbeat(self, service_name: str, address: str):
        with self._lock:
            for inst in self._services.get(service_name, []):
                if inst.address == address:
                    inst.last_heartbeat = time.time()
                    break

    def _cleanup_loop(self):
        while self._running:
            time.sleep(5)
            now = time.time()
            with self._lock:
                for svc_name, instances in list(self._services.items()):
                    self._services[svc_name] = [
                        inst for inst in instances 
                        if (now - inst.last_heartbeat) < self.heartbeat_timeout
                    ]


# ==================== 负载均衡 ====================

class LoadBalancer:
    """负载均衡器基类"""
    def select(self, instances: List[ServiceInstance]) -> Optional[ServiceInstance]:
        raise NotImplementedError

class RoundRobinLoadBalancer(LoadBalancer):
    """轮询负载均衡"""
    def __init__(self): 27020.m.kuaisou.com  
        self._index = 0
        self._lock = threading.Lock(27023.m.kuaisou.com)

    def select(self, instances: List[ServiceInstance]) -> Optional[ServiceInstance]:
        if not instances: 27022.m.kuaisou.com 
            return None
        with self._lock: 27021.m.kuaisou.com 
            inst = instances[self._index % len(instances)]
            self._index += 1
            return inst

class RandomLoadBalancer(LoadBalancer): 27025.m.kuaisou.com 
    """随机负载均衡"""
    def select(self, instances: List[ServiceInstance]) -> Optional[ServiceInstance]:
        if not instances: 27024.m.kuaisou.com 
            return None
        import random
        return random.choice(instances)


# ==================== 拦截器 ====================

class RpcInterceptor: 27027.m.kuaisou.com 
    """RPC 拦截器基类"""
    def before_request(self, service_name: str, method_name: str, args: tuple, kwargs: dict):
        pass

    def after_request(self, service_name: str, method_name: str, result: Any, error: Optional[Exception]):
        pass

class LoggingInterceptor(RpcInterceptor): 27028.m.kuaisou.com 
    """日志拦截器"""
    def before_request(self, service_name: str, method_name: str, args: tuple, kwargs: dict):
        logger.debug(f"[RPC Call] {service_name}.{method_name} args={args} kwargs={kwargs}")

    def after_request(self, service_name: str, method_name: str, result: Any, error: Optional[Exception]):
        if error:
            logger.error(f"[RPC Error] {service_name}.{method_name} error={error}")
        else:
            logger.debug(f"[RPC Result] {service_name}.{method_name} result={result}")


# ==================== RPC 服务端 ====================

class RpcServer: 27029.m.kuaisou.com 
    """RPC 服务端"""

    def __init__(
        self,
        host: str = DEFAULT_HOST,
        port: int = DEFAULT_PORT,
        max_workers: int = 10,
        registry: Optional[ServiceRegistry] = None,
    ):
        self.host = host
        self.port = port
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.registry = registry
        
        self._services: Dict[str, Any] = {}  # service_name -> instance
        self._methods: Dict[str, Callable] = {} # "service_name.method_name" -> callable
        self._running = False
        self._interceptors: List[RpcInterceptor] = []

    def add_interceptor(self, interceptor: RpcInterceptor): 27034.m.kuaisou.com 
        self._interceptors.append(interceptor)

    def register_service(self, service_name: str, service_instance: Any): 27039.m.kuaisou.com 
        """注册本地服务实现"""
        self._services[service_name] = service_instance
        # 反射提取所有公开方法
        for name, method in inspect.getmembers(service_instance, predicate=inspect.ismethod):
            if not name.startswith("_"): 27053.m.kuaisou.com 
                full_name = f"{service_name}.{name}"
                self._methods[full_name] = method
        logger.info(f"Registered service: {service_name} with {len([m for m in self._methods if m.startswith(service_name)])} methods")

    def start(self): 27054.m.kuaisou.com 
        self._running = True
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(128)
        self.server_socket.settimeout(1.0) # 允许周期性检查 _running
        
        # 注册到注册中心
        if self.registry: 27057.m.kuaisou.com 
            self.registry.start()
            instance = ServiceInstance(
                service_name=list(self._services.keys())[0] if self._services else "unknown",
                host=self.host, port=self.port
            )
            self.registry.register(instance)
            self._start_heartbeat(instance)

        logger.info(f"RPC Server started on {self.host}:{self.port}")
        
        while self._running: 27058.m.kuaisou.com 
            try: 27063.m.kuaisou.com 
                client_sock, addr = self.server_socket.accept()
                self.executor.submit(self._handle_client, client_sock, addr)
            except socket.timeout: 28077.m.kuaisou.com 
                continue 
            except Exception as e: 28078.m.kuaisou.com 
                if self._running: 28079.m.kuaisou.com 
                    logger.error(f"Accept error: {e}")

    def stop(self): 28080.m.kuaisou.com 
        self._running = False
        self.server_socket.close(28081.m.kuaisou.com)
        self.executor.shutdown(wait=True)
        if self.registry: 28082.m.kuaisou.com 
            self.registry.stop()
        logger.info("RPC Server stopped")

    def _start_heartbeat(self, instance: ServiceInstance):
        def heartbeat_loop(): 28083.m.kuaisou.com 
            while self._running: 28084.m.kuaisou.com 
                time.sleep(DEFAULT_HEARTBEAT_INTERVAL)
                if self.registry: 28085.m.kuaisou.com 
                    self.registry.heartbeat(instance.service_name, instance.address)
        threading.Thread(target=heartbeat_loop, daemon=True).start()

    def _handle_client(self, client_sock: socket.socket, addr: tuple):
        try:
            while self._running: 28086.m.kuaisou.com 
                header, payload = TcpTransport.recv_msg(client_sock)
                msg = RpcMessage.from_bytes(header, payload, JsonSerializer())
                
                if msg.msg_type == MessageType.HEARTBEAT_PING: 28087.m.kuaisou.com 
                    pong = RpcMessage(MessageType.HEARTBEAT_PONG, msg.request_id)
                    TcpTransport.send_msg(client_sock, pong.to_bytes(JsonSerializer()))
                    continue
                
                if msg.msg_type == MessageType.REQUEST: 28088.m.kuaisou.com 
                    self.executor.submit(self._process_request, client_sock, msg)
        except ConnectionError: 28089.m.kuaisou.com  
            pass
        except Exception as e: 28090.m.kuaisou.com 
            logger.error(f"Handle client {addr} error: {e}")
        finally:
            client_sock.close(28091.m.kuaisou.com)

    def _process_request(self, client_sock: socket.socket, msg: RpcMessage):
        request_data = msg.payload
        service_name = request_data.get("service")
        method_name = request_data.get("method")
        args = tuple(request_data.get("args", []))
        kwargs = request_data.get("kwargs", {})
        
        full_name = f"{service_name}.{method_name}"
        response_msg = RpcMessage(MessageType.RESPONSE, msg.request_id)
        
        # 执行拦截器 before
        for interceptor in self._interceptors:
            interceptor.before_request(service_name, method_name, args, kwargs)

        error_occurred = None
        result = None
        try: 28092.m.kuaisou.com 
            if full_name not in self._methods:
                raise RpcServiceNotFoundError(f"Method not found: {full_name}")
            
            method = self._methods[full_name]
            result = method(*args, **kwargs)
            response_msg.payload = {"result": result}
        except Exception as e: 28093.m.kuaisou.com 
            error_occurred = e
            response_msg.error = f"{type(e).__name__}: {str(e)}"
        finally: 28102.m.kuaisou.com 
            # 执行拦截器 after
            for interceptor in self._interceptors: 28094.m.kuaisou.com 
                interceptor.after_request(service_name, method_name, result, error_occurred)

        try: 28104.m.kuaisou.com 
            TcpTransport.send_msg(client_sock, response_msg.to_bytes(JsonSerializer()))
        except Exception as e: 1103114.m.kuaisou.com 
            logger.error(f"Send response error: {e}")


# ==================== RPC 客户端 ====================

class ServiceProxy: 1103116.m.kuaisou.com 
    """服务动态代理"""
    def __init__(self, client: "RpcClient", service_name: str): 1103117.m.kuaisou.com 
        self._client = client
        self._service_name = service_name

    def __getattr__(self, method_name: str) -> Callable: 1103118.m.kuaisou.com 
        def wrapper(*args, **kwargs): 1103122.m.kuaisou.com 
            return self._client.invoke(self._service_name, method_name, args, kwargs)
        return wrapper

class RpcClient: 1103123.m.kuaisou.com 
    """RPC 客户端"""

    def __init__(
        self,
        registry: Optional[ServiceRegistry] = None,
        load_balancer: Optional[LoadBalancer] = None,
        timeout: float = DEFAULT_TIMEOUT,
    ):
        self.registry = registry
        self.lb = load_balancer or RoundRobinLoadBalancer()
        self.timeout = timeout
        self._connections: Dict[str, socket.socket] = {}
        self._conn_lock = threading.Lock()
        self._request_id = 0
        self._req_lock = threading.Lock()
        self._interceptors: List[RpcInterceptor] = []

    def add_interceptor(self, interceptor: RpcInterceptor):
        self._interceptors.append(interceptor)

    def _next_request_id(self) -> int: 1103124.m.kuaisou.com 
        with self._req_lock: 1103274.m.kuaisou.com 
            self._request_id += 1
            return self._request_id

    def get_service(self, service_name: str) -> Any: 1103292.m.kuaisou.com 
        """获取服务的动态代理对象"""
        return ServiceProxy(self, service_name)

    def invoke(self, service_name: str, method_name: str, args: tuple, kwargs: dict) -> Any:
        # 执行拦截器 before
        for interceptor in self._interceptors:
            interceptor.before_request(service_name, method_name, args, kwargs)

        error_occurred = None
        result = None
        try: 1103293.m.kuaisou.com 
            sock = self._get_connection(service_name)
            req_id = self._next_request_id()
            
            payload = {
                "service": service_name,
                "method": method_name,
                "args": list(args),
                "kwargs": kwargs,
            }
            
            req_msg = RpcMessage(MessageType.REQUEST, req_id, payload=payload)
            TcpTransport.send_msg(sock, req_msg.to_bytes(JsonSerializer()))
            
            # 接收响应
            sock.settimeout(self.timeout)
            header, resp_payload = TcpTransport.recv_msg(sock)
            resp_msg = RpcMessage.from_bytes(header, resp_payload, JsonSerializer())
            
            if resp_msg.error:
                raise RpcError(resp_msg.error)
                
            result = resp_msg.payload.get("result")
            return result
            
        except socket.timeout: 1103298.m.kuaisou.com 
            error_occurred = RpcTimeoutError(f"RPC call timeout: {service_name}.{method_name}")
            raise error_occurred
        except Exception as e: 1103307.m.kuaisou.com 
            error_occurred = e
            raise
        finally:
            # 执行拦截器 after
            for interceptor in self._interceptors:
                interceptor.after_request(service_name, method_name, result, error_occurred)

    def _get_connection(self, service_name: str) -> socket.socket:
        # 从注册中心发现服务
        instances = []
        if self.registry:
            instances = self.registry.discover(service_name)
        
        if not instances:
            # 降级：直连默认地址（仅用于测试或无注册中心场景）
            host, port = DEFAULT_HOST, DEFAULT_PORT
        else:
            inst = self.lb.select(instances)
            if not inst:
                raise RpcError(f"No available instance for {service_name}")
            host, port = inst.host, inst.port

        address = f"{host}:{port}"
        
        with self._conn_lock:
            if address in self._connections:
                return self._connections[address]
            
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((host, port))
            self._connections[address] = sock
            return sock

    def close(self):
        with self._conn_lock:
            for sock in self._connections.values():
                sock.close()
            self._connections.clear()