Ensure expectedOffset advances immediately after a claim

Author: sjvanrossum

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

@@ -504,6 +504,7 @@ public ProcessContinuation processElement(
           if (!tracker.tryClaim(rawRecord.offset())) {
             return ProcessContinuation.stop();
           }
+          expectedOffset = rawRecord.offset() + 1;
           try {
             KafkaRecord<K, V> kafkaRecord =
                 new KafkaRecord<>(
@@ -520,7 +521,6 @@ public ProcessContinuation processElement(
                     + (rawRecord.value() == null ? 0 : rawRecord.value().length);
             avgRecordSize.update(recordSize);
             rawSizes.update(recordSize);
-            expectedOffset = rawRecord.offset() + 1;
             Instant outputTimestamp;
             // The outputTimestamp and watermark will be computed by timestampPolicy, where the
             // WatermarkEstimator should be a manual one.