【分布式系统百科】数据再平衡：固定分区、动态分区与 TiKV 的调度策略

class FixedPartitionRebalancer:
    def __init__(self, num_partitions, num_replicas):
        self.num_partitions = num_partitions
        self.num_replicas = num_replicas
        
    def assign_partitions(self, nodes):
        """
        将分区均匀分配给节点
        返回: {partition_id: [node_list]}
        """
        if len(nodes) == 0:
            return {}
            
        assignments = {}
        nodes_list = sorted(nodes)
        
        for partition_id in range(self.num_partitions):
            # 为每个分区选择 num_replicas 个节点
            replicas = []
            for replica_idx in range(self.num_replicas):
                # 使用一致性的方式选择节点
                node_idx = (partition_id * self.num_replicas + replica_idx) % len(nodes_list)
                replicas.append(nodes_list[node_idx])
            assignments[partition_id] = replicas
            
        return assignments
    
    def compute_moves(self, old_assignment, new_assignment):
        """
        计算需要移动的分区
        返回: [(partition_id, from_node, to_node)]
        """
        moves = []
        
        for partition_id in new_assignment:
            old_nodes = set(old_assignment.get(partition_id, []))
            new_nodes = set(new_assignment[partition_id])
            
            # 找出需要添加的副本
            to_add = new_nodes - old_nodes
            # 找出需要删除的副本
            to_remove = old_nodes - new_nodes
            
            # 配对删除和添加，形成移动操作
            to_add_list = list(to_add)
            to_remove_list = list(to_remove)
            
            for i in range(min(len(to_add_list), len(to_remove_list))):
                moves.append((partition_id, to_remove_list[i], to_add_list[i]))
                
        return moves

# 使用示例
rebalancer = FixedPartitionRebalancer(num_partitions=256, num_replicas=3)

# 初始 4 个节点
old_nodes = ['node1', 'node2', 'node3', 'node4']
old_assignment = rebalancer.assign_partitions(old_nodes)

# 扩容到 6 个节点
new_nodes = ['node1', 'node2', 'node3', 'node4', 'node5', 'node6']
new_assignment = rebalancer.assign_partitions(new_nodes)

# 计算需要迁移的分区
moves = rebalancer.compute_moves(old_assignment, new_assignment)
print(f"需要迁移 {len(moves)} 个分区副本")

PUT /my_index
{
  "settings": {
    "number_of_shards": 5,
    "number_of_replicas": 1
  }
}

class DynamicPartition:
    def __init__(self, start_key, end_key, size_mb, qps):
        self.start_key = start_key
        self.end_key = end_key
        self.size_mb = size_mb
        self.qps = qps
        
class DynamicPartitionManager:
    def __init__(self, split_threshold_mb=96, merge_threshold_mb=32):
        self.split_threshold_mb = split_threshold_mb
        self.merge_threshold_mb = merge_threshold_mb
        self.partitions = []
        
    def check_split(self, partition):
        """
        检查是否需要分裂
        """
        if partition.size_mb > self.split_threshold_mb:
            # 找到中点 key
            mid_key = self.find_median_key(partition)
            
            # 创建两个新分区
            left = DynamicPartition(
                partition.start_key, 
                mid_key, 
                partition.size_mb / 2,
                partition.qps / 2
            )
            right = DynamicPartition(
                mid_key, 
                partition.end_key, 
                partition.size_mb / 2,
                partition.qps / 2
            )
            
            return (left, right)
        return None
    
    def check_merge(self, partition1, partition2):
        """
        检查两个相邻分区是否可以合并
        """
        if partition1.end_key != partition2.start_key:
            return None  # 不相邻
            
        total_size = partition1.size_mb + partition2.size_mb
        
        # 只有两个都很小，或者合并后不超过阈值时才合并
        if (partition1.size_mb < self.merge_threshold_mb and 
            partition2.size_mb < self.merge_threshold_mb and
            total_size < self.split_threshold_mb):
            
            merged = DynamicPartition(
                partition1.start_key,
                partition2.end_key,
                total_size,
                partition1.qps + partition2.qps
            )
            return merged
            
        return None
    
    def find_median_key(self, partition):
        """
        找到分区的中位数 key (简化实现)
        """
        # 实际实现需要扫描分区数据
        # 这里简化为中点
        if isinstance(partition.start_key, str) and isinstance(partition.end_key, str):
            # 字符串的中点
            return partition.start_key + (partition.end_key - partition.start_key) // 2
        return (partition.start_key + partition.end_key) // 2
    
    def rebalance(self):
        """
        执行再平衡：检查所有分区，执行必要的分裂和合并
        """
        i = 0
        while i < len(self.partitions):
            partition = self.partitions[i]
            
            # 检查分裂
            split_result = self.check_split(partition)
            if split_result:
                left, right = split_result
                print(f"分裂分区 [{partition.start_key}, {partition.end_key}) "
                      f"-> [{left.start_key}, {left.end_key}) 和 "
                      f"[{right.start_key}, {right.end_key})")
                self.partitions[i] = left
                self.partitions.insert(i + 1, right)
                i += 2
                continue
            
            # 检查合并
            if i + 1 < len(self.partitions):
                next_partition = self.partitions[i + 1]
                merge_result = self.check_merge(partition, next_partition)
                if merge_result:
                    print(f"合并分区 [{partition.start_key}, {partition.end_key}) 和 "
                          f"[{next_partition.start_key}, {next_partition.end_key}) "
                          f"-> [{merge_result.start_key}, {merge_result.end_key})")
                    self.partitions[i] = merge_result
                    self.partitions.pop(i + 1)
                    continue
            
            i += 1

