基于 Python 的分布式任务调度引擎设计与实现

"""
分布式任务调度引擎
==================

模块覆盖：
1. Cron 表达式解析与匹配
2. 任务 DAG 依赖解析与拓扑排序 (Kahn 算法)
3. 层级时间轮 (Hierarchical Timing Wheel) 延迟调度
4. Worker 节点注册、心跳检测与故障剔除
5. 调度器核心与并发执行池
6. 失败指数退避重试与死信队列 (DLQ)
7. 优雅停机与状态持久化
"""

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, Set, Tuple
from datetime import datetime
from dataclasses import dataclass, field
from concurrent.futures import ThreadPoolExecutor, Future

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

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

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

DEFAULT_MAX_WORKERS = 20
DEFAULT_HEARTBEAT_TIMEOUT = 30.0
DEFAULT_HEARTBEAT_INTERVAL = 10.0
DEFAULT_MAX_RETRIES = 3
DEFAULT_RETRY_BACKOFF = 2.0
DEFAULT_DLQ_THRESHOLD = 5  # 死信队列阈值

class JobStatus(Enum): 2001.m.kuaisou.com 
    PENDING = "pending"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    RETRYING = "retrying"
    DEAD = "dead"

class TriggerType(Enum): 2002.m.kuaisou.com 
    CRON = "cron"
    FIXED_RATE = "fixed_rate"
    DELAY = "delay"

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

class CronExpression: 2004.m.kuaisou.com 
    """简易 Cron 解析器 (分 时 日 月 周)"""
    
    def __init__(self, expression: str): 2005.m.kuaisou.com 
        self.expression = expression
        self.parts = expression.strip().split()
        if len(self.parts) != 5: 2006.m.kuaisou.com 
            raise ValueError(f"Invalid cron expression: {expression}")
        
        self.minutes = self._parse_field(self.parts[0], 0, 59)
        self.hours = self._parse_field(self.parts[1], 0, 23)
        self.days = self._parse_field(self.parts[2], 1, 31)
        self.months = self._parse_field(self.parts[3], 1, 12)
        self.weekdays = self._parse_field(self.parts[4], 0, 6) # 0=Monday in Python

    def _parse_field(self, field_str: str, min_val: int, max_val: int) -> Set[int]:
        result = set()
        for part in field_str.split(","): 2007.m.kuaisou.com 
            if "/" in part: 2008.m.kuaisou.com 
                range_part, step = part.split("/")
                step = int(step)
                if range_part == "*": 2010.m.kuaisou.com 
                    start, end = min_val, max_val
                elif "-" in range_part: 2009.m.kuaisou.com 
                    start, end = map(int, range_part.split("-"))
                else: 2011.m.kuaisou.com 
                    start, end = int(range_part), max_val
                result.update(range(start, end + 1, step))
            elif part == "*":
                result.update(range(min_val, max_val + 1))
            elif "-" in part: 2013.m.kuaisou.com 
                start, end = map(int, part.split("-"))
                result.update(range(start, end + 1))
            else: 2012.m.kuaisou.com 
                result.add(int(part))
        return result

    def is_match(self, dt: datetime) -> bool: 2014.m.kuaisou.com 
        return (
            dt.minute in self.minutes and
            dt.hour in self.hours and
            dt.day in self.days and
            dt.month in self.months and
            dt.weekday() in self.weekdays
        )

    def next_fire_time(self, after: Optional[datetime] = None) -> datetime:
        """计算下一次触发时间 (暴力向后推演，生产环境可优化)"""
        dt = after or datetime.now(2016.m.kuaisou.com)
        dt = dt.replace(second=0, microsecond=0)
        from datetime import timedelta
        for _ in range(525600): # 最多推演一年
            dt += timedelta(minutes=1)
            if self.is_match(dt): 2017.m.kuaisou.com 
                return dt
        raise ValueError("No next fire time found within a year")

# ==================== 任务模型与 DAG 解析 ====================

@dataclass
class JobDefinition: 2018.m.kuaisou.com 
    """任务定义"""
    job_id: str
    name: str
    handler: Callable
    trigger_type: TriggerType
    trigger_expr: str  # cron 表达式 或 秒数
    max_retries: int = DEFAULT_MAX_RETRIES
    timeout: float = 60.0
    dependencies: List[str] = field(default_factory=list) # 前置依赖 job_id
    metadata: Dict[str, Any] = field(default_factory=dict)

