Add Error Handling to Kafka IO

CHANGES.md

@@ -66,6 +66,7 @@
 * TextIO now supports skipping multiple header lines (Java) ([#17990](https://github.com/apache/beam/issues/17990)).
 * Python GCSIO is now implemented with GCP GCS Client instead of apitools ([#25676](https://github.com/apache/beam/issues/25676))
 * Adding support for LowCardinality DataType in ClickHouse (Java) ([#29533](https://github.com/apache/beam/pull/29533)).
+* Added support for handling bad records to KafkaIO (Java) ([#29546](https://github.com/apache/beam/pull/29546))
 
 ## New Features / Improvements
 

buildSrc/src/main/groovy/org/apache/beam/gradle/kafka/KafkaTestUtilities.groovy

@@ -40,7 +40,7 @@ class KafkaTestUtilities {
         '"keySizeBytes": "10",' +
         '"valueSizeBytes": "90"' +
         '}',
-        "--readTimeout=120",
+        "--readTimeout=60",
         "--kafkaTopic=beam",
         "--withTestcontainers=true",
         "--kafkaContainerVersion=5.5.2",
@@ -56,6 +56,7 @@ class KafkaTestUtilities {
           excludeTestsMatching "*SDFResumesCorrectly" //Kafka SDF does not work for kafka versions <2.0.1
           excludeTestsMatching "*StopReadingFunction" //Kafka SDF does not work for kafka versions <2.0.1
           excludeTestsMatching "*WatermarkUpdateWithSparseMessages" //Kafka SDF does not work for kafka versions <2.0.1
+          excludeTestsMatching "*KafkaIOSDFReadWithErrorHandler"
         }
       }
     }

examples/java/src/main/java/org/apache/beam/examples/KafkaStreaming.java

@@ -49,8 +49,11 @@
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.transforms.Combine;
 import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.transforms.errorhandling.BadRecord;
+import org.apache.beam.sdk.transforms.errorhandling.ErrorHandler.BadRecordErrorHandler;
 import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime;
 import org.apache.beam.sdk.transforms.windowing.FixedWindows;
 import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
@@ -60,6 +63,8 @@
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
+import org.apache.kafka.common.errors.SerializationException;
+import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.IntegerDeserializer;
 import org.apache.kafka.common.serialization.IntegerSerializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
@@ -97,7 +102,7 @@ public interface KafkaStreamingOptions extends PipelineOptions {
      * to use your own Kafka server.
      */
     @Description("Kafka server host")
-    @Default.String("kafka_server:9092")
+    @Default.String("localhost:9092")
     String getKafkaHost();
 
     void setKafkaHost(String value);
@@ -208,15 +213,22 @@ public void run() {
       // Start reading form Kafka with the latest offset
       consumerConfig.put("auto.offset.reset", "latest");
 
-      PCollection<KV<String, Integer>> pCollection =
-          pipeline.apply(
-              KafkaIO.<String, Integer>read()
-                  .withBootstrapServers(options.getKafkaHost())
-                  .withTopic(TOPIC_NAME)
-                  .withKeyDeserializer(StringDeserializer.class)
-                  .withValueDeserializer(IntegerDeserializer.class)
-                  .withConsumerConfigUpdates(consumerConfig)
-                  .withoutMetadata());
+      // Register an error handler for any deserialization errors.
+      // Errors are simulated with an intentionally failing deserializer
+      PCollection<KV<String, Integer>> pCollection;
+      try (BadRecordErrorHandler<PCollection<BadRecord>> errorHandler =
+          pipeline.registerBadRecordErrorHandler(new LogErrors())) {
+        pCollection =
+            pipeline.apply(
+                KafkaIO.<String, Integer>read()
+                    .withBootstrapServers(options.getKafkaHost())
+                    .withTopic(TOPIC_NAME)
+                    .withKeyDeserializer(StringDeserializer.class)
+                    .withValueDeserializer(IntermittentlyFailingIntegerDeserializer.class)
+                    .withConsumerConfigUpdates(consumerConfig)
+                    .withBadRecordErrorHandler(errorHandler)
+                    .withoutMetadata());
+      }
 
       pCollection
           // Apply a window and a trigger ourput repeatedly.
@@ -317,4 +329,39 @@ public void processElement(ProcessContext c, IntervalWindow w) throws Exception
       c.output(c.element());
     }
   }
+
+  // Simple PTransform to log Error information
+  static class LogErrors extends PTransform<PCollection<BadRecord>, PCollection<BadRecord>> {
+
+    @Override
+    public PCollection<BadRecord> expand(PCollection<BadRecord> input) {
+      return input.apply("Log Errors", ParDo.of(new LogErrorFn()));
+    }
+
+    static class LogErrorFn extends DoFn<BadRecord, BadRecord> {
+      @ProcessElement
+      public void processElement(@Element BadRecord record, OutputReceiver<BadRecord> receiver) {
+        System.out.println(record);
+        receiver.output(record);
+      }
+    }
+  }
+
+  // Intentionally failing deserializer to simulate bad data from Kafka
+  public static class IntermittentlyFailingIntegerDeserializer implements Deserializer<Integer> {
+
+    public static final IntegerDeserializer INTEGER_DESERIALIZER = new IntegerDeserializer();
+    public int deserializeCount = 0;
+
+    public IntermittentlyFailingIntegerDeserializer() {}
+
+    @Override
+    public Integer deserialize(String topic, byte[] data) {
+      deserializeCount++;
+      if (deserializeCount % 10 == 0) {
+        throw new SerializationException("Expected Serialization Exception");
+      }
+      return INTEGER_DESERIALIZER.deserialize(topic, data);
+    }
+  }
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/errorhandling/ErrorHandler.java

@@ -17,6 +17,9 @@
  */
 package org.apache.beam.sdk.transforms.errorhandling;
 
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.Serializable;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
@@ -49,22 +52,24 @@
  *     <p>Simple usage with one DLQ
  *     <pre>{@code
  * PCollection<?> records = ...;
- * try (ErrorHandler<E,T> errorHandler = pipeline.registerErrorHandler(SomeSink.write())) {
- *  PCollection<?> results = records.apply(SomeIO.write().withDeadLetterQueue(errorHandler));
+ * try (BadRecordErrorHandler<T> errorHandler = pipeline.registerBadRecordErrorHandler(SomeSink.write())) {
+ *  PCollection<?> results = records.apply(SomeIO.write().withErrorHandler(errorHandler));
  * }
  * results.apply(SomeOtherTransform);
  * }</pre>
  *     Usage with multiple DLQ stages
  *     <pre>{@code
  * PCollection<?> records = ...;
- * try (ErrorHandler<E,T> errorHandler = pipeline.registerErrorHandler(SomeSink.write())) {
- *  PCollection<?> results = records.apply(SomeIO.write().withDeadLetterQueue(errorHandler))
- *                        .apply(OtherTransform.builder().withDeadLetterQueue(errorHandler));
+ * try (BadRecordErrorHandler<T> errorHandler = pipeline.registerBadRecordErrorHandler(SomeSink.write())) {
+ *  PCollection<?> results = records.apply(SomeIO.write().withErrorHandler(errorHandler))
+ *                        .apply(OtherTransform.builder().withErrorHandler(errorHandler));
  * }
  * results.apply(SomeOtherTransform);
  * }</pre>
+ *     This is marked as serializable despite never being needed on the runner, to enable it to be a
+ *     parameter of an Autovalue configured PTransform.
  */
-public interface ErrorHandler<ErrorT, OutputT extends POutput> extends AutoCloseable {
+public interface ErrorHandler<ErrorT, OutputT extends POutput> extends AutoCloseable, Serializable {
 
   void addErrorCollection(PCollection<ErrorT> errorCollection);
 
@@ -79,13 +84,16 @@ class PTransformErrorHandler<ErrorT, OutputT extends POutput>
     private static final Logger LOG = LoggerFactory.getLogger(PTransformErrorHandler.class);
     private final PTransform<PCollection<ErrorT>, OutputT> sinkTransform;
 
-    private final Pipeline pipeline;
+    // transient as Pipelines are not serializable
+    private final transient Pipeline pipeline;
 
     private final Coder<ErrorT> coder;
 
-    private final List<PCollection<ErrorT>> errorCollections = new ArrayList<>();
+    // transient as PCollections are not serializable
+    private transient List<PCollection<ErrorT>> errorCollections = new ArrayList<>();
 
-    private @Nullable OutputT sinkOutput = null;
+    // transient as PCollections are not serializable
+    private transient @Nullable OutputT sinkOutput = null;
 
     private boolean closed = false;
 
@@ -103,6 +111,12 @@ public PTransformErrorHandler(
       this.coder = coder;
     }
 
+    private void readObject(ObjectInputStream aInputStream)
+        throws ClassNotFoundException, IOException {
+      aInputStream.defaultReadObject();
+      errorCollections = new ArrayList<>();
+    }
+
     @Override
     public void addErrorCollection(PCollection<ErrorT> errorCollection) {
       errorCollections.add(errorCollection);

sdks/java/io/kafka/kafka-01103/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="0.11.0.3"
     undelimited="01103"
+    sdfCompatible=false
 }
 
 apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-100/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="1.0.0"
     undelimited="100"
+    sdfCompatible=false
 }
 
-apply from: "../kafka-integration-test.gradle"
+apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-111/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="1.1.1"
     undelimited="111"
+    sdfCompatible=false
 }
 
 apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-201/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="2.0.1"
     undelimited="201"
+    sdfCompatible=true
 }
 
 apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-211/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="2.1.1"
     undelimited="211"
+    sdfCompatible=true
 }
 
 apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-222/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="2.2.2"
     undelimited="222"
+    sdfCompatible=true
 }
 
 apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-231/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="2.3.1"
     undelimited="231"
+    sdfCompatible=true
 }
 
 apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-241/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="2.4.1"
     undelimited="241"
+    sdfCompatible=true
 }
 
 apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-251/build.gradle

@@ -18,6 +18,7 @@
 project.ext {
     delimited="2.5.1"
     undelimited="251"
+    sdfCompatible=true
 }
 
 apply from: "../kafka-integration-test.gradle"

sdks/java/io/kafka/kafka-integration-test.gradle

@@ -39,4 +39,4 @@ dependencies {
 
 configurations.create("kafkaVersion$undelimited")
 
-tasks.register("kafkaVersion${undelimited}BatchIT",KafkaTestUtilities.KafkaBatchIT, project.ext.delimited, project.ext.undelimited, false, configurations, project)
+tasks.register("kafkaVersion${undelimited}BatchIT",KafkaTestUtilities.KafkaBatchIT, project.ext.delimited, project.ext.undelimited, project.ext.sdfCompatible, configurations, project)