From c02cc1b4a885a4ce53baf32e72fce95593f6e481 Mon Sep 17 00:00:00 2001 From: yumoqing Date: Wed, 9 Mar 2022 16:18:54 +0800 Subject: [PATCH] bugfix --- appPublic/zmq_reqrep.py | 135 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 135 insertions(+) create mode 100755 appPublic/zmq_reqrep.py diff --git a/appPublic/zmq_reqrep.py b/appPublic/zmq_reqrep.py new file mode 100755 index 0000000..29e0886 --- /dev/null +++ b/appPublic/zmq_reqrep.py @@ -0,0 +1,135 @@ +# zmq_reqresp.py +import asyncio +import zmq +import zmq.asyncio +from .background import Background +from inspect import iscoroutinefunction + +class ZmqRequester(object): + def __init__(self, url, async_mode=False, timeout=0): + super().__init__() + self.async_mode = async_mode + self.url = url + self.timeout = timeout + self._connect() + + def __del__(self): + self._close() + + def _connect(self): + if self.async_mode: + self.ctx = zmq.asyncio.Context() + else: + self.ctx = zmq.Context() + + # Socket to talk to server + self.sock = self.ctx.socket(zmq.REQ) + self.sock.connect(self.url) + if self.timeout > 0: + self.sock.setsockopt(zmq.LINGER, 0) + self.poller = zmq.Poller() + self.poller.register(self.sock, zmq.POLLIN) + + def _close(self): + self.sock.close() + self.ctx.term() + + def send(self, msg): + """ + send s string to responser, and return a string + """ + if self.async_mode: + raise Exception('ZMQ_Requester: in async mode, use asend instead') + b = msg.encode('utf-8') + r = self.send_b(b) + if r is not None: + return r.decode('utf-8') + return None + + def send_b(self, b): + """ + send a bytes and return a bytes + """ + if self.async_mode: + raise Exception('ZMQ_Requester: in async mode, use asend_b instead') + self.sock.send(b) + if self.timeout > 0: + if self.poller.poll(self.timeout * 1000): + return self.sock.recv() + else: + self._close() + self._connect() + return None + else: + return self.sock.recv() + + async def asend_b(self, b): + if not self.async_mode: + raise Exception('ZMQ_Requester: not in async mode, use send_b instead') + await self.sock.send_multipart([b]) + if self.timeout > 0: + if self.poller.poll(self.timeout * 1000): + r = await self.sock.recv_multipart() + return r + else: + self._close() + self._connect() + return None + r = await self.sock.recv_multipart() + return r[0] + + async def asend(self, msg): + if not self.async_mode: + raise Exception('ZMQ_Requester: not in async mode, use send instead') + b = msg.encode('utf-8') + r = await self.asend_b(b) + if r is None: + return None + return r.decode('utf-8') + +class ZmqReplier(object): + def __init__(self, url, handler, async_mode=False): + self.async_mode = async_mode + self.url = url + if not self.async_mode and iscoroutinefunction(handler): + raise('not in async mode, handler can not be a coroutine') + + self.handler = handler + if self.async_mode: + self.ctx = zmq.asyncio.Context() + else: + self.ctx = zmq.Context() + self.sock = self.ctx.socket(zmq.REP) + self.sock.bind(self.url) + self.keep_running = True + + async def async_run(self): + while self.keep_running: + bs = await self.sock.recv_multipart() + b = b[0] + if iscoroutinefunction(self.handler): + rb = await self.handler(b) + else: + rb =self.self.handler(b) + + if isinstance(rb, str): + rb = rb.encode('utf-8') + await self.sock.send_multipart([rb]) + + def run(self): + self.background = Background(self._run) + self.background.daemon = True + self.background.start() + + def _run(self): + while self.keep_running: + b = self.sock.recv() + + rb = self.handler(b) + if isinstance(rb, str): + rb = rb.encode('utf-8') + self.sock.send(rb) + + def stop(self): + self.keep_running = False + self.join()