@Override public JobModel getJobModel() { SystemAdmins systemAdmins = new SystemAdmins(config); StreamMetadataCache streamMetadataCache = new StreamMetadataCache(systemAdmins, 5000, SystemClock.instance()); systemAdmins.start(); try { String containerId = Integer.toString(config.getInt(JobConfig.PROCESSOR_ID())); GrouperMetadata grouperMetadata = new GrouperMetadataImpl(ImmutableMap.of(String.valueOf(containerId), locationId), Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap()); return JobModelManager.readJobModel(this.config, Collections.emptyMap(), streamMetadataCache, grouperMetadata); } finally { systemAdmins.stop(); } }
@Test public void testGetGrouperMetadata() { // Mocking setup. LocalityManager mockLocalityManager = mock(LocalityManager.class); TaskAssignmentManager mockTaskAssignmentManager = Mockito.mock(TaskAssignmentManager.class); Map<String, Map<String, String>> localityMappings = new HashMap<>(); localityMappings.put("0", ImmutableMap.of(SetContainerHostMapping.HOST_KEY, "abc-affinity")); Map<String, String> taskAssignment = ImmutableMap.of("task-0", "0"); // Mock the container locality assignment. when(mockLocalityManager.readContainerLocality()).thenReturn(localityMappings); // Mock the container to task assignment. when(mockTaskAssignmentManager.readTaskAssignment()).thenReturn(taskAssignment); when(mockTaskAssignmentManager.readTaskModes()).thenReturn(Collections.singletonMap(new TaskName("task-0"), TaskMode.Active)); GrouperMetadataImpl grouperMetadata = JobModelManager.getGrouperMetadata(new MapConfig(), mockLocalityManager, mockTaskAssignmentManager); Mockito.verify(mockLocalityManager).readContainerLocality(); Mockito.verify(mockTaskAssignmentManager).readTaskAssignment(); Assert.assertEquals(ImmutableMap.of("0", new LocationId("abc-affinity"), "1", new LocationId("ANY_HOST")), grouperMetadata.getProcessorLocality()); Assert.assertEquals(ImmutableMap.of(new TaskName("task-0"), new LocationId("abc-affinity")), grouperMetadata.getTaskLocality()); }
when(mockGrouperMetadata.getPreviousTaskToProcessorAssignment()).thenReturn(new HashMap<>()); Mockito.doNothing().when(mockTaskAssignmentManager).writeTaskContainerMapping(Mockito.any(), Mockito.any(), Mockito.any());
@Test(expected = IllegalArgumentException.class) public void testShouldFailWhenProcessorLocalityIsEmpty() { TaskNameGrouper taskNameGrouper = buildSimpleGrouper(3); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), new HashMap<>()); taskNameGrouper.group(new HashSet<>(), grouperMetadata); }
/** * Builds the {@link GrouperMetadataImpl} based upon provided {@param jobModelVersion} * and {@param processorNodes}. * @param jobModelVersion the most recent jobModelVersion available in the zookeeper. * @param processorNodes the list of live processors in the zookeeper. * @return the built grouper metadata. */ private GrouperMetadataImpl getGrouperMetadata(String jobModelVersion, List<ProcessorNode> processorNodes) { Map<TaskName, String> taskToProcessorId = new HashMap<>(); Map<TaskName, List<SystemStreamPartition>> taskToSSPs = new HashMap<>(); if (jobModelVersion != null) { JobModel jobModel = zkUtils.getJobModel(jobModelVersion); for (ContainerModel containerModel : jobModel.getContainers().values()) { for (TaskModel taskModel : containerModel.getTasks().values()) { taskToProcessorId.put(taskModel.getTaskName(), containerModel.getId()); for (SystemStreamPartition partition : taskModel.getSystemStreamPartitions()) { taskToSSPs.computeIfAbsent(taskModel.getTaskName(), k -> new ArrayList<>()); taskToSSPs.get(taskModel.getTaskName()).add(partition); } } } } Map<String, LocationId> processorLocality = new HashMap<>(); for (ProcessorNode processorNode : processorNodes) { ProcessorData processorData = processorNode.getProcessorData(); processorLocality.put(processorData.getProcessorId(), processorData.getLocationId()); } Map<TaskName, LocationId> taskLocality = zkUtils.readTaskLocality(); return new GrouperMetadataImpl(processorLocality, taskLocality, taskToSSPs, taskToProcessorId); }
GrouperMetadata grouperMetadata = new GrouperMetadataImpl(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap()); JobModel newJobModel = JobModelManager.readJobModel(this.config, Collections.emptyMap(), streamMetadataCache, grouperMetadata); LOG.info("pid=" + processorId + "Generated new Job Model. Version = " + nextJMVersion);
@Test public void testBalancerEmptyTaskMapping() { Set<TaskModel> taskModels = generateTaskModels(3); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), new HashMap<>()); Set<ContainerModel> groupContainers = new GroupByContainerCount(1).group(taskModels); Set<ContainerModel> balanceContainers = new GroupByContainerCount(1).group(taskModels, grouperMetadata); // Results should be the same as calling group() assertEquals(groupContainers, balanceContainers); }
@Test(expected = UnsupportedOperationException.class) public void testBalancerResultImmutable() { Set<TaskModel> taskModels = generateTaskModels(3); Set<ContainerModel> prevContainers = new GroupByContainerCount(3).group(taskModels); Map<TaskName, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), prevTaskToContainerMapping); Set<ContainerModel> containers = new GroupByContainerCount(2).group(taskModels, grouperMetadata); containers.remove(containers.iterator().next()); }
public static JobModelManager getJobModelManagerUsingReadModel(Config config, StreamMetadataCache streamMetadataCache, HttpServer server, LocalityManager localityManager, Map<String, LocationId> processorLocality) { JobModel jobModel = JobModelManager.readJobModel(config, new HashMap<>(), streamMetadataCache, new GrouperMetadataImpl(processorLocality, new HashMap<>(), new HashMap<>(), new HashMap<>())); return new JobModelManager(new JobModel(jobModel.getConfig(), jobModel.getContainers(), localityManager), server, localityManager); } }
@Test(expected = IllegalArgumentException.class) public void testBalancerEmptyTasks() { Set<TaskModel> taskModels = generateTaskModels(3); Set<ContainerModel> prevContainers = new GroupByContainerCount(3).group(taskModels); Map<TaskName, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), prevTaskToContainerMapping); new GroupByContainerCount(5).group(new HashSet<>(), grouperMetadata); }
@Test(expected = IllegalArgumentException.class) public void testBalancerNewContainerCountGreaterThanTasks() { Set<TaskModel> taskModels = generateTaskModels(3); Set<ContainerModel> prevContainers = new GroupByContainerCount(3).group(taskModels); Map<TaskName, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), prevTaskToContainerMapping); new GroupByContainerCount(5).group(taskModels, grouperMetadata); // Should throw }
@Test public void testBalancerNewContainerCountOne() { Set<TaskModel> taskModels = generateTaskModels(3); Set<ContainerModel> prevContainers = new GroupByContainerCount(3).group(taskModels); Map<TaskName, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), prevTaskToContainerMapping); Set<ContainerModel> groupContainers = new GroupByContainerCount(1).group(taskModels); Set<ContainerModel> balanceContainers = new GroupByContainerCount(1).group(taskModels, grouperMetadata); // Results should be the same as calling group() assertEquals(groupContainers, balanceContainers); }
@Test public void testGroupTaskCountIncrease() { int taskCount = 3; Set<TaskModel> taskModels = generateTaskModels(taskCount); Set<ContainerModel> prevContainers = new GroupByContainerCount(2).group(generateTaskModels(taskCount - 1)); // Here's the key step Map<TaskName, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), prevTaskToContainerMapping); Set<ContainerModel> groupContainers = new GroupByContainerCount(1).group(taskModels); Set<ContainerModel> balanceContainers = new GroupByContainerCount(1).group(taskModels, grouperMetadata); // Results should be the same as calling group() assertEquals(groupContainers, balanceContainers); }
@Test public void testBalancerOldContainerCountOne() { Set<TaskModel> taskModels = generateTaskModels(3); Set<ContainerModel> prevContainers = new GroupByContainerCount(1).group(taskModels); Map<TaskName, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), prevTaskToContainerMapping); Set<ContainerModel> groupContainers = new GroupByContainerCount(3).group(taskModels); Set<ContainerModel> balanceContainers = new GroupByContainerCount(3).group(taskModels, grouperMetadata); // Results should be the same as calling group() assertEquals(groupContainers, balanceContainers); }
@Test public void testGroupTaskCountDecrease() { int taskCount = 3; Set<TaskModel> taskModels = generateTaskModels(taskCount); Set<ContainerModel> prevContainers = new GroupByContainerCount(3).group(generateTaskModels(taskCount + 1)); // Here's the key step Map<TaskName, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(new HashMap<>(), new HashMap<>(), new HashMap<>(), prevTaskToContainerMapping); Set<ContainerModel> groupContainers = new GroupByContainerCount(1).group(taskModels); Set<ContainerModel> balanceContainers = new GroupByContainerCount(1).group(taskModels, grouperMetadata); // Results should be the same as calling group() assertEquals(groupContainers, balanceContainers); }
testTaskName3, testLocationId3); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(processorLocality, taskLocality, new HashMap<>(), new HashMap<>());
testTaskName3, testLocationId3); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(processorLocality, taskLocality, new HashMap<>(), new HashMap<>()); testProcessorId2, testLocationId2); grouperMetadata = new GrouperMetadataImpl(processorLocality, taskLocality, new HashMap<>(), new HashMap<>());
@Test public void testGenerateContainerModelForSingleContainer() { TaskNameGrouper taskNameGrouper = buildSimpleGrouper(1); String testProcessorId1 = "testProcessorId1"; LocationId testLocationId1 = new LocationId("testLocationId1"); LocationId testLocationId2 = new LocationId("testLocationId2"); LocationId testLocationId3 = new LocationId("testLocationId3"); TaskName testTaskName1 = new TaskName("testTasKId1"); TaskName testTaskName2 = new TaskName("testTaskId2"); TaskName testTaskName3 = new TaskName("testTaskId3"); TaskModel testTaskModel1 = new TaskModel(testTaskName1, new HashSet<>(), new Partition(0)); TaskModel testTaskModel2 = new TaskModel(testTaskName2, new HashSet<>(), new Partition(1)); TaskModel testTaskModel3 = new TaskModel(testTaskName3, new HashSet<>(), new Partition(2)); Map<String, LocationId> processorLocality = ImmutableMap.of(testProcessorId1, testLocationId1); Map<TaskName, LocationId> taskLocality = ImmutableMap.of(testTaskName1, testLocationId1, testTaskName2, testLocationId2, testTaskName3, testLocationId3); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(processorLocality, taskLocality, new HashMap<>(), new HashMap<>()); Set<TaskModel> taskModels = ImmutableSet.of(testTaskModel1, testTaskModel2, testTaskModel3); Set<ContainerModel> expectedContainerModels = ImmutableSet.of(new ContainerModel(testProcessorId1, ImmutableMap.of(testTaskName1, testTaskModel1, testTaskName2, testTaskModel2, testTaskName3, testTaskModel3))); Set<ContainerModel> actualContainerModels = taskNameGrouper.group(taskModels, grouperMetadata); assertEquals(expectedContainerModels, actualContainerModels); }
@Test public void testShouldGenerateCorrectContainerModelWhenTaskLocalityIsEmpty() { TaskNameGrouper taskNameGrouper = buildSimpleGrouper(3); String testProcessorId1 = "testProcessorId1"; String testProcessorId2 = "testProcessorId2"; String testProcessorId3 = "testProcessorId3"; LocationId testLocationId1 = new LocationId("testLocationId1"); LocationId testLocationId2 = new LocationId("testLocationId2"); LocationId testLocationId3 = new LocationId("testLocationId3"); TaskName testTaskName1 = new TaskName("testTasKId1"); TaskName testTaskName2 = new TaskName("testTaskId2"); TaskName testTaskName3 = new TaskName("testTaskId3"); TaskModel testTaskModel1 = new TaskModel(testTaskName1, new HashSet<>(), new Partition(0)); TaskModel testTaskModel2 = new TaskModel(testTaskName2, new HashSet<>(), new Partition(1)); TaskModel testTaskModel3 = new TaskModel(testTaskName3, new HashSet<>(), new Partition(2)); Map<String, LocationId> processorLocality = ImmutableMap.of(testProcessorId1, testLocationId1, testProcessorId2, testLocationId2, testProcessorId3, testLocationId3); Map<TaskName, LocationId> taskLocality = ImmutableMap.of(testTaskName1, testLocationId1); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(processorLocality, taskLocality, new HashMap<>(), new HashMap<>()); Set<TaskModel> taskModels = ImmutableSet.of(testTaskModel1, testTaskModel2, testTaskModel3); Set<ContainerModel> expectedContainerModels = ImmutableSet.of(new ContainerModel(testProcessorId1, ImmutableMap.of(testTaskName1, testTaskModel1)), new ContainerModel(testProcessorId2, ImmutableMap.of(testTaskName2, testTaskModel2)), new ContainerModel(testProcessorId3, ImmutableMap.of(testTaskName3, testTaskModel3))); Set<ContainerModel> actualContainerModels = taskNameGrouper.group(taskModels, grouperMetadata); assertEquals(expectedContainerModels, actualContainerModels); }
@Test public void testShouldUseTaskLocalityWhenGeneratingContainerModels() { TaskNameGrouper taskNameGrouper = buildSimpleGrouper(3); String testProcessorId1 = "testProcessorId1"; String testProcessorId2 = "testProcessorId2"; String testProcessorId3 = "testProcessorId3"; LocationId testLocationId1 = new LocationId("testLocationId1"); LocationId testLocationId2 = new LocationId("testLocationId2"); LocationId testLocationId3 = new LocationId("testLocationId3"); TaskName testTaskName1 = new TaskName("testTasKId1"); TaskName testTaskName2 = new TaskName("testTaskId2"); TaskName testTaskName3 = new TaskName("testTaskId3"); TaskModel testTaskModel1 = new TaskModel(testTaskName1, new HashSet<>(), new Partition(0)); TaskModel testTaskModel2 = new TaskModel(testTaskName2, new HashSet<>(), new Partition(1)); TaskModel testTaskModel3 = new TaskModel(testTaskName3, new HashSet<>(), new Partition(2)); Map<String, LocationId> processorLocality = ImmutableMap.of(testProcessorId1, testLocationId1, testProcessorId2, testLocationId2, testProcessorId3, testLocationId3); Map<TaskName, LocationId> taskLocality = ImmutableMap.of(testTaskName1, testLocationId1, testTaskName2, testLocationId2, testTaskName3, testLocationId3); GrouperMetadataImpl grouperMetadata = new GrouperMetadataImpl(processorLocality, taskLocality, new HashMap<>(), new HashMap<>()); Set<TaskModel> taskModels = ImmutableSet.of(testTaskModel1, testTaskModel2, testTaskModel3); Set<ContainerModel> expectedContainerModels = ImmutableSet.of(new ContainerModel(testProcessorId1, ImmutableMap.of(testTaskName1, testTaskModel1)), new ContainerModel(testProcessorId2, ImmutableMap.of(testTaskName2, testTaskModel2)), new ContainerModel(testProcessorId3, ImmutableMap.of(testTaskName3, testTaskModel3))); Set<ContainerModel> actualContainerModels = taskNameGrouper.group(taskModels, grouperMetadata); assertEquals(expectedContainerModels, actualContainerModels); }