SchedulerThriftInterface implements AnnotatedAuroraAdmin { private Optional<String> rewriteInstance( IInstanceConfigRewrite instanceRewrite, MutableStoreProvider storeProvider) { IInstanceKey instanceKey = instanceRewrite.getInstanceKey(); Optional<String> error = Optional.absent(); Iterable<IScheduledTask> tasks = storeProvider.getTaskStore().fetchTasks( Query.instanceScoped(instanceKey.getJobKey(), instanceKey.getInstanceId()) .active()); Optional<IAssignedTask> task = Optional.fromNullable(Iterables.getOnlyElement(tasks, null)) .transform(IScheduledTask::getAssignedTask); if (task.isPresent()) { if (task.get().getTask().equals(instanceRewrite.getOldTask())) { ITaskConfig newConfiguration = instanceRewrite.getRewrittenTask(); boolean changed = storeProvider.getUnsafeTaskStore().unsafeModifyInPlace( task.get().getTaskId(), newConfiguration); if (!changed) { error = Optional.of("Did not change " + task.get().getTaskId()); } } else { error = Optional.of("CAS compare failed for " + instanceKey); } } else { error = Optional.of("No active task found for " + instanceKey); } return error; } @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 testRewriteInstance() throws Exception { TaskConfig storedConfig = productionTask(); ITaskConfig modifiedConfig = ITaskConfig.build( storedConfig.deepCopy().setExecutorConfig(new ExecutorConfig("aurora", "rewritten"))); String taskId = "task_id"; IScheduledTask storedTask = IScheduledTask.build(new ScheduledTask().setAssignedTask( new AssignedTask() .setTaskId(taskId) .setTask(storedConfig))); InstanceKey instanceKey = new InstanceKey(storedConfig.getJob(), 0); storageUtil.expectTaskFetch( Query.instanceScoped(IInstanceKey.build(instanceKey)).active(), storedTask); expect(storageUtil.taskStore.unsafeModifyInPlace(taskId, modifiedConfig)) .andReturn(true); control.replay(); RewriteConfigsRequest request = new RewriteConfigsRequest( ImmutableList.of(ConfigRewrite.instanceRewrite( new InstanceConfigRewrite(instanceKey, storedConfig, modifiedConfig.newBuilder())))); assertOkResponse(thrift.rewriteConfigs(request)); }