/*
 * Copyright 2019 Amazon.com, Inc. or its affiliates.
 * Licensed 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 software.amazon.kinesis.lifecycle;

import java.time.Instant;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.BrokenBarrierException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;
import org.junit.After;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestName;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.runners.MockitoJUnitRunner;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;
import software.amazon.kinesis.common.InitialPositionInStreamExtended;
import software.amazon.kinesis.common.RequestDetails;
import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.lifecycle.ConsumerStates.ShardConsumerState;
import software.amazon.kinesis.lifecycle.events.ProcessRecordsInput;
import software.amazon.kinesis.lifecycle.events.TaskExecutionListenerInput;
import software.amazon.kinesis.retrieval.RecordsDeliveryAck;
import software.amazon.kinesis.retrieval.RecordsPublisher;
import software.amazon.kinesis.retrieval.RecordsRetrieved;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;

import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyLong;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.verifyZeroInteractions;
import static org.mockito.Mockito.when;

/**
 * Unit tests of {@link ShardConsumer}.
 */
@RunWith(MockitoJUnitRunner.class)
@Slf4j
public class ShardConsumerTest {

    private final String shardId = "shardId-0-0";
    private final String concurrencyToken = "TestToken";
    private final RequestDetails lastSuccessfulRequestDetails = new RequestDetails();
    private ShardInfo shardInfo;
    private TaskExecutionListenerInput initialTaskInput;
    private TaskExecutionListenerInput processTaskInput;
    private TaskExecutionListenerInput shutdownTaskInput;
    private TaskExecutionListenerInput shutdownRequestedTaskInput;
    private TaskExecutionListenerInput shutdownRequestedAwaitTaskInput;

    private ExecutorService executorService;

    @Mock
    private RecordsPublisher recordsPublisher;

    @Mock
    private ShutdownNotification shutdownNotification;

    @Mock
    private ConsumerState blockedOnParentsState;

    @Mock
    private ConsumerTask blockedOnParentsTask;

    @Mock
    private ConsumerState initialState;

    @Mock
    private ConsumerTask initializeTask;

    @Mock
    private ConsumerState processingState;

    @Mock
    private ConsumerTask processingTask;

    @Mock
    private ConsumerState shutdownState;

    @Mock
    private ConsumerTask shutdownTask;

    @Mock
    private TaskResult initializeTaskResult;

    @Mock
    private TaskResult processingTaskResult;

    @Mock
    private TaskResult blockOnParentsTaskResult;

    @Mock
    private ConsumerState shutdownCompleteState;

    @Mock
    private ShardConsumerArgument shardConsumerArgument;

    @Mock
    private ConsumerState shutdownRequestedState;

    @Mock
    private ConsumerTask shutdownRequestedTask;

    @Mock
    private ConsumerState shutdownRequestedAwaitState;

    @Mock
    private TaskExecutionListener taskExecutionListener;

    private ProcessRecordsInput processRecordsInput;

    private Optional<Long> logWarningForTaskAfterMillis = Optional.empty();

    @Rule
    public TestName testName = new TestName();

    @Before
    public void before() {
        MockitoAnnotations.initMocks(this);
        shardInfo = new ShardInfo(shardId, concurrencyToken, null, ExtendedSequenceNumber.TRIM_HORIZON);
        ThreadFactory factory = new ThreadFactoryBuilder()
                .setNameFormat("test-" + testName.getMethodName() + "-%04d")
                .setDaemon(true)
                .build();
        executorService = new ThreadPoolExecutor(4, 4, 1, TimeUnit.SECONDS, new LinkedBlockingQueue<>(), factory);

        processRecordsInput = ProcessRecordsInput.builder()
                .isAtShardEnd(false)
                .cacheEntryTime(Instant.now())
                .millisBehindLatest(1000L)
                .records(Collections.emptyList())
                .build();
        initialTaskInput = TaskExecutionListenerInput.builder()
                .shardInfo(shardInfo)
                .taskType(TaskType.INITIALIZE)
                .build();
        processTaskInput = TaskExecutionListenerInput.builder()
                .shardInfo(shardInfo)
                .taskType(TaskType.PROCESS)
                .build();
        shutdownRequestedTaskInput = TaskExecutionListenerInput.builder()
                .shardInfo(shardInfo)
                .taskType(TaskType.SHUTDOWN_NOTIFICATION)
                .build();
        shutdownRequestedAwaitTaskInput = TaskExecutionListenerInput.builder()
                .shardInfo(shardInfo)
                .taskType(TaskType.SHUTDOWN_COMPLETE)
                .build();
        shutdownTaskInput = TaskExecutionListenerInput.builder()
                .shardInfo(shardInfo)
                .taskType(TaskType.SHUTDOWN)
                .build();
    }

    @After
    public void after() {
        List<Runnable> remainder = executorService.shutdownNow();
        assertTrue(remainder.isEmpty());
    }

    private class TestPublisher implements RecordsPublisher {

        final CyclicBarrier barrier = new CyclicBarrier(2);
        final CyclicBarrier requestBarrier = new CyclicBarrier(2);

