Examples with KafkaTopicPartition org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition used on opensource projects

Example 26 with KafkaTopicPartition

use of org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition in project flink by apache.

the class Kafka09Fetcher method commitInternalOffsetsToKafka.

@Override
public void commitInternalOffsetsToKafka(Map<KafkaTopicPartition, Long> offsets) throws Exception {
    KafkaTopicPartitionState<TopicPartition>[] partitions = subscribedPartitionStates();
    Map<TopicPartition, OffsetAndMetadata> offsetsToCommit = new HashMap<>(partitions.length);
    for (KafkaTopicPartitionState<TopicPartition> partition : partitions) {
        Long lastProcessedOffset = offsets.get(partition.getKafkaTopicPartition());
        if (lastProcessedOffset != null) {
            // committed offsets through the KafkaConsumer need to be 1 more than the last processed offset.
            // This does not affect Flink's checkpoints/saved state.
            long offsetToCommit = lastProcessedOffset + 1;
            offsetsToCommit.put(partition.getKafkaPartitionHandle(), new OffsetAndMetadata(offsetToCommit));
            partition.setCommittedOffset(offsetToCommit);
        }
    }
    // record the work to be committed by the main consumer thread and make sure the consumer notices that
    consumerThread.setOffsetsToCommit(offsetsToCommit);
}

Example 27 with KafkaTopicPartition

use of org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition in project flink by apache.

the class Kafka09FetcherTest method ensureOffsetsGetCommitted.

@Test
public void ensureOffsetsGetCommitted() throws Exception {
    // test data
    final KafkaTopicPartition testPartition1 = new KafkaTopicPartition("test", 42);
    final KafkaTopicPartition testPartition2 = new KafkaTopicPartition("another", 99);
    final Map<KafkaTopicPartition, Long> testCommitData1 = new HashMap<>();
    testCommitData1.put(testPartition1, 11L);
    testCommitData1.put(testPartition2, 18L);
    final Map<KafkaTopicPartition, Long> testCommitData2 = new HashMap<>();
    testCommitData2.put(testPartition1, 19L);
    testCommitData2.put(testPartition2, 28L);
    final BlockingQueue<Map<TopicPartition, OffsetAndMetadata>> commitStore = new LinkedBlockingQueue<>();
    // ----- the mock consumer with poll(), wakeup(), and commit(A)sync calls ----
    final MultiShotLatch blockerLatch = new MultiShotLatch();
    KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
    when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {

        @Override
        public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) throws InterruptedException {
            blockerLatch.await();
            return ConsumerRecords.empty();
        }
    });
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) {
            blockerLatch.trigger();
            return null;
        }
    }).when(mockConsumer).wakeup();
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) {
            @SuppressWarnings("unchecked") Map<TopicPartition, OffsetAndMetadata> offsets = (Map<TopicPartition, OffsetAndMetadata>) invocation.getArguments()[0];
            OffsetCommitCallback callback = (OffsetCommitCallback) invocation.getArguments()[1];
            commitStore.add(offsets);
            callback.onComplete(offsets, null);
            return null;
        }
    }).when(mockConsumer).commitAsync(Mockito.<Map<TopicPartition, OffsetAndMetadata>>any(), any(OffsetCommitCallback.class));
    // make sure the fetcher creates the mock consumer
    whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);
    // ----- create the test fetcher -----
    @SuppressWarnings("unchecked") SourceContext<String> sourceContext = mock(SourceContext.class);
    Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets = Collections.singletonMap(new KafkaTopicPartition("test", 42), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
    KeyedDeserializationSchema<String> schema = new KeyedDeserializationSchemaWrapper<>(new SimpleStringSchema());
    final Kafka09Fetcher<String> fetcher = new Kafka09Fetcher<>(sourceContext, partitionsWithInitialOffsets, null, /* periodic watermark extractor */
    null, /* punctuated watermark extractor */
    new TestProcessingTimeService(), 10, /* watermark interval */
    this.getClass().getClassLoader(), "task_name", new UnregisteredMetricsGroup(), schema, new Properties(), 0L, false);
    // ----- run the fetcher -----
    final AtomicReference<Throwable> error = new AtomicReference<>();
    final Thread fetcherRunner = new Thread("fetcher runner") {

        @Override
        public void run() {
            try {
                fetcher.runFetchLoop();
            } catch (Throwable t) {
                error.set(t);
            }
        }
    };
    fetcherRunner.start();
    // ----- trigger the first offset commit -----
    fetcher.commitInternalOffsetsToKafka(testCommitData1);
    Map<TopicPartition, OffsetAndMetadata> result1 = commitStore.take();
    for (Entry<TopicPartition, OffsetAndMetadata> entry : result1.entrySet()) {
        TopicPartition partition = entry.getKey();
        if (partition.topic().equals("test")) {
            assertEquals(42, partition.partition());
            assertEquals(12L, entry.getValue().offset());
        } else if (partition.topic().equals("another")) {
            assertEquals(99, partition.partition());
            assertEquals(17L, entry.getValue().offset());
        }
    }
    // ----- trigger the second offset commit -----
    fetcher.commitInternalOffsetsToKafka(testCommitData2);
    Map<TopicPartition, OffsetAndMetadata> result2 = commitStore.take();
    for (Entry<TopicPartition, OffsetAndMetadata> entry : result2.entrySet()) {
        TopicPartition partition = entry.getKey();
        if (partition.topic().equals("test")) {
            assertEquals(42, partition.partition());
            assertEquals(20L, entry.getValue().offset());
        } else if (partition.topic().equals("another")) {
            assertEquals(99, partition.partition());
            assertEquals(27L, entry.getValue().offset());
        }
    }
    // ----- test done, wait till the fetcher is done for a clean shutdown -----
    fetcher.cancel();
    fetcherRunner.join();
    // check that there were no errors in the fetcher
    final Throwable caughtError = error.get();
    if (caughtError != null && !(caughtError instanceof Handover.ClosedException)) {
        throw new Exception("Exception in the fetcher", caughtError);
    }
}

