private IScheduledTask makeFlappyTaskWithStates(
String taskId, Iterable<ScheduleStatus> states, @Nullable String ancestorId) {
Amount<Long, Time> timeInState = Amount.of(10L, Time.SECONDS);
ScheduledTask base = makeTask(taskId, INIT).newBuilder();
for (ScheduleStatus status : states) {
base.addToTaskEvents(new TaskEvent(clock.nowMillis(), status));
clock.advance(timeInState);
}
base.setAncestorId(ancestorId);
final IScheduledTask result = IScheduledTask.build(base);
// Insert the task if it doesn't already exist.
storage.write(
new MutateWork.NoResult.Quiet() {
@Override
protected void execute(MutableStoreProvider storeProvider) {
TaskStore.Mutable taskStore = storeProvider.getUnsafeTaskStore();
if (taskStore.fetchTasks(Query.taskScoped(Tasks.id(result))).isEmpty()) {
taskStore.saveTasks(ImmutableSet.of(result));
}
}
});
return result;
}
@Test
public void testTaskDeleted() {
expectAnyMaintenanceCalls();
expectOfferDeclineIn(10);
final IScheduledTask task = makeTask("a", PENDING);
Capture<Runnable> timeoutCapture = expectTaskGroupBackoff(10);
expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.<TaskInfo>absent());
expectTaskGroupBackoff(10, 20);
expect(preemptor.findPreemptionSlotFor("a")).andReturn(Optional.<String>absent());
replayAndCreateScheduler();
offerQueue.addOffer(OFFER_A);
changeState(task, INIT, PENDING);
timeoutCapture.getValue().run();
// Ensure the offer was consumed.
changeState(task, INIT, PENDING);
storage.write(
new MutateWork.NoResult.Quiet() {
@Override
protected void execute(MutableStoreProvider storeProvider) {
storeProvider.getUnsafeTaskStore().deleteTasks(Tasks.ids(task));
}
});
taskGroups.tasksDeleted(new TasksDeleted(ImmutableSet.of(task)));
timeoutCapture.getValue().run();
}
private void changeState(IScheduledTask task, ScheduleStatus oldState, ScheduleStatus newState) {
final IScheduledTask copy = IScheduledTask.build(task.newBuilder().setStatus(newState));
// Insert the task if it doesn't already exist.
storage.write(
new MutateWork.NoResult.Quiet() {
@Override
protected void execute(MutableStoreProvider storeProvider) {
TaskStore.Mutable taskStore = storeProvider.getUnsafeTaskStore();
if (taskStore.fetchTasks(Query.taskScoped(Tasks.id(copy))).isEmpty()) {
taskStore.saveTasks(ImmutableSet.of(copy));
}
}
});
taskGroups.taskChangedState(new TaskStateChange(copy, oldState));
}
@Test
public void testLoadFromStorage() {
expectTaskGroupBackoff(10);
replayAndCreateScheduler();
final IScheduledTask c = makeTask("c", RUNNING);
storage.write(
new MutateWork.NoResult.Quiet() {
@Override
protected void execute(MutableStoreProvider store) {
store
.getUnsafeTaskStore()
.saveTasks(ImmutableSet.of(makeTask("a", KILLED), makeTask("b", PENDING), c));
}
});
taskGroups.storageStarted(new StorageStarted());
changeState(c, RUNNING, FINISHED);
}
