[FLINK-37282] Add ProducerPool

Add first class producer pool that self-manages all resources and allows to recycle producers by transactional ids.

Author: AHeise

flink-connector-kafka/archunit-violations/86dfd459-67a9-4b26-9b5c-0b0bbf22681a

@@ -26,6 +26,12 @@ org.apache.flink.connector.kafka.sink.KafkaWriterITCase does not satisfy: only o
 * reside in a package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type InternalMiniClusterExtension and annotated with @RegisterExtension\
 * reside outside of package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type MiniClusterExtension and annotated with @RegisterExtension or are , and of type MiniClusterTestEnvironment and annotated with @TestEnv\
 * reside in a package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class InternalMiniClusterExtension\
+* reside outside of package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class MiniClusterExtension\
+ or contain any fields that are public, static, and of type MiniClusterWithClientResource and final and annotated with @ClassRule or contain any fields that is of type MiniClusterWithClientResource and public and final and not static and annotated with @Rule
+org.apache.flink.connector.kafka.sink.internal.ProducerPoolImplITCase does not satisfy: only one of the following predicates match:\
+* reside in a package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type InternalMiniClusterExtension and annotated with @RegisterExtension\
+* reside outside of package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type MiniClusterExtension and annotated with @RegisterExtension or are , and of type MiniClusterTestEnvironment and annotated with @TestEnv\
+* reside in a package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class InternalMiniClusterExtension\
 * reside outside of package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class MiniClusterExtension\
  or contain any fields that are public, static, and of type MiniClusterWithClientResource and final and annotated with @ClassRule or contain any fields that is of type MiniClusterWithClientResource and public and final and not static and annotated with @Rule
 org.apache.flink.connector.kafka.source.KafkaSourceITCase does not satisfy: only one of the following predicates match:\

flink-connector-kafka/archunit-violations/c0d94764-76a0-4c50-b617-70b1754c4612

@@ -50,6 +50,7 @@ Method <org.apache.flink.connector.kafka.sink.ExactlyOnceKafkaWriter.getProducer
 Method <org.apache.flink.connector.kafka.sink.KafkaSink.getKafkaProducerConfig()> is annotated with <org.apache.flink.annotation.VisibleForTesting> in (KafkaSink.java:0)
 Method <org.apache.flink.connector.kafka.sink.KafkaSinkBuilder.setRecordSerializer(org.apache.flink.connector.kafka.sink.KafkaRecordSerializationSchema)> calls method <org.apache.flink.api.java.ClosureCleaner.clean(java.lang.Object, org.apache.flink.api.common.ExecutionConfig$ClosureCleanerLevel, boolean)> in (KafkaSinkBuilder.java:152)
 Method <org.apache.flink.connector.kafka.sink.KafkaWriter.getCurrentProducer()> is annotated with <org.apache.flink.annotation.VisibleForTesting> in (KafkaWriter.java:0)
+Method <org.apache.flink.connector.kafka.sink.internal.ProducerPoolImpl.getProducers()> is annotated with <org.apache.flink.annotation.VisibleForTesting> in (ProducerPoolImpl.java:0)
 Method <org.apache.flink.connector.kafka.source.KafkaSource.createReader(org.apache.flink.api.connector.source.SourceReaderContext, java.util.function.Consumer)> is annotated with <org.apache.flink.annotation.VisibleForTesting> in (KafkaSource.java:0)
 Method <org.apache.flink.connector.kafka.source.KafkaSource.getConfiguration()> is annotated with <org.apache.flink.annotation.VisibleForTesting> in (KafkaSource.java:0)
 Method <org.apache.flink.connector.kafka.source.KafkaSource.getKafkaSubscriber()> is annotated with <org.apache.flink.annotation.VisibleForTesting> in (KafkaSource.java:0)

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

@@ -0,0 +1,78 @@
+/*
+ * 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.sink.internal;
+
+import org.apache.flink.annotation.Internal;
+
+import java.util.Objects;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * An immutable class that represents a transactional id and a checkpoint id. It's used inside the
+ * {@link ProducerPoolImpl} to keep track of the transactions that are currently in-flight. The
+ * checkpoint id is used to subsume committed transactions wrt to recycling producers.
+ */
+@Internal
+public class CheckpointTransaction {
+    private final String transactionalId;
+    private final long checkpointId;
+
+    public CheckpointTransaction(String transactionalId, long checkpointId) {
+        this.transactionalId = checkNotNull(transactionalId, "transactionalId must not be null");
+        this.checkpointId = checkpointId;
+    }
+
+    public long getCheckpointId() {
+        return checkpointId;
+    }
+
+    public String getTransactionalId() {
+        return transactionalId;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        CheckpointTransaction that = (CheckpointTransaction) o;
+        return checkpointId == that.checkpointId
+                && Objects.equals(transactionalId, that.transactionalId);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(transactionalId, checkpointId);
+    }
+
+    @Override
+    public String toString() {
+        return "CheckpointTransaction{"
+                + "transactionalId='"
+                + transactionalId
+                + '\''
+                + ", checkpointId="
+                + checkpointId
+                + '}';
+    }
+}

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

