/*
 * 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.errors.ConnectException;
import org.apache.kafka.connect.util.Callback;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Future;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * <p>
 * OffsetStorageWriter is a buffered writer that wraps the simple {@link OffsetBackingStore} interface.
 * It maintains a copy of the key-value data in memory and buffers writes. It allows you to take
 * a snapshot, which can then be asynchronously flushed to the backing store while new writes
 * continue to be processed. This allows Kafka Connect to process offset commits in the background
 * while continuing to process messages.
 * </p>
 * <p>
 * Connect uses an OffsetStorage implementation to save state about the current progress of
 * source (import to Kafka) connectors, which may have many input partitions and "offsets" may not be as
 * simple as they are for Kafka partitions or files. Offset storage is not required for sink connectors
 * because they can use Kafka's native offset storage (or the sink data store can handle offset
 * storage to achieve exactly once semantics).
 * </p>
 * <p>
 * Both partitions and offsets are generic data objects. This allows different connectors to use
 * whatever representation they need, even arbitrarily complex records. These are translated
 * internally into the serialized form the OffsetBackingStore uses.
 * </p>
 * <p>
 * Note that this only provides write functionality. This is intentional to ensure stale data is
 * never read. Offset data should only be read during startup or reconfiguration of a task. By
 * always serving those requests by reading the values from the backing store, we ensure we never
 * accidentally use stale data. (One example of how this can occur: a task is processing input
 * partition A, writing offsets; reconfiguration causes partition A to be reassigned elsewhere;
 * reconfiguration causes partition A to be reassigned to this node, but now the offset data is out
 * of date). Since these offsets are created and managed by the connector itself, there's no way
 * for the offset management layer to know which keys are "owned" by which tasks at any given
 * time.
 * </p>
 * <p>
 * This class is thread-safe.
 * </p>
 */
public class OffsetStorageWriter {
    private static final Logger log = LoggerFactory.getLogger(OffsetStorageWriter.class);

    private final OffsetBackingStore backingStore;
    private final Converter keyConverter;
    private final Converter valueConverter;
    private final String namespace;
    // Offset data in Connect format
    private Map<Map<String, Object>, Map<String, Object>> data = new HashMap<>();

    private Map<Map<String, Object>, Map<String, Object>> toFlush = null;
    private final Semaphore flushInProgress = new Semaphore(1);
    // Unique ID for each flush request to handle callbacks after timeouts
    private long currentFlushId = 0;

    public OffsetStorageWriter(OffsetBackingStore backingStore,
                               String namespace, Converter keyConverter, Converter valueConverter) {
        this.backingStore = backingStore;
        this.namespace = namespace;
        this.keyConverter = keyConverter;
        this.valueConverter = valueConverter;
    }

    /**
     * Set an offset for a partition using Connect data values
     * @param partition the partition to store an offset for
     * @param offset the offset
     */
    @SuppressWarnings("unchecked")
    public synchronized void offset(Map<String, ?> partition, Map<String, ?> offset) {
        data.put((Map<String, Object>) partition, (Map<String, Object>) offset);
    }

    private boolean flushing() {
        return toFlush != null;
    }

    /**
     * Performs the first step of a flush operation, snapshotting the current state. This does not
     * actually initiate the flush with the underlying storage. Ensures that any previous flush operations
     * have finished before beginning a new flush.
     *
     * @return true if a flush was initiated, false if no data was available
     * @throws ConnectException if the previous flush is not complete before this method is called
     */
    public boolean beginFlush() {
        try {
            return beginFlush(0, TimeUnit.NANOSECONDS);
        } catch (InterruptedException | TimeoutException e) {
            log.error("Invalid call to OffsetStorageWriter beginFlush() while already flushing, the "
                    + "framework should not allow this");
            throw new ConnectException("OffsetStorageWriter is already flushing");
        }
    }

