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 io.netty.util.HashedWheelTimer;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentMap;
  import org.apache.bookkeeper.conf.ClientConfiguration;
  import org.apache.bookkeeper.feature.Feature;
  import org.apache.bookkeeper.feature.FeatureProvider;
  import org.apache.bookkeeper.net.BookieId;
  import org.apache.bookkeeper.net.BookieNode;
  import org.apache.bookkeeper.net.DNSToSwitchMapping;
  import org.apache.bookkeeper.net.NetworkTopology;
  import org.apache.bookkeeper.net.NetworkTopologyImpl;
  import org.apache.bookkeeper.net.Node;
  import org.apache.bookkeeper.net.NodeBase;
  import org.apache.bookkeeper.proto.BookieAddressResolver;
  import org.apache.bookkeeper.stats.StatsLogger;
  import org.apache.bookkeeper.util.BookKeeperConstants;
  import org.apache.commons.lang3.tuple.Pair;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * A placement policy use region information in the network topology for placing ensembles.
   *
   * @see EnsemblePlacementPolicy
   */
  public class RegionAwareEnsemblePlacementPolicy extends RackawareEnsemblePlacementPolicy {
      static final Logger LOG = LoggerFactory.getLogger(RegionAwareEnsemblePlacementPolicy.class);
  
      public static final String REPP_REGIONS_TO_WRITE = "reppRegionsToWrite";
      public static final String REPP_MINIMUM_REGIONS_FOR_DURABILITY = "reppMinimumRegionsForDurability";
      public static final String REPP_ENABLE_DURABILITY_ENFORCEMENT_IN_REPLACE =
          "reppEnableDurabilityEnforcementInReplace";
      public static final String REPP_DISABLE_DURABILITY_FEATURE_NAME = "reppDisableDurabilityFeatureName";
      public static final String REPP_DISALLOW_BOOKIE_PLACEMENT_IN_REGION_FEATURE_NAME =
          "reppDisallowBookiePlacementInRegionFeatureName";
      public static final String REPP_DISABLE_DURABILITY_ENFORCEMENT_FEATURE = "reppDisableDurabilityEnforcementFeature";
      public static final String REPP_ENABLE_VALIDATION = "reppEnableValidation";
      public static final String REGION_AWARE_ANOMALOUS_ENSEMBLE = "region_aware_anomalous_ensemble";
      static final int MINIMUM_REGIONS_FOR_DURABILITY_DEFAULT = 2;
      static final int REGIONID_DISTANCE_FROM_LEAVES = 2;
      static final String UNKNOWN_REGION = "UnknownRegion";
      static final int REMOTE_NODE_IN_REORDER_SEQUENCE = 2;
  
      protected final Map<String, TopologyAwareEnsemblePlacementPolicy> perRegionPlacement;
      protected final ConcurrentMap<BookieId, String> address2Region;
      protected FeatureProvider featureProvider;
      protected String disallowBookiePlacementInRegionFeatureName;
      protected String myRegion = null;
      protected int minRegionsForDurability = 0;
      protected boolean enableValidation = true;
      protected boolean enforceDurabilityInReplace = false;
      protected Feature disableDurabilityFeature;
      private int lastRegionIndex = 0;
  
      RegionAwareEnsemblePlacementPolicy() {
          super();
          perRegionPlacement = new HashMap<String, TopologyAwareEnsemblePlacementPolicy>();
          address2Region = new ConcurrentHashMap<BookieId, String>();
      }
  
      protected String getLocalRegion(BookieNode node) {
          if (null == node || null == node.getAddr()) {
              return UNKNOWN_REGION;
          }
          return getRegion(node.getAddr());
      }
  
      protected String getRegion(BookieId addr) {
          String region = address2Region.get(addr);
          if (null == region) {
              region = parseBookieRegion(addr);
              address2Region.putIfAbsent(addr, region);
          }
          return region;
      }
  
      protected String parseBookieRegion(BookieId addr) {
          String networkLocation = resolveNetworkLocation(addr);
          if (NetworkTopology.DEFAULT_REGION_AND_RACK.equals(networkLocation)) {
              return UNKNOWN_REGION;
          } else {
              String[] parts = networkLocation.split(NodeBase.PATH_SEPARATOR_STR);
              if (parts.length <= 1) {
                  return UNKNOWN_REGION;
              } else {
                  return parts[1];
              }
          }
      }
  
      @Override
      public void handleBookiesThatLeft(Set<BookieId> leftBookies) {
          // case 1: (In some situation, eg, Broker and bookie restart concurrently.)
          //1. Bookie X join cluster for the first time, encounters a region exception, and `address2Region` record X's
          // region as default-region.
          //2. Bookie X left cluster and is removed from knownBookies, but address2Region retains the information of
          // bookie X.
          //3. update Bookie X's rack info, and calling `onBookieRackChange` will only update address2Region for
          // addresses present in knownBookies; therefore, bookie X's region info is not updated.
          //4. Bookie X join cluster again, since address2Region contains the previous default-region information,
          // getRegion will directly use cached data, resulting of an incorrect region.
  
          // The bookie region is initialized to "default-region" in address2Region.
          // We should ensure that when a bookie leaves the cluster,
          // we also clean up the corresponding region information for that bookie in address2Region,
          // so that it can update the correct region for the bookie during onBookieRackChange and
          // handleBookiesThatJoined.
          // to avoid traffic skew in ensemble selection.
          leftBookies.forEach(address2Region::remove);
          super.handleBookiesThatLeft(leftBookies);
  
          for (TopologyAwareEnsemblePlacementPolicy policy: perRegionPlacement.values()) {
              policy.handleBookiesThatLeft(leftBookies);
          }
      }
  
      @Override
      public void handleBookiesThatJoined(Set<BookieId> joinedBookies) {
          Map<String, Set<BookieId>> perRegionClusterChange = new HashMap<String, Set<BookieId>>();
  
          // node joined
          for (BookieId addr : joinedBookies) {
              BookieNode node = createBookieNode(addr);
              topology.add(node);
              knownBookies.put(addr, node);
              historyBookies.put(addr, node);
              String region = getLocalRegion(node);
              if (null == perRegionPlacement.get(region)) {
                  perRegionPlacement.put(region, new RackawareEnsemblePlacementPolicy()
                          .initialize(dnsResolver, timer, this.reorderReadsRandom, this.stabilizePeriodSeconds,
                                  this.reorderThresholdPendingRequests, this.isWeighted, this.maxWeightMultiple,
                                  this.minNumRacksPerWriteQuorum, this.enforceMinNumRacksPerWriteQuorum,
                                  this.ignoreLocalNodeInPlacementPolicy,
                                  this.useHostnameResolveLocalNodePlacementPolicy, statsLogger, bookieAddressResolver)
                          .withDefaultRack(NetworkTopology.DEFAULT_REGION_AND_RACK));
              }
  
              Set<BookieId> regionSet = perRegionClusterChange.get(region);
              if (null == regionSet) {
                  regionSet = new HashSet<BookieId>();
                  regionSet.add(addr);
                  perRegionClusterChange.put(region, regionSet);
              } else {
                  regionSet.add(addr);
              }
  
              if (LOG.isDebugEnabled()) {
                  LOG.debug("Cluster changed : bookie {} joined the cluster.", addr);
              }
          }
  
          for (Map.Entry<String, TopologyAwareEnsemblePlacementPolicy> regionEntry : perRegionPlacement.entrySet()) {
              Set<BookieId> regionSet = perRegionClusterChange.get(regionEntry.getKey());
              if (null == regionSet) {
                  regionSet = new HashSet<BookieId>();
              }
              regionEntry.getValue().handleBookiesThatJoined(regionSet);
          }
      }
  
      @Override
      public void onBookieRackChange(List<BookieId> bookieAddressList) {
          rwLock.writeLock().lock();
          try {
              bookieAddressList.forEach(bookieAddress -> {
                  try {
                      BookieNode node = knownBookies.get(bookieAddress);
                      if (node != null) {
                          // refresh the rack info if its a known bookie
                          BookieNode newNode = createBookieNode(bookieAddress);
                          if (!newNode.getNetworkLocation().equals(node.getNetworkLocation())) {
                              topology.remove(node);
                              topology.add(newNode);
                              knownBookies.put(bookieAddress, newNode);
                              historyBookies.put(bookieAddress, newNode);
                          }
                          //Handle per region placement policy.
                          String oldRegion = getRegion(bookieAddress);
                          String newRegion = parseBookieRegion(newNode.getAddr());
                          if (oldRegion.equals(newRegion)) {
                              TopologyAwareEnsemblePlacementPolicy regionPlacement = perRegionPlacement.get(oldRegion);
                              regionPlacement.onBookieRackChange(Collections.singletonList(bookieAddress));
                          } else {
                              address2Region.put(bookieAddress, newRegion);
                              TopologyAwareEnsemblePlacementPolicy oldRegionPlacement = perRegionPlacement.get(oldRegion);
                              oldRegionPlacement.handleBookiesThatLeft(Collections.singleton(bookieAddress));
                              TopologyAwareEnsemblePlacementPolicy newRegionPlacement = perRegionPlacement.get(
                                      newRegion);
                              if (newRegionPlacement == null) {
                                  newRegionPlacement = new RackawareEnsemblePlacementPolicy()
                                          .initialize(dnsResolver, timer, this.reorderReadsRandom,
                                                  this.stabilizePeriodSeconds, this.reorderThresholdPendingRequests,
                                                  this.isWeighted, this.maxWeightMultiple,
                                                  this.minNumRacksPerWriteQuorum, this.enforceMinNumRacksPerWriteQuorum,
                                                  this.ignoreLocalNodeInPlacementPolicy,
                                                  this.useHostnameResolveLocalNodePlacementPolicy, statsLogger,
                                                  bookieAddressResolver)
                                          .withDefaultRack(NetworkTopology.DEFAULT_REGION_AND_RACK);
                                  perRegionPlacement.put(newRegion, newRegionPlacement);
                              }
                              newRegionPlacement.handleBookiesThatJoined(Collections.singleton(bookieAddress));
                          }
                      }
                  } catch (IllegalArgumentException | NetworkTopologyImpl.InvalidTopologyException e) {
                      LOG.error("Failed to update bookie rack info: {} ", bookieAddress, e);
                  }
              });
          } finally {
              rwLock.writeLock().unlock();
          }
      }
  
      @Override
      public RegionAwareEnsemblePlacementPolicy initialize(ClientConfiguration conf,
                                                           Optional<DNSToSwitchMapping> optionalDnsResolver,
                                                           HashedWheelTimer timer,
                                                           FeatureProvider featureProvider,
                                                           StatsLogger statsLogger,
                                                           BookieAddressResolver bookieAddressResolver) {
          super.initialize(conf, optionalDnsResolver, timer, featureProvider, statsLogger, bookieAddressResolver)
                  .withDefaultRack(NetworkTopology.DEFAULT_REGION_AND_RACK);
          myRegion = getLocalRegion(localNode);
          enableValidation = conf.getBoolean(REPP_ENABLE_VALIDATION, true);
          // We have to statically provide regions we want the writes to go through and how many regions
          // are required for durability. This decision cannot be driven by the active bookies as the
          // current topology will not be indicative of constraints that must be enforced for durability
          String regionsString = conf.getString(REPP_REGIONS_TO_WRITE, null);
          if (null != regionsString) {
              // Regions are specified as
              // R1;R2;...
              String[] regions = regionsString.split(";");
              for (String region : regions) {
                  perRegionPlacement.put(region, new RackawareEnsemblePlacementPolicy(true)
                          .initialize(dnsResolver, timer, this.reorderReadsRandom, this.stabilizePeriodSeconds,
                                  this.reorderThresholdPendingRequests, this.isWeighted, this.maxWeightMultiple,
                                  this.minNumRacksPerWriteQuorum, this.enforceMinNumRacksPerWriteQuorum,
                                  this.ignoreLocalNodeInPlacementPolicy, this.useHostnameResolveLocalNodePlacementPolicy,
                                  statsLogger, bookieAddressResolver)
                          .withDefaultRack(NetworkTopology.DEFAULT_REGION_AND_RACK));
              }
              minRegionsForDurability = conf.getInt(REPP_MINIMUM_REGIONS_FOR_DURABILITY,
                      MINIMUM_REGIONS_FOR_DURABILITY_DEFAULT);
              if (minRegionsForDurability > 0) {
                  enforceDurability = true;
                  enforceDurabilityInReplace = conf.getBoolean(REPP_ENABLE_DURABILITY_ENFORCEMENT_IN_REPLACE, true);
              }
              if (regions.length < minRegionsForDurability) {
                  throw new IllegalArgumentException(
                          "Regions provided are insufficient to meet the durability constraints");
              }
          }
          this.featureProvider = featureProvider;
          this.disallowBookiePlacementInRegionFeatureName =
              conf.getString(REPP_DISALLOW_BOOKIE_PLACEMENT_IN_REGION_FEATURE_NAME);
          this.disableDurabilityFeature = conf.getFeature(REPP_DISABLE_DURABILITY_ENFORCEMENT_FEATURE, null);
          if (null == disableDurabilityFeature) {
              this.disableDurabilityFeature =
                      featureProvider.getFeature(
                          conf.getString(REPP_DISABLE_DURABILITY_FEATURE_NAME,
                                  BookKeeperConstants.FEATURE_REPP_DISABLE_DURABILITY_ENFORCEMENT));
          }
          return this;
      }
  
      protected List<BookieNode> selectRandomFromRegions(Set<String> availableRegions,
                                              int numBookies,
                                              Set<Node> excludeBookies,
                                              Predicate<BookieNode> predicate,
                                              Ensemble<BookieNode> ensemble)
          throws BKException.BKNotEnoughBookiesException {
          List<BookieNode> availableBookies = new ArrayList<BookieNode>();
          for (BookieNode bookieNode: knownBookies.values()) {
              if (availableRegions.contains(getLocalRegion(bookieNode))) {
                  availableBookies.add(bookieNode);
              }
          }
  
          return selectRandomInternal(availableBookies,  numBookies, excludeBookies, predicate, ensemble);
      }
  
  
      @Override
      public PlacementResult<List<BookieId>> newEnsemble(int ensembleSize, int writeQuorumSize,
              int ackQuorumSize, Map<String, byte[]> customMetadata, Set<BookieId> excludedBookies)
              throws BKException.BKNotEnoughBookiesException {
  
          int effectiveMinRegionsForDurability = disableDurabilityFeature.isAvailable() ? 1 : minRegionsForDurability;
  
          // All of these conditions indicate bad configuration
          if (ackQuorumSize < effectiveMinRegionsForDurability) {
              throw new IllegalArgumentException(
                      "Ack Quorum size provided are insufficient to meet the durability constraints");
          } else if (ensembleSize < writeQuorumSize) {
              throw new IllegalArgumentException(
                      "write quorum (" + writeQuorumSize + ") cannot exceed ensemble size (" + ensembleSize + ")");
          } else if (writeQuorumSize < ackQuorumSize) {
              throw new IllegalArgumentException(
                      "ack quorum (" + ackQuorumSize + ") cannot exceed write quorum size (" + writeQuorumSize + ")");
          } else if (effectiveMinRegionsForDurability > 0) {
              // We must survive the failure of numRegions - effectiveMinRegionsForDurability. When these
              // regions have failed we would spread the replicas over the remaining
              // effectiveMinRegionsForDurability regions; we have to make sure that the ack quorum is large
              // enough such that there is a configuration for spreading the replicas across
              // effectiveMinRegionsForDurability - 1 regions
              if (ackQuorumSize <= (writeQuorumSize - (writeQuorumSize / effectiveMinRegionsForDurability))) {
                  throw new IllegalArgumentException("ack quorum (" + ackQuorumSize + ") "
                      + "violates the requirement to satisfy durability constraints when running in degraded mode");
              }
          }
  
          rwLock.readLock().lock();
          try {
              Set<BookieId> comprehensiveExclusionBookiesSet = addDefaultRackBookiesIfMinNumRacksIsEnforced(
                      excludedBookies);
              Set<Node> excludeNodes = convertBookiesToNodes(comprehensiveExclusionBookiesSet);
              List<String> availableRegions = new ArrayList<>();
              for (Map.Entry<String, TopologyAwareEnsemblePlacementPolicy> entry : perRegionPlacement.entrySet()) {
                  String region = entry.getKey();
                  if (!entry.getValue().knownBookies.isEmpty() && (null == disallowBookiePlacementInRegionFeatureName
                          || !featureProvider.scope(region).getFeature(disallowBookiePlacementInRegionFeatureName)
                              .isAvailable())) {
                      availableRegions.add(region);
                  }
              }
              int numRegionsAvailable = availableRegions.size();
  
              // If we were unable to get region information or all regions are disallowed which is
              // an invalid configuration; default to random selection from the set of nodes
              if (numRegionsAvailable < 1) {
                  // We cant disallow all regions; if we did, raise an alert to draw attention
                  if (perRegionPlacement.keySet().size() >= 1) {
                      LOG.error("No regions available, invalid configuration");
                  }
                  List<BookieNode> bns = selectRandom(ensembleSize, excludeNodes, TruePredicate.INSTANCE,
                      EnsembleForReplacementWithNoConstraints.INSTANCE);
                  ArrayList<BookieId> addrs = new ArrayList<BookieId>(ensembleSize);
                  for (BookieNode bn : bns) {
                      addrs.add(bn.getAddr());
                  }
                  return PlacementResult.of(addrs,
                                            isEnsembleAdheringToPlacementPolicy(
                                                    addrs, writeQuorumSize, ackQuorumSize));
              }
  
              // Single region, fall back to RackAwareEnsemblePlacement
              if (numRegionsAvailable < 2) {
                  RRTopologyAwareCoverageEnsemble ensemble = new RRTopologyAwareCoverageEnsemble(ensembleSize,
                          writeQuorumSize, ackQuorumSize, REGIONID_DISTANCE_FROM_LEAVES,
                          effectiveMinRegionsForDurability > 0 ? new HashSet<>(perRegionPlacement.keySet()) : null,
                          effectiveMinRegionsForDurability, minNumRacksPerWriteQuorum);
                  TopologyAwareEnsemblePlacementPolicy nextPolicy = perRegionPlacement.get(
                          availableRegions.iterator().next());
                  return nextPolicy.newEnsemble(ensembleSize, writeQuorumSize, writeQuorumSize,
                          comprehensiveExclusionBookiesSet, ensemble, ensemble);
              }
  
              int remainingEnsemble = ensembleSize;
              int remainingWriteQuorum = writeQuorumSize;
  
              // Equally distribute the nodes across all regions to whatever extent possible
              // with the hierarchy in mind
              // Try and place as many nodes in a region as possible, the ones that cannot be
              // accommodated are placed on other regions
              // Within each region try and follow rack aware placement
              Map<String, Pair<Integer, Integer>> regionsWiseAllocation = new HashMap<>();
              for (String region: availableRegions) {
                  regionsWiseAllocation.put(region, Pair.of(0, 0));
              }
              int remainingEnsembleBeforeIteration;
              int numRemainingRegions;
              Set<String> regionsReachedMaxAllocation = new HashSet<String>();
              RRTopologyAwareCoverageEnsemble ensemble;
              do {
                  numRemainingRegions = numRegionsAvailable - regionsReachedMaxAllocation.size();
                  ensemble = new RRTopologyAwareCoverageEnsemble(ensembleSize, writeQuorumSize, ackQuorumSize,
                          REGIONID_DISTANCE_FROM_LEAVES,
                          // We pass all regions we know off to the coverage ensemble as
                          // regardless of regions that are available; constraints are
                          // always applied based on all possible regions
                          effectiveMinRegionsForDurability > 0 ? new HashSet<>(perRegionPlacement.keySet()) : null,
                          effectiveMinRegionsForDurability, minNumRacksPerWriteQuorum);
                  remainingEnsembleBeforeIteration = remainingEnsemble;
                  int regionsToAllocate = numRemainingRegions;
                  int startRegionIndex = lastRegionIndex % numRegionsAvailable;
                  int localRegionIndex = -1;
                  if (myRegion != null && !UNKNOWN_REGION.equals(myRegion)) {
                      localRegionIndex = availableRegions.indexOf(myRegion);
                  }
                  String region = myRegion;
                  for (int i = 0; i < numRegionsAvailable; ++i) {
                      // select the local region first, and for the rest region select, use round-robin selection.
                      if (i > 0 || localRegionIndex == -1) {
                          if (startRegionIndex % numRegionsAvailable == localRegionIndex) {
                              startRegionIndex++;
                          }
                          region = availableRegions.get(startRegionIndex % numRegionsAvailable);
                          startRegionIndex++;
                      }
                      final Pair<Integer, Integer> currentAllocation = regionsWiseAllocation.get(region);
                      TopologyAwareEnsemblePlacementPolicy policyWithinRegion = perRegionPlacement.get(region);
                      if (!regionsReachedMaxAllocation.contains(region)) {
                          if (numRemainingRegions <= 0) {
                              LOG.error("Inconsistent State: This should never happen");
                              throw new BKException.BKNotEnoughBookiesException();
                          }
                          // try to place the bookies as balance as possible across all the regions
                          int addToEnsembleSize = Math.min(remainingEnsemble, remainingEnsemble / regionsToAllocate
                                  + (remainingEnsemble % regionsToAllocate == 0 ? 0 : 1));
                          boolean success = false;
                          while (addToEnsembleSize > 0) {
                              int addToWriteQuorum = Math.max(1, Math.min(remainingWriteQuorum,
                                          Math.round(1.0f * writeQuorumSize * addToEnsembleSize / ensembleSize)));
                              // Temp ensemble will be merged back into the ensemble only if we are able to successfully
                              // allocate the target number of bookies in this region; if we fail because we dont have
                              // enough bookies; then we retry the process with a smaller target
                              RRTopologyAwareCoverageEnsemble tempEnsemble =
                                  new RRTopologyAwareCoverageEnsemble(ensemble);
                              int newEnsembleSize = currentAllocation.getLeft() + addToEnsembleSize;
                              int newWriteQuorumSize = currentAllocation.getRight() + addToWriteQuorum;
                              try {
                                  List<BookieId> allocated = policyWithinRegion
                                          .newEnsemble(newEnsembleSize, newWriteQuorumSize, newWriteQuorumSize,
                                                  comprehensiveExclusionBookiesSet, tempEnsemble, tempEnsemble)
                                          .getResult();
                                  ensemble = tempEnsemble;
                                  remainingEnsemble -= addToEnsembleSize;
                                  remainingWriteQuorum -= addToWriteQuorum;
                                  regionsWiseAllocation.put(region, Pair.of(newEnsembleSize, newWriteQuorumSize));
                                  success = true;
                                  regionsToAllocate--;
                                  lastRegionIndex = startRegionIndex;
                                  LOG.info("Region {} allocating bookies with ensemble size {} "
                                          + "and write quorum size {} : {}",
                                          region, newEnsembleSize, newWriteQuorumSize, allocated);
                                  break;
                              } catch (BKException.BKNotEnoughBookiesException exc) {
                                  LOG.warn("Could not allocate {} bookies in region {}, try allocating {} bookies",
                                          newEnsembleSize, region, (newEnsembleSize - 1));
                                  addToEnsembleSize--;
                              }
                          }
  
                          // we couldn't allocate additional bookies from the region,
                          // it should have reached its max allocation.
                          if (!success) {
                              regionsReachedMaxAllocation.add(region);
                          }
                      }
  
                      if (regionsReachedMaxAllocation.contains(region)) {
                          if (currentAllocation.getLeft() > 0) {
                              LOG.info("Allocating {} bookies in region {} : ensemble {} exclude {}",
                                      currentAllocation.getLeft(), region, comprehensiveExclusionBookiesSet, ensemble);
                              policyWithinRegion.newEnsemble(
                                      currentAllocation.getLeft(),
                                      currentAllocation.getRight(),
                                      currentAllocation.getRight(),
                                      comprehensiveExclusionBookiesSet,
                                      ensemble,
                                      ensemble);
                              LOG.info("Allocated {} bookies in region {} : {}",
                                      currentAllocation.getLeft(), region, ensemble);
                          }
                      }
                  }
  
                  if (regionsReachedMaxAllocation.containsAll(regionsWiseAllocation.keySet())) {
                      break;
                  }
              } while ((remainingEnsemble > 0) && (remainingEnsemble < remainingEnsembleBeforeIteration));
  
              List<BookieId> bookieList = ensemble.toList();
              if (ensembleSize != bookieList.size()) {
                  LOG.error("Not enough {} bookies are available to form an ensemble : {}.",
                            ensembleSize, bookieList);
                  throw new BKException.BKNotEnoughBookiesException();
              }
  
              if (enableValidation && !ensemble.validate()) {
                  LOG.error("Not enough {} bookies are available to form a valid ensemble : {}.",
                      ensembleSize, bookieList);
                  throw new BKException.BKNotEnoughBookiesException();
              }
              LOG.info("Bookies allocated successfully {}", ensemble);
              List<BookieId> ensembleList = ensemble.toList();
              return PlacementResult.of(ensembleList,
                      isEnsembleAdheringToPlacementPolicy(ensembleList, writeQuorumSize, ackQuorumSize));
          } finally {
              rwLock.readLock().unlock();
          }
      }
  
      @Override
      public PlacementResult<BookieId> replaceBookie(int ensembleSize, int writeQuorumSize, int ackQuorumSize,
              Map<String, byte[]> customMetadata, List<BookieId> currentEnsemble,
              BookieId bookieToReplace, Set<BookieId> excludeBookies)
              throws BKException.BKNotEnoughBookiesException {
          rwLock.readLock().lock();
          try {
              boolean enforceDurability = enforceDurabilityInReplace && !disableDurabilityFeature.isAvailable();
              int effectiveMinRegionsForDurability = enforceDurability ? minRegionsForDurability : 1;
              Set<BookieId> comprehensiveExclusionBookiesSet = addDefaultRackBookiesIfMinNumRacksIsEnforced(
                      excludeBookies);
              Set<Node> excludeNodes = convertBookiesToNodes(comprehensiveExclusionBookiesSet);
              RRTopologyAwareCoverageEnsemble ensemble = new RRTopologyAwareCoverageEnsemble(ensembleSize,
                  writeQuorumSize,
                  ackQuorumSize,
                  REGIONID_DISTANCE_FROM_LEAVES,
                  effectiveMinRegionsForDurability > 0 ? new HashSet<String>(perRegionPlacement.keySet()) : null,
                  effectiveMinRegionsForDurability, minNumRacksPerWriteQuorum);
  
              BookieNode bookieNodeToReplace = knownBookies.get(bookieToReplace);
              if (null == bookieNodeToReplace) {
                  bookieNodeToReplace = createBookieNode(bookieToReplace);
              }
              excludeNodes.add(bookieNodeToReplace);
  
              for (BookieId bookieAddress: currentEnsemble) {
                  if (bookieAddress.equals(bookieToReplace)) {
                      continue;
                  }
  
                  BookieNode bn = knownBookies.get(bookieAddress);
                  if (null == bn) {
                      bn = createBookieNode(bookieAddress);
                  }
  
                  excludeNodes.add(bn);
  
                  if (!ensemble.apply(bn, ensemble)) {
                      LOG.warn("Anomalous ensemble detected");
                      if (null != statsLogger) {
                          statsLogger.getCounter(REGION_AWARE_ANOMALOUS_ENSEMBLE).inc();
                      }
                      enforceDurability = false;
                  }
  
                  ensemble.addNode(bn);
              }
  
              if (LOG.isDebugEnabled()) {
                  LOG.debug("Try to choose a new bookie to replace {}, excluding {}.", bookieToReplace,
                      excludeNodes);
              }
              // pick a candidate from same rack to replace
              BookieNode candidate = replaceFromRack(bookieNodeToReplace, excludeNodes,
                  ensemble, ensemble, enforceDurability);
              if (LOG.isDebugEnabled()) {
                  LOG.debug("Bookie {} is chosen to replace bookie {}.", candidate, bookieNodeToReplace);
              }
              BookieId candidateAddr = candidate.getAddr();
              List<BookieId> newEnsemble = new ArrayList<BookieId>(currentEnsemble);
              if (currentEnsemble.isEmpty()) {
                  /*
                   * in testing code there are test cases which would pass empty
                   * currentEnsemble
                   */
                  newEnsemble.add(candidateAddr);
              } else {
                  newEnsemble.set(currentEnsemble.indexOf(bookieToReplace), candidateAddr);
              }
              return PlacementResult.of(candidateAddr,
                      isEnsembleAdheringToPlacementPolicy(newEnsemble, writeQuorumSize, ackQuorumSize));
          } finally {
              rwLock.readLock().unlock();
          }
      }
  
      protected BookieNode replaceFromRack(BookieNode bookieNodeToReplace,
                                           Set<Node> excludeBookies,
                                           Predicate<BookieNode> predicate,
                                           Ensemble<BookieNode> ensemble,
                                           boolean enforceDurability)
          throws BKException.BKNotEnoughBookiesException {
          Set<String> availableRegions = new HashSet<String>();
          for (String region: perRegionPlacement.keySet()) {
              if ((null == disallowBookiePlacementInRegionFeatureName)
                      || !featureProvider.scope(region).getFeature(disallowBookiePlacementInRegionFeatureName)
                          .isAvailable()) {
                  availableRegions.add(region);
              }
          }
          String regionForBookieToReplace = getLocalRegion(bookieNodeToReplace);
          if (availableRegions.contains(regionForBookieToReplace)) {
              TopologyAwareEnsemblePlacementPolicy regionPolicy = perRegionPlacement.get(regionForBookieToReplace);
              if (null != regionPolicy) {
                  try {
                      // select one from local rack => it falls back to selecting a node from the region
                      // if the rack does not have an available node, selecting from the same region
                      // should not violate durability constraints so we can simply not have to check
                      // for that.
                      return regionPolicy.selectFromNetworkLocation(
                          bookieNodeToReplace.getNetworkLocation(),
                          excludeBookies,
                          TruePredicate.INSTANCE,
                          EnsembleForReplacementWithNoConstraints.INSTANCE,
                          true);
                  } catch (BKException.BKNotEnoughBookiesException e) {
                      LOG.warn("Failed to choose a bookie from {} : "
                              + "excluded {}, fallback to choose bookie randomly from the cluster.",
                          bookieNodeToReplace.getNetworkLocation(), excludeBookies);
                  }
              }
          }
  
          // randomly choose one from all the regions that are available, ignore the provided predicate if we are not
          // enforcing durability.
          return selectRandomFromRegions(availableRegions, 1,
              excludeBookies,
              enforceDurability ? predicate : TruePredicate.INSTANCE,
              enforceDurability ? ensemble : EnsembleForReplacementWithNoConstraints.INSTANCE).get(0);
      }
  
      @Override
      public final DistributionSchedule.WriteSet reorderReadSequence(
              List<BookieId> ensemble,
              BookiesHealthInfo bookiesHealthInfo,
              DistributionSchedule.WriteSet writeSet) {
          if (UNKNOWN_REGION.equals(myRegion)) {
              return super.reorderReadSequence(ensemble, bookiesHealthInfo, writeSet);
          } else {
              Map<Integer, String> writeSetWithRegion = new HashMap<>();
              for (int i = 0; i < writeSet.size(); i++) {
                  int idx = writeSet.get(i);
                  writeSetWithRegion.put(idx, getRegion(ensemble.get(idx)));
              }
              return super.reorderReadSequenceWithRegion(ensemble, writeSet, writeSetWithRegion,
                  bookiesHealthInfo, true, myRegion, REMOTE_NODE_IN_REORDER_SEQUENCE);
          }
      }
  
      @Override
      public final DistributionSchedule.WriteSet reorderReadLACSequence(
              List<BookieId> ensemble,
              BookiesHealthInfo bookiesHealthInfo,
              DistributionSchedule.WriteSet writeSet) {
          if (UNKNOWN_REGION.equals(myRegion)) {
              return super.reorderReadLACSequence(ensemble, bookiesHealthInfo, writeSet);
          }
          DistributionSchedule.WriteSet finalList = reorderReadSequence(ensemble, bookiesHealthInfo, writeSet);
          finalList.addMissingIndices(ensemble.size());
          return finalList;
      }
  
      @Override
      public PlacementPolicyAdherence isEnsembleAdheringToPlacementPolicy(List<BookieId> ensembleList,
              int writeQuorumSize, int ackQuorumSize) {
          /**
           * TODO: have to implement actual logic for this method for
           * RegionAwareEnsemblePlacementPolicy. For now return true value.
           *
           * - https://github.com/apache/bookkeeper/issues/1898
           */
          return PlacementPolicyAdherence.MEETS_STRICT;
      }
  
      @Override
      public void updateBookieInfo(Map<BookieId, BookieInfoReader.BookieInfo> bookieInfoMap) {
          super.updateBookieInfo(bookieInfoMap);
          for (TopologyAwareEnsemblePlacementPolicy policy : perRegionPlacement.values()) {
              policy.updateBookieInfo(bookieInfoMap);
          }
      }
  }
  