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.image.loader;
  
  import org.apache.kafka.common.metadata.MetadataRecordType;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.image.MetadataDelta;
  import org.apache.kafka.image.MetadataImage;
  import org.apache.kafka.image.MetadataProvenance;
  import org.apache.kafka.raft.Batch;
  import org.apache.kafka.raft.LeaderAndEpoch;
  import org.apache.kafka.server.common.ApiMessageAndVersion;
  import org.apache.kafka.server.fault.FaultHandler;
  
  import org.slf4j.Logger;
  
  import static java.util.concurrent.TimeUnit.NANOSECONDS;
  
  /**
   * Loads batches of metadata updates from Raft commits into MetadataDelta-s. Multiple batches from a commit
   * are buffered into a MetadataDelta to achieve batching of records and reduce the number of times
   * MetadataPublishers must be updated. This class also supports metadata transactions (KIP-866).
   *
   *
   */
  public class MetadataBatchLoader {
  
      enum TransactionState {
          NO_TRANSACTION,
          STARTED_TRANSACTION,
          CONTINUED_TRANSACTION,
          ENDED_TRANSACTION,
          ABORTED_TRANSACTION
      }
  
      @FunctionalInterface
      public interface MetadataUpdater {
          void update(MetadataDelta delta, MetadataImage image, LogDeltaManifest manifest);
      }
  
      private final Logger log;
      private final Time time;
      private final FaultHandler faultHandler;
      private final MetadataUpdater callback;
  
      private MetadataImage image;
      private MetadataDelta delta;
      private long lastOffset;
      private int lastEpoch;
      private long lastContainedLogTimeMs;
      private long numBytes;
      private int numBatches;
      private long totalBatchElapsedNs;
      private TransactionState transactionState;
      private boolean hasSeenRecord;
  
      public MetadataBatchLoader(
          LogContext logContext,
          Time time,
          FaultHandler faultHandler,
          MetadataUpdater callback
      ) {
          this.log = logContext.logger(MetadataBatchLoader.class);
          this.time = time;
          this.faultHandler = faultHandler;
          this.callback = callback;
          this.resetToImage(MetadataImage.EMPTY);
          this.hasSeenRecord = false;
      }
  
      /**
       * @return True if this batch loader has seen at least one record.
       */
      public boolean hasSeenRecord() {
          return hasSeenRecord;
      }
  
      /**
       * Reset the state of this batch loader to the given image. Any un-flushed state will be
       * discarded. This is called after applying a delta and passing it back to MetadataLoader, or
       * when MetadataLoader loads a snapshot.
       *
       * @param image     Metadata image to reset this batch loader's state to.
       */
      public final void resetToImage(MetadataImage image) {
          this.image = image;
          this.hasSeenRecord = !image.isEmpty();
          this.delta = new MetadataDelta.Builder().setImage(image).build();
          this.transactionState = TransactionState.NO_TRANSACTION;
          this.lastOffset = image.provenance().lastContainedOffset();
          this.lastEpoch = image.provenance().lastContainedEpoch();
          this.lastContainedLogTimeMs = image.provenance().lastContainedLogTimeMs();
          this.numBytes = 0;
          this.numBatches = 0;
          this.totalBatchElapsedNs = 0;
      }
  
      /**
       * Load a batch of records from the log. We have to do some bookkeeping here to
       * translate between batch offsets and record offsets, and track the number of bytes we
       * have read. Additionally, there is the chance that one of the records is a metadata
       * version change which needs to be handled differently.
       * </p>
       * If this batch starts a transaction, any records preceding the transaction in this
       * batch will be implicitly added to the transaction.
       *
       * @param batch    The reader which yields the batches.
       * @return         The time in nanoseconds that elapsed while loading this batch
       */
  
      public long loadBatch(Batch<ApiMessageAndVersion> batch, LeaderAndEpoch leaderAndEpoch) {
          long startNs = time.nanoseconds();
          int indexWithinBatch = 0;
  
          lastContainedLogTimeMs = batch.appendTimestamp();
          lastEpoch = batch.epoch();
  
          for (ApiMessageAndVersion record : batch.records()) {
              try {
                  replay(record);
              } catch (Throwable e) {
                  faultHandler.handleFault("Error loading metadata log record from offset " +
                      (batch.baseOffset() + indexWithinBatch), e);
              }
  
              // Emit the accumulated delta if a new transaction has been started and one of the following is true
              // 1) this is not the first record in this batch
              // 2) this is not the first batch since last emitting a delta
              if (transactionState == TransactionState.STARTED_TRANSACTION && (indexWithinBatch > 0 || numBatches > 0)) {
                  // Accumulated delta is aligned on batch boundaries iff the BeginTransactionRecord is the first record in a batch
                  MetadataProvenance provenance = new MetadataProvenance(lastOffset, lastEpoch, lastContainedLogTimeMs, indexWithinBatch == 0);
                  LogDeltaManifest manifest = LogDeltaManifest.newBuilder()
                      .provenance(provenance)
                      .leaderAndEpoch(leaderAndEpoch)
                      .numBatches(numBatches) // This will be zero if we have not yet read a batch
                      .elapsedNs(totalBatchElapsedNs)
                      .numBytes(numBytes)     // This will be zero if we have not yet read a batch
                      .build();
                  if (log.isDebugEnabled()) {
                      log.debug("handleCommit: Generated a metadata delta between {} and {} from {} batch(es) in {} us. The delta is {}.",
                              image.offset(), manifest.provenance().lastContainedOffset(),
                              manifest.numBatches(), NANOSECONDS.toMicros(manifest.elapsedNs()),
                              provenance.isOffsetBatchAligned() ? "batch aligned" : "not batch aligned");
                  }
                  applyDeltaAndUpdate(delta, manifest);
                  transactionState = TransactionState.STARTED_TRANSACTION;
              }
  
              lastOffset = batch.baseOffset() + indexWithinBatch;
              indexWithinBatch++;
          }
          long elapsedNs = time.nanoseconds() - startNs;
  
          // Update state for the manifest. The actual byte count will only be included in the last delta emitted for
          // a given batch or transaction.
          lastOffset = batch.lastOffset();
          numBytes += batch.sizeInBytes();
          numBatches += 1;
          totalBatchElapsedNs += elapsedNs;
          return totalBatchElapsedNs;
      }
  
      /**
       * Flush the metadata accumulated in this batch loader if not in the middle of a transaction. The
       * flushed metadata will be passed to the {@link MetadataUpdater} configured for this class.
       */
      public void maybeFlushBatches(LeaderAndEpoch leaderAndEpoch, boolean isOffsetBatchAligned) {
          MetadataProvenance provenance = new MetadataProvenance(lastOffset, lastEpoch, lastContainedLogTimeMs, isOffsetBatchAligned);
          LogDeltaManifest manifest = LogDeltaManifest.newBuilder()
              .provenance(provenance)
              .leaderAndEpoch(leaderAndEpoch)
              .numBatches(numBatches)
              .elapsedNs(totalBatchElapsedNs)
              .numBytes(numBytes)
              .build();
          switch (transactionState) {
              case STARTED_TRANSACTION:
              case CONTINUED_TRANSACTION:
                  log.debug("handleCommit: not publishing since a transaction starting at {} is still in progress. " +
                      "{} batch(es) processed so far.", image.offset(), numBatches);
                  break;
              case ABORTED_TRANSACTION:
                  log.debug("handleCommit: publishing empty delta between {} and {} from {} batch(es) " +
                      "since a transaction was aborted", image.offset(), manifest.provenance().lastContainedOffset(),
                      manifest.numBatches());
                  applyDeltaAndUpdate(new MetadataDelta.Builder().setImage(image).build(), manifest);
                  break;
              case ENDED_TRANSACTION:
              case NO_TRANSACTION:
                  if (log.isDebugEnabled()) {
                      log.debug("handleCommit: Generated a metadata delta between {} and {} from {} batch(es) in {} us.",
                          image.offset(), manifest.provenance().lastContainedOffset(),
                          manifest.numBatches(), NANOSECONDS.toMicros(manifest.elapsedNs()));
                  }
                  applyDeltaAndUpdate(delta, manifest);
                  break;
          }
      }
  
      private void replay(ApiMessageAndVersion record) {
          MetadataRecordType type = MetadataRecordType.fromId(record.message().apiKey());
          switch (type) {
              case BEGIN_TRANSACTION_RECORD:
                  if (transactionState == TransactionState.STARTED_TRANSACTION ||
                      transactionState == TransactionState.CONTINUED_TRANSACTION) {
                      throw new RuntimeException("Encountered BeginTransactionRecord while already in a transaction");
                  } else {
                      transactionState = TransactionState.STARTED_TRANSACTION;
                  }
                  break;
              case END_TRANSACTION_RECORD:
                  if (transactionState == TransactionState.CONTINUED_TRANSACTION ||
                      transactionState == TransactionState.STARTED_TRANSACTION) {
                      transactionState = TransactionState.ENDED_TRANSACTION;
                  } else {
                      throw new RuntimeException("Encountered EndTransactionRecord without having seen a BeginTransactionRecord");
                  }
                  break;
              case ABORT_TRANSACTION_RECORD:
                  if (transactionState == TransactionState.CONTINUED_TRANSACTION ||
                      transactionState == TransactionState.STARTED_TRANSACTION) {
                      transactionState = TransactionState.ABORTED_TRANSACTION;
                  } else {
                      throw new RuntimeException("Encountered AbortTransactionRecord without having seen a BeginTransactionRecord");
                  }
                  break;
              default:
                  switch (transactionState) {
                      case STARTED_TRANSACTION:
                          // If we see a non-transaction record after starting a transaction, transition to CONTINUED_TRANSACTION
                          transactionState = TransactionState.CONTINUED_TRANSACTION;
                          break;
                      case ENDED_TRANSACTION:
                      case ABORTED_TRANSACTION:
                          // If we see a non-transaction record after ending a transaction, transition back to NO_TRANSACTION
                          transactionState = TransactionState.NO_TRANSACTION;
                          break;
                      case CONTINUED_TRANSACTION:
                      case NO_TRANSACTION:
                      default:
                          break;
                  }
                  hasSeenRecord = true;
                  delta.replay(record.message());
          }
      }
  
      private void applyDeltaAndUpdate(MetadataDelta delta, LogDeltaManifest manifest) {
          try {
              image = delta.apply(manifest.provenance());
          } catch (Throwable e) {
              faultHandler.handleFault("Error generating new metadata image from " +
                  "metadata delta between offset " + image.offset() +
                  " and " + manifest.provenance().lastContainedOffset(), e);
          }
  
          // Whether we can apply the delta or not, we need to make sure the batch loader gets reset
          // to the image known to MetadataLoader
          callback.update(delta, image, manifest);
          resetToImage(image);
      }
  }
  