TaskHistoryPrunerTest.java

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

import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.io.Closer;

import org.apache.aurora.common.application.Lifecycle;
import org.apache.aurora.common.base.Command;
import org.apache.aurora.common.quantity.Amount;
import org.apache.aurora.common.quantity.Time;
import org.apache.aurora.common.testing.easymock.EasyMockTest;
import org.apache.aurora.common.util.testing.FakeClock;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
import org.apache.aurora.scheduler.async.DelayExecutor;
import org.apache.aurora.scheduler.base.JobKeys;
import org.apache.aurora.scheduler.base.TaskTestUtil;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
import org.apache.aurora.scheduler.pruning.TaskHistoryPruner.HistoryPrunnerSettings;
import org.apache.aurora.scheduler.state.StateManager;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
import org.apache.aurora.scheduler.testing.FakeStatsProvider;
import org.easymock.Capture;
import org.easymock.EasyMock;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import static org.apache.aurora.gen.ScheduleStatus.FINISHED;
import static org.apache.aurora.gen.ScheduleStatus.KILLED;
import static org.apache.aurora.gen.ScheduleStatus.LOST;
import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
import static org.apache.aurora.gen.ScheduleStatus.STARTING;
import static org.apache.aurora.scheduler.pruning.TaskHistoryPruner.TASKS_PRUNED;
import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expectLastCall;
import static org.junit.Assert.assertEquals;

public class TaskHistoryPrunerTest extends EasyMockTest {
  private static final String SLAVE_HOST = "HOST_A";
  private static final Amount<Long, Time> ONE_MS = Amount.of(1L, Time.MILLISECONDS);
  private static final Amount<Long, Time> ONE_MINUTE = Amount.of(1L, Time.MINUTES);
  private static final Amount<Long, Time> ONE_DAY = Amount.of(1L, Time.DAYS);
  private static final Amount<Long, Time> ONE_HOUR = Amount.of(1L, Time.HOURS);
  private static final int PER_JOB_HISTORY = 2;

  private DelayExecutor executor;
  private FakeClock clock;
  private StateManager stateManager;
  private StorageTestUtil storageUtil;
  private TaskHistoryPruner pruner;
  private Closer closer;
  private Command shutdownCommand;
  private FakeStatsProvider statsProvider;

  @Before
  public void setUp() throws Exception {
    executor = createMock(DelayExecutor.class);
    clock = new FakeClock();
    stateManager = createMock(StateManager.class);
    storageUtil = new StorageTestUtil(this);
    storageUtil.expectOperations();
    shutdownCommand = createMock(Command.class);
    TaskEventBatchWorker batchWorker = createMock(TaskEventBatchWorker.class);
    statsProvider = new FakeStatsProvider();
    expectBatchExecute(batchWorker, storageUtil.storage, control).anyTimes();

    pruner = new TaskHistoryPruner(
        executor,
        stateManager,
        clock,
        new HistoryPrunnerSettings(ONE_DAY, ONE_MINUTE, PER_JOB_HISTORY),
        storageUtil.storage,
        new Lifecycle(shutdownCommand),
        batchWorker,
        statsProvider);
    closer = Closer.create();
  }

  @After
  public void tearDownCloser() throws Exception {
    closer.close();
  }

  @Test
  public void testNoPruning() {
    long taskATimestamp = clock.nowMillis();
    IScheduledTask a = makeTask("a", FINISHED);

    clock.advance(ONE_MS);
    long taskBTimestamp = clock.nowMillis();
    IScheduledTask b = makeTask("b", LOST);

    expectNoImmediatePrune(ImmutableSet.of(a));
    expectOneDelayedPrune(taskATimestamp);
    expectNoImmediatePrune(ImmutableSet.of(a, b));
    expectOneDelayedPrune(taskBTimestamp);

    control.replay();

    pruner.recordStateChange(TaskStateChange.initialized(a));
    pruner.recordStateChange(TaskStateChange.initialized(b));
  }

