src/test/java/org/apache/aurora/scheduler/preemptor/PendingTaskProcessorTest.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.preemptor;

import java.util.Arrays;
import java.util.Optional;

import javax.annotation.Nullable;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMultiset;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;

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.AssignedTask;
import org.apache.aurora.gen.HostAttributes;
import org.apache.aurora.gen.MaintenanceMode;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.gen.TaskConfig;
import org.apache.aurora.gen.TaskEvent;
import org.apache.aurora.scheduler.base.JobKeys;
import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.base.TaskGroupKey;
import org.apache.aurora.scheduler.filter.AttributeAggregate;
import org.apache.aurora.scheduler.offers.HostOffer;
import org.apache.aurora.scheduler.offers.OfferManager;
import org.apache.aurora.scheduler.state.ClusterState;
import org.apache.aurora.scheduler.stats.CachedCounters;
import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
import org.apache.aurora.scheduler.testing.FakeStatsProvider;
import org.apache.mesos.v1.Protos;
import org.junit.Before;
import org.junit.Test;

import static org.apache.aurora.gen.ScheduleStatus.PENDING;
import static org.apache.aurora.scheduler.preemptor.PreemptorMetrics.TASK_PROCESSOR_RUN_NAME;
import static org.apache.aurora.scheduler.preemptor.PreemptorMetrics.UNMATCHED_TASKS;
import static org.apache.aurora.scheduler.preemptor.PreemptorMetrics.attemptsByJobStatName;
import static org.apache.aurora.scheduler.preemptor.PreemptorMetrics.attemptsStatName;
import static org.apache.aurora.scheduler.preemptor.PreemptorMetrics.slotSearchByJobStatName;
import static org.apache.aurora.scheduler.preemptor.PreemptorMetrics.slotSearchStatName;
import static org.easymock.EasyMock.anyObject;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.expectLastCall;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

public class PendingTaskProcessorTest extends EasyMockTest {
  private static final String CACHE_NAME = "TEST";
  private static final String CACHE_SIZE_STAT_NAME = "TEST_cache_size";
  private static final String SLAVE_ID_1 = "slave_id_1";
  private static final String SLAVE_ID_2 = "slave_id_2";
  private static final IJobKey JOB_A = JobKeys.from("role_a", "env", "job_a");
  private static final IJobKey JOB_B = JobKeys.from("role_b", "env", "job_b");
  private static final IJobKey JOB_C = JobKeys.from("role_c", "env", "job_c");
  private static final IScheduledTask TASK_A = makeTask(JOB_A, SLAVE_ID_1, "id1");
  private static final IScheduledTask TASK_B = makeTask(JOB_B, SLAVE_ID_2, "id2");
  private static final IScheduledTask TASK_C = makeTask(JOB_C, SLAVE_ID_2, "id3");
  private static final PreemptionProposal SLOT_A = createPreemptionProposal(TASK_A, SLAVE_ID_1);
  private static final Amount<Long, Time> PREEMPTION_DELAY = Amount.of(30L, Time.SECONDS);
  private static final Amount<Long, Time> EXPIRATION = Amount.of(10L, Time.MINUTES);
  private static final Integer RESERVATION_BATCH_SIZE = 5;

  private StorageTestUtil storageUtil;
  private OfferManager offerManager;
  private FakeStatsProvider statsProvider;
  private PreemptionVictimFilter preemptionVictimFilter;
  private PendingTaskProcessor slotFinder;
  private BiCache<PreemptionProposal, TaskGroupKey> slotCache;
  private ClusterState clusterState;
  private FakeClock clock;