# 使用示例
manager = DynamicPartitionManager()
manager.partitions = [
    DynamicPartition(0, 100, 120, 1000),    # 太大，需要分裂
    DynamicPartition(100, 200, 20, 100),    # 正常
    DynamicPartition(200, 250, 15, 50),     # 太小
    DynamicPartition(250, 280, 18, 60),     # 太小，可以与前一个合并
]

manager.rebalance()

def pre_split(key_range, num_partitions):
    """
    预先将 key 空间分裂成多个分区
    """
    start, end = key_range
    step = (end - start) // num_partitions
    
    partitions = []
    for i in range(num_partitions):
        partition_start = start + i * step
        partition_end = start + (i + 1) * step if i < num_partitions - 1 else end
        partitions.append(DynamicPartition(partition_start, partition_end, 0, 0))
    
    return partitions

# 将 key 空间 [0, 1000000) 预分裂成 10 个分区
initial_partitions = pre_split((0, 1000000), 10)

byte[][] splits = new byte[][] {
    Bytes.toBytes("row_100"),
    Bytes.toBytes("row_200"),
    Bytes.toBytes("row_300"),
    // ...
};
HTableDescriptor desc = new HTableDescriptor(TableName.valueOf("mytable"));
admin.createTable(desc, splits);

import hashlib
import bisect

class VirtualNode:
    def __init__(self, node_id, token, vnode_index):
        self.node_id = node_id
        self.token = token
        self.vnode_index = vnode_index
        
    def __repr__(self):
        return f"VNode({self.node_id}, token={self.token}, idx={self.vnode_index})"

class ConsistentHashRing:
    def __init__(self, vnodes_per_node=256):
        self.vnodes_per_node = vnodes_per_node
        self.ring = []  # 有序的 vnode 列表
        self.nodes = set()
        
    def add_node(self, node_id):
        """
        添加新节点：为其生成 vnodes_per_node 个 token
        """
        if node_id in self.nodes:
            return
            
        self.nodes.add(node_id)
        
        for i in range(self.vnodes_per_node):
            # 生成确定性的 token
            token_str = f"{node_id}:{i}"
            token = int(hashlib.md5(token_str.encode()).hexdigest(), 16)
            
            vnode = VirtualNode(node_id, token, i)
            bisect.insort(self.ring, (token, vnode))
        
        print(f"节点 {node_id} 加入，创建了 {self.vnodes_per_node} 个 vnode")
        print(f"总 vnode 数: {len(self.ring)}")
    
    def remove_node(self, node_id):
        """
        移除节点：删除其所有 vnode
        """
        if node_id not in self.nodes:
            return
            
        self.nodes.remove(node_id)
        self.ring = [(token, vnode) for token, vnode in self.ring 
                     if vnode.node_id != node_id]
        
        print(f"节点 {node_id} 移除，删除了 {self.vnodes_per_node} 个 vnode")
    
    def get_node(self, key):
        """
        根据 key 找到负责的节点
        """
        if not self.ring:
            return None
            
        key_hash = int(hashlib.md5(key.encode()).hexdigest(), 16)
        
        # 找到第一个 token >= key_hash 的 vnode
        idx = bisect.bisect_left([token for token, _ in self.ring], key_hash)
        
        if idx == len(self.ring):
            idx = 0  # 环形结构
            
        return self.ring[idx][1].node_id
    
    def get_distribution(self):
        """
        统计每个节点负责的 vnode 数量
        """
        distribution = {}
        for _, vnode in self.ring:
            distribution[vnode.node_id] = distribution.get(vnode.node_id, 0) + 1
        return distribution

# 使用示例
ring = ConsistentHashRing(vnodes_per_node=256)

# 初始 3 个节点
ring.add_node("node1")
ring.add_node("node2")
ring.add_node("node3")

print("\n初始分布:", ring.get_distribution())

# 添加第 4 个节点
ring.add_node("node4")
print("\n添加 node4 后分布:", ring.get_distribution())

# 测试 key 分布
keys = [f"key_{i}" for i in range(10000)]
node_counts = {}
for key in keys:
    node = ring.get_node(key)
    node_counts[node] = node_counts.get(node, 0) + 1

print("\n10000 个 key 的分布:")
for node, count in sorted(node_counts.items()):
    print(f"  {node}: {count} keys ({count/100:.1f}%)")