    /**
     * Performs the first step of a flush operation, snapshotting the current state. This does not
     * actually initiate the flush with the underlying storage. Ensures that any previous flush operations
     * have finished before beginning a new flush.
     * <p>If and only if this method returns true, the caller must call {@link #doFlush(Callback)}
     * or {@link #cancelFlush()} to finish the flush operation and allow later calls to complete.
     *
     * @param timeout A maximum duration to wait for previous flushes to finish before giving up on waiting
     * @param timeUnit Units of the timeout argument
     * @return true if a flush was initiated, false if no data was available
     * @throws InterruptedException if this thread was interrupted while waiting for the previous flush to complete
     * @throws TimeoutException if the {@code timeout} elapses before previous flushes are complete.
     */
    public boolean beginFlush(long timeout, TimeUnit timeUnit) throws InterruptedException, TimeoutException {
        if (flushInProgress.tryAcquire(Math.max(0, timeout), timeUnit)) {
            synchronized (this) {
                if (data.isEmpty()) {
                    flushInProgress.release();
                    return false;
                } else {
                    toFlush = data;
                    data = new HashMap<>();
                    return true;
                }
            }
        } else {
            throw new TimeoutException("Timed out waiting for previous flush to finish");
        }
    }

    /**
     * Flush the current offsets and clear them from this writer. This is non-blocking: it
     * moves the current set of offsets out of the way, serializes the data, and asynchronously
     * writes the data to the backing store. If no offsets need to be written, the callback is
     * still invoked, but no Future is returned.
     *
     * @return a Future, or null if there are no offsets to commit
     */
    public Future<Void> doFlush(final Callback<Void> callback) {

        final long flushId;
        // Serialize
        final Map<ByteBuffer, ByteBuffer> offsetsSerialized;

        synchronized (this) {
            flushId = currentFlushId;

            try {
                offsetsSerialized = new HashMap<>(toFlush.size());
                for (Map.Entry<Map<String, Object>, Map<String, Object>> entry : toFlush.entrySet()) {
                    // Offsets are specified as schemaless to the converter, using whatever internal schema is appropriate
                    // for that data. The only enforcement of the format is here.
                    OffsetUtils.validateFormat(entry.getKey());
                    OffsetUtils.validateFormat(entry.getValue());
                    // When serializing the key, we add in the namespace information so the key is [namespace, real key]
                    byte[] key = keyConverter.fromConnectData(namespace, null, Arrays.asList(namespace, entry.getKey()));
                    ByteBuffer keyBuffer = (key != null) ? ByteBuffer.wrap(key) : null;
                    byte[] value = valueConverter.fromConnectData(namespace, null, entry.getValue());
                    ByteBuffer valueBuffer = (value != null) ? ByteBuffer.wrap(value) : null;
                    offsetsSerialized.put(keyBuffer, valueBuffer);
                }
            } catch (Throwable t) {
                // Must handle errors properly here or the writer will be left mid-flush forever and be
                // unable to make progress.
                log.error("CRITICAL: Failed to serialize offset data, making it impossible to commit "
                        + "offsets under namespace {}. This likely won't recover unless the "
                        + "unserializable partition or offset information is overwritten.", namespace);
                log.error("Cause of serialization failure:", t);
                callback.onCompletion(t, null);
                return null;
            }

            // And submit the data
            log.debug("Submitting {} entries to backing store. The offsets are: {}", offsetsSerialized.size(), toFlush);
        }

        return backingStore.set(offsetsSerialized, (error, result) -> {
            boolean isCurrent = handleFinishWrite(flushId, error, result);
            if (isCurrent && callback != null) {
                callback.onCompletion(error, result);
            }
        });
    }

    /**
     * Cancel a flush that has been initiated by {@link #beginFlush}. This should not be called if
     * {@link #doFlush} has already been invoked. It should be used if an operation performed
     * between beginFlush and doFlush failed.
     */
    public synchronized void cancelFlush() {
        // Verify we're still flushing data to handle a race between cancelFlush() calls from up the
        // call stack and callbacks from the write request to underlying storage
        if (flushing()) {
            // Just recombine the data and place it back in the primary storage
            toFlush.putAll(data);
            data = toFlush;
            currentFlushId++;
            flushInProgress.release();
            toFlush = null;
        }
    }

    /**
     * Handle completion of a write. Returns true if this callback is for the current flush
     * operation, false if it's for an old operation that should now be ignored.
     */
    private synchronized boolean handleFinishWrite(long flushId, Throwable error, Void result) {
        // Callbacks need to be handled carefully since the flush operation may have already timed
        // out and been cancelled.
        if (flushId != currentFlushId)
            return false;

        if (error != null) {
            cancelFlush();
        } else {
            currentFlushId++;
            flushInProgress.release();
            toFlush = null;
        }
        return true;
    }
}
