DefaultDeschedulerService implements DeschedulerService { @Override public List<DeschedulingResult> deschedule(Map<String, TaskRelocationPlan> plannedAheadTaskRelocationPlans) { List<Pair<Job, List<Task>>> allJobsAndTasks = jobOperations.getJobsAndTasks(); Map<String, Job<?>> jobs = allJobsAndTasks.stream().map(Pair::getLeft).collect(Collectors.toMap(Job::getId, j -> j)); Map<String, Task> tasksById = allJobsAndTasks.stream() .flatMap(p -> p.getRight().stream()) .collect(Collectors.toMap(Task::getId, t -> t)); EvacuatedAgentsAllocationTracker evacuatedAgentsAllocationTracker = new EvacuatedAgentsAllocationTracker(nodeDataResolver.resolve(), tasksById); EvictionQuotaTracker evictionQuotaTracker = new EvictionQuotaTracker(evictionOperations, jobs); TaskMigrationDescheduler taskMigrationDescheduler = new TaskMigrationDescheduler( plannedAheadTaskRelocationPlans, evacuatedAgentsAllocationTracker, evictionQuotaTracker, jobs, tasksById, titusRuntime ); Map<String, DeschedulingResult> requestedImmediateEvictions = taskMigrationDescheduler.findAllImmediateEvictions(); Map<String, DeschedulingResult> requestedEvictions = taskMigrationDescheduler.findRequestedJobOrTaskMigrations(); Map<String, DeschedulingResult> allRequestedEvictions = CollectionsExt.merge(requestedImmediateEvictions, requestedEvictions); Map<String, DeschedulingResult> regularEvictions = new HashMap<>(); Optional<Pair<Node, List<Task>>> bestMatch; while ((bestMatch = taskMigrationDescheduler.nextBestMatch()).isPresent()) { Node agent = bestMatch.get().getLeft(); List<Task> tasks = bestMatch.get().getRight(); tasks.forEach(task -> { if (!allRequestedEvictions.containsKey(task.getId())) { TaskRelocationPlan relocationPlan = plannedAheadTaskRelocationPlans.get(task.getId()); if (relocationPlan == null) { relocationPlan = newNotDelayedRelocationPlan(task, true); } regularEvictions.put( task.getId(), DeschedulingResult.newBuilder() .withTask(task) .withAgentInstance(agent) .withTaskRelocationPlan(relocationPlan) .build() ); } }); } for (Task task : tasksById.values()) { if (allRequestedEvictions.containsKey(task.getId()) || regularEvictions.containsKey(task.getId())) { continue; } if (evacuatedAgentsAllocationTracker.isEvacuated(task)) { DeschedulingFailure failure = taskMigrationDescheduler.getDeschedulingFailure(task); TaskRelocationPlan relocationPlan = plannedAheadTaskRelocationPlans.get(task.getId()); if (relocationPlan == null) { relocationPlan = newNotDelayedRelocationPlan(task, false); } Node agent = evacuatedAgentsAllocationTracker.getRemovableAgent(task); regularEvictions.put( task.getId(), DeschedulingResult.newBuilder() .withTask(task) .withAgentInstance(agent) .withTaskRelocationPlan(relocationPlan) .withFailure(failure) .build() ); } } return CollectionsExt.merge(new ArrayList<>(allRequestedEvictions.values()), new ArrayList<>(regularEvictions.values())); } @Inject  DefaultDeschedulerService(ReadOnlyJobOperations jobOperations,
                                     ReadOnlyEvictionOperations evictionOperations,
                                     NodeDataResolver nodeDataResolver,
                                     TitusRuntime titusRuntime); @Override List<DeschedulingResult> deschedule(Map<String, TaskRelocationPlan> plannedAheadTaskRelocationPlans);  }
@Test public void testAllExpectedJobMigrationsAreFound() { List<Task> tasksOfJob1 = jobOperations.getTasks("job1"); dataGenerator.place("active1", tasksOfJob1.get(0), tasksOfJob1.get(1)); dataGenerator.place("removable1", tasksOfJob1.get(2), tasksOfJob1.get(3)); dataGenerator.setQuota("job1", 2); List<Task> tasksOfJob2 = jobOperations.getTasks("job2"); dataGenerator.place("active1", tasksOfJob2.get(0), tasksOfJob2.get(1)); dataGenerator.place("removable1", tasksOfJob2.get(2), tasksOfJob2.get(3)); dataGenerator.setQuota("job2", 2); dataGenerator.addJobAttribute("jobImmediate", RelocationAttributes.RELOCATION_REQUIRED_BY_IMMEDIATELY, "" + (clock.wallTime() + 1)); Task taskImmediate = jobOperations.getTasks("jobImmediate").get(0); dataGenerator.place("active1", taskImmediate); List<DeschedulingResult> results = deschedulerService.deschedule(Collections.emptyMap()); assertThat(results).hasSize(5); for (DeschedulingResult result : results) { boolean isImmediateJobMigration = result.getTask().getId().equals(taskImmediate.getId()); if (isImmediateJobMigration) { assertThat(result.getAgentInstance().getServerGroupId()).isEqualTo("active1"); } else { assertThat(result.getAgentInstance().getServerGroupId()).isEqualTo("removable1"); } TaskRelocationPlan plan = result.getTaskRelocationPlan(); assertThat(plan.getReason()).isEqualTo(TaskRelocationReason.TaskMigration); if (isImmediateJobMigration) { assertThat(plan.getReasonMessage()).containsSequence("Job marked for immediate eviction"); } else { assertThat(plan.getReasonMessage()).containsSequence("Enough quota to migrate the task"); } } }