Java tutorial
/** * 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; } } }