FakeMaster.java example

Explorer
aurora-master
/**
 * 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.app.local;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import javax.inject.Inject;

import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;
import com.google.common.eventbus.EventBus;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.SettableFuture;

import org.apache.aurora.scheduler.app.local.simulator.events.Started;
import org.apache.aurora.scheduler.mesos.DriverFactory;
import org.apache.mesos.Protos.ExecutorID;
import org.apache.mesos.Protos.Filters;
import org.apache.mesos.Protos.FrameworkID;
import org.apache.mesos.Protos.MasterInfo;
import org.apache.mesos.Protos.Offer;
import org.apache.mesos.Protos.OfferID;
import org.apache.mesos.Protos.Request;
import org.apache.mesos.Protos.SlaveID;
import org.apache.mesos.Protos.Status;
import org.apache.mesos.Protos.TaskID;
import org.apache.mesos.Protos.TaskInfo;
import org.apache.mesos.Protos.TaskState;
import org.apache.mesos.Protos.TaskStatus;
import org.apache.mesos.Scheduler;
import org.apache.mesos.SchedulerDriver;
import org.apache.mesos.v1.Protos;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.util.Objects.requireNonNull;

/**
 * A simulated master for use in scheduler testing.
 */
@SuppressWarnings("deprecation")
public class FakeMaster implements SchedulerDriver, DriverFactory {

  private static final Logger LOG = LoggerFactory.getLogger(FakeMaster.class);

  private final Map<TaskID, Task> activeTasks = Collections.synchronizedMap(Maps.newHashMap());
  private final Map<OfferID, Offer> idleOffers = Collections.synchronizedMap(Maps.newHashMap());
  private final Map<OfferID, Offer> sentOffers = Collections.synchronizedMap(Maps.newHashMap());

  private final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
  private final CountDownLatch stopped = new CountDownLatch(1);

  private final SettableFuture<Scheduler> schedulerFuture = SettableFuture.create();
  private final EventBus eventBus;

  @Inject
  FakeMaster(EventBus eventBus) {
    this.eventBus = requireNonNull(eventBus);
  }

  public void addResources(Iterable<Offer> offers) {
    assertNotStopped();

    synchronized (idleOffers) {
      for (Offer offer : offers) {
        checkState(!idleOffers.containsKey(offer.getId()), "Duplicate offer id " + offer.getId());
        idleOffers.put(offer.getId(), offer);
      }
    }
  }

  public void changeState(TaskID task, TaskState state) {
    assertNotStopped();

    checkState(activeTasks.containsKey(task), "Task " + task + " does not exist.");
    Futures.getUnchecked(schedulerFuture).statusUpdate(this, TaskStatus.newBuilder()
        .setTaskId(task)
        .setState(state)
        .build());
  }

  @Override
  public SchedulerDriver create(
      Scheduler scheduler,
      Optional<Protos.Credential> credentials,
      Protos.FrameworkInfo frameworkInfo,
      String master) {

    schedulerFuture.set(scheduler);
    return this;
  }

  @Override
  public Status start() {
    assertNotStopped();

    Futures.getUnchecked(schedulerFuture).registered(this,
        FrameworkID.newBuilder().setValue("local").build(),
        MasterInfo.getDefaultInstance());

    eventBus.post(new Started());

    executor.scheduleAtFixedRate(
        () -> {
          List<Offer> allOffers;
          synchronized (sentOffers) {
            synchronized (idleOffers) {
              sentOffers.putAll(idleOffers);
              allOffers = ImmutableList.copyOf(idleOffers.values());
              idleOffers.clear();
            }
          }

          if (allOffers.isEmpty()) {
            LOG.info("All offers consumed, suppressing offer cycle.");
          } else {
            Futures.getUnchecked(schedulerFuture).resourceOffers(this, allOffers);
          }
        },
        1,
        5,
        TimeUnit.SECONDS);

    return Status.DRIVER_RUNNING;
  }

  @Override
  public Status stop(boolean failover) {
    return stop();
  }

  @Override
  public Status stop() {
    stopped.countDown();
    MoreExecutors.shutdownAndAwaitTermination(executor, 1, TimeUnit.SECONDS);
    return Status.DRIVER_RUNNING;
  }

