[FLINK-40912][Connectors/Kafka] Close the exactly-once writer's AdminClient

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ExactlyOnceKafkaWriter.java

@@ -48,6 +48,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.annotation.Nullable;
+
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -106,7 +108,7 @@ class ExactlyOnceKafkaWriter<IN> extends KafkaWriter<IN> {
     private final int totalNumberOfOwnedSubtasks;
     private final int[] ownedSubtaskIds;
 
-    /** Lazily created admin client for {@link TransactionAbortStrategyImpl}. */
+    /** Lazily created admin client for listing transactions and resolving topic patterns. */
     private AdminClient adminClient;
 
     /**
@@ -273,7 +275,8 @@ public void close() throws Exception {
                 this::abortCurrentProducer,
                 () -> closeAll(producerPool),
                 backchannel,
-                super::close);
+                super::close,
+                adminClient);
     }
 
     private void abortCurrentProducer() {
@@ -297,6 +300,11 @@ ProducerPool getProducerPool() {
         return producerPool;
     }
 
+    @Nullable
+    Admin getAdminClient() {
+        return adminClient;
+    }
+
     @VisibleForTesting
     public String getTransactionalIdPrefix() {
         return transactionalIdPrefix;
@@ -361,7 +369,7 @@ private TransactionAbortStrategyContextImpl getTransactionAbortStrategyContext(
                 prefixesToAbort,
                 startCheckpointId,
                 aborter,
-                this::getAdminClient,
+                this::getOrCreateAdminClient,
                 precommittedTransactionalIds);
     }
 
@@ -375,7 +383,8 @@ private Collection<String> getTopicNames() {
         if (identifier.getTopics() != null) {
             return identifier.getTopics();
         }
-        return AdminUtils.getTopicsByPattern(getAdminClient(), identifier.getTopicPattern());
+        return AdminUtils.getTopicsByPattern(
+                getOrCreateAdminClient(), identifier.getTopicPattern());
     }
 
     private Optional<KafkaDatasetIdentifier> getDatasetIdentifier() {
@@ -388,7 +397,7 @@ private Optional<KafkaDatasetIdentifier> getDatasetIdentifier() {
         return Optional.empty();
     }
 
-    private Admin getAdminClient() {
+    private Admin getOrCreateAdminClient() {
         if (adminClient == null) {
             adminClient = AdminClient.create(kafkaProducerConfig);
         }

flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/ExactlyOnceKafkaWriterITCase.java

@@ -33,13 +33,15 @@
 import org.apache.flink.test.junit5.MiniClusterExtension;
 
 import com.google.common.collect.Iterables;
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.errors.ProducerFencedException;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.ValueSource;
 
 import java.io.IOException;
@@ -48,11 +50,13 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 import static org.apache.flink.connector.kafka.sink.internal.TransactionalIdFactory.buildTransactionalId;
 import static org.apache.flink.connector.kafka.testutils.KafkaUtil.drainAllRecordsFromTopic;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatCode;
 
 /** Tests for the standalone KafkaWriter. */
@@ -374,6 +378,46 @@ void testAbortOnClose() throws Exception {
         }
     }
 
+    @Test
+    void testPoolingClosesAdminClient() throws Exception {
+        final ExactlyOnceKafkaWriter<Integer> writer =
+                createExactlyOnceWriter(TransactionNamingStrategy.POOLING);
+        // The outer try provides fallback cleanup so a failing test does not leak the client.
+        // The closure assertion runs before that fallback cleanup.
+        try (final Admin adminClient = writer.getAdminClient()) {
+            try (writer) {
+                assertThat(adminClient)
+                        .as("POOLING lists transactions through an admin client")
+                        .isNotNull();
+                assertThat(adminClient.listTopics().names().get(10, TimeUnit.SECONDS))
+                        .contains(topic);
+            }
+
+            assertThatThrownBy(() -> adminClient.listTopics().names().get(10, TimeUnit.SECONDS))
+                    .as("The writer must close its admin client")
+                    .hasRootCauseExactlyInstanceOf(IllegalStateException.class)
+                    .hasRootCauseMessage("Cannot accept new calls when AdminClient is closing.");
+        }
+    }
+
+    @ParameterizedTest
+    @EnumSource(
+            value = TransactionNamingStrategy.class,
+            names = "POOLING",
+            mode = EnumSource.Mode.EXCLUDE)
+    void testNonPoolingStrategiesNeverCreateAdminClient(TransactionNamingStrategy namingStrategy)
+            throws Exception {
+        final ExactlyOnceKafkaWriter<Integer> writer = createExactlyOnceWriter(namingStrategy);
+        try (writer) {
+            assertThat(writer.getAdminClient())
+                    .as("Naming strategy %s should not create an admin client", namingStrategy)
+                    .isNull();
+        }
+        assertThat(writer.getAdminClient())
+                .as("Closing the writer must not create an admin client")
+                .isNull();
+    }
+
     /** Test that producers are reused when committed. */
     @Test
     void shouldSkipIdsOfCommitterForPooledTransactions() throws Exception {
@@ -423,6 +467,15 @@ void shouldSkipIdsOfCommitterForPooledTransactions() throws Exception {
         }
     }
 
+    private ExactlyOnceKafkaWriter<Integer> createExactlyOnceWriter(
+            TransactionNamingStrategy namingStrategy) throws IOException {
+        return createWriter(
+                builder ->
+                        builder.setDeliveryGuarantee(DeliveryGuarantee.EXACTLY_ONCE)
+                                .setTransactionNamingStrategy(namingStrategy),
+                createInitContext());
+    }
+
     private static Collection<FlinkKafkaInternalProducer<byte[], byte[]>> getProducers(
             ExactlyOnceKafkaWriter<Integer> writer) {
         return ((ProducerPoolImpl) writer.getProducerPool()).getProducers();