Code review comment: `checkpointStore.claimOwnership()` method takes a list of partition ownership instances.

Code before review:
private Mono<Void> loadBalance(final Tuple2<Map<String, PartitionOwnership>, List<String>> tuple) {
return Mono.fromRunnable(() -> {
logger.info("Starting load balancer");
Map<String, PartitionOwnership> partitionOwnershipMap = tuple.getT1();
List<String> partitionIds = tuple.getT2();
if (CoreUtils.isNullOrEmpty(partitionIds)) {
throw logger.logExceptionAsError(Exceptions.propagate(
new IllegalStateException("There are no partitions in Event Hub " + eventHubName)));
}
int numberOfPartitions = partitionIds.size();
logger.info("Partition manager returned {} ownership records", partitionOwnershipMap.size());
logger.info("EventHubAsyncClient returned {} partitions", numberOfPartitions);
if (!isValid(partitionOwnershipMap)) {
throw logger.logExceptionAsError(Exceptions.propagate(
new IllegalStateException("Invalid partitionOwnership data from CheckpointStore")));
}
/*
* Remove all partitions' ownership that have not been modified for a configuration period of time. This
* means that the previous EventProcessor that owned the partition is probably down and the partition is now
* eligible to be claimed by other EventProcessors.
*/
Map<String, PartitionOwnership> activePartitionOwnershipMap = removeInactivePartitionOwnerships(
partitionOwnershipMap);
logger.info("Number of active ownership records {}", activePartitionOwnershipMap.size());
/*
* Create a map of owner id and a list of partitions it owns
*/
Map<String, List<PartitionOwnership>> ownerPartitionMap = activePartitionOwnershipMap.values()
.stream()
.collect(
Collectors.groupingBy(PartitionOwnership::getOwnerId, mapping(Function.identity(), toList())));
ownerPartitionMap.putIfAbsent(this.ownerId, new ArrayList<>());
if (CoreUtils.isNullOrEmpty(activePartitionOwnershipMap)) {
/*
* If the active partition ownership map is empty, this is the first time an event processor is
* running or all Event Processors are down for this Event Hub, consumer group combination. All
* partitions in this Event Hub are available to claim. Choose a random partition to claim ownership.
*/
claimOwnership(partitionOwnershipMap, ownerPartitionMap,
partitionIds.get(RANDOM.nextInt(numberOfPartitions)));
return;
}
/*
* Find the minimum number of partitions every event processor should own when the load is
* evenly distributed.
*/
int numberOfActiveEventProcessors = ownerPartitionMap.size();
logger.info("Number of active event processors {}", ownerPartitionMap.size());
int minPartitionsPerEventProcessor = numberOfPartitions / numberOfActiveEventProcessors;
/*
* If the number of partitions in Event Hub is not evenly divisible by number of active event processors,
* a few Event Processors may own 1 additional partition than the minimum when the load is balanced.
* Calculate the number of event processors that can own additional partition.
*/
int numberOfEventProcessorsWithAdditionalPartition = numberOfPartitions % numberOfActiveEventProcessors;
logger.info("Expected min partitions per event processor = {}, expected number of event "
+ "processors with additional partition = {}", minPartitionsPerEventProcessor,
numberOfEventProcessorsWithAdditionalPartition);
if (isLoadBalanced(minPartitionsPerEventProcessor, numberOfEventProcessorsWithAdditionalPartition,
ownerPartitionMap)) {
logger.info("Load is balanced");
checkpointStore.claimOwnership(partitionPumpManager.getPartitionPumps().keySet()
.stream()
.map(partitionId -> createPartitionOwnershipRequest(partitionOwnershipMap, partitionId))
.collect(Collectors.toList()))
.subscribe();
return;
}
if (!shouldOwnMorePartitions(minPartitionsPerEventProcessor, ownerPartitionMap)) {
logger.info("This event processor owns {} partitions and shouldn't own more",
ownerPartitionMap.get(ownerId).size());
checkpointStore.claimOwnership(partitionPumpManager.getPartitionPumps().keySet()
.stream()
.map(partitionId -> createPartitionOwnershipRequest(partitionOwnershipMap, partitionId))
.collect(Collectors.toList()))
.subscribe();
return;
}
logger.info(
"Load is unbalanced and this event processor should own more partitions");
/*
* If some partitions are unclaimed, this could be because an event processor is down and
* it's partitions are now available for others to own or because event processors are just
* starting up and gradually claiming partitions to own or new partitions were added to Event Hub.
* Find any partition that is not actively owned and claim it.
*
* OR
*
* Find a partition to steal from another event processor. Pick the event processor that has owns the
* highest number of partitions.
*/
String partitionToClaim = partitionIds.parallelStream()
.filter(partitionId -> !activePartitionOwnershipMap.containsKey(partitionId))
.findAny()
.orElseGet(() -> {
logger.info("No unclaimed partitions, stealing from another event processor");
return findPartitionToSteal(ownerPartitionMap);
});
claimOwnership(partitionOwnershipMap, ownerPartitionMap, partitionToClaim);
});
}

