PartitionManager { static PartitionInfo getGroupInfo(String me, SortedSet<String> children, Collection<TableRange> tablets, int groupSize) { int numGroups = Math.max(1, children.size() / groupSize); int[] groupSizes = new int[numGroups]; int count = 0; int myGroupId = -1; int myId = -1; for (String child : children) { if (child.equals(me)) { myGroupId = count; myId = groupSizes[count]; } groupSizes[count]++; count = (count + 1) % numGroups; } List<TableRange> rangesCopy = new ArrayList<>(tablets); Collections.sort(rangesCopy); Collections.shuffle(rangesCopy, new Random(42)); List<TableRange> groupsTablets = new ArrayList<>(); count = 0; for (TableRange tr : rangesCopy) { if (count == myGroupId) { groupsTablets.add(tr); } count = (count + 1) % numGroups; } return new PartitionInfo(myId, myGroupId, groupSizes[myGroupId], numGroups, children.size(), groupsTablets); } PartitionManager(Environment env, long minSleepTime, long maxSleepTime); void stop(); boolean shouldProcess(Notification notification); static final String ZK_FINDER_PREFIX; }
@Test public void testGrouping() { IntFunction<String> nff = i -> String.format("f-%04d", i); for (int numSplits : new int[] {1, 10, 100, 1000}) { for (int numWorkers : new int[] {1, 5, 10, 11, 100}) { for (int groupSize : new int[] {1, 2, 3, 5, 7, 13, 17, 19, 43, 73, 97}) { int expectedGroups = Math.max(1, numWorkers / groupSize); int maxGroupSize = Math.min(numWorkers, groupSize + (int) Math.ceil((numWorkers % groupSize) / (double) expectedGroups)); TreeSet<String> children = new TreeSet<>(); IntStream.range(0, numWorkers).mapToObj(nff).forEach(children::add); Collection<Bytes> rows = IntStream.iterate(0, i -> i + 1000).limit(numSplits) .mapToObj(i -> String.format("r%06d", i)).map(Bytes::of).collect(toList()); Collection<TableRange> tablets = TableRange.fromBytes(rows); Set<String> idCombos = new HashSet<>(); Map<Integer, RangeSet> groupTablets = new HashMap<>(); for (int i = 0; i < numWorkers; i++) { String me = nff.apply(i); PartitionInfo pi = PartitionManager.getGroupInfo(me, children, tablets, groupSize); Assert.assertEquals(expectedGroups, pi.getNumGroups()); Assert.assertTrue(pi.getMyGroupSize() >= Math.min(numWorkers, groupSize) && pi.getMyGroupSize() <= maxGroupSize); Assert.assertEquals(numWorkers, pi.getNumWorkers()); Assert.assertTrue(pi.getMyIdInGroup() >= 0 && pi.getMyIdInGroup() < maxGroupSize); Assert.assertTrue(pi.getMyGroupId() >= 0 && pi.getMyGroupId() < expectedGroups); Assert.assertFalse(idCombos.contains(pi.getMyGroupId() + ":" + pi.getMyIdInGroup())); idCombos.add(pi.getMyGroupId() + ":" + pi.getMyIdInGroup()); if (!groupTablets.containsKey(pi.getMyGroupId())) { groupTablets.put(pi.getMyGroupId(), pi.getMyGroupsRanges()); } else { Assert.assertEquals(groupTablets.get(pi.getMyGroupId()), pi.getMyGroupsRanges()); } } Assert.assertEquals(numWorkers, idCombos.size()); HashSet<TableRange> allTabletsFromGroups = new HashSet<>(); for (RangeSet tabletSet : groupTablets.values()) { tabletSet.forEach(tr -> { Assert.assertFalse(allTabletsFromGroups.contains(tr)); allTabletsFromGroups.add(tr); }); } Assert.assertEquals(new HashSet<>(tablets), allTabletsFromGroups); IntSummaryStatistics summaryStats = groupTablets.values().stream().mapToInt(RangeSet::size).summaryStatistics(); Assert.assertTrue(summaryStats.getMax() - summaryStats.getMin() < 2); } } } }