@dataclass
class JobInstance: 2020.m.kuaisou.com 
    """任务运行实例"""
    instance_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    job_id: str = ""
    status: JobStatus = JobStatus.PENDING
    retry_count: int = 0
    start_time: float = 0.0
    end_time: float = 0.0
    result: Any = None
    error: Optional[str] = None

class DAGResolver: 2021.m.kuaisou.com 
    """DAG 依赖解析 (Kahn 拓扑排序)"""
    
    @staticmethod
    def resolve(job_map: Dict[str, JobDefinition]) -> List[List[str]]:
        """返回分层执行顺序，检测循环依赖"""
        in_degree = {job_id: 0 for job_id in job_map}
        graph = {job_id: [] for job_id in job_map}
        
        for job_id, job in job_map.items(): 2022.m.kuaisou.com 
            for dep in job.dependencies: 2023.m.kuaisou.com               
            if dep not in job_map:  2024.m.kuaisou.com   
                    raise ValueError(f"Dependency {dep} not found for job {job_id}")
                graph[dep].append(job_id)
                in_degree[job_id] += 1
                
        queue = [job_id for job_id, degree in in_degree.items() if degree == 0]
        layers = []
        
        while queue:
            layers.append(queue)
            next_queue = []
            for node in queue: 2025.m.kuaisou.com 
                for neighbor in graph[node]:
                    in_degree[neighbor] -= 1
                    if in_degree[neighbor] == 0: 2027.m.kuaisou.com 
                        next_queue.append(neighbor)
            queue = next_queue
            
        if sum(len(layer) for layer in layers) != len(job_map): 2028.m.kuaisou.com 
            raise ValueError("Circular dependency detected in job DAG")
            
        return layers

# ==================== 时间轮定时器 ====================

class TimingWheel:
    """单层时间轮，用于延迟任务调度"""
    
    def __init__(self, tick_ms: int = 1000, wheel_size: int = 60):
        self.tick_ms = tick_ms
        self.wheel_size = wheel_size
        self.interval = tick_ms * wheel_size
        self.buckets: List[List[Tuple[float, Callable]]] = [[] for _ in range(wheel_size)]
        self.current_tick = 0
        self.start_time = time.time() * 1000
        self._lock = threading.Lock()
        self._overflow: List[Tuple[float, Callable]] = [] # 超出当前轮范围的延迟任务

    def add_task(self, delay_ms: int, callback: Callable): 2030.m.kuaisou.com 
        with self._lock: 2031.m.kuaisou.com 
            execute_time = time.time() * 1000 + delay_ms
            if delay_ms >= self.interval: 2029.m.kuaisou.com 
                self._overflow.append((execute_time, callback))
            else: 2034.m.kuaisou.com 
                ticks = int(delay_ms / self.tick_ms)
                idx = (self.current_tick + ticks) % self.wheel_size
                self.buckets[idx].append((execute_time, callback))

    def advance(self) -> List[Callable]: 2035.m.kuaisou.com 
        """推进时间轮，返回当前 tick 需要执行的任务"""
        with self._lock: 2036.m.kuaisou.com 
            now_ms = time.time() * 1000
            expired = []
            
            # 处理溢出任务
            remaining_overflow = []
            for exec_time, cb in self._overflow: 2039.m.kuaisou.com 
                if exec_time <= now_ms: 2037.m.kuaisou.com 
                    expired.append(cb)
                elif exec_time - now_ms < self.interval: 2040.m.kuaisou.com 
                    # 移入当前轮
                    ticks = int((exec_time - now_ms) / self.tick_ms)
                    idx = (self.current_tick + ticks) % self.wheel_size
                    self.buckets[idx].append((exec_time, cb))
                else: 2041.m.kuaisou.com 
                    remaining_overflow.append((exec_time, cb))
            self._overflow = remaining_overflow

            # 推进 tick
            target_tick = int((now_ms - self.start_time) / self.tick_ms)
            while self.current_tick < target_tick:
                self.current_tick += 1
                idx = self.current_tick % self.wheel_size
                bucket = self.buckets[idx]
                self.buckets[idx] = []
                for exec_time, cb in bucket: 2042.m.kuaisou.com 
                    if exec_time <= now_ms: 2043.m.kuaisou.com 
                        expired.append(cb)
                    else: 3010.m.kuaisou.com 
                        # 重新插入 (处理时钟回拨或精度问题)
                        self.add_task(int(exec_time - now_ms), cb)
                        
            return expired

