/*
 * 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.runtime;

import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.errors.InvalidProducerEpochException;
import org.apache.kafka.common.internals.Plugin;
import org.apache.kafka.common.metrics.Sensor;
import org.apache.kafka.common.metrics.stats.Avg;
import org.apache.kafka.common.metrics.stats.Max;
import org.apache.kafka.common.metrics.stats.Min;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.runtime.errors.ErrorHandlingMetrics;
import org.apache.kafka.connect.runtime.errors.ErrorReporter;
import org.apache.kafka.connect.runtime.errors.ProcessingContext;
import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
import org.apache.kafka.connect.runtime.isolation.LoaderSwap;
import org.apache.kafka.connect.source.SourceRecord;
import org.apache.kafka.connect.source.SourceTask;
import org.apache.kafka.connect.source.SourceTask.TransactionBoundary;
import org.apache.kafka.connect.storage.CloseableOffsetStorageReader;
import org.apache.kafka.connect.storage.ClusterConfigState;
import org.apache.kafka.connect.storage.ConnectorOffsetBackingStore;
import org.apache.kafka.connect.storage.Converter;
import org.apache.kafka.connect.storage.HeaderConverter;
import org.apache.kafka.connect.storage.OffsetStorageWriter;
import org.apache.kafka.connect.storage.StatusBackingStore;
import org.apache.kafka.connect.util.ConnectorTaskId;
import org.apache.kafka.connect.util.LoggingContext;
import org.apache.kafka.connect.util.TopicAdmin;
import org.apache.kafka.connect.util.TopicCreationGroup;

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

import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Supplier;


/**
 * WorkerTask that uses a {@link SourceTask} to ingest data into Kafka, with support for exactly-once semantics.
 */
class ExactlyOnceWorkerSourceTask extends AbstractWorkerSourceTask {
    private static final Logger log = LoggerFactory.getLogger(ExactlyOnceWorkerSourceTask.class);

    private boolean transactionOpen;
    private final LinkedHashMap<SourceRecord, RecordMetadata> committableRecords;

    private final TransactionBoundaryManager transactionBoundaryManager;
    private final TransactionMetricsGroup transactionMetrics;

    private final Runnable preProducerCheck;
    private final Runnable postProducerCheck;

    public ExactlyOnceWorkerSourceTask(ConnectorTaskId id,
                                       SourceTask task,
                                       TaskStatus.Listener statusListener,
                                       TargetState initialState,
                                       Plugin<Converter> keyConverterPlugin,
                                       Plugin<Converter> valueConverterPlugin,
                                       Plugin<HeaderConverter> headerConverterPlugin,
                                       TransformationChain<SourceRecord, SourceRecord> transformationChain,
                                       Producer<byte[], byte[]> producer,
                                       TopicAdmin admin,
                                       Map<String, TopicCreationGroup> topicGroups,
                                       CloseableOffsetStorageReader offsetReader,
                                       OffsetStorageWriter offsetWriter,
                                       ConnectorOffsetBackingStore offsetStore,
                                       WorkerConfig workerConfig,
                                       ClusterConfigState configState,
                                       ConnectMetrics connectMetrics,
                                       ErrorHandlingMetrics errorMetrics,
                                       ClassLoader loader,
                                       Time time,
                                       RetryWithToleranceOperator<SourceRecord> retryWithToleranceOperator,
                                       StatusBackingStore statusBackingStore,
                                       SourceConnectorConfig sourceConfig,
                                       Executor closeExecutor,
                                       Runnable preProducerCheck,
                                       Runnable postProducerCheck,
                                       Supplier<List<ErrorReporter<SourceRecord>>> errorReportersSupplier,
                                       TaskPluginsMetadata pluginsMetadata,
                                       Function<ClassLoader, LoaderSwap> pluginLoaderSwapper) {
        super(id, task, statusListener, initialState, configState, keyConverterPlugin, valueConverterPlugin, headerConverterPlugin, transformationChain,
                buildTransactionContext(sourceConfig),
                producer, admin, topicGroups, offsetReader, offsetWriter, offsetStore, workerConfig, connectMetrics, errorMetrics,
                loader, time, retryWithToleranceOperator, statusBackingStore, closeExecutor, errorReportersSupplier, pluginsMetadata, pluginLoaderSwapper);

        this.transactionOpen = false;
        this.committableRecords = new LinkedHashMap<>();

        this.preProducerCheck = preProducerCheck;
        this.postProducerCheck = postProducerCheck;

        this.transactionBoundaryManager = buildTransactionManager(workerConfig, sourceConfig, sourceTaskContext.transactionContext());
        this.transactionMetrics = new TransactionMetricsGroup(id, connectMetrics);
    }

