[improve][broker] Avoid topic-wide deduplication lock

microbench/src/main/java/org/apache/pulsar/broker/service/MessageDeduplicationSequenceCheckBenchmark.java

@@ -0,0 +1,153 @@
+/*
+ * 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.pulsar.broker.service;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.TimeUnit;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.ThreadParams;
+
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@Fork(2)
+public class MessageDeduplicationSequenceCheckBenchmark {
+
+    @State(Scope.Benchmark)
+    public static class SequenceMaps {
+        private final ConcurrentMap<String, Long> monitorDistinct = createMap();
+        private final ConcurrentMap<String, Long> atomicDistinct = createMap();
+        private final ConcurrentMap<String, Long> monitorShared = createMap();
+        private final ConcurrentMap<String, Long> atomicShared = createMap();
+
+        private static ConcurrentMap<String, Long> createMap() {
+            ConcurrentMap<String, Long> map = new ConcurrentHashMap<>();
+            for (int i = 0; i < 64; i++) {
+                map.put("producer-" + i, 0L);
+            }
+            return map;
+        }
+    }
+
+    @State(Scope.Thread)
+    public static class ProducerAccess {
+        private String producerName;
+        private long sequenceId;
+
+        @org.openjdk.jmh.annotations.Setup
+        public void setup(ThreadParams threadParams) {
+            producerName = "producer-" + threadParams.getThreadIndex();
+        }
+
+        long nextSequenceId() {
+            return ++sequenceId;
+        }
+    }
+
+    @Benchmark
+    @Threads(1)
+    public void monitorDistinctSingle(SequenceMaps maps, ProducerAccess access) {
+        updateWithMonitor(maps.monitorDistinct, access.producerName, access.nextSequenceId());
+    }
+
+    @Benchmark
+    @Threads(1)
+    public void atomicDistinctSingle(SequenceMaps maps, ProducerAccess access) {
+        updateAtomically(maps.atomicDistinct, access.producerName, access.nextSequenceId());
+    }
+
+    @Benchmark
+    @Threads(16)
+    public void monitorDistinctConcurrent(SequenceMaps maps, ProducerAccess access) {
+        updateWithMonitor(maps.monitorDistinct, access.producerName, access.nextSequenceId());
+    }
+
+    @Benchmark
+    @Threads(16)
+    public void atomicDistinctConcurrent(SequenceMaps maps, ProducerAccess access) {
+        updateAtomically(maps.atomicDistinct, access.producerName, access.nextSequenceId());
+    }
+
+    @Benchmark
+    @Threads(16)
+    public void monitorSharedProducer(SequenceMaps maps) {
+        incrementWithMonitor(maps.monitorShared, "producer-0");
+    }
+
+    @Benchmark
+    @Threads(16)
+    public void atomicSharedProducer(SequenceMaps maps) {
+        incrementAtomically(maps.atomicShared, "producer-0");
+    }
+
+    private static void updateWithMonitor(ConcurrentMap<String, Long> sequenceIds,
+                                          String producerName, long sequenceId) {
+        synchronized (sequenceIds) {
+            Long previous = sequenceIds.get(producerName);
+            if (previous == null || sequenceId > previous) {
+                sequenceIds.put(producerName, sequenceId);
+            }
+        }
+    }
+
+    private static void updateAtomically(ConcurrentMap<String, Long> sequenceIds,
+                                         String producerName, long sequenceId) {
+        Long boxedSequenceId = sequenceId;
+        while (true) {
+            Long previous = sequenceIds.get(producerName);
+            if (previous != null && sequenceId <= previous) {
+                return;
+            }
+            if (previous == null) {
+                if (sequenceIds.putIfAbsent(producerName, boxedSequenceId) == null) {
+                    return;
+                }
+            } else if (sequenceIds.replace(producerName, previous, boxedSequenceId)) {
+                return;
+            }
+        }
+    }
+
+    private static void incrementWithMonitor(ConcurrentMap<String, Long> sequenceIds, String producerName) {
+        synchronized (sequenceIds) {
+            sequenceIds.put(producerName, sequenceIds.get(producerName) + 1);
+        }
+    }
+
+    private static void incrementAtomically(ConcurrentMap<String, Long> sequenceIds, String producerName) {
+        while (true) {
+            Long previous = sequenceIds.get(producerName);
+            Long next = previous + 1;
+            if (sequenceIds.replace(producerName, previous, next)) {
+                return;
+            }
+        }
+    }
+}

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/MessageDeduplication.java

@@ -31,6 +31,7 @@
 import java.util.TreeMap;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.bookkeeper.mledger.AsyncCallbacks.DeleteCursorCallback;
@@ -109,12 +110,12 @@ public MessageDupUnknownException(String topicName, String producerName) {
     // Map that contains the highest sequenceId that have been sent by each producers. The map will be updated before
     // the messages are persisted
     @VisibleForTesting
-    final Map<String, Long> highestSequencedPushed = new ConcurrentHashMap<>();
+    final ConcurrentMap<String, Long> highestSequencedPushed = new ConcurrentHashMap<>();
 
     // Map that contains the highest sequenceId that have been persistent by each producers. The map will be updated
     // after the messages are persisted
     @VisibleForTesting
-    final Map<String, Long> highestSequencedPersisted = new ConcurrentHashMap<>();
+    final ConcurrentMap<String, Long> highestSequencedPersisted = new ConcurrentHashMap<>();
 
     // Number of persisted entries after which to store a snapshot of the sequence ids map
     private final int snapshotInterval;
@@ -479,9 +480,10 @@ public MessageDupStatus isDuplicateNormal(PublishContext publishContext, ByteBuf
             publishContext.setProperty(IS_LAST_CHUNK, Boolean.FALSE);
             return MessageDupStatus.NotDup;
         }
-        // Synchronize the get() and subsequent put() on the map. This would only be relevant if the producer
-        // disconnects and re-connects very quickly. At that point the call can be coming from a different thread
-        synchronized (highestSequencedPushed) {
+        // A producer that disconnects and reconnects quickly can publish from two threads. Update its sequence
+        // atomically without serializing unrelated producers on one topic-wide monitor.
+        Long newHighestSequenceId = null;
+        while (true) {
             Long lastSequenceIdPushed = highestSequencedPushed.get(producerName);
             if (lastSequenceIdPushed != null && sequenceId <= lastSequenceIdPushed) {
                 log.debug()
@@ -503,7 +505,17 @@ public MessageDupStatus isDuplicateNormal(PublishContext publishContext, ByteBuf
                     return MessageDupStatus.Unknown;
                 }
             }
-            highestSequencedPushed.put(producerName, highestSequenceId);
+            if (newHighestSequenceId == null) {
+                newHighestSequenceId = highestSequenceId;
+            }
+            if (lastSequenceIdPushed == null) {
+                if (highestSequencedPushed.putIfAbsent(producerName, newHighestSequenceId) == null) {
+                    break;
+                }
+            } else if (highestSequencedPushed.replace(producerName, lastSequenceIdPushed,
+                    newHighestSequenceId)) {
+                break;
+            }
         }
         // Only put sequence ID into highestSequencedPushed and
         // highestSequencedPersisted until receive and persistent the last chunk.

pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/MessageDuplicationTest.java

@@ -38,10 +38,16 @@
 import io.netty.buffer.Unpooled;
 import io.netty.channel.EventLoopGroup;
 import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import lombok.CustomLog;
 import org.apache.bookkeeper.mledger.ManagedCursor;
 import org.apache.bookkeeper.mledger.ManagedLedger;
@@ -175,6 +181,52 @@ public void testIsDuplicate() {
         assertEquals(lastSequenceIdPushed.longValue(), 5);
     }
 
+    @Test
+    public void testConcurrentDuplicateCheckForSameProducer() throws Exception {
+        PulsarService pulsarService = mock(PulsarService.class);
+        ServiceConfiguration serviceConfiguration = new ServiceConfiguration();
+        serviceConfiguration.setBrokerDeduplicationEntriesInterval(BROKER_DEDUPLICATION_ENTRIES_INTERVAL);
+        serviceConfiguration.setBrokerDeduplicationMaxNumberOfProducers(BROKER_DEDUPLICATION_MAX_NUMBER_PRODUCERS);
+        serviceConfiguration.setReplicatorPrefix(REPLICATOR_PREFIX);
+
+        doReturn(serviceConfiguration).when(pulsarService).getConfiguration();
+        MessageDeduplication messageDeduplication = spyWithClassAndConstructorArgs(MessageDeduplication.class,
+                pulsarService, mock(PersistentTopic.class), mock(ManagedLedger.class));
+        doReturn(true).when(messageDeduplication).isEnabled();
+
+        int threads = 32;
+        ExecutorService executor = Executors.newFixedThreadPool(threads);
+        CountDownLatch start = new CountDownLatch(1);
+        ByteBuf message = getMessage("producer", 1);
+        List<Future<MessageDeduplication.MessageDupStatus>> results = new ArrayList<>();
+        try {
+            for (int i = 0; i < threads; i++) {
+                Topic.PublishContext context = getPublishContext("producer", 1);
+                results.add(executor.submit(() -> {
+                    start.await();
+                    return messageDeduplication.isDuplicate(context, message);
+                }));
+            }
+            start.countDown();
+
+            int accepted = 0;
+            int unknown = 0;
+            for (Future<MessageDeduplication.MessageDupStatus> result : results) {
+                switch (result.get()) {
+                    case NotDup -> accepted++;
+                    case Unknown -> unknown++;
+                    default -> throw new AssertionError("A sequence cannot be known persisted before completion");
+                }
+            }
+            assertEquals(accepted, 1);
+            assertEquals(unknown, threads - 1);
+            assertEquals(messageDeduplication.highestSequencedPushed.get("producer").longValue(), 1L);
+        } finally {
+            message.release();
+            executor.shutdownNow();
+        }
+    }
+
     @Test
     @SuppressWarnings("unchecked")
     public void testInactiveProducerRemove() throws Exception {