@Override
  public void restartShards(IJobKey jobKey, final Set<Integer> shards, final String requestingUser)
      throws ScheduleException {

    if (!JobKeys.isValid(jobKey)) {
      throw new ScheduleException("Invalid job key: " + jobKey);
    }

    if (shards.isEmpty()) {
      throw new ScheduleException("At least one shard must be specified.");
    }

    final Query.Builder query = Query.instanceScoped(jobKey, shards).active();
    storage.write(
        new MutateWork.NoResult<ScheduleException>() {
          @Override
          protected void execute(MutableStoreProvider storeProvider) throws ScheduleException {

            Set<IScheduledTask> matchingTasks = storeProvider.getTaskStore().fetchTasks(query);
            if (matchingTasks.size() != shards.size()) {
              throw new ScheduleException("Not all requested shards are active.");
            }
            LOG.info("Restarting shards matching " + query);
            for (String taskId : Tasks.ids(matchingTasks)) {
              stateManager.changeState(
                  taskId,
                  Optional.<ScheduleStatus>absent(),
                  RESTARTING,
                  Optional.of("Restarted by " + requestingUser));
            }
          }
        });
  }
  private static Optional<IScheduledTask> getActiveInstance(
      TaskStore taskStore, IJobKey job, int instanceId) {

    return Optional.fromNullable(
        Iterables.getOnlyElement(
            taskStore.fetchTasks(Query.instanceScoped(job, instanceId).active()), null));
  }
  @Test
  public void testGetJobUpdateDiffWithUpdateRemove() throws Exception {
    TaskConfig task1 = defaultTask(false).setNumCpus(1.0);
    TaskConfig task2 = defaultTask(false).setNumCpus(2.0);
    TaskConfig task3 = defaultTask(false).setNumCpus(3.0);

    ImmutableSet.Builder<IScheduledTask> tasks = ImmutableSet.builder();
    makeTasks(0, 10, task1, tasks);
    makeTasks(10, 20, task2, tasks);
    makeTasks(20, 30, task3, tasks);

    expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.absent());
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), tasks.build());

    control.replay();

    JobUpdateRequest request =
        new JobUpdateRequest()
            .setTaskConfig(defaultTask(false).setNumCpus(6.0))
            .setInstanceCount(20)
            .setSettings(new JobUpdateSettings());

    GetJobUpdateDiffResult expected =
        new GetJobUpdateDiffResult()
            .setRemove(ImmutableSet.of(group(task3, new Range(20, 29))))
            .setUpdate(
                ImmutableSet.of(group(task1, new Range(0, 9)), group(task2, new Range(10, 19))))
            .setAdd(ImmutableSet.of())
            .setUnchanged(ImmutableSet.of());

    Response response = assertOkResponse(thrift.getJobUpdateDiff(request));
    assertEquals(expected, response.getResult().getGetJobUpdateDiffResult());
  }
  @Test
  public void testGetPendingReasonFailsStatusSet() throws Exception {
    Builder query = Query.unscoped().byStatus(ScheduleStatus.ASSIGNED);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.getPendingReason(query.get()));
  }
  @Test
  public void testGetPendingReasonFailsSlavesSet() throws Exception {
    Builder query = Query.unscoped().bySlave("host1");

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.getPendingReason(query.get()));
  }
  @Test
  public void testGetTasksStatus() throws Exception {
    Builder query = Query.unscoped();
    Iterable<IScheduledTask> tasks = makeDefaultScheduledTasks(10);
    storageUtil.expectTaskFetch(query, ImmutableSet.copyOf(tasks));

    control.replay();

    ImmutableList<ScheduledTask> expected = IScheduledTask.toBuildersList(tasks);
    Response response = assertOkResponse(thrift.getTasksStatus(new TaskQuery()));
    assertEquals(expected, response.getResult().getScheduleStatusResult().getTasks());
  }
Exemple #7
0
  /**
   * Computes totals for each of the {@link MetricType}s.
   *
   * @return aggregates for each global metric type.
   * @throws StorageException if there was a problem fetching tasks from storage.
   */
  public List<GlobalMetric> computeConsumptionTotals() throws StorageException {
    List<GlobalMetric> counts =
        FluentIterable.from(Arrays.asList(MetricType.values()))
            .transform(TO_GLOBAL_METRIC)
            .toList();

    for (ITaskConfig task : getTasks(Query.unscoped().active())) {
      for (GlobalMetric count : counts) {
        count.accumulate(task);
      }
    }
    return counts;
  }
  @Override
  public synchronized void killTasks(Query.Builder query, final String user)
      throws ScheduleException {

    checkNotNull(query);
    LOG.info("Killing tasks matching " + query);

    boolean jobDeleted = false;

    if (Query.isOnlyJobScoped(query)) {
      // If this looks like a query for all tasks in a job, instruct the scheduler modules to
      // delete the job.
      IJobKey jobKey = JobKeys.from(query).get();
      for (JobManager manager : jobManagers) {
        if (manager.deleteJob(jobKey)) {
          jobDeleted = true;
        }
      }
    }

    // Unless statuses were specifically supplied, only attempt to kill active tasks.
    final Query.Builder taskQuery =
        query.get().isSetStatuses() ? query.byStatus(ACTIVE_STATES) : query;

    int tasksAffected =
        storage.write(
            new MutateWork.Quiet<Integer>() {
              @Override
              public Integer apply(MutableStoreProvider storeProvider) {
                int total = 0;
                for (String taskId :
                    Tasks.ids(storeProvider.getTaskStore().fetchTasks(taskQuery))) {
                  boolean changed =
                      stateManager.changeState(
                          taskId,
                          Optional.<ScheduleStatus>absent(),
                          KILLING,
                          Optional.of("Killed by " + user));

                  if (changed) {
                    total++;
                  }
                }
                return total;
              }
            });

    if (!jobDeleted && (tasksAffected == 0)) {
      throw new ScheduleException("No jobs to kill");
    }
  }
  @Test
  public void testGetPendingReason() throws Exception {
    Builder query = Query.unscoped().byJob(JOB_KEY);
    Builder filterQuery = Query.unscoped().byJob(JOB_KEY).byStatus(ScheduleStatus.PENDING);
    String taskId1 = "task_id_test1";
    String taskId2 = "task_id_test2";
    ImmutableSet<Veto> result =
        ImmutableSet.of(Veto.constraintMismatch("first"), Veto.constraintMismatch("second"));

    ITaskConfig taskConfig = ITaskConfig.build(defaultTask(true));
    IScheduledTask pendingTask1 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(
                    new AssignedTask().setTaskId(taskId1).setTask(taskConfig.newBuilder()))
                .setStatus(ScheduleStatus.PENDING));

    IScheduledTask pendingTask2 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(
                    new AssignedTask().setTaskId(taskId2).setTask(taskConfig.newBuilder()))
                .setStatus(ScheduleStatus.PENDING));

    storageUtil.expectTaskFetch(filterQuery, pendingTask1, pendingTask2);
    expect(nearestFit.getNearestFit(TaskGroupKey.from(taskConfig))).andReturn(result).times(2);

    control.replay();

    String reason = "Constraint not satisfied: first,Constraint not satisfied: second";
    Set<PendingReason> expected =
        ImmutableSet.of(
            new PendingReason().setTaskId(taskId1).setReason(reason),
            new PendingReason().setTaskId(taskId2).setReason(reason));

    Response response = assertOkResponse(thrift.getPendingReason(query.get()));
    assertEquals(expected, response.getResult().getGetPendingReasonResult().getReasons());
  }
  @Test
  public void testModifySnapshotBeforeCommit() throws Exception {
    expect(snapshotStore.createSnapshot()).andReturn(SNAPSHOT1);
    Snapshot modified = SNAPSHOT1.deepCopy().setTasks(ImmutableSet.of(TASK1));
    Capture<MutateWork<Object, Exception>> transaction = createCapture();
    expect(primaryStorage.write(capture(transaction))).andReturn(null);
    distributedStore.persist(modified);
    shutDownNow.execute();

    control.replay();

    clock.advance(INTERVAL);
    storageBackup.createSnapshot();
    String backup1 = storageBackup.createBackupName();
    recovery.stage(backup1);
    assertEquals(
        IScheduledTask.setFromBuilders(SNAPSHOT1.getTasks()), recovery.query(Query.unscoped()));
    recovery.deleteTasks(Query.taskScoped(Tasks.id(TASK2)));
    assertEquals(
        IScheduledTask.setFromBuilders(modified.getTasks()), recovery.query(Query.unscoped()));
    recovery.commit();
    transaction.getValue().apply(storeProvider);
  }
  @Test
  public void testGetTasksWithoutConfigs() throws Exception {
    Builder query = Query.unscoped();
    storageUtil.expectTaskFetch(query, ImmutableSet.copyOf(makeDefaultScheduledTasks(10)));

    control.replay();

    ImmutableList<ScheduledTask> expected =
        IScheduledTask.toBuildersList(
            makeDefaultScheduledTasks(10, defaultTask(true).setExecutorConfig(null)));

    Response response = assertOkResponse(thrift.getTasksWithoutConfigs(new TaskQuery()));
    assertEquals(expected, response.getResult().getScheduleStatusResult().getTasks());
  }
  @Test
  public void testKillExistingCollisionFailedKill() throws Exception {
    IJobConfiguration killExisting =
        IJobConfiguration.build(
            job.newBuilder().setCronCollisionPolicy(CronCollisionPolicy.KILL_EXISTING));
    Capture<Runnable> jobTriggerCapture = expectJobAccepted(killExisting);
    expectActiveTaskFetch(TASK);
    scheduler.killTasks(Query.jobScoped(killExisting.getKey()).active(), CRON_USER);
    expectLastCall().andThrow(new ScheduleException("injected"));

    control.replay();

    cron.receiveJob(new SanitizedConfiguration(killExisting));
    jobTriggerCapture.getValue().run();
  }
  @Test
  public void testEmptyConfigSummary() throws Exception {
    IJobKey key = JobKeys.from("test", "test", "test");

    storageUtil.expectTaskFetch(Query.jobScoped(key).active(), ImmutableSet.of());

    ConfigSummary summary =
        new ConfigSummary().setKey(key.newBuilder()).setGroups(Sets.newHashSet());

    ConfigSummaryResult expected = new ConfigSummaryResult().setSummary(summary);

    control.replay();

    Response response = assertOkResponse(thrift.getConfigSummary(key.newBuilder()));
    assertEquals(expected, response.getResult().getConfigSummaryResult());
  }
  @Test
  public void testGetConfigSummary() throws Exception {
    IJobKey key = JobKeys.from("test", "test", "test");

    TaskConfig firstGroupTask = defaultTask(true);
    TaskConfig secondGroupTask = defaultTask(true).setNumCpus(2);

    IScheduledTask first1 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(firstGroupTask).setInstanceId(0)));

    IScheduledTask first2 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(firstGroupTask).setInstanceId(1)));

    IScheduledTask second =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(secondGroupTask).setInstanceId(2)));

    storageUtil.expectTaskFetch(Query.jobScoped(key).active(), first1, first2, second);

    ConfigGroup group1 =
        new ConfigGroup()
            .setConfig(firstGroupTask)
            .setInstances(IRange.toBuildersSet(convertRanges(toRanges(ImmutableSet.of(0, 1)))));
    ConfigGroup group2 =
        new ConfigGroup()
            .setConfig(secondGroupTask)
            .setInstances(IRange.toBuildersSet(convertRanges(toRanges(ImmutableSet.of(2)))));

    ConfigSummary summary =
        new ConfigSummary().setKey(key.newBuilder()).setGroups(Sets.newHashSet(group1, group2));

    ConfigSummaryResult expected = new ConfigSummaryResult().setSummary(summary);

    control.replay();

    Response response = assertOkResponse(thrift.getConfigSummary(key.newBuilder()));
    assertEquals(
        IConfigSummaryResult.build(expected),
        IConfigSummaryResult.build(response.getResult().getConfigSummaryResult()));
  }
Exemple #15
0
  /**
   * Displays the aggregate utilization for roles within a metric type.
   *
   * @param metric Metric id.
   * @return HTML-formatted utilization within the metric type.
   */
  @GET
  @Path("/{metric}")
  @Produces(MediaType.TEXT_HTML)
  public Response aggregateRoles(@PathParam("metric") final String metric) {
    final MetricType type = getTypeByName(metric);

    Function<ITaskConfig, Display> toKey =
        new Function<ITaskConfig, Display>() {
          @Override
          public Display apply(ITaskConfig task) {
            String role = task.getJob().getRole();
            return new Display(role, metric + "/" + role);
          }
        };
    Map<Display, Metric> byRole =
        counter.computeAggregates(Query.unscoped().active(), type.filter, toKey);
    return Response.ok(fillTemplate(byRole)).build();
  }
Exemple #16
0
  /**
   * Displays the aggregate utilization for jobs within a role.
   *
   * @param metric Metric id.
   * @param role Role for jobs to aggregate.
   * @return HTML-formatted utilization within the metric/role.
   */
  @GET
  @Path("/{metric}/{role}")
  @Produces(MediaType.TEXT_HTML)
  public Response aggregateJobs(
      @PathParam("metric") String metric, @PathParam("role") String role) {

    MetricType type = getTypeByName(metric);
    Function<ITaskConfig, Display> toKey =
        new Function<ITaskConfig, Display>() {
          @Override
          public Display apply(ITaskConfig task) {
            return new Display(task.getJobName(), null);
          }
        };
    Map<Display, Metric> byJob =
        counter.computeAggregates(Query.roleScoped(role).active(), type.filter, toKey);
    return Response.ok(fillTemplate(byJob)).build();
  }
  @Test
  public void testExecution() {
    expect(statsProvider.makeCounter(EXPLICIT_STAT_NAME)).andReturn(explicitRuns);
    expect(statsProvider.makeCounter(IMPLICIT_STAT_NAME)).andReturn(implicitRuns);
    clock = FakeScheduledExecutor.scheduleAtFixedRateExecutor(executorService, 2, 5);

    IScheduledTask task = makeTask("id1", TaskTestUtil.makeConfig(TaskTestUtil.JOB));
    storageUtil.expectOperations();
    storageUtil
        .expectTaskFetch(Query.unscoped().byStatus(Tasks.SLAVE_ASSIGNED_STATES), task)
        .times(5);

    driver.reconcileTasks(ImmutableSet.of(TASK_TO_PROTO.apply(task)));
    expectLastCall().times(5);

    driver.reconcileTasks(ImmutableSet.of());
    expectLastCall().times(2);

    control.replay();

    TaskReconciler reconciler =
        new TaskReconciler(SETTINGS, storageUtil.storage, driver, executorService, statsProvider);

    reconciler.startAsync().awaitRunning();

    clock.advance(INITIAL_DELAY);
    assertEquals(1L, explicitRuns.get());
    assertEquals(0L, implicitRuns.get());

    clock.advance(SPREAD);
    assertEquals(1L, explicitRuns.get());
    assertEquals(1L, implicitRuns.get());

    clock.advance(EXPLICIT_SCHEDULE);
    assertEquals(2L, explicitRuns.get());
    assertEquals(1L, implicitRuns.get());

    clock.advance(IMPLICT_SCHEDULE);
    assertEquals(5L, explicitRuns.get());
    assertEquals(2L, implicitRuns.get());
  }
  @Test
  public void testGetAllJobs() throws Exception {
    JobConfiguration cronJobOne =
        makeJob()
            .setCronSchedule("1 * * * *")
            .setKey(JOB_KEY.newBuilder())
            .setTaskConfig(nonProductionTask());
    JobKey jobKey2 = JOB_KEY.newBuilder().setRole("other_role");
    JobConfiguration cronJobTwo =
        makeJob().setCronSchedule("2 * * * *").setKey(jobKey2).setTaskConfig(nonProductionTask());
    TaskConfig immediateTaskConfig =
        defaultTask(false)
            .setJob(JOB_KEY.newBuilder().setName("immediate"))
            .setJobName("immediate")
            .setOwner(ROLE_IDENTITY);
    IScheduledTask immediateTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(immediateTaskConfig))
                .setStatus(ScheduleStatus.ASSIGNED));
    JobConfiguration immediateJob =
        new JobConfiguration()
            .setKey(JOB_KEY.newBuilder().setName("immediate"))
            .setOwner(ROLE_IDENTITY)
            .setInstanceCount(1)
            .setTaskConfig(immediateTaskConfig);

    Set<JobConfiguration> crons = ImmutableSet.of(cronJobOne, cronJobTwo);
    expect(storageUtil.jobStore.fetchJobs()).andReturn(IJobConfiguration.setFromBuilders(crons));
    storageUtil.expectTaskFetch(Query.unscoped().active(), immediateTask);

    control.replay();

    Set<JobConfiguration> allJobs =
        ImmutableSet.<JobConfiguration>builder().addAll(crons).add(immediateJob).build();
    assertEquals(
        IJobConfiguration.setFromBuilders(allJobs),
        IJobConfiguration.setFromBuilders(
            thrift.getJobs(null).getResult().getGetJobsResult().getConfigs()));
  }
    @Override
    public Amount<Long, Time> getReevaluationDelay(
        IInstanceKey instance,
        IJobUpdateInstructions instructions,
        MutableStoreProvider storeProvider,
        StateManager stateManager,
        JobUpdateStatus status) {

      String taskId =
          Tasks.id(
              Iterables.getOnlyElement(
                  storeProvider
                      .getTaskStore()
                      .fetchTasks(Query.instanceScoped(instance).active())));
      LOG.info("Killing " + instance + " while " + status);
      stateManager.changeState(
          storeProvider,
          taskId,
          Optional.absent(),
          ScheduleStatus.KILLING,
          Optional.of("Killed for job update."));
      return Amount.of(
          (long) instructions.getSettings().getMaxWaitToInstanceRunningMs(), Time.MILLISECONDS);
    }
  @Test
  public void testGetJobSummaryWithoutNextRun() throws Exception {
    // 31st of February, there is no such day.
    String cronSchedule = "* * 31 2 *";

    TaskConfig task =
        nonProductionTask()
            .setJobName(JOB_KEY.getName())
            .setOwner(ROLE_IDENTITY)
            .setEnvironment(JOB_KEY.getEnvironment());
    JobConfiguration job = makeJob().setCronSchedule(cronSchedule).setTaskConfig(task);
    expect(cronPredictor.predictNextRun(CrontabEntry.parse(cronSchedule)))
        .andReturn(Optional.absent())
        .anyTimes();
    storageUtil.expectTaskFetch(Query.roleScoped(ROLE));
    Set<JobConfiguration> jobOnly = ImmutableSet.of(job);
    expect(storageUtil.jobStore.fetchJobs()).andReturn(IJobConfiguration.setFromBuilders(jobOnly));

    control.replay();

    JobSummaryResult result = thrift.getJobSummary(ROLE).getResult().getJobSummaryResult();
    assertEquals(1, result.getSummaries().size());
    assertFalse(result.getSummariesIterator().next().isSetNextCronRunMs());
  }
  @Test
  public void testGetJobUpdateDiffWithUnchanged() throws Exception {
    expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.absent());
    storageUtil.expectTaskFetch(
        Query.jobScoped(JOB_KEY).active(), ImmutableSet.copyOf(makeDefaultScheduledTasks(10)));

    control.replay();

    JobUpdateRequest request =
        new JobUpdateRequest()
            .setTaskConfig(defaultTask(true))
            .setInstanceCount(10)
            .setSettings(new JobUpdateSettings());

    GetJobUpdateDiffResult expected =
        new GetJobUpdateDiffResult()
            .setUnchanged(ImmutableSet.of(group(defaultTask(true), new Range(0, 9))))
            .setRemove(ImmutableSet.of())
            .setUpdate(ImmutableSet.of())
            .setAdd(ImmutableSet.of());

    Response response = assertOkResponse(thrift.getJobUpdateDiff(request));
    assertEquals(expected, response.getResult().getGetJobUpdateDiffResult());
  }
Exemple #22
0
 @VisibleForTesting
 static Query.Builder jobHistoryQuery(IJobKey jobKey) {
   return Query.jobScoped(jobKey).byStatus(apiConstants.TERMINAL_STATES);
 }
  @Test
  public void testCreateAndRestoreNewSnapshot() {
    ImmutableSet<IScheduledTask> tasks =
        ImmutableSet.of(
            IScheduledTask.build(new ScheduledTask().setStatus(ScheduleStatus.PENDING)));
    Set<QuotaConfiguration> quotas =
        ImmutableSet.of(new QuotaConfiguration("steve", ResourceAggregates.none().newBuilder()));
    IHostAttributes attribute =
        IHostAttributes.build(
            new HostAttributes(
                "host", ImmutableSet.of(new Attribute("attr", ImmutableSet.of("value")))));
    StoredJob job =
        new StoredJob(
            "jobManager", new JobConfiguration().setKey(new JobKey("owner", "env", "name")));
    String frameworkId = "framework_id";
    ILock lock =
        ILock.build(
            new Lock()
                .setKey(LockKey.job(JobKeys.from("testRole", "testEnv", "testJob").newBuilder()))
                .setToken("lockId")
                .setUser("testUser")
                .setTimestampMs(12345L));
    SchedulerMetadata metadata =
        new SchedulerMetadata().setFrameworkId(frameworkId).setVersion(CURRENT_API_VERSION);

    storageUtil.expectOperations();
    expect(storageUtil.taskStore.fetchTasks(Query.unscoped())).andReturn(tasks);
    expect(storageUtil.quotaStore.fetchQuotas())
        .andReturn(ImmutableMap.of("steve", ResourceAggregates.none()));
    expect(storageUtil.attributeStore.getHostAttributes()).andReturn(ImmutableSet.of(attribute));
    expect(storageUtil.jobStore.fetchManagerIds()).andReturn(ImmutableSet.of("jobManager"));
    expect(storageUtil.jobStore.fetchJobs("jobManager"))
        .andReturn(ImmutableSet.of(IJobConfiguration.build(job.getJobConfiguration())));
    expect(storageUtil.schedulerStore.fetchFrameworkId()).andReturn(frameworkId);
    expect(storageUtil.lockStore.fetchLocks()).andReturn(ImmutableSet.of(lock));

    expectDataWipe();
    storageUtil.taskStore.saveTasks(tasks);
    storageUtil.quotaStore.saveQuota("steve", ResourceAggregates.none());
    storageUtil.attributeStore.saveHostAttributes(attribute);
    storageUtil.jobStore.saveAcceptedJob(
        job.getJobManagerId(), IJobConfiguration.build(job.getJobConfiguration()));
    storageUtil.schedulerStore.saveFrameworkId(frameworkId);
    storageUtil.lockStore.saveLock(lock);

    control.replay();

    Snapshot expected =
        new Snapshot()
            .setTimestamp(NOW)
            .setTasks(IScheduledTask.toBuildersSet(tasks))
            .setQuotaConfigurations(quotas)
            .setHostAttributes(ImmutableSet.of(attribute.newBuilder()))
            .setJobs(ImmutableSet.of(job))
            .setSchedulerMetadata(metadata)
            .setLocks(ILock.toBuildersSet(ImmutableSet.of(lock)));

    assertEquals(expected, snapshotStore.createSnapshot());

    snapshotStore.applySnapshot(expected);
  }
 private IExpectationSetters<?> expectActiveTaskFetch(IScheduledTask... activeTasks) {
   return storageUtil.expectTaskFetch(Query.jobScoped(job.getKey()).active(), activeTasks);
 }
  @Test
  public void testGetJobSummary() throws Exception {
    long nextCronRunMs = 100;
    TaskConfig ownedCronJobTask =
        nonProductionTask()
            .setJob(JOB_KEY.newBuilder())
            .setJobName(JOB_KEY.getName())
            .setOwner(ROLE_IDENTITY)
            .setEnvironment(JOB_KEY.getEnvironment());
    JobConfiguration ownedCronJob =
        makeJob().setCronSchedule(CRON_SCHEDULE).setTaskConfig(ownedCronJobTask);
    IScheduledTask ownedCronJobScheduledTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(ownedCronJobTask))
                .setStatus(ScheduleStatus.ASSIGNED));
    Identity otherOwner = new Identity("other", "other");
    JobConfiguration unownedCronJob =
        makeJob()
            .setOwner(otherOwner)
            .setCronSchedule(CRON_SCHEDULE)
            .setKey(JOB_KEY.newBuilder().setRole("other"))
            .setTaskConfig(ownedCronJobTask.deepCopy().setOwner(otherOwner));
    TaskConfig ownedImmediateTaskInfo =
        defaultTask(false)
            .setJob(JOB_KEY.newBuilder().setName("immediate"))
            .setJobName("immediate")
            .setOwner(ROLE_IDENTITY);
    Set<JobConfiguration> ownedCronJobOnly = ImmutableSet.of(ownedCronJob);
    Set<JobSummary> ownedCronJobSummaryOnly =
        ImmutableSet.of(
            new JobSummary()
                .setJob(ownedCronJob)
                .setStats(new JobStats())
                .setNextCronRunMs(nextCronRunMs));
    Set<JobSummary> ownedCronJobSummaryWithRunningTask =
        ImmutableSet.of(
            new JobSummary()
                .setJob(ownedCronJob)
                .setStats(new JobStats().setActiveTaskCount(1))
                .setNextCronRunMs(nextCronRunMs));
    Set<JobConfiguration> unownedCronJobOnly = ImmutableSet.of(unownedCronJob);
    Set<JobConfiguration> bothCronJobs = ImmutableSet.of(ownedCronJob, unownedCronJob);

    IScheduledTask ownedImmediateTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(ownedImmediateTaskInfo))
                .setStatus(ScheduleStatus.ASSIGNED));
    JobConfiguration ownedImmediateJob =
        new JobConfiguration()
            .setKey(JOB_KEY.newBuilder().setName("immediate"))
            .setOwner(ROLE_IDENTITY)
            .setInstanceCount(1)
            .setTaskConfig(ownedImmediateTaskInfo);
    Builder query = Query.roleScoped(ROLE);

    Set<JobSummary> ownedImmediateJobSummaryOnly =
        ImmutableSet.of(
            new JobSummary()
                .setJob(ownedImmediateJob)
                .setStats(new JobStats().setActiveTaskCount(1)));

    expect(cronPredictor.predictNextRun(CrontabEntry.parse(CRON_SCHEDULE)))
        .andReturn(Optional.of(new Date(nextCronRunMs)))
        .anyTimes();

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(ownedCronJobOnly));

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(bothCronJobs));

    storageUtil.expectTaskFetch(query, ownedImmediateTask);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(unownedCronJobOnly));

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs()).andReturn(ImmutableSet.of());

    // Handle the case where a cron job has a running task (same JobKey present in both stores).
    storageUtil.expectTaskFetch(query, ownedCronJobScheduledTask);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(ImmutableSet.of(ownedCronJob)));

    control.replay();

    assertEquals(jobSummaryResponse(ownedCronJobSummaryOnly), thrift.getJobSummary(ROLE));

    assertEquals(jobSummaryResponse(ownedCronJobSummaryOnly), thrift.getJobSummary(ROLE));

    Response jobSummaryResponse = thrift.getJobSummary(ROLE);
    assertEquals(
        jobSummaryResponse(ownedImmediateJobSummaryOnly),
        IResponse.build(jobSummaryResponse).newBuilder());

    assertEquals(jobSummaryResponse(ImmutableSet.of()), thrift.getJobSummary(ROLE));

    assertEquals(
        jobSummaryResponse(ownedCronJobSummaryWithRunningTask), thrift.getJobSummary(ROLE));
  }
 @Benchmark
 public Iterable<IScheduledTask> run() {
   return storage.read(store -> store.getTaskStore().fetchTasks(Query.unscoped()));
 }
  @Test
  public void testGetJobs() throws Exception {
    TaskConfig ownedCronJobTask =
        nonProductionTask()
            .setJobName(JOB_KEY.getName())
            .setOwner(ROLE_IDENTITY)
            .setEnvironment(JOB_KEY.getEnvironment());
    JobConfiguration ownedCronJob =
        makeJob().setCronSchedule(CRON_SCHEDULE).setTaskConfig(ownedCronJobTask);
    IScheduledTask ownedCronJobScheduledTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(ownedCronJobTask))
                .setStatus(ScheduleStatus.ASSIGNED));
    Identity otherOwner = new Identity("other", "other");
    JobConfiguration unownedCronJob =
        makeJob()
            .setOwner(otherOwner)
            .setCronSchedule(CRON_SCHEDULE)
            .setKey(JOB_KEY.newBuilder().setRole("other"))
            .setTaskConfig(ownedCronJobTask.deepCopy().setOwner(otherOwner));
    TaskConfig ownedImmediateTaskInfo =
        defaultTask(false)
            .setJob(JOB_KEY.newBuilder().setName("immediate"))
            .setJobName("immediate")
            .setOwner(ROLE_IDENTITY);
    Set<JobConfiguration> ownedCronJobOnly = ImmutableSet.of(ownedCronJob);
    Set<JobConfiguration> unownedCronJobOnly = ImmutableSet.of(unownedCronJob);
    Set<JobConfiguration> bothCronJobs = ImmutableSet.of(ownedCronJob, unownedCronJob);
    IScheduledTask ownedImmediateTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(ownedImmediateTaskInfo))
                .setStatus(ScheduleStatus.ASSIGNED));
    JobConfiguration ownedImmediateJob =
        new JobConfiguration()
            .setKey(JOB_KEY.newBuilder().setName("immediate"))
            .setOwner(ROLE_IDENTITY)
            .setInstanceCount(1)
            .setTaskConfig(ownedImmediateTaskInfo);
    Query.Builder query = Query.roleScoped(ROLE).active();

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(ownedCronJobOnly));

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(bothCronJobs));

    storageUtil.expectTaskFetch(query, ownedImmediateTask);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(unownedCronJobOnly));

    expect(storageUtil.jobStore.fetchJobs()).andReturn(ImmutableSet.of());
    storageUtil.expectTaskFetch(query);

    // Handle the case where a cron job has a running task (same JobKey present in both stores).
    storageUtil.expectTaskFetch(query, ownedCronJobScheduledTask);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(ImmutableSet.of(ownedCronJob)));

    control.replay();

    assertJobsEqual(
        ownedCronJob,
        Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));

    assertJobsEqual(
        ownedCronJob,
        Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));

    Set<JobConfiguration> queryResult3 =
        thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs();
    assertJobsEqual(ownedImmediateJob, Iterables.getOnlyElement(queryResult3));
    assertEquals(
        ITaskConfig.build(ownedImmediateTaskInfo),
        ITaskConfig.build(Iterables.getOnlyElement(queryResult3).getTaskConfig()));

    assertTrue(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs().isEmpty());

    assertJobsEqual(
        ownedCronJob,
        Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));
  }
 private TaskQuery setupPaginatedQuery(Iterable<IScheduledTask> tasks, int offset, int limit) {
   TaskQuery query = new TaskQuery().setOffset(offset).setLimit(limit);
   Builder builder = Query.arbitrary(query);
   storageUtil.expectTaskFetch(builder, ImmutableSet.copyOf(tasks));
   return query;
 }