    private static WorkerTransactionContext buildTransactionContext(SourceConnectorConfig sourceConfig) {
        return TransactionBoundary.CONNECTOR.equals(sourceConfig.transactionBoundary())
                ? new WorkerTransactionContext()
                : null;
    }

    @Override
    protected void prepareToInitializeTask() {
        preProducerCheck.run();

        // Try not to initialize the transactional producer (which may accidentally fence out other, later task generations) if we've already
        // been shut down at this point
        if (isStopping())
            return;
        producer.initTransactions();

        postProducerCheck.run();
    }

    @Override
    protected void prepareToEnterSendLoop() {
        transactionBoundaryManager.initialize();
    }

    @Override
    protected void beginSendIteration() {
        // No-op
    }

    @Override
    protected void prepareToPollTask() {
        // No-op
    }

    @Override
    protected void recordDropped(SourceRecord record) {
        synchronized (committableRecords) {
            committableRecords.put(record, null);
        }
        transactionBoundaryManager.maybeCommitTransactionForRecord(record);
    }

    @Override
    protected Optional<SubmittedRecords.SubmittedRecord> prepareToSendRecord(
            SourceRecord sourceRecord,
            ProducerRecord<byte[], byte[]> producerRecord
    ) {
        if (offsetStore.primaryOffsetsTopic().equals(producerRecord.topic())) {
            // This is to prevent deadlock that occurs when:
            //     1. A task provides a record whose topic is the task's offsets topic
            //     2. That record is dispatched to the task's producer in a transaction that remains open
            //        at least until the worker polls the task again
            //     3. In the subsequent call to SourceTask::poll, the task requests offsets from the worker
            //        (which requires a read to the end of the offsets topic, and will block until any open
            //        transactions on the topic are either committed or aborted)
            throw new ConnectException("Source tasks may not produce to their own offsets topics when exactly-once support is enabled");
        }
        maybeBeginTransaction();
        return Optional.empty();
    }

    @Override
    protected void recordDispatched(SourceRecord record) {
        // Offsets are converted & serialized in the OffsetWriter
        // Important: we only save offsets for the record after it has been accepted by the producer; this way,
        // we commit those offsets if and only if the record is sent successfully.
        offsetWriter.offset(record.sourcePartition(), record.sourceOffset());
        transactionMetrics.addRecord();
        transactionBoundaryManager.maybeCommitTransactionForRecord(record);
    }

    @Override
    protected void batchDispatched() {
        transactionBoundaryManager.maybeCommitTransactionForBatch();
    }

    @Override
    protected void recordSent(
            SourceRecord sourceRecord,
            ProducerRecord<byte[], byte[]> producerRecord,
            RecordMetadata recordMetadata
    ) {
        synchronized (committableRecords) {
            committableRecords.put(sourceRecord, recordMetadata);
        }
    }

    @Override
    protected void producerSendFailed(
            ProcessingContext<SourceRecord> context,
            boolean synchronous,
            ProducerRecord<byte[], byte[]> producerRecord,
            SourceRecord preTransformRecord,
            Exception e
    ) {
        if (synchronous) {
            throw maybeWrapProducerSendException(
                    "Unrecoverable exception trying to send",
                    e
            );
        } else {
            // No-op; all asynchronously-reported producer exceptions should be bubbled up again by Producer::commitTransaction
        }
    }

    @Override
    protected void finalOffsetCommit(boolean failed) {
        if (failed) {
            log.debug("Skipping final offset commit as task has failed");
            return;
        }

        // It should be safe to commit here even if we were in the middle of retrying on RetriableExceptions in the
        // send loop since we only track source offsets for records that have been successfully dispatched to the
        // producer.
        // Any records that we were retrying on (and any records after them in the batch) won't be included in the
        // transaction and their offsets won't be committed, but (unless the user has requested connector-defined
        // transaction boundaries), it's better to commit some data than none.
        transactionBoundaryManager.maybeCommitFinalTransaction();
    }

    @Override
    public void removeMetrics() {
        Utils.closeQuietly(transactionMetrics, "source task transaction metrics tracker");
        super.removeMetrics();
    }

    @Override
    protected void onPause() {
        super.onPause();
        // Commit the transaction now so that we don't end up with a hanging transaction, or worse, get fenced out
        // and fail the task once unpaused
        transactionBoundaryManager.maybeCommitFinalTransaction();
    }

    private void maybeBeginTransaction() {
        if (!transactionOpen) {
            producer.beginTransaction();
            transactionOpen = true;
        }
    }