  @Test
  public void testStorageStartedWithPruning() {
    long taskATimestamp = clock.nowMillis();
    IScheduledTask a = makeTask("a", FINISHED);

    clock.advance(ONE_MINUTE);
    long taskBTimestamp = clock.nowMillis();
    IScheduledTask b = makeTask("b", LOST);

    clock.advance(ONE_MINUTE);
    long taskCTimestamp = clock.nowMillis();
    IScheduledTask c = makeTask("c", FINISHED);

    clock.advance(ONE_MINUTE);
    IScheduledTask d = makeTask("d", FINISHED);
    IScheduledTask e = makeTask(JobKeys.from("role", "env", "job-x"), "e", FINISHED);

    expectNoImmediatePrune(ImmutableSet.of(a));
    expectOneDelayedPrune(taskATimestamp);
    expectNoImmediatePrune(ImmutableSet.of(a, b));
    expectOneDelayedPrune(taskBTimestamp);
    expectImmediatePrune(ImmutableSet.of(a, b, c), a);
    expectOneDelayedPrune(taskCTimestamp);
    expectImmediatePrune(ImmutableSet.of(b, c, d), b);
    expectDefaultDelayedPrune();
    expectNoImmediatePrune(ImmutableSet.of(e));
    expectDefaultDelayedPrune();

    control.replay();

    assertEquals(0L, statsProvider.getValue(TASKS_PRUNED));
    for (IScheduledTask task : ImmutableList.of(a, b, c, d, e)) {
      pruner.recordStateChange(TaskStateChange.initialized(task));
    }
    assertEquals(2L, statsProvider.getValue(TASKS_PRUNED));
  }

  @Test
  public void testStateChange() {
    IScheduledTask starting = makeTask("a", STARTING);
    IScheduledTask running = copy(starting, RUNNING);
    IScheduledTask killed = copy(starting, KILLED);

    expectNoImmediatePrune(ImmutableSet.of(killed));
    expectDefaultDelayedPrune();

    control.replay();

    // No future set for non-terminal state transition.
    changeState(starting, running);

    // Future set for terminal state transition.
    changeState(running, killed);
  }

  @Test
  public void testActivateFutureAndExceedHistoryGoal() {
    IScheduledTask running = makeTask("a", RUNNING);
    IScheduledTask killed = copy(running, KILLED);
    expectNoImmediatePrune(ImmutableSet.of(running));
    Capture<Runnable> delayedDelete = expectDefaultDelayedPrune();

    // Expect task "a" to be pruned when future is activated.
    expectDeleteTasks("a");

    control.replay();

    // Capture future for inactive task "a"
    changeState(running, killed);
    clock.advance(ONE_HOUR);
    assertEquals(0L, statsProvider.getValue(TASKS_PRUNED));
    // Execute future to prune task "a" from the system.
    delayedDelete.getValue().run();
    assertEquals(1L, statsProvider.getValue(TASKS_PRUNED));
  }

  @Test
  public void testSuppressEmptyDelete() {
    IScheduledTask running = makeTask("a", RUNNING);
    IScheduledTask killed = copy(running, KILLED);
    expectImmediatePrune(
        ImmutableSet.of(makeTask("b", KILLED), makeTask("c", KILLED), makeTask("d", KILLED)));
    expectDefaultDelayedPrune();

    control.replay();

    changeState(running, killed);
  }

  @Test
  public void testJobHistoryExceeded() {
    IScheduledTask a = makeTask("a", RUNNING);
    clock.advance(ONE_MS);
    IScheduledTask aKilled = copy(a, KILLED);

    IScheduledTask b = makeTask("b", RUNNING);
    clock.advance(ONE_MS);
    IScheduledTask bKilled = copy(b, KILLED);

    IScheduledTask c = makeTask("c", RUNNING);
    clock.advance(ONE_MS);
    IScheduledTask cLost = copy(c, LOST);

    IScheduledTask d = makeTask("d", RUNNING);
    clock.advance(ONE_MS);
    IScheduledTask dLost = copy(d, LOST);

    expectNoImmediatePrune(ImmutableSet.of(a));
    expectDefaultDelayedPrune();
    expectNoImmediatePrune(ImmutableSet.of(a, b));
    expectDefaultDelayedPrune();
    expectNoImmediatePrune(ImmutableSet.of(a, b)); // no pruning yet due to min threshold
    expectDefaultDelayedPrune();
    clock.advance(ONE_HOUR);
    expectImmediatePrune(ImmutableSet.of(a, b, c, d), a, b); // now prune 2 tasks
    expectDefaultDelayedPrune();

    control.replay();

    changeState(a, aKilled);
    changeState(b, bKilled);
    changeState(c, cLost);
    assertEquals(0L, statsProvider.getValue(TASKS_PRUNED));
    changeState(d, dLost);
    assertEquals(2L, statsProvider.getValue(TASKS_PRUNED));
  }

