/*
 * 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.flink.streaming.connectors.kafka.internal;

import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.core.testutils.MultiShotLatch;
import org.apache.flink.core.testutils.OneShotLatch;
import org.apache.flink.metrics.groups.UnregisteredMetricsGroup;
import org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.connectors.kafka.internals.KafkaCommitCallback;
import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition;
import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartitionStateSentinel;
import org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService;
import org.apache.flink.streaming.util.serialization.KeyedDeserializationSchema;
import org.apache.flink.streaming.util.serialization.KeyedDeserializationSchemaWrapper;

import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.consumer.OffsetCommitCallback;
import org.apache.kafka.common.TopicPartition;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.ReentrantLock;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.anyLong;
import static org.powermock.api.mockito.PowerMockito.doAnswer;
import static org.powermock.api.mockito.PowerMockito.mock;
import static org.powermock.api.mockito.PowerMockito.when;
import static org.powermock.api.mockito.PowerMockito.whenNew;

/**
 * Unit tests for the {@link Kafka09Fetcher}.
 */
@RunWith(PowerMockRunner.class)
@PrepareForTest(KafkaConsumerThread.class)
public class Kafka09FetcherTest {

	@Test
	public void testCommitDoesNotBlock() throws Exception {

		// test data
		final KafkaTopicPartition testPartition = new KafkaTopicPartition("test", 42);
		final Map<KafkaTopicPartition, Long> testCommitData = new HashMap<>();
		testCommitData.put(testPartition, 11L);

		// to synchronize when the consumer is in its blocking method
		final OneShotLatch sync = new OneShotLatch();

		// ----- the mock consumer with blocking poll calls ----
		final MultiShotLatch blockerLatch = new MultiShotLatch();

		KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
		when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {

			@Override
			public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) throws InterruptedException {
				sync.trigger();
				blockerLatch.await();
				return ConsumerRecords.empty();
			}
		});

		doAnswer(new Answer<Void>() {
			@Override
			public Void answer(InvocationOnMock invocation) {
				blockerLatch.trigger();
				return null;
			}
		}).when(mockConsumer).wakeup();

		// make sure the fetcher creates the mock consumer
		whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

		// ----- create the test fetcher -----

		@SuppressWarnings("unchecked")
		SourceContext<String> sourceContext = mock(SourceContext.class);
		Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
			Collections.singletonMap(new KafkaTopicPartition("test", 42), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
		KeyedDeserializationSchema<String> schema = new KeyedDeserializationSchemaWrapper<>(new SimpleStringSchema());

		final Kafka09Fetcher<String> fetcher = new Kafka09Fetcher<>(
				sourceContext,
				partitionsWithInitialOffsets,
				null,
				null, /* periodic watermark extractor */
				null, /* punctuated watermark extractor */
				new TestProcessingTimeService(),
				10, /* watermark interval */
				this.getClass().getClassLoader(),
				"task_name",
				schema,
				new Properties(),
				0L,
				new UnregisteredMetricsGroup(),
				new UnregisteredMetricsGroup(),
				false);

		// ----- run the fetcher -----

		final AtomicReference<Throwable> error = new AtomicReference<>();
		final Thread fetcherRunner = new Thread("fetcher runner") {

			@Override
			public void run() {
				try {
					fetcher.runFetchLoop(false);
				} catch (Throwable t) {
					error.set(t);
				}
			}
		};
		fetcherRunner.start();

		// wait until the fetcher has reached the method of interest
		sync.await();

		// ----- trigger the offset commit -----

		final AtomicReference<Throwable> commitError = new AtomicReference<>();
		final Thread committer = new Thread("committer runner") {
			@Override
			public void run() {
				try {
					fetcher.commitInternalOffsetsToKafka(testCommitData, mock(KafkaCommitCallback.class));
				} catch (Throwable t) {
					commitError.set(t);
				}
			}
		};
		committer.start();

		// ----- ensure that the committer finishes in time  -----
		committer.join(30000);
		assertFalse("The committer did not finish in time", committer.isAlive());

		// ----- test done, wait till the fetcher is done for a clean shutdown -----
		fetcher.cancel();
		fetcherRunner.join();

		// check that there were no errors in the fetcher
		final Throwable fetcherError = error.get();
		if (fetcherError != null && !(fetcherError instanceof Handover.ClosedException)) {
			throw new Exception("Exception in the fetcher", fetcherError);
		}
		final Throwable committerError = commitError.get();
		if (committerError != null) {
			throw new Exception("Exception in the committer", committerError);
		}
	}

	@Test
	public void ensureOffsetsGetCommitted() throws Exception {

		// test data
		final KafkaTopicPartition testPartition1 = new KafkaTopicPartition("test", 42);
		final KafkaTopicPartition testPartition2 = new KafkaTopicPartition("another", 99);

		final Map<KafkaTopicPartition, Long> testCommitData1 = new HashMap<>();
		testCommitData1.put(testPartition1, 11L);
		testCommitData1.put(testPartition2, 18L);

		final Map<KafkaTopicPartition, Long> testCommitData2 = new HashMap<>();
		testCommitData2.put(testPartition1, 19L);
		testCommitData2.put(testPartition2, 28L);

		final BlockingQueue<Map<TopicPartition, OffsetAndMetadata>> commitStore = new LinkedBlockingQueue<>();

		// ----- the mock consumer with poll(), wakeup(), and commit(A)sync calls ----

		final MultiShotLatch blockerLatch = new MultiShotLatch();

		KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);

		when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {
			@Override
			public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) throws InterruptedException {
				blockerLatch.await();
				return ConsumerRecords.empty();
			}
		});

		doAnswer(new Answer<Void>() {
			@Override
			public Void answer(InvocationOnMock invocation) {
				blockerLatch.trigger();
				return null;
			}
		}).when(mockConsumer).wakeup();

		doAnswer(new Answer<Void>() {
			@Override
			public Void answer(InvocationOnMock invocation) {
				@SuppressWarnings("unchecked")
				Map<TopicPartition, OffsetAndMetadata> offsets =
						(Map<TopicPartition, OffsetAndMetadata>) invocation.getArguments()[0];

				OffsetCommitCallback callback = (OffsetCommitCallback) invocation.getArguments()[1];

				commitStore.add(offsets);
				callback.onComplete(offsets, null);

				return null;
			}
		}).when(mockConsumer).commitAsync(
				Mockito.<Map<TopicPartition, OffsetAndMetadata>>any(), any(OffsetCommitCallback.class));

		// make sure the fetcher creates the mock consumer
		whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

		// ----- create the test fetcher -----

		@SuppressWarnings("unchecked")
		SourceContext<String> sourceContext = mock(SourceContext.class);
		Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
			Collections.singletonMap(new KafkaTopicPartition("test", 42), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
		KeyedDeserializationSchema<String> schema = new KeyedDeserializationSchemaWrapper<>(new SimpleStringSchema());

		final Kafka09Fetcher<String> fetcher = new Kafka09Fetcher<>(
				sourceContext,
				partitionsWithInitialOffsets,
				null,
				null, /* periodic watermark extractor */
				null, /* punctuated watermark extractor */
				new TestProcessingTimeService(),
				10, /* watermark interval */
				this.getClass().getClassLoader(),
				"task_name",
				schema,
				new Properties(),
				0L,
				new UnregisteredMetricsGroup(),
				new UnregisteredMetricsGroup(),
				false);

		// ----- run the fetcher -----

		final AtomicReference<Throwable> error = new AtomicReference<>();
		final Thread fetcherRunner = new Thread("fetcher runner") {

			@Override
			public void run() {
				try {
					fetcher.runFetchLoop(false);
				} catch (Throwable t) {
					error.set(t);
				}
			}
		};
		fetcherRunner.start();

		// ----- trigger the first offset commit -----

		fetcher.commitInternalOffsetsToKafka(testCommitData1, mock(KafkaCommitCallback.class));
		Map<TopicPartition, OffsetAndMetadata> result1 = commitStore.take();

		for (Entry<TopicPartition, OffsetAndMetadata> entry : result1.entrySet()) {
			TopicPartition partition = entry.getKey();
			if (partition.topic().equals("test")) {
				assertEquals(42, partition.partition());
				assertEquals(12L, entry.getValue().offset());
			}
			else if (partition.topic().equals("another")) {
				assertEquals(99, partition.partition());
				assertEquals(17L, entry.getValue().offset());
			}
		}

		// ----- trigger the second offset commit -----

		fetcher.commitInternalOffsetsToKafka(testCommitData2, mock(KafkaCommitCallback.class));
		Map<TopicPartition, OffsetAndMetadata> result2 = commitStore.take();

		for (Entry<TopicPartition, OffsetAndMetadata> entry : result2.entrySet()) {
			TopicPartition partition = entry.getKey();
			if (partition.topic().equals("test")) {
				assertEquals(42, partition.partition());
				assertEquals(20L, entry.getValue().offset());
			}
			else if (partition.topic().equals("another")) {
				assertEquals(99, partition.partition());
				assertEquals(27L, entry.getValue().offset());
			}
		}

		// ----- test done, wait till the fetcher is done for a clean shutdown -----
		fetcher.cancel();
		fetcherRunner.join();

		// check that there were no errors in the fetcher
		final Throwable caughtError = error.get();
		if (caughtError != null && !(caughtError instanceof Handover.ClosedException)) {
			throw new Exception("Exception in the fetcher", caughtError);
		}
	}

	@Test
	public void testCancellationWhenEmitBlocks() throws Exception {

		// ----- some test data -----

		final String topic = "test-topic";
		final int partition = 3;
		final byte[] payload = new byte[] {1, 2, 3, 4};

		final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(
				new ConsumerRecord<>(topic, partition, 15, payload, payload),
				new ConsumerRecord<>(topic, partition, 16, payload, payload),
				new ConsumerRecord<>(topic, partition, 17, payload, payload));

		final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
		data.put(new TopicPartition(topic, partition), records);

		final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);

		// ----- the test consumer -----

		final KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
		when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {
			@Override
			public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) {
				return consumerRecords;
			}
		});

		whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);

		// ----- build a fetcher -----

		BlockingSourceContext<String> sourceContext = new BlockingSourceContext<>();
		Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
			Collections.singletonMap(new KafkaTopicPartition(topic, partition), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
		KeyedDeserializationSchema<String> schema = new KeyedDeserializationSchemaWrapper<>(new SimpleStringSchema());

		final Kafka09Fetcher<String> fetcher = new Kafka09Fetcher<>(
				sourceContext,
				partitionsWithInitialOffsets,
				null,
				null, /* periodic watermark extractor */
				null, /* punctuated watermark extractor */
				new TestProcessingTimeService(),
				10, /* watermark interval */
				this.getClass().getClassLoader(),
				"task_name",
				schema,
				new Properties(),
				0L,
				new UnregisteredMetricsGroup(),
				new UnregisteredMetricsGroup(),
				false);

		// ----- run the fetcher -----

		final AtomicReference<Throwable> error = new AtomicReference<>();
		final Thread fetcherRunner = new Thread("fetcher runner") {

			@Override
			public void run() {
				try {
					fetcher.runFetchLoop(false);
				} catch (Throwable t) {
					error.set(t);
				}
			}
		};
		fetcherRunner.start();

		// wait until the thread started to emit records to the source context
		sourceContext.waitTillHasBlocker();

		// now we try to cancel the fetcher, including the interruption usually done on the task thread
		// once it has finished, there must be no more thread blocked on the source context
		fetcher.cancel();
		fetcherRunner.interrupt();
		fetcherRunner.join();

		assertFalse("fetcher threads did not properly finish", sourceContext.isStillBlocking());
	}

	// ------------------------------------------------------------------------
	//  test utilities
	// ------------------------------------------------------------------------

	private static final class BlockingSourceContext<T> implements SourceContext<T> {

		private final ReentrantLock lock = new ReentrantLock();
		private final OneShotLatch inBlocking = new OneShotLatch();

		@Override
		public void collect(T element) {
			block();
		}

		@Override
		public void collectWithTimestamp(T element, long timestamp) {
			block();
		}

		@Override
		public void emitWatermark(Watermark mark) {
			block();
		}

		@Override
		public void markAsTemporarilyIdle() {
			throw new UnsupportedOperationException();
		}

		@Override
		public Object getCheckpointLock() {
			return new Object();
		}

		@Override
		public void close() {}

		void waitTillHasBlocker() throws InterruptedException {
			inBlocking.await();
		}

		boolean isStillBlocking() {
			return lock.isLocked();
		}

		@SuppressWarnings({"InfiniteLoopStatement", "SynchronizationOnLocalVariableOrMethodParameter"})
		private void block() {
			lock.lock();
			try {
				inBlocking.trigger();

				// put this thread to sleep indefinitely
				final Object o = new Object();
				while (true) {
					synchronized (o) {
						o.wait();
					}
				}
			}
			catch (InterruptedException e) {
				// exit cleanly, simply reset the interruption flag
				Thread.currentThread().interrupt();
			}
			finally {
				lock.unlock();
			}
		}
	}
}