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.connect.mirror;
  
  import org.apache.kafka.clients.consumer.Consumer;
  import org.apache.kafka.clients.consumer.ConsumerRecord;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.common.utils.Utils;
  import org.apache.kafka.connect.util.KafkaBasedLog;
  import org.apache.kafka.connect.util.TopicAdmin;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.Arrays;
  import java.util.Map;
  import java.util.Optional;
  import java.util.OptionalLong;
  import java.util.concurrent.ConcurrentHashMap;
  
  /**
   * Used internally by MirrorMaker. Stores offset syncs and performs offset translation.
   * <p>A limited number of offset syncs can be stored per TopicPartition, in a way which provides better translation
   * later in the topic, closer to the live end of the topic.
   * This maintains the following invariants for each topic-partition in the in-memory sync storage:
   * <ul>
   *     <li>Invariant A: syncs[0] is the latest offset sync from the syncs topic</li>
   *     <li>Invariant B: For each i,j, i < j, syncs[i] != syncs[j]: syncs[i].upstream <= syncs[j].upstream + 2^j - 2^i</li>
   *     <li>Invariant C: For each i,j, i < j, syncs[i] != syncs[j]: syncs[i].upstream >= syncs[j].upstream + 2^(i-2)</li>
   *     <li>Invariant D: syncs[63] is the earliest offset sync from the syncs topic usable for translation</li>
   * </ul>
   * <p>The above invariants ensure that the store is kept updated upon receipt of each sync, and that distinct
   * offset syncs are separated by approximately exponential space. They can be checked locally (by comparing all adjacent
   * indexes) but hold globally (for all pairs of any distance). This allows updates to the store in linear time.
   * <p>Offset translation uses the syncs[i] which most closely precedes the upstream consumer group's current offset.
   * For a fixed in-memory state, translation of variable upstream offsets will be monotonic.
   * For variable in-memory state, translation of a fixed upstream offset will not be monotonic.
   * <p>Translation will be unavailable for all topic-partitions before an initial read-to-end of the offset syncs topic
   * is complete. Translation will be unavailable after that if no syncs are present for a topic-partition, if replication
   * started after the position of the consumer group, or if relevant offset syncs for the topic were potentially used as
   * for translation in an earlier generation of the sync store.
   */
  public class OffsetSyncStore implements AutoCloseable {
  
      private static final Logger log = LoggerFactory.getLogger(OffsetSyncStore.class);
      // Store one offset sync for each bit of the topic offset.
      // Visible for testing
      static final int SYNCS_PER_PARTITION = Long.SIZE;
      private final KafkaBasedLog<byte[], byte[]> backingStore;
      private final Map<TopicPartition, OffsetSync[]> offsetSyncs = new ConcurrentHashMap<>();
      private final TopicAdmin admin;
      protected volatile boolean initializationMustReadToEnd = true;
      protected volatile boolean readToEnd = false;
  
      // package access to avoid Java 21 "this-escape" warning
      OffsetSyncStore(MirrorCheckpointConfig config) {
          Consumer<byte[], byte[]> consumer = null;
          TopicAdmin admin = null;
          KafkaBasedLog<byte[], byte[]> store;
          try {
              consumer = MirrorUtils.newConsumer(config.offsetSyncsTopicConsumerConfig());
              admin = new TopicAdmin(
                      config.offsetSyncsTopicAdminConfig(),
                      config.forwardingAdmin(config.offsetSyncsTopicAdminConfig()));
              store = createBackingStore(config, consumer, admin);
          } catch (Throwable t) {
              Utils.closeQuietly(consumer, "consumer for offset syncs");
              Utils.closeQuietly(admin, "admin client for offset syncs");
              throw t;
          }
          this.admin = admin;
          this.backingStore = store;
      }
  
      private KafkaBasedLog<byte[], byte[]> createBackingStore(MirrorCheckpointConfig config, Consumer<byte[], byte[]> consumer, TopicAdmin admin) {
          return KafkaBasedLog.withExistingClients(
                  config.offsetSyncsTopic(),
                  consumer,
                  null,
                  admin,
                  (error, record) -> this.handleRecord(record),
                  Time.SYSTEM,
                  ignored -> { },
                  topicPartition -> topicPartition.partition() == 0
          );
      }
  
      // for testing
      OffsetSyncStore() {
          this.admin = null;
          this.backingStore = null;
      }
  
      /**
       * Start the OffsetSyncStore, blocking until all previous Offset Syncs have been read from backing storage.
       */
      public void start(boolean initializationMustReadToEnd) {
          this.initializationMustReadToEnd = initializationMustReadToEnd;
          log.debug("OffsetSyncStore starting - must read to OffsetSync end = {}", initializationMustReadToEnd);
          backingStoreStart();
          readToEnd = true;
      }
  
      // overridable for testing
      void backingStoreStart() {
          backingStore.start(false);
      }
  
      public OptionalLong translateDownstream(String group, TopicPartition sourceTopicPartition, long upstreamOffset) {
          if (!readToEnd) {
              // If we have not read to the end of the syncs topic at least once, decline to translate any offsets.
              // This prevents emitting stale offsets while initially reading the offset syncs topic.
              log.debug("translateDownstream({},{},{}): Skipped (initial offset syncs read still in progress)",
                      group, sourceTopicPartition, upstreamOffset);
              return OptionalLong.empty();
          }
          Optional<OffsetSync> offsetSync = latestOffsetSync(sourceTopicPartition, upstreamOffset);
          if (offsetSync.isPresent()) {
              if (offsetSync.get().upstreamOffset() > upstreamOffset) {
                  // Offset is too far in the past to translate accurately
                  log.debug("translateDownstream({},{},{}): Skipped ({} is ahead of upstream consumer group {})",
                          group, sourceTopicPartition, upstreamOffset,
                          offsetSync.get(), upstreamOffset);
                  return OptionalLong.of(-1L);
              }
              // If the consumer group is ahead of the offset sync, we can translate the upstream offset only 1
              // downstream offset past the offset sync itself. This is because we know that future records must appear
              // ahead of the offset sync, but we cannot estimate how many offsets from the upstream topic
              // will be written vs dropped. If we overestimate, then we may skip the correct offset and have data loss.
              // This also handles consumer groups at the end of a topic whose offsets point past the last valid record.
              // This may cause re-reading of records depending on the age of the offset sync.
              // s=offset sync pair, ?=record may or may not be replicated, g=consumer group offset, r=re-read record
              // source |-s?????r???g-|
              //          |  ______/
              //          | /
              //          vv
              // target |-sg----r-----|
              long upstreamStep = upstreamOffset == offsetSync.get().upstreamOffset() ? 0 : 1;
              log.debug("translateDownstream({},{},{}): Translated {} (relative to {})",
                      group, sourceTopicPartition, upstreamOffset,
                      offsetSync.get().downstreamOffset() + upstreamStep,
                      offsetSync.get()
              );
              return OptionalLong.of(offsetSync.get().downstreamOffset() + upstreamStep);
          } else {
              log.debug("translateDownstream({},{},{}): Skipped (offset sync not found)",
                      group, sourceTopicPartition, upstreamOffset);
              return OptionalLong.empty();
          }
      }
  
      @Override
      public void close() {
          Utils.closeQuietly(backingStore != null ? backingStore::stop : null, "backing store for offset syncs");
          Utils.closeQuietly(admin, "admin client for offset syncs");
      }
  
      protected void handleRecord(ConsumerRecord<byte[], byte[]> record) {
          OffsetSync offsetSync = OffsetSync.deserializeRecord(record);
          TopicPartition sourceTopicPartition = offsetSync.topicPartition();
          offsetSyncs.compute(sourceTopicPartition, (ignored, syncs) ->
                  syncs == null ? createInitialSyncs(offsetSync) : updateExistingSyncs(syncs, offsetSync)
          );
      }
  
      private OffsetSync[] updateExistingSyncs(OffsetSync[] syncs, OffsetSync offsetSync) {
          // Make a copy of the array before mutating it, so that readers do not see inconsistent data
          // TODO: consider batching updates so that this copy can be performed less often for high-volume sync topics.
          OffsetSync[] mutableSyncs = Arrays.copyOf(syncs, SYNCS_PER_PARTITION);
          updateSyncArray(mutableSyncs, syncs, offsetSync);
          if (log.isTraceEnabled()) {
              log.trace("New sync {} applied, new state is {}", offsetSync, offsetArrayToString(mutableSyncs));
          }
          return mutableSyncs;
      }
  
      private String offsetArrayToString(OffsetSync[] syncs) {
          StringBuilder stateString = new StringBuilder();
          stateString.append("[");
          for (int i = 0; i < SYNCS_PER_PARTITION; i++) {
              if (i == 0 || syncs[i] != syncs[i - 1]) {
                  if (i != 0) {
                      stateString.append(",");
                  }
                  // Print only if the sync is interesting, a series of repeated syncs will be elided
                  stateString.append(syncs[i].upstreamOffset());
                  stateString.append(":");
                  stateString.append(syncs[i].downstreamOffset());
              }
          }
          stateString.append("]");
          return stateString.toString();
      }
  
      private OffsetSync[] createInitialSyncs(OffsetSync firstSync) {
          OffsetSync[] syncs = new OffsetSync[SYNCS_PER_PARTITION];
          clearSyncArray(syncs, firstSync);
          return syncs;
      }
  
      private void clearSyncArray(OffsetSync[] syncs, OffsetSync offsetSync) {
          // If every element of the store is the same, then it satisfies invariants B and C trivially.
          for (int i = 0; i < SYNCS_PER_PARTITION; i++) {
              syncs[i] = offsetSync;
          }
      }
  
      private void updateSyncArray(OffsetSync[] syncs, OffsetSync[] original, OffsetSync offsetSync) {
          long upstreamOffset = offsetSync.upstreamOffset();
          // While reading to the end of the topic, ensure that our earliest sync is later than
          // any earlier sync that could have been used for translation, to preserve monotonicity
          // If the upstream offset rewinds, all previous offsets are invalid, so overwrite them all.
          boolean onlyLoadLastOffset = !readToEnd && initializationMustReadToEnd;
          boolean upstreamRewind = upstreamOffset < syncs[0].upstreamOffset();
          if (onlyLoadLastOffset || upstreamRewind) {
              clearSyncArray(syncs, offsetSync);
              return;
          }
          OffsetSync replacement = offsetSync;
          // Index into the original syncs array to the replacement we'll consider next.
          int replacementIndex = 0;
          // Invariant A is always violated once a new sync appears.
          // Repair Invariant A: the latest sync must always be updated
          syncs[0] = replacement;
          for (int current = 1; current < SYNCS_PER_PARTITION; current++) {
              int previous = current - 1;
  
              // Try to choose a value from the old array as the replacement
              // This allows us to keep more distinct values stored overall, improving translation.
              boolean skipOldValue;
              do {
                  OffsetSync oldValue = original[replacementIndex];
                  // If oldValue is not recent enough, then it is not valid to use at the current index.
                  // It may still be valid when used in a later index where values are allowed to be older.
                  boolean isRecent = invariantB(syncs[previous], oldValue, previous, current);
                  // Ensure that this value is sufficiently separated from the previous value
                  // We prefer to keep more recent syncs of similar precision (i.e. the value in replacement)
                  // If this value is too close to the previous value, it will never be valid in a later position.
                  boolean separatedFromPrevious = invariantC(syncs[previous], oldValue, previous);
                  // Ensure that this value is sufficiently separated from the next value
                  // We prefer to keep existing syncs of lower precision (i.e. the value in syncs[next])
                  int next = current + 1;
                  boolean separatedFromNext = next >= SYNCS_PER_PARTITION || invariantC(oldValue, syncs[next], current);
                  // If the next value in the old array is a duplicate of the current one, then they are equivalent
                  // This value will not need to be considered again
                  int nextReplacement = replacementIndex + 1;
                  boolean duplicate = nextReplacement < SYNCS_PER_PARTITION && oldValue == original[nextReplacement];
  
                  // Promoting the oldValue to the replacement only happens if it satisfies all invariants.
                  boolean promoteOldValueToReplacement = isRecent && separatedFromPrevious && separatedFromNext;
                  if (promoteOldValueToReplacement) {
                      replacement = oldValue;
                  }
                  // The index should be skipped without promoting if we know that it will not be used at a later index
                  // based only on the observed part of the array so far.
                  skipOldValue = duplicate || !separatedFromPrevious;
                  if (promoteOldValueToReplacement || skipOldValue) {
                      replacementIndex++;
                  }
                  // We may need to skip past multiple indices, so keep looping until we're done skipping forward.
                  // The index must not get ahead of the current index, as we only promote from low index to high index.
              } while (replacementIndex < current && skipOldValue);
  
              // The replacement variable always contains a value which satisfies the invariants for this index.
              // This replacement may or may not be used, since the invariants could already be satisfied,
              // and in that case, prefer to keep the existing tail of the syncs array rather than updating it.
              assert invariantB(syncs[previous], replacement, previous, current);
              assert invariantC(syncs[previous], replacement, previous);
  
              // Test if changes to the previous index affected the invariant for this index
              if (invariantB(syncs[previous], syncs[current], previous, current)) {
                  // Invariant B holds for syncs[current]: it must also hold for all later values
                  break;
              } else {
                  // Invariant B violated for syncs[current]: sync is now too old and must be updated
                  // Repair Invariant B: swap in replacement
                  syncs[current] = replacement;
  
                  assert invariantB(syncs[previous], syncs[current], previous, current);
                  assert invariantC(syncs[previous], syncs[current], previous);
              }
          }
      }
  
      private boolean invariantB(OffsetSync iSync, OffsetSync jSync, int i, int j) {
          long bound = jSync.upstreamOffset() + (1L << j) - (1L << i);
          return iSync == jSync || bound < 0 || iSync.upstreamOffset() <= bound;
      }
  
      private boolean invariantC(OffsetSync iSync, OffsetSync jSync, int i) {
          long bound = jSync.upstreamOffset() + (1L << Math.max(i - 2, 0));
          return iSync == jSync || (bound >= 0 && iSync.upstreamOffset() >= bound);
      }
  
      private Optional<OffsetSync> latestOffsetSync(TopicPartition topicPartition, long upstreamOffset) {
          return Optional.ofNullable(offsetSyncs.get(topicPartition))
                  .map(syncs -> lookupLatestSync(syncs, upstreamOffset));
      }
  
  
      private OffsetSync lookupLatestSync(OffsetSync[] syncs, long upstreamOffset) {
          // linear search the syncs, effectively a binary search over the topic offsets
          // Search from latest to earliest to find the sync that gives the best accuracy
          for (int i = 0; i < SYNCS_PER_PARTITION; i++) {
              OffsetSync offsetSync = syncs[i];
              if (offsetSync.upstreamOffset() <= upstreamOffset) {
                  return offsetSync;
              }
          }
          return syncs[SYNCS_PER_PARTITION - 1];
      }
  
      // For testing
      OffsetSync syncFor(TopicPartition topicPartition, int syncIdx) {
          OffsetSync[] syncs = offsetSyncs.get(topicPartition);
          if (syncs == null)
              throw new IllegalArgumentException("No syncs present for " + topicPartition);
          if (syncIdx >= syncs.length)
              throw new IllegalArgumentException(
                      "Requested sync " + (syncIdx + 1) + " for " + topicPartition
                              + " but there are only " + syncs.length + " syncs available for that topic partition"
              );
          return syncs[syncIdx];
      }
  }
  