GroupByContainerIds implements TaskNameGrouper { @Override public Set<ContainerModel> group(Set<TaskModel> tasks) { List<String> containerIds = new ArrayList<>(startContainerCount); for (int i = 0; i < startContainerCount; i++) { containerIds.add(String.valueOf(i)); } return group(tasks, containerIds); } GroupByContainerIds(int count); @Override Set<ContainerModel> group(Set<TaskModel> tasks); @Override Set<ContainerModel> group(Set<TaskModel> tasks, List<String> containerIds); @Override Set<ContainerModel> group(Set<TaskModel> taskModels, GrouperMetadata grouperMetadata);  }
@Test(expected = IllegalArgumentException.class) public void testGroupEmptyTasks() { buildSimpleGrouper(1).group(new HashSet()); }
@Test(expected = UnsupportedOperationException.class) public void testGrouperResultImmutable() { Set<TaskModel> taskModels = generateTaskModels(3); Set<ContainerModel> containers = buildSimpleGrouper(2).group(taskModels); containers.remove(containers.iterator().next()); }
@Test public void testGroupHappyPath() { Set<TaskModel> taskModels = generateTaskModels(5); Set<ContainerModel> containers = buildSimpleGrouper(2).group(taskModels); Map<String, ContainerModel> containersMap = new HashMap<>(); for (ContainerModel container : containers) { containersMap.put(container.getId(), container); } assertEquals(2, containers.size()); ContainerModel container0 = containersMap.get("0"); ContainerModel container1 = containersMap.get("1"); assertNotNull(container0); assertNotNull(container1); assertEquals("0", container0.getId()); assertEquals("1", container1.getId()); assertEquals(3, container0.getTasks().size()); assertEquals(2, container1.getTasks().size()); assertTrue(container0.getTasks().containsKey(getTaskName(0))); assertTrue(container0.getTasks().containsKey(getTaskName(2))); assertTrue(container0.getTasks().containsKey(getTaskName(4))); assertTrue(container1.getTasks().containsKey(getTaskName(1))); assertTrue(container1.getTasks().containsKey(getTaskName(3))); }
@Test public void testGroupWithNullContainerIds() { Set<TaskModel> taskModels = generateTaskModels(5); List<String> containerIds = null; Set<ContainerModel> containers = buildSimpleGrouper(2).group(taskModels, containerIds); Map<String, ContainerModel> containersMap = new HashMap<>(); for (ContainerModel container : containers) { containersMap.put(container.getId(), container); } assertEquals(2, containers.size()); ContainerModel container0 = containersMap.get("0"); ContainerModel container1 = containersMap.get("1"); assertNotNull(container0); assertNotNull(container1); assertEquals("0", container0.getId()); assertEquals("1", container1.getId()); assertEquals(3, container0.getTasks().size()); assertEquals(2, container1.getTasks().size()); assertTrue(container0.getTasks().containsKey(getTaskName(0))); assertTrue(container0.getTasks().containsKey(getTaskName(2))); assertTrue(container0.getTasks().containsKey(getTaskName(4))); assertTrue(container1.getTasks().containsKey(getTaskName(1))); assertTrue(container1.getTasks().containsKey(getTaskName(3))); }
@Test(expected = IllegalArgumentException.class) public void testGroupWithEmptyContainerIds() { Set<TaskModel> taskModels = generateTaskModels(5); buildSimpleGrouper(2).group(taskModels, Collections.emptyList()); }
@Test public void testGroupHappyPathWithListOfContainers() { Set<TaskModel> taskModels = generateTaskModels(5); List<String> containerIds = new ArrayList<String>() { { add("4"); add("2"); } }; Set<ContainerModel> containers = buildSimpleGrouper().group(taskModels, containerIds); Map<String, ContainerModel> containersMap = new HashMap<>(); for (ContainerModel container : containers) { containersMap.put(container.getId(), container); } assertEquals(2, containers.size()); ContainerModel container0 = containersMap.get("4"); ContainerModel container1 = containersMap.get("2"); assertNotNull(container0); assertNotNull(container1); assertEquals("4", container0.getId()); assertEquals("2", container1.getId()); assertEquals(3, container0.getTasks().size()); assertEquals(2, container1.getTasks().size()); assertTrue(container0.getTasks().containsKey(getTaskName(0))); assertTrue(container0.getTasks().containsKey(getTaskName(2))); assertTrue(container0.getTasks().containsKey(getTaskName(4))); assertTrue(container1.getTasks().containsKey(getTaskName(1))); assertTrue(container1.getTasks().containsKey(getTaskName(3))); }
@Test public void testGroupManyTasks() { Set<TaskModel> taskModels = generateTaskModels(21); List<String> containerIds = new ArrayList<String>() { { add("4"); add("2"); } }; Set<ContainerModel> containers = buildSimpleGrouper().group(taskModels, containerIds); Map<String, ContainerModel> containersMap = new HashMap<>(); for (ContainerModel container : containers) { containersMap.put(container.getId(), container); } assertEquals(2, containers.size()); ContainerModel container0 = containersMap.get("4"); ContainerModel container1 = containersMap.get("2"); assertNotNull(container0); assertNotNull(container1); assertEquals("4", container0.getId()); assertEquals("2", container1.getId()); assertEquals(11, container0.getTasks().size()); assertEquals(10, container1.getTasks().size()); assertTrue(container0.getTasks().containsKey(getTaskName(0))); assertTrue(container0.getTasks().containsKey(getTaskName(10))); assertTrue(container0.getTasks().containsKey(getTaskName(12))); assertTrue(container0.getTasks().containsKey(getTaskName(14))); assertTrue(container0.getTasks().containsKey(getTaskName(16))); assertTrue(container0.getTasks().containsKey(getTaskName(18))); assertTrue(container0.getTasks().containsKey(getTaskName(2))); assertTrue(container0.getTasks().containsKey(getTaskName(3))); assertTrue(container0.getTasks().containsKey(getTaskName(5))); assertTrue(container0.getTasks().containsKey(getTaskName(7))); assertTrue(container0.getTasks().containsKey(getTaskName(9))); assertTrue(container1.getTasks().containsKey(getTaskName(1))); assertTrue(container1.getTasks().containsKey(getTaskName(11))); assertTrue(container1.getTasks().containsKey(getTaskName(13))); assertTrue(container1.getTasks().containsKey(getTaskName(15))); assertTrue(container1.getTasks().containsKey(getTaskName(17))); assertTrue(container1.getTasks().containsKey(getTaskName(19))); assertTrue(container1.getTasks().containsKey(getTaskName(20))); assertTrue(container1.getTasks().containsKey(getTaskName(4))); assertTrue(container1.getTasks().containsKey(getTaskName(6))); assertTrue(container1.getTasks().containsKey(getTaskName(8))); }
@Test public void testFewerTasksThanContainers() { final String testContainerId1 = "1"; final String testContainerId2 = "2"; Set<TaskModel> taskModels = generateTaskModels(1); List<String> containerIds = ImmutableList.of(testContainerId1, testContainerId2); Map<TaskName, TaskModel> expectedTasks = taskModels.stream() .collect(Collectors.toMap(TaskModel::getTaskName, x -> x)); ContainerModel expectedContainerModel = new ContainerModel(testContainerId1, expectedTasks); Set<ContainerModel> actualContainerModels = buildSimpleGrouper().group(taskModels, containerIds); assertEquals(1, actualContainerModels.size()); assertEquals(ImmutableSet.of(expectedContainerModel), 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); }
@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(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); }
@Test public void testShouldGenerateIdenticalTaskDistributionWhenNoChangeInProcessorGroup() { 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); actualContainerModels = taskNameGrouper.group(taskModels, grouperMetadata); assertEquals(expectedContainerModels, actualContainerModels); }
@Test public void testShouldMinimizeTaskShuffleWhenAvailableProcessorInGroupChanges() { 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); processorLocality = ImmutableMap.of(testProcessorId1, testLocationId1, testProcessorId2, testLocationId2); grouperMetadata = new GrouperMetadataImpl(processorLocality, taskLocality, new HashMap<>(), new HashMap<>()); actualContainerModels = taskNameGrouper.group(taskModels, grouperMetadata); expectedContainerModels = ImmutableSet.of(new ContainerModel(testProcessorId1, ImmutableMap.of(testTaskName1, testTaskModel1, testTaskName3, testTaskModel3)), new ContainerModel(testProcessorId2, ImmutableMap.of(testTaskName2, testTaskModel2))); assertEquals(expectedContainerModels, actualContainerModels); }
@Test public void testMoreTasksThanProcessors() { String testProcessorId1 = "testProcessorId1"; String testProcessorId2 = "testProcessorId2"; 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"); Map<String, LocationId> processorLocality = ImmutableMap.of(testProcessorId1, testLocationId1, testProcessorId2, testLocationId2); 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 = generateTaskModels(1); List<String> containerIds = ImmutableList.of(testProcessorId1, testProcessorId2); Map<TaskName, TaskModel> expectedTasks = taskModels.stream() .collect(Collectors.toMap(TaskModel::getTaskName, x -> x)); ContainerModel expectedContainerModel = new ContainerModel(testProcessorId1, expectedTasks); Set<ContainerModel> actualContainerModels = buildSimpleGrouper().group(taskModels, grouperMetadata); assertEquals(1, actualContainerModels.size()); assertEquals(ImmutableSet.of(expectedContainerModel), actualContainerModels); }