TaskMigrationDescheduler { Map<String, DeschedulingResult> findRequestedJobOrTaskMigrations() { long now = clock.wallTime(); Map<String, DeschedulingResult> result = new HashMap<>(); tasksById.values().forEach(task -> { Job<?> job = jobsById.get(task.getJobId()); Node instance = evacuatedAgentsAllocationTracker.getAgent(task); if (job != null && instance != null) { RelocationPredicates.checkIfRelocationRequired(job, task).ifPresent(reason -> { if (evictionQuotaTracker.getSystemEvictionQuota() > 0 && canTerminate(task)) { long quota = evictionQuotaTracker.getJobEvictionQuota(task.getJobId()); if (quota > 0) { evictionQuotaTracker.consumeQuota(task.getJobId()); result.put(task.getId(), newDeschedulingResultForRequestedRelocation(now, task, instance, reason.getRight())); } } }); } }); return result; } TaskMigrationDescheduler(Map<String, TaskRelocationPlan> plannedAheadTaskRelocationPlans,
                             EvacuatedAgentsAllocationTracker evacuatedAgentsAllocationTracker,
                             EvictionQuotaTracker evictionQuotaTracker,
                             Map<String, Job<?>> jobsById,
                             Map<String, Task> tasksById,
                             TitusRuntime titusRuntime);   }
@Test public void testJobRequiredMigrationBy() { Task job1Task0 = jobOperations.getTasks("job1").get(0); dataGenerator.place("active1", job1Task0); dataGenerator.setQuota("job1", 1); dataGenerator.addJobAttribute("job1", RelocationAttributes.RELOCATION_REQUIRED_BY, "" + clock.wallTime()); clock.advanceTime(Duration.ofSeconds(1)); Map<String, DeschedulingResult> results = newDescheduler(Collections.emptyMap()).findRequestedJobOrTaskMigrations(); assertThat(results).isNotEmpty(); }
@Test public void testTaskRequiredMigration() { Task job1Task0 = jobOperations.getTasks("job1").get(0); dataGenerator.place("active1", job1Task0); dataGenerator.setQuota("job1", 1); dataGenerator.addTaskAttribute(job1Task0.getId(), RelocationAttributes.RELOCATION_REQUIRED, "true"); clock.advanceTime(Duration.ofSeconds(1)); Map<String, DeschedulingResult> results = newDescheduler(Collections.emptyMap()).findRequestedJobOrTaskMigrations(); assertThat(results).isNotEmpty(); }