    private void commitTransaction() {
        log.debug("{} Committing offsets", this);

        long started = time.milliseconds();

        AtomicReference<Throwable> flushError = new AtomicReference<>();
        boolean shouldFlush = false;
        try {
            // Begin the flush without waiting, as there should not be any concurrent flushes.
            // This is because commitTransaction is always called on the same thread, and should always block until
            // the flush is complete, or cancel the flush if an error occurs.
            shouldFlush = offsetWriter.beginFlush();
        } catch (Throwable e) {
            flushError.compareAndSet(null, e);
        }
        if (flushError.get() == null && !transactionOpen && !shouldFlush) {
            // There is no contents on the framework side to commit, so skip the offset flush and producer commit
            long durationMillis = time.milliseconds() - started;
            recordCommitSuccess(durationMillis);
            log.debug("{} Finished commitOffsets successfully in {} ms", this, durationMillis);

            commitSourceTask();
            return;
        }

        // We might have just aborted a transaction, in which case we'll have to begin a new one
        // in order to commit offsets
        maybeBeginTransaction();

        if (shouldFlush) {
            // Now we can actually write the offsets to the internal topic.
            // No need to track the flush future here since it's guaranteed to complete by the time
            // Producer::commitTransaction completes
            // We do have to track failures for that callback though, since they may originate from outside
            // the producer (i.e., the offset writer or the backing offset store), and would not cause
            // Producer::commitTransaction to fail
            offsetWriter.doFlush((error, result) -> {
                if (error != null) {
                    log.error("{} Failed to flush offsets to storage: ", ExactlyOnceWorkerSourceTask.this, error);
                    flushError.compareAndSet(null, error);
                } else {
                    log.trace("{} Finished flushing offsets to storage", ExactlyOnceWorkerSourceTask.this);
                }
            });
        }

        // Only commit the transaction if we were able to serialize the offsets.
        // Otherwise, we may commit source records without committing their offsets
        Throwable error = flushError.get();
        if (error == null) {
            try {
                // Commit the transaction
                // Blocks until all outstanding records have been sent and ack'd
                producer.commitTransaction();
            } catch (Throwable t) {
                log.error("{} Failed to commit producer transaction", ExactlyOnceWorkerSourceTask.this, t);
                flushError.compareAndSet(null, t);
            }
            transactionOpen = false;
        }

        error = flushError.get();
        if (error != null) {
            recordCommitFailure(time.milliseconds() - started);
            offsetWriter.cancelFlush();
            throw maybeWrapProducerSendException(
                    "Failed to flush offsets and/or records for task " + id,
                    error
            );
        }

        transactionMetrics.commitTransaction();

        long durationMillis = time.milliseconds() - started;
        recordCommitSuccess(durationMillis);
        log.debug("{} Finished commitOffsets successfully in {} ms", this, durationMillis);

        // Synchronize in order to guarantee that writes on other threads are picked up by this one
        synchronized (committableRecords) {
            committableRecords.forEach(this::commitTaskRecord);
            committableRecords.clear();
        }
        commitSourceTask();
    }

    private RuntimeException maybeWrapProducerSendException(String message, Throwable error) {
        if (isPossibleTransactionTimeoutError(error)) {
            return wrapTransactionTimeoutError(error);
        } else {
            return new ConnectException(message, error);
        }
    }

    private static boolean isPossibleTransactionTimeoutError(Throwable error) {
        return error instanceof InvalidProducerEpochException
            || error.getCause() instanceof InvalidProducerEpochException;
    }

    private ConnectException wrapTransactionTimeoutError(Throwable error) {
        return new ConnectException(
            "The task " + id + " was unable to finish writing records to Kafka before its producer transaction expired. "
                + "It may be necessary to reconfigure this connector in order for it to run healthily with exactly-once support. "
                + "Options for this include: tune the connector's producer configuration for higher throughput, "
                + "increase the transaction timeout for the connector's producers, "
                + "decrease the offset commit interval (if using interval-based transaction boundaries), "
                + "or use the 'poll' transaction boundary (if the connector is not already configured to use it).",
            error
        );
    }

    @Override
    public String toString() {
        return "ExactlyOnceWorkerSourceTask{" +
            "id=" + id +
            '}';
    }

    private abstract class TransactionBoundaryManager {
        protected boolean shouldCommitTransactionForRecord(SourceRecord record) {
            return false;
        }

        protected boolean shouldCommitTransactionForBatch(long currentTimeMs) {
            return false;
        }

        protected boolean shouldCommitFinalTransaction() {
            return false;
        }

        /**
         * Hook to signal that a new transaction cycle has been started. May be invoked
         * multiple times if the task is paused and then resumed. It can be assumed that
         * a new transaction is created at least every time an existing transaction is
         * committed; this is just a hook to notify that a new transaction may have been
         * created outside of that flow as well.
         */
        protected void initialize() {
        }

        public void maybeCommitTransactionForRecord(SourceRecord record) {
            maybeCommitTransaction(shouldCommitTransactionForRecord(record));
        }

        public void maybeCommitTransactionForBatch() {
            maybeCommitTransaction(shouldCommitTransactionForBatch(time.milliseconds()));
        }