class PartitionMigration:
    def __init__(self, partition_id, source_node, target_node):
        self.partition_id = partition_id
        self.source_node = source_node
        self.target_node = target_node
        self.state = "init"
        self.copied_keys = set()
        self.pending_writes = []
        
    def phase1_copy(self):
        """
        阶段 1: 复制数据
        """
        print(f"开始复制分区 {self.partition_id} 从 {self.source_node} 到 {self.target_node}")
        self.state = "copying"
        
        # 获取源节点的所有 key
        keys = self.source_node.get_keys(self.partition_id)
        
        # 批量复制数据
        batch_size = 1000
        for i in range(0, len(keys), batch_size):
            batch = keys[i:i+batch_size]
            
            # 读取源节点数据
            data = self.source_node.read_batch(batch)
            
            # 写入目标节点
            self.target_node.write_batch(data)
            
            # 记录已复制的 key
            self.copied_keys.update(batch)
            
            print(f"  已复制 {len(self.copied_keys)}/{len(keys)} 个 key")
        
        print(f"分区 {self.partition_id} 数据复制完成")
        
    def phase2_switchover(self):
        """
        阶段 2: 切换
        """
        print(f"开始切换分区 {self.partition_id}")
        self.state = "switching"
        
        # 1. 短暂阻塞写操作（设置只读模式）
        self.source_node.set_readonly(self.partition_id)
        
        # 2. 复制增量数据（在只读期间产生的写请求）
        incremental_data = self.source_node.get_incremental_writes(self.partition_id)
        if incremental_data:
            self.target_node.write_batch(incremental_data)
            print(f"  复制了 {len(incremental_data)} 个增量写入")
        
        # 3. 更新元数据服务器：分区所有权转移
        metadata_service.update_partition_owner(self.partition_id, self.target_node)
        
        # 4. 目标节点开始服务
        self.target_node.set_active(self.partition_id)
        
        # 5. 源节点停止服务
        self.source_node.remove_partition(self.partition_id)
        
        self.state = "completed"
        print(f"分区 {self.partition_id} 切换完成")
        
    def execute(self):
        """
        执行完整的迁移流程
        """
        try:
            self.phase1_copy()
            self.phase2_switchover()
            return True
        except Exception as e:
            print(f"迁移失败: {e}")
            self.state = "failed"
            return False

class PartitionRouter:
    def __init__(self):
        self.routing_table = {}  # {partition_id: (node, epoch)}
        self.current_epoch = 0
        
    def get_node(self, partition_id):
        """
        获取分区的当前负责节点
        """
        return self.routing_table.get(partition_id)
    
    def start_migration(self, partition_id, source_node, target_node):
        """
        开始迁移：增加 epoch，设置双节点模式
        """
        self.current_epoch += 1
        self.routing_table[partition_id] = {
            'source': source_node,
            'target': target_node,
            'epoch': self.current_epoch,
            'state': 'migrating'
        }
        
    def finish_migration(self, partition_id, target_node):
        """
        完成迁移：切换到目标节点
        """
        self.current_epoch += 1
        self.routing_table[partition_id] = {
            'source': None,
            'target': target_node,
            'epoch': self.current_epoch,
            'state': 'active'
        }
    
    def read(self, partition_id, key):
        """
        读取数据，处理迁移期间的情况
        """
        info = self.routing_table.get(partition_id)
        if not info:
            raise Exception(f"分区 {partition_id} 不存在")
        
        if info['state'] == 'active':
            # 正常状态：读目标节点
            return info['target'].read(key)
        
        elif info['state'] == 'migrating':
            # 迁移中：优先读源节点，如果源节点没有则读目标节点
            value = info['source'].read(key)
            if value is None:
                value = info['target'].read(key)
            return value
        
    def write(self, partition_id, key, value):
        """
        写入数据，处理迁移期间的情况
        """
        info = self.routing_table.get(partition_id)
        if not info:
            raise Exception(f"分区 {partition_id} 不存在")
        
        if info['state'] == 'active':
            # 正常状态：写目标节点
            info['target'].write(key, value)
        
        elif info['state'] == 'migrating':
            # 迁移中：双写
            info['source'].write(key, value)
            info['target'].write(key, value)

class IdempotentMigration:
    def copy_key_idempotent(self, key, value, version):
        """
        幂等的数据复制
        """
        existing = self.target_node.read_with_version(key)
        
        if existing is None:
            # 目标节点没有这个 key，直接写入
            self.target_node.write_with_version(key, value, version)
        elif existing['version'] < version:
            # 目标节点的版本较旧，覆盖
            self.target_node.write_with_version(key, value, version)
        else:
            # 目标节点的版本更新或相同，跳过
            pass

type StoreHeartbeat struct {
    StoreID       uint64
    Capacity      uint64  // 总容量
    Available     uint64  // 可用容量
    RegionCount   uint32  // Region 数量
    LeaderCount   uint32  // Leader Region 数量
    BytesWritten  uint64  // 写入字节数
    BytesRead     uint64  // 读取字节数
    KeysWritten   uint64  // 写入 key 数
    KeysRead      uint64  // 读取 key 数
    IsBusy        bool    // 是否繁忙
    StartTime     uint64  // 启动时间
}