        Subscriber<? super RecordsRetrieved> subscriber;
        final Subscription subscription = mock(Subscription.class);

        TestPublisher() {
            this(false);
        }

        TestPublisher(boolean enableCancelAwait) {
            doAnswer(a -> {
                        requestBarrier.await();
                        return null;
                    })
                    .when(subscription)
                    .request(anyLong());
            doAnswer(a -> {
                        if (enableCancelAwait) {
                            requestBarrier.await();
                        }
                        return null;
                    })
                    .when(subscription)
                    .cancel();
        }

        @Override
        public void start(
                ExtendedSequenceNumber extendedSequenceNumber,
                InitialPositionInStreamExtended initialPositionInStreamExtended) {}

        @Override
        public void notify(RecordsDeliveryAck ack) {}

        @Override
        public void shutdown() {}

        @Override
        public RequestDetails getLastSuccessfulRequestDetails() {
            return lastSuccessfulRequestDetails;
        }

        @Override
        public void subscribe(Subscriber<? super RecordsRetrieved> s) {
            subscriber = s;
            subscriber.onSubscribe(subscription);
            try {
                barrier.await();
            } catch (Exception ex) {
                throw new RuntimeException(ex);
            }
        }

        @Override
        public void restartFrom(RecordsRetrieved recordsRetrieved) {}

        public void awaitSubscription() throws InterruptedException, BrokenBarrierException {
            barrier.await();
            barrier.reset();
        }

        public void awaitRequest() throws InterruptedException, BrokenBarrierException {
            requestBarrier.await();
            requestBarrier.reset();
        }

        public void awaitInitialSetup() throws InterruptedException, BrokenBarrierException {
            awaitRequest();
            awaitSubscription();
        }

        public void publish() {
            publish(() -> processRecordsInput);
        }

        public void publish(RecordsRetrieved input) {
            subscriber.onNext(input);
        }
    }

    @Test
    public void simpleTest() throws Exception {
        CyclicBarrier taskCallBarrier = new CyclicBarrier(2);

        mockSuccessfulInitialize(null);

        mockSuccessfulProcessing(taskCallBarrier);

        mockSuccessfulShutdown(null);

        TestPublisher cache = new TestPublisher();
        final ShardConsumer consumer = createShardConsumer(cache);

        boolean initComplete = false;
        while (!initComplete) {
            initComplete = consumer.initializeComplete().get();
        }

        consumer.subscribe();
        cache.awaitInitialSetup();

        cache.publish();
        awaitAndResetBarrier(taskCallBarrier);
        cache.awaitRequest();

        cache.publish();
        awaitAndResetBarrier(taskCallBarrier);
        cache.awaitRequest();

        consumer.leaseLost();
        boolean shutdownComplete = consumer.shutdownComplete().get();
        while (!shutdownComplete) {
            shutdownComplete = consumer.shutdownComplete().get();
        }

        verify(cache.subscription, times(3)).request(anyLong());
        verify(cache.subscription).cancel();
        verify(processingState, times(2)).createTask(eq(shardConsumerArgument), eq(consumer), any());
        verify(taskExecutionListener, times(1)).beforeTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(2)).beforeTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(shutdownTaskInput);

        initialTaskInput =
                initialTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        processTaskInput =
                processTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        shutdownTaskInput = shutdownTaskInput.toBuilder()
                .taskOutcome(TaskOutcome.SUCCESSFUL)
                .build();

