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 d43f3ae6 [FLINK-39940] DynamicKafkaSource should refresh restored
cluster properties from KafkaMetadataService (#266)
d43f3ae6 is described below
commit d43f3ae66531d39a12102cac18123603ca73eb0e
Author: bowenli86 <[email protected]>
AuthorDate: Wed Jun 17 11:38:00 2026 -0700
[FLINK-39940] DynamicKafkaSource should refresh restored cluster properties
from KafkaMetadataService (#266)
DynamicKafkaSource restores KafkaStream metadata from the checkpointed
enumerator state when a job is resumed from a checkpoint or savepoint. Before
this change, restored sub-enumerators were rebuilt from the checkpointed
cluster properties without first refreshing them from KafkaMetadataService.
That means non-bootstrap Kafka client properties such as authentication or
security settings can be stale at restore time, and the job can fail before the
next periodic metadata refresh has a [...]
The restore path should use the latest metadata-service properties
immediately, while still preserving the checkpointed bootstrap.servers for
restored clusters so the enumerator reconnects to the same restored cluster
identity. Restore-time metadata lookup failures should also fail fast instead
of being deferred until a later discovery cycle.
This PR also adds DEBUG-level startup-offset logs for checkpoint snapshot
and restore to make restore issues easier to diagnose.
---
.../enumerator/DynamicKafkaSourceEnumerator.java | 130 ++++++-
.../DynamicKafkaSourceEnumeratorTest.java | 395 +++++++++++++++++++--
2 files changed, 488 insertions(+), 37 deletions(-)
diff --git
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumerator.java
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumerator.java
index 73c0a143..ba20fa98 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
@@ -41,6 +41,7 @@ import
org.apache.flink.connector.kafka.source.enumerator.subscriber.KafkaSubscr
import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit;
import org.apache.flink.util.Preconditions;
+import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.common.KafkaException;
import org.slf4j.Logger;
@@ -50,7 +51,9 @@ import javax.annotation.Nullable;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.Collections;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
@@ -158,10 +161,18 @@ public class DynamicKafkaSourceEnumerator
this.stoppableKafkaEnumContextProxyFactory =
stoppableKafkaEnumContextProxyFactory;
this.splitAssignmentStrategy =
createSplitAssignmentStrategy(properties);
+ if
(!dynamicKafkaSourceEnumState.getClusterEnumeratorStates().isEmpty()) {
+ logger.info("Dynamic Kafka source restored from checkpointed
enumerator state");
+ }
+
// handle checkpoint state and rebuild contexts
this.clusterEnumeratorMap = new HashMap<>();
this.clusterEnumContextMap = new HashMap<>();
this.latestKafkaStreams =
dynamicKafkaSourceEnumState.getKafkaStreams();
+ if (!this.latestKafkaStreams.isEmpty()) {
+ this.latestKafkaStreams =
+
refreshRestoredClusterPropertiesFromMetadataService(this.latestKafkaStreams);
+ }
Map<String, Properties> clusterProperties = new HashMap<>();
Map<String, OffsetsInitializer> clusterStartingOffsets = new
HashMap<>();
@@ -186,9 +197,18 @@ public class DynamicKafkaSourceEnumerator
Set<String> activeSplitIds = new HashSet<>();
for (Entry<String, KafkaSourceEnumState> clusterEnumState :
dynamicKafkaSourceEnumState.getClusterEnumeratorStates().entrySet()) {
+ String clusterId = clusterEnumState.getKey();
+ KafkaSourceEnumState state = clusterEnumState.getValue();
+ if (!state.assignedSplits().isEmpty() ||
!state.unassignedSplits().isEmpty()) {
+ logger.debug(
+ "Restored enumerator startup offsets for cluster {}
assigned={} unassigned={}",
+ clusterId,
+ summarizeSplitOffsets(state.assignedSplits()),
+ summarizeSplitOffsets(state.unassignedSplits()));
+ }
this.latestClusterTopicsMap.put(
- clusterEnumState.getKey(),
- clusterEnumState.getValue().assignedSplits().stream()
+ clusterId,
+ state.assignedSplits().stream()
.map(KafkaPartitionSplit::getTopic)
.collect(Collectors.toSet()));
clusterEnumState
@@ -197,21 +217,81 @@ public class DynamicKafkaSourceEnumerator
.forEach(
splitStatus ->
activeSplitIds.add(
- toDynamicSplitId(
- clusterEnumState.getKey(),
- splitStatus.split())));
+ toDynamicSplitId(clusterId,
splitStatus.split())));
createEnumeratorWithAssignedTopicPartitions(
- clusterEnumState.getKey(),
- this.latestClusterTopicsMap.get(clusterEnumState.getKey()),
- clusterEnumState.getValue(),
- clusterProperties.get(clusterEnumState.getKey()),
- clusterStartingOffsets.get(clusterEnumState.getKey()),
- clusterStoppingOffsets.get(clusterEnumState.getKey()));
+ clusterId,
+ this.latestClusterTopicsMap.get(clusterId),
+ state,
+ clusterProperties.get(clusterId),
+ clusterStartingOffsets.get(clusterId),
+ clusterStoppingOffsets.get(clusterId));
}
splitAssignmentStrategy.onMetadataRefresh(activeSplitIds);
}
+ private Set<KafkaStream>
refreshRestoredClusterPropertiesFromMetadataService(
+ Set<KafkaStream> restoredKafkaStreams) {
+ Set<KafkaStream> fetchedKafkaStreams =
+
kafkaStreamSubscriber.getSubscribedStreams(kafkaMetadataService);
+
+ Map<String, Properties> fetchedClusterPropertiesById =
+ extractClusterPropertiesById(fetchedKafkaStreams);
+ Set<KafkaStream> mergedKafkaStreams = new HashSet<>();
+ for (KafkaStream restoredKafkaStream : restoredKafkaStreams) {
+ Map<String, ClusterMetadata> mergedClusterMetadataMap = new
HashMap<>();
+ for (Entry<String, ClusterMetadata> restoredClusterEntry :
+ restoredKafkaStream.getClusterMetadataMap().entrySet()) {
+ String kafkaClusterId = restoredClusterEntry.getKey();
+ ClusterMetadata restoredClusterMetadata =
restoredClusterEntry.getValue();
+
+ Properties mergedProperties = new Properties();
+ Properties fetchedProperties =
fetchedClusterPropertiesById.get(kafkaClusterId);
+ if (fetchedProperties != null) {
+ KafkaPropertiesUtil.copyProperties(fetchedProperties,
mergedProperties);
+ }
+
+ String restoredBootstrapServers =
+ restoredClusterMetadata
+ .getProperties()
+
.getProperty(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG);
+ if (restoredBootstrapServers != null) {
+ mergedProperties.setProperty(
+ CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
restoredBootstrapServers);
+ }
+ if (mergedProperties.isEmpty()) {
+ KafkaPropertiesUtil.copyProperties(
+ restoredClusterMetadata.getProperties(),
mergedProperties);
+ }
+
+ mergedClusterMetadataMap.put(
+ kafkaClusterId,
+ new ClusterMetadata(
+ restoredClusterMetadata.getTopics(),
+ mergedProperties,
+
restoredClusterMetadata.getStartingOffsetsInitializer(),
+
restoredClusterMetadata.getStoppingOffsetsInitializer()));
+ }
+ mergedKafkaStreams.add(
+ new KafkaStream(restoredKafkaStream.getStreamId(),
mergedClusterMetadataMap));
+ }
+
+ return mergedKafkaStreams;
+ }
+
+ private static Map<String, Properties> extractClusterPropertiesById(
+ Set<KafkaStream> kafkaStreams) {
+ Map<String, Properties> clusterPropertiesById = new HashMap<>();
+ for (KafkaStream kafkaStream : kafkaStreams) {
+ for (Entry<String, ClusterMetadata> clusterEntry :
+ kafkaStream.getClusterMetadataMap().entrySet()) {
+ clusterPropertiesById.put(
+ clusterEntry.getKey(),
clusterEntry.getValue().getProperties());
+ }
+ }
+ return clusterPropertiesById;
+ }
+
/**
* Discover Kafka clusters and initialize sub enumerators. Bypass kafka
metadata service
* discovery if there exists prior state. Exceptions with initializing
Kafka source are treated
@@ -563,18 +643,40 @@ public class DynamicKafkaSourceEnumerator
@Override
public DynamicKafkaSourceEnumState snapshotState(long checkpointId) throws
Exception {
Map<String, KafkaSourceEnumState> subEnumeratorStateByCluster = new
HashMap<>();
+ boolean isCheckpointSnapshot = checkpointId >= 0;
// populate map for all assigned splits
for (Entry<String, SplitEnumerator<KafkaPartitionSplit,
KafkaSourceEnumState>>
clusterEnumerator : clusterEnumeratorMap.entrySet()) {
- subEnumeratorStateByCluster.put(
- clusterEnumerator.getKey(),
- clusterEnumerator.getValue().snapshotState(checkpointId));
+ KafkaSourceEnumState state =
clusterEnumerator.getValue().snapshotState(checkpointId);
+ subEnumeratorStateByCluster.put(clusterEnumerator.getKey(), state);
+ if (isCheckpointSnapshot) {
+ logger.debug(
+ "Checkpoint {} cluster {} enumerator startup offsets
for assigned splits {}",
+ checkpointId,
+ clusterEnumerator.getKey(),
+ summarizeSplitOffsets(state.assignedSplits()));
+ logger.debug(
+ "Checkpoint {} cluster {} enumerator startup offsets
for unassigned splits {}",
+ checkpointId,
+ clusterEnumerator.getKey(),
+ summarizeSplitOffsets(state.unassignedSplits()));
+ }
}
return new DynamicKafkaSourceEnumState(latestKafkaStreams,
subEnumeratorStateByCluster);
}
+ private static String
summarizeSplitOffsets(Collection<KafkaPartitionSplit> splits) {
+ if (splits.isEmpty()) {
+ return "[]";
+ }
+ return splits.stream()
+ .sorted(Comparator.comparing(split ->
split.getTopicPartition().toString()))
+ .map(split -> split.getTopicPartition() + "=" +
split.getStartingOffset())
+ .collect(Collectors.joining(",", "[", "]"));
+ }
+
@Override
public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent) {
Preconditions.checkArgument(
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 40e0dd00..e6867d32 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
@@ -29,6 +29,7 @@ import
org.apache.flink.connector.kafka.dynamic.metadata.KafkaMetadataService;
import org.apache.flink.connector.kafka.dynamic.metadata.KafkaStream;
import
org.apache.flink.connector.kafka.dynamic.source.DynamicKafkaSourceOptions;
import org.apache.flink.connector.kafka.dynamic.source.GetMetadataUpdateEvent;
+import org.apache.flink.connector.kafka.dynamic.source.MetadataUpdateEvent;
import
org.apache.flink.connector.kafka.dynamic.source.enumerator.subscriber.KafkaStreamSetSubscriber;
import
org.apache.flink.connector.kafka.dynamic.source.split.DynamicKafkaSourceSplit;
import org.apache.flink.connector.kafka.source.KafkaSourceOptions;
@@ -44,6 +45,7 @@ import org.apache.flink.runtime.state.OperatorStateHandle;
import org.apache.flink.runtime.state.OperatorStreamStateHandle;
import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
import
org.apache.flink.streaming.connectors.kafka.DynamicKafkaSourceTestHelper;
+import org.apache.flink.testutils.logging.LoggerAuditingExtension;
import com.google.common.collect.ImmutableSet;
import org.apache.kafka.clients.CommonClientConfigs;
@@ -52,6 +54,7 @@ import org.apache.kafka.common.TopicPartition;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
import java.util.ArrayList;
import java.util.Collection;
@@ -61,6 +64,7 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
+import java.util.Objects;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.Callable;
@@ -72,6 +76,7 @@ import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatCode;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.slf4j.event.Level.DEBUG;
/** A test for {@link DynamicKafkaSourceEnumerator}. */
public class DynamicKafkaSourceEnumeratorTest {
@@ -80,6 +85,12 @@ public class DynamicKafkaSourceEnumeratorTest {
private static final int NUM_SPLITS_PER_CLUSTER = 3;
private static final int NUM_RECORDS_PER_SPLIT = 5;
private static final String SOURCE_READER_SPLIT_STATE_NAME =
"source-reader-splits";
+ private static final String REFRESHED_CLUSTER_PROPERTY_KEY =
"refreshed.cluster.property";
+ private static final String REFRESHED_CLUSTER_PROPERTY_VALUE =
"from-metadata-service";
+
+ @RegisterExtension
+ public final LoggerAuditingExtension dynamicEnumeratorLogger =
+ new LoggerAuditingExtension(DynamicKafkaSourceEnumerator.class,
DEBUG);
@BeforeAll
public static void beforeAll() throws Throwable {
@@ -175,9 +186,7 @@ public class DynamicKafkaSourceEnumeratorTest {
}
@Test
- public void
-
testStartupWithKafkaMetadataServiceFailure_withContinuousDiscoveryAndCheckpointState()
- throws Throwable {
+ public void testRestoreWithKafkaMetadataServiceFailure_surfacesFailure()
throws Throwable {
// init enumerator with checkpoint state
final DynamicKafkaSourceEnumState dynamicKafkaSourceEnumState =
getCheckpointState();
Properties properties = new Properties();
@@ -185,26 +194,22 @@ public class DynamicKafkaSourceEnumeratorTest {
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "1");
properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
"0");
try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
- new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
- DynamicKafkaSourceEnumerator enumerator =
- new DynamicKafkaSourceEnumerator(
- new
KafkaStreamSetSubscriber(Collections.singleton(TOPIC)),
- new MockKafkaMetadataService(true),
- context,
- OffsetsInitializer.committedOffsets(),
- new NoStoppingOffsetsInitializer(),
- properties,
- Boundedness.CONTINUOUS_UNBOUNDED,
- dynamicKafkaSourceEnumState,
- new TestKafkaEnumContextProxyFactory())) {
- enumerator.start();
-
- assertThat(context.getPeriodicCallables()).hasSize(1);
- // no exception
- context.runPeriodicCallable(0);
-
- assertThatThrownBy(() -> context.runPeriodicCallable(0))
- .hasRootCause(new RuntimeException("Mock exception"));
+ new MockSplitEnumeratorContext<>(NUM_SUBTASKS)) {
+ assertThatThrownBy(
+ () ->
+ new DynamicKafkaSourceEnumerator(
+ new KafkaStreamSetSubscriber(
+
Collections.singleton(TOPIC)),
+ new MockKafkaMetadataService(true),
+ context,
+
OffsetsInitializer.committedOffsets(),
+ new NoStoppingOffsetsInitializer(),
+ properties,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ dynamicKafkaSourceEnumState,
+ new
TestKafkaEnumContextProxyFactory()))
+ .isInstanceOf(RuntimeException.class)
+ .hasMessage("Mock exception");
}
}
@@ -475,6 +480,124 @@ public class DynamicKafkaSourceEnumeratorTest {
}
}
+ @Test
+ public void testSnapshotStateLogsAssignedAndUnassignedOffsets() throws
Throwable {
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
+ DynamicKafkaSourceEnumerator enumerator =
createEnumerator(context)) {
+ enumerator.start();
+
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
0);
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
1);
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
2);
+ runAllOneTimeCallables(context);
+
+ enumerator.snapshotState(7L);
+
+ // Snapshot logs in dynamic Kafka source encode special offset
sentinels:
+ // -2 = EARLIEST_OFFSET in KafkaPartitionSplit (not a real
partition offset).
+ assertThat(dynamicEnumeratorLogger.getMessages())
+ .as(
+ "checkpoint snapshot should label enumerator
startup offsets for assigned splits")
+ .anyMatch(
+ message ->
+ message.equals(
+ "Checkpoint 7 cluster
kafka-cluster-0 enumerator startup offsets for assigned splits "
+ + "["
+ + TOPIC
+ + "-0=-2,"
+ + TOPIC
+ + "-1=-2,"
+ + TOPIC
+ + "-2=-2]"))
+ .as(
+ "checkpoint snapshot should label enumerator
startup offsets for unassigned splits")
+ .anyMatch(
+ message ->
+ message.equals(
+ "Checkpoint 7 cluster
kafka-cluster-0 enumerator startup offsets for unassigned splits []"))
+ .as(
+ "checkpoint snapshot should include equivalent
logs for second cluster in the same format")
+ .anyMatch(
+ message ->
+ message.equals(
+ "Checkpoint 7 cluster
kafka-cluster-1 enumerator startup offsets for assigned splits "
+ + "["
+ + TOPIC
+ + "-0=-2,"
+ + TOPIC
+ + "-1=-2,"
+ + TOPIC
+ + "-2=-2]"))
+ .as(
+ "checkpoint snapshot should include empty
unassigned offsets for second cluster")
+ .anyMatch(
+ message ->
+ message.equals(
+ "Checkpoint 7 cluster
kafka-cluster-1 enumerator startup offsets for unassigned splits []"));
+ }
+ }
+
+ @Test
+ public void testRestoreLogsCheckpointedOffsets() throws Throwable {
+ Properties properties = new Properties();
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "0");
+
properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
"0");
+
+ DynamicKafkaSourceEnumState checkpointState = getCheckpointState();
+
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(NUM_SUBTASKS)) {
+ try (DynamicKafkaSourceEnumerator restoredEnumerator =
+ new DynamicKafkaSourceEnumerator(
+ new
KafkaStreamSetSubscriber(Collections.singleton(TOPIC)),
+ new MockKafkaMetadataService(
+ Collections.singleton(
+
DynamicKafkaSourceTestHelper.getKafkaStream(TOPIC))),
+ context,
+ OffsetsInitializer.earliest(),
+ new NoStoppingOffsetsInitializer(),
+ properties,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ checkpointState,
+ new TestKafkaEnumContextProxyFactory())) {
+ restoredEnumerator.start();
+
+ assertThat(dynamicEnumeratorLogger.getMessages())
+ .as("restore path should emit checkpoint state logs")
+ .anyMatch(
+ message ->
+ message.contains(
+ "Dynamic Kafka source restored
from checkpointed enumerator state"))
+ .as(
+ "restore path should label restored enumerator
startup offsets for first active cluster")
+ .anyMatch(
+ message ->
+ message.equals(
+ "Restored enumerator startup
offsets for cluster kafka-cluster-0 assigned=["
+ + TOPIC
+ + "-0=-2,"
+ + TOPIC
+ + "-1=-2,"
+ + TOPIC
+ + "-2=-2]
unassigned=[]"))
+ .as(
+ "restore path should label restored enumerator
startup offsets for second active cluster")
+ .anyMatch(
+ message ->
+ message.equals(
+ "Restored enumerator startup
offsets for cluster kafka-cluster-1 assigned=["
+ + TOPIC
+ + "-0=-2,"
+ + TOPIC
+ + "-1=-2,"
+ + TOPIC
+ + "-2=-2]
unassigned=[]"));
+ }
+ }
+ }
+
@Test
public void testEnumeratorStateDoesNotContainStaleTopicPartitions() throws
Throwable {
final String topic2 = TOPIC + "_2";
@@ -653,6 +776,201 @@ public class DynamicKafkaSourceEnumeratorTest {
}
}
+ @Test
+ public void
testRestoreRefreshesNonBootstrapClusterPropertiesFromMetadataService()
+ throws Throwable {
+ KafkaStream baseStream =
DynamicKafkaSourceTestHelper.getKafkaStream(TOPIC);
+ DynamicKafkaSourceEnumState checkpointState =
getCheckpointState(baseStream);
+
+ DynamicKafkaSourceEnumStateSerializer serializer =
+ new DynamicKafkaSourceEnumStateSerializer();
+ DynamicKafkaSourceEnumState serializedRestoredState =
+ serializer.deserialize(
+ serializer.getVersion(),
serializer.serialize(checkpointState));
+
+ KafkaStream restoredStream =
+ serializedRestoredState.getKafkaStreams().stream()
+ .findFirst()
+ .orElseThrow(() -> new AssertionError("Missing
restored stream"));
+ Map<String, String> restoredBootstrapByCluster = new HashMap<>();
+ for (Entry<String, ClusterMetadata> entry :
+ restoredStream.getClusterMetadataMap().entrySet()) {
+ restoredBootstrapByCluster.put(
+ entry.getKey(),
+ entry.getValue()
+ .getProperties()
+
.getProperty(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG));
+ }
+
+ Map<String, ClusterMetadata> refreshedClusterMetadataMap = new
HashMap<>();
+ for (Entry<String, ClusterMetadata> entry :
+ restoredStream.getClusterMetadataMap().entrySet()) {
+ Properties refreshedProperties = new Properties();
+ refreshedProperties.setProperty(
+ CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
+ "placeholder-from-managed-config");
+ refreshedProperties.setProperty(
+ REFRESHED_CLUSTER_PROPERTY_KEY,
REFRESHED_CLUSTER_PROPERTY_VALUE);
+ refreshedClusterMetadataMap.put(
+ entry.getKey(),
+ new ClusterMetadata(
+ entry.getValue().getTopics(),
+ refreshedProperties,
+ entry.getValue().getStartingOffsetsInitializer(),
+ entry.getValue().getStoppingOffsetsInitializer()));
+ }
+ KafkaStream refreshedStream =
+ new KafkaStream(restoredStream.getStreamId(),
refreshedClusterMetadataMap);
+
+ Properties properties = new Properties();
+
properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
"0");
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "0");
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
+ DynamicKafkaSourceEnumerator enumerator =
+ new DynamicKafkaSourceEnumerator(
+ new
KafkaStreamSetSubscriber(Collections.singleton(TOPIC)),
+ new MockKafkaMetadataService(
+
Collections.singleton(refreshedStream)),
+ context,
+ OffsetsInitializer.committedOffsets(),
+ new NoStoppingOffsetsInitializer(),
+ properties,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ serializedRestoredState,
+ new TestKafkaEnumContextProxyFactory())) {
+ enumerator.start();
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
0);
+
+ MetadataUpdateEvent metadataUpdateEvent =
getLatestMetadataUpdateEvent(context, 0);
+ KafkaStream updatedStream =
+ metadataUpdateEvent.getKafkaStreams().stream()
+ .findFirst()
+ .orElseThrow(
+ () -> new AssertionError("Missing metadata
update stream"));
+ for (Entry<String, ClusterMetadata> entry :
+ updatedStream.getClusterMetadataMap().entrySet()) {
+ assertThat(
+ entry.getValue()
+ .getProperties()
+
.getProperty(REFRESHED_CLUSTER_PROPERTY_KEY))
+ .as(
+ "restored reader metadata should include
refreshed non-bootstrap properties")
+ .isEqualTo(REFRESHED_CLUSTER_PROPERTY_VALUE);
+ assertThat(
+ entry.getValue()
+ .getProperties()
+
.getProperty(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG))
+ .as("restored bootstrap should remain checkpoint
value")
+
.isEqualTo(restoredBootstrapByCluster.get(entry.getKey()));
+ }
+
+ Map<String, ?> clusterEnumeratorMap =
+ (Map<String, ?>) Whitebox.getInternalState(enumerator,
"clusterEnumeratorMap");
+ for (Entry<String, ?> clusterEnumeratorEntry :
clusterEnumeratorMap.entrySet()) {
+ Properties clusterEnumeratorProperties =
+ (Properties)
+ Whitebox.getInternalState(
+ clusterEnumeratorEntry.getValue(),
"properties");
+
assertThat(clusterEnumeratorProperties.getProperty(REFRESHED_CLUSTER_PROPERTY_KEY))
+ .as("restored sub-enumerator should use refreshed
non-bootstrap properties")
+ .isEqualTo(REFRESHED_CLUSTER_PROPERTY_VALUE);
+ assertThat(
+ clusterEnumeratorProperties.getProperty(
+
CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG))
+ .as("restored sub-enumerator should preserve
checkpoint bootstrap servers")
+
.isEqualTo(restoredBootstrapByCluster.get(clusterEnumeratorEntry.getKey()));
+ }
+ }
+ }
+
+ @Test
+ public void
testRestoreReconcilesRemovedAndAddedClustersFromLatestMetadata() throws
Throwable {
+ KafkaStream baseStream =
DynamicKafkaSourceTestHelper.getKafkaStream(TOPIC);
+ DynamicKafkaSourceEnumState checkpointState =
getCheckpointState(baseStream);
+
+ DynamicKafkaSourceEnumStateSerializer serializer =
+ new DynamicKafkaSourceEnumStateSerializer();
+ DynamicKafkaSourceEnumState restoredState =
+ serializer.deserialize(
+ serializer.getVersion(),
serializer.serialize(checkpointState));
+
+ String cluster0 = DynamicKafkaSourceTestHelper.getKafkaClusterId(0);
+ String removedCluster =
DynamicKafkaSourceTestHelper.getKafkaClusterId(1);
+ String addedCluster = removedCluster + "-replacement";
+
+ Map<String, ClusterMetadata> latestClusterMetadataMap = new
HashMap<>();
+ latestClusterMetadataMap.put(
+ cluster0,
+ copyClusterMetadataWithOverrides(
+
Objects.requireNonNull(baseStream.getClusterMetadataMap().get(cluster0)),
+ props ->
+ props.setProperty(
+ REFRESHED_CLUSTER_PROPERTY_KEY,
+ REFRESHED_CLUSTER_PROPERTY_VALUE)));
+ latestClusterMetadataMap.put(
+ addedCluster,
+ copyClusterMetadataWithOverrides(
+ Objects.requireNonNull(
+
baseStream.getClusterMetadataMap().get(removedCluster)),
+ props ->
+ props.setProperty(
+ REFRESHED_CLUSTER_PROPERTY_KEY,
+ REFRESHED_CLUSTER_PROPERTY_VALUE)));
+
+ KafkaStream latestMetadataStream =
+ new KafkaStream(baseStream.getStreamId(),
latestClusterMetadataMap);
+
+ Properties properties = new Properties();
+
properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
"0");
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "1");
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
+ DynamicKafkaSourceEnumerator enumerator =
+ new DynamicKafkaSourceEnumerator(
+ new
KafkaStreamSetSubscriber(Collections.singleton(TOPIC)),
+ new MockKafkaMetadataService(
+
Collections.singleton(latestMetadataStream)),
+ context,
+ OffsetsInitializer.committedOffsets(),
+ new NoStoppingOffsetsInitializer(),
+ properties,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ restoredState,
+ new TestKafkaEnumContextProxyFactory())) {
+ enumerator.start();
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
0);
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
1);
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
2);
+ runAllOneTimeCallables(context);
+
+ context.runPeriodicCallable(0);
+ runAllOneTimeCallables(context);
+
+ enumerator.handleSourceEvent(0, new GetMetadataUpdateEvent());
+ MetadataUpdateEvent metadataUpdateEvent =
getLatestMetadataUpdateEvent(context, 0);
+ KafkaStream latestReaderStream =
+ metadataUpdateEvent.getKafkaStreams().stream()
+ .findFirst()
+ .orElseThrow(
+ () -> new AssertionError("Missing metadata
update stream"));
+ assertThat(latestReaderStream.getClusterMetadataMap().keySet())
+ .as("restored readers should converge to latest cluster
set after discovery")
+ .containsExactlyInAnyOrder(cluster0, addedCluster);
+ assertThat(latestReaderStream.getClusterMetadataMap())
+ .doesNotContainKey(removedCluster);
+
+ DynamicKafkaSourceEnumState postRefreshState =
enumerator.snapshotState(-1);
+ assertThat(postRefreshState.getClusterEnumeratorStates().keySet())
+ .as("restored enumerator state should drop removed cluster
and add new cluster")
+ .containsExactlyInAnyOrder(cluster0, addedCluster);
+ assertThat(postRefreshState.getClusterEnumeratorStates())
+ .doesNotContainKey(removedCluster);
+ }
+ }
+
@Test
public void testAddSplitsBack() throws Throwable {
try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
@@ -1690,6 +2008,37 @@ public class DynamicKafkaSourceEnumeratorTest {
}
}
+ private MetadataUpdateEvent getLatestMetadataUpdateEvent(
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context, int
readerId)
+ throws Exception {
+ List<SourceEvent> sourceEvents =
context.getSentSourceEvent().get(readerId);
+ assertThat(sourceEvents)
+ .as("reader %s should have received source events", readerId)
+ .isNotNull();
+ return sourceEvents.stream()
+ .filter(MetadataUpdateEvent.class::isInstance)
+ .map(MetadataUpdateEvent.class::cast)
+ .reduce((first, second) -> second)
+ .orElseThrow(
+ () ->
+ new AssertionError(
+ String.format(
+ "reader %s did not receive
metadata update event",
+ readerId)));
+ }
+
+ private ClusterMetadata copyClusterMetadataWithOverrides(
+ ClusterMetadata baseClusterMetadata, Consumer<Properties>
overrideConsumer) {
+ Properties copiedProperties = new Properties();
+ copiedProperties.putAll(baseClusterMetadata.getProperties());
+ overrideConsumer.accept(copiedProperties);
+ return new ClusterMetadata(
+ baseClusterMetadata.getTopics(),
+ copiedProperties,
+ baseClusterMetadata.getStartingOffsetsInitializer(),
+ baseClusterMetadata.getStoppingOffsetsInitializer());
+ }
+
private static class TestKafkaEnumContextProxyFactory
implements
StoppableKafkaEnumContextProxy.StoppableKafkaEnumContextProxyFactory {
