You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by re...@apache.org on 2017/10/17 15:17:36 UTC
[01/13] beam git commit: Kafka exactly-once sink. Tested manually
with direct runner and on dataflow.
Repository: beam
Updated Branches:
refs/heads/master 9524cbd88 -> 0d2a39bd8
Kafka exactly-once sink.
Tested manually with direct runner and on dataflow.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/d745cc9d
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/d745cc9d
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/d745cc9d
Branch: refs/heads/master
Commit: d745cc9d8cc1735d3bc3c67ba3e2617cb7f11a8c
Parents: 9524cbd
Author: Raghu Angadi <ra...@google.com>
Authored: Thu Jul 20 23:26:49 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:01:22 2017 -0700
----------------------------------------------------------------------
sdks/java/io/kafka/pom.xml | 12 +-
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 563 ++++++++++++++++++-
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 9 +-
3 files changed, 579 insertions(+), 5 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/d745cc9d/sdks/java/io/kafka/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/pom.xml b/sdks/java/io/kafka/pom.xml
index a94d958..f1ddb51 100644
--- a/sdks/java/io/kafka/pom.xml
+++ b/sdks/java/io/kafka/pom.xml
@@ -80,7 +80,17 @@
<artifactId>auto-value</artifactId>
<scope>provided</scope>
</dependency>
-
+
+ <dependency>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-annotations</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-databind</artifactId>
+ </dependency>
+
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-expression</artifactId>
http://git-wip-us.apache.org/repos/asf/beam/blob/d745cc9d/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index aab99c3..5813203 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -21,10 +21,13 @@ import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.auto.value.AutoValue;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.Joiner;
+import com.google.common.base.MoreObjects;
import com.google.common.collect.ComparisonChain;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
@@ -59,6 +62,7 @@ import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.AvroCoder;
+import org.apache.beam.sdk.coders.BigEndianLongCoder;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
@@ -75,14 +79,23 @@ import org.apache.beam.sdk.metrics.SinkMetrics;
import org.apache.beam.sdk.metrics.SourceMetrics;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.ValueProvider;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.windowing.AfterPane;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
+import org.apache.beam.sdk.transforms.windowing.Repeatedly;
+import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
@@ -92,21 +105,29 @@ import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.producer.Callback;
import org.apache.kafka.clients.producer.KafkaProducer;
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.clients.producer.RecordMetadata;
+import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.AuthorizationException;
+import org.apache.kafka.common.errors.OutOfOrderSequenceException;
+import org.apache.kafka.common.errors.ProducerFencedException;
import org.apache.kafka.common.errors.WakeupException;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.AppInfoParser;
+import org.joda.time.DateTimeUtils;
+import org.joda.time.DateTimeZone;
import org.joda.time.Duration;
import org.joda.time.Instant;
+import org.joda.time.format.DateTimeFormat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -283,6 +304,9 @@ public class KafkaIO {
public static <K, V> Write<K, V> write() {
return new AutoValue_KafkaIO_Write.Builder<K, V>()
.setProducerConfig(Write.DEFAULT_PRODUCER_PROPERTIES)
+ .setEOS(false)
+ .setNumShards(0)
+ .setConsumerFactoryFn(Read.KAFKA_CONSUMER_FACTORY_FN)
.build();
}
@@ -1407,6 +1431,15 @@ public class KafkaIO {
@Nullable abstract Class<? extends Serializer<K>> getKeySerializer();
@Nullable abstract Class<? extends Serializer<V>> getValueSerializer();
+ // Configuration for EOS sink
+ abstract boolean isEOS();
+ @Nullable abstract String getSinkGroupId();
+ abstract int getNumShards();
+ @Nullable abstract Coder<K> getKeyCoder();
+ @Nullable abstract Coder<V> getValueCoder();
+ @Nullable abstract
+ SerializableFunction<Map<String, Object>, ? extends Consumer<?, ?>> getConsumerFactoryFn();
+
abstract Builder<K, V> toBuilder();
@AutoValue.Builder
@@ -1417,6 +1450,13 @@ public class KafkaIO {
SerializableFunction<Map<String, Object>, Producer<K, V>> fn);
abstract Builder<K, V> setKeySerializer(Class<? extends Serializer<K>> serializer);
abstract Builder<K, V> setValueSerializer(Class<? extends Serializer<V>> serializer);
+ abstract Builder<K, V> setEOS(boolean eosEnabled);
+ abstract Builder<K, V> setSinkGroupId(String sinkGroupId);
+ abstract Builder<K, V> setNumShards(int numShards);
+ abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
+ abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
+ abstract Builder<K, V> setConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, ? extends Consumer<?, ?>> fn);
abstract Write<K, V> build();
}
@@ -1472,6 +1512,35 @@ public class KafkaIO {
return toBuilder().setProducerFactoryFn(producerFactoryFn).build();
}
+ public Write<K, V> withEOS() {
+ return toBuilder().setEOS(true).build();
+ }
+
+ /**
+ * Should be unique for the job. This is also used in naming the producers used in EOS sink.
+ * TODO: expand javaDoc.
+ */
+ public Write<K, V> withSinkGroupId(String sinkGroupId) {
+ return toBuilder().setSinkGroupId(sinkGroupId).build();
+ }
+
+ public Write<K, V> withNumShards(int numShards) {
+ return toBuilder().setNumShards(numShards).build();
+ }
+
+ public Write<K, V> withKeyCoder(Coder<K> keyCoder) {
+ return toBuilder().setKeyCoder(keyCoder).build();
+ }
+
+ public Write<K, V> withValueCoder(Coder<V> valueCoder) {
+ return toBuilder().setValueCoder(valueCoder).build();
+ }
+
+ public Write<K, V> withConsumerFactoryFn(
+ SerializableFunction<Map<String, Object>, ? extends Consumer<?, ?>> consumerFactoryFn) {
+ return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
+ }
+
/**
* Writes just the values to Kafka. This is useful for writing collections of values rather
* thank {@link KV}s.
@@ -1483,11 +1552,23 @@ public class KafkaIO {
@Override
public PDone expand(PCollection<KV<K, V>> input) {
checkArgument(
- getProducerConfig().get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG) != null,
- "withBootstrapServers() is required");
+ getProducerConfig().get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG) != null,
+ "withBootstrapServers() is required");
checkArgument(getTopic() != null, "withTopic() is required");
- input.apply(ParDo.of(new KafkaWriter<>(this)));
+ if (isEOS()) {
+ EOSWrite.ensureEOSSupport();
+ checkArgument(getSinkGroupId() != null,
+ "withSinkGroupId() is required for exactly-once sink");
+
+ // TODO: Verify that the group_id does not have existing state stored on Kafka unless
+ // this is an upgrade. This avoids issues with simple mistake of reusing group_id
+ // across multiple runs or across multiple jobs.
+
+ input.apply(new EOSWrite<>(this));
+ } else {
+ input.apply(ParDo.of(new KafkaWriter<>(this)));
+ }
return PDone.in(input.getPipeline());
}
@@ -1552,6 +1633,7 @@ public class KafkaIO {
}
}
+
private static class NullOnlyCoder<T> extends AtomicCoder<T> {
@Override
public void encode(T value, OutputStream outStream) {
@@ -1696,4 +1778,479 @@ public class KafkaIO {
throw new RuntimeException(String.format(
"Could not extract the Kafka Deserializer type from %s", deserializer));
}
+
+ ////////////////////////////////// Exactly-Once Sink \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+
+ /**
+ * Exactly-once sink transform.
+ */
+ private static class EOSWrite<K, V> extends PTransform<PCollection<KV<K, V>>, PCollection<Void>> {
+ //
+ // Dataflow ensures at-least once processing for side effects like sinks. In order to provide
+ // exactly-once semantics, a sink needs to be idempotent or it should avoid writing records
+ // that have already been written. This snk does the latter. All the the records are ordered
+ // across a fixed number of shards and records in each shard are written in order. It drops
+ // any records that are already written and buffers those arriving out of order.
+ //
+ // // Exactly once sink involves two shuffles of the records:
+ // A -- GBK --> B -- GBK --> C
+ //
+ // Processing guarantees also require deterministic processing within user transforms.
+ // in this case that implies the order of the records seen by C should not be affected by
+ // restarts in upstream stages link B & A.
+ //
+ // A : Assigns a random shard for message. Note that there are no ordering guarantees for
+ // writing user records to Kafka. User can still control partitioning among topic
+ // partitions as with regular sink (of course, there are no ordering guarantees in
+ // regular Kafka sink either).
+ // B : Assigns an id sequentially for each messages within a shard.
+ // C : Writes each shard to Kafka in sequential id order. In Dataflow, when C sees a record
+ // and id, it implies that record and the associated id are checkpointed to persistent
+ // storage and this record will always have same id, even in retries.
+ // Exactly-once semantics are achieved by writing records in the strict order of
+ // these checkpointed sequence ids.
+ //
+ // Parallelism for B and C is fixed to 'numShards', which defaults to number of partitions
+ // for the topic. A few reasons for that:
+ // - B & C implement their functionality using per-key state. Shard id makes it independent
+ // of cardinality of user key.
+ // - We create one producer per shard, and its 'transactional id' is based on shard id. This
+ // requires that number of shards to be finite. This also helps with batching. and avoids
+ // initializing producers and transactions.
+ // - Most importantly, each of sharded writers stores 'next message id' in partition
+ // metadata, which is committed atomically with Kafka transactions. This is critical
+ // to handle retries of C correctly. Initial testing showed number of shards could be
+ // larger than number of partitions for the topic.
+ //
+ // Number of shards can change across multiple runs of a pipeline (job upgrade in Dataflow).
+ //
+
+ private final Write<K, V> spec;
+
+ EOSWrite(Write<K, V> spec) {
+ this.spec = spec;
+ }
+
+ @Override
+ public PCollection<Void> expand(PCollection<KV<K, V>> input) {
+ return input
+ .apply(Window.<KV<K, V>>into(new GlobalWindows()) // Everything into global window.
+ .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
+ .discardingFiredPanes())
+ .apply(String.format("Shuffle across %d shards", spec.getNumShards()),
+ ParDo.of(new EOSReshard<K, V>(DateTimeUtils.currentTimeMillis(),
+ spec.getNumShards())))
+ .apply("Persist sharding", GroupByKey.<Integer, KV<K, V>>create())
+ .apply("Assign sequential ids", ParDo.of(new EOSSequencer<K, V>()))
+ .apply("Persist ids", GroupByKey.<Integer, KV<Long, KV<K, V>>>create())
+ .apply(String.format("Write to Kafka topic '%s'", spec.getTopic()),
+ ParDo.of(new KafkaEOWriter<>(spec)));
+
+ // TODO: add metrics.
+ }
+ }
+
+ /**
+ * Shuffle messages assigning each randomly to a shard.
+ */
+ private static class EOSReshard<K, V> extends DoFn<KV<K, V>, KV<Integer, KV<K, V>>> {
+ private final Random random;
+ private final int numShards;
+
+ EOSReshard(long seed, int numShards) {
+ this.random = new Random(seed);
+ this.numShards = numShards;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext ctx) {
+ ctx.output(KV.of(random.nextInt(numShards), ctx.element()));
+ }
+ }
+
+ private static class EOSSequencer<K, V>
+ extends DoFn<KV<Integer, Iterable<KV<K, V>>>, KV<Integer, KV<Long, KV<K, V>>>> {
+ private static final String NEXT_ID = "nextId";
+ @StateId(NEXT_ID)
+ private final StateSpec<ValueState<Long>> nextIdSpec = StateSpecs.value();
+
+ @ProcessElement
+ public void processElement(@StateId(NEXT_ID) ValueState<Long> nextIdState, ProcessContext ctx) {
+ long nextId = MoreObjects.firstNonNull(nextIdState.read(), 0L);
+ int shard = ctx.element().getKey();
+ for (KV<K, V> value : ctx.element().getValue()) {
+ ctx.output(KV.of(shard, KV.of(nextId, value)));
+ nextId++;
+ }
+ nextIdState.write(nextId);
+ }
+ }
+
+ private static class KafkaEOWriter<K, V>
+ extends DoFn<KV<Integer, Iterable<KV<Long, KV<K, V>>>>, Void> {
+
+ private static final String NEXT_ID = "nextId";
+ private static final String MIN_BUFFERED_ID = "minBufferedId";
+ private static final String OUT_OF_ORDER_BUFFER = "outOfOrderBuffer";
+ private static final String WRITER_ID = "writerId";
+
+ private static final ObjectMapper JSON_MAPPER = new ObjectMapper();
+
+ @StateId(NEXT_ID)
+ private final StateSpec<ValueState<Long>> sequenceIdSpec = StateSpecs.value();
+ @StateId(MIN_BUFFERED_ID)
+ private final StateSpec<ValueState<Long>> minBufferedId = StateSpecs.value();
+ @StateId(OUT_OF_ORDER_BUFFER)
+ private final StateSpec<BagState<KV<Long, KV<K, V>>>> outOfOrderBuffer;
+ // A random id assigned to each shard.
+ // Helps with detecting when multiple jobs are mistakenly started with same 'groupId' used for
+ // storing state on Kafka side. This also include the case where a job is restarted with same
+ // groupId, but the state is not explicitly cleared.
+ // Better to be safe and error out with a clear message.
+ @StateId(WRITER_ID)
+ private final StateSpec<ValueState<String>> writerIdSpec = StateSpecs.value();
+
+ private final Write<K, V> spec;
+
+ // One cache for each sink (usually there is only one sink per pipeline
+ private static final Map<String, Map<Integer, ShardWriter<?, ?>>> CACHE_BY_GROUP_ID =
+ new HashMap<>();
+ // TODO: Need a way to close producers that are no longer relevant (may be have a timeout?).
+
+ // This would mainly matter only in batch pipelines. Could be configurable.
+ private static final int MAX_RECORDS_PER_TXN = 1000;
+
+
+ KafkaEOWriter(Write<K, V> spec) {
+ this.spec = spec;
+ this.outOfOrderBuffer = StateSpecs.bag(KvCoder.of(
+ BigEndianLongCoder.of(), KvCoder.of(spec.getKeyCoder(), spec.getValueCoder())));
+ }
+
+ @ProcessElement
+ public void processElement(@StateId(NEXT_ID) ValueState<Long> nextIdState,
+ @StateId(MIN_BUFFERED_ID) ValueState<Long> minBufferedIdState,
+ @StateId(OUT_OF_ORDER_BUFFER)
+ BagState<KV<Long, KV<K, V>>> oooBufferState,
+ @StateId(WRITER_ID) ValueState<String> writerIdState,
+ ProcessContext ctx)
+ throws IOException {
+
+ int shard = ctx.element().getKey();
+
+ minBufferedIdState.readLater();
+ long nextId = MoreObjects.firstNonNull(nextIdState.read(), 0L);
+ long minBufferedId = MoreObjects.firstNonNull(minBufferedIdState.read(), Long.MAX_VALUE);
+
+ ShardWriter<K, V> writer = getShardWriter(shard, writerIdState, nextId);
+ long committedId = writer.committedId;
+
+ if (committedId >= nextId) {
+ // This is a retry of an already committed batch.
+ LOG.info("{}: committed id {} is ahead of expected {}. {} records will be dropped "
+ + "(these are already written).",
+ shard, committedId, nextId - 1, committedId - nextId + 1);
+ nextId = committedId + 1;
+ }
+
+ try {
+ writer.beginTxn();
+ int txnSize = 0;
+
+ // Iterate in recordId order. The input iterator could be mostly sorted.
+ // There might be out of order messages buffered in earlier iterations. These
+ // will get merged if and when minBufferedId matches nextId.
+
+ Iterator<KV<Long, KV<K, V>>> iter = ctx.element().getValue().iterator();
+
+ while (iter.hasNext()) {
+ KV<Long, KV<K, V>> kv = iter.next();
+ long recordId = kv.getKey();
+
+ if (recordId < nextId) {
+ LOG.info("{}: dropping older record {}. Already committed till {}",
+ shard, recordId, committedId);
+ continue;
+ }
+
+ if (recordId > nextId) {
+ // Out of order delivery. Should be pretty rare (what about in a batch pipeline?)
+
+ LOG.info("{}: Saving out of order record {}, next record id to be written is {}",
+ shard, recordId, nextId);
+
+ // checkState(recordId - nextId < 10000, "records are way out of order");
+
+ oooBufferState.add(kv);
+ minBufferedId = Math.min(minBufferedId, recordId);
+ minBufferedIdState.write(minBufferedId);
+ continue;
+ }
+
+ // recordId and nextId match. Finally write record.
+
+ writer.sendRecord(kv.getValue());
+ nextId++;
+
+ if (++txnSize >= MAX_RECORDS_PER_TXN) {
+ writer.commitTxn(recordId);
+ txnSize = 0;
+ writer.beginTxn();
+ }
+
+ if (minBufferedId == nextId) {
+ // One or more of the buffered records can be committed now.
+ // Read all the buffered records in to memory and sort them. Reading into memory
+ // might be problematic in extreme cases. Might need to improve it in future.
+
+ List<KV<Long, KV<K, V>>> buffered = Lists.newArrayList(oooBufferState.read());
+ Collections.sort(buffered, new KV.OrderByKey<Long, KV<K, V>>());
+
+ LOG.info("{} : merging {} buffered records with min buffered id",
+ shard, buffered.size(), minBufferedId);
+
+ oooBufferState.clear();
+ minBufferedIdState.clear();
+ minBufferedId = Long.MAX_VALUE;
+
+ iter = Iterators.mergeSorted(ImmutableList.of(iter, buffered.iterator()),
+ new KV.OrderByKey<Long, KV<K, V>>());
+ }
+ }
+
+ writer.commitTxn(nextId - 1);
+ nextIdState.write(nextId);
+ } catch (ProducerFencedException | OutOfOrderSequenceException | AuthorizationException e) {
+ // JavaDoc says these are not recoverable errors and producer should be closed.
+
+ // Close the producer and and remove it from the cache. A new producer will be created
+ // in retry. It is possible that a rough worker keeps retrying and ends up fencing off
+ // active producers. How likely this might be or how well such a scenario is handled
+ // depends on the runner. For now we will leave it to upper layers, will need to revisit.
+
+ LOG.warn("{} : closing producer {} after unrecoverable error. The work might be migrated. "
+ + "committed id {}, current id {}.",
+ writer.shard, writer.producerName, writer.committedId, nextId - 1, e);
+
+ CACHE_BY_GROUP_ID.get(spec.getSinkGroupId()).remove(writer.shard);
+ writer.producer.close();
+ throw e;
+ }
+ }
+
+ private static class ShardMetadata {
+ @JsonProperty("seq")
+ public final long sequenceId;
+ @JsonProperty("id")
+ public final String writerId;
+
+ private ShardMetadata() { // for json
+ sequenceId = -1;
+ writerId = null;
+ }
+
+ ShardMetadata(long sequenceId, String writerId) {
+ this.sequenceId = sequenceId;
+ this.writerId = writerId;
+ }
+ }
+
+ private static class ShardWriter<K, V> {
+ private final int shard;
+ private final String writerId;
+ private final Producer<K, V> producer;
+ private final String producerName;
+ private final Write<K, V> spec;
+ private long committedId;
+
+ ShardWriter(int shard,
+ String writerId,
+ Producer<K, V> producer,
+ String producerName,
+ Write<K, V> spec,
+ long committedId) {
+ this.shard = shard;
+ this.writerId = writerId;
+ this.producer = producer;
+ this.producerName = producerName;
+ this.spec = spec;
+ this.committedId = committedId;
+ }
+
+ void beginTxn() {
+ producer.beginTransaction();
+ }
+
+ void sendRecord(KV<K, V> record) {
+ try {
+ producer.send(
+ new ProducerRecord<>(spec.getTopic(), record.getKey(), record.getValue()));
+ } catch (KafkaException e) {
+ producer.abortTransaction();
+ throw e;
+ }
+ }
+
+ void commitTxn(long lastRecordId) throws IOException {
+ try {
+ // Store id in consumer group metadata for the partition
+ producer.sendOffsetsToTransaction(
+ ImmutableMap.of(new TopicPartition(spec.getTopic(), shard),
+ new OffsetAndMetadata(
+ Long.MAX_VALUE, // So that consumer group does not expire.
+ JSON_MAPPER.writeValueAsString(new ShardMetadata(lastRecordId,
+ writerId)
+ ))),
+ spec.getSinkGroupId());
+ producer.commitTransaction();
+
+ LOG.info("{} : committed {} records", shard, lastRecordId - committedId);
+
+ committedId = lastRecordId;
+ } catch (KafkaException e) {
+ producer.abortTransaction();
+ throw e;
+ }
+ }
+ }
+
+ private ShardWriter<K, V> getShardWriter(int shard,
+ ValueState<String> writerIdState,
+ long nextId) throws IOException {
+
+ Map<Integer, ShardWriter<?, ?>> cache;
+
+ synchronized (CACHE_BY_GROUP_ID) {
+ cache = CACHE_BY_GROUP_ID.get(spec.getSinkGroupId());
+ if (cache == null) {
+ cache = new HashMap<>();
+ CACHE_BY_GROUP_ID.put(spec.getSinkGroupId(), cache);
+ }
+ }
+
+ synchronized (cache) {
+ @SuppressWarnings("unchecked")
+ ShardWriter<K, V> shardWriter = (ShardWriter<K, V>) cache.get(shard);
+ if (shardWriter != null) {
+ // any sanity checks?
+ return shardWriter;
+ }
+ }
+
+ // initialize new shard
+
+ String producerName = String.format("producer_%d_for_%s", shard, spec.getSinkGroupId());
+ Producer<K, V> producer = initializeEosProducer(spec, producerName);
+
+ String writerId = writerIdState.read();
+
+ // Fetch latest committed metadata for the partition (if any). Checks committed sequence ids.
+ try {
+
+ Consumer<?, ?> consumer = spec.getConsumerFactoryFn().apply((ImmutableMap.of(
+ ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, spec
+ .getProducerConfig().get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),
+ ConsumerConfig.GROUP_ID_CONFIG, spec.getSinkGroupId(),
+ ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class,
+ ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class
+ )));
+
+ OffsetAndMetadata committed;
+ try {
+ committed = consumer.committed(new TopicPartition(spec.getTopic(), shard));
+
+ } finally {
+ consumer.close();
+ }
+
+ long committedSeqId = -1;
+
+ if (committed == null || committed.metadata() == null || committed.metadata().isEmpty()) {
+ checkState(nextId == 0 && writerId == null,
+ "State exists for shard %d (nextId %s, writerId '%s'), but there is no state "
+ + "stored with Kafka topic '%s' group id '%s'",
+ shard, nextId, writerId, spec.getTopic(), spec.getSinkGroupId());
+
+ writerId = String.format("%X - %s",
+ new Random().nextInt(Integer.MAX_VALUE),
+ DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss")
+ .withZone(DateTimeZone.UTC)
+ .print(DateTimeUtils.currentTimeMillis()));
+ writerIdState.write(writerId);
+ LOG.info("Assigned writer id '{}' to shard {}", writerId, shard);
+
+ } else {
+ ShardMetadata metadata = JSON_MAPPER.readValue(committed.metadata(),
+ ShardMetadata.class);
+
+ checkNotNull(metadata.writerId);
+
+ if (writerId == null) {
+ // a) This might be a restart of the job from scratch, in which case metatdata
+ // should be ignored and overwritten with new one.
+ // b) This job might be started with an incorrect group id which is an error.
+ // c) There is an extremely small chance that this is a retry of the first bundle
+ // where metatdate was committed to Kafka but the bundle results were not committed
+ // in Beam, in which case it should be treated as correct metadata.
+ // How can we tell these three cases apart? Be safe and throw an exception.
+ //
+ // We could let users explicitly an option to override the existing metadata.
+ //
+ throw new IllegalStateException(String.format(
+ "Kafka metadata exists for shard %d, but there is no stored state for it. "
+ + "This mostly indicates groupId '%s' is already used else where or in earlier runs. "
+ + "Try another group id. Metadata : '%s'.",
+ shard, spec.getSinkGroupId(), committed.metadata()));
+ }
+
+ checkState(writerId.equals(metadata.writerId),
+ "Writer ids don't match. This is mostly a unintended misuse of groupId('%s')."
+ + "Beam '%s', Kafka '%s'",
+ spec.getSinkGroupId(), writerId, metadata.writerId);
+
+ committedSeqId = metadata.sequenceId;
+
+ checkState(committedSeqId >= (nextId - 1),
+ "Committed sequence id can not be lower than %s, partition metadata : %s",
+ nextId - 1, committed.metadata());
+ }
+
+ ShardWriter<K, V> shardWriter = new ShardWriter<>(shard, writerId, producer, producerName,
+ spec, committedSeqId);
+
+ synchronized (cache) {
+ checkState(cache.get(shard) == null,
+ "Unexpected concurrent execution of shard %s", shard);
+ cache.put(shard, shardWriter);
+ }
+
+ LOG.info("{} : initialized producer {} with committed sequence id {}",
+ shard, producerName, committedSeqId);
+
+ return shardWriter;
+
+ } catch (Exception e) {
+ producer.close();
+ throw e;
+ }
+ }
+ }
+
+ private static <K, V> Producer<K, V> initializeEosProducer(Write<K, V> spec,
+ String producerName) {
+
+ Map<String, Object> producerConfig = new HashMap<>(spec.getProducerConfig());
+ producerConfig.putAll(ImmutableMap.of(
+ ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, spec.getKeySerializer(),
+ ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, spec.getValueSerializer(),
+ ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, true,
+ ProducerConfig.TRANSACTIONAL_ID_CONFIG, producerName));
+
+ Producer<K, V> producer = spec.getProducerFactoryFn() != null
+ ? spec.getProducerFactoryFn().apply((producerConfig))
+ : new KafkaProducer<K, V>(producerConfig);
+
+ producer.initTransactions();
+ return producer;
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/d745cc9d/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 482f5a2..48d4922 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -740,6 +740,13 @@ public class KafkaIOTest {
.withTopic(topic)
.withKeySerializer(IntegerSerializer.class)
.withValueSerializer(LongSerializer.class)
+ // .withEOS()
+ // .withSinkGroupId("test")
+ // .withNumShards(1)
+ // .withKeyCoder(BigEndianIntegerCoder.of())
+ // .withValueCoder(BigEndianLongCoder.of())
+ // .withConsumerFactoryFn(new ConsumerFactoryFn(
+ // Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST))
.withProducerFactoryFn(new ProducerFactoryFn()));
p.run();
@@ -1070,7 +1077,7 @@ public class KafkaIOTest {
/**
* Singleton MockProudcer. Using a singleton here since we need access to the object to fetch
* the actual records published to the producer. This prohibits running the tests using
- * the producer in parallel, but there are only one or two tests.
+ * the producer in parallel. These tests only take a few millis each.
*/
private static final MockProducer<Integer, Long> MOCK_PRODUCER =
new MockProducer<Integer, Long>(
[06/13] beam git commit: Improve how we use MockProducer in tests.
Remove global instance and use a map to use same producer instance inside
beam and test verifier.
Posted by re...@apache.org.
Improve how we use MockProducer in tests.
Remove global instance and use a map to use same producer
instance inside beam and test verifier.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/3bd4085d
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/3bd4085d
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/3bd4085d
Branch: refs/heads/master
Commit: 3bd4085db154241c477f8cc318162530029456fb
Parents: 236484b
Author: Raghu Angadi <ra...@google.com>
Authored: Tue Aug 22 14:37:59 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:05 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 2 +-
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 208 ++++++++++++-------
2 files changed, 139 insertions(+), 71 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/3bd4085d/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 78227a0..1d9560b 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -1811,7 +1811,7 @@ public class KafkaIO {
// and id, it implies that record and the associated id are checkpointed to persistent
// storage and this record will always have same id, even in retries.
// Exactly-once semantics are achieved by writing records in the strict order of
- // these checkpointed sequence ids.
+ // these check-pointed sequence ids.
//
// Parallelism for B and C is fixed to 'numShards', which defaults to number of partitions
// for the topic. A few reasons for that:
http://git-wip-us.apache.org/repos/asf/beam/blob/3bd4085d/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index eaf30d6..88c9d21 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -34,13 +34,17 @@ import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.nio.ByteBuffer;
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.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;
@@ -180,7 +184,7 @@ public class KafkaIOTest {
//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) {
+ public void assign(final Collection<TopicPartition> assigned) {
super.assign(assigned);
assignedPartitions.set(ImmutableList.copyOf(assigned));
for (TopicPartition tp : assigned) {
@@ -724,11 +728,10 @@ public class KafkaIOTest {
int numElements = 1000;
- synchronized (MOCK_PRODUCER_LOCK) {
+ try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) {
- MOCK_PRODUCER.clear();
-
- ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread().start();
+ ProducerSendCompletionThread completionThread =
+ new ProducerSendCompletionThread(producerWrapper.mockProducer).start();
String topic = "test";
@@ -740,32 +743,26 @@ public class KafkaIOTest {
.withTopic(topic)
.withKeySerializer(IntegerSerializer.class)
.withValueSerializer(LongSerializer.class)
- .withEOS()
- .withSinkGroupId("test")
- .withNumShards(1)
- .withConsumerFactoryFn(new ConsumerFactoryFn(
- Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST))
- .withProducerFactoryFn(new ProducerFactoryFn()));
+ .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)));
p.run();
completionThread.shutdown();
- verifyProducerRecords(topic, numElements, false);
+ verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, false);
}
}
- //@Test (TODO: reenable. global MOCK_PRODUCER does not work with 0.11 anymore for multiple tests)
+ @Test
public void testValuesSink() throws Exception {
// similar to testSink(), but use values()' interface.
int numElements = 1000;
- synchronized (MOCK_PRODUCER_LOCK) {
-
- MOCK_PRODUCER.clear();
+ try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) {
- ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread().start();
+ ProducerSendCompletionThread completionThread =
+ new ProducerSendCompletionThread(producerWrapper.mockProducer).start();
String topic = "test";
@@ -777,18 +774,56 @@ public class KafkaIOTest {
.withBootstrapServers("none")
.withTopic(topic)
.withValueSerializer(LongSerializer.class)
- .withProducerFactoryFn(new ProducerFactoryFn())
+ .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))
.values());
p.run();
completionThread.shutdown();
- verifyProducerRecords(topic, numElements, true);
+ verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, true);
+ }
+ }
+
+ @Test
+ public void testEOSink() throws Exception {
+ // similar to testSink(), enables to EOS.
+ // This does not actually test exactly-once-semantics. Mainly exercises the code.
+
+ 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()
+ .withSinkGroupId("test")
+ .withNumShards(1)
+ .withConsumerFactoryFn(new ConsumerFactoryFn(
+ Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST))
+ .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)));
+
+ p.run();
+
+ completionThread.shutdown();
+
+ verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, false);
}
}
- //@Test (TODO: reenable. global MOCK_PRODUCER does not work with 0.11 anymore for multiple tests)
+
+ @Test
public void testSinkWithSendErrors() throws Throwable {
// similar to testSink(), except that up to 10 of the send calls to producer will fail
// asynchronously.
@@ -802,15 +837,13 @@ public class KafkaIOTest {
int numElements = 1000;
- synchronized (MOCK_PRODUCER_LOCK) {
+ try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) {
- MOCK_PRODUCER.clear();
+ ProducerSendCompletionThread completionThreadWithErrors =
+ new ProducerSendCompletionThread(producerWrapper.mockProducer, 10, 100).start();
String topic = "test";
- ProducerSendCompletionThread completionThreadWithErrors =
- new ProducerSendCompletionThread(10, 100).start();
-
p
.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn())
.withoutMetadata())
@@ -819,7 +852,7 @@ public class KafkaIOTest {
.withTopic(topic)
.withKeySerializer(IntegerSerializer.class)
.withValueSerializer(LongSerializer.class)
- .withProducerFactoryFn(new ProducerFactoryFn()));
+ .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)));
try {
p.run();
@@ -910,19 +943,21 @@ public class KafkaIOTest {
assertThat(displayData, hasDisplayItem("receive.buffer.bytes", 524288));
}
- //@Test (TODO: reenable. global MOCK_PRODUCER does not work with 0.11 anymore for multiple tests)
+ @Test
public void testSinkDisplayData() {
- KafkaIO.Write<Integer, Long> write = KafkaIO.<Integer, Long>write()
+ 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());
+ .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey));
- DisplayData displayData = DisplayData.from(write);
+ 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("topic", "myTopic"));
+ assertThat(displayData, hasDisplayItem("bootstrap.servers", "myServerA:9092,myServerB:9092"));
+ assertThat(displayData, hasDisplayItem("retries", 3));
+ }
}
// interface for testing coder inference
@@ -1002,17 +1037,16 @@ public class KafkaIOTest {
KafkaIO.inferCoder(registry, NonInferableObjectDeserializer.class);
}
- // @Test (TODO: reenable. global MOCK_PRODUCER does not work with 0.11 anymore for multiple tests)
+ @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;
- synchronized (MOCK_PRODUCER_LOCK) {
+ try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) {
- MOCK_PRODUCER.clear();
-
- ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread().start();
+ ProducerSendCompletionThread completionThread =
+ new ProducerSendCompletionThread(producerWrapper.mockProducer).start();
String topic = "test";
@@ -1024,7 +1058,7 @@ public class KafkaIOTest {
.withTopic(topic)
.withKeySerializer(IntegerSerializer.class)
.withValueSerializer(LongSerializer.class)
- .withProducerFactoryFn(new ProducerFactoryFn()));
+ .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)));
PipelineResult result = p.run();
@@ -1047,10 +1081,11 @@ public class KafkaIOTest {
}
}
- private static void verifyProducerRecords(String topic, int numElements, boolean keyIsAbsent) {
+ private static void verifyProducerRecords(MockProducer<Integer, Long> mockProducer,
+ String topic, int numElements, boolean keyIsAbsent) {
// verify that appropriate messages are written to kafka
- List<ProducerRecord<Integer, Long>> sent = MOCK_PRODUCER.history();
+ List<ProducerRecord<Integer, Long>> sent = mockProducer.history();
// sort by values
Collections.sort(sent, new Comparator<ProducerRecord<Integer, Long>>() {
@@ -1073,44 +1108,67 @@ public class KafkaIOTest {
}
/**
- * Singleton MockProudcer. Using a singleton here since we need access to the object to fetch
- * the actual records published to the producer. This prohibits running the tests using
- * the producer in parallel. These tests only take a few millis each.
+ * This wrapper over MockProducer 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. inject errors).
*/
- private static final MockProducer<Integer, Long> MOCK_PRODUCER =
- new MockProducer<Integer, Long>(
- false, // disable synchronous completion of send. see ProducerSendCompletionThread below.
- new IntegerSerializer(),
- new LongSerializer()) {
+ private static class MockProducerWrapper implements AutoCloseable {
- // override flush() so that it does not complete all the waiting sends, giving a chance to
- // ProducerCompletionThread to inject errors.
+ final String producerKey;
+ final MockProducer<Integer, Long> mockProducer;
- @Override
- public void flush() {
- while (completeNext()) {
- // there are some uncompleted records. let the completion thread handle them.
- try {
- Thread.sleep(10);
- } catch (InterruptedException e) {
+ 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 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));
+ }
+
+ public void close() {
+ MOCK_PRODUCER_MAP.remove(producerKey);
+ if (!mockProducer.closed()) {
+ mockProducer.close();
}
- };
+ }
+ }
- // use a separate object serialize tests using MOCK_PRODUCER so that we don't interfere
- // with Kafka MockProducer locking itself.
- private static final Object MOCK_PRODUCER_LOCK = new Object();
+ 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.
-
// There may not be a key serializer if we're interested only in values.
if (config.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG) != null) {
Utils.newInstance(
@@ -1124,12 +1182,17 @@ public class KafkaIOTest {
.asSubclass(Serializer.class)
).configure(config, false);
- return MOCK_PRODUCER;
+ // 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 {
- public InjectedErrorException(String message) {
+ InjectedErrorException(String message) {
super(message);
}
}
@@ -1142,18 +1205,22 @@ public class KafkaIOTest {
*/
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() {
+ ProducerSendCompletionThread(MockProducer<Integer, Long> mockProducer) {
// complete everything successfully
- this(0, 0);
+ this(mockProducer, 0, 0);
}
- ProducerSendCompletionThread(final int maxErrors, final int errorFrequency) {
+ ProducerSendCompletionThread(MockProducer<Integer, Long> mockProducer,
+ int maxErrors,
+ int errorFrequency) {
+ this.mockProducer = mockProducer;
this.maxErrors = maxErrors;
this.errorFrequency = errorFrequency;
injectorThread = Executors.newSingleThreadExecutor();
@@ -1169,14 +1236,14 @@ public class KafkaIOTest {
boolean successful;
if (errorsInjected < maxErrors && ((numCompletions + 1) % errorFrequency) == 0) {
- successful = MOCK_PRODUCER.errorNext(
+ successful = mockProducer.errorNext(
new InjectedErrorException("Injected Error #" + (errorsInjected + 1)));
if (successful) {
errorsInjected++;
}
} else {
- successful = MOCK_PRODUCER.completeNext();
+ successful = mockProducer.completeNext();
}
if (successful) {
@@ -1186,6 +1253,7 @@ public class KafkaIOTest {
try {
Thread.sleep(1);
} catch (InterruptedException e) {
+ // ok to retry.
}
}
}
[05/13] beam git commit: assign a random shard id once in setup(),
rather than each time in processElement().
Posted by re...@apache.org.
assign a random shard id once in setup(), rather than each time in processElement().
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/236484b0
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/236484b0
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/236484b0
Branch: refs/heads/master
Commit: 236484b010a8e7b8bb6e6bc60c20ab7fced2b964
Parents: f4f6105
Author: Raghu Angadi <ra...@google.com>
Authored: Tue Aug 8 14:18:54 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:05 2017 -0700
----------------------------------------------------------------------
.../main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java | 12 +++++++++---
1 file changed, 9 insertions(+), 3 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/236484b0/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 63dc734..78227a0 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -1866,15 +1866,21 @@ public class KafkaIO {
*/
private static class EOSReshard<K, V> extends DoFn<KV<K, V>, KV<Integer, KV<K, V>>> {
private final int numShards;
+ private transient int shardId;
EOSReshard(int numShards) {
this.numShards = numShards;
}
+ @Setup
+ public void setup() {
+ shardId = ThreadLocalRandom.current().nextInt(numShards);
+ }
+
@ProcessElement
public void processElement(ProcessContext ctx) {
- int shard = ThreadLocalRandom.current().nextInt(numShards);
- ctx.output(KV.of(shard, ctx.element()));
+ shardId = (shardId + 1) % numShards; // round-robin among shards.
+ ctx.output(KV.of(shardId, ctx.element()));
}
}
@@ -2196,7 +2202,7 @@ public class KafkaIO {
//
throw new IllegalStateException(String.format(
"Kafka metadata exists for shard %s, but there is no stored state for it. "
- + "This mostly indicates groupId '%s' is already used else where or in earlier runs. "
+ + "This mostly indicates groupId '%s' is used else where or in earlier runs. "
+ "Try another group id. Metadata for this shard on Kafka : '%s'",
shard, spec.getSinkGroupId(), committed.metadata()));
}
[07/13] beam git commit: Input PCollection has the coder for
elements. No need to ask for them explicitly.
Posted by re...@apache.org.
Input PCollection has the coder for elements. No need to ask for them explicitly.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/5767156f
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/5767156f
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/5767156f
Branch: refs/heads/master
Commit: 5767156fe7810a2c2156d9644687c5768bc75ad1
Parents: 99f5541
Author: Raghu Angadi <ra...@google.com>
Authored: Thu Jul 27 13:02:26 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:05 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 30 ++++----------------
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 2 --
2 files changed, 6 insertions(+), 26 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/5767156f/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 654a137..5d50cf7 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -1437,8 +1437,6 @@ public class KafkaIO {
abstract boolean isEOS();
@Nullable abstract String getSinkGroupId();
abstract int getNumShards();
- @Nullable abstract Coder<K> getKeyCoder();
- @Nullable abstract Coder<V> getValueCoder();
@Nullable abstract
SerializableFunction<Map<String, Object>, ? extends Consumer<?, ?>> getConsumerFactoryFn();
@@ -1455,8 +1453,6 @@ public class KafkaIO {
abstract Builder<K, V> setEOS(boolean eosEnabled);
abstract Builder<K, V> setSinkGroupId(String sinkGroupId);
abstract Builder<K, V> setNumShards(int numShards);
- abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
- abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
abstract Builder<K, V> setConsumerFactoryFn(
SerializableFunction<Map<String, Object>, ? extends Consumer<?, ?>> fn);
abstract Write<K, V> build();
@@ -1530,14 +1526,6 @@ public class KafkaIO {
return toBuilder().setNumShards(numShards).build();
}
- public Write<K, V> withKeyCoder(Coder<K> keyCoder) {
- return toBuilder().setKeyCoder(keyCoder).build();
- }
-
- public Write<K, V> withValueCoder(Coder<V> valueCoder) {
- return toBuilder().setValueCoder(valueCoder).build();
- }
-
public Write<K, V> withConsumerFactoryFn(
SerializableFunction<Map<String, Object>, ? extends Consumer<?, ?>> consumerFactoryFn) {
return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
@@ -1561,8 +1549,6 @@ public class KafkaIO {
if (isEOS()) {
EOSWrite.ensureEOSSupport();
checkNotNull(getSinkGroupId(), "A group id is required for exactly-once sink");
- checkNotNull(getKeyCoder(), "Key coder is required for exact-once sink");
- checkNotNull(getValueCoder(), "Value coder is required for exactly-once sink");
// TODO: Verify that the group_id does not have existing state stored on Kafka unless
// this is an upgrade. This avoids issues with simple mistake of reusing group_id
@@ -1857,9 +1843,7 @@ public class KafkaIO {
.apply("Assign sequential ids", ParDo.of(new EOSSequencer<K, V>()))
.apply("Persist ids", GroupByKey.<Integer, KV<Long, KV<K, V>>>create())
.apply(String.format("Write to Kafka topic '%s'", spec.getTopic()),
- ParDo.of(new KafkaEOWriter<>(spec)));
-
- // TODO: add metrics.
+ ParDo.of(new KafkaEOWriter<>(spec, input.getCoder())));
}
}
@@ -1918,10 +1902,9 @@ public class KafkaIO {
private final StateSpec<ValueState<Long>> minBufferedId = StateSpecs.value();
@StateId(OUT_OF_ORDER_BUFFER)
private final StateSpec<BagState<KV<Long, KV<K, V>>>> outOfOrderBuffer;
- // A random id assigned to each shard.
- // Helps with detecting when multiple jobs are mistakenly started with same 'groupId' used for
- // storing state on Kafka side. This also include the case where a job is restarted with same
- // groupId, but the state is not explicitly cleared.
+ // A random id assigned to each shard. Helps with detecting when multiple jobs are mistakenly
+ // started with same groupId used for storing state on Kafka side including the case where
+ // a job is restarted with same groupId, but the metadata from previous run is not removed.
// Better to be safe and error out with a clear message.
@StateId(WRITER_ID)
private final StateSpec<ValueState<String>> writerIdSpec = StateSpecs.value();
@@ -1939,10 +1922,9 @@ public class KafkaIO {
private final Counter elementsBuffered = Metrics.counter(METRIC_NAMESPACE, "elementsBuffered");
private final Counter numTransactions = Metrics.counter(METRIC_NAMESPACE, "numTransactions");
- KafkaEOWriter(Write<K, V> spec) {
+ KafkaEOWriter(Write<K, V> spec, Coder<KV<K, V>> elemCoder) {
this.spec = spec;
- this.outOfOrderBuffer = StateSpecs.bag(KvCoder.of(
- BigEndianLongCoder.of(), KvCoder.of(spec.getKeyCoder(), spec.getValueCoder())));
+ this.outOfOrderBuffer = StateSpecs.bag(KvCoder.of(BigEndianLongCoder.of(), elemCoder));
}
@ProcessElement
http://git-wip-us.apache.org/repos/asf/beam/blob/5767156f/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 3a6c974..eaf30d6 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -743,8 +743,6 @@ public class KafkaIOTest {
.withEOS()
.withSinkGroupId("test")
.withNumShards(1)
- .withKeyCoder(BigEndianIntegerCoder.of())
- .withValueCoder(BigEndianLongCoder.of())
.withConsumerFactoryFn(new ConsumerFactoryFn(
Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST))
.withProducerFactoryFn(new ProducerFactoryFn()));
[03/13] beam git commit: Add basic metrics.
Posted by re...@apache.org.
Add basic metrics.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/99f55419
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/99f55419
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/99f55419
Branch: refs/heads/master
Commit: 99f554197ada52a8d628ade64523bd761a61075a
Parents: d9c95f0
Author: Raghu Angadi <ra...@google.com>
Authored: Wed Jul 26 14:29:11 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:04 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 30 +++++++++++++-------
1 file changed, 20 insertions(+), 10 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/99f55419/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 9f5f493..654a137 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -76,6 +76,7 @@ import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;
import org.apache.beam.sdk.io.kafka.KafkaCheckpointMark.PartitionMark;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Gauge;
+import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.metrics.SinkMetrics;
import org.apache.beam.sdk.metrics.SourceMetrics;
import org.apache.beam.sdk.options.PipelineOptions;
@@ -1905,6 +1906,10 @@ public class KafkaIO {
private static final String OUT_OF_ORDER_BUFFER = "outOfOrderBuffer";
private static final String WRITER_ID = "writerId";
+ private static final String METRIC_NAMESPACE = "KafkaEOSink";
+
+ // Not sure of a good limit. This applies only for large bundles.
+ private static final int MAX_RECORDS_PER_TXN = 1000;
private static final ObjectMapper JSON_MAPPER = new ObjectMapper();
@StateId(NEXT_ID)
@@ -1928,9 +1933,11 @@ public class KafkaIO {
new HashMap<>();
// TODO: Need a way to close producers that are no longer relevant (may be have a timeout?).
- // This would mainly matter only in batch pipelines. Could be configurable.
- private static final int MAX_RECORDS_PER_TXN = 1000;
-
+ // Metrics
+ private final Counter elementsWritten = SinkMetrics.elementsWritten();
+ // Elements buffered due to out of order arrivals.
+ private final Counter elementsBuffered = Metrics.counter(METRIC_NAMESPACE, "elementsBuffered");
+ private final Counter numTransactions = Metrics.counter(METRIC_NAMESPACE, "numTransactions");
KafkaEOWriter(Write<K, V> spec) {
this.spec = spec;
@@ -1995,23 +2002,24 @@ public class KafkaIO {
oooBufferState.add(kv);
minBufferedId = Math.min(minBufferedId, recordId);
minBufferedIdState.write(minBufferedId);
+ elementsBuffered.inc();
continue;
}
// recordId and nextId match. Finally write record.
- writer.sendRecord(kv.getValue());
+ writer.sendRecord(kv.getValue(), elementsWritten);
nextId++;
if (++txnSize >= MAX_RECORDS_PER_TXN) {
- writer.commitTxn(recordId);
+ writer.commitTxn(recordId, numTransactions);
txnSize = 0;
writer.beginTxn();
}
if (minBufferedId == nextId) {
// One or more of the buffered records can be committed now.
- // Read all the buffered records in to memory and sort them. Reading into memory
+ // Read all of them in to memory and sort them. Reading into memory
// might be problematic in extreme cases. Might need to improve it in future.
List<KV<Long, KV<K, V>>> buffered = Lists.newArrayList(oooBufferState.read());
@@ -2029,7 +2037,7 @@ public class KafkaIO {
}
}
- writer.commitTxn(nextId - 1);
+ writer.commitTxn(nextId - 1, numTransactions);
nextIdState.write(nextId);
} catch (ProducerFencedException | OutOfOrderSequenceException | AuthorizationException e) {
// JavaDoc says these are not recoverable errors and producer should be closed.
@@ -2055,7 +2063,7 @@ public class KafkaIO {
@JsonProperty("id")
public final String writerId;
- private ShardMetadata() { // for json
+ private ShardMetadata() { // for json deserializer
sequenceId = -1;
writerId = null;
}
@@ -2092,17 +2100,18 @@ public class KafkaIO {
producer.beginTransaction();
}
- void sendRecord(KV<K, V> record) {
+ void sendRecord(KV<K, V> record, Counter sendCounter) {
try {
producer.send(
new ProducerRecord<>(spec.getTopic(), record.getKey(), record.getValue()));
+ sendCounter.inc();
} catch (KafkaException e) {
producer.abortTransaction();
throw e;
}
}
- void commitTxn(long lastRecordId) throws IOException {
+ void commitTxn(long lastRecordId, Counter numTransactions) throws IOException {
try {
// Store id in consumer group metadata for the partition.
// NOTE: Kafka keeps this metadata for 24 hours since the last update. This limits
@@ -2116,6 +2125,7 @@ public class KafkaIO {
spec.getSinkGroupId());
producer.commitTransaction();
+ numTransactions.inc();
LOG.info("{} : committed {} records", shard, lastRecordId - committedId);
committedId = lastRecordId;
[02/13] beam git commit: Disable other sink tests. Need to fix how
tests use global MOCK_PRODUCER. A few more other minor changes.
Posted by re...@apache.org.
Disable other sink tests. Need to fix how tests use global MOCK_PRODUCER. A few more other minor changes.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/6bf618d9
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/6bf618d9
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/6bf618d9
Branch: refs/heads/master
Commit: 6bf618d97920bba357c9d592216757d8409db90f
Parents: d745cc9
Author: Raghu Angadi <ra...@google.com>
Authored: Fri Jul 21 13:49:19 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:01:44 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 8 +++----
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 22 ++++++++++----------
2 files changed, 15 insertions(+), 15 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/6bf618d9/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 5813203..cb1c287 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -1792,12 +1792,12 @@ public class KafkaIO {
// across a fixed number of shards and records in each shard are written in order. It drops
// any records that are already written and buffers those arriving out of order.
//
- // // Exactly once sink involves two shuffles of the records:
- // A -- GBK --> B -- GBK --> C
+ // Exactly once sink involves two shuffles of the records:
+ // A : Assign a shard ---> B : Assign sequential ID ---> C : Write to Kafka in order
//
// Processing guarantees also require deterministic processing within user transforms.
- // in this case that implies the order of the records seen by C should not be affected by
- // restarts in upstream stages link B & A.
+ // Here, that requires order of the records committed to Kafka by C should not be affected by
+ // restarts in C and its upstream stages.
//
// A : Assigns a random shard for message. Note that there are no ordering guarantees for
// writing user records to Kafka. User can still control partitioning among topic
http://git-wip-us.apache.org/repos/asf/beam/blob/6bf618d9/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 48d4922..3a6c974 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -740,13 +740,13 @@ public class KafkaIOTest {
.withTopic(topic)
.withKeySerializer(IntegerSerializer.class)
.withValueSerializer(LongSerializer.class)
- // .withEOS()
- // .withSinkGroupId("test")
- // .withNumShards(1)
- // .withKeyCoder(BigEndianIntegerCoder.of())
- // .withValueCoder(BigEndianLongCoder.of())
- // .withConsumerFactoryFn(new ConsumerFactoryFn(
- // Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST))
+ .withEOS()
+ .withSinkGroupId("test")
+ .withNumShards(1)
+ .withKeyCoder(BigEndianIntegerCoder.of())
+ .withValueCoder(BigEndianLongCoder.of())
+ .withConsumerFactoryFn(new ConsumerFactoryFn(
+ Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST))
.withProducerFactoryFn(new ProducerFactoryFn()));
p.run();
@@ -757,7 +757,7 @@ public class KafkaIOTest {
}
}
- @Test
+ //@Test (TODO: reenable. global MOCK_PRODUCER does not work with 0.11 anymore for multiple tests)
public void testValuesSink() throws Exception {
// similar to testSink(), but use values()' interface.
@@ -790,7 +790,7 @@ public class KafkaIOTest {
}
}
- @Test
+ //@Test (TODO: reenable. global MOCK_PRODUCER does not work with 0.11 anymore for multiple tests)
public void testSinkWithSendErrors() throws Throwable {
// similar to testSink(), except that up to 10 of the send calls to producer will fail
// asynchronously.
@@ -912,7 +912,7 @@ public class KafkaIOTest {
assertThat(displayData, hasDisplayItem("receive.buffer.bytes", 524288));
}
- @Test
+ //@Test (TODO: reenable. global MOCK_PRODUCER does not work with 0.11 anymore for multiple tests)
public void testSinkDisplayData() {
KafkaIO.Write<Integer, Long> write = KafkaIO.<Integer, Long>write()
.withBootstrapServers("myServerA:9092,myServerB:9092")
@@ -1004,7 +1004,7 @@ public class KafkaIOTest {
KafkaIO.inferCoder(registry, NonInferableObjectDeserializer.class);
}
- @Test
+ // @Test (TODO: reenable. global MOCK_PRODUCER does not work with 0.11 anymore for multiple tests)
public void testSinkMetrics() throws Exception {
// Simply read from kafka source and write to kafka sink. Then verify the metrics are reported.
[11/13] beam git commit: update JavaDoc for withEOS().
Posted by re...@apache.org.
update JavaDoc for withEOS().
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/0d5a3e26
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/0d5a3e26
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/0d5a3e26
Branch: refs/heads/master
Commit: 0d5a3e263205eabb26f269ac27d5ae009123ea26
Parents: 5e280fa
Author: Raghu Angadi <ra...@google.com>
Authored: Thu Oct 5 18:28:52 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:25 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 79 +++++++++++++++-----
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 4 +-
2 files changed, 61 insertions(+), 22 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/0d5a3e26/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index b9b6bd4..5291606 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -174,11 +174,14 @@ import org.slf4j.LoggerFactory;
* .updateConsumerProperties(ImmutableMap.of("receive.buffer.bytes", 1024 * 1024))
*
* // custom function for calculating record timestamp (default is processing time)
- * .withTimestampFn(new MyTypestampFunction())
+ * .withTimestampFn(new MyTimestampFunction())
*
* // custom function for watermark (default is record timestamp)
* .withWatermarkFn(new MyWatermarkFunction())
*
+ * // restrict reader to committed messages on Kafka (see method documentation).
+ * .withReadCommitted()
+ *
* // finally, if you don't need Kafka metadata, you can drop it
* .withoutMetadata() // PCollection<KV<Long, String>>
* )
@@ -235,6 +238,9 @@ import org.slf4j.LoggerFactory;
* // you can further customize KafkaProducer used to write the records by adding more
* // settings for ProducerConfig. e.g, to enable compression :
* .updateProducerProperties(ImmutableMap.of("compression.type", "gzip"))
+ *
+ * // Optionally enable exactly-once sink (on supported runners). See JavaDoc for withEOS().
+ * .withEOS(20, "eos-sink-group-id");
* );
* }</pre>
*
@@ -1530,29 +1536,62 @@ public class KafkaIO {
}
/**
- * TODO: User friendly javadoc.
- * Note on performance: Exactly-once sink involves two shuffles of input records in order to
- * provide the right semantics. As a result, the input records go through 2
- * serialization-deserialization cycles. Depending on volume and cost of serialization, the
- * CPU cost might be noticeable. The cost could be minimized by writing byte arrays.
+ * Provides exactly-once semantics while writing to Kafka, which enables applications with
+ * end-to-end exactly-once guarantees on top of exactly-once semantics <i>within</i> Beam
+ * pipelines. It ensures that records written to sink are committed on Kafka exactly once,
+ * even in the case of retries during pipeline execution even when some processing is retried.
+ * Retries typically occur when workers restart (as in failure recovery), or when the work is
+ * redistributed (as in an autoscaling event).
+ *
+ * <p>Beam runners typically provide exactly-once semantics for results of a pipeline, but not
+ * for side effects from user code in transform. If a transform such as Kafka sink writes
+ * to an external system, those writes might occur more than once. When EOS is enabled here,
+ * the sink transform ties checkpointing semantics in compatible Beam runners and transactions
+ * in Kafka (version 0.11+) to ensure a record is written only once. As the implementation
+ * relies on runners checkpoint semantics, not all the runners are compatible. The sink throws
+ * an exception during initialization if the runner is not whitelisted. Flink runner is
+ * one of the runners whose checkpoint semantics are not compatible with current
+ * implementation (hope to provide a solution in near future). Dataflow runner and Spark
+ * runners are whitelisted as compatible.
+ *
+ * <p>Note on performance: Exactly-once sink involves two shuffles of the records. In addition
+ * to cost of shuffling the records among workers, the records go through 2
+ * serialization-deserialization cycles. Depending on volume and cost of serialization,
+ * the CPU cost might be noticeable. The CPU cost can be reduced by writing byte arrays
+ * (i.e. serializing them to byte before writing to Kafka sink).
+ *
+ * @param numShards Sets sink parallelism. The state metadata stored on Kafka is stored across
+ * this many virtual partitions using {@code sinkGroupId}. A good rule of thumb is to set
+ * this to be around number of partitions in Kafka topic.
+ *
+ * @param sinkGroupId The <i>group id</i> used to store small amount of state as metadata on
+ * Kafka. It is similar to <i>consumer group id</i> used with a {@link KafkaConsumer}. Each
+ * job should use a unique group id so that restarts/updates of job preserve the state to
+ * ensure exactly-once semantics. The state is committed atomically with sink transactions
+ * on Kafka. See {@link KafkaProducer#sendOffsetsToTransaction(Map, String)} for more
+ * information. The sink performs multiple sanity checks during initialization to catch
+ * common mistakes so that it does not end up using state that does not <i>seem</i> to
+ * be written by the same job.
*/
- public Write<K, V> withEOS() {
+ public Write<K, V> withEOS(int numShards, String sinkGroupId) {
EOSWrite.ensureEOSSupport();
- return toBuilder().setEOS(true).build();
+ checkArgument(numShards >= 1, "numShards should be >= 1");
+ checkArgument(sinkGroupId != null, "sinkGroupId is required for exactly-once sink");
+ return toBuilder()
+ .setEOS(true)
+ .setNumShards(numShards)
+ .setSinkGroupId(sinkGroupId)
+ .build();
}
/**
- * Should be unique for the job. This is also used in naming the producers used in EOS sink.
- * TODO: expand javaDoc.
+ * When exactly-once semantics are enabled (see {@link #withEOS(int, String)}), the sink needs
+ * to fetch previously stored state with Kafka topic. Fetching the metadata requires a
+ * consumer. Similar to {@link Read#withConsumerFactoryFn(SerializableFunction)}, a factory
+ * function can be supplied if required in a specific case.
+ * The default is {@link KafkaConsumer}.
+ * @param consumerFactoryFn
*/
- public Write<K, V> withSinkGroupId(String sinkGroupId) {
- return toBuilder().setSinkGroupId(sinkGroupId).build();
- }
-
- public Write<K, V> withNumShards(int numShards) {
- return toBuilder().setNumShards(numShards).build();
- }
-
public Write<K, V> withConsumerFactoryFn(
SerializableFunction<Map<String, Object>, ? extends Consumer<?, ?>> consumerFactoryFn) {
return toBuilder().setConsumerFactoryFn(consumerFactoryFn).build();
@@ -1579,7 +1618,9 @@ public class KafkaIO {
// TODO: Verify that the group_id does not have existing state stored on Kafka unless
// this is an upgrade. This avoids issues with simple mistake of reusing group_id
- // across multiple runs or across multiple jobs.
+ // across multiple runs or across multiple jobs. This is checked when the sink
+ // transform initializes while processing the output. It might be better to
+ // check here to catch common mistake.
input.apply(new EOSWrite<>(this));
} else {
http://git-wip-us.apache.org/repos/asf/beam/blob/0d5a3e26/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index b1e4a4d..89748dd 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -821,9 +821,7 @@ public class KafkaIOTest {
.withTopic(topic)
.withKeySerializer(IntegerSerializer.class)
.withValueSerializer(LongSerializer.class)
- .withEOS()
- .withSinkGroupId("test")
- .withNumShards(1)
+ .withEOS(1, "test")
.withConsumerFactoryFn(new ConsumerFactoryFn(
Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST))
.withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)));
[09/13] beam git commit: Support version 0.11 through ProducerSpEL
(reflection). Tests work with both 0.11 and older versions.
Posted by re...@apache.org.
Support version 0.11 through ProducerSpEL (reflection).
Tests work with both 0.11 and older versions.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/265e7c74
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/265e7c74
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/265e7c74
Branch: refs/heads/master
Commit: 265e7c743daf8c17cd1064c740290e99878ffb79
Parents: 3bd4085
Author: Raghu Angadi <ra...@google.com>
Authored: Tue Sep 19 13:50:05 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:05 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 44 ++++---
.../apache/beam/sdk/io/kafka/ProducerSpEL.java | 118 +++++++++++++++++++
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 44 +++++--
3 files changed, 181 insertions(+), 25 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/265e7c74/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 1d9560b..76e8a1e 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -124,9 +124,6 @@ import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.errors.AuthorizationException;
-import org.apache.kafka.common.errors.OutOfOrderSequenceException;
-import org.apache.kafka.common.errors.ProducerFencedException;
import org.apache.kafka.common.errors.WakeupException;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.Deserializer;
@@ -1525,6 +1522,7 @@ public class KafkaIO {
* CPU cost might be noticeable. The cost could be minimized by writing byte arrays.
*/
public Write<K, V> withEOS() {
+ EOSWrite.ensureEOSSupport();
return toBuilder().setEOS(true).build();
}
@@ -1830,6 +1828,13 @@ public class KafkaIO {
private final Write<K, V> spec;
+ static void ensureEOSSupport() {
+ checkArgument(
+ ProducerSpEL.supportsTransactions(), "%s %s",
+ "This version of Kafka client does not support transactions required to support",
+ "exactly-once semantics. Please use Kafka client version 0.11 or newer.");
+ }
+
EOSWrite(Write<K, V> spec) {
this.spec = spec;
}
@@ -1942,6 +1947,12 @@ public class KafkaIO {
this.outOfOrderBuffer = StateSpecs.bag(KvCoder.of(BigEndianLongCoder.of(), elemCoder));
}
+ @Setup
+ public void setup() {
+ // This is on the worker. Ensure the runtime version is till compatible.
+ EOSWrite.ensureEOSSupport();
+ }
+
@ProcessElement
public void processElement(@StateId(NEXT_ID) ValueState<Long> nextIdState,
@StateId(MIN_BUFFERED_ID) ValueState<Long> minBufferedIdState,
@@ -2043,7 +2054,7 @@ public class KafkaIO {
writer.commitTxn(nextId - 1, numTransactions);
nextIdState.write(nextId);
- } catch (ProducerFencedException | OutOfOrderSequenceException | AuthorizationException e) {
+ } catch (ProducerSpEL.UnrecoverableProducerException e) {
// Producer JavaDoc says these are not recoverable errors and producer should be closed.
// Close the producer and a new producer will be initialized in retry.
@@ -2110,7 +2121,7 @@ public class KafkaIO {
}
void beginTxn() {
- producer.beginTransaction();
+ ProducerSpEL.beginTransaction(producer);
}
void sendRecord(KV<K, V> record, Counter sendCounter) {
@@ -2119,7 +2130,7 @@ public class KafkaIO {
new ProducerRecord<>(spec.getTopic(), record.getKey(), record.getValue()));
sendCounter.inc();
} catch (KafkaException e) {
- producer.abortTransaction();
+ ProducerSpEL.abortTransaction(producer);
throw e;
}
}
@@ -2130,20 +2141,21 @@ public class KafkaIO {
// NOTE: Kafka keeps this metadata for 24 hours since the last update. This limits
// how long the pipeline could be down before resuming it. It does not look like
// this TTL can be adjusted (asked about it on Kafka users list).
- producer.sendOffsetsToTransaction(
+ ProducerSpEL.sendOffsetsToTransaction(
+ producer,
ImmutableMap.of(new TopicPartition(spec.getTopic(), shard),
- new OffsetAndMetadata(
- 0L, JSON_MAPPER.writeValueAsString(new ShardMetadata(lastRecordId,
- writerId)))),
+ new OffsetAndMetadata(0L,
+ JSON_MAPPER.writeValueAsString(
+ new ShardMetadata(lastRecordId, writerId)))),
spec.getSinkGroupId());
- producer.commitTransaction();
+ ProducerSpEL.commitTransaction(producer);
numTransactions.inc();
- LOG.info("{} : committed {} records", shard, lastRecordId - committedId);
+ LOG.debug("{} : committed {} records", shard, lastRecordId - committedId);
committedId = lastRecordId;
} catch (KafkaException e) {
- producer.abortTransaction();
+ ProducerSpEL.abortTransaction(producer);
throw e;
}
}
@@ -2315,14 +2327,14 @@ public class KafkaIO {
producerConfig.putAll(ImmutableMap.of(
ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, spec.getKeySerializer(),
ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, spec.getValueSerializer(),
- ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, true,
- ProducerConfig.TRANSACTIONAL_ID_CONFIG, producerName));
+ ProducerSpEL.ENABLE_IDEMPOTENCE_CONFIG, true,
+ ProducerSpEL.TRANSACTIONAL_ID_CONFIG, producerName));
Producer<K, V> producer = spec.getProducerFactoryFn() != null
? spec.getProducerFactoryFn().apply((producerConfig))
: new KafkaProducer<K, V>(producerConfig);
- producer.initTransactions();
+ ProducerSpEL.initTransactions(producer);
return producer;
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/265e7c74/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ProducerSpEL.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ProducerSpEL.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ProducerSpEL.java
new file mode 100644
index 0000000..ec9bb64
--- /dev/null
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ProducerSpEL.java
@@ -0,0 +1,118 @@
+package org.apache.beam.sdk.io.kafka;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.producer.Producer;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.ApiException;
+import org.apache.kafka.common.errors.AuthorizationException;
+
+/**
+ * ProducerSpEL to handle newer versions Producer API. The API is updated in Kafka 0.11
+ * to support exactly-once semantics.
+ */
+class ProducerSpEL {
+
+ private static boolean supportsTransactions;
+
+ private static Method initTransactionsMethod;
+ private static Method beginTransactionMethod;
+ private static Method commitTransactionMethod;
+ private static Method abortTransactionMethod;
+ private static Method sendOffsetsToTransactionMethod;
+
+ static final String ENABLE_IDEMPOTENCE_CONFIG = "enable.idempotence";
+ static final String TRANSACTIONAL_ID_CONFIG = "transactional.id";
+
+ private static Class producerFencedExceptionClass;
+ private static Class outOfOrderSequenceExceptionClass;
+
+ static {
+ try {
+ initTransactionsMethod = Producer.class.getMethod("initTransactions");
+ beginTransactionMethod = Producer.class.getMethod("beginTransaction");
+ commitTransactionMethod = Producer.class.getMethod("commitTransaction");
+ abortTransactionMethod = Producer.class.getMethod("abortTransaction");
+ sendOffsetsToTransactionMethod = Producer.class.getMethod(
+ "sendOffsetsToTransaction", Map.class, String.class);
+
+ producerFencedExceptionClass = Class.forName(
+ "org.apache.kafka.common.errors.ProducerFencedException");
+ outOfOrderSequenceExceptionClass = Class.forName(
+ "org.apache.kafka.common.errors.OutOfOrderSequenceException");
+
+ supportsTransactions = true;
+ } catch (ClassNotFoundException | NoSuchMethodException e) {
+ supportsTransactions = false;
+ }
+ }
+
+ /**
+ * Wraps an unrecoverable producer exceptions, including the ones related transactions
+ * introduced in 0.11 (as described in documentation for {@link Producer}). The calller should
+ * close the producer when this exception is thrown.
+ */
+ static class UnrecoverableProducerException extends ApiException {
+ UnrecoverableProducerException(ApiException cause) {
+ super(cause);
+ }
+ }
+
+ static boolean supportsTransactions() {
+ return supportsTransactions;
+ }
+
+ private static void ensureTransactionsSupport() {
+ checkArgument(supportsTransactions(),
+ "This version of Kafka client library does not support transactions. ",
+ "Please used version 0.11 or later.");
+ }
+
+ private static Object invoke(Method method, Object obj, Object... args) {
+ try {
+ return method.invoke(obj, args);
+ } catch (IllegalAccessException | InvocationTargetException e) {
+ return new RuntimeException(e);
+ } catch (ApiException e) {
+ Class eClass = e.getClass();
+ if (producerFencedExceptionClass.isAssignableFrom(eClass)
+ || outOfOrderSequenceExceptionClass.isAssignableFrom(eClass)
+ || AuthorizationException.class.isAssignableFrom(eClass)) {
+ throw new UnrecoverableProducerException(e);
+ }
+ throw e;
+ }
+ }
+
+ static void initTransactions(Producer<?, ?> producer) {
+ ensureTransactionsSupport();
+ invoke(initTransactionsMethod, producer);
+ }
+
+ static void beginTransaction(Producer<?, ?> producer) {
+ ensureTransactionsSupport();
+ invoke(beginTransactionMethod, producer);
+ }
+
+ static void commitTransaction(Producer<?, ?> producer) {
+ ensureTransactionsSupport();
+ invoke(commitTransactionMethod, producer);
+ }
+
+ static void abortTransaction(Producer<?, ?> producer) {
+ ensureTransactionsSupport();
+ invoke(abortTransactionMethod, producer);
+ }
+
+ static void sendOffsetsToTransaction(Producer<?, ?> producer,
+ Map<TopicPartition, OffsetAndMetadata> offsets,
+ String consumerGroupId) {
+ ensureTransactionsSupport();
+ invoke(sendOffsetsToTransactionMethod, producer, offsets, consumerGroupId);
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/265e7c74/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 88c9d21..1532788 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -32,9 +32,9 @@ import com.google.common.collect.Lists;
import java.io.IOException;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
import java.nio.ByteBuffer;
import java.util.ArrayList;
-import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
@@ -114,14 +114,19 @@ 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',
- * or 'mvn test -Dkafka.clients.version=0.9.0.1' for either Kafka client version
+ * or 'mvn test -Dkafka.clients.version=0.9.0.1' for either Kafka client 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.
@@ -184,7 +189,7 @@ public class KafkaIOTest {
//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 Collection<TopicPartition> assigned) {
+ public void assign(final List<TopicPartition> assigned) {
super.assign(assigned);
assignedPartitions.set(ImmutableList.copyOf(assigned));
for (TopicPartition tp : assigned) {
@@ -790,6 +795,11 @@ public class KafkaIOTest {
// similar to testSink(), enables to EOS.
// This does not actually test exactly-once-semantics. Mainly exercises the code.
+ if (!ProducerSpEL.supportsTransactions()) {
+ LOG.warn("testEOSink() is disabled as Kafka client version does not support transactions.");
+ return;
+ }
+
int numElements = 1000;
try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) {
@@ -1108,16 +1118,28 @@ public class KafkaIOTest {
}
/**
- * This wrapper over MockProducer 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. inject errors).
+ * 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>(
@@ -1147,8 +1169,12 @@ public class KafkaIOTest {
public void close() {
MOCK_PRODUCER_MAP.remove(producerKey);
- if (!mockProducer.closed()) {
- mockProducer.close();
+ try {
+ if (closedMethod == null || !((Boolean) closedMethod.invoke(mockProducer))) {
+ mockProducer.close();
+ }
+ } catch (Exception e) { // Not expected.
+ throw new RuntimeException(e);
}
}
}
[12/13] beam git commit: update JavaDoc for withEOS().
Posted by re...@apache.org.
update JavaDoc for withEOS().
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/686b22f3
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/686b22f3
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/686b22f3
Branch: refs/heads/master
Commit: 686b22f36e9ad487d99a184c9d9c6218b7f938d4
Parents: 0d5a3e2
Author: Raghu Angadi <ra...@google.com>
Authored: Thu Oct 5 18:28:52 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:03:17 2017 -0700
----------------------------------------------------------------------
.../kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java | 1 -
1 file changed, 1 deletion(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/686b22f3/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 5291606..603e62f 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -1614,7 +1614,6 @@ public class KafkaIO {
if (isEOS()) {
EOSWrite.ensureEOSSupport();
- checkNotNull(getSinkGroupId(), "A group id is required for exactly-once sink");
// TODO: Verify that the group_id does not have existing state stored on Kafka unless
// this is an upgrade. This avoids issues with simple mistake of reusing group_id
[04/13] beam git commit: 1) Add a note about metadata expiration. 2)
Fetch number of partitions if numShards is not set.
Posted by re...@apache.org.
1) Add a note about metadata expiration.
2) Fetch number of partitions if numShards is not set.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/d9c95f0e
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/d9c95f0e
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/d9c95f0e
Branch: refs/heads/master
Commit: d9c95f0ea40be1b5d4fdef9e123321367186fc9d
Parents: 6bf618d
Author: Raghu Angadi <ra...@google.com>
Authored: Wed Jul 26 14:02:27 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:04 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 68 ++++++++++++--------
1 file changed, 42 insertions(+), 26 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/d9c95f0e/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index cb1c287..9f5f493 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -55,6 +55,7 @@ import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
@@ -1558,8 +1559,9 @@ public class KafkaIO {
if (isEOS()) {
EOSWrite.ensureEOSSupport();
- checkArgument(getSinkGroupId() != null,
- "withSinkGroupId() is required for exactly-once sink");
+ checkNotNull(getSinkGroupId(), "A group id is required for exactly-once sink");
+ checkNotNull(getKeyCoder(), "Key coder is required for exact-once sink");
+ checkNotNull(getValueCoder(), "Value coder is required for exactly-once sink");
// TODO: Verify that the group_id does not have existing state stored on Kafka unless
// this is an upgrade. This avoids issues with simple mistake of reusing group_id
@@ -1833,13 +1835,23 @@ public class KafkaIO {
@Override
public PCollection<Void> expand(PCollection<KV<K, V>> input) {
+
+ int numShards = spec.getNumShards();
+ if (numShards <= 0) {
+ try (Consumer<?, ?> consumer = openConsumer(spec)) {
+ numShards = consumer.partitionsFor(spec.getTopic()).size();
+ LOG.info("Using {} shards for exactly-once, matching number of partitions for topic '{}'",
+ numShards, spec.getTopic());
+ }
+ }
+ checkState(numShards > 0, "Could not set number of shards");
+
return input
.apply(Window.<KV<K, V>>into(new GlobalWindows()) // Everything into global window.
.triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
.discardingFiredPanes())
- .apply(String.format("Shuffle across %d shards", spec.getNumShards()),
- ParDo.of(new EOSReshard<K, V>(DateTimeUtils.currentTimeMillis(),
- spec.getNumShards())))
+ .apply(String.format("Shuffle across %d shards", numShards),
+ ParDo.of(new EOSReshard<K, V>(numShards)))
.apply("Persist sharding", GroupByKey.<Integer, KV<K, V>>create())
.apply("Assign sequential ids", ParDo.of(new EOSSequencer<K, V>()))
.apply("Persist ids", GroupByKey.<Integer, KV<Long, KV<K, V>>>create())
@@ -1854,17 +1866,16 @@ public class KafkaIO {
* Shuffle messages assigning each randomly to a shard.
*/
private static class EOSReshard<K, V> extends DoFn<KV<K, V>, KV<Integer, KV<K, V>>> {
- private final Random random;
private final int numShards;
- EOSReshard(long seed, int numShards) {
- this.random = new Random(seed);
+ EOSReshard(int numShards) {
this.numShards = numShards;
}
@ProcessElement
public void processElement(ProcessContext ctx) {
- ctx.output(KV.of(random.nextInt(numShards), ctx.element()));
+ int shard = ThreadLocalRandom.current().nextInt(numShards);
+ ctx.output(KV.of(shard, ctx.element()));
}
}
@@ -2093,14 +2104,15 @@ public class KafkaIO {
void commitTxn(long lastRecordId) throws IOException {
try {
- // Store id in consumer group metadata for the partition
+ // Store id in consumer group metadata for the partition.
+ // NOTE: Kafka keeps this metadata for 24 hours since the last update. This limits
+ // how long the pipeline could be down before resuming it. It does not look like
+ // this TTL can be adjusted (asked about it on Kafka users list).
producer.sendOffsetsToTransaction(
ImmutableMap.of(new TopicPartition(spec.getTopic(), shard),
new OffsetAndMetadata(
- Long.MAX_VALUE, // So that consumer group does not expire.
- JSON_MAPPER.writeValueAsString(new ShardMetadata(lastRecordId,
- writerId)
- ))),
+ 0L, JSON_MAPPER.writeValueAsString(new ShardMetadata(lastRecordId,
+ writerId)))),
spec.getSinkGroupId());
producer.commitTransaction();
@@ -2147,20 +2159,10 @@ public class KafkaIO {
// Fetch latest committed metadata for the partition (if any). Checks committed sequence ids.
try {
- Consumer<?, ?> consumer = spec.getConsumerFactoryFn().apply((ImmutableMap.of(
- ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, spec
- .getProducerConfig().get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),
- ConsumerConfig.GROUP_ID_CONFIG, spec.getSinkGroupId(),
- ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class,
- ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class
- )));
-
OffsetAndMetadata committed;
- try {
- committed = consumer.committed(new TopicPartition(spec.getTopic(), shard));
- } finally {
- consumer.close();
+ try (Consumer<?, ?> consumer = openConsumer(spec)) {
+ committed = consumer.committed(new TopicPartition(spec.getTopic(), shard));
}
long committedSeqId = -1;
@@ -2236,6 +2238,20 @@ public class KafkaIO {
}
}
+ /**
+ * Opens a generic consumer that is mainly meant for metadata operations like fetching
+ * number of partitions for a topic rather than for fetching messages.
+ */
+ private static Consumer<?, ?> openConsumer(Write<?, ?> spec) {
+ return spec.getConsumerFactoryFn().apply((ImmutableMap.of(
+ ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, spec
+ .getProducerConfig().get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),
+ ConsumerConfig.GROUP_ID_CONFIG, spec.getSinkGroupId(),
+ ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class,
+ ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class
+ )));
+ }
+
private static <K, V> Producer<K, V> initializeEosProducer(Write<K, V> spec,
String producerName) {
[10/13] beam git commit: Close producers after one minute of
inactivity. This closes producers when ranges move to other workers. Misc
fixes around.
Posted by re...@apache.org.
Close producers after one minute of inactivity.
This closes producers when ranges move to other workers.
Misc fixes around.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/f4f6105e
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/f4f6105e
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/f4f6105e
Branch: refs/heads/master
Commit: f4f6105e9ca7da0af938ba65544398be2f6fc4e8
Parents: 5767156
Author: Raghu Angadi <ra...@google.com>
Authored: Fri Jul 28 17:39:17 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:05 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 172 +++++++++++++------
1 file changed, 115 insertions(+), 57 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/f4f6105e/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 5d50cf7..63dc734 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -28,6 +28,13 @@ import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.base.MoreObjects;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+import com.google.common.cache.RemovalCause;
+import com.google.common.cache.RemovalListener;
+import com.google.common.cache.RemovalNotification;
import com.google.common.collect.ComparisonChain;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
@@ -1510,6 +1517,13 @@ public class KafkaIO {
return toBuilder().setProducerFactoryFn(producerFactoryFn).build();
}
+ /**
+ * TODO: User friendly javadoc.
+ * Note on performance: Exactly-once sink involves two shuffles of input records in order to
+ * provide the right semantics. As a result, the input records go through 2
+ * serialization-deserialization cycles. Depending on volume and cost of serialization, the
+ * CPU cost might be noticeable. The cost could be minimized by writing byte arrays.
+ */
public Write<K, V> withEOS() {
return toBuilder().setEOS(true).build();
}
@@ -1827,8 +1841,8 @@ public class KafkaIO {
if (numShards <= 0) {
try (Consumer<?, ?> consumer = openConsumer(spec)) {
numShards = consumer.partitionsFor(spec.getTopic()).size();
- LOG.info("Using {} shards for exactly-once, matching number of partitions for topic '{}'",
- numShards, spec.getTopic());
+ LOG.info("Using {} shards for exactly-once writer, matching number of partitions "
+ + "for topic '{}'", numShards, spec.getTopic());
}
}
checkState(numShards > 0, "Could not set number of shards");
@@ -1911,11 +1925,6 @@ public class KafkaIO {
private final Write<K, V> spec;
- // One cache for each sink (usually there is only one sink per pipeline
- private static final Map<String, Map<Integer, ShardWriter<?, ?>>> CACHE_BY_GROUP_ID =
- new HashMap<>();
- // TODO: Need a way to close producers that are no longer relevant (may be have a timeout?).
-
// Metrics
private final Counter elementsWritten = SinkMetrics.elementsWritten();
// Elements buffered due to out of order arrivals.
@@ -1942,13 +1951,19 @@ public class KafkaIO {
long nextId = MoreObjects.firstNonNull(nextIdState.read(), 0L);
long minBufferedId = MoreObjects.firstNonNull(minBufferedIdState.read(), Long.MAX_VALUE);
- ShardWriter<K, V> writer = getShardWriter(shard, writerIdState, nextId);
+ ShardWriterCache<K, V> cache =
+ (ShardWriterCache<K, V>) CACHE_BY_GROUP_ID.getUnchecked(spec.getSinkGroupId());
+ ShardWriter<K, V> writer = cache.removeIfPresent(shard);
+ if (writer == null) {
+ writer = initShardWriter(shard, writerIdState, nextId);
+ }
+
long committedId = writer.committedId;
if (committedId >= nextId) {
// This is a retry of an already committed batch.
LOG.info("{}: committed id {} is ahead of expected {}. {} records will be dropped "
- + "(these are already written).",
+ + "(these are already written).",
shard, committedId, nextId - 1, committedId - nextId + 1);
nextId = committedId + 1;
}
@@ -2007,7 +2022,7 @@ public class KafkaIO {
List<KV<Long, KV<K, V>>> buffered = Lists.newArrayList(oooBufferState.read());
Collections.sort(buffered, new KV.OrderByKey<Long, KV<K, V>>());
- LOG.info("{} : merging {} buffered records with min buffered id",
+ LOG.info("{} : merging {} buffered records (min buffered id is {}).",
shard, buffered.size(), minBufferedId);
oooBufferState.clear();
@@ -2021,25 +2036,31 @@ public class KafkaIO {
writer.commitTxn(nextId - 1, numTransactions);
nextIdState.write(nextId);
+
} catch (ProducerFencedException | OutOfOrderSequenceException | AuthorizationException e) {
- // JavaDoc says these are not recoverable errors and producer should be closed.
+ // Producer JavaDoc says these are not recoverable errors and producer should be closed.
- // Close the producer and and remove it from the cache. A new producer will be created
- // in retry. It is possible that a rough worker keeps retrying and ends up fencing off
+ // Close the producer and a new producer will be initialized in retry.
+ // It is possible that a rough worker keeps retrying and ends up fencing off
// active producers. How likely this might be or how well such a scenario is handled
// depends on the runner. For now we will leave it to upper layers, will need to revisit.
- LOG.warn("{} : closing producer {} after unrecoverable error. The work might be migrated. "
- + "committed id {}, current id {}.",
+ LOG.warn("{} : closing producer {} after unrecoverable error. The work might have migrated."
+ + " Committed id {}, current id {}.",
writer.shard, writer.producerName, writer.committedId, nextId - 1, e);
- CACHE_BY_GROUP_ID.get(spec.getSinkGroupId()).remove(writer.shard);
writer.producer.close();
+ writer = null; // No need to cache it.
throw e;
+ } finally {
+ if (writer != null) {
+ cache.insert(shard, writer);
+ }
}
}
private static class ShardMetadata {
+
@JsonProperty("seq")
public final long sequenceId;
@JsonProperty("id")
@@ -2056,7 +2077,11 @@ public class KafkaIO {
}
}
+ /**
+ * A wrapper around Kafka producer. One for each of the shards.
+ */
private static class ShardWriter<K, V> {
+
private final int shard;
private final String writerId;
private final Producer<K, V> producer;
@@ -2102,8 +2127,8 @@ public class KafkaIO {
producer.sendOffsetsToTransaction(
ImmutableMap.of(new TopicPartition(spec.getTopic(), shard),
new OffsetAndMetadata(
- 0L, JSON_MAPPER.writeValueAsString(new ShardMetadata(lastRecordId,
- writerId)))),
+ 0L, JSON_MAPPER.writeValueAsString(new ShardMetadata(lastRecordId,
+ writerId)))),
spec.getSinkGroupId());
producer.commitTransaction();
@@ -2118,39 +2143,18 @@ public class KafkaIO {
}
}
- private ShardWriter<K, V> getShardWriter(int shard,
- ValueState<String> writerIdState,
- long nextId) throws IOException {
-
- Map<Integer, ShardWriter<?, ?>> cache;
-
- synchronized (CACHE_BY_GROUP_ID) {
- cache = CACHE_BY_GROUP_ID.get(spec.getSinkGroupId());
- if (cache == null) {
- cache = new HashMap<>();
- CACHE_BY_GROUP_ID.put(spec.getSinkGroupId(), cache);
- }
- }
-
- synchronized (cache) {
- @SuppressWarnings("unchecked")
- ShardWriter<K, V> shardWriter = (ShardWriter<K, V>) cache.get(shard);
- if (shardWriter != null) {
- // any sanity checks?
- return shardWriter;
- }
- }
-
- // initialize new shard
+ private ShardWriter<K, V> initShardWriter(int shard,
+ ValueState<String> writerIdState,
+ long nextId) throws IOException {
String producerName = String.format("producer_%d_for_%s", shard, spec.getSinkGroupId());
Producer<K, V> producer = initializeEosProducer(spec, producerName);
- String writerId = writerIdState.read();
-
// Fetch latest committed metadata for the partition (if any). Checks committed sequence ids.
try {
+ String writerId = writerIdState.read();
+
OffsetAndMetadata committed;
try (Consumer<?, ?> consumer = openConsumer(spec)) {
@@ -2161,7 +2165,7 @@ public class KafkaIO {
if (committed == null || committed.metadata() == null || committed.metadata().isEmpty()) {
checkState(nextId == 0 && writerId == null,
- "State exists for shard %d (nextId %s, writerId '%s'), but there is no state "
+ "State exists for shard %s (nextId %s, writerId '%s'), but there is no state "
+ "stored with Kafka topic '%s' group id '%s'",
shard, nextId, writerId, spec.getTopic(), spec.getSinkGroupId());
@@ -2191,9 +2195,9 @@ public class KafkaIO {
// We could let users explicitly an option to override the existing metadata.
//
throw new IllegalStateException(String.format(
- "Kafka metadata exists for shard %d, but there is no stored state for it. "
+ "Kafka metadata exists for shard %s, but there is no stored state for it. "
+ "This mostly indicates groupId '%s' is already used else where or in earlier runs. "
- + "Try another group id. Metadata : '%s'.",
+ + "Try another group id. Metadata for this shard on Kafka : '%s'",
shard, spec.getSinkGroupId(), committed.metadata()));
}
@@ -2209,25 +2213,79 @@ public class KafkaIO {
nextId - 1, committed.metadata());
}
- ShardWriter<K, V> shardWriter = new ShardWriter<>(shard, writerId, producer, producerName,
- spec, committedSeqId);
-
- synchronized (cache) {
- checkState(cache.get(shard) == null,
- "Unexpected concurrent execution of shard %s", shard);
- cache.put(shard, shardWriter);
- }
-
LOG.info("{} : initialized producer {} with committed sequence id {}",
shard, producerName, committedSeqId);
- return shardWriter;
+ return new ShardWriter<>(shard, writerId, producer, producerName, spec, committedSeqId);
} catch (Exception e) {
producer.close();
throw e;
}
}
+
+ /**
+ * A wrapper around guava cache to provide insert()/remove() semantics. A ShardWriter will
+ * be closed if it is stays in cache for more than 1 minute, i.e. not used inside EOSWrite
+ * DoFn for a minute or more.
+ */
+ private static class ShardWriterCache<K, V> {
+
+ static final ScheduledExecutorService SCHEDULED_CLEAN_UP_THREAD =
+ Executors.newSingleThreadScheduledExecutor();
+
+ static final int CLEAN_UP_CHECK_INTERVAL_MS = 10 * 1000;
+ static final int IDLE_TIMEOUT_MS = 60 * 1000;
+
+ private final Cache<Integer, ShardWriter<K, V>> cache;
+
+ ShardWriterCache() {
+ this.cache = CacheBuilder
+ .newBuilder()
+ .expireAfterWrite(IDLE_TIMEOUT_MS, TimeUnit.MILLISECONDS)
+ .removalListener(new RemovalListener<Integer, ShardWriter<K, V>>() {
+ @Override
+ public void onRemoval(RemovalNotification<Integer, ShardWriter<K, V>> notification) {
+ if (notification.getCause() != RemovalCause.EXPLICIT) {
+ ShardWriter writer = notification.getValue();
+ LOG.info("{} : Closing idle shard writer {} after 1 minute of idle time.",
+ writer.shard, writer.producerName);
+ writer.producer.close();
+ }
+ }
+ }).build();
+
+ // run cache.cleanUp() every 10 seconds.
+ SCHEDULED_CLEAN_UP_THREAD.scheduleAtFixedRate(
+ new Runnable() {
+ @Override
+ public void run() {
+ cache.cleanUp();
+ }
+ },
+ CLEAN_UP_CHECK_INTERVAL_MS, CLEAN_UP_CHECK_INTERVAL_MS, TimeUnit.MILLISECONDS);
+ }
+
+ ShardWriter<K, V> removeIfPresent(int shard) {
+ return cache.asMap().remove(shard);
+ }
+
+ void insert(int shard, ShardWriter<K, V> writer) {
+ ShardWriter<K, V> existing = cache.asMap().putIfAbsent(shard, writer);
+ checkState(existing == null,
+ "Unexpected multiple instances of writers for shard %s", shard);
+ }
+ }
+
+ // One cache for each sink (usually there is only one sink per pipeline)
+ private static final LoadingCache<String, ShardWriterCache<?, ?>> CACHE_BY_GROUP_ID =
+ CacheBuilder.newBuilder()
+ .build(new CacheLoader<String, ShardWriterCache<?, ?>>() {
+ @Override
+ public ShardWriterCache<?, ?> load(String key) throws Exception {
+ return new ShardWriterCache<>();
+ }
+ });
}
/**
[08/13] beam git commit: Add runner check. Add withReadCommit() to
reader.
Posted by re...@apache.org.
Add runner check. Add withReadCommit() to reader.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/5e280fad
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/5e280fad
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/5e280fad
Branch: refs/heads/master
Commit: 5e280fad28b97c26f27bfb037ef7f45b8fe034e5
Parents: 265e7c7
Author: Raghu Angadi <ra...@google.com>
Authored: Wed Oct 4 17:43:47 2017 -0700
Committer: Raghu Angadi <ra...@google.com>
Committed: Tue Oct 17 00:02:05 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 31 ++++++++++++++++++++
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 11 ++++---
2 files changed, 38 insertions(+), 4 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/5e280fad/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 76e8a1e..b9b6bd4 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -549,6 +549,18 @@ public class KafkaIO {
}
/**
+ * Sets "isolation_level" to "read_committed" in Kafka consumer configuration. This is
+ * ensures that the consumer does not read uncommitted messages. Kafka version 0.11
+ * introduced transactional writes. Applications requiring end-to-end exactly-once
+ * semantics should only read committed messages. See JavaDoc for {@link KafkaConsumer}
+ * for more description.
+ */
+ public Read<K, V> withReadCommitted() {
+ return updateConsumerProperties(
+ ImmutableMap.<String, Object>of("isolation.level", "read_committed"));
+ }
+
+ /**
* Returns a {@link PTransform} for PCollection of {@link KV}, dropping Kafka metatdata.
*/
public PTransform<PBegin, PCollection<KV<K, V>>> withoutMetadata() {
@@ -1275,6 +1287,9 @@ public class KafkaIO {
private void updateLatestOffsets() {
for (PartitionState p : partitionStates) {
try {
+ // If "read_committed" is enabled in the config, this seeks to 'Last Stable Offset'.
+ // As a result uncommitted messages are not counted in backlog. It is correct since
+ // the reader can not read them anyway.
consumerSpEL.evaluateSeek2End(offsetConsumer, p.topicPartition);
long offset = offsetConsumer.position(p.topicPartition);
p.setLatestOffset(offset);
@@ -1573,6 +1588,22 @@ public class KafkaIO {
return PDone.in(input.getPipeline());
}
+ public void validate(PipelineOptions options) {
+ if (isEOS()) {
+ String runner = options.getRunner().getName();
+ if (runner.equals("org.apache.beam.runners.direct.DirectRunner")
+ || runner.startsWith("org.apache.beam.runners.dataflow.")
+ || runner.startsWith("org.apache.beam.runners.spark.")) {
+ return;
+ }
+ throw new UnsupportedOperationException(
+ runner + " is not whitelisted among runners compatible with Kafka exactly-once sink. "
+ + "This implementation of exactly-once sink relies on specific checkpoint guarantees. "
+ + "Only the runners with known to have compatible checkpoint semantics are whitelisted."
+ );
+ }
+ }
+
// set config defaults
private static final Map<String, Object> DEFAULT_PRODUCER_PROPERTIES =
ImmutableMap.<String, Object>of(
http://git-wip-us.apache.org/repos/asf/beam/blob/5e280fad/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 1532788..b1e4a4d 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -791,9 +791,13 @@ public class KafkaIOTest {
}
@Test
- public void testEOSink() throws Exception {
- // similar to testSink(), enables to EOS.
- // This does not actually test exactly-once-semantics. Mainly exercises the code.
+ public void testEOSink() {
+ // 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 KafkaEOWriter DoFn directly to ensure it handles retries correctly.
if (!ProducerSpEL.supportsTransactions()) {
LOG.warn("testEOSink() is disabled as Kafka client version does not support transactions.");
@@ -832,7 +836,6 @@ public class KafkaIOTest {
}
}
-
@Test
public void testSinkWithSendErrors() throws Throwable {
// similar to testSink(), except that up to 10 of the send calls to producer will fail
[13/13] beam git commit: This closes #3612
Posted by re...@apache.org.
This closes #3612
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/0d2a39bd
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/0d2a39bd
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/0d2a39bd
Branch: refs/heads/master
Commit: 0d2a39bd89af0a1fa1c9d908daed5f1f64335b32
Parents: 9524cbd 686b22f
Author: Reuven Lax <re...@relax-macbookpro2.roam.corp.google.com>
Authored: Tue Oct 17 07:25:52 2017 -0700
Committer: Reuven Lax <re...@relax-macbookpro2.roam.corp.google.com>
Committed: Tue Oct 17 07:25:52 2017 -0700
----------------------------------------------------------------------
sdks/java/io/kafka/pom.xml | 12 +-
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 720 ++++++++++++++++++-
.../apache/beam/sdk/io/kafka/ProducerSpEL.java | 118 +++
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 222 ++++--
4 files changed, 1006 insertions(+), 66 deletions(-)
----------------------------------------------------------------------