[feature][transaction] Purge old aborted txn (#1879)

gaoran10 · eolivelli · web-flow · commit 1dda87903fb2 · 2023-07-13T18:10:31.000+08:00
This PR needs producer state manager recovery. 

### Motivation

Purge old aborted txns.

### Modifications

Add a timer task to purge useless old aborted data in the cache.

Co-authored-by: Enrico Olivelli &lt;enrico.olivelli@datastax.com&gt;
diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaProtocolHandler.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaProtocolHandler.java
@@ -47,6 +47,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Function;
@@ -117,6 +118,8 @@ public class KafkaProtocolHandler implements ProtocolHandler, TenantContextManag
     private MigrationManager migrationManager;
     private ReplicaManager replicaManager;
 
+    private ScheduledFuture<?> txUpdatedPurgeAbortedTxOffsetsTimeHandle;
+
     private final Map<String, GroupCoordinator> groupCoordinatorsByTenant = new ConcurrentHashMap<>();
     private final Map<String, TransactionCoordinator> transactionCoordinatorByTenant = new ConcurrentHashMap<>();
 
@@ -309,6 +312,16 @@ private void invalidatePartitionLog(TopicName topicName) {
         schemaRegistryManager = new SchemaRegistryManager(kafkaConfig, brokerService.getPulsar(),
                 brokerService.getAuthenticationService());
         migrationManager = new MigrationManager(kafkaConfig, brokerService.getPulsar());
+
+        if (kafkaConfig.isKafkaTransactionCoordinatorEnabled()
+                && kafkaConfig.getKafkaTxnPurgeAbortedTxnIntervalSeconds() > 0) {
+            txUpdatedPurgeAbortedTxOffsetsTimeHandle = service.getPulsar().getExecutor().scheduleWithFixedDelay(() -> {
+                        getReplicaManager().updatePurgeAbortedTxnsOffsets();
+                    },
+                    kafkaConfig.getKafkaTxnPurgeAbortedTxnIntervalSeconds(),
+                    kafkaConfig.getKafkaTxnPurgeAbortedTxnIntervalSeconds(),
+                    TimeUnit.SECONDS);
+        }
     }
 
     private TransactionCoordinator createAndBootTransactionCoordinator(String tenant) {
@@ -522,6 +535,10 @@ public Map<InetSocketAddress, ChannelInitializer<SocketChannel>> newChannelIniti
 
     @Override
     public void close() {
+        if (txUpdatedPurgeAbortedTxOffsetsTimeHandle != null) {
+            txUpdatedPurgeAbortedTxOffsetsTimeHandle.cancel(false);
+        }
+
         if (producePurgatory != null) {
             producePurgatory.shutdown();
         }
diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaServiceConfiguration.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaServiceConfiguration.java
@@ -451,6 +451,12 @@ public class KafkaServiceConfiguration extends ServiceConfiguration {
     )
     private int kafkaTransactionRecoveryNumThreads = 8;
 
+    @FieldContext(
+            category = CATEGORY_KOP_TRANSACTION,
+            doc = "Interval for purging aborted transactions from memory (requires reads from storage)"
+    )
+    private int kafkaTxnPurgeAbortedTxnIntervalSeconds = 60 * 60;
+
     @FieldContext(
             category = CATEGORY_KOP_TRANSACTION,
             doc = "The interval in milliseconds at which to rollback transactions that have timed out."
diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/scala/Either.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/scala/Either.java
@@ -16,6 +16,7 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import lombok.Getter;
+import lombok.ToString;
 
 /**
  * A simple Java migration of <a href="https://www.scala-lang.org/api/2.13.6/scala/util/Either.html">Scala Either</a>.
@@ -45,6 +46,7 @@
  * @param <W> the type of the 2nd possible value (the right side)
  */
 @Getter
+@ToString
 public class Either<V, W> {
 
     private final V left;
diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/storage/PartitionLog.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/storage/PartitionLog.java
@@ -203,9 +203,12 @@ private CompletableFuture<Void> loadTopicProperties() {
                     this.entryFormatter = buildEntryFormatter(topicProperties);
                     this.kafkaTopicUUID = properties.get("kafkaTopicUUID");
                     this.producerStateManager =
-                            new ProducerStateManager(fullPartitionName, kafkaTopicUUID,
+                            new ProducerStateManager(
+                                    fullPartitionName,
+                                    kafkaTopicUUID,
                                     producerStateManagerSnapshotBuffer,
-                                    kafkaConfig.getKafkaTxnProducerStateTopicSnapshotIntervalSeconds());
+                                    kafkaConfig.getKafkaTxnProducerStateTopicSnapshotIntervalSeconds(),
+                                    kafkaConfig.getKafkaTxnPurgeAbortedTxnIntervalSeconds());
                 });
     }
 
@@ -1064,6 +1067,25 @@ private MemoryRecords trimInvalidBytes(MemoryRecords records, LogAppendInfo info
         }
     }
 
+    /**
+     * Remove all the AbortedTxn that are no more referred by existing data on the topic.
+     * @return
+     */
+    public CompletableFuture<?> updatePurgeAbortedTxnsOffset() {
+        if (!kafkaConfig.isKafkaTransactionCoordinatorEnabled()) {
+            // no need to scan the topic, because transactions are disabled
+            return CompletableFuture.completedFuture(null);
+        }
+        if (!producerStateManager.hasSomeAbortedTransactions()) {
+            // nothing to do
+            return CompletableFuture.completedFuture(null);
+        }
+        return fetchOldestAvailableIndexFromTopic()
+                .thenAccept(offset ->
+                    producerStateManager.updateAbortedTxnsPurgeOffset(offset));
+
+    }
+
     public CompletableFuture<Long> fetchOldestAvailableIndexFromTopic() {
         final CompletableFuture<Long> future = new CompletableFuture<>();
 
@@ -1130,6 +1152,19 @@ public CompletableFuture<?> takeProducerSnapshot() {
         });
     }
 
+    public CompletableFuture<Long> forcePurgeAbortTx() {
+        return initFuture.thenCompose((___)  -> {
+            // purge can be taken only on the same thread that is used for writes
+            ManagedLedgerImpl ml = (ManagedLedgerImpl) getPersistentTopic().getManagedLedger();
+            ExecutorService executorService = ml.getScheduledExecutor().chooseThread(ml.getName());
+
+            return updatePurgeAbortedTxnsOffset()
+                    .thenApplyAsync((____) -> {
+                        return getProducerStateManager().executePurgeAbortedTx();
+                    }, executorService);
+        });
+    }
+
     public CompletableFuture<Long> recoverTxEntries(
             long offset,
             Executor executor) {
@@ -1206,7 +1241,6 @@ public CompletableFuture<Long> recoverTxEntries(
         }));
     }
 
