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

import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Executor;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.TypeLiteral;

import org.apache.aurora.common.quantity.Amount;
import org.apache.aurora.common.quantity.Time;
import org.apache.aurora.common.stats.StatsProvider;
import org.apache.aurora.common.testing.easymock.EasyMockTest;
import org.apache.aurora.gen.CountSlaPolicy;
import org.apache.aurora.gen.HostAttributes;
import org.apache.aurora.gen.HostMaintenanceRequest;
import org.apache.aurora.gen.HostStatus;
import org.apache.aurora.gen.MaintenanceMode;
import org.apache.aurora.gen.PercentageSlaPolicy;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.gen.SlaPolicy;
import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
import org.apache.aurora.scheduler.async.AsyncModule.AsyncExecutor;
import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.base.TaskTestUtil;
import org.apache.aurora.scheduler.config.types.TimeAmount;
import org.apache.aurora.scheduler.events.EventSink;
import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
import org.apache.aurora.scheduler.events.PubsubEventModule;
import org.apache.aurora.scheduler.sla.SlaManager;
import org.apache.aurora.scheduler.state.PubsubTestUtil;
import org.apache.aurora.scheduler.state.StateManager;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
import org.apache.aurora.scheduler.storage.entities.IHostMaintenanceRequest;
import org.apache.aurora.scheduler.storage.entities.IHostStatus;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.entities.ISlaPolicy;
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.MaintenanceMode.DRAINED;
import static org.apache.aurora.gen.MaintenanceMode.DRAINING;
import static org.apache.aurora.gen.MaintenanceMode.NONE;
import static org.apache.aurora.gen.MaintenanceMode.SCHEDULED;
import static org.apache.aurora.gen.ScheduleStatus.KILLED;
import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
import static org.easymock.EasyMock.anyObject;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.junit.Assert.assertEquals;

public class MaintenanceControllerImplTest extends EasyMockTest {

  private static final String HOST_A = "a";
  private static final Set<String> A = ImmutableSet.of(HOST_A);
  private static final Protos.OfferID OFFER_ID = Protos.OfferID.newBuilder()
      .setValue("offer-id")
      .build();
  private static final Protos.AgentID AGENT_ID = Protos.AgentID.newBuilder()
      .setValue("agent-id")
      .build();
  private static final Protos.FrameworkID FRAMEWORK_ID = Protos.FrameworkID.newBuilder()
      .setValue("framework-id")
      .build();
  private static final Protos.URL AGENT_URL = Protos.URL.newBuilder()
      .setAddress(Protos.Address.newBuilder()
          .setHostname(HOST_A)
          .setPort(5051))
      .setScheme("http")
      .build();
  private static final Protos.Unavailability UNAVAILABILITY = Protos.Unavailability.newBuilder()
      .setStart(Protos.TimeInfo.newBuilder()
          .setNanoseconds(Amount.of(1L, Time.MINUTES).as(Time.NANOSECONDS)))
      .build();
  private static final SlaPolicy SLA_POLICY = SlaPolicy.percentageSlaPolicy(
      new PercentageSlaPolicy(95, 1800));

  private static final Protos.InverseOffer INVERSE_OFFER = Protos.InverseOffer.newBuilder()
      .setId(OFFER_ID)
      .setAgentId(AGENT_ID)
      .setUrl(AGENT_URL)
      .setFrameworkId(FRAMEWORK_ID)
      .setUnavailability(UNAVAILABILITY)
      .build();

  private static final SlaPolicy COUNT_SLA_POLICY = SlaPolicy.countSlaPolicy(
      new CountSlaPolicy()
          .setCount(2)
          .setDurationSecs(1800)
  );

  private StorageTestUtil storageUtil;
  private StateManager stateManager;
  private SlaManager slaManager;
  private MaintenanceController.MaintenanceControllerImpl maintenance;
  private EventSink eventSink;

