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.storage;
  
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.connect.runtime.WorkerConfig;
  import org.apache.kafka.connect.util.Callback;
  import org.apache.kafka.connect.util.LoggingContext;
  import org.apache.kafka.connect.util.TopicAdmin;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.nio.ByteBuffer;
  import java.time.Duration;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.Future;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.function.Supplier;
  
  /**
   * An {@link OffsetBackingStore} with support for reading from and writing to a worker-global
   * offset backing store and/or a connector-specific offset backing store.
   */
  public class ConnectorOffsetBackingStore implements OffsetBackingStore {
  
      private static final Logger log = LoggerFactory.getLogger(ConnectorOffsetBackingStore.class);
  
      /**
       * Builds an offset store that uses a connector-specific offset topic as the primary store and
       * the worker-global offset store as the secondary store.
       *
       * @param loggingContext a {@link Supplier} for the {@link LoggingContext} that should be used
       *                       for messages logged by this offset store; may not be null, and may never return null
       * @param workerStore the worker-global offset store; may not be null
       * @param connectorStore the connector-specific offset store; may not be null
       * @param connectorOffsetsTopic the name of the connector-specific offset topic; may not be null
       * @param connectorStoreAdmin the topic admin to use for the connector-specific offset topic; may not be null
       * @return an offset store backed primarily by the connector-specific offset topic and secondarily
       * by the worker-global offset store; never null
       */
      public static ConnectorOffsetBackingStore withConnectorAndWorkerStores(
              Supplier<LoggingContext> loggingContext,
              OffsetBackingStore workerStore,
              KafkaOffsetBackingStore connectorStore,
              String connectorOffsetsTopic,
              TopicAdmin connectorStoreAdmin
      ) {
          Objects.requireNonNull(loggingContext);
          Objects.requireNonNull(workerStore);
          Objects.requireNonNull(connectorStore);
          Objects.requireNonNull(connectorOffsetsTopic);
          Objects.requireNonNull(connectorStoreAdmin);
          return new ConnectorOffsetBackingStore(
                  Time.SYSTEM,
                  loggingContext,
                  connectorOffsetsTopic,
                  workerStore,
                  connectorStore,
                  connectorStoreAdmin
          );
      }
  
      /**
       * Builds an offset store that uses the worker-global offset store as the primary store, and no secondary store.
       *
       * @param loggingContext a {@link Supplier} for the {@link LoggingContext} that should be used
       *                       for messages logged by this offset store; may not be null, and may never return null
       * @param workerStore the worker-global offset store; may not be null
       * @param workerOffsetsTopic the name of the worker-global offset topic; may be null if the worker
       *                           does not use an offset topic for its offset store
       * @return an offset store for the connector backed solely by the worker-global offset store; never null
       */
      public static ConnectorOffsetBackingStore withOnlyWorkerStore(
              Supplier<LoggingContext> loggingContext,
              OffsetBackingStore workerStore,
              String workerOffsetsTopic
      ) {
          Objects.requireNonNull(loggingContext);
          Objects.requireNonNull(workerStore);
          return new ConnectorOffsetBackingStore(Time.SYSTEM, loggingContext, workerOffsetsTopic, workerStore, null, null);
      }
  
      /**
       * Builds an offset store that uses a connector-specific offset topic as the primary store, and no secondary store.
       *
       * @param loggingContext a {@link Supplier} for the {@link LoggingContext} that should be used
       *                       for messages logged by this offset store; may not be null, and may never return null
       * @param connectorStore the connector-specific offset store; may not be null
       * @param connectorOffsetsTopic the name of the connector-specific offset topic; may not be null
       * @param connectorStoreAdmin the topic admin to use for the connector-specific offset topic; may not be null
       * @return an offset store for the connector backed solely by the connector-specific offset topic; never null
       */
      public static ConnectorOffsetBackingStore withOnlyConnectorStore(
              Supplier<LoggingContext> loggingContext,
              KafkaOffsetBackingStore connectorStore,
              String connectorOffsetsTopic,
              TopicAdmin connectorStoreAdmin
      ) {
          Objects.requireNonNull(loggingContext);
          Objects.requireNonNull(connectorOffsetsTopic);
          Objects.requireNonNull(connectorStoreAdmin);
          return new ConnectorOffsetBackingStore(
                  Time.SYSTEM,
                  loggingContext,
                  connectorOffsetsTopic,
                  null,
                  connectorStore,
                  connectorStoreAdmin
          );
      }
  
      private final Time time;
      private final Supplier<LoggingContext> loggingContext;
      private final String primaryOffsetsTopic;
      private final Optional<OffsetBackingStore> workerStore;
      private final Optional<KafkaOffsetBackingStore> connectorStore;
      private final Optional<TopicAdmin> connectorStoreAdmin;
  
      ConnectorOffsetBackingStore(
              Time time,
              Supplier<LoggingContext> loggingContext,
              String primaryOffsetsTopic,
              OffsetBackingStore workerStore,
              KafkaOffsetBackingStore connectorStore,
              TopicAdmin connectorStoreAdmin
      ) {
          if (workerStore == null && connectorStore == null) {
              throw new IllegalArgumentException("At least one non-null offset store must be provided");
          }
          this.time = time;
          this.loggingContext = loggingContext;
          this.primaryOffsetsTopic = primaryOffsetsTopic;
          this.workerStore = Optional.ofNullable(workerStore);
          this.connectorStore = Optional.ofNullable(connectorStore);
          this.connectorStoreAdmin = Optional.ofNullable(connectorStoreAdmin);
      }
  
      public String primaryOffsetsTopic() {
          return primaryOffsetsTopic;
      }
  
      /**
       * If configured to use a connector-specific offset store, {@link OffsetBackingStore#start() start} that store.
       *
       * <p>The worker-global offset store is not modified; it is the caller's responsibility to ensure that it is started
       * before calls to {@link #get(Collection)} and {@link #set(Map, Callback)} take place.
       */
      @Override
      public void start() {
          // Worker offset store should already be started
          connectorStore.ifPresent(OffsetBackingStore::start);
      }
  
      /**
       * If configured to use a connector-specific offset store, {@link OffsetBackingStore#stop() stop} that store,
       * and {@link TopicAdmin#close(Duration) close} the topic admin used by that store.
       *
       * <p>The worker-global offset store is not modified as it may be used for other connectors that either already exist,
       * or will be created, on this worker.
       */
      @Override
      public void stop() {
          // Worker offset store should not be stopped as it may be used for multiple connectors
          connectorStore.ifPresent(OffsetBackingStore::stop);
          connectorStoreAdmin.ifPresent(TopicAdmin::close);
      }
  
      /**
       * Get the offset values for the specified keys.
       *
       * <p>If configured to use a connector-specific offset store, priority is given to the values contained in that store,
       * and the values in the worker-global offset store (if one is provided) are used as a fallback for keys that are not
       * present in the connector-specific store.
       *
       * <p>If not configured to use a connector-specific offset store, only the values contained in the worker-global
       * offset store are returned.
  
       * @param keys list of keys to look up
       * @return future for the resulting map from key to value
       */
      @Override
      public Future<Map<ByteBuffer, ByteBuffer>> get(Collection<ByteBuffer> keys) {
          Future<Map<ByteBuffer, ByteBuffer>> workerGetFuture = getFromStore(workerStore, keys);
          Future<Map<ByteBuffer, ByteBuffer>> connectorGetFuture = getFromStore(connectorStore, keys);
  
          return new Future<>() {
              @Override
              public boolean cancel(boolean mayInterruptIfRunning) {
                  // Note the use of | instead of || here; this causes cancel to be invoked on both futures,
                  // even if the first call to cancel returns true
                  return workerGetFuture.cancel(mayInterruptIfRunning)
                          | connectorGetFuture.cancel(mayInterruptIfRunning);
              }
  
              @Override
              public boolean isCancelled() {
                  return workerGetFuture.isCancelled()
                          || connectorGetFuture.isCancelled();
              }
  
              @Override
              public boolean isDone() {
                  return workerGetFuture.isDone()
                          && connectorGetFuture.isDone();
              }
  
              @Override
              public Map<ByteBuffer, ByteBuffer> get() throws InterruptedException, ExecutionException {
                  Map<ByteBuffer, ByteBuffer> result = new HashMap<>(workerGetFuture.get());
                  result.putAll(connectorGetFuture.get());
                  return result;
              }
  
              @Override
              public Map<ByteBuffer, ByteBuffer> get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
                  long timeoutMs = unit.toMillis(timeout);
                  long endTime = time.milliseconds() + timeoutMs;
                  Map<ByteBuffer, ByteBuffer> result = new HashMap<>(workerGetFuture.get(timeoutMs, unit));
                  timeoutMs = Math.max(1, endTime - time.milliseconds());
                  result.putAll(connectorGetFuture.get(timeoutMs, TimeUnit.MILLISECONDS));
                  return result;
              }
          };
      }
  
      /**
       * Store the specified offset key/value pairs.
       *
       * <p>If configured to use a connector-specific offset store, the returned {@link Future} corresponds to a
       * write to that store, and the passed-in {@link Callback} is invoked once that write completes. If a worker-global
       * store is provided, a secondary write is made to that store if the write to the connector-specific store
       * succeeds.
       * <p>
       * Normally, errors with this secondary write are not reflected in the returned {@link Future} or the passed-in
       * {@link Callback}; they are only logged as a warning to users. The only exception to this rule is when the
       * offsets that need to be committed contain tombstone records.
       * <p>When the to-be-committed offsets contain tombstones, offset commits take place in three phases:
       * <ol>
       *     <li>First, only the tombstone offsets are written to the worker-global store. Failures during this step will
       *     be reflected in the returned {@link Future} and reported to the passed-in {@link Callback}.</li>
       *     <li>If and only if the previous write to the worker-global store succeeded, all offsets (both tombstones and
       *     non-tombstones) are written to the connector-specific store. Failures during this step will also be
       *     reflected in the returned {@link Future} and reported to the passed-in {@link Callback}.</li>
       *     <li>Finally, if and only if the previous write to the connector-specific store succeeded, all offsets with
       *     non-tombstone values are written to the worker-global store. Failures during this step will only be reported
       *     as warning log messages, and will not be reflected in the returned {@link Future} or reported to the
       *     passed-in {@link Callback}.</li>
       * </ol>
       *
       * <p>If not configured to use a connector-specific offset store, the returned {@link Future} corresponds to a
       * write to the worker-global offset store, and the passed-in {@link Callback} is invoked once that write completes.
  
       * @param values map from key to value
       * @param callback callback to invoke on completion of the primary write
       * @return void future for the primary write
      *
       * @see <a href="https://issues.apache.org/jira/browse/KAFKA-15018">KAFKA-15018</a> for context on the three-step
       * write sequence
       *
       */
      @Override
      public Future<Void> set(Map<ByteBuffer, ByteBuffer> values, Callback<Void> callback) {
          final OffsetBackingStore primaryStore;
          final OffsetBackingStore secondaryStore;
          if (connectorStore.isPresent()) {
              primaryStore = connectorStore.get();
              secondaryStore = workerStore.orElse(null);
          } else if (workerStore.isPresent()) {
              primaryStore = workerStore.get();
              secondaryStore = null;
          } else {
              // Should never happen since we check for this case in the constructor, but just in case, this should
              // be more informative than the NPE that would otherwise be thrown
              throw new IllegalStateException("At least one non-null offset store must be provided");
          }
  
          Map<ByteBuffer, ByteBuffer> regularOffsets = new HashMap<>();
          Map<ByteBuffer, ByteBuffer> tombstoneOffsets = new HashMap<>();
          values.forEach((partition, offset) -> {
              if (offset == null) {
                  tombstoneOffsets.put(partition, null);
              } else {
                  regularOffsets.put(partition, offset);
              }
          });
  
          if (secondaryStore != null && !tombstoneOffsets.isEmpty()) {
              return new ChainedOffsetWriteFuture(
                  primaryStore,
                  secondaryStore,
                  values,
                  regularOffsets,
                  tombstoneOffsets,
                  callback
              );
          } else {
              return setPrimaryThenSecondary(primaryStore, secondaryStore, values, regularOffsets, callback);
          }
      }
  
      private Future<Void> setPrimaryThenSecondary(
          OffsetBackingStore primaryStore,
          OffsetBackingStore secondaryStore,
          Map<ByteBuffer, ByteBuffer> completeOffsets,
          Map<ByteBuffer, ByteBuffer> nonTombstoneOffsets,
          Callback<Void> callback
      ) {
          return primaryStore.set(completeOffsets, (primaryWriteError, ignored) -> {
              if (secondaryStore != null) {
                  if (primaryWriteError != null) {
                      log.trace("Skipping offsets write to secondary store because primary write has failed", primaryWriteError);
                  } else {
                      try {
                          // Invoke OffsetBackingStore::set but ignore the resulting future; we don't block on writes to this
                          // backing store.
                          secondaryStore.set(nonTombstoneOffsets, (secondaryWriteError, ignored2) -> {
                              try (LoggingContext context = loggingContext()) {
                                  if (secondaryWriteError != null) {
                                      log.warn("Failed to write offsets to secondary backing store", secondaryWriteError);
                                  } else {
                                      log.debug("Successfully flushed offsets to secondary backing store");
                                  }
                              }
                          });
                      } catch (Exception e) {
                          log.warn("Failed to write offsets to secondary backing store", e);
                      }
                  }
              }
              try (LoggingContext context = loggingContext()) {
                  callback.onCompletion(primaryWriteError, ignored);
              }
          });
      }
  
      @Override
      public Set<Map<String, Object>> connectorPartitions(String connectorName) {
          Set<Map<String, Object>> partitions = new HashSet<>();
          workerStore.ifPresent(offsetBackingStore -> partitions.addAll(offsetBackingStore.connectorPartitions(connectorName)));
          connectorStore.ifPresent(offsetBackingStore -> partitions.addAll(offsetBackingStore.connectorPartitions(connectorName)));
          return partitions;
      }
  
      /**
       * If configured to use a connector-specific offset store,
       * {@link OffsetBackingStore#configure(WorkerConfig) configure} that store.
       *
       * <p>The worker-global offset store is not modified; it is the caller's responsibility to ensure that it is configured
       * before calls to {@link #start()}, {@link #get(Collection)} and {@link #set(Map, Callback)} take place.
       */
      @Override
      public void configure(WorkerConfig config) {
          // Worker offset store should already be configured
          connectorStore.ifPresent(store -> store.configure(config));
      }
  
      // For testing
      public boolean hasConnectorSpecificStore() {
          return connectorStore.isPresent();
      }
  
      // For testing
      public boolean hasWorkerGlobalStore() {
          return workerStore.isPresent();
      }
  
      private LoggingContext loggingContext() {
          LoggingContext result = loggingContext.get();
          Objects.requireNonNull(result);
          return result;
      }
  
      private static Future<Map<ByteBuffer, ByteBuffer>> getFromStore(Optional<? extends OffsetBackingStore> store, Collection<ByteBuffer> keys) {
          return store.map(s -> s.get(keys)).orElseGet(() -> CompletableFuture.completedFuture(Collections.emptyMap()));
      }
  
      private class ChainedOffsetWriteFuture implements Future<Void> {
  
          private final OffsetBackingStore primaryStore;
          private final OffsetBackingStore secondaryStore;
          private final Map<ByteBuffer, ByteBuffer> completeOffsets;
          private final Map<ByteBuffer, ByteBuffer> regularOffsets;
          private final Callback<Void> callback;
          private final AtomicReference<Throwable> writeError;
          private final CountDownLatch completed;
  
          public ChainedOffsetWriteFuture(
              OffsetBackingStore primaryStore,
              OffsetBackingStore secondaryStore,
              Map<ByteBuffer, ByteBuffer> completeOffsets,
              Map<ByteBuffer, ByteBuffer> regularOffsets,
              Map<ByteBuffer, ByteBuffer> tombstoneOffsets,
              Callback<Void> callback
          ) {
              this.primaryStore = primaryStore;
              this.secondaryStore = secondaryStore;
              this.completeOffsets = completeOffsets;
              this.regularOffsets = regularOffsets;
              this.callback = callback;
              this.writeError = new AtomicReference<>();
              this.completed = new CountDownLatch(1);
  
              secondaryStore.set(tombstoneOffsets, this::onFirstWrite);
          }
  
          private void onFirstWrite(Throwable error, Void ignored) {
              if (error != null) {
                  log.trace("Skipping offsets write to primary store because secondary tombstone write has failed", error);
                  try (LoggingContext context = loggingContext()) {
                      callback.onCompletion(error, ignored);
                      writeError.compareAndSet(null, error);
                      completed.countDown();
                  }
                  return;
              }
              setPrimaryThenSecondary(primaryStore, secondaryStore, completeOffsets, regularOffsets, this::onSecondWrite);
          }
  
          private void onSecondWrite(Throwable error, Void ignored) {
              callback.onCompletion(error, ignored);
              writeError.compareAndSet(null, error);
              completed.countDown();
          }
  
          @Override
          public boolean cancel(boolean mayInterruptIfRunning) {
              return false;
          }
  
          @Override
          public boolean isCancelled() {
              return false;
          }
  
          @Override
          public boolean isDone() {
              return completed.getCount() == 0;
          }
  
          @Override
          public Void get() throws InterruptedException, ExecutionException {
              completed.await();
              if (writeError.get() != null) {
                  throw new ExecutionException(writeError.get());
              }
              return null;
          }
  
          @Override
          public Void get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
              if (!completed.await(timeout, unit)) {
                  throw new TimeoutException("Failed to complete offset write in time");
              }
              if (writeError.get() != null) {
                  throw new ExecutionException(writeError.get());
              }
              return null;
          }
      }
  
  }
  