You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 

130 lines
3.9 KiB

import asyncio
import threading
import zmq
import zmq.asyncio
from abc import ABC, abstractmethod
# This const defines Thread <-> Actor zmq one-to-one connection
# We create a seperate zmq context, so zqm address 'inproc://xxx' doesn't matter
# It is default address and you may want to use AsyncZmqThread another way
ZMQ_THREAD_ACTOR_ADDR = 'inproc://xxx'
# Inherience order (threading.Thread, ABC) is essential. Otherwise it's a MRO error.
class AsyncZmqThread(threading.Thread, ABC):
"""Class for wrapping zmq socket into a thread with it's own asyncio event loop
"""
def __init__(self,
zmq_context: zmq.asyncio.Context,
zmq_socket_addr: str,
zmq_socket_type = zmq.PAIR
):
super(AsyncZmqThread, self).__init__()
self._zmq_context = zmq_context # you can use it in child classes
self.__zmq_socket_addr = zmq_socket_addr
self.__zmq_socket_type = zmq_socket_type
self.__asyncio_loop = None
self.__zmq_socket = None
async def __message_recv_loop(self):
while True:
text = await self.__zmq_socket.recv_string()
asyncio.ensure_future(self._on_message_to_thread(text))
async def _send_message_from_thread(self, message: str):
await self.__zmq_socket.send_string(message)
@abstractmethod
async def _on_message_to_thread(self, message: str):
"""Override this method to receive messages"""
@abstractmethod
async def _run_thread(self):
"""Override this method to do some async work.
This method uses a separate thread.
You can block yourself here if you don't do any await.
Example:
```
async def _run_thread(self):
i = 0
while True:
await asyncio.sleep(1)
i += 1
await self._send_message_from_thread(f'{self.name}: ping {i}')
```
"""
def run(self):
self.__asyncio_loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.__asyncio_loop)
self.__zmq_socket = self._zmq_context.socket(self.__zmq_socket_type)
self.__zmq_socket.connect(self.__zmq_socket_addr)
asyncio.ensure_future(self.__message_recv_loop())
self.__asyncio_loop.run_until_complete(self._run_thread())
# TODO: implement stop signal handling
class AsyncZmqActor(AsyncZmqThread):
"""Threaded and Async Actor model based on ZMQ inproc communication
override following:
```
async def _run_thread(self)
async def _on_message_to_thread(self, message: str)
```
both methods run in actor's thread
you can call `self._send_message_from_thread('txt')`
to receive it later in `self._recv_message_from_thread()`.
Example:
```
class MyActor(AsyncZmqActor):
async def _run_thread(self):
self.counter = 0
# runs in a different thread
await self._send_message_from_thread('some_txt_message_to_actor')
def async _on_message_to_thread(self, message):
# runs in Thread-actor
self.counter++
asyncZmqActor = MyActor()
asyncZmqActor.start()
```
"""
def __init__(self):
super(AsyncZmqActor, self).__init__(zmq.asyncio.Context(), ZMQ_THREAD_ACTOR_ADDR)
self.__actor_socket = self._zmq_context.socket(zmq.PAIR)
self.__actor_socket.bind(ZMQ_THREAD_ACTOR_ADDR)
async def _put_message_to_thread(self, message: str):
"""It "sends" `message` to thread,
but we can't await it's `AsyncZmqThread._on_message_to_thread()`
so it's "put", not "send"
"""
await self.__actor_socket.send_string(message)
async def _recv_message_from_thread(self) -> str:
"""Returns next message ``'txt'`` from thread sent by
``AsyncZmqActor._send_message_from_thread('txt')``
"""
return await self.__actor_socket.recv_string()
# TODO: implement graceful stopping