type RegionHeartbeat struct {
    RegionID      uint64
    Leader        *Peer      // 当前 Leader
    Peers         []*Peer    // 所有副本
    StartKey      []byte     // 起始 key
    EndKey        []byte     // 结束 key
    ApproximateSize uint64   // 大小估算
    ApproximateKeys uint64   // key 数量估算
    BytesWritten  uint64     // 写入速率
    BytesRead     uint64     // 读取速率
    KeysWritten   uint64
    KeysRead      uint64
}

class BalanceLeaderScheduler:
    def __init__(self):
        self.name = "balance-leader"
        
    def schedule(self, cluster_info):
        """
        生成 Leader 平衡的调度操作
        """
        stores = cluster_info.get_stores()
        
        # 统计每个节点的 Leader 数量
        leader_counts = {}
        for store in stores:
            leader_counts[store.id] = store.leader_count
        
        # 找到 Leader 最多和最少的节点
        max_store = max(stores, key=lambda s: s.leader_count)
        min_store = min(stores, key=lambda s: s.leader_count)
        
        # 如果差距太小，不需要调度
        if max_store.leader_count - min_store.leader_count <= 1:
            return None
        
        # 从 max_store 选择一个 Region，将其 Leader 转移到 min_store
        candidate_region = self.select_region(cluster_info, max_store, min_store)
        
        if candidate_region:
            return {
                'type': 'transfer-leader',
                'region_id': candidate_region.id,
                'from_store': max_store.id,
                'to_store': min_store.id
            }
        
        return None
    
    def select_region(self, cluster_info, from_store, to_store):
        """
        选择合适的 Region 进行 Leader 转移
        """
        # 找到在 from_store 上是 Leader，且在 to_store 上有副本的 Region
        for region in cluster_info.get_regions():
            if region.leader_store_id == from_store.id:
                for peer in region.peers:
                    if peer.store_id == to_store.id:
                        return region
        return None

class BalanceRegionScheduler:
    def __init__(self):
        self.name = "balance-region"
        
    def schedule(self, cluster_info):
        """
        生成 Region 平衡的调度操作
        """
        stores = cluster_info.get_stores()
        
        # 按数据量排序
        stores_by_size = sorted(stores, key=lambda s: s.used_size, reverse=True)
        
        source_store = stores_by_size[0]
        target_store = stores_by_size[-1]
        
        # 检查是否需要平衡
        avg_size = sum(s.used_size for s in stores) / len(stores)
        
        if source_store.used_size < avg_size * 1.1:
            return None  # 已经足够均衡
        
        # 选择一个 Region 从 source_store 移动到 target_store
        candidate_region = self.select_region(cluster_info, source_store, target_store)
        
        if candidate_region:
            return {
                'type': 'move-peer',
                'region_id': candidate_region.id,
                'from_store': source_store.id,
                'to_store': target_store.id
            }
        
        return None
    
    def select_region(self, cluster_info, from_store, to_store):
        """
        选择合适的 Region 进行迁移
        """
        # 找到在 from_store 上有副本的 Region
        candidates = []
        for region in cluster_info.get_regions():
            if any(peer.store_id == from_store.id for peer in region.peers):
                # 检查 to_store 上是否还没有这个 Region 的副本
                if not any(peer.store_id == to_store.id for peer in region.peers):
                    candidates.append(region)
        
        # 优先选择较大的 Region，加快平衡速度
        if candidates:
            return max(candidates, key=lambda r: r.approximate_size)
        
        return None

class HotRegionScheduler:
    def __init__(self, hot_threshold_bytes=10 * 1024 * 1024):  # 10MB/s
        self.name = "hot-region"
        self.hot_threshold = hot_threshold_bytes
        self.hot_cache = {}  # {region_id: hot_info}
        
    def update_hot_cache(self, region_heartbeat):
        """
        更新热点缓存
        """
        region_id = region_heartbeat.region_id
        
        # 计算读写速率 (字节/秒)
        write_rate = region_heartbeat.bytes_written
        read_rate = region_heartbeat.bytes_read
        
        total_rate = write_rate + read_rate
        
        if total_rate > self.hot_threshold:
            self.hot_cache[region_id] = {
                'region_id': region_id,
                'store_id': region_heartbeat.leader.store_id,
                'write_rate': write_rate,
                'read_rate': read_rate,
                'total_rate': total_rate,
                'last_update': time.time()
            }
        else:
            # 不再热，从缓存中移除
            self.hot_cache.pop(region_id, None)
    
    def schedule(self, cluster_info):
        """
        生成热点平衡的调度操作
        """
        # 统计每个节点的热点负载
        store_hot_load = {}
        for store in cluster_info.get_stores():
            store_hot_load[store.id] = 0
        
        for hot_info in self.hot_cache.values():
            store_id = hot_info['store_id']
            store_hot_load[store_id] += hot_info['total_rate']
        
        # 找到最热和最冷的节点
        if not store_hot_load:
            return None
        
        hottest_store = max(store_hot_load.items(), key=lambda x: x[1])
        coldest_store = min(store_hot_load.items(), key=lambda x: x[1])
        
        # 检查是否需要调度
        if hottest_store[1] < coldest_store[1] * 2:
            return None  # 热点分布相对均衡
        
        # 从 hottest_store 选择一个热点 Region，转移 Leader 到 coldest_store
        hot_region = self.select_hot_region(cluster_info, hottest_store[0], coldest_store[0])
        
        if hot_region:
            return {
                'type': 'transfer-leader',
                'region_id': hot_region['region_id'],
                'from_store': hottest_store[0],
                'to_store': coldest_store[0],
                'reason': 'hot-region'
            }
        
        return None
    
    def select_hot_region(self, cluster_info, from_store, to_store):
        """
        选择热点 Region
        """
        for hot_info in self.hot_cache.values():
            if hot_info['store_id'] == from_store:
                # 检查这个 Region 在 to_store 上是否有副本
                region = cluster_info.get_region(hot_info['region_id'])
                if any(peer.store_id == to_store for peer in region.peers):
                    return hot_info
        return None

