human/utils/async_task_queue.py
2024-10-10 19:01:13 +08:00

49 lines
1.4 KiB
Python

#encoding = utf8
import asyncio
import threading
class AsyncTaskQueue:
def __init__(self):
self._queue = asyncio.Queue()
self._loop = asyncio.new_event_loop()
self._thread = threading.Thread(target=self._run_loop)
self._worker_task = None
self._loop_running = threading.Event()
self._thread.start()
def _run_loop(self):
print('_run_loop')
self._loop_running.set()
asyncio.set_event_loop(self._loop)
self._loop.run_forever()
async def _worker(self):
print('_worker')
while self._loop_running.is_set():
task = await self._queue.get()
if task is None:
break
print('run task')
await task
self._queue.task_done()
print('_worker finish')
def add_task(self, coro):
print('add_task')
asyncio.run_coroutine_threadsafe(self._queue.put(coro), self._loop)
def start_worker(self):
if not self._worker_task:
self._worker_task = asyncio.run_coroutine_threadsafe(self._worker(), self._loop)
def stop(self):
self._loop_running.clear()
asyncio.run_coroutine_threadsafe(self._queue.put(None), self._loop).result()
if self._worker_task:
self._worker_task.result()
self._loop.call_soon_threadsafe(self._loop.stop)
self._thread.join()
self._loop.close()