【编程难度第一名】深入理解调度器设计：编程领域最难的挑战之一

func updateTaskStatus(task *Task, newStatus Status) error {
    if task.Status == Pending {
        if newStatus == Scheduling {
            // 执行调度逻辑
            scheduleTask(task)
            task.Status = Scheduling
        } else {
            return errors.New("invalid transition")
        }
    } else if task.Status == Scheduling {
        if newStatus == Running {
            // 启动任务
            startTask(task)
            task.Status = Running
        } else if newStatus == Failed {
            // 调度失败
            handleScheduleFailure(task)
            task.Status = Failed
        } else {
            return errors.New("invalid transition")
        }
    } else if task.Status == Running {
        if newStatus == Success {
            // 成功
            onSuccess(task)
            task.Status = Success
        } else if newStatus == Failed {
            // 失败
            onFailed(task)
            task.Status = Failed
        } else if newStatus == Lost {
            // 失联
            onLost(task)
            task.Status = Lost
        } else {
            return errors.New("invalid transition")
        }
    } else if task.Status == Failed {
        if newStatus == Retry {
            // 重试
            retryTask(task)
            task.Status = Retry
        } else {
            return errors.New("invalid transition")
        }
    }
    // ... 还有更多状态

    return updateDB(task)
}

type Transition struct {
    From   []Status
    To     Status
    Action func(*Task) error
}

var transitions = []Transition{
    {
        From:   []Status{Pending, Retry},
        To:     Scheduling,
        Action: scheduleTask,
    },
    {
        From:   []Status{Scheduling},
        To:     Running,
        Action: startTask,
    },
    {
        From:   []Status{Running},
        To:     Success,
        Action: onSuccess,
    },
    {
        From:   []Status{Running},
        To:     Failed,
        Action: onFailed,
    },
    {
        From:   []Status{Running},
        To:     Lost,
        Action: onLost,
    },
    {
        From:   []Status{Failed},
        To:     Retry,
        Action: retryTask,
    },
}

func ApplyTransition(task *Task, to Status) error {
    for _, t := range transitions {
        if t.To == to && contains(t.From, task.Status) {
            // 执行副作用
            if err := t.Action(task); err != nil {
                return err
            }
            // 更新数据库（CAS）
            return updateStatus(task.ID, task.Version, to)
        }
    }
    return fmt.Errorf("invalid transition %s → %s", task.Status, to)
}

{
    From:   []Status{Lost},
    To:     Recovering,
    Action: startRecovery,
},

func updateStatus(task *Task) error {
    if task.Status == Pending {
        task.Status = Scheduling
        // 直接在这里做副作用
        selectNode(task)  // 如果这里失败了怎么办？
        startContainer(task)  // 如果这里失败了怎么办？
        // 状态已经改了，但副作用失败了，数据不一致！
    }
}

func ApplyTransition(task *Task, to Status) error {
    // 1. 先执行副作用
    if err := executeAction(task, to); err != nil {
        return err  // 副作用失败，状态不变
    }

    // 2. 副作用成功，才更新状态
    return updateStatusCAS(task, to)
}

func scheduleTask(ctx context.Context, task *Task) error {
    // 1. 资源过滤（Filter）
    candidateNodes := filterNodes(task.Requirements)
    if len(candidateNodes) == 0 {
        return errors.New("no suitable node")
    }

    // 2. 节点打分（Score）
    scoredNodes := scoreNodes(candidateNodes, task)

    // 3. 选择最优节点
    bestNode := selectBestNode(scoredNodes)

    // 4. 写入调度结果
    task.ScheduledNode = bestNode.Name
    task.ScheduledAt = time.Now()

    return nil
}

func filterNodes(requirements Requirements) []Node {
    var result []Node

    for _, node := range allNodes {
        // 检查内存
        if node.AvailableMemory < requirements.Memory {
            continue
        }

        // 检查 CPU
        if node.AvailableCPU < requirements.CPU {
            continue
        }

        // 检查 GPU
        if requirements.GPU > 0 && node.AvailableGPU < requirements.GPU {
            continue
        }

        // 检查节点健康状态
        if !node.IsReady {
            continue
        }

        // 检查标签匹配
        if !matchLabels(node.Labels, requirements.NodeSelector) {
            continue
        }

        result = append(result, node)
    }

    return result
}

