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.coordinator.share;
  
  import org.apache.kafka.server.share.persister.PersisterStateBatch;
  
  import java.util.ArrayList;
  import java.util.Iterator;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Objects;
  import java.util.TreeSet;
  
  public class PersisterStateBatchCombiner {
      private List<PersisterStateBatch> combinedBatchList;    // link between pruning and merging
      private final long startOffset;
      private TreeSet<PersisterStateBatch> sortedBatches;
      private List<PersisterStateBatch> finalBatchList;   // final list is built here
  
      public PersisterStateBatchCombiner(
          List<PersisterStateBatch> batchesSoFar,
          List<PersisterStateBatch> newBatches,
          long startOffset
      ) {
          initializeCombinedList(batchesSoFar, newBatches);
          int estimatedResultSize = (combinedBatchList.size() * 3) / 2;   // heuristic size - 50% overallocation
          finalBatchList = new ArrayList<>(estimatedResultSize);
          this.startOffset = startOffset;
      }
  
      private void initializeCombinedList(List<PersisterStateBatch> batchesSoFar, List<PersisterStateBatch> newBatches) {
          boolean soFarEmpty = batchesSoFar == null || batchesSoFar.isEmpty();
          boolean newBatchesEmpty = newBatches == null || newBatches.isEmpty();
  
          if (soFarEmpty && newBatchesEmpty) {
              combinedBatchList = new ArrayList<>();
          } else if (soFarEmpty) {
              combinedBatchList = new ArrayList<>(newBatches);    // new list as the original one could be unmodifiable
          } else if (newBatchesEmpty) {
              combinedBatchList = new ArrayList<>(batchesSoFar);  // new list as the original one could be unmodifiable
          } else {
              combinedBatchList = new ArrayList<>(batchesSoFar.size() + newBatches.size());
              combinedBatchList.addAll(batchesSoFar);
              combinedBatchList.addAll(newBatches);
          }
      }
  
      /**
       * Algorithm: Merge current state batches and new batches into a single non-overlapping batch list.
       * Input: batchesSoFar, newBatches, startOffset
       * Output: combined list with non-overlapping batches (finalBatchList)
       * <p>
       * - Add both currentBatches and newBatches into a single list combinedBatchList
       * - if combinedBatchList.size() <= 1 return combinedBatchList
       * <p>
       * - Remove/prune any batches from the combinedBatchList:
       * -    if batch.lastOffset < startOffset then remove batch from combinedBatchList
       * -    else if batch.firstOffset > startOffset then we will keep the batch
       * -    else if batch.firstOffset <= startOffset <= batch.lastOffset then keep [startOffset, batch.lastOffset] part only and discard rest.
       * <p>
       * - create a treeset sortedBatches using pruned combinedBatchList
       * - find first 2 mergeable batches in sortedBatches set, say, prev and candidate.
       * - remove any non-overlapping batches from sortedBatches encountered during the find operation and add them to a finalBatchList
       * - do repeat until a mergeable pair is not found:
       * -    based on various conditions of offset overlap and batch state differences combine the batches or
       *      create new batches, if required, and add to the sortedBatches.
       * -    find first 2 mergeable batches in sortedBatches set, say, prev and candidate.
       * -    remove any non-mergeable batches from sortedBatches encountered during the find operation and add them to a finalBatchList
       * - done
       * - return the finalBatchList
       *
       * @return list of {@link PersisterStateBatch} representing non-overlapping combined batches
       */
      public List<PersisterStateBatch> combineStateBatches() {
          pruneBatches();
          mergeBatches();
          return finalBatchList;
      }
  
      private void mergeBatches() {
          if (combinedBatchList.size() < 2) {
              finalBatchList = combinedBatchList;
              return;
          }
  
          sortedBatches = new TreeSet<>(combinedBatchList);
  
          MergeCandidatePair overlapState = getMergeCandidatePair();
  
          while (overlapState != MergeCandidatePair.EMPTY) {
              PersisterStateBatch prev = overlapState.prev();
              PersisterStateBatch candidate = overlapState.candidate();
  
              // remove both previous and candidate for easier
              // assessment about adding batches to sortedBatches
              sortedBatches.remove(prev);
              sortedBatches.remove(candidate);
  
              if (compareBatchDeliveryInfo(candidate, prev) == 0) {  // same state and overlap or contiguous
                  // overlap and same state (prev.firstOffset <= candidate.firstOffset) due to sort
                  // covers:
                  // case:        1        2          3            4          5           6          7 (contiguous)
                  // prev:        ------   -------    -------      -------   -------   --------    -------
                  // candidate:   ------   ----       ----------     ---        ----       -------        -------
                  handleSameStateMerge(prev, candidate);  // pair can be contiguous or overlapping
              } else {
                  // If we reach here then it is guaranteed that the batch pair is overlapping and
                  // non-contiguous because getMergeCandidatePair only returns contiguous pair if
                  // the constituents have the same delivery count and state.
                  // covers:
                  // case:        1        2*          3            4          5           6             7*
                  // prev:        ------   -------    -------      -------    -------     --------      ------
                  // candidate:   ------   ----       ---------      ----        ----        -------   -------
                  // max batches: 1           2       2                3          2            2          2
                  // min batches: 1           1       1                1          1            2          1
                  // * not possible with treeset
                  handleDiffStateOverlap(prev, candidate);
              }
              overlapState = getMergeCandidatePair();
          }
          finalBatchList.addAll(sortedBatches);   // some non overlapping batches might have remained
      }
  
      /**
       * Compares the non-offset state of 2 batches i.e. the deliveryCount and deliverState.
       * <p>
       * Uses standard compareTo contract x < y => +int, x > y => -int, x == y => 0
       *
       * @param b1 - {@link PersisterStateBatch} to compare
       * @param b2 - {@link PersisterStateBatch} to compare
       * @return int representing comparison result.
       */
      private int compareBatchDeliveryInfo(PersisterStateBatch b1, PersisterStateBatch b2) {
          int deltaCount = Short.compare(b1.deliveryCount(), b2.deliveryCount());
  
          // Delivery state could be:
          // 0 - AVAILABLE (non-terminal)
          // 1 - ACQUIRED - should not be persisted yet
          // 2 - ACKNOWLEDGED (terminal)
          // 3 - ARCHIVING - not implemented in KIP-932 - non-terminal - leads only to ARCHIVED
          // 4 - ARCHIVED (terminal)
  
          if (deltaCount == 0) {   // same delivery count
              return Byte.compare(b1.deliveryState(), b2.deliveryState());
          }
          return deltaCount;
      }
  
      /**
       * Accepts a sorted set of state batches and finds the first 2 batches which can be merged.
       * Merged implies that they have some offsets in common or, they are contiguous with the same state.
       * <p>
       * Any non-mergeable batches prefixing a good mergeable pair are removed from the sortedBatches.
       * For example:
       * ----- ----  ----- -----      -----
       *                      ------
       * <---------------> <-------->
       * non-overlapping   1st overlapping pair
       *
       * @return object representing the overlap state
       */
      private MergeCandidatePair getMergeCandidatePair() {
          if (sortedBatches == null || sortedBatches.isEmpty()) {
              return MergeCandidatePair.EMPTY;
          }
          Iterator<PersisterStateBatch> iter = sortedBatches.iterator();
          PersisterStateBatch prev = iter.next();
          List<PersisterStateBatch> nonOverlapping = new LinkedList<>();
          while (iter.hasNext()) {
              PersisterStateBatch candidate = iter.next();
              if (candidate.firstOffset() <= prev.lastOffset() || // overlap
                  prev.lastOffset() + 1 == candidate.firstOffset() && compareBatchDeliveryInfo(prev, candidate) == 0) {  // contiguous
                  updateBatchContainers(nonOverlapping);
                  return new MergeCandidatePair(
                      prev,
                      candidate
                  );
              }
              nonOverlapping.add(prev);
              prev = candidate;
          }
  
          updateBatchContainers(nonOverlapping);
          return MergeCandidatePair.EMPTY;
      }
  
      private void updateBatchContainers(List<PersisterStateBatch> nonOverlappingBatches) {
          nonOverlappingBatches.forEach(sortedBatches::remove);
          finalBatchList.addAll(nonOverlappingBatches);
      }
  
      /**
       * Accepts a list of {@link PersisterStateBatch} and checks:
       * - last offset is < start offset => batch is removed
       * - first offset > start offset => batch is preserved
       * - start offset intersects the batch => part of batch before start offset is removed and
       * the part after it is preserved.
       */
      private void pruneBatches() {
          if (startOffset != -1) {
              List<PersisterStateBatch> retainedBatches = new ArrayList<>(combinedBatchList.size());
              combinedBatchList.forEach(batch -> {
                  if (batch.lastOffset() < startOffset) {
                      // batch is expired, skip current iteration
                      // -------
                      //         | -> start offset
                      return;
                  }
  
                  if (batch.firstOffset() >= startOffset) {
                      // complete batch is valid
                      //    ---------
                      //  | -> start offset
                      retainedBatches.add(batch);
                  } else {
                      // start offset intersects batch
                      //   ---------
                      //       |     -> start offset
                      retainedBatches.add(new PersisterStateBatch(startOffset, batch.lastOffset(), batch.deliveryState(), batch.deliveryCount()));
                  }
              });
              // update the instance variable
              combinedBatchList = retainedBatches;
          }
      }
  
      private void handleSameStateMerge(PersisterStateBatch prev, PersisterStateBatch candidate) {
          sortedBatches.add(new PersisterStateBatch(
              prev.firstOffset(),
              // cover cases
              // prev:      ------   --------       ---------
              // candidate:   ---       ----------           -----
              Math.max(candidate.lastOffset(), prev.lastOffset()),
              prev.deliveryState(),
              prev.deliveryCount()
          ));
      }
  
      private void handleDiffStateOverlap(PersisterStateBatch prev, PersisterStateBatch candidate) {
          if (candidate.firstOffset() == prev.firstOffset()) {
              handleDiffStateOverlapFirstOffsetAligned(prev, candidate);
          } else {    // candidate.firstOffset() > prev.firstOffset()
              handleDiffStateOverlapFirstOffsetNotAligned(prev, candidate);
          }
      }
  
      private void handleDiffStateOverlapFirstOffsetAligned(PersisterStateBatch prev, PersisterStateBatch candidate) {
          if (candidate.lastOffset() == prev.lastOffset()) {  // case 1
              // candidate can never have lower or equal priority
              // since sortedBatches order takes that into account.
              // -------
              // -------
              sortedBatches.add(candidate);
          } else {
              // case 2 is not possible with TreeSet. It is symmetric to case 3.
              // case 3
              // --------
              // -----------
              if (compareBatchDeliveryInfo(candidate, prev) < 0) {
                  sortedBatches.add(prev);
                  sortedBatches.add(new PersisterStateBatch(
                      prev.lastOffset() + 1,
                      candidate.lastOffset(),
                      candidate.deliveryState(),
                      candidate.deliveryCount()
                  ));
              } else {
                  // candidate priority is >= prev
                  sortedBatches.add(candidate);
              }
          }
      }
  
      private void handleDiffStateOverlapFirstOffsetNotAligned(PersisterStateBatch prev, PersisterStateBatch candidate) {
          if (candidate.lastOffset() < prev.lastOffset()) {    // case 4
              handleDiffStateOverlapPrevSwallowsCandidate(prev, candidate);
          } else if (candidate.lastOffset() == prev.lastOffset()) {    // case 5
              handleDiffStateOverlapLastOffsetAligned(prev, candidate);
          } else {    // case 6
              handleDiffStateOverlapCandidateOffsetsLarger(prev, candidate);
          }
      }
  
      private void handleDiffStateOverlapPrevSwallowsCandidate(PersisterStateBatch prev, PersisterStateBatch candidate) {
          // --------
          //   ----
          if (compareBatchDeliveryInfo(candidate, prev) < 0) {
              sortedBatches.add(prev);
          } else {
              sortedBatches.add(new PersisterStateBatch(
                  prev.firstOffset(),
                  candidate.firstOffset() - 1,
                  prev.deliveryState(),
                  prev.deliveryCount()
              ));
  
              sortedBatches.add(candidate);
  
              sortedBatches.add(new PersisterStateBatch(
                  candidate.lastOffset() + 1,
                  prev.lastOffset(),
                  prev.deliveryState(),
                  prev.deliveryCount()
              ));
          }
      }
  
      private void handleDiffStateOverlapLastOffsetAligned(PersisterStateBatch prev, PersisterStateBatch candidate) {
          // --------
          //    -----
          if (compareBatchDeliveryInfo(candidate, prev) < 0) {
              sortedBatches.add(prev);
          } else {
              sortedBatches.add(new PersisterStateBatch(
                  prev.firstOffset(),
                  candidate.firstOffset() - 1,
                  prev.deliveryState(),
                  prev.deliveryCount()
              ));
  
              sortedBatches.add(candidate);
          }
      }
  
      private void handleDiffStateOverlapCandidateOffsetsLarger(PersisterStateBatch prev, PersisterStateBatch candidate) {
          //   -------
          //      -------
          if (compareBatchDeliveryInfo(candidate, prev) < 0) {
              sortedBatches.add(prev);
  
              sortedBatches.add(new PersisterStateBatch(
                  prev.lastOffset() + 1,
                  candidate.lastOffset(),
                  candidate.deliveryState(),
                  candidate.deliveryCount()
              ));
          } else {
              // candidate has higher priority
              sortedBatches.add(new PersisterStateBatch(
                  prev.firstOffset(),
                  candidate.firstOffset() - 1,
                  prev.deliveryState(),
                  prev.deliveryCount()
              ));
  
              sortedBatches.add(candidate);
          }
      }
  
      /**
       * Holder class for intermediate state
       * used in the batch merge algorithm.
       */
      static class MergeCandidatePair {
          private final PersisterStateBatch prev;
          private final PersisterStateBatch candidate;
          public static final MergeCandidatePair EMPTY = new MergeCandidatePair(null, null);
  
          public MergeCandidatePair(
              PersisterStateBatch prev,
              PersisterStateBatch candidate
          ) {
              this.prev = prev;
              this.candidate = candidate;
          }
  
          public PersisterStateBatch prev() {
              return prev;
          }
  
          public PersisterStateBatch candidate() {
              return candidate;
          }
  
          @Override
          public boolean equals(Object o) {
              if (this == o) return true;
              if (!(o instanceof MergeCandidatePair that)) return false;
              return Objects.equals(prev, that.prev) && Objects.equals(candidate, that.candidate);
          }
  
          @Override
          public int hashCode() {
              return Objects.hash(prev, candidate);
          }
      }
  }
  