  @Before
  public void setUp() {
    storageUtil = new StorageTestUtil(this);
    storageUtil.expectOperations();
    offerManager = createMock(OfferManager.class);
    preemptionVictimFilter = createMock(PreemptionVictimFilter.class);
    statsProvider = new FakeStatsProvider();
    clusterState = createMock(ClusterState.class);
    clock = new FakeClock();
    slotCache = new BiCache<>(
        statsProvider,
        new BiCache.BiCacheSettings(EXPIRATION, CACHE_NAME),
        clock);

    slotFinder = new PendingTaskProcessor(
        storageUtil.storage,
        offerManager,
        preemptionVictimFilter,
        new PreemptorMetrics(new CachedCounters(statsProvider)),
        PREEMPTION_DELAY,
        slotCache,
        clusterState,
        clock,
        RESERVATION_BATCH_SIZE);
  }

  @Test
  public void testSearchSlotSuccessful() throws Exception {
    expectGetPendingTasks(TASK_A, TASK_B);
    expectGetClusterState(TASK_A, TASK_B);
    HostOffer offer1 = makeOffer(SLAVE_ID_1);
    HostOffer offer2 = makeOffer(SLAVE_ID_2);
    expectOffers(offer1, offer2);
    expectSlotSearch(TASK_A.getAssignedTask().getTask(), TASK_A);
    expectSlotSearch(TASK_B.getAssignedTask().getTask(), TASK_B);

    control.replay();

    clock.advance(PREEMPTION_DELAY);

    slotFinder.run();
    assertEquals(1L, statsProvider.getLongValue(TASK_PROCESSOR_RUN_NAME));
    assertEquals(2L, statsProvider.getLongValue(attemptsStatName(true)));
    assertEquals(1L, statsProvider.getLongValue(attemptsByJobStatName(JOB_A)));
    assertEquals(1L, statsProvider.getLongValue(attemptsByJobStatName(JOB_B)));
    assertEquals(2L, statsProvider.getLongValue(slotSearchStatName(true, true)));
    assertEquals(1L, statsProvider.getLongValue(slotSearchByJobStatName(true, JOB_A)));
    assertEquals(1L, statsProvider.getLongValue(slotSearchByJobStatName(true, JOB_B)));
    assertEquals(0L, statsProvider.getLongValue(slotSearchStatName(false, true)));
    assertEquals(0L, statsProvider.getLongValue(UNMATCHED_TASKS));
    assertEquals(2L, statsProvider.getLongValue(CACHE_SIZE_STAT_NAME));
  }

  @Test
  public void testSearchSlotFailed() throws Exception {
    expectGetPendingTasks(TASK_A);
    expectGetClusterState(TASK_A);
    HostOffer offer1 = makeOffer(SLAVE_ID_1);
    expectOffers(offer1);
    expectSlotSearch(TASK_A.getAssignedTask().getTask());

    control.replay();

    clock.advance(PREEMPTION_DELAY);

    slotFinder.run();
    assertEquals(1L, statsProvider.getLongValue(TASK_PROCESSOR_RUN_NAME));
    assertEquals(1L, statsProvider.getLongValue(attemptsStatName(true)));
    assertEquals(1L, statsProvider.getLongValue(attemptsByJobStatName(JOB_A)));
    assertEquals(0L, statsProvider.getLongValue(slotSearchStatName(true, true)));
    assertEquals(1L, statsProvider.getLongValue(slotSearchStatName(false, true)));
    assertEquals(1L, statsProvider.getLongValue(slotSearchByJobStatName(false, JOB_A)));
    assertEquals(1L, statsProvider.getLongValue(UNMATCHED_TASKS));
  }

  @Test
  public void testHasCachedSlots() throws Exception {
    slotCache.put(SLOT_A, group(TASK_A));
    expectGetPendingTasks(TASK_A);
    expectGetClusterState(TASK_A);
    HostOffer offer1 = makeOffer(SLAVE_ID_1);
    expectOffers(offer1);

    control.replay();

    clock.advance(PREEMPTION_DELAY);

    slotFinder.run();
    assertEquals(1L, statsProvider.getLongValue(TASK_PROCESSOR_RUN_NAME));
    assertEquals(0L, statsProvider.getLongValue(attemptsStatName(true)));
    assertEquals(0L, statsProvider.getLongValue(slotSearchStatName(true, true)));
    assertEquals(0L, statsProvider.getLongValue(slotSearchStatName(false, true)));
    assertEquals(0L, statsProvider.getLongValue(UNMATCHED_TASKS));
  }