  @Before
  public void setUp() throws Exception {
    storageUtil = new StorageTestUtil(this);
    storageUtil.expectOperations();
    stateManager = createMock(StateManager.class);
    slaManager = createMock(SlaManager.class);
    TaskEventBatchWorker batchWorker = createMock(TaskEventBatchWorker.class);
    expectBatchExecute(batchWorker, storageUtil.storage, control).anyTimes();

    Injector injector = Guice.createInjector(
        new PubsubEventModule(),
        new MaintenanceModule(new MaintenanceModule.Options()),
        new AbstractModule() {
          @Override
          protected void configure() {
            bind(Storage.class).toInstance(storageUtil.storage);
            bind(StateManager.class).toInstance(stateManager);
            bind(SlaManager.class).toInstance(slaManager);
            bind(StatsProvider.class).toInstance(new FakeStatsProvider());
            bind(Executor.class).annotatedWith(AsyncExecutor.class)
                .toInstance(MoreExecutors.directExecutor());
            bind(TaskEventBatchWorker.class).toInstance(batchWorker);
            bind(new TypeLiteral<Amount<Long, Time>>() { })
                .annotatedWith(
                    MaintenanceController.MaintenanceControllerImpl.PollingInterval.class)
                .toInstance(new TimeAmount(1, Time.MINUTES));
          }
        });
    maintenance = injector.getInstance(MaintenanceController.MaintenanceControllerImpl.class);
    eventSink = PubsubTestUtil.startPubsub(injector);
  }

  private static IScheduledTask makeTask(String host, String taskId) {
    ScheduledTask builder = TaskTestUtil.addStateTransition(
        TaskTestUtil.makeTask(taskId, TaskTestUtil.JOB),
        RUNNING,
        1000).newBuilder();
    builder.getAssignedTask().setSlaveHost(host);
    return IScheduledTask.build(builder);
  }

  @Test
  public void testMaintenanceCycle() {
    IScheduledTask task1 = makeTask(HOST_A, "taskA");
    IScheduledTask task2 = makeTask(HOST_A, "taskB");

    expectMaintenanceModeChange(HOST_A, SCHEDULED);
    expectFetchTasksByHost(HOST_A, ImmutableSet.of(task1, task2));
    expectTaskDraining(task1);
    expectTaskDraining(task2);
    expectMaintenanceModeChange(HOST_A, DRAINING);
    IHostAttributes attributes =
        IHostAttributes.build(new HostAttributes().setHost(HOST_A).setMode(DRAINING));
    IHostMaintenanceRequest maintenanceRequest =
        IHostMaintenanceRequest.build(new HostMaintenanceRequest()
            .setHost(HOST_A)
            .setCreatedTimestampMs(System.currentTimeMillis())
            .setTimeoutSecs(7200)
            .setDefaultSlaPolicy(SLA_POLICY));

    storageUtil.hostMaintenanceStore.saveHostMaintenanceRequest(
        anyObject(IHostMaintenanceRequest.class));
    expect(storageUtil.attributeStore.getHostAttributes(HOST_A))
        .andReturn(Optional.of(attributes)).times(2);
    expect(storageUtil.hostMaintenanceStore.getHostMaintenanceRequest(HOST_A))
        .andReturn(Optional.of(maintenanceRequest)).times(2);

    expect(storageUtil.attributeStore.getHostAttributes()).andReturn(ImmutableSet.of(attributes));
    expectFetchTasksByHost(HOST_A, ImmutableSet.of(task2));
    // TaskA is KILLED and therefore no longer active
    expectFetchTasksByHost(HOST_A, ImmutableSet.of());
    expectMaintenanceModeChange(HOST_A, DRAINED);
    expectMaintenanceModeChange(HOST_A, NONE);
    storageUtil.hostMaintenanceStore.removeHostMaintenanceRequest(HOST_A);
    // end maintenance
    storageUtil.hostMaintenanceStore.removeHostMaintenanceRequest(HOST_A);

    control.replay();

    assertStatus(HOST_A, SCHEDULED, maintenance.startMaintenance(A));
    assertStatus(HOST_A, DRAINING, maintenance.drain(A));
    assertStatus(HOST_A, DRAINING, maintenance.getStatus(A));
    eventSink.post(
        TaskStateChange.transition(
            IScheduledTask.build(task1.newBuilder().setStatus(KILLED)), RUNNING));
    eventSink.post(
        TaskStateChange.transition(
            IScheduledTask.build(task2.newBuilder().setStatus(KILLED)), RUNNING));
    assertStatus(HOST_A, NONE, maintenance.endMaintenance(A));
  }

  @Test
  public void testUnknownHost() {
    expect(storageUtil.attributeStore.getHostAttributes("b"))
        .andReturn(Optional.empty());

    control.replay();

    assertEquals(ImmutableSet.of(),
        maintenance.startMaintenance(ImmutableSet.of("b")));
  }

