Examples with ConsumerRecord org.apache.kafka.clients.consumer.ConsumerRecord used on opensource projects

Example 6 with ConsumerRecord

use of org.apache.kafka.clients.consumer.ConsumerRecord in project kafka by apache.

the class ProcessorTopologyTestDriver method process.

/**
     * Send an input message with the given key, value and timestamp on the specified topic to the topology, and then commit the messages.
     *
     * @param topicName the name of the topic on which the message is to be sent
     * @param key the raw message key
     * @param value the raw message value
     * @param timestamp the raw message timestamp
     */
private void process(String topicName, byte[] key, byte[] value, long timestamp) {
    TopicPartition tp = partitionsByTopic.get(topicName);
    if (tp != null) {
        // Add the record ...
        long offset = offsetsByTopicPartition.get(tp).incrementAndGet();
        task.addRecords(tp, records(new ConsumerRecord<>(tp.topic(), tp.partition(), offset, timestamp, TimestampType.CREATE_TIME, 0L, 0, 0, key, value)));
        producer.clear();
        // Process the record ...
        task.process();
        ((InternalProcessorContext) task.context()).setRecordContext(new ProcessorRecordContext(timestamp, offset, tp.partition(), topicName));
        task.commit();
        // Capture all the records sent to the producer ...
        for (ProducerRecord<byte[], byte[]> record : producer.history()) {
            Queue<ProducerRecord<byte[], byte[]>> outputRecords = outputRecordsByTopic.get(record.topic());
            if (outputRecords == null) {
                outputRecords = new LinkedList<>();
                outputRecordsByTopic.put(record.topic(), outputRecords);
            }
            outputRecords.add(record);
            // Forward back into the topology if the produced record is to an internal or a source topic ...
            if (internalTopics.contains(record.topic()) || topology.sourceTopics().contains(record.topic())) {
                process(record.topic(), record.key(), record.value(), record.timestamp());
            }
        }
    } else {
        final TopicPartition global = globalPartitionsByTopic.get(topicName);
        if (global == null) {
            throw new IllegalArgumentException("Unexpected topic: " + topicName);
        }
        final long offset = offsetsByTopicPartition.get(global).incrementAndGet();
        globalStateTask.update(new ConsumerRecord<>(global.topic(), global.partition(), offset, timestamp, TimestampType.CREATE_TIME, 0L, 0, 0, key, value));
        globalStateTask.flushState();
    }
}

Example 7 with ConsumerRecord

use of org.apache.kafka.clients.consumer.ConsumerRecord in project kafka by apache.

the class FetcherTest method testFetchMaxPollRecords.

@Test
public void testFetchMaxPollRecords() {
    Fetcher<byte[], byte[]> fetcher = createFetcher(subscriptions, new Metrics(time), 2);
    List<ConsumerRecord<byte[], byte[]>> records;
    subscriptions.assignFromUser(singleton(tp));
    subscriptions.seek(tp, 1);
    client.prepareResponse(matchesOffset(tp, 1), fetchResponse(this.records, Errors.NONE, 100L, 0));
    client.prepareResponse(matchesOffset(tp, 4), fetchResponse(this.nextRecords, Errors.NONE, 100L, 0));
    assertEquals(1, fetcher.sendFetches());
    consumerClient.poll(0);
    records = fetcher.fetchedRecords().get(tp);
    assertEquals(2, records.size());
    assertEquals(3L, subscriptions.position(tp).longValue());
    assertEquals(1, records.get(0).offset());
    assertEquals(2, records.get(1).offset());
    assertEquals(0, fetcher.sendFetches());
    consumerClient.poll(0);
    records = fetcher.fetchedRecords().get(tp);
    assertEquals(1, records.size());
    assertEquals(4L, subscriptions.position(tp).longValue());
    assertEquals(3, records.get(0).offset());
    assertTrue(fetcher.sendFetches() > 0);
    consumerClient.poll(0);
    records = fetcher.fetchedRecords().get(tp);
    assertEquals(2, records.size());
    assertEquals(6L, subscriptions.position(tp).longValue());
    assertEquals(4, records.get(0).offset());
    assertEquals(5, records.get(1).offset());
}

Example 8 with ConsumerRecord

use of org.apache.kafka.clients.consumer.ConsumerRecord in project samza by apache.

