39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
|
from concurrent.futures import ThreadPoolExecutor
|
||
|
|
from itertools import count
|
||
|
|
from threading import Lock
|
||
|
|
from typing import Any, Callable, Dict
|
||
|
|
import traceback
|
||
|
|
|
||
|
|
|
||
|
|
class KeyedBackgroundTaskRunner:
|
||
|
|
"""按 key 去重的后台任务执行器，避免同类长任务重复堆积。"""
|
||
|
|
|
||
|
|
def __init__(self, max_workers: int, thread_name_prefix: str):
|
||
|
|
self._executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix=thread_name_prefix)
|
||
|
|
self._lock = Lock()
|
||
|
|
self._seq = count(1)
|
||
|
|
self._tasks: Dict[str, Dict[str, Any]] = {}
|
||
|
|
|
||
|
|
def submit(self, key: str, func: Callable[..., Any], *args: Any, **kwargs: Any) -> bool:
|
||
|
|
with self._lock:
|
||
|
|
current = self._tasks.get(key)
|
||
|
|
if current and not current["future"].done():
|
||
|
|
return False
|
||
|
|
|
||
|
|
token = next(self._seq)
|
||
|
|
future = self._executor.submit(self._run_task, key, token, func, *args, **kwargs)
|
||
|
|
self._tasks[key] = {"token": token, "future": future}
|
||
|
|
return True
|
||
|
|
|
||
|
|
def _run_task(self, key: str, token: int, func: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
|
||
|
|
try:
|
||
|
|
func(*args, **kwargs)
|
||
|
|
except Exception:
|
||
|
|
print(f"[BackgroundTaskRunner] Task failed, key={key}")
|
||
|
|
traceback.print_exc()
|
||
|
|
finally:
|
||
|
|
with self._lock:
|
||
|
|
current = self._tasks.get(key)
|
||
|
|
if current and current["token"] == token:
|
||
|
|
self._tasks.pop(key, None)
|