[server]
labels = "zone=z1,rack=r1,host=h1"

{
  "group_id": "pd",
  "id": "default",
  "role": "voter",
  "count": 3,
  "location_labels": ["zone", "rack", "host"],
  "constraints": [
    {"key": "zone", "op": "in", "values": ["z1", "z2", "z3"]}
  ]
}

class PlacementChecker:
    def check_placement(self, region, placement_rule):
        """
        检查 Region 的副本是否满足放置规则
        """
        peers = region.peers
        
        # 检查副本数量
        if len(peers) != placement_rule['count']:
            return False, "副本数量不匹配"
        
        # 检查位置隔离
        location_labels = placement_rule['location_labels']
        
        for label in location_labels:
            label_values = [self.get_store_label(peer.store_id, label) for peer in peers]
            
            # 同一层级的标签值应该各不相同（隔离）
            if len(label_values) != len(set(label_values)):
                return False, f"标签 {label} 没有隔离"
        
        # 检查约束条件
        for constraint in placement_rule.get('constraints', []):
            if not self.check_constraint(peers, constraint):
                return False, f"违反约束 {constraint}"
        
        return True, "满足放置规则"
    
    def check_constraint(self, peers, constraint):
        """
        检查约束条件
        """
        key = constraint['key']
        op = constraint['op']
        values = constraint['values']
        
        for peer in peers:
            label_value = self.get_store_label(peer.store_id, key)
            
            if op == 'in':
                if label_value not in values:
                    return False
            elif op == 'not_in':
                if label_value in values:
                    return False
        
        return True

class SchedulerManager:
    def __init__(self):
        self.schedulers = []
        self.max_pending_operators = 10  # 最多同时执行 10 个操作
        self.max_snapshot_count = 3      # 最多同时传输 3 个快照
        self.pending_operators = []
        
    def run_schedulers(self, cluster_info):
        """
        运行所有调度器，生成调度操作
        """
        if len(self.pending_operators) >= self.max_pending_operators:
            return  # 已达到并发上限
        
        for scheduler in self.schedulers:
            operator = scheduler.schedule(cluster_info)
            
            if operator:
                # 检查是否可以添加这个操作
                if self.can_add_operator(operator):
                    self.pending_operators.append(operator)
                    print(f"添加调度操作: {operator}")
                    
                    if len(self.pending_operators) >= self.max_pending_operators:
                        break
    
    def can_add_operator(self, operator):
        """
        检查是否可以添加新的调度操作
        """
        # 检查是否涉及正在迁移的 Region
        for pending in self.pending_operators:
            if pending['region_id'] == operator['region_id']:
                return False  # 同一个 Region 不能同时有多个操作
        
        # 检查快照传输数量
        if operator['type'] == 'move-peer':
            snapshot_count = sum(1 for op in self.pending_operators if op['type'] == 'move-peer')
            if snapshot_count >= self.max_snapshot_count:
                return False
        
        return True

type Operator struct {
    RegionID uint64
    Kind     OperatorKind  // transfer-leader, add-peer, remove-peer, etc.
    Steps    []OperatorStep
    Status   OperatorStatus
}

type OperatorStep struct {
    Type     StepType  // add_learner, promote_learner, transfer_leader, remove_peer
    PeerID   uint64
    StoreID  uint64
}

// 示例：将 Region 1 从 Store 1 迁移到 Store 4
operator := Operator{
    RegionID: 1,
    Kind: OpRegion,
    Steps: []OperatorStep{
        {Type: AddLearner, StoreID: 4},     // 在 Store 4 上添加 Learner
        {Type: PromoteLearner, StoreID: 4}, // 提升为 Follower
        {Type: RemovePeer, StoreID: 1},     // 移除 Store 1 的副本
    },
}

import time