the class TestStreamProcessor method verifyNumMessages.

/**
   * Consumes data from the topic until there are no new messages for a while
   * and asserts that the number of consumed messages is as expected.
   */
private void verifyNumMessages(String topic, int expectedNumMessages) {
    KafkaConsumer consumer = getKafkaConsumer();
    consumer.subscribe(Collections.singletonList(topic));
    int count = 0;
    int emptyPollCount = 0;
    while (count < expectedNumMessages && emptyPollCount < 5) {
        ConsumerRecords records = consumer.poll(5000);
        if (!records.isEmpty()) {
            Iterator<ConsumerRecord> iterator = records.iterator();
            while (iterator.hasNext()) {
                ConsumerRecord record = iterator.next();
                Assert.assertEquals(new String((byte[]) record.value()), String.valueOf(count));
                count++;
            }
        } else {
            emptyPollCount++;
        }
    }
    Assert.assertEquals(count, expectedNumMessages);
}

Example 9 with ConsumerRecord

use of org.apache.kafka.clients.consumer.ConsumerRecord in project beam by apache.

the class KafkaIOTest method mkMockConsumer.

// Update mock consumer with records distributed among the given topics, each with given number
// of partitions. Records are assigned in round-robin order among the partitions.
private static MockConsumer<byte[], byte[]> mkMockConsumer(List<String> topics, int partitionsPerTopic, int numElements, OffsetResetStrategy offsetResetStrategy) {
    final List<TopicPartition> partitions = new ArrayList<>();
    final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> records = new HashMap<>();
    Map<String, List<PartitionInfo>> partitionMap = new HashMap<>();
    for (String topic : topics) {
        List<PartitionInfo> partIds = new ArrayList<>(partitionsPerTopic);
        for (int i = 0; i < partitionsPerTopic; i++) {
            TopicPartition tp = new TopicPartition(topic, i);
            partitions.add(tp);
            partIds.add(new PartitionInfo(topic, i, null, null, null));
            records.put(tp, new ArrayList<ConsumerRecord<byte[], byte[]>>());
        }
        partitionMap.put(topic, partIds);
    }
    int numPartitions = partitions.size();
    final long[] offsets = new long[numPartitions];
    for (int i = 0; i < numElements; i++) {
        int pIdx = i % numPartitions;
        TopicPartition tp = partitions.get(pIdx);
        records.get(tp).add(new ConsumerRecord<>(tp.topic(), tp.partition(), offsets[pIdx]++, // key is 4 byte record id
        ByteBuffer.wrap(new byte[4]).putInt(i).array(), // value is 8 byte record id
        ByteBuffer.wrap(new byte[8]).putLong(i).array()));
    }
    // This is updated when reader assigns partitions.
    final AtomicReference<List<TopicPartition>> assignedPartitions = new AtomicReference<>(Collections.<TopicPartition>emptyList());
    final MockConsumer<byte[], byte[]> consumer = new MockConsumer<byte[], byte[]>(offsetResetStrategy) {

        // override assign() in order to set offset limits & to save assigned partitions.
        //remove keyword '@Override' here, it can work with Kafka client 0.9 and 0.10 as:
        //1. SpEL can find this function, either input is List or Collection;
        //2. List extends Collection, so super.assign() could find either assign(List)
        //  or assign(Collection).
        public void assign(final List<TopicPartition> assigned) {
            super.assign(assigned);
            assignedPartitions.set(ImmutableList.copyOf(assigned));
            for (TopicPartition tp : assigned) {
                updateBeginningOffsets(ImmutableMap.of(tp, 0L));
                updateEndOffsets(ImmutableMap.of(tp, (long) records.get(tp).size()));
            }
        }

        // Override offsetsForTimes() in order to look up the offsets by timestamp.
        // Remove keyword '@Override' here, Kafka client 0.10.1.0 previous versions does not have
        // this method.
        // Should return Map<TopicPartition, OffsetAndTimestamp>, but 0.10.1.0 previous versions
        // does not have the OffsetAndTimestamp class. So return a raw type and use reflection
        // here.
        @SuppressWarnings("unchecked")
        public Map offsetsForTimes(Map<TopicPartition, Long> timestampsToSearch) {
            HashMap<TopicPartition, Object> result = new HashMap<>();
            try {
                Class<?> cls = Class.forName("org.apache.kafka.clients.consumer.OffsetAndTimestamp");
                // OffsetAndTimestamp(long offset, long timestamp)
                Constructor constructor = cls.getDeclaredConstructor(long.class, long.class);
                // In test scope, timestamp == offset.
                for (Map.Entry<TopicPartition, Long> entry : timestampsToSearch.entrySet()) {
                    long maxOffset = offsets[partitions.indexOf(entry.getKey())];
                    Long offset = entry.getValue();
                    if (offset >= maxOffset) {
                        offset = null;
                    }
                    result.put(entry.getKey(), constructor.newInstance(entry.getValue(), offset));
                }
                return result;
            } catch (ClassNotFoundException | IllegalAccessException | InstantiationException | NoSuchMethodException | InvocationTargetException e) {
                throw new RuntimeException(e);
            }
        }
    };
    for (String topic : topics) {
        consumer.updatePartitions(topic, partitionMap.get(topic));
    }
    // MockConsumer does not maintain any relationship between partition seek position and the
    // records added. e.g. if we add 10 records to a partition and then seek to end of the
    // partition, MockConsumer is still going to return the 10 records in next poll. It is
    // our responsibility to make sure currently enqueued records sync with partition offsets.
    // The following task will be called inside each invocation to MockConsumer.poll().
    // We enqueue only the records with the offset >= partition's current position.
    Runnable recordEnqueueTask = new Runnable() {

        @Override
        public void run() {
            // add all the records with offset >= current partition position.
            for (TopicPartition tp : assignedPartitions.get()) {
                long curPos = consumer.position(tp);
                for (ConsumerRecord<byte[], byte[]> r : records.get(tp)) {
                    if (r.offset() >= curPos) {
                        consumer.addRecord(r);
                    }
                }
            }
            consumer.schedulePollTask(this);
        }
    };
    consumer.schedulePollTask(recordEnqueueTask);
    return consumer;
}

