@Test
public void testStorageException() {
IScheduledTask task = makeTask("a", PENDING);
TaskInfo mesosTask =
TaskInfo.newBuilder()
.setName(Tasks.id(task))
.setTaskId(TaskID.newBuilder().setValue(Tasks.id(task)))
.setSlaveId(SlaveID.newBuilder().setValue("slaveId"))
.build();
Capture<Runnable> timeoutCapture = expectTaskGroupBackoff(10);
expectAnyMaintenanceCalls();
expectOfferDeclineIn(10);
expect(assigner.maybeAssign(OFFER_A, task)).andThrow(new StorageException("Injected failure."));
Capture<Runnable> timeoutCapture2 = expectTaskGroupBackoff(10, 20);
expect(preemptor.findPreemptionSlotFor("a")).andReturn(Optional.<String>absent());
expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.of(mesosTask));
driver.launchTask(OFFER_A.getId(), mesosTask);
expectLastCall();
replayAndCreateScheduler();
changeState(task, INIT, PENDING);
offerQueue.addOffer(OFFER_A);
timeoutCapture.getValue().run();
timeoutCapture2.getValue().run();
}
@Test
public void testDriverNotReady() {
IScheduledTask task = makeTask("a", PENDING);
TaskInfo mesosTask =
TaskInfo.newBuilder()
.setName(Tasks.id(task))
.setTaskId(TaskID.newBuilder().setValue(Tasks.id(task)))
.setSlaveId(SlaveID.newBuilder().setValue("slaveId"))
.build();
Capture<Runnable> timeoutCapture = expectTaskGroupBackoff(10);
expectAnyMaintenanceCalls();
expectOfferDeclineIn(10);
expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.of(mesosTask));
driver.launchTask(OFFER_A.getId(), mesosTask);
expectLastCall().andThrow(new IllegalStateException("Driver not ready."));
expect(
stateManager.changeState(
Query.taskScoped("a").byStatus(PENDING), LOST, TaskScheduler.LAUNCH_FAILED_MSG))
.andReturn(1);
replayAndCreateScheduler();
changeState(task, INIT, PENDING);
offerQueue.addOffer(OFFER_A);
timeoutCapture.getValue().run();
}
@Test
public void testTaskAssigned() {
expectAnyMaintenanceCalls();
expectOfferDeclineIn(10);
IScheduledTask task = makeTask("a", PENDING);
TaskInfo mesosTask = makeTaskInfo(task);
Capture<Runnable> timeoutCapture = expectTaskGroupBackoff(10);
expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.<TaskInfo>absent());
expect(preemptor.findPreemptionSlotFor("a")).andReturn(Optional.<String>absent());
Capture<Runnable> timeoutCapture2 = expectTaskGroupBackoff(10, 20);
expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.of(mesosTask));
driver.launchTask(OFFER_A.getId(), mesosTask);
Capture<Runnable> timeoutCapture3 = expectTaskGroupBackoff(10);
expectTaskGroupBackoff(10, 20);
expect(preemptor.findPreemptionSlotFor("b")).andReturn(Optional.<String>absent());
replayAndCreateScheduler();
offerQueue.addOffer(OFFER_A);
changeState(task, INIT, PENDING);
timeoutCapture.getValue().run();
timeoutCapture2.getValue().run();
// Ensure the offer was consumed.
changeState(makeTask("b"), INIT, PENDING);
timeoutCapture3.getValue().run();
}
private Capture<IScheduledTask> expectTaskScheduled(IScheduledTask task) {
TaskInfo mesosTask = makeTaskInfo(task);
Capture<IScheduledTask> taskScheduled = createCapture();
expect(assigner.maybeAssign(EasyMock.<Offer>anyObject(), capture(taskScheduled)))
.andReturn(Optional.of(mesosTask));
driver.launchTask(EasyMock.<OfferID>anyObject(), eq(mesosTask));
return taskScheduled;
}
@Test
public void testChangingMaintenancePreferences() {
expectOffer();
expect(maintenance.getMode("HOST_A")).andReturn(MaintenanceMode.NONE);
expectOffer();
expect(maintenance.getMode("HOST_B")).andReturn(MaintenanceMode.SCHEDULED);
expectOffer();
expect(maintenance.getMode("HOST_C")).andReturn(MaintenanceMode.DRAINED);
IScheduledTask taskA = makeTask("A", PENDING);
TaskInfo mesosTaskA = makeTaskInfo(taskA);
expect(assigner.maybeAssign(OFFER_B, taskA)).andReturn(Optional.of(mesosTaskA));
driver.launchTask(OFFER_B.getId(), mesosTaskA);
Capture<Runnable> captureA = expectTaskGroupBackoff(10);
IScheduledTask taskB = makeTask("B", PENDING);
TaskInfo mesosTaskB = makeTaskInfo(taskB);
expect(assigner.maybeAssign(OFFER_C, taskB)).andReturn(Optional.of(mesosTaskB));
driver.launchTask(OFFER_C.getId(), mesosTaskB);
Capture<Runnable> captureB = expectTaskGroupBackoff(10);
replayAndCreateScheduler();
offerQueue.addOffer(OFFER_A);
offerQueue.addOffer(OFFER_B);
offerQueue.addOffer(OFFER_C);
// Initially, we'd expect the offers to be consumed in order (A, B), with (C) unschedulable
// Expected order now (B), with (C, A) unschedulable
changeHostMaintenanceState("HOST_A", MaintenanceMode.DRAINING);
changeState(taskA, INIT, PENDING);
captureA.getValue().run();
// Expected order now (C), with (A) unschedulable and (B) already consumed
changeHostMaintenanceState("HOST_C", MaintenanceMode.NONE);
changeState(taskB, INIT, PENDING);
captureB.getValue().run();
}
@Test
public void testBasicMaintenancePreferences() {
expectOffer();
expect(maintenance.getMode("HOST_D")).andReturn(MaintenanceMode.DRAINED);
expectOffer();
expect(maintenance.getMode("HOST_C")).andReturn(MaintenanceMode.DRAINING);
expectOffer();
expect(maintenance.getMode("HOST_B")).andReturn(MaintenanceMode.SCHEDULED);
expectOffer();
expect(maintenance.getMode("HOST_A")).andReturn(MaintenanceMode.NONE);
IScheduledTask taskA = makeTask("A", PENDING);
TaskInfo mesosTaskA = makeTaskInfo(taskA);
expect(assigner.maybeAssign(OFFER_A, taskA)).andReturn(Optional.of(mesosTaskA));
driver.launchTask(OFFER_A.getId(), mesosTaskA);
Capture<Runnable> captureA = expectTaskGroupBackoff(10);
IScheduledTask taskB = makeTask("B", PENDING);
TaskInfo mesosTaskB = makeTaskInfo(taskB);
expect(assigner.maybeAssign(OFFER_B, taskB)).andReturn(Optional.of(mesosTaskB));
driver.launchTask(OFFER_B.getId(), mesosTaskB);
Capture<Runnable> captureB = expectTaskGroupBackoff(10);
replayAndCreateScheduler();
offerQueue.addOffer(OFFER_D);
offerQueue.addOffer(OFFER_C);
offerQueue.addOffer(OFFER_B);
offerQueue.addOffer(OFFER_A);
changeState(taskA, INIT, PENDING);
captureA.getValue().run();
changeState(taskB, INIT, PENDING);
captureB.getValue().run();
}
@Test
public void testDontDeclineAcceptedOffer() throws OfferQueue.LaunchException {
expectAnyMaintenanceCalls();
Capture<Runnable> offerExpirationCapture = expectOfferDeclineIn(10);
Function<Offer, Optional<TaskInfo>> offerAcceptor =
createMock(new Clazz<Function<Offer, Optional<TaskInfo>>>() {});
final TaskInfo taskInfo = TaskInfo.getDefaultInstance();
expect(offerAcceptor.apply(OFFER_A)).andReturn(Optional.of(taskInfo));
driver.launchTask(OFFER_A.getId(), taskInfo);
replayAndCreateScheduler();
offerQueue.addOffer(OFFER_A);
offerQueue.launchFirst(offerAcceptor);
offerExpirationCapture.getValue().run();
}