[FLINK-34554] Introduce transaction strategies

Move existing way of naming and aborting transactions into specific strategies. Next commit will add new strategies.

Author: AHeise

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

@@ -26,8 +26,11 @@
 import org.apache.flink.connector.kafka.sink.internal.ProducerPool;
 import org.apache.flink.connector.kafka.sink.internal.ProducerPoolImpl;
 import org.apache.flink.connector.kafka.sink.internal.ReadableBackchannel;
+import org.apache.flink.connector.kafka.sink.internal.TransactionAbortStrategyContextImpl;
+import org.apache.flink.connector.kafka.sink.internal.TransactionAbortStrategyImpl;
 import org.apache.flink.connector.kafka.sink.internal.TransactionFinished;
-import org.apache.flink.connector.kafka.sink.internal.TransactionalIdFactory;
+import org.apache.flink.connector.kafka.sink.internal.TransactionNamingStrategyContextImpl;
+import org.apache.flink.connector.kafka.sink.internal.TransactionNamingStrategyImpl;
 import org.apache.flink.runtime.checkpoint.CheckpointIDCounter;
 import org.apache.flink.util.FlinkRuntimeException;
 
@@ -54,7 +57,17 @@
  */
 class ExactlyOnceKafkaWriter<IN> extends KafkaWriter<IN> {
     private static final Logger LOG = LoggerFactory.getLogger(ExactlyOnceKafkaWriter.class);
+    /**
+     * Prefix for the transactional id. Must be unique across all sinks writing to the same broker.
+     */
     private final String transactionalIdPrefix;
+    /**
+     * Strategy to abort lingering transactions from previous executions during writer
+     * initialization.
+     */
+    private final TransactionAbortStrategyImpl transactionAbortStrategy;
+    /** Strategy to name transactions. */
+    private final TransactionNamingStrategyImpl transactionNamingStrategy;
 
     private final KafkaWriterState kafkaWriterState;
     private final Collection<KafkaWriterState> recoveredStates;
@@ -72,6 +85,8 @@ class ExactlyOnceKafkaWriter<IN> extends KafkaWriter<IN> {
      * related methods.
      */
     private final ReadableBackchannel<TransactionFinished> backchannel;
+    /** The context used to name transactions. */
+    private final TransactionNamingStrategyContextImpl namingContext;
 
     /**
      * Constructor creating a kafka writer.
@@ -95,6 +110,8 @@ class ExactlyOnceKafkaWriter<IN> extends KafkaWriter<IN> {
             WriterInitContext sinkInitContext,
             KafkaRecordSerializationSchema<IN> recordSerializer,
             SerializationSchema.InitializationContext schemaContext,
+            TransactionAbortStrategyImpl transactionAbortStrategy,
+            TransactionNamingStrategyImpl transactionNamingStrategy,
             Collection<KafkaWriterState> recoveredStates) {
         super(
                 deliveryGuarantee,
@@ -104,6 +121,11 @@ class ExactlyOnceKafkaWriter<IN> extends KafkaWriter<IN> {
                 schemaContext);
         this.transactionalIdPrefix =
                 checkNotNull(transactionalIdPrefix, "transactionalIdPrefix must not be null");
+        this.transactionAbortStrategy =
+                checkNotNull(transactionAbortStrategy, "transactionAbortStrategy must not be null");
+        this.transactionNamingStrategy =
+                checkNotNull(
+                        transactionNamingStrategy, "transactionNamingStrategy must not be null");
 
         try {
             recordSerializer.open(schemaContext, kafkaSinkContext);
@@ -127,6 +149,9 @@ class ExactlyOnceKafkaWriter<IN> extends KafkaWriter<IN> {
                                 subtaskId,
                                 sinkInitContext.getTaskInfo().getAttemptNumber(),
                                 transactionalIdPrefix);
+        this.namingContext =
+                new TransactionNamingStrategyContextImpl(
+                        transactionalIdPrefix, subtaskId, restoredCheckpointId, producerPool);
     }
 
     @Override
@@ -147,13 +172,10 @@ public void initialize() {
     }
 
     private FlinkKafkaInternalProducer<byte[], byte[]> startTransaction(long checkpointId) {
+        namingContext.setNextCheckpointId(checkpointId);
         FlinkKafkaInternalProducer<byte[], byte[]> producer =
-                producerPool.getTransactionalProducer(
-                        TransactionalIdFactory.buildTransactionalId(
-                                transactionalIdPrefix,
-                                kafkaSinkContext.getParallelInstanceId(),
-                                checkpointId),
-                        checkpointId);
+                transactionNamingStrategy.getTransactionalProducer(namingContext);
+        namingContext.setLastCheckpointId(checkpointId);
         producer.beginTransaction();
         return producer;
     }
@@ -236,13 +258,34 @@ private void abortLingeringTransactions(
             }
         }
 
-        try (TransactionAborter transactionAborter =
-                new TransactionAborter(
-                        kafkaSinkContext.getParallelInstanceId(),
-                        kafkaSinkContext.getNumberOfParallelInstances(),
-                        id -> producerPool.getTransactionalProducer(id, startCheckpointId),
-                        producerPool::recycle)) {
-            transactionAborter.abortLingeringTransactions(prefixesToAbort, startCheckpointId);
-        }
+        LOG.info(
+                "Aborting lingering transactions with prefixes {} using {}",
+                prefixesToAbort,
+                transactionAbortStrategy);
+        TransactionAbortStrategyContextImpl context =
+                getTransactionAbortStrategyContext(startCheckpointId, prefixesToAbort);
+        transactionAbortStrategy.abortTransactions(context);
+    }
+
+    private TransactionAbortStrategyContextImpl getTransactionAbortStrategyContext(
+            long startCheckpointId, List<String> prefixesToAbort) {
+        TransactionAbortStrategyImpl.TransactionAborter aborter =
+                transactionalId -> {
+                    // getTransactionalProducer already calls initTransactions, which cancels the
+                    // transaction
+                    FlinkKafkaInternalProducer<byte[], byte[]> producer =
+                            producerPool.getTransactionalProducer(transactionalId, 0);
+                    LOG.debug("Aborting transaction {}", transactionalId);
+                    producer.flush();
+                    short epoch = producer.getEpoch();
+                    producerPool.recycle(producer);
+                    return epoch;
+                };
+        return new TransactionAbortStrategyContextImpl(
+                kafkaSinkContext.getParallelInstanceId(),
+                kafkaSinkContext.getNumberOfParallelInstances(),
+                prefixesToAbort,
+                startCheckpointId,
+                aborter);
     }
 }

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

@@ -82,16 +82,19 @@ public class KafkaSink<IN>
     private final KafkaRecordSerializationSchema<IN> recordSerializer;
     private final Properties kafkaProducerConfig;
     private final String transactionalIdPrefix;
+    private final TransactionNamingStrategy transactionNamingStrategy;
 
     KafkaSink(
             DeliveryGuarantee deliveryGuarantee,
             Properties kafkaProducerConfig,
             String transactionalIdPrefix,
-            KafkaRecordSerializationSchema<IN> recordSerializer) {
+            KafkaRecordSerializationSchema<IN> recordSerializer,
+            TransactionNamingStrategy transactionNamingStrategy) {
         this.deliveryGuarantee = deliveryGuarantee;
         this.kafkaProducerConfig = kafkaProducerConfig;
         this.transactionalIdPrefix = transactionalIdPrefix;
         this.recordSerializer = recordSerializer;
+        this.transactionNamingStrategy = transactionNamingStrategy;
     }
 
     /**
@@ -141,6 +144,8 @@ public KafkaWriter<IN> restoreWriter(
                             context,
                             recordSerializer,
                             context.asSerializationSchemaInitializationContext(),
+                            transactionNamingStrategy.getAbortImpl(),
+                            transactionNamingStrategy.getImpl(),
                             recoveredState);
         } else {
             writer =

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

@@ -74,6 +74,7 @@ public class KafkaSinkBuilder<IN> {
 
     private final Properties kafkaProducerConfig;
     private KafkaRecordSerializationSchema<IN> recordSerializer;
+    private TransactionNamingStrategy transactionNamingStrategy = TransactionNamingStrategy.DEFAULT;
 
     KafkaSinkBuilder() {
         kafkaProducerConfig = new Properties();
@@ -125,6 +126,20 @@ public KafkaSinkBuilder<IN> setProperty(String key, String value) {
         return this;
     }
 
+    /**
+     * Sets the {@link TransactionNamingStrategy} that is used to name the transactions.
+     *
+     * <p>By default {@link TransactionNamingStrategy#DEFAULT} is used. It's recommended to change
+     * the strategy only if specific issues occur.
+     */
+    public KafkaSinkBuilder<IN> setTransactionNamingStrategy(
+            TransactionNamingStrategy transactionNamingStrategy) {
+        this.transactionNamingStrategy =
+                checkNotNull(
+                        transactionNamingStrategy, "transactionNamingStrategy must not be null");
+        return this;
+    }
+
     /**
      * Sets the {@link KafkaRecordSerializationSchema} that transforms incoming records to {@link
      * org.apache.kafka.clients.producer.ProducerRecord}s.
@@ -162,7 +177,7 @@ public KafkaSinkBuilder<IN> setTransactionalIdPrefix(String transactionalIdPrefi
         checkState(
                 transactionalIdPrefix.getBytes(StandardCharsets.UTF_8).length
                         <= MAXIMUM_PREFIX_BYTES,
-                "The configured prefix is too long and the resulting transactionalId might exceed Kafka's transactionalIds size.");
+                "The configured prefix is too long and the resulting transactionalIdPrefix might exceed Kafka's transactionalIdPrefix size.");
         return this;
     }
 
@@ -180,12 +195,12 @@ private void sanityCheck() {
         checkNotNull(
                 kafkaProducerConfig.getProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),
                 "bootstrapServers");
+        checkNotNull(recordSerializer, "recordSerializer");
         if (deliveryGuarantee == DeliveryGuarantee.EXACTLY_ONCE) {
             checkState(
                     transactionalIdPrefix != null,
-                    "EXACTLY_ONCE delivery guarantee requires a transactionIdPrefix to be set to provide unique transaction names across multiple KafkaSinks writing to the same Kafka cluster.");
+                    "EXACTLY_ONCE delivery guarantee requires a transactionalIdPrefix to be set to provide unique transaction names across multiple KafkaSinks writing to the same Kafka cluster.");
         }
-        checkNotNull(recordSerializer, "recordSerializer");
     }
 
     /**
@@ -196,6 +211,10 @@ private void sanityCheck() {
     public KafkaSink<IN> build() {
         sanityCheck();
         return new KafkaSink<>(
-                deliveryGuarantee, kafkaProducerConfig, transactionalIdPrefix, recordSerializer);
+                deliveryGuarantee,
+                kafkaProducerConfig,
+                transactionalIdPrefix,
+                recordSerializer,
+                transactionNamingStrategy);
     }
 }