        public void maybeCommitFinalTransaction() {
            maybeCommitTransaction(shouldCommitFinalTransaction());
        }

        private void maybeCommitTransaction(boolean shouldCommit) {
            if (shouldCommit) {
                try (LoggingContext loggingContext = LoggingContext.forOffsets(id)) {
                    commitTransaction();
                }
            }
        }
    }

    private TransactionBoundaryManager buildTransactionManager(
            WorkerConfig workerConfig,
            SourceConnectorConfig sourceConfig,
            WorkerTransactionContext transactionContext) {
        TransactionBoundary boundary = sourceConfig.transactionBoundary();
        return switch (boundary) {
            case POLL -> new TransactionBoundaryManager() {
                @Override
                protected boolean shouldCommitTransactionForBatch(long currentTimeMs) {
                    return true;
                }

                @Override
                protected boolean shouldCommitFinalTransaction() {
                    return true;
                }
            };
            case INTERVAL -> {
                long transactionBoundaryInterval = Optional.ofNullable(sourceConfig.transactionBoundaryInterval())
                        .orElse(workerConfig.offsetCommitInterval());
                yield new TransactionBoundaryManager() {
                    private final long commitInterval = transactionBoundaryInterval;
                    private long lastCommit;

                    @Override
                    public void initialize() {
                        this.lastCommit = time.milliseconds();
                    }

                    @Override
                    protected boolean shouldCommitTransactionForBatch(long currentTimeMs) {
                        if (time.milliseconds() >= lastCommit + commitInterval) {
                            lastCommit = time.milliseconds();
                            return true;
                        } else {
                            return false;
                        }
                    }

                    @Override
                    protected boolean shouldCommitFinalTransaction() {
                        return true;
                    }
                };
            }
            case CONNECTOR -> {
                Objects.requireNonNull(transactionContext, "Transaction context must be provided when using connector-defined transaction boundaries");
                yield new TransactionBoundaryManager() {
                    @Override
                    protected boolean shouldCommitFinalTransaction() {
                        return shouldCommitTransactionForBatch(time.milliseconds());
                    }

                    @Override
                    protected boolean shouldCommitTransactionForBatch(long currentTimeMs) {
                        if (transactionContext.shouldAbortBatch()) {
                            log.info("Aborting transaction for batch as requested by connector");
                            maybeAbortTransaction();
                            // We abort the transaction, which causes all the records up to this point to be dropped, but we still want to
                            // commit offsets so that the task doesn't see the same records all over again
                            return true;
                        }
                        return transactionContext.shouldCommitBatch();
                    }

                    @Override
                    protected boolean shouldCommitTransactionForRecord(SourceRecord record) {
                        if (transactionContext.shouldAbortOn(record)) {
                            log.info("Aborting transaction for record on topic {} as requested by connector", record.topic());
                            log.trace("Last record in aborted transaction: {}", record);
                            maybeAbortTransaction();
                            // We abort the transaction, which causes all the records up to this point to be dropped, but we still want to
                            // commit offsets so that the task doesn't see the same records all over again
                            return true;
                        }
                        return transactionContext.shouldCommitOn(record);
                    }

                    private void maybeAbortTransaction() {
                        if (!transactionOpen) {
                            log.warn("Ignoring request by task to abort transaction as the current transaction is empty");
                            return;
                        }
                        producer.abortTransaction();
                        transactionMetrics.abortTransaction();
                        transactionOpen = false;
                    }
                };
            }
        };
    }

    TransactionMetricsGroup transactionMetricsGroup() {
        return transactionMetrics;
    }


    static class TransactionMetricsGroup implements AutoCloseable {
        private final Sensor transactionSize;
        private int size;
        private final ConnectMetrics.MetricGroup metricGroup;

        public TransactionMetricsGroup(ConnectorTaskId id, ConnectMetrics connectMetrics) {
            ConnectMetricsRegistry registry = connectMetrics.registry();
            metricGroup = connectMetrics.group(registry.sourceTaskGroupName(),
                    registry.connectorTagName(), id.connector(),
                    registry.taskTagName(), Integer.toString(id.task()));

            transactionSize = metricGroup.sensor("transaction-size");
            transactionSize.add(metricGroup.metricName(registry.transactionSizeAvg), new Avg());
            transactionSize.add(metricGroup.metricName(registry.transactionSizeMin), new Min());
            transactionSize.add(metricGroup.metricName(registry.transactionSizeMax), new Max());
        }

        @Override
        public void close() {
            metricGroup.close();
        }

        void addRecord() {
            size++;
        }

        void abortTransaction() {
            size = 0;
        }

        void commitTransaction() {
            transactionSize.record(size);
            size = 0;
        }

        protected ConnectMetrics.MetricGroup metricGroup() {
            return metricGroup;
        }

    }

}
