@Test
  public void testResistsStarvation() {
    // TODO(wfarner): This test requires intimate knowledge of the way futures are used inside
    // TaskScheduler.  It's time to test using a real ScheduledExecutorService.

    expectAnyMaintenanceCalls();

    IScheduledTask jobA0 = makeTask("a0", PENDING);

    ScheduledTask jobA1Builder = jobA0.newBuilder();
    jobA1Builder.getAssignedTask().setTaskId("a1");
    jobA1Builder.getAssignedTask().setInstanceId(1);
    IScheduledTask jobA1 = IScheduledTask.build(jobA1Builder);

    ScheduledTask jobA2Builder = jobA0.newBuilder();
    jobA2Builder.getAssignedTask().setTaskId("a2");
    jobA2Builder.getAssignedTask().setInstanceId(2);
    IScheduledTask jobA2 = IScheduledTask.build(jobA2Builder);

    IScheduledTask jobB0 = makeTask("b0", PENDING);

    expectOfferDeclineIn(10);
    expectOfferDeclineIn(10);
    expectOfferDeclineIn(10);
    expectOfferDeclineIn(10);

    Capture<Runnable> timeoutA = expectTaskGroupBackoff(10);
    Capture<Runnable> timeoutB = expectTaskGroupBackoff(10);

    Capture<IScheduledTask> firstScheduled = expectTaskScheduled(jobA0);
    Capture<IScheduledTask> secondScheduled = expectTaskScheduled(jobB0);

    // Expect another watch of the task group for job A.
    expectTaskGroupBackoff(10);

    replayAndCreateScheduler();

    offerQueue.addOffer(OFFER_A);
    offerQueue.addOffer(OFFER_B);
    offerQueue.addOffer(OFFER_C);
    offerQueue.addOffer(OFFER_D);
    changeState(jobA0, INIT, PENDING);
    changeState(jobA1, INIT, PENDING);
    changeState(jobA2, INIT, PENDING);
    changeState(jobB0, INIT, PENDING);
    timeoutA.getValue().run();
    timeoutB.getValue().run();
    assertEquals(
        ImmutableSet.of(jobA0, jobB0),
        ImmutableSet.of(firstScheduled.getValue(), secondScheduled.getValue()));
  }
 private TaskInfo makeTaskInfo(IScheduledTask task) {
   return TaskInfo.newBuilder()
       .setName(Tasks.id(task))
       .setTaskId(TaskID.newBuilder().setValue(Tasks.id(task)))
       .setSlaveId(SlaveID.newBuilder().setValue("slave-id" + task.toString()))
       .build();
 }
  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 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();
  }
  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));
  }
 private IScheduledTask makeTask(String taskId) {
   return IScheduledTask.build(
       new ScheduledTask()
           .setAssignedTask(
               new AssignedTask()
                   .setInstanceId(0)
                   .setTaskId(taskId)
                   .setTask(
                       new TaskConfig()
                           .setJobName("job-" + taskId)
                           .setOwner(
                               new Identity().setRole("role-" + taskId).setUser("user-" + taskId))
                           .setEnvironment("env-" + taskId))));
 }
  @Test
  public void testNoPenaltyForInterruptedTasks() {
    expectAnyMaintenanceCalls();

    makeFlappyTaskWithStates("a0", EnumSet.of(INIT, PENDING, ASSIGNED, RESTARTING, FAILED), null);
    IScheduledTask taskA1 =
        IScheduledTask.build(makeTask("a1", INIT).newBuilder().setAncestorId("a0"));

    expectOfferDeclineIn(10);
    Capture<Runnable> first = expectTaskGroupBackoff(10);

    expectTaskScheduled(taskA1);

    replayAndCreateScheduler();

    offerQueue.addOffer(OFFER_A);

    changeState(taskA1, INIT, PENDING);

    first.getValue().run();
  }
 private IScheduledTask makeTask(String taskId, ScheduleStatus status) {
   return IScheduledTask.build(makeTask(taskId).newBuilder().setStatus(status));
 }