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

import org.apache.kafka.common.internals.FatalExitError;
import org.apache.kafka.common.utils.Exit;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.test.TestUtils;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

public class ShutdownableThreadTest {

    @AfterEach
    public void tearDown() {
        Exit.resetExitProcedure();
    }

    @Test
    public void testShutdownWhenCalledAfterThreadStart() throws InterruptedException {
        AtomicReference<Optional<Integer>> statusCodeOption = new AtomicReference<>(Optional.empty());
        Exit.setExitProcedure((statusCode, ignored) -> {
            statusCodeOption.set(Optional.of(statusCode));
            // Sleep until interrupted to emulate the fact that `System.exit()` never returns
            Utils.sleep(Long.MAX_VALUE);
            throw new AssertionError();
        });
        CountDownLatch latch = new CountDownLatch(1);
        ShutdownableThread thread = new ShutdownableThread("shutdownable-thread-test") {
            @Override
            public void doWork() {
                latch.countDown();
                throw new FatalExitError();
            }
        };
        thread.start();
        assertTrue(latch.await(10, TimeUnit.SECONDS), "doWork was not invoked");

        thread.shutdown();
        TestUtils.waitForCondition(() -> statusCodeOption.get().isPresent(), "Status code was not set by exit procedure");
        assertEquals(1, statusCodeOption.get().get());
    }

    @Test
    public void testIsThreadStarted() throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(1);
        ShutdownableThread thread = new ShutdownableThread("shutdownable-thread-test") {
            @Override
            public void doWork() {
                latch.countDown();
            }
        };
        assertFalse(thread.isStarted());
        thread.start();
        assertTrue(latch.await(10, TimeUnit.SECONDS), "doWork was not invoked");
        assertTrue(thread.isStarted());

        thread.shutdown();
    }

    @ParameterizedTest
    @ValueSource(booleans = {true, false})
    public void testShutdownWhenTestTimesOut(boolean isInterruptible) {
        // Mask the exit procedure so the contents of the "test" can't shut down the JVM
        Exit.setExitProcedure((statusCode, ignored) -> {
            throw new FatalExitError();
        });
        // This latch will be triggered only after the "test" finishes
        CountDownLatch afterTest = new CountDownLatch(1);
        try {
            // This is the "test", which uses a ShutdownableThread with exit procedures masked.
            CountDownLatch startupLatch = new CountDownLatch(1);
            ShutdownableThread thread = new ShutdownableThread("shutdownable-thread-timeout", isInterruptible) {
                @Override
                public void doWork() {
                    // Tell the test that we finished starting, and are ready to shut down.
                    startupLatch.countDown();
                    if (isInterruptible) {
                        // Swallow the interruption sent by the thread
                        try {
                            afterTest.await();
                        } catch (InterruptedException ignored) {
                        }
                    }
                    // Trigger a fatal exit, after the test has completed and the exit procedure masking has been cleaned up
                    try {
                        afterTest.await();
                    } catch (InterruptedException ignored) {
                    }
                    throw new FatalExitError();
                }
            };
            thread.start();
            startupLatch.await();
            // Interrupt ourselves, as if the test was interrupted for timing out
            Thread.currentThread().interrupt();
            thread.shutdown();
            fail("Shutdown should have been interrupted");
        } catch (InterruptedException ignored) {
            // Swallow the interruption, so that the surrounding test passes
        } finally {
            Exit.resetExitProcedure();
        }
        // After the test has stopped waiting for the thread and cleaned up the Exit procedure, let the thread continue
        afterTest.countDown();
    }
}
