[fix][txn] Fix deadlock when loading transaction buffer snapshot (#24401)

Author: BewareMyPower

pulsar-broker/src/main/java/org/apache/pulsar/broker/PulsarService.java

@@ -301,6 +301,7 @@ public class PulsarService implements AutoCloseable, ShutdownService {
 
     private TransactionPendingAckStoreProvider transactionPendingAckStoreProvider;
     private final ExecutorProvider transactionExecutorProvider;
+    private final ExecutorProvider transactionSnapshotRecoverExecutorProvider;
     private final MonotonicClock monotonicClock;
     private String brokerId;
     private final CompletableFuture<Void> readyForIncomingRequestsFuture = new CompletableFuture<>();
@@ -371,8 +372,11 @@ public PulsarService(ServiceConfiguration config,
         if (config.isTransactionCoordinatorEnabled()) {
             this.transactionExecutorProvider = new ExecutorProvider(this.getConfiguration()
                     .getNumTransactionReplayThreadPoolSize(), "pulsar-transaction-executor");
+            this.transactionSnapshotRecoverExecutorProvider = new ExecutorProvider(this.getConfiguration()
+                    .getNumTransactionReplayThreadPoolSize(), "pulsar-transaction-snapshot-recover");
         } else {
             this.transactionExecutorProvider = null;
+            this.transactionSnapshotRecoverExecutorProvider = null;
         }
 
         this.ioEventLoopGroup = EventLoopUtil.newEventLoopGroup(config.getNumIOThreads(), config.isEnableBusyWait(),
@@ -661,6 +665,9 @@ public CompletableFuture<Void> closeAsync() {
             if (transactionExecutorProvider != null) {
                 transactionExecutorProvider.shutdownNow();
             }
+            if (transactionSnapshotRecoverExecutorProvider != null) {
+                transactionSnapshotRecoverExecutorProvider.shutdownNow();
+            }
             if (transactionTimer != null) {
                 transactionTimer.stop();
             }

pulsar-broker/src/main/java/org/apache/pulsar/broker/transaction/buffer/impl/SingleSnapshotAbortedTxnProcessorImpl.java

@@ -87,7 +87,7 @@ public boolean checkAbortedTransaction(TxnID txnID) {
     public CompletableFuture<Position> recoverFromSnapshot() {
         final var future = new CompletableFuture<Position>();
         final var pulsar = topic.getBrokerService().getPulsar();
-        pulsar.getTransactionExecutorProvider().getExecutor(this).execute(() -> {
+        pulsar.getTransactionSnapshotRecoverExecutorProvider().getExecutor(this).execute(() -> {
             try {
                 final var snapshot = pulsar.getTransactionBufferSnapshotServiceFactory().getTxnBufferSnapshotService()
                         .getTableView().readLatest(topic.getName());

pulsar-broker/src/main/java/org/apache/pulsar/broker/transaction/buffer/impl/SnapshotSegmentAbortedTxnProcessorImpl.java

@@ -228,7 +228,7 @@ public CompletableFuture<Void> takeAbortedTxnsSnapshot(Position maxReadPosition)
     public CompletableFuture<Position> recoverFromSnapshot() {
         final var pulsar = topic.getBrokerService().getPulsar();
         final var future = new CompletableFuture<Position>();
-        pulsar.getTransactionExecutorProvider().getExecutor(this).execute(() -> {
+        pulsar.getTransactionSnapshotRecoverExecutorProvider().getExecutor(this).execute(() -> {
             try {
                 final var indexes = pulsar.getTransactionBufferSnapshotServiceFactory()
                         .getTxnBufferSnapshotIndexService().getTableView().readLatest(topic.getName());

pulsar-broker/src/main/java/org/apache/pulsar/broker/transaction/buffer/impl/TableView.java

@@ -58,6 +58,16 @@ public TableView(Function<TopicName, CompletableFuture<Reader<T>>> readerCreator
     }
 
     public T readLatest(String topic) throws Exception {
+        try {
+            return internalReadLatest(topic);
+        } catch (Exception e) {
+            final var namespace = TopicName.get(topic).getNamespaceObject();
+            readers.remove(namespace);
+            throw e;
+        }
+    }
+
+    private T internalReadLatest(String topic) throws Exception {
         final var reader = getReader(topic);
         while (wait(reader.hasMoreEventsAsync(), "has more events")) {
             final var msg = wait(reader.readNextAsync(), "read message");

pulsar-broker/src/main/java/org/apache/pulsar/utils/SimpleCache.java

@@ -41,7 +41,7 @@ private class ExpirableValue<V> {
 
         boolean tryExpire() {
             if (System.currentTimeMillis() >= deadlineMs) {
-                expireCallback.accept(value);
+                cancel();
                 return true;
             } else {
                 return false;
@@ -51,6 +51,10 @@ boolean tryExpire() {
         void updateDeadline() {
             deadlineMs = System.currentTimeMillis() + timeoutMs;
         }
+
+        void cancel() {
+            expireCallback.accept(value);
+        }
     }
 
     public SimpleCache(final ScheduledExecutorService scheduler, final long timeoutMs, final long frequencyMs) {
@@ -80,4 +84,14 @@ public synchronized V get(final K key, final Supplier<V> valueSupplier, final Co
         cache.put(key, newValue);
         return newValue.value;
     }
+
+    public void remove(final K key) {
+        final ExpirableValue<V> value;
+        synchronized (this) {
+            value = cache.remove(key);
+        }
+        if (value != null) {
+            value.cancel();
+        }
+    }
 }

pulsar-broker/src/test/java/org/apache/pulsar/broker/transaction/TransactionTest.java

@@ -178,7 +178,10 @@ public class TransactionTest extends TransactionTestBase {
 
     @BeforeClass
     protected void setup() throws Exception {
-       setUpBase(NUM_BROKERS, NUM_PARTITIONS, NAMESPACE1 + "/test", 0);
+        // Use a single transaction thread to reproduce possible deadlock easily
+        conf.setNumTransactionReplayThreadPoolSize(1);
+        conf.setManagedLedgerNumSchedulerThreads(1);
+        setUpBase(NUM_BROKERS, NUM_PARTITIONS, NAMESPACE1 + "/test", 0);
     }
 
     @AfterClass(alwaysRun = true)

pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java

@@ -241,6 +241,8 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
     private final AtomicReference<ClientCnx> clientCnxUsedForConsumerRegistration = new AtomicReference<>();
     private final AtomicInteger previousExceptionCount = new AtomicInteger();
     private volatile boolean hasSoughtByTimestamp = false;
+    // This field will be set after the state becomes Failed, then the following operations will fail immediately
+    private volatile Throwable failReason = null;
 
     static <T> ConsumerImpl<T> newConsumerImpl(PulsarClientImpl client,
                                                String topic,
@@ -980,11 +982,13 @@ public CompletableFuture<Void> connectionOpened(final ClientCnx cnx) {
                 } else if (!subscribeFuture.isDone()) {
                     // unable to create new consumer, fail operation
                     setState(State.Failed);
+                    final Throwable throwable = PulsarClientException.wrap(e, String.format("Failed to subscribe the "
+                            + "topic %s with subscription name %s when connecting to the broker", topicName.toString(),
+                            subscription));
+                    fail(throwable);
+
                     closeConsumerTasks();
-                    subscribeFuture.completeExceptionally(
-                            PulsarClientException.wrap(e, String.format("Failed to subscribe the topic %s "
-                                            + "with subscription name %s when connecting to the broker",
-                                    topicName.toString(), subscription)));
+                    subscribeFuture.completeExceptionally(throwable);
                     client.cleanupConsumer(this);
                 } else if (isUnrecoverableError(e.getCause())) {
                     closeWhenReceivedUnrecoverableError(e.getCause(), cnx);
@@ -1025,7 +1029,7 @@ protected void closeWhenReceivedUnrecoverableError(Throwable t, ClientCnx cnx) {
                 topic, subscription, cnxStr, t.getClass().getName(), t.getMessage());
         closeAsync().whenComplete((__, ex) -> {
             if (ex == null) {
-                setState(State.Failed);
+                fail(t);
                 return;
             }
             log.error("[{}][{}] {} Failed to close consumer after got an error that does not support to retry: {} {}",
@@ -1119,7 +1123,7 @@ public boolean connectionFailed(PulsarClientException exception) {
         if (nonRetriableError || timeout) {
             exception.setPreviousExceptionCount(previousExceptionCount);
             if (subscribeFuture.completeExceptionally(exception)) {
-                setState(State.Failed);
+                fail(exception);
                 if (nonRetriableError) {
                     log.info("[{}] Consumer creation failed for consumer {} with unretriableError {}",
                             topic, consumerId, exception.getMessage());
@@ -2870,6 +2874,10 @@ private void internalGetLastMessageIdAsync(final Backoff backoff,
                 return null;
             });
         } else {
+            if (failReason != null) {
+                future.completeExceptionally(failReason);
+                return;
+            }
             long nextDelay = Math.min(backoff.next(), remainingTime.get());
             if (nextDelay <= 0) {
                 future.completeExceptionally(
@@ -3301,4 +3309,9 @@ public Producer<byte[]> getRetryLetterProducer() {
     public Producer<byte[]> getDeadLetterProducer() throws ExecutionException, InterruptedException {
         return (deadLetterProducer == null || !deadLetterProducer.isDone()) ? null : deadLetterProducer.get();
     }
+
+    private void fail(Throwable throwable) {
+        setState(State.Failed);
+        failReason = throwable;
+    }
 }