This is an automated email from the ASF dual-hosted git repository.
MartijnVisser pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-connector-kafka.git
The following commit(s) were added to refs/heads/main by this push:
new 91e01a56 [FLINK-40137][Connectors/Kafka] Extract ReaderRecoveryGate
from DynamicKafkaSourceEnumerator
91e01a56 is described below
commit 91e01a56ef6f5ea216a24c854a5a218616613f0b
Author: Sylwester Lachiewicz <[email protected]>
AuthorDate: Mon Jul 13 20:44:35 2026 +0200
[FLINK-40137][Connectors/Kafka] Extract ReaderRecoveryGate from
DynamicKafkaSourceEnumerator
The enumerator accumulated three interacting pieces of recovery gating
state (initial registration pending, reported splits pending
redistribution, deferred metadata update readers). Extract them into a
dedicated, unit-tested ReaderRecoveryGate so the gating state and its
lifecycle have an explicit name and structure. No behavior change.
Generated-by: Claude Fable 5
---
.../enumerator/DynamicKafkaSourceEnumerator.java | 42 ++----
.../source/enumerator/ReaderRecoveryGate.java | 132 +++++++++++++++++
.../source/enumerator/ReaderRecoveryGateTest.java | 156 +++++++++++++++++++++
3 files changed, 301 insertions(+), 29 deletions(-)
diff --git
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumerator.java
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumerator.java
index 32bd3b51..ea38c9e6 100644
---
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumerator.java
+++
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumerator.java
@@ -114,9 +114,7 @@ public class DynamicKafkaSourceEnumerator
private Map<String, DynamicKafkaSourceEnumState.RetainedClusterState>
retainedClusterEnumeratorStates;
private boolean firstDiscoveryComplete;
- private boolean initialReaderRegistrationPending;
- private final Map<Integer, List<DynamicKafkaSourceSplit>>
pendingReportedSplitsByReader;
- private final Set<Integer> pendingMetadataUpdateReaders;
+ private final ReaderRecoveryGate readerRecoveryGate;
public DynamicKafkaSourceEnumerator(
KafkaStreamSubscriber kafkaStreamSubscriber,
@@ -218,10 +216,8 @@ public class DynamicKafkaSourceEnumerator
runnable,
"dynamic-kafka-enumerator-closing-worker"));
this.asynchronousEnumeratorCloseFailure = new AtomicReference<>();
this.splitAssignmentStrategy =
createSplitAssignmentStrategy(properties);
- this.initialReaderRegistrationPending =
- hasRestoredEnumeratorState(dynamicKafkaSourceEnumState);
- this.pendingReportedSplitsByReader = new HashMap<>();
- this.pendingMetadataUpdateReaders = new HashSet<>();
+ this.readerRecoveryGate =
+ new
ReaderRecoveryGate(hasRestoredEnumeratorState(dynamicKafkaSourceEnumState));
if
(!dynamicKafkaSourceEnumState.getClusterEnumeratorStates().isEmpty()) {
logger.info("Dynamic Kafka source restored from checkpointed
enumerator state");
@@ -559,7 +555,7 @@ public class DynamicKafkaSourceEnumerator
/** NOTE: Must run on coordinator thread. */
private void sendMetadataUpdateEventToAvailableReaders() {
if (shouldDeferMetadataUpdateEvents()) {
-
pendingMetadataUpdateReaders.addAll(enumContext.registeredReaders().keySet());
+
readerRecoveryGate.deferMetadataUpdates(enumContext.registeredReaders().keySet());
return;
}
@@ -769,11 +765,8 @@ public class DynamicKafkaSourceEnumerator
logger.debug("Adding reader {}", subtaskId);
ReaderInfo readerInfo = enumContext.registeredReaders().get(subtaskId);
if (readerInfo != null) {
- List<DynamicKafkaSourceSplit> reportedSplits =
- readerInfo.getReportedSplitsOnRegistration();
- if (!reportedSplits.isEmpty()) {
- pendingReportedSplitsByReader.put(subtaskId, new
ArrayList<>(reportedSplits));
- }
+ readerRecoveryGate.recordReportedSplits(
+ subtaskId, readerInfo.getReportedSplitsOnRegistration());
}
if (tryCompletePendingReaderRegistration()) {
@@ -785,19 +778,15 @@ public class DynamicKafkaSourceEnumerator
}
private boolean tryCompletePendingReaderRegistration() {
- boolean hasPendingRecovery =
- initialReaderRegistrationPending ||
!pendingReportedSplitsByReader.isEmpty();
- if (!hasPendingRecovery) {
+ if (!readerRecoveryGate.hasPendingRecovery()) {
return false;
}
if (!firstDiscoveryComplete || !allReadersRegistered()) {
return true;
}
- if (initialReaderRegistrationPending) {
- initialReaderRegistrationPending = false;
- }
- if (!pendingReportedSplitsByReader.isEmpty()) {
+ readerRecoveryGate.markInitialRegistrationComplete();
+ if (readerRecoveryGate.hasReportedSplits()) {
reassignReportedSplits();
} else {
flushPendingSplitAssignmentsForRegisteredReaders();
@@ -831,7 +820,7 @@ public class DynamicKafkaSourceEnumerator
long currentTimeMillis = System.currentTimeMillis();
for (Entry<Integer, List<DynamicKafkaSourceSplit>> readerSplits :
- new TreeMap<>(pendingReportedSplitsByReader).entrySet()) {
+ readerRecoveryGate.drainReportedSplits().entrySet()) {
int readerId = readerSplits.getKey();
for (DynamicKafkaSourceSplit split : readerSplits.getValue()) {
if (isSplitActive(split)) {
@@ -876,19 +865,14 @@ public class DynamicKafkaSourceEnumerator
if (!retainedSplitsByReader.isEmpty()) {
enumContext.assignSplits(new
SplitsAssignment<>(retainedSplitsByReader));
}
- pendingReportedSplitsByReader.clear();
}
private boolean shouldDeferMetadataUpdateEvents() {
- return initialReaderRegistrationPending
- || (!pendingReportedSplitsByReader.isEmpty() &&
!allReadersRegistered());
+ return
readerRecoveryGate.shouldDeferMetadataUpdateEvents(allReadersRegistered());
}
private void flushPendingMetadataUpdateEvents() {
- List<Integer> readers = new ArrayList<>(pendingMetadataUpdateReaders);
- Collections.sort(readers);
- pendingMetadataUpdateReaders.clear();
- for (int readerId : readers) {
+ for (int readerId :
readerRecoveryGate.drainDeferredMetadataUpdateReaders()) {
if (enumContext.registeredReaders().containsKey(readerId)) {
sendMetadataUpdateEvent(readerId);
}
@@ -1044,7 +1028,7 @@ public class DynamicKafkaSourceEnumerator
if (enumContext.registeredReaders().containsKey(subtaskId)) {
if (shouldDeferMetadataUpdateEvents()) {
- pendingMetadataUpdateReaders.add(subtaskId);
+ readerRecoveryGate.deferMetadataUpdate(subtaskId);
} else {
sendMetadataUpdateEvent(subtaskId);
}
diff --git
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/ReaderRecoveryGate.java
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/ReaderRecoveryGate.java
new file mode 100644
index 00000000..5707a078
--- /dev/null
+++
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/ReaderRecoveryGate.java
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.kafka.dynamic.source.enumerator;
+
+import org.apache.flink.annotation.Internal;
+import
org.apache.flink.connector.kafka.dynamic.source.split.DynamicKafkaSourceSplit;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.Set;
+import java.util.TreeMap;
+
+/**
+ * Tracks the recovery-time reader registration state of the {@link
DynamicKafkaSourceEnumerator}.
+ *
+ * <p>The gate is armed by two recovery triggers:
+ *
+ * <ul>
+ * <li><b>Enumerator restore from checkpoint</b>: armed at construction
({@code
+ * restoredFromCheckpoint = true}). Split assignment and metadata update
events must be
+ * deferred until the first metadata discovery has completed and every
reader has
+ * (re-)registered, so that restored reader splits can be redistributed
consistently.
+ * <li><b>Reader re-registration after partial failover</b>: armed on a
running enumerator when a
+ * re-registering reader reports checkpointed splits ({@link
#recordReportedSplits} with
+ * non-empty splits while initial reader registration is already
complete). Metadata update
+ * events are deferred until all readers have registered again and
reported splits are
+ * redistributed.
+ * </ul>
+ *
+ * <p>This class owns that gating state; the enumerator remains responsible
for acting on it.
+ *
+ * <p>This class is not thread-safe and must only be used from the coordinator
thread: the
+ * enumerator methods that touch it are called by the {@code
SourceCoordinator}, and metadata
+ * discovery results are handed back through {@code runInCoordinatorThread}.
+ */
+@Internal
+class ReaderRecoveryGate {
+
+ /** Set on restore; cleared once all readers have registered after the
first discovery. */
+ private boolean initialReaderRegistrationPending;
+
+ /** Splits reported by readers on registration, pending redistribution. */
+ private final Map<Integer, List<DynamicKafkaSourceSplit>>
pendingReportedSplitsByReader =
+ new HashMap<>();
+
+ /** Readers whose metadata update events were deferred during recovery. */
+ private final Set<Integer> pendingMetadataUpdateReaders = new HashSet<>();
+
+ ReaderRecoveryGate(boolean restoredFromCheckpoint) {
+ this.initialReaderRegistrationPending = restoredFromCheckpoint;
+ }
+
+ /** Records splits a reader reported on registration; an empty report is
ignored. */
+ void recordReportedSplits(int subtaskId, List<DynamicKafkaSourceSplit>
reportedSplits) {
+ if (!reportedSplits.isEmpty()) {
+ pendingReportedSplitsByReader.put(subtaskId, new
ArrayList<>(reportedSplits));
+ }
+ }
+
+ /** Whether recovery gating is active and registrations must be deferred.
*/
+ boolean hasPendingRecovery() {
+ return initialReaderRegistrationPending ||
!pendingReportedSplitsByReader.isEmpty();
+ }
+
+ /**
+ * Whether metadata update events must be deferred instead of sent, given
the current reader
+ * registration completeness.
+ */
+ boolean shouldDeferMetadataUpdateEvents(boolean allReadersRegistered) {
+ return initialReaderRegistrationPending
+ || (!pendingReportedSplitsByReader.isEmpty() &&
!allReadersRegistered);
+ }
+
+ void deferMetadataUpdate(int readerId) {
+ pendingMetadataUpdateReaders.add(readerId);
+ }
+
+ void deferMetadataUpdates(Collection<Integer> readerIds) {
+ pendingMetadataUpdateReaders.addAll(readerIds);
+ }
+
+ /** Returns the deferred metadata update readers in ascending order and
clears them. */
+ List<Integer> drainDeferredMetadataUpdateReaders() {
+ List<Integer> readers = new ArrayList<>(pendingMetadataUpdateReaders);
+ Collections.sort(readers);
+ pendingMetadataUpdateReaders.clear();
+ return readers;
+ }
+
+ void markInitialRegistrationComplete() {
+ initialReaderRegistrationPending = false;
+ }
+
+ boolean hasReportedSplits() {
+ return !pendingReportedSplitsByReader.isEmpty();
+ }
+
+ /**
+ * Returns the reported splits ordered by reader id and clears the pending
state.
+ *
+ * <p>Note: the pending state is cleared eagerly, so the gate must not be
consulted for pending
+ * reported splits while reassigning (e.g. from {@code
handleNoMoreSplits}).
+ */
+ NavigableMap<Integer, List<DynamicKafkaSourceSplit>> drainReportedSplits()
{
+ NavigableMap<Integer, List<DynamicKafkaSourceSplit>>
reportedSplitsByReader =
+ new TreeMap<>(pendingReportedSplitsByReader);
+ pendingReportedSplitsByReader.clear();
+ return reportedSplitsByReader;
+ }
+}
diff --git
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/ReaderRecoveryGateTest.java
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/ReaderRecoveryGateTest.java
new file mode 100644
index 00000000..e6456ed4
--- /dev/null
+++
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/ReaderRecoveryGateTest.java
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.kafka.dynamic.source.enumerator;
+
+import
org.apache.flink.connector.kafka.dynamic.source.split.DynamicKafkaSourceSplit;
+import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit;
+
+import org.apache.kafka.common.TopicPartition;
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.NavigableMap;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link ReaderRecoveryGate}. */
+class ReaderRecoveryGateTest {
+
+ @Test
+ void testFreshStartHasNoPendingRecovery() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+
+ assertThat(gate.hasPendingRecovery()).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+ assertThat(gate.hasReportedSplits()).isFalse();
+ }
+
+ @Test
+ void testRestoredStartGatesUntilInitialRegistrationCompletes() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
+
+ assertThat(gate.hasPendingRecovery()).isTrue();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isTrue();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isTrue();
+
+ gate.markInitialRegistrationComplete();
+
+ assertThat(gate.hasPendingRecovery()).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+ }
+
+ @Test
+ void
testCompletingInitialRegistrationDoesNotReleaseGateWithPendingReportedSplits() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
+ gate.recordReportedSplits(1, Collections.singletonList(split("topic",
0)));
+
+ gate.markInitialRegistrationComplete();
+
+ assertThat(gate.hasPendingRecovery()).isTrue();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isTrue();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+
+ gate.drainReportedSplits();
+
+ assertThat(gate.hasPendingRecovery()).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+ }
+
+ @Test
+ void testReportedSplitsGateUntilAllReadersRegistered() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ gate.recordReportedSplits(1, Collections.singletonList(split("topic",
0)));
+
+ assertThat(gate.hasPendingRecovery()).isTrue();
+ assertThat(gate.hasReportedSplits()).isTrue();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isTrue();
+ assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+ }
+
+ @Test
+ void testEmptyReportedSplitsAreIgnored() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ gate.recordReportedSplits(1, Collections.emptyList());
+
+ assertThat(gate.hasPendingRecovery()).isFalse();
+ assertThat(gate.hasReportedSplits()).isFalse();
+ }
+
+ @Test
+ void testRepeatedReportForSameReaderReplacesPreviousReport() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ DynamicKafkaSourceSplit firstReport = split("topic", 0);
+ DynamicKafkaSourceSplit secondReport = split("topic", 1);
+ gate.recordReportedSplits(1, Collections.singletonList(firstReport));
+ gate.recordReportedSplits(1, Collections.singletonList(secondReport));
+
+ // A reader that registers again keeps one entry holding its latest
report, so the
+ // duplicate-owner check in reassignReportedSplits is not tripped by
its own re-report.
+ NavigableMap<Integer, List<DynamicKafkaSourceSplit>> drained =
gate.drainReportedSplits();
+
+ assertThat(drained.keySet()).containsExactly(1);
+ assertThat(drained.get(1)).containsExactly(secondReport);
+ }
+
+ @Test
+ void testEmptyRepeatedReportRetainsPreviousReport() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ DynamicKafkaSourceSplit firstReport = split("topic", 0);
+ gate.recordReportedSplits(1, Collections.singletonList(firstReport));
+ gate.recordReportedSplits(1, Collections.emptyList());
+
+
assertThat(gate.drainReportedSplits().get(1)).containsExactly(firstReport);
+ }
+
+ @Test
+ void testDrainReportedSplitsReturnsReaderOrderAndClears() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ DynamicKafkaSourceSplit splitReader2 = split("topic", 2);
+ DynamicKafkaSourceSplit splitReader0 = split("topic", 0);
+ gate.recordReportedSplits(2, Collections.singletonList(splitReader2));
+ gate.recordReportedSplits(0, Collections.singletonList(splitReader0));
+
+ NavigableMap<Integer, List<DynamicKafkaSourceSplit>> drained =
gate.drainReportedSplits();
+
+ assertThat(drained.keySet()).containsExactly(0, 2);
+ assertThat(drained.get(0)).containsExactly(splitReader0);
+ assertThat(drained.get(2)).containsExactly(splitReader2);
+ assertThat(gate.hasReportedSplits()).isFalse();
+ assertThat(gate.drainReportedSplits()).isEmpty();
+ }
+
+ @Test
+ void testDrainDeferredMetadataUpdateReadersReturnsSortedAndClears() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
+ gate.deferMetadataUpdate(3);
+ gate.deferMetadataUpdates(Arrays.asList(1, 2, 3));
+
+
assertThat(gate.drainDeferredMetadataUpdateReaders()).containsExactly(1, 2, 3);
+ assertThat(gate.drainDeferredMetadataUpdateReaders()).isEmpty();
+ }
+
+ private static DynamicKafkaSourceSplit split(String topic, int partition) {
+ return new DynamicKafkaSourceSplit(
+ "cluster0", new KafkaPartitionSplit(new TopicPartition(topic,
partition), 0L));
+ }
+}
