content: /*
   * 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> extends TokenAllocatorBase<Unit>
  {
      final Multimap<Unit, Token> unitToTokens;
      final int replicas;
  
      ReplicationAwareTokenAllocator(NavigableMap<Token, Unit> sortedTokens, ReplicationStrategy<Unit> strategy, IPartitioner partitioner)
      {
          super(sortedTokens, strategy, partitioner);
          unitToTokens = HashMultimap.create();
          for (Map.Entry<Token, Unit> en : sortedTokens.entrySet())
              unitToTokens.put(en.getValue(), en.getKey());
          this.replicas = strategy.replicas();
      }
  
      public int getReplicas()
      {
          return replicas;
      }
  
      public Collection<Token> addUnit(Unit newUnit, int numTokens)
      {
          assert !unitToTokens.containsKey(newUnit);
  
          if (unitCount() < replicas)
              // Allocation does not matter for now; everything replicates everywhere. However, at this point it is
              // important to start the cluster/datacenter with suitably varied token range sizes so that the algorithm
              // can maintain good balance for any number of nodes.
              return generateSplits(newUnit, numTokens);
          if (numTokens > sortedTokens.size())
              // Some of the heuristics below can't deal with this very unlikely case. Use splits for now,
              // later allocations can fix any problems this may cause.
              return generateSplits(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 splits as above.
              // 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 generateSplits(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));
              }
          }
  
          ImmutableList<Token> newTokens = ImmutableList.copyOf(unitToTokens.get(newUnit));
          TokenAllocatorDiagnostics.unitedAdded(this, numTokens, unitToTokens, sortedTokens, newTokens, newUnit);
          return newTokens;
      }
  
      /**
       * Selects tokens by repeatedly splitting the largest range in the ring at the given ratio.
       * This is used to choose tokens for the first nodes in the ring where the algorithm cannot be applied (e.g. when
       * number of nodes < RF). It generates a reasonably chaotic initial token split, after which the algorithm behaves
       * well for an unbounded number of nodes.
       */
  
      @Override
      Collection<Token> generateSplits(Unit newUnit, int numTokens)
      {
          Collection<Token> tokens = super.generateSplits(newUnit, numTokens);
          unitToTokens.putAll(newUnit, tokens);
          TokenAllocatorDiagnostics.splitsGenerated(this, numTokens, unitToTokens, sortedTokens, newUnit, tokens);
          return tokens;
      }
  
      /**
       * 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);
  
          TokenAllocatorDiagnostics.tokenInfosCreated(this, unitToTokens, 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 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);
          TokenAllocatorDiagnostics.unitRemoved(this, n, unitToTokens, sortedTokens);
      }
  
      public int unitCount()
      {
          return unitToTokens.asMap().size();
      }
  
      public String toString()
      {
          return getClass().getSimpleName();
      }
  
      /**
       * 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;
          }
      }
  }
  
  