[feat][txn] Transaction buffer snapshot writer reuse

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SystemTopicTxnBufferSnapshotService.java

@@ -21,63 +21,127 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.broker.systopic.NamespaceEventsSystemTopicFactory;
 import org.apache.pulsar.broker.systopic.SystemTopicClient;
 import org.apache.pulsar.broker.systopic.SystemTopicClientBase;
 import org.apache.pulsar.client.api.PulsarClient;
 import org.apache.pulsar.client.api.PulsarClientException;
 import org.apache.pulsar.common.events.EventType;
+import org.apache.pulsar.common.naming.NamespaceName;
 import org.apache.pulsar.common.naming.TopicName;
-import org.apache.pulsar.common.util.FutureUtil;
 
+@Slf4j
 public class SystemTopicTxnBufferSnapshotService<T> {
 
-    protected final Map<TopicName, SystemTopicClient<T>> clients;
+    protected final ConcurrentHashMap<NamespaceName, SystemTopicClient<T>> clients;
     protected final NamespaceEventsSystemTopicFactory namespaceEventsSystemTopicFactory;
 
     protected final Class<T> schemaType;
     protected final EventType systemTopicType;
 
+    private final ConcurrentHashMap<NamespaceName, ReferenceCountedWriter<T>> refCountedWriterMap;
+
+    // The class ReferenceCountedWriter will maintain the reference count,
+    // when the reference count decrement to 0, it will be removed from writerFutureMap, the writer will be closed.
+    public static class ReferenceCountedWriter<T> {
+
+        private final AtomicLong referenceCount;
+        private final NamespaceName namespaceName;
+        private final CompletableFuture<SystemTopicClient.Writer<T>> future;
+        private final SystemTopicTxnBufferSnapshotService<T> snapshotService;
+
+        public ReferenceCountedWriter(NamespaceName namespaceName,
+                                      CompletableFuture<SystemTopicClient.Writer<T>> future,
+                                      SystemTopicTxnBufferSnapshotService<T> snapshotService) {
+            this.referenceCount = new AtomicLong(1);
+            this.namespaceName = namespaceName;
+            this.snapshotService = snapshotService;
+            this.future = future;
+            this.future.exceptionally(t -> {
+                        log.error("[{}] Failed to create TB snapshot writer.", namespaceName, t);
+                snapshotService.refCountedWriterMap.remove(namespaceName, this);
+                return null;
+            });
+        }
+
+        public CompletableFuture<SystemTopicClient.Writer<T>> getFuture() {
+            return future;
+        }
+
+        private synchronized boolean retain() {
+            return this.referenceCount.incrementAndGet() > 0;
+        }
+
+        public synchronized void release() {
+            if (this.referenceCount.decrementAndGet() == 0) {
+                snapshotService.refCountedWriterMap.remove(namespaceName, this);
+                future.thenAccept(writer -> {
+                    final String topicName = writer.getSystemTopicClient().getTopicName().toString();
+                    writer.closeAsync().exceptionally(t -> {
+                        if (t != null) {
+                            log.error("[{}] Failed to close TB snapshot writer.", topicName, t);
+                        } else {
+                            if (log.isDebugEnabled()) {
+                                log.debug("[{}] Success to close TB snapshot writer.", topicName);
+                            }
+                        }
+                        return null;
+                    });
+                });
+            }
+        }
+
+    }
+
     public SystemTopicTxnBufferSnapshotService(PulsarClient client, EventType systemTopicType,
                                                Class<T> schemaType) {
         this.namespaceEventsSystemTopicFactory = new NamespaceEventsSystemTopicFactory(client);
         this.systemTopicType = systemTopicType;
         this.schemaType = schemaType;
         this.clients = new ConcurrentHashMap<>();
-    }
-
-    public CompletableFuture<SystemTopicClient.Writer<T>> createWriter(TopicName topicName) {
-        return getTransactionBufferSystemTopicClient(topicName).thenCompose(SystemTopicClient::newWriterAsync);
+        this.refCountedWriterMap = new ConcurrentHashMap<>();
     }
 
     public CompletableFuture<SystemTopicClient.Reader<T>> createReader(TopicName topicName) {
-        return getTransactionBufferSystemTopicClient(topicName).thenCompose(SystemTopicClient::newReaderAsync);
+        return getTransactionBufferSystemTopicClient(topicName.getNamespaceObject()).newReaderAsync();
     }
 
     public void removeClient(TopicName topicName, SystemTopicClientBase<T> transactionBufferSystemTopicClient) {
         if (transactionBufferSystemTopicClient.getReaders().size() == 0
                 && transactionBufferSystemTopicClient.getWriters().size() == 0) {
-            clients.remove(topicName);
+            clients.remove(topicName.getNamespaceObject());
         }
     }
 
-    protected CompletableFuture<SystemTopicClient<T>> getTransactionBufferSystemTopicClient(TopicName topicName) {
+    public ReferenceCountedWriter<T> getReferenceWriter(NamespaceName namespaceName) {
+        return refCountedWriterMap.compute(namespaceName, (k, v) -> {
+            if (v != null && v.retain()) {
+                return v;
+            } else {
+                return new ReferenceCountedWriter<>(namespaceName,
+                        getTransactionBufferSystemTopicClient(namespaceName).newWriterAsync(), this);
+            }
+        });
+    }
+
+    private SystemTopicClient<T> getTransactionBufferSystemTopicClient(NamespaceName namespaceName) {
         TopicName systemTopicName = NamespaceEventsSystemTopicFactory
-                .getSystemTopicName(topicName.getNamespaceObject(), systemTopicType);
+                .getSystemTopicName(namespaceName, systemTopicType);
         if (systemTopicName == null) {
-            return FutureUtil.failedFuture(
-                    new PulsarClientException
-                            .InvalidTopicNameException("Can't create SystemTopicBaseTxnBufferSnapshotIndexService, "
-                            + "because the topicName is null!"));
+            throw new RuntimeException(new PulsarClientException.InvalidTopicNameException(
+                    "Can't get the TB system topic client for namespace " + namespaceName
+                            + " with type " + systemTopicType + "."));
         }
-        return CompletableFuture.completedFuture(clients.computeIfAbsent(systemTopicName,
+
+        return clients.computeIfAbsent(namespaceName,
                 (v) -> namespaceEventsSystemTopicFactory
-                        .createTransactionBufferSystemTopicClient(systemTopicName,
-                                this, schemaType)));
+                        .createTransactionBufferSystemTopicClient(systemTopicName, this, schemaType));
     }
 
     public void close() throws Exception {
-        for (Map.Entry<TopicName, SystemTopicClient<T>> entry : clients.entrySet()) {
+        for (Map.Entry<NamespaceName, SystemTopicClient<T>> entry : clients.entrySet()) {
             entry.getValue().close();
         }
     }

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

@@ -28,6 +28,7 @@
 import org.apache.bookkeeper.mledger.impl.PositionImpl;
 import org.apache.commons.collections4.map.LinkedMap;
 import org.apache.pulsar.broker.service.BrokerServiceException;
+import org.apache.pulsar.broker.service.SystemTopicTxnBufferSnapshotService.ReferenceCountedWriter;
 import org.apache.pulsar.broker.service.persistent.PersistentTopic;
 import org.apache.pulsar.broker.systopic.SystemTopicClient;
 import org.apache.pulsar.broker.transaction.buffer.AbortedTxnProcessor;
@@ -42,7 +43,7 @@
 @Slf4j
 public class SingleSnapshotAbortedTxnProcessorImpl implements AbortedTxnProcessor {
     private final PersistentTopic topic;
-    private final CompletableFuture<SystemTopicClient.Writer<TransactionBufferSnapshot>> takeSnapshotWriter;
+    private final ReferenceCountedWriter<TransactionBufferSnapshot> takeSnapshotWriter;
     /**
      * Aborts, map for jude message is aborted, linked for remove abort txn in memory when this
      * position have been deleted.
@@ -51,12 +52,14 @@ public class SingleSnapshotAbortedTxnProcessorImpl implements AbortedTxnProcesso
 
     private volatile long lastSnapshotTimestamps;
 
+    private volatile boolean isClosed = false;
+
     public SingleSnapshotAbortedTxnProcessorImpl(PersistentTopic topic) {
         this.topic = topic;
         this.takeSnapshotWriter = this.topic.getBrokerService().getPulsar()
                 .getTransactionBufferSnapshotServiceFactory()
-                .getTxnBufferSnapshotService().createWriter(TopicName.get(topic.getName()));
-        this.takeSnapshotWriter.exceptionally((ex) -> {
+                .getTxnBufferSnapshotService().getReferenceWriter(TopicName.get(topic.getName()).getNamespaceObject());
+        this.takeSnapshotWriter.getFuture().exceptionally((ex) -> {
                     log.error("{} Failed to create snapshot writer", topic.getName());
                     topic.close();
                     return null;
@@ -132,7 +135,7 @@ public CompletableFuture<PositionImpl> recoverFromSnapshot() {
 
     @Override
     public CompletableFuture<Void> clearAbortedTxnSnapshot() {
-        return this.takeSnapshotWriter.thenCompose(writer -> {
+        return this.takeSnapshotWriter.getFuture().thenCompose(writer -> {
             TransactionBufferSnapshot snapshot = new TransactionBufferSnapshot();
             snapshot.setTopicName(topic.getName());
             return writer.deleteAsync(snapshot.getTopicName(), snapshot);
@@ -141,7 +144,7 @@ public CompletableFuture<Void> clearAbortedTxnSnapshot() {
 
     @Override
     public CompletableFuture<Void> takeAbortedTxnsSnapshot(PositionImpl maxReadPosition) {
-        return takeSnapshotWriter.thenCompose(writer -> {
+        return takeSnapshotWriter.getFuture().thenCompose(writer -> {
             TransactionBufferSnapshot snapshot = new TransactionBufferSnapshot();
             snapshot.setTopicName(topic.getName());
             snapshot.setMaxReadPositionLedgerId(maxReadPosition.getLedgerId());
@@ -175,8 +178,12 @@ public long getLastSnapshotTimestamps() {
     }
 
     @Override
-    public CompletableFuture<Void> closeAsync() {
-        return takeSnapshotWriter.thenCompose(SystemTopicClient.Writer::closeAsync);
+    public synchronized CompletableFuture<Void> closeAsync() {
+        if (!isClosed) {
+            isClosed = true;
+            takeSnapshotWriter.release();
+        }
+        return CompletableFuture.completedFuture(null);
     }
 
     private void closeReader(SystemTopicClient.Reader<TransactionBufferSnapshot> reader) {

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

@@ -43,6 +43,7 @@
 import org.apache.commons.lang3.tuple.MutablePair;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pulsar.broker.service.BrokerServiceException;
+import org.apache.pulsar.broker.service.SystemTopicTxnBufferSnapshotService.ReferenceCountedWriter;
 import org.apache.pulsar.broker.service.persistent.PersistentTopic;
 import org.apache.pulsar.broker.systopic.SystemTopicClient;
 import org.apache.pulsar.broker.transaction.buffer.AbortedTxnProcessor;
@@ -442,10 +443,10 @@ public class PersistentWorker {
         private final PersistentTopic topic;
 
         //Persistent snapshot segment and index at the single thread.
-        private final CompletableFuture<SystemTopicClient.Writer<TransactionBufferSnapshotSegment>>
-                snapshotSegmentsWriterFuture;
-        private final CompletableFuture<SystemTopicClient.Writer<TransactionBufferSnapshotIndexes>>
-                snapshotIndexWriterFuture;
+        private final ReferenceCountedWriter<TransactionBufferSnapshotSegment> snapshotSegmentsWriter;
+        private final ReferenceCountedWriter<TransactionBufferSnapshotIndexes> snapshotIndexWriter;
+
+        private volatile boolean closed = false;
 
         private enum OperationState {
             None,
@@ -470,18 +471,20 @@ public enum OperationType {
 
         public PersistentWorker(PersistentTopic topic) {
             this.topic = topic;
-            this.snapshotSegmentsWriterFuture =  this.topic.getBrokerService().getPulsar()
+            this.snapshotSegmentsWriter = this.topic.getBrokerService().getPulsar()
                     .getTransactionBufferSnapshotServiceFactory()
-                    .getTxnBufferSnapshotSegmentService().createWriter(TopicName.get(topic.getName()));
-            this.snapshotSegmentsWriterFuture.exceptionally(ex -> {
+                    .getTxnBufferSnapshotSegmentService()
+                    .getReferenceWriter(TopicName.get(topic.getName()).getNamespaceObject());
+            this.snapshotSegmentsWriter.getFuture().exceptionally(ex -> {
                         log.error("{} Failed to create snapshot index writer", topic.getName());
                         topic.close();
                         return null;
                     });
-            this.snapshotIndexWriterFuture =  this.topic.getBrokerService().getPulsar()
+            this.snapshotIndexWriter =  this.topic.getBrokerService().getPulsar()
                     .getTransactionBufferSnapshotServiceFactory()
-                    .getTxnBufferSnapshotIndexService().createWriter(TopicName.get(topic.getName()));
-            this.snapshotIndexWriterFuture.exceptionally((ex) -> {
+                    .getTxnBufferSnapshotIndexService()
+                    .getReferenceWriter(TopicName.get(topic.getName()).getNamespaceObject());
+            this.snapshotIndexWriter.getFuture().exceptionally((ex) -> {
                         log.error("{} Failed to create snapshot writer", topic.getName());
                         topic.close();
                         return null;
@@ -631,7 +634,7 @@ private CompletableFuture<Void> writeSnapshotSegmentAsync(LinkedList<TxnID> segm
             transactionBufferSnapshotSegment.setPersistentPositionLedgerId(
                     abortedMarkerPersistentPosition.getLedgerId());
 
-            return snapshotSegmentsWriterFuture.thenCompose(segmentWriter -> {
+            return snapshotSegmentsWriter.getFuture().thenCompose(segmentWriter -> {
                 transactionBufferSnapshotSegment.setSequenceId(this.sequenceID.get());
                 return segmentWriter.writeAsync(buildKey(this.sequenceID.get()), transactionBufferSnapshotSegment);
             }).thenCompose((messageId) -> {
@@ -668,7 +671,7 @@ private CompletableFuture<Void> deleteSnapshotSegment(List<PositionImpl> positio
             List<CompletableFuture<Void>> results = new ArrayList<>();
             for (PositionImpl positionNeedToDelete : positionNeedToDeletes) {
                 long sequenceIdNeedToDelete = indexes.get(positionNeedToDelete).getSequenceID();
-                CompletableFuture<Void> res = snapshotSegmentsWriterFuture
+                CompletableFuture<Void> res = snapshotSegmentsWriter.getFuture()
                         .thenCompose(writer -> writer.deleteAsync(buildKey(sequenceIdNeedToDelete), null))
                         .thenCompose(messageId -> {
                             if (log.isDebugEnabled()) {
@@ -695,7 +698,7 @@ private CompletableFuture<Void> deleteSnapshotSegment(List<PositionImpl> positio
 
         private CompletableFuture<Void> updateSnapshotIndex(TransactionBufferSnapshotIndexesMetadata snapshotSegment) {
             TransactionBufferSnapshotIndexes snapshotIndexes = new TransactionBufferSnapshotIndexes();
-            CompletableFuture<Void> res = snapshotIndexWriterFuture
+            CompletableFuture<Void> res = snapshotIndexWriter.getFuture()
                     .thenCompose((indexesWriter) -> {
                         snapshotIndexes.setIndexList(indexes.values().stream().toList());
                         snapshotIndexes.setSnapshot(snapshotSegment);
@@ -712,7 +715,7 @@ private CompletableFuture<Void> updateSnapshotIndex(TransactionBufferSnapshotInd
 
         private CompletableFuture<Void> clearSnapshotSegmentAndIndexes() {
             CompletableFuture<Void> res = persistentWorker.clearAllSnapshotSegments()
-                    .thenCompose((ignore) -> snapshotIndexWriterFuture
+                    .thenCompose((ignore) -> snapshotIndexWriter.getFuture()
                             .thenCompose(indexesWriter -> indexesWriter.writeAsync(topic.getName(), null)))
                     .thenRun(() ->
                             log.debug("Successes to clear the snapshot segment and indexes for the topic [{}]",
@@ -747,7 +750,7 @@ private CompletableFuture<Void> clearAllSnapshotSegments() {
                                 Message<TransactionBufferSnapshotSegment> message = reader.readNextAsync()
                                         .get(getSystemClientOperationTimeoutMs(), TimeUnit.MILLISECONDS);
                                 if (topic.getName().equals(message.getValue().getTopicName())) {
-                                   snapshotSegmentsWriterFuture.get().write(message.getKey(), null);
+                                   snapshotSegmentsWriter.getFuture().get().write(message.getKey(), null);
                                 }
                             }
                             return CompletableFuture.completedFuture(null);
@@ -760,11 +763,13 @@ private CompletableFuture<Void> clearAllSnapshotSegments() {
            });
         }
 
-
-        CompletableFuture<Void> closeAsync() {
-            return CompletableFuture.allOf(
-                    this.snapshotIndexWriterFuture.thenCompose(SystemTopicClient.Writer::closeAsync),
-                    this.snapshotSegmentsWriterFuture.thenCompose(SystemTopicClient.Writer::closeAsync));
+        synchronized CompletableFuture<Void> closeAsync() {
+            if (!closed) {
+                closed = true;
+                snapshotSegmentsWriter.release();
+                snapshotIndexWriter.release();
+            }
+            return CompletableFuture.completedFuture(null);
         }
     }