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.connect.data.SchemaAndValue;
  import org.apache.kafka.connect.errors.ConnectException;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.nio.ByteBuffer;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.CancellationException;
  import java.util.concurrent.Future;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  /**
   * Implementation of OffsetStorageReader. Unlike {@link OffsetStorageWriter} which is implemented
   * directly, the interface is only separate from this implementation because it needs to be
   * included in the public API package.
   */
  public class OffsetStorageReaderImpl implements CloseableOffsetStorageReader {
      private static final Logger log = LoggerFactory.getLogger(OffsetStorageReaderImpl.class);
  
      private final OffsetBackingStore backingStore;
      private final String namespace;
      private final Converter keyConverter;
      private final Converter valueConverter;
      private final AtomicBoolean closed;
      private final Set<Future<Map<ByteBuffer, ByteBuffer>>> offsetReadFutures;
  
      public OffsetStorageReaderImpl(OffsetBackingStore backingStore, String namespace,
                                     Converter keyConverter, Converter valueConverter) {
          this.backingStore = backingStore;
          this.namespace = namespace;
          this.keyConverter = keyConverter;
          this.valueConverter = valueConverter;
          this.closed = new AtomicBoolean(false);
          this.offsetReadFutures = new HashSet<>();
      }
  
      @Override
      public <T> Map<String, Object> offset(Map<String, T> partition) {
          return offsets(Collections.singletonList(partition)).get(partition);
      }
  
      @Override
      @SuppressWarnings("unchecked")
      public <T> Map<Map<String, T>, Map<String, Object>> offsets(Collection<Map<String, T>> partitions) {
          // Serialize keys so backing store can work with them
          Map<ByteBuffer, Map<String, T>> serializedToOriginal = new HashMap<>(partitions.size());
          for (Map<String, T> key : partitions) {
              try {
                  // Offsets are treated as schemaless, their format is only validated here (and the returned value below)
                  OffsetUtils.validateFormat(key);
                  byte[] keySerialized = keyConverter.fromConnectData(namespace, null, Arrays.asList(namespace, key));
                  ByteBuffer keyBuffer = (keySerialized != null) ? ByteBuffer.wrap(keySerialized) : null;
                  serializedToOriginal.put(keyBuffer, key);
              } catch (Throwable t) {
                  log.error("CRITICAL: Failed to serialize partition key when getting offsets for task with "
                          + "namespace {}. No value for this data will be returned, which may break the "
                          + "task or cause it to skip some data.", namespace, t);
              }
          }
  
          // Get serialized key -> serialized value from backing store
          Map<ByteBuffer, ByteBuffer> raw;
          try {
              Future<Map<ByteBuffer, ByteBuffer>> offsetReadFuture;
              synchronized (offsetReadFutures) {
                  if (closed.get()) {
                      throw new ConnectException(
                          "Offset reader is closed. This is likely because the task has already been "
                              + "scheduled to stop but has taken longer than the graceful shutdown "
                              + "period to do so.");
                  }
                  offsetReadFuture = backingStore.get(serializedToOriginal.keySet());
                  offsetReadFutures.add(offsetReadFuture);
              }
  
              try {
                  raw = offsetReadFuture.get();
              } catch (CancellationException e) {
                  throw new ConnectException(
                      "Offset reader closed while attempting to read offsets. This is likely because "
                          + "the task was been scheduled to stop but has taken longer than the "
                          + "graceful shutdown period to do so.");
              } finally {
                  synchronized (offsetReadFutures) {
                      offsetReadFutures.remove(offsetReadFuture);
                  }
              }
          } catch (Exception e) {
              log.error("Failed to fetch offsets from namespace {}: ", namespace, e);
              throw new ConnectException("Failed to fetch offsets.", e);
          }
  
          // Deserialize all the values and map back to the original keys
          Map<Map<String, T>, Map<String, Object>> result = new HashMap<>(partitions.size());
          for (Map.Entry<ByteBuffer, ByteBuffer> rawEntry : raw.entrySet()) {
              try {
                  // Since null could be a valid key, explicitly check whether map contains the key
                  if (!serializedToOriginal.containsKey(rawEntry.getKey())) {
                      log.error("Should be able to map {} back to a requested partition-offset key, backing "
                              + "store may have returned invalid data", rawEntry.getKey());
                      continue;
                  }
                  Map<String, T> origKey = serializedToOriginal.get(rawEntry.getKey());
                  SchemaAndValue deserializedSchemaAndValue = valueConverter.toConnectData(namespace, rawEntry.getValue() != null ? rawEntry.getValue().array() : null);
                  Object deserializedValue = deserializedSchemaAndValue.value();
                  OffsetUtils.validateFormat(deserializedValue);
  
                  result.put(origKey, (Map<String, Object>) deserializedValue);
              } catch (Throwable t) {
                  log.error("CRITICAL: Failed to deserialize offset data when getting offsets for task with"
                          + " namespace {}. No value for this data will be returned, which may break the "
                          + "task or cause it to skip some data. This could either be due to an error in "
                          + "the connector implementation or incompatible schema.", namespace, t);
              }
          }
  
          return result;
      }
  
      @Override
      public void close() {
          if (!closed.getAndSet(true)) {
              synchronized (offsetReadFutures) {
                  for (Future<Map<ByteBuffer, ByteBuffer>> offsetReadFuture : offsetReadFutures) {
                      try {
                          offsetReadFuture.cancel(true);
                      } catch (Throwable t) {
                          log.error("Failed to cancel offset read future", t);
                      }
                  }
                  offsetReadFutures.clear();
              }
          }
      }
  }
  