func scoreNodes(nodes []Node, task *Task) []ScoredNode {
    var result []ScoredNode

    for _, node := range nodes {
        score := 0

        // 资源充足度（0-100 分）
        memScore := (node.AvailableMemory / node.TotalMemory) * 100
        cpuScore := (node.AvailableCPU / node.TotalCPU) * 100
        score += int(memScore + cpuScore) / 2

        // 负载均衡（0-100 分）
        loadScore := (1 - node.CurrentLoad) * 100
        score += int(loadScore)

        // 亲和性（0-100 分）
        if hasAffinity(node, task) {
            score += 100
        }

        result = append(result, ScoredNode{
            Node:  node,
            Score: score,
        })
    }

    // 按分数排序
    sort.Slice(result, func(i, j int) bool {
        return result[i].Score > result[j].Score
    })

    return result
}

func startTask(ctx context.Context, task *Task) error {
    // 1. 构建容器配置
    containerConfig := buildContainerConfig(task)

    // 2. 调用底层系统启动容器（Kubernetes/Volcano/Ray）
    // 注意：这里必须幂等！
    jobID := fmt.Sprintf("task-%d", task.ID)
    err := clusterClient.CreateJob(ctx, jobID, containerConfig)
    if err != nil {
        // 如果是"已存在"错误，认为成功
        if isAlreadyExistsError(err) {
            // 验证任务确实在运行
            job, err := clusterClient.GetJob(ctx, jobID)
            if err != nil {
                return err
            }
            if job.Status != "Running" {
                return fmt.Errorf("job exists but not running: %s", job.Status)
            }
        } else {
            return err
        }
    }

    // 3. 记录启动时间
    task.StartedAt = time.Now()

    // 4. 初始化心跳
    task.HeartbeatVersion = 1
    task.LastHeartbeatAt = time.Now()

    // 5. 注册监控
    registerMetrics(task)

    return nil
}

func onSuccess(ctx context.Context, task *Task) error {
    // 1. 记录完成时间
    task.FinishedAt = time.Now()

    // 2. 清理资源
    err := cleanupResources(ctx, task)
    if err != nil {
        log.Error("cleanup failed", "taskID", task.ID, "err", err)
        // 注意：清理失败不影响任务成功
    }

    // 3. 发送通知
    notifyUser(task, "Task completed successfully")

    // 4. 更新 Workflow 状态
    updateWorkflowProgress(task.WorkflowID)

    // 5. 记录 metrics
    recordMetrics(task, "success")

    // 6. 触发依赖任务
    triggerDependentTasks(task)

    return nil
}

func onFailed(ctx context.Context, task *Task) error {
    // 1. 记录失败时间
    task.FinishedAt = time.Now()

    // 2. 收集错误信息
    errorInfo := collectErrorInfo(task)
    task.ErrorReason = errorInfo

    // 3. 判断是否需要重试
    if task.RetryCount < task.MaxRetries {
        // 转到 RetryWait 状态
        task.RetryCount++
        task.Status = StatusRetryWait

        // 计算下次重试时间（指数退避）
        backoffDuration := calculateBackoff(task.RetryCount)
        task.NextRetryAt = time.Now().Add(backoffDuration)

        log.Info("task will retry",
            "taskID", task.ID,
            "retryCount", task.RetryCount,
            "nextRetryAt", task.NextRetryAt)
    } else {
        // 已达重试上限，真正失败
        log.Error("task failed permanently",
            "taskID", task.ID,
            "error", errorInfo)

        // 发送告警
        sendAlert(task, "Task failed after max retries")
    }

    // 4. 清理资源
    cleanupResources(ctx, task)

    // 5. 记录 metrics
    recordMetrics(task, "failed")

    return nil
}

func calculateBackoff(retryCount int) time.Duration {
    // 指数退避：1s, 2s, 4s, 8s, 16s, ...
    seconds := math.Pow(2, float64(retryCount))
    // 最大 5 分钟
    if seconds > 300 {
        seconds = 300
    }
    return time.Duration(seconds) * time.Second
}