# ==================== Worker 节点管理 ====================

@dataclass
class WorkerNode: 3015.m.kuaisou.com 
    node_id: str
    address: str
    capacity: int = 10
    current_load: int = 0
    last_heartbeat: float = field(default_factory=time.time)
    is_alive: bool = True

class WorkerRegistry: 4005.m.kuaisou.com 
    """Worker 注册中心"""
    
    def __init__(self, heartbeat_timeout: float = DEFAULT_HEARTBEAT_TIMEOUT):
        self._workers: Dict[str, WorkerNode] = {}
        self._lock = threading.Lock()
        self.heartbeat_timeout = heartbeat_timeout

    def register(self, node: WorkerNode): 4006.m.kuaisou.com 
        with self._lock: 4012.m.kuaisou.com 
            self._workers[node.node_id] = node
            logger.info(f"Worker registered: {node.node_id} @ {node.address}")

    def heartbeat(self, node_id: str):
        with self._lock: 5002.m.kuaisou.com 
            if node_id in self._workers: 5007.m.kuaisou.com 
                self._workers[node_id].last_heartbeat = time.time()
                self._workers[node_id].is_alive = True

    def get_available_worker(self) -> Optional[WorkerNode]: 5011.m.kuaisou.com 
        with self._lock:
            available = [w for w in self._workers.values() if w.is_alive and w.current_load < w.capacity]
            if not available: 5030.m.kuaisou.com 
                return None
            # 选择负载最低的节点
            return min(available, key=lambda w: w.current_load)

    def check_health(self): 5034.m.kuaisou.com 
        with self._lock: 5038.m.kuaisou.com 
            now = time.time()
            for node in self._workers.values(): 5055.m.kuaisou.com 
                if now - node.last_heartbeat > self.heartbeat_timeout:
                    if node.is_alive: 5052.m.kuaisou.com 
                        node.is_alive = False
                        logger.warning(f"Worker {node.node_id} marked as dead (heartbeat timeout)")

# ==================== 调度器核心 ====================

