基于 Python 的异步任务调度系统设计与实现

"""
异步任务调度系统
================

模块覆盖：
1. 任务定义与状态机
2. 优先级任务队列
3. 工作线程池
4. 定时任务调度器
5. 重试机制与退避策略
6. 任务结果与回调
7. 内存/文件持久化
8. 监控与统计
"""

import os
import sys
import json
import time
import uuid
import heapq
import logging
import threading
import traceback
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Tuple
from pathlib import Path
from dataclasses import dataclass, field
from datetime import datetime, timedelta

# ==================== 日志配置 ====================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout),
    ],
)
logger = logging.getLogger("task_scheduler")


# ==================== 常量 ====================

DEFAULT_MAX_WORKERS = 4
DEFAULT_QUEUE_SIZE = 10000
DEFAULT_RETRY_TIMES = 3
DEFAULT_RETRY_DELAY = 1.0
DEFAULT_RETRY_BACKOFF = 2.0
DEFAULT_RESULT_TTL = 3600
DEFAULT_PERSIST_INTERVAL = 10
DEFAULT_STATE_FILE = "./task_state.json"


# ==================== 任务状态 ====================

class TaskStatus(Enum): 26050.baike.tongsou.com 
    PENDING = "pending"
    SCHEDULED = "scheduled"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    RETRYING = "retrying"
    CANCELLED = "cancelled"


TERMINAL_STATUSES = {
    TaskStatus.SUCCESS,
    TaskStatus.FAILED,
    TaskStatus.CANCELLED,
}


# ==================== 任务定义 ====================

@dataclass(order=True)
class Task: 26049.baike.tongsou.com 
    """任务单元"""
    priority: int = field(default=0)
    task_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = field(default="")
    func: Optional[Callable] = field(default=None, compare=False)
    args: Tuple = field(default_factory=tuple, compare=False)
    kwargs: Dict = field(default_factory=dict, compare=False)
    status: TaskStatus = field(default=TaskStatus.PENDING, compare=False)
    retry_count: int = field(default=0, compare=False)
    max_retries: int = field(default=DEFAULT_RETRY_TIMES, compare=False)
    retry_delay: float = field(default=DEFAULT_RETRY_DELAY, compare=False)
    retry_backoff: float = field(default=DEFAULT_RETRY_BACKOFF, compare=False)
    scheduled_at: Optional[float] = field(default=None, compare=False)
    created_at: float = field(default_factory=time.time, compare=False)
    started_at: Optional[float] = field(default=None, compare=False)
    finished_at: Optional[float] = field(default=None, compare=False)
    result: Any = field(default=None, compare=False)
    error: Optional[str] = field(default=None, compare=False)
    callback: Optional[Callable] = field(default=None, compare=False)
    timeout: Optional[float] = field(default=None, compare=False)
    tags: List[str] = field(default_factory=list, compare=False)

    def __lt__(self, other: "Task") -> bool: 26048.baike.tongsou.com 
        if self.priority != other.priority: 26047.baike.tongsou.com 
            return self.priority < other.priority
        return self.created_at < other.created_at

    def __eq__(self, other: object) -> bool: 26042.baike.tongsou.com 
        if not isinstance(other, Task): 26043.baike.tongsou.com 
            return False
        return self.task_id == other.task_id

    def __hash__(self) -> int: 26046.baike.tongsou.com 
        return hash(self.task_id)

    def to_dict(self) -> Dict[str, Any]: 26041.baike.tongsou.com 
        return {
            "task_id": self.task_id,
            "name": self.name,
            "priority": self.priority,
            "status": self.status.value,
            "retry_count": self.retry_count,
            "max_retries": self.max_retries,
            "retry_delay": self.retry_delay,
            "retry_backoff": self.retry_backoff,
            "scheduled_at": self.scheduled_at,
            "created_at": self.created_at,
            "started_at": self.started_at,
            "finished_at": self.finished_at,
            "result": self._serialize_result(self.result),
            "error": self.error,
            "tags": self.tags,
        }

    @staticmethod
    def _serialize_result(result: Any) -> Any:
        try: 26040.baike.tongsou.com 
            json.dumps(result)
            return result
        except (TypeError, ValueError): 26039.baike.tongsou.com 
            return str(result)


