DefaultNodeConditionController implements NodeConditionController { @VisibleForTesting Mono<Void> handleNodesWithBadCondition(ExecutionContext executionContext) { int iterationCount = executionContext.getExecutionId().getTotal(); logger.debug("Starting node condition controller iteration {} ...", iterationCount); if (hasStaleData()) { logger.info("Stale data. Skipping the node condition control loop iteration- {} ", iterationCount); return Mono.empty(); } return handleNodesWithBadCondition(); } DefaultNodeConditionController(RelocationConfiguration relocationConfiguration,
                                          NodeDataResolver nodeDataResolver,
                                          JobDataReplicator jobDataReplicator,
                                          ReadOnlyJobOperations jobOperations,
                                          JobManagementClient jobManagementClient,
                                          TitusRuntime titusRuntime); @Override void activate(); @Override void deactivate(); Mono<Void> handleNodesWithBadCondition();  }
@Test public void noTerminationsOnDataStaleness() { TitusRuntime titusRuntime = mock(TitusRuntime.class); when(titusRuntime.getRegistry()).thenReturn(new DefaultRegistry()); RelocationConfiguration configuration = mock(RelocationConfiguration.class); when(configuration.getBadNodeConditionPattern()).thenReturn(".*Problem"); when(configuration.isTaskTerminationOnBadNodeConditionEnabled()).thenReturn(true); when(configuration.getDataStalenessThresholdMs()).thenReturn(8000L); NodeDataResolver nodeDataResolver = mock(NodeDataResolver.class); when(nodeDataResolver.getStalenessMs()).thenReturn(5L); JobDataReplicator jobDataReplicator = mock(JobDataReplicator.class); when(jobDataReplicator.getStalenessMs()).thenReturn(10L); ReadOnlyJobOperations readOnlyJobOperations = mock(ReadOnlyJobOperations.class); JobManagementClient jobManagementClient = mock(JobManagementClient.class); Set<String> terminatedTaskIds = new HashSet<>(); when(jobManagementClient.killTask(anyString(), anyBoolean(), any())).thenAnswer(invocation -> { String taskIdToBeTerminated = invocation.getArgument(0); terminatedTaskIds.add(taskIdToBeTerminated); return Mono.empty(); }); DefaultNodeConditionController nodeConditionCtrl = new DefaultNodeConditionController(configuration, nodeDataResolver, jobDataReplicator, readOnlyJobOperations, jobManagementClient, titusRuntime); ExecutionContext executionContext = ExecutionContext.newBuilder().withIteration(ExecutionId.initial()).build(); StepVerifier.create(nodeConditionCtrl.handleNodesWithBadCondition(executionContext)) .verifyComplete(); assertThat(terminatedTaskIds).isEmpty(); }
@Test public void checkTasksTerminatedDueToBadNodeConditions() { Map<String, Node> nodeMap = buildNodes(); List<Job<BatchJobExt>> jobs = getJobs(true); Map<String, List<Task>> tasksByJobIdMap = buildTasksForJobAndNodeAssignment(new ArrayList<>(nodeMap.values()), jobs); TitusRuntime titusRuntime = mock(TitusRuntime.class); when(titusRuntime.getRegistry()).thenReturn(new DefaultRegistry()); RelocationConfiguration configuration = mock(RelocationConfiguration.class); when(configuration.getBadNodeConditionPattern()).thenReturn(".*Failure"); when(configuration.isTaskTerminationOnBadNodeConditionEnabled()).thenReturn(true); NodeDataResolver nodeDataResolver = mock(NodeDataResolver.class); when(nodeDataResolver.resolve()).thenReturn(nodeMap); JobDataReplicator jobDataReplicator = mock(JobDataReplicator.class); when(jobDataReplicator.getStalenessMs()).thenReturn(0L); ReadOnlyJobOperations readOnlyJobOperations = mock(ReadOnlyJobOperations.class); when(readOnlyJobOperations.getJobs()).thenReturn(new ArrayList<>(jobs)); tasksByJobIdMap.forEach((key, value) -> when(readOnlyJobOperations.getTasks(key)).thenReturn(value)); JobManagementClient jobManagementClient = mock(JobManagementClient.class); Set<String> terminatedTaskIds = new HashSet<>(); when(jobManagementClient.killTask(anyString(), anyBoolean(), any())).thenAnswer(invocation -> { String taskIdToBeTerminated = invocation.getArgument(0); terminatedTaskIds.add(taskIdToBeTerminated); return Mono.empty(); }); DefaultNodeConditionController nodeConditionCtrl = new DefaultNodeConditionController(configuration, nodeDataResolver, jobDataReplicator, readOnlyJobOperations, jobManagementClient, titusRuntime); ExecutionContext executionContext = ExecutionContext.newBuilder().withIteration(ExecutionId.initial()).build(); StepVerifier.create(nodeConditionCtrl.handleNodesWithBadCondition(executionContext)).verifyComplete(); assertThat(terminatedTaskIds).isNotEmpty(); assertThat(terminatedTaskIds.size()).isEqualTo(2); verifyTerminatedTasksOnBadNodes(terminatedTaskIds, tasksByJobIdMap, nodeMap); }
@Test public void badNodeConditionsIgnoredForJobsNotOptingIn() { Map<String, Node> nodeMap = buildNodes(); List<Job<BatchJobExt>> jobs = getJobs(false); Map<String, List<Task>> stringListMap = buildTasksForJobAndNodeAssignment(new ArrayList<>(nodeMap.values()), jobs); TitusRuntime titusRuntime = mock(TitusRuntime.class); when(titusRuntime.getRegistry()).thenReturn(new DefaultRegistry()); RelocationConfiguration configuration = mock(RelocationConfiguration.class); when(configuration.getBadNodeConditionPattern()).thenReturn(".*Failure"); when(configuration.isTaskTerminationOnBadNodeConditionEnabled()).thenReturn(true); NodeDataResolver nodeDataResolver = mock(NodeDataResolver.class); when(nodeDataResolver.resolve()).thenReturn(nodeMap); JobDataReplicator jobDataReplicator = mock(JobDataReplicator.class); when(jobDataReplicator.getStalenessMs()).thenReturn(0L); ReadOnlyJobOperations readOnlyJobOperations = mock(ReadOnlyJobOperations.class); when(readOnlyJobOperations.getJobs()).thenReturn(new ArrayList<>(jobs)); stringListMap.forEach((key, value) -> when(readOnlyJobOperations.getTasks(key)).thenReturn(value)); JobManagementClient jobManagementClient = mock(JobManagementClient.class); Set<String> terminatedTaskIds = new HashSet<>(); when(jobManagementClient.killTask(anyString(), anyBoolean(), any())).thenAnswer(invocation -> { String taskIdToBeTerminated = invocation.getArgument(0); terminatedTaskIds.add(taskIdToBeTerminated); return Mono.empty(); }); DefaultNodeConditionController nodeConditionController = new DefaultNodeConditionController(configuration, nodeDataResolver, jobDataReplicator, readOnlyJobOperations, jobManagementClient, titusRuntime); ExecutionContext executionContext = ExecutionContext.newBuilder().withIteration(ExecutionId.initial()).build(); StepVerifier.create(nodeConditionController.handleNodesWithBadCondition(executionContext)) .verifyComplete(); assertThat(terminatedTaskIds).isEmpty(); }