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 70fb6311 [FLINK-39944] Make dynamic Kafka metadata refresh resilient
(#274)
70fb6311 is described below
commit 70fb6311c1d5483cb8fc049909be5058b98cd296
Author: lnbest0707 <[email protected]>
AuthorDate: Wed Jun 24 11:48:42 2026 -0700
[FLINK-39944] Make dynamic Kafka metadata refresh resilient (#274)
This change isolates dynamic stream metadata discovery on a dedicated
stoppable worker and
marshals result handling back to the coordinator thread. Because metadata
discovery and
stale-cluster activity checks can now run from different workers, metadata
service calls are
serialized. Metadata updates are sent to readers before stale enumerator
close waits complete,
and shutdown closes the metadata service before waiting for discovery so
in-flight calls can
unblock. Asynchronous stale-enumerator close failures are retained and
surfaced from close.
---
.../enumerator/DynamicKafkaSourceEnumerator.java | 146 +++++++-
.../enumerator/StoppableKafkaEnumContextProxy.java | 23 +-
...ppableKafkaMetadataServiceDiscoveryContext.java | 149 ++++++++
.../SynchronizedKafkaMetadataService.java | 63 ++++
.../dynamic/source/DynamicKafkaSourceITTest.java | 19 +-
.../DynamicKafkaSourceEnumeratorTest.java | 383 +++++++++++++++++++++
...leKafkaMetadataServiceDiscoveryContextTest.java | 190 ++++++++++
.../SynchronizedKafkaMetadataServiceTest.java | 196 +++++++++++
8 files changed, 1148 insertions(+), 21 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 cc3a34f7..a31840db 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
@@ -62,6 +62,10 @@ import java.util.Map.Entry;
import java.util.Properties;
import java.util.Set;
import java.util.TreeMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
/**
@@ -90,6 +94,10 @@ public class DynamicKafkaSourceEnumerator
private final Boundedness boundedness;
private final
StoppableKafkaEnumContextProxy.StoppableKafkaEnumContextProxyFactory
stoppableKafkaEnumContextProxyFactory;
+ private final StoppableKafkaMetadataServiceDiscoveryContext
+ kafkaMetadataServiceDiscoveryContext;
+ private final ExecutorService enumeratorClosingExecutor;
+ private final AtomicReference<Throwable>
asynchronousEnumeratorCloseFailure;
// options
private final long kafkaMetadataServiceDiscoveryIntervalMs;
@@ -123,7 +131,9 @@ public class DynamicKafkaSourceEnumerator
boundedness,
dynamicKafkaSourceEnumState,
StoppableKafkaEnumContextProxy.StoppableKafkaEnumContextProxyFactory
- .getDefaultFactory());
+ .getDefaultFactory(),
+ StoppableKafkaMetadataServiceDiscoveryContext
+
.StoppableKafkaMetadataServiceDiscoveryContextFactory.getDefaultFactory());
}
@VisibleForTesting
@@ -138,6 +148,35 @@ public class DynamicKafkaSourceEnumerator
DynamicKafkaSourceEnumState dynamicKafkaSourceEnumState,
StoppableKafkaEnumContextProxy.StoppableKafkaEnumContextProxyFactory
stoppableKafkaEnumContextProxyFactory) {
+ this(
+ kafkaStreamSubscriber,
+ kafkaMetadataService,
+ enumContext,
+ startingOffsetsInitializer,
+ stoppingOffsetInitializer,
+ properties,
+ boundedness,
+ dynamicKafkaSourceEnumState,
+ stoppableKafkaEnumContextProxyFactory,
+ StoppableKafkaMetadataServiceDiscoveryContext
+ .StoppableKafkaMetadataServiceDiscoveryContextFactory
+ .getSplitEnumeratorContextFactory());
+ }
+
+ DynamicKafkaSourceEnumerator(
+ KafkaStreamSubscriber kafkaStreamSubscriber,
+ KafkaMetadataService kafkaMetadataService,
+ SplitEnumeratorContext<DynamicKafkaSourceSplit> enumContext,
+ OffsetsInitializer startingOffsetsInitializer,
+ OffsetsInitializer stoppingOffsetInitializer,
+ Properties properties,
+ Boundedness boundedness,
+ DynamicKafkaSourceEnumState dynamicKafkaSourceEnumState,
+
StoppableKafkaEnumContextProxy.StoppableKafkaEnumContextProxyFactory
+ stoppableKafkaEnumContextProxyFactory,
+ StoppableKafkaMetadataServiceDiscoveryContext
+
.StoppableKafkaMetadataServiceDiscoveryContextFactory
+ kafkaMetadataServiceDiscoveryContextFactory) {
this.kafkaStreamSubscriber = kafkaStreamSubscriber;
this.boundedness = boundedness;
@@ -162,8 +201,16 @@ public class DynamicKafkaSourceEnumerator
this.kafkaMetadataServiceDiscoveryFailureCount = 0;
this.firstDiscoveryComplete = false;
- this.kafkaMetadataService = kafkaMetadataService;
+ this.kafkaMetadataService = new
SynchronizedKafkaMetadataService(kafkaMetadataService);
this.stoppableKafkaEnumContextProxyFactory =
stoppableKafkaEnumContextProxyFactory;
+ this.kafkaMetadataServiceDiscoveryContext =
+
kafkaMetadataServiceDiscoveryContextFactory.create(enumContext);
+ this.enumeratorClosingExecutor =
+ Executors.newSingleThreadExecutor(
+ runnable ->
+ createDaemonThread(
+ runnable,
"dynamic-kafka-enumerator-closing-worker"));
+ this.asynchronousEnumeratorCloseFailure = new AtomicReference<>();
this.splitAssignmentStrategy =
createSplitAssignmentStrategy(properties);
if
(!dynamicKafkaSourceEnumState.getClusterEnumeratorStates().isEmpty()) {
@@ -313,12 +360,15 @@ public class DynamicKafkaSourceEnumerator
}
if (kafkaMetadataServiceDiscoveryIntervalMs <= 0) {
- enumContext.callAsync(
- () ->
kafkaStreamSubscriber.getSubscribedStreams(kafkaMetadataService),
- this::onHandleSubscribedStreamsFetch);
+ logger.info("Scheduling one-time dynamic Kafka metadata refresh");
+ kafkaMetadataServiceDiscoveryContext.callAsync(
+ this::fetchSubscribedKafkaStreams,
this::onHandleSubscribedStreamsFetch);
} else {
- enumContext.callAsync(
- () ->
kafkaStreamSubscriber.getSubscribedStreams(kafkaMetadataService),
+ logger.info(
+ "Scheduling dynamic Kafka metadata refresh every {} ms",
+ kafkaMetadataServiceDiscoveryIntervalMs);
+ kafkaMetadataServiceDiscoveryContext.callAsync(
+ this::fetchSubscribedKafkaStreams,
this::onHandleSubscribedStreamsFetch,
0,
kafkaMetadataServiceDiscoveryIntervalMs);
@@ -346,7 +396,24 @@ public class DynamicKafkaSourceEnumerator
// --------------- private methods for metadata discovery ---------------
+ private Set<KafkaStream> fetchSubscribedKafkaStreams() {
+ logger.debug("Fetching subscribed Kafka streams for metadata refresh");
+ Set<KafkaStream> fetchedKafkaStreams =
+
kafkaStreamSubscriber.getSubscribedStreams(kafkaMetadataService);
+ logger.debug(
+ "Fetched {} subscribed Kafka streams for metadata refresh",
+ fetchedKafkaStreams.size());
+ return fetchedKafkaStreams;
+ }
+
+ private static Thread createDaemonThread(Runnable runnable, String
threadName) {
+ Thread thread = new Thread(runnable, threadName);
+ thread.setDaemon(true);
+ return thread;
+ }
+
private void onHandleSubscribedStreamsFetch(Set<KafkaStream>
fetchedKafkaStreams, Throwable t) {
+ logger.debug("Handling subscribed Kafka streams fetched by metadata
refresh");
firstDiscoveryComplete = true;
Set<KafkaStream> handledFetchKafkaStreams =
handleFetchSubscribedStreamsError(fetchedKafkaStreams, t);
@@ -397,13 +464,13 @@ public class DynamicKafkaSourceEnumerator
throw new RuntimeException("unable to snapshot state in metadata
change", e);
}
- logger.info("Closing enumerators due to metadata change");
-
- closeAllEnumeratorsAndContexts();
latestClusterTopicsMap = newClustersTopicsMap;
latestKafkaStreams = handledFetchKafkaStreams;
sendMetadataUpdateEventToAvailableReaders();
+ logger.info("Closing enumerators due to metadata change");
+
+ closeAllEnumeratorsAndContexts();
retainRemovedClusterEnumeratorStates(
dynamicKafkaSourceEnumState.getClusterEnumeratorStates(),
latestClusterTopicsMap.keySet());
@@ -587,17 +654,50 @@ public class DynamicKafkaSourceEnumerator
}
private void closeAllEnumeratorsAndContexts() {
- clusterEnumeratorMap.forEach(
+ Map<String, StoppableKafkaEnumContextProxy>
closingClusterEnumContextMap =
+ new HashMap<>(clusterEnumContextMap);
+ Map<String, SplitEnumerator<KafkaPartitionSplit, KafkaSourceEnumState>>
+ closingClusterEnumeratorMap = new
HashMap<>(clusterEnumeratorMap);
+ closingClusterEnumContextMap
+ .values()
+ .forEach(StoppableKafkaEnumContextProxy::prepareForClose);
+ clusterEnumContextMap.clear();
+ clusterEnumeratorMap.clear();
+
+ enumeratorClosingExecutor.execute(
+ () ->
+ closeEnumeratorsAndContexts(
+ closingClusterEnumContextMap,
closingClusterEnumeratorMap));
+ }
+
+ private void closeEnumeratorsAndContexts(
+ Map<String, StoppableKafkaEnumContextProxy>
closingClusterEnumContextMap,
+ Map<String, SplitEnumerator<KafkaPartitionSplit,
KafkaSourceEnumState>>
+ closingClusterEnumeratorMap) {
+ closingClusterEnumeratorMap.forEach(
(cluster, subEnumerator) -> {
try {
- clusterEnumContextMap.get(cluster).close();
+ closingClusterEnumContextMap.get(cluster).close();
subEnumerator.close();
} catch (Exception e) {
- throw new RuntimeException(e);
+ handleAsynchronousEnumeratorCloseFailure(e);
}
});
- clusterEnumContextMap.clear();
- clusterEnumeratorMap.clear();
+ }
+
+ private void handleAsynchronousEnumeratorCloseFailure(Exception e) {
+ asynchronousEnumeratorCloseFailure.compareAndSet(null, e);
+ try {
+ enumContext.runInCoordinatorThread(
+ () -> {
+ throw new RuntimeException(e);
+ });
+ } catch (Throwable coordinatorFailure) {
+ logger.warn(
+ "Unable to propagate asynchronous dynamic Kafka enumerator
close failure to "
+ + "the coordinator thread. The failure will be
rethrown during close.",
+ coordinatorFailure);
+ }
}
/**
@@ -754,6 +854,13 @@ public class DynamicKafkaSourceEnumerator
@Override
public void close() throws IOException {
try {
+ kafkaMetadataServiceDiscoveryContext.prepareForClose();
+
clusterEnumContextMap.values().forEach(StoppableKafkaEnumContextProxy::prepareForClose);
+
+ // Metadata service close may unblock an in-flight metadata
discovery call.
+ kafkaMetadataService.close();
+ kafkaMetadataServiceDiscoveryContext.close();
+
// close contexts first since they may have running tasks
for (StoppableKafkaEnumContextProxy subEnumContext :
clusterEnumContextMap.values()) {
subEnumContext.close();
@@ -764,7 +871,14 @@ public class DynamicKafkaSourceEnumerator
clusterEnumerator.getValue().close();
}
- kafkaMetadataService.close();
+ enumeratorClosingExecutor.shutdown();
+ enumeratorClosingExecutor.awaitTermination(Long.MAX_VALUE,
TimeUnit.MILLISECONDS);
+
+ Throwable asynchronousCloseFailure =
asynchronousEnumeratorCloseFailure.get();
+ if (asynchronousCloseFailure != null) {
+ throw new RuntimeException(
+ "Failed to close stale dynamic Kafka enumerator",
asynchronousCloseFailure);
+ }
} catch (Exception e) {
throw new RuntimeException(e);
}
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 752a5d6b..89843837 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
@@ -211,12 +211,19 @@ public class StoppableKafkaEnumContextProxy
@Override
public void close() throws Exception {
logger.info("Closing enum context for {}", kafkaClusterId);
+ prepareForClose();
if (subEnumeratorWorker != null) {
- // KafkaSubscriber worker thread will fail if admin client is
closed in the middle.
- // Swallow the error and set the context to closed state.
+ subEnumeratorWorker.awaitTermination(Long.MAX_VALUE,
TimeUnit.MILLISECONDS);
+ }
+ }
+
+ void prepareForClose() {
+ if (subEnumeratorWorker != null) {
+ // KafkaSubscriber worker thread may fail if KafkaSourceEnumerator
closes its Admin
+ // client while discovery is in flight. Mark the context closing
before stale
+ // enumerators are closed so those failures are swallowed and
callbacks are skipped.
isClosing = true;
subEnumeratorWorker.shutdown();
- subEnumeratorWorker.awaitTermination(Long.MAX_VALUE,
TimeUnit.MILLISECONDS);
}
}
@@ -226,6 +233,11 @@ public class StoppableKafkaEnumContextProxy
*/
protected <T> Callable<T> wrapCallAsyncCallable(Callable<T> callable) {
return () -> {
+ if (isClosing) {
+ throw new HandledFlinkKafkaException(
+ new KafkaException("Enumerator context is closing"),
kafkaClusterId);
+ }
+
try {
return callable.call();
} catch (Exception e) {
@@ -254,6 +266,11 @@ public class StoppableKafkaEnumContextProxy
protected <T> BiConsumer<T, Throwable> wrapCallAsyncCallableHandler(
BiConsumer<T, Throwable> mainHandler) {
return (result, t) -> {
+ if (isClosing) {
+ logger.debug("Skipping callback for closing enum context {}",
kafkaClusterId);
+ return;
+ }
+
// check if exception is handled
Optional<HandledFlinkKafkaException> throwable =
ExceptionUtils.findThrowable(t,
HandledFlinkKafkaException.class);
diff --git
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaMetadataServiceDiscoveryContext.java
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaMetadataServiceDiscoveryContext.java
new file mode 100644
index 00000000..4410862d
--- /dev/null
+++
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaMetadataServiceDiscoveryContext.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.dynamic.source.enumerator;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+import
org.apache.flink.connector.kafka.dynamic.source.split.DynamicKafkaSourceSplit;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
+
+/**
+ * Runs dynamic Kafka stream metadata discovery outside Flink's shared source
coordinator worker.
+ *
+ * <p>The source coordinator worker is also used by per-cluster Kafka
partition discovery. Removed
+ * clusters can block that shared worker in Kafka Admin calls, so stream
metadata discovery needs an
+ * isolated worker to keep cluster removal reconciliation responsive.
+ */
+@Internal
+public class StoppableKafkaMetadataServiceDiscoveryContext implements
AutoCloseable {
+ private final SplitEnumeratorContext<DynamicKafkaSourceSplit> enumContext;
+ private final ScheduledExecutorService metadataDiscoveryWorker;
+ private volatile boolean isClosing;
+
+ public StoppableKafkaMetadataServiceDiscoveryContext(
+ SplitEnumeratorContext<DynamicKafkaSourceSplit> enumContext) {
+ this.enumContext = enumContext;
+ this.metadataDiscoveryWorker =
+ Executors.newScheduledThreadPool(
+ 1,
+ runnable ->
+ createDaemonThread(
+ runnable,
"dynamic-kafka-metadata-discovery-worker"));
+ this.isClosing = false;
+ }
+
+ public <T> void callAsync(Callable<T> callable, BiConsumer<T, Throwable>
handler) {
+ metadataDiscoveryWorker.execute(() -> invokeCallable(callable,
handler));
+ }
+
+ public <T> void callAsync(
+ Callable<T> callable,
+ BiConsumer<T, Throwable> handler,
+ long initialDelay,
+ long period) {
+ metadataDiscoveryWorker.scheduleAtFixedRate(
+ () -> invokeCallable(callable, handler),
+ initialDelay,
+ period,
+ TimeUnit.MILLISECONDS);
+ }
+
+ private <T> void invokeCallable(Callable<T> callable, BiConsumer<T,
Throwable> handler) {
+ T result = null;
+ Throwable throwable = null;
+ try {
+ result = callable.call();
+ } catch (Throwable t) {
+ throwable = t;
+ }
+
+ T fetchedResult = result;
+ Throwable fetchError = throwable;
+ runInCoordinatorThreadIfOpen(handler, fetchedResult, fetchError);
+ }
+
+ private synchronized <T> void runInCoordinatorThreadIfOpen(
+ BiConsumer<T, Throwable> handler, T result, Throwable t) {
+ if (!isClosing) {
+ enumContext.runInCoordinatorThread(() -> handleResult(handler,
result, t));
+ }
+ }
+
+ private <T> void handleResult(BiConsumer<T, Throwable> handler, T result,
Throwable t) {
+ if (!isClosing) {
+ handler.accept(result, t);
+ }
+ }
+
+ private static Thread createDaemonThread(Runnable runnable, String
threadName) {
+ Thread thread = new Thread(runnable, threadName);
+ thread.setDaemon(true);
+ return thread;
+ }
+
+ void prepareForClose() {
+ synchronized (this) {
+ isClosing = true;
+ metadataDiscoveryWorker.shutdownNow();
+ }
+ }
+
+ @Override
+ public void close() throws InterruptedException {
+ prepareForClose();
+ metadataDiscoveryWorker.awaitTermination(Long.MAX_VALUE,
TimeUnit.MILLISECONDS);
+ }
+
+ /** Factory for metadata discovery contexts. */
+ @Internal
+ public interface StoppableKafkaMetadataServiceDiscoveryContextFactory {
+ StoppableKafkaMetadataServiceDiscoveryContext create(
+ SplitEnumeratorContext<DynamicKafkaSourceSplit> enumContext);
+
+ static StoppableKafkaMetadataServiceDiscoveryContextFactory
getDefaultFactory() {
+ return StoppableKafkaMetadataServiceDiscoveryContext::new;
+ }
+
+ static StoppableKafkaMetadataServiceDiscoveryContextFactory
+ getSplitEnumeratorContextFactory() {
+ return enumContext ->
+ new
StoppableKafkaMetadataServiceDiscoveryContext(enumContext) {
+ @Override
+ public <T> void callAsync(
+ Callable<T> callable, BiConsumer<T, Throwable>
handler) {
+ enumContext.callAsync(callable, handler);
+ }
+
+ @Override
+ public <T> void callAsync(
+ Callable<T> callable,
+ BiConsumer<T, Throwable> handler,
+ long initialDelay,
+ long period) {
+ enumContext.callAsync(callable, handler,
initialDelay, period);
+ }
+ };
+ }
+ }
+}
diff --git
a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/SynchronizedKafkaMetadataService.java
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/SynchronizedKafkaMetadataService.java
new file mode 100644
index 00000000..2cf09cf0
--- /dev/null
+++
b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/SynchronizedKafkaMetadataService.java
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.kafka.dynamic.source.enumerator;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.connector.kafka.dynamic.metadata.KafkaMetadataService;
+import org.apache.flink.connector.kafka.dynamic.metadata.KafkaStream;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Serializes Kafka metadata service calls across dynamic Kafka source worker
threads.
+ *
+ * <p>Dynamic stream metadata refresh and stale-cluster error handling run on
separate workers.
+ * Metadata service implementations predate that split and are not required to
be thread-safe.
+ */
+@Internal
+final class SynchronizedKafkaMetadataService implements KafkaMetadataService {
+ private final KafkaMetadataService delegate;
+
+ SynchronizedKafkaMetadataService(KafkaMetadataService delegate) {
+ this.delegate = delegate;
+ }
+
+ @Override
+ public synchronized Set<KafkaStream> getAllStreams() {
+ return delegate.getAllStreams();
+ }
+
+ @Override
+ public synchronized Map<String, KafkaStream>
describeStreams(Collection<String> streamIds) {
+ return delegate.describeStreams(streamIds);
+ }
+
+ @Override
+ public synchronized boolean isClusterActive(String kafkaClusterId) {
+ return delegate.isClusterActive(kafkaClusterId);
+ }
+
+ @Override
+ public void close() throws Exception {
+ // Closing must be able to unblock an in-flight metadata call during
source shutdown.
+ delegate.close();
+ }
+}
diff --git
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/DynamicKafkaSourceITTest.java
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/DynamicKafkaSourceITTest.java
index a28e4412..3bb77e7c 100644
---
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/DynamicKafkaSourceITTest.java
+++
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/DynamicKafkaSourceITTest.java
@@ -277,7 +277,7 @@ class DynamicKafkaSourceITTest {
for (int readerId = 0; readerId < numSubtasks; readerId++) {
registerReader(context, enumerator, readerId);
}
- runAllOneTimeCallables(context);
+ waitForInitialSplitAssignments(context);
verifyAllSplitsAssignedOnce(
context.getSplitsAssignmentSequence(),
metadataService.getAllStreams());
@@ -571,7 +571,7 @@ class DynamicKafkaSourceITTest {
enumerator.start();
registerReader(context, enumerator, 0);
registerReader(context, enumerator, 1);
- runAllOneTimeCallables(context);
+ waitForInitialSplitAssignments(context);
List<DynamicKafkaSourceSplit> assignedSplits =
context.getSplitsAssignmentSequence().stream()
@@ -1415,6 +1415,21 @@ class DynamicKafkaSourceITTest {
}
}
+ private void waitForInitialSplitAssignments(
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context)
throws Exception {
+ CommonTestUtils.waitUtil(
+ () -> {
+ try {
+ runAllOneTimeCallables(context);
+ } catch (Throwable t) {
+ throw new RuntimeException(t);
+ }
+ return
!context.getSplitsAssignmentSequence().isEmpty();
+ },
+ Duration.ofSeconds(10),
+ "Initial dynamic Kafka split assignment did not complete");
+ }
+
private void verifyAllSplitsAssignedOnce(
List<SplitsAssignment<DynamicKafkaSourceSplit>> assignments,
Set<KafkaStream> kafkaStreams) {
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 4683037d..25f748de 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
@@ -39,6 +39,7 @@ import
org.apache.flink.connector.kafka.source.enumerator.initializer.NoStopping
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.core.testutils.CommonTestUtils;
import org.apache.flink.mock.Whitebox;
import
org.apache.flink.runtime.checkpoint.RoundRobinOperatorStateRepartitioner;
import org.apache.flink.runtime.state.OperatorStateHandle;
@@ -56,6 +57,7 @@ import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
+import java.time.Duration;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
@@ -68,6 +70,11 @@ import java.util.Objects;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
@@ -1128,6 +1135,185 @@ public class DynamicKafkaSourceEnumeratorTest {
}
}
+ @Test
+ public void
testMetadataRefreshSendsReaderMetadataBeforeClosingStaleEnumerators()
+ throws Throwable {
+ KafkaStream initialStream =
DynamicKafkaSourceTestHelper.getKafkaStream(TOPIC);
+ KafkaStream shrunkStream =
DynamicKafkaSourceTestHelper.getKafkaStream(TOPIC);
+ String removedCluster =
DynamicKafkaSourceTestHelper.getKafkaClusterId(1);
+ shrunkStream.getClusterMetadataMap().remove(removedCluster);
+
+ MockKafkaMetadataService metadataService =
+ new
MockKafkaMetadataService(Collections.singleton(initialStream));
+ BlockingCloseKafkaEnumContextProxyFactory enumContextProxyFactory =
+ new BlockingCloseKafkaEnumContextProxyFactory();
+
+ Properties properties = new Properties();
+
properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
"0");
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "1");
+ ExecutorService refreshExecutor = Executors.newSingleThreadExecutor();
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
+ DynamicKafkaSourceEnumerator enumerator =
+ new DynamicKafkaSourceEnumerator(
+ new
KafkaStreamSetSubscriber(Collections.singleton(TOPIC)),
+ metadataService,
+ context,
+ OffsetsInitializer.committedOffsets(),
+ new NoStoppingOffsetsInitializer(),
+ properties,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ new DynamicKafkaSourceEnumState(),
+ enumContextProxyFactory)) {
+ enumerator.start();
+ context.runPeriodicCallable(0);
+ runAllOneTimeCallables(context);
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
0);
+
+
metadataService.setKafkaStreams(Collections.singleton(shrunkStream));
+ Future<?> metadataRefresh =
+ refreshExecutor.submit(
+ () -> {
+ try {
+ context.runPeriodicCallable(0);
+ } catch (Throwable t) {
+ throw new RuntimeException(t);
+ }
+ });
+
+ assertThat(enumContextProxyFactory.awaitCloseStarted())
+ .as("metadata refresh should start closing stale
enumerators")
+ .isTrue();
+
+ MetadataUpdateEvent metadataUpdateEvent =
+ getLatestMetadataUpdateEventWithoutContextSync(context, 0);
+ KafkaStream latestReaderStream =
+ metadataUpdateEvent.getKafkaStreams().stream()
+ .findFirst()
+ .orElseThrow(
+ () -> new AssertionError("Missing metadata
update stream"));
+ assertThat(latestReaderStream.getClusterMetadataMap())
+ .as("metadata event sent to reader should remove stale
clusters before close")
+ .doesNotContainKey(removedCluster);
+
+ try {
+ assertThatCode(() -> metadataRefresh.get(10, TimeUnit.SECONDS))
+ .as("metadata refresh should not wait for stale
enumerator close")
+ .doesNotThrowAnyException();
+ } finally {
+ enumContextProxyFactory.allowClose();
+ }
+ } finally {
+ refreshExecutor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testCloseSurfacesAsynchronousStaleEnumeratorCloseFailure()
throws Throwable {
+ KafkaStream initialStream =
DynamicKafkaSourceTestHelper.getKafkaStream(TOPIC);
+ KafkaStream shrunkStream =
DynamicKafkaSourceTestHelper.getKafkaStream(TOPIC);
+ String removedCluster =
DynamicKafkaSourceTestHelper.getKafkaClusterId(1);
+ shrunkStream.getClusterMetadataMap().remove(removedCluster);
+
+ MockKafkaMetadataService metadataService =
+ new
MockKafkaMetadataService(Collections.singleton(initialStream));
+ ThrowingCloseKafkaEnumContextProxyFactory enumContextProxyFactory =
+ new ThrowingCloseKafkaEnumContextProxyFactory(removedCluster);
+
+ Properties properties = new Properties();
+
properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
"0");
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "1");
+
+ try (DroppingCoordinatorThreadContext context =
+ new DroppingCoordinatorThreadContext(NUM_SUBTASKS)) {
+ DynamicKafkaSourceEnumerator enumerator =
+ new DynamicKafkaSourceEnumerator(
+ new
KafkaStreamSetSubscriber(Collections.singleton(TOPIC)),
+ metadataService,
+ context,
+ OffsetsInitializer.committedOffsets(),
+ new NoStoppingOffsetsInitializer(),
+ properties,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ new DynamicKafkaSourceEnumState(),
+ enumContextProxyFactory);
+
+ enumerator.start();
+ context.runPeriodicCallable(0);
+ runAllOneTimeCallables(context);
+
+
metadataService.setKafkaStreams(Collections.singleton(shrunkStream));
+ context.runPeriodicCallable(0);
+ runAllOneTimeCallables(context);
+
+ assertThatThrownBy(enumerator::close)
+ .hasMessageContaining("Failed to close stale dynamic Kafka
enumerator")
+ .hasRootCauseMessage("test close failure");
+ }
+ }
+
+ @Test
+ public void
testProductionMetadataRefreshBypassesBlockedSourceCoordinatorAsyncCallable()
+ throws Throwable {
+ KafkaStream kafkaStream =
DynamicKafkaSourceTestHelper.getKafkaStream(TOPIC);
+ BlockingDescribeStreamsKafkaMetadataService metadataService =
+ new
BlockingDescribeStreamsKafkaMetadataService(Collections.singleton(kafkaStream));
+ CountDownLatch sourceCoordinatorCallableStarted = new
CountDownLatch(1);
+ CountDownLatch allowSourceCoordinatorCallableToFinish = new
CountDownLatch(1);
+ ExecutorService sourceCoordinatorWorker =
Executors.newSingleThreadExecutor();
+
+ try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
+ new MockSplitEnumeratorContext<>(NUM_SUBTASKS);
+ DynamicKafkaSourceEnumerator enumerator =
+ createProductionEnumerator(context, metadataService)) {
+ context.callAsync(
+ () -> {
+ sourceCoordinatorCallableStarted.countDown();
+
awaitUninterruptibly(allowSourceCoordinatorCallableToFinish);
+ return null;
+ },
+ (result, t) -> {});
+ Future<?> blockedSourceCoordinatorCallable =
+ sourceCoordinatorWorker.submit(
+ () -> {
+ try {
+ context.runNextOneTimeCallable();
+ } catch (Throwable t) {
+ throw new RuntimeException(t);
+ }
+ });
+ assertThat(sourceCoordinatorCallableStarted.await(10,
TimeUnit.SECONDS))
+ .as("source coordinator async callable should start")
+ .isTrue();
+
+ enumerator.start();
+ assertThat(metadataService.awaitDescribeStreamsStarted())
+ .as("production metadata discovery worker should fetch
streams")
+ .isTrue();
+ mockRegisterReaderAndSendReaderStartupEvent(context, enumerator,
0);
+
+ metadataService.allowDescribeStreams();
+ CommonTestUtils.waitUtil(
+ () -> hasLatestMetadataUpdateEvent(context, 0,
kafkaStream),
+ Duration.ofSeconds(10),
+ "Metadata refresh did not bypass blocked source
coordinator async callable");
+ assertThat(blockedSourceCoordinatorCallable.isDone())
+ .as("source coordinator async callable should still be
blocked")
+ .isFalse();
+
+ allowSourceCoordinatorCallableToFinish.countDown();
+ assertThatCode(() -> blockedSourceCoordinatorCallable.get(10,
TimeUnit.SECONDS))
+ .as("source coordinator async callable should finish after
release")
+ .doesNotThrowAnyException();
+ } finally {
+ metadataService.allowDescribeStreams();
+ allowSourceCoordinatorCallableToFinish.countDown();
+ sourceCoordinatorWorker.shutdownNow();
+ }
+ }
+
@Test
public void testAddSplitsBack() throws Throwable {
try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
@@ -1853,6 +2039,24 @@ public class DynamicKafkaSourceEnumeratorTest {
new TestKafkaEnumContextProxyFactory());
}
+ private DynamicKafkaSourceEnumerator createProductionEnumerator(
+ SplitEnumeratorContext<DynamicKafkaSourceSplit> context,
+ KafkaMetadataService kafkaMetadataService) {
+ Properties properties = new Properties();
+
properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(),
"0");
+ properties.setProperty(
+
DynamicKafkaSourceOptions.STREAM_METADATA_DISCOVERY_INTERVAL_MS.key(), "0");
+ return new DynamicKafkaSourceEnumerator(
+ new KafkaStreamSetSubscriber(Collections.singleton(TOPIC)),
+ kafkaMetadataService,
+ context,
+ OffsetsInitializer.earliest(),
+ new NoStoppingOffsetsInitializer(),
+ properties,
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ new DynamicKafkaSourceEnumState());
+ }
+
private void mockRegisterReaderAndSendReaderStartupEvent(
MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context,
DynamicKafkaSourceEnumerator enumerator,
@@ -2120,6 +2324,21 @@ public class DynamicKafkaSourceEnumeratorTest {
}
}
+ private static void awaitUninterruptibly(CountDownLatch latch) {
+ boolean interrupted = false;
+ while (true) {
+ try {
+ latch.await();
+ break;
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+
private DynamicKafkaSourceEnumState getCheckpointState(KafkaStream
kafkaStream)
throws Throwable {
try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context =
@@ -2169,6 +2388,28 @@ public class DynamicKafkaSourceEnumeratorTest {
MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context, int
readerId)
throws Exception {
List<SourceEvent> sourceEvents =
context.getSentSourceEvent().get(readerId);
+ assertThat(sourceEvents)
+ .as("source events should have been sent to reader %s",
readerId)
+ .isNotNull();
+ return sourceEvents.stream()
+ .filter(MetadataUpdateEvent.class::isInstance)
+ .map(MetadataUpdateEvent.class::cast)
+ .reduce((first, second) -> second)
+ .orElseThrow(
+ () ->
+ new AssertionError(
+ String.format(
+ "metadata update event was not
sent to reader %s",
+ readerId)));
+ }
+
+ @SuppressWarnings("unchecked")
+ private MetadataUpdateEvent getLatestMetadataUpdateEventWithoutContextSync(
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context, int
readerId) {
+ Map<Integer, List<SourceEvent>> sentSourceEvents =
+ (Map<Integer, List<SourceEvent>>)
+ Whitebox.getInternalState(context, "sentSourceEvent");
+ List<SourceEvent> sourceEvents = sentSourceEvents.get(readerId);
assertThat(sourceEvents)
.as("reader %s should have received source events", readerId)
.isNotNull();
@@ -2184,6 +2425,19 @@ public class DynamicKafkaSourceEnumeratorTest {
readerId)));
}
+ private boolean hasLatestMetadataUpdateEvent(
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context,
+ int readerId,
+ KafkaStream expectedKafkaStream) {
+ try {
+ return getLatestMetadataUpdateEventWithoutContextSync(context,
readerId)
+ .getKafkaStreams()
+ .equals(Collections.singleton(expectedKafkaStream));
+ } catch (AssertionError e) {
+ return false;
+ }
+ }
+
private ClusterMetadata copyClusterMetadataWithOverrides(
ClusterMetadata baseClusterMetadata, Consumer<Properties>
overrideConsumer) {
Properties copiedProperties = new Properties();
@@ -2242,4 +2496,133 @@ public class DynamicKafkaSourceEnumeratorTest {
period);
}
}
+
+ private static class BlockingCloseKafkaEnumContextProxyFactory
+ implements
StoppableKafkaEnumContextProxy.StoppableKafkaEnumContextProxyFactory {
+ private final CountDownLatch closeStarted = new CountDownLatch(1);
+ private final CountDownLatch allowClose = new CountDownLatch(1);
+
+ @Override
+ public StoppableKafkaEnumContextProxy create(
+ SplitEnumeratorContext<DynamicKafkaSourceSplit> enumContext,
+ String kafkaClusterId,
+ KafkaMetadataService kafkaMetadataService,
+ Runnable signalNoMoreSplitsCallback) {
+ return new BlockingCloseKafkaEnumContextProxy(
+ kafkaClusterId,
+ kafkaMetadataService,
+ (MockSplitEnumeratorContext<DynamicKafkaSourceSplit>)
enumContext,
+ closeStarted,
+ allowClose);
+ }
+
+ private boolean awaitCloseStarted() throws InterruptedException {
+ return closeStarted.await(10, TimeUnit.SECONDS);
+ }
+
+ private void allowClose() {
+ allowClose.countDown();
+ }
+ }
+
+ private static class BlockingCloseKafkaEnumContextProxy extends
TestKafkaEnumContextProxy {
+ private final CountDownLatch closeStarted;
+ private final CountDownLatch allowClose;
+
+ public BlockingCloseKafkaEnumContextProxy(
+ String kafkaClusterId,
+ KafkaMetadataService kafkaMetadataService,
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit>
enumContext,
+ CountDownLatch closeStarted,
+ CountDownLatch allowClose) {
+ super(kafkaClusterId, kafkaMetadataService, enumContext);
+ this.closeStarted = closeStarted;
+ this.allowClose = allowClose;
+ }
+
+ @Override
+ public void close() throws Exception {
+ closeStarted.countDown();
+ assertThat(allowClose.await(10, TimeUnit.SECONDS))
+ .as("test should allow stale enumerator close to complete")
+ .isTrue();
+ super.close();
+ }
+ }
+
+ private static class ThrowingCloseKafkaEnumContextProxyFactory
+ implements
StoppableKafkaEnumContextProxy.StoppableKafkaEnumContextProxyFactory {
+ private final String failingClusterId;
+
+ private ThrowingCloseKafkaEnumContextProxyFactory(String
failingClusterId) {
+ this.failingClusterId = failingClusterId;
+ }
+
+ @Override
+ public StoppableKafkaEnumContextProxy create(
+ SplitEnumeratorContext<DynamicKafkaSourceSplit> enumContext,
+ String kafkaClusterId,
+ KafkaMetadataService kafkaMetadataService,
+ Runnable signalNoMoreSplitsCallback) {
+ if (failingClusterId.equals(kafkaClusterId)) {
+ return new ThrowingCloseKafkaEnumContextProxy(
+ kafkaClusterId,
+ kafkaMetadataService,
+ (MockSplitEnumeratorContext<DynamicKafkaSourceSplit>)
enumContext);
+ }
+ return new TestKafkaEnumContextProxy(
+ kafkaClusterId,
+ kafkaMetadataService,
+ (MockSplitEnumeratorContext<DynamicKafkaSourceSplit>)
enumContext);
+ }
+ }
+
+ private static class ThrowingCloseKafkaEnumContextProxy extends
TestKafkaEnumContextProxy {
+ private ThrowingCloseKafkaEnumContextProxy(
+ String kafkaClusterId,
+ KafkaMetadataService kafkaMetadataService,
+ MockSplitEnumeratorContext<DynamicKafkaSourceSplit>
enumContext) {
+ super(kafkaClusterId, kafkaMetadataService, enumContext);
+ }
+
+ @Override
+ public void close() throws Exception {
+ throw new Exception("test close failure");
+ }
+ }
+
+ private static class DroppingCoordinatorThreadContext
+ extends MockSplitEnumeratorContext<DynamicKafkaSourceSplit> {
+ private DroppingCoordinatorThreadContext(int parallelism) {
+ super(parallelism);
+ }
+
+ @Override
+ public void runInCoordinatorThread(Runnable runnable) {}
+ }
+
+ private static class BlockingDescribeStreamsKafkaMetadataService
+ extends MockKafkaMetadataService {
+ private final CountDownLatch describeStreamsStarted = new
CountDownLatch(1);
+ private final CountDownLatch allowDescribeStreams = new
CountDownLatch(1);
+
+ private BlockingDescribeStreamsKafkaMetadataService(Set<KafkaStream>
kafkaStreams) {
+ super(kafkaStreams);
+ }
+
+ @Override
+ public Map<String, KafkaStream> describeStreams(Collection<String>
streamIds) {
+ describeStreamsStarted.countDown();
+ awaitUninterruptibly(allowDescribeStreams);
+ return super.describeStreams(streamIds);
+ }
+
+ private boolean awaitDescribeStreamsStarted() throws
InterruptedException {
+ return describeStreamsStarted.await(10, TimeUnit.SECONDS);
+ }
+
+ private void allowDescribeStreams() {
+ allowDescribeStreams.countDown();
+ }
+ }
}
diff --git
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaMetadataServiceDiscoveryContextTest.java
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaMetadataServiceDiscoveryContextTest.java
new file mode 100644
index 00000000..45ab57e9
--- /dev/null
+++
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/StoppableKafkaMetadataServiceDiscoveryContextTest.java
@@ -0,0 +1,190 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.kafka.dynamic.source.enumerator;
+
+import org.apache.flink.api.connector.source.ReaderInfo;
+import org.apache.flink.api.connector.source.SourceEvent;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+import org.apache.flink.api.connector.source.SplitsAssignment;
+import
org.apache.flink.connector.kafka.dynamic.source.split.DynamicKafkaSourceSplit;
+import org.apache.flink.core.testutils.CommonTestUtils;
+import org.apache.flink.metrics.groups.SplitEnumeratorMetricGroup;
+
+import org.junit.jupiter.api.Test;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** A test for {@link StoppableKafkaMetadataServiceDiscoveryContext}. */
+public class StoppableKafkaMetadataServiceDiscoveryContextTest {
+
+ @Test
+ public void testMetadataDiscoveryDoesNotUseSourceCoordinatorWorker()
throws Exception {
+ TrackingSplitEnumeratorContext enumContext = new
TrackingSplitEnumeratorContext();
+
+ AtomicReference<String> callbackResult = new AtomicReference<>();
+ try (StoppableKafkaMetadataServiceDiscoveryContext context =
+ new
StoppableKafkaMetadataServiceDiscoveryContext(enumContext)) {
+ context.callAsync(
+ () -> "metadata-refresh-complete",
+ (result, t) -> {
+ assertThat(t).isNull();
+ callbackResult.set(result);
+ });
+
+ CommonTestUtils.waitUtil(
+ () ->
"metadata-refresh-complete".equals(callbackResult.get()),
+ Duration.ofSeconds(5),
+ "Metadata refresh callback did not complete");
+ assertThat(callbackResult).hasValue("metadata-refresh-complete");
+ assertThat(enumContext.coordinatorThreadCallCount()).isEqualTo(1);
+ assertThat(enumContext.sourceCoordinatorAsyncCallCount()).isZero();
+ }
+ }
+
+ @Test
+ public void testCloseWaitsForMetadataDiscoveryWorker() throws Exception {
+ TrackingSplitEnumeratorContext enumContext = new
TrackingSplitEnumeratorContext();
+ CountDownLatch callableStarted = new CountDownLatch(1);
+ CountDownLatch allowCallableToFinish = new CountDownLatch(1);
+ ExecutorService closeExecutor = Executors.newSingleThreadExecutor();
+
+ try (StoppableKafkaMetadataServiceDiscoveryContext context =
+ new
StoppableKafkaMetadataServiceDiscoveryContext(enumContext)) {
+ context.callAsync(
+ () -> {
+ callableStarted.countDown();
+ awaitUninterruptibly(allowCallableToFinish);
+ return "metadata-refresh-complete";
+ },
+ (result, t) -> {});
+
+ assertThat(callableStarted.await(10, TimeUnit.SECONDS))
+ .as("metadata discovery callable should start")
+ .isTrue();
+
+ Future<?> closeFuture =
+ closeExecutor.submit(
+ () -> {
+ context.close();
+ return null;
+ });
+ assertThatThrownBy(() -> closeFuture.get(100,
TimeUnit.MILLISECONDS))
+ .as("close should wait for the metadata discovery worker")
+ .isInstanceOf(TimeoutException.class);
+
+ allowCallableToFinish.countDown();
+ assertThatCode(() -> closeFuture.get(10, TimeUnit.SECONDS))
+ .as("close should finish after the metadata discovery
worker exits")
+ .doesNotThrowAnyException();
+ assertThat(enumContext.coordinatorThreadCallCount()).isZero();
+ } finally {
+ allowCallableToFinish.countDown();
+ closeExecutor.shutdownNow();
+ }
+ }
+
+ private static void awaitUninterruptibly(CountDownLatch latch) {
+ boolean interrupted = false;
+ while (true) {
+ try {
+ latch.await();
+ break;
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+
+ private static class TrackingSplitEnumeratorContext
+ implements SplitEnumeratorContext<DynamicKafkaSourceSplit> {
+ private final AtomicInteger sourceCoordinatorAsyncCallCount = new
AtomicInteger();
+ private final AtomicInteger coordinatorThreadCallCount = new
AtomicInteger();
+
+ @Override
+ public SplitEnumeratorMetricGroup metricGroup() {
+ return null;
+ }
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, SourceEvent event)
{}
+
+ @Override
+ public int currentParallelism() {
+ return 1;
+ }
+
+ @Override
+ public Map<Integer, ReaderInfo> registeredReaders() {
+ return Collections.emptyMap();
+ }
+
+ @Override
+ public void assignSplits(SplitsAssignment<DynamicKafkaSourceSplit>
newSplitAssignments) {}
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {}
+
+ @Override
+ public <T> void callAsync(Callable<T> callable, BiConsumer<T,
Throwable> handler) {
+ sourceCoordinatorAsyncCallCount.incrementAndGet();
+ }
+
+ @Override
+ public <T> void callAsync(
+ Callable<T> callable,
+ BiConsumer<T, Throwable> handler,
+ long initialDelay,
+ long period) {
+ sourceCoordinatorAsyncCallCount.incrementAndGet();
+ }
+
+ @Override
+ public void runInCoordinatorThread(Runnable runnable) {
+ coordinatorThreadCallCount.incrementAndGet();
+ runnable.run();
+ }
+
+ private int sourceCoordinatorAsyncCallCount() {
+ return sourceCoordinatorAsyncCallCount.get();
+ }
+
+ private int coordinatorThreadCallCount() {
+ return coordinatorThreadCallCount.get();
+ }
+ }
+}
diff --git
a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/SynchronizedKafkaMetadataServiceTest.java
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/SynchronizedKafkaMetadataServiceTest.java
new file mode 100644
index 00000000..7896967e
--- /dev/null
+++
b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/SynchronizedKafkaMetadataServiceTest.java
@@ -0,0 +1,196 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.kafka.dynamic.source.enumerator;
+
+import org.apache.flink.connector.kafka.dynamic.metadata.KafkaMetadataService;
+import org.apache.flink.connector.kafka.dynamic.metadata.KafkaStream;
+import
org.apache.flink.connector.kafka.dynamic.source.enumerator.subscriber.KafkaStreamSetSubscriber;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+
+/** A test for {@link SynchronizedKafkaMetadataService}. */
+public class SynchronizedKafkaMetadataServiceTest {
+
+ @Test
+ public void testSubscribedStreamsAndClusterActivityChecksAreSerialized()
throws Exception {
+ BlockingKafkaMetadataService delegate = new
BlockingKafkaMetadataService();
+ KafkaMetadataService metadataService = new
SynchronizedKafkaMetadataService(delegate);
+ KafkaStreamSetSubscriber kafkaStreamSubscriber =
+ new KafkaStreamSetSubscriber(Collections.singleton("stream"));
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ Future<Set<KafkaStream>> subscribedStreams =
+ executor.submit(
+ () ->
kafkaStreamSubscriber.getSubscribedStreams(metadataService));
+ assertThat(delegate.awaitDescribeStreamsStarted())
+ .as("subscribed stream metadata fetch should start")
+ .isTrue();
+
+ Future<Boolean> isClusterActive =
+ executor.submit(() ->
metadataService.isClusterActive("cluster"));
+ assertThat(delegate.awaitIsClusterActiveStarted(100,
TimeUnit.MILLISECONDS))
+ .as("cluster activity check should wait for subscribed
stream fetch")
+ .isFalse();
+
+ delegate.allowDescribeStreams();
+ assertThat(subscribedStreams.get(10, TimeUnit.SECONDS)).isEmpty();
+ assertThat(isClusterActive.get(10, TimeUnit.SECONDS)).isTrue();
+ assertThat(delegate.awaitIsClusterActiveStarted())
+ .as("cluster activity check should run after subscribed
stream fetch")
+ .isTrue();
+ assertThat(delegate.maxConcurrentCalls()).isEqualTo(1);
+ } finally {
+ delegate.allowDescribeStreams();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testCloseCanUnblockInFlightMetadataCall() throws Exception {
+ BlockingKafkaMetadataService delegate = new
BlockingKafkaMetadataService();
+ KafkaMetadataService metadataService = new
SynchronizedKafkaMetadataService(delegate);
+ KafkaStreamSetSubscriber kafkaStreamSubscriber =
+ new KafkaStreamSetSubscriber(Collections.singleton("stream"));
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ Future<Set<KafkaStream>> subscribedStreams =
+ executor.submit(
+ () ->
kafkaStreamSubscriber.getSubscribedStreams(metadataService));
+ assertThat(delegate.awaitDescribeStreamsStarted())
+ .as("subscribed stream metadata fetch should start")
+ .isTrue();
+
+ Future<?> close =
+ executor.submit(
+ () -> {
+ metadataService.close();
+ return null;
+ });
+ assertThatCode(() -> close.get(10, TimeUnit.SECONDS))
+ .as("metadata service close should unblock the in-flight
fetch")
+ .doesNotThrowAnyException();
+ assertThat(subscribedStreams.get(10, TimeUnit.SECONDS)).isEmpty();
+ assertThat(delegate.awaitCloseCalled())
+ .as("delegate close should run while metadata fetch is in
flight")
+ .isTrue();
+ } finally {
+ delegate.allowDescribeStreams();
+ executor.shutdownNow();
+ }
+ }
+
+ private static class BlockingKafkaMetadataService implements
KafkaMetadataService {
+ private final CountDownLatch describeStreamsStarted = new
CountDownLatch(1);
+ private final CountDownLatch allowDescribeStreams = new
CountDownLatch(1);
+ private final CountDownLatch isClusterActiveStarted = new
CountDownLatch(1);
+ private final CountDownLatch closeCalled = new CountDownLatch(1);
+ private final AtomicInteger concurrentCalls = new AtomicInteger();
+ private final AtomicInteger maxConcurrentCalls = new AtomicInteger();
+
+ @Override
+ public Set<KafkaStream> getAllStreams() {
+ return runMetadataCall(Collections::emptySet);
+ }
+
+ @Override
+ public Map<String, KafkaStream> describeStreams(Collection<String>
streamIds) {
+ return runMetadataCall(
+ () -> {
+ describeStreamsStarted.countDown();
+ assertThat(allowDescribeStreams.await(10,
TimeUnit.SECONDS))
+ .as("test should release subscribed stream
metadata fetch")
+ .isTrue();
+ return Collections.emptyMap();
+ });
+ }
+
+ @Override
+ public boolean isClusterActive(String kafkaClusterId) {
+ return runMetadataCall(
+ () -> {
+ isClusterActiveStarted.countDown();
+ return true;
+ });
+ }
+
+ private <T> T runMetadataCall(CheckedSupplier<T> callable) {
+ int activeCalls = concurrentCalls.incrementAndGet();
+ maxConcurrentCalls.updateAndGet(currentMax -> Math.max(currentMax,
activeCalls));
+ try {
+ return callable.get();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ } finally {
+ concurrentCalls.decrementAndGet();
+ }
+ }
+
+ private boolean awaitDescribeStreamsStarted() throws
InterruptedException {
+ return describeStreamsStarted.await(10, TimeUnit.SECONDS);
+ }
+
+ private boolean awaitIsClusterActiveStarted() throws
InterruptedException {
+ return awaitIsClusterActiveStarted(10, TimeUnit.SECONDS);
+ }
+
+ private boolean awaitIsClusterActiveStarted(long timeout, TimeUnit
timeUnit)
+ throws InterruptedException {
+ return isClusterActiveStarted.await(timeout, timeUnit);
+ }
+
+ private void allowDescribeStreams() {
+ allowDescribeStreams.countDown();
+ }
+
+ private int maxConcurrentCalls() {
+ return maxConcurrentCalls.get();
+ }
+
+ @Override
+ public void close() {
+ closeCalled.countDown();
+ allowDescribeStreams();
+ }
+
+ private boolean awaitCloseCalled() throws InterruptedException {
+ return closeCalled.await(10, TimeUnit.SECONDS);
+ }
+ }
+
+ @FunctionalInterface
+ private interface CheckedSupplier<T> {
+ T get() throws Exception;
+ }
+}
