code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.workers;
  
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.mockito.Mockito.doAnswer;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.verifyNoInteractions;
  import static org.mockito.Mockito.when;
  
  import com.fasterxml.jackson.databind.JsonNode;
  import io.airbyte.config.Configs;
  import io.airbyte.config.EnvConfigs;
  import io.airbyte.config.StandardSync;
  import io.airbyte.config.StandardSyncInput;
  import io.airbyte.workers.internal.HeartbeatMonitor;
  import java.time.Duration;
  import java.time.temporal.ChronoUnit;
  import java.util.Map;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.function.BiConsumer;
  import org.apache.commons.lang3.tuple.ImmutablePair;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.DisplayName;
  import org.junit.jupiter.api.Nested;
  import org.junit.jupiter.api.Test;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  @SuppressWarnings("PMD.JUnitTestsShouldIncludeAssert")
  class WorkerUtilsTest {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(GentleCloseWithHeartbeat.class);
  
    @Nested
    class GentleCloseWithHeartbeat {
  
      private final Duration CHECK_HEARTBEAT_DURATION = Duration.of(10, ChronoUnit.MILLIS);
  
      private final Duration SHUTDOWN_TIME_DURATION = Duration.of(100, ChronoUnit.MILLIS);
  
      private Process process;
      private HeartbeatMonitor heartbeatMonitor;
      private BiConsumer<Process, Duration> forceShutdown;
  
      @SuppressWarnings("unchecked")
      @BeforeEach
      void setup() {
        process = mock(Process.class);
        heartbeatMonitor = mock(HeartbeatMonitor.class);
        forceShutdown = mock(BiConsumer.class);
      }
  
      private void runShutdown() {
        gentleCloseWithHeartbeat(
            new WorkerConfigs(new EnvConfigs()),
            process,
            heartbeatMonitor,
            SHUTDOWN_TIME_DURATION,
            CHECK_HEARTBEAT_DURATION,
            SHUTDOWN_TIME_DURATION,
            forceShutdown);
      }
  
      @SuppressWarnings("BusyWait")
      @DisplayName("Verify that shutdown waits indefinitely when heartbeat and process are healthy.")
      @Test
      void testStartsWait() throws InterruptedException {
        when(process.isAlive()).thenReturn(true);
        final AtomicInteger recordedBeats = new AtomicInteger(0);
        doAnswer((ignored) -> {
          recordedBeats.incrementAndGet();
          return true;
        }).when(heartbeatMonitor).isBeating();
  
        final Thread thread = new Thread(this::runShutdown);
  
        thread.start();
  
        // block until the loop is running.
        while (recordedBeats.get() < 3) {
          Thread.sleep(10);
        }
  
        thread.stop();
      }
  
      @Test
      @DisplayName("Test heartbeat ends and graceful shutdown.")
      void testGracefulShutdown() {
        when(heartbeatMonitor.isBeating()).thenReturn(false);
        when(process.isAlive()).thenReturn(false);
  
        runShutdown();
  
        verifyNoInteractions(forceShutdown);
      }
  
      @Test
      @DisplayName("Test heartbeat ends and shutdown is forced.")
      void testForcedShutdown() {
        when(heartbeatMonitor.isBeating()).thenReturn(false);
        when(process.isAlive()).thenReturn(true);
  
        runShutdown();
  
        verify(forceShutdown).accept(process, SHUTDOWN_TIME_DURATION);
      }
  
      @Test
      @DisplayName("Test process dies.")
      void testProcessDies() {
        when(heartbeatMonitor.isBeating()).thenReturn(true);
        when(process.isAlive()).thenReturn(false);
        runShutdown();
  
        verifyNoInteractions(forceShutdown);
      }
  
    }
  
    @Test
    void testMapStreamNamesToSchemasWithNullNamespace() {
      final ImmutablePair<StandardSync, StandardSyncInput> syncPair = TestConfigHelpers.createSyncConfig();
      final StandardSyncInput syncInput = syncPair.getValue();
      final Map<String, JsonNode> mapOutput = WorkerUtils.mapStreamNamesToSchemas(syncInput);
      assertNotNull(mapOutput.get("user_preferences"));
    }
  
    @Test
    void testMapStreamNamesToSchemasWithMultipleNamespaces() {
      final ImmutablePair<StandardSync, StandardSyncInput> syncPair = TestConfigHelpers.createSyncConfig(true);
      final StandardSyncInput syncInput = syncPair.getValue();
      final Map<String, JsonNode> mapOutput = WorkerUtils.mapStreamNamesToSchemas(syncInput);
      assertNotNull(mapOutput.get("namespaceuser_preferences"));
      assertNotNull(mapOutput.get("namespace2user_preferences"));
    }
  
    /**
     * As long as the the heartbeatMonitor detects a heartbeat, the process will be allowed to continue.
     * This method checks the heartbeat once every minute. Once there is no heartbeat detected, if the
     * process has ended, then the method returns. If the process is still running it is given a grace
     * period of the timeout arguments passed into the method. Once those expire the process is killed
     * forcibly. If the process cannot be killed, this method will log that this is the case, but then
     * returns.
     *
     * @param process - process to monitor.
     * @param heartbeatMonitor - tracks if the heart is still beating for the given process.
     * @param gracefulShutdownDuration - grace period to give the process to die after its heart stops
     *        beating.
     * @param checkHeartbeatDuration - frequency with which the heartbeat of the process is checked.
     * @param forcedShutdownDuration - amount of time to wait if a process needs to be destroyed
     *        forcibly.
     */
    static void gentleCloseWithHeartbeat(final WorkerConfigs workerConfigs,
                                         final Process process,
                                         final HeartbeatMonitor heartbeatMonitor,
                                         final Duration gracefulShutdownDuration,
                                         final Duration checkHeartbeatDuration,
                                         final Duration forcedShutdownDuration,
                                         final BiConsumer<Process, Duration> forceShutdown) {
      while (process.isAlive() && heartbeatMonitor.isBeating()) {
        try {
          if (workerConfigs.getWorkerEnvironment().equals(Configs.WorkerEnvironment.KUBERNETES)) {
            LOGGER.debug("Gently closing process {} with heartbeat..", process.info().commandLine().get());
          }
  
          process.waitFor(checkHeartbeatDuration.toMillis(), TimeUnit.MILLISECONDS);
        } catch (final InterruptedException e) {
          LOGGER.error("Exception while waiting for process to finish", e);
        }
      }
  
      if (process.isAlive()) {
        try {
          if (workerConfigs.getWorkerEnvironment().equals(Configs.WorkerEnvironment.KUBERNETES)) {
            LOGGER.debug("Gently closing process {} without heartbeat..", process.info().commandLine().get());
          }
  
          process.waitFor(gracefulShutdownDuration.toMillis(), TimeUnit.MILLISECONDS);
        } catch (final InterruptedException e) {
          LOGGER.error("Exception during grace period for process to finish. This can happen when cancelling jobs.");
        }
      }
  
      // if we were unable to exist gracefully, force shutdown...
      if (process.isAlive()) {
        if (workerConfigs.getWorkerEnvironment().equals(Configs.WorkerEnvironment.KUBERNETES)) {
          LOGGER.debug("Force shutdown process {}..", process.info().commandLine().get());
        }
  
        forceShutdown.accept(process, forcedShutdownDuration);
      }
    }
  
  }
  