org.apache.bookkeeper.client.DynamicWeightedRandomSelectionImpl Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.bookkeeper.client;
import com.google.common.math.Quantiles;
import com.google.common.math.Quantiles.ScaleAndIndex;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* DynamicWeightedRandomSelectionImpl class implements both getNextRandom
* overloaded methods. Where getNextRandom() method considers all bookies it
* knows of as candidates, but getNextRandom(Collection selectedNodes) method
* considers only 'selectedNodes' as candidates.
*/
class DynamicWeightedRandomSelectionImpl implements WeightedRandomSelection {
static final Logger LOG = LoggerFactory.getLogger(DynamicWeightedRandomSelectionImpl.class);
int maxProbabilityMultiplier;
final Map weightMap;
final ReadWriteLock rwLock = new ReentrantReadWriteLock(true);
Random rand;
DynamicWeightedRandomSelectionImpl() {
this(-1);
}
DynamicWeightedRandomSelectionImpl(int maxMultiplier) {
this.maxProbabilityMultiplier = maxMultiplier;
this.weightMap = new HashMap();
rand = new Random(System.currentTimeMillis());
}
@Override
public void updateMap(Map updatedMap) {
rwLock.writeLock().lock();
try {
weightMap.clear();
weightMap.putAll(updatedMap);
} finally {
rwLock.writeLock().unlock();
}
}
@Override
public T getNextRandom() {
rwLock.readLock().lock();
try {
return getNextRandom(weightMap.keySet());
} finally {
rwLock.readLock().unlock();
}
}
@Override
public T getNextRandom(Collection selectedNodes) {
rwLock.readLock().lock();
try {
/*
* calculate minWeight and actual total weight.
*/
long minWeight = Long.MAX_VALUE;
long actTotalWeight = 0;
for (T node : selectedNodes) {
long weight = 0;
if ((weightMap.containsKey(node))) {
weight = weightMap.get(node).getWeight();
}
actTotalWeight += weight;
if (weight > 0 && minWeight > weight) {
minWeight = weight;
}
}
long medianWeight;
/*
* if actTotalWeight is 0, then assign 1 to minWeight and
* medianWeight.
*/
if (actTotalWeight == 0) {
minWeight = 1L;
medianWeight = 1L;
} else {
/*
* calculate medianWeight.
*/
Function weightFunc = (node) -> {
long weight = 0;
if ((weightMap.containsKey(node))) {
weight = weightMap.get(node).getWeight();
}
return weight;
};
ArrayList weightList = selectedNodes.stream().map(weightFunc)
.collect(Collectors.toCollection(ArrayList::new));
ScaleAndIndex median = Quantiles.median();
medianWeight = (long) median.compute(weightList);
}
/*
* initialize maxWeight value based on maxProbabilityMultiplier.
*/
long maxWeight = maxProbabilityMultiplier * medianWeight;
/*
* apply weighted random selection to select an element randomly
* based on weight.
*/
long cumTotalWeight = 0;
T nextRandomNode = null;
for (T node : selectedNodes) {
long weight = 0;
if ((weightMap.containsKey(node))) {
weight = weightMap.get(node).getWeight();
}
if (weight <= 0) {
weight = minWeight;
} else if (maxWeight > 0 && weight > maxWeight) {
weight = maxWeight;
}
long tmpRandLong = rand.nextLong();
if (tmpRandLong == Long.MIN_VALUE) {
tmpRandLong++;
}
long randValue = Math.abs(tmpRandLong) % (cumTotalWeight + weight);
if (randValue >= cumTotalWeight) {
nextRandomNode = node;
}
cumTotalWeight += weight;
}
return nextRandomNode;
} finally {
rwLock.readLock().unlock();
}
}
@Override
public void setMaxProbabilityMultiplier(int max) {
this.maxProbabilityMultiplier = max;
}
}