108 lines
3.5 KiB
Python
108 lines
3.5 KiB
Python
import asyncio
|
|
import logging
|
|
from typing import Any, Callable
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AsyncTaskQueue:
|
|
def __init__(self, worker_count: int = 3):
|
|
self.queue = asyncio.Queue()
|
|
self.worker_count = worker_count
|
|
self.workers = []
|
|
self._stop_event = asyncio.Event()
|
|
|
|
async def start_workers(self):
|
|
"""Start the worker pool."""
|
|
logger.debug(f"Starting {self.worker_count} workers...")
|
|
for i in range(self.worker_count):
|
|
worker = asyncio.create_task(self.worker_loop(f"Worker-{i+1}"))
|
|
self.workers.append(worker)
|
|
logger.debug("All workers started.")
|
|
|
|
async def stop_workers(self):
|
|
"""Stop the worker pool and wait for all tasks to complete."""
|
|
logger.debug("Stopping workers...")
|
|
self._stop_event.set() # Signal workers to stop
|
|
await self.queue.join() # Wait for all tasks to be processed
|
|
for worker in self.workers:
|
|
worker.cancel()
|
|
await asyncio.gather(*self.workers, return_exceptions=True)
|
|
logger.debug("All workers stopped.")
|
|
|
|
|
|
async def worker_loop(self, name: str):
|
|
"""Worker loop: Process tasks from the queue."""
|
|
logger.debug(f"{name} started.")
|
|
while not self._stop_event.is_set():
|
|
task = None
|
|
try:
|
|
task = await self.queue.get()
|
|
logger.info(f"{name} processing: {task['name']}")
|
|
await task['func'](*task['args'], **task['kwargs'])
|
|
except asyncio.CancelledError:
|
|
logger.debug(f"{name} received cancellation.")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error in {name}: {e}", exc_info=True)
|
|
finally:
|
|
if task is not None:
|
|
self.queue.task_done()
|
|
logger.debug(f"{name} exited.")
|
|
|
|
|
|
|
|
async def enqueue(self, name: str, func: Callable, *args: Any, **kwargs: Any):
|
|
"""Add a task to the queue."""
|
|
logger.debug(f"Enqueuing task: {name}")
|
|
await self.queue.put({'name': name, 'func': func, 'args': args, 'kwargs': kwargs})
|
|
|
|
|
|
|
|
async def run_sync_task_in_thread(func: Callable, *args: Any, **kwargs: Any):
|
|
"""Run a synchronous function in a separate thread.
|
|
|
|
Args:
|
|
func: The synchronous function to run.
|
|
*args: Arguments to pass to the function.
|
|
"""
|
|
await asyncio.to_thread(func, *args, **kwargs)
|
|
|
|
"""
|
|
from asyncTaskQueue import AsyncTaskQueue, run_sync_task_in_thread
|
|
import asyncio
|
|
|
|
# Example async task
|
|
async def your_async_function(name: str, duration: int):
|
|
print(f"{name} started, will sleep for {duration} seconds")
|
|
await asyncio.sleep(duration)
|
|
print(f"{name} finished")
|
|
|
|
# Example sync task
|
|
def your_sync_function(name: str, duration: int):
|
|
print(f"{name} started, will sleep for {duration} seconds")
|
|
import time
|
|
time.sleep(duration)
|
|
print(f"{name} finished")
|
|
|
|
async def main():
|
|
queue = AsyncTaskQueue(worker_count=2)
|
|
await queue.start_workers()
|
|
|
|
# Enqueue async tasks directly
|
|
await queue.enqueue("AsyncTask1", your_async_function, "AsyncTask1", 2)
|
|
await queue.enqueue("AsyncTask2", your_async_function, "AsyncTask2", 1)
|
|
|
|
# Enqueue sync tasks using the wrapper
|
|
await queue.enqueue("SyncTask1", run_sync_task_in_thread, your_sync_function, "SyncTask1", 2)
|
|
|
|
# Let tasks run for a while
|
|
await asyncio.sleep(5)
|
|
|
|
# Stop workers
|
|
await queue.stop_workers()
|
|
|
|
asyncio.run(main())
|
|
|
|
|
|
""" |