Add algorithmic token allocation

patch by branimir; reviewed by benedict for CASSANDRA-7032

CHANGES.txt

@@ -1,4 +1,5 @@
 3.0:
+ * Add algorithmic token allocation (CASSANDRA-7032)
  * Add nodetool command to replay batchlog (CASSANDRA-9547)
  * Make file buffer cache independent of paths being read (CASSANDRA-8897)
  * Remove deprecated legacy Hadoop code (CASSANDRA-9353)

conf/cassandra.yaml

@@ -24,6 +24,17 @@ cluster_name: 'Test Cluster'
 # multiple tokens per node, see http://wiki.apache.org/cassandra/Operations
 num_tokens: 256
 
+# Triggers automatic allocation of num_tokens tokens for this node. The allocation
+# algorithm attempts to choose tokens in a way that optimizes replicated load over
+# the nodes in the datacenter for the replication strategy used by the specified
+# keyspace.
+#
+# The load assigned to each node will be close to proportional to its number of
+# vnodes.
+#
+# Only supported with the Murmur3Partitioner.
+# allocate_tokens_keyspace: KEYSPACE
+
 # initial_token allows you to specify tokens manually.  While you can use # it with
 # vnodes (num_tokens > 1, above) -- in which case you should provide a 
 # comma-separated list -- it's primarily used when adding nodes # to legacy clusters 

src/java/org/apache/cassandra/config/Config.java

@@ -73,6 +73,8 @@ public class Config
     /* initial token in the ring */
     public String initial_token;
     public Integer num_tokens = 1;
+    /** Triggers automatic allocation of tokens if set, using the replication strategy of the referenced keyspace */
+    public String allocate_tokens_for_keyspace = null;
 
     public volatile Long request_timeout_in_ms = 10000L;
 

src/java/org/apache/cassandra/config/DatabaseDescriptor.java

@@ -811,6 +811,11 @@ public class DatabaseDescriptor
         return tokensFromString(System.getProperty("cassandra.initial_token", conf.initial_token));
     }
 