func onLost(ctx context.Context, task *Task) error {
    // 1. 记录失联时间
    task.LostAt = time.Now()

    // 2. 尝试查询底层系统状态
    job, err := clusterClient.GetJob(ctx, fmt.Sprintf("task-%d", task.ID))
    if err == nil && job != nil {
        // 任务其实还在运行，只是心跳丢了
        if job.Status == "Running" {
            log.Warn("task lost but still running", "taskID", task.ID)
            // 恢复心跳
            task.LastHeartbeatAt = time.Now()
            task.Status = StatusRunning
            return nil
        }

        // 任务已经结束
        if job.Status == "Success" {
            task.Status = StatusSuccess
            return onSuccess(ctx, task)
        }

        if job.Status == "Failed" {
            task.Status = StatusFailed
            return onFailed(ctx, task)
        }
    }

    // 3. 真的失联了，决定是否重试
    if task.RetryCount < task.MaxRetries {
        task.RetryCount++
        task.Status = StatusRetryWait
        task.NextRetryAt = time.Now().Add(calculateBackoff(task.RetryCount))
    } else {
        task.Status = StatusFailed
        task.ErrorReason = "Task lost and exceeded max retries"
    }

    // 4. 清理资源
    cleanupResources(ctx, task)

    // 5. 发送告警
    sendAlert(task, "Task lost")

    return nil
}

type TaskRepo interface {
    UpdateStatusCAS(ctx context.Context,
                    taskID int64,
                    oldVersion int64,
                    newStatus Status,
                    newVersion int64) (bool, error)
}

func (r *PostgresTaskRepo) UpdateStatusCAS(
    ctx context.Context,
    taskID int64,
    oldVersion int64,
    newStatus Status,
    newVersion int64,
) (bool, error) {
    result, err := r.db.ExecContext(ctx, `
        UPDATE task_runs
        SET status = $1,
            version = $2,
            updated_at = NOW()
        WHERE id = $3
          AND version = $4
    `, newStatus, newVersion, taskID, oldVersion)

    if err != nil {
        return false, err
    }

    rowsAffected, err := result.RowsAffected()
    if err != nil {
        return false, err
    }

    // 如果影响行数为 0，说明 version 不匹配
    return rowsAffected > 0, nil
}

func (s *Scheduler) handleTask(ctx context.Context, task *Task) error {
    oldVersion := task.Version
    newVersion := oldVersion + 1

    // 执行副作用
    if err := scheduleTask(ctx, task); err != nil {
        return err
    }

    // CAS 更新
    success, err := s.repo.UpdateStatusCAS(
        ctx,
        task.ID,
        oldVersion,
        StatusScheduling,
        newVersion,
    )

    if err != nil {
        return err
    }

    if !success {
        // CAS 失败，说明任务已被别的调度器处理
        log.Info("task already processed by another scheduler",
                 "taskID", task.ID)
        return nil
    }

    // 更新本地任务对象
    task.Status = StatusScheduling
    task.Version = newVersion

    return nil
}

func (r *PostgresTaskRepo) PickRunnableTasks(
    ctx context.Context,
    limit int,
) ([]*Task, error) {
    rows, err := r.db.QueryContext(ctx, `
        SELECT id, status, version, ...
        FROM task_runs
        WHERE status IN ('Pending', 'RetryWait')
          AND (next_retry_at IS NULL OR next_retry_at <= NOW())
          AND is_paused = FALSE
        ORDER BY priority DESC, id
        FOR UPDATE SKIP LOCKED
        LIMIT $1
    `, limit)

    if err != nil {
        return nil, err
    }
    defer rows.Close()

    var tasks []*Task
    for rows.Next() {
        var task Task
        if err := rows.Scan(&task.ID, &task.Status, ...); err != nil {
            return nil, err
        }
        tasks = append(tasks, &task)
    }

    return tasks, nil
}

func (c *K8sClient) CreateJob(ctx context.Context, jobName string, config JobConfig) error {
    // 1. 先检查是否已存在
    existingJob, err := c.clientset.BatchV1().Jobs(namespace).Get(
        ctx,
        jobName,
        metav1.GetOptions{},
    )

    if err == nil {
        // 任务已存在
        if existingJob.Status.Active > 0 {
            // 正在运行，直接返回成功
            return nil
        }

        if existingJob.Status.Succeeded > 0 {
            // 已经成功，直接返回成功
            return nil
        }

        // 失败了，删除重建
        if existingJob.Status.Failed > 0 {
            err = c.clientset.BatchV1().Jobs(namespace).Delete(
                ctx,
                jobName,
                metav1.DeleteOptions{},
            )
            if err != nil {
                return err
            }
        }
    } else if !errors.IsNotFound(err) {
        // 其他错误
        return err
    }

    // 2. 创建任务
    job := buildK8sJob(jobName, config)
    _, err = c.clientset.BatchV1().Jobs(namespace).Create(
        ctx,
        job,
        metav1.CreateOptions{},
    )

    if err != nil && errors.IsAlreadyExists(err) {
        // 并发创建，已存在，认为成功
        return nil
    }

    return err
}

