SchedulerThriftInterface implements AnnotatedAuroraAdmin { @Override public Response startJobUpdate(JobUpdateRequest mutableRequest, @Nullable String message) { requireNonNull(mutableRequest); if (!mutableRequest.getTaskConfig().isIsService()) { return invalidRequest(NON_SERVICE_TASK); } JobUpdateSettings settings = requireNonNull(mutableRequest.getSettings()); if (settings.getUpdateGroupSize() <= 0) { return invalidRequest(INVALID_GROUP_SIZE); } if (settings.getMaxPerInstanceFailures() < 0) { return invalidRequest(INVALID_MAX_INSTANCE_FAILURES); } if (settings.getMaxFailedInstances() < 0) { return invalidRequest(INVALID_MAX_FAILED_INSTANCES); } if (settings.getMaxPerInstanceFailures() * mutableRequest.getInstanceCount() > thresholds.getMaxUpdateInstanceFailures()) { return invalidRequest(TOO_MANY_POTENTIAL_FAILED_INSTANCES); } if (settings.getMinWaitInInstanceRunningMs() < 0) { return invalidRequest(INVALID_MIN_WAIT_TO_RUNNING); } if (settings.getBlockIfNoPulsesAfterMs() < 0) { return invalidRequest(INVALID_PULSE_TIMEOUT); } 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); } return storage.write(storeProvider -> { IJobKey job = request.getTaskConfig().getJob(); if (getCronJob(storeProvider, job).isPresent()) { return invalidRequest(NO_CRON); } String updateId = uuidGenerator.createNew().toString(); IJobUpdateSettings settings1 = request.getSettings(); JobDiff diff = JobDiff.compute( storeProvider.getTaskStore(), job, JobDiff.asMap(request.getTaskConfig(), request.getInstanceCount()), settings1.getUpdateOnlyTheseInstances()); Set<Integer> invalidScope = diff.getOutOfScopeInstances( Numbers.rangesToInstanceIds(settings1.getUpdateOnlyTheseInstances())); if (!invalidScope.isEmpty()) { return invalidRequest( "The update request attempted to update specific instances," + " but some are irrelevant to the update and current job state: " + invalidScope); } if (diff.isNoop()) { return addMessage(empty(), OK, NOOP_JOB_UPDATE_MESSAGE); } JobUpdateInstructions instructions = new JobUpdateInstructions() .setSettings(settings1.newBuilder()) .setInitialState(buildInitialState(diff.getReplacedInstances())); Set<Integer> replacements = diff.getReplacementInstances(); if (!replacements.isEmpty()) { instructions.setDesiredState( new InstanceTaskConfig() .setTask(request.getTaskConfig().newBuilder()) .setInstances(IRange.toBuildersSet(convertRanges(toRanges(replacements))))); } String remoteUserName = auditMessages.getRemoteUserName(); IJobUpdate update = IJobUpdate.build(new JobUpdate() .setSummary(new JobUpdateSummary() .setKey(new JobUpdateKey(job.newBuilder(), updateId)) .setUser(remoteUserName)) .setInstructions(instructions)); Response response = empty(); try { validateTaskLimits( request.getTaskConfig(), request.getInstanceCount(), quotaManager.checkJobUpdate(update, storeProvider)); jobUpdateController.start( update, new AuditData(remoteUserName, Optional.fromNullable(message))); return response.setResponseCode(OK) .setResult(Result.startJobUpdateResult( new StartJobUpdateResult(update.getSummary().getKey().newBuilder()))); } catch (UpdateStateException | TaskValidationException e) { return error(INVALID_REQUEST, e); } }); } @Inject  SchedulerThriftInterface(
      ConfigurationManager configurationManager,
      Thresholds thresholds,
      NonVolatileStorage storage,
      LockManager lockManager,
      StorageBackup backup,
      Recovery recovery,
      CronJobManager cronJobManager,
      MaintenanceController maintenance,
      QuotaManager quotaManager,
      StateManager stateManager,
      TaskIdGenerator taskIdGenerator,
      UUIDGenerator uuidGenerator,
      JobUpdateController jobUpdateController,
      ReadOnlyScheduler.Iface readOnlyScheduler,
      AuditMessages auditMessages); @Override Response createJob(JobConfiguration mutableJob); @Override Response scheduleCronJob(JobConfiguration mutableJob); @Override Response replaceCronTemplate(JobConfiguration mutableJob); @Override Response descheduleCronJob(JobKey mutableJobKey); @Override Response populateJobConfig(JobConfiguration description); @Override Response startCronJob(JobKey mutableJobKey); @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 getJobUpdateDiff(JobUpdateRequest request); @Override Response killTasks(@Nullable JobKey mutableJob, @Nullable Set<Integer> instances); @Override Response restartShards(JobKey mutableJobKey, Set<Integer> shardIds); @Override Response getQuota(String ownerRole); @Override Response setQuota(String ownerRole, ResourceAggregate resourceAggregate); @Override Response startMaintenance(Hosts hosts); @Override Response drainHosts(Hosts hosts); @Override Response maintenanceStatus(Hosts hosts); @Override Response endMaintenance(Hosts hosts); @Override Response forceTaskState(String taskId, ScheduleStatus status); @Override Response performBackup(); @Override Response listBackups(); @Override Response stageRecovery(String backupId); @Override Response queryRecovery(TaskQuery query); @Override Response deleteRecoveryTasks(TaskQuery query); @Override Response commitRecovery(); @Override Response unloadRecovery(); @Override Response snapshot(); @Override Response rewriteConfigs(RewriteConfigsRequest request); @Override Response addInstances(InstanceKey key, int count); Optional<IJobConfiguration> getCronJob(StoreProvider storeProvider, IJobKey jobKey); @Override Response startJobUpdate(JobUpdateRequest mutableRequest, @Nullable String message); @Override Response pauseJobUpdate(JobUpdateKey mutableKey, @Nullable String message); @Override Response resumeJobUpdate(JobUpdateKey mutableKey, @Nullable String message); @Override Response abortJobUpdate(JobUpdateKey mutableKey, @Nullable String message); @Override Response pulseJobUpdate(JobUpdateKey mutableUpdateKey); @Override Response getJobUpdateSummaries(JobUpdateQuery mutableQuery); @Override Response getJobUpdateDetails(JobUpdateKey key);  }
