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.kafka.metadata.placement;
  
  import org.apache.kafka.common.errors.InvalidReplicationFactorException;
  import org.apache.kafka.metadata.OptionalStringComparator;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Random;
  
  
  /**
   * The striped replica placer.
   * <p>
   * <h3>Goals</h3>
   * <p>The design of this placer attempts to satisfy a few competing goals. Firstly, we want
   * to spread the replicas as evenly as we can across racks. In the simple case where
   * broker racks have not been configured, this goal is a no-op, of course. But it is the
   * highest priority goal in multi-rack clusters.
   *
   * <p>Our second goal is to spread the replicas evenly across brokers. Since we are placing
   * multiple partitions, we try to avoid putting each partition on the same set of
   * replicas, even if it does satisfy the rack placement goal. If any specific broker is
   * fenced, we would like the new leaders to distributed evenly across the remaining
   * brokers.
   *
   * <p>However, we treat the rack placement goal as higher priority than this goal-- if you
   * configure 10 brokers in rack A and B, and 1 broker in rack C, you will end up with a
   * lot of partitions on that one broker in rack C.  If you were to place a lot of
   * partitions with replication factor 3, each partition would try to get a replica there.
   * In general racks are supposed to be about the same size -- if they aren't, this is a
   * user error.
   *
   * <p>Finally, we would prefer to place replicas on unfenced brokers, rather than on fenced
   * brokers.
   * <p>
   * <h3>Constraints</h3>
   * In addition to these goals, we have two constraints. Unlike the goals, these are not
   * optional -- they are mandatory. Placement will fail if a constraint cannot be
   * satisfied. The first constraint is that we can't place more than one replica on the
   * same broker. This imposes an upper limit on replication factor-- for example, a 3-node
   * cluster can't have any topics with replication factor 4. This constraint comes from
   * Kafka's internal design.
   *
   * <p>The second constraint is that the leader of each partition must be an unfenced broker.
   * This constraint is a bit arbitrary. In theory, we could allow people to create
   * new topics even if every broker were fenced. However, this would be confusing for
   * users.
   * <p>
   * <h3>Algorithm</h3>
   * <p>The StripedReplicaPlacer constructor loads the broker data into rack objects. Each
   * rack object contains a sorted list of fenced brokers, and a separate sorted list of
   * unfenced brokers. The racks themselves are organized into a sorted list, stored inside
   * the top-level RackList object.
   *
   * <p>The general idea is that we place replicas on to racks in a round-robin fashion. So if
   * we had racks A, B, C, and D, and we were creating a new partition with replication
   * factor 3, our first replica might come from A, our second from B, and our third from C.
   * Of course our placement would not be very fair if we always started with rack A.
   * Therefore, we generate a random starting offset when the RackList is created. So one
   * time we might go B, C, D. Another time we might go C, D, A. And so forth.
   *
   * <p>Note that each partition we generate advances the starting offset by one.
   * So in our 4-rack cluster, with 3 partitions, we might choose these racks:
   * <pre>
   * partition 1: A, B, C
   * partition 2: B, C, A
   * partition 3: C, A, B
   * </pre>
   * This is what generates the characteristic "striped" pattern of this placer.
   *
   * <p>So far I haven't said anything about how we choose a replica from within a rack.  In
   * fact, this is also done in a round-robin fashion. So if rack A had replica A0, A1, A2,
   * and A3, we might return A0 the first time, A1, the second, A2 the third, and so on.
   * Just like with the racks, we add a random starting offset to mix things up a bit.
   *
   * <p>So let's say you had a cluster with racks A, B, and C, and each rack had 3 replicas,
   * for 9 nodes in total.
   * If all the offsets were 0, you'd get placements like this:
   * <pre>
   * partition 1: A0, B0, C0
   * partition 2: B1, C1, A1
   * partition 3: C2, A2, B2
   * </pre>
   * <p>One additional complication with choosing a replica within a rack is that we want to
   * choose the unfenced replicas first. In a big cluster with lots of nodes available,
   * we'd prefer not to place a new partition on a node that is fenced. Therefore, we
   * actually maintain two lists, rather than the single list I described above.
   * We only start using the fenced node list when the unfenced node list is totally
   * exhausted.
   *
   * <p>Furthermore, we cannot place the first replica (the leader) of a new partition on a
   * fenced replica. Therefore, we have some special logic to ensure that this doesn't
   * happen.
   */
  public class StripedReplicaPlacer implements ReplicaPlacer {
      /**
       * A list of brokers that we can iterate through.
       */
      static class BrokerList {
          static final BrokerList EMPTY = new BrokerList();
          private final List<Integer> brokers = new ArrayList<>(0);
  
          /**
           * How many brokers we have retrieved from the list during the current iteration epoch.
           */
          private int index = 0;
  
          /**
           * The offset to add to the index in order to calculate the list entry to fetch.  The
           * addition is done modulo the list size.
           */
          private int offset = 0;
  
          /**
           * The last known iteration epoch. If we call next with a different epoch than this, the
           * index and offset will be reset.
           */
          private int epoch = 0;
  
          BrokerList add(int broker) {
              this.brokers.add(broker);
              return this;
          }
  
          /**
           * Initialize this broker list by sorting it and randomizing the start offset.
           *
           * @param random    The random number generator.
           */
          void initialize(Random random) {
              if (!brokers.isEmpty()) {
                  brokers.sort(Integer::compareTo);
                  this.offset = random.nextInt(brokers.size());
              }
          }
  
          /**
           * Randomly shuffle the brokers in this list.
           */
          void shuffle(Random random) {
              Collections.shuffle(brokers, random);
          }
  
          /**
           * @return          The number of brokers in this list.
           */
          int size() {
              return brokers.size();
          }
  
          /**
           * Get the next broker in this list, or -1 if there are no more elements to be
           * returned.
           *
           * @param epoch     The current iteration epoch.
           *
           * @return          The broker ID, or -1 if there are no more brokers to be
           *                  returned in this epoch.
           */
          int next(int epoch) {
              if (brokers.isEmpty()) return -1;
              if (this.epoch != epoch) {
                  this.epoch = epoch;
                  this.index = 0;
                  this.offset = (offset + 1) % brokers.size();
              }
              if (index >= brokers.size()) return -1;
              int broker = brokers.get((index + offset) % brokers.size());
              index++;
              return broker;
          }
      }
  
      /**
       * A rack in the cluster, which contains brokers.
       */
      static class Rack {
          private final BrokerList fenced = new BrokerList();
          private final BrokerList unfenced = new BrokerList();
  
          /**
           * Initialize this rack.
           *
           * @param random    The random number generator.
           */
          void initialize(Random random) {
              fenced.initialize(random);
              unfenced.initialize(random);
          }
  
          void shuffle(Random random) {
              fenced.shuffle(random);
              unfenced.shuffle(random);
          }
  
          BrokerList fenced() {
              return fenced;
          }
  
          BrokerList unfenced() {
              return unfenced;
          }
  
          /**
           * Get the next unfenced broker in this rack, or -1 if there are no more brokers
           * to be returned.
           *
           * @param epoch     The current iteration epoch.
           *
           * @return          The broker ID, or -1 if there are no more brokers to be
           *                  returned in this epoch.
           */
          int nextUnfenced(int epoch) {
              return unfenced.next(epoch);
          }
  
          /**
           * Get the next broker in this rack, or -1 if there are no more brokers to be
           * returned.
           *
           * @param epoch     The current iteration epoch.
           *
           * @return          The broker ID, or -1 if there are no more brokers to be
           *                  returned in this epoch.
           */
          int next(int epoch) {
              int result = unfenced.next(epoch);
              if (result >= 0) return result;
              return fenced.next(epoch);
          }
      }
  
      /**
       * A list of racks that we can iterate through.
       */
      static class RackList {
          /**
           * The random number generator.
           */
          private final Random random;
  
          /**
           * A map from rack names to the brokers contained within them.
           */
          private final Map<Optional<String>, Rack> racks = new HashMap<>();
  
          /**
           * The names of all the racks in the cluster.
           * Racks which have at least one unfenced broker come first (in sorted order),
           * followed by racks which have only fenced brokers (also in sorted order).
           */
          private final List<Optional<String>> rackNames = new ArrayList<>();
  
          /**
           * The total number of brokers in the cluster, both fenced and unfenced.
           */
          private final int numTotalBrokers;
  
          /**
           * The total number of unfenced brokers in the cluster.
           */
          private final int numUnfencedBrokers;
  
          /**
           * The iteration epoch.
           */
          private int epoch = 0;
  
          /**
           * The offset we use to determine which rack is returned first.
           */
          private int offset;
  
          RackList(Random random, Iterator<UsableBroker> iterator) {
              this.random = random;
              int numTotalBrokersCount = 0, numUnfencedBrokersCount = 0;
              while (iterator.hasNext()) {
                  UsableBroker broker = iterator.next();
                  Rack rack = racks.get(broker.rack());
                  if (rack == null) {
                      rackNames.add(broker.rack());
                      rack = new Rack();
                      racks.put(broker.rack(), rack);
                  }
                  if (broker.fenced()) {
                      rack.fenced().add(broker.id());
                  } else {
                      numUnfencedBrokersCount++;
                      rack.unfenced().add(broker.id());
                  }
                  numTotalBrokersCount++;
              }
              for (Rack rack : racks.values()) {
                  rack.initialize(random);
              }
              this.rackNames.sort(OptionalStringComparator.INSTANCE);
              this.numTotalBrokers = numTotalBrokersCount;
              this.numUnfencedBrokers = numUnfencedBrokersCount;
              this.offset = rackNames.isEmpty() ? 0 : random.nextInt(rackNames.size());
          }
  
          int numTotalBrokers() {
              return numTotalBrokers;
          }
  
          int numUnfencedBrokers() {
              return numUnfencedBrokers;
          }
  
          // VisibleForTesting
          List<Optional<String>> rackNames() {
              return rackNames;
          }
  
          List<Integer> place(int replicationFactor) {
              throwInvalidReplicationFactorIfNonPositive(replicationFactor);
              throwInvalidReplicationFactorIfTooFewBrokers(replicationFactor, numTotalBrokers());
              throwInvalidReplicationFactorIfZero(numUnfencedBrokers());
              // If we have returned as many assignments as there are unfenced brokers in
              // the cluster, shuffle the rack list and broker lists to try to avoid
              // repeating the same assignments again.
              // But don't reset the iteration epoch for a single unfenced broker -- otherwise we would loop forever
              if (epoch == numUnfencedBrokers && numUnfencedBrokers > 1) {
                  shuffle();
                  epoch = 0;
              }
              if (offset == rackNames.size()) {
                  offset = 0;
              }
              List<Integer> brokers = new ArrayList<>(replicationFactor);
              int firstRackIndex = offset;
              while (true) {
                  Optional<String> name = rackNames.get(firstRackIndex);
                  Rack rack = racks.get(name);
                  int result = rack.nextUnfenced(epoch);
                  if (result >= 0) {
                      brokers.add(result);
                      break;
                  }
                  firstRackIndex++;
                  if (firstRackIndex == rackNames.size()) {
                      firstRackIndex = 0;
                  }
              }
              int rackIndex = offset;
              for (int replica = 1; replica < replicationFactor; replica++) {
                  int result = -1;
                  do {
                      if (rackIndex == firstRackIndex) {
                          firstRackIndex = -1;
                      } else {
                          Optional<String> rackName = rackNames.get(rackIndex);
                          Rack rack = racks.get(rackName);
                          result = rack.next(epoch);
                      }
                      rackIndex++;
                      if (rackIndex == rackNames.size()) {
                          rackIndex = 0;
                      }
                  } while (result < 0);
                  brokers.add(result);
              }
              epoch++;
              offset++;
              return brokers;
          }
  
          void shuffle() {
              Collections.shuffle(rackNames, random);
              for (Rack rack : racks.values()) {
                  rack.shuffle(random);
              }
          }
      }
  
      private static void throwInvalidReplicationFactorIfNonPositive(int replicationFactor) {
          if (replicationFactor <= 0) {
              throw new InvalidReplicationFactorException("Invalid replication factor " +
                      replicationFactor + ": the replication factor must be positive.");
          }
      }
  
      private static void throwInvalidReplicationFactorIfZero(int numUnfenced) {
          if (numUnfenced == 0) {
              throw new InvalidReplicationFactorException("All brokers are currently fenced.");
          }
      }
  
      private static void throwInvalidReplicationFactorIfTooFewBrokers(int replicationFactor, int numTotalBrokers) {
          if (replicationFactor > numTotalBrokers) {
              throw new InvalidReplicationFactorException("The target replication factor " +
                      "of " + replicationFactor + " cannot be reached because only " +
                      numTotalBrokers + " broker(s) are registered.");
          }
      }
  
      private final Random random;
  
      public StripedReplicaPlacer(Random random) {
          this.random = random;
      }
  
      @Override
      public TopicAssignment place(
          PlacementSpec placement,
          ClusterDescriber cluster
      ) throws InvalidReplicationFactorException {
          RackList rackList = new RackList(random, cluster.usableBrokers());
          throwInvalidReplicationFactorIfNonPositive(placement.numReplicas());
          throwInvalidReplicationFactorIfZero(rackList.numUnfencedBrokers());
          throwInvalidReplicationFactorIfTooFewBrokers(placement.numReplicas(),
              rackList.numTotalBrokers());
          List<List<Integer>> placements = new ArrayList<>(placement.numPartitions());
          for (int partition = 0; partition < placement.numPartitions(); partition++) {
              placements.add(rackList.place(placement.numReplicas()));
          }
          return new TopicAssignment(
              placements.stream().map(replicas -> new PartitionAssignment(replicas, cluster)).toList()
          );
      }
  }
  