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.bookkeeper.client;
  
  import com.google.common.math.Quantiles;
  import com.google.common.math.Quantiles.ScaleAndIndex;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Random;
  import java.util.concurrent.locks.ReadWriteLock;
  import java.util.concurrent.locks.ReentrantReadWriteLock;
  import java.util.function.Function;
  import java.util.stream.Collectors;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * DynamicWeightedRandomSelectionImpl class implements both getNextRandom
   * overloaded methods. Where getNextRandom() method considers all bookies it
   * knows of as candidates, but getNextRandom(Collection selectedNodes) method
   * considers only 'selectedNodes' as candidates.
   */
  class DynamicWeightedRandomSelectionImpl<T> implements WeightedRandomSelection<T> {
      static final Logger LOG = LoggerFactory.getLogger(DynamicWeightedRandomSelectionImpl.class);
  
      int maxProbabilityMultiplier;
      final Map<T, WeightedObject> weightMap;
      final ReadWriteLock rwLock = new ReentrantReadWriteLock(true);
      Random rand;
  
      DynamicWeightedRandomSelectionImpl() {
          this(-1);
      }
  
      DynamicWeightedRandomSelectionImpl(int maxMultiplier) {
          this.maxProbabilityMultiplier = maxMultiplier;
          this.weightMap = new HashMap<T, WeightedObject>();
          rand = new Random(System.currentTimeMillis());
      }
  
      @Override
      public void updateMap(Map<T, WeightedObject> updatedMap) {
          rwLock.writeLock().lock();
          try {
              weightMap.clear();
              weightMap.putAll(updatedMap);
          } finally {
              rwLock.writeLock().unlock();
          }
      }
  
      @Override
      public T getNextRandom() {
          rwLock.readLock().lock();
          try {
              return getNextRandom(weightMap.keySet());
          } finally {
              rwLock.readLock().unlock();
          }
      }
  
      @Override
      public T getNextRandom(Collection<T> selectedNodes) {
          rwLock.readLock().lock();
          try {
              /*
               * calculate minWeight and actual total weight.
               */
              long minWeight = Long.MAX_VALUE;
              long actTotalWeight = 0;
              for (T node : selectedNodes) {
                  long weight = 0;
                  if ((weightMap.containsKey(node))) {
                      weight = weightMap.get(node).getWeight();
                  }
                  actTotalWeight += weight;
                  if (weight > 0 && minWeight > weight) {
                      minWeight = weight;
                  }
              }
  
              long medianWeight;
              /*
               * if actTotalWeight is 0, then assign 1 to minWeight and
               * medianWeight.
               */
              if (actTotalWeight == 0) {
                  minWeight = 1L;
                  medianWeight = 1L;
              } else {
                  /*
                   * calculate medianWeight.
                   */
                  Function<? super T, ? extends Long> weightFunc = (node) -> {
                      long weight = 0;
                      if ((weightMap.containsKey(node))) {
                          weight = weightMap.get(node).getWeight();
                      }
                      return weight;
                  };
                  ArrayList<Long> weightList = selectedNodes.stream().map(weightFunc)
                          .collect(Collectors.toCollection(ArrayList::new));
                  ScaleAndIndex median = Quantiles.median();
                  medianWeight = (long) median.compute(weightList);
              }
  
              /*
               * initialize maxWeight value based on maxProbabilityMultiplier.
               */
              long maxWeight = maxProbabilityMultiplier * medianWeight;
  
              /*
               * apply weighted random selection to select an element randomly
               * based on weight.
               */
              long cumTotalWeight = 0;
              T nextRandomNode = null;
              for (T node : selectedNodes) {
                  long weight = 0;
                  if ((weightMap.containsKey(node))) {
                      weight = weightMap.get(node).getWeight();
                  }
                  if (weight <= 0) {
                      weight = minWeight;
                  } else if (maxWeight > 0 && weight > maxWeight) {
                      weight = maxWeight;
                  }
                  long tmpRandLong = rand.nextLong();
                  if (tmpRandLong == Long.MIN_VALUE) {
                      tmpRandLong++;
                  }
                  long randValue = Math.abs(tmpRandLong) % (cumTotalWeight + weight);
                  if (randValue >= cumTotalWeight) {
                      nextRandomNode = node;
                  }
                  cumTotalWeight += weight;
              }
              return nextRandomNode;
          } finally {
              rwLock.readLock().unlock();
          }
      }
  
      @Override
      public void setMaxProbabilityMultiplier(int max) {
          this.maxProbabilityMultiplier = max;
      }
  }
  