  @Test
  public void testDrainEmptyHost() {
    storageUtil.hostMaintenanceStore.saveHostMaintenanceRequest(
        anyObject(IHostMaintenanceRequest.class));
    expectMaintenanceModeChange(HOST_A, SCHEDULED);
    expectFetchTasksByHost(HOST_A, ImmutableSet.of());
    expectMaintenanceModeChange(HOST_A, DRAINED);

    control.replay();

    assertStatus(HOST_A, SCHEDULED, maintenance.startMaintenance(A));
    assertStatus(HOST_A, DRAINED, maintenance.drain(A));
  }

  @Test
  public void testEndEarly() {
    expectMaintenanceModeChange(HOST_A, SCHEDULED);
    expectMaintenanceModeChange(HOST_A, NONE);
    expect(storageUtil.attributeStore.getHostAttributes(HOST_A)).andReturn(Optional.of(
        IHostAttributes.build(new HostAttributes().setHost(HOST_A).setMode(NONE))));
    storageUtil.hostMaintenanceStore.removeHostMaintenanceRequest(HOST_A);

    control.replay();

    assertStatus(HOST_A, SCHEDULED, maintenance.startMaintenance(A));

    // End maintenance without DRAINING.
    assertStatus(HOST_A, NONE, maintenance.endMaintenance(A));

    // Make sure a later transition on the host does not cause any ill effects that could surface
    // from stale internal state.
    eventSink.post(TaskStateChange.transition(
        IScheduledTask.build(makeTask(HOST_A, "taskA").newBuilder().setStatus(KILLED)), RUNNING));
  }

  @Test
  public void testSlaDrain() {
    storageUtil.hostMaintenanceStore.saveHostMaintenanceRequest(
        anyObject(IHostMaintenanceRequest.class));
    expectMaintenanceModeChange(HOST_A, DRAINING);

    control.replay();

    assertStatus(
        HOST_A,
        DRAINING,
        maintenance.slaDrain(ImmutableSet.of(HOST_A), COUNT_SLA_POLICY, 1800));
  }

  @Test
  public void testSlaDrainUnknownHost() {
    storageUtil.hostMaintenanceStore.saveHostMaintenanceRequest(
        anyObject(IHostMaintenanceRequest.class));
    expect(storageUtil.attributeStore.getHostAttributes("unknown"))
        .andReturn(Optional.empty());

    control.replay();

    assertEquals(ImmutableSet.of(),
        maintenance.slaDrain(ImmutableSet.of("unknown"), COUNT_SLA_POLICY, 1800));
  }

  @Test
  public void testIteration() {
    IScheduledTask task1 = makeTask(HOST_A, "taskA");
    IScheduledTask task2 = makeTask(HOST_A, "taskB");

    IHostAttributes attributes =
        IHostAttributes.build(new HostAttributes().setHost(HOST_A).setMode(DRAINING));
    expect(storageUtil.attributeStore.getHostAttributes())
        .andReturn(ImmutableSet.of(attributes));
    expectFetchTasksByHost(HOST_A, ImmutableSet.of(task1, task2));
    IHostMaintenanceRequest maintenanceRequest =
        IHostMaintenanceRequest.build(new HostMaintenanceRequest()
            .setHost(HOST_A)
            .setCreatedTimestampMs(System.currentTimeMillis())
            .setTimeoutSecs(7200)
            .setDefaultSlaPolicy(SLA_POLICY));
    expect(storageUtil.hostMaintenanceStore.getHostMaintenanceRequest(HOST_A))
        .andReturn(Optional.of(maintenanceRequest)).times(2);
    expectTaskDraining(task1);
    expectTaskDraining(task2);

    control.replay();

    maintenance.runForTest();
  }

  @Test
  public void testIterationEmptyHost() {
    IHostAttributes attributes =
        IHostAttributes.build(new HostAttributes().setHost(HOST_A).setMode(DRAINING));
    expect(storageUtil.attributeStore.getHostAttributes())
        .andReturn(ImmutableSet.of(attributes));
    expectFetchTasksByHost(HOST_A, ImmutableSet.of());
    expectMaintenanceModeChange(HOST_A, DRAINED);

    control.replay();

    maintenance.runForTest();
  }