        verify(taskExecutionListener, times(1)).afterTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(2)).afterTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).afterTaskExecution(shutdownTaskInput);
        verifyNoMoreInteractions(taskExecutionListener);
    }

    @Test(timeout = 1000L)
    public void testLeaseLossIsNonBlocking() throws Exception {
        CyclicBarrier taskCallBarrier = new CyclicBarrier(2);
        CyclicBarrier processingTaskInterlock = new CyclicBarrier(2);

        mockSuccessfulInitialize(null);

        mockSuccessfulProcessing(taskCallBarrier, processingTaskInterlock);

        mockSuccessfulShutdown(null);

        TestPublisher cache = new TestPublisher();
        final ShardConsumer consumer = createShardConsumer(cache);

        boolean initComplete = false;
        while (!initComplete) {
            initComplete = consumer.initializeComplete().get();
        }

        consumer.subscribe();
        cache.awaitInitialSetup();

        log.debug("Setup complete publishing entry");
        cache.publish();
        awaitAndResetBarrier(taskCallBarrier);
        consumer.leaseLost();

        //
        // This will block if a lock is held on ShardConsumer#this
        //
        consumer.executeLifecycle();
        assertFalse(consumer.isShutdown());

        log.debug("Release processing task interlock");
        awaitAndResetBarrier(processingTaskInterlock);

        while (!consumer.isShutdown()) {
            consumer.executeLifecycle();
            Thread.yield();
        }

        verify(cache.subscription, times(1)).request(anyLong());
        verify(cache.subscription).cancel();
        verify(processingState, times(1)).createTask(eq(shardConsumerArgument), eq(consumer), any());
        verify(taskExecutionListener, times(1)).beforeTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(shutdownTaskInput);

        initialTaskInput =
                initialTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        processTaskInput =
                processTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        shutdownTaskInput = shutdownTaskInput.toBuilder()
                .taskOutcome(TaskOutcome.SUCCESSFUL)
                .build();

        verify(taskExecutionListener, times(1)).afterTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(1)).afterTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).afterTaskExecution(shutdownTaskInput);
        verifyNoMoreInteractions(taskExecutionListener);
    }

    @Test
    public void testDataArrivesAfterProcessing2() throws Exception {
        CyclicBarrier taskCallBarrier = new CyclicBarrier(2);

        mockSuccessfulInitialize(null);

        mockSuccessfulProcessing(taskCallBarrier);

        mockSuccessfulShutdown(null);

        TestPublisher cache = new TestPublisher();
        final ShardConsumer consumer = createShardConsumer(cache);

        boolean initComplete = false;
        while (!initComplete) {
            initComplete = consumer.initializeComplete().get();
        }

        consumer.subscribe();
        cache.awaitInitialSetup();

        cache.publish();
        awaitAndResetBarrier(taskCallBarrier);

        verify(processingState).createTask(any(), any(), any());
        verify(processingTask).call();

        cache.awaitRequest();

        cache.publish();
        awaitAndResetBarrier(taskCallBarrier);
        verify(processingState, times(2)).createTask(any(), any(), any());
        verify(processingTask, times(2)).call();

        cache.awaitRequest();

        cache.publish();
        awaitAndResetBarrier(taskCallBarrier);
        cache.awaitRequest();

        log.info("Starting shutdown");
        consumer.leaseLost();
        boolean shutdownComplete;
        do {
            shutdownComplete = consumer.shutdownComplete().get();
        } while (!shutdownComplete);

        verify(processingState, times(3)).createTask(any(), any(), any());
        verify(processingTask, times(3)).call();
        verify(processingState).shutdownTransition(eq(ShutdownReason.LEASE_LOST));
        verify(shutdownState).shutdownTransition(eq(ShutdownReason.LEASE_LOST));
        verify(taskExecutionListener, times(1)).beforeTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(3)).beforeTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(shutdownTaskInput);

        initialTaskInput =
                initialTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        processTaskInput =
                processTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        shutdownTaskInput = shutdownTaskInput.toBuilder()
                .taskOutcome(TaskOutcome.SUCCESSFUL)
                .build();

        verify(taskExecutionListener, times(1)).afterTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(3)).afterTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).afterTaskExecution(shutdownTaskInput);
        verifyNoMoreInteractions(taskExecutionListener);
    }

    @Test
    @Ignore
    public final void testInitializationStateUponFailure() throws Exception {
        final ShardConsumer consumer = createShardConsumer(recordsPublisher);

        when(initialState.createTask(eq(shardConsumerArgument), eq(consumer), any()))
                .thenReturn(initializeTask);
        when(initializeTask.call()).thenReturn(new TaskResult(new Exception("Bad")));
        when(initializeTask.taskType()).thenReturn(TaskType.INITIALIZE);
        when(initialState.failureTransition()).thenReturn(initialState);

        CyclicBarrier taskBarrier = new CyclicBarrier(2);

        when(initialState.requiresDataAvailability()).thenAnswer(i -> {
            taskBarrier.await();
            return false;
        });

        consumer.executeLifecycle();
        for (int i = 0; i < 4; ++i) {
            awaitAndResetBarrier(taskBarrier);
        }

        verify(initialState, times(5)).createTask(eq(shardConsumerArgument), eq(consumer), any());
        verify(initialState, never()).successTransition();
        verify(initialState, never()).shutdownTransition(any());
    }

    /**
     * Test method to verify consumer undergoes the transition WAITING_ON_PARENT_SHARDS -> INITIALIZING -> PROCESSING
     */
    @Test
    public final void testSuccessfulConsumerStateTransition() {
        ExecutorService directExecutorService = spy(executorService);

        doAnswer(this::directlyExecuteRunnable).when(directExecutorService).execute(any());

        final ShardConsumer consumer = createShardConsumer(directExecutorService, blockedOnParentsState);

        mockSuccessfulUnblockOnParents();
        mockSuccessfulInitializeWithFailureTransition();
        mockSuccessfulProcessing(null);

        int arbitraryExecutionCount = 3;
        do {
            try {
                consumer.executeLifecycle();
            } catch (Exception e) {
                // Suppress any exception like the scheduler.
                fail("Unexpected exception while executing consumer lifecycle");
            }
        } while (--arbitraryExecutionCount > 0);

        assertEquals(
                ShardConsumerState.PROCESSING.consumerState().state(),
                consumer.currentState().state());
        verify(directExecutorService, times(2)).execute(any());
    }

    /**
     * Test method to verify consumer does not transition to PROCESSING from WAITING_ON_PARENT_SHARDS when
     * INITIALIZING tasks gets rejected.
     */
    @Test
    public final void testConsumerNotTransitionsToProcessingWhenInitializationFails() {
        ExecutorService failingService = spy(executorService);
        final ShardConsumer consumer = createShardConsumer(failingService, blockedOnParentsState);

        mockSuccessfulUnblockOnParents();
        mockSuccessfulInitializeWithFailureTransition();
        mockSuccessfulProcessing(null);

        // Failing the initialization task and all other attempts after that.
        doAnswer(this::directlyExecuteRunnable)
                .doThrow(new RejectedExecutionException())
                .when(failingService)
                .execute(any());

        int arbitraryExecutionCount = 5;
        do {
            try {
                consumer.executeLifecycle();
            } catch (Exception e) {
                // Suppress any exception like the scheduler.
                fail("Unexpected exception while executing consumer lifecycle");
            }
        } while (--arbitraryExecutionCount > 0);

        assertEquals(
                ShardConsumerState.INITIALIZING.consumerState().state(),
                consumer.currentState().state());
        verify(failingService, times(5)).execute(any());
    }

    /**
     * Test method to verify consumer transition to PROCESSING from WAITING_ON_PARENT_SHARDS with
     * intermittent INITIALIZING task rejections.
     */
    @Test
    public final void testConsumerTransitionsToProcessingWithIntermittentInitializationFailures() {
        ExecutorService failingService = spy(executorService);
        final ShardConsumer consumer = createShardConsumer(failingService, blockedOnParentsState);

        mockSuccessfulUnblockOnParents();
        mockSuccessfulInitializeWithFailureTransition();
        mockSuccessfulProcessing(null);

        // Failing the initialization task and few other attempts after that.
        doAnswer(this::directlyExecuteRunnable)
                .doThrow(new RejectedExecutionException())
                .doThrow(new RejectedExecutionException())
                .doThrow(new RejectedExecutionException())
                .doAnswer(this::directlyExecuteRunnable)
                .when(failingService)
                .execute(any());

        int arbitraryExecutionCount = 6;
        do {
            try {
                consumer.executeLifecycle();
            } catch (Exception e) {
                // Suppress any exception like the scheduler.
                fail("Unexpected exception while executing consumer lifecycle");
            }
        } while (--arbitraryExecutionCount > 0);

        assertEquals(
                ShardConsumerState.PROCESSING.consumerState().state(),
                consumer.currentState().state());
        verify(failingService, times(5)).execute(any());
    }

    /**
     * Test method to verify consumer does not transition to INITIALIZING when WAITING_ON_PARENT_SHARDS task rejected.
     */
    @Test
    public final void testConsumerNotTransitionsToInitializingWhenWaitingOnParentsFails() {
        ExecutorService failingService = spy(executorService);
        final ShardConsumer consumer = createShardConsumer(failingService, blockedOnParentsState);

        mockSuccessfulUnblockOnParentsWithFailureTransition();
        mockSuccessfulInitializeWithFailureTransition();

        // Failing the waiting_on_parents task and few other attempts after that.
        doThrow(new RejectedExecutionException()).when(failingService).execute(any());

        int arbitraryExecutionCount = 5;
        do {
            try {
                consumer.executeLifecycle();
            } catch (Exception e) {
                // Suppress any exception like the scheduler.
                fail("Unexpected exception while executing consumer lifecycle");
            }
        } while (--arbitraryExecutionCount > 0);

        assertEquals(
                ShardConsumerState.WAITING_ON_PARENT_SHARDS.consumerState().state(),
                consumer.currentState().state());
        verify(failingService, times(5)).execute(any());
    }

    /**
     * Test method to verify consumer stays in INITIALIZING state when InitializationTask fails.
     */
    @Test(expected = RejectedExecutionException.class)
    public final void testInitializationStateUponSubmissionFailure() throws Exception {
        ExecutorService failingService = mock(ExecutorService.class);
        final ShardConsumer consumer = createShardConsumer(failingService, initialState);

        doThrow(new RejectedExecutionException()).when(failingService).execute(any());

        boolean initComplete;
        do {
            initComplete = consumer.initializeComplete().get();
        } while (!initComplete);
        verifyZeroInteractions(taskExecutionListener);
    }

    @Test
    public void testErrorThrowableInInitialization() throws Exception {
        final ShardConsumer consumer = createShardConsumer(recordsPublisher);

        when(initialState.createTask(any(), any(), any())).thenReturn(initializeTask);
        when(initialState.taskType()).thenReturn(TaskType.INITIALIZE);
        when(initializeTask.call()).thenAnswer(i -> {
            throw new Error("Error");
        });

        try {
            consumer.initializeComplete().get();
        } catch (ExecutionException ee) {
            assertThat(ee.getCause(), instanceOf(Error.class));
        }
        verify(taskExecutionListener, times(1)).beforeTaskExecution(initialTaskInput);
        verifyNoMoreInteractions(taskExecutionListener);
    }

    @Test
    public void testRequestedShutdownWhileQuiet() throws Exception {
        CyclicBarrier taskBarrier = new CyclicBarrier(2);

        TestPublisher cache = new TestPublisher();
        final ShardConsumer consumer = createShardConsumer(cache);

        mockSuccessfulInitialize(null);

        mockSuccessfulProcessing(taskBarrier);

        when(processingState.shutdownTransition(eq(ShutdownReason.REQUESTED))).thenReturn(shutdownRequestedState);
        when(shutdownRequestedState.createTask(any(), any(), any())).thenReturn(shutdownRequestedTask);
        when(shutdownRequestedState.taskType()).thenReturn(TaskType.SHUTDOWN_NOTIFICATION);
        when(shutdownRequestedTask.call()).thenReturn(new TaskResult(null));

        when(shutdownRequestedState.shutdownTransition(eq(ShutdownReason.REQUESTED)))
                .thenReturn(shutdownRequestedAwaitState);
        when(shutdownRequestedAwaitState.createTask(any(), any(), any())).thenReturn(null);
        when(shutdownRequestedAwaitState.shutdownTransition(eq(ShutdownReason.LEASE_LOST)))
                .thenReturn(shutdownState);
        when(shutdownRequestedAwaitState.taskType()).thenReturn(TaskType.SHUTDOWN_COMPLETE);

        mockSuccessfulShutdown(null);

        boolean init = consumer.initializeComplete().get();
        while (!init) {
            init = consumer.initializeComplete().get();
        }

        consumer.subscribe();
        cache.awaitInitialSetup();

        cache.publish();
        awaitAndResetBarrier(taskBarrier);
        cache.awaitRequest();

        cache.publish();
        awaitAndResetBarrier(taskBarrier);
        cache.awaitRequest();

        consumer.gracefulShutdown(shutdownNotification);
        boolean shutdownComplete = consumer.shutdownComplete().get();
        assertFalse(shutdownComplete);
        shutdownComplete = consumer.shutdownComplete().get();
        assertFalse(shutdownComplete);

        consumer.leaseLost();
        shutdownComplete = consumer.shutdownComplete().get();
        assertFalse(shutdownComplete);
        shutdownComplete = consumer.shutdownComplete().get();
        assertTrue(shutdownComplete);

        verify(processingState, times(2)).createTask(any(), any(), any());
        verify(shutdownRequestedState, never()).shutdownTransition(eq(ShutdownReason.LEASE_LOST));
        verify(shutdownRequestedState).createTask(any(), any(), any());
        verify(shutdownRequestedState).shutdownTransition(eq(ShutdownReason.REQUESTED));
        verify(shutdownRequestedAwaitState).createTask(any(), any(), any());
        verify(shutdownRequestedAwaitState).shutdownTransition(eq(ShutdownReason.LEASE_LOST));
        verify(taskExecutionListener, times(1)).beforeTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(2)).beforeTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(shutdownRequestedTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(shutdownRequestedAwaitTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(shutdownTaskInput);

        initialTaskInput =
                initialTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        processTaskInput =
                processTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        shutdownRequestedTaskInput = shutdownRequestedTaskInput.toBuilder()
                .taskOutcome(TaskOutcome.SUCCESSFUL)
                .build();
        shutdownTaskInput = shutdownTaskInput.toBuilder()
                .taskOutcome(TaskOutcome.SUCCESSFUL)
                .build();
        // No task is created/run for this shutdownRequestedAwaitState, so there's no task outcome.

        // shutdownNotification.shutdownComplete() should only be called for gracefulShutdown
        verify(shutdownNotification, times(1)).shutdownComplete();

        verify(taskExecutionListener, times(1)).afterTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(2)).afterTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).afterTaskExecution(shutdownRequestedTaskInput);
        verify(taskExecutionListener, times(1)).afterTaskExecution(shutdownRequestedAwaitTaskInput);
        verify(taskExecutionListener, times(1)).afterTaskExecution(shutdownTaskInput);
        verifyNoMoreInteractions(taskExecutionListener);
    }

    @Test
    public void testExceptionInProcessingStopsRequests() throws Exception {
        TestPublisher cache = new TestPublisher();

        ShardConsumer consumer = new ShardConsumer(
                cache,
                executorService,
                shardInfo,
                Optional.of(1L),
                shardConsumerArgument,
                initialState,
                1,
                taskExecutionListener,
                0);

        mockSuccessfulInitialize(null);
        mockSuccessfulProcessing(null);

        CyclicBarrier taskCallBarrier = new CyclicBarrier(2);
        final RuntimeException expectedException = new RuntimeException("Whee");
        when(processingTask.call()).thenAnswer(a -> {
            try {
                throw expectedException;
            } finally {
                taskCallBarrier.await();
            }
        });

        boolean initComplete;
        do {
            initComplete = consumer.initializeComplete().get();
        } while (!initComplete);

        consumer.subscribe();
        cache.awaitInitialSetup();

        cache.publish();
        awaitAndResetBarrier(taskCallBarrier);
        cache.awaitRequest();

        Throwable healthCheckOutcome = consumer.healthCheck();

        assertThat(healthCheckOutcome, equalTo(expectedException));

        verify(cache.subscription, times(2)).request(anyLong());
        verify(taskExecutionListener, times(1)).beforeTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(processTaskInput);

        initialTaskInput =
                initialTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();

        verify(taskExecutionListener, times(1)).afterTaskExecution(initialTaskInput);
        verifyNoMoreInteractions(taskExecutionListener);
    }

    @Test
    public void testLongRunningTasks() throws Exception {
        TestPublisher cache = new TestPublisher();

        ShardConsumer consumer = new ShardConsumer(
                cache,
                executorService,
                shardInfo,
                Optional.of(1L),
                shardConsumerArgument,
                initialState,
                1,
                taskExecutionListener,
                0);

        CyclicBarrier taskArriveBarrier = new CyclicBarrier(2);
        CyclicBarrier taskDepartBarrier = new CyclicBarrier(2);

        mockSuccessfulInitialize(taskArriveBarrier, taskDepartBarrier);
        mockSuccessfulProcessing(taskArriveBarrier, taskDepartBarrier);
        mockSuccessfulShutdown(taskArriveBarrier, taskDepartBarrier);

        CompletableFuture<Boolean> initSuccess = consumer.initializeComplete();

        awaitAndResetBarrier(taskArriveBarrier);
        assertNotNull(consumer.taskRunningTime());
        consumer.healthCheck();
        awaitAndResetBarrier(taskDepartBarrier);

        assertFalse(initSuccess.get());
        verify(initializeTask).call();

        initSuccess = consumer.initializeComplete();
        verify(initializeTask).call();
        assertTrue(initSuccess.get());
        consumer.healthCheck();

        assertNull(consumer.taskRunningTime());

        consumer.subscribe();
        cache.awaitInitialSetup();

        cache.publish();

        awaitAndResetBarrier(taskArriveBarrier);
        Instant previousTaskStartTime = consumer.taskDispatchedAt();
        assertNotNull(consumer.taskRunningTime());
        consumer.healthCheck();
        awaitAndResetBarrier(taskDepartBarrier);

        consumer.healthCheck();

        cache.requestBarrier.await();
        assertNull(consumer.taskRunningTime());
        cache.requestBarrier.reset();

        // Sleep for 10 millis before processing next task. If we don't; then the following
        // assertion on time fails. This happens if cache.publish() is executed in the
        // same millisecond as that of previousTaskStartTime, resulting in the unit test failure.
        Thread.sleep(10);
        cache.publish();

        awaitAndResetBarrier(taskArriveBarrier);
        Instant currentTaskStartTime = consumer.taskDispatchedAt();
        assertNotEquals(currentTaskStartTime, previousTaskStartTime);
        awaitAndResetBarrier(taskDepartBarrier);

        cache.requestBarrier.await();
        assertNull(consumer.taskRunningTime());
        cache.requestBarrier.reset();

        consumer.leaseLost();

        assertTrue(consumer.isShutdownRequested());
        CompletableFuture<Boolean> shutdownComplete = consumer.shutdownComplete();

        awaitAndResetBarrier(taskArriveBarrier);
        assertNotNull(consumer.taskRunningTime());
        awaitAndResetBarrier(taskDepartBarrier);

        assertFalse(shutdownComplete.get());

        shutdownComplete = consumer.shutdownComplete();
        assertTrue(shutdownComplete.get());

        assertNull(consumer.taskRunningTime());
        consumer.healthCheck();

        verify(taskExecutionListener, times(1)).beforeTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(2)).beforeTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).beforeTaskExecution(shutdownTaskInput);

        initialTaskInput =
                initialTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        processTaskInput =
                processTaskInput.toBuilder().taskOutcome(TaskOutcome.SUCCESSFUL).build();
        shutdownTaskInput = shutdownTaskInput.toBuilder()
                .taskOutcome(TaskOutcome.SUCCESSFUL)
                .build();

        verify(taskExecutionListener, times(1)).afterTaskExecution(initialTaskInput);
        verify(taskExecutionListener, times(2)).afterTaskExecution(processTaskInput);
        verify(taskExecutionListener, times(1)).afterTaskExecution(shutdownTaskInput);
        verifyNoMoreInteractions(taskExecutionListener);
    }

    @Test
    public void testEmptyShardProcessingRaceCondition() throws Exception {
        final RecordsPublisher mockPublisher = mock(RecordsPublisher.class);
        final ExecutorService mockExecutor = mock(ExecutorService.class);
        final ConsumerState mockState = mock(ConsumerState.class);
        final ShardConsumer consumer = new ShardConsumer(
                mockPublisher,
                mockExecutor,
                shardInfo,
                Optional.of(1L),
                shardConsumerArgument,
                mockState,
                1,
                taskExecutionListener,
                0);

        when(mockState.state()).thenReturn(ShardConsumerState.WAITING_ON_PARENT_SHARDS);
        when(mockState.taskType()).thenReturn(TaskType.BLOCK_ON_PARENT_SHARDS);
        final ConsumerTask mockTask = mock(ConsumerTask.class);
        when(mockState.createTask(any(), any(), any())).thenReturn(mockTask);
        // Simulate successful BlockedOnParent task execution
        // and successful Initialize task execution
        when(mockTask.call()).thenReturn(new TaskResult(false));

        log.info("Scheduler Thread: Invoking ShardConsumer.executeLifecycle() to initiate async"
                + " processing of blocked on parent task");
        consumer.executeLifecycle();
        final ArgumentCaptor<Runnable> taskToExecute = ArgumentCaptor.forClass(Runnable.class);
        verify(mockExecutor, timeout(100)).execute(taskToExecute.capture());
        taskToExecute.getValue().run();
        log.info("RecordProcessor Thread: Simulated successful execution of Blocked on parent task");
        reset(mockExecutor);

        log.info("Scheduler Thread: Invoking ShardConsumer.executeLifecycle() to move to InitializingState"
                + " and initiate async processing of initialize task");
        when(mockState.successTransition()).thenReturn(mockState);
        when(mockState.state()).thenReturn(ShardConsumerState.INITIALIZING);
        when(mockState.taskType()).thenReturn(TaskType.INITIALIZE);
        consumer.executeLifecycle();
        verify(mockExecutor, timeout(100)).execute(taskToExecute.capture());
        log.info("RecordProcessor Thread: Simulated successful execution of Initialize task");
        taskToExecute.getValue().run();

        log.info("Scheduler Thread: Invoking ShardConsumer.executeLifecycle() to move to ProcessingState"
                + " and mark initialization future as complete");
        when(mockState.state()).thenReturn(ShardConsumerState.PROCESSING);
        consumer.executeLifecycle();

        // Simulate the race where
        // scheduler invokes executeLifecycle which performs Publisher.subscribe(subscriber)
        // on recordProcessor thread
        // but before scheduler thread finishes initialization, handleInput is invoked
        // on record processor thread.

        // Since ShardConsumer creates its own instance of subscriber that cannot be mocked
        // this test sequence will appear a little odd.
        // In order to control the order in which execution occurs, lets first invoke
        // handleInput, although this will never happen, since there isn't a way
        // to control the precise timing of the thread execution, this is the best way
        final CountDownLatch processTaskLatch = new CountDownLatch(1);
        new Thread(() -> {
                    reset(mockState);
                    when(mockState.taskType()).thenReturn(TaskType.PROCESS);
                    final ConsumerTask mockProcessTask = mock(ConsumerTask.class);
                    when(mockState.createTask(any(), any(), any())).thenReturn(mockProcessTask);
                    when(mockProcessTask.call()).then(input -> {
                        // first we want to wait for subscribe to be called,
                        // but we cannot control the timing, so wait for 10 seconds
                        // to let the main thread invoke executeLifecyle which
                        // will perform subscribe
                        processTaskLatch.countDown();
                        log.info("Record Processor Thread: Holding shardConsumer lock, waiting for 10 seconds to"
                                + " let subscribe be called by scheduler thread");
                        Thread.sleep(10 * 1000);
                        log.info("RecordProcessor Thread: Done waiting");
                        // then return shard end result
                        log.info(
                                "RecordProcessor Thread: Simulating execution of ProcessTask and returning shard-end result");
                        return new TaskResult(true);
                    });
                    final Subscription mockSubscription = mock(Subscription.class);
                    consumer.handleInput(
                            ProcessRecordsInput.builder().isAtShardEnd(true).build(), mockSubscription);
                })
                .start();

        processTaskLatch.await();

        // invoke executeLifecycle, which should invoke subscribe
        // meanwhile if scheduler tries to acquire the ShardConsumer lock it will
        // be blocked during initialization processing because handleInput was
        // already invoked and will be holding the lock. Thereby creating the
        // race condition we want.
        reset(mockState);
        AtomicBoolean successTransitionCalled = new AtomicBoolean(false);
        AtomicBoolean shutdownTransitionCalled = new AtomicBoolean(false);
        when(mockState.state()).then(input -> {
            if (successTransitionCalled.get() && shutdownTransitionCalled.get()) {
                return ShardConsumerState.SHUTTING_DOWN;
            }
            return ShardConsumerState.PROCESSING;
        });
        log.info("Scheduler Thread: Invoking ShardConsumer.executeLifecycle() to invoke subscribe and"
                + " complete initialization");
        consumer.executeLifecycle();
        log.info("Scheduler Thread: Done initializing the ShardConsumer");

        log.info("Verifying scheduler did not perform shutdown transition during initialization");
        verify(mockState, times(0)).shutdownTransition(any());
    }

    private void mockSuccessfulShutdown(CyclicBarrier taskCallBarrier) {
        mockSuccessfulShutdown(taskCallBarrier, null);
    }

    private void mockSuccessfulShutdown(CyclicBarrier taskArriveBarrier, CyclicBarrier taskDepartBarrier) {
        when(shutdownState.createTask(eq(shardConsumerArgument), any(), any())).thenReturn(shutdownTask);
        when(shutdownState.taskType()).thenReturn(TaskType.SHUTDOWN);
        when(shutdownTask.call()).thenAnswer(i -> {
            awaitBarrier(taskArriveBarrier);
            awaitBarrier(taskDepartBarrier);
            return new TaskResult(null);
        });
        when(shutdownState.shutdownTransition(any())).thenReturn(shutdownCompleteState);
        when(shutdownState.state()).thenReturn(ConsumerStates.ShardConsumerState.SHUTTING_DOWN);

        when(shutdownCompleteState.isTerminal()).thenReturn(true);
    }

    private void mockSuccessfulProcessing(CyclicBarrier taskCallBarrier) {
        mockSuccessfulProcessing(taskCallBarrier, null);
    }

    private void mockSuccessfulProcessing(CyclicBarrier taskCallBarrier, CyclicBarrier taskInterlockBarrier) {
        when(processingState.createTask(eq(shardConsumerArgument), any(), any()))
                .thenReturn(processingTask);
        when(processingState.taskType()).thenReturn(TaskType.PROCESS);
        when(processingTask.taskType()).thenReturn(TaskType.PROCESS);
        when(processingTask.call()).thenAnswer(i -> {
            awaitBarrier(taskCallBarrier);
            awaitBarrier(taskInterlockBarrier);
            return processingTaskResult;
        });
        when(processingTaskResult.getException()).thenReturn(null);
        when(processingState.successTransition()).thenReturn(processingState);
        when(processingState.shutdownTransition(any())).thenReturn(shutdownState);
        when(processingState.state()).thenReturn(ConsumerStates.ShardConsumerState.PROCESSING);
    }

    private void mockSuccessfulInitializeWithFailureTransition() {
        mockSuccessfulInitialize(null, null);
        when(initialState.failureTransition()).thenReturn(initialState);
    }

    private void mockSuccessfulInitialize(CyclicBarrier taskCallBarrier) {
        mockSuccessfulInitialize(taskCallBarrier, null);
    }

    private void mockSuccessfulInitialize(CyclicBarrier taskCallBarrier, CyclicBarrier taskInterlockBarrier) {
        when(initialState.createTask(eq(shardConsumerArgument), any(), any())).thenReturn(initializeTask);
        when(initialState.taskType()).thenReturn(TaskType.INITIALIZE);
        when(initializeTask.taskType()).thenReturn(TaskType.INITIALIZE);
        when(initializeTask.call()).thenAnswer(i -> {
            awaitBarrier(taskCallBarrier);
            awaitBarrier(taskInterlockBarrier);
            return initializeTaskResult;
        });
        when(initializeTaskResult.getException()).thenReturn(null);
        when(initialState.successTransition()).thenReturn(processingState);
        when(initialState.state()).thenReturn(ConsumerStates.ShardConsumerState.INITIALIZING);
    }

    private void mockSuccessfulUnblockOnParentsWithFailureTransition() {
        mockSuccessfulUnblockOnParents();
        when(blockedOnParentsState.failureTransition()).thenReturn(blockedOnParentsState);
    }

    private void mockSuccessfulUnblockOnParents() {
        when(blockedOnParentsState.createTask(eq(shardConsumerArgument), any(), any()))
                .thenReturn(blockedOnParentsTask);
        when(blockedOnParentsState.taskType()).thenReturn(TaskType.BLOCK_ON_PARENT_SHARDS);
        when(blockedOnParentsTask.call()).thenAnswer(i -> blockOnParentsTaskResult);
        when(blockOnParentsTaskResult.getException()).thenReturn(null);
        when(blockedOnParentsState.successTransition()).thenReturn(initialState);
        when(blockedOnParentsState.state()).thenReturn(ShardConsumerState.WAITING_ON_PARENT_SHARDS);
    }

    private void awaitBarrier(CyclicBarrier barrier) throws Exception {
        if (barrier != null) {
            barrier.await();
        }
    }

    private void awaitAndResetBarrier(CyclicBarrier barrier) throws Exception {
        barrier.await();
        barrier.reset();
    }

    private Object directlyExecuteRunnable(InvocationOnMock invocation) {
        Object[] args = invocation.getArguments();
        Runnable runnable = (Runnable) args[0];
        runnable.run();
        return null;
    }

    private ShardConsumer createShardConsumer(final RecordsPublisher publisher) {
        return createShardConsumer(publisher, executorService, initialState);
    }

    private ShardConsumer createShardConsumer(final ExecutorService executorService, final ConsumerState state) {
        return createShardConsumer(recordsPublisher, executorService, state);
    }

    private ShardConsumer createShardConsumer(
            final RecordsPublisher publisher, final ExecutorService executorService, final ConsumerState state) {
        return new ShardConsumer(
                publisher,
                executorService,
                shardInfo,
                logWarningForTaskAfterMillis,
                shardConsumerArgument,
                state,
                1,
                taskExecutionListener,
                0);
    }
}