  @Test
  public void serviceShutdownOnFailure() {
    IScheduledTask running = makeTask("a", RUNNING);
    IScheduledTask killed = copy(running, KILLED);
    expectNoImmediatePrune(ImmutableSet.of(running));
    Capture<Runnable> delayedDelete = expectDefaultDelayedPrune();

    expectDeleteTasks("a");
    expectLastCall().andThrow(new RuntimeException("oops"));

    shutdownCommand.execute();

    control.replay();

    changeState(running, killed);
    clock.advance(ONE_HOUR);
    assertEquals(0L, statsProvider.getValue(TASKS_PRUNED));
    delayedDelete.getValue().run();
    assertEquals(0L, statsProvider.getValue(TASKS_PRUNED));
  }

  private void expectDeleteTasks(String... tasks) {
    stateManager.deleteTasks(storageUtil.mutableStoreProvider, ImmutableSet.copyOf(tasks));
  }

  private Capture<Runnable> expectDefaultDelayedPrune() {
    return expectDelayedPrune(ONE_DAY.as(Time.MILLISECONDS));
  }

  private Capture<Runnable> expectOneDelayedPrune(long timestampMillis) {
    return expectDelayedPrune(timestampMillis);
  }

  private void expectNoImmediatePrune(ImmutableSet<IScheduledTask> tasksInJob) {
    expectImmediatePrune(tasksInJob);
  }

  private void expectImmediatePrune(
      ImmutableSet<IScheduledTask> tasksInJob,
      IScheduledTask... pruned) {

    // Expect a deferred prune operation when a new task is being watched.
    executor.execute(EasyMock.<Runnable>anyObject());
    expectLastCall().andAnswer(
        () -> {
          Runnable work = (Runnable) EasyMock.getCurrentArguments()[0];
          work.run();
          return null;
        }
    );

    IJobKey jobKey = Iterables.getOnlyElement(
        FluentIterable.from(tasksInJob).transform(Tasks::getJob).toSet());
    storageUtil.expectTaskFetch(TaskHistoryPruner.jobHistoryQuery(jobKey), tasksInJob);
    if (pruned.length > 0) {
      stateManager.deleteTasks(storageUtil.mutableStoreProvider, Tasks.ids(pruned));
    }
  }

  private Capture<Runnable> expectDelayedPrune(long timestampMillis) {
    Capture<Runnable> capture = createCapture();
    executor.execute(
        EasyMock.capture(capture),
        eq(Amount.of(pruner.calculateTimeout(timestampMillis), Time.MILLISECONDS)));
    return capture;
  }

  private void changeState(IScheduledTask oldStateTask, IScheduledTask newStateTask) {
    pruner.recordStateChange(TaskStateChange.transition(newStateTask, oldStateTask.getStatus()));
  }

  private IScheduledTask copy(IScheduledTask task, ScheduleStatus status) {
    return IScheduledTask.build(task.newBuilder().setStatus(status));
  }

  private IScheduledTask makeTask(
      IJobKey job,
      String taskId,
      ScheduleStatus status) {

    ScheduledTask builder = TaskTestUtil.addStateTransition(
        TaskTestUtil.makeTask(taskId, job), status, clock.nowMillis())
        .newBuilder();
    builder.getAssignedTask().setSlaveHost(SLAVE_HOST);
    return IScheduledTask.build(builder);
  }

  private IScheduledTask makeTask(String taskId, ScheduleStatus status) {
    return makeTask(TaskTestUtil.JOB, taskId, status);
  }

}