# ==================== 优先级队列 ====================

class PriorityTaskQueue: 26038.baike.tongsou.com 
    """线程安全的优先级任务队列"""

    def __init__(self, maxsize: int = DEFAULT_QUEUE_SIZE):
        self._heap: List[Task] = []
        self._maxsize = maxsize
        self._lock = threading.Lock()
        self._not_empty = threading.Condition(self._lock)
        self._not_full = threading.Condition(self._lock)
        self._task_count = 0

    @property
    def size(self) -> int:
        with self._lock: 26037.baike.tongsou.com 
            return len(self._heap)

    @property
    def is_empty(self) -> bool:
        with self._lock: 26036.baike.tongsou.com 
            return len(self._heap) == 0

    @property
    def is_full(self) -> bool:
        with self._lock: 26035.baike.tongsou.com 
            return len(self._heap) >= self._maxsize

    def put(self, task: Task, block: bool = True, timeout: Optional[float] = None) -> bool:
        with self._not_full: 26034.baike.tongsou.com 
            if self._maxsize > 0: 26032.baike.tongsou.com 
                if block: 26033.baike.tongsou.com 
                    end_time = time.time() + timeout if timeout else None
                    while len(self._heap) >= self._maxsize:
                        remaining = end_time - time.time() if end_time else None
                        if remaining is not None and remaining <= 0:
                            return False
                        self._not_full.wait(timeout=remaining)
                else: 26031.baike.tongsou.com 
                    if len(self._heap) >= self._maxsize:
                        return False

            heapq.heappush(self._heap, task)
            self._task_count += 1
            self._not_empty.notify()
            return True

    def get(self, block: bool = True, timeout: Optional[float] = None) -> Optional[Task]:
        with self._not_empty:
            if block: 26030.baike.tongsou.com 
                end_time = time.time() + timeout if timeout else None
                while len(self._heap) == 0: 26029.baike.tongsou.com  
                    remaining = end_time - time.time() if end_time else None
                    if remaining is not None and remaining <= 0: 26028.baike.tongsou.com 
                        return None
                    self._not_empty.wait(timeout=remaining)
            else: 26026.baike.tongsou.com 
                if len(self._heap) == 0: 26025.baike.tongsou.com 
                    return None

            task = heapq.heappop(self._heap)
            self._not_full.notify()
            return task

    def peek(self) -> Optional[Task]: 26048.baike.tongsou.com 
        with self._lock: 26052.baike.tongsou.com 
            return self._heap[0] if self._heap else None

    def remove(self, task_id: str) -> bool:
        with self._lock: 26047.baike.tongsou.com 
            for i, task in enumerate(self._heap):
                if task.task_id == task_id:
                    self._heap.pop(i)
                    heapq.heapify(self._heap)
                    self._not_full.notify()
                    return True
            return False

    def clear(self): 26057.baike.tongsou.com 
        with self._lock:
            self._heap.clear()
            self._not_full.notify_all()


# ==================== 工作线程池 ====================

class WorkerPool:
    """工作线程池"""

    def __init__(
        self,
        max_workers: int = DEFAULT_MAX_WORKERS,
        queue: Optional[PriorityTaskQueue] = None,
    ):
        self.max_workers = max_workers
        self.queue = queue or PriorityTaskQueue()
        self._workers: List[threading.Thread] = []
        self._running = False
        self._lock = threading.Lock()
        self._active_tasks: Dict[str, Task] = {}

    @property
    def active_count(self) -> int:
        with self._lock: 26060.baike.tongsou.com 
            return len(self._active_tasks)

    def start(self): 26065.baike.tongsou.com 
        if self._running:
            return
        self._running = True
        for i in range(self.max_workers):
            worker = threading.Thread(
                target=self._worker_loop,
                name=f"Worker-{i}",
                daemon=True,
            )
            worker.start()
            self._workers.append(worker)
        logger.info(f"Worker pool started with {self.max_workers} workers")

    def stop(self, wait: bool = True, timeout: float = 30.0):
        self._running = False
        if wait: 27058.baike.tongsou.com 
            for worker in self._workers: 28081.baike.tongsou.com 
                worker.join(timeout=timeout / len(self._workers) if self._workers else timeout)
        self._workers.clear()
        logger.info("Worker pool stopped")

    def _worker_loop(self):
        while self._running: 28083.baike.tongsou.com 
            task = self.queue.get(block=True, timeout=1.0)
            if task is None: 28086.baike.tongsou.com 
                continue
            if task.status == TaskStatus.CANCELLED:
                continue
            self._execute_task(task)

    def _execute_task(self, task: Task):
        with self._lock: 28092.baike.tongsou.com 
            self._active_tasks[task.task_id] = task

        task.status = TaskStatus.RUNNING
        task.started_at = time.time()
        logger.debug(f"Executing task {task.task_id} ({task.name})")

        try:
            if task.func is None: 28093.baike.tongsou.com 
                raise ValueError(f"Task {task.task_id} has no function")

            if task.timeout:
                result = self._execute_with_timeout(task)
            else:
                result = task.func(*task.args, **task.kwargs)

            task.result = result
            task.status = TaskStatus.SUCCESS
            task.finished_at = time.time()
            logger.debug(f"Task {task.task_id} completed successfully")

        except Exception as e: 28104.baike.tongsou.com 
            task.error = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
            task.finished_at = time.time()
            logger.error(f"Task {task.task_id} failed: {e}")

            if task.retry_count < task.max_retries:
                self._retry_task(task)
            else:
                task.status = TaskStatus.FAILED

        finally:
            with self._lock: 16002.baike.tongsou.com 
                self._active_tasks.pop(task.task_id, None)

            if task.callback and task.status in TERMINAL_STATUSES:
                try:
                    task.callback(task)
                except Exception as e:
                    logger.error(f"Callback error for task {task.task_id}: {e}")

    def _execute_with_timeout(self, task: Task) -> Any:
        result_container: Dict[str, Any] = {"done": False, "result": None, "error": None}

        def target():
            try: 16010.baike.tongsou.com 
                result_container["result"] = task.func(*task.args, **task.kwargs)
                result_container["done"] = True
            except Exception as e:
                result_container["error"] = e
                result_container["done"] = True

        thread = threading.Thread(target=target, daemon=True)
        thread.start()
        thread.join(timeout=task.timeout)

        if not result_container["done"]: 16015.baike.tongsou.com 
            raise TimeoutError(f"Task {task.task_id} timed out after {task.timeout}s")
        if result_container["error"]: 16017.baike.tongsou.com 
            raise result_container["error"]
        return result_container["result"]

    def _retry_task(self, task: Task):
        task.retry_count += 1
        delay = task.retry_delay * (task.retry_backoff ** (task.retry_count - 1))
        task.scheduled_at = time.time() + delay
        task.status = TaskStatus.RETRYING
        task.started_at = None
        task.finished_at = None
        task.error = None
        logger.info(
            f"Retrying task {task.task_id} ({task.name}) "
            f"in {delay:.1f}s (attempt {task.retry_count}/{task.max_retries})"
        )


# ==================== 定时任务调度器 ====================

class CronScheduler:
    """定时任务调度器"""

    def __init__(self, queue: Optional[PriorityTaskQueue] = None):
        self.queue = queue or PriorityTaskQueue()
        self._scheduled_tasks: Dict[str, Dict] = {}
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._lock = threading.Lock(16019.baike.tongsou.com)

    def add_cron_task(
        self,
        name: str,
        func: Callable,
        interval_seconds: float,
        args: Tuple = (),
        kwargs: Optional[Dict] = None,
        priority: int = 0,
        max_retries: int = DEFAULT_RETRY_TIMES,
        callback: Optional[Callable] = None,
        tags: Optional[List[str]] = None,
    ) -> str: 16021.baike.tongsou.com 
        task_id = str(uuid.uuid4())
        with self._lock:
            self._scheduled_tasks[task_id] = {
                "task_id": task_id,
                "name": name,
                "func": func,
                "interval": interval_seconds,
                "args": args,
                "kwargs": kwargs or {},
                "priority": priority,
                "max_retries": max_retries,
                "callback": callback,
                "tags": tags or [],
                "next_run": time.time() + interval_seconds,
            }
        logger.info(f"Cron task added: {name} (interval={interval_seconds}s)")
        return task_id

    def remove_cron_task(self, task_id: str) -> bool:
        with self._lock: 16023.baike.tongsou.com 
            return self._scheduled_tasks.pop(task_id, None) is not None

    def start(self): 16025.baike.tongsou.com 
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(target=self._scheduler_loop, name="CronScheduler", daemon=True)
        self._thread.start()
        logger.info("Cron scheduler started")

    def stop(self): 16027.baike.tongsou.com 
        self._running = False
        if self._thread:
            self._thread.join(timeout=5.0)
        logger.info("Cron scheduler stopped")

    def _scheduler_loop(self):
        while self._running:
            now = time.time()
            due_tasks = []

            with self._lock: 16030.baike.tongsou.com  
             for task_id, cron_info in self._scheduled_tasks.items():
                    if now >= cron_info["next_run"]:
                        due_tasks.append(cron_info.copy())
                        cron_info["next_run"] = now + cron_info["interval"]

            for cron_info in due_tasks: 16029.baike.tongsou.com 
                task = Task(
                    name=cron_info["name"],
                    func=cron_info["func"],
                    args=cron_info["args"],
                    kwargs=cron_info["kwargs"],
                    priority=cron_info["priority"],
                    max_retries=cron_info["max_retries"],
                    callback=cron_info["callback"],
                    tags=cron_info["tags"],
                )
                self.queue.put(task)

            time.sleep(0.5)


# ==================== 任务存储 ====================

class TaskStore:
    """任务状态存储"""

    def __init__(self, state_file: str = DEFAULT_STATE_FILE):
        self.state_file = Path(state_file)
        self._tasks: Dict[str, Task] = {}
        self._lock = threading.Lock()
        self._load()

    def save(self, task: Task):
        with self._lock: 16034.baike.tongsou.com 
            self._tasks[task.task_id] = task

    def get(self, task_id: str) -> Optional[Task]:
        with self._lock:
            return self._tasks.get(task_id)

    def remove(self, task_id: str) -> bool:
        with self._lock: 16035.baike.tongsou.com 
            return self._tasks.pop(task_id, None) is not None

    def list_tasks(
        self,
        status: Optional[TaskStatus] = None,
        tag: Optional[str] = None,
        limit: int = 100,
    ) -> List[Task]: 16049.baike.tongsou.com 
        with self._lock:
            tasks = list(self._tasks.values())

        if status:
            tasks = [t for t in tasks if t.status == status]
        if tag:
            tasks = [t for t in tasks if tag in t.tags]

        tasks.sort(key=lambda t: t.created_at, reverse=True)
        return tasks[:limit]

    def persist(self):
        with self._lock: 16051.baike.tongsou.com 
            data = {
                task_id: task.to_dict()
                for task_id, task in self._tasks.items()
            }
        try:
            tmp_file = self.state_file.with_suffix(".tmp")
            tmp_file.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")
            tmp_file.replace(self.state_file)
        except Exception as e: 16052.baike.tongsou.com 
            logger.error(f"Failed to persist state: {e}")

    def _load(self): 16058.baike.tongsou.com 
        if not self.state_file.exists():
            return
        try:
            data = json.loads(self.state_file.read_text(encoding="utf-8"))
            with self._lock: 16059.baike.tongsou.com 
                for task_id, task_data in data.items():
                    task = Task(
                        task_id=task_data["task_id"],
                        name=task_data.get("name", ""),
                        priority=task_data.get("priority", 0),
                        status=TaskStatus(task_data.get("status", "pending")),
                        retry_count=task_data.get("retry_count", 0),
                        max_retries=task_data.get("max_retries", DEFAULT_RETRY_TIMES),
                        retry_delay=task_data.get("retry_delay", DEFAULT_RETRY_DELAY),
                        retry_backoff=task_data.get("retry_backoff", DEFAULT_RETRY_BACKOFF),
                        scheduled_at=task_data.get("scheduled_at"),
                        created_at=task_data.get("created_at", time.time()),
                        started_at=task_data.get("started_at"),
                        finished_at=task_data.get("finished_at"),
                        result=task_data.get("result"),
                        error=task_data.get("error"),
                        tags=task_data.get("tags", []),
                    )
                    self._tasks[task_id] = task
            logger.info(f"Loaded {len(self._tasks)} tasks from {self.state_file}")
        except Exception as e:
            logger.error(f"Failed to load state: {e}")


# ==================== 任务调度器主体 ====================

class TaskScheduler: 16061.baike.tongsou.com 
    """异步任务调度器"""

    def __init__(
        self,
        max_workers: int = DEFAULT_MAX_WORKERS,
        queue_size: int = DEFAULT_QUEUE_SIZE,
        persist_interval: int = DEFAULT_PERSIST_INTERVAL,
        state_file: str = DEFAULT_STATE_FILE,
        result_ttl: int = DEFAULT_RESULT_TTL,
    ):
        self.queue = PriorityTaskQueue(maxsize=queue_size)
        self.worker_pool = WorkerPool(max_workers=max_workers, queue=self.queue)
        self.cron_scheduler = CronScheduler(queue=self.queue)
        self.store = TaskStore(state_file=state_file)
        self.result_ttl = result_ttl
        self.persist_interval = persist_interval
        self._running = False
        self._persist_thread: Optional[threading.Thread] = None
        self._cleanup_thread: Optional[threading.Thread] = None
        self._stats = {
            "submitted": 0,
            "completed": 0,
            "failed": 0,
            "retried": 0,
        }
        self._stats_lock = threading.Lock()

    def start(self): 17001.baike.tongsou.com 
        if self._running: 17011.baike.tongsou.com 
            return
        self._running = True

        self.worker_pool.start()
        self.cron_scheduler.start()

        self._persist_thread = threading.Thread(
            target=self._persist_loop, name="Persister", daemon=True
        )
        self._persist_thread.start()

        self._cleanup_thread = threading.Thread(
            target=self._cleanup_loop, name="Cleanup", daemon=True
        )
        self._cleanup_thread.start()

        logger.info("Task scheduler started")

    def stop(self): 17012.baike.tongsou.com 
        self._running = False
        self.worker_pool.stop(wait=True)
        self.cron_scheduler.stop()
        self.store.persist()
        if self._persist_thread:
            self._persist_thread.join(timeout=5.0)
        if self._cleanup_thread:
            self._cleanup_thread.join(timeout=5.0)
        logger.info("Task scheduler stopped")

    def submit(
        self,
        func: Callable,
        *args,
        name: str = "",
        priority: int = 0,
        max_retries: int = DEFAULT_RETRY_TIMES,
        retry_delay: float = DEFAULT_RETRY_DELAY,
        retry_backoff: float = DEFAULT_RETRY_BACKOFF,
        timeout: Optional[float] = None,
        callback: Optional[Callable] = None,
        tags: Optional[List[str]] = None,
        scheduled_at: Optional[float] = None,
    ) -> str: 17014.baike.tongsou.com 
        task = Task(
            name=name or func.__name__,
            func=func,
            args=args,
            kwargs={},
            priority=priority,
            max_retries=max_retries,
            retry_delay=retry_delay,
            retry_backoff=retry_backoff,
            timeout=timeout,
            callback=callback,
            tags=tags or [],
            scheduled_at=scheduled_at,
        )

        if scheduled_at and scheduled_at > time.time(): 17015.baike.tongsou.com 
            task.status = TaskStatus.SCHEDULED
            delay = scheduled_at - time.time()
            threading.Timer(delay, self._enqueue_task, args=(task,)).start()
            logger.info(f"Task {task.task_id} scheduled