+    public static String getAllocateTokensKeyspace()
+    {
+        return System.getProperty("cassandra.allocate_tokens_keyspace", conf.allocate_tokens_for_keyspace);
+    }
+
     public static Collection<String> tokensFromString(String tokenString)
     {
         List<String> tokens = new ArrayList<String>();

src/java/org/apache/cassandra/db/commitlog/CommitLog.java

@@ -74,7 +74,7 @@ public class CommitLog implements CommitLogMBean
     public ParameterizedClass compressorClass;
     final public String location;
 
-    static private CommitLog construct()
+    private static CommitLog construct()
     {
         CommitLog log = new CommitLog(DatabaseDescriptor.getCommitLogLocation(), new CommitLogArchiver());
 

src/java/org/apache/cassandra/db/commitlog/CompressedSegment.java

@@ -35,7 +35,7 @@ import org.apache.cassandra.utils.SyncUtil;
  */
 public class CompressedSegment extends CommitLogSegment
 {
-    static private final ThreadLocal<ByteBuffer> compressedBufferHolder = new ThreadLocal<ByteBuffer>() {
+    private static final ThreadLocal<ByteBuffer> compressedBufferHolder = new ThreadLocal<ByteBuffer>() {
         protected ByteBuffer initialValue()
         {
             return ByteBuffer.allocate(0);

src/java/org/apache/cassandra/dht/BootStrapper.java

@@ -31,6 +31,7 @@ import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.config.Schema;
 import org.apache.cassandra.db.Keyspace;
 import org.apache.cassandra.db.TypeSizes;
+import org.apache.cassandra.dht.tokenallocator.TokenAllocation;
 import org.apache.cassandra.exceptions.ConfigurationException;
 import org.apache.cassandra.gms.FailureDetector;
 import org.apache.cassandra.io.IVersionedSerializer;
@@ -150,37 +151,61 @@ public class BootStrapper extends ProgressEventNotifierSupport
 
     /**
      * if initialtoken was specified, use that (split on comma).
-     * otherwise, if num_tokens == 1, pick a token to assume half the load of the most-loaded node.
+     * otherwise, if allocationKeyspace is specified use the token allocation algorithm to generate suitable tokens
      * else choose num_tokens tokens at random
      */
-    public static Collection<Token> getBootstrapTokens(final TokenMetadata metadata) throws ConfigurationException
+    public static Collection<Token> getBootstrapTokens(final TokenMetadata metadata, InetAddress address) throws ConfigurationException
     {
+        String allocationKeyspace = DatabaseDescriptor.getAllocateTokensKeyspace();
         Collection<String> initialTokens = DatabaseDescriptor.getInitialTokens();
+        if (initialTokens.size() > 0 && allocationKeyspace != null)
+            logger.warn("manually specified tokens override automatic allocation");
+
         // if user specified tokens, use those
         if (initialTokens.size() > 0)
-        {
-            logger.debug("tokens manually specified as {}",  initialTokens);
-            List<Token> tokens = new ArrayList<>(initialTokens.size());
-            for (String tokenString : initialTokens)
-            {
-                Token token = StorageService.getPartitioner().getTokenFactory().fromString(tokenString);
-                if (metadata.getEndpoint(token) != null)
-                    throw new ConfigurationException("Bootstrapping to existing token " + tokenString + " is not allowed (decommission/removenode the old node first).");
-                tokens.add(token);
-            }
-            return tokens;
-        }
-
+            return getSpecifiedTokens(metadata, initialTokens);
+        
         int numTokens = DatabaseDescriptor.getNumTokens();
         if (numTokens < 1)
             throw new ConfigurationException("num_tokens must be >= 1");
 
+        if (allocationKeyspace != null)
+            return allocateTokens(metadata, address, allocationKeyspace, numTokens);
+
         if (numTokens == 1)
-            logger.warn("Picking random token for a single vnode.  You should probably add more vnodes; failing that, you should probably specify the token manually");
+            logger.warn("Picking random token for a single vnode.  You should probably add more vnodes and/or use the automatic token allocation mechanism.");
 
         return getRandomTokens(metadata, numTokens);
     }
 
+    private static Collection<Token> getSpecifiedTokens(final TokenMetadata metadata,
+                                                                Collection<String> initialTokens)
+    {
+        logger.debug("tokens manually specified as {}",  initialTokens);
+        List<Token> tokens = new ArrayList<>(initialTokens.size());
+        for (String tokenString : initialTokens)
+        {
+            Token token = StorageService.getPartitioner().getTokenFactory().fromString(tokenString);
+            if (metadata.getEndpoint(token) != null)
+                throw new ConfigurationException("Bootstrapping to existing token " + tokenString + " is not allowed (decommission/removenode the old node first).");
+            tokens.add(token);
+        }
+        return tokens;
+    }
+
+    static Collection<Token> allocateTokens(final TokenMetadata metadata,
+                                            InetAddress address,
+                                            String allocationKeyspace,
+                                            int numTokens)
+    {
+        Keyspace ks = Keyspace.open(allocationKeyspace);
+        if (ks == null)
+            throw new ConfigurationException("Problem opening token allocation keyspace " + allocationKeyspace);
+        AbstractReplicationStrategy rs = ks.getReplicationStrategy();
+        
+        return TokenAllocation.allocateTokens(metadata, rs, StorageService.getPartitioner(), address, numTokens);
+    }
+
     public static Collection<Token> getRandomTokens(TokenMetadata metadata, int numTokens)
     {
         Set<Token> tokens = new HashSet<>(numTokens);

src/java/org/apache/cassandra/dht/ByteOrderedPartitioner.java

@@ -116,6 +116,20 @@ public class ByteOrderedPartitioner implements IPartitioner
         {
             return token;
         }
+
+        @Override
+        public double size(Token next)
+        {
+            throw new UnsupportedOperationException(String.format("Token type %s does not support token allocation.",
+                                                                  getClass().getSimpleName()));
+        }
+
+        @Override
+        public Token increaseSlightly()
+        {
+            throw new UnsupportedOperationException(String.format("Token type %s does not support token allocation.",
+                                                                  getClass().getSimpleName()));
+        }
     }
 
     public BytesToken getToken(ByteBuffer key)

src/java/org/apache/cassandra/dht/ComparableObjectToken.java

@@ -66,4 +66,18 @@ abstract class ComparableObjectToken<C extends Comparable<C>> extends Token
 
         return token.compareTo(((ComparableObjectToken<C>) o).token);
     }
+
+    @Override
+    public double size(Token next)
+    {
+        throw new UnsupportedOperationException(String.format("Token type %s does not support token allocation.",
+                                                              getClass().getSimpleName()));
+    }
+
+    @Override
+    public Token increaseSlightly()
+    {
+        throw new UnsupportedOperationException(String.format("Token type %s does not support token allocation.",
+                                                              getClass().getSimpleName()));
+    }
 }

src/java/org/apache/cassandra/dht/Murmur3Partitioner.java

@@ -20,10 +20,7 @@ package org.apache.cassandra.dht;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.concurrent.ThreadLocalRandom;
 
 import org.apache.cassandra.db.DecoratedKey;
@@ -140,6 +137,21 @@ public class Murmur3Partitioner implements IPartitioner
         {
             return token;
         }
+
+        @Override
+        public double size(Token next)
+        {
+            LongToken n = (LongToken) next;
+            long v = n.token - token;  // Overflow acceptable and desired.
+            double d = Math.scalb((double) v, -Long.SIZE); // Scale so that the full range is 1.
+            return d > 0.0 ? d : (d + 1.0); // Adjust for signed long, also making sure t.size(t) == 1.
+        }
+
+        @Override
+        public Token increaseSlightly()
+        {
+            return new LongToken(token + 1);
+        }
     }
 
     /**
@@ -170,7 +182,12 @@ public class Murmur3Partitioner implements IPartitioner
 
     public LongToken getRandomToken()
     {
-        return new LongToken(normalize(ThreadLocalRandom.current().nextLong()));
+        return getRandomToken(ThreadLocalRandom.current());
+    }
+
+    public LongToken getRandomToken(Random r)
+    {
+        return new LongToken(normalize(r.nextLong()));
     }
 
     private long normalize(long v)

src/java/org/apache/cassandra/dht/Token.java

@@ -73,6 +73,18 @@ public abstract class Token implements RingPosition<Token>, Serializable
     abstract public long getHeapSize();
     abstract public Object getTokenValue();
 
+    /**
+     * Returns a measure for the token space covered between this token and next.
+     * Used by the token allocation algorithm (see CASSANDRA-7032).
+     */
+    abstract public double size(Token next);
+    /**
+     * Returns a token that is slightly greater than this. Used to avoid clashes
+     * between nodes in separate datacentres trying to use the same token via
+     * the token allocation algorithm.
+     */
+    abstract public Token increaseSlightly();
+
     public Token getToken()
     {
         return this;

src/java/org/apache/cassandra/dht/tokenallocator/ReplicationAwareTokenAllocator.java

@@ -0,0 +1,805 @@
+/*
+ * 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.HashMultimap;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Multimap;
+
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.dht.Token;
+
+/**
+ * A Replication Aware allocator for tokens, that attempts to ensure an even distribution of ownership across
+ * the known cluster for the provided replication strategy.
+ *
+ * A unit is shorthand for a "unit of ownership" which translates roughly to a node, or a disk on the node,
+ * a CPU on the node, or some other relevant unit of ownership. These units should be the lowest rung over which
+ * ownership needs to be evenly distributed. At the moment only nodes as a whole are treated as units, but that
+ * will change with the introduction of token ranges per disk.
+ */
+class ReplicationAwareTokenAllocator<Unit> implements TokenAllocator<Unit>
+{
+    final NavigableMap<Token, Unit> sortedTokens;
+    final Multimap<Unit, Token> unitToTokens;
+    final ReplicationStrategy<Unit> strategy;
+    final IPartitioner partitioner;
+    final int replicas;
+
+    ReplicationAwareTokenAllocator(NavigableMap<Token, Unit> sortedTokens, ReplicationStrategy<Unit> strategy, IPartitioner partitioner)
+    {
+        this.sortedTokens = sortedTokens;
+        unitToTokens = HashMultimap.create();
+        for (Map.Entry<Token, Unit> en : sortedTokens.entrySet())
+            unitToTokens.put(en.getValue(), en.getKey());
+        this.strategy = strategy;
+        this.replicas = strategy.replicas();
+        this.partitioner = partitioner;
+    }
+
+    public Collection<Token> addUnit(Unit newUnit, int numTokens)
+    {
+        assert !unitToTokens.containsKey(newUnit);
+
+        if (unitCount() < replicas)
+            // Allocation does not matter; everything replicates everywhere.
+            return generateRandomTokens(newUnit, numTokens);
+        if (numTokens > sortedTokens.size())
+            // Some of the heuristics below can't deal with this case. Use random for now, later allocations can fix any problems this may cause.
+            return generateRandomTokens(newUnit, numTokens);
+
+        // ============= construct our initial token ring state =============
+
+        double optTokenOwnership = optimalTokenOwnership(numTokens);
+        Map<Object, GroupInfo> groups = Maps.newHashMap();
+        Map<Unit, UnitInfo<Unit>> unitInfos = createUnitInfos(groups);
+        if (groups.size() < replicas)
+        {
+            // We need at least replicas groups to do allocation correctly. If there aren't enough, 
+            // use random allocation.
+            // This part of the code should only be reached via the RATATest. StrategyAdapter should disallow
+            // token allocation in this case as the algorithm is not able to cover the behavior of NetworkTopologyStrategy.
+            return generateRandomTokens(newUnit, numTokens);
+        }
+
+        // initialise our new unit's state (with an idealised ownership)
+        // strategy must already know about this unit
+        UnitInfo<Unit> newUnitInfo = new UnitInfo<>(newUnit, numTokens * optTokenOwnership, groups, strategy);
+
+        // build the current token ring state
+        TokenInfo<Unit> tokens = createTokenInfos(unitInfos, newUnitInfo.group);
+        newUnitInfo.tokenCount = numTokens;
+
+        // ============= construct and rank our candidate token allocations =============
+
+        // walk the token ring, constructing the set of candidates in ring order
+        // as the midpoints between all existing tokens
+        CandidateInfo<Unit> candidates = createCandidates(tokens, newUnitInfo, optTokenOwnership);
+
+        // Evaluate the expected improvements from all candidates and form a priority queue.
+        PriorityQueue<Weighted<CandidateInfo<Unit>>> improvements = new PriorityQueue<>(sortedTokens.size());
+        CandidateInfo<Unit> candidate = candidates;
+        do
+        {
+            double impr = evaluateImprovement(candidate, optTokenOwnership, 1.0 / numTokens);
+            improvements.add(new Weighted<>(impr, candidate));
+            candidate = candidate.next;
+        } while (candidate != candidates);
+
+        // ============= iteratively take the best candidate, and re-rank =============
+
+        CandidateInfo<Unit> bestToken = improvements.remove().value;
+        for (int vn = 1; ; ++vn)
+        {
+            candidates = bestToken.removeFrom(candidates);
+            confirmCandidate(bestToken);
+
+            if (vn == numTokens)
+                break;
+
+            while (true)
+            {
+                // Get the next candidate in the queue. Its improvement may have changed (esp. if multiple tokens
+                // were good suggestions because they could improve the same problem)-- evaluate it again to check
+                // if it is still a good candidate.
+                bestToken = improvements.remove().value;
+                double impr = evaluateImprovement(bestToken, optTokenOwnership, (vn + 1.0) / numTokens);
+                Weighted<CandidateInfo<Unit>> next = improvements.peek();
+
+                // If it is better than the next in the queue, it is good enough. This is a heuristic that doesn't
+                // get the best results, but works well enough and on average cuts search time by a factor of O(vnodes).
+                if (next == null || impr >= next.weight)
+                    break;
+                improvements.add(new Weighted<>(impr, bestToken));
+            }
+        }
+
+        return ImmutableList.copyOf(unitToTokens.get(newUnit));
+    }
+
+    private Collection<Token> generateRandomTokens(Unit newUnit, int numTokens)
+    {
+        Set<Token> tokens = new HashSet<>(numTokens);
+        while (tokens.size() < numTokens)
+        {
+            Token token = partitioner.getRandomToken();
+            if (!sortedTokens.containsKey(token))
+            {
+                tokens.add(token);
+                sortedTokens.put(token, newUnit);
+                unitToTokens.put(newUnit, token);
+            }
+        }
+        return tokens;
+    }
+
+    private Map<Unit, UnitInfo<Unit>> createUnitInfos(Map<Object, GroupInfo> groups)
+    {
+        Map<Unit, UnitInfo<Unit>> map = Maps.newHashMap();
+        for (Unit n : sortedTokens.values())
+        {
+            UnitInfo<Unit> ni = map.get(n);
+            if (ni == null)
+                map.put(n, ni = new UnitInfo<>(n, 0, groups, strategy));
+            ni.tokenCount++;
+        }
+        return map;
+    }
+
+    /**
+     * Construct the token ring as a CircularList of TokenInfo,
+     * and populate the ownership of the UnitInfo's provided
+     */
+    private TokenInfo<Unit> createTokenInfos(Map<Unit, UnitInfo<Unit>> units, GroupInfo newUnitGroup)
+    {
+        // build the circular list
+        TokenInfo<Unit> prev = null;
+        TokenInfo<Unit> first = null;
+        for (Map.Entry<Token, Unit> en : sortedTokens.entrySet())
+        {
+            Token t = en.getKey();
+            UnitInfo<Unit> ni = units.get(en.getValue());
+            TokenInfo<Unit> ti = new TokenInfo<>(t, ni);
+            first = ti.insertAfter(first, prev);
+            prev = ti;
+        }
+
+        TokenInfo<Unit> curr = first;
+        do
+        {
+            populateTokenInfoAndAdjustUnit(curr, newUnitGroup);
+            curr = curr.next;
+        } while (curr != first);
+
+        return first;
+    }
+
+    private CandidateInfo<Unit> createCandidates(TokenInfo<Unit> tokens, UnitInfo<Unit> newUnitInfo, double initialTokenOwnership)
+    {
+        TokenInfo<Unit> curr = tokens;
+        CandidateInfo<Unit> first = null;
+        CandidateInfo<Unit> prev = null;
+        do
+        {
+            CandidateInfo<Unit> candidate = new CandidateInfo<Unit>(partitioner.midpoint(curr.prev.token, curr.token), curr, newUnitInfo);
+            first = candidate.insertAfter(first, prev);
+
+            candidate.replicatedOwnership = initialTokenOwnership;
+            populateCandidate(candidate);
+
+            prev = candidate;
+            curr = curr.next;
+        } while (curr != tokens);
+        prev.next = first;
+        return first;
+    }
+
+    private void populateCandidate(CandidateInfo<Unit> candidate)
+    {
+        // Only finding replication start would do.
+        populateTokenInfo(candidate, candidate.owningUnit.group);
+    }
+
+    /**
+     * Incorporates the selected candidate into the ring, adjusting ownership information and calculated token
+     * information.
+     */
+    private void confirmCandidate(CandidateInfo<Unit> candidate)
+    {
+        // This process is less efficient than it could be (loops through each vnode's replication span instead
+        // of recalculating replicationStart, replicationThreshold from existing data + new token data in an O(1)
+        // case analysis similar to evaluateImprovement). This is fine as the method does not dominate processing
+        // time.
+
+        // Put the accepted candidate in the token list.
+        UnitInfo<Unit> newUnit = candidate.owningUnit;
+        Token newToken = candidate.token;
+        sortedTokens.put(newToken, newUnit.unit);
+        unitToTokens.put(newUnit.unit, newToken);
+
+        TokenInfo<Unit> prev = candidate.prevInRing();
+        TokenInfo<Unit> newTokenInfo = new TokenInfo<>(newToken, newUnit);
+        newTokenInfo.replicatedOwnership = candidate.replicatedOwnership;
+        newTokenInfo.insertAfter(prev, prev);   // List is not empty so this won't need to change head of list.
+
+        // Update data for candidate.
+        populateTokenInfoAndAdjustUnit(newTokenInfo, newUnit.group);
+
+        ReplicationVisitor replicationVisitor = new ReplicationVisitor();
+        assert newTokenInfo.next == candidate.split;
+        for (TokenInfo<Unit> curr = newTokenInfo.next; !replicationVisitor.visitedAll(); curr = curr.next)
+        {
+            // update the candidate between curr and next
+            candidate = candidate.next;
+            populateCandidate(candidate);
+
+            if (!replicationVisitor.add(curr.owningUnit.group))
+                continue;    // If we've already seen this group, the token cannot be affected.
+
+            populateTokenInfoAndAdjustUnit(curr, newUnit.group);
+        }
+
+        replicationVisitor.clean();
+    }
+
+    /**
+     * Calculates the {@code replicationStart} of a token, as well as {@code replicationThreshold} which is chosen in a way
+     * that permits {@code findUpdatedReplicationStart} to quickly identify changes in ownership.
+     */
+    private Token populateTokenInfo(BaseTokenInfo<Unit, ?> token, GroupInfo newUnitGroup)
+    {
+        GroupInfo tokenGroup = token.owningUnit.group;
+        PopulateVisitor visitor = new PopulateVisitor();
+
+        // Replication start = the end of a token from the RF'th different group seen before the token.
+        Token replicationStart;
+        // The end of a token from the RF-1'th different group seen before the token.
+        Token replicationThreshold = token.token;
+        GroupInfo currGroup;
+        for (TokenInfo<Unit> curr = token.prevInRing(); ; curr = curr.prev)
+        {
+            replicationStart = curr.token;
+            currGroup = curr.owningUnit.group;
+            if (!visitor.add(currGroup))
+                continue; // Group is already seen.
+            if (visitor.visitedAll())
+                break;
+
+            replicationThreshold = replicationStart;
+            // Another instance of the same group precedes us in the replication range of the ring,
+            // so this is where our replication range begins
+            if (currGroup == tokenGroup)
+                break;
+        }
+        if (newUnitGroup == tokenGroup)
+            // new token is always a boundary (as long as it's closer than replicationStart)
+            replicationThreshold = token.token;
+        else if (newUnitGroup != currGroup && visitor.seen(newUnitGroup))
+            // already has new group in replication span before last seen. cannot be affected
+            replicationThreshold = replicationStart;
+        visitor.clean();
+
+        token.replicationThreshold = replicationThreshold;
+        token.replicationStart = replicationStart;
+        return replicationStart;
+    }
+
+    private void populateTokenInfoAndAdjustUnit(TokenInfo<Unit> populate, GroupInfo newUnitGroup)
+    {
+        Token replicationStart = populateTokenInfo(populate, newUnitGroup);
+        double newOwnership = replicationStart.size(populate.token);
+        double oldOwnership = populate.replicatedOwnership;
+        populate.replicatedOwnership = newOwnership;
+        populate.owningUnit.ownership += newOwnership - oldOwnership;
+    }
+
+    /**
+     * Evaluates the improvement in variance for both units and individual tokens when candidate is inserted into the
+     * ring.
+     */
+    private double evaluateImprovement(CandidateInfo<Unit> candidate, double optTokenOwnership, double newUnitMult)
+    {
+        double tokenChange = 0;
+
+        UnitInfo<Unit> candidateUnit = candidate.owningUnit;
+        Token candidateEnd = candidate.token;
+
+        // Form a chain of units affected by the insertion to be able to qualify change of unit ownership.
+        // A unit may be affected more than once.
+        UnitAdjustmentTracker<Unit> unitTracker = new UnitAdjustmentTracker<>(candidateUnit);
+
+        // Reflect change in ownership of the splitting token (candidate).
+        tokenChange += applyOwnershipAdjustment(candidate, candidateUnit, candidate.replicationStart, candidateEnd, optTokenOwnership, unitTracker);
+
+        // Loop through all vnodes that replicate candidate or split and update their ownership.
+        ReplicationVisitor replicationVisitor = new ReplicationVisitor();
+        for (TokenInfo<Unit> curr = candidate.split; !replicationVisitor.visitedAll(); curr = curr.next)
+        {
+            UnitInfo<Unit> currUnit = curr.owningUnit;
+
+            if (!replicationVisitor.add(currUnit.group))
+                continue;    // If this group is already seen, the token cannot be affected.
+
+            Token replicationEnd = curr.token;
+            Token replicationStart = findUpdatedReplicationStart(curr, candidate);
+            tokenChange += applyOwnershipAdjustment(curr, currUnit, replicationStart, replicationEnd, optTokenOwnership, unitTracker);
+        }
+        replicationVisitor.clean();
+
+        double nodeChange = unitTracker.calculateUnitChange(newUnitMult, optTokenOwnership);
+        return -(tokenChange + nodeChange);
+    }
+
+    /**
+     * Returns the start of the replication span for the token {@code curr} when {@code candidate} is inserted into the
+     * ring.
+     */
+    private Token findUpdatedReplicationStart(TokenInfo<Unit> curr, CandidateInfo<Unit> candidate)
+    {
+        return furtherStartToken(curr.replicationThreshold, candidate.token, curr.token);
+    }
+
+    /**
+     * Applies the ownership adjustment for the given element, updating tracked unit ownership and returning the change
+     * of variance.
+     */
+    private double applyOwnershipAdjustment(BaseTokenInfo<Unit, ?> curr, UnitInfo<Unit> currUnit,
+            Token replicationStart, Token replicationEnd,
+            double optTokenOwnership, UnitAdjustmentTracker<Unit> unitTracker)
+    {
+        double oldOwnership = curr.replicatedOwnership;
+        double newOwnership = replicationStart.size(replicationEnd);
+        double tokenCount = currUnit.tokenCount;
+        assert tokenCount > 0;
+        unitTracker.add(currUnit, newOwnership - oldOwnership);
+        return (sq(newOwnership - optTokenOwnership) - sq(oldOwnership - optTokenOwnership)) / sq(tokenCount);
+    }
+
+    /**
+     * Tracker for unit ownership changes. The changes are tracked by a chain of UnitInfos where the adjustedOwnership
+     * field is being updated as we see changes in token ownership.
+     *
+     * The chain ends with an element that points to itself; this element must be specified as argument to the
+     * constructor as well as be the first unit with which 'add' is called; when calculating the variance change
+     * a separate multiplier is applied to it (used to permit more freedom in choosing the first tokens of a unit).
+     */
+    private static class UnitAdjustmentTracker<Unit>
+    {
+        UnitInfo<Unit> unitsChain;
+
+        UnitAdjustmentTracker(UnitInfo<Unit> newUnit)
+        {
+            unitsChain = newUnit;
+        }
+
+        void add(UnitInfo<Unit> currUnit, double diff)
+        {
+            if (currUnit.prevUsed == null)
+            {
+                assert unitsChain.prevUsed != null || currUnit == unitsChain;
+
+                currUnit.adjustedOwnership = currUnit.ownership + diff;
+                currUnit.prevUsed = unitsChain;
+                unitsChain = currUnit;
+            }
+            else
+            {
+                currUnit.adjustedOwnership += diff;
+            }
+        }
+
+        double calculateUnitChange(double newUnitMult, double optTokenOwnership)
+        {
+            double unitChange = 0;
+            UnitInfo<Unit> unitsChain = this.unitsChain;
+            // Now loop through the units chain and add the unit-level changes. Also clear the groups' seen marks.
+            while (true)
+            {
+                double newOwnership = unitsChain.adjustedOwnership;
+                double oldOwnership = unitsChain.ownership;
+                double tokenCount = unitsChain.tokenCount;
+                double diff = (sq(newOwnership / tokenCount - optTokenOwnership) - sq(oldOwnership / tokenCount - optTokenOwnership));
+                UnitInfo<Unit> prev = unitsChain.prevUsed;
+                unitsChain.prevUsed = null;
+                if (unitsChain != prev)
+                    unitChange += diff;
+                else
+                {
+                    unitChange += diff * newUnitMult;
+                    break;
+                }
+                unitsChain = prev;
+            }
+            this.unitsChain = unitsChain;
+            return unitChange;
+        }
+    }
+
+
+    /**
+     * Helper class for marking/unmarking visited a chain of groups
+     */
+    private abstract class GroupVisitor
+    {
+        GroupInfo groupChain = GroupInfo.TERMINATOR;
+        int seen = 0;
+
+        abstract GroupInfo prevSeen(GroupInfo group);
+        abstract void setPrevSeen(GroupInfo group, GroupInfo prevSeen);
+
+        // true iff this is the first time we've visited this group
+        boolean add(GroupInfo group)
+        {
+            if (prevSeen(group) != null)
+                return false;
+            ++seen;
+            setPrevSeen(group, groupChain);
+            groupChain = group;
+            return true;
+        }
+
+        boolean visitedAll()
+        {
+            return seen >= replicas;
+        }
+
+        boolean seen(GroupInfo group)
+        {
+            return prevSeen(group) != null;
+        }
+
+        // Clean group seen markers.
+        void clean()
+        {
+            GroupInfo groupChain = this.groupChain;
+            while (groupChain != GroupInfo.TERMINATOR)
+            {
+                GroupInfo prev = prevSeen(groupChain);
+                setPrevSeen(groupChain, null);
+                groupChain = prev;
+            }
+            this.groupChain = GroupInfo.TERMINATOR;
+        }
+    }
+
+    private class ReplicationVisitor extends GroupVisitor
+    {
+        GroupInfo prevSeen(GroupInfo group)
+        {
+            return group.prevSeen;
+        }
+
+        void setPrevSeen(GroupInfo group, GroupInfo prevSeen)
+        {
+            group.prevSeen = prevSeen;
+        }
+    }
+
+    private class PopulateVisitor extends GroupVisitor
+    {
+        GroupInfo prevSeen(GroupInfo group)
+        {
+            return group.prevPopulate;
+        }
+
+        void setPrevSeen(GroupInfo group, GroupInfo prevSeen)
+        {
+            group.prevPopulate = prevSeen;
+        }
+    }
+
+    private Map.Entry<Token, Unit> mapEntryFor(Token t)
+    {
+        Map.Entry<Token, Unit> en = sortedTokens.floorEntry(t);
+        if (en == null)
+            en = sortedTokens.lastEntry();
+        return en;
+    }
+
+    Unit unitFor(Token t)
+    {
+        return mapEntryFor(t).getValue();
+    }
+
+    private double optimalTokenOwnership(int tokensToAdd)
+    {
+        return 1.0 * replicas / (sortedTokens.size() + tokensToAdd);
+    }
+
+    /**
+     * Selects from {@code t1}, {@code t2} the token that forms a bigger range with {@code towards} as the upper bound,
+     * taking into account wrapping.
+     * Unlike Token.size(), equality is taken to mean "same as" rather than covering the whole range.
+     */
+    private static Token furtherStartToken(Token t1, Token t2, Token towards)
+    {
+        if (t1.equals(towards))
+            return t2;
+        if (t2.equals(towards))
+            return t1;
+
+        return t1.size(towards) > t2.size(towards) ? t1 : t2;
+    }
+
+    private static double sq(double d)
+    {
+        return d * d;
+    }
+
+
+    /**
+     * For testing, remove the given unit preserving correct state of the allocator.
+     */
+    void removeUnit(Unit n)
+    {
+        Collection<Token> tokens = unitToTokens.removeAll(n);
+        sortedTokens.keySet().removeAll(tokens);
+    }
+
+    int unitCount()
+    {
+        return unitToTokens.asMap().size();
+    }
+
+    public String toString()
+    {
+        return getClass().getSimpleName();
+    }
+
+    // get or initialise the shared GroupInfo associated with the unit
+    private static <Unit> GroupInfo getGroup(Unit unit, Map<Object, GroupInfo> groupMap, ReplicationStrategy<Unit> strategy)
+    {
+        Object groupClass = strategy.getGroup(unit);
+        GroupInfo group = groupMap.get(groupClass);
+        if (group == null)
+            groupMap.put(groupClass, group = new GroupInfo(groupClass));
+        return group;
+    }
+
+    /**
+     * Unique group object that one or more UnitInfo objects link to.
+     */
+    private static class GroupInfo
+    {
+        /**
+         * Group identifier given by ReplicationStrategy.getGroup(Unit).
+         */
+        final Object group;
+
+        /**
+         * Seen marker. When non-null, the group is already seen in replication walks.
+         * Also points to previous seen group to enable walking the seen groups and clearing the seen markers.
+         */
+        GroupInfo prevSeen = null;
+        /**
+         * Same marker/chain used by populateTokenInfo.
+         */
+        GroupInfo prevPopulate = null;
+
+        /**
+         * Value used as terminator for seen chains.
+         */
+        static GroupInfo TERMINATOR = new GroupInfo(null);
+
+        public GroupInfo(Object group)
+        {
+            this.group = group;
+        }
+
+        public String toString()
+        {
+            return group.toString() + (prevSeen != null ? "*" : "");
+        }
+    }
+
+    /**
+     * Unit information created and used by ReplicationAwareTokenDistributor. Contained vnodes all point to the same
+     * instance.
+     */
+    static class UnitInfo<Unit>
+    {
+        final Unit unit;
+        final GroupInfo group;
+        double ownership;
+        int tokenCount;
+
+        /**
+         * During evaluateImprovement this is used to form a chain of units affected by the candidate insertion.
+         */
+        UnitInfo<Unit> prevUsed;
+        /**
+         * During evaluateImprovement this holds the ownership after the candidate insertion.
+         */
+        double adjustedOwnership;
+
+        private UnitInfo(Unit unit, GroupInfo group)
+        {
+            this.unit = unit;
+            this.group = group;
+            this.tokenCount = 0;
+        }
+
+        public UnitInfo(Unit unit, double ownership, Map<Object, GroupInfo> groupMap, ReplicationStrategy<Unit> strategy)
+        {
+            this(unit, getGroup(unit, groupMap, strategy));
+            this.ownership = ownership;
+        }
+
+        public String toString()
+        {
+            return String.format("%s%s(%.2e)%s",
+                    unit, unit == group.group ? (group.prevSeen != null ? "*" : "") : ":" + group.toString(),
+                    ownership, prevUsed != null ? (prevUsed == this ? "#" : "->" + prevUsed.toString()) : "");
+        }
+    }
+
+    private static class CircularList<T extends CircularList<T>>
+    {
+        T prev;
+        T next;
+
+        /**
+         * Inserts this after unit in the circular list which starts at head. Returns the new head of the list, which
+         * only changes if head was null.
+         */
+        @SuppressWarnings("unchecked")
+        T insertAfter(T head, T unit)
+        {
+            if (head == null)
+            {
+                return prev = next = (T) this;
+            }
+            assert unit != null;
+            assert unit.next != null;
+            prev = unit;
+            next = unit.next;
+            prev.next = (T) this;
+            next.prev = (T) this;
+            return head;
+        }
+
+        /**
+         * Removes this from the list that starts at head. Returns the new head of the list, which only changes if the
+         * head was removed.
+         */
+        T removeFrom(T head)
+        {
+            next.prev = prev;
+            prev.next = next;
+            return this == head ? (this == next ? null : next) : head;
+        }
+    }
+
+    private static class BaseTokenInfo<Unit, T extends BaseTokenInfo<Unit, T>> extends CircularList<T>
+    {
+        final Token token;
+        final UnitInfo<Unit> owningUnit;
+
+        /**
+         * Start of the replication span for the vnode, i.e. the first token of the RF'th group seen before the token.
+         * The replicated ownership of the unit is the range between {@code replicationStart} and {@code token}.
+         */
+        Token replicationStart;
+        /**
+         * The closest position that the new candidate can take to become the new replication start. If candidate is
+         * closer, the start moves to this position. Used to determine replicationStart after insertion of new token.
+         *
+         * Usually the RF minus one boundary, i.e. the first token of the RF-1'th group seen before the token.
+         */
+        Token replicationThreshold;
+        /**
+         * Current replicated ownership. This number is reflected in the owning unit's ownership.
+         */
+        double replicatedOwnership = 0;
+
+        public BaseTokenInfo(Token token, UnitInfo<Unit> owningUnit)
+        {
+            this.token = token;
+            this.owningUnit = owningUnit;
+        }
+
+        public String toString()
+        {
+            return String.format("%s(%s)", token, owningUnit);
+        }
+
+        /**
+         * Previous unit in the token ring. For existing tokens this is prev,
+         * for candidates it's "split".
+         */
+        TokenInfo<Unit> prevInRing()
+        {
+            return null;
+        }
+    }
+
+    /**
+     * TokenInfo about existing tokens/vnodes.
+     */
+    private static class TokenInfo<Unit> extends BaseTokenInfo<Unit, TokenInfo<Unit>>
+    {
+        public TokenInfo(Token token, UnitInfo<Unit> owningUnit)
+        {
+            super(token, owningUnit);
+        }
+
+        TokenInfo<Unit> prevInRing()
+        {
+            return prev;
+        }
+    }
+
+    /**
+     * TokenInfo about candidate new tokens/vnodes.
+     */
+    private static class CandidateInfo<Unit> extends BaseTokenInfo<Unit, CandidateInfo<Unit>>
+    {
+        // directly preceding token in the current token ring
+        final TokenInfo<Unit> split;
+
+        public CandidateInfo(Token token, TokenInfo<Unit> split, UnitInfo<Unit> owningUnit)
+        {
+            super(token, owningUnit);
+            this.split = split;
+        }
+
+        TokenInfo<Unit> prevInRing()
+        {
+            return split.prev;
+        }
+    }
+
+    static void dumpTokens(String lead, BaseTokenInfo<?, ?> tokens)
+    {
+        BaseTokenInfo<?, ?> token = tokens;
+        do
+        {
+            System.out.format("%s%s: rs %s rt %s size %.2e\n", lead, token, token.replicationStart, token.replicationThreshold, token.replicatedOwnership);
+            token = token.next;
+        } while (token != null && token != tokens);
+    }
+
+    static class Weighted<T> implements Comparable<Weighted<T>>
+    {
+        final double weight;
+        final T value;
+
+        public Weighted(double weight, T value)
+        {
+            this.weight = weight;
+            this.value = value;
+        }
+
+        @Override
+        public int compareTo(Weighted<T> o)
+        {
+            int cmp = Double.compare(o.weight, this.weight);
+            return cmp;
+        }
+
+        @Override
+        public String toString()
+        {
+            return String.format("%s<%s>", value, weight);
+        }
+    }
+}
+

src/java/org/apache/cassandra/dht/tokenallocator/ReplicationStrategy.java

@@ -0,0 +1,29 @@
+/*
+ * 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;
+
+interface ReplicationStrategy<Unit>
+{
+    int replicas();
+
+    /**
+     * Returns a group identifier. getGroup(a) == getGroup(b) iff a and b are on the same group.
+     * @return Some hashable object.
+     */
+    Object getGroup(Unit unit);
+}

src/java/org/apache/cassandra/dht/tokenallocator/TokenAllocation.java

@@ -0,0 +1,259 @@
+/*
+ * 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.net.InetAddress;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.TreeMap;
+
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+
+import org.apache.commons.math.stat.descriptive.SummaryStatistics;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.exceptions.ConfigurationException;
+import org.apache.cassandra.locator.AbstractReplicationStrategy;
+import org.apache.cassandra.locator.IEndpointSnitch;
+import org.apache.cassandra.locator.NetworkTopologyStrategy;
+import org.apache.cassandra.locator.SimpleStrategy;
+import org.apache.cassandra.locator.TokenMetadata;
+import org.apache.cassandra.locator.TokenMetadata.Topology;
+
+public class TokenAllocation
+{
+    private static final Logger logger = LoggerFactory.getLogger(TokenAllocation.class);
+
+    public static Collection<Token> allocateTokens(final TokenMetadata tokenMetadata,
+                                                   final AbstractReplicationStrategy rs,
+                                                   final IPartitioner partitioner,
+                                                   final InetAddress endpoint,
+                                                   int numTokens)
+    {
+        StrategyAdapter strategy = getStrategy(tokenMetadata, rs, endpoint);
+        Collection<Token> tokens = create(tokenMetadata, strategy, partitioner).addUnit(endpoint, numTokens);
+        tokens = adjustForCrossDatacenterClashes(tokenMetadata, strategy, tokens);
+
+        if (logger.isWarnEnabled())
+        {
+            logger.warn("Selected tokens {}", tokens);
+            SummaryStatistics os = replicatedOwnershipStats(tokenMetadata, rs, endpoint);
+            TokenMetadata tokenMetadataCopy = tokenMetadata.cloneOnlyTokenMap();
+            tokenMetadataCopy.updateNormalTokens(tokens, endpoint);
+            SummaryStatistics ns = replicatedOwnershipStats(tokenMetadataCopy, rs, endpoint);
+            logger.warn("Replicated node load in datacentre before allocation " + statToString(os));
+            logger.warn("Replicated node load in datacentre after allocation " + statToString(ns));
+
+            // TODO: Is it worth doing the replicated ownership calculation always to be able to raise this alarm?
+            if (ns.getStandardDeviation() > os.getStandardDeviation())
+                logger.warn("Unexpected growth in standard deviation after allocation.");
+        }
+        return tokens;
+    }
+
+    private static Collection<Token> adjustForCrossDatacenterClashes(final TokenMetadata tokenMetadata,
+                                                                     StrategyAdapter strategy, Collection<Token> tokens)
+    {
+        List<Token> filtered = Lists.newArrayListWithCapacity(tokens.size());
+
+        for (Token t : tokens)
+        {
+            while (tokenMetadata.getEndpoint(t) != null)
+            {
+                InetAddress other = tokenMetadata.getEndpoint(t);
+                if (strategy.inAllocationRing(other))
+                    throw new ConfigurationException(String.format("Allocated token %s already assigned to node %s. Is another node also allocating tokens?", t, other));
+                t = t.increaseSlightly();
+            }
+            filtered.add(t);
+        }
+        return filtered;
+    }
+
+    // return the ratio of ownership for each endpoint
+    public static Map<InetAddress, Double> evaluateReplicatedOwnership(TokenMetadata tokenMetadata, AbstractReplicationStrategy rs)
+    {
+        Map<InetAddress, Double> ownership = Maps.newHashMap();
+        List<Token> sortedTokens = tokenMetadata.sortedTokens();
+        Iterator<Token> it = sortedTokens.iterator();
+        Token current = it.next();
+        while (it.hasNext())
+        {
+            Token next = it.next();
+            addOwnership(tokenMetadata, rs, current, next, ownership);
+            current = next;
+        }
+        addOwnership(tokenMetadata, rs, current, sortedTokens.get(0), ownership);
+
+        return ownership;
+    }
+
+    static void addOwnership(final TokenMetadata tokenMetadata, final AbstractReplicationStrategy rs, Token current, Token next, Map<InetAddress, Double> ownership)
+    {
+        double size = current.size(next);
+        Token representative = current.getPartitioner().midpoint(current, next);
+        for (InetAddress n : rs.calculateNaturalEndpoints(representative, tokenMetadata))
+        {
+            Double v = ownership.get(n);
+            ownership.put(n, v != null ? v + size : size);
+        }
+    }
+
+    public static String statToString(SummaryStatistics stat)
+    {
+        return String.format("max %.2f min %.2f stddev %.4f", stat.getMax() / stat.getMean(), stat.getMin() / stat.getMean(), stat.getStandardDeviation());
+    }
+
+    public static SummaryStatistics replicatedOwnershipStats(TokenMetadata tokenMetadata,
+                                                             AbstractReplicationStrategy rs, InetAddress endpoint)
+    {
+        SummaryStatistics stat = new SummaryStatistics();
+        StrategyAdapter strategy = getStrategy(tokenMetadata, rs, endpoint);
+        for (Map.Entry<InetAddress, Double> en : evaluateReplicatedOwnership(tokenMetadata, rs).entrySet())
+        {
+            // Filter only in the same datacentre.
+            if (strategy.inAllocationRing(en.getKey()))
+                stat.addValue(en.getValue() / tokenMetadata.getTokens(en.getKey()).size());
+        }
+        return stat;
+    }
+
+    static TokenAllocator<InetAddress> create(TokenMetadata tokenMetadata, StrategyAdapter strategy, IPartitioner partitioner)
+    {
+        NavigableMap<Token, InetAddress> sortedTokens = new TreeMap<>();
+        for (Map.Entry<Token, InetAddress> en : tokenMetadata.getNormalAndBootstrappingTokenToEndpointMap().entrySet())
+        {
+            if (strategy.inAllocationRing(en.getValue()))
+                sortedTokens.put(en.getKey(), en.getValue());
+        }
+        return new ReplicationAwareTokenAllocator<>(sortedTokens, strategy, partitioner);
+    }
+
+    interface StrategyAdapter extends ReplicationStrategy<InetAddress>
+    {
+        // return true iff the provided endpoint occurs in the same virtual token-ring we are allocating for
+        // i.e. the set of the nodes that share ownership with the node we are allocating
+        // alternatively: return false if the endpoint's ownership is independent of the node we are allocating tokens for
+        boolean inAllocationRing(InetAddress other);
+    }
+
+    static StrategyAdapter getStrategy(final TokenMetadata tokenMetadata, final AbstractReplicationStrategy rs, final InetAddress endpoint)
+    {
+        if (rs instanceof NetworkTopologyStrategy)
+            return getStrategy(tokenMetadata, (NetworkTopologyStrategy) rs, rs.snitch, endpoint);
+        if (rs instanceof SimpleStrategy)
+            return getStrategy(tokenMetadata, (SimpleStrategy) rs, endpoint);
+        throw new ConfigurationException("Token allocation does not support replication strategy " + rs.getClass().getSimpleName());
+    }
+
+    static StrategyAdapter getStrategy(final TokenMetadata tokenMetadata, final SimpleStrategy rs, final InetAddress endpoint)
+    {
+        final int replicas = rs.getReplicationFactor();
+
+        return new StrategyAdapter()
+        {
+            @Override
+            public int replicas()
+            {
+                return replicas;
+            }
+
+            @Override
+            public Object getGroup(InetAddress unit)
+            {
+                return unit;
+            }
+
+            @Override
+            public boolean inAllocationRing(InetAddress other)
+            {
+                return true;
+            }
+        };
+    }
+
+    static StrategyAdapter getStrategy(final TokenMetadata tokenMetadata, final NetworkTopologyStrategy rs, final IEndpointSnitch snitch, final InetAddress endpoint)
+    {
+        final String dc = snitch.getDatacenter(endpoint);
+        final int replicas = rs.getReplicationFactor(dc);
+
+        Topology topology = tokenMetadata.getTopology();
+        int racks = topology.getDatacenterRacks().get(dc).size();
+
+        if (replicas >= racks)
+        {
+            return new StrategyAdapter()
+            {
+                @Override
+                public int replicas()
+                {
+                    return replicas;
+                }
+
+                @Override
+                public Object getGroup(InetAddress unit)
+                {
+                    return snitch.getRack(unit);
+                }
+
+                @Override
+                public boolean inAllocationRing(InetAddress other)
+                {
+                    return dc.equals(snitch.getDatacenter(other));
+                }
+            };
+        }
+        else if (racks == 1)
+        {
+            // One rack, each node treated as separate.
+            return new StrategyAdapter()
+            {
+                @Override
+                public int replicas()
+                {
+                    return replicas;
+                }
+
+                @Override
+                public Object getGroup(InetAddress unit)
+                {
+                    return unit;
+                }
+
+                @Override
+                public boolean inAllocationRing(InetAddress other)
+                {
+                    return dc.equals(snitch.getDatacenter(other));
+                }
+            };
+        }
+        else
+            throw new ConfigurationException(
+                                            String.format("Token allocation failed: the number of racks %d in datacentre %s is lower than its replication factor %d.",
+                                                          replicas, dc, racks));
+    }
+}
+

src/java/org/apache/cassandra/dht/tokenallocator/TokenAllocator.java

@@ -0,0 +1,27 @@
+/*
+ * 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.Collection;
+
+import org.apache.cassandra.dht.Token;
+
+public interface TokenAllocator<Unit>
+{
+    public Collection<Token> addUnit(Unit newUnit, int numTokens);
+}

src/java/org/apache/cassandra/locator/TokenMetadata.java

@@ -781,7 +781,8 @@ public class TokenMetadata
     {
         List tokens = sortedTokens();
         int index = Collections.binarySearch(tokens, token);
-        assert index >= 0 : token + " not found in " + StringUtils.join(tokenToEndpointMap.keySet(), ", ");
+//        assert index >= 0 : token + " not found in " + StringUtils.join(tokenToEndpointMap.keySet(), ", ");
+        if (index < 0) index = -index-1;
         return (Token) (index == 0 ? tokens.get(tokens.size() - 1) : tokens.get(index - 1));
     }
 
@@ -789,7 +790,8 @@ public class TokenMetadata
     {
         List tokens = sortedTokens();
         int index = Collections.binarySearch(tokens, token);
-        assert index >= 0 : token + " not found in " + StringUtils.join(tokenToEndpointMap.keySet(), ", ");
+//        assert index >= 0 : token + " not found in " + StringUtils.join(tokenToEndpointMap.keySet(), ", ");
+        if (index < 0) return (Token) tokens.get(-index-1);
         return (Token) ((index == (tokens.size() - 1)) ? tokens.get(0) : tokens.get(index + 1));
     }
 

src/java/org/apache/cassandra/service/StorageService.java

@@ -876,7 +876,7 @@ public class StorageService extends NotificationBroadcasterSupport implements IE
                     throw new UnsupportedOperationException(s);
                 }
                 setMode(Mode.JOINING, "getting bootstrap token", true);
-                bootstrapTokens = BootStrapper.getBootstrapTokens(tokenMetadata);
+                bootstrapTokens = BootStrapper.getBootstrapTokens(tokenMetadata, FBUtilities.getBroadcastAddress());
             }
             else
             {

test/conf/logback-test.xml

@@ -66,7 +66,6 @@
 
   <root level="DEBUG">
     <appender-ref ref="ASYNCFILE" />
-    <appender-ref ref="STDERR" />
     <appender-ref ref="STDOUT" />
   </root>
 </configuration>

test/long/org/apache/cassandra/db/commitlog/CommitLogStressTest.java

@@ -127,7 +127,7 @@ public class CommitLogStressTest
     ReplayPosition discardedPos;
     
     @BeforeClass
-    static public void initialize() throws FileNotFoundException, IOException, InterruptedException
+    public static void initialize() throws FileNotFoundException, IOException, InterruptedException
     {
         try (FileInputStream fis = new FileInputStream("CHANGES.txt"))
         {

test/long/org/apache/cassandra/dht/tokenallocator/ReplicationAwareTokenAllocatorTest.java

@@ -0,0 +1,700 @@
+/*
+ * 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 junit.framework.Assert;
+
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
+
+import org.apache.commons.math.stat.descriptive.SummaryStatistics;
+
+import org.junit.Test;
+
+import org.apache.cassandra.dht.Murmur3Partitioner;
+import org.apache.cassandra.dht.Token;
+
+public class ReplicationAwareTokenAllocatorTest
+{
+    private static final int MAX_VNODE_COUNT = 64;
+
+    private static final int TARGET_CLUSTER_SIZE = 250;
+
+    interface TestReplicationStrategy extends ReplicationStrategy<Unit>
+    {
+        void addUnit(Unit n);
+
+        void removeUnit(Unit n);
+
+        /**
+         * Returns a list of all replica units for given token.
+         */
+        List<Unit> getReplicas(Token token, NavigableMap<Token, Unit> sortedTokens);
+
+        /**
+         * Returns the start of the token span that is replicated in this token.
+         * Note: Though this is not trivial to see, the replicated span is always contiguous. A token in the same
+         * group acts as a barrier; if one is not found the token replicates everything up to the replica'th distinct
+         * group seen in front of it.
+         */
+        Token replicationStart(Token token, Unit unit, NavigableMap<Token, Unit> sortedTokens);
+
+        /**
+         * Multiplier for the acceptable disbalance in the cluster. With some strategies it is harder to achieve good
+         * results.
+         */
+        public double spreadExpectation();
+    }
+
+    static class NoReplicationStrategy implements TestReplicationStrategy
+    {
+        public List<Unit> getReplicas(Token token, NavigableMap<Token, Unit> sortedTokens)
+        {
+            return Collections.singletonList(sortedTokens.ceilingEntry(token).getValue());
+        }
+
+        public Token replicationStart(Token token, Unit unit, NavigableMap<Token, Unit> sortedTokens)
+        {
+            return sortedTokens.lowerKey(token);
+        }
+
+        public String toString()
+        {
+            return "No replication";
+        }
+
+        public void addUnit(Unit n)
+        {
+        }
+
+        public void removeUnit(Unit n)
+        {
+        }
+
+        public int replicas()
+        {
+            return 1;
+        }
+
+        public boolean sameGroup(Unit n1, Unit n2)
+        {
+            return false;
+        }
+
+        public Object getGroup(Unit unit)
+        {
+            return unit;
+        }
+
+        public double spreadExpectation()
+        {
+            return 1;
+        }
+    }
+
+    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);
+    }
+
+    static interface TokenCount
+    {
+        int tokenCount(int perUnitCount, Random rand);
+
+        double spreadExpectation();
+    }
+
+    static TokenCount fixedTokenCount = new TokenCount()
+    {
+        public int tokenCount(int perUnitCount, Random rand)
+        {
+            return perUnitCount;
+        }
+
+        public double spreadExpectation()
+        {
+            return 4;  // High tolerance to avoid flakiness.
+        }
+    };
+
+    static TokenCount varyingTokenCount = new TokenCount()
+    {
+        public int tokenCount(int perUnitCount, Random rand)
+        {
+            if (perUnitCount == 1) return 1;
+            // 25 to 175%
+            return rand.nextInt(perUnitCount * 3 / 2) + (perUnitCount + 3) / 4;
+        }
+
+        public double spreadExpectation()
+        {
+            return 8;  // High tolerance to avoid flakiness.
+        }
+    };
+
+    Murmur3Partitioner partitioner = new Murmur3Partitioner();
+    Random seededRand = new Random(2);
+
+    private void random(Map<Token, Unit> map, TestReplicationStrategy rs, int unitCount, TokenCount tc, int perUnitCount)
+    {
+        System.out.format("\nRandom generation of %d units with %d tokens each\n", unitCount, perUnitCount);
+        Random rand = seededRand;
+        for (int i = 0; i < unitCount; i++)
+        {
+            Unit unit = new Unit();
+            rs.addUnit(unit);
+            int tokens = tc.tokenCount(perUnitCount, rand);
+            for (int j = 0; j < tokens; j++)
+            {
+                map.put(partitioner.getRandomToken(rand), unit);
+            }
+        }
+    }
+
+    @Test
+    public void testExistingCluster()
+    {
+        for (int rf = 1; rf <= 5; ++rf)
+        {
+            for (int perUnitCount = 1; perUnitCount <= MAX_VNODE_COUNT; perUnitCount *= 4)
+            {
+                testExistingCluster(perUnitCount, fixedTokenCount, new SimpleReplicationStrategy(rf));
+                testExistingCluster(perUnitCount, varyingTokenCount, new SimpleReplicationStrategy(rf));
+                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));
+                    testExistingCluster(perUnitCount, varyingTokenCount, new UnbalancedGroupReplicationStrategy(rf, groupSize / 2, groupSize * 2, seededRand));
+                }
+                testExistingCluster(perUnitCount, fixedTokenCount, new FixedGroupCountReplicationStrategy(rf, rf * 2));
+            }
+        }
+    }
+
+    public void testExistingCluster(int perUnitCount, TokenCount tc, TestReplicationStrategy rs)
+    {
+        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);
+
+        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);
+        System.out.println();
+    }
+
+    @Test
+    public void testNewCluster()
+    {
+        for (int rf = 2; rf <= 5; ++rf)
+        {
+            for (int perUnitCount = 1; perUnitCount <= MAX_VNODE_COUNT; perUnitCount *= 4)
+            {
+                testNewCluster(perUnitCount, fixedTokenCount, new SimpleReplicationStrategy(rf));
+                testNewCluster(perUnitCount, varyingTokenCount, new SimpleReplicationStrategy(rf));
+                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));
+                    testNewCluster(perUnitCount, varyingTokenCount, new UnbalancedGroupReplicationStrategy(rf, groupSize / 2, groupSize * 2, seededRand));
+                }
+                testNewCluster(perUnitCount, fixedTokenCount, new FixedGroupCountReplicationStrategy(rf, rf * 2));
+            }
+        }
+    }
+
+    public void testNewCluster(int perUnitCount, TokenCount tc, TestReplicationStrategy rs)
+    {
+        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);
+        System.out.println();
+    }
+
+    private void loseAndReplace(ReplicationAwareTokenAllocator<Unit> t, int howMany, TokenCount tc, int perUnitCount)
+    {
+        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);
+    }
+
+    static class Summary
+    {
+        double min = 1;
+        double max = 1;
+        double stddev = 0;
+
+        void update(SummaryStatistics stat)
+        {
+            min = Math.min(min, stat.getMin());
+            max = Math.max(max, stat.getMax());
+            stddev = Math.max(stddev, stat.getStandardDeviation());
+        }
+
+        public String toString()
+        {
+            return String.format("max %.2f min %.2f stddev %.4f", max, min, stddev);
+        }
+    }
+
+    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);
+        }
+    }
+
+
+    private static String mms(SummaryStatistics s)
+    {
+        return String.format("max %.2f min %.2f stddev %.4f", s.getMax(), s.getMin(), s.getStandardDeviation());
+    }
+
+
+    int nextUnitId = 0;
+
+    final class Unit implements Comparable<Unit>
+    {
+        int unitId = nextUnitId++;
+
+        public String toString()
+        {
+            return Integer.toString(unitId);
+        }
+
+        @Override
+        public int compareTo(Unit o)
+        {
+            return Integer.compare(unitId, o.unitId);
+        }
+    }
+}

