Examples with TopicAssigner org.apache.storm.kafka.spout.subscription.TopicAssigner used on opensource projects

Example 1 with TopicAssigner

use of org.apache.storm.kafka.spout.subscription.TopicAssigner in project storm by apache.

the class KafkaSpoutRebalanceTest method spoutMustIgnoreFailsForTuplesItIsNotAssignedAfterRebalance.

@Test
public void spoutMustIgnoreFailsForTuplesItIsNotAssignedAfterRebalance() throws Exception {
    // Failing tuples for partitions that are no longer assigned is useless since the spout will not be allowed to commit them if they later pass
    TopicAssigner assignerMock = mock(TopicAssigner.class);
    KafkaSpoutRetryService retryServiceMock = mock(KafkaSpoutRetryService.class);
    KafkaSpout<String, String> spout = new KafkaSpout<>(createKafkaSpoutConfigBuilder(topicFilterMock, partitionerMock, -1).setOffsetCommitPeriodMs(10).setRetry(retryServiceMock).build(), consumerFactory, assignerMock);
    String topic = SingleTopicKafkaSpoutConfiguration.TOPIC;
    TopicPartition partitionThatWillBeRevoked = new TopicPartition(topic, 1);
    TopicPartition assignedPartition = new TopicPartition(topic, 2);
    when(retryServiceMock.getMessageId(any(TopicPartition.class), anyLong())).thenReturn(new KafkaSpoutMessageId(partitionThatWillBeRevoked, 0)).thenReturn(new KafkaSpoutMessageId(assignedPartition, 0));
    // Emit a message on each partition and revoke the first partition
    List<KafkaSpoutMessageId> emittedMessageIds = emitOneMessagePerPartitionThenRevokeOnePartition(spout, partitionThatWillBeRevoked, assignedPartition, assignerMock);
    // Check that only two message ids were generated
    verify(retryServiceMock, times(2)).getMessageId(any(TopicPartition.class), anyLong());
    // Fail both emitted tuples
    spout.fail(emittedMessageIds.get(0));
    spout.fail(emittedMessageIds.get(1));
    // Check that only the tuple on the currently assigned partition is retried
    verify(retryServiceMock, never()).schedule(emittedMessageIds.get(0));
    verify(retryServiceMock).schedule(emittedMessageIds.get(1));
}

Example 2 with TopicAssigner

use of org.apache.storm.kafka.spout.subscription.TopicAssigner in project storm by apache.

the class KafkaSpoutRebalanceTest method testReassignPartitionSeeksForOnlyNewPartitions.

@Test
public void testReassignPartitionSeeksForOnlyNewPartitions() {
    /*
         * When partitions are reassigned, the spout should seek with the first poll offset strategy for new partitions.
         * Previously assigned partitions should be left alone, since the spout keeps the emitted and acked state for those.
         */
    TopicAssigner assignerMock = mock(TopicAssigner.class);
    KafkaSpout<String, String> spout = new KafkaSpout<>(createKafkaSpoutConfigBuilder(topicFilterMock, partitionerMock, -1).setFirstPollOffsetStrategy(FirstPollOffsetStrategy.UNCOMMITTED_EARLIEST).build(), consumerFactory, assignerMock);
    String topic = SingleTopicKafkaSpoutConfiguration.TOPIC;
    TopicPartition assignedPartition = new TopicPartition(topic, 1);
    TopicPartition newPartition = new TopicPartition(topic, 2);
    // Setup spout with mock consumer so we can get at the rebalance listener
    spout.open(conf, contextMock, collectorMock);
    spout.activate();
    ArgumentCaptor<ConsumerRebalanceListener> rebalanceListenerCapture = ArgumentCaptor.forClass(ConsumerRebalanceListener.class);
    verify(assignerMock).assignPartitions(any(), any(), rebalanceListenerCapture.capture());
    // Assign partitions to the spout
    ConsumerRebalanceListener consumerRebalanceListener = rebalanceListenerCapture.getValue();
    Set<TopicPartition> assignedPartitions = new HashSet<>();
    assignedPartitions.add(assignedPartition);
    consumerRebalanceListener.onPartitionsAssigned(assignedPartitions);
    reset(consumerMock);
    // Set up committed so it looks like some messages have been committed on each partition
    long committedOffset = 500;
    when(consumerMock.committed(assignedPartition)).thenReturn(new OffsetAndMetadata(committedOffset));
    when(consumerMock.committed(newPartition)).thenReturn(new OffsetAndMetadata(committedOffset));
    // Now rebalance and add a new partition
    consumerRebalanceListener.onPartitionsRevoked(assignedPartitions);
    Set<TopicPartition> newAssignedPartitions = new HashSet<>();
    newAssignedPartitions.add(assignedPartition);
    newAssignedPartitions.add(newPartition);
    consumerRebalanceListener.onPartitionsAssigned(newAssignedPartitions);
    // This partition was previously assigned, so the consumer position shouldn't change
    verify(consumerMock, never()).seek(eq(assignedPartition), anyLong());
    // This partition is new, and should start at the committed offset
    verify(consumerMock).seek(newPartition, committedOffset);
}

Example 3 with TopicAssigner