@@ -18,8 +18,6 @@
 
 package org.apache.flink.connector.kafka.sink.internal;
 
-import org.apache.flink.annotation.Internal;
-
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerConfig;
@@ -46,9 +44,7 @@
 /**
  * A {@link KafkaProducer} that exposes private fields to allow resume producing from a given state.
  */
-@Internal
 public class FlinkKafkaInternalProducer<K, V> extends KafkaProducer<K, V> {
-
     private static final Logger LOG = LoggerFactory.getLogger(FlinkKafkaInternalProducer.class);
     private static final String TRANSACTION_MANAGER_FIELD_NAME = "transactionManager";
     private static final String TRANSACTION_MANAGER_STATE_ENUM =
@@ -173,28 +169,51 @@ public void initTransactionId(String transactionalId) {
         }
     }
 
+    /**
+     * Sets the transaction manager state to uninitialized.
+     *
+     * <p>Can only be called if the producer is in a transaction. Its main purpose is to resolve the
+     * split brain scenario between writer and committer.
+     */
+    public void transactionCompletedExternally() {
+        checkState(inTransaction, "Not in transactional state");
+        this.inTransaction = false;
+        this.hasRecordsInTransaction = false;
+        Object transactionManager = getTransactionManager();
+        synchronized (transactionManager) {
+            setField(transactionManager, "transactionalId", transactionalId);
+            setField(
+                    transactionManager,
+                    "currentState",
+                    getTransactionManagerState("UNINITIALIZED"));
+        }
+    }
+
+    /**
+     * Sets the transactional id and sets the transaction manager state to uninitialized.
+     *
+     * <p>Can only be called if the producer is not in a transaction.
+     */
     public void setTransactionId(String transactionalId) {
-        if (!transactionalId.equals(this.transactionalId)) {
-            checkState(
-                    !inTransaction,
-                    String.format("Another transaction %s is still open.", transactionalId));
-            LOG.debug("Change transaction id from {} to {}", this.transactionalId, transactionalId);
-            Object transactionManager = getTransactionManager();
-            synchronized (transactionManager) {
-                setField(transactionManager, "transactionalId", transactionalId);
-                setField(
-                        transactionManager,
-                        "currentState",
-                        getTransactionManagerState("UNINITIALIZED"));
-                this.transactionalId = transactionalId;
-            }
+        checkState(
+                !inTransaction,
+                String.format("Another transaction %s is still open.", transactionalId));
+        LOG.debug("Change transaction id from {} to {}", this.transactionalId, transactionalId);
+        this.transactionalId = transactionalId;
+        Object transactionManager = getTransactionManager();
+        synchronized (transactionManager) {
+            setField(transactionManager, "transactionalId", transactionalId);
+            setField(
+                    transactionManager,
+                    "currentState",
+                    getTransactionManagerState("UNINITIALIZED"));
         }
     }
 
     /**
      * Besides committing {@link org.apache.kafka.clients.producer.KafkaProducer#commitTransaction}
      * is also adding new partitions to the transaction. flushNewPartitions method is moving this
-     * logic to pre-commit/flush, to make resumeTransaction simpler. Otherwise resumeTransaction
+     * logic to pre-commit/flush, to make resumeTransaction simpler. Otherwise, resumeTransaction
      * would require to restore state of the not yet added/"in-flight" partitions.
      */
     private void flushNewPartitions() {

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

@@ -0,0 +1,43 @@
+/*
+ * 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.sink.internal;
+
+/** A pool of producers that can be recycled. */
+public interface ProducerPool extends AutoCloseable {
+    /**
+     * Notify the pool that a transaction has finished. The producer with the given transactional id
+     * can be recycled.
+     */
+    void recycleByTransactionId(String transactionalId);
+
+    /**
+     * Get a producer for the given transactional id and checkpoint id. The producer is not recycled
+     * until it is passed to the committer, the committer commits the transaction, and {@link
+     * #recycleByTransactionId(String)} is called. Alternatively, the producer can be recycled by
+     * {@link #recycle(FlinkKafkaInternalProducer)}.
+     */
+    FlinkKafkaInternalProducer<byte[], byte[]> getTransactionalProducer(
+            String transactionalId, long checkpointId);
+
+    /**
+     * Explicitly recycle a producer. This is useful when the producer has not been passed to the
+     * committer.
+     */
+    void recycle(FlinkKafkaInternalProducer<byte[], byte[]> producer);
+}