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;
}
@Before
public void setUp() {
FakeClock clock = new FakeClock();
clock.setNowMillis(NOW);
storageUtil = new StorageTestUtil(this);
snapshotStore = new SnapshotStoreImpl(clock, storageUtil.storage);
}
@Test
public void testFlappingTasks() {
expectAnyMaintenanceCalls();
makeFlappyTask("a0", null);
IScheduledTask taskA1 =
IScheduledTask.build(makeTask("a1", INIT).newBuilder().setAncestorId("a0"));
expectOfferDeclineIn(10);
Capture<Runnable> first = expectTaskGroupBackoff(10);
expect(flappingStrategy.calculateBackoffMs(0)).andReturn(5L);
// Since A1 has been penalized, the task has to wait for another 10 ms until the penalty has
// expired.
Capture<Runnable> flapping = expectTaskRetryIn(10);
expectTaskScheduled(taskA1);
replayAndCreateScheduler();
offerQueue.addOffer(OFFER_A);
changeState(taskA1, INIT, PENDING);
first.getValue().run();
clock.waitFor(10);
flapping.getValue().run();
}
@Test
public void testFlappingTasksBackoffTruncation() {
expectAnyMaintenanceCalls();
makeFlappyTask("a0", null);
makeFlappyTask("a1", "a0");
makeFlappyTask("a2", "a1");
IScheduledTask taskA3 =
IScheduledTask.build(makeTask("a3", INIT).newBuilder().setAncestorId("a2"));
expectOfferDeclineIn(10);
Capture<Runnable> first = expectTaskGroupBackoff(10);
// The ancestry chain is 3 long, but if the backoff strategy truncates, we don't traverse the
// entire history.
expect(flappingStrategy.calculateBackoffMs(0)).andReturn(5L);
expect(flappingStrategy.calculateBackoffMs(5L)).andReturn(5L);
Capture<Runnable> flapping = expectTaskRetryIn(10);
expectTaskScheduled(taskA3);
replayAndCreateScheduler();
offerQueue.addOffer(OFFER_A);
changeState(taskA3, INIT, PENDING);
first.getValue().run();
clock.waitFor(10);
flapping.getValue().run();
}
@Override
public TemporaryStorage apply(Snapshot snapshot) {
final Storage storage = DbUtil.createFlaggedStorage();
FakeClock clock = new FakeClock();
clock.setNowMillis(snapshot.getTimestamp());
final SnapshotStore<Snapshot> snapshotStore = new SnapshotStoreImpl(clock, storage);
snapshotStore.applySnapshot(snapshot);
return new TemporaryStorage() {
@Override
public void deleteTasks(final Query.Builder query) {
storage.write(
new MutateWork.NoResult.Quiet() {
@Override
public void execute(MutableStoreProvider storeProvider) {
Set<String> ids =
FluentIterable.from(storeProvider.getTaskStore().fetchTasks(query))
.transform(Tasks::id)
.toSet();
storeProvider.getUnsafeTaskStore().deleteTasks(ids);
}
});
}
@Override
public Iterable<IScheduledTask> fetchTasks(final Query.Builder query) {
return storage.read(
new Work.Quiet<Iterable<IScheduledTask>>() {
@Override
public Iterable<IScheduledTask> apply(StoreProvider storeProvider) {
return storeProvider.getTaskStore().fetchTasks(query);
}
});
}
@Override
public Snapshot toSnapshot() {
return snapshotStore.createSnapshot();
}
};
}
@Before
public void setUp() {
storage = MemStorage.newEmptyStorage();
maintenance = createMock(MaintenanceController.class);
stateManager = createMock(StateManager.class);
assigner = createMock(TaskAssigner.class);
retryStrategy = createMock(BackoffStrategy.class);
driver = createMock(Driver.class);
executor = createMock(ScheduledExecutorService.class);
future = createMock(ScheduledFuture.class);
returnDelay = createMock(OfferReturnDelay.class);
clock = new FakeClock();
clock.setNowMillis(0);
flappingStrategy = createMock(BackoffStrategy.class);
preemptor = createMock(Preemptor.class);
}
@Test
public void testRecover() throws Exception {
expect(snapshotStore.createSnapshot()).andReturn(SNAPSHOT1);
Capture<MutateWork<Object, Exception>> transaction = createCapture();
expect(primaryStorage.write(capture(transaction))).andReturn(null);
distributedStore.persist(SNAPSHOT1);
shutDownNow.execute();
control.replay();
assertEquals(ImmutableSet.<String>of(), recovery.listBackups());
clock.advance(INTERVAL);
storageBackup.createSnapshot();
String backup1 = storageBackup.createBackupName();
assertEquals(ImmutableSet.of(backup1), recovery.listBackups());
recovery.stage(backup1);
assertEquals(
IScheduledTask.setFromBuilders(SNAPSHOT1.getTasks()), recovery.query(Query.unscoped()));
recovery.commit();
transaction.getValue().apply(storeProvider);
}