  @Test
  public void testIterationMaintenanceTimeout() {
    IScheduledTask task1 = makeTask(HOST_A, "taskA");
    IScheduledTask task2 = makeTask(HOST_A, "taskB");

    IHostAttributes attributes =
        IHostAttributes.build(new HostAttributes().setHost(HOST_A).setMode(DRAINING));
    expect(storageUtil.attributeStore.getHostAttributes())
        .andReturn(ImmutableSet.of(attributes));
    expectFetchTasksByHost(HOST_A, ImmutableSet.of(task1, task2));
    IHostMaintenanceRequest maintenanceRequest =
        IHostMaintenanceRequest.build(new HostMaintenanceRequest()
            .setHost(HOST_A)
            .setCreatedTimestampMs(0)
            .setDefaultSlaPolicy(SLA_POLICY));
    expect(storageUtil.hostMaintenanceStore.getHostMaintenanceRequest(HOST_A))
        .andReturn(Optional.of(maintenanceRequest)).times(2);
    expectTaskDraining(task1, true);
    expectTaskDraining(task2, true);

    control.replay();

    maintenance.runForTest();
  }

  @Test
  public void testIterationNoMaintenanceRequest() {
    IScheduledTask task1 = makeTask(HOST_A, "taskA");
    IScheduledTask task2 = makeTask(HOST_A, "taskB");

    IHostAttributes attributes =
        IHostAttributes.build(new HostAttributes().setHost(HOST_A).setMode(DRAINING));
    expect(storageUtil.attributeStore.getHostAttributes())
        .andReturn(ImmutableSet.of(attributes));
    expectFetchTasksByHost(HOST_A, ImmutableSet.of(task1, task2));
    expect(storageUtil.hostMaintenanceStore.getHostMaintenanceRequest(HOST_A))
        .andReturn(Optional.empty()).times(2);

    control.replay();

    maintenance.runForTest();
  }

  @Test
  public void testGetMode() {
    expect(storageUtil.attributeStore.getHostAttributes(HOST_A)).andReturn(Optional.of(
        IHostAttributes.build(new HostAttributes().setHost(HOST_A).setMode(DRAINING))));
    expect(storageUtil.attributeStore.getHostAttributes("unknown")).andReturn(Optional.empty());

    control.replay();

    assertEquals(DRAINING, maintenance.getMode(HOST_A));
    assertEquals(NONE, maintenance.getMode("unknown"));
  }

  @Test
  public void testInverseOfferDrain() {
    IScheduledTask task1 = makeTask(HOST_A, "taskA");
    storageUtil.hostMaintenanceStore.saveHostMaintenanceRequest(
        anyObject(IHostMaintenanceRequest.class));
    expectFetchTasksByHost(HOST_A, ImmutableSet.of(task1));
    expectTaskDraining(task1);
    IHostMaintenanceRequest maintenanceRequest =
        IHostMaintenanceRequest.build(new HostMaintenanceRequest()
            .setHost(HOST_A)
            .setCreatedTimestampMs(System.currentTimeMillis())
            .setTimeoutSecs(7200)
            .setDefaultSlaPolicy(SLA_POLICY));
    expect(storageUtil.hostMaintenanceStore.getHostMaintenanceRequest(HOST_A))
        .andReturn(Optional.of(maintenanceRequest)).times(1);

    control.replay();
    maintenance.drainForInverseOffer(INVERSE_OFFER);
  }

  private void expectTaskDraining(IScheduledTask task) {
    expectTaskDraining(task, false);
  }

  private void expectTaskDraining(IScheduledTask task, boolean force) {
    slaManager.checkSlaThenAct(
        eq(task),
        eq(ISlaPolicy.build(SLA_POLICY)),
        anyObject(Storage.MutateWork.class),
        anyObject(ImmutableMap.class),
        eq(force));
  }

  private void expectFetchTasksByHost(String hostName, Set<IScheduledTask> tasks) {
    expect(storageUtil.taskStore.fetchTasks(Query.slaveScoped(hostName).active())).andReturn(tasks);
  }

  private void expectMaintenanceModeChange(String hostName, MaintenanceMode mode) {
    IHostAttributes attributes = IHostAttributes.build(new HostAttributes().setHost(hostName));

    expect(storageUtil.attributeStore.getHostAttributes(hostName))
        .andReturn(Optional.of(attributes));
    IHostAttributes updated = IHostAttributes.build(attributes.newBuilder().setMode(mode));
    expect(storageUtil.attributeStore.saveHostAttributes(updated)).andReturn(true);
  }

  private void assertStatus(String host, MaintenanceMode mode, Set<IHostStatus> statuses) {
    assertEquals(ImmutableSet.of(IHostStatus.build(new HostStatus(host, mode))), statuses);
  }
}