/*
* 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.cassandra.dht.tokenallocator;
import java.util.*;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import org.apache.commons.math3.stat.descriptive.SummaryStatistics;
import org.junit.Assert;
import org.junit.Ignore;
import org.apache.cassandra.dht.IPartitioner;
import org.apache.cassandra.dht.Token;
/**
* Base class for {@link Murmur3ReplicationAwareTokenAllocatorTest} and {@link RandomReplicationAwareTokenAllocatorTest},
* we need to separate classes to avoid timeous in case flaky tests need to be repeated, see CASSANDRA-12784.
*/
@Ignore
abstract class AbstractReplicationAwareTokenAllocatorTest extends TokenAllocatorTestBase
{
static class SimpleReplicationStrategy implements TestReplicationStrategy
{
int replicas;
public SimpleReplicationStrategy(int replicas)
{
super();
this.replicas = replicas;
}
public List<Unit> getReplicas(Token token, NavigableMap<Token, Unit> sortedTokens)
{
List<Unit> endpoints = new ArrayList<Unit>(replicas);
token = sortedTokens.ceilingKey(token);
if (token == null)
token = sortedTokens.firstKey();
Iterator<Unit> iter = Iterables.concat(sortedTokens.tailMap(token, true).values(), sortedTokens.values()).iterator();
while (endpoints.size() < replicas)
{
if (!iter.hasNext())
return endpoints;
Unit ep = iter.next();
if (!endpoints.contains(ep))
endpoints.add(ep);
}
return endpoints;
}
public Token replicationStart(Token token, Unit unit, NavigableMap<Token, Unit> sortedTokens)
{
Set<Unit> seenUnits = Sets.newHashSet();
int unitsFound = 0;
for (Map.Entry<Token, Unit> en : Iterables.concat(
sortedTokens.headMap(token, false).descendingMap().entrySet(),
sortedTokens.descendingMap().entrySet()))
{
Unit n = en.getValue();
// Same group as investigated unit is a break; anything that could replicate in it replicates there.
if (n == unit)
break;
if (seenUnits.add(n))
{
if (++unitsFound == replicas)
break;
}
token = en.getKey();
}
return token;
}
public void addUnit(Unit n)
{
}
public void removeUnit(Unit n)
{
}
public String toString()
{
return String.format("Simple %d replicas", replicas);
}
public int replicas()
{
return replicas;
}
public boolean sameGroup(Unit n1, Unit n2)
{
return false;
}
public Unit getGroup(Unit unit)
{
// The unit is the group.
return unit;
}
public double spreadExpectation()
{
return 1;
}
}
static abstract class GroupReplicationStrategy implements TestReplicationStrategy
{
final int replicas;
final Map<Unit, Integer> groupMap;
public GroupReplicationStrategy(int replicas)
{
this.replicas = replicas;
this.groupMap = Maps.newHashMap();
}
public List<Unit> getReplicas(Token token, NavigableMap<Token, Unit> sortedTokens)
{
List<Unit> endpoints = new ArrayList<Unit>(replicas);
BitSet usedGroups = new BitSet();
if (sortedTokens.isEmpty())
return endpoints;
token = sortedTokens.ceilingKey(token);
if (token == null)
token = sortedTokens.firstKey();
Iterator<Unit> iter = Iterables.concat(sortedTokens.tailMap(token, true).values(), sortedTokens.values()).iterator();
while (endpoints.size() < replicas)
{
// For simlicity assuming list can't be exhausted before finding all replicas.
Unit ep = iter.next();
int group = groupMap.get(ep);
if (!usedGroups.get(group))
{
endpoints.add(ep);
usedGroups.set(group);
}
}
return endpoints;
}
public Token lastReplicaToken(Token token, NavigableMap<Token, Unit> sortedTokens)
{
BitSet usedGroups = new BitSet();
int groupsFound = 0;
token = sortedTokens.ceilingKey(token);
if (token == null)
token = sortedTokens.firstKey();
for (Map.Entry<Token, Unit> en :
Iterables.concat(sortedTokens.tailMap(token, true).entrySet(),
sortedTokens.entrySet()))
{
Unit ep = en.getValue();
int group = groupMap.get(ep);
if (!usedGroups.get(group))
{
usedGroups.set(group);
if (++groupsFound >= replicas)
return en.getKey();
}
}
return token;
}
public Token replicationStart(Token token, Unit unit, NavigableMap<Token, Unit> sortedTokens)
{
// replicated ownership
int unitGroup = groupMap.get(unit); // unit must be already added
BitSet seenGroups = new BitSet();
int groupsFound = 0;
for (Map.Entry<Token, Unit> en : Iterables.concat(
sortedTokens.headMap(token, false).descendingMap().entrySet(),
sortedTokens.descendingMap().entrySet()))
{
Unit n = en.getValue();
int ngroup = groupMap.get(n);
// Same group as investigated unit is a break; anything that could replicate in it replicates there.
if (ngroup == unitGroup)
break;
if (!seenGroups.get(ngroup))
{
if (++groupsFound == replicas)
break;
seenGroups.set(ngroup);
}
token = en.getKey();
}
return token;
}
public String toString()
{
Map<Integer, Integer> idToSize = instanceToCount(groupMap);
Map<Integer, Integer> sizeToCount = Maps.newTreeMap();
sizeToCount.putAll(instanceToCount(idToSize));
return String.format("%s strategy, %d replicas, group size to count %s", getClass().getSimpleName(), replicas, sizeToCount);
}
@Override
public int replicas()
{
return replicas;
}
public boolean sameGroup(Unit n1, Unit n2)
{
return groupMap.get(n1).equals(groupMap.get(n2));
}
public void removeUnit(Unit n)
{
groupMap.remove(n);
}
public Integer getGroup(Unit unit)
{
return groupMap.get(unit);
}
public double spreadExpectation()
{
return 1.5; // Even balanced racks get disbalanced when they lose nodes.
}
}
private static <T> Map<T, Integer> instanceToCount(Map<?, T> map)
{
Map<T, Integer> idToCount = Maps.newHashMap();
for (Map.Entry<?, T> en : map.entrySet())
{
Integer old = idToCount.get(en.getValue());
idToCount.put(en.getValue(), old != null ? old + 1 : 1);
}
return idToCount;
}
/**
* Group strategy spreading units into a fixed number of groups.
*/
static class FixedGroupCountReplicationStrategy extends GroupReplicationStrategy
{
int groupId;
int groupCount;
public FixedGroupCountReplicationStrategy(int replicas, int groupCount)
{
super(replicas);
assert groupCount >= replicas;
groupId = 0;
this.groupCount = groupCount;
}
public void addUnit(Unit n)
{
groupMap.put(n, groupId++ % groupCount);
}
}
/**
* Group strategy with a fixed number of units per group.
*/
static class BalancedGroupReplicationStrategy extends GroupReplicationStrategy
{
int groupId;
int groupSize;
public BalancedGroupReplicationStrategy(int replicas, int groupSize)
{
super(replicas);
groupId = 0;
this.groupSize = groupSize;
}
public void addUnit(Unit n)
{
groupMap.put(n, groupId++ / groupSize);
}
}
static class UnbalancedGroupReplicationStrategy extends GroupReplicationStrategy
{
int groupId;
int nextSize;
int num;
int minGroupSize;
int maxGroupSize;
Random rand;
public UnbalancedGroupReplicationStrategy(int replicas, int minGroupSize, int maxGroupSize, Random rand)
{
super(replicas);
groupId = -1;
nextSize = 0;
num = 0;
this.maxGroupSize = maxGroupSize;
this.minGroupSize = minGroupSize;
this.rand = rand;
}
public void addUnit(Unit n)
{
if (++num > nextSize)
{
nextSize = minGroupSize + rand.nextInt(maxGroupSize - minGroupSize + 1);
++groupId;
num = 0;
}
groupMap.put(n, groupId);
}
public double spreadExpectation()
{
return 2;
}
}
static Map<Unit, Double> evaluateReplicatedOwnership(ReplicationAwareTokenAllocator<Unit> t)
{
Map<Unit, Double> ownership = Maps.newHashMap();
Iterator<Token> it = t.sortedTokens.keySet().iterator();
if (!it.hasNext())
return ownership;
Token current = it.next();
while (it.hasNext())
{
Token next = it.next();
addOwnership(t, current, next, ownership);
current = next;
}
addOwnership(t, current, t.sortedTokens.firstKey(), ownership);
return ownership;
}
private static void addOwnership(ReplicationAwareTokenAllocator<Unit> t, Token current, Token next, Map<Unit, Double> ownership)
{
TestReplicationStrategy ts = (TestReplicationStrategy) t.strategy;
double size = current.size(next);
Token representative = t.partitioner.midpoint(current, next);
for (Unit n : ts.getReplicas(representative, t.sortedTokens))
{
Double v = ownership.get(n);
ownership.put(n, v != null ? v + size : size);
}
}
private static double replicatedTokenOwnership(Token token, NavigableMap<Token, Unit> sortedTokens, ReplicationStrategy<Unit> strategy)
{
TestReplicationStrategy ts = (TestReplicationStrategy) strategy;
Token next = sortedTokens.higherKey(token);
if (next == null)
next = sortedTokens.firstKey();
return ts.replicationStart(token, sortedTokens.get(token), sortedTokens).size(next);
}
protected void testExistingCluster(IPartitioner partitioner, int maxVNodeCount)
{
for (int rf = 1; rf <= 5; ++rf)
{
for (int perUnitCount = 1; perUnitCount <= maxVNodeCount; perUnitCount *= 4)
{
testExistingCluster(perUnitCount, fixedTokenCount, new SimpleReplicationStrategy(rf), partitioner);
testExistingCluster(perUnitCount, varyingTokenCount, new SimpleReplicationStrategy(rf), partitioner);
if (rf == 1) continue; // Replication strategy doesn't matter for RF = 1.
for (int groupSize = 4; groupSize <= 64 && groupSize * rf * 4 < TARGET_CLUSTER_SIZE; groupSize *= 4)
{
testExistingCluster(perUnitCount, fixedTokenCount,
new BalancedGroupReplicationStrategy(rf, groupSize), partitioner);
testExistingCluster(perUnitCount, varyingTokenCount,
new UnbalancedGroupReplicationStrategy(rf, groupSize / 2, groupSize * 2, seededRand),
partitioner);
}
testExistingCluster(perUnitCount, fixedTokenCount,
new FixedGroupCountReplicationStrategy(rf, rf * 2), partitioner);
}
}
}
private void testExistingCluster(int perUnitCount, TokenCount tc, TestReplicationStrategy rs, IPartitioner partitioner)
{
System.out.println("Testing existing cluster, target " + perUnitCount + " vnodes, replication " + rs);
final int targetClusterSize = TARGET_CLUSTER_SIZE;
NavigableMap<Token, Unit> tokenMap = Maps.newTreeMap();
random(tokenMap, rs, targetClusterSize / 2, tc, perUnitCount, partitioner);
ReplicationAwareTokenAllocator<Unit> t = new ReplicationAwareTokenAllocator<>(tokenMap, rs, partitioner);
grow(t, targetClusterSize * 9 / 10, tc, perUnitCount, false);
grow(t, targetClusterSize, tc, perUnitCount, true);
loseAndReplace(t, targetClusterSize / 10, tc, perUnitCount, partitioner);
System.out.println();
}
protected void testNewCluster(IPartitioner partitioner, int maxVNodeCount)
{
// This test is flaky because the selection of the tokens for the first RF nodes (which is random, with an
// uncontrolled seed) can sometimes cause a pathological situation where the algorithm will find a (close to)
// ideal distribution of tokens for some number of nodes, which in turn will inevitably cause it to go into a
// bad (unacceptable to the test criteria) distribution after adding one more node.
// This should happen very rarely, unless something is broken in the token allocation code.
for (int rf = 2; rf <= 5; ++rf)
{
for (int perUnitCount = 1; perUnitCount <= maxVNodeCount; perUnitCount *= 4)
{
testNewCluster(perUnitCount, fixedTokenCount, new SimpleReplicationStrategy(rf), partitioner);
testNewCluster(perUnitCount, varyingTokenCount, new SimpleReplicationStrategy(rf), partitioner);
if (rf == 1) continue; // Replication strategy doesn't matter for RF = 1.
for (int groupSize = 4; groupSize <= 64 && groupSize * rf * 8 < TARGET_CLUSTER_SIZE; groupSize *= 4)
{
testNewCluster(perUnitCount, fixedTokenCount,
new BalancedGroupReplicationStrategy(rf, groupSize), partitioner);
testNewCluster(perUnitCount, varyingTokenCount,
new UnbalancedGroupReplicationStrategy(rf, groupSize / 2, groupSize * 2, seededRand),
partitioner);
}
testNewCluster(perUnitCount, fixedTokenCount,
new FixedGroupCountReplicationStrategy(rf, rf * 2), partitioner);
}
}
}
private void testNewCluster(int perUnitCount, TokenCount tc, TestReplicationStrategy rs, IPartitioner partitioner)
{
System.out.println("Testing new cluster, target " + perUnitCount + " vnodes, replication " + rs);
final int targetClusterSize = TARGET_CLUSTER_SIZE;
NavigableMap<Token, Unit> tokenMap = Maps.newTreeMap();
ReplicationAwareTokenAllocator<Unit> t = new ReplicationAwareTokenAllocator<>(tokenMap, rs, partitioner);
grow(t, targetClusterSize * 2 / 5, tc, perUnitCount, false);
grow(t, targetClusterSize, tc, perUnitCount, true);
loseAndReplace(t, targetClusterSize / 5, tc, perUnitCount, partitioner);
System.out.println();
}
private void loseAndReplace(ReplicationAwareTokenAllocator<Unit> t, int howMany,
TokenCount tc, int perUnitCount, IPartitioner partitioner)
{
int fullCount = t.unitCount();
System.out.format("Losing %d units. ", howMany);
for (int i = 0; i < howMany; ++i)
{
Unit u = t.unitFor(partitioner.getRandomToken(seededRand));
t.removeUnit(u);
((TestReplicationStrategy) t.strategy).removeUnit(u);
}
// Grow half without verifying.
grow(t, (t.unitCount() + fullCount * 3) / 4, tc, perUnitCount, false);
// Metrics should be back to normal by now. Check that they remain so.
grow(t, fullCount, tc, perUnitCount, true);
}
public void grow(ReplicationAwareTokenAllocator<Unit> t, int targetClusterSize, TokenCount tc, int perUnitCount, boolean verifyMetrics)
{
int size = t.unitCount();
Summary su = new Summary();
Summary st = new Summary();
Random rand = new Random(targetClusterSize + perUnitCount);
TestReplicationStrategy strategy = (TestReplicationStrategy) t.strategy;
if (size < targetClusterSize)
{
System.out.format("Adding %d unit(s) using %s...", targetClusterSize - size, t.toString());
long time = System.currentTimeMillis();
while (size < targetClusterSize)
{
int tokens = tc.tokenCount(perUnitCount, rand);
Unit unit = new Unit();
strategy.addUnit(unit);
t.addUnit(unit, tokens);
++size;
if (verifyMetrics)
updateSummary(t, su, st, false);
}
System.out.format(" Done in %.3fs\n", (System.currentTimeMillis() - time) / 1000.0);
if (verifyMetrics)
{
updateSummary(t, su, st, true);
double maxExpected = 1.0 + tc.spreadExpectation() * strategy.spreadExpectation() / (perUnitCount * t.replicas);
if (su.max > maxExpected)
{
Assert.fail(String.format("Expected max unit size below %.4f, was %.4f", maxExpected, su.max));
}
// We can't verify lower side range as small loads can't always be fixed.
}
}
}
private void updateSummary(ReplicationAwareTokenAllocator<Unit> t, Summary su, Summary st, boolean print)
{
int size = t.sortedTokens.size();
double inverseAverage = 1.0 * size / t.strategy.replicas();
Map<Unit, Double> ownership = evaluateReplicatedOwnership(t);
SummaryStatistics unitStat = new SummaryStatistics();
for (Map.Entry<Unit, Double> en : ownership.entrySet())
unitStat.addValue(en.getValue() * inverseAverage / t.unitToTokens.get(en.getKey()).size());
su.update(unitStat);
SummaryStatistics tokenStat = new SummaryStatistics();
for (Token tok : t.sortedTokens.keySet())
tokenStat.addValue(replicatedTokenOwnership(tok, t.sortedTokens, t.strategy) * inverseAverage);
st.update(tokenStat);
if (print)
{
System.out.format("Size %d(%d) \tunit %s token %s %s\n",
t.unitCount(), size,
mms(unitStat),
mms(tokenStat),
t.strategy);
System.out.format("Worst intermediate unit\t%s token %s\n", su, st);
}
}
}