JobQuotaController implements QuotaController<Job<?>> { @Override public JobQuotaController update(Job<?> updatedJob) { int currentDesired = JobFunctions.getJobDesiredSize(job); int newDesired = JobFunctions.getJobDesiredSize(updatedJob); if (currentDesired == newDesired) { if (job.getJobDescriptor().getDisruptionBudget().equals(updatedJob.getJobDescriptor().getDisruptionBudget())) { return this; } } return new JobQuotaController( updatedJob, jobOperations, effectiveDisruptionBudgetResolver, containerHealthService, this, titusRuntime ); } JobQuotaController(Job<?> job,
                              V3JobOperations jobOperations,
                              EffectiveJobDisruptionBudgetResolver effectiveDisruptionBudgetResolver,
                              ContainerHealthService containerHealthService,
                              TitusRuntime titusRuntime); private  JobQuotaController(Job<?> newJob,
                               V3JobOperations jobOperations,
                               EffectiveJobDisruptionBudgetResolver effectiveDisruptionBudgetResolver,
                               ContainerHealthService containerHealthService,
                               JobQuotaController previousJobQuotaController,
                               TitusRuntime titusRuntime); Job<?> getJob(); @Override EvictionQuota getQuota(Reference reference); @Override ConsumptionResult consume(String taskId); @Override void giveBackConsumedQuota(String taskId); @Override JobQuotaController update(Job<?> updatedJob);  }
@Test public void testUpdate() { Job<BatchJobExt> job = newBatchJob(10, budget(perTaskRelocationLimitPolicy(2), hourlyRatePercentage(20), Collections.emptyList())); com.netflix.titus.api.model.reference.Reference jobReference = com.netflix.titus.api.model.reference.Reference.job(job.getId()); scheduleJob(job, 10); JobQuotaController jobController = new JobQuotaController(job, jobOperations, SelfJobDisruptionBudgetResolver.getInstance(), containerHealthService, titusRuntime); assertThat(jobController.getQuota(jobReference).getQuota()).isEqualTo(2); Task task = jobOperations.getTasks(job.getId()).get(0); assertThat(jobController.consume(task.getId()).isApproved()).isTrue(); jobComponentStub.killTask(task, false, false, V3JobOperations.Trigger.Eviction); Task replacement1 = jobComponentStub.getJobOperations().getTasks().stream().filter(t -> t.getOriginalId().equals(task.getId())).findFirst().get(); assertThat(jobController.consume(replacement1.getId()).isApproved()).isTrue(); jobComponentStub.moveTaskToState(replacement1, TaskState.Started); jobComponentStub.killTask(replacement1, false, false, V3JobOperations.Trigger.Eviction); Task replacement2 = jobComponentStub.getJobOperations().getTasks().stream().filter(t -> t.getOriginalId().equals(task.getId())).findFirst().get(); assertThat(jobController.consume(replacement2.getId()).isApproved()).isFalse(); assertThat(jobController.getQuota(jobReference).getQuota()).isEqualTo(0); Job<BatchJobExt> updatedJob = jobComponentStub.changeJob( exceptBudget(job, budget(perTaskRelocationLimitPolicy(3), hourlyRatePercentage(80), Collections.emptyList())) ); JobQuotaController updatedController = jobController.update(updatedJob); assertThat(updatedController.consume(replacement2.getId()).isApproved()).isTrue(); jobComponentStub.moveTaskToState(replacement2, TaskState.Started); jobComponentStub.killTask(replacement2, false, false, V3JobOperations.Trigger.Eviction); Task replacement3 = jobComponentStub.getJobOperations().getTasks().stream().filter(t -> t.getOriginalId().equals(task.getId())).findFirst().get(); assertThat(updatedController.consume(replacement3.getId()).isApproved()).isFalse(); assertThat(updatedController.getQuota(jobReference).getQuota()).isEqualTo(5); Job<BatchJobExt> scaledJob = jobComponentStub.changeJob(JobFunctions.changeBatchJobSize(updatedJob, 20)); jobComponentStub.createDesiredTasks(scaledJob); JobQuotaController updatedController2 = jobController.update(scaledJob); assertThat(updatedController2.getQuota(jobReference).getQuota()).isEqualTo(13); }