  @Override
  public Status abort() {
    throw new UnsupportedOperationException();
  }

  @Override
  public Status join() {
    try {
      stopped.await();
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
    }
    return Status.DRIVER_RUNNING;
  }

  @Override
  public Status run() {
    assertNotStopped();
    return Status.DRIVER_RUNNING;
  }

  @Override
  public Status requestResources(Collection<Request> requests) {
    throw new UnsupportedOperationException();
  }

  private void assertNotStopped() {
    Preconditions.checkState(stopped.getCount() == 1, "Driver is stopped.");
  }

  private void checkState(boolean assertion, String failureMessage) {
    if (!assertion) {
      Futures.getUnchecked(schedulerFuture).error(this, failureMessage);
      stop();
      throw new IllegalStateException(failureMessage);
    }
  }

  @Override
  public Status launchTasks(
      Collection<OfferID> offerIds,
      Collection<TaskInfo> tasks,
      Filters filters) {

    throw new UnsupportedOperationException();
  }

  @Override
  public Status launchTasks(Collection<OfferID> offerIds, Collection<TaskInfo> tasks) {
    assertNotStopped();

    OfferID id = Iterables.getOnlyElement(offerIds);
    Offer offer = sentOffers.remove(id);
    checkState(offer != null, "Offer " + id + " is invalid.");

    final TaskInfo task = Iterables.getOnlyElement(tasks);
    synchronized (activeTasks) {
      checkState(
          !activeTasks.containsKey(task.getTaskId()),
          "Task " + task.getTaskId() + " already exists.");
      activeTasks.put(task.getTaskId(), new Task(offer, task));
    }

    executor.schedule(
        () -> Futures.getUnchecked(schedulerFuture).statusUpdate(
            this,
            TaskStatus.newBuilder()
                .setTaskId(task.getTaskId())
                .setState(TaskState.TASK_RUNNING)
                .build()),
        1,
        TimeUnit.SECONDS);

    return Status.DRIVER_RUNNING;
  }

  @Override
  public Status launchTasks(OfferID offerId, Collection<TaskInfo> tasks, Filters filters) {
    throw new UnsupportedOperationException();
  }

  @Override
  public Status launchTasks(OfferID offerId, Collection<TaskInfo> tasks) {
    throw new UnsupportedOperationException();
  }

  @Override
  public Status killTask(TaskID taskId) {
    assertNotStopped();

    Task task = activeTasks.remove(taskId);
    checkState(task != null, "Task " + taskId + " not found.");
    idleOffers.put(task.getOffer().getId(), task.getOffer());

    Futures.getUnchecked(schedulerFuture).statusUpdate(this, TaskStatus.newBuilder()
        .setTaskId(taskId)
        .setState(TaskState.TASK_FINISHED)
        .build());

    return Status.DRIVER_RUNNING;
  }

  @Override
  public Status declineOffer(OfferID offerId, Filters filters) {
    assertNotStopped();

    throw new UnsupportedOperationException();
  }

  @Override
  public Status declineOffer(OfferID offerId) {
    assertNotStopped();
    return null;
  }

  @Override
  public Status reviveOffers() {
    assertNotStopped();
    throw new UnsupportedOperationException();
  }

  @Override
  public Status acknowledgeStatusUpdate(TaskStatus status) {
    assertNotStopped();
    throw new UnsupportedOperationException();
  }

  @Override
  public Status sendFrameworkMessage(ExecutorID executorId, SlaveID slaveId, byte[] data) {
    throw new UnsupportedOperationException();
  }

  @Override
  public Status reconcileTasks(Collection<TaskStatus> statuses) {
    throw new UnsupportedOperationException();
  }

  @Override
  public Status acceptOffers(
      Collection<OfferID> offerIds,
      Collection<Offer.Operation> operations,
      Filters filters) {
    throw new UnsupportedOperationException();
  }

  @Override
  public Status suppressOffers() {
    throw new UnsupportedOperationException();
  }

  private static final class Task {
    private final Offer offer;
    private final TaskInfo taskInfo;

    private Task(Offer offer, TaskInfo taskInfo) {
      this.offer = offer;
      this.taskInfo = taskInfo;
    }

    Offer getOffer() {
      return offer;
    }

    TaskInfo getTask() {
      return taskInfo;
    }
  }
}