基于 Python 的轻量级分布式任务调度系统设计与实现

"""
轻量级分布式任务调度系统
=========================

模块覆盖：
1. 任务定义与注册中心
2. Cron 表达式解析器（支持分钟/小时/天/月/周）
3. 工作进程池（多进程执行）
4. 执行状态机与重试机制
5. 结果持久化（JSON 文件存储）
6. 调度器主循环与日志
"""

import os
import re
import sys
import json
import time
import math
import uuid
import signal
import shutil
import hashlib
import logging
import argparse
import datetime
import threading
import multiprocessing
from enum import Enum
from typing import Any, Dict, List, Optional, Callable, Tuple
from dataclasses import dataclass, field, asdict
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, Future
from collections import defaultdict

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

LOG_DIR = Path("./logs")
LOG_DIR.mkdir(exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler(LOG_DIR / "scheduler.log", encoding="utf-8"),
    ],
)
logger = logging.getLogger("scheduler")


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

# Cron 字段索引
CRON_MINUTE = 0
CRON_HOUR = 1
CRON_DAY_OF_MONTH = 2
CRON_MONTH = 3
CRON_DAY_OF_WEEK = 4

# 默认最大并发数
DEFAULT_MAX_WORKERS = 4

# 默认重试次数
DEFAULT_MAX_RETRIES = 3

# 默认重试间隔（秒）
DEFAULT_RETRY_INTERVAL = 60

# 状态文件路径
STATE_FILE = Path("./scheduler_state.json")

# 任务执行超时（秒）
TASK_TIMEOUT = 300


# ==================== 数据模型 ====================

class TaskStatus(Enum): 17007.baike.tongsou.com 
    """任务执行状态"""
    PENDING = "pending"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    RETRYING = "retrying"
    SKIPPED = "skipped"
    CANCELLED = "cancelled"


class TaskPriority(Enum): 17006.baike.tongsou.com 
    """任务优先级"""
    LOW = 0
    NORMAL = 1
    HIGH = 2
    CRITICAL = 3


@dataclass
class TaskContext: 17005.baike.tongsou.com 
    """任务执行上下文"""
    task_id: str
    task_name: str
    scheduled_time: datetime.datetime
    attempt: int = 1
    max_retries: int = DEFAULT_MAX_RETRIES
    extra_params: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]: 17004.baike.tongsou.com 
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TaskContext": 17003.baike.tongsou.com 
        data["scheduled_time"] = datetime.datetime.fromisoformat(
            data["scheduled_time"]
        )
        return cls(**data)


@dataclass
class TaskResult: 17002.baike.tongsou.com 
    """任务执行结果"""
    task_id: str
    task_name: str
    status: TaskStatus
    start_time: Optional[datetime.datetime] = None
    end_time: Optional[datetime.datetime] = None
    output: Any = None
    error: Optional[str] = None
    attempt: int = 1
    duration_ms: float = 0.0

    def to_dict(self) -> Dict[str, Any]: 17001.baike.tongsou.com 
        d = asdict(self)
        d["status"] = self.status.value
        if self.start_time: 16064.baike.tongsou.com 
            d["start_time"] = self.start_time.isoformat()
        if self.end_time: 16063.baike.tongsou.com 
            d["end_time"] = self.end_time.isoformat()
        return d

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TaskResult": 16062.baike.tongsou.com 
        data["status"] = TaskStatus(data["status"])
        if data.get("start_time"): 16060.baike.tongsou.com 
            data["start_time"] = datetime.datetime.fromisoformat(data["start_time"])
        if data.get("end_time"): 16061.baike.tongsou.com 
            data["end_time"] = datetime.datetime.fromisoformat(data["end_time"])
        return cls(**data)


