Code review comment: When maxScaleCount is set to a non-zero value, it's ok to be greedy because there's a cap to how much it can be assigned to one CFP instance. I agree that the system might still go though couple rounds of load balancing itself but much less than when no cap is set. The idea here is that as new instances are introduced into the system, they should start the work immediately on all the partitions that are available rather than only taking one at the time.

Code before review:
public List<Lease> selectLeasesToTake(List<Lease> allLeases) {
Map<String, Integer> workerToPartitionCount = new HashMap<>();
List<Lease> expiredLeases = new ArrayList<>();
Map<String, Lease> allPartitions = new HashMap<>();
this.categorizeLeases(allLeases, allPartitions, expiredLeases, workerToPartitionCount);
int partitionCount = allPartitions.size();
int workerCount = workerToPartitionCount.size();
if (partitionCount <= 0) {
return new ArrayList<Lease>();
}
int target = this.calculateTargetPartitionCount(partitionCount, workerCount);
int myCount = workerToPartitionCount.get(this.hostName);
int partitionsNeededForMe = target - myCount;
if (expiredLeases.size() > 0) {
if (this.maxPartitionCount == 0 && (partitionsNeededForMe <= 0 || (partitionsNeededForMe > 1 && workerToPartitionCount.size() > 1))) {
partitionsNeededForMe = 1;
}
for (Lease lease : expiredLeases) {
this.logger.info("Found unused or expired lease {}; previous lease count for instance owner {} is {}, count of leases to target is {} and maxScaleCount {} ",
lease.getLeaseToken(), this.hostName, myCount, partitionsNeededForMe, this.maxPartitionCount);
}
return expiredLeases.subList(0, Math.min(partitionsNeededForMe, expiredLeases.size()));
}
if (partitionsNeededForMe <= 0)
return new ArrayList<Lease>();
Lease stolenLease = getLeaseToSteal(workerToPartitionCount, target, partitionsNeededForMe, allPartitions);
List<Lease> stolenLeases = new ArrayList<>();
if (stolenLease != null) {
stolenLeases.add(stolenLease);
}
return stolenLeases;
}

Line under review: if (this.maxPartitionCount == 0 && (partitionsNeededForMe <= 0 || (partitionsNeededForMe > 1 && workerToPartitionCount.size() > 1))) {

Code after addressing the review:
public List<Lease> selectLeasesToTake(List<Lease> allLeases) {
Map<String, Integer> workerToPartitionCount = new HashMap<>();
List<Lease> expiredLeases = new ArrayList<>();
Map<String, Lease> allPartitions = new HashMap<>();
this.categorizeLeases(allLeases, allPartitions, expiredLeases, workerToPartitionCount);
int partitionCount = allPartitions.size();
int workerCount = workerToPartitionCount.size();
if (partitionCount <= 0) {
return new ArrayList<Lease>();
}
int target = this.calculateTargetPartitionCount(partitionCount, workerCount);
int myCount = workerToPartitionCount.get(this.hostName);
int partitionsNeededForMe = target - myCount;
if (expiredLeases.size() > 0) {
if ((this.maxPartitionCount == 0 && partitionsNeededForMe <= 0) || (partitionsNeededForMe > 1 && workerToPartitionCount.size() > 1)) {
partitionsNeededForMe = 1;
}
if (partitionsNeededForMe == 1) {
Random random = new Random();
Lease expiredLease = expiredLeases.get(random.nextInt(expiredLeases.size()));
this.logger.info("Found unused or expired lease {} (owner was {}); previous lease count for instance owner {} is {}, count of leases to target is {} and maxScaleCount {} ",
expiredLease.getLeaseToken(), expiredLease.getOwner(), this.hostName, myCount, partitionsNeededForMe, this.maxPartitionCount);
return Collections.singletonList(expiredLease);
} else {
for (Lease lease : expiredLeases) {
this.logger.info("Found unused or expired lease {} (owner was {}); previous lease count for instance owner {} is {} and maxScaleCount {} ",
lease.getLeaseToken(), lease.getOwner(), this.hostName, myCount, this.maxPartitionCount);
}
}
return expiredLeases.subList(0, Math.min(partitionsNeededForMe, expiredLeases.size()));
}
if (partitionsNeededForMe <= 0)
return new ArrayList<Lease>();
Lease stolenLease = getLeaseToSteal(workerToPartitionCount, target, partitionsNeededForMe, allPartitions);
List<Lease> stolenLeases = new ArrayList<>();
if (stolenLease != null) {
stolenLeases.add(stolenLease);
}
return stolenLeases;
}