/** * Find the next element in the current batch OR schedule {@link KafkaRecordIterator#pollRecords()} (hasMore = false). */ private void findNext() { if (consumerRecordIterator.hasNext()) { nextRecord = consumerRecordIterator.next(); hasMore = nextRecord.offset() < endOffset; } else { hasMore = false; nextRecord = null; } }
@Override public ConsumerRecord<byte[], byte[]> next() { ConsumerRecord<byte[], byte[]> value = nextRecord; Preconditions.checkState(value.offset() < endOffset); findNext(); return value; }
@Nonnull @Override public List<OrderedPartitionableRecord<Integer, Long>> poll(long timeout) { List<OrderedPartitionableRecord<Integer, Long>> polledRecords = new ArrayList<>(); for (ConsumerRecord<byte[], byte[]> record : consumer.poll(timeout)) { polledRecords.add(new OrderedPartitionableRecord<>( record.topic(), record.partition(), record.offset(), record.value() == null ? null : ImmutableList.of(record.value()) )); } return polledRecords; }
private <T> List<Long> collectRecordOffsets(List<ConsumerRecord<T, T>> records) { List<Long> res = new ArrayList<>(records.size()); for (ConsumerRecord<?, ?> record : records) res.add(record.offset()); return res; } }
public KafkaSpoutMessageId(ConsumerRecord<?, ?> consumerRecord, boolean nullTuple) { this(new TopicPartition(consumerRecord.topic(), consumerRecord.partition()), consumerRecord.offset(), nullTuple); }
void set(ConsumerRecord<byte[], byte[]> consumerRecord) { this.partition = consumerRecord.partition(); this.timestamp = consumerRecord.timestamp(); this.offset = consumerRecord.offset(); this.value = consumerRecord.value(); this.recordKey = consumerRecord.key(); }
@Override public List<Object> apply(ConsumerRecord<K, V> record) { return new Values(record.topic(), record.partition(), record.offset(), record.key(), record.value()); }
@Test public void testHasNextNoEnd() { long lastOffset = (long) expectedRecords.size(); this.kafkaRecordIterator = new KafkaRecordIterator(this.consumer, topicPartition, 5L, null, POLL_TIMEOUT_MS); this.compareIterator(expectedRecords.stream() .filter((consumerRecord) -> consumerRecord.offset() >= 5L && consumerRecord.offset() < lastOffset) .collect(Collectors.toList()), this.kafkaRecordIterator); }
@Test public void testHasNextLastRecord() { long startOffset = (long) (expectedRecords.size() - 1); long lastOffset = (long) expectedRecords.size(); this.kafkaRecordIterator = new KafkaRecordIterator(this.consumer, topicPartition, startOffset, lastOffset, POLL_TIMEOUT_MS); this.compareIterator(expectedRecords.stream() .filter((consumerRecord) -> consumerRecord.offset() >= startOffset && consumerRecord.offset() < lastOffset) .collect(Collectors.toList()), this.kafkaRecordIterator); }
@Test public void testHasNextGivenStartEnd() { this.kafkaRecordIterator = new KafkaRecordIterator(this.consumer, topicPartition, 2L, 4L, POLL_TIMEOUT_MS); this.compareIterator(expectedRecords.stream() .filter((consumerRecord) -> consumerRecord.offset() >= 2L && consumerRecord.offset() < 4L) .collect(Collectors.toList()), this.kafkaRecordIterator); }
@Test public void testHasNextFirstRecord() { this.kafkaRecordIterator = new KafkaRecordIterator(this.consumer, topicPartition, 0L, 1L, POLL_TIMEOUT_MS); this.compareIterator(expectedRecords.stream() .filter((consumerRecord) -> consumerRecord.offset() >= 0L && consumerRecord.offset() < 1L) .collect(Collectors.toList()), this.kafkaRecordIterator); }
@Test public void testHasNextNoStart() { this.kafkaRecordIterator = new KafkaRecordIterator(this.consumer, topicPartition, null, 10L, POLL_TIMEOUT_MS); this.compareIterator(expectedRecords.stream() .filter((consumerRecord) -> consumerRecord.offset() >= 0L && consumerRecord.offset() < 10L) .collect(Collectors.toList()), this.kafkaRecordIterator); }
@Test(expected = KafkaRecordIterator.PollTimeoutException.class) public void testPullingBeyondLimit() { //FYI In the Tx world Commits can introduce offset gaps therefore //thus(RECORD_NUMBER + 1) as beyond limit offset is only true if the topic has not Tx or any Control msg. long increment = readUncommitted ? 1 : 2; long requestedEnd = expectedRecords.size() + increment; long requestedStart = expectedRecords.size() - 1; this.kafkaRecordIterator = new KafkaRecordIterator(this.consumer, topicPartition, requestedStart, requestedEnd, POLL_TIMEOUT_MS); this.compareIterator(expectedRecords.stream() .filter((consumerRecord) -> consumerRecord.offset() >= requestedStart) .collect(Collectors.toList()), this.kafkaRecordIterator); }
@Test public void iterator() throws Exception { Map<TopicPartition, List<ConsumerRecord<Integer, String>>> records = new LinkedHashMap<>(); String topic = "topic"; records.put(new TopicPartition(topic, 0), new ArrayList<ConsumerRecord<Integer, String>>()); ConsumerRecord<Integer, String> record1 = new ConsumerRecord<>(topic, 1, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, 1, "value1"); ConsumerRecord<Integer, String> record2 = new ConsumerRecord<>(topic, 1, 1, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, 2, "value2"); records.put(new TopicPartition(topic, 1), Arrays.asList(record1, record2)); records.put(new TopicPartition(topic, 2), new ArrayList<ConsumerRecord<Integer, String>>()); ConsumerRecords<Integer, String> consumerRecords = new ConsumerRecords<>(records); Iterator<ConsumerRecord<Integer, String>> iter = consumerRecords.iterator(); int c = 0; for (; iter.hasNext(); c++) { ConsumerRecord<Integer, String> record = iter.next(); assertEquals(1, record.partition()); assertEquals(topic, record.topic()); assertEquals(c, record.offset()); } assertEquals(2, c); } }
@Test public void testFetchNonContinuousRecords() { // if we are fetching from a compacted topic, there may be gaps in the returned records // this test verifies the fetcher updates the current fetched/consumed positions correctly for this case MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE, TimestampType.CREATE_TIME, 0L); builder.appendWithOffset(15L, 0L, "key".getBytes(), "value-1".getBytes()); builder.appendWithOffset(20L, 0L, "key".getBytes(), "value-2".getBytes()); builder.appendWithOffset(30L, 0L, "key".getBytes(), "value-3".getBytes()); MemoryRecords records = builder.build(); List<ConsumerRecord<byte[], byte[]>> consumerRecords; subscriptions.assignFromUser(singleton(tp0)); subscriptions.seek(tp0, 0); // normal fetch assertEquals(1, fetcher.sendFetches()); client.prepareResponse(fullFetchResponse(tp0, records, Errors.NONE, 100L, 0)); consumerClient.poll(time.timer(0)); consumerRecords = fetcher.fetchedRecords().get(tp0); assertEquals(3, consumerRecords.size()); assertEquals(31L, subscriptions.position(tp0).longValue()); // this is the next fetching position assertEquals(15L, consumerRecords.get(0).offset()); assertEquals(20L, consumerRecords.get(1).offset()); assertEquals(30L, consumerRecords.get(2).offset()); }
@Override @SuppressWarnings("deprecation") public ConsumerRecords<String, String> onConsume(ConsumerRecords<String, String> records) { // This will ensure that we get the cluster metadata when onConsume is called for the first time // as subsequent compareAndSet operations will fail. CLUSTER_ID_BEFORE_ON_CONSUME.compareAndSet(NO_CLUSTER_ID, CLUSTER_META.get()); Map<TopicPartition, List<ConsumerRecord<String, String>>> recordMap = new HashMap<>(); for (TopicPartition tp : records.partitions()) { List<ConsumerRecord<String, String>> lst = new ArrayList<>(); for (ConsumerRecord<String, String> record: records.records(tp)) { lst.add(new ConsumerRecord<>(record.topic(), record.partition(), record.offset(), record.timestamp(), record.timestampType(), record.checksum(), record.serializedKeySize(), record.serializedValueSize(), record.key(), record.value().toUpperCase(Locale.ROOT))); } recordMap.put(tp, lst); } return new ConsumerRecords<String, String>(recordMap); }
private void seekAndConsumeRecord(ByteBuffer responseBuffer, long toOffset) { // Seek to skip the bad record and fetch again. subscriptions.seek(tp0, toOffset); // Should not throw exception after the seek. fetcher.fetchedRecords(); assertEquals(1, fetcher.sendFetches()); client.prepareResponse(fullFetchResponse(tp0, MemoryRecords.readableRecords(responseBuffer), Errors.NONE, 100L, 0)); consumerClient.poll(time.timer(0)); List<ConsumerRecord<byte[], byte[]>> records = fetcher.fetchedRecords().get(tp0); assertEquals(1, records.size()); assertEquals(toOffset, records.get(0).offset()); assertEquals(toOffset + 1, subscriptions.position(tp0).longValue()); }
@Test public void testFetchNormal() { subscriptions.assignFromUser(singleton(tp0)); subscriptions.seek(tp0, 0); // normal fetch assertEquals(1, fetcher.sendFetches()); assertFalse(fetcher.hasCompletedFetches()); client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0)); consumerClient.poll(time.timer(0)); assertTrue(fetcher.hasCompletedFetches()); Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> partitionRecords = fetcher.fetchedRecords(); assertTrue(partitionRecords.containsKey(tp0)); List<ConsumerRecord<byte[], byte[]>> records = partitionRecords.get(tp0); assertEquals(3, records.size()); assertEquals(4L, subscriptions.position(tp0).longValue()); // this is the next fetching position long offset = 1; for (ConsumerRecord<byte[], byte[]> record : records) { assertEquals(offset, record.offset()); offset += 1; } }
private KafkaTridentSpoutEmitter<String, String> createEmitter(FirstPollOffsetStrategy firstPollOffsetStrategy) { return new KafkaTridentSpoutEmitter<>( SingleTopicKafkaTridentSpoutConfiguration.createKafkaSpoutConfigBuilder(-1) .setRecordTranslator(r -> new Values(r.offset()), new Fields("offset")) .setFirstPollOffsetStrategy(firstPollOffsetStrategy) .setPollTimeoutMs(1) .build(), topologyContextMock, config -> consumer, new TopicAssigner()); }
@Test @SuppressWarnings("deprecation") public void testOldConstructor() { String topic = "topic"; int partition = 0; long offset = 23; String key = "key"; String value = "value"; ConsumerRecord<String, String> record = new ConsumerRecord<>(topic, partition, offset, key, value); assertEquals(topic, record.topic()); assertEquals(partition, record.partition()); assertEquals(offset, record.offset()); assertEquals(key, record.key()); assertEquals(value, record.value()); assertEquals(TimestampType.NO_TIMESTAMP_TYPE, record.timestampType()); assertEquals(ConsumerRecord.NO_TIMESTAMP, record.timestamp()); assertEquals(ConsumerRecord.NULL_CHECKSUM, record.checksum()); assertEquals(ConsumerRecord.NULL_SIZE, record.serializedKeySize()); assertEquals(ConsumerRecord.NULL_SIZE, record.serializedValueSize()); assertEquals(Optional.empty(), record.leaderEpoch()); assertEquals(new RecordHeaders(), record.headers()); }