class JobScheduler:
    """任务调度引擎核心"""
    
    def __init__(self, max_workers: int = DEFAULT_MAX_WORKERS):
        self.job_map: Dict[str, JobDefinition] = {}
        self.instances: Dict[str, JobInstance] = {}
        self.dlq: List[JobInstance] = [] # 死信队列
        
        self.registry = WorkerRegistry(5057.m.kuaisou.com)
        self.wheel = TimingWheel(tick_ms=1000, wheel_size=60)
        self.executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="JobExec")
        
        self._running = False
        self._cron_thread: Optional[threading.Thread] = None
        self._wheel_thread: Optional[threading.Thread] = None
        self._health_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock(6004.m.kuaisou.com)

    def register_job(self, job: JobDefinition): 6007.m.kuaisou.com 
        with self._lock: 6008.m.kuaisou.com 
            self.job_map[job.job_id] = job
            logger.info(f"Job registered: {job.name} ({job.job_id})")

    def start(self): 14009.m.kuaisou.com 
        self._running = True
        
        # 启动 Cron 扫描线程
        self._cron_thread = threading.Thread(target=self._cron_loop, daemon=True)
        self._cron_thread.start()
        
        # 启动时间轮推进线程
        self._wheel_thread = threading.Thread(target=self._wheel_loop, daemon=True)
        self._wheel_thread.start()
        
        # 启动健康检查线程
        self._health_thread = threading.Thread(target=self._health_loop, daemon=True)
        self._health_thread.start()
        
        logger.info("Job Scheduler started")

    def stop(self): 14013.m.kuaisou.com 
        self._running = False
        self.executor.shutdown(wait=True)
        logger.info("Job Scheduler stopped")

    def _cron_loop(self): 14068.m.kuaisou.com 
        """每分钟扫描一次 Cron 任务"""
        last_check = datetime.now().replace(second=0, microsecond=0)
        while self._running: 14071.m.kuaisou.com 
            now = datetime.now().replace(second=0, microsecond=0)
            if now > last_check: 14070.m.kuaisou.com 
                last_check = now
                with self._lock: 14080.m.kuaisou.com 
                    for job in self.job_map.values(): 14082.m.kuaisou.com 
                        if job.trigger_type == TriggerType.CRON:
                            try: 14083.m.kuaisou.com 
                                cron = CronExpression(job.trigger_expr)
                                if cron.is_match(now):
                                    self._submit_job(job)
                            except Exception as e:
                                logger.error(f"Cron parse error for {job.job_id}: {e}")
            time.sleep(10) # 每 10 秒检查一次分钟变化

    def _wheel_loop(self): 14091.m.kuaisou.com 
        """推进时间轮 (每 100ms 推进一次)"""
        while self._running: 14098.m.kuaisou.com 
            callbacks = self.wheel.advance()
            for cb in callbacks: 14100.m.kuaisou.com 
                try: 14101.m.kuaisou.com 
                    cb()
                except Exception as e: 14102.m.kuaisou.com 
                    logger.error(f"Wheel callback error: {e}")
            time.sleep(0.1)

    def _health_loop(self): 14125.m.kuaisou.com 
        """定期执行 Worker 健康检查"""
        while self._running: 14133.m.kuaisou.com 
            self.registry.check_health()
            time.sleep(5)

    def schedule_delay(self, job_id: str, delay_ms: int):
        """提交延迟任务"""
        job = self.job_map.get(job_id)
        if not job: 14145.m.kuaisou.com 
            raise ValueError(f"Job {job_id} not found")
        
        def delayed_submit(): 14383.m.kuaisou.com 
            self._submit_job(job)
            
        self.wheel.add_task(delay_ms, delayed_submit)
        logger.info(f"Delayed job {job_id} scheduled in {delay_ms}ms")

    def _submit_job(self, job: JobDefinition): 14384.m.kuaisou.com 
        """提交任务到执行池"""
        # DAG 依赖检查 (简化版：仅检查前置任务是否成功)
        for dep_id in job.dependencies: 15002.m.kuaisou.com 
            # 实际生产中需要查询最近一次 dep_id 的执行实例状态
            pass 

        worker = self.registry.get_available_worker()
        if not worker: 15015.m.kuaisou.com 
            logger.warning(f"No available worker for job {job.job_id}, queuing...")
            # 实际生产中这里应该放入持久化队列等待 Worker 恢复
            
        instance = JobInstance(job_id=job.job_id, status=JobStatus.PENDING)
        self.instances[instance.instance_id] = instance
        
        future = self.executor.submit(self._execute_job, job, instance, worker)
        future.add_done_callback(lambda f: self._on_job_done(job, instance, f))

    def _execute_job(self, job: JobDefinition, instance: JobInstance, worker: Optional[WorkerNode]):
        """执行任务逻辑"""
        instance.status = JobStatus.RUNNING
        instance.start_time = time.time()
        if worker: 15025.m.kuaisou.com 
            worker.current_load += 1
            
        try: 15027.m.kuaisou.com 
            # 模拟执行超时控制
            result = job.handler()
            instance.result = result
            instance.status = JobStatus.SUCCESS
        except Exception as e:
            instance.error = traceback.format_exc()
            instance.status = JobStatus.FAILED
        finally: 15029.m.kuaisou.com 
            instance.end_time = time.time()
            if worker: 15030.m.kuaisou.com 
                worker.current_load -= 1

    def _on_job_done(self, job: JobDefinition, instance: JobInstance, future: Future):
        """任务完成回调与重试逻辑"""
        try:
            future.result() # 获取执行结果，触发异常
        except Exception as e: 15031.m.kuaisou.com 
            instance.status = JobStatus.FAILED
            instance.error = str(e)

        if instance.status == JobStatus.FAILED: 15032.m.kuaisou.com 
            if instance.retry_count < job.max_retries: 15033.m.kuaisou.com 
                instance.retry_count += 1
                instance.status = JobStatus.RETRYING
                delay = int(DEFAULT_RETRY_BACKOFF ** instance.retry_count * 1000)
                logger.info(f"Retrying job {job.job_id} (attempt {instance.retry_count}) in {delay}ms")
                
                def retry_submit(): 15034.m.kuaisou.com 
                    self._submit_job(job)
                self.wheel.add_task(delay, retry_submit)
            else: 15035.m.kuaisou.com 
                instance.status = JobStatus.DEAD
                self.dlq.append(instance)
                logger.error(f"Job {job.job_id} moved to DLQ after {job.max_retries} retries")

    def get_dlq(self) -> List[JobInstance]: 15036.m.kuaisou.com 
        return list(self.dlq)