| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.sdk.io.kafka; |
| |
| import static org.apache.beam.sdk.metrics.MetricResultsMatchers.attemptedMetricsResult; |
| import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem; |
| import static org.hamcrest.Matchers.containsString; |
| import static org.hamcrest.Matchers.greaterThan; |
| import static org.hamcrest.Matchers.hasItem; |
| import static org.hamcrest.Matchers.isA; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertNull; |
| import static org.junit.Assert.assertThat; |
| import static org.junit.Assert.assertTrue; |
| import static org.junit.Assume.assumeTrue; |
| import static org.junit.internal.matchers.ThrowableCauseMatcher.hasCause; |
| import static org.junit.internal.matchers.ThrowableMessageMatcher.hasMessage; |
| |
| import java.io.IOException; |
| import java.lang.reflect.Method; |
| import java.nio.ByteBuffer; |
| import java.util.AbstractMap.SimpleEntry; |
| import java.util.ArrayList; |
| import java.util.Collection; |
| import java.util.Collections; |
| import java.util.Comparator; |
| import java.util.HashMap; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Optional; |
| import java.util.concurrent.ConcurrentHashMap; |
| import java.util.concurrent.ConcurrentMap; |
| import java.util.concurrent.ExecutorService; |
| import java.util.concurrent.Executors; |
| import java.util.concurrent.ThreadLocalRandom; |
| import java.util.concurrent.TimeUnit; |
| import java.util.concurrent.atomic.AtomicBoolean; |
| import java.util.concurrent.atomic.AtomicReference; |
| import java.util.stream.Collectors; |
| import javax.annotation.Nullable; |
| import org.apache.beam.sdk.Pipeline.PipelineExecutionException; |
| import org.apache.beam.sdk.PipelineResult; |
| import org.apache.beam.sdk.coders.BigEndianIntegerCoder; |
| import org.apache.beam.sdk.coders.BigEndianLongCoder; |
| import org.apache.beam.sdk.coders.CoderRegistry; |
| import org.apache.beam.sdk.coders.InstantCoder; |
| import org.apache.beam.sdk.coders.StringUtf8Coder; |
| import org.apache.beam.sdk.coders.VarIntCoder; |
| import org.apache.beam.sdk.coders.VarLongCoder; |
| import org.apache.beam.sdk.io.Read; |
| import org.apache.beam.sdk.io.UnboundedSource; |
| import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; |
| import org.apache.beam.sdk.io.kafka.serialization.InstantDeserializer; |
| import org.apache.beam.sdk.metrics.MetricName; |
| import org.apache.beam.sdk.metrics.MetricNameFilter; |
| import org.apache.beam.sdk.metrics.MetricQueryResults; |
| import org.apache.beam.sdk.metrics.MetricResult; |
| import org.apache.beam.sdk.metrics.MetricsFilter; |
| import org.apache.beam.sdk.metrics.SinkMetrics; |
| import org.apache.beam.sdk.metrics.SourceMetrics; |
| import org.apache.beam.sdk.options.PipelineOptionsFactory; |
| import org.apache.beam.sdk.testing.PAssert; |
| import org.apache.beam.sdk.testing.TestPipeline; |
| import org.apache.beam.sdk.transforms.Count; |
| import org.apache.beam.sdk.transforms.Distinct; |
| import org.apache.beam.sdk.transforms.DoFn; |
| import org.apache.beam.sdk.transforms.Flatten; |
| import org.apache.beam.sdk.transforms.MapElements; |
| import org.apache.beam.sdk.transforms.Max; |
| import org.apache.beam.sdk.transforms.Min; |
| import org.apache.beam.sdk.transforms.ParDo; |
| import org.apache.beam.sdk.transforms.SerializableFunction; |
| import org.apache.beam.sdk.transforms.Values; |
| import org.apache.beam.sdk.transforms.display.DisplayData; |
| import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| import org.apache.beam.sdk.transforms.windowing.FixedWindows; |
| import org.apache.beam.sdk.transforms.windowing.Window; |
| import org.apache.beam.sdk.util.CoderUtils; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.PCollectionList; |
| import org.apache.beam.sdk.values.TypeDescriptors; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Lists; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.Uninterruptibles; |
| import org.apache.kafka.clients.consumer.Consumer; |
| import org.apache.kafka.clients.consumer.ConsumerConfig; |
| import org.apache.kafka.clients.consumer.ConsumerRecord; |
| import org.apache.kafka.clients.consumer.MockConsumer; |
| import org.apache.kafka.clients.consumer.OffsetAndTimestamp; |
| import org.apache.kafka.clients.consumer.OffsetResetStrategy; |
| import org.apache.kafka.clients.producer.MockProducer; |
| import org.apache.kafka.clients.producer.Producer; |
| import org.apache.kafka.clients.producer.ProducerConfig; |
| import org.apache.kafka.clients.producer.ProducerRecord; |
| import org.apache.kafka.common.KafkaException; |
| import org.apache.kafka.common.PartitionInfo; |
| import org.apache.kafka.common.TopicPartition; |
| import org.apache.kafka.common.record.TimestampType; |
| import org.apache.kafka.common.serialization.ByteArrayDeserializer; |
| import org.apache.kafka.common.serialization.Deserializer; |
| import org.apache.kafka.common.serialization.IntegerDeserializer; |
| import org.apache.kafka.common.serialization.IntegerSerializer; |
| import org.apache.kafka.common.serialization.LongDeserializer; |
| import org.apache.kafka.common.serialization.LongSerializer; |
| import org.apache.kafka.common.serialization.Serializer; |
| import org.apache.kafka.common.serialization.StringDeserializer; |
| import org.apache.kafka.common.utils.Utils; |
| import org.hamcrest.collection.IsIterableContainingInAnyOrder; |
| import org.hamcrest.collection.IsIterableWithSize; |
| import org.joda.time.Duration; |
| import org.joda.time.Instant; |
| import org.junit.Ignore; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.rules.ExpectedException; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** |
| * Tests of {@link KafkaIO}. Run with 'mvn test -Dkafka.clients.version=0.10.1.1', to test with a |
| * specific Kafka version. |
| */ |
| @RunWith(JUnit4.class) |
| public class KafkaIOTest { |
| |
| private static final Logger LOG = LoggerFactory.getLogger(KafkaIOTest.class); |
| |
| /* |
| * The tests below borrow code and structure from CountingSourceTest. In addition verifies |
| * the reader interleaves the records from multiple partitions. |
| * |
| * Other tests to consider : |
| * - test KafkaRecordCoder |
| */ |
| |
| @Rule public final transient TestPipeline p = TestPipeline.create(); |
| |
| @Rule public ExpectedException thrown = ExpectedException.none(); |
| |
| private static final Instant LOG_APPEND_START_TIME = new Instant(600 * 1000); |
| private static final String TIMESTAMP_START_MILLIS_CONFIG = "test.timestamp.start.millis"; |
| private static final String TIMESTAMP_TYPE_CONFIG = "test.timestamp.type"; |
| |
| // 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, |
| Map<String, Object> config) { |
| |
| 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<>()); |
| } |
| partitionMap.put(topic, partIds); |
| } |
| |
| int numPartitions = partitions.size(); |
| final long[] offsets = new long[numPartitions]; |
| |
| long timestampStartMillis = |
| (Long) |
| config.getOrDefault(TIMESTAMP_START_MILLIS_CONFIG, LOG_APPEND_START_TIME.getMillis()); |
| TimestampType timestampType = |
| TimestampType.forName( |
| (String) |
| config.getOrDefault( |
| TIMESTAMP_TYPE_CONFIG, TimestampType.LOG_APPEND_TIME.toString())); |
| |
| for (int i = 0; i < numElements; i++) { |
| int pIdx = i % numPartitions; |
| TopicPartition tp = partitions.get(pIdx); |
| |
| byte[] key = ByteBuffer.wrap(new byte[4]).putInt(i).array(); // key is 4 byte record id |
| byte[] value = ByteBuffer.wrap(new byte[8]).putLong(i).array(); // value is 8 byte record id |
| |
| records |
| .get(tp) |
| .add( |
| new ConsumerRecord<>( |
| tp.topic(), |
| tp.partition(), |
| offsets[pIdx]++, |
| timestampStartMillis + Duration.standardSeconds(i).getMillis(), |
| timestampType, |
| 0, |
| key.length, |
| value.length, |
| key, |
| value)); |
| } |
| |
| // 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 |
| public synchronized void assign(final Collection<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. |
| @Override |
| public synchronized Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes( |
| Map<TopicPartition, Long> timestampsToSearch) { |
| return timestampsToSearch.entrySet().stream() |
| .map( |
| e -> { |
| // In test scope, timestamp == offset. |
| long maxOffset = offsets[partitions.indexOf(e.getKey())]; |
| long offset = e.getValue(); |
| OffsetAndTimestamp value = |
| (offset >= maxOffset) ? null : new OffsetAndTimestamp(offset, offset); |
| return new SimpleEntry<>(e.getKey(), value); |
| }) |
| .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue)); |
| } |
| }; |
| |
| 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. |
| int recordsAdded = 0; |
| 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); |
| recordsAdded++; |
| } |
| } |
| } |
| if (recordsAdded == 0) { |
| if (config.get("inject.error.at.eof") != null) { |
| consumer.setException(new KafkaException("Injected error in consumer.poll()")); |
| } |
| // MockConsumer.poll(timeout) does not actually wait even when there aren't any |
| // records. |
| // Add a small wait here in order to avoid busy looping in the reader. |
| Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS); |
| // TODO: BEAM-4086: testUnboundedSourceWithoutBoundedWrapper() occasionally hangs |
| // without this wait. Need to look into it. |
| } |
| consumer.schedulePollTask(this); |
| } |
| }; |
| |
| consumer.schedulePollTask(recordEnqueueTask); |
| return consumer; |
| } |
| |
| private static class ConsumerFactoryFn |
| implements SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> { |
| private final List<String> topics; |
| private final int partitionsPerTopic; |
| private final int numElements; |
| private final OffsetResetStrategy offsetResetStrategy; |
| |
| ConsumerFactoryFn( |
| List<String> topics, |
| int partitionsPerTopic, |
| int numElements, |
| OffsetResetStrategy offsetResetStrategy) { |
| this.topics = topics; |
| this.partitionsPerTopic = partitionsPerTopic; |
| this.numElements = numElements; |
| this.offsetResetStrategy = offsetResetStrategy; |
| } |
| |
| @Override |
| public Consumer<byte[], byte[]> apply(Map<String, Object> config) { |
| return mkMockConsumer(topics, partitionsPerTopic, numElements, offsetResetStrategy, config); |
| } |
| } |
| |
| private static KafkaIO.Read<Integer, Long> mkKafkaReadTransform( |
| int numElements, @Nullable SerializableFunction<KV<Integer, Long>, Instant> timestampFn) { |
| return mkKafkaReadTransform(numElements, numElements, timestampFn); |
| } |
| |
| /** |
| * Creates a consumer with two topics, with 10 partitions each. numElements are (round-robin) |
| * assigned all the 20 partitions. |
| */ |
| private static KafkaIO.Read<Integer, Long> mkKafkaReadTransform( |
| int numElements, |
| int maxNumRecords, |
| @Nullable SerializableFunction<KV<Integer, Long>, Instant> timestampFn) { |
| |
| List<String> topics = ImmutableList.of("topic_a", "topic_b"); |
| |
| KafkaIO.Read<Integer, Long> reader = |
| KafkaIO.<Integer, Long>read() |
| .withBootstrapServers("myServer1:9092,myServer2:9092") |
| .withTopics(topics) |
| .withConsumerFactoryFn( |
| new ConsumerFactoryFn( |
| topics, 10, numElements, OffsetResetStrategy.EARLIEST)) // 20 partitions |
| .withKeyDeserializer(IntegerDeserializer.class) |
| .withValueDeserializer(LongDeserializer.class) |
| .withMaxNumRecords(maxNumRecords); |
| |
| if (timestampFn != null) { |
| return reader.withTimestampFn(timestampFn); |
| } else { |
| return reader; |
| } |
| } |
| |
| private static class AssertMultipleOf implements SerializableFunction<Iterable<Long>, Void> { |
| private final int num; |
| |
| AssertMultipleOf(int num) { |
| this.num = num; |
| } |
| |
| @Override |
| public Void apply(Iterable<Long> values) { |
| for (Long v : values) { |
| assertEquals(0, v % num); |
| } |
| return null; |
| } |
| } |
| |
| public static void addCountingAsserts(PCollection<Long> input, long numElements) { |
| // Count == numElements |
| // Unique count == numElements |
| // Min == 0 |
| // Max == numElements-1 |
| addCountingAsserts(input, numElements, numElements, 0L, numElements - 1); |
| } |
| |
| public static void addCountingAsserts( |
| PCollection<Long> input, long count, long uniqueCount, long min, long max) { |
| |
| PAssert.thatSingleton(input.apply("Count", Count.globally())).isEqualTo(count); |
| |
| PAssert.thatSingleton(input.apply(Distinct.create()).apply("UniqueCount", Count.globally())) |
| .isEqualTo(uniqueCount); |
| |
| PAssert.thatSingleton(input.apply("Min", Min.globally())).isEqualTo(min); |
| |
| PAssert.thatSingleton(input.apply("Max", Max.globally())).isEqualTo(max); |
| } |
| |
| @Test |
| public void testUnboundedSource() { |
| int numElements = 1000; |
| |
| PCollection<Long> input = |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply(Values.create()); |
| |
| addCountingAsserts(input, numElements); |
| p.run(); |
| } |
| |
| @Test |
| public void testUnreachableKafkaBrokers() { |
| // Expect an exception when the Kafka brokers are not reachable on the workers. |
| // We specify partitions explicitly so that splitting does not involve server interaction. |
| // Set request timeout to 10ms so that test does not take long. |
| |
| thrown.expect(Exception.class); |
| thrown.expectMessage("Reader-0: Timeout while initializing partition 'test-0'"); |
| |
| int numElements = 1000; |
| PCollection<Long> input = |
| p.apply( |
| KafkaIO.<Integer, Long>read() |
| .withBootstrapServers("8.8.8.8:9092") // Google public DNS ip. |
| .withTopicPartitions(ImmutableList.of(new TopicPartition("test", 0))) |
| .withKeyDeserializer(IntegerDeserializer.class) |
| .withValueDeserializer(LongDeserializer.class) |
| .withConsumerConfigUpdates( |
| ImmutableMap.of( |
| ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, 10, |
| ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 5, |
| ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 8, |
| ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, 8)) |
| .withMaxNumRecords(10) |
| .withoutMetadata()) |
| .apply(Values.create()); |
| |
| addCountingAsserts(input, numElements); |
| p.run(); |
| } |
| |
| @Test |
| public void testUnboundedSourceWithSingleTopic() { |
| // same as testUnboundedSource, but with single topic |
| |
| int numElements = 1000; |
| String topic = "my_topic"; |
| |
| KafkaIO.Read<Integer, Long> reader = |
| KafkaIO.<Integer, Long>read() |
| .withBootstrapServers("none") |
| .withTopic("my_topic") |
| .withConsumerFactoryFn( |
| new ConsumerFactoryFn( |
| ImmutableList.of(topic), 10, numElements, OffsetResetStrategy.EARLIEST)) |
| .withMaxNumRecords(numElements) |
| .withKeyDeserializer(IntegerDeserializer.class) |
| .withValueDeserializer(LongDeserializer.class); |
| |
| PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create()); |
| |
| addCountingAsserts(input, numElements); |
| p.run(); |
| } |
| |
| @Test |
| public void testUnboundedSourceWithExplicitPartitions() { |
| int numElements = 1000; |
| |
| List<String> topics = ImmutableList.of("test"); |
| |
| KafkaIO.Read<byte[], Long> reader = |
| KafkaIO.<byte[], Long>read() |
| .withBootstrapServers("none") |
| .withTopicPartitions(ImmutableList.of(new TopicPartition("test", 5))) |
| .withConsumerFactoryFn( |
| new ConsumerFactoryFn( |
| topics, 10, numElements, OffsetResetStrategy.EARLIEST)) // 10 partitions |
| .withKeyDeserializer(ByteArrayDeserializer.class) |
| .withValueDeserializer(LongDeserializer.class) |
| .withMaxNumRecords(numElements / 10); |
| |
| PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create()); |
| |
| // assert that every element is a multiple of 5. |
| PAssert.that(input).satisfies(new AssertMultipleOf(5)); |
| |
| PAssert.thatSingleton(input.apply(Count.globally())).isEqualTo(numElements / 10L); |
| |
| p.run(); |
| } |
| |
| private static class ElementValueDiff extends DoFn<Long, Long> { |
| @ProcessElement |
| public void processElement(ProcessContext c) throws Exception { |
| c.output(c.element() - c.timestamp().getMillis()); |
| } |
| } |
| |
| @Test |
| public void testUnboundedSourceTimestamps() { |
| |
| int numElements = 1000; |
| |
| PCollection<Long> input = |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply(Values.create()); |
| |
| addCountingAsserts(input, numElements); |
| |
| PCollection<Long> diffs = |
| input |
| .apply("TimestampDiff", ParDo.of(new ElementValueDiff())) |
| .apply("DistinctTimestamps", Distinct.create()); |
| // This assert also confirms that diffs only has one unique value. |
| PAssert.thatSingleton(diffs).isEqualTo(0L); |
| |
| p.run(); |
| } |
| |
| @Test |
| public void testUnboundedSourceLogAppendTimestamps() { |
| // LogAppendTime (server side timestamp) for records is set based on record index |
| // in MockConsumer above. Ensure that those exact timestamps are set by the source. |
| int numElements = 1000; |
| |
| PCollection<Long> input = |
| p.apply(mkKafkaReadTransform(numElements, null).withLogAppendTime().withoutMetadata()) |
| .apply(Values.create()); |
| |
| addCountingAsserts(input, numElements); |
| |
| PCollection<Long> diffs = |
| input |
| .apply( |
| MapElements.into(TypeDescriptors.longs()) |
| .via(t -> LOG_APPEND_START_TIME.plus(Duration.standardSeconds(t)).getMillis())) |
| .apply("TimestampDiff", ParDo.of(new ElementValueDiff())) |
| .apply("DistinctTimestamps", Distinct.create()); |
| |
| // This assert also confirms that diff only has one unique value. |
| PAssert.thatSingleton(diffs).isEqualTo(0L); |
| |
| p.run(); |
| } |
| |
| @Test |
| public void testUnboundedSourceCustomTimestamps() { |
| // The custom timestamps is set to customTimestampStartMillis + value. |
| // Tests basic functionality of custom timestamps. |
| |
| final int numElements = 1000; |
| final long customTimestampStartMillis = 80000L; |
| |
| PCollection<Long> input = |
| p.apply( |
| mkKafkaReadTransform(numElements, null) |
| .withTimestampPolicyFactory( |
| (tp, prevWatermark) -> |
| new CustomTimestampPolicyWithLimitedDelay<Integer, Long>( |
| (record -> |
| new Instant( |
| TimeUnit.SECONDS.toMillis(record.getKV().getValue()) |
| + customTimestampStartMillis)), |
| Duration.ZERO, |
| prevWatermark)) |
| .withoutMetadata()) |
| .apply(Values.create()); |
| |
| addCountingAsserts(input, numElements); |
| |
| PCollection<Long> diffs = |
| input |
| .apply( |
| MapElements.into(TypeDescriptors.longs()) |
| .via(t -> TimeUnit.SECONDS.toMillis(t) + customTimestampStartMillis)) |
| .apply("TimestampDiff", ParDo.of(new ElementValueDiff())) |
| .apply("DistinctTimestamps", Distinct.create()); |
| |
| // This assert also confirms that diff only has one unique value. |
| PAssert.thatSingleton(diffs).isEqualTo(0L); |
| |
| p.run(); |
| } |
| |
| @Test |
| public void testUnboundedSourceCreateTimestamps() { |
| // Same as testUnboundedSourceCustomTimestamps with create timestamp. |
| |
| final int numElements = 1000; |
| final long createTimestampStartMillis = 50000L; |
| |
| PCollection<Long> input = |
| p.apply( |
| mkKafkaReadTransform(numElements, null) |
| .withCreateTime(Duration.ZERO) |
| .withConsumerConfigUpdates( |
| ImmutableMap.of( |
| TIMESTAMP_TYPE_CONFIG, |
| "CreateTime", |
| TIMESTAMP_START_MILLIS_CONFIG, |
| createTimestampStartMillis)) |
| .withoutMetadata()) |
| .apply(Values.create()); |
| |
| addCountingAsserts(input, numElements); |
| |
| PCollection<Long> diffs = |
| input |
| .apply( |
| MapElements.into(TypeDescriptors.longs()) |
| .via(t -> TimeUnit.SECONDS.toMillis(t) + createTimestampStartMillis)) |
| .apply("TimestampDiff", ParDo.of(new ElementValueDiff())) |
| .apply("DistinctTimestamps", Distinct.create()); |
| |
| // This assert also confirms that diff only has one unique value. |
| PAssert.thatSingleton(diffs).isEqualTo(0L); |
| |
| p.run(); |
| } |
| |
| // Returns TIMESTAMP_MAX_VALUE for watermark when all the records are read from a partition. |
| static class TimestampPolicyWithEndOfSource<K, V> implements TimestampPolicyFactory<K, V> { |
| private final long maxOffset; |
| |
| TimestampPolicyWithEndOfSource(long maxOffset) { |
| this.maxOffset = maxOffset; |
| } |
| |
| @Override |
| public TimestampPolicy<K, V> createTimestampPolicy( |
| TopicPartition tp, Optional<Instant> previousWatermark) { |
| return new TimestampPolicy<K, V>() { |
| long lastOffset = 0; |
| Instant lastTimestamp = BoundedWindow.TIMESTAMP_MIN_VALUE; |
| |
| @Override |
| public Instant getTimestampForRecord(PartitionContext ctx, KafkaRecord<K, V> record) { |
| lastOffset = record.getOffset(); |
| lastTimestamp = new Instant(record.getTimestamp()); |
| return lastTimestamp; |
| } |
| |
| @Override |
| public Instant getWatermark(PartitionContext ctx) { |
| if (lastOffset < maxOffset) { |
| return lastTimestamp; |
| } else { // EOF |
| return BoundedWindow.TIMESTAMP_MAX_VALUE; |
| } |
| } |
| }; |
| } |
| } |
| |
| @Test |
| public void testUnboundedSourceWithExceptionInKafkaFetch() { |
| // Similar testUnboundedSource, but with an injected exception inside Kafk Consumer poll. |
| |
| // The reader should throw an IOException: |
| thrown.expectCause(isA(IOException.class)); |
| thrown.expectCause(hasMessage(containsString("Exception while reading from Kafka"))); |
| // The original exception is from MockConsumer.poll(): |
| thrown.expectCause(hasCause(isA(KafkaException.class))); |
| thrown.expectCause(hasCause(hasMessage(containsString("Injected error in consumer.poll()")))); |
| |
| int numElements = 1000; |
| String topic = "my_topic"; |
| |
| KafkaIO.Read<Integer, Long> reader = |
| KafkaIO.<Integer, Long>read() |
| .withBootstrapServers("none") |
| .withTopic("my_topic") |
| .withConsumerFactoryFn( |
| new ConsumerFactoryFn( |
| ImmutableList.of(topic), 10, numElements, OffsetResetStrategy.EARLIEST)) |
| .withMaxNumRecords(2 * numElements) // Try to read more messages than available. |
| .withConsumerConfigUpdates(ImmutableMap.of("inject.error.at.eof", true)) |
| .withKeyDeserializer(IntegerDeserializer.class) |
| .withValueDeserializer(LongDeserializer.class); |
| |
| PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create()); |
| |
| addCountingAsserts(input, numElements); |
| p.run(); |
| } |
| |
| @Test |
| @Ignore // TODO : BEAM-4086 : enable once flakiness is fixed. |
| public void testUnboundedSourceWithoutBoundedWrapper() { |
| // This is same as testUnboundedSource() without the BoundedSource wrapper. |
| // Most of the tests in this file set 'maxNumRecords' on the source, which wraps |
| // the unbounded source in a bounded source. As a result, the test pipeline run as |
| // bounded/batch pipelines under direct-runner. |
| // This tests runs without such a wrapper over unbounded wrapper, and depends on watermark |
| // progressing to infinity to end the test (see TimestampPolicyWithEndOfSource above). |
| |
| final int numElements = 1000; |
| final int numPartitions = 10; |
| String topic = "testUnboundedSourceWithoutBoundedWrapper"; |
| |
| KafkaIO.Read<byte[], Long> reader = |
| KafkaIO.<byte[], Long>read() |
| .withBootstrapServers(topic) |
| .withTopic(topic) |
| .withConsumerFactoryFn( |
| new ConsumerFactoryFn( |
| ImmutableList.of(topic), |
| numPartitions, |
| numElements, |
| OffsetResetStrategy.EARLIEST)) |
| .withKeyDeserializer(ByteArrayDeserializer.class) |
| .withValueDeserializer(LongDeserializer.class) |
| .withTimestampPolicyFactory( |
| new TimestampPolicyWithEndOfSource<>(numElements / numPartitions - 1)); |
| |
| p.apply("readFromKafka", reader.withoutMetadata()) |
| .apply(Values.create()) |
| .apply(Window.into(FixedWindows.of(Duration.standardDays(100)))); |
| |
| PipelineResult result = p.run(); |
| |
| MetricName elementsRead = SourceMetrics.elementsRead().getName(); |
| |
| MetricQueryResults metrics = |
| result |
| .metrics() |
| .queryMetrics( |
| MetricsFilter.builder() |
| .addNameFilter(MetricNameFilter.inNamespace(elementsRead.getNamespace())) |
| .build()); |
| |
| assertThat( |
| metrics.getCounters(), |
| hasItem( |
| attemptedMetricsResult( |
| elementsRead.getNamespace(), |
| elementsRead.getName(), |
| "readFromKafka", |
| (long) numElements))); |
| } |
| |
| private static class RemoveKafkaMetadata<K, V> extends DoFn<KafkaRecord<K, V>, KV<K, V>> { |
| @ProcessElement |
| public void processElement(ProcessContext ctx) { |
| ctx.output(ctx.element().getKV()); |
| } |
| } |
| |
| @Test |
| public void testUnboundedSourceSplits() throws Exception { |
| |
| int numElements = 1000; |
| int numSplits = 10; |
| |
| // Coders must be specified explicitly here due to the way the transform |
| // is used in the test. |
| UnboundedSource<KafkaRecord<Integer, Long>, ?> initial = |
| mkKafkaReadTransform(numElements, null) |
| .withKeyDeserializerAndCoder(IntegerDeserializer.class, BigEndianIntegerCoder.of()) |
| .withValueDeserializerAndCoder(LongDeserializer.class, BigEndianLongCoder.of()) |
| .makeSource(); |
| |
| List<? extends UnboundedSource<KafkaRecord<Integer, Long>, ?>> splits = |
| initial.split(numSplits, p.getOptions()); |
| assertEquals("Expected exact splitting", numSplits, splits.size()); |
| |
| long elementsPerSplit = numElements / numSplits; |
| assertEquals("Expected even splits", numElements, elementsPerSplit * numSplits); |
| PCollectionList<Long> pcollections = PCollectionList.empty(p); |
| for (int i = 0; i < splits.size(); ++i) { |
| pcollections = |
| pcollections.and( |
| p.apply("split" + i, Read.from(splits.get(i)).withMaxNumRecords(elementsPerSplit)) |
| .apply("Remove Metadata " + i, ParDo.of(new RemoveKafkaMetadata<>())) |
| .apply("collection " + i, Values.create())); |
| } |
| PCollection<Long> input = pcollections.apply(Flatten.pCollections()); |
| |
| addCountingAsserts(input, numElements); |
| p.run(); |
| } |
| |
| /** A timestamp function that uses the given value as the timestamp. */ |
| private static class ValueAsTimestampFn |
| implements SerializableFunction<KV<Integer, Long>, Instant> { |
| @Override |
| public Instant apply(KV<Integer, Long> input) { |
| return new Instant(input.getValue()); |
| } |
| } |
| |
| // Kafka records are read in a separate thread inside the reader. As a result advance() might not |
| // read any records even from the mock consumer, especially for the first record. |
| // This is a helper method to loop until we read a record. |
| private static void advanceOnce(UnboundedReader<?> reader, boolean isStarted) throws IOException { |
| if (!isStarted && reader.start()) { |
| return; |
| } |
| while (!reader.advance()) { |
| // very rarely will there be more than one attempts. |
| // In case of a bug we might end up looping forever, and test will fail with a timeout. |
| |
| // Avoid hard cpu spinning in case of a test failure. |
| try { |
| Thread.sleep(1); |
| } catch (InterruptedException e) { |
| throw new RuntimeException(e); |
| } |
| } |
| } |
| |
| @Test |
| public void testUnboundedSourceCheckpointMark() throws Exception { |
| int numElements = 85; // 85 to make sure some partitions have more records than other. |
| |
| // create a single split: |
| UnboundedSource<KafkaRecord<Integer, Long>, KafkaCheckpointMark> source = |
| mkKafkaReadTransform(numElements, new ValueAsTimestampFn()) |
| .makeSource() |
| .split(1, PipelineOptionsFactory.create()) |
| .get(0); |
| |
| UnboundedReader<KafkaRecord<Integer, Long>> reader = source.createReader(null, null); |
| final int numToSkip = 20; // one from each partition. |
| |
| // advance numToSkip elements |
| for (int i = 0; i < numToSkip; ++i) { |
| advanceOnce(reader, i > 0); |
| } |
| |
| // Confirm that we get the expected element in sequence before checkpointing. |
| assertEquals(numToSkip - 1, (long) reader.getCurrent().getKV().getValue()); |
| assertEquals(numToSkip - 1, reader.getCurrentTimestamp().getMillis()); |
| |
| // Checkpoint and restart, and confirm that the source continues correctly. |
| KafkaCheckpointMark mark = |
| CoderUtils.clone( |
| source.getCheckpointMarkCoder(), (KafkaCheckpointMark) reader.getCheckpointMark()); |
| reader = source.createReader(null, mark); |
| |
| // Confirm that we get the next elements in sequence. |
| // This also confirms that Reader interleaves records from each partitions by the reader. |
| |
| for (int i = numToSkip; i < numElements; i++) { |
| advanceOnce(reader, i > numToSkip); |
| assertEquals(i, (long) reader.getCurrent().getKV().getValue()); |
| assertEquals(i, reader.getCurrentTimestamp().getMillis()); |
| } |
| } |
| |
| @Test |
| public void testUnboundedSourceCheckpointMarkWithEmptyPartitions() throws Exception { |
| // Similar to testUnboundedSourceCheckpointMark(), but verifies that source resumes |
| // properly from empty partitions, without missing messages added since checkpoint. |
| |
| // Initialize consumer with fewer elements than number of partitions so that some are empty. |
| int initialNumElements = 5; |
| UnboundedSource<KafkaRecord<Integer, Long>, KafkaCheckpointMark> source = |
| mkKafkaReadTransform(initialNumElements, new ValueAsTimestampFn()) |
| .makeSource() |
| .split(1, PipelineOptionsFactory.create()) |
| .get(0); |
| |
| UnboundedReader<KafkaRecord<Integer, Long>> reader = source.createReader(null, null); |
| |
| for (int l = 0; l < initialNumElements; ++l) { |
| advanceOnce(reader, l > 0); |
| } |
| |
| // Checkpoint and restart, and confirm that the source continues correctly. |
| KafkaCheckpointMark mark = |
| CoderUtils.clone( |
| source.getCheckpointMarkCoder(), (KafkaCheckpointMark) reader.getCheckpointMark()); |
| |
| // Create another source with MockConsumer with OffsetResetStrategy.LATEST. This insures that |
| // the reader need to explicitly need to seek to first offset for partitions that were empty. |
| |
| int numElements = 100; // all the 20 partitions will have elements |
| List<String> topics = ImmutableList.of("topic_a", "topic_b"); |
| |
| source = |
| KafkaIO.<Integer, Long>read() |
| .withBootstrapServers("none") |
| .withTopics(topics) |
| .withConsumerFactoryFn( |
| new ConsumerFactoryFn(topics, 10, numElements, OffsetResetStrategy.LATEST)) |
| .withKeyDeserializer(IntegerDeserializer.class) |
| .withValueDeserializer(LongDeserializer.class) |
| .withMaxNumRecords(numElements) |
| .withTimestampFn(new ValueAsTimestampFn()) |
| .makeSource() |
| .split(1, PipelineOptionsFactory.create()) |
| .get(0); |
| |
| reader = source.createReader(null, mark); |
| |
| // Verify in any order. As the partitions are unevenly read, the returned records are not in a |
| // simple order. Note that testUnboundedSourceCheckpointMark() verifies round-robin oder. |
| |
| List<Long> expected = new ArrayList<>(); |
| List<Long> actual = new ArrayList<>(); |
| for (long i = initialNumElements; i < numElements; i++) { |
| advanceOnce(reader, i > initialNumElements); |
| expected.add(i); |
| actual.add(reader.getCurrent().getKV().getValue()); |
| } |
| assertThat(actual, IsIterableContainingInAnyOrder.containsInAnyOrder(expected.toArray())); |
| } |
| |
| @Test |
| public void testUnboundedSourceMetrics() { |
| int numElements = 1000; |
| |
| String readStep = "readFromKafka"; |
| |
| p.apply( |
| readStep, |
| mkKafkaReadTransform(numElements, new ValueAsTimestampFn()) |
| .withConsumerConfigUpdates( |
| ImmutableMap.of(ConsumerConfig.GROUP_ID_CONFIG, "test.group")) |
| .commitOffsetsInFinalize() |
| .withoutMetadata()); |
| |
| PipelineResult result = p.run(); |
| |
| String splitId = "0"; |
| |
| MetricName elementsRead = SourceMetrics.elementsRead().getName(); |
| MetricName elementsReadBySplit = SourceMetrics.elementsReadBySplit(splitId).getName(); |
| MetricName bytesRead = SourceMetrics.bytesRead().getName(); |
| MetricName bytesReadBySplit = SourceMetrics.bytesReadBySplit(splitId).getName(); |
| MetricName backlogElementsOfSplit = SourceMetrics.backlogElementsOfSplit(splitId).getName(); |
| MetricName backlogBytesOfSplit = SourceMetrics.backlogBytesOfSplit(splitId).getName(); |
| |
| MetricQueryResults metrics = result.metrics().allMetrics(); |
| |
| Iterable<MetricResult<Long>> counters = metrics.getCounters(); |
| |
| assertThat( |
| counters, |
| hasItem( |
| attemptedMetricsResult( |
| elementsRead.getNamespace(), elementsRead.getName(), readStep, 1000L))); |
| |
| assertThat( |
| counters, |
| hasItem( |
| attemptedMetricsResult( |
| elementsReadBySplit.getNamespace(), |
| elementsReadBySplit.getName(), |
| readStep, |
| 1000L))); |
| |
| assertThat( |
| counters, |
| hasItem( |
| attemptedMetricsResult( |
| bytesRead.getNamespace(), bytesRead.getName(), readStep, 12000L))); |
| |
| assertThat( |
| counters, |
| hasItem( |
| attemptedMetricsResult( |
| bytesReadBySplit.getNamespace(), bytesReadBySplit.getName(), readStep, 12000L))); |
| |
| MetricQueryResults backlogElementsMetrics = |
| result |
| .metrics() |
| .queryMetrics( |
| MetricsFilter.builder() |
| .addNameFilter( |
| MetricNameFilter.named( |
| backlogElementsOfSplit.getNamespace(), |
| backlogElementsOfSplit.getName())) |
| .build()); |
| |
| // since gauge values may be inconsistent in some environments assert only on their existence. |
| assertThat(backlogElementsMetrics.getGauges(), IsIterableWithSize.iterableWithSize(1)); |
| |
| MetricQueryResults backlogBytesMetrics = |
| result |
| .metrics() |
| .queryMetrics( |
| MetricsFilter.builder() |
| .addNameFilter( |
| MetricNameFilter.named( |
| backlogBytesOfSplit.getNamespace(), backlogBytesOfSplit.getName())) |
| .build()); |
| |
| // since gauge values may be inconsistent in some environments assert only on their existence. |
| assertThat(backlogBytesMetrics.getGauges(), IsIterableWithSize.iterableWithSize(1)); |
| |
| // Check checkpointMarkCommitsEnqueued metric. |
| MetricQueryResults commitsEnqueuedMetrics = |
| result |
| .metrics() |
| .queryMetrics( |
| MetricsFilter.builder() |
| .addNameFilter( |
| MetricNameFilter.named( |
| KafkaUnboundedReader.METRIC_NAMESPACE, |
| KafkaUnboundedReader.CHECKPOINT_MARK_COMMITS_ENQUEUED_METRIC)) |
| .build()); |
| |
| assertThat(commitsEnqueuedMetrics.getCounters(), IsIterableWithSize.iterableWithSize(1)); |
| assertThat( |
| commitsEnqueuedMetrics.getCounters().iterator().next().getAttempted(), greaterThan(0L)); |
| } |
| |
| @Test |
| public void testSink() throws Exception { |
| // Simply read from kafka source and write to kafka sink. Then verify the records |
| // are correctly published to mock kafka producer. |
| |
| int numElements = 1000; |
| |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| |
| ProducerSendCompletionThread completionThread = |
| new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); |
| |
| String topic = "test"; |
| |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply( |
| KafkaIO.<Integer, Long>write() |
| .withBootstrapServers("none") |
| .withTopic(topic) |
| .withKeySerializer(IntegerSerializer.class) |
| .withValueSerializer(LongSerializer.class) |
| .withInputTimestamp() |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); |
| |
| p.run(); |
| |
| completionThread.shutdown(); |
| |
| verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, false, true); |
| } |
| } |
| |
| @Test |
| public void testValuesSink() throws Exception { |
| // similar to testSink(), but use values()' interface. |
| |
| int numElements = 1000; |
| |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| |
| ProducerSendCompletionThread completionThread = |
| new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); |
| |
| String topic = "test"; |
| |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply(Values.create()) // there are no keys |
| .apply( |
| KafkaIO.<Integer, Long>write() |
| .withBootstrapServers("none") |
| .withTopic(topic) |
| .withValueSerializer(LongSerializer.class) |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)) |
| .values()); |
| |
| p.run(); |
| |
| completionThread.shutdown(); |
| |
| verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, true, false); |
| } |
| } |
| |
| @Test |
| public void testRecordsSink() throws Exception { |
| // Simply read from kafka source and write to kafka sink using ProducerRecord transform. Then |
| // verify the records are correctly published to mock kafka producer. |
| |
| int numElements = 1000; |
| |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| |
| ProducerSendCompletionThread completionThread = |
| new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); |
| |
| String topic = "test"; |
| |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply(ParDo.of(new KV2ProducerRecord(topic))) |
| .setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of())) |
| .apply( |
| KafkaIO.<Integer, Long>writeRecords() |
| .withBootstrapServers("none") |
| .withTopic(topic) |
| .withKeySerializer(IntegerSerializer.class) |
| .withValueSerializer(LongSerializer.class) |
| .withInputTimestamp() |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); |
| |
| p.run(); |
| |
| completionThread.shutdown(); |
| |
| verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, false, true); |
| } |
| } |
| |
| @Test |
| public void testSinkToMultipleTopics() throws Exception { |
| // Set different output topic names |
| int numElements = 1000; |
| |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| |
| ProducerSendCompletionThread completionThread = |
| new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); |
| |
| String defaultTopic = "test"; |
| |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply(ParDo.of(new KV2ProducerRecord(defaultTopic, false))) |
| .setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of())) |
| .apply( |
| KafkaIO.<Integer, Long>writeRecords() |
| .withBootstrapServers("none") |
| .withKeySerializer(IntegerSerializer.class) |
| .withValueSerializer(LongSerializer.class) |
| .withInputTimestamp() |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); |
| |
| p.run(); |
| |
| completionThread.shutdown(); |
| |
| // Verify that appropriate messages are written to different Kafka topics |
| List<ProducerRecord<Integer, Long>> sent = producerWrapper.mockProducer.history(); |
| |
| for (int i = 0; i < numElements; i++) { |
| ProducerRecord<Integer, Long> record = sent.get(i); |
| if (i % 2 == 0) { |
| assertEquals("test_2", record.topic()); |
| } else { |
| assertEquals("test_1", record.topic()); |
| } |
| assertEquals(i, record.key().intValue()); |
| assertEquals(i, record.value().longValue()); |
| assertEquals(i, record.timestamp().intValue()); |
| } |
| } |
| } |
| |
| @Test |
| public void testSinkProducerRecordsWithCustomTS() throws Exception { |
| int numElements = 1000; |
| |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| |
| ProducerSendCompletionThread completionThread = |
| new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); |
| |
| final String defaultTopic = "test"; |
| final Long ts = System.currentTimeMillis(); |
| |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply(ParDo.of(new KV2ProducerRecord(defaultTopic, ts))) |
| .setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of())) |
| .apply( |
| KafkaIO.<Integer, Long>writeRecords() |
| .withBootstrapServers("none") |
| .withKeySerializer(IntegerSerializer.class) |
| .withValueSerializer(LongSerializer.class) |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); |
| |
| p.run(); |
| |
| completionThread.shutdown(); |
| |
| // Verify that messages are written with user-defined timestamp |
| List<ProducerRecord<Integer, Long>> sent = producerWrapper.mockProducer.history(); |
| |
| for (int i = 0; i < numElements; i++) { |
| ProducerRecord<Integer, Long> record = sent.get(i); |
| assertEquals(defaultTopic, record.topic()); |
| assertEquals(i, record.key().intValue()); |
| assertEquals(i, record.value().longValue()); |
| assertEquals(ts, record.timestamp()); |
| } |
| } |
| } |
| |
| private static class KV2ProducerRecord |
| extends DoFn<KV<Integer, Long>, ProducerRecord<Integer, Long>> { |
| final String topic; |
| final boolean isSingleTopic; |
| final Long ts; |
| |
| KV2ProducerRecord(String topic) { |
| this(topic, true); |
| } |
| |
| KV2ProducerRecord(String topic, Long ts) { |
| this(topic, true, ts); |
| } |
| |
| KV2ProducerRecord(String topic, boolean isSingleTopic) { |
| this(topic, isSingleTopic, null); |
| } |
| |
| KV2ProducerRecord(String topic, boolean isSingleTopic, Long ts) { |
| this.topic = topic; |
| this.isSingleTopic = isSingleTopic; |
| this.ts = ts; |
| } |
| |
| @ProcessElement |
| public void processElement(ProcessContext ctx) { |
| KV<Integer, Long> kv = ctx.element(); |
| if (isSingleTopic) { |
| ctx.output(new ProducerRecord<>(topic, null, ts, kv.getKey(), kv.getValue())); |
| } else { |
| if (kv.getKey() % 2 == 0) { |
| ctx.output(new ProducerRecord<>(topic + "_2", null, ts, kv.getKey(), kv.getValue())); |
| } else { |
| ctx.output(new ProducerRecord<>(topic + "_1", null, ts, kv.getKey(), kv.getValue())); |
| } |
| } |
| } |
| } |
| |
| @Test |
| public void testExactlyOnceSink() { |
| // testSink() with EOS enabled. |
| // This does not actually inject retries in a stage to test exactly-once-semantics. |
| // It mainly exercises the code in normal flow without retries. |
| // Ideally we should test EOS Sink by triggering replays of a messages between stages. |
| // It is not feasible to test such retries with direct runner. When DoFnTester supports |
| // state, we can test ExactlyOnceWriter DoFn directly to ensure it handles retries correctly. |
| |
| if (!ProducerSpEL.supportsTransactions()) { |
| LOG.warn( |
| "testExactlyOnceSink() is disabled as Kafka client version does not support transactions."); |
| return; |
| } |
| |
| int numElements = 1000; |
| |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| |
| ProducerSendCompletionThread completionThread = |
| new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); |
| |
| String topic = "test"; |
| |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply( |
| KafkaIO.<Integer, Long>write() |
| .withBootstrapServers("none") |
| .withTopic(topic) |
| .withKeySerializer(IntegerSerializer.class) |
| .withValueSerializer(LongSerializer.class) |
| .withEOS(1, "test") |
| .withConsumerFactoryFn( |
| new ConsumerFactoryFn( |
| Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST)) |
| .withPublishTimestampFunction((e, ts) -> ts) |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); |
| |
| p.run(); |
| |
| completionThread.shutdown(); |
| |
| verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, false, true); |
| } |
| } |
| |
| @Test |
| public void testSinkWithSendErrors() throws Throwable { |
| // similar to testSink(), except that up to 10 of the send calls to producer will fail |
| // asynchronously. |
| |
| // TODO: Ideally we want the pipeline to run to completion by retrying bundles that fail. |
| // We limit the number of errors injected to 10 below. This would reflect a real streaming |
| // pipeline. But I am sure how to achieve that. For now expect an exception: |
| |
| thrown.expect(InjectedErrorException.class); |
| thrown.expectMessage("Injected Error #1"); |
| |
| int numElements = 1000; |
| |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| |
| ProducerSendCompletionThread completionThreadWithErrors = |
| new ProducerSendCompletionThread(producerWrapper.mockProducer, 10, 100).start(); |
| |
| String topic = "test"; |
| |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply( |
| KafkaIO.<Integer, Long>write() |
| .withBootstrapServers("none") |
| .withTopic(topic) |
| .withKeySerializer(IntegerSerializer.class) |
| .withValueSerializer(LongSerializer.class) |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); |
| |
| try { |
| p.run(); |
| } catch (PipelineExecutionException e) { |
| // throwing inner exception helps assert that first exception is thrown from the Sink |
| throw e.getCause().getCause(); |
| } finally { |
| completionThreadWithErrors.shutdown(); |
| } |
| } |
| } |
| |
| @Test |
| public void testUnboundedSourceStartReadTime() { |
| |
| assumeTrue(new ConsumerSpEL().hasOffsetsForTimes()); |
| |
| int numElements = 1000; |
| // In this MockConsumer, we let the elements of the time and offset equal and there are 20 |
| // partitions. So set this startTime can read half elements. |
| int startTime = numElements / 20 / 2; |
| int maxNumRecords = numElements / 2; |
| |
| PCollection<Long> input = |
| p.apply( |
| mkKafkaReadTransform(numElements, maxNumRecords, new ValueAsTimestampFn()) |
| .withStartReadTime(new Instant(startTime)) |
| .withoutMetadata()) |
| .apply(Values.create()); |
| |
| addCountingAsserts(input, maxNumRecords, maxNumRecords, maxNumRecords, numElements - 1); |
| p.run(); |
| } |
| |
| @Rule public ExpectedException noMessagesException = ExpectedException.none(); |
| |
| @Test |
| public void testUnboundedSourceStartReadTimeException() { |
| |
| assumeTrue(new ConsumerSpEL().hasOffsetsForTimes()); |
| |
| noMessagesException.expect(RuntimeException.class); |
| |
| int numElements = 1000; |
| // In this MockConsumer, we let the elements of the time and offset equal and there are 20 |
| // partitions. So set this startTime can not read any element. |
| int startTime = numElements / 20; |
| |
| p.apply( |
| mkKafkaReadTransform(numElements, numElements, new ValueAsTimestampFn()) |
| .withStartReadTime(new Instant(startTime)) |
| .withoutMetadata()) |
| .apply(Values.create()); |
| |
| p.run(); |
| } |
| |
| @Test |
| public void testSourceDisplayData() { |
| KafkaIO.Read<Integer, Long> read = mkKafkaReadTransform(10, null); |
| |
| DisplayData displayData = DisplayData.from(read); |
| |
| assertThat(displayData, hasDisplayItem("topics", "topic_a,topic_b")); |
| assertThat(displayData, hasDisplayItem("enable.auto.commit", false)); |
| assertThat(displayData, hasDisplayItem("bootstrap.servers", "myServer1:9092,myServer2:9092")); |
| assertThat(displayData, hasDisplayItem("auto.offset.reset", "latest")); |
| assertThat(displayData, hasDisplayItem("receive.buffer.bytes", 524288)); |
| } |
| |
| @Test |
| public void testSourceWithExplicitPartitionsDisplayData() { |
| KafkaIO.Read<byte[], byte[]> read = |
| KafkaIO.readBytes() |
| .withBootstrapServers("myServer1:9092,myServer2:9092") |
| .withTopicPartitions( |
| ImmutableList.of(new TopicPartition("test", 5), new TopicPartition("test", 6))) |
| .withConsumerFactoryFn( |
| new ConsumerFactoryFn( |
| Lists.newArrayList("test"), |
| 10, |
| 10, |
| OffsetResetStrategy.EARLIEST)); // 10 partitions |
| |
| DisplayData displayData = DisplayData.from(read); |
| |
| assertThat(displayData, hasDisplayItem("topicPartitions", "test-5,test-6")); |
| assertThat(displayData, hasDisplayItem("enable.auto.commit", false)); |
| assertThat(displayData, hasDisplayItem("bootstrap.servers", "myServer1:9092,myServer2:9092")); |
| assertThat(displayData, hasDisplayItem("auto.offset.reset", "latest")); |
| assertThat(displayData, hasDisplayItem("receive.buffer.bytes", 524288)); |
| } |
| |
| @Test |
| public void testSinkDisplayData() { |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| KafkaIO.Write<Integer, Long> write = |
| KafkaIO.<Integer, Long>write() |
| .withBootstrapServers("myServerA:9092,myServerB:9092") |
| .withTopic("myTopic") |
| .withValueSerializer(LongSerializer.class) |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)) |
| .withProducerConfigUpdates(ImmutableMap.of("retry.backoff.ms", 100)); |
| |
| DisplayData displayData = DisplayData.from(write); |
| |
| assertThat(displayData, hasDisplayItem("topic", "myTopic")); |
| assertThat(displayData, hasDisplayItem("bootstrap.servers", "myServerA:9092,myServerB:9092")); |
| assertThat(displayData, hasDisplayItem("retries", 3)); |
| assertThat(displayData, hasDisplayItem("retry.backoff.ms", 100)); |
| } |
| } |
| |
| // interface for testing coder inference |
| private interface DummyInterface<T> {} |
| |
| // interface for testing coder inference |
| private interface DummyNonparametricInterface {} |
| |
| // class for testing coder inference |
| private static class DeserializerWithInterfaces |
| implements DummyInterface<String>, DummyNonparametricInterface, Deserializer<Long> { |
| |
| @Override |
| public void configure(Map<String, ?> configs, boolean isKey) {} |
| |
| @Override |
| public Long deserialize(String topic, byte[] bytes) { |
| return 0L; |
| } |
| |
| @Override |
| public void close() {} |
| } |
| |
| // class for which a coder cannot be infered |
| private static class NonInferableObject {} |
| |
| // class for testing coder inference |
| private static class NonInferableObjectDeserializer implements Deserializer<NonInferableObject> { |
| |
| @Override |
| public void configure(Map<String, ?> configs, boolean isKey) {} |
| |
| @Override |
| public NonInferableObject deserialize(String topic, byte[] bytes) { |
| return new NonInferableObject(); |
| } |
| |
| @Override |
| public void close() {} |
| } |
| |
| @Test |
| public void testInferKeyCoder() { |
| CoderRegistry registry = CoderRegistry.createDefault(); |
| |
| assertTrue( |
| KafkaIO.inferCoder(registry, LongDeserializer.class).getValueCoder() |
| instanceof VarLongCoder); |
| |
| assertTrue( |
| KafkaIO.inferCoder(registry, StringDeserializer.class).getValueCoder() |
| instanceof StringUtf8Coder); |
| |
| assertTrue( |
| KafkaIO.inferCoder(registry, InstantDeserializer.class).getValueCoder() |
| instanceof InstantCoder); |
| |
| assertTrue( |
| KafkaIO.inferCoder(registry, DeserializerWithInterfaces.class).getValueCoder() |
| instanceof VarLongCoder); |
| } |
| |
| @Rule public ExpectedException cannotInferException = ExpectedException.none(); |
| |
| @Test |
| public void testInferKeyCoderFailure() throws Exception { |
| cannotInferException.expect(RuntimeException.class); |
| |
| CoderRegistry registry = CoderRegistry.createDefault(); |
| KafkaIO.inferCoder(registry, NonInferableObjectDeserializer.class); |
| } |
| |
| @Test |
| public void testSinkMetrics() throws Exception { |
| // Simply read from kafka source and write to kafka sink. Then verify the metrics are reported. |
| |
| int numElements = 1000; |
| |
| try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { |
| |
| ProducerSendCompletionThread completionThread = |
| new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); |
| |
| String topic = "test"; |
| |
| p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) |
| .apply( |
| "writeToKafka", |
| KafkaIO.<Integer, Long>write() |
| .withBootstrapServers("none") |
| .withTopic(topic) |
| .withKeySerializer(IntegerSerializer.class) |
| .withValueSerializer(LongSerializer.class) |
| .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); |
| |
| PipelineResult result = p.run(); |
| |
| MetricName elementsWritten = SinkMetrics.elementsWritten().getName(); |
| |
| MetricQueryResults metrics = |
| result |
| .metrics() |
| .queryMetrics( |
| MetricsFilter.builder() |
| .addNameFilter(MetricNameFilter.inNamespace(elementsWritten.getNamespace())) |
| .build()); |
| |
| assertThat( |
| metrics.getCounters(), |
| hasItem( |
| attemptedMetricsResult( |
| elementsWritten.getNamespace(), |
| elementsWritten.getName(), |
| "writeToKafka", |
| 1000L))); |
| |
| completionThread.shutdown(); |
| } |
| } |
| |
| @Test |
| public void testOffsetConsumerConfigOverrides() throws Exception { |
| KafkaUnboundedReader reader1 = |
| new KafkaUnboundedReader( |
| new KafkaUnboundedSource( |
| KafkaIO.read() |
| .withBootstrapServers("broker_1:9092,broker_2:9092") |
| .withTopic("my_topic") |
| .withOffsetConsumerConfigOverrides(null), |
| 0), |
| null); |
| assertTrue( |
| reader1 |
| .getOffsetConsumerConfig() |
| .get(ConsumerConfig.GROUP_ID_CONFIG) |
| .toString() |
| .matches(".*_offset_consumer_\\d+_none")); |
| assertEquals( |
| false, reader1.getOffsetConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG)); |
| assertEquals( |
| "read_uncommitted", |
| reader1.getOffsetConsumerConfig().get(ConsumerConfig.ISOLATION_LEVEL_CONFIG)); |
| |
| String offsetGroupId = "group.offsetConsumer"; |
| KafkaUnboundedReader reader2 = |
| new KafkaUnboundedReader( |
| new KafkaUnboundedSource( |
| KafkaIO.read() |
| .withBootstrapServers("broker_1:9092,broker_2:9092") |
| .withTopic("my_topic") |
| .withOffsetConsumerConfigOverrides( |
| ImmutableMap.of(ConsumerConfig.GROUP_ID_CONFIG, offsetGroupId)), |
| 0), |
| null); |
| assertEquals( |
| offsetGroupId, reader2.getOffsetConsumerConfig().get(ConsumerConfig.GROUP_ID_CONFIG)); |
| assertEquals( |
| false, reader2.getOffsetConsumerConfig().get(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG)); |
| assertEquals( |
| "read_uncommitted", |
| reader2.getOffsetConsumerConfig().get(ConsumerConfig.ISOLATION_LEVEL_CONFIG)); |
| } |
| |
| private static void verifyProducerRecords( |
| MockProducer<Integer, Long> mockProducer, |
| String topic, |
| int numElements, |
| boolean keyIsAbsent, |
| boolean verifyTimestamp) { |
| |
| // verify that appropriate messages are written to kafka |
| List<ProducerRecord<Integer, Long>> sent = mockProducer.history(); |
| |
| // sort by values |
| sent.sort(Comparator.comparingLong(ProducerRecord::value)); |
| |
| for (int i = 0; i < numElements; i++) { |
| ProducerRecord<Integer, Long> record = sent.get(i); |
| assertEquals(topic, record.topic()); |
| if (keyIsAbsent) { |
| assertNull(record.key()); |
| } else { |
| assertEquals(i, record.key().intValue()); |
| } |
| assertEquals(i, record.value().longValue()); |
| if (verifyTimestamp) { |
| assertEquals(i, record.timestamp().intValue()); |
| } |
| } |
| } |
| |
| /** |
| * This wrapper over MockProducer. It also places the mock producer in global MOCK_PRODUCER_MAP. |
| * The map is needed so that the producer returned by ProducerFactoryFn during pipeline can be |
| * used in verification after the test. We also override {@code flush()} method in MockProducer so |
| * that test can control behavior of {@code send()} method (e.g. to inject errors). |
| */ |
| private static class MockProducerWrapper implements AutoCloseable { |
| |
| final String producerKey; |
| final MockProducer<Integer, Long> mockProducer; |
| |
| // MockProducer has "closed" method starting version 0.11. |
| private static Method closedMethod; |
| |
| static { |
| try { |
| closedMethod = MockProducer.class.getMethod("closed"); |
| } catch (NoSuchMethodException e) { |
| closedMethod = null; |
| } |
| } |
| |
| MockProducerWrapper() { |
| producerKey = String.valueOf(ThreadLocalRandom.current().nextLong()); |
| mockProducer = |
| new MockProducer<Integer, Long>( |
| false, // disable synchronous completion of send. see ProducerSendCompletionThread |
| // below. |
| new IntegerSerializer(), |
| new LongSerializer()) { |
| |
| // override flush() so that it does not complete all the waiting sends, giving a chance |
| // to |
| // ProducerCompletionThread to inject errors. |
| |
| @Override |
| public synchronized void flush() { |
| while (completeNext()) { |
| // there are some uncompleted records. let the completion thread handle them. |
| try { |
| Thread.sleep(10); |
| } catch (InterruptedException e) { |
| // ok to retry. |
| } |
| } |
| } |
| }; |
| |
| // Add the producer to the global map so that producer factory function can access it. |
| assertNull(MOCK_PRODUCER_MAP.putIfAbsent(producerKey, mockProducer)); |
| } |
| |
| @Override |
| public void close() { |
| MOCK_PRODUCER_MAP.remove(producerKey); |
| try { |
| if (closedMethod == null || !((Boolean) closedMethod.invoke(mockProducer))) { |
| mockProducer.close(); |
| } |
| } catch (Exception e) { // Not expected. |
| throw new RuntimeException(e); |
| } |
| } |
| } |
| |
| private static final ConcurrentMap<String, MockProducer<Integer, Long>> MOCK_PRODUCER_MAP = |
| new ConcurrentHashMap<>(); |
| |
| private static class ProducerFactoryFn |
| implements SerializableFunction<Map<String, Object>, Producer<Integer, Long>> { |
| final String producerKey; |
| |
| ProducerFactoryFn(String producerKey) { |
| this.producerKey = producerKey; |
| } |
| |
| @SuppressWarnings("unchecked") |
| @Override |
| public Producer<Integer, Long> apply(Map<String, Object> config) { |
| |
| // Make sure the config is correctly set up for serializers. |
| Utils.newInstance( |
| (Class<? extends Serializer<?>>) |
| ((Class<?>) config.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG)) |
| .asSubclass(Serializer.class)) |
| .configure(config, true); |
| |
| Utils.newInstance( |
| (Class<? extends Serializer<?>>) |
| ((Class<?>) config.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG)) |
| .asSubclass(Serializer.class)) |
| .configure(config, false); |
| |
| // Returning same producer in each instance in a pipeline seems to work fine currently. |
| // If DirectRunner creates multiple DoFn instances for sinks, we might need to handle |
| // it appropriately. I.e. allow multiple producers for each producerKey and concatenate |
| // all the messages written to each producer for verification after the pipeline finishes. |
| |
| return MOCK_PRODUCER_MAP.get(producerKey); |
| } |
| } |
| |
| private static class InjectedErrorException extends RuntimeException { |
| InjectedErrorException(String message) { |
| super(message); |
| } |
| } |
| |
| /** |
| * We start MockProducer with auto-completion disabled. That implies a record is not marked sent |
| * until #completeNext() is called on it. This class starts a thread to asynchronously 'complete' |
| * the the sends. During completion, we can also make those requests fail. This error injection is |
| * used in one of the tests. |
| */ |
| private static class ProducerSendCompletionThread { |
| |
| private final MockProducer<Integer, Long> mockProducer; |
| private final int maxErrors; |
| private final int errorFrequency; |
| private final AtomicBoolean done = new AtomicBoolean(false); |
| private final ExecutorService injectorThread; |
| private int numCompletions = 0; |
| |
| ProducerSendCompletionThread(MockProducer<Integer, Long> mockProducer) { |
| // complete everything successfully |
| this(mockProducer, 0, 0); |
| } |
| |
| ProducerSendCompletionThread( |
| MockProducer<Integer, Long> mockProducer, int maxErrors, int errorFrequency) { |
| this.mockProducer = mockProducer; |
| this.maxErrors = maxErrors; |
| this.errorFrequency = errorFrequency; |
| injectorThread = Executors.newSingleThreadExecutor(); |
| } |
| |
| @SuppressWarnings("FutureReturnValueIgnored") |
| ProducerSendCompletionThread start() { |
| injectorThread.submit( |
| () -> { |
| int errorsInjected = 0; |
| |
| while (!done.get()) { |
| boolean successful; |
| |
| if (errorsInjected < maxErrors && ((numCompletions + 1) % errorFrequency) == 0) { |
| successful = |
| mockProducer.errorNext( |
| new InjectedErrorException("Injected Error #" + (errorsInjected + 1))); |
| |
| if (successful) { |
| errorsInjected++; |
| } |
| } else { |
| successful = mockProducer.completeNext(); |
| } |
| |
| if (successful) { |
| numCompletions++; |
| } else { |
| // wait a bit since there are no unsent records |
| try { |
| Thread.sleep(1); |
| } catch (InterruptedException e) { |
| // ok to retry. |
| } |
| } |
| } |
| }); |
| |
| return this; |
| } |
| |
| void shutdown() { |
| done.set(true); |
| injectorThread.shutdown(); |
| try { |
| assertTrue(injectorThread.awaitTermination(10, TimeUnit.SECONDS)); |
| } catch (InterruptedException e) { |
| Thread.currentThread().interrupt(); |
| throw new RuntimeException(e); |
| } |
| } |
| } |
| } |
