@Test public void testAssignment() { int maxNumPartitionsForTopic = -1; int totalNumPartitions = 0; Set<PartitionInfo> partitions = new HashSet<>(); for (int i = 0; i < NUM_TOPICS; i++) { int randomNumPartitions = 4 * (RANDOM.nextInt(100) + 1); maxNumPartitionsForTopic = Math.max(randomNumPartitions, maxNumPartitionsForTopic); totalNumPartitions += randomNumPartitions; for (int j = 0; j < randomNumPartitions; j++) { partitions.add(new PartitionInfo(TOPIC_PREFIX + i, j, NODE_0, nodes(), nodes())); } } Cluster cluster = new Cluster("cluster", Arrays.asList(nodes()), partitions, Collections.emptySet(), Collections.emptySet()); Metadata metadata = new Metadata(METADATA_REFRESH_BACKOFF, METADATA_EXPIRY_MS, new LogContext(), new ClusterResourceListeners()); Map<String, Set<PartitionInfo>> topicToTopicPartitions = new HashMap<>(partitions.size()); for (PartitionInfo tp : partitions) { topicToTopicPartitions.putIfAbsent(tp.topic(), new HashSet<>()); topicToTopicPartitions.get(tp.topic()).add(tp); } List<MetadataResponse.TopicMetadata> topicMetadata = new ArrayList<>(partitions.size()); for (Map.Entry<String, Set<PartitionInfo>> entry : topicToTopicPartitions.entrySet()) { List<MetadataResponse.PartitionMetadata> partitionMetadata = new ArrayList<>(entry.getValue().size()); for (PartitionInfo tp : entry.getValue()) { partitionMetadata.add(new MetadataResponse.PartitionMetadata(Errors.NONE, tp.partition(), NODE_0, Optional.of(RecordBatch.NO_PARTITION_LEADER_EPOCH), Arrays.asList(nodes()), Arrays.asList(nodes()), Collections.emptyList())); } topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.NONE, entry.getKey(), false, partitionMetadata)); } MetadataResponse metadataResponse = KafkaCruiseControlUtils.prepareMetadataResponse(cluster.nodes(), cluster.clusterResource().clusterId(), MetadataResponse.NO_CONTROLLER_ID, topicMetadata); metadata.update(KafkaCruiseControlUtils.REQUEST_VERSION_UPDATE, metadataResponse, 0); MetricSamplerPartitionAssignor assignor = new DefaultMetricSamplerPartitionAssignor(); Set<TopicPartition> assignment = assignor.assignPartitions(metadata.fetch()); int maxAssignedNumPartitionsForFetcher = -1; int minAssignedNumPartitionsForFetcher = Integer.MAX_VALUE; int totalAssignedNumPartitions = 0; maxAssignedNumPartitionsForFetcher = Math.max(maxAssignedNumPartitionsForFetcher, assignment.size()); minAssignedNumPartitionsForFetcher = Math.min(minAssignedNumPartitionsForFetcher, assignment.size()); Set<TopicPartition> uniqueAssignedPartitions = new HashSet<>(assignment); totalAssignedNumPartitions += assignment.size(); assertEquals("Total assigned number of partitions should be " + totalNumPartitions, totalNumPartitions, totalAssignedNumPartitions); assertEquals("Total number of unique assigned partitions should be " + totalNumPartitions, totalNumPartitions, uniqueAssignedPartitions.size()); int avgAssignedPartitionsPerFetcher = totalNumPartitions; assertTrue("In the worst case the max number of partitions assigned to a metric fetchers should not differ by " + "more than the partition number of the biggest topic, which is " + maxNumPartitionsForTopic, maxAssignedNumPartitionsForFetcher - avgAssignedPartitionsPerFetcher <= maxNumPartitionsForTopic); assertTrue("In the worst case the min number of partitions assigned to a metric fetchers should not differ by " + "more than the partition number of the biggest topic, which is " + maxNumPartitionsForTopic, avgAssignedPartitionsPerFetcher - minAssignedNumPartitionsForFetcher <= maxNumPartitionsForTopic); }
@Override public List<Set<TopicPartition>> assignPartitions(Cluster cluster, int numMetricFetchers) { if (numMetricFetchers != SUPPORTED_NUM_METRIC_FETCHER) { throw new IllegalArgumentException("DefaultMetricSamplerPartitionAssignor supports only a single metric fetcher."); } List<Set<TopicPartition>> assignments = new ArrayList<>(); assignments.add(assignPartitions(cluster)); return assignments; }
DefaultMetricSamplerPartitionAssignor implements MetricSamplerPartitionAssignor { @Override public List<Set<TopicPartition>> assignPartitions(Cluster cluster, int numMetricFetchers) { if (numMetricFetchers != SUPPORTED_NUM_METRIC_FETCHER) { throw new IllegalArgumentException("DefaultMetricSamplerPartitionAssignor supports only a single metric fetcher."); } List<Set<TopicPartition>> assignments = new ArrayList<>(); assignments.add(assignPartitions(cluster)); return assignments; } }
DefaultMetricSamplerPartitionAssignor implements MetricSamplerPartitionAssignor { @Override public List<Set<TopicPartition>> assignPartitions(Cluster cluster, int numMetricFetchers) { if (numMetricFetchers != SUPPORTED_NUM_METRIC_FETCHER) { throw new IllegalArgumentException("DefaultMetricSamplerPartitionAssignor supports only a single metric fetcher."); } List<Set<TopicPartition>> assignments = new ArrayList<>(); assignments.add(assignPartitions(cluster)); return assignments; }  }
DefaultMetricSamplerPartitionAssignor implements MetricSamplerPartitionAssignor { @Override public List<Set<TopicPartition>> assignPartitions(Cluster cluster, int numMetricFetchers) { if (numMetricFetchers != SUPPORTED_NUM_METRIC_FETCHER) { throw new IllegalArgumentException("DefaultMetricSamplerPartitionAssignor supports only a single metric fetcher."); } List<Set<TopicPartition>> assignments = new ArrayList<>(); assignments.add(assignPartitions(cluster)); return assignments; }  @Override void configure(Map<String, ?> configs); @Override List<Set<TopicPartition>> assignPartitions(Cluster cluster, int numMetricFetchers); @Override Set<TopicPartition> assignPartitions(Cluster cluster); }
DefaultMetricSamplerPartitionAssignor implements MetricSamplerPartitionAssignor { @Override public List<Set<TopicPartition>> assignPartitions(Cluster cluster, int numMetricFetchers) { if (numMetricFetchers != SUPPORTED_NUM_METRIC_FETCHER) { throw new IllegalArgumentException("DefaultMetricSamplerPartitionAssignor supports only a single metric fetcher."); } List<Set<TopicPartition>> assignments = new ArrayList<>(); assignments.add(assignPartitions(cluster)); return assignments; }  @Override void configure(Map<String, ?> configs); @Override List<Set<TopicPartition>> assignPartitions(Cluster cluster, int numMetricFetchers); @Override Set<TopicPartition> assignPartitions(Cluster cluster);  }