class RateLimiter:
    def __init__(self, rate_bytes_per_sec):
        self.rate = rate_bytes_per_sec
        self.allowance = rate_bytes_per_sec
        self.last_check = time.time()
        
    def acquire(self, bytes_count):
        """
        获取指定字节数的配额，如果超限则阻塞
        """
        current_time = time.time()
        elapsed = current_time - self.last_check
        
        # 补充配额
        self.allowance += elapsed * self.rate
        self.allowance = min(self.allowance, self.rate)  # 不超过上限
        
        self.last_check = current_time
        
        if self.allowance >= bytes_count:
            # 有足够配额，直接消耗
            self.allowance -= bytes_count
            return
        else:
            # 配额不足，需要等待
            wait_time = (bytes_count - self.allowance) / self.rate
            time.sleep(wait_time)
            self.allowance = 0

class ThrottledMigration:
    def __init__(self, rate_limit_mbps=100):
        self.rate_limiter = RateLimiter(rate_limit_mbps * 1024 * 1024)
        
    def migrate_data(self, source, target, partition_id):
        """
        带限流的数据迁移
        """
        keys = source.get_keys(partition_id)
        batch_size = 1000
        
        for i in range(0, len(keys), batch_size):
            batch = keys[i:i+batch_size]
            
            # 读取数据
            data = source.read_batch(batch)
            data_size = sum(len(v) for v in data.values())
            
            # 申请配额
            self.rate_limiter.acquire(data_size)
            
            # 写入目标节点
            target.write_batch(data)

class AdaptiveThrottle:
    def __init__(self, max_rate_mbps=200):
        self.max_rate = max_rate_mbps * 1024 * 1024
        self.current_rate = self.max_rate * 0.5  # 从 50% 开始
        
    def adjust_rate(self, system_metrics):
        """
        根据系统指标调整迁移速率
        """
        cpu_usage = system_metrics['cpu_usage']
        disk_util = system_metrics['disk_util']
        p99_latency = system_metrics['p99_latency_ms']
        
        # 如果系统压力大，降低迁移速率
        if cpu_usage > 0.8 or disk_util > 0.8 or p99_latency > 100:
            self.current_rate *= 0.9  # 降低 10%
        # 如果系统压力小，提高迁移速率
        elif cpu_usage < 0.5 and disk_util < 0.5 and p99_latency < 50:
            self.current_rate *= 1.1  # 提高 10%
        
        # 限制在合理范围内
        self.current_rate = max(10 * 1024 * 1024, self.current_rate)  # 至少 10MB/s
        self.current_rate = min(self.max_rate, self.current_rate)      # 不超过上限
        
        return self.current_rate

class TimeWindowScheduler:
    def __init__(self, allowed_hours):
        self.allowed_hours = allowed_hours  # 例如 [0, 1, 2, 3, 4, 5] 表示凌晨 0-6 点
        
    def is_allowed(self):
        """
        检查当前是否在允许的时间窗口内
        """
        current_hour = time.localtime().tm_hour
        return current_hour in self.allowed_hours
    
    def schedule_if_allowed(self, scheduler, cluster_info):
        """
        只在允许的时间窗口内执行调度
        """
        if self.is_allowed():
            return scheduler.schedule(cluster_info)
        else:
            return None

class ImbalanceDetector:
    def __init__(self, load_threshold=0.3, min_gap_abs=100):
        self.load_threshold = load_threshold
        self.min_gap_abs = min_gap_abs

    def should_rebalance(self, node_loads: dict) -> bool:
        """
        检测是否需要再平衡。
        node_loads: {node_id: load_score}
        load_threshold: 最大/最小负载比超过此值时触发。
        min_gap_abs: 最大与最小负载的绝对差值下限，避免低负载时误触发。
        """
        if len(node_loads) < 2:
            return False
        max_load = max(node_loads.values())
        min_load = min(node_loads.values())
        if min_load == 0:
            return max_load > self.min_gap_abs
        ratio = (max_load - min_load) / min_load
        return ratio > self.load_threshold and (max_load - min_load) > self.min_gap_abs

class MigrationBackpressure:
    def __init__(self, latency_threshold_ms=100, cpu_threshold=0.85):
        self.latency_threshold_ms = latency_threshold_ms
        self.cpu_threshold = cpu_threshold
        self.paused = False

    def check_and_adjust(self, metrics: dict, rate_limiter) -> str:
        """
        根据业务指标动态调整迁移速率。
        metrics: {"p99_latency_ms": float, "cpu_usage": float, "error_rate": float}
        """
        if metrics["error_rate"] > 0.01:
            self.paused = True
            return "暂停迁移：错误率过高"

        if metrics["p99_latency_ms"] > self.latency_threshold_ms:
            rate_limiter.reduce_rate(0.5)
            return "降速50%：P99延迟超标"

        if metrics["cpu_usage"] > self.cpu_threshold:
            rate_limiter.reduce_rate(0.7)
            return "降速30%：CPU使用率过高"

        if not self.paused and metrics["p99_latency_ms"] < self.latency_threshold_ms * 0.5:
            rate_limiter.increase_rate(1.2)
            return "提速20%：系统负载健康"

        return "维持当前速率"

import random
import hashlib
from typing import List, Dict, Set
from dataclasses import dataclass

