content: /**
   * 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 org.apache.aurora.scheduler.updater;
  
  import java.util.Optional;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.function.Consumer;
  
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.ImmutableMap;
  
  import org.apache.aurora.common.quantity.Time;
  import org.apache.aurora.common.testing.easymock.EasyMockTest;
  import org.apache.aurora.common.util.BackoffStrategy;
  import org.apache.aurora.gen.InstanceTaskConfig;
  import org.apache.aurora.gen.JobInstanceUpdateEvent;
  import org.apache.aurora.gen.JobUpdate;
  import org.apache.aurora.gen.JobUpdateAction;
  import org.apache.aurora.gen.JobUpdateDetails;
  import org.apache.aurora.gen.JobUpdateEvent;
  import org.apache.aurora.gen.JobUpdateInstructions;
  import org.apache.aurora.gen.JobUpdateKey;
  import org.apache.aurora.gen.JobUpdateSettings;
  import org.apache.aurora.gen.JobUpdateStatus;
  import org.apache.aurora.gen.PercentageSlaPolicy;
  import org.apache.aurora.gen.ScheduleStatus;
  import org.apache.aurora.gen.SlaPolicy;
  import org.apache.aurora.scheduler.base.InstanceKeys;
  import org.apache.aurora.scheduler.base.JobKeys;
  import org.apache.aurora.scheduler.base.TaskTestUtil;
  import org.apache.aurora.scheduler.config.types.TimeAmount;
  import org.apache.aurora.scheduler.sla.SlaManager;
  import org.apache.aurora.scheduler.storage.Storage;
  import org.apache.aurora.scheduler.storage.entities.IAssignedTask;
  import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
  import org.apache.aurora.scheduler.storage.entities.IJobInstanceUpdateEvent;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateDetails;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.apache.aurora.scheduler.storage.entities.ISlaPolicy;
  import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
  import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
  import org.apache.aurora.scheduler.testing.FakeScheduledExecutor;
  import org.apache.aurora.scheduler.testing.FakeStatsProvider;
  import org.apache.aurora.scheduler.updater.UpdaterModule.UpdateActionBatchWorker;
  import org.easymock.Capture;
  import org.junit.Before;
  import org.junit.Test;
  
  import static org.apache.aurora.scheduler.base.TaskTestUtil.JOB;
  import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask;
  import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
  import static org.easymock.EasyMock.anyObject;
  import static org.easymock.EasyMock.capture;
  import static org.easymock.EasyMock.eq;
  import static org.easymock.EasyMock.expect;
  import static org.easymock.EasyMock.expectLastCall;
  import static org.easymock.EasyMock.newCapture;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  
  public class SlaKillControllerTest extends EasyMockTest {
  
    private static final ITaskConfig OLD_CONFIG = TaskTestUtil.makeConfig(JOB);
    private static final SlaPolicy TEST_SLA_POLICY = SlaPolicy.percentageSlaPolicy(
        new PercentageSlaPolicy()
            .setPercentage(0)
            .setDurationSecs(0));
    private static final ITaskConfig NEW_CONFIG = ITaskConfig.build(
        TaskTestUtil.makeConfig(JOB).newBuilder().setSlaPolicy(TEST_SLA_POLICY));
    private static final IScheduledTask TASK = IScheduledTask.build(
        makeTask("id", OLD_CONFIG).newBuilder().setStatus(ScheduleStatus.RUNNING));
    private static final IAssignedTask ASSIGNED_TASK = TASK.getAssignedTask();
    private static final IInstanceKey INSTANCE_KEY = InstanceKeys.from(
        JOB,
        ASSIGNED_TASK.getInstanceId());
    private static final IJobUpdate UPDATE = IJobUpdate.build(
        new JobUpdate()
            .setInstructions(new JobUpdateInstructions()
                .setDesiredState(new InstanceTaskConfig()
                    .setTask(NEW_CONFIG.newBuilder()))
                .setSettings(new JobUpdateSettings()
                    .setSlaAware(true))));
    private static final IJobUpdateKey UPDATE_ID =
        IJobUpdateKey.build(new JobUpdateKey(JOB.newBuilder(), "update_id"));
    private static final String KILL_ATTEMPTS_STAT_NAME = SlaKillController.SLA_KILL_ATTEMPT
        + JobKeys.canonicalString(JOB);
    private static final String KILL_SUCCESSES_STAT_NAME = SlaKillController.SLA_KILL_SUCCESS
        + JobKeys.canonicalString(JOB);
  
    private StorageTestUtil storageUtil;
    private UpdateActionBatchWorker batchWorker;
    private SlaManager slaManager;
    private FakeScheduledExecutor clock;
    private BackoffStrategy backoffStrategy;
    private FakeStatsProvider statsProvider;
    private SlaKillController slaKillController;
    private CountDownLatch killCommandHasExecuted;
    private Consumer<Storage.MutableStoreProvider> fakeKillCommand;
  
    @Before
    public void setUp() {
      storageUtil = new StorageTestUtil(this);
      storageUtil.expectOperations();
      batchWorker = createMock(UpdateActionBatchWorker.class);
      slaManager = createMock(SlaManager.class);
      ScheduledExecutorService executor = createMock(ScheduledExecutorService.class);
      clock = FakeScheduledExecutor.scheduleExecutor(executor);
      backoffStrategy = createMock(BackoffStrategy.class);
      statsProvider = new FakeStatsProvider();
      slaKillController = new SlaKillController(
          executor,
          batchWorker,
          slaManager,
          clock,
          backoffStrategy,
          statsProvider);
      killCommandHasExecuted = new CountDownLatch(2);
      fakeKillCommand = mutableStoreProvider -> killCommandHasExecuted.countDown();
    }
  
    @Test
    public <T, E extends Exception> void testSlaKill() throws Exception {
      IJobUpdateDetails updateDetails = IJobUpdateDetails.build(
          new JobUpdateDetails(
              UPDATE.newBuilder(),
              ImmutableList.of(new JobUpdateEvent(JobUpdateStatus.ROLLING_FORWARD, 123L)),
              ImmutableList.of()));
      expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_ID))
          .andReturn(Optional.of(updateDetails))
          .anyTimes();
      Capture<IJobInstanceUpdateEvent> instanceUpdateEventCapture = newCapture();
      storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(
          eq(UPDATE_ID),
          capture(instanceUpdateEventCapture));
      expectLastCall().times(2);
      storageUtil.expectTaskFetch(TASK.getAssignedTask().getTaskId(), TASK);
      Capture<Storage.MutateWork<T, E>> workCapture = createCapture();
      slaManager.checkSlaThenAct(
          eq(TASK),
          eq(ISlaPolicy.build(TEST_SLA_POLICY)),
          capture(workCapture),
          eq(ImmutableMap.of()),
          eq(false));
      expectBatchExecute(batchWorker, storageUtil.storage, control);
      expect(backoffStrategy.calculateBackoffMs(0)).andReturn(42L);
  
      control.replay();
  
      // Kill command has not been executed yet
      assertEquals(2, killCommandHasExecuted.getCount());
  
      // Start an SLA-aware kill
      slaKillController.slaKill(
          storageUtil.mutableStoreProvider,
          INSTANCE_KEY,
          TASK,
          UPDATE_ID,
          ISlaPolicy.build(TEST_SLA_POLICY),
          JobUpdateStatus.ROLLING_FORWARD,
          fakeKillCommand);
  
      // Ensure the SLA_CHECKING_MESSAGE message is added
      assertTrue(
          checkInstanceEventMatches(
              instanceUpdateEventCapture.getValue(),
              INSTANCE_KEY,
              JobUpdateAction.INSTANCE_UPDATING,
              SlaKillController.SLA_CHECKING_MESSAGE));
      instanceUpdateEventCapture.reset();
      assertFalse(instanceUpdateEventCapture.hasCaptured());
  
      // Pretend SLA passes, executes work
      workCapture.getValue().apply(storageUtil.mutableStoreProvider);
      assertEquals(1, killCommandHasExecuted.getCount());
  
      // Ensure the SLA_PASSED_MESSAGE message is added
      assertTrue(
          checkInstanceEventMatches(
              instanceUpdateEventCapture.getValue(),
              INSTANCE_KEY,
              JobUpdateAction.INSTANCE_UPDATING,
              SlaKillController.SLA_PASSED_MESSAGE));
    }
  
    /**
     * Test that SLA kills are retried in case the SLA check does not pass.
     */
    @Test
    public <T, E extends Exception> void testSlaKillRetry() throws Exception {
      IJobUpdateDetails updateDetails = IJobUpdateDetails.build(
          new JobUpdateDetails(
              UPDATE.newBuilder(),
              ImmutableList.of(new JobUpdateEvent(JobUpdateStatus.ROLLING_FORWARD, 123L)),
              ImmutableList.of()));
      expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_ID))
          .andReturn(Optional.of(updateDetails))
          .anyTimes();
      storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(eq(UPDATE_ID), anyObject());
      expectLastCall().times(2);
      storageUtil.expectTaskFetch(TASK.getAssignedTask().getTaskId(), TASK).times(2);
      storageUtil.expectTaskFetch(
          TASK.getAssignedTask().getTaskId(),
          IScheduledTask.build(TASK.newBuilder().setStatus(ScheduleStatus.KILLING)));
      Capture<Storage.MutateWork<T, E>> workCapture = createCapture();
      slaManager.checkSlaThenAct(
          eq(TASK),
          eq(ISlaPolicy.build(TEST_SLA_POLICY)),
          capture(workCapture),
          eq(ImmutableMap.of()),
          eq(false));
      expectLastCall().times(2);
      expectBatchExecute(batchWorker, storageUtil.storage, control).times(3);
      expect(backoffStrategy.calculateBackoffMs(0L)).andReturn(42L);
      expect(backoffStrategy.calculateBackoffMs(42L)).andReturn(84L);
  
      control.replay();
  
      // Kill command has not been executed yet
      assertFalse(statsProvider.getAllValues().keySet().contains(KILL_ATTEMPTS_STAT_NAME));
      assertFalse(statsProvider.getAllValues().keySet().contains(KILL_SUCCESSES_STAT_NAME));
      assertFalse(workCapture.hasCaptured());
      assertEquals(killCommandHasExecuted.getCount(), 2);
  
      // Start an SLA-aware kill
      slaKillController.slaKill(
          storageUtil.mutableStoreProvider,
          INSTANCE_KEY,
          TASK,
          UPDATE_ID,
          ISlaPolicy.build(TEST_SLA_POLICY),
          JobUpdateStatus.ROLLING_FORWARD,
          fakeKillCommand);
  
      // SLA check is called and discarded, pretending it failed
      assertEquals(1, statsProvider.getLongValue(KILL_ATTEMPTS_STAT_NAME));
      assertFalse(statsProvider.getAllValues().keySet().contains(KILL_SUCCESSES_STAT_NAME));
      assertTrue(workCapture.hasCaptured());
      workCapture.reset();
      assertEquals(2, killCommandHasExecuted.getCount());
      assertFalse(workCapture.hasCaptured());
  
      // Another SLA kill is scheduled assuming the previous attempt failed
      assertEquals(1, clock.countDeferredWork());
      clock.advance(TimeAmount.of(42L, Time.MILLISECONDS));
  
      // The second SLA check passes and the kill function is called
      assertEquals(2, killCommandHasExecuted.getCount());
      workCapture.getValue().apply(storageUtil.mutableStoreProvider);
      assertEquals(1, killCommandHasExecuted.getCount());
      assertEquals(2, statsProvider.getLongValue(KILL_ATTEMPTS_STAT_NAME));
      assertEquals(1, statsProvider.getLongValue(KILL_SUCCESSES_STAT_NAME));
  
      // One more SLA kill is scheduled assuming the previous attempt failed. Since the previous
      // attempt did not fail, we do a NOOP since the task is already KILLING
      clock.advance(TimeAmount.of(84L, Time.MILLISECONDS));
      assertEquals(2, statsProvider.getLongValue(KILL_ATTEMPTS_STAT_NAME));
      assertEquals(1, statsProvider.getLongValue(KILL_SUCCESSES_STAT_NAME));
      assertEquals(1, killCommandHasExecuted.getCount());
      assertEquals(0, clock.countDeferredWork());
    }
  
    @Test
    public <T, E extends Exception> void testSlaKillRollingBack() throws Exception {
      IJobUpdateDetails updateDetails = IJobUpdateDetails.build(
          new JobUpdateDetails(
              UPDATE.newBuilder(),
              ImmutableList.of(new JobUpdateEvent(JobUpdateStatus.ROLLING_BACK, 123L)),
              ImmutableList.of()));
      expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_ID))
          .andReturn(Optional.of(updateDetails))
          .anyTimes();
      Capture<IJobInstanceUpdateEvent> instanceUpdateEventCapture = newCapture();
      storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(
          eq(UPDATE_ID),
          capture(instanceUpdateEventCapture));
      expectLastCall().times(2);
      storageUtil.expectTaskFetch(TASK.getAssignedTask().getTaskId(), TASK);
      Capture<Storage.MutateWork<T, E>> workCapture = createCapture();
      slaManager.checkSlaThenAct(
          eq(TASK),
          eq(ISlaPolicy.build(TEST_SLA_POLICY)),
          capture(workCapture),
          eq(ImmutableMap.of()),
          eq(false));
      expectBatchExecute(batchWorker, storageUtil.storage, control);
      expect(backoffStrategy.calculateBackoffMs(0)).andReturn(42L);
  
      control.replay();
  
      // Kill command has not been executed yet
      assertEquals(2, killCommandHasExecuted.getCount());
  
      // Start an SLA-aware kill
      slaKillController.slaKill(
          storageUtil.mutableStoreProvider,
          INSTANCE_KEY,
          TASK,
          UPDATE_ID,
          ISlaPolicy.build(TEST_SLA_POLICY),
          JobUpdateStatus.ROLLING_BACK,
          fakeKillCommand);
  
      // Ensure the SLA_CHECKING_MESSAGE message is added with ROLLING_BACK action
      assertTrue(
          checkInstanceEventMatches(
              instanceUpdateEventCapture.getValue(),
              INSTANCE_KEY,
              JobUpdateAction.INSTANCE_ROLLING_BACK,
              SlaKillController.SLA_CHECKING_MESSAGE));
      instanceUpdateEventCapture.reset();
      assertFalse(instanceUpdateEventCapture.hasCaptured());
  
      // Pretend SLA passes, executes work
      workCapture.getValue().apply(storageUtil.mutableStoreProvider);
      assertEquals(1, killCommandHasExecuted.getCount());
  
      // Ensure the SLA_PASSED_MESSAGE message is added with ROLLING_BACK action
      assertTrue(
          checkInstanceEventMatches(
              instanceUpdateEventCapture.getValue(),
              INSTANCE_KEY,
              JobUpdateAction.INSTANCE_ROLLING_BACK,
              SlaKillController.SLA_PASSED_MESSAGE));
    }
  
    @Test
    public void testSlaKillFailOnPause() throws Exception {
      IJobUpdateDetails updateDetails = IJobUpdateDetails.build(
          new JobUpdateDetails(
              UPDATE.newBuilder(),
              ImmutableList.of(new JobUpdateEvent(JobUpdateStatus.ROLLING_FORWARD, 123L)),
              ImmutableList.of()));
      IJobUpdateDetails pausedUpdateDetails = IJobUpdateDetails.build(
          new JobUpdateDetails(
              UPDATE.newBuilder(),
              ImmutableList.of(new JobUpdateEvent(JobUpdateStatus.ROLL_FORWARD_PAUSED, 123L)),
              ImmutableList.of()));
      expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_ID))
          .andReturn(Optional.of(updateDetails));
      storageUtil.jobUpdateStore.saveJobInstanceUpdateEvent(eq(UPDATE_ID), anyObject());
      expectLastCall();
      expectBatchExecute(batchWorker, storageUtil.storage, control);
      expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_ID))
          .andReturn(Optional.of(pausedUpdateDetails));
  
      control.replay();
  
      // Kill command has not been executed yet
      assertEquals(2, killCommandHasExecuted.getCount());
  
      // Start an SLA-aware kill
      slaKillController.slaKill(
          storageUtil.mutableStoreProvider,
          INSTANCE_KEY,
          TASK,
          UPDATE_ID,
          ISlaPolicy.build(TEST_SLA_POLICY),
          JobUpdateStatus.ROLLING_FORWARD,
          fakeKillCommand);
  
      // Nothing should happen since status has changed
      assertEquals(2, killCommandHasExecuted.getCount());
      assertEquals(0, clock.countDeferredWork());
      assertFalse(statsProvider.getAllValues().keySet().contains(KILL_ATTEMPTS_STAT_NAME));
    }
  
    @Test
    public void testSlaKillNoDuplicateEvents() throws Exception {
      IJobUpdateDetails updateDetails = IJobUpdateDetails.build(
          new JobUpdateDetails(
              UPDATE.newBuilder(),
              ImmutableList.of(new JobUpdateEvent(JobUpdateStatus.ROLLING_FORWARD, 123L)),
              ImmutableList.of(
                  new JobInstanceUpdateEvent()
                      .setInstanceId(INSTANCE_KEY.getInstanceId())
                      .setAction(JobUpdateAction.INSTANCE_UPDATING)
                      .setMessage(SlaKillController.SLA_CHECKING_MESSAGE))));
      IJobUpdateDetails pausedUpdateDetails = IJobUpdateDetails.build(
          new JobUpdateDetails(
              UPDATE.newBuilder(),
              ImmutableList.of(new JobUpdateEvent(JobUpdateStatus.ROLL_FORWARD_PAUSED, 123L)),
              ImmutableList.of()));
      expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_ID))
          .andReturn(Optional.of(updateDetails));
      expectBatchExecute(batchWorker, storageUtil.storage, control);
      expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_ID))
          .andReturn(Optional.of(pausedUpdateDetails));
  
      control.replay();
  
      // Start an SLA-aware kill, update already contains event so we don't expect a save
      slaKillController.slaKill(
          storageUtil.mutableStoreProvider,
          INSTANCE_KEY,
          TASK,
          UPDATE_ID,
          ISlaPolicy.build(TEST_SLA_POLICY),
          JobUpdateStatus.ROLLING_FORWARD,
          fakeKillCommand);
    }
  
    @Test
    public void testSlaKillInvalidStatus() {
      control.replay();
  
      // Start an SLA-aware kill, throws an exception since the kill was called while the update
      // was not active
      try {
        slaKillController.slaKill(
            storageUtil.mutableStoreProvider,
            INSTANCE_KEY,
            TASK,
            UPDATE_ID,
            ISlaPolicy.build(TEST_SLA_POLICY),
            JobUpdateStatus.ROLL_FORWARD_PAUSED,
            fakeKillCommand);
      } catch (RuntimeException e) {
        return;
      }
  
      fail();
    }
  
    private boolean checkInstanceEventMatches(IJobInstanceUpdateEvent event,
                                              IInstanceKey instance,
                                              JobUpdateAction action,
                                              String message) {
  
      return event.getInstanceId() == instance.getInstanceId()
          && event.getAction() == action
          && event.getMessage().equals(message);
    }
  }
  