func (s *Scheduler) watchTaskChanges(ctx context.Context) {
    listener := pq.NewListener(
        s.dbURL,
        10*time.Second,
        time.Minute,
        func(ev pq.ListenerEventType, err error) {
            if err != nil {
                log.Error("listener error", "err", err)
            }
        },
    )

    err := listener.Listen("task_change")
    if err != nil {
        log.Fatal("listen failed", "err", err)
    }

    for {
        select {
        case notification := <-listener.Notify:
            var event TaskChangeEvent
            if err := json.Unmarshal([]byte(notification.Extra), &event); err != nil {
                log.Error("unmarshal failed", "err", err)
                continue
            }

            // 处理事件
            s.handleTaskChange(ctx, event)

        case <-ctx.Done():
            return
        }
    }
}

func (s *Scheduler) runRecoveryLoop(ctx context.Context) {
    ticker := time.NewTicker(30 * time.Second)
    defer ticker.Stop()

    for {
        select {
        case <-ticker.C:
            s.recoverLostTasks(ctx)
        case <-ctx.Done():
            return
        }
    }
}

func (s *Scheduler) recoverLostTasks(ctx context.Context) {
    // 查找所有 Lost 状态的任务
    lostTasks, err := s.repo.FindTasksByStatus(ctx, StatusLost)
    if err != nil {
        log.Error("find lost tasks failed", "err", err)
        return
    }

    for _, task := range lostTasks {
        // 查询底层系统状态
        job, err := s.clusterClient.GetJob(
            ctx,
            fmt.Sprintf("task-%d", task.ID),
        )

        if err != nil {
            log.Error("get job failed", "taskID", task.ID, "err", err)
            continue
        }

        // 根据实际状态恢复
        if job.Status == "Running" {
            // 任务还在运行，恢复心跳
            err = s.applyTransition(ctx, task, StatusRunning)
        } else if job.Status == "Success" {
            // 任务已成功
            err = s.applyTransition(ctx, task, StatusSuccess)
        } else if job.Status == "Failed" {
            // 任务已失败
            err = s.applyTransition(ctx, task, StatusFailed)
        } else {
            // 任务真的丢了，重试
            if task.RetryCount < task.MaxRetries {
                err = s.applyTransition(ctx, task, StatusRetryWait)
            } else {
                err = s.applyTransition(ctx, task, StatusFailed)
            }
        }

        if err != nil {
            log.Error("recover task failed", "taskID", task.ID, "err", err)
        }
    }
}

type Task struct {
    // ...
    IsPaused     bool      `db:"is_paused"`
    PauseReason  string    `db:"pause_reason"`
    PausedAt     time.Time `db:"paused_at"`
    PausedBy     string    `db:"paused_by"`
}

任何状态 → Pause （允许）
Pause → 其他状态（不允许，除非 Resume）

func (s *Scheduler) PauseTask(
    ctx context.Context,
    taskID int64,
    reason string,
    operator string,
) error {
    tx, err := s.db.BeginTx(ctx, nil)
    if err != nil {
        return err
    }
    defer tx.Rollback()

    // 更新暂停标志
    _, err = tx.ExecContext(ctx, `
        UPDATE task_runs
        SET is_paused = TRUE,
            pause_reason = $1,
            paused_at = NOW(),
            paused_by = $2,
            updated_at = NOW()
        WHERE id = $3
    `, reason, operator, taskID)

    if err != nil {
        return err
    }

    // 如果任务正在运行，尝试停止
    task, err := s.repo.GetTask(ctx, taskID)
    if err != nil {
        return err
    }

    if task.Status == StatusRunning {
        // 调用底层系统停止任务
        err = s.clusterClient.StopJob(
            ctx,
            fmt.Sprintf("task-%d", taskID),
        )
        if err != nil {
            log.Error("stop job failed", "taskID", taskID, "err", err)
            // 继续，因为暂停标志已设置，下次调度会跳过
        }
    }

    return tx.Commit()
}

