apppublic/appPublic/zmq_reqrep.py
2023-02-23 15:17:24 +08:00

136 lines
3.3 KiB
Python
Executable File

# 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()