Examples with PartitionInfo org.apache.kafka.common.PartitionInfo used on opensource projects

Example 36 with PartitionInfo

use of org.apache.kafka.common.PartitionInfo in project kafka by apache.

the class DefaultPartitionGrouper method partitionGroups.

/**
     * Generate tasks with the assigned topic partitions.
     *
     * @param topicGroups   group of topics that need to be joined together
     * @param metadata      metadata of the consuming cluster
     * @return The map from generated task ids to the assigned partitions
     */
public Map<TaskId, Set<TopicPartition>> partitionGroups(Map<Integer, Set<String>> topicGroups, Cluster metadata) {
    Map<TaskId, Set<TopicPartition>> groups = new HashMap<>();
    for (Map.Entry<Integer, Set<String>> entry : topicGroups.entrySet()) {
        Integer topicGroupId = entry.getKey();
        Set<String> topicGroup = entry.getValue();
        int maxNumPartitions = maxNumPartitions(metadata, topicGroup);
        for (int partitionId = 0; partitionId < maxNumPartitions; partitionId++) {
            Set<TopicPartition> group = new HashSet<>(topicGroup.size());
            for (String topic : topicGroup) {
                List<PartitionInfo> partitions = metadata.partitionsForTopic(topic);
                if (partitionId < partitions.size()) {
                    group.add(new TopicPartition(topic, partitionId));
                }
            }
            groups.put(new TaskId(topicGroupId, partitionId), Collections.unmodifiableSet(group));
        }
    }
    return Collections.unmodifiableMap(groups);
}

Example 37 with PartitionInfo

use of org.apache.kafka.common.PartitionInfo in project storm by apache.

the class KafkaOffsetLagUtil method getOffsetLags.

/**
     *
     * @param newKafkaSpoutOffsetQuery represents the information needed to query kafka for log head and spout offsets
     * @return log head offset, spout offset and lag for each partition
     */
public static List<KafkaOffsetLagResult> getOffsetLags(NewKafkaSpoutOffsetQuery newKafkaSpoutOffsetQuery) {
    KafkaConsumer<String, String> consumer = null;
    List<KafkaOffsetLagResult> result = new ArrayList<>();
    try {
        Properties props = new Properties();
        props.put("bootstrap.servers", newKafkaSpoutOffsetQuery.getBootStrapBrokers());
        props.put("group.id", newKafkaSpoutOffsetQuery.getConsumerGroupId());
        props.put("enable.auto.commit", "false");
        props.put("session.timeout.ms", "30000");
        props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        if (newKafkaSpoutOffsetQuery.getSecurityProtocol() != null) {
            props.put("security.protocol", newKafkaSpoutOffsetQuery.getSecurityProtocol());
        }
        List<TopicPartition> topicPartitionList = new ArrayList<>();
        consumer = new KafkaConsumer<>(props);
        for (String topic : newKafkaSpoutOffsetQuery.getTopics().split(",")) {
            List<PartitionInfo> partitionInfoList = consumer.partitionsFor(topic);
            if (partitionInfoList != null) {
                for (PartitionInfo partitionInfo : partitionInfoList) {
                    topicPartitionList.add(new TopicPartition(partitionInfo.topic(), partitionInfo.partition()));
                }
            }
        }
        consumer.assign(topicPartitionList);
        for (TopicPartition topicPartition : topicPartitionList) {
            OffsetAndMetadata offsetAndMetadata = consumer.committed(topicPartition);
            long committedOffset = offsetAndMetadata != null ? offsetAndMetadata.offset() : -1;
            consumer.seekToEnd(toArrayList(topicPartition));
            result.add(new KafkaOffsetLagResult(topicPartition.topic(), topicPartition.partition(), committedOffset, consumer.position(topicPartition)));
        }
    } finally {
        if (consumer != null) {
            consumer.close();
        }
    }
    return result;
}

Example 38 with PartitionInfo

use of org.apache.kafka.common.PartitionInfo in project storm by apache.

the class ManualPartitionNamedSubscription method refreshAssignment.

@Override
public void refreshAssignment() {
    List<TopicPartition> allPartitions = new ArrayList<>();
    for (String topic : topics) {
        for (PartitionInfo partitionInfo : consumer.partitionsFor(topic)) {
            allPartitions.add(new TopicPartition(partitionInfo.topic(), partitionInfo.partition()));
        }
    }
    Collections.sort(allPartitions, TopicPartitionComparator.INSTANCE);
    Set<TopicPartition> newAssignment = new HashSet<>(partitioner.partition(allPartitions, context));
    if (!newAssignment.equals(currentAssignment)) {
        if (currentAssignment != null) {
            listener.onPartitionsRevoked(currentAssignment);
            listener.onPartitionsAssigned(newAssignment);
        }
        currentAssignment = newAssignment;
        consumer.assign(currentAssignment);
    }
}