func (s *Scheduler) ResumeTask(
    ctx context.Context,
    taskID int64,
    operator string,
) error {
    _, err := s.db.ExecContext(ctx, `
        UPDATE task_runs
        SET is_paused = FALSE,
            pause_reason = NULL,
            paused_at = NULL,
            paused_by = NULL,
            updated_at = NOW()
        WHERE id = $1
    `, taskID)

    return err
}

type Scheduler struct {
    // ...
    isDegraded   bool
    degradeMode  DegradeMode
    degradeUntil time.Time
}

type DegradeMode int

const (
    DegradeModeNormal DegradeMode = iota
    DegradeModePartial  // 部分降级：只调度高优先级任务
    DegraduModeFull     // 完全降级：停止所有调度
)

func (s *Scheduler) enterDegradeMode(
    mode DegradeMode,
    duration time.Duration,
    reason string,
) {
    s.mu.Lock()
    defer s.mu.Unlock()

    s.isDegraded = true
    s.degradeMode = mode
    s.degradeUntil = time.Now().Add(duration)

    log.Warn("entering degrade mode",
        "mode", mode,
        "duration", duration,
        "reason", reason)

    // 发送告警
    s.alertManager.Send(Alert{
        Level:   "warning",
        Title:   "Scheduler Degraded",
        Message: fmt.Sprintf("Reason: %s, Mode: %v, Duration: %v",
                            reason, mode, duration),
    })
}

func (s *Scheduler) checkAutoDegrade(ctx context.Context) {
    // 检查 GPU 资源
    gpuUsage, err := s.getGPUUsage(ctx)
    if err == nil && gpuUsage > 0.95 {
        s.enterDegradeMode(
            DegradeModePartial,
            10*time.Minute,
            "GPU usage > 95%",
        )
        return
    }

    // 检查 Node 健康度
    healthyNodes, totalNodes := s.getNodeHealth(ctx)
    if float64(healthyNodes)/float64(totalNodes) < 0.5 {
        s.enterDegradeMode(
            DegraduModeFull,
            30*time.Minute,
            "Less than 50% nodes healthy",
        )
        return
    }

    // 检查调度器负载
    if s.metrics.PendingTasks > 10000 {
        s.enterDegradeMode(
            DegradeModePartial,
            5*time.Minute,
            "Too many pending tasks",
        )
        return
    }
}

func (s *Scheduler) RunOnce(ctx context.Context) error {
    // 检查是否需要降级
    s.checkAutoDegrade(ctx)

    // 检查是否还在降级期
    if s.isDegraded && time.Now().After(s.degradeUntil) {
        s.exitDegradeMode()
    }

    // 如果在降级模式
    if s.isDegraded {
        if s.degradeMode == DegraduModeFull {
            // 完全降级：不调度任何任务
            log.Info("in full degrade mode, skipping scheduling")
            return nil
        }

        if s.degradeMode == DegradeModePartial {
            // 部分降级：只调度高优先级任务
            tasks, err := s.repo.PickHighPriorityTasks(ctx, 10)
            if err != nil {
                return err
            }
            return s.processTasks(ctx, tasks)
        }
    }

    // 正常调度
    tasks, err := s.repo.PickRunnableTasks(ctx, 100)
    if err != nil {
        return err
    }
    return s.processTasks(ctx, tasks)
}

func (s *Server) HandleHeartbeat(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()

    var req HeartbeatRequest
    if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
        http.Error(w, "invalid request", 400)
        return
    }

    // 查询任务
    task, err := s.repo.GetTask(ctx, req.TaskID)
    if err != nil {
        http.Error(w, "task not found", 404)
        return
    }

    // 关键：检查版本号
    if req.Version < task.HeartbeatVersion {
        // 旧心跳，丢弃
        log.Warn("outdated heartbeat",
            "taskID", req.TaskID,
            "receivedVersion", req.Version,
            "currentVersion", task.HeartbeatVersion)

        // 返回 200，但不更新（避免节点重试）
        w.WriteHeader(200)
        return
    }

    // 更新心跳
    err = s.repo.UpdateHeartbeat(ctx, req.TaskID, req.Version, time.Now())
    if err != nil {
        http.Error(w, "update failed", 500)
        return
    }

    // 可选：更新 metrics
    if req.Metrics != nil {
        s.metricsStore.Update(req.TaskID, req.Metrics)
    }

    w.WriteHeader(200)
}

