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_alter():
result = tuple_alter(
(1, 2, 3, 2, 1),
(2, lambda x: x + 1),
(3, lambda x: x + 2),
(4, lambda x: x + 3),
)
assert result == (1, 2, 4, 4, 4)
def test_tuple_alter():
result = tuple_alter(
(1, 2, 3, 2, 1),
(2, lambda x: x + 1),
(3, lambda x: x + 2),
(4, lambda x: x + 3),
)
assert result == (1, 2, 4, 4, 4)
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
