TaskRelocationDataInjector { Observable<TaskQueryResult> injectIntoTaskQueryResult(Observable<TaskQueryResult> tasksObservable) { if (!featureActivationConfiguration.isMergingTaskMigrationPlanInGatewayEnabled()) { return tasksObservable; } return tasksObservable.flatMap(queryResult -> { Set<String> taskIds = queryResult.getItemsList().stream().map(Task::getId).collect(Collectors.toSet()); if (shouldUseRelocationCache()) { Map<String, TaskRelocationPlan> plans = relocationDataReplicator.getCurrent().getPlans(); List<Task> newTaskList = queryResult.getItemsList().stream() .map(task -> { TaskRelocationPlan plan = plans.get(task.getId()); return plan != null ? newTaskWithRelocationPlan(task, plan) : task; }) .collect(Collectors.toList()); return Observable.just(queryResult.toBuilder().clearItems().addAllItems(newTaskList).build()); } return ReactorExt.toObservable(relocationServiceClient.findTaskRelocationPlans(taskIds)) .timeout(getTaskRelocationTimeout(), TimeUnit.MILLISECONDS, scheduler) .doOnError(error -> logger.info("Could not resolve task relocation status for tasks: taskIds={}, error={}", taskIds, ExceptionExt.toMessageChain(error))) .onErrorReturn(e -> Collections.emptyList()) .map(relocationPlans -> { Map<String, TaskRelocationPlan> plansById = relocationPlans.stream().collect(Collectors.toMap(TaskRelocationPlan::getTaskId, p -> p)); if (plansById.isEmpty()) { return queryResult; } List<Task> newTaskList = queryResult.getItemsList().stream() .map(task -> { TaskRelocationPlan plan = plansById.get(task.getId()); return plan != null ? newTaskWithRelocationPlan(task, plan) : task; }) .collect(Collectors.toList()); return queryResult.toBuilder().clearItems().addAllItems(newTaskList).build(); }); }); } @Inject  TaskRelocationDataInjector(
            GrpcClientConfiguration configuration,
            JobManagerConfiguration jobManagerConfiguration,
            FeatureActivationConfiguration featureActivationConfiguration,
            RelocationServiceClient relocationServiceClient,
            RelocationDataReplicator relocationDataReplicator); @VisibleForTesting  TaskRelocationDataInjector(
            GrpcClientConfiguration configuration,
            JobManagerConfiguration jobManagerConfiguration,
            FeatureActivationConfiguration featureActivationConfiguration,
            RelocationServiceClient relocationServiceClient,
            RelocationDataReplicator relocationDataReplicator,
            Scheduler scheduler);   }
@Test public void testFindTasksWithRelocationDeadline() { long deadline1 = titusRuntime.getClock().wallTime() + 1_000; long deadline2 = titusRuntime.getClock().wallTime() + 2_000; TaskQueryResult queryResult = TaskQueryResult.newBuilder() .addItems(TASK1) .addItems(TASK2) .build(); List<TaskRelocationPlan> relocationPlans = Arrays.asList( newRelocationPlan(TASK1, deadline1), newRelocationPlan(TASK2, deadline2) ); when(relocationServiceClient.findTaskRelocationPlans(asSet(TASK1.getId(), TASK2.getId()))).thenReturn(Mono.just(relocationPlans)); TaskQueryResult merged = taskRelocationDataInjector.injectIntoTaskQueryResult(Observable.just(queryResult)).toBlocking().first(); assertThat(merged.getItems(0).getMigrationDetails().getNeedsMigration()).isTrue(); assertThat(merged.getItems(0).getMigrationDetails().getDeadline()).isEqualTo(deadline1); assertThat(merged.getItems(1).getMigrationDetails().getNeedsMigration()).isTrue(); assertThat(merged.getItems(1).getMigrationDetails().getDeadline()).isEqualTo(deadline2); }
@Test public void testFindTasksWithoutRelocationDeadline() { long deadline1 = titusRuntime.getClock().wallTime() + 1_000; TaskQueryResult queryResult = TaskQueryResult.newBuilder() .addItems(TASK1) .addItems(TASK2) .build(); List<TaskRelocationPlan> relocationPlans = Collections.singletonList(newRelocationPlan(TASK1, deadline1)); when(relocationServiceClient.findTaskRelocationPlans(asSet(TASK1.getId(), TASK2.getId()))).thenReturn(Mono.just(relocationPlans)); TaskQueryResult merged = taskRelocationDataInjector.injectIntoTaskQueryResult(Observable.just(queryResult)).toBlocking().first(); assertThat(merged.getItems(0).getMigrationDetails().getNeedsMigration()).isTrue(); assertThat(merged.getItems(0).getMigrationDetails().getDeadline()).isEqualTo(deadline1); assertThat(merged.getItems(1).getMigrationDetails().getNeedsMigration()).isFalse(); }
@Test public void testFindTasksWithLegacyMigration() { long deadline1 = titusRuntime.getClock().wallTime() + 1_000; long legacyDeadline2 = titusRuntime.getClock().wallTime() + 2_000; Task legacyTask = toLegacyTask(TASK2, legacyDeadline2); TaskQueryResult queryResult = TaskQueryResult.newBuilder() .addItems(TASK1) .addItems(legacyTask) .build(); List<TaskRelocationPlan> relocationPlans = Collections.singletonList(newRelocationPlan(TASK1, deadline1)); when(relocationServiceClient.findTaskRelocationPlans(asSet(TASK1.getId(), TASK2.getId()))).thenReturn(Mono.just(relocationPlans)); TaskQueryResult merged = taskRelocationDataInjector.injectIntoTaskQueryResult(Observable.just(queryResult)).toBlocking().first(); assertThat(merged.getItems(0).getMigrationDetails().getNeedsMigration()).isTrue(); assertThat(merged.getItems(0).getMigrationDetails().getDeadline()).isEqualTo(deadline1); assertThat(merged.getItems(1)).isEqualTo(legacyTask); }
@Test public void testFindTasksWithRelocationDataFetchTimeout() { when(relocationServiceClient.findTaskRelocationPlans(asSet(TASK1.getId(), TASK2.getId()))).thenReturn(Mono.never()); TaskQueryResult queryResult = TaskQueryResult.newBuilder() .addItems(TASK1) .addItems(TASK2) .build(); ExtTestSubscriber<TaskQueryResult> testSubscriber = new ExtTestSubscriber<>(); taskRelocationDataInjector.injectIntoTaskQueryResult(Observable.just(queryResult)).subscribe(testSubscriber); testSubscriber.assertOpen(); testScheduler.advanceTimeBy(RELOCATION_TIMEOUT_MS, TimeUnit.MILLISECONDS); testSubscriber.assertOnCompleted(); TaskQueryResult result = testSubscriber.takeNext(); assertThat(result.getItemsList()).contains(TASK1, TASK2); }