def move_leadership(self, partition, new_leader): """Return a new state that is the result of changing the leadership of a single partition. :param partition: The partition index of the partition to change the leadership of. :param new_leader: The broker index of the new leader replica. """ new_state = copy(self) # Update the partition replica tuple source = new_state.replicas[partition][0] new_leader_index = self.replicas[partition].index(new_leader) new_state.replicas = tuple_alter( self.replicas, (partition, lambda replicas: tuple_replace( replicas, (0, replicas[new_leader_index]), (new_leader_index, replicas[0]), )), ) # Update the broker leader weights partition_weight = self.partition_weights[partition] new_state.broker_leader_weights = tuple_alter( self.broker_leader_weights, (source, lambda leader_weight: leader_weight - partition_weight), (new_leader, lambda leader_weight: leader_weight + partition_weight), ) # Update the total leader movement size new_state.leader_movement_count += 1 return new_state
def test_tuple_replace(): result = tuple_replace( (1, 2, 3, 2, 1), (2, 1), (3, 2), (4, 3), ) assert result == (1, 2, 1, 2, 3)
def move_leadership(self, partition, new_leader): """Return a new state that is the result of changing the leadership of a single partition. :param partition: The partition index of the partition to change the leadership of. :param new_leader: The broker index of the new leader replica. """ new_state = copy(self) # Update the partition replica tuple source = new_state.replicas[partition][0] new_leader_index = self.replicas[partition].index(new_leader) new_state.replicas = tuple_alter( self.replicas, (partition, lambda replicas: tuple_replace( replicas, (0, replicas[new_leader_index]), (new_leader_index, replicas[0]), )), ) new_state.pending_partitions = self.pending_partitions + (partition, ) # Update the leader count new_state.broker_leader_counts = tuple_alter( self.broker_leader_counts, (source, lambda leader_count: leader_count - 1), (new_leader, lambda leader_count: leader_count + 1), ) # Update the broker leader weights partition_weight = self.partition_weights[partition] new_state.broker_leader_weights = tuple_alter( self.broker_leader_weights, (source, lambda leader_weight: leader_weight - partition_weight), (new_leader, lambda leader_weight: leader_weight + partition_weight), ) # Update the total leader movement size new_state.leader_movement_count += 1 return new_state
def add_replica(self, partition, broker): new_state = copy(self) # Add replica to partition replica tuple new_state.replicas = tuple_alter( self.replicas, (partition, lambda replicas: replicas + (broker, )), ) new_state.pending_partitions = self.pending_partitions + (partition, ) # Update the broker partition count new_state.broker_partition_counts = tuple_alter( self.broker_partition_counts, (broker, lambda partition_count: partition_count + 1), ) # Update the broker weight partition_weight = self.partition_weights[partition] new_state.broker_weights = tuple_alter( self.broker_weights, (broker, lambda broker_weight: broker_weight + partition_weight), ) # Update the topic weights topic = new_state.partition_topic[partition] new_state.topic_weights = tuple_alter( self.topic_weights, (topic, lambda topic_weight: topic_weight + partition_weight) ) # Update the total weight new_state.total_weight = self.total_weight + partition_weight # Update the topic broker counts topic = self.partition_topic[partition] new_state.topic_broker_count = tuple_alter( self.topic_broker_count, (topic, lambda broker_counts: tuple_alter( broker_counts, (broker, lambda count: count + 1), )), ) # Update topic replica count new_state.topic_replica_count = tuple_alter( self.topic_replica_count, (topic, lambda replica_count: replica_count + 1), ) # Update the topic broker imbalance new_state.topic_broker_imbalance = tuple_replace( self.topic_broker_imbalance, (topic, new_state._calculate_topic_imbalance(topic)), ) new_state._weighted_topic_broker_imbalance = ( self._weighted_topic_broker_imbalance - self.topic_weights[topic] * self.topic_broker_imbalance[topic] + new_state.topic_weights[topic] * new_state.topic_broker_imbalance[topic] ) # Update the replication group replica counts rg = self.broker_rg[broker] new_state.rg_replicas = tuple_alter( self.rg_replicas, (rg, lambda replica_counts: tuple_alter( replica_counts, (partition, lambda count: count + 1), )), ) return new_state
def move(self, partition, source, dest): """Return a new state that is the result of moving a single partition. :param partition: The partition index of the partition to move. :param source: The broker index of the broker to move the partition from. :param dest: The broker index of the broker to move the partition to. """ new_state = copy(self) # Update the partition replica tuple source_index = self.replicas[partition].index(source) new_state.replicas = tuple_alter( self.replicas, (partition, lambda replicas: tuple_replace( replicas, (source_index, dest), )), ) new_state.pending_partitions = self.pending_partitions + (partition, ) # Update the broker weights partition_weight = self.partition_weights[partition] new_state.broker_weights = tuple_alter( self.broker_weights, (source, lambda broker_weight: broker_weight - partition_weight), (dest, lambda broker_weight: broker_weight + partition_weight), ) # Update the broker partition count new_state.broker_partition_counts = tuple_alter( self.broker_partition_counts, (source, lambda partition_count: partition_count - 1), (dest, lambda partition_count: partition_count + 1), ) # Update the broker leader weights if source_index == 0: new_state.broker_leader_weights = tuple_alter( self.broker_leader_weights, (source, lambda lw: lw - partition_weight), (dest, lambda lw: lw + partition_weight), ) new_state.broker_leader_counts = tuple_alter( self.broker_leader_counts, (source, lambda leader_count: leader_count - 1), (dest, lambda leader_count: leader_count + 1), ) new_state.leader_movement_count += 1 # Update the topic broker counts topic = self.partition_topic[partition] new_state.topic_broker_count = tuple_alter( self.topic_broker_count, (topic, lambda broker_count: tuple_alter( broker_count, (source, lambda count: count - 1), (dest, lambda count: count + 1), )), ) # Update the topic broker imbalance new_state.topic_broker_imbalance = tuple_replace( self.topic_broker_imbalance, (topic, new_state._calculate_topic_imbalance(topic)), ) new_state._weighted_topic_broker_imbalance = ( self._weighted_topic_broker_imbalance + self.topic_weights[topic] * ( new_state.topic_broker_imbalance[topic] - self.topic_broker_imbalance[topic] ) ) # Update the replication group replica counts source_rg = self.broker_rg[source] dest_rg = self.broker_rg[dest] if source_rg != dest_rg: new_state.rg_replicas = tuple_alter( self.rg_replicas, (source_rg, lambda replica_counts: tuple_alter( replica_counts, (partition, lambda replica_count: replica_count - 1), )), (dest_rg, lambda replica_counts: tuple_alter( replica_counts, (partition, lambda replica_count: replica_count + 1), )), ) # Update the movement sizes new_state.movement_size += self.partition_sizes[partition] new_state.movement_count += 1 return new_state