@Test public void testPartition() throws InterruptedException { List<PartitionInfo> partitions = Arrays.asList( new PartitionInfo("topix", 0, new Node(0, "0", 9092, "us-east-1a"), null, null), new PartitionInfo("topix", 1, new Node(1, "0", 9092, "us-east-1e"), null, null), new PartitionInfo("topix", 3, new Node(2, "0", 9092, "us-east-1d"), null, null), new PartitionInfo("topix", 2, new Node(3, "0", 9092, "us-east-1e"), null, null)); LocalityAwareRandomPartitioner partitioner = new LocalityAwareRandomPartitioner("us-east-1e", 10000); assertEquals("us-east-1e", partitioner.getRack()); partitioner.partition(String.valueOf(0), partitions); List<Integer> filteredPartitions = partitioner.getLocalPartitions().stream() .map(p -> p.partition()).collect(Collectors.toList()); assertEquals(Arrays.asList(new Integer[] { 1, 2 }), filteredPartitions); Set<Integer> outputPartitionIds = new HashSet<>(Arrays.asList(1, 2)); for (int i = 0; i < 10000; i++) { int partition = partitioner.partition(String.valueOf(i), partitions); assertTrue("Bad partition:" + partition, outputPartitionIds.contains(partition)); } partitioner = new LocalityAwareRandomPartitioner("us-east-1c", 1000); partitioner.partition(String.valueOf(0), partitions); filteredPartitions = partitioner.getLocalPartitions().stream().map(p -> p.partition()) .collect(Collectors.toList()); assertEquals(4, filteredPartitions.size()); assertEquals(Arrays.asList(new Integer[] { 0, 1, 3, 2 }), filteredPartitions); Thread.sleep(1100); partitions = Arrays.asList( new PartitionInfo("topix", 0, new Node(0, "0", 9092, "us-east-1c"), null, null), new PartitionInfo("topix", 1, new Node(1, "0", 9092, "us-east-1e"), null, null), new PartitionInfo("topix", 3, new Node(2, "0", 9092, "us-east-1c"), null, null), new PartitionInfo("topix", 2, new Node(3, "0", 9092, "us-east-1d"), null, null)); outputPartitionIds = new HashSet<>(Arrays.asList(0, 3)); for (int i = 0; i < 10000; i++) { int partition = partitioner.partition(String.valueOf(i), partitions); assertTrue("Bad partition:" + partition, outputPartitionIds.contains(partition)); } Thread.sleep(1100); partitions = Arrays.asList( new PartitionInfo("topix", 0, new Node(0, "0", 9092, "us-east-1d"), null, null), new PartitionInfo("topix", 1, new Node(1, "0", 9092, "us-east-1e"), null, null), new PartitionInfo("topix", 3, new Node(2, "0", 9092, "us-east-1c"), null, null), new PartitionInfo("topix", 2, new Node(3, "0", 9092, "us-east-1d"), null, null)); outputPartitionIds = new HashSet<>(Arrays.asList(3)); for (int i = 0; i < 10000; i++) { int partition = partitioner.partition(String.valueOf(i), partitions); assertTrue("Bad partition:" + partition, outputPartitionIds.contains(partition)); } }
@Override public int partition(Object messageKey, List<PartitionInfo> partitions) { if (localPartitions == null || isTimeToRefresh()) { checkAndAssignLocalPartitions(partitions); updateNextRefreshTime(); } return localPartitions.get(Math.abs(random.nextInt() % localPartitions.size())).partition(); }
LocalityAwareRandomPartitioner extends LocalityAwarePartitioner { @Override public int partition(Object messageKey, List<PartitionInfo> partitions) { if (localPartitions == null || isTimeToRefresh()) { checkAndAssignLocalPartitions(partitions); updateNextRefreshTime(); } return localPartitions.get(Math.abs(random.nextInt() % localPartitions.size())).partition(); } }
LocalityAwareRandomPartitioner extends LocalityAwarePartitioner { @Override public int partition(Object messageKey, List<PartitionInfo> partitions) { if (localPartitions == null || isTimeToRefresh()) { checkAndAssignLocalPartitions(partitions); updateNextRefreshTime(); } return localPartitions.get(Math.abs(random.nextInt() % localPartitions.size())).partition(); } LocalityAwareRandomPartitioner(); protected  LocalityAwareRandomPartitioner(String rack, long refreshIntervalMs); }
LocalityAwareRandomPartitioner extends LocalityAwarePartitioner { @Override public int partition(Object messageKey, List<PartitionInfo> partitions) { if (localPartitions == null || isTimeToRefresh()) { checkAndAssignLocalPartitions(partitions); updateNextRefreshTime(); } return localPartitions.get(Math.abs(random.nextInt() % localPartitions.size())).partition(); } LocalityAwareRandomPartitioner(); protected  LocalityAwareRandomPartitioner(String rack, long refreshIntervalMs); @Override int partition(Object messageKey, List<PartitionInfo> partitions); }
LocalityAwareRandomPartitioner extends LocalityAwarePartitioner { @Override public int partition(Object messageKey, List<PartitionInfo> partitions) { if (localPartitions == null || isTimeToRefresh()) { checkAndAssignLocalPartitions(partitions); updateNextRefreshTime(); } return localPartitions.get(Math.abs(random.nextInt() % localPartitions.size())).partition(); } LocalityAwareRandomPartitioner(); protected  LocalityAwareRandomPartitioner(String rack, long refreshIntervalMs); @Override int partition(Object messageKey, List<PartitionInfo> partitions); static final long REFRESH_INTERVAL_MS; }