-
     private void readNextEntriesForRecovery(ManagedCursor cursor, AtomicLong cursorOffset,
                                             KafkaTopicConsumerManager tcm,
                                             AtomicLong entryCounter,
@@ -1310,6 +1344,7 @@ private void updateProducerStateManager(long lastOffset, AnalyzeResult analyzeRe
 
         // do system clean up stuff in this thread
         producerStateManager.maybeTakeSnapshot(recoveryExecutor);
+        producerStateManager.maybePurgeAbortedTx();
     }
 
     private void decodeEntriesForRecovery(final CompletableFuture<DecodeResult> future,
diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/storage/PartitionLogManager.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/storage/PartitionLogManager.java
@@ -107,17 +107,15 @@ public int size() {
         return logMap.size();
     }
 
-    public CompletableFuture<Void> takeProducerStateSnapshots() {
-        List<CompletableFuture<Void>> handles = new ArrayList<>();
+    public CompletableFuture<?> updatePurgeAbortedTxnsOffsets() {
+        List<CompletableFuture<?>> handles = new ArrayList<>();
         logMap.values().forEach(log -> {
             if (log.isInitialised()) {
-                handles.add(log
-                        .getProducerStateManager()
-                        .takeSnapshot(recoveryExecutor)
-                        .thenApply(___ -> null));
+                handles.add(log.updatePurgeAbortedTxnsOffset());
             }
         });
-        return FutureUtil.waitForAll(handles);
+        return FutureUtil
+                .waitForAll(handles);
     }
 }
 
diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/storage/ProducerStateManager.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/storage/ProducerStateManager.java
@@ -13,6 +13,7 @@
  */
 package io.streamnative.pulsar.handlers.kop.storage;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Maps;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -22,6 +23,7 @@
 import java.util.TreeMap;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicLong;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.bookkeeper.util.SafeRunnable;
@@ -47,21 +49,27 @@ public class ProducerStateManager {
     private final ProducerStateManagerSnapshotBuffer producerStateManagerSnapshotBuffer;
 
     private final int kafkaTxnProducerStateTopicSnapshotIntervalSeconds;
+    private final int kafkaTxnPurgeAbortedTxnIntervalSeconds;
 
     private volatile long mapEndOffset = -1;
 
     private long lastSnapshotTime;
+    private long lastPurgeAbortedTxnTime;
 
+    private volatile long abortedTxnsPurgeOffset = -1;
 
     public ProducerStateManager(String topicPartition,
                                 String kafkaTopicUUID,
                                 ProducerStateManagerSnapshotBuffer producerStateManagerSnapshotBuffer,
-                                int kafkaTxnProducerStateTopicSnapshotIntervalSeconds) {
+                                int kafkaTxnProducerStateTopicSnapshotIntervalSeconds,
+                                int kafkaTxnPurgeAbortedTxnIntervalSeconds) {
         this.topicPartition = topicPartition;
         this.kafkaTopicUUID = kafkaTopicUUID;
         this.producerStateManagerSnapshotBuffer = producerStateManagerSnapshotBuffer;
         this.kafkaTxnProducerStateTopicSnapshotIntervalSeconds = kafkaTxnProducerStateTopicSnapshotIntervalSeconds;
+        this.kafkaTxnPurgeAbortedTxnIntervalSeconds = kafkaTxnPurgeAbortedTxnIntervalSeconds;
         this.lastSnapshotTime = System.currentTimeMillis();
+        this.lastPurgeAbortedTxnTime = System.currentTimeMillis();
     }
 
     public CompletableFuture<Void> recover(PartitionLog partitionLog, Executor executor) {
@@ -159,6 +167,34 @@ void maybeTakeSnapshot(Executor executor) {
         takeSnapshot(executor);
     }
 
+    void updateAbortedTxnsPurgeOffset(long abortedTxnsPurgeOffset) {
+        if (log.isDebugEnabled()) {
+            log.debug("{} updateAbortedTxnsPurgeOffset {}", topicPartition, abortedTxnsPurgeOffset);
+        }
+        if (abortedTxnsPurgeOffset < 0) {
+            return;
+        }
+        this.abortedTxnsPurgeOffset = abortedTxnsPurgeOffset;
+    }
+
+    long maybePurgeAbortedTx() {
+        if (mapEndOffset == -1) {
+            return 0;
+        }
+        long now = System.currentTimeMillis();
+        long deltaFromLast = (now - lastPurgeAbortedTxnTime) / 1000;
+        if (deltaFromLast / 1000 <= kafkaTxnPurgeAbortedTxnIntervalSeconds) {
+            return 0;
+        }
+        lastPurgeAbortedTxnTime = now;
+        return executePurgeAbortedTx();
+    }
+
+    @VisibleForTesting
+    long executePurgeAbortedTx() {
+        return purgeAbortedTxns(abortedTxnsPurgeOffset);
+    }
+
     private ProducerStateManagerSnapshot getProducerStateManagerSnapshot() {
         ProducerStateManagerSnapshot snapshot;
         synchronized (abortedIndexList) {
@@ -270,15 +306,39 @@ public void completeTxn(CompletedTxn completedTxn) {
         }
     }
 
+    public boolean hasSomeAbortedTransactions() {
+        return !abortedIndexList.isEmpty();
+    }
+
+    public long purgeAbortedTxns(long offset) {
+        AtomicLong count = new AtomicLong();
+        synchronized (abortedIndexList) {
+            abortedIndexList.removeIf(tx -> {
+                boolean toRemove = tx.lastOffset() < offset;
+                if (toRemove) {
+                    log.info("Transaction {} can be removed (lastOffset {} < {})", tx, tx.lastOffset(), offset);
+                    count.incrementAndGet();
+                }
+                return toRemove;
+            });
+            if (!abortedIndexList.isEmpty()) {
+                log.info("There are still {} aborted tx on {}", abortedIndexList.size(), topicPartition);
+            }
+        }
+        return count.get();
+    }
+
     public List<FetchResponse.AbortedTransaction> getAbortedIndexList(long fetchOffset) {
-        List<FetchResponse.AbortedTransaction> abortedTransactions = new ArrayList<>();
-        for (AbortedTxn abortedTxn : abortedIndexList) {
-            if (abortedTxn.lastOffset() >= fetchOffset) {
-                abortedTransactions.add(
-                        new FetchResponse.AbortedTransaction(abortedTxn.producerId(), abortedTxn.firstOffset()));
+        synchronized (abortedIndexList) {
+            List<FetchResponse.AbortedTransaction> abortedTransactions = new ArrayList<>();
+            for (AbortedTxn abortedTxn : abortedIndexList) {
+                if (abortedTxn.lastOffset() >= fetchOffset) {
+                    abortedTransactions.add(
+                            new FetchResponse.AbortedTransaction(abortedTxn.producerId(), abortedTxn.firstOffset()));
+                }
             }
+            return abortedTransactions;
         }
-        return abortedTransactions;
     }
 
     public void handleMissingDataBeforeRecovery(long minOffset, long snapshotOffset) {
diff --git a/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/storage/ReplicaManager.java b/kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/storage/ReplicaManager.java
@@ -337,4 +337,8 @@ public void tryCompleteDelayedFetch(DelayedOperationKey key) {
         }
     }
 
+    public CompletableFuture<?> updatePurgeAbortedTxnsOffsets() {
+        return logManager.updatePurgeAbortedTxnsOffsets();
+    }
+
 }
diff --git a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionTest.java b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionTest.java
diff --git a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionWithOAuthBearerAuthTest.java b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionWithOAuthBearerAuthTest.java
diff --git a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/storage/ProducerStateManagerTest.java b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/storage/ProducerStateManagerTest.java

Original file line number	Diff line number	Diff line change
`@@ -107,17 +107,15 @@ public int size() {`
`107`	`107`	`return logMap.size();`
`108`	`108`	`}`
`109`	`109`
`110`		`- public CompletableFuture<Void> takeProducerStateSnapshots() {`
`111`		`- List<CompletableFuture<Void>> handles = new ArrayList<>();`
	`110`	`+ public CompletableFuture<?> updatePurgeAbortedTxnsOffsets() {`
	`111`	`+ List<CompletableFuture<?>> handles = new ArrayList<>();`
`112`	`112`	`logMap.values().forEach(log -> {`
`113`	`113`	`if (log.isInitialised()) {`
`114`		`- handles.add(log`
`115`		`- .getProducerStateManager()`
`116`		`- .takeSnapshot(recoveryExecutor)`
`117`		`- .thenApply(___ -> null));`
	`114`	`+ handles.add(log.updatePurgeAbortedTxnsOffset());`
`118`	`115`	`}`
`119`	`116`	`});`
`120`		`- return FutureUtil.waitForAll(handles);`
	`117`	`+ return FutureUtil`
	`118`	`+ .waitForAll(handles);`
`121`	`119`	`}`
`122`	`120`	`}`
`123`	`121`
Original file line number	Diff line number	Diff line change
`@@ -337,4 +337,8 @@ public void tryCompleteDelayedFetch(DelayedOperationKey key) {`
`337`	`337`	`}`
`338`	`338`	`}`
`339`	`339`
	`340`	`+ public CompletableFuture<?> updatePurgeAbortedTxnsOffsets() {`
	`341`	`+ return logManager.updatePurgeAbortedTxnsOffsets();`
	`342`	`+ }`
	`343`	`+`
`340`	`344`	`}`