com.azure.cosmos.implementation.changefeed.implementation.EqualPartitionsBalancingStrategy Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of azure-cosmos Show documentation
Show all versions of azure-cosmos Show documentation
This Package contains Microsoft Azure Cosmos SDK (with Reactive Extension Reactor support) for Azure Cosmos DB SQL API
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.changefeed.implementation;
import com.azure.cosmos.implementation.changefeed.Lease;
import com.azure.cosmos.implementation.changefeed.PartitionLoadBalancingStrategy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/**
* Implementation for {@link PartitionLoadBalancingStrategy}.
*/
class EqualPartitionsBalancingStrategy implements PartitionLoadBalancingStrategy {
private final Logger logger = LoggerFactory.getLogger(EqualPartitionsBalancingStrategy.class);
private final String hostName;
private final int minPartitionCount;
private final int maxPartitionCount;
private final Duration leaseExpirationInterval;
public EqualPartitionsBalancingStrategy(String hostName, int minPartitionCount, int maxPartitionCount, Duration leaseExpirationInterval) {
if (hostName == null) {
throw new IllegalArgumentException("hostName");
}
this.hostName = hostName;
this.minPartitionCount = minPartitionCount;
this.maxPartitionCount = maxPartitionCount;
this.leaseExpirationInterval = leaseExpirationInterval;
}
@Override
public List selectLeasesToTake(List allLeases) {
Map workerToPartitionCount = new HashMap<>();
List expiredLeases = new ArrayList<>();
Map allPartitions = new HashMap<>();
this.categorizeLeases(allLeases, allPartitions, expiredLeases, workerToPartitionCount);
int partitionCount = allPartitions.size();
int workerCount = workerToPartitionCount.size();
if (partitionCount <= 0) {
return new ArrayList();
}
int target = this.calculateTargetPartitionCount(partitionCount, workerCount);
int myCount = workerToPartitionCount.get(this.hostName);
int partitionsNeededForMe = target - myCount;
/*
Logger.InfoFormat(
"Host '{0}' {1} partitions, {2} hosts, {3} available leases, target = {4}, min = {5}, max = {6}, mine = {7}, will try to take {8} lease(s) for myself'.",
this.hostName,
partitionCount,
workerCount,
expiredLeases.Count,
target,
this.minScaleCount,
this.maxScaleCount,
myCount,
Math.Max(partitionsNeededForMe, 0));
*/
if (partitionsNeededForMe <= 0)
return new ArrayList();
if (expiredLeases.size() > 0) {
return expiredLeases.subList(0, Math.min(partitionsNeededForMe, expiredLeases.size()));
}
Lease stolenLease = getLeaseToSteal(workerToPartitionCount, target, partitionsNeededForMe, allPartitions);
List stolenLeases = new ArrayList<>();
if (stolenLease != null) {
stolenLeases.add(stolenLease);
}
return stolenLeases;
}
private static Lease getLeaseToSteal(
Map workerToPartitionCount,
int target,
int partitionsNeededForMe,
Map allPartitions) {
Map.Entry workerToStealFrom = findWorkerWithMostPartitions(workerToPartitionCount);
if (workerToStealFrom.getValue() > target - (partitionsNeededForMe > 1 ? 1 : 0)) {
for (Map.Entry entry : allPartitions.entrySet()) {
if (entry.getValue().getOwner().equalsIgnoreCase(workerToStealFrom.getKey())) {
return entry.getValue();
}
}
}
return null;
}
private static Map.Entry findWorkerWithMostPartitions(Map workerToPartitionCount) {
Map.Entry workerToStealFrom = new ChangeFeedHelper.KeyValuePair<>("", 0);
for (Map.Entry entry : workerToPartitionCount.entrySet()) {
if (workerToStealFrom.getValue() <= entry.getValue()) {
workerToStealFrom = entry;
}
}
return workerToStealFrom;
}
private int calculateTargetPartitionCount(int partitionCount, int workerCount) {
int target = 1;
if (partitionCount > workerCount) {
target = (int)Math.ceil((double)partitionCount / workerCount);
}
if (this.maxPartitionCount > 0 && target > this.maxPartitionCount) {
target = this.maxPartitionCount;
}
if (this.minPartitionCount > 0 && target < this.minPartitionCount) {
target = this.minPartitionCount;
}
return target;
}
private void categorizeLeases(
List allLeases,
Map allPartitions,
List expiredLeases,
Map workerToPartitionCount) {
for (Lease lease : allLeases) {
// Debug.Assert(lease.LeaseToken != null, "TakeLeasesAsync: lease.LeaseToken cannot be null.");
allPartitions.put(lease.getLeaseToken(), lease);
if (lease.getOwner() == null || lease.getOwner().isEmpty() || this.isExpired(lease)) {
this.logger.info("Found unused or expired lease {}", lease.getLeaseToken());
expiredLeases.add(lease);
} else {
String assignedTo = lease.getOwner();
Integer count = workerToPartitionCount.get(assignedTo);
if (count != null) {
workerToPartitionCount.replace(assignedTo, count + 1);
} else {
workerToPartitionCount.put(assignedTo, 1);
}
}
}
if (!workerToPartitionCount.containsKey(this.hostName)) {
workerToPartitionCount.put(this.hostName, 0);
}
}
private boolean isExpired(Lease lease) {
if (lease.getOwner() == null || lease.getOwner().isEmpty() || lease.getTimestamp() == null) {
return true;
}
Instant leaseExpireTime = Instant.parse(lease.getTimestamp()).plus(this.leaseExpirationInterval);
this.logger.debug("Current lease timestamp: {}, current time: {}", leaseExpireTime, Instant.now());
return leaseExpireTime.isBefore(Instant.now());
}
}