Agent源码解析：从零构建一个可扩展的任务调度框架

1次阅读

没有评论

共计 2219 个字符，预计需要花费 6 分钟才能阅读完成。

在分布式系统中，任务调度是一个核心组件。传统的任务调度框架往往面临以下问题：

单点故障：集中式调度器一旦崩溃，整个系统将瘫痪
扩展性差：难以动态增减 Worker 节点应对负载波动
缺乏弹性：无法根据任务优先级动态调整资源分配

这些问题在微服务架构下尤为突出，亟需一个更健壮的解决方案。

与 Celery、Airflow 等流行框架相比，Agent 在架构上有显著差异：

Celery
基于消息队列的推模型
强依赖中间件 (RabbitMQ/Redis)
静态 Worker 配置
Airflow
面向批处理的 DAG 调度
中心化调度器
复杂依赖管理
Agent
混合推拉模型
去中心化调度
动态 Worker 注册

Agent 采用双层队列设计：

全局优先级队列
基于 Redis 的 Sorted Set 实现
按任务优先级和提交时间排序
O(logN) 的插入 / 删除复杂度
本地执行队列
每个 Worker 维护的内存队列
采用双端队列结构
支持任务抢占和暂停

# 任务队列实现示例
class TaskQueue:
    def __init__(self, redis_conn):
        self.redis = redis_conn
        self.local_queue = deque()

    def push(self, task_id, priority):
        # Redis ZADD 操作
        self.redis.zadd('global_queue', {task_id: priority})

    def pull(self, worker_id):
        # 原子获取最高优先级任务
        script = """local task = redis.call('ZPOPMAX', KEYS[1])
        if task then
            redis.call('HSET', KEYS[2], task[1], ARGV[1])
            return task
        end
        return nil
        """
        return self.redis.eval(
            script, 2, 
            'global_queue', 
            'assigned_tasks',
            worker_id
        )

Agent 使用 Redlock 算法保证分布式一致性：

获取锁时向所有 Redis 节点发送 SETNX 命令
当获得多数节点响应时视为成功
锁自动过期防止死锁
采用令牌续期机制处理长任务

时间复杂度分析：
– 获取锁：O(N) N 为 Redis 节点数
– 释放锁：O(1)

通过心跳机制实现：

Worker 定期上报负载指标
调度器根据负载动态调整任务分配
支持优雅下线：
停止接收新任务
完成存量任务
注销注册信息

完整任务定义和调度示例：

from dataclasses import dataclass
from enum import Enum
import time
import random

class TaskStatus(Enum):
    PENDING = 1
    RUNNING = 2
    COMPLETED = 3
    FAILED = 4

@dataclass
class Task:
    id: str
    priority: int
    max_retries: int = 3
    retry_delay: float = 1.0

    def execute(self):
        # 模拟任务执行
        if random.random() < 0.2:
            raise ValueError("Simulated failure")
        time.sleep(0.5)
        return True

class TaskScheduler:
    def __init__(self):
        self.queue = TaskQueue(redis_conn)
        self.workers = {}

    def schedule(self, task: Task):
        self.queue.push(task.id, task.priority)

    def run_worker(self, worker_id):
        while True:
            task_id = self.queue.pull(worker_id)
            if not task_id:
                time.sleep(0.1)
                continue

            task = self.get_task(task_id)
            for attempt in range(task.max_retries):
                try:
                    result = task.execute()
                    self.mark_completed(task_id)
                    break
                except Exception as e:
                    if attempt == task.max_retries - 1:
                        self.mark_failed(task_id)
                    time.sleep(task.retry_delay)

根据业务特点可调整以下参数：