func (s *Scheduler) CheckLostTasks(ctx context.Context) error {
    now := time.Now()

    // 查找所有 Running 任务
    tasks, err := s.repo.FindTasksByStatus(ctx, StatusRunning)
    if err != nil {
        return err
    }

    for _, task := range tasks {
        // 条件 1：心跳超时
        heartbeatTimeout := task.LastHeartbeatAt.Add(task.HeartbeatTimeout)
        if !now.After(heartbeatTimeout) {
            // 心跳正常
            continue
        }

        // 心跳超时了，进一步检查

        // 条件 2：查询节点状态
        node, err := s.clusterClient.GetNode(ctx, task.ScheduledNode)
        if err != nil {
            log.Error("get node failed", "node", task.ScheduledNode, "err", err)
        } else {
            if node.IsReady {
                // 节点是健康的，可能只是任务心跳丢了
                // 给一次宽限期
                if now.Sub(task.LastHeartbeatAt) < 2*task.HeartbeatTimeout {
                    continue
                }
            }
        }

        // 条件 3：查询底层任务状态
        job, err := s.clusterClient.GetJob(
            ctx,
            fmt.Sprintf("task-%d", task.ID),
        )
        if err != nil {
            log.Error("get job failed", "taskID", task.ID, "err", err)
        } else {
            if job.Status == "Running" {
                // 任务还在运行，只是心跳丢了
                // 恢复心跳
                log.Warn("task heartbeat lost but job still running",
                    "taskID", task.ID)
                task.LastHeartbeatAt = now
                s.repo.UpdateHeartbeat(ctx, task.ID, task.HeartbeatVersion, now)
                continue
            }

            if job.Status == "Success" {
                // 任务已成功，心跳没上报
                log.Info("task success but heartbeat lost",
                    "taskID", task.ID)
                s.applyTransition(ctx, task, StatusSuccess)
                continue
            }

            if job.Status == "Failed" {
                // 任务已失败
                log.Info("task failed but heartbeat lost",
                    "taskID", task.ID)
                s.applyTransition(ctx, task, StatusFailed)
                continue
            }
        }

        // 条件 4：检查日志更新时间
        lastLogTime, err := s.logStore.GetLastUpdateTime(task.ID)
        if err == nil {
            if now.Sub(lastLogTime) < task.HeartbeatTimeout {
                // 日志还在更新，任务可能还在跑
                log.Info("task heartbeat lost but logs still updating",
                    "taskID", task.ID)
                continue
            }
        }

        // 所有条件都确认：任务真的失联了
        log.Warn("task confirmed lost",
            "taskID", task.ID,
            "lastHeartbeat", task.LastHeartbeatAt,
            "timeout", task.HeartbeatTimeout)

        err = s.applyTransition(ctx, task, StatusLost)
        if err != nil {
            log.Error("mark task lost failed", "taskID", task.ID, "err", err)
        }
    }

    return nil
}

Pending      ：待调度（新任务 / 等待首跑）
Scheduling   ：正在选节点 / 提交到底层系统
Running      ：底层已接收，任务在跑
Success      ：执行成功
Failed       ：执行失败（且不再重试）
RetryWait    ：失败后等待重试时间窗
Lost         ：失联（心跳超时 / 节点炸了）
Paused       ：人工或系统暂停
Canceled     ：人工取消

package fsm

import (
    "context"
    "fmt"
)

type Status string

const (
    StatusPending    Status = "Pending"
    StatusScheduling Status = "Scheduling"
    StatusRunning    Status = "Running"
    StatusSuccess    Status = "Success"
    StatusFailed     Status = "Failed"
    StatusRetryWait  Status = "RetryWait"
    StatusLost       Status = "Lost"
    StatusPaused     Status = "Paused"
    StatusCanceled   Status = "Canceled"
)

type Task struct {
    ID         int64
    Status     Status
    Version    int64
    RetryCount int
    MaxRetries int
    // ... 其他字段
}

type Transition struct {
    From   []Status
    To     Status
    Action func(ctx context.Context, task *Task) error
}