  @Test
  public void testMultipleTaskGroups() throws Exception {
    IScheduledTask task1 = makeTask(JOB_A, "1");
    IScheduledTask task2 = makeTask(JOB_A, "2");
    IScheduledTask task3 = makeTask(JOB_A, "3");
    IScheduledTask task4 = makeTask(JOB_B, "4");
    IScheduledTask task5 = makeTask(JOB_B, "5");

    expectGetPendingTasks(task1, task4, task2, task5, task3);
    expectGetClusterState(TASK_A, TASK_B);

    HostOffer offer1 = makeOffer(SLAVE_ID_1);
    HostOffer offer2 = makeOffer(SLAVE_ID_2);
    expectOffers(offer1, offer2);
    expectSlotSearch(task1.getAssignedTask().getTask());
    expectSlotSearch(task4.getAssignedTask().getTask(), TASK_B);
    PreemptionProposal proposal1 = createPreemptionProposal(TASK_B, SLAVE_ID_1);
    PreemptionProposal proposal2 = createPreemptionProposal(TASK_B, SLAVE_ID_2);

    control.replay();

    clock.advance(PREEMPTION_DELAY);

    slotFinder.run();
    assertEquals(slotCache.get(proposal1), Optional.of(group(task4)));
    assertEquals(slotCache.get(proposal2), Optional.of(group(task5)));
    assertEquals(1L, statsProvider.getLongValue(TASK_PROCESSOR_RUN_NAME));
    assertEquals(3L, statsProvider.getLongValue(attemptsStatName(true)));
    assertEquals(1L, statsProvider.getLongValue(attemptsByJobStatName(JOB_A)));
    assertEquals(2L, statsProvider.getLongValue(attemptsByJobStatName(JOB_B)));
    assertEquals(2L, statsProvider.getLongValue(slotSearchStatName(true, true)));
    assertEquals(2L, statsProvider.getLongValue(slotSearchByJobStatName(true, JOB_B)));

    // TODO(wfarner): This test depends on the iteration order of a hash set (the set containing
    // task groups), and as a result this stat could be 0 or 2 depending on which group is
    // evaluated first.
    assertTrue(ImmutableSet.of(0L, 2L).contains(
        statsProvider.getLongValue(slotSearchStatName(false, true))));
    assertEquals(1L, statsProvider.getLongValue(UNMATCHED_TASKS));
    assertEquals(2L, statsProvider.getLongValue(CACHE_SIZE_STAT_NAME));
  }

  @Test
  public void testNoVictims() throws Exception {
    expectGetClusterState();
    control.replay();

    clock.advance(PREEMPTION_DELAY);

    slotFinder.run();
    assertEquals(1L, statsProvider.getLongValue(TASK_PROCESSOR_RUN_NAME));
    assertEquals(0L, statsProvider.getLongValue(attemptsStatName(true)));
    assertEquals(0L, statsProvider.getLongValue(attemptsStatName(false)));
    assertEquals(0L, statsProvider.getLongValue(slotSearchStatName(true, true)));
    assertEquals(0L, statsProvider.getLongValue(slotSearchStatName(false, true)));
    assertEquals(0L, statsProvider.getLongValue(UNMATCHED_TASKS));
  }