@Test public void testStartUpdate() throws Exception { expectGetRemoteUser(); expectNoCronJob(); ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask(); expect(taskIdGenerator.generate(newTask, 6)).andReturn(TASK_ID); IScheduledTask oldTask1 = buildTaskForJobUpdate(0, "old"); IScheduledTask oldTask2 = buildTaskForJobUpdate(1, "old"); IScheduledTask oldTask3 = buildTaskForJobUpdate(2, "old2"); IScheduledTask oldTask4 = buildTaskForJobUpdate(3, "old2"); IScheduledTask oldTask5 = buildTaskForJobUpdate(4, "old"); IScheduledTask oldTask6 = buildTaskForJobUpdate(5, "old"); IScheduledTask oldTask7 = buildTaskForJobUpdate(6, "old"); IJobUpdate update = buildJobUpdate(6, newTask, ImmutableMap.of( oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 1), new Range(4, 6)), oldTask3.getAssignedTask().getTask(), ImmutableSet.of(new Range(2, 3)) )); expect(quotaManager.checkJobUpdate( update, storageUtil.mutableStoreProvider)).andReturn(ENOUGH_QUOTA); expect(uuidGenerator.createNew()).andReturn(UU_ID); storageUtil.expectTaskFetch( Query.unscoped().byJob(JOB_KEY).active(), oldTask1, oldTask2, oldTask3, oldTask4, oldTask5, oldTask6, oldTask7); jobUpdateController.start(update, AUDIT); control.replay(); Response response = assertOkResponse(thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE)); assertEquals( new StartJobUpdateResult(UPDATE_KEY.newBuilder()), response.getResult().getStartJobUpdateResult()); }
@Test public void testStartUpdateEmptyDesired() throws Exception { expectGetRemoteUser(); expectNoCronJob(); ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask(); expect(taskIdGenerator.generate(newTask, 1)).andReturn(TASK_ID); IScheduledTask oldTask1 = buildTaskForJobUpdate(0); IScheduledTask oldTask2 = buildTaskForJobUpdate(1); IJobUpdate update = buildJobUpdate( 1, newTask, ImmutableMap.of(oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 1)))); expectJobUpdateQuotaCheck(ENOUGH_QUOTA); JobUpdate expected = update.newBuilder(); expected.getInstructions().setInitialState(ImmutableSet.of( new InstanceTaskConfig(newTask.newBuilder(), ImmutableSet.of(new Range(1, 1))))); expected.getInstructions().unsetDesiredState(); expect(uuidGenerator.createNew()).andReturn(UU_ID); storageUtil.expectTaskFetch( Query.unscoped().byJob(JOB_KEY).active(), oldTask1, oldTask2); jobUpdateController.start(IJobUpdate.build(expected), AUDIT); control.replay(); Response response = assertOkResponse( thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE)); assertEquals( new StartJobUpdateResult(UPDATE_KEY.newBuilder()), response.getResult().getStartJobUpdateResult()); }
@Test(expected = NullPointerException.class) public void testStartUpdateFailsNullRequest() throws Exception { control.replay(); thrift.startJobUpdate(null, AUDIT_MESSAGE); }
@Test(expected = NullPointerException.class) public void testStartUpdateFailsNullTaskConfig() throws Exception { control.replay(); thrift.startJobUpdate( new JobUpdateRequest(null, 5, buildJobUpdateSettings()), AUDIT_MESSAGE); }
@Test public void testStartUpdateFailsInvalidGroupSize() throws Exception { control.replay(); JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); updateRequest.getSettings().setUpdateGroupSize(0); assertEquals( invalidResponse(SchedulerThriftInterface.INVALID_GROUP_SIZE), thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsInvalidMaxInstanceFailures() throws Exception { control.replay(); JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); updateRequest.getSettings().setMaxPerInstanceFailures(-1); assertEquals( invalidResponse(SchedulerThriftInterface.INVALID_MAX_INSTANCE_FAILURES), thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsTooManyPerInstanceFailures() throws Exception { control.replay(); JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); updateRequest.getSettings() .setMaxPerInstanceFailures(THRESHOLDS.getMaxUpdateInstanceFailures() + 10); assertEquals( invalidResponse(SchedulerThriftInterface.TOO_MANY_POTENTIAL_FAILED_INSTANCES), thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsInvalidMaxFailedInstances() throws Exception { control.replay(); JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); updateRequest.getSettings().setMaxFailedInstances(-1); assertEquals( invalidResponse(SchedulerThriftInterface.INVALID_MAX_FAILED_INSTANCES), thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsInvalidMinWaitInRunning() throws Exception { control.replay(); JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); updateRequest.getSettings().setMinWaitInInstanceRunningMs(-1); assertEquals( invalidResponse(SchedulerThriftInterface.INVALID_MIN_WAIT_TO_RUNNING), thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsNonServiceTask() throws Exception { control.replay(); JobUpdateRequest request = buildJobUpdateRequest(populatedTask().setIsService(false)); assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsInvalidPulseTimeout() throws Exception { control.replay(); JobUpdateRequest updateRequest = buildServiceJobUpdateRequest(); updateRequest.getSettings().setBlockIfNoPulsesAfterMs(-1); assertEquals( invalidResponse(SchedulerThriftInterface.INVALID_PULSE_TIMEOUT), thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsForCronJob() throws Exception { JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask()); expectCronJob(); control.replay(); assertEquals(invalidResponse(NO_CRON), thrift.startJobUpdate(request, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsConfigValidation() throws Exception { JobUpdateRequest request = buildJobUpdateRequest(INVALID_TASK_CONFIG.setIsService(true)); control.replay(); assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); }
@Test public void testStartNoopUpdate() throws Exception { expectNoCronJob(); expect(uuidGenerator.createNew()).andReturn(UU_ID); ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask(); IScheduledTask oldTask = buildTaskForJobUpdate(0); storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); IJobUpdate update = buildJobUpdate( 1, newTask, ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0)))); control.replay(); JobUpdateRequest request = buildJobUpdateRequest(update); Response response = thrift.startJobUpdate(request, AUDIT_MESSAGE); assertResponse(OK, response); assertEquals( NOOP_JOB_UPDATE_MESSAGE, Iterables.getOnlyElement(response.getDetails()).getMessage()); }
@Test public void testStartUpdateInvalidScope() throws Exception { expectNoCronJob(); expect(uuidGenerator.createNew()).andReturn(UU_ID); IScheduledTask oldTask = buildTaskForJobUpdate(0); storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask(); JobUpdate builder = buildJobUpdate( 1, newTask, ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0)))) .newBuilder(); builder.getInstructions().getSettings() .setUpdateOnlyTheseInstances(ImmutableSet.of(new Range(100, 100))); control.replay(); JobUpdateRequest request = buildJobUpdateRequest(IJobUpdate.build(builder)); assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); }
@Test public void testStartScopeIncludesNoop() throws Exception { expectGetRemoteUser(); expectNoCronJob(); expect(uuidGenerator.createNew()).andReturn(UU_ID); ScheduledTask taskBuilder = buildTaskForJobUpdate(0).newBuilder(); taskBuilder.getAssignedTask().getTask().setNumCpus(100); IScheduledTask newTask = IScheduledTask.build(taskBuilder); IScheduledTask oldTask1 = buildTaskForJobUpdate(1); IScheduledTask oldTask2 = buildTaskForJobUpdate(2); storageUtil.expectTaskFetch( Query.unscoped().byJob(JOB_KEY).active(), newTask, oldTask1, oldTask2); expectJobUpdateQuotaCheck(ENOUGH_QUOTA); expect(taskIdGenerator.generate(EasyMock.<ITaskConfig>anyObject(), anyInt())) .andReturn(TASK_ID); jobUpdateController.start(EasyMock.<IJobUpdate>anyObject(), eq(AUDIT)); ITaskConfig newConfig = newTask.getAssignedTask().getTask(); JobUpdate builder = buildJobUpdate( 3, newConfig, ImmutableMap.of( newConfig, ImmutableSet.of(new Range(0, 0)), oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(1, 2)))) .newBuilder(); builder.getInstructions().getSettings() .setUpdateOnlyTheseInstances(ImmutableSet.of(new Range(0, 2))); control.replay(); JobUpdateRequest request = buildJobUpdateRequest(IJobUpdate.build(builder)); assertResponse(OK, thrift.startJobUpdate(request, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsInstanceCountCheck() throws Exception { JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask()); request.setInstanceCount(THRESHOLDS.getMaxTasksPerJob() + 1); expectGetRemoteUser(); expectNoCronJob(); expect(uuidGenerator.createNew()).andReturn(UU_ID); storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active()); expectJobUpdateQuotaCheck(ENOUGH_QUOTA); control.replay(); assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsTaskIdLength() throws Exception { JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask()); expectGetRemoteUser(); expectNoCronJob(); expect(uuidGenerator.createNew()).andReturn(UU_ID); storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active()); expectJobUpdateQuotaCheck(ENOUGH_QUOTA); expect(taskIdGenerator.generate(ITaskConfig.build(request.getTaskConfig()), 6)) .andReturn(Strings.repeat("a", MAX_TASK_ID_LENGTH + 1)); control.replay(); assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsQuotaCheck() throws Exception { JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask()); expectGetRemoteUser(); expectNoCronJob(); expect(uuidGenerator.createNew()).andReturn(UU_ID); IScheduledTask oldTask = buildTaskForJobUpdate(0); storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); expect(taskIdGenerator.generate(ITaskConfig.build(request.getTaskConfig()), 6)) .andReturn(TASK_ID); expectJobUpdateQuotaCheck(NOT_ENOUGH_QUOTA); control.replay(); assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE)); }
@Test public void testStartUpdateFailsInController() throws Exception { expectGetRemoteUser(); expectNoCronJob(); IScheduledTask oldTask = buildTaskForJobUpdate(0, "old"); ITaskConfig newTask = buildTaskForJobUpdate(0, "new").getAssignedTask().getTask(); IJobUpdate update = buildJobUpdate( 1, newTask, ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0)))); expect(uuidGenerator.createNew()).andReturn(UU_ID); expect(taskIdGenerator.generate(newTask, 1)).andReturn(TASK_ID); expectJobUpdateQuotaCheck(ENOUGH_QUOTA); storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask); jobUpdateController.start(update, AUDIT); expectLastCall().andThrow(new UpdateStateException("failed")); control.replay(); assertResponse( INVALID_REQUEST, thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE)); }