/*
 * 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.RecordMetadata;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.InvalidTopicException;
import org.apache.kafka.common.errors.RecordTooLargeException;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.connect.data.Schema;
import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.integration.MonitorableSourceConnector;
import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
import org.apache.kafka.connect.runtime.errors.ErrorHandlingMetrics;
import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest;
import org.apache.kafka.connect.runtime.isolation.Plugins;
import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
import org.apache.kafka.connect.source.SourceRecord;
import org.apache.kafka.connect.source.SourceTask;
import org.apache.kafka.connect.source.TransactionContext;
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.storage.StringConverter;
import org.apache.kafka.connect.test.util.ConcurrencyUtils;
import org.apache.kafka.connect.test.util.MockitoUtils;
import org.apache.kafka.connect.util.Callback;
import org.apache.kafka.connect.util.ConnectorTaskId;
import org.apache.kafka.connect.util.ParameterizedTest;
import org.apache.kafka.connect.util.TopicAdmin;
import org.apache.kafka.connect.util.TopicCreationGroup;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnit;
import org.mockito.junit.MockitoRule;
import org.mockito.stubbing.OngoingStubbing;
import org.mockito.verification.VerificationMode;

import java.time.Duration;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import static java.util.Collections.emptySet;
import static org.apache.kafka.connect.integration.MonitorableSourceConnector.TOPIC_CONFIG;
import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
import static org.apache.kafka.connect.runtime.SourceConnectorConfig.TOPIC_CREATION_GROUPS_CONFIG;
import static org.apache.kafka.connect.runtime.SourceConnectorConfig.TRANSACTION_BOUNDARY_CONFIG;
import static org.apache.kafka.connect.runtime.SourceConnectorConfig.TRANSACTION_BOUNDARY_INTERVAL_CONFIG;
import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
import static org.apache.kafka.connect.runtime.TopicCreationConfig.EXCLUDE_REGEX_CONFIG;
import static org.apache.kafka.connect.runtime.TopicCreationConfig.INCLUDE_REGEX_CONFIG;
import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_CREATION_ENABLE_CONFIG;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.atLeast;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

@RunWith(Parameterized.class)
public class ExactlyOnceWorkerSourceTaskTest {
    private static final String TOPIC = "topic";
    private static final Map<String, byte[]> PARTITION = Collections.singletonMap("key", "partition".getBytes());
    private static final Map<String, ?> OFFSET = offset(12);

    // Connect-format data
    private static final Schema KEY_SCHEMA = Schema.INT32_SCHEMA;
    private static final Integer KEY = -1;
    private static final Schema RECORD_SCHEMA = Schema.INT64_SCHEMA;
    private static final Long VALUE_1 = 12L;
    private static final Long VALUE_2 = 13L;
    // Serialized data. The actual format of this data doesn't matter -- we just want to see that the right version
    // is used in the right place.
    private static final byte[] SERIALIZED_KEY = "converted-key".getBytes();
    private static final byte[] SERIALIZED_RECORD = "converted-record".getBytes();

    private final ExecutorService executor = Executors.newSingleThreadExecutor();
    private final ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
    private WorkerConfig config;
    private SourceConnectorConfig sourceConfig;
    private Plugins plugins;
    private MockConnectMetrics metrics;
    @Mock private ErrorHandlingMetrics errorHandlingMetrics;
    private Time time;
    private ExactlyOnceWorkerSourceTask workerTask;
    @Mock private SourceTask sourceTask;
    @Mock private Converter keyConverter;
    @Mock private Converter valueConverter;
    @Mock private HeaderConverter headerConverter;
    @Mock private TransformationChain<SourceRecord, SourceRecord> transformationChain;
    @Mock private Producer<byte[], byte[]> producer;
    @Mock private TopicAdmin admin;
    @Mock private CloseableOffsetStorageReader offsetReader;
    @Mock private OffsetStorageWriter offsetWriter;
    @Mock private ClusterConfigState clusterConfigState;
    @Mock private TaskStatus.Listener statusListener;
    @Mock private StatusBackingStore statusBackingStore;
    @Mock private ConnectorOffsetBackingStore offsetStore;
    @Mock private Runnable preProducerCheck;
    @Mock private Runnable postProducerCheck;

    @Rule public MockitoRule rule = MockitoJUnit.rule();

    private static final Map<String, String> TASK_PROPS = new HashMap<>();
    static {
        TASK_PROPS.put(TaskConfig.TASK_CLASS_CONFIG, TestSourceTask.class.getName());
    }
    private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);

    private static final SourceRecord SOURCE_RECORD_1 =
            new SourceRecord(PARTITION, OFFSET, TOPIC, null, KEY_SCHEMA, KEY, RECORD_SCHEMA, VALUE_1);
    private static final SourceRecord SOURCE_RECORD_2 =
            new SourceRecord(PARTITION, OFFSET, TOPIC, null, KEY_SCHEMA, KEY, RECORD_SCHEMA, VALUE_2);

    private static final List<SourceRecord> RECORDS = Arrays.asList(SOURCE_RECORD_1, SOURCE_RECORD_2);
    private final AtomicReference<CountDownLatch> pollLatch = new AtomicReference<>(new CountDownLatch(0));
    private final AtomicReference<List<SourceRecord>> pollRecords = new AtomicReference<>(RECORDS);

    private final boolean enableTopicCreation;

    private boolean taskStarted;
    private Future<?> workerTaskFuture;

    @ParameterizedTest.Parameters
    public static Collection<Boolean> parameters() {
        return Arrays.asList(false, true);
    }

    public ExactlyOnceWorkerSourceTaskTest(boolean enableTopicCreation) {
        this.enableTopicCreation = enableTopicCreation;
        this.taskStarted = false;
    }

    @Before
    public void setup() throws Exception {
        Map<String, String> workerProps = workerProps();
        plugins = new Plugins(workerProps);
        config = new StandaloneConfig(workerProps);
        sourceConfig = new SourceConnectorConfig(plugins, sourceConnectorProps(), true);
        metrics = new MockConnectMetrics();
        time = Time.SYSTEM;
        when(offsetStore.primaryOffsetsTopic()).thenReturn("offsets-topic");
        when(sourceTask.poll()).thenAnswer(invocation -> {
            // Capture the records we'll return before decrementing the poll latch,
            // since test cases may mutate the poll records field immediately
            // after awaiting the latch
            List<SourceRecord> result = pollRecords.get();
            pollLatch.get().countDown();
            Thread.sleep(10);
            return result;
        });
    }

    @After
    public void teardown() throws Exception {
        // In most tests, we don't really care about how many times the task got polled,
        // how many times we prepared to write offsets, etc.
        // We add these verify(..., anyTimes())... calls to make sure that those interactions
        // with our mocks don't cause the verifyNoMoreInteractions(...) call to fail.
        // Individual tests can add more fine-grained verification logic if necessary; the
        // verifications here will not conflict with them

        verify(sourceTask, MockitoUtils.anyTimes()).poll();
        verify(sourceTask, MockitoUtils.anyTimes()).commit();
        verify(sourceTask, MockitoUtils.anyTimes()).commitRecord(any(), any());

        verify(offsetWriter, MockitoUtils.anyTimes()).offset(PARTITION, OFFSET);
        verify(offsetWriter, MockitoUtils.anyTimes()).beginFlush();
        verify(offsetWriter, MockitoUtils.anyTimes()).doFlush(any());

        if (enableTopicCreation) {
            verify(admin, MockitoUtils.anyTimes()).describeTopics(TOPIC);
        }

        verify(statusBackingStore, MockitoUtils.anyTimes()).getTopic(any(), any());

        verify(offsetStore, MockitoUtils.anyTimes()).primaryOffsetsTopic();

        verifyNoMoreInteractions(statusListener, producer, sourceTask, admin, offsetWriter, statusBackingStore, offsetStore, preProducerCheck, postProducerCheck);
        if (metrics != null) metrics.stop();
    }

    private Map<String, String> workerProps() {
        Map<String, String> props = new HashMap<>();
        props.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
        props.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
        props.put("internal.key.converter", "org.apache.kafka.connect.json.JsonConverter");
        props.put("internal.value.converter", "org.apache.kafka.connect.json.JsonConverter");
        props.put("internal.key.converter.schemas.enable", "false");
        props.put("internal.value.converter.schemas.enable", "false");
        props.put("offset.storage.file.filename", "/tmp/connect.offsets");
        props.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(enableTopicCreation));
        return props;
    }

    private Map<String, String> sourceConnectorProps() {
        return sourceConnectorProps(SourceTask.TransactionBoundary.DEFAULT);
    }

    private Map<String, String> sourceConnectorProps(SourceTask.TransactionBoundary transactionBoundary) {
        // setup up props for the source connector
        Map<String, String> props = new HashMap<>();
        props.put("name", "foo-connector");
        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
        props.put(TASKS_MAX_CONFIG, String.valueOf(1));
        props.put(TOPIC_CONFIG, TOPIC);
        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
        props.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", "foo", "bar"));
        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
        props.put(TRANSACTION_BOUNDARY_CONFIG, transactionBoundary.toString());
        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + "foo" + "." + INCLUDE_REGEX_CONFIG, TOPIC);
        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + "bar" + "." + INCLUDE_REGEX_CONFIG, ".*");
        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + "bar" + "." + EXCLUDE_REGEX_CONFIG, TOPIC);
        return props;
    }

    private static Map<String, ?> offset(int n) {
        return Collections.singletonMap("key", n);
    }

    private void createWorkerTask() {
        createWorkerTask(TargetState.STARTED);
    }

    private void createWorkerTask(TargetState initialState) {
        createWorkerTask(initialState, keyConverter, valueConverter, headerConverter);
    }

    private void createWorkerTask(TargetState initialState, Converter keyConverter, Converter valueConverter, HeaderConverter headerConverter) {
        workerTask = new ExactlyOnceWorkerSourceTask(taskId, sourceTask, statusListener, initialState, keyConverter, valueConverter, headerConverter,
                transformationChain, producer, admin, TopicCreationGroup.configuredGroups(sourceConfig), offsetReader, offsetWriter, offsetStore,
                config, clusterConfigState, metrics, errorHandlingMetrics, plugins.delegatingLoader(), time, RetryWithToleranceOperatorTest.noopOperator(), statusBackingStore,
                sourceConfig, Runnable::run, preProducerCheck, postProducerCheck, Collections::emptyList);
    }

    @Test
    public void testRemoveMetrics() {
        createWorkerTask();

        workerTask.removeMetrics();

        assertEquals(emptySet(), filterToTaskMetrics(metrics.metrics().metrics().keySet()));
    }

    private Set<MetricName> filterToTaskMetrics(Set<MetricName> metricNames) {
        return metricNames
                .stream()
                .filter(m -> metrics.registry().taskGroupName().equals(m.group())
                        || metrics.registry().sourceTaskGroupName().equals(m.group()))
                .collect(Collectors.toSet());
    }

    @Test
    public void testStartPaused() throws Exception {
        createWorkerTask(TargetState.PAUSED);

        final CountDownLatch pauseLatch = new CountDownLatch(1);
        doAnswer(invocation -> {
            pauseLatch.countDown();
            return null;
        }).when(statusListener).onPause(eq(taskId));

        startTaskThread();

        assertTrue(pauseLatch.await(5, TimeUnit.SECONDS));

        awaitShutdown();

        verify(statusListener).onPause(taskId);
        verify(sourceTask, never()).start(any());
        verify(sourceTask, never()).poll();
        verifyCleanShutdown();
        assertPollMetrics(0);
    }

    @Test
    public void testPause() throws Exception {
        createWorkerTask();

        expectSuccessfulSends();
        expectSuccessfulFlushes();

        final CountDownLatch pauseLatch = new CountDownLatch(1);
        doAnswer(invocation -> {
            pauseLatch.countDown();
            return null;
        }).when(statusListener).onPause(eq(taskId));

        startTaskThread();

        awaitPolls(2);
        workerTask.transitionTo(TargetState.PAUSED);

        long pollsBeforePause = pollCount();
        ConcurrencyUtils.awaitLatch(pauseLatch, "task did not pause in time");

        // since the transition is observed asynchronously, the count could be off by one loop iteration
        assertTrue(pollCount() - pollsBeforePause <= 1);

        awaitShutdown();

        verify(statusListener).onPause(taskId);
        // Task should have flushed offsets for every record poll, once on pause, and once for end-of-life offset commit
        verifyTransactions(pollCount() + 2, pollCount() + 2);
        verifySends();
        verifyPossibleTopicCreation();
        // make sure we didn't poll again after triggering shutdown
        assertTrue(pollCount() - pollsBeforePause <= 1);

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
    }

    @Test
    public void testFailureInPreProducerCheck() throws Exception {
        createWorkerTask();

        Exception exception = new ConnectException("Failed to perform zombie fencing");
        doThrow(exception).when(preProducerCheck).run();

        workerTask.initialize(TASK_CONFIG);
        // No need to execute on a separate thread; preflight checks should all take place before the poll-send loop starts
        workerTask.run();
        verify(preProducerCheck).run();
        verifyShutdown(true, false);
    }

    @Test
    public void testFailureInProducerInitialization() throws Exception {
        createWorkerTask();

        Exception exception = new ConnectException("You can't do that!");
        doThrow(exception).when(producer).initTransactions();

        workerTask.initialize(TASK_CONFIG);
        // No need to execute on a separate thread; preflight checks should all take place before the poll-send loop starts
        workerTask.run();
        verify(preProducerCheck).run();
        verify(producer).initTransactions();
        verifyShutdown(true, false);
    }

    @Test
    public void testFailureInPostProducerCheck() throws Exception {
        createWorkerTask();

        Exception exception = new ConnectException("New task configs for the connector have already been generated");
        doThrow(exception).when(postProducerCheck).run();

        workerTask.initialize(TASK_CONFIG);
        // No need to execute on a separate thread; preflight checks should all take place before the poll-send loop starts
        workerTask.run();
        verify(preProducerCheck).run();
        verify(producer).initTransactions();
        verify(postProducerCheck).run();
        verifyShutdown(true, false);
    }

    @Test
    public void testFailureInOffsetStoreStart() throws Exception {
        createWorkerTask();

        Exception exception = new ConnectException("No soup for you!");
        doThrow(exception).when(offsetStore).start();

        workerTask.initialize(TASK_CONFIG);
        // No need to execute on a separate thread; preflight checks should all take place before the poll-send loop starts
        workerTask.run();
        verify(preProducerCheck).run();
        verify(producer).initTransactions();
        verify(postProducerCheck).run();
        verify(offsetStore).start();
        verifyShutdown(true, false);
    }

    @Test
    public void testPollsInBackground() throws Exception {
        createWorkerTask();

        expectSuccessfulSends();
        expectSuccessfulFlushes();

        startTaskThread();

        int minPolls = 10;
        awaitPolls(minPolls);
        awaitShutdown();

        // Task should have flushed offsets for every record poll and for end-of-life offset commit
        verifyTransactions(pollCount() + 1, pollCount() + 1);
        verifySends();
        verifyPossibleTopicCreation();

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
        assertPollMetrics(minPolls);
        assertTransactionMetrics(RECORDS.size());
    }

    @Test
    public void testFailureInPoll() throws Exception {
        createWorkerTask();

        final CountDownLatch pollLatch = new CountDownLatch(1);
        final RuntimeException exception = new RuntimeException();
        when(sourceTask.poll()).thenAnswer(invocation -> {
            pollLatch.countDown();
            throw exception;
        });

        startTaskThread();

        ConcurrencyUtils.awaitLatch(pollLatch, "task was not polled in time");
        //Failure in poll should trigger automatic stop of the worker
        awaitShutdown(false);

        verifyPreflight();
        verifyStartup();
        verifyShutdown(true, false);
        assertPollMetrics(0);
    }

    @Test
    public void testFailureInPollAfterCancel() throws Exception {
        createWorkerTask();

        final CountDownLatch pollLatch = new CountDownLatch(1);
        final CountDownLatch workerCancelLatch = new CountDownLatch(1);
        final RuntimeException exception = new RuntimeException();
        when(sourceTask.poll()).thenAnswer(invocation -> {
            pollLatch.countDown();
            ConcurrencyUtils.awaitLatch(workerCancelLatch, "task was not cancelled in time");
            throw exception;
        });

        startTaskThread();

        ConcurrencyUtils.awaitLatch(pollLatch, "task was not polled in time");
        workerTask.cancel();
        workerCancelLatch.countDown();
        awaitShutdown(false);

        verifyPreflight();
        verifyStartup();
        verifyShutdown(false, true);
        assertPollMetrics(0);
    }

    @Test
    public void testFailureInPollAfterStop() throws Exception {
        createWorkerTask();

        final CountDownLatch pollLatch = new CountDownLatch(1);
        final CountDownLatch workerStopLatch = new CountDownLatch(1);
        final RuntimeException exception = new RuntimeException();
        // We check if there are offsets that need to be committed before shutting down the task
        when(sourceTask.poll()).thenAnswer(invocation -> {
            pollLatch.countDown();
            ConcurrencyUtils.awaitLatch(workerStopLatch, "task was not stopped in time");
            throw exception;
        });

        startTaskThread();

        ConcurrencyUtils.awaitLatch(pollLatch, "task was not polled in time");
        workerTask.stop();
        workerStopLatch.countDown();
        awaitShutdown(false);

        verifyTransactions(0, 0);

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
        assertPollMetrics(0);
    }

    @Test
    public void testPollReturnsNoRecords() throws Exception {
        // Test that the task handles an empty list of records
        createWorkerTask();

        // Make sure the task returns empty batches from poll before we start polling it
        pollRecords.set(Collections.emptyList());

        when(offsetWriter.beginFlush()).thenReturn(false);

        startTaskThread();

        awaitEmptyPolls(10);

        awaitShutdown();

        // task commits for each empty poll, plus the final commit
        verifyTransactions(pollCount() + 1, 0);

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
        assertPollMetrics(0);
    }

    @Test
    public void testPollBasedCommit() throws Exception {
        Map<String, String> connectorProps = sourceConnectorProps(SourceTask.TransactionBoundary.POLL);
        sourceConfig = new SourceConnectorConfig(plugins, connectorProps, enableTopicCreation);

        createWorkerTask();

        expectSuccessfulSends();
        expectSuccessfulFlushes();

        startTaskThread();

        awaitPolls(10);

        awaitShutdown();

        // Task should have flushed offsets for every record poll, and for end-of-life offset commit
        verifyTransactions(pollCount() + 1, pollCount() + 1);
        verifySends();
        verifyPossibleTopicCreation();

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
        assertPollMetrics(1);
        assertTransactionMetrics(RECORDS.size());
    }

    @Test
    public void testIntervalBasedCommit() throws Exception {
        long commitInterval = 618;
        Map<String, String> connectorProps = sourceConnectorProps(SourceTask.TransactionBoundary.INTERVAL);
        connectorProps.put(TRANSACTION_BOUNDARY_INTERVAL_CONFIG, Long.toString(commitInterval));
        sourceConfig = new SourceConnectorConfig(plugins, connectorProps, enableTopicCreation);

        time = new MockTime();

        createWorkerTask();

        expectSuccessfulSends();
        expectSuccessfulFlushes();

        startTaskThread();

        awaitPolls(2);
        assertEquals("No flushes should have taken place before offset commit interval has elapsed", 0, flushCount());
        time.sleep(commitInterval);

        awaitPolls(2);
        assertEquals("One flush should have taken place after offset commit interval has elapsed", 1, flushCount());
        time.sleep(commitInterval * 2);

        awaitPolls(2);
        assertEquals("Two flushes should have taken place after offset commit interval has elapsed again", 2, flushCount());

        awaitShutdown();

        // Task should have flushed offsets twice based on offset commit interval, and performed final end-of-life offset commit
        verifyTransactions(3, 3);
        verifySends();
        verifyPossibleTopicCreation();

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
        assertPollMetrics(2);
        assertTransactionMetrics(RECORDS.size() * 2);
    }

    @Test
    public void testConnectorCommitOnBatch() throws Exception {
        testConnectorBasedCommit(TransactionContext::commitTransaction, false);
    }

    @Test
    public void testConnectorCommitOnRecord() throws Exception {
        testConnectorBasedCommit(ctx -> ctx.commitTransaction(SOURCE_RECORD_2), false);
    }

    @Test
    public void testConnectorAbortOnBatch() throws Exception {
        testConnectorBasedCommit(TransactionContext::abortTransaction, true);
    }

    @Test
    public void testConnectorAbortOnRecord() throws Exception {
        testConnectorBasedCommit(ctx -> ctx.abortTransaction(SOURCE_RECORD_2), true);
    }

    private void testConnectorBasedCommit(Consumer<TransactionContext> requestCommit, boolean abort) throws Exception {
        Map<String, String> connectorProps = sourceConnectorProps(SourceTask.TransactionBoundary.CONNECTOR);
        sourceConfig = new SourceConnectorConfig(plugins, connectorProps, enableTopicCreation);
        createWorkerTask();

        expectSuccessfulSends();
        expectSuccessfulFlushes();

        TransactionContext transactionContext = workerTask.sourceTaskContext.transactionContext();

        startTaskThread();

        awaitPolls(3);
        assertEquals("No flushes should have taken place without connector requesting transaction commit",
                0, flushCount());

        requestCommit.accept(transactionContext);
        awaitPolls(3);
        assertEquals("One flush should have taken place after transaction commit/abort was requested",
                1, flushCount());

        awaitPolls(3);
        assertEquals("Only one flush should still have taken place without connector re-requesting commit/abort, even on identical records",
                1, flushCount());

        awaitShutdown();

        assertEquals("Task should have flushed offsets once based on connector-defined boundaries, and skipped final end-of-life offset commit",
                1, flushCount());
        // We begin a new transaction after connector-requested aborts so that we can still write offsets for the source records that were aborted
        verify(producer, times(abort ? 3 : 2)).beginTransaction();
        verifySends();
        if (abort) {
            verify(producer).abortTransaction();
        }
        verify(producer).commitTransaction();

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
        verifyPossibleTopicCreation();
        assertPollMetrics(1);
        assertTransactionMetrics(abort ? 0 : (3 * RECORDS.size()));
    }

    @Test
    public void testConnectorAbortsEmptyTransaction() throws Exception {
        Map<String, String> connectorProps = sourceConnectorProps(SourceTask.TransactionBoundary.CONNECTOR);
        sourceConfig = new SourceConnectorConfig(plugins, connectorProps, enableTopicCreation);
        createWorkerTask();

        expectPossibleTopicCreation();
        expectTaskGetTopic();
        expectApplyTransformationChain();
        expectConvertHeadersAndKeyValue();
        TransactionContext transactionContext = workerTask.sourceTaskContext.transactionContext();

        startTaskThread();

        // Ensure the task produces at least one non-empty batch
        awaitPolls(1);
        // Start returning empty batches from SourceTask::poll
        awaitEmptyPolls(1);
        // The task commits the active transaction, which should not be empty and should
        // result in a single call to Producer::commitTransaction
        transactionContext.commitTransaction();

        // Ensure the task produces at least one empty batch after the transaction has been committed
        awaitEmptyPolls(1);
        // The task aborts the active transaction, which has no records in it and should not trigger a call to
        // Producer::abortTransaction
        transactionContext.abortTransaction();
        // Make sure we go through at least one more poll, to give the framework a chance to handle the
        // transaction abort request
        awaitEmptyPolls(1);

        awaitShutdown(true);

        verify(producer, times(1)).beginTransaction();
        verify(producer, times(1)).commitTransaction();
        verify(producer, atLeast(RECORDS.size())).send(any(), any());
        // We never abort transactions in this test!
        verify(producer, never()).abortTransaction();

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
        verifyPossibleTopicCreation();
    }

    @Test
    public void testMixedConnectorTransactionBoundaryCommitLastRecordAbortBatch() throws Exception {
        // We fail tasks that try to abort and commit a transaction for the same record or same batch
        // But we don't fail if they try to commit the last record of a batch and abort the entire batch
        // Instead, we give precedence to the record-based operation

        Map<String, String> connectorProps = sourceConnectorProps(SourceTask.TransactionBoundary.CONNECTOR);
        sourceConfig = new SourceConnectorConfig(plugins, connectorProps, enableTopicCreation);
        createWorkerTask();

        expectPossibleTopicCreation();
        expectTaskGetTopic();
        expectApplyTransformationChain();
        expectConvertHeadersAndKeyValue();
        when(offsetWriter.beginFlush())
                .thenReturn(true)
                .thenReturn(false);

        workerTask.initialize(TASK_CONFIG);

        TransactionContext transactionContext = workerTask.sourceTaskContext.transactionContext();

        // Request that the batch be aborted
        transactionContext.abortTransaction();
        // Request that the last record in the batch be committed
        transactionContext.commitTransaction(RECORDS.get(RECORDS.size() - 1));

        workerTask.toSend = RECORDS;
        assertTrue(workerTask.sendRecords());

        verifyTransactions(2, 1);
        verifySends(RECORDS.size());
        // We never abort transactions in this test!
        verify(producer, never()).abortTransaction();

        verifyPossibleTopicCreation();
    }

    @Test
    public void testMixedConnectorTransactionBoundaryAbortLastRecordCommitBatch() throws Exception {
        // We fail tasks that try to abort and commit a transaction for the same record or same batch
        // But we don't fail if they try to abort the last record of a batch and commit the entire batch
        // Instead, we give precedence to the record-based operation

        Map<String, String> connectorProps = sourceConnectorProps(SourceTask.TransactionBoundary.CONNECTOR);
        sourceConfig = new SourceConnectorConfig(plugins, connectorProps, enableTopicCreation);
        createWorkerTask();

        expectPossibleTopicCreation();
        expectTaskGetTopic();
        expectApplyTransformationChain();
        expectConvertHeadersAndKeyValue();
        when(offsetWriter.beginFlush())
                .thenReturn(true)
                .thenReturn(false);

        workerTask.initialize(TASK_CONFIG);

        TransactionContext transactionContext = workerTask.sourceTaskContext.transactionContext();

        // Request that the last record in the batch be aborted
        transactionContext.abortTransaction(RECORDS.get(RECORDS.size() - 1));
        // Request that the batch be committed
        transactionContext.commitTransaction();

        workerTask.toSend = RECORDS;
        assertTrue(workerTask.sendRecords());

        verify(offsetWriter, times(2)).beginFlush();
        verify(sourceTask, times(2)).commit();
        // We open a transaction once for the aborted batch, and once to commit offsets for that batch
        verify(producer, times(2)).beginTransaction();
        verify(producer, times(1)).commitTransaction();
        verify(offsetWriter, times(1)).doFlush(any());
        verify(producer, times(1)).abortTransaction();

        verifySends(RECORDS.size());

        verifyPossibleTopicCreation();
    }

    @Test
    public void testCommitFlushSyncCallbackFailure() throws Exception {
        Exception failure = new RecordTooLargeException();
        when(offsetWriter.beginFlush()).thenReturn(true);
        when(offsetWriter.doFlush(any())).thenAnswer(invocation -> {
            Callback<Void> callback = invocation.getArgument(0);
            callback.onCompletion(failure, null);
            return null;
        });
        testCommitFailure(failure, false);
    }

    @Test
    public void testCommitFlushAsyncCallbackFailure() throws Exception {
        Exception failure = new RecordTooLargeException();
        when(offsetWriter.beginFlush()).thenReturn(true);
        // doFlush delegates its callback to the producer,
        // which delays completing the callback until commitTransaction
        AtomicReference<Callback<Void>> callback = new AtomicReference<>();
        when(offsetWriter.doFlush(any())).thenAnswer(invocation -> {
            callback.set(invocation.getArgument(0));
            return null;
        });
        doAnswer(invocation -> {
            callback.get().onCompletion(failure, null);
            return null;
        }).when(producer).commitTransaction();
        testCommitFailure(failure, true);
    }

    @Test
    public void testCommitTransactionFailure() throws Exception {
        Exception failure = new RecordTooLargeException();
        when(offsetWriter.beginFlush()).thenReturn(true);
        doThrow(failure).when(producer).commitTransaction();
        testCommitFailure(failure, true);
    }

    private void testCommitFailure(Exception commitException, boolean executeCommit) throws Exception {
        createWorkerTask();

        // Unlike the standard WorkerSourceTask class, this one fails permanently when offset commits don't succeed
        final CountDownLatch taskFailure = new CountDownLatch(1);
        final AtomicReference<Exception> taskException = new AtomicReference<>();
        doAnswer(invocation -> {
            taskFailure.countDown();
            taskException.set(invocation.getArgument(1));
            return null;
        }).when(statusListener).onFailure(eq(taskId), any());

        expectSuccessfulSends();

        startTaskThread();

        ConcurrencyUtils.awaitLatch(taskFailure, "task did not fail in time");

        awaitShutdown();

        assertEquals(commitException, taskException.get().getCause());
        verifySends();
        verifyPossibleTopicCreation();
        verify(producer).beginTransaction();
        if (executeCommit) {
            verify(producer).commitTransaction();
        }
        verify(offsetWriter).cancelFlush();

        verifyPreflight();
        verifyStartup();
        verifyShutdown(true, false);
    }

    @Test
    public void testSendRecordsRetries() {
        createWorkerTask();

        // Differentiate only by Kafka partition so we can reuse conversion expectations
        SourceRecord record1 = new SourceRecord(PARTITION, offset(1), TOPIC, 1, KEY_SCHEMA, KEY, RECORD_SCHEMA, VALUE_1);
        SourceRecord record2 = new SourceRecord(PARTITION, offset(2), TOPIC, 2, KEY_SCHEMA, KEY, RECORD_SCHEMA, VALUE_1);
        SourceRecord record3 = new SourceRecord(PARTITION, offset(3), TOPIC, 3, KEY_SCHEMA, KEY, RECORD_SCHEMA, VALUE_1);

        expectPossibleTopicCreation();
        expectTaskGetTopic();
        expectApplyTransformationChain();
        expectConvertHeadersAndKeyValue();

        // We're trying to send three records
        workerTask.toSend = Arrays.asList(record1, record2, record3);
        OngoingStubbing<Future<RecordMetadata>> producerSend = when(producer.send(any(), any()));
        // The first one is sent successfully
        producerSend = expectSuccessfulSend(producerSend);
        // The second one encounters a retriable failure
        producerSend = expectSynchronousFailedSend(producerSend, new org.apache.kafka.common.errors.TimeoutException("retriable sync failure"));
        // All subsequent sends (which should include a retry of the second record and an initial try for the third) succeed
        expectSuccessfulSend(producerSend);

        assertFalse(workerTask.sendRecords());
        assertEquals(Arrays.asList(record2, record3), workerTask.toSend);
        verify(producer).beginTransaction();
        // When using poll-based transaction boundaries, we do not commit transactions while retrying delivery for a batch
        verify(producer, never()).commitTransaction();
        verifySends(2);
        verifyPossibleTopicCreation();

        // Next they all succeed
        assertTrue(workerTask.sendRecords());
        assertNull(workerTask.toSend);
        // After the entire batch is successfully dispatched to the producer, we can finally commit the transaction
        verify(producer, times(1)).commitTransaction();
        verifySends(4);

        verify(offsetWriter).offset(PARTITION, offset(1));
        verify(offsetWriter).offset(PARTITION, offset(2));
        verify(offsetWriter).offset(PARTITION, offset(3));
    }

    @Test
    public void testSendRecordsProducerSendFailsImmediately() {
        createWorkerTask();

        SourceRecord record1 = new SourceRecord(PARTITION, OFFSET, TOPIC, 1, KEY_SCHEMA, KEY, RECORD_SCHEMA, VALUE_1);
        SourceRecord record2 = new SourceRecord(PARTITION, OFFSET, TOPIC, 2, KEY_SCHEMA, KEY, RECORD_SCHEMA, VALUE_1);

        expectPossibleTopicCreation();
        expectConvertHeadersAndKeyValue();
        expectApplyTransformationChain();

        when(producer.send(any(), any()))
                .thenThrow(new KafkaException("Producer closed while send in progress", new InvalidTopicException(TOPIC)));

        workerTask.toSend = Arrays.asList(record1, record2);
        assertThrows(ConnectException.class, workerTask::sendRecords);

        verify(producer).beginTransaction();
        verify(producer).send(any(), any());
        verify(offsetWriter, never()).offset(any(), any());
        verifyPossibleTopicCreation();
    }

    @Test
    public void testSlowTaskStart() throws Exception {
        final CountDownLatch startupLatch = new CountDownLatch(1);
        final CountDownLatch finishStartupLatch = new CountDownLatch(1);

        createWorkerTask();

        doAnswer(invocation -> {
            startupLatch.countDown();
            ConcurrencyUtils.awaitLatch(finishStartupLatch, "task was not allowed to finish startup in time");
            return null;
        }).when(sourceTask).start(eq(TASK_PROPS));

        when(offsetWriter.beginFlush()).thenReturn(false);
        startTaskThread();

        // Stopping immediately while the other thread has work to do should result in no polling, no offset commits,
        // exiting the work thread immediately, and the stop() method will be invoked in the background thread since it
        // cannot be invoked immediately in the thread trying to stop the task.
        ConcurrencyUtils.awaitLatch(startupLatch, "task did not start up in time");
        workerTask.stop();
        finishStartupLatch.countDown();
        awaitShutdown(false);

        verify(sourceTask, never()).poll();
        // task commit called on shutdown
        verifyTransactions(1, 0);
        verifySends(0);

        verifyPreflight();
        verifyStartup();
        verifyCleanShutdown();
    }

    @Test
    public void testCancel() {
        createWorkerTask();

        // workerTask said something dumb on twitter
        workerTask.cancel();

        verify(offsetReader).close();
        verify(producer).close(Duration.ZERO);
    }

    /**
     * @return how many times the source task was {@link SourceTask#poll() polled}
     */
    private int pollCount() {
        return (int) MockitoUtils.countInvocations(sourceTask, "poll");
    }

    /**
     * @return how many times an {@link OffsetStorageWriter#doFlush(Callback) offset flush} took place
     */
    private int flushCount() {
        return (int) MockitoUtils.countInvocations(offsetWriter, "doFlush", Callback.class);
    }

    private void awaitPolls(int minimum, List<SourceRecord> records) {
        pollRecords.set(records);
        pollLatch.set(new CountDownLatch(minimum));
        ConcurrencyUtils.awaitLatch(pollLatch.get(), "task was not polled " + minimum + " time(s) quickly enough");
    }

    private void awaitEmptyPolls(int minimum) {
        awaitPolls(minimum, Collections.emptyList());
    }

    private void awaitPolls(int minimum) {
        awaitPolls(minimum, RECORDS);
    }

    private void awaitShutdown() throws InterruptedException, ExecutionException {
        awaitShutdown(true);
    }

    private void awaitShutdown(boolean triggerStop) throws InterruptedException, ExecutionException {
        if (triggerStop)
            workerTask.stop();
        assertTrue(workerTask.awaitStop(1000));
        workerTaskFuture.get();
    }

    private void startTaskThread() {
        workerTask.initialize(TASK_CONFIG);
        workerTaskFuture = executor.submit(workerTask);
    }

    private void expectSuccessfulFlushes() {
        when(offsetWriter.beginFlush()).thenReturn(true);
        when(offsetWriter.doFlush(any())).thenAnswer(invocation -> {
            Callback<Void> cb = invocation.getArgument(0);
            cb.onCompletion(null, null);
            return null;
        });
    }

    private void expectSuccessfulSends() {
        expectConvertHeadersAndKeyValue();
        expectApplyTransformationChain();
        expectSuccessfulSend(when(producer.send(any(), any())));
        expectTaskGetTopic();
        expectPossibleTopicCreation();
    }

    private OngoingStubbing<Future<RecordMetadata>> expectSuccessfulSend(OngoingStubbing<Future<RecordMetadata>> whenSend) {
        return whenSend.thenAnswer(invocation -> {
            org.apache.kafka.clients.producer.Callback cb = invocation.getArgument(1);
            cb.onCompletion(
                    new RecordMetadata(
                            new TopicPartition(TOPIC, 0),
                            0,
                            0,
                            0L,
                            0,
                            0),
                    null);
            return null;
        });
    }

    private OngoingStubbing<Future<RecordMetadata>> expectSynchronousFailedSend(OngoingStubbing<Future<RecordMetadata>> whenSend, Exception failure) {
        return whenSend.thenThrow(failure);
    }

    private void expectConvertHeadersAndKeyValue() {
        Headers headers = new RecordHeaders();
        for (Header header : headers) {
            when(headerConverter.fromConnectHeader(eq(TOPIC), eq(header.key()), eq(Schema.STRING_SCHEMA), eq(new String(header.value()))))
                    .thenReturn(header.value());
        }
        when(keyConverter.fromConnectData(eq(TOPIC), eq(headers), eq(KEY_SCHEMA), eq(KEY)))
                .thenReturn(SERIALIZED_KEY);
        when(valueConverter.fromConnectData(eq(TOPIC), eq(headers), eq(RECORD_SCHEMA), eq(VALUE_1)))
                .thenReturn(SERIALIZED_RECORD);
    }

    private void expectApplyTransformationChain() {
        when(transformationChain.apply(any(), any()))
                .thenAnswer(invocation -> invocation.getArgument(1));
    }

    private void expectTaskGetTopic() {
        when(statusBackingStore.getTopic(any(), eq(TOPIC)))
                .thenAnswer(invocation -> new TopicStatus(
                        invocation.getArgument(1),
                        new ConnectorTaskId(invocation.getArgument(0), 0),
                        time.milliseconds()
                ));
    }

    private void expectPossibleTopicCreation() {
        if (config.topicCreationEnable()) {
            Set<String> created = Collections.singleton(TOPIC);
            Set<String> existing = Collections.emptySet();
            TopicAdmin.TopicCreationResponse creationResponse = new TopicAdmin.TopicCreationResponse(created, existing);
            when(admin.createOrFindTopics(any())).thenReturn(creationResponse);
        }
    }

    private void verifyPreflight() {
        verify(preProducerCheck).run();
        verify(producer).initTransactions();
        verify(postProducerCheck).run();
        verify(offsetStore).start();
    }

    private void verifyStartup() {
        verify(sourceTask).initialize(any());
        verify(sourceTask).start(TASK_PROPS);
        taskStarted = true;
        verify(statusListener).onStartup(taskId);
    }

    private void verifyCleanShutdown() throws Exception {
        verifyShutdown(false, false);
    }

    private void verifyShutdown(boolean taskFailed, boolean cancelled) throws Exception {
        if (cancelled) {
            verify(producer).close(Duration.ZERO);
            verify(producer, times(2)).close(any());
            verify(offsetReader, times(2)).close();
        } else {
            verify(producer).close(any());
            verify(offsetReader).close();
        }

        verify(offsetStore).stop();
        verify(admin).close(any());
        verify(transformationChain).close();
        verify(headerConverter).close();

        if (taskStarted) {
            verify(sourceTask).stop();
        }

        if (taskFailed) {
            verify(statusListener).onFailure(eq(taskId), any());
        } else if (!cancelled) {
            verify(statusListener).onShutdown(taskId);
        }
    }

    private void verifyPossibleTopicCreation() {
        if (enableTopicCreation) {
            verify(admin).createOrFindTopics(any());
        } else {
            verify(admin, never()).createOrFindTopics(any());
        }
    }

    private void verifySends() {
        verifySends(pollCount() * RECORDS.size());
    }

    private void verifySends(int count) {
        verify(producer, times(count)).send(any(), any());
    }

    private void verifyTransactions(int numTaskCommits, int numProducerCommits) throws InterruptedException {
        // these operations happen on every commit opportunity
        VerificationMode commitOpportunities = times(numTaskCommits);
        verify(offsetWriter, commitOpportunities).beginFlush();
        verify(sourceTask, commitOpportunities).commit();
        // these operations only happen on non-empty commits
        VerificationMode commits = times(numProducerCommits);
        verify(producer, commits).beginTransaction();
        verify(producer, commits).commitTransaction();
        verify(offsetWriter, commits).doFlush(any());
    }

    private void assertTransactionMetrics(int minimumMaxSizeExpected) {
        MetricGroup transactionGroup = workerTask.transactionMetricsGroup().metricGroup();
        double actualMin = metrics.currentMetricValueAsDouble(transactionGroup, "transaction-size-min");
        double actualMax = metrics.currentMetricValueAsDouble(transactionGroup, "transaction-size-max");
        double actualAvg = metrics.currentMetricValueAsDouble(transactionGroup, "transaction-size-avg");
        assertTrue(actualMin >= 0);
        assertTrue(actualMax >= minimumMaxSizeExpected);

        if (actualMax - actualMin <= 0.000001d) {
            assertEquals(actualMax, actualAvg, 0.000002d);
        } else {
            assertTrue("Average transaction size should be greater than minimum transaction size", actualAvg > actualMin);
            assertTrue("Average transaction size should be less than maximum transaction size", actualAvg < actualMax);
        }
    }

    private void assertPollMetrics(int minimumPollCountExpected) {
        MetricGroup sourceTaskGroup = workerTask.sourceTaskMetricsGroup().metricGroup();
        MetricGroup taskGroup = workerTask.taskMetricsGroup().metricGroup();
        double pollRate = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-poll-rate");
        double pollTotal = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-poll-total");
        if (minimumPollCountExpected > 0) {
            assertEquals(RECORDS.size(), metrics.currentMetricValueAsDouble(taskGroup, "batch-size-max"), 0.000001d);
            assertEquals(RECORDS.size(), metrics.currentMetricValueAsDouble(taskGroup, "batch-size-avg"), 0.000001d);
            assertTrue(pollRate > 0.0d);
        } else {
            assertEquals(0.0d, pollRate, 0.0);
        }
        assertTrue(pollTotal >= minimumPollCountExpected);

        double writeRate = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-write-rate");
        double writeTotal = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-write-total");
        if (minimumPollCountExpected > 0) {
            assertTrue(writeRate > 0.0d);
        } else {
            assertEquals(0.0d, writeRate, 0.0);
        }
        assertTrue(writeTotal >= minimumPollCountExpected);

        double pollBatchTimeMax = metrics.currentMetricValueAsDouble(sourceTaskGroup, "poll-batch-max-time-ms");
        double pollBatchTimeAvg = metrics.currentMetricValueAsDouble(sourceTaskGroup, "poll-batch-avg-time-ms");
        if (minimumPollCountExpected > 0) {
            assertTrue(pollBatchTimeMax >= 0.0d);
        }
        assertTrue(Double.isNaN(pollBatchTimeAvg) || pollBatchTimeAvg > 0.0d);
        double activeCount = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-active-count");
        double activeCountMax = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-active-count-max");
        assertEquals(0, activeCount, 0.000001d);
        if (minimumPollCountExpected > 0) {
            assertEquals(RECORDS.size(), activeCountMax, 0.000001d);
        }
    }

    private abstract static class TestSourceTask extends SourceTask {
    }

}
