This is an automated email from the ASF dual-hosted git repository.
bowenli86 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 9d114deb [FLINK-40710][Connectors/Kafka] Preserve Dynamic Kafka
partition ownership during reader recovery (#323)
9d114deb is described below
commit 9d114deb30c6bfe7f074f657ef3ef964b028d7e2
Author: bowenli86 <[email protected]>
AuthorDate: Thu Oct 22 19:22:15 2026 -0700
[FLINK-40710][Connectors/Kafka] Preserve Dynamic Kafka partition ownership
during reader recovery (#323)
* [FLINK-40710][Connectors/Kafka] Preserve Dynamic Kafka partition
ownership during reader recovery
Use cluster and topic-partition keys for current ownership so retained
splits cannot collide with active partitions from another cluster. Batch
initial recovery returns once per cluster to avoid quadratic pending scans.
---
.../c0d94764-76a0-4c50-b617-70b1754c4612 | 1 -
.../enumerator/DynamicKafkaSourceEnumerator.java | 177 +++-
.../source/enumerator/ReaderRecoveryGate.java | 75 +-
.../enumerator/StoppableKafkaEnumContextProxy.java | 20 +-
.../source/enumerator/KafkaSourceEnumerator.java | 34 +-
.../DynamicKafkaSourceEnumeratorRecoveryTest.java | 1087 ++++++++++++++++++++
.../DynamicKafkaSourceEnumeratorTest.java | 11 +-
.../source/enumerator/ReaderRecoveryGateTest.java | 102 +-
...ceEnumeratorPendingOwnershipRegressionTest.java | 149 +++
...aSourceEnumeratorRecoveryDiscoveryRaceTest.java | 230 +++++
10 files changed, 1804 insertions(+), 82 deletions(-)
diff --git
a/flink-connector-kafka/archunit-violations/c0d94764-76a0-4c50-b617-70b1754c4612
b/flink-connector-kafka/archunit-violations/c0d94764-76a0-4c50-b617-70b1754c4612
index 0a64cb1b..2429b474 100644
---
a/flink-connector-kafka/archunit-violations/c0d94764-76a0-4c50-b617-70b1754c4612
+++
b/flink-connector-kafka/archunit-violations/c0d94764-76a0-4c50-b617-70b1754c4612
@@ -46,7 +46,6 @@ Method
<org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumStateS
Method
<org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumerator.deepCopyProperties(java.util.Properties,
java.util.Properties)> is annotated with
<org.apache.flink.annotation.VisibleForTesting> in
(KafkaSourceEnumerator.java:0)
Method
<org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumerator.getPartitionChange(java.util.Set,
boolean)> is annotated with <org.apache.flink.annotation.VisibleForTesting> in
(KafkaSourceEnumerator.java:0)
Method
<org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumerator.getPendingPartitionSplitAssignment()>
is annotated with <org.apache.flink.annotation.VisibleForTesting> in
(KafkaSourceEnumerator.java:0)
-Method
<org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumerator.getSplitOwner(org.apache.kafka.common.TopicPartition,
int)> is annotated with <org.apache.flink.annotation.VisibleForTesting> in
(KafkaSourceEnumerator.java:0)
Method
<org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumerator.topicIntegrityCheckEnabled()>
is annotated with <org.apache.flink.annotation.VisibleForTesting> in
(KafkaSourceEnumerator.java:0)
Method
<org.apache.flink.connector.kafka.source.reader.KafkaPartitionSplitReader.consumer()>
is annotated with <org.apache.flink.annotation.VisibleForTesting> in
(KafkaPartitionSplitReader.java:0)
Method
<org.apache.flink.connector.kafka.source.reader.KafkaPartitionSplitReader.getPollTimeout()>
is annotated with <org.apache.flink.annotation.VisibleForTesting> in
(KafkaPartitionSplitReader.java:0)
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 60a2fe9b..487a7400 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
@@ -26,6 +26,7 @@ import org.apache.flink.api.connector.source.SourceEvent;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.apache.flink.api.connector.source.SplitsAssignment;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.connector.kafka.dynamic.metadata.ClusterMetadata;
import org.apache.flink.connector.kafka.dynamic.metadata.KafkaMetadataService;
import org.apache.flink.connector.kafka.dynamic.metadata.KafkaStream;
@@ -47,6 +48,7 @@ import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.OffsetResetStrategy;
import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -117,6 +119,8 @@ public class DynamicKafkaSourceEnumerator
private final ReaderRecoveryGate readerRecoveryGate;
private final Set<Integer> readersWithNoMoreSplits = new HashSet<>();
private boolean splitAssignmentInProgress;
+ // Reader checkpoint reports may refer to owners from before recovery
redistribution.
+ private final Map<Tuple2<String, TopicPartition>, Integer>
currentSplitOwners = new HashMap<>();
public DynamicKafkaSourceEnumerator(
KafkaStreamSubscriber kafkaStreamSubscriber,
@@ -413,7 +417,7 @@ public class DynamicKafkaSourceEnumerator
private void handleNoMoreSplits() {
// A cluster callback may run before other clusters have assigned
their splits.
- if (splitAssignmentInProgress ||
readerRecoveryGate.hasPendingRecovery()) {
+ if (splitAssignmentInProgress ||
readerRecoveryGate.isReconciliationPending()) {
return;
}
if (Boundedness.BOUNDED.equals(boundedness)) {
@@ -520,6 +524,7 @@ public class DynamicKafkaSourceEnumerator
// create enumerators
Set<String> activeSplitIds = new HashSet<>();
+ Set<Tuple2<String, TopicPartition>> knownPartitions = new HashSet<>();
for (Entry<String, Set<String>> activeClusterTopics :
latestClusterTopicsMap.entrySet()) {
String kafkaClusterId = activeClusterTopics.getKey();
KafkaSourceEnumState kafkaSourceEnumState =
@@ -539,9 +544,14 @@ public class DynamicKafkaSourceEnumerator
Set<SplitAndAssignmentStatus> partitions =
filterStateByTopics(kafkaSourceEnumState,
activeClusterTopics.getValue());
partitions.forEach(
- splitStatus ->
- activeSplitIds.add(
- toDynamicSplitId(kafkaClusterId,
splitStatus.split())));
+ splitStatus -> {
+ activeSplitIds.add(
+ toDynamicSplitId(kafkaClusterId,
splitStatus.split()));
+ knownPartitions.add(
+ Tuple2.of(
+ kafkaClusterId,
+
splitStatus.split().getTopicPartition()));
+ });
newKafkaSourceEnumState =
new KafkaSourceEnumState(
@@ -564,6 +574,21 @@ public class DynamicKafkaSourceEnumerator
}
splitAssignmentStrategy.onMetadataRefresh(activeSplitIds);
+ // Readers retain removed-cluster splits and reactivate them without a
new assignment.
+ retainedClusterEnumeratorStates.forEach(
+ (clusterId, retainedState) ->
+ retainedState
+ .getKafkaSourceEnumState()
+ .splits()
+ .forEach(
+ splitStatus ->
+ knownPartitions.add(
+ Tuple2.of(
+ clusterId,
+ splitStatus
+
.split()
+
.getTopicPartition()))));
+ currentSplitOwners.keySet().retainAll(knownPartitions);
startAllEnumerators();
tryCompletePendingReaderRegistration();
}
@@ -638,8 +663,21 @@ public class DynamicKafkaSourceEnumerator
kafkaClusterId,
kafkaMetadataService,
signalNoMoreSplitsCallback);
- KafkaSourceEnumerator.SplitOwnerSelector splitOwnerSelector =
+
context.setReaderReadyForAssignment(readerRecoveryGate::isReaderReadyForAssignment);
+ context.setSplitAssignmentListener(
+ (reader, split) ->
currentSplitOwners.put(toPartitionKey(split), reader));
+ KafkaSourceEnumerator.SplitOwnerSelector strategyOwnerSelector =
splitAssignmentStrategy.createSplitOwnerSelector(kafkaClusterId);
+ KafkaSourceEnumerator.SplitOwnerSelector splitOwnerSelector =
+ (split, numReaders) ->
+ currentSplitOwners.computeIfAbsent(
+ Tuple2.of(kafkaClusterId,
split.getTopicPartition()),
+ ignored ->
+ strategyOwnerSelector == null
+ ?
KafkaSourceEnumerator.getSplitOwner(
+
split.getTopicPartition(), numReaders)
+ :
strategyOwnerSelector.getSplitOwner(
+ split, numReaders));
SplitEnumeratorContext<KafkaPartitionSplit> assignmentContext =
splitAssignmentStrategy.createAssignmentContext(kafkaClusterId, context);
@@ -759,10 +797,26 @@ public class DynamicKafkaSourceEnumerator
@Override
public void addSplitsBack(List<DynamicKafkaSourceSplit> splits, int
subtaskId) {
logger.debug("Adding splits back for {}", subtaskId);
+ List<DynamicKafkaSourceSplit> ownedSplits =
+ splits.stream()
+ .filter(split -> isCurrentSplitOwner(split, subtaskId))
+ .collect(Collectors.toList());
+ List<DynamicKafkaSourceSplit> returnedRetainedSplits = new
ArrayList<>();
+ long currentTimeMillis = System.currentTimeMillis();
+ for (DynamicKafkaSourceSplit split : ownedSplits) {
+ if (!isClusterActive(split.getKafkaClusterId())) {
+ DynamicKafkaSourceSplit retainedSplit =
+ getRetainedReportedSplit(split, currentTimeMillis);
+ if (retainedSplit != null) {
+ returnedRetainedSplits.add(retainedSplit);
+ }
+ }
+ }
+ readerRecoveryGate.recordReturnedSplits(subtaskId,
returnedRetainedSplits);
runWithSplitAssignmentInProgress(
() -> {
- splitAssignmentStrategy.onSplitsBack(splits, subtaskId);
- addSplitsBackToClusterEnumerators(splits, subtaskId,
false);
+ splitAssignmentStrategy.onSplitsBack(ownedSplits,
subtaskId);
+ addSplitsBackToClusterEnumerators(ownedSplits, subtaskId,
false);
});
handleNoMoreSplits();
}
@@ -803,17 +857,24 @@ public class DynamicKafkaSourceEnumerator
public void addReader(int subtaskId) {
logger.debug("Adding reader {}", subtaskId);
readersWithNoMoreSplits.remove(subtaskId);
+ readerRecoveryGate.startReaderRegistration(subtaskId);
ReaderInfo readerInfo = enumContext.registeredReaders().get(subtaskId);
- if (readerInfo != null) {
- readerRecoveryGate.recordReportedSplits(
- subtaskId, readerInfo.getReportedSplitsOnRegistration());
- }
+ List<DynamicKafkaSourceSplit> restoredSplits =
+ readerInfo == null
+ ? Collections.emptyList()
+ : readerInfo.getReportedSplitsOnRegistration();
+ readerRecoveryGate.recordReportedSplits(subtaskId, restoredSplits);
if (tryCompletePendingReaderRegistration()) {
return;
}
- runWithSplitAssignmentInProgress(() ->
addReaderToClusterEnumerators(subtaskId));
+ runWithSplitAssignmentInProgress(
+ () -> {
+ reassignReportedSplitsFromRestartedReader(subtaskId,
restoredSplits);
+ readerRecoveryGate.completeReaderRegistration(subtaskId);
+ addReaderToClusterEnumerators(subtaskId);
+ });
handleNoMoreSplits();
}
@@ -828,12 +889,13 @@ public class DynamicKafkaSourceEnumerator
// Draining the gate clears its pending state before reassignment
finishes.
runWithSplitAssignmentInProgress(
() -> {
- readerRecoveryGate.markInitialRegistrationComplete();
if (readerRecoveryGate.hasReportedSplits()) {
reassignReportedSplits();
- } else {
- flushPendingSplitAssignmentsForRegisteredReaders();
}
+ // Returned splits must replace pending copies before any
child can flush
+ // assignments.
+ readerRecoveryGate.markInitialRegistrationComplete();
+ flushPendingSplitAssignmentsForRegisteredReaders();
});
handleNoMoreSplits();
flushPendingMetadataUpdateEvents();
@@ -907,21 +969,87 @@ public class DynamicKafkaSourceEnumerator
.map(reportedSplit -> reportedSplit.split)
.collect(Collectors.toList());
splitAssignmentStrategy.onRecoveredSplits(activeSplits,
enumContext.currentParallelism());
+ activeSplits.forEach(split ->
currentSplitOwners.remove(toPartitionKey(split)));
+
+ if (!activeSplits.isEmpty()) {
+ // With the initial gate closed, registeredReaders() is empty, so
the reader ID
+ // has no effect here. The flush after reconciliation covers all
readers. Enqueue
+ // one batch per cluster to avoid scanning pending copies once per
reported split.
+ int reportingReader =
activeReportedSplits.values().iterator().next().readerId;
+ addSplitsBackToClusterEnumerators(activeSplits, reportingReader,
true);
+ }
+
+ assignRetainedSplits(retainedSplitsByReader);
+ }
- for (ReportedSplit reportedSplit : activeReportedSplits.values()) {
- addSplitsBackToClusterEnumerators(
- Collections.singletonList(reportedSplit.split),
reportedSplit.readerId, true);
+ private void reassignReportedSplitsFromRestartedReader(
+ int readerId, List<DynamicKafkaSourceSplit> restoredSplits) {
+ List<DynamicKafkaSourceSplit> activeSplits = new ArrayList<>();
+ List<DynamicKafkaSourceSplit> retainedSplits = new ArrayList<>();
+ long currentTimeMillis = System.currentTimeMillis();
+ Map<Tuple2<String, TopicPartition>, DynamicKafkaSourceSplit>
reportedSplits =
+ new HashMap<>();
+ restoredSplits.forEach(split ->
reportedSplits.put(toPartitionKey(split), split));
+ // Returned assignments may have moved to this reader since its
restored checkpoint.
+ readerRecoveryGate
+ .drainReportedSplits(readerId)
+ .forEach(split -> reportedSplits.put(toPartitionKey(split),
split));
+ List<DynamicKafkaSourceSplit> orderedSplits = new
ArrayList<>(reportedSplits.values());
+ orderedSplits.sort(
+ Comparator.comparing(DynamicKafkaSourceSplit::splitId)
+
.thenComparing(DynamicKafkaSourceSplit::getKafkaClusterId));
+ for (DynamicKafkaSourceSplit split : orderedSplits) {
+ if (!isCurrentSplitOwner(split, readerId)) {
+ continue;
+ }
+ if (isSplitActive(split)) {
+ // Active metadata also reactivates retained reader reports.
Reconcile their
+ // physical ownership before sending a historical shadow back
to a reader.
+ activeSplits.add(split.clearRetention());
+ } else {
+ DynamicKafkaSourceSplit retainedSplit =
+ getRetainedReportedSplit(split, currentTimeMillis);
+ if (retainedSplit != null) {
+ retainedSplits.add(retainedSplit);
+ }
+ }
}
+ splitAssignmentStrategy.onSplitsBack(activeSplits, readerId);
+ addSplitsBackToClusterEnumerators(activeSplits, readerId, true);
+ if (!retainedSplits.isEmpty()) {
+ assignRetainedSplits(Collections.singletonMap(readerId,
retainedSplits));
+ }
+ }
- flushPendingSplitAssignmentsForRegisteredReaders();
+ private void assignRetainedSplits(Map<Integer,
List<DynamicKafkaSourceSplit>> splitsByReader) {
+ if (splitsByReader.isEmpty()) {
+ return;
+ }
+ // Direct retained assignments bypass the sub-enumerator's assignment
listener.
+ splitsByReader.forEach(
+ (reader, splits) ->
+ splits.forEach(
+ split ->
currentSplitOwners.put(toPartitionKey(split), reader)));
+ enumContext.assignSplits(new SplitsAssignment<>(splitsByReader));
+ }
- if (!retainedSplitsByReader.isEmpty()) {
- enumContext.assignSplits(new
SplitsAssignment<>(retainedSplitsByReader));
+ private boolean isCurrentSplitOwner(DynamicKafkaSourceSplit split, int
readerId) {
+ Integer currentOwner = currentSplitOwners.get(toPartitionKey(split));
+ if (currentOwner == null || currentOwner == readerId) {
+ return true;
}
+ // Historical ownership is not a transfer request. The current owner's
returned
+ // assignments or its checkpoint report restore the split at the
correct position.
+ logger.debug(
+ "Ignoring split {} from historical reader {}; current owner is
{}",
+ split.splitId(),
+ readerId,
+ currentOwner);
+ return false;
}
private boolean shouldDeferMetadataUpdateEvents() {
- return
readerRecoveryGate.shouldDeferMetadataUpdateEvents(allReadersRegistered());
+ return readerRecoveryGate.shouldDeferMetadataUpdateEvents();
}
private void flushPendingMetadataUpdateEvents() {
@@ -1150,6 +1278,11 @@ public class DynamicKafkaSourceEnumerator
}
}
+ private static Tuple2<String, TopicPartition>
toPartitionKey(DynamicKafkaSourceSplit split) {
+ return Tuple2.of(
+ split.getKafkaClusterId(),
split.getKafkaPartitionSplit().getTopicPartition());
+ }
+
private static String toDynamicSplitId(String kafkaClusterId,
KafkaPartitionSplit split) {
return kafkaClusterId + "-" + split.splitId();
}
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
index 14bc6930..4cc8656e 100644
---
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
@@ -38,19 +38,9 @@ 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>After enumerator restore, assignments and metadata updates wait for the
complete reader cohort
+ * and the first metadata discovery. A local reader restart only waits for
that reader's checkpoint
+ * report to be reconciled with the surviving enumerator's current ownership.
*
* <p>This class owns that gating state; the enumerator remains responsible
for acting on it.
*
@@ -64,10 +54,13 @@ 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. */
+ /** Reported or returned splits awaiting recovery reassignment. */
private final Map<Integer, List<DynamicKafkaSourceSplit>>
pendingReportedSplitsByReader =
new HashMap<>();
+ /** Readers whose checkpoint reports are being reconciled before pending
assignments flush. */
+ private final Set<Integer> readersAwaitingSplitReconciliation = new
HashSet<>();
+
/** Readers whose metadata update events were deferred during recovery. */
private final Set<Integer> pendingMetadataUpdateReaders = new HashSet<>();
@@ -75,9 +68,9 @@ class ReaderRecoveryGate {
this.initialReaderRegistrationPending = restoredFromCheckpoint;
}
- /** Merges a reader's registration report with pending splits; an empty
report is ignored. */
+ /** Merges non-empty initial recovery reports with pending splits; empty
reports are ignored. */
void recordReportedSplits(int subtaskId, List<DynamicKafkaSourceSplit>
reportedSplits) {
- if (!reportedSplits.isEmpty()) {
+ if (initialReaderRegistrationPending && !reportedSplits.isEmpty()) {
pendingReportedSplitsByReader.put(
subtaskId,
mergeReportedSplits(
@@ -85,6 +78,16 @@ class ReaderRecoveryGate {
}
}
+ /** Retains splits returned while their reader is inactive until recovery
reassignment. */
+ void recordReturnedSplits(int readerId, List<DynamicKafkaSourceSplit>
returnedSplits) {
+ if (!returnedSplits.isEmpty()) {
+ pendingReportedSplitsByReader.put(
+ readerId,
+ mergeReportedSplits(
+ pendingReportedSplitsByReader.get(readerId),
returnedSplits));
+ }
+ }
+
/** Returns a copy of the pending reports without draining them. */
Map<Integer, List<DynamicKafkaSourceSplit>> snapshotReportedSplits() {
Map<Integer, List<DynamicKafkaSourceSplit>> snapshot = new HashMap<>();
@@ -116,16 +119,12 @@ class ReaderRecoveryGate {
/** Whether recovery gating is active and registrations must be deferred.
*/
boolean hasPendingRecovery() {
- return initialReaderRegistrationPending ||
!pendingReportedSplitsByReader.isEmpty();
+ return initialReaderRegistrationPending;
}
- /**
- * 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);
+ /** Whether metadata updates must wait for checkpoint report
reconciliation. */
+ boolean shouldDeferMetadataUpdateEvents() {
+ return initialReaderRegistrationPending ||
!readersAwaitingSplitReconciliation.isEmpty();
}
void deferMetadataUpdate(int readerId) {
@@ -144,19 +143,43 @@ class ReaderRecoveryGate {
return readers;
}
+ void startReaderRegistration(int readerId) {
+ readersAwaitingSplitReconciliation.add(readerId);
+ }
+
+ void completeReaderRegistration(int readerId) {
+ readersAwaitingSplitReconciliation.remove(readerId);
+ }
+
+ boolean isReaderReadyForAssignment(int readerId) {
+ return !initialReaderRegistrationPending
+ && !readersAwaitingSplitReconciliation.contains(readerId);
+ }
+
+ boolean isReconciliationPending() {
+ return initialReaderRegistrationPending ||
!readersAwaitingSplitReconciliation.isEmpty();
+ }
+
void markInitialRegistrationComplete() {
initialReaderRegistrationPending = false;
+ readersAwaitingSplitReconciliation.clear();
}
boolean hasReportedSplits() {
return !pendingReportedSplitsByReader.isEmpty();
}
+ /** Returns and removes only the pending splits for the recovering reader.
*/
+ List<DynamicKafkaSourceSplit> drainReportedSplits(int readerId) {
+ List<DynamicKafkaSourceSplit> reportedSplits =
+ pendingReportedSplitsByReader.remove(readerId);
+ return reportedSplits == null ? Collections.emptyList() :
reportedSplits;
+ }
+
/**
* 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}).
+ * <p>The initial registration gate remains closed while these reports are
reassigned.
*/
NavigableMap<Integer, List<DynamicKafkaSourceSplit>> drainReportedSplits()
{
NavigableMap<Integer, List<DynamicKafkaSourceSplit>>
reportedSplitsByReader =
diff --git
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaEnumContextProxy.java
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaEnumContextProxy.java
index 89843837..a591b322 100644
---
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaEnumContextProxy.java
+++
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaEnumContextProxy.java
@@ -46,6 +46,7 @@ import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
+import java.util.function.IntPredicate;
import java.util.stream.Collectors;
/**
@@ -74,6 +75,9 @@ public class StoppableKafkaEnumContextProxy
private final Runnable signalNoMoreSplitsCallback;
private boolean noMoreSplits = false;
private volatile boolean isClosing;
+ private IntPredicate readerReadyForAssignment = reader -> true;
+ private BiConsumer<Integer, DynamicKafkaSourceSplit>
splitAssignmentListener =
+ (reader, split) -> {};
/**
* Constructor for the enumerator context.
@@ -116,7 +120,18 @@ public class StoppableKafkaEnumContextProxy
@Override
public Map<Integer, ReaderInfo> registeredReaders() {
- return enumContext.registeredReaders();
+ return enumContext.registeredReaders().entrySet().stream()
+ .filter(entry -> readerReadyForAssignment.test(entry.getKey()))
+ .collect(Collectors.toMap(Map.Entry::getKey,
Map.Entry::getValue));
+ }
+
+ void setReaderReadyForAssignment(IntPredicate readerReadyForAssignment) {
+ this.readerReadyForAssignment = readerReadyForAssignment;
+ }
+
+ void setSplitAssignmentListener(
+ BiConsumer<Integer, DynamicKafkaSourceSplit>
splitAssignmentListener) {
+ this.splitAssignmentListener = splitAssignmentListener;
}
/** Wrap splits with cluster metadata. */
@@ -147,6 +162,9 @@ public class StoppableKafkaEnumContextProxy
.collect(Collectors.toList())));
if (!readerToSplitsMap.isEmpty()) {
+ readerToSplitsMap.forEach(
+ (reader, splits) ->
+ splits.forEach(split ->
splitAssignmentListener.accept(reader, split)));
enumContext.assignSplits(new
SplitsAssignment<>(readerToSplitsMap));
}
}
diff --git
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java
index 2f936706..e66a73c3 100644
---
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java
+++
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java
@@ -459,12 +459,23 @@ public class KafkaSourceEnumerator
LOG.debug("Partition discovery is disabled.");
noMoreNewPartitionSplits = true;
}
+ List<KafkaPartitionSplit> newPartitionSplits = new ArrayList<>();
for (KafkaPartitionSplit split :
partitionSplitChange.newPartitionSplits) {
- unassignedSplits.put(split.getTopicPartition(), split);
+ TopicPartition partition = split.getTopicPartition();
+ KafkaPartitionSplit pendingSplit = unassignedSplits.get(partition);
+ // A checkpoint split can arrive while initialization runs on a
worker. Its
+ // position takes precedence over a late discovery result; only
MIGRATED needs init.
+ if (assignedSplits.containsKey(partition)
+ || (pendingSplit != null && !pendingSplit.isMigrated())) {
+ LOG.debug("Ignoring stale initialization for known partition
{}", partition);
+ continue;
+ }
+ unassignedSplits.put(partition, split);
+ newPartitionSplits.add(split);
}
LOG.info("Partition split changes: {}", partitionSplitChange);
// TODO: Handle removed partitions.
-
addPartitionSplitChangeToPendingAssignments(partitionSplitChange.newPartitionSplits);
+ addPartitionSplitChangeToPendingAssignments(newPartitionSplits);
assignPendingPartitionSplits(context.registeredReaders().keySet());
}
@@ -472,7 +483,21 @@ public class KafkaSourceEnumerator
private void addPartitionSplitChangeToPendingAssignments(
Collection<KafkaPartitionSplit> newPartitionSplits) {
int numReaders = context.currentParallelism();
- List<KafkaPartitionSplit> sortedSplits = new
ArrayList<>(newPartitionSplits);
+ Map<TopicPartition, KafkaPartitionSplit> splitsByPartition = new
HashMap<>();
+ newPartitionSplits.forEach(
+ split -> splitsByPartition.put(split.getTopicPartition(),
split));
+ // Split equality includes offsets. Replace pending physical copies
across all owners
+ // once per batch; if a batch repeats a partition, its last supplied
position wins.
+ pendingPartitionSplitAssignment
+ .values()
+ .forEach(
+ pending ->
+ pending.removeIf(
+ previous ->
+ splitsByPartition.containsKey(
+
previous.getTopicPartition())));
+ pendingPartitionSplitAssignment.values().removeIf(Set::isEmpty);
+ List<KafkaPartitionSplit> sortedSplits = new
ArrayList<>(splitsByPartition.values());
sortedSplits.sort(
Comparator.comparing(
(KafkaPartitionSplit split) ->
split.getTopicPartition().topic())
@@ -626,8 +651,7 @@ public class KafkaSourceEnumerator
* @param numReaders the total number of readers.
* @return the id of the subtask that owns the split.
*/
- @VisibleForTesting
- static int getSplitOwner(TopicPartition tp, int numReaders) {
+ public static int getSplitOwner(TopicPartition tp, int numReaders) {
int startIndex = ((tp.topic().hashCode() * 31) & 0x7FFFFFFF) %
numReaders;
// here, the assumption is that the id of Kafka partitions are always
ascending
diff --git
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorRecoveryTest.java
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorRecoveryTest.java
index 88adb63f..d20c2116 100644
---
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorRecoveryTest.java
+++
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorRecoveryTest.java
@@ -35,15 +35,23 @@ import
org.apache.flink.connector.kafka.dynamic.source.enumerator.subscriber.Kaf
import
org.apache.flink.connector.kafka.dynamic.source.split.DynamicKafkaSourceSplit;
import org.apache.flink.connector.kafka.source.KafkaSourceOptions;
import org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumState;
+import
org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumerator;
import
org.apache.flink.connector.kafka.source.enumerator.initializer.NoStoppingOffsetsInitializer;
import
org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit;
import org.apache.flink.connector.kafka.testutils.MockKafkaMetadataService;
+import org.apache.flink.mock.Whitebox;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.common.TopicPartition;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
@@ -54,6 +62,8 @@ import java.util.Properties;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.function.BiConsumer;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
import static org.assertj.core.api.Assertions.assertThat;
@@ -624,6 +634,1083 @@ public class DynamicKafkaSourceEnumeratorRecoveryTest {
}
}
+ @Test
+ public void
testLocalFailureBeforeInitialDiscoveryPreservesReportedSplits() throws
Throwable {
+ String streamId = "stream";
+ String clusterId = "cluster-0";
+ String topic = "topic";
+ List<DynamicKafkaSourceSplit> checkpointSplits =
+ List.of(
+ new DynamicKafkaSourceSplit(
+ clusterId,
+ new KafkaPartitionSplit(new
TopicPartition(topic, 0), 42L)),
+ new DynamicKafkaSourceSplit(
+ clusterId,
+ new KafkaPartitionSplit(new
TopicPartition(topic, 1), 43L)));
+ KafkaStream stream = createKafkaStream(streamId, clusterId, topic);
+
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(2);
+ DynamicKafkaSourceEnumerator enumerator =
+ createEnumerator(
+ streamId,
+ new
MockKafkaMetadataService(Collections.singleton(stream)),
+ context,
+ createGlobalModeProperties(),
+ createRestoredState(stream, clusterId,
checkpointSplits))) {
+ enumerator.start();
+ registerReaderAndReportRestoredSplits(context, enumerator, 0,
checkpointSplits);
+ registerReaderAndReportRestoredSplits(context, enumerator, 1,
Collections.emptyList());
+ assertThat(context.getSplitsAssignmentSequence()).isEmpty();
+
+ // No post-checkpoint assignment exists to return before metadata
discovery completes.
+ context.unregisterReader(0);
+ enumerator.addSplitsBack(Collections.emptyList(), 0);
+ registerReaderAndReportRestoredSplits(context, enumerator, 0,
Collections.emptyList());
+ assertThat(context.getSplitsAssignmentSequence()).isEmpty();
+
+ context.runNextOneTimeCallable();
+
+ assertExactCheckpointOffsets(context, 0, checkpointSplits);
+ assertSinglePhysicalOwners(context, 0, checkpointSplits.size());
+ }
+ }
+
+ @Test
+ public void testInitialRestoreKeepsOneOwnerBeforeAnyLocalFailure() throws
Throwable {
+ verifyNumericCheckpointRecovery(false, false);
+ }
+
+ @Test
+ public void testInitialRestoreDiscoveryBeforeRegistrationKeepsOneOwner()
throws Throwable {
+ verifyNumericCheckpointRecovery(true, false);
+ }
+
+ @Test
+ public void testLocalFailureBeforeNewCheckpointKeepsOneOwner() throws
Throwable {
+ verifyNumericCheckpointRecovery(false, true);
+ }
+
+ @Test
+ public void testLocalFailureWithReverseRegistrationKeepsOneOwner() throws
Throwable {
+ verifyNumericCheckpointRecovery(false, true, 3, true, false, false,
false);
+ }
+
+ @Test
+ public void testLocalFailureWithInterleavedReturnsKeepsOneOwner() throws
Throwable {
+ verifyNumericCheckpointRecovery(false, true, 3, true, true, false,
false);
+ }
+
+ @Test
+ public void testOnlyReaderZeroFailsWithoutMovingSurvivingSplits() throws
Throwable {
+ verifyNumericCheckpointRecovery(false, true, 1, false, false, false,
false);
+ }
+
+ @Test
+ public void testOnlyReaderOneFailsWithoutMovingSurvivingSplits() throws
Throwable {
+ verifyNumericCheckpointRecovery(false, true, 2, false, false, false,
false);
+ }
+
+ @Test
+ public void
testPerClusterLocalFailureDoesNotReclaimSurvivingReaderSplits() throws
Throwable {
+ verifyNumericCheckpointRecovery(false, true, 1, false, false, false,
false, "per_cluster");
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {1, 2})
+ public void
testLocalFailureOfReactivatedRetainedSplitsKeepsCurrentOwners(int
failedReaderMask)
+ throws Throwable {
+ verifyNumericCheckpointRecovery(
+ false, true, failedReaderMask, false, false, false, false,
"global", true);
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {true, false})
+ public void
testPerClusterLocalFailureAfterRetainedClusterReactivationKeepsRestoredOwner(
+ boolean bounded) throws Throwable {
+ String clusterId = "cluster-0";
+ KafkaStream inactiveStream = new KafkaStream("stream",
Collections.emptyMap());
+ KafkaStream activeStream = createKafkaStream("stream", clusterId,
"topic");
+ // After rescaling, the retained reader can differ from the
per-cluster hash owner.
+ DynamicKafkaSourceSplit checkpointSplit =
+ new DynamicKafkaSourceSplit(
+ clusterId,
+ new KafkaPartitionSplit(new TopicPartition("topic",
0), 123L, 1000L));
+ DynamicKafkaSourceSplit retainedSplit =
checkpointSplit.retainUntil(Long.MAX_VALUE);
+ assertThat(
+ KafkaSourceEnumerator.getSplitOwner(
+
checkpointSplit.getKafkaPartitionSplit().getTopicPartition(), 2))
+ .isEqualTo(1);
+
+ Properties properties = createBaseProperties();
+ properties.setProperty(
+ DynamicKafkaSourceOptions.STREAM_ENUMERATOR_MODE.key(),
"per_cluster");
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "1");
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_REMOVED_CLUSTER_RETENTION_MS.key(),
+ "60000");
+ DynamicKafkaSourceEnumState restoredState =
+ new DynamicKafkaSourceEnumState(
+ Collections.singleton(inactiveStream),
+ Collections.emptyMap(),
+ Collections.singletonMap(
+ clusterId,
+ new
DynamicKafkaSourceEnumState.RetainedClusterState(
+ new KafkaSourceEnumState(
+
List.of(checkpointSplit.getKafkaPartitionSplit()),
+ Collections.emptyList(),
+ true),
+ Long.MAX_VALUE)),
+ Collections.emptyMap());
+ List<String> events = new ArrayList<>();
+ try (MockKafkaMetadataService metadataService =
+ new
MockKafkaMetadataService(Collections.singleton(inactiveStream));
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new
MockSplitEnumeratorContext<DynamicKafkaSourceSplit>(2) {
+ @Override
+ public void assignSplits(
+ SplitsAssignment<DynamicKafkaSourceSplit>
assignment) {
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+ splits.forEach(
+ split ->
+
events.add(
+
"assign:"
+
+ reader
+
+ ":"
+
+ split
+
.splitId())));
+ super.assignSplits(assignment);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int reader) {
+ events.add("done:" + reader);
+ super.signalNoMoreSplits(reader);
+ }
+ };
+ DynamicKafkaSourceEnumerator enumerator =
+ new DynamicKafkaSourceEnumerator(
+ new
KafkaStreamSetSubscriber(Collections.singleton("stream")),
+ metadataService,
+ context,
+ OffsetsInitializer.earliest(),
+ bounded
+ ? OffsetsInitializer.latest()
+ : new NoStoppingOffsetsInitializer(),
+ properties,
+ bounded ? Boundedness.BOUNDED :
Boundedness.CONTINUOUS_UNBOUNDED,
+ restoredState,
+ new NoOpKafkaEnumContextProxyFactory())) {
+ enumerator.start();
+ registerReaderAndReportRestoredSplits(context, enumerator, 0,
List.of(retainedSplit));
+ registerReaderAndReportRestoredSplits(context, enumerator, 1,
Collections.emptyList());
+ context.runPeriodicCallable(0);
+
+ assertExactCheckpointOffsets(context, 0, List.of(retainedSplit));
+ context.getSplitsAssignmentSequence()
+ .forEach(assignment ->
assertThat(assignment.assignment()).containsOnlyKeys(0));
+ if (bounded) {
+ assertBoundedAssignmentsBeforeCompletion(events, 0, 1);
+ assertThat(events.stream().filter(event ->
event.equals("done:0"))).hasSize(1);
+ assertThat(events.stream().filter(event ->
event.equals("done:1"))).hasSize(1);
+ }
+
+ int assignmentsBeforeReactivation =
context.getSplitsAssignmentSequence().size();
+
metadataService.setKafkaStreams(Collections.singleton(activeStream));
+ context.runPeriodicCallable(0);
+ KafkaSourceEnumerator subEnumerator =
+ (KafkaSourceEnumerator)
+ ((Map<?, ?>)
+ Whitebox.getInternalState(
+ enumerator,
"clusterEnumeratorMap"))
+ .get(clusterId);
+ completePartitionDiscovery(subEnumerator);
+ assertThat(context.getSplitsAssignmentSequence())
+ .hasSize(assignmentsBeforeReactivation);
+
+ for (int attempt = 0; attempt < 2; attempt++) {
+ events.clear();
+ int priorAssignments =
context.getSplitsAssignmentSequence().size();
+ context.unregisterReader(0);
+ enumerator.addSplitsBack(List.of(checkpointSplit), 0);
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, 0, List.of(retainedSplit));
+
+ assertExactCheckpointOffsets(context, priorAssignments,
List.of(checkpointSplit));
+ context.getSplitsAssignmentSequence().stream()
+ .skip(priorAssignments)
+ .forEach(
+ assignment ->
+
assertThat(assignment.assignment()).containsOnlyKeys(0));
+ if (bounded) {
+ assertBoundedAssignmentsBeforeCompletion(events, 0, 1);
+ assertThat(events.stream().filter(event ->
event.equals("done:0"))).hasSize(1);
+ assertThat(events).doesNotContain("done:1");
+ } else {
+ assertThat(events).noneMatch(event ->
event.startsWith("done:"));
+ }
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @MethodSource("collidingSplitOwners")
+ public void
testLocalFailureKeepsDistinctOwnersWhenActiveAndRetainedSplitIdsCollide(
+ boolean bounded, int retainedReader) throws Throwable {
+ DynamicKafkaSourceSplit activeSplit =
+ new DynamicKafkaSourceSplit(
+ "a", new KafkaPartitionSplit(new TopicPartition("b-c",
0), 123L, 1000L));
+ DynamicKafkaSourceSplit retainedSplit =
+ new DynamicKafkaSourceSplit(
+ "a-b",
+ new KafkaPartitionSplit(new
TopicPartition("c", 0), 456L, 2000L))
+ .retainUntil(Long.MAX_VALUE);
+ assertThat(activeSplit.splitId()).isEqualTo(retainedSplit.splitId());
+ assertThat(
+ KafkaSourceEnumerator.getSplitOwner(
+
activeSplit.getKafkaPartitionSplit().getTopicPartition(), 2))
+ .isZero();
+ KafkaStream activeStream = createKafkaStream("stream", "a", "b-c");
+ DynamicKafkaSourceEnumState restoredState =
+ createRestoredState(activeStream, "a", List.of(activeSplit));
+ Properties properties = createBaseProperties();
+ properties.setProperty(
+ DynamicKafkaSourceOptions.STREAM_ENUMERATOR_MODE.key(),
"per_cluster");
+ List<String> events = new ArrayList<>();
+ Map<Integer, List<DynamicKafkaSourceSplit>> reports =
+ Map.of(
+ 0,
+ retainedReader == 0
+ ? List.of(activeSplit, retainedSplit)
+ : List.of(activeSplit),
+ 1,
+ retainedReader == 1 ? List.of(retainedSplit) :
Collections.emptyList());
+
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new
MockSplitEnumeratorContext<DynamicKafkaSourceSplit>(2) {
+ @Override
+ public void assignSplits(
+ SplitsAssignment<DynamicKafkaSourceSplit>
assignment) {
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+ splits.forEach(
+ split ->
+
events.add(
+
"assign:"
+
+ reader
+
+ ":"
+
+ split
+
.splitId())));
+ super.assignSplits(assignment);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int reader) {
+ events.add("done:" + reader);
+ super.signalNoMoreSplits(reader);
+ }
+ };
+ DynamicKafkaSourceEnumerator enumerator =
+ new DynamicKafkaSourceEnumerator(
+ new
KafkaStreamSetSubscriber(Collections.singleton("stream")),
+ new
MockKafkaMetadataService(Collections.singleton(activeStream)),
+ context,
+ OffsetsInitializer.earliest(),
+ bounded
+ ? OffsetsInitializer.latest()
+ : new NoStoppingOffsetsInitializer(),
+ properties,
+ bounded ? Boundedness.BOUNDED :
Boundedness.CONTINUOUS_UNBOUNDED,
+ restoredState,
+ new NoOpKafkaEnumContextProxyFactory())) {
+ enumerator.start();
+ registerReaderAndReportRestoredSplits(context, enumerator, 0,
reports.get(0));
+ registerReaderAndReportRestoredSplits(context, enumerator, 1,
reports.get(1));
+ context.runNextOneTimeCallable();
+ completePartitionDiscovery(
+ (KafkaSourceEnumerator)
+ ((Map<?, ?>)
+ Whitebox.getInternalState(
+ enumerator,
"clusterEnumeratorMap"))
+ .get("a"));
+
+ assertExactCheckpointOffsets(context, 0, List.of(activeSplit,
retainedSplit));
+ reports.forEach(
+ (reader, expected) ->
+ assertThat(
+
context.getSplitsAssignmentSequence().stream()
+ .flatMap(
+ assignment ->
+ assignment
+
.assignment()
+
.getOrDefault(
+
reader,
+
Collections
+
.emptyList())
+
.stream()))
+
.containsExactlyInAnyOrderElementsOf(expected));
+
+ int priorAssignments =
context.getSplitsAssignmentSequence().size();
+ events.clear();
+ context.unregisterReader(retainedReader);
+ enumerator.addSplitsBack(Collections.emptyList(), retainedReader);
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, retainedReader,
reports.get(retainedReader));
+
+ assertExactCheckpointOffsets(context, priorAssignments,
reports.get(retainedReader));
+ context.getSplitsAssignmentSequence().stream()
+ .skip(priorAssignments)
+ .forEach(
+ assignment ->
+ assertThat(assignment.assignment())
+ .containsOnlyKeys(retainedReader));
+ if (bounded) {
+ assertBoundedAssignmentsBeforeCompletion(
+ events, retainedReader,
reports.get(retainedReader).size());
+ assertThat(events.stream().filter(event ->
event.equals("done:" + retainedReader)))
+ .hasSize(1);
+ assertThat(events).doesNotContain("done:" + (1 -
retainedReader));
+ } else {
+ assertThat(events).noneMatch(event ->
event.startsWith("done:"));
+ }
+ }
+ }
+
+ private static Stream<Arguments> collidingSplitOwners() {
+ return Stream.of(
+ Arguments.of(true, 0),
+ Arguments.of(true, 1),
+ Arguments.of(false, 0),
+ Arguments.of(false, 1));
+ }
+
+ @Test
+ public void testLocalFailurePreservesStillInactiveRetainedSplit() throws
Throwable {
+ String cluster = "cluster-0";
+ KafkaStream stream = createKafkaStream("stream", cluster,
"active-topic");
+ DynamicKafkaSourceSplit retained =
+ new DynamicKafkaSourceSplit(
+ cluster,
+ new KafkaPartitionSplit(
+ new TopicPartition("inactive-topic",
0), 123L))
+ .retainUntil(Long.MAX_VALUE);
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(2);
+ DynamicKafkaSourceEnumerator enumerator =
+ createEnumerator(
+ "stream",
+ new
MockKafkaMetadataService(Collections.singleton(stream)),
+ context,
+ createGlobalModeProperties(),
+ createRestoredState(stream, cluster,
Collections.emptyList()))) {
+ enumerator.start();
+ registerReaderAndReportRestoredSplits(context, enumerator, 0,
List.of(retained));
+ registerReaderAndReportRestoredSplits(context, enumerator, 1,
Collections.emptyList());
+ context.runNextOneTimeCallable();
+ assertExactCheckpointOffsets(context, 0, List.of(retained));
+
+ int priorAssignments =
context.getSplitsAssignmentSequence().size();
+ context.unregisterReader(0);
+ enumerator.addSplitsBack(Collections.emptyList(), 0);
+ registerReaderAndReportRestoredSplits(context, enumerator, 0,
List.of(retained));
+
+ assertExactCheckpointOffsets(context, priorAssignments,
List.of(retained));
+ context.getSplitsAssignmentSequence().stream()
+ .skip(priorAssignments)
+ .forEach(assignment ->
assertThat(assignment.assignment()).containsOnlyKeys(0));
+ }
+ }
+
+ @Test
+ public void testLocalFailureAfterNewCheckpointUsesReportedOffsets() throws
Throwable {
+ verifyNumericCheckpointRecovery(false, true, 1, false, false, true,
false);
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {true, false})
+ public void testLocalFailurePreservesCurrentOwnersAcrossClusterRetention(
+ boolean failWhileInactive) throws Throwable {
+ String streamId = "stream";
+ String clusterId = "cluster-0";
+ String topic = "topic";
+ List<DynamicKafkaSourceSplit> checkpointSplits =
+ createSplits(clusterId, topic, 11).stream()
+ .map(
+ split ->
+ withOffset(
+ split,
+ 1000L
+ +
split.getKafkaPartitionSplit()
+
.getPartition()))
+ .collect(Collectors.toList());
+ List<DynamicKafkaSourceSplit> historicalReaderZeroSplits =
+ checkpointSplits.stream()
+ .filter(split ->
split.getKafkaPartitionSplit().getPartition() % 2 == 0)
+ .collect(Collectors.toList());
+ List<DynamicKafkaSourceSplit> historicalReaderOneSplits =
+ checkpointSplits.stream()
+ .filter(split ->
split.getKafkaPartitionSplit().getPartition() % 2 == 1)
+ .collect(Collectors.toList());
+ KafkaStream stream = createKafkaStream(streamId, clusterId, topic);
+ Properties properties = createGlobalModeProperties();
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "1");
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_REMOVED_CLUSTER_RETENTION_MS.key(),
+ "60000");
+
+ try (MockKafkaMetadataService metadataService =
+ new
MockKafkaMetadataService(Collections.singleton(stream));
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(2);
+ DynamicKafkaSourceEnumerator enumerator =
+ createEnumerator(
+ streamId,
+ metadataService,
+ context,
+ properties,
+ createRestoredState(stream, clusterId,
checkpointSplits))) {
+ enumerator.start();
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, 0, historicalReaderZeroSplits);
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, 1, historicalReaderOneSplits);
+ context.runPeriodicCallable(0);
+ assertExactCheckpointOffsets(context, 0, checkpointSplits);
+
+ Map<Integer, List<DynamicKafkaSourceSplit>> currentAssignments =
new HashMap<>();
+ context.getSplitsAssignmentSequence()
+ .forEach(
+ assignment ->
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+ currentAssignments
+
.computeIfAbsent(
+
reader,
+
ignored ->
+
new ArrayList<>())
+
.addAll(splits)));
+ assertThat(historicalReaderZeroSplits.stream())
+ .anyMatch(currentAssignments.get(1)::contains);
+
+ metadataService.setKafkaStreams(
+ Collections.singleton(new KafkaStream(streamId,
Collections.emptyMap())));
+ context.runPeriodicCallable(0);
+
assertThat(enumerator.snapshotState(1L).getRetainedClusterEnumeratorStates())
+ .containsKey(clusterId);
+
+ if (!failWhileInactive) {
+ metadataService.setKafkaStreams(Collections.singleton(stream));
+ context.runPeriodicCallable(0);
+ }
+ int assignmentsBeforeFailure =
context.getSplitsAssignmentSequence().size();
+ context.unregisterReader(0);
+ enumerator.addSplitsBack(currentAssignments.get(0), 0);
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, 0, historicalReaderZeroSplits);
+
+ List<DynamicKafkaSourceSplit> recoveredSplits = new ArrayList<>();
+ for (SplitsAssignment<DynamicKafkaSourceSplit> assignment :
+ context.getSplitsAssignmentSequence()
+ .subList(
+ assignmentsBeforeFailure,
+
context.getSplitsAssignmentSequence().size())) {
+ assertThat(assignment.assignment()).containsOnlyKeys(0);
+ recoveredSplits.addAll(assignment.assignment().get(0));
+ }
+ assertThat(recoveredSplits)
+ .allSatisfy(
+ split ->
assertThat(split.isRetained()).isEqualTo(failWhileInactive));
+ assertThat(
+ recoveredSplits.stream()
+
.map(DynamicKafkaSourceSplit::clearRetention)
+ .collect(Collectors.toList()))
+
.containsExactlyInAnyOrderElementsOf(currentAssignments.get(0));
+
+ if (failWhileInactive) {
+ int assignmentsBeforeReactivation =
context.getSplitsAssignmentSequence().size();
+ metadataService.setKafkaStreams(Collections.singleton(stream));
+ context.runPeriodicCallable(0);
+ assertThat(context.getSplitsAssignmentSequence())
+ .hasSize(assignmentsBeforeReactivation);
+ }
+ }
+ }
+
+ @Test
+ public void
testCheckpointReportsReplacePreinitializedPendingOwnersAndOffsets()
+ throws Throwable {
+ verifyNumericCheckpointRecovery(true, true, 3, true, true, false,
true);
+ }
+
+ @Test
+ public void
testNewPartitionDiscoveryWaitsForRecoveryAndSurvivesRepeatedLocalFailure()
+ throws Throwable {
+ String cluster = "cluster-0";
+ KafkaStream stream = createKafkaStream("stream", cluster, "topic");
+ List<DynamicKafkaSourceSplit> checkpointSplits = createSplits(cluster,
"topic", 2);
+ DynamicKafkaSourceSplit discovered =
+ withOffset(createSplits(cluster, "topic", 3).get(2), 2500L);
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(2);
+ DynamicKafkaSourceEnumerator enumerator =
+ createEnumerator(
+ "stream",
+ new
MockKafkaMetadataService(Collections.singleton(stream)),
+ context,
+ createGlobalModeProperties(),
+ createRestoredState(stream, cluster,
checkpointSplits))) {
+ enumerator.start();
+ context.runNextOneTimeCallable();
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, 0,
Collections.singletonList(checkpointSplits.get(0)));
+ // Invoke the real coordinator callback with a completed
broker-discovery result.
+ // No broker or mock owner-selection implementation participates
in the regression.
+ Object subEnumerator =
+ ((Map<?, ?>) Whitebox.getInternalState(enumerator,
"clusterEnumeratorMap"))
+ .get(cluster);
+ Class<?> changeClass =
+ Class.forName(KafkaSourceEnumerator.class.getName() +
"$PartitionSplitChange");
+ Constructor<?> constructor =
changeClass.getDeclaredConstructor(Set.class, Set.class);
+ constructor.setAccessible(true);
+ Method handler =
+ KafkaSourceEnumerator.class.getDeclaredMethod(
+ "handlePartitionSplitChanges", changeClass,
Throwable.class);
+ handler.setAccessible(true);
+ handler.invoke(
+ subEnumerator,
+ constructor.newInstance(
+
Collections.singleton(discovered.getKafkaPartitionSplit()),
+ Collections.emptySet()),
+ null);
+ assertThat(context.getSplitsAssignmentSequence()).isEmpty();
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, 1,
Collections.singletonList(checkpointSplits.get(1)));
+ List<DynamicKafkaSourceSplit> expected = new
ArrayList<>(checkpointSplits);
+ expected.add(discovered);
+ assertExactCheckpointOffsets(context, 0, expected);
+ Map<Integer, List<DynamicKafkaSourceSplit>> actualOwners = new
HashMap<>();
+ context.getSplitsAssignmentSequence()
+ .forEach(
+ assignment ->
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+ actualOwners
+
.computeIfAbsent(
+
reader,
+
ignored ->
+
new ArrayList<>())
+
.addAll(splits)));
+ int failedReader =
+ actualOwners.entrySet().stream()
+ .filter(entry ->
entry.getValue().contains(discovered))
+ .findFirst()
+ .get()
+ .getKey();
+ for (int attempt = 0; attempt < 2; attempt++) {
+ int priorAssignments =
context.getSplitsAssignmentSequence().size();
+ context.unregisterReader(failedReader);
+ enumerator.addSplitsBack(actualOwners.get(failedReader),
failedReader);
+ registerReaderAndReportRestoredSplits(
+ context,
+ enumerator,
+ failedReader,
+
Collections.singletonList(checkpointSplits.get(failedReader)));
+ assertExactCheckpointOffsets(
+ context, priorAssignments,
actualOwners.get(failedReader));
+ context.getSplitsAssignmentSequence().stream()
+ .skip(priorAssignments)
+ .forEach(
+ assignment ->
+ assertThat(assignment.assignment())
+
.containsOnlyKeys(failedReader));
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @MethodSource("boundedRecoveryModes")
+ public void testBoundedRecoveryAssignsAllClustersBeforeNoMoreSplits(
+ String mode, boolean discoverFirst) throws Throwable {
+ verifyBoundedRecoveryAssignmentOrder(mode, discoverFirst, true);
+ }
+
+ @ParameterizedTest
+ @MethodSource("boundedRecoveryModes")
+ public void testBoundedLocalRecoveryAssignsAllClustersBeforeNoMoreSplits(
+ String mode, boolean discoverFirst) throws Throwable {
+ verifyBoundedRecoveryAssignmentOrder(mode, discoverFirst, false);
+ }
+
+ private void verifyBoundedRecoveryAssignmentOrder(
+ String mode, boolean discoverFirst, boolean checkInitialOrder)
throws Throwable {
+ Map<String, ClusterMetadata> metadata = new HashMap<>();
+ Map<String, KafkaSourceEnumState> states = new HashMap<>();
+ Map<Integer, List<DynamicKafkaSourceSplit>> reports = new HashMap<>();
+ for (String cluster : List.of("cluster-0", "cluster-1")) {
+ metadata.putAll(createKafkaStream("stream", cluster,
"topic").getClusterMetadataMap());
+ List<DynamicKafkaSourceSplit> splits = new ArrayList<>();
+ for (int reader = 0; reader < 2; reader++) {
+ DynamicKafkaSourceSplit split =
+ new DynamicKafkaSourceSplit(
+ cluster,
+ new KafkaPartitionSplit(
+ new TopicPartition("topic", reader),
100L + reader, 1000L));
+ splits.add(split);
+ reports.computeIfAbsent(reader, ignored -> new
ArrayList<>()).add(split);
+ }
+ states.put(
+ cluster,
+ new KafkaSourceEnumState(unwrapSplits(splits),
Collections.emptyList(), true));
+ }
+ KafkaStream stream = new KafkaStream("stream", metadata);
+ Properties properties = createBaseProperties();
+
properties.setProperty(DynamicKafkaSourceOptions.STREAM_ENUMERATOR_MODE.key(),
mode);
+ List<String> events = new ArrayList<>();
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new
MockSplitEnumeratorContext<DynamicKafkaSourceSplit>(2) {
+ @Override
+ public void assignSplits(
+ SplitsAssignment<DynamicKafkaSourceSplit>
assignment) {
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+ splits.forEach(
+ split ->
+
events.add(
+
"assign:"
+
+ reader
+
+ ":"
+
+ split
+
.splitId())));
+ super.assignSplits(assignment);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int reader) {
+ events.add("done:" + reader);
+ super.signalNoMoreSplits(reader);
+ }
+ };
+ DynamicKafkaSourceEnumerator enumerator =
+ new DynamicKafkaSourceEnumerator(
+ new
KafkaStreamSetSubscriber(Collections.singleton("stream")),
+ new
MockKafkaMetadataService(Collections.singleton(stream)),
+ context,
+ OffsetsInitializer.earliest(),
+ OffsetsInitializer.latest(),
+ properties,
+ Boundedness.BOUNDED,
+ new DynamicKafkaSourceEnumState(
+ Collections.singleton(stream), states),
+ new NoOpKafkaEnumContextProxyFactory())) {
+ enumerator.start();
+ // Complete actual sub-enumerator discovery callbacks without
broker I/O. Each
+ // sub-enumerator must subsequently signal completion during its
addReader flush.
+ Map<?, ?> subEnumerators =
+ (Map<?, ?>) Whitebox.getInternalState(enumerator,
"clusterEnumeratorMap");
+ Class<?> changeClass =
+ Class.forName(KafkaSourceEnumerator.class.getName() +
"$PartitionSplitChange");
+ Constructor<?> constructor =
changeClass.getDeclaredConstructor(Set.class, Set.class);
+ constructor.setAccessible(true);
+ Method handler =
+ KafkaSourceEnumerator.class.getDeclaredMethod(
+ "handlePartitionSplitChanges", changeClass,
Throwable.class);
+ handler.setAccessible(true);
+ for (Object subEnumerator : subEnumerators.values()) {
+ handler.invoke(
+ subEnumerator,
+ constructor.newInstance(Collections.emptySet(),
Collections.emptySet()),
+ null);
+ }
+ if (discoverFirst) {
+ context.runNextOneTimeCallable();
+ }
+ for (int reader = 0; reader < 2; reader++) {
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, reader, reports.get(reader));
+ }
+ if (!discoverFirst) {
+ context.runNextOneTimeCallable();
+ }
+ if (checkInitialOrder) {
+ assertBoundedAssignmentsBeforeCompletion(events, 0, 2);
+ assertBoundedAssignmentsBeforeCompletion(events, 1, 2);
+ }
+ Map<Integer, List<DynamicKafkaSourceSplit>> assigned = new
HashMap<>();
+ context.getSplitsAssignmentSequence()
+ .forEach(
+ assignment ->
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+
assigned.computeIfAbsent(
+
reader,
+
ignored ->
+
new ArrayList<>())
+
.addAll(splits)));
+ events.clear();
+ int priorAssignments =
context.getSplitsAssignmentSequence().size();
+ context.unregisterReader(0);
+ context.unregisterReader(1);
+ List<DynamicKafkaSourceSplit> readerZeroReports = new
ArrayList<>(assigned.get(0));
+ readerZeroReports.add(
+ new DynamicKafkaSourceSplit(
+ "removed-cluster",
+ new KafkaPartitionSplit(
+ new
TopicPartition("removed-topic", 0), 5L, 1000L))
+ .retainUntil(Long.MAX_VALUE));
+ enumerator.addSplitsBack(Collections.emptyList(), 0);
+ registerReaderAndReportRestoredSplits(context, enumerator, 0,
readerZeroReports);
+
+ assertBoundedAssignmentsBeforeCompletion(events, 0, 3);
+ assertThat(events).noneMatch(event ->
event.startsWith("assign:1:"));
+ assertThat(events).doesNotContain("done:1");
+ assertExactCheckpointOffsets(context, priorAssignments,
readerZeroReports);
+ context.getSplitsAssignmentSequence().stream()
+ .skip(priorAssignments)
+ .flatMap(assignment ->
assignment.assignment().values().stream())
+ .flatMap(List::stream)
+ .forEach(
+ split ->
+
assertThat(split.getKafkaPartitionSplit().getStoppingOffset())
+ .contains(1000L));
+ registerReaderAndReportRestoredSplits(context, enumerator, 1,
assigned.get(1));
+ assertBoundedAssignmentsBeforeCompletion(events, 1, 2);
+ assertThat(events.stream().filter(event ->
event.equals("done:0"))).hasSize(1);
+ assertThat(events.stream().filter(event ->
event.equals("done:1"))).hasSize(1);
+ }
+ }
+
+ private static Stream<Arguments> boundedRecoveryModes() {
+ return Stream.of(
+ Arguments.of("per_cluster", true),
+ Arguments.of("per_cluster", false),
+ Arguments.of("global", true),
+ Arguments.of("global", false));
+ }
+
+ private static void assertBoundedAssignmentsBeforeCompletion(
+ List<String> events, int reader, int expectedSplits) {
+ String assignmentPrefix = "assign:" + reader + ":";
+ int completion = events.indexOf("done:" + reader);
+ assertThat(completion).as("completion must be signaled: %s",
events).isNotNegative();
+ assertThat(events.stream().filter(event ->
event.startsWith(assignmentPrefix)))
+ .hasSize(expectedSplits);
+ assertThat(
+ events.subList(0, completion).stream()
+ .filter(event ->
event.startsWith(assignmentPrefix)))
+ .as("all restored assignments must precede completion: %s",
events)
+ .hasSize(expectedSplits);
+ }
+
+ private void verifyNumericCheckpointRecovery(boolean discoverFirst,
boolean localFailure)
+ throws Throwable {
+ verifyNumericCheckpointRecovery(discoverFirst, localFailure, 3, false,
false, false, false);
+ }
+
+ private void verifyNumericCheckpointRecovery(
+ boolean discoverFirst,
+ boolean localFailure,
+ int failedReaderMask,
+ boolean reverseRegistration,
+ boolean interleavedReturns,
+ boolean checkpointAfterRedistribution,
+ boolean restoredPendingCopies)
+ throws Throwable {
+ verifyNumericCheckpointRecovery(
+ discoverFirst,
+ localFailure,
+ failedReaderMask,
+ reverseRegistration,
+ interleavedReturns,
+ checkpointAfterRedistribution,
+ restoredPendingCopies,
+ "global");
+ }
+
+ private void verifyNumericCheckpointRecovery(
+ boolean discoverFirst,
+ boolean localFailure,
+ int failedReaderMask,
+ boolean reverseRegistration,
+ boolean interleavedReturns,
+ boolean checkpointAfterRedistribution,
+ boolean restoredPendingCopies,
+ String mode)
+ throws Throwable {
+ verifyNumericCheckpointRecovery(
+ discoverFirst,
+ localFailure,
+ failedReaderMask,
+ reverseRegistration,
+ interleavedReturns,
+ checkpointAfterRedistribution,
+ restoredPendingCopies,
+ mode,
+ false);
+ }
+
+ private void verifyNumericCheckpointRecovery(
+ boolean discoverFirst,
+ boolean localFailure,
+ int failedReaderMask,
+ boolean reverseRegistration,
+ boolean interleavedReturns,
+ boolean checkpointAfterRedistribution,
+ boolean restoredPendingCopies,
+ String mode,
+ boolean retainedCheckpointReports)
+ throws Throwable {
+ String streamId = "stream";
+ String clusterId = "cluster-0";
+ String topic = "topic";
+ int parallelism = 2;
+ List<DynamicKafkaSourceSplit> checkpointSplits =
+ createSplits(clusterId, topic, 11).stream()
+ .map(
+ split ->
+ new DynamicKafkaSourceSplit(
+ clusterId,
+ new KafkaPartitionSplit(
+
split.getKafkaPartitionSplit()
+
.getTopicPartition(),
+ 1000L
+ +
split.getKafkaPartitionSplit()
+
.getTopicPartition()
+
.partition())))
+ .collect(Collectors.toList());
+ KafkaStream stream = createKafkaStream(streamId, clusterId, topic);
+ Properties properties = createBaseProperties();
+
properties.setProperty(DynamicKafkaSourceOptions.STREAM_ENUMERATOR_MODE.key(),
mode);
+ Map<Integer, List<DynamicKafkaSourceSplit>> checkpointOwners = new
HashMap<>();
+ for (DynamicKafkaSourceSplit split : checkpointSplits) {
+ int owner =
+
split.getKafkaPartitionSplit().getTopicPartition().partition() % parallelism;
+ checkpointOwners
+ .computeIfAbsent(owner, ignored -> new ArrayList<>())
+ .add(retainedCheckpointReports ?
split.retainUntil(Long.MAX_VALUE) : split);
+ }
+ DynamicKafkaSourceEnumState restoredState =
+ createRestoredState(stream, clusterId, checkpointSplits);
+ if (restoredPendingCopies) {
+ restoredState =
+ new DynamicKafkaSourceEnumState(
+ Collections.singleton(stream),
+ Collections.singletonMap(
+ clusterId,
+ new KafkaSourceEnumState(
+ Collections.emptyList(),
+ checkpointSplits.stream()
+ .map(
+ split ->
+
withOffset(split, 9000L)
+
.getKafkaPartitionSplit())
+
.collect(Collectors.toList()),
+ true)));
+ }
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(parallelism);
+ DynamicKafkaSourceEnumerator enumerator =
+ createEnumerator(
+ streamId,
+ new
MockKafkaMetadataService(Collections.singleton(stream)),
+ context,
+ properties,
+ restoredState)) {
+ enumerator.start();
+ if (discoverFirst) {
+ context.runNextOneTimeCallable();
+ }
+ for (int reader = 0; reader < parallelism; reader++) {
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, reader,
checkpointOwners.get(reader));
+ }
+ if (!discoverFirst) {
+ context.runNextOneTimeCallable();
+ }
+ assertSinglePhysicalOwners(context, 0, 11);
+ assertExactCheckpointOffsets(context, 0, checkpointSplits);
+ if (!localFailure) {
+ return;
+ }
+ // Model SourceCoordinator local recovery: return all
+ // assignments since the restored checkpoint, then register
original checkpoint reports.
+ Map<Integer, List<DynamicKafkaSourceSplit>>
uncheckpointedAssignments = new HashMap<>();
+ for (SplitsAssignment<DynamicKafkaSourceSplit> assignment :
+ context.getSplitsAssignmentSequence()) {
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+ uncheckpointedAssignments
+ .computeIfAbsent(
+ reader, ignored -> new
ArrayList<>())
+ .addAll(splits));
+ }
+ int assignmentsBeforeFailure =
context.getSplitsAssignmentSequence().size();
+ Map<Integer, List<DynamicKafkaSourceSplit>> currentAssignments =
new HashMap<>();
+ uncheckpointedAssignments.forEach(
+ (reader, splits) -> currentAssignments.put(reader, new
ArrayList<>(splits)));
+ if (checkpointAfterRedistribution) {
+ // Runtime no longer returns assignments covered by a
completed checkpoint.
+ // Reader reports now use the current owner and the
checkpointed next offset.
+ checkpointOwners.clear();
+ currentAssignments.forEach(
+ (reader, splits) ->
+ checkpointOwners.put(
+ reader,
+ splits.stream()
+ .map(
+ split ->
+ withOffset(
+ split,
+
split.getKafkaPartitionSplit()
+
.getStartingOffset()
+
+ 10000L))
+
.collect(Collectors.toList())));
+ uncheckpointedAssignments.replaceAll((reader, splits) ->
Collections.emptyList());
+ }
+ for (int reader = 0; reader < parallelism; reader++) {
+ if ((failedReaderMask & (1 << reader)) != 0) {
+ context.unregisterReader(reader);
+ }
+ }
+ if (!interleavedReturns) {
+ for (int reader = 0; reader < parallelism; reader++) {
+ if ((failedReaderMask & (1 << reader)) != 0) {
+
enumerator.addSplitsBack(uncheckpointedAssignments.get(reader), reader);
+ }
+ }
+ }
+ for (int index = 0; index < parallelism; index++) {
+ int reader = reverseRegistration ? parallelism - 1 - index :
index;
+ if ((failedReaderMask & (1 << reader)) != 0) {
+ if (interleavedReturns) {
+
enumerator.addSplitsBack(uncheckpointedAssignments.get(reader), reader);
+ }
+ registerReaderAndReportRestoredSplits(
+ context, enumerator, reader,
checkpointOwners.get(reader));
+ }
+ }
+ Map<Integer, List<DynamicKafkaSourceSplit>>
assignmentsAfterFailure = new HashMap<>();
+ context.getSplitsAssignmentSequence().stream()
+ .skip(assignmentsBeforeFailure)
+ .forEach(
+ assignment ->
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+
assignmentsAfterFailure
+
.computeIfAbsent(
+
reader,
+
ignored ->
+
new ArrayList<>())
+
.addAll(splits)));
+ for (int reader = 0; reader < parallelism; reader++) {
+ if ((failedReaderMask & (1 << reader)) == 0) {
+
assertThat(assignmentsAfterFailure).doesNotContainKey(reader);
+ } else {
+ List<DynamicKafkaSourceSplit> expected =
+ checkpointAfterRedistribution
+ ? checkpointOwners.get(reader)
+ : currentAssignments.get(reader);
+ assertThat(assignmentsAfterFailure.get(reader))
+ .as(
+ "failed reader %s keeps its actual owner
and checkpoint offsets",
+ reader)
+ .containsExactlyInAnyOrderElementsOf(expected);
+ }
+ }
+ }
+ }
+
+ private static DynamicKafkaSourceSplit withOffset(DynamicKafkaSourceSplit
split, long offset) {
+ return new DynamicKafkaSourceSplit(
+ split.getKafkaClusterId(),
+ new KafkaPartitionSplit(
+ split.getKafkaPartitionSplit().getTopicPartition(),
offset));
+ }
+
+ private static void completePartitionDiscovery(KafkaSourceEnumerator
enumerator)
+ throws Exception {
+ Class<?> changeClass =
+ Class.forName(KafkaSourceEnumerator.class.getName() +
"$PartitionSplitChange");
+ Constructor<?> constructor =
changeClass.getDeclaredConstructor(Set.class, Set.class);
+ constructor.setAccessible(true);
+ Method handler =
+ KafkaSourceEnumerator.class.getDeclaredMethod(
+ "handlePartitionSplitChanges", changeClass,
Throwable.class);
+ handler.setAccessible(true);
+ handler.invoke(
+ enumerator,
+ constructor.newInstance(Collections.emptySet(),
Collections.emptySet()),
+ null);
+ }
+
+ private static void assertExactCheckpointOffsets(
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context,
+ int skipAssignments,
+ List<DynamicKafkaSourceSplit> expected) {
+ assertThat(
+ context.getSplitsAssignmentSequence().stream()
+ .skip(skipAssignments)
+ .flatMap(assignment ->
assignment.assignment().values().stream())
+ .flatMap(List::stream)
+ .collect(Collectors.toList()))
+ .containsExactlyInAnyOrderElementsOf(expected);
+ }
+
+ private static void assertSinglePhysicalOwners(
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context,
+ int skipAssignments,
+ int expectedPartitions) {
+ Map<String, Set<Integer>> owners = new HashMap<>();
+ context.getSplitsAssignmentSequence().stream()
+ .skip(skipAssignments)
+ .forEach(
+ assignment ->
+ assignment
+ .assignment()
+ .forEach(
+ (reader, splits) ->
+ splits.forEach(
+ split ->
+
owners.computeIfAbsent(
+
split
+
.splitId(),
+
ignored ->
+
new HashSet<>())
+
.add(reader))));
+ assertThat(owners).hasSize(expectedPartitions);
+ assertThat(owners.entrySet())
+ .as("every physical split must have one active reader: %s",
owners)
+ .allSatisfy(entry -> assertThat(entry.getValue()).hasSize(1));
+ }
+
+ private static void registerReaderAndReportRestoredSplits(
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context,
+ DynamicKafkaSourceEnumerator enumerator,
+ int readerId,
+ List<DynamicKafkaSourceSplit> restoredSplits) {
+ context.registerReader(
+ ReaderInfo.createReaderInfo(readerId, "location-" + readerId,
restoredSplits));
+ enumerator.addReader(readerId);
+ }
+
private static void registerReadersWithoutReportedSplits(
MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context,
DynamicKafkaSourceEnumerator enumerator,
diff --git
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorTest.java
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorTest.java
index 389beb68..e100aedf 100644
---
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorTest.java
+++
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorTest.java
@@ -257,8 +257,15 @@ public class DynamicKafkaSourceEnumeratorTest {
context.registerReader(
ReaderInfo.createReaderInfo(0, "restarted",
reportedSplitsByReader.get(0)));
enumerator.addReader(0);
- assertThat(context.getSplitsAssignmentSequence()).isEmpty();
- assertThat(context.splitsAtCompletion).isEmpty();
+ assertThat(context.getAssignedSplits(0))
+
.containsExactlyInAnyOrderElementsOf(reportedSplitsByReader.get(0));
+ assertThat(context.splitsAtCompletion).containsOnlyKeys(0);
+ assertThat(context.splitsAtCompletion.get(0))
+ .singleElement()
+ .isEqualTo(
+ reportedSplitsByReader.get(0).stream()
+ .map(DynamicKafkaSourceSplit::splitId)
+ .collect(Collectors.toSet()));
context.registerReader(
ReaderInfo.createReaderInfo(1, "restarted",
reportedSplitsByReader.get(1)));
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
index cf3d5376..4b2cfa51 100644
---
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
@@ -40,8 +40,7 @@ class ReaderRecoveryGateTest {
ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
assertThat(gate.hasPendingRecovery()).isFalse();
- assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isFalse();
- assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents()).isFalse();
assertThat(gate.hasReportedSplits()).isFalse();
}
@@ -50,56 +49,66 @@ class ReaderRecoveryGateTest {
ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
assertThat(gate.hasPendingRecovery()).isTrue();
- assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isTrue();
- assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isTrue();
+ assertThat(gate.shouldDeferMetadataUpdateEvents()).isTrue();
gate.markInitialRegistrationComplete();
assertThat(gate.hasPendingRecovery()).isFalse();
- assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents()).isFalse();
}
@Test
- void
testCompletingInitialRegistrationDoesNotReleaseGateWithPendingReportedSplits() {
+ void testDrainingReportsDoesNotReleaseInitialAssignmentGate() {
ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
gate.recordReportedSplits(1, Collections.singletonList(split("topic",
0)));
+ gate.startReaderRegistration(1);
- gate.markInitialRegistrationComplete();
+ gate.drainReportedSplits();
assertThat(gate.hasPendingRecovery()).isTrue();
- assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isTrue();
- assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+ assertThat(gate.isReaderReadyForAssignment(1)).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents()).isTrue();
- gate.drainReportedSplits();
+ gate.markInitialRegistrationComplete();
assertThat(gate.hasPendingRecovery()).isFalse();
- assertThat(gate.shouldDeferMetadataUpdateEvents(false)).isFalse();
- assertThat(gate.shouldDeferMetadataUpdateEvents(true)).isFalse();
+ assertThat(gate.isReaderReadyForAssignment(1)).isTrue();
+ assertThat(gate.isReconciliationPending()).isFalse();
}
@Test
- void testReportedSplitsGateUntilAllReadersRegistered() {
+ void testLocalRegistrationOnlyBlocksTheRecoveringReader() {
ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ gate.startReaderRegistration(1);
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();
+ assertThat(gate.hasPendingRecovery()).isFalse();
+ assertThat(gate.hasReportedSplits()).isFalse();
+ assertThat(gate.isReaderReadyForAssignment(0)).isTrue();
+ assertThat(gate.isReaderReadyForAssignment(1)).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents()).isTrue();
+
+ gate.completeReaderRegistration(1);
+
+ assertThat(gate.isReaderReadyForAssignment(1)).isTrue();
+ assertThat(gate.isReconciliationPending()).isFalse();
+ assertThat(gate.shouldDeferMetadataUpdateEvents()).isFalse();
}
@Test
- void testEmptyReportedSplitsAreIgnored() {
- ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ void testEmptyReportDuringInitialRecovery() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
gate.recordReportedSplits(1, Collections.emptyList());
- assertThat(gate.hasPendingRecovery()).isFalse();
+ assertThat(gate.hasPendingRecovery()).isTrue();
assertThat(gate.hasReportedSplits()).isFalse();
+ assertThat(gate.drainReportedSplits()).isEmpty();
+ assertThat(gate.hasPendingRecovery()).isTrue();
}
@Test
void testRepeatedReportMergesSplitsAndPrefersLatestOffsets() {
- ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
DynamicKafkaSourceSplit firstReport = split("topic", 0);
DynamicKafkaSourceSplit secondReport = split("topic", 1);
DynamicKafkaSourceSplit updatedSplit =
@@ -118,7 +127,7 @@ class ReaderRecoveryGateTest {
@Test
void testSnapshotReportedSplitsIsIndependentAndDoesNotDrain() {
- ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
DynamicKafkaSourceSplit firstReport = split("topic", 0);
DynamicKafkaSourceSplit secondReport = split("topic", 1);
gate.recordReportedSplits(1, Collections.singletonList(firstReport));
@@ -139,8 +148,8 @@ class ReaderRecoveryGateTest {
}
@Test
- void testEmptyRepeatedReportRetainsPreviousReport() {
- ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ void testEmptyRepeatedReportPreservesPendingSplits() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
DynamicKafkaSourceSplit firstReport = split("topic", 0);
gate.recordReportedSplits(1, Collections.singletonList(firstReport));
gate.recordReportedSplits(1, Collections.emptyList());
@@ -149,8 +158,51 @@ class ReaderRecoveryGateTest {
}
@Test
- void testDrainReportedSplitsReturnsReaderOrderAndClears() {
+ void testReturnedSplitsRemainPendingWithoutInitialRecoveryGate() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ DynamicKafkaSourceSplit firstReturnedSplit = split("topic", 0);
+ DynamicKafkaSourceSplit secondReturnedSplit = split("topic", 1);
+ DynamicKafkaSourceSplit updatedSplit =
+ new DynamicKafkaSourceSplit(
+ "cluster0", new KafkaPartitionSplit(new
TopicPartition("topic", 0), 42L));
+ gate.recordReturnedSplits(1,
Collections.singletonList(firstReturnedSplit));
+ gate.recordReturnedSplits(1, Arrays.asList(updatedSplit,
secondReturnedSplit));
+
+ assertThat(gate.hasPendingRecovery()).isFalse();
+ assertThat(gate.hasReportedSplits()).isTrue();
+ assertThat(gate.snapshotReportedSplits().get(1))
+ .containsExactlyInAnyOrder(updatedSplit, secondReturnedSplit);
+
assertThat(gate.snapshotReportedSplits().get(1).get(0).getStartingOffset()).isEqualTo(42L);
+ }
+
+ @Test
+ void testDrainReportedSplitsForOneReaderPreservesOtherReaders() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ DynamicKafkaSourceSplit splitReader1 = split("topic", 1);
+ DynamicKafkaSourceSplit splitReader2 = split("topic", 2);
+ gate.recordReturnedSplits(1, Collections.singletonList(splitReader1));
+ gate.recordReturnedSplits(2, Collections.singletonList(splitReader2));
+
+ assertThat(gate.drainReportedSplits(1)).containsExactly(splitReader1);
+ assertThat(gate.drainReportedSplits(1)).isEmpty();
+ assertThat(gate.snapshotReportedSplits()).containsOnlyKeys(2);
+
assertThat(gate.snapshotReportedSplits().get(2)).containsExactly(splitReader2);
+ }
+
+ @Test
+ void testEmptyReturnedSplitsAreIgnored() {
ReaderRecoveryGate gate = new ReaderRecoveryGate(false);
+ gate.recordReturnedSplits(1, Collections.emptyList());
+
+ assertThat(gate.hasPendingRecovery()).isFalse();
+ assertThat(gate.hasReportedSplits()).isFalse();
+ assertThat(gate.snapshotReportedSplits()).isEmpty();
+ assertThat(gate.drainReportedSplits(1)).isEmpty();
+ }
+
+ @Test
+ void testDrainReportedSplitsReturnsReaderOrderAndClears() {
+ ReaderRecoveryGate gate = new ReaderRecoveryGate(true);
DynamicKafkaSourceSplit splitReader2 = split("topic", 2);
DynamicKafkaSourceSplit splitReader0 = split("topic", 0);
gate.recordReportedSplits(2, Collections.singletonList(splitReader2));
diff --git
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorPendingOwnershipRegressionTest.java
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorPendingOwnershipRegressionTest.java
new file mode 100644
index 00000000..7468db74
--- /dev/null
+++
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorPendingOwnershipRegressionTest.java
@@ -0,0 +1,149 @@
+/*
+ * 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.source.enumerator;
+
+import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.api.connector.source.ReaderInfo;
+import org.apache.flink.api.connector.source.mocks.MockSplitEnumeratorContext;
+import
org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
+import
org.apache.flink.connector.kafka.source.enumerator.subscriber.KafkaSubscriber;
+import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit;
+
+import org.apache.kafka.common.TopicPartition;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Properties;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Offline regressions for physical-partition ownership during returned-split
reconciliation. */
+class KafkaSourceEnumeratorPendingOwnershipRegressionTest {
+
+ @Test
+ void sameOwnerSameOffsetIsIdempotentControl() throws Exception {
+ assertSinglePendingPhysicalSplit(0, 0, 100L, 100L);
+ }
+
+ @Test
+ void changedOwnerMustNotRetainPreviousPendingOwner() throws Exception {
+ assertSinglePendingPhysicalSplit(0, 1, 100L, 100L);
+ }
+
+ @Test
+ void changedOwnerAndOffsetMustNotRetainPreviousPendingCopy() throws
Exception {
+ assertSinglePendingPhysicalSplit(0, 1, 100L, 200L);
+ }
+
+ @Test
+ void changedOffsetMustNotRetainTwoCopiesOnSameOwner() throws Exception {
+ assertSinglePendingPhysicalSplit(0, 0, 100L, 200L);
+ }
+
+ @Test
+ void lowerCheckpointOffsetReplacesHigherPendingOffset() throws Exception {
+ assertSinglePendingPhysicalSplit(0, 1, 200L, 100L);
+ }
+
+ @ParameterizedTest
+ @CsvSource({"100, 200", "200, 100"})
+ void sameBatchKeepsLastPhysicalSplitAndPreservesOtherPartitions(
+ long firstOffset, long lastOffset) throws Exception {
+ TopicPartition partition = new TopicPartition("topic", 0);
+ KafkaPartitionSplit unrelated =
+ new KafkaPartitionSplit(new TopicPartition("topic", 1), 500L);
+ KafkaPartitionSplit expected = new KafkaPartitionSplit(partition,
lastOffset);
+ try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
+ new MockSplitEnumeratorContext<>(1);
+ KafkaSourceEnumerator enumerator =
+ new KafkaSourceEnumerator(
+ client -> Collections.emptySet(),
+ OffsetsInitializer.earliest(),
+ OffsetsInitializer.earliest(),
+ new Properties(),
+ context,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ new KafkaSourceEnumState(
+ Collections.emptyList(),
Collections.emptyList(), true))) {
+ enumerator.addSplitsBack(List.of(unrelated), 0);
+ enumerator.addSplitsBack(
+ List.of(new KafkaPartitionSplit(partition, firstOffset),
expected), 0);
+ assertThat(enumerator.snapshotState(1L).unassignedSplits())
+ .containsExactlyInAnyOrder(unrelated, expected);
+
+ context.registerReader(new ReaderInfo(0, "reader-0"));
+ enumerator.addReader(0);
+
+ assertThat(context.getSplitsAssignmentSequence()).hasSize(1);
+
assertThat(context.getSplitsAssignmentSequence().get(0).assignment().get(0))
+ .containsExactlyInAnyOrder(unrelated, expected);
+ assertThat(enumerator.snapshotState(2L).assignedSplits())
+ .containsExactlyInAnyOrder(unrelated, expected);
+ }
+ }
+
+ private void assertSinglePendingPhysicalSplit(
+ int firstOwner, int secondOwner, long firstOffset, long
secondOffset) throws Exception {
+ AtomicInteger owner = new AtomicInteger(firstOwner);
+ try (MockSplitEnumeratorContext<KafkaPartitionSplit> context =
+ new MockSplitEnumeratorContext<>(2)) {
+ KafkaSubscriber subscriber = client -> Collections.emptySet();
+ KafkaSourceEnumerator enumerator =
+ new KafkaSourceEnumerator(
+ subscriber,
+ OffsetsInitializer.earliest(),
+ OffsetsInitializer.earliest(),
+ new Properties(),
+ context,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ new KafkaSourceEnumState(
+ Collections.emptySet(),
Collections.emptySet(), true),
+ (split, readers) -> owner.get());
+ TopicPartition partition = new TopicPartition("topic", 2);
+ enumerator.addSplitsBack(
+ Collections.singletonList(new
KafkaPartitionSplit(partition, firstOffset)),
+ firstOwner);
+ owner.set(secondOwner);
+ enumerator.addSplitsBack(
+ Collections.singletonList(new
KafkaPartitionSplit(partition, secondOffset)),
+ secondOwner);
+ for (int reader = 0; reader < 2; reader++) {
+ context.registerReader(new ReaderInfo(reader, "reader-" +
reader));
+ enumerator.addReader(reader);
+ }
+ List<String> emitted =
+ context.getSplitsAssignmentSequence().stream()
+ .flatMap(a -> a.assignment().entrySet().stream())
+ .flatMap(e -> e.getValue().stream().map(s ->
e.getKey() + ":" + s))
+ .collect(Collectors.toList());
+ assertThat(emitted)
+ .as("one pending assignment for physical topic partition
%s", partition)
+ .hasSize(1);
+
assertThat(context.getSplitsAssignmentSequence().get(0).assignment())
+ .containsOnlyKeys(secondOwner);
+
assertThat(context.getSplitsAssignmentSequence().get(0).assignment().get(secondOwner))
+ .containsExactly(new KafkaPartitionSplit(partition,
secondOffset));
+ }
+ }
+}
diff --git
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorRecoveryDiscoveryRaceTest.java
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorRecoveryDiscoveryRaceTest.java
new file mode 100644
index 00000000..06e6ba41
--- /dev/null
+++
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorRecoveryDiscoveryRaceTest.java
@@ -0,0 +1,230 @@
+/*
+ * 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.source.enumerator;
+
+import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.api.connector.source.ReaderInfo;
+import org.apache.flink.api.connector.source.mocks.MockSplitEnumeratorContext;
+import org.apache.flink.connector.kafka.source.KafkaSourceOptions;
+import
org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
+import
org.apache.flink.connector.kafka.source.enumerator.subscriber.KafkaSubscriber;
+import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit;
+import org.apache.flink.mock.Whitebox;
+
+import org.apache.kafka.clients.admin.MockAdminClient;
+import org.apache.kafka.common.TopicPartition;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.lang.reflect.Method;
+import java.util.ArrayDeque;
+import java.util.Collections;
+import java.util.Deque;
+import java.util.List;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.function.BiConsumer;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.SoftAssertions.assertSoftly;
+
+/** Replays delayed initialization callbacks against newer authoritative
checkpoint splits. */
+class KafkaSourceEnumeratorRecoveryDiscoveryRaceTest {
+ private static final TopicPartition PARTITION = new
TopicPartition("topic", 0);
+ private static final long CHECKPOINT_STOP = 1500L;
+ private static final long DISCOVERY_STOP = 2000L;
+
+ @ParameterizedTest
+ @CsvSource({
+ "120, 900, false",
+ "900, 120, false",
+ "120, 900, true",
+ "900, 120, true",
+ "-2, 900, false",
+ "-3, 900, false"
+ })
+ void checkpointSplitWinsOverDelayedDiscoveryRegardlessOfOffsetOrder(
+ long checkpointOffset, long discoveryOffset, boolean
alreadyRegistered)
+ throws Throwable {
+ KafkaPartitionSplit checkpoint =
+ new KafkaPartitionSplit(PARTITION, checkpointOffset,
CHECKPOINT_STOP);
+ try (HoldingCompletionContext context = new HoldingCompletionContext();
+ KafkaSourceEnumerator enumerator =
+ newEnumerator(context, Collections.emptyList(),
discoveryOffset)) {
+ discover(enumerator);
+ // Compute the real worker result first, but hold its real
completion callback.
+ context.runNextOneTimeCallable();
+ assertThat(context.completions).hasSize(1);
+ assertThat(context.getSplitsAssignmentSequence()).isEmpty();
+ if (alreadyRegistered) {
+ registerReader(context, enumerator);
+ }
+ enumerator.addSplitsBack(Collections.singletonList(checkpoint), 0);
+ KafkaSourceEnumState before = enumerator.snapshotState(1L);
+ assertThat(alreadyRegistered ? before.assignedSplits() :
before.unassignedSplits())
+ .containsExactly(checkpoint);
+ assertThat(before.initialDiscoveryFinished()).isFalse();
+
+ context.deliverCompletion();
+ KafkaSourceEnumState afterCallback = enumerator.snapshotState(2L);
+ List<KafkaPartitionSplit> pendingAfterCallback =
+
enumerator.getPendingPartitionSplitAssignment().values().stream()
+ .flatMap(Set::stream)
+ .collect(Collectors.toList());
+ if (!alreadyRegistered) {
+ registerReader(context, enumerator);
+ }
+ KafkaSourceEnumState afterRegistration =
enumerator.snapshotState(3L);
+ // Collect both persisted-state and emitted-assignment failures in
the red baseline.
+ assertSoftly(
+ softly -> {
+
softly.assertThat(afterCallback.initialDiscoveryFinished()).isTrue();
+ softly.assertThat(
+ alreadyRegistered
+ ?
afterCallback.assignedSplits()
+ :
afterCallback.unassignedSplits())
+ .as(
+ "checkpoint split survives the delayed
callback in snapshot state")
+ .containsExactly(checkpoint);
+ softly.assertThat(
+ alreadyRegistered
+ ?
afterCallback.unassignedSplits()
+ :
afterCallback.assignedSplits())
+ .isEmpty();
+ if (alreadyRegistered) {
+ softly.assertThat(pendingAfterCallback).isEmpty();
+ } else {
+
softly.assertThat(pendingAfterCallback).containsExactly(checkpoint);
+ }
+ softly.assertThat(emittedSplits(context))
+ .as(
+ "exactly one checkpoint assignment;
never the stale discovery copy")
+ .containsExactly(checkpoint);
+ softly.assertThat(afterRegistration.assignedSplits())
+ .containsExactly(checkpoint);
+
softly.assertThat(afterRegistration.unassignedSplits()).isEmpty();
+ });
+ assertCompletionBookkeeping(context, enumerator);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void absentOrUnresolvedMigratedPartitionStillAcceptsInitialization(boolean
migrated)
+ throws Throwable {
+ List<KafkaPartitionSplit> restored =
+ migrated
+ ? Collections.singletonList(
+ new KafkaPartitionSplit(PARTITION,
KafkaPartitionSplit.MIGRATED))
+ : Collections.emptyList();
+ KafkaPartitionSplit initialized = new KafkaPartitionSplit(PARTITION,
900L, DISCOVERY_STOP);
+ try (HoldingCompletionContext context = new HoldingCompletionContext();
+ KafkaSourceEnumerator enumerator = newEnumerator(context,
restored, 900L)) {
+ discover(enumerator);
+ context.runNextOneTimeCallable();
+ assertThat(context.completions).hasSize(1);
+ context.deliverCompletion();
+ KafkaSourceEnumState pending = enumerator.snapshotState(1L);
+
assertThat(pending.unassignedSplits()).containsExactly(initialized);
+ assertThat(pending.assignedSplits()).isEmpty();
+ registerReader(context, enumerator);
+ assertThat(emittedSplits(context)).containsExactly(initialized);
+ KafkaSourceEnumState assigned = enumerator.snapshotState(2L);
+ assertThat(assigned.assignedSplits()).containsExactly(initialized);
+ assertThat(assigned.unassignedSplits()).isEmpty();
+ assertCompletionBookkeeping(context, enumerator);
+ }
+ }
+
+ private static KafkaSourceEnumerator newEnumerator(
+ HoldingCompletionContext context,
+ List<KafkaPartitionSplit> restored,
+ long discoveryOffset) {
+ Properties properties = new Properties();
+
properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
"0");
+ KafkaSubscriber subscriber = client ->
Collections.singleton(PARTITION);
+ KafkaSourceEnumerator enumerator =
+ new KafkaSourceEnumerator(
+ subscriber,
+ OffsetsInitializer.offsets(
+ Collections.singletonMap(PARTITION,
discoveryOffset)),
+ OffsetsInitializer.offsets(
+ Collections.singletonMap(PARTITION,
DISCOVERY_STOP)),
+ properties,
+ context,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ new KafkaSourceEnumState(Collections.emptyList(),
restored, false));
+ // No broker is involved; the production initializer resolves the
specified offsets.
+ Whitebox.setInternalState(enumerator, "adminClient", new
MockAdminClient());
+ return enumerator;
+ }
+
+ private static void discover(KafkaSourceEnumerator enumerator) throws
Exception {
+ Method discover =
+ KafkaSourceEnumerator.class.getDeclaredMethod(
+ "checkPartitionChanges", Set.class, Throwable.class);
+ discover.setAccessible(true);
+ discover.invoke(enumerator, Collections.singleton(PARTITION), null);
+ }
+
+ private static void registerReader(
+ HoldingCompletionContext context, KafkaSourceEnumerator
enumerator) {
+ context.registerReader(new ReaderInfo(0, "reader-0"));
+ enumerator.addReader(0);
+ }
+
+ private static List<KafkaPartitionSplit>
emittedSplits(HoldingCompletionContext context) {
+ return context.getSplitsAssignmentSequence().stream()
+ .flatMap(assignment ->
assignment.assignment().values().stream())
+ .flatMap(List::stream)
+ .collect(Collectors.toList());
+ }
+
+ private static void assertCompletionBookkeeping(
+ HoldingCompletionContext context, KafkaSourceEnumerator
enumerator) throws Exception {
+
assertThat(enumerator.snapshotState(4L).initialDiscoveryFinished()).isTrue();
+ discover(enumerator);
+ assertThat(context.getOneTimeCallables()).isEmpty();
+ assertThat(context.completions).isEmpty();
+ }
+
+ /** Holds only delivery of the actual initializer result, not its
computation or split state. */
+ private static class HoldingCompletionContext
+ extends MockSplitEnumeratorContext<KafkaPartitionSplit> {
+ private final Deque<Runnable> completions = new ArrayDeque<>();
+
+ private HoldingCompletionContext() {
+ super(1);
+ }
+
+ @Override
+ public <T> void callAsync(Callable<T> callable, BiConsumer<T,
Throwable> handler) {
+ super.callAsync(
+ callable,
+ (result, error) -> completions.addLast(() ->
handler.accept(result, error)));
+ }
+
+ private void deliverCompletion() {
+ completions.removeFirst().run();
+ }
+ }
+}
