ReadFromKafkaDoFn: handle failed seek (#32456)

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java

@@ -441,10 +441,11 @@ public ProcessContinuation processElement(
  ConsumerSpEL.evaluateAssign(
  consumer, ImmutableList.of(kafkaSourceDescriptor.getTopicPartition()));
  long startOffset = tracker.currentRestriction().getFrom();
-
  long expectedOffset = startOffset;
  consumer.seek(kafkaSourceDescriptor.getTopicPartition(), startOffset);
  ConsumerRecords<byte[], byte[]> rawRecords = ConsumerRecords.empty();
+ long skippedRecords = 0L;
+ final Stopwatch sw = Stopwatch.createStarted();
 
  while (true) {
  rawRecords = poll(consumer, kafkaSourceDescriptor.getTopicPartition());
@@ -461,6 +462,36 @@ public ProcessContinuation processElement(
  return ProcessContinuation.resume();
  }
  for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+ // If the Kafka consumer returns a record with an offset that is already processed
+ // the record can be safely skipped. This is needed because there is a possibility
+ // that the seek() above fails to move the offset to the desired position. In which
+ // case poll() would return records that are already cnsumed.
+ if (rawRecord.offset() < startOffset) {
+ // If the start offset is not reached even after skipping the records for 10 seconds
+ // then the processing is stopped with a backoff to give the Kakfa server some time
+ // catch up.
+ if (sw.elapsed().getSeconds() > 10L) {
+ LOG.error(
+ "The expected offset ({}) was not reached even after"
+ + " skipping consumed records for 10 seconds. The offset we could"
+ + " reach was {}. The processing of this bundle will be attempted"
+ + " at a later time.",
+ expectedOffset,
+ rawRecord.offset());
+ return ProcessContinuation.resume()
+ .withResumeDelay(org.joda.time.Duration.standardSeconds(10L));
+ }
+ skippedRecords++;
+ continue;
+ }
+ if (skippedRecords > 0L) {
+ LOG.warn(
+ "{} records were skipped due to seek returning an"
+ + " earlier position than requested position of {}",
+ skippedRecords,
+ expectedOffset);
+ skippedRecords = 0L;
+ }
  if (!tracker.tryClaim(rawRecord.offset())) {
  return ProcessContinuation.stop();
  }

sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFnTest.java

@@ -110,6 +110,12 @@ public class ReadFromKafkaDoFnTest {
  private final ReadFromKafkaDoFn<String, String> exceptionDofnInstance =
  ReadFromKafkaDoFn.create(makeReadSourceDescriptor(exceptionConsumer), RECORDS);
 
+ private final SimpleMockKafkaConsumerWithBrokenSeek consumerWithBrokenSeek =
+ new SimpleMockKafkaConsumerWithBrokenSeek(OffsetResetStrategy.NONE, topicPartition);
+
+ private final ReadFromKafkaDoFn<String, String> dofnInstanceWithBrokenSeek =
+ ReadFromKafkaDoFn.create(makeReadSourceDescriptor(consumerWithBrokenSeek), RECORDS);
+
  private ReadSourceDescriptors<String, String> makeReadSourceDescriptor(
  Consumer<byte[], byte[]> kafkaMockConsumer) {
  return ReadSourceDescriptors.<String, String>read()
@@ -295,6 +301,17 @@ public synchronized long position(TopicPartition partition) {
  }
  }
 
+ private static class SimpleMockKafkaConsumerWithBrokenSeek extends SimpleMockKafkaConsumer {
+
+ public SimpleMockKafkaConsumerWithBrokenSeek(
+ OffsetResetStrategy offsetResetStrategy, TopicPartition topicPartition) {
+ super(offsetResetStrategy, topicPartition);
+ }
+
+ @Override
+ public synchronized void seek(TopicPartition partition, long offset) {}
+ }
+
  private static class MockMultiOutputReceiver implements MultiOutputReceiver {
 
  MockOutputReceiver<KV<KafkaSourceDescriptor, KafkaRecord<String, String>>> mockOutputReceiver =
@@ -377,6 +394,7 @@ private List<KV<KafkaSourceDescriptor, KafkaRecord<String, String>>> createExpec
  public void setUp() throws Exception {
  dofnInstance.setup();
  exceptionDofnInstance.setup();
+ dofnInstanceWithBrokenSeek.setup();
  consumer.reset();
  }
 
@@ -475,6 +493,24 @@ public void testProcessElement() throws Exception {
  receiver.getGoodRecords());
  }
 
+ @Test
+ public void testProcessElementWithEarlierOffset() throws Exception {
+ MockMultiOutputReceiver receiver = new MockMultiOutputReceiver();
+ consumerWithBrokenSeek.setNumOfRecordsPerPoll(6L);
+ consumerWithBrokenSeek.setCurrentPos(0L);
+ long startOffset = 3L;
+ OffsetRangeTracker tracker =
+ new OffsetRangeTracker(new OffsetRange(startOffset, startOffset + 3));
+ KafkaSourceDescriptor descriptor =
+ KafkaSourceDescriptor.of(topicPartition, null, null, null, null, null);
+ ProcessContinuation result =
+ dofnInstanceWithBrokenSeek.processElement(descriptor, tracker, null, receiver);
+ assertEquals(ProcessContinuation.stop(), result);
+ assertEquals(
+ createExpectedRecords(descriptor, startOffset, 3, "key", "value"),
+ receiver.getGoodRecords());
+ }
+
  @Test
  public void testRawSizeMetric() throws Exception {
  final int numElements = 1000;