@dataclass
class Node:
    id: str
    capacity_mb: int
    used_mb: int = 0
    partitions: Set[int] = None
    
    def __post_init__(self):
        if self.partitions is None:
            self.partitions = set()
    
    def available_mb(self):
        return self.capacity_mb - self.used_mb
    
    def usage_ratio(self):
        return self.used_mb / self.capacity_mb if self.capacity_mb > 0 else 0

@dataclass
class Partition:
    id: int
    size_mb: int
    qps: int
    key_range: tuple  # (start, end)

class RebalanceSimulator:
    def __init__(self, strategy='fixed'):
        self.strategy = strategy
        self.nodes: Dict[str, Node] = {}
        self.partitions: Dict[int, Partition] = {}
        self.partition_to_node: Dict[int, str] = {}
        
    def add_node(self, node: Node):
        self.nodes[node.id] = node
        print(f"添加节点 {node.id}, 容量 {node.capacity_mb}MB")
        
    def remove_node(self, node_id: str):
        if node_id not in self.nodes:
            return
        
        # 将该节点的分区迁移到其他节点
        partitions_to_move = list(self.nodes[node_id].partitions)
        for partition_id in partitions_to_move:
            self.move_partition(partition_id, node_id, self.select_target_node(partition_id))
        
        del self.nodes[node_id]
        print(f"移除节点 {node_id}")
        
    def add_partition(self, partition: Partition):
        self.partitions[partition.id] = partition
        # 选择负载最轻的节点
        target_node = self.select_target_node(partition.id)
        self.assign_partition(partition.id, target_node)
        
    def assign_partition(self, partition_id: int, node_id: str):
        partition = self.partitions[partition_id]
        node = self.nodes[node_id]
        
        self.partition_to_node[partition_id] = node_id
        node.partitions.add(partition_id)
        node.used_mb += partition.size_mb
        
    def move_partition(self, partition_id: int, from_node_id: str, to_node_id: str):
        if from_node_id == to_node_id:
            return
            
        partition = self.partitions[partition_id]
        from_node = self.nodes[from_node_id]
        to_node = self.nodes[to_node_id]
        
        # 从源节点移除
        from_node.partitions.remove(partition_id)
        from_node.used_mb -= partition.size_mb
        
        # 添加到目标节点
        to_node.partitions.add(partition_id)
        to_node.used_mb += partition.size_mb
        
        self.partition_to_node[partition_id] = to_node_id
        
        print(f"迁移分区 {partition_id} ({partition.size_mb}MB) "
              f"从 {from_node_id} 到 {to_node_id}")
        
    def select_target_node(self, partition_id: int) -> str:
        """
        选择分区的目标节点
        """
        # 排除已经有这个分区的节点
        current_node = self.partition_to_node.get(partition_id)
        
        # 选择使用率最低的节点
        min_usage = float('inf')
        target = None
        
        for node_id, node in self.nodes.items():
            if node_id == current_node:
                continue
            if node.usage_ratio() < min_usage:
                min_usage = node.usage_ratio()
                target = node_id
        
        return target
        
    def rebalance(self):
        """
        执行再平衡
        """
        print(f"\n=== 开始再平衡 (策略: {self.strategy}) ===")
        
        if self.strategy == 'fixed':
            self.rebalance_fixed()
        elif self.strategy == 'dynamic':
            self.rebalance_dynamic()
        elif self.strategy == 'proportional':
            self.rebalance_proportional()
            
        self.print_status()
        
    def rebalance_fixed(self):
        """
        固定分区策略：移动整个分区以平衡负载
        """
        # 计算平均使用率
        total_used = sum(node.used_mb for node in self.nodes.values())
        total_capacity = sum(node.capacity_mb for node in self.nodes.values())
        avg_usage = total_used / total_capacity
        
        # 找到过载和空闲的节点
        overloaded = [node for node in self.nodes.values() if node.usage_ratio() > avg_usage * 1.2]
        underloaded = [node for node in self.nodes.values() if node.usage_ratio() < avg_usage * 0.8]
        
        # 从过载节点迁移分区到空闲节点
        for over_node in overloaded:
            for under_node in underloaded:
                # 选择合适大小的分区
                for partition_id in list(over_node.partitions):
                    partition = self.partitions[partition_id]
                    
                    if under_node.available_mb() >= partition.size_mb:
                        self.move_partition(partition_id, over_node.id, under_node.id)
                        
                        # 检查是否已经平衡
                        if over_node.usage_ratio() <= avg_usage * 1.1:
                            break
                            
                if over_node.usage_ratio() <= avg_usage * 1.1:
                    break
                    
    def rebalance_dynamic(self):
        """
        动态分区策略：分裂大分区，合并小分区
        """
        split_threshold = 100  # MB
        merge_threshold = 20   # MB
        
        # 检查分裂
        to_split = []
        for partition_id, partition in self.partitions.items():
            if partition.size_mb > split_threshold:
                to_split.append(partition_id)
        
        for partition_id in to_split:
            self.split_partition(partition_id)
        
        # 检查合并
        sorted_partitions = sorted(self.partitions.values(), key=lambda p: p.key_range[0])
        i = 0
        while i < len(sorted_partitions) - 1:
            p1 = sorted_partitions[i]
            p2 = sorted_partitions[i + 1]
            
            if (p1.size_mb < merge_threshold and p2.size_mb < merge_threshold and
                p1.key_range[1] == p2.key_range[0]):
                self.merge_partitions(p1.id, p2.id)
                sorted_partitions = sorted(self.partitions.values(), key=lambda p: p.key_range[0])
            else:
                i += 1
                
    def split_partition(self, partition_id: int):
        """
        分裂分区
        """
        partition = self.partitions[partition_id]
        node_id = self.partition_to_node[partition_id]
        
        # 找到分裂点
        start, end = partition.key_range
        mid = (start + end) // 2
        
        # 创建两个新分区
        new_id_1 = max(self.partitions.keys()) + 1
        new_id_2 = new_id_1 + 1
        
        p1 = Partition(new_id_1, partition.size_mb // 2, partition.qps // 2, (start, mid))
        p2 = Partition(new_id_2, partition.size_mb // 2, partition.qps // 2, (mid, end))
        
        # 移除旧分区
        node = self.nodes[node_id]
        node.partitions.remove(partition_id)
        node.used_mb -= partition.size_mb
        del self.partitions[partition_id]
        del self.partition_to_node[partition_id]
        
        # 添加新分区
        self.partitions[new_id_1] = p1
        self.partitions[new_id_2] = p2
        self.assign_partition(new_id_1, node_id)
        
        # 第二个分区可能分配到其他节点以平衡负载
        target = self.select_target_node(new_id_2)
        self.assign_partition(new_id_2, target)
        
        print(f"分裂分区 {partition_id} ({partition.size_mb}MB) "
              f"-> {new_id_1} 和 {new_id_2} (各 {partition.size_mb//2}MB)")
        
    def merge_partitions(self, partition_id_1: int, partition_id_2: int):
        """
        合并分区
        """
        p1 = self.partitions[partition_id_1]
        p2 = self.partitions[partition_id_2]
        
        # 创建合并后的分区
        new_id = max(self.partitions.keys()) + 1
        merged = Partition(
            new_id,
            p1.size_mb + p2.size_mb,
            p1.qps + p2.qps,
            (p1.key_range[0], p2.key_range[1])
        )
        
        # 移除旧分区
        for pid in [partition_id_1, partition_id_2]:
            node_id = self.partition_to_node[pid]
            node = self.nodes[node_id]
            node.partitions.remove(pid)
            node.used_mb -= self.partitions[pid].size_mb
            del self.partition_to_node[pid]
            del self.partitions[pid]
        
        # 添加新分区
        self.partitions[new_id] = merged
        target = self.select_target_node(new_id)
        self.assign_partition(new_id, target)
        
        print(f"合并分区 {partition_id_1} 和 {partition_id_2} -> {new_id} ({merged.size_mb}MB)")
        
    def rebalance_proportional(self):
        """
        按比例分区策略：保持每个节点的分区数量一致
        """
        target_partitions_per_node = len(self.partitions) // len(self.nodes)
        
        overloaded = [node for node in self.nodes.values() 
                      if len(node.partitions) > target_partitions_per_node * 1.2]
        underloaded = [node for node in self.nodes.values() 
                       if len(node.partitions) < target_partitions_per_node * 0.8]
        
        for over_node in overloaded:
            for under_node in underloaded:
                while (len(over_node.partitions) > target_partitions_per_node and
                       len(under_node.partitions) < target_partitions_per_node):
                    # 随机选择一个分区迁移
                    partition_id = random.choice(list(over_node.partitions))
                    self.move_partition(partition_id, over_node.id, under_node.id)
                    
    def print_status(self):
        """
        打印集群状态
        """
        print("\n=== 集群状态 ===")
        for node_id, node in sorted(self.nodes.items()):
            print(f"{node_id}: {node.used_mb}/{node.capacity_mb}MB "
                  f"({node.usage_ratio()*100:.1f}%), "
                  f"{len(node.partitions)} 个分区")

# 测试不同策略
def test_rebalance():
    for strategy in ['fixed', 'dynamic', 'proportional']:
        print(f"\n{'='*60}")
        print(f"测试策略: {strategy}")
        print(f"{'='*60}")
        
        sim = RebalanceSimulator(strategy=strategy)
        
        # 初始 3 个节点
        sim.add_node(Node('node1', 1000))
        sim.add_node(Node('node2', 1000))
        sim.add_node(Node('node3', 1000))
        
        # 添加一些分区
        for i in range(20):
            size = random.randint(30, 150)
            qps = random.randint(100, 1000)
            partition = Partition(i, size, qps, (i*1000, (i+1)*1000))
            sim.add_partition(partition)
        
        print("\n初始状态:")
        sim.print_status()
        
        # 添加新节点触发再平衡
        sim.add_node(Node('node4', 1000))
        sim.rebalance()

if __name__ == '__main__':
    test_rebalance()