brige/human
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases 2 Wiki Activity

modify async task queue

Browse Source
This commit is contained in:
brige 2024-10-10 11:31:09 +08:00
parent f9d45e6d44
commit 8173d5356e
1 changed files with 12 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
utils/async_task_queue.py
View File

@ -9,17 +9,17 @@ class AsyncTaskQueue:
self._queue = asyncio.Queue() self._queue = asyncio.Queue()
self._loop = asyncio.new_event_loop() self._loop = asyncio.new_event_loop()
self._thread = threading.Thread(target=self._run_loop) self._thread = threading.Thread(target=self._run_loop)
self.worker_task = None self._worker_task = None
self.loop_running = threading.Event() self._loop_running = threading.Event()
self._thread.start() self._thread.start()
def _run_loop(self): def _run_loop(self):
asyncio.set_event_loop(self._loop) asyncio.set_event_loop(self._loop)
self.loop_running.set() # 设置事件，表明事件循环正在运行 self._loop_running.set()
self._loop.run_forever() # 启动事件循环 self._loop.run_forever()
async def _worker(self): async def _worker(self):
while True: while self._loop_running.is_set():
task = await self._queue.get() task = await self._queue.get()
if task is None: if task is None:
break break
@ -30,13 +30,14 @@ class AsyncTaskQueue:
asyncio.run_coroutine_threadsafe(self._queue.put(coro), self._loop) asyncio.run_coroutine_threadsafe(self._queue.put(coro), self._loop)
def start_worker(self): def start_worker(self):
if not self.worker_task: if not self._worker_task:
self.worker_task = asyncio.run_coroutine_threadsafe(self._worker(), self._loop) self._worker_task = asyncio.run_coroutine_threadsafe(self._worker(), self._loop)
def stop(self): def stop(self):
self._loop_running.clear()
asyncio.run_coroutine_threadsafe(self._queue.put(None), self._loop).result() asyncio.run_coroutine_threadsafe(self._queue.put(None), self._loop).result()
if self.worker_task: if self._worker_task:
self.worker_task.result() self._worker_task.result()
self._loop.call_soon_threadsafe(self._loop.stop) self._loop.call_soon_threadsafe(self._loop.stop)
self._thread.join() # 等待线程结束 self._thread.join()
self._loop.close() # 关闭事件循环 self._loop.close()