Example 28 with KafkaTopicPartition

use of org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition in project flink by apache.

the class FlinkKafkaConsumerBaseMigrationTest method testRestoreFromFlink11WithEmptyStateWithPartitions.

/** Test restoring from an empty state taken using Flink 1.1, when some partitions could be found for topics. */
@Test
public void testRestoreFromFlink11WithEmptyStateWithPartitions() throws Exception {
    final List<KafkaTopicPartition> partitions = new ArrayList<>();
    partitions.add(new KafkaTopicPartition("abc", 13));
    partitions.add(new KafkaTopicPartition("def", 7));
    final DummyFlinkKafkaConsumer<String> consumerFunction = new DummyFlinkKafkaConsumer<>(partitions);
    StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator = new StreamSource<>(consumerFunction);
    final AbstractStreamOperatorTestHarness<String> testHarness = new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);
    testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);
    testHarness.setup();
    // restore state from binary snapshot file using legacy method
    testHarness.initializeStateFromLegacyCheckpoint(getResourceFilename("kafka-consumer-migration-test-flink1.1-snapshot-empty-state"));
    testHarness.open();
    // the expected state in "kafka-consumer-migration-test-flink1.1-snapshot-empty-state";
    // since the state is empty, the consumer should reflect on the startup mode to determine start offsets.
    final HashMap<KafkaTopicPartition, Long> expectedSubscribedPartitionsWithStartOffsets = new HashMap<>();
    expectedSubscribedPartitionsWithStartOffsets.put(new KafkaTopicPartition("abc", 13), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
    expectedSubscribedPartitionsWithStartOffsets.put(new KafkaTopicPartition("def", 7), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
    // assert that there are partitions and is identical to expected list
    Assert.assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
    Assert.assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());
    Assert.assertEquals(expectedSubscribedPartitionsWithStartOffsets, consumerFunction.getSubscribedPartitionsToStartOffsets());
    // assert that no state was restored
    Assert.assertTrue(consumerFunction.getRestoredState() == null);
    consumerOperator.close();
    consumerOperator.cancel();
}

Example 29 with KafkaTopicPartition

use of org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition in project flink by apache.

the class KafkaConsumerPartitionAssignmentTest method testMultiplePartitionsPerConsumers.

@Test
public void testMultiplePartitionsPerConsumers() {
    try {
        final int[] partitionIDs = { 4, 52, 17, 1, 2, 3, 89, 42, 31, 127, 14 };
        final List<KafkaTopicPartition> partitions = new ArrayList<>();
        final Set<KafkaTopicPartition> allPartitions = new HashSet<>();
        for (int p : partitionIDs) {
            KafkaTopicPartition part = new KafkaTopicPartition("test-topic", p);
            partitions.add(part);
            allPartitions.add(part);
        }
        final int numConsumers = 3;
        final int minPartitionsPerConsumer = partitions.size() / numConsumers;
        final int maxPartitionsPerConsumer = partitions.size() / numConsumers + 1;
        for (int i = 0; i < numConsumers; i++) {
            Map<KafkaTopicPartition, Long> subscribedPartitionsToStartOffsets = new HashMap<>();
            FlinkKafkaConsumerBase.initializeSubscribedPartitionsToStartOffsets(subscribedPartitionsToStartOffsets, partitions, i, numConsumers, StartupMode.GROUP_OFFSETS, null);
            List<KafkaTopicPartition> subscribedPartitions = new ArrayList<>(subscribedPartitionsToStartOffsets.keySet());
            assertTrue(subscribedPartitions.size() >= minPartitionsPerConsumer);
            assertTrue(subscribedPartitions.size() <= maxPartitionsPerConsumer);
            for (KafkaTopicPartition p : subscribedPartitions) {
                // check that the element was actually contained
                assertTrue(allPartitions.remove(p));
            }
        }
        // all partitions must have been assigned
        assertTrue(allPartitions.isEmpty());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 30 with KafkaTopicPartition

use of org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition in project flink by apache.

the class KafkaConsumerPartitionAssignmentTest method testPartitionsFewerThanConsumers.

@Test
public void testPartitionsFewerThanConsumers() {
    try {
        List<KafkaTopicPartition> inPartitions = Arrays.asList(new KafkaTopicPartition("test-topic", 4), new KafkaTopicPartition("test-topic", 52), new KafkaTopicPartition("test-topic", 17), new KafkaTopicPartition("test-topic", 1));
        final Set<KafkaTopicPartition> allPartitions = new HashSet<>();
        allPartitions.addAll(inPartitions);
        final int numConsumers = 2 * inPartitions.size() + 3;
        for (int i = 0; i < numConsumers; i++) {
            Map<KafkaTopicPartition, Long> subscribedPartitionsToStartOffsets = new HashMap<>();
            FlinkKafkaConsumerBase.initializeSubscribedPartitionsToStartOffsets(subscribedPartitionsToStartOffsets, inPartitions, i, numConsumers, StartupMode.GROUP_OFFSETS, null);
            List<KafkaTopicPartition> subscribedPartitions = new ArrayList<>(subscribedPartitionsToStartOffsets.keySet());
            assertTrue(subscribedPartitions.size() <= 1);
            for (KafkaTopicPartition p : subscribedPartitions) {
                // check that the element was actually contained
                assertTrue(allPartitions.remove(p));
            }
        }
        // all partitions must have been assigned
        assertTrue(allPartitions.isEmpty());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}