Example 39 with PartitionInfo

use of org.apache.kafka.common.PartitionInfo in project storm by apache.

the class ManualPartitionPatternSubscription method refreshAssignment.

@Override
public void refreshAssignment() {
    List<TopicPartition> allPartitions = new ArrayList<>();
    for (Map.Entry<String, List<PartitionInfo>> entry : consumer.listTopics().entrySet()) {
        if (pattern.matcher(entry.getKey()).matches()) {
            for (PartitionInfo partitionInfo : entry.getValue()) {
                allPartitions.add(new TopicPartition(partitionInfo.topic(), partitionInfo.partition()));
            }
        }
    }
    Collections.sort(allPartitions, TopicPartitionComparator.INSTANCE);
    Set<TopicPartition> newAssignment = new HashSet<>(parter.partition(allPartitions, context));
    if (!newAssignment.equals(currentAssignment)) {
        if (currentAssignment != null) {
            listener.onPartitionsRevoked(currentAssignment);
            listener.onPartitionsAssigned(newAssignment);
        }
        currentAssignment = newAssignment;
        consumer.assign(currentAssignment);
    }
}

Example 40 with PartitionInfo

use of org.apache.kafka.common.PartitionInfo in project kafka by apache.

the class KafkaProducerTest method testMetadataFetchOnStaleMetadata.

@PrepareOnlyThisForTest(Metadata.class)
@Test
public void testMetadataFetchOnStaleMetadata() throws Exception {
    Properties props = new Properties();
    props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
    KafkaProducer<String, String> producer = new KafkaProducer<>(props, new StringSerializer(), new StringSerializer());
    Metadata metadata = PowerMock.createNiceMock(Metadata.class);
    MemberModifier.field(KafkaProducer.class, "metadata").set(producer, metadata);
    String topic = "topic";
    ProducerRecord<String, String> initialRecord = new ProducerRecord<>(topic, "value");
    // Create a record with a partition higher than the initial (outdated) partition range
    ProducerRecord<String, String> extendedRecord = new ProducerRecord<>(topic, 2, null, "value");
    Collection<Node> nodes = Collections.singletonList(new Node(0, "host1", 1000));
    final Cluster emptyCluster = new Cluster(null, nodes, Collections.<PartitionInfo>emptySet(), Collections.<String>emptySet(), Collections.<String>emptySet());
    final Cluster initialCluster = new Cluster("dummy", Collections.singletonList(new Node(0, "host1", 1000)), Arrays.asList(new PartitionInfo(topic, 0, null, null, null)), Collections.<String>emptySet(), Collections.<String>emptySet());
    final Cluster extendedCluster = new Cluster("dummy", Collections.singletonList(new Node(0, "host1", 1000)), Arrays.asList(new PartitionInfo(topic, 0, null, null, null), new PartitionInfo(topic, 1, null, null, null), new PartitionInfo(topic, 2, null, null, null)), Collections.<String>emptySet(), Collections.<String>emptySet());
    // Expect exactly one fetch for each attempt to refresh while topic metadata is not available
    final int refreshAttempts = 5;
    EasyMock.expect(metadata.fetch()).andReturn(emptyCluster).times(refreshAttempts - 1);
    EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
    EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
    PowerMock.replay(metadata);
    producer.send(initialRecord);
    PowerMock.verify(metadata);
    // Expect exactly one fetch if topic metadata is available and records are still within range
    PowerMock.reset(metadata);
    EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
    EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
    PowerMock.replay(metadata);
    producer.send(initialRecord, null);
    PowerMock.verify(metadata);
    // Expect exactly two fetches if topic metadata is available but metadata response still returns
    // the same partition size (either because metadata are still stale at the broker too or because
    // there weren't any partitions added in the first place).
    PowerMock.reset(metadata);
    EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
    EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
    EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
    PowerMock.replay(metadata);
    try {
        producer.send(extendedRecord, null);
        fail("Expected KafkaException to be raised");
    } catch (KafkaException e) {
    // expected
    }
    PowerMock.verify(metadata);
    // Expect exactly two fetches if topic metadata is available but outdated for the given record
    PowerMock.reset(metadata);
    EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
    EasyMock.expect(metadata.fetch()).andReturn(extendedCluster).once();
    EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
    PowerMock.replay(metadata);
    producer.send(extendedRecord, null);
    PowerMock.verify(metadata);
}