use of org.apache.storm.kafka.spout.subscription.TopicAssigner in project storm by apache.

the class KafkaTridentSpoutEmitterPartitioningTest method testAssignPartitions.

@Test
public void testAssignPartitions() {
    // Verify correct unwrapping of partitions and delegation of assignment
    TopicAssigner assignerMock = mock(TopicAssigner.class);
    KafkaTridentSpoutEmitter<String, String> emitter = new KafkaTridentSpoutEmitter<>(SingleTopicKafkaTridentSpoutConfiguration.createKafkaSpoutConfigBuilder(-1).build(), topologyContextMock, config -> consumer, assignerMock);
    List<KafkaTridentSpoutTopicPartition> allPartitions = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
        allPartitions.add(new KafkaTridentSpoutTopicPartition(SingleTopicKafkaSpoutConfiguration.TOPIC, i));
    }
    Set<TopicPartition> unwrappedPartitions = allPartitions.stream().map(kttp -> kttp.getTopicPartition()).collect(Collectors.toSet());
    emitter.refreshPartitions(allPartitions);
    verify(assignerMock).assignPartitions(eq(consumer), eq(unwrappedPartitions), any(ConsumerRebalanceListener.class));
}

Example 4 with TopicAssigner

use of org.apache.storm.kafka.spout.subscription.TopicAssigner in project storm by apache.

the class KafkaTridentSpoutEmitterPartitioningTest method testGetOrderedPartitionsIsConsistent.

@Test
public void testGetOrderedPartitionsIsConsistent() {
    KafkaTridentSpoutEmitter<String, String> emitter = new KafkaTridentSpoutEmitter<>(SingleTopicKafkaTridentSpoutConfiguration.createKafkaSpoutConfigBuilder(-1).build(), topologyContextMock, config -> consumer, new TopicAssigner());
    Set<TopicPartition> allPartitions = new HashSet<>();
    int numPartitions = 10;
    for (int i = 0; i < numPartitions; i++) {
        allPartitions.add(new TopicPartition(SingleTopicKafkaSpoutConfiguration.TOPIC, i));
    }
    List<Map<String, Object>> serializedPartitions = allPartitions.stream().map(tp -> tpSerializer.toMap(tp)).collect(Collectors.toList());
    List<KafkaTridentSpoutTopicPartition> orderedPartitions = emitter.getOrderedPartitions(serializedPartitions);
    assertThat("Should contain all partitions", orderedPartitions.size(), is(allPartitions.size()));
    Collections.shuffle(serializedPartitions);
    List<KafkaTridentSpoutTopicPartition> secondGetOrderedPartitions = emitter.getOrderedPartitions(serializedPartitions);
    assertThat("Ordering must be consistent", secondGetOrderedPartitions, is(orderedPartitions));
    serializedPartitions.add(tpSerializer.toMap(new TopicPartition(SingleTopicKafkaSpoutConfiguration.TOPIC, numPartitions)));
    List<KafkaTridentSpoutTopicPartition> orderedPartitionsWithNewPartition = emitter.getOrderedPartitions(serializedPartitions);
    orderedPartitionsWithNewPartition.remove(orderedPartitionsWithNewPartition.size() - 1);
    assertThat("Adding new partitions should not shuffle the existing ordering", orderedPartitionsWithNewPartition, is(orderedPartitions));
}

Example 5 with TopicAssigner

use of org.apache.storm.kafka.spout.subscription.TopicAssigner in project storm by apache.

the class KafkaTridentSpoutEmitterPartitioningTest method testGetPartitionsForTask.

@Test
public void testGetPartitionsForTask() {
    // Verify correct wrapping/unwrapping of partition and delegation of partition assignment
    ManualPartitioner partitionerMock = mock(ManualPartitioner.class);
    when(partitionerMock.getPartitionsForThisTask(any(), any())).thenAnswer(invocation -> {
        List<TopicPartition> partitions = new ArrayList<>(invocation.getArgument(0));
        partitions.remove(0);
        return new HashSet<>(partitions);
    });
    KafkaTridentSpoutEmitter<String, String> emitter = new KafkaTridentSpoutEmitter<>(SingleTopicKafkaTridentSpoutConfiguration.createKafkaSpoutConfigBuilder(mock(TopicFilter.class), partitionerMock, -1).build(), topologyContextMock, config -> consumer, new TopicAssigner());
    List<KafkaTridentSpoutTopicPartition> allPartitions = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
        allPartitions.add(new KafkaTridentSpoutTopicPartition(SingleTopicKafkaSpoutConfiguration.TOPIC, i));
    }
    List<TopicPartition> unwrappedPartitions = allPartitions.stream().map(kttp -> kttp.getTopicPartition()).collect(Collectors.toList());
    List<KafkaTridentSpoutTopicPartition> partitionsForTask = emitter.getPartitionsForTask(0, 2, allPartitions);
    verify(partitionerMock).getPartitionsForThisTask(eq(unwrappedPartitions), any(TopologyContext.class));
    allPartitions.remove(0);
    assertThat("Should have assigned all except the first partition to this task", new HashSet<>(partitionsForTask), is(new HashSet<>(allPartitions)));
}