分布式任务调度平台架构设计实战

// 获取分布式锁
func (l *RedisLock) Acquire(ctx context.Context, key, value string, ttl time.Duration) (bool, error) {
    result, err := l.redis.SetNX(ctx, key, value, ttl).Result()
    if err != nil {
        return false, err
    }
    return result, nil
}

// Lua 脚本：原子性释放锁（只有锁持有者才能释放）
const releaseLockScript = `
if redis.call("get", KEYS[1]) == ARGV[1] then
    return redis.call("del", KEYS[1])
else
    return 0
end
`

// 续期锁 TTL（防止执行时间超过预估时间导致锁提前释放）
func (l *RedisLock) Extend(ctx context.Context, key, value string, ttl time.Duration) (bool, error) {
    script := redis.NewScript(`
        if redis.call("get", KEYS[1]) == ARGV[1] then
            return redis.call("pexpire", KEYS[1], ARGV[2])
        else
            return 0
        end
    `)
    result, err := script.Run(ctx, l.redis, []string{key}, value, ttl.Milliseconds()).Int()
    return result == 1, err
}

type ConsistentHash struct {
    nodes        []uint32           // 排序后的虚拟节点哈希值数组
    nodeMap      map[uint32]string   // hash -> physicalNode
    virtualCount int                // 每个物理节点的虚拟节点数
    mu           sync.RWMutex
}

func (c *ConsistentHash) AddNode(node string) {
    c.mu.Lock()
    defer c.mu.Unlock()
    for i := 0; i < c.virtualCount; i++ {
        vnodeKey := fmt.Sprintf("%s-vn-%d", node, i)
        hashVal := crc32.ChecksumIEEE([]byte(vnodeKey))
        c.nodes = append(c.nodes, hashVal)
        c.nodeMap[hashVal] = node
    }
    sort.Slice(c.nodes, func(i, j int) bool { return c.nodes[i] < c.nodes[j] })
}

// GetShard 返回任务ID对应的分片执行器节点
func (c *ConsistentHash) GetShard(taskId string) string {
    c.mu.RLock()
    defer c.mu.RUnlock()
    hashVal := crc32.ChecksumIEEE([]byte(taskId))
    idx := sort.Search(len(c.nodes), func(i int) bool {
        return c.nodes[i] >= hashVal
    })
    if idx == len(c.nodes) {
        idx = 0
    }
    return c.nodeMap[c.nodes[idx]]
}

// 任务实例状态机
const (
    InstanceStatusPending  = "PENDING"  // 待执行
    InstanceStatusRunning  = "RUNNING"  // 执行中
    InstanceStatusSuccess  = "SUCCESS"  // 执行成功
    InstanceStatusFailed  = "FAILED"   // 执行失败（已重试完毕）
    InstanceStatusRetry   = "RETRY"    // 等待重试
    InstanceStatusDLQ     = "DLQ"      // 死信队列（人工干预）
)

// 幂等插入：只有 PENDING 状态才能转为 RUNNING
func (s *TaskInstanceService) TryStart(ctx context.Context, instanceId string) (bool, error) {
    sql := `UPDATE task_instance
            SET status = ?, start_time = NOW(), executor_id = ?
            WHERE id = ? AND status = ?`
    result, err := s.db.ExecContext(ctx, sql,
        InstanceStatusRunning, s.executorId, instanceId, InstanceStatusPending)
    if err != nil {
        return false, err
    }
    affected, _ := result.RowsAffected()
    return affected == 1, nil // affected=1 表示抢到了执行权
}

const (
    leaderKey         = "scheduler:leader"
    leaderChannel    = "scheduler:leader_change"
    leaderTTL        = 6 * time.Second
    heartbeatInterval = 3 * time.Second
)

type SchedulerNode struct {
    nodeId       string
    redis        *redis.Client
    isLeader     atomic.Bool
}

func (n *SchedulerNode) Run(ctx context.Context) {
    if n.isLeader.Load() {
        n.renewLeader(ctx)
    } else {
        n.tryAcquireLeader(ctx)
    }
}

func (n *SchedulerNode) renewLeader(ctx context.Context) {
    ticker := time.NewTicker(heartbeatInterval)
    defer ticker.Stop()
    for {
        select {
        case <-ticker.C:
            script := redis.NewScript(`
                if redis.call("get", KEYS[1]) == ARGV[1] then
                    return redis.call("pexpire", KEYS[1], ARGV[2])
                else
                    return 0
                end
            `)
            ok, _ := script.Run(ctx, n.redis, []string{leaderKey}, n.nodeId, leaderTTL.Milliseconds()).Int()
            if ok == 0 {
                n.isLeader.Store(false)
                n.onLeadershipLost()
                go n.tryAcquireLeader(ctx)
                return
            }
        case <-ctx.Done():
            n.redis.Del(ctx, leaderKey)
            return
        }
    }
}

func (n *SchedulerNode) tryAcquireLeader(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            return
        default:
        }
        ok, _ := n.redis.SetNX(ctx, leaderKey, n.nodeId, leaderTTL).Result()
        if ok {
            n.isLeader.Store(true)
            n.redis.Publish(ctx, leaderChannel, n.nodeId)
            go n.renewLeader(ctx)
            return
        }
        time.Sleep(1 * time.Second)
    }
}