@Test public void testFindTaskWithRelocationDeadline() { long deadlineTimestamp = titusRuntime.getClock().wallTime() + 1_000; when(relocationServiceClient.findTaskRelocationPlan(TASK1.getId())).thenReturn( Mono.just(Optional.of(newRelocationPlan(TASK1, deadlineTimestamp))) ); Task merged = taskRelocationDataInjector.injectIntoTask(TASK1.getId(), Observable.just(TASK1)).toBlocking().first(); assertThat(merged.getMigrationDetails().getNeedsMigration()).isTrue(); assertThat(merged.getMigrationDetails().getDeadline()).isEqualTo(deadlineTimestamp); }
Observable<Task> injectIntoTask(String taskId, Observable<Task> taskObservable) { if (!featureActivationConfiguration.isMergingTaskMigrationPlanInGatewayEnabled()) { return taskObservable; } if (shouldUseRelocationCache()) { return taskObservable.map(task -> newTaskWithRelocationPlan(task, relocationDataReplicator.getCurrent().getPlans().get(taskId))); } Observable<Optional<TaskRelocationPlan>> relocationPlanResolver = ReactorExt.toObservable(relocationServiceClient.findTaskRelocationPlan(taskId)) .timeout(getTaskRelocationTimeout(), TimeUnit.MILLISECONDS, scheduler) .doOnError(error -> logger.info("Could not resolve task relocation status for task: taskId={}, error={}", taskId, ExceptionExt.toMessageChain(error))) .onErrorReturn(e -> Optional.empty()); return Observable.zip( taskObservable, relocationPlanResolver, (task, planOpt) -> planOpt.map(plan -> newTaskWithRelocationPlan(task, plan)).orElse(task) ); }
TaskRelocationDataInjector { Observable<Task> injectIntoTask(String taskId, Observable<Task> taskObservable) { if (!featureActivationConfiguration.isMergingTaskMigrationPlanInGatewayEnabled()) { return taskObservable; } if (shouldUseRelocationCache()) { return taskObservable.map(task -> newTaskWithRelocationPlan(task, relocationDataReplicator.getCurrent().getPlans().get(taskId))); } Observable<Optional<TaskRelocationPlan>> relocationPlanResolver = ReactorExt.toObservable(relocationServiceClient.findTaskRelocationPlan(taskId)) .timeout(getTaskRelocationTimeout(), TimeUnit.MILLISECONDS, scheduler) .doOnError(error -> logger.info("Could not resolve task relocation status for task: taskId={}, error={}", taskId, ExceptionExt.toMessageChain(error))) .onErrorReturn(e -> Optional.empty()); return Observable.zip( taskObservable, relocationPlanResolver, (task, planOpt) -> planOpt.map(plan -> newTaskWithRelocationPlan(task, plan)).orElse(task) ); } }
TaskRelocationDataInjector { Observable<Task> injectIntoTask(String taskId, Observable<Task> taskObservable) { if (!featureActivationConfiguration.isMergingTaskMigrationPlanInGatewayEnabled()) { return taskObservable; } if (shouldUseRelocationCache()) { return taskObservable.map(task -> newTaskWithRelocationPlan(task, relocationDataReplicator.getCurrent().getPlans().get(taskId))); } Observable<Optional<TaskRelocationPlan>> relocationPlanResolver = ReactorExt.toObservable(relocationServiceClient.findTaskRelocationPlan(taskId)) .timeout(getTaskRelocationTimeout(), TimeUnit.MILLISECONDS, scheduler) .doOnError(error -> logger.info("Could not resolve task relocation status for task: taskId={}, error={}", taskId, ExceptionExt.toMessageChain(error))) .onErrorReturn(e -> Optional.empty()); return Observable.zip( taskObservable, relocationPlanResolver, (task, planOpt) -> planOpt.map(plan -> newTaskWithRelocationPlan(task, plan)).orElse(task) ); } @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); }
TaskRelocationDataInjector { Observable<Task> injectIntoTask(String taskId, Observable<Task> taskObservable) { if (!featureActivationConfiguration.isMergingTaskMigrationPlanInGatewayEnabled()) { return taskObservable; } if (shouldUseRelocationCache()) { return taskObservable.map(task -> newTaskWithRelocationPlan(task, relocationDataReplicator.getCurrent().getPlans().get(taskId))); } Observable<Optional<TaskRelocationPlan>> relocationPlanResolver = ReactorExt.toObservable(relocationServiceClient.findTaskRelocationPlan(taskId)) .timeout(getTaskRelocationTimeout(), TimeUnit.MILLISECONDS, scheduler) .doOnError(error -> logger.info("Could not resolve task relocation status for task: taskId={}, error={}", taskId, ExceptionExt.toMessageChain(error))) .onErrorReturn(e -> Optional.empty()); return Observable.zip( taskObservable, relocationPlanResolver, (task, planOpt) -> planOpt.map(plan -> newTaskWithRelocationPlan(task, plan)).orElse(task) ); } @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);  }
TaskRelocationDataInjector { Observable<Task> injectIntoTask(String taskId, Observable<Task> taskObservable) { if (!featureActivationConfiguration.isMergingTaskMigrationPlanInGatewayEnabled()) { return taskObservable; } if (shouldUseRelocationCache()) { return taskObservable.map(task -> newTaskWithRelocationPlan(task, relocationDataReplicator.getCurrent().getPlans().get(taskId))); } Observable<Optional<TaskRelocationPlan>> relocationPlanResolver = ReactorExt.toObservable(relocationServiceClient.findTaskRelocationPlan(taskId)) .timeout(getTaskRelocationTimeout(), TimeUnit.MILLISECONDS, scheduler) .doOnError(error -> logger.info("Could not resolve task relocation status for task: taskId={}, error={}", taskId, ExceptionExt.toMessageChain(error))) .onErrorReturn(e -> Optional.empty()); return Observable.zip( taskObservable, relocationPlanResolver, (task, planOpt) -> planOpt.map(plan -> newTaskWithRelocationPlan(task, plan)).orElse(task) ); } @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);   }