# encoding: UTF-8 import threading import traceback import signal import zmq from msgpack import packb, unpackb from json import dumps, loads import cPickle pDumps = cPickle.dumps pLoads = cPickle.loads # 实现Ctrl-c中断recv signal.signal(signal.SIGINT, signal.SIG_DFL) ######################################################################## class RpcObject(object): """ RPC对象 提供对数据的序列化打包和解包接口，目前提供了json、msgpack、cPickle三种工具。 msgpack：性能更高，但通常需要安装msgpack相关工具； json：性能略低但通用性更好，大部分编程语言都内置了相关的库。 cPickle：性能一般且仅能用于Python，但是可以直接传送Python对象，非常方便。 因此建议尽量使用msgpack，如果要和某些语言通讯没有提供msgpack时再使用json， 当传送的数据包含很多自定义的Python对象时建议使用cPickle。 如果希望使用其他的序列化工具也可以在这里添加。 """ #---------------------------------------------------------------------- def __init__(self): """Constructor""" # 默认使用msgpack作为序列化工具 #self.useMsgpack() self.usePickle() #---------------------------------------------------------------------- def pack(self, data): """打包""" pass #---------------------------------------------------------------------- def unpack(self, data): """解包""" pass #---------------------------------------------------------------------- def __jsonPack(self, data): """使用json打包""" return dumps(data) #---------------------------------------------------------------------- def __jsonUnpack(self, data): """使用json解包""" return loads(data) #---------------------------------------------------------------------- def __msgpackPack(self, data): """使用msgpack打包""" return packb(data) #---------------------------------------------------------------------- def __msgpackUnpack(self, data): """使用msgpack解包""" return unpackb(data) #---------------------------------------------------------------------- def __picklePack(self, data): """使用cPickle打包""" return pDumps(data) #---------------------------------------------------------------------- def __pickleUnpack(self, data): """使用cPickle解包""" return pLoads(data) #---------------------------------------------------------------------- def useJson(self): """使用json作为序列化工具""" self.pack = self.__jsonPack self.unpack = self.__jsonUnpack #---------------------------------------------------------------------- def useMsgpack(self): """使用msgpack作为序列化工具""" self.pack = self.__msgpackPack self.unpack = self.__msgpackUnpack #---------------------------------------------------------------------- def usePickle(self): """使用cPickle作为序列化工具""" self.pack = self.__picklePack self.unpack = self.__pickleUnpack ######################################################################## class RpcServer(RpcObject): """RPC服务器""" #---------------------------------------------------------------------- def __init__(self, repAddress, pubAddress): """Constructor""" super(RpcServer, self).__init__() # 保存功能函数的字典，key是函数名，value是函数对象 self.__functions = {} # zmq端口相关 self.__context = zmq.Context() self.__socketREP = self.__context.socket(zmq.REP) # 请求回应socket self.__socketREP.bind(repAddress) self.__socketPUB = self.__context.socket(zmq.PUB) # 数据广播socket self.__socketPUB.bind(pubAddress) # 工作线程相关 self.__active = False # 服务器的工作状态 self.__thread = threading.Thread(target=self.__run) # 服务器的工作线程 #---------------------------------------------------------------------- def start(self): """启动服务器""" # 将服务器设为启动 self.__active = True # 启动工作线程 self.__thread.start() #---------------------------------------------------------------------- def stop(self): """停止服务器""" # 将服务器设为停止 self.__active = False # 等待工作线程退出 self.__thread.join() #---------------------------------------------------------------------- def __run(self): """连续运行函数""" while self.__active: # 从请求响应socket收取请求数据 reqb = self.__socketREP.recv() # 序列化解包 req = self.unpack(reqb) # 获取函数名和参数 name, args, kwargs = req # 获取引擎中对应的函数对象，并执行调用，如果有异常则捕捉后返回 try: func = self.__functions[name] r = func(*args, **kwargs) rep = [True, r] except Exception as e: rep = [False, traceback.format_exc()] # 序列化打包 repb = self.pack(rep) # 通过请求响应socket返回调用结果 self.__socketREP.send(repb) #---------------------------------------------------------------------- def publish(self, topic, data): """ 广播推送数据 topic：主题内容 data：具体的数据 """ # 序列化数据 datab = self.pack(data) # 通过广播socket发送数据 self.__socketPUB.send_multipart([topic, datab]) #---------------------------------------------------------------------- def register(self, func): """注册函数""" self.__functions[func.__name__] = func ######################################################################## class RpcClient(RpcObject): """RPC客户端""" #---------------------------------------------------------------------- def __init__(self, reqAddress, subAddress): """Constructor""" super(RpcClient, self).__init__() # zmq端口相关 self.__reqAddress = reqAddress self.__subAddress = subAddress self.__context = zmq.Context() self.__socketREQ = self.__context.socket(zmq.REQ) # 请求发出socket self.__socketSUB = self.__context.socket(zmq.SUB) # 广播订阅socket # 工作线程相关，用于处理服务器推送的数据 self.__active = False # 客户端的工作状态 self.__thread = threading.Thread(target=self.__run) # 客户端的工作线程 #---------------------------------------------------------------------- def __getattr__(self, name): """实现远程调用功能""" # 执行远程调用任务 def dorpc(*args, **kwargs): # 生成请求 req = [name, args, kwargs] # 序列化打包请求 reqb = self.pack(req) # 发送请求并等待回应 self.__socketREQ.send(reqb) repb = self.__socketREQ.recv() # 序列化解包回应 rep = self.unpack(repb) # 若正常则返回结果，调用失败则触发异常 if rep[0]: return rep[1] else: raise RemoteException(rep[1]) return dorpc #---------------------------------------------------------------------- def start(self): """启动客户端""" # 连接端口 self.__socketREQ.connect(self.__reqAddress) self.__socketSUB.connect(self.__subAddress) # 将服务器设为启动 self.__active = True # 启动工作线程 self.__thread.start() #---------------------------------------------------------------------- def stop(self): """停止客户端""" # 将客户端设为停止 self.__active = False # 等待工作线程退出 self.__thread.join() #---------------------------------------------------------------------- def __run(self): """连续运行函数""" while self.__active: # 从订阅socket收取广播数据 topic, datab = self.__socketSUB.recv_multipart() # 序列化解包 data = self.unpack(datab) # 调用回调函数处理 self.callback(topic, data) #---------------------------------------------------------------------- def callback(self, topic, data): """回调函数，必须由用户实现""" raise NotImplementedError #---------------------------------------------------------------------- def subscribe(self, topic): """ 订阅特定主题的广播数据 可以使用topic=''来订阅所有的主题 """ self.__socketSUB.setsockopt(zmq.SUBSCRIBE, topic) ######################################################################## class RemoteException(Exception): """RPC远程异常""" #---------------------------------------------------------------------- def __init__(self, value): """Constructor""" self.__value = value #---------------------------------------------------------------------- def __str__(self): """输出错误信息""" return self.__value