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.controller;
  
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.metadata.Replicas;
  import org.apache.kafka.server.common.TopicIdPartition;
  import org.apache.kafka.timeline.SnapshotRegistry;
  import org.apache.kafka.timeline.TimelineHashMap;
  
  import java.util.Arrays;
  import java.util.Iterator;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.NoSuchElementException;
  
  import static org.apache.kafka.metadata.LeaderConstants.NO_LEADER;
  import static org.apache.kafka.metadata.Replicas.NONE;
  
  
  /**
   * Associates brokers with their in-sync partitions.
   *
   * This is useful when we need to remove a broker from all the ISRs, or move all leaders
   * away from a broker.
   *
   * We also track all the partitions that currently have no leader.
   *
   * The core data structure is a map from broker IDs to topic maps.  Each topic map relates
   * topic UUIDs to arrays of partition IDs.
   *
   * Each entry in the array has a high bit which indicates that the broker is the leader
   * for the given partition, as well as 31 low bits which contain the partition id.  This
   * works because partition IDs cannot be negative.
   */
  public class BrokersToIsrs {
      private static final int LEADER_FLAG = 0x8000_0000;
  
      private static final int REPLICA_MASK = 0x7fff_ffff;
  
      static class PartitionsOnReplicaIterator implements Iterator<TopicIdPartition> {
          private final Iterator<Entry<Uuid, int[]>> iterator;
          private final boolean leaderOnly;
          private int offset = 0;
          Uuid uuid = Uuid.ZERO_UUID;
          int[] replicas = NONE;
          private TopicIdPartition next = null;
  
          PartitionsOnReplicaIterator(Map<Uuid, int[]> topicMap, boolean leaderOnly) {
              this.iterator = topicMap.entrySet().iterator();
              this.leaderOnly = leaderOnly;
          }
  
          @Override
          public boolean hasNext() {
              if (next != null) return true;
              while (true) {
                  if (offset >= replicas.length) {
                      if (!iterator.hasNext()) return false;
                      offset = 0;
                      Entry<Uuid, int[]> entry = iterator.next();
                      uuid = entry.getKey();
                      replicas = entry.getValue();
                  }
                  int replica = replicas[offset++];
                  if ((!leaderOnly) || (replica & LEADER_FLAG) != 0) {
                      next = new TopicIdPartition(uuid, replica & REPLICA_MASK);
                      return true;
                  }
              }
          }
  
          @Override
          public TopicIdPartition next() {
              if (!hasNext()) {
                  throw new NoSuchElementException();
              }
              TopicIdPartition result = next;
              next = null;
              return result;
          }
      }
  
      private final SnapshotRegistry snapshotRegistry;
  
      /**
       * A map of broker IDs to the partitions that the broker is in the ISR for.
       * Partitions with no isr members appear in this map under id NO_LEADER.
       */
      private final TimelineHashMap<Integer, TimelineHashMap<Uuid, int[]>> isrMembers;
      
      BrokersToIsrs(SnapshotRegistry snapshotRegistry) {
          this.snapshotRegistry = snapshotRegistry;
          this.isrMembers = new TimelineHashMap<>(snapshotRegistry, 0);
      }
  
      /**
       * Update our records of a partition's ISR.
       *
       * @param topicId       The topic ID of the partition.
       * @param partitionId   The partition ID of the partition.
       * @param prevIsr       The previous ISR, or null if the partition is new.
       * @param nextIsr       The new ISR, or null if the partition is being removed.
       * @param prevLeader    The previous leader, or NO_LEADER if the partition had no leader.
       * @param nextLeader    The new leader, or NO_LEADER if the partition now has no leader.
       */
      void update(Uuid topicId, int partitionId, int[] prevIsr, int[] nextIsr,
                  int prevLeader, int nextLeader) {
          int[] prev;
          if (prevIsr == null) {
              prev = NONE;
          } else {
              if (prevLeader == NO_LEADER) {
                  prev = Replicas.copyWith(prevIsr, NO_LEADER);
              } else {
                  prev = Replicas.clone(prevIsr);
              }
              Arrays.sort(prev);
          }
          int[] next;
          if (nextIsr == null) {
              next = NONE;
          } else {
              if (nextLeader == NO_LEADER) {
                  next = Replicas.copyWith(nextIsr, NO_LEADER);
              } else {
                  next = Replicas.clone(nextIsr);
              }
              Arrays.sort(next);
          }
          int i = 0, j = 0;
          while (true) {
              if (i == prev.length) {
                  if (j == next.length) {
                      break;
                  }
                  int newReplica = next[j];
                  add(newReplica, topicId, partitionId, newReplica == nextLeader);
                  j++;
              } else if (j == next.length) {
                  int prevReplica = prev[i];
                  remove(prevReplica, topicId, partitionId, prevReplica == prevLeader);
                  i++;
              } else {
                  int prevReplica = prev[i];
                  int newReplica = next[j];
                  if (prevReplica < newReplica) {
                      remove(prevReplica, topicId, partitionId, prevReplica == prevLeader);
                      i++;
                  } else if (prevReplica > newReplica) {
                      add(newReplica, topicId, partitionId, newReplica == nextLeader);
                      j++;
                  } else {
                      boolean wasLeader = prevReplica == prevLeader;
                      boolean isLeader = prevReplica == nextLeader;
                      if (wasLeader != isLeader) {
                          change(prevReplica, topicId, partitionId, wasLeader, isLeader);
                      }
                      i++;
                      j++;
                  }
              }
          }
      }
  
      void removeTopicEntryForBroker(Uuid topicId, int brokerId) {
          Map<Uuid, int[]> topicMap = isrMembers.get(brokerId);
          if (topicMap != null) {
              topicMap.remove(topicId);
          }
      }
  
      private void add(int brokerId, Uuid topicId, int newPartition, boolean leader) {
          if (leader) {
              newPartition = newPartition | LEADER_FLAG;
          }
          TimelineHashMap<Uuid, int[]> topicMap = isrMembers.get(brokerId);
          if (topicMap == null) {
              topicMap = new TimelineHashMap<>(snapshotRegistry, 0);
              isrMembers.put(brokerId, topicMap);
          }
          int[] partitions = topicMap.get(topicId);
          int[] newPartitions;
          if (partitions == null) {
              newPartitions = new int[1];
          } else {
              newPartitions = new int[partitions.length + 1];
              System.arraycopy(partitions, 0, newPartitions, 0, partitions.length);
          }
          newPartitions[newPartitions.length - 1] = newPartition;
          topicMap.put(topicId, newPartitions);
      }
  
      private void change(int brokerId, Uuid topicId, int partition,
                          boolean wasLeader, boolean isLeader) {
          TimelineHashMap<Uuid, int[]> topicMap = isrMembers.get(brokerId);
          if (topicMap == null) {
              throw new RuntimeException("Broker " + brokerId + " has no isrMembers " +
                  "entry, so we can't change " + topicId + ":" + partition);
          }
          int[] partitions = topicMap.get(topicId);
          if (partitions == null) {
              throw new RuntimeException("Broker " + brokerId + " has no " +
                  "entry in isrMembers for topic " + topicId);
          }
          int[] newPartitions = new int[partitions.length];
          int target = wasLeader ? partition | LEADER_FLAG : partition;
          for (int i = 0; i < partitions.length; i++) {
              int cur = partitions[i];
              if (cur == target) {
                  newPartitions[i] = isLeader ? partition | LEADER_FLAG : partition;
              } else {
                  newPartitions[i] = cur;
              }
          }
          topicMap.put(topicId, newPartitions);
      }
  
      private void remove(int brokerId, Uuid topicId, int removedPartition, boolean leader) {
          if (leader) {
              removedPartition = removedPartition | LEADER_FLAG;
          }
          TimelineHashMap<Uuid, int[]> topicMap = isrMembers.get(brokerId);
          if (topicMap == null) {
              throw new RuntimeException("Broker " + brokerId + " has no isrMembers " +
                  "entry, so we can't remove " + topicId + ":" + removedPartition);
          }
          int[] partitions = topicMap.get(topicId);
          if (partitions == null) {
              throw new RuntimeException("Broker " + brokerId + " has no " +
                  "entry in isrMembers for topic " + topicId);
          }
          if (partitions.length == 1) {
              if (partitions[0] != removedPartition) {
                  throw new RuntimeException("Broker " + brokerId + " has no " +
                      "entry in isrMembers for " + topicId + ":" + removedPartition);
              }
              topicMap.remove(topicId);
              if (topicMap.isEmpty()) {
                  isrMembers.remove(brokerId);
              }
          } else {
              int[] newPartitions = new int[partitions.length - 1];
              int i = 0;
              for (int partition : partitions) {
                  if (partition != removedPartition) {
                      newPartitions[i++] = partition;
                  }
              }
              topicMap.put(topicId, newPartitions);
          }
      }
  
      PartitionsOnReplicaIterator iterator(int brokerId, boolean leadersOnly) {
          Map<Uuid, int[]> topicMap = isrMembers.get(brokerId);
          if (topicMap == null) {
              topicMap = Map.of();
          }
          return new PartitionsOnReplicaIterator(topicMap, leadersOnly);
      }
  
      PartitionsOnReplicaIterator partitionsWithNoLeader() {
          return iterator(NO_LEADER, true);
      }
  
      PartitionsOnReplicaIterator partitionsWithBrokerInIsr(int brokerId) {
          return iterator(brokerId, false);
      }
  
      boolean hasLeaderships(int brokerId) {
          return iterator(brokerId, true).hasNext();
      }
  }
  