type ActionFunc func(ctx context.Context, task *Task) error

package fsm

var transitions = []Transition{
    {
        From:   []Status{StatusPending, StatusRetryWait},
        To:     StatusScheduling,
        Action: scheduleTask, // 选节点、写 scheduled_node、下发到底层
    },
    {
        From:   []Status{StatusScheduling},
        To:     StatusRunning,
        Action: startTask, // 已成功提交，记录 start_time / 注册心跳
    },
    {
        From:   []Status{StatusRunning},
        To:     StatusSuccess,
        Action: onSuccess, // 清理资源、发通知、更新 metrics
    },
    {
        From:   []Status{StatusRunning},
        To:     StatusFailed,
        Action: onFailed,
    },
    {
        From:   []Status{StatusRunning},
        To:     StatusLost,
        Action: onLost,
    },
    {
        From:   []Status{StatusFailed, StatusLost},
        To:     StatusRetryWait,
        Action: onRetryWait,
    },
    {
        From:   []Status{StatusRunning, StatusPending, StatusRetryWait},
        To:     StatusPaused,
        Action: onPaused, // 记录原因，不再被调度
    },
    {
        From:   []Status{StatusPending, StatusRetryWait, StatusPaused},
        To:     StatusCanceled,
        Action: onCanceled,
    },
}

package fsm

func ApplyTransition(
    ctx context.Context,
    repo TaskRepo,
    task *Task,
    to Status,
) error {
    // 1. 查找匹配的转换
    var found *Transition
    for i := range transitions {
        t := &transitions[i]
        if t.To != to {
            continue
        }
        for _, from := range t.From {
            if from == task.Status {
                found = t
                break
            }
        }
        if found != nil {
            break
        }
    }

    if found == nil {
        return fmt.Errorf("invalid transition %s -> %s", task.Status, to)
    }

    // 2. 先执行副作用（一般在一个 DB 事务里更安全）
    if err := found.Action(ctx, task); err != nil {
        return fmt.Errorf("action failed: %w", err)
    }

    // 3. CAS 更新状态 + version（单写者）
    oldVersion := task.Version
    task.Status = to
    task.Version++

    ok, err := repo.UpdateStatusCAS(
        ctx,
        task.ID,
        oldVersion,
        to,
        task.Version,
    )
    if err != nil {
        return fmt.Errorf("update status failed: %w", err)
    }

    if !ok {
        return fmt.Errorf("cas conflict when updating task %d", task.ID)
    }

    return nil
}

func contains(statuses []Status, status Status) bool {
    for _, s := range statuses {
        if s == status {
            return true
        }
    }
    return false
}

package repo

import (
    "context"
    "database/sql"
)

type TaskRepo interface {
    PickRunnableTasks(ctx context.Context, limit int) ([]*Task, error)
    UpdateStatusCAS(ctx context.Context, taskID int64, oldVersion int64,
                    newStatus Status, newVersion int64) (bool, error)
    GetTask(ctx context.Context, taskID int64) (*Task, error)
    // ... 其他方法
}

type PostgresTaskRepo struct {
    db *sql.DB
}

func (r *PostgresTaskRepo) PickRunnableTasks(
    ctx context.Context,
    limit int,
) ([]*Task, error) {
    query := `
        SELECT
            id, workflow_run_id, name, status, version,
            priority, retry_count, max_retries,
            heartbeat_version, last_heartbeat_at, heartbeat_timeout,
            payload, created_at, updated_at
        FROM task_runs
        WHERE status IN ('Pending', 'RetryWait')
          AND (next_retry_at IS NULL OR next_retry_at <= NOW())
          AND is_paused = FALSE
        ORDER BY priority DESC, id
        FOR UPDATE SKIP LOCKED
        LIMIT $1
    `

    rows, err := r.db.QueryContext(ctx, query, limit)
    if err != nil {
        return nil, err
    }
    defer rows.Close()

    var tasks []*Task
    for rows.Next() {
        task := &Task{}
        err := rows.Scan(
            &task.ID, &task.WorkflowRunID, &task.Name,
            &task.Status, &task.Version,
            &task.Priority, &task.RetryCount, &task.MaxRetries,
            &task.HeartbeatVersion, &task.LastHeartbeatAt,
            &task.HeartbeatTimeout,
            &task.Payload, &task.CreatedAt, &task.UpdatedAt,
        )
        if err != nil {
            return nil, err
        }
        tasks = append(tasks, task)
    }

    return tasks, rows.Err()
}