Line under review: .collect(Collectors.toList()))

Code after addressing the review:
private Mono<Void> loadBalance(final Tuple2<Map<String, PartitionOwnership>, List<String>> tuple) {
return Mono.fromRunnable(() -> {
logger.info("Starting load balancer");
Map<String, PartitionOwnership> partitionOwnershipMap = tuple.getT1();
List<String> partitionIds = tuple.getT2();
if (CoreUtils.isNullOrEmpty(partitionIds)) {
throw logger.logExceptionAsError(Exceptions.propagate(
new IllegalStateException("There are no partitions in Event Hub " + eventHubName)));
}
int numberOfPartitions = partitionIds.size();
logger.info("Partition manager returned {} ownership records", partitionOwnershipMap.size());
logger.info("EventHubAsyncClient returned {} partitions", numberOfPartitions);
if (!isValid(partitionOwnershipMap)) {
throw logger.logExceptionAsError(Exceptions.propagate(
new IllegalStateException("Invalid partitionOwnership data from CheckpointStore")));
}
/*
* Remove all partitions' ownership that have not been modified for a configuration period of time. This
* means that the previous EventProcessor that owned the partition is probably down and the partition is now
* eligible to be claimed by other EventProcessors.
*/
Map<String, PartitionOwnership> activePartitionOwnershipMap = removeInactivePartitionOwnerships(
partitionOwnershipMap);
logger.info("Number of active ownership records {}", activePartitionOwnershipMap.size());
/*
* Create a map of owner id and a list of partitions it owns
*/
Map<String, List<PartitionOwnership>> ownerPartitionMap = activePartitionOwnershipMap.values()
.stream()
.collect(
Collectors.groupingBy(PartitionOwnership::getOwnerId, mapping(Function.identity(), toList())));
ownerPartitionMap.putIfAbsent(this.ownerId, new ArrayList<>());
if (CoreUtils.isNullOrEmpty(activePartitionOwnershipMap)) {
/*
* If the active partition ownership map is empty, this is the first time an event processor is
* running or all Event Processors are down for this Event Hub, consumer group combination. All
* partitions in this Event Hub are available to claim. Choose a random partition to claim ownership.
*/
claimOwnership(partitionOwnershipMap, ownerPartitionMap,
partitionIds.get(RANDOM.nextInt(numberOfPartitions)));
return;
}
/*
* Find the minimum number of partitions every event processor should own when the load is
* evenly distributed.
*/
int numberOfActiveEventProcessors = ownerPartitionMap.size();
logger.info("Number of active event processors {}", ownerPartitionMap.size());
int minPartitionsPerEventProcessor = numberOfPartitions / numberOfActiveEventProcessors;
/*
* If the number of partitions in Event Hub is not evenly divisible by number of active event processors,
* a few Event Processors may own 1 additional partition than the minimum when the load is balanced.
* Calculate the number of event processors that can own additional partition.
*/
int numberOfEventProcessorsWithAdditionalPartition = numberOfPartitions % numberOfActiveEventProcessors;
logger.info("Expected min partitions per event processor = {}, expected number of event "
+ "processors with additional partition = {}", minPartitionsPerEventProcessor,
numberOfEventProcessorsWithAdditionalPartition);
if (isLoadBalanced(minPartitionsPerEventProcessor, numberOfEventProcessorsWithAdditionalPartition,
ownerPartitionMap)) {
logger.info("Load is balanced");
checkpointStore.claimOwnership(partitionPumpManager.getPartitionPumps().keySet()
.stream()
.map(partitionId -> createPartitionOwnershipRequest(partitionOwnershipMap, partitionId))
.collect(Collectors.toList()))
.subscribe();
return;
}
if (!shouldOwnMorePartitions(minPartitionsPerEventProcessor, ownerPartitionMap)) {
logger.info("This event processor owns {} partitions and shouldn't own more",
ownerPartitionMap.get(ownerId).size());
checkpointStore.claimOwnership(partitionPumpManager.getPartitionPumps().keySet()
.stream()
.map(partitionId -> createPartitionOwnershipRequest(partitionOwnershipMap, partitionId))
.collect(Collectors.toList()))
.subscribe();
return;
}
logger.info(
"Load is unbalanced and this event processor should own more partitions");
/*
* If some partitions are unclaimed, this could be because an event processor is down and
* it's partitions are now available for others to own or because event processors are just
* starting up and gradually claiming partitions to own or new partitions were added to Event Hub.
* Find any partition that is not actively owned and claim it.
*
* OR
*
* Find a partition to steal from another event processor. Pick the event processor that has owns the
* highest number of partitions.
*/
String partitionToClaim = partitionIds.parallelStream()
.filter(partitionId -> !activePartitionOwnershipMap.containsKey(partitionId))
.findAny()
.orElseGet(() -> {
logger.info("No unclaimed partitions, stealing from another event processor");
return findPartitionToSteal(ownerPartitionMap);
});
claimOwnership(partitionOwnershipMap, ownerPartitionMap, partitionToClaim);
});
}