You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by yh...@apache.org on 2023/05/28 01:00:45 UTC
[beam] branch master updated: [Flink][KafkaIO] Finalize checkpoints during stop/drain scenario (#26883)
This is an automated email from the ASF dual-hosted git repository.
yhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 68d991aef26 [Flink][KafkaIO] Finalize checkpoints during stop/drain scenario (#26883)
68d991aef26 is described below
commit 68d991aef2643d1da63654bad6e10c49118edcd4
Author: Kanishk Karanawat <kk...@gmail.com>
AuthorDate: Sat May 27 21:00:34 2023 -0400
[Flink][KafkaIO] Finalize checkpoints during stop/drain scenario (#26883)
* add fix for flink runner and kafkaio where offsets are not committed during stop
* update UnboundedSourceWrapperTest unit test
---------
Co-authored-by: Kanishk Karanawat <kk...@twitter.com>
---
.../streaming/io/UnboundedSourceWrapper.java | 59 +++++++++++-----------
.../streaming/io/UnboundedSourceWrapperTest.java | 6 +++
.../beam/sdk/io/kafka/KafkaUnboundedReader.java | 33 +++++++-----
3 files changed, 55 insertions(+), 43 deletions(-)
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
index 92d0652e11f..c0e196e8130 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
@@ -383,44 +383,46 @@ public class UnboundedSourceWrapper<OutputT, CheckpointMarkT extends UnboundedSo
@Override
public void snapshotState(FunctionSnapshotContext functionSnapshotContext) throws Exception {
if (!isRunning) {
+ // This implies that stop/drain is invoked and final checkpoint is triggered. This method
+ // should not be skipped in this scenario so that the notifyCheckpointComplete method is still
+ // invoked and performs the finalization step after commit is complete.
LOG.debug("snapshotState() called on closed source");
- } else {
+ }
- if (checkpointCoder == null) {
- // no checkpoint coder available in this source
- return;
- }
+ if (checkpointCoder == null) {
+ // no checkpoint coder available in this source
+ return;
+ }
- stateForCheckpoint.clear();
+ stateForCheckpoint.clear();
- long checkpointId = functionSnapshotContext.getCheckpointId();
+ long checkpointId = functionSnapshotContext.getCheckpointId();
- // we checkpoint the sources along with the CheckpointMarkT to ensure
- // than we have a correct mapping of checkpoints to sources when
- // restoring
- List<CheckpointMarkT> checkpointMarks = new ArrayList<>(localSplitSources.size());
+ // we checkpoint the sources along with the CheckpointMarkT to ensure
+ // than we have a correct mapping of checkpoints to sources when
+ // restoring
+ List<CheckpointMarkT> checkpointMarks = new ArrayList<>(localSplitSources.size());
- for (int i = 0; i < localSplitSources.size(); i++) {
- UnboundedSource<OutputT, CheckpointMarkT> source = localSplitSources.get(i);
- UnboundedSource.UnboundedReader<OutputT> reader = localReaders.get(i);
+ for (int i = 0; i < localSplitSources.size(); i++) {
+ UnboundedSource<OutputT, CheckpointMarkT> source = localSplitSources.get(i);
+ UnboundedSource.UnboundedReader<OutputT> reader = localReaders.get(i);
- @SuppressWarnings("unchecked")
- CheckpointMarkT mark = (CheckpointMarkT) reader.getCheckpointMark();
- checkpointMarks.add(mark);
- KV<UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT> kv = KV.of(source, mark);
- stateForCheckpoint.add(kv);
- }
+ @SuppressWarnings("unchecked")
+ CheckpointMarkT mark = (CheckpointMarkT) reader.getCheckpointMark();
+ checkpointMarks.add(mark);
+ KV<UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT> kv = KV.of(source, mark);
+ stateForCheckpoint.add(kv);
+ }
- // cleanup old pending checkpoints and add new checkpoint
- int diff = pendingCheckpoints.size() - MAX_NUMBER_PENDING_CHECKPOINTS;
- if (diff >= 0) {
- for (Iterator<Long> iterator = pendingCheckpoints.keySet().iterator(); diff >= 0; diff--) {
- iterator.next();
- iterator.remove();
- }
+ // cleanup old pending checkpoints and add new checkpoint
+ int diff = pendingCheckpoints.size() - MAX_NUMBER_PENDING_CHECKPOINTS;
+ if (diff >= 0) {
+ for (Iterator<Long> iterator = pendingCheckpoints.keySet().iterator(); diff >= 0; diff--) {
+ iterator.next();
+ iterator.remove();
}
- pendingCheckpoints.put(checkpointId, checkpointMarks);
}
+ pendingCheckpoints.put(checkpointId, checkpointMarks);
}
@Override
@@ -527,7 +529,6 @@ public class UnboundedSourceWrapper<OutputT, CheckpointMarkT extends UnboundedSo
@Override
public void notifyCheckpointComplete(long checkpointId) throws Exception {
-
List<CheckpointMarkT> checkpointMarks = pendingCheckpoints.get(checkpointId);
if (checkpointMarks != null) {
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapperTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapperTest.java
index 2b6a8639142..3eb1ef69836 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapperTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapperTest.java
@@ -422,6 +422,9 @@ public class UnboundedSourceWrapperTest {
assertTrue("Did not successfully read first batch of elements.", readFirstBatchOfElements);
+ // simulate pipeline stop/drain scenario, where sources are closed first.
+ sourceOperator.cancel();
+
// draw a snapshot
OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
@@ -431,6 +434,9 @@ public class UnboundedSourceWrapperTest {
testHarness.notifyOfCompletedCheckpoint(0);
assertEquals(flinkWrapper.getLocalSplitSources().size(), finalizeList.size());
+ // stop the pipeline
+ testHarness.close();
+
// create a completely new source but restore from the snapshot
TestCountingSource restoredSource = new TestCountingSource(numElements);
UnboundedSourceWrapper<KV<Integer, Integer>, TestCountingSource.CounterMark>
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
index 019a7c746d7..ab92370a357 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
@@ -550,10 +550,8 @@ class KafkaUnboundedReader<K, V> extends UnboundedReader<KafkaRecord<K, V>> {
records, RECORDS_ENQUEUE_POLL_TIMEOUT.getMillis(), TimeUnit.MILLISECONDS)) {
records = ConsumerRecords.empty();
}
- KafkaCheckpointMark checkpointMark = finalizedCheckpointMark.getAndSet(null);
- if (checkpointMark != null) {
- commitCheckpointMark(checkpointMark);
- }
+
+ commitCheckpointMark();
} catch (InterruptedException e) {
LOG.warn("{}: consumer thread is interrupted", this, e); // not expected
break;
@@ -569,17 +567,21 @@ class KafkaUnboundedReader<K, V> extends UnboundedReader<KafkaRecord<K, V>> {
}
}
- private void commitCheckpointMark(KafkaCheckpointMark checkpointMark) {
- LOG.debug("{}: Committing finalized checkpoint {}", this, checkpointMark);
- Consumer<byte[], byte[]> consumer = Preconditions.checkStateNotNull(this.consumer);
+ private void commitCheckpointMark() {
+ KafkaCheckpointMark checkpointMark = finalizedCheckpointMark.getAndSet(null);
- consumer.commitSync(
- checkpointMark.getPartitions().stream()
- .filter(p -> p.getNextOffset() != UNINITIALIZED_OFFSET)
- .collect(
- Collectors.toMap(
- p -> new TopicPartition(p.getTopic(), p.getPartition()),
- p -> new OffsetAndMetadata(p.getNextOffset()))));
+ if (checkpointMark != null) {
+ LOG.debug("{}: Committing finalized checkpoint {}", this, checkpointMark);
+ Consumer<byte[], byte[]> consumer = Preconditions.checkStateNotNull(this.consumer);
+
+ consumer.commitSync(
+ checkpointMark.getPartitions().stream()
+ .filter(p -> p.getNextOffset() != UNINITIALIZED_OFFSET)
+ .collect(
+ Collectors.toMap(
+ p -> new TopicPartition(p.getTopic(), p.getPartition()),
+ p -> new OffsetAndMetadata(p.getNextOffset()))));
+ }
}
/**
@@ -733,6 +735,9 @@ class KafkaUnboundedReader<K, V> extends UnboundedReader<KafkaRecord<K, V>> {
}
}
+ // Commit any pending finalized checkpoint before shutdown.
+ commitCheckpointMark();
+
Closeables.close(keyDeserializerInstance, true);
Closeables.close(valueDeserializerInstance, true);