func (r *PostgresTaskRepo) UpdateStatusCAS(
    ctx context.Context,
    taskID int64,
    oldVersion int64,
    newStatus Status,
    newVersion int64,
) (bool, error) {
    result, err := r.db.ExecContext(ctx, `
        UPDATE task_runs
        SET status = $1,
            version = $2,
            updated_at = NOW()
        WHERE id = $3
          AND version = $4
    `, newStatus, newVersion, taskID, oldVersion)

    if err != nil {
        return false, err
    }

    rowsAffected, err := result.RowsAffected()
    if err != nil {
        return false, err
    }

    return rowsAffected > 0, nil
}

package scheduler

import (
    "context"
    "sync"
    "time"
)

type Scheduler struct {
    repo         TaskRepo
    clusterClient ClusterClient
    workerPool   *WorkerPool
    logger       Logger

    // 降级相关
    mu           sync.RWMutex
    isDegraded   bool
    degradeMode  DegradeMode
    degradeUntil time.Time
}

func (s *Scheduler) Run(ctx context.Context) error {
    // 启动多个后台循环
    g, ctx := errgroup.WithContext(ctx)

    // 主调度循环
    g.Go(func() error {
        return s.runMainLoop(ctx)
    })

    // 心跳检查循环
    g.Go(func() error {
        return s.runHeartbeatCheckLoop(ctx)
    })

    // 恢复循环
    g.Go(func() error {
        return s.runRecoveryLoop(ctx)
    })

    // 降级检查循环
    g.Go(func() error {
        return s.runDegradeCheckLoop(ctx)
    })

    return g.Wait()
}

func (s *Scheduler) runMainLoop(ctx context.Context) error {
    ticker := time.NewTicker(5 * time.Second)
    defer ticker.Stop()

    for {
        select {
        case <-ticker.C:
            if err := s.RunOnce(ctx); err != nil {
                s.logger.Error("schedule once failed", "err", err)
            }
        case <-ctx.Done():
            return ctx.Err()
        }
    }
}

func (s *Scheduler) RunOnce(ctx context.Context) error {
    // 检查降级
    if s.shouldSkipScheduling() {
        return nil
    }

    // 拉取任务
    limit := s.getTaskLimit()
    tasks, err := s.repo.PickRunnableTasks(ctx, limit)
    if err != nil {
        return err
    }

    if len(tasks) == 0 {
        return nil
    }

    s.logger.Info("picked tasks", "count", len(tasks))

    // 并发处理任务
    for _, t := range tasks {
        task := t // 避免闭包问题
        s.workerPool.Submit(func() {
            if err := s.handleTask(ctx, task); err != nil {
                s.logger.Error("handleTask failed",
                    "taskID", task.ID, "err", err)
            }
        })
    }

    return nil
}

func (s *Scheduler) handleTask(ctx context.Context, task *Task) error {
    s.logger.Info("handling task",
        "taskID", task.ID,
        "status", task.Status)

    // Pending / RetryWait -> Scheduling
    if err := ApplyTransition(ctx, s.repo, task, StatusScheduling); err != nil {
        return fmt.Errorf("transition to Scheduling failed: %w", err)
    }

    s.logger.Info("task scheduled", "taskID", task.ID)

    // Scheduling -> Running
    if err := ApplyTransition(ctx, s.repo, task, StatusRunning); err != nil {
        return fmt.Errorf("transition to Running failed: %w", err)
    }

    s.logger.Info("task started", "taskID", task.ID)

    return nil
}

func (s *Scheduler) shouldSkipScheduling() bool {
    s.mu.RLock()
    defer s.mu.RUnlock()

    if !s.isDegraded {
        return false
    }

    // 检查是否过期
    if time.Now().After(s.degradeUntil) {
        return false
    }

    return s.degradeMode == DegraduModeFull
}

func (s *Scheduler) getTaskLimit() int {
    s.mu.RLock()
    defer s.mu.RUnlock()

    if s.isDegraded && s.degradeMode == DegradeModePartial {
        return 10  // 降级时只拉取 10 个
    }

    return 100  // 正常时拉取 100 个
}