test/long/org/apache/cassandra/io/compress/CompressorPerformance.java

@@ -8,7 +8,7 @@ import java.util.concurrent.ThreadLocalRandom;
 public class CompressorPerformance
 {
 
-    static public void testPerformances() throws IOException
+    public static void testPerformances() throws IOException
     {
         for (ICompressor compressor: new ICompressor[] {
                 SnappyCompressor.instance,  // warm up
@@ -38,7 +38,7 @@ public class CompressorPerformance
     static ByteBuffer dataSource;
     static int bufLen;
 
-    static private void testPerformance(ICompressor compressor, BufferType in, BufferType out) throws IOException
+    private static void testPerformance(ICompressor compressor, BufferType in, BufferType out) throws IOException
     {
         int len = dataSource.capacity();
         int bufLen = compressor.initialCompressedBufferLength(len);

test/unit/org/apache/cassandra/db/commitlog/CommitLogTestReplayer.java

@@ -36,7 +36,7 @@ import org.apache.cassandra.io.util.FastByteArrayInputStream;
  */
 public class CommitLogTestReplayer extends CommitLogReplayer
 {
-    static public void examineCommitLog(Predicate<Mutation> processor) throws IOException
+    public static void examineCommitLog(Predicate<Mutation> processor) throws IOException
     {
         CommitLog.instance.sync(true);
 

test/unit/org/apache/cassandra/dht/BootStrapperTest.java

@@ -18,13 +18,22 @@
 */
 package org.apache.cassandra.dht;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.util.Collection;
 import java.util.HashSet;
-import java.util.Set;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
+import com.google.common.collect.Lists;
+
+import org.apache.commons.math.stat.descriptive.SummaryStatistics;
+
+import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -34,25 +43,35 @@ import org.apache.cassandra.SchemaLoader;
 import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.config.Schema;
 import org.apache.cassandra.db.Keyspace;
+import org.apache.cassandra.dht.tokenallocator.TokenAllocation;
 import org.apache.cassandra.exceptions.ConfigurationException;
 import org.apache.cassandra.gms.IFailureDetectionEventListener;
 import org.apache.cassandra.gms.IFailureDetector;
 import org.apache.cassandra.locator.TokenMetadata;
 import org.apache.cassandra.service.StorageService;
-
-import static org.junit.Assert.*;
+import org.apache.cassandra.utils.FBUtilities;
 
 @RunWith(OrderedJUnit4ClassRunner.class)
 public class BootStrapperTest
 {
+    static IPartitioner oldPartitioner;
+
     @BeforeClass
     public static void setup() throws ConfigurationException
     {
+        oldPartitioner = DatabaseDescriptor.getPartitioner();
+        DatabaseDescriptor.setPartitioner(Murmur3Partitioner.instance);
         SchemaLoader.startGossiper();
         SchemaLoader.prepareServer();
         SchemaLoader.schemaDefinition("BootStrapperTest");
     }
 
+    @AfterClass
+    public static void tearDown()
+    {
+        DatabaseDescriptor.setPartitioner(oldPartitioner);
+    }
+
     @Test
     public void testSourceTargetComputation() throws UnknownHostException
     {
@@ -112,14 +131,81 @@ public class BootStrapperTest
 
     private void generateFakeEndpoints(int numOldNodes) throws UnknownHostException
     {
-        TokenMetadata tmd = StorageService.instance.getTokenMetadata();
+        generateFakeEndpoints(StorageService.instance.getTokenMetadata(), numOldNodes, 1);
+    }
+
+    private void generateFakeEndpoints(TokenMetadata tmd, int numOldNodes, int numVNodes) throws UnknownHostException
+    {
         tmd.clearUnsafe();
         IPartitioner p = StorageService.getPartitioner();
 
         for (int i = 1; i <= numOldNodes; i++)
         {
             // leave .1 for myEndpoint
-            tmd.updateNormalToken(p.getRandomToken(), InetAddress.getByName("127.0.0." + (i + 1)));
+            InetAddress addr = InetAddress.getByName("127.0.0." + (i + 1));
+            List<Token> tokens = Lists.newArrayListWithCapacity(numVNodes);
+            for (int j = 0; j < numVNodes; ++j)
+                tokens.add(p.getRandomToken());
+            
+            tmd.updateNormalTokens(tokens, addr);
         }
     }
+    
+    @Test
+    public void testAllocateTokens() throws UnknownHostException
+    {
+        int vn = 16;
+        String ks = "BootStrapperTestKeyspace3";
+        TokenMetadata tm = new TokenMetadata();
+        generateFakeEndpoints(tm, 10, vn);
+        InetAddress addr = FBUtilities.getBroadcastAddress();
+        allocateTokensForNode(vn, ks, tm, addr);
+    }
+
+    private void allocateTokensForNode(int vn, String ks, TokenMetadata tm, InetAddress addr)
+    {
+        SummaryStatistics os = TokenAllocation.replicatedOwnershipStats(tm, Keyspace.open(ks).getReplicationStrategy(), addr);
+        Collection<Token> tokens = BootStrapper.allocateTokens(tm, addr, ks, vn);
+        assertEquals(vn, tokens.size());
+        tm.updateNormalTokens(tokens, addr);
+        SummaryStatistics ns = TokenAllocation.replicatedOwnershipStats(tm, Keyspace.open(ks).getReplicationStrategy(), addr);
+        verifyImprovement(os, ns);
+    }
+
+    private void verifyImprovement(SummaryStatistics os, SummaryStatistics ns)
+    {
+        if (ns.getStandardDeviation() > os.getStandardDeviation())
+        {
+            fail(String.format("Token allocation unexpectedly increased standard deviation.\nStats before:\n%s\nStats after:\n%s", os, ns));
+        }
+    }
+
+    
+    @Test
+    public void testAllocateTokensMultipleKeyspaces() throws UnknownHostException
+    {
+        // TODO: This scenario isn't supported very well. Investigate a multi-keyspace version of the algorithm.
+        int vn = 16;
+        String ks3 = "BootStrapperTestKeyspace4"; // RF = 3
+        String ks2 = "BootStrapperTestKeyspace5"; // RF = 2
+
+        TokenMetadata tm = new TokenMetadata();
+        generateFakeEndpoints(tm, 10, vn);
+        
+        InetAddress dcaddr = FBUtilities.getBroadcastAddress();
+        SummaryStatistics os3 = TokenAllocation.replicatedOwnershipStats(tm, Keyspace.open(ks3).getReplicationStrategy(), dcaddr);
+        SummaryStatistics os2 = TokenAllocation.replicatedOwnershipStats(tm, Keyspace.open(ks2).getReplicationStrategy(), dcaddr);
+        String cks = ks3;
+        String nks = ks2;
+        for (int i=11; i<=20; ++i)
+        {
+            allocateTokensForNode(vn, cks, tm, InetAddress.getByName("127.0.0." + (i + 1)));
+            String t = cks; cks = nks; nks = t;
+        }
+        
+        SummaryStatistics ns3 = TokenAllocation.replicatedOwnershipStats(tm, Keyspace.open(ks3).getReplicationStrategy(), dcaddr);
+        SummaryStatistics ns2 = TokenAllocation.replicatedOwnershipStats(tm, Keyspace.open(ks2).getReplicationStrategy(), dcaddr);
+        verifyImprovement(os3, ns3);
+        verifyImprovement(os2, ns2);
+    }
 }