86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
import logging
|
|
import queue
|
|
import threading
|
|
from typing import Callable, Optional
|
|
|
|
from .config import settings
|
|
|
|
|
|
class SimpleTaskQueue:
|
|
def __init__(self, name: str, maxsize: int = 0):
|
|
self._name = name
|
|
self._queue: "queue.Queue[tuple[Optional[str], Callable, tuple, dict]]" = queue.Queue(maxsize=maxsize)
|
|
self._threads: list[threading.Thread] = []
|
|
self._started = False
|
|
self._lock = threading.Lock()
|
|
self._inflight: set[str] = set()
|
|
self._inflight_lock = threading.Lock()
|
|
|
|
def start(self, workers: int):
|
|
import logging
|
|
logger = logging.getLogger(self._name)
|
|
with self._lock:
|
|
if self._started:
|
|
logger.warning(f"{self._name} 队列已经启动，跳过")
|
|
return
|
|
self._started = True
|
|
logger.info(f"启动 {workers} 个 {self._name} worker 线程")
|
|
for idx in range(max(1, workers)):
|
|
thread = threading.Thread(
|
|
target=self._worker,
|
|
name=f"{self._name}-worker-{idx + 1}",
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
self._threads.append(thread)
|
|
logger.info(f"Worker 线程 {self._name}-worker-{idx + 1} 已启动")
|
|
logger.info(f"{self._name} 队列启动完成，共 {len(self._threads)} 个 worker")
|
|
|
|
def enqueue(self, func: Callable, *args, **kwargs) -> None:
|
|
self._queue.put((None, func, args, kwargs))
|
|
|
|
def enqueue_unique(self, key: str, func: Callable, *args, **kwargs) -> bool:
|
|
if not key:
|
|
self.enqueue(func, *args, **kwargs)
|
|
return True
|
|
with self._inflight_lock:
|
|
if key in self._inflight:
|
|
return False
|
|
self._inflight.add(key)
|
|
self._queue.put((key, func, args, kwargs))
|
|
return True
|
|
|
|
def _worker(self):
|
|
import logging
|
|
logger = logging.getLogger(self._name)
|
|
logger.info(f"Worker 线程 {threading.current_thread().name} 开始运行")
|
|
while True:
|
|
try:
|
|
logger.debug(f"Worker {threading.current_thread().name} 等待任务...")
|
|
key, func, args, kwargs = self._queue.get()
|
|
logger.info(f"Worker {threading.current_thread().name} 获取到任务: {func.__name__}, key={key}")
|
|
try:
|
|
func(*args, **kwargs)
|
|
logger.info(f"Worker {threading.current_thread().name} 任务完成: {func.__name__}")
|
|
except Exception as e:
|
|
logger.exception(f"{self._name} worker 任务失败: {e}")
|
|
finally:
|
|
if key:
|
|
with self._inflight_lock:
|
|
self._inflight.discard(key)
|
|
self._queue.task_done()
|
|
except Exception as e:
|
|
logger.exception(f"Worker {threading.current_thread().name} 发生异常: {e}")
|
|
|
|
|
|
sing_task_queue = SimpleTaskQueue("sing")
|
|
|
|
|
|
def start_sing_workers():
|
|
import logging
|
|
logger = logging.getLogger("sing")
|
|
workers = max(1, settings.SING_MERGE_MAX_CONCURRENCY or 1)
|
|
logger.info(f"启动 {workers} 个唱歌任务 worker 线程")
|
|
sing_task_queue.start(workers)
|
|
logger.info("唱歌任务 worker 线程启动完成")
|