  @Test
  public void testGetPreemptionSequence() {
    TaskGroupKey a = group(TASK_A);
    TaskGroupKey b = group(TASK_B);
    TaskGroupKey c = group(TASK_C);

    control.replay();

    assertEquals(
        ImmutableList.of(),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(), 2));
    assertEquals(
        ImmutableList.of(a),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a), 3));
    assertEquals(
        ImmutableList.of(a, a, a),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, a, a), 5));
    assertEquals(
        ImmutableList.of(a, a, a),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, a, a), 2));
    assertEquals(
        ImmutableList.of(a, b, a, b, a),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, a, a, b, b), 1));
    assertEquals(
        ImmutableList.of(a, a, b, b, a),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, a, a, b, b), 2));
    assertEquals(
        ImmutableList.of(a, a, a, b, b),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, a, a, b, b), 3));
    assertEquals(
        ImmutableList.of(a, a, a, b, b),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, a, a, b, b), 5));
    assertEquals(
        ImmutableList.of(a, a, b, b, a, a, a),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, a, a, a, a, b, b), 2));
    assertEquals(
        ImmutableList.of(a, b, c, b, c),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, b, b, c, c), 1));
    assertEquals(
        ImmutableList.of(a, b, b, c, c),
        PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(a, b, b, c, c), 2));
  }

  @Test(expected = IllegalArgumentException.class)
  public void testGetPreemptionSequenceInvalidArgument() {
    control.replay();

    PendingTaskProcessor.getPreemptionSequence(ImmutableMultiset.of(), 0);
  }

  private Multimap<String, PreemptionVictim> getVictims(IScheduledTask... tasks) {
    return Multimaps.transformValues(
        Multimaps.index(Arrays.asList(tasks), task -> task.getAssignedTask().getSlaveId()),
        task -> PreemptionVictim.fromTask(task.getAssignedTask())
    );
  }

  private HostOffer makeOffer(String slaveId) {
    Protos.Offer.Builder builder = Protos.Offer.newBuilder();
    builder.getIdBuilder().setValue("id");
    builder.getFrameworkIdBuilder().setValue("framework-id");
    builder.getAgentIdBuilder().setValue(slaveId);
    builder.setHostname(slaveId);
    return new HostOffer(
        builder.build(),
        IHostAttributes.build(new HostAttributes().setMode(MaintenanceMode.NONE)));
  }

  private void expectOffers(HostOffer... offers) {
    expect(offerManager.getAll()).andReturn(ImmutableSet.copyOf(offers));
  }

  private void expectGetClusterState(IScheduledTask... returnedTasks) {
    expect(clusterState.getSlavesToActiveTasks()).andReturn(getVictims(returnedTasks));
  }

  private void expectSlotSearch(ITaskConfig config, IScheduledTask... victims) {
    expect(preemptionVictimFilter.filterPreemptionVictims(
        eq(config),
        anyObject(),
        anyObject(AttributeAggregate.class),
        anyObject(),
        eq(storageUtil.storeProvider)));
    expectLastCall().andReturn(
        victims.length == 0
            ? Optional.empty()
            : Optional.of(ImmutableSet.copyOf(getVictims(victims).values())))
        .anyTimes();
  }

  private static PreemptionProposal createPreemptionProposal(IScheduledTask task, String slaveId) {
    return new PreemptionProposal(
        ImmutableSet.of(PreemptionVictim.fromTask(task.getAssignedTask())),
        slaveId);
  }

  private static IScheduledTask makeTask(IJobKey key, String taskId) {
    return makeTask(key, null, taskId);
  }

  private static TaskGroupKey group(IScheduledTask task) {
    return TaskGroupKey.from(task.getAssignedTask().getTask());
  }

  private static IScheduledTask makeTask(IJobKey key, @Nullable String slaveId, String taskId) {
    ScheduledTask task = new ScheduledTask()
        .setAssignedTask(new AssignedTask()
            .setSlaveId(slaveId)
            .setTaskId(taskId)
            .setTask(new TaskConfig()
                .setPriority(1)
                .setProduction(true)
                .setJob(key.newBuilder())));
    task.addToTaskEvents(new TaskEvent(0, PENDING));
    return IScheduledTask.build(task);
  }

  private void expectGetPendingTasks(IScheduledTask... returnedTasks) {
    storageUtil.expectTaskFetch(Query.statusScoped(PENDING), ImmutableSet.copyOf(returnedTasks));
  }
}