Example 10 with ConsumerRecord

use of org.apache.kafka.clients.consumer.ConsumerRecord in project flink by apache.

the class Kafka010FetcherTest method testCancellationWhenEmitBlocks.

@Test
public void testCancellationWhenEmitBlocks() throws Exception {
    // ----- some test data -----
    final String topic = "test-topic";
    final int partition = 3;
    final byte[] payload = new byte[] { 1, 2, 3, 4 };
    final List<ConsumerRecord<byte[], byte[]>> records = Arrays.asList(new ConsumerRecord<byte[], byte[]>(topic, partition, 15, payload, payload), new ConsumerRecord<byte[], byte[]>(topic, partition, 16, payload, payload), new ConsumerRecord<byte[], byte[]>(topic, partition, 17, payload, payload));
    final Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> data = new HashMap<>();
    data.put(new TopicPartition(topic, partition), records);
    final ConsumerRecords<byte[], byte[]> consumerRecords = new ConsumerRecords<>(data);
    // ----- the test consumer -----
    final KafkaConsumer<?, ?> mockConsumer = mock(KafkaConsumer.class);
    when(mockConsumer.poll(anyLong())).thenAnswer(new Answer<ConsumerRecords<?, ?>>() {

        @Override
        public ConsumerRecords<?, ?> answer(InvocationOnMock invocation) {
            return consumerRecords;
        }
    });
    whenNew(KafkaConsumer.class).withAnyArguments().thenReturn(mockConsumer);
    // ----- build a fetcher -----
    BlockingSourceContext<String> sourceContext = new BlockingSourceContext<>();
    Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets = Collections.singletonMap(new KafkaTopicPartition(topic, partition), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
    KeyedDeserializationSchema<String> schema = new KeyedDeserializationSchemaWrapper<>(new SimpleStringSchema());
    final Kafka010Fetcher<String> fetcher = new Kafka010Fetcher<>(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();
    // wait until the thread started to emit records to the source context
    sourceContext.waitTillHasBlocker();
    // now we try to cancel the fetcher, including the interruption usually done on the task thread
    // once it has finished, there must be no more thread blocked on the source context
    fetcher.cancel();
    fetcherRunner.interrupt();
    fetcherRunner.join();
    assertFalse("fetcher threads did not properly finish", sourceContext.isStillBlocking());
}