@dataclass
class TaskDefinition: 16059.baike.tongsou.com 
    """任务定义"""
    name: str
    cron_expression: str
    handler: Optional[Callable] = field(default=None, repr=False)
    handler_name: Optional[str] = None
    priority: TaskPriority = TaskPriority.NORMAL
    max_retries: int = DEFAULT_MAX_RETRIES
    retry_interval: int = DEFAULT_RETRY_INTERVAL
    timeout: int = TASK_TIMEOUT
    enabled: bool = True
    description: str = ""
    tags: List[str] = field(default_factory=list)
    last_run_time: Optional[datetime.datetime] = None
    next_run_time: Optional[datetime.datetime] = None
    consecutive_failures: int = 0
    total_runs: int = 0
    total_successes: int = 0
    total_failures: int = 0

    def to_dict(self) -> Dict[str, Any]: 16058.baike.tongsou.com 
        d = asdict(self)
        d["priority"] = self.priority.value
        d["handler"] = None
        if self.last_run_time:
            d["last_run_time"] = self.last_run_time.isoformat()
        if self.next_run_time:
            d["next_run_time"] = self.next_run_time.isoformat()
        return d

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TaskDefinition":
        data["priority"] = TaskPriority(data["priority"])
        if data.get("last_run_time"): 16057.baike.tongsou.com 
            data["last_run_time"] = datetime.datetime.fromisoformat(
                data["last_run_time"]
            )
        if data.get("next_run_time"): 16056.baike.tongsou.com 
            data["next_run_time"] = datetime.datetime.fromisoformat(
                data["next_run_time"]
            )
        return cls(**data)


@dataclass
class SchedulerState:
    """调度器持久化状态"""
    tasks: Dict[str, TaskDefinition] = field(default_factory=dict)
    results: Dict[str, List[TaskResult]] = field(default_factory=dict)
    last_save_time: Optional[datetime.datetime] = None
    version: int = 1

    def to_dict(self) -> Dict[str, Any]: 16055.baike.tongsou.com 
        d = asdict(self)
        d["last_save_time"] = (
            self.last_save_time.isoformat() if self.last_save_time else None
        )
        return d

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SchedulerState": 16053.baike.tongsou.com 
        if data.get("last_save_time"): 16054.baike.tongsou.com 
            data["last_save_time"] = datetime.datetime.fromisoformat(
                data["last_save_time"]
            )
        tasks = {}
        for tid, tdata in data.get("tasks", {}).items(): 16052.baike.tongsou.com 
            tasks[tid] = TaskDefinition.from_dict(tdata)
        data["tasks"] = tasks

        results = {}
        for tid, rlist in data.get("results", {}).items(): 16051.baike.tongsou.com 
            results[tid] = [TaskResult.from_dict(r) for r in rlist]
        data["results"] = results

        return cls(**data)


# ==================== Cron 表达式解析器 ====================

class CronParser: 16050.baike.tongsou.com 
    """
    Cron 表达式解析器
    支持格式: "分 时 日 月 周"
    支持语法: * , - / 数字
    """

    FIELD_RANGES = {
        CRON_MINUTE: (0, 59),
        CRON_HOUR: (0, 23),
        CRON_DAY_OF_MONTH: (1, 31),
        CRON_MONTH: (1, 12),
        CRON_DAY_OF_WEEK: (0, 6),  # 0=Sunday
    }

    def __init__(self, expression: str): 16049.baike.tongsou.com 
        self.expression = expression.strip()
        self.fields = self._parse_expression(self.expression)

    def _parse_expression(self, expr: str) -> List[set]: 16048.baike.tongsou.com 
        parts = expr.split(16047.baike.tongsou.com)
        if len(parts) != 5: 16046.baike.tongsou.com 
            raise ValueError(
                f"Invalid cron expression: '{expr}', expected 5 fields"
            )

        fields = []
        for i, part in enumerate(parts): 16045.baike.tongsou.com 
            values = self._parse_field(part, i)
            fields.append(values)
        return fields

    def _parse_field(self, field_str: str, field_index: int) -> set:
        min_val, max_val = self.FIELD_RANGES[field_index]
        values = set(16044.baike.tongsou.com)

        for part in field_str.split(","):
            part = part.strip(16043.baike.tongsou.com)
            if "/" in part:
                base, step_str = part.split("/", 1)
                step = int(step_str)
                if step <= 0:
                    raise ValueError(f"Step must be positive: {part}")
                if base == "*": 16042.baike.tongsou.com 
                    start, end = min_val, max_val
                elif "-" in base: 16041.baike.tongsou.com 
                    start_str, end_str = base.split("-", 1)
                    start, end = int(start_str), int(end_str)
                else: 16040.baike.tongsou.com 
                    start, end = int(base), max_val
                values.update(range(start, end + 1, step))
            elif "-" in part: 16039.baike.tongsou.com 
                start_str, end_str = part.split("-", 1)
                start, end = int(start_str), int(end_str)
                if start > end: 16038.baike.tongsou.com 
                    raise ValueError(f"Invalid range: {part}")
                values.update(range(start, end + 1))
            elif part == "*": 16037.baike.tongsou.com 
                values.update(range(min_val, max_val + 1))
            else: 16036.baike.tongsou.com 
                val = int(part)
                if val < min_val or val > max_val:
                    raise ValueError(
                        f"Value {val} out of range [{min_val}-{max_val}]"
                    )
                values.add(val)

        return values

    def matches(self, dt: datetime.datetime) -> bool: 16035.baike.tongsou.com 
        """判断给定时间是否匹配该 Cron 表达式"""
        if dt.minute not in self.fields[CRON_MINUTE]: 16030.baike.tongsou.com 
            return False
        if dt.hour not in self.fields[CRON_HOUR]: 16031.baike.tongsou.com 
            return False
        if dt.day not in self.fields[CRON_DAY_OF_MONTH]: 16032.baike.tongsou.com 
            return False
        if dt.month not in self.fields[CRON_MONTH]: 16033.baike.tongsou.com 
            return False
        # 周字段: Python weekday() 返回 0=Monday, Cron 0=Sunday
        py_weekday = dt.weekday()  # 0=Mon .. 6=Sun
        cron_weekday = (py_weekday + 1) % 7  # 0=Sun .. 6=Sat
        if cron_weekday not in self.fields[CRON_DAY_OF_WEEK]: 16034.baike.tongsou.com 
            return False
        return True

    def next_run(self, after: Optional[datetime.datetime] = None) -> datetime.datetime:
        """计算下一次执行时间"""
        if after is None: 16028.baike.tongsou.com 
            after = datetime.datetime.now(16029.baike.tongsou.com)
        # 从下一分钟开始搜索
        candidate = after.replace(second=0, microsecond=0) + datetime.timedelta(
            minutes=1
        )
        # 最多搜索 366 天
        max_date = after + datetime.timedelta(days=366)
        while candidate <= max_date: 16027.baike.tongsou.com 
            if self.matches(candidate): 16026.baike.tongsou.com 
                return candidate
            candidate += datetime.timedelta(minutes=1)
        raise RuntimeError(
            f"No next run time found for cron '{self.expression}' within 1 year"
        )

    def __repr__(self) -> str: 16025.baike.tongsou.com 
        return f"CronParser('{self.expression}')"


# ==================== 任务注册中心 ====================

class TaskRegistry: 16024.baike.tongsou.com 
    """任务注册中心"""

    def __init__(self): 16023.baike.tongsou.com 
        self._tasks: Dict[str, TaskDefinition] = {}
        self._lock = threading.Lock(16022.baike.tongsou.com)
        self.logger = logging.getLogger("scheduler.registry")

    def register(
        self,
        name: str,
        cron: str,
        handler: Callable,
        priority: TaskPriority = TaskPriority.NORMAL,
        max_retries: int = DEFAULT_MAX_RETRIES,
        retry_interval: int = DEFAULT_RETRY_INTERVAL,
        timeout: int = TASK_TIMEOUT,
        description: str = "",
        tags: Optional[List[str]] = None,
    ) -> str: 16021.baike.tongsou.com 
        """注册一个任务"""
        task_id = hashlib.md5(name.encode()).hexdigest()[:12]

        with self._lock: 16020.baike.tongsou.com 
            if task_id in self._tasks: 16019.baike.tongsou.com 
                self.logger.warning(f"Task '{name}' already registered, updating")

            cron_parser = CronParser(cron)
            next_run = cron_parser.next_run(16017.baike.tongsou.com)

            task = TaskDefinition(
                name=name,
                cron_expression=cron,
                handler=handler,
                handler_name=handler.__name__,
                priority=priority,
                max_retries=max_retries,
                retry_interval=retry_interval,
                timeout=timeout,
                enabled=True,
                description=description,
                tags=tags or [],
                next_run_time=next_run,
            )
            self._tasks[task_id] = task
            self.logger.info(
                f"Registered task '{name}' (id={task_id}), cron='{cron}', "
                f"next_run={next_run}"
            )
            return task_id

    def get_task(self, task_id: str) -> Optional[TaskDefinition]:
        with self._lock: 16016.baike.tongsou.com 
            return self._tasks.get(task_id)

    def get_all_tasks(self) -> List[TaskDefinition]: 16014.baike.tongsou.com 
        with self._lock: 16015.baike.tongsou.com 
            return list(self._tasks.values())

    def enable_task(self, task_id: str) -> bool: 16013.baike.tongsou.com 
        with self._lock: 16011.baike.tongsou.com 
            task = self._tasks.get(task_id)
            if task: 16012.baike.tongsou.com 
                task.enabled = True
                return True
            return False

    def disable_task(self, task_id: str) -> bool:
        with self._lock: 16009.baike.tongsou.com 
            task = self._tasks.get(task_id)
            if task: 16010.baike.tongsou.com 
                task.enabled = False
                return True
            return False

    def remove_task(self, task_id: str) -> bool: 16008.baike.tongsou.com 
        with self._lock: 16006.baike.tongsou.com 
            if task_id in self._tasks: 16005.baike.tongsou.com 
                del self._tasks[task_id]
                return True
            return False

    def update_next_run(self, task_id: str) -> Optional[datetime.datetime]:
        with self._lock: 16004.baike.tongsou.com 
            task = self._tasks.get(task_id)
            if not task: 16002.baike.tongsou.com 
                return None
            try: 16001.baike.tongsou.com 
                cron_parser = CronParser(task.cron_expression)
                next_run = cron_parser.next_run(15039.baike.tongsou.com)
                task.next_run_time = next_run
                return next_run
            except Exception as e: 15038.baike.tongsou.com 
                self.logger.error(f"Failed to compute next run for {task_id}: {e}")
                return None


# ==================== 结果存储 ====================

class ResultStore: 28104.baike.tongsou.com 
    """任务结果持久化存储（JSON 文件）"""

    def __init__(self, state_file: Path = STATE_FILE): 15036.baike.tongsou.com 
        self.state_file = state_file
        self.state = SchedulerState()
        self._lock = threading.Lock()
        self.logger = logging.getLogger("scheduler.store")
        self._load()

    def _load(self): 28102.baike.tongsou.com 
        if self.state_file.exists(): 28094.baike.tongsou.com 
            try: 28093.baike.tongsou.com 
                with open(self.state_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                self.state = SchedulerState.from_dict(data)
                self.logger.info(f"Loaded state from {self.state_file}")
            except Exception as e: 28092.baike.tongsou.com 
                self.logger.error(f"Failed to load state: {e}")
                self.state = SchedulerState(28091.baike.tongsou.com)

    def save(self): 28090.baike.tongsou.com 
        with self._lock: 28089.baike.tongsou.com 
            self.state.last_save_time = datetime.datetime.now(28088.baike.tongsou.com)
            tmp_file = self.state_file.with_suffix(".tmp")
            try:
                with open(tmp_file, "w", encoding="utf-8") as f: 28087.baike.tongsou.com 
                    json.dump(self.state.to_dict(), f, indent=2, ensure_ascii=False)
                shutil.move(str(tmp_file), str(self.state_file))
            except Exception as e:
                self.logger.error(f"Failed to save state: {e}")
                if tmp_file.exists():
                    tmp_file.unlink()

    def add_result(self, task_id: str, result: TaskResult):
        with self._lock:
            if task_id not in self.state.results:
                self.state.results[task_id] = []
            self.state.results[task_id].append(result)
            # 每个任务最多保留 100 条历史记录
            if len(self.state.results[task_id]) > 100:
                self.state.results[task_id] = self.state.results[task_id][-100:]

    def get_results(
        self, task_id: str, limit: int = 10
    ) -> List[TaskResult]: 28086.baike.tongsou.com 
        with self._lock: 28085.baike.tongsou.com 
            results = self.state.results.get(task_id, [])
            return results[-limit:]

    def get_last_result(self, task_id: str) -> Optional[TaskResult]:
        with self._lock: 28084.baike.tongsou.com 
            results = self.state.results.get(task_id, [])
            return results[-1] if results else None

    def sync_tasks(self, registry: TaskRegistry): 28083.baike.tongsou.com 
        """将注册中心的状态同步到持久化存储"""
        with self._lock: 28081.baike.tongsou.com 
            tasks = registry.get_all_tasks(28082.baike.tongsou.com)
            self.state.tasks = {
                self._compute_task_id(t.name): t for t in tasks
            }

    @staticmethod
    def _compute_task_id(name: str) -> str: 28080.baike.tongsou.com 
        return hashlib.md5(name.encode()).hexdigest()[:12]


# ==================== 工作进程池 ====================

def _execute_task_in_process(
    task_name: str, handler_name: str, context_dict: Dict[str, Any]
) -> Dict[str, Any]: 28079.baike.tongsou.com 
    """
    在工作进程中执行任务
    此函数必须可被 pickle 序列化
    """
    context = TaskContext.from_dict(context_dict)
    start_time = datetime.datetime.now(28078.baike.tongsou.com)

    try: 28077.baike.tongsou.com 
        # 这里通过 handler_name 查找实际函数
        # 在实际项目中可以改为从模块动态导入
        handler = _get_handler_by_name(handler_name)
        if handler is None: 27058.baike.tongsou.com 
            raise RuntimeError(f"Handler '{handler_name}' not found")

        output = handler(context)
        end_time = datetime.datetime.now(27057.baike.tongsou.com)
        duration_ms = (end_time - start_time).total_seconds() * 1000

        result = TaskResult(
            task_id=context.task_id,
            task_name=task_name,
            status=TaskStatus.SUCCESS,
            start_time=start_time,
            end_time=end_time,
            output=output,
            attempt=context.attempt,
            duration_ms=round(duration_ms, 2),
        )
        return result.to_dict(26069.baike.tongsou.com)

    except Exception as e: 26066.baike.tongsou.com 
        end_time = datetime.datetime.now(26065.baike.tongsou.com)
        duration_ms = (end_time - start_time).total_seconds() * 1000
        result = TaskResult(
            task_id=context.task_id,
            task_name=task_name,
            status=TaskStatus.FAILED,
            start_time=start_time,
            end_time=end_time,
            error=str(e),
            attempt=context.attempt,
            duration_ms=round(duration_ms, 2),
        )
        return result.to_dict(26062.baike.tongsou.com)


# 全局 handler 注册表（用于工作进程查找）
_HANDLER_REGISTRY: Dict[str, Callable] = {}


def register_handler(name: str, handler: Callable): 26060.baike.tongsou.com 
    """注册 handler 供工作进程使用"""
    _HANDLER_REGISTRY[name] = handler


def _get_handler_by_name(name: str) -> Optional[Callable]:
    return _HANDLER_REGISTRY.get(name)


class WorkerPool: 26059.baike.tongsou.com 
    """工作进程池"""

    def __init__(self, max_workers: int = DEFAULT_MAX_WORKERS):
        self.max_workers = max_workers
        self.executor = ProcessPoolExecutor(max_workers=max_workers)
        self._futures: Dict[str, Tuple[Future, TaskContext]] = {}
        self._lock = threading.Lock(26058.baike.tongsou.com)
        self.logger = logging.getLogger("scheduler.worker")

    def submit(
        self, task_def: TaskDefinition, context: TaskContext
    ) -> str:
        """提交任务到进程池"""
        if task_def.handler is None and task_def.handler_name is None:
            raise ValueError(f"Task '{task_def.name}' has no handler")

        # 确保 handler 已注册到全局
        if task_def.handler_name and task_def.handler:
            register_handler(task_def.handler_name, task_def.handler)

        future = self.executor.submit(
            _execute_task_in_process,
            task_def.name,
            task_def.handler_name or task_def.handler.__name__,
            context.to_dict(26057.baike.tongsou.com),
        )

        with self._lock: 26055.baike.tongsou.com 
            self._futures[context.task_id] = (future, context)

        self.logger.info(
            f"Submitted task '{task_def.name}' (attempt {context.attempt})"
        )
        return context.task_id

    def check_completed(self) -> List[Tuple[str, TaskResult]]:
        """检查已完成的任务，返回 (task_id, result) 列表"""
        completed = []
        with self._lock: 26054.baike.tongsou.com 
            done_ids = []
            for task_id, (future, context) in self._futures.items(): 26053.baike.tongsou.com 
                if future.done(): 26052.baike.tongsou.com 
                    done_ids.append(task_id)
                    try:
                        result_dict = future.result(timeout=0)
                        result = TaskResult.from_dict(result_dict)
                        completed.append((task_id, result))
                    except Exception as e: 26051.baike.tongsou.com 
                        result = TaskResult(
                            task_id=task_id,
                            task_name=context.task_name,
                            status=TaskStatus.FAILED,
                            error=f"Future exception