ReadOnlySchedulerImpl implements ReadOnlyScheduler.Iface { @Override public Response getJobUpdateDiff(JobUpdateRequest mutableRequest) { IJobUpdateRequest request; try { request = IJobUpdateRequest.build(new JobUpdateRequest(mutableRequest).setTaskConfig( configurationManager.validateAndPopulate( ITaskConfig.build(mutableRequest.getTaskConfig())).newBuilder())); } catch (TaskDescriptionException e) { return error(INVALID_REQUEST, e); } IJobKey job = request.getTaskConfig().getJob(); return storage.read(storeProvider -> { if (storeProvider.getCronJobStore().fetchJob(job).isPresent()) { return invalidRequest(NO_CRON); } JobDiff diff = JobDiff.compute( storeProvider.getTaskStore(), job, JobDiff.asMap(request.getTaskConfig(), request.getInstanceCount()), request.getSettings().getUpdateOnlyTheseInstances()); Map<Integer, ITaskConfig> replaced = diff.getReplacedInstances(); Map<Integer, ITaskConfig> replacements = Maps.asMap( diff.getReplacementInstances(), Functions.constant(request.getTaskConfig())); Map<Integer, ITaskConfig> add = Maps.filterKeys( replacements, Predicates.in(Sets.difference(replacements.keySet(), replaced.keySet()))); Map<Integer, ITaskConfig> remove = Maps.filterKeys( replaced, Predicates.in(Sets.difference(replaced.keySet(), replacements.keySet()))); Map<Integer, ITaskConfig> update = Maps.filterKeys( replaced, Predicates.in(Sets.intersection(replaced.keySet(), replacements.keySet()))); return ok(Result.getJobUpdateDiffResult(new GetJobUpdateDiffResult() .setAdd(instancesToConfigGroups(add)) .setRemove(instancesToConfigGroups(remove)) .setUpdate(instancesToConfigGroups(update)) .setUnchanged(instancesToConfigGroups(diff.getUnchangedInstances())))); }); } @Inject  ReadOnlySchedulerImpl(
      ConfigurationManager configurationManager,
      Storage storage,
      NearestFit nearestFit,
      CronPredictor cronPredictor,
      QuotaManager quotaManager); @Override Response populateJobConfig(JobConfiguration description); @Override Response getTasksStatus(TaskQuery query); @Override Response getTasksWithoutConfigs(TaskQuery query); @Override Response getPendingReason(TaskQuery query); @Override Response getConfigSummary(JobKey job); @Override Response getRoleSummary(); @Override Response getJobSummary(@Nullable String maybeNullRole); @Override Response getJobs(@Nullable String maybeNullRole); @Override Response getQuota(String ownerRole); @Override Response getJobUpdateSummaries(JobUpdateQuery mutableQuery); @Override Response getJobUpdateDetails(JobUpdateKey mutableKey); @Override Response getJobUpdateDiff(JobUpdateRequest mutableRequest);  }
@Test public void testGetJobUpdateDiffWithUpdateAdd() throws Exception { TaskConfig task1 = defaultTask(false).setNumCpus(1.0); TaskConfig task2 = defaultTask(false).setNumCpus(2.0); TaskConfig task3 = defaultTask(false).setNumCpus(3.0); TaskConfig task4 = defaultTask(false).setNumCpus(4.0); TaskConfig task5 = defaultTask(false).setNumCpus(5.0); ImmutableSet.Builder<IScheduledTask> tasks = ImmutableSet.builder(); makeTasks(0, 10, task1, tasks); makeTasks(10, 20, task2, tasks); makeTasks(20, 30, task3, tasks); makeTasks(30, 40, task4, tasks); makeTasks(40, 50, task5, tasks); expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.absent()); storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), tasks.build()); control.replay(); TaskConfig newTask = defaultTask(false).setNumCpus(6.0); JobUpdateRequest request = new JobUpdateRequest() .setTaskConfig(newTask) .setInstanceCount(60) .setSettings(new JobUpdateSettings() .setUpdateOnlyTheseInstances(ImmutableSet.of(new Range(10, 59)))); GetJobUpdateDiffResult expected = new GetJobUpdateDiffResult() .setAdd(ImmutableSet.of(group(newTask, new Range(50, 59)))) .setUpdate(ImmutableSet.of( group(task2, new Range(10, 19)), group(task3, new Range(20, 29)), group(task4, new Range(30, 39)), group(task5, new Range(40, 49)))) .setUnchanged(ImmutableSet.of(group(task1, new Range(0, 9)))) .setRemove(ImmutableSet.of()); Response response = assertOkResponse(thrift.getJobUpdateDiff(request)); assertEquals(expected, response.getResult().getGetJobUpdateDiffResult()); }
@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 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()); }
@Test public void testGetJobUpdateDiffNoCron() throws Exception { expect(storageUtil.jobStore.fetchJob(JOB_KEY)) .andReturn(Optional.of(IJobConfiguration.build(CRON_JOB))); control.replay(); JobUpdateRequest request = new JobUpdateRequest().setTaskConfig(defaultTask(false)); Response expected = Responses.empty() .setResponseCode(INVALID_REQUEST) .setDetails(ImmutableList.of(new ResponseDetail(NO_CRON))); assertEquals(expected, thrift.getJobUpdateDiff(request)); }
@Test public void testGetJobUpdateDiffInvalidConfig() throws Exception { control.replay(); JobUpdateRequest request = new JobUpdateRequest().setTaskConfig(defaultTask(false).setNumCpus(-1)); assertResponse(INVALID_REQUEST, thrift.getJobUpdateDiff(request)); }