[improve][broker] Speed up auto-split Key_Shared consumer selection

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

@@ -0,0 +1,85 @@
+/*
+ * 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 static org.mockito.Mockito.mock;
+import java.util.Map;
+import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.TimeUnit;
+import org.apache.pulsar.client.api.Range;
+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.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@Fork(2)
+@State(Scope.Benchmark)
+public class AutoSplitSelectorBenchmark {
+
+    @Param({"2", "10", "50"})
+    private int consumerCount;
+
+    private HashRangeAutoSplitStickyKeyConsumerSelector selector;
+    private ConcurrentSkipListMap<Integer, Consumer> previousLookup;
+
+    @Setup
+    public void setup() throws Exception {
+        selector = new HashRangeAutoSplitStickyKeyConsumerSelector(true);
+        for (int i = 0; i < consumerCount; i++) {
+            selector.addConsumer(mock(Consumer.class)).join();
+        }
+        previousLookup = new ConcurrentSkipListMap<>();
+        for (Map.Entry<Consumer, java.util.List<Range>> entry : selector.getConsumerKeyHashRanges().entrySet()) {
+            for (Range range : entry.getValue()) {
+                previousLookup.put(range.getEnd(), entry.getKey());
+            }
+        }
+    }
+
+    @Benchmark
+    public Consumer skipListLookup(Cursor cursor) {
+        return previousLookup.ceilingEntry(cursor.nextHash()).getValue();
+    }
+
+    @Benchmark
+    public Consumer snapshotLookup(Cursor cursor) {
+        return selector.select(cursor.nextHash());
+    }
+
+    @State(Scope.Thread)
+    public static class Cursor {
+        private int hash;
+
+        int nextHash() {
+            hash = (hash + 40503) & 0xffff;
+            return hash;
+        }
+    }
+}

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

@@ -50,7 +50,8 @@
  *
  * 0 -< 65536(consumer-1)
  *
- * In this approach use skip list map to maintain the hash range and consumers.
+ * In this approach use skip list map to maintain the hash range and consumers. An immutable array snapshot is
+ * published after membership changes for lookups.
  *
  * Select consumer will return the ceiling key of message key hashcode % range size.
  *
@@ -62,6 +63,7 @@ public class HashRangeAutoSplitStickyKeyConsumerSelector implements StickyKeyCon
     private final Map<Consumer, Integer> consumerRange;
     private final boolean addOrRemoveReturnsImpactedConsumersResult;
     private ConsumerHashAssignmentsSnapshot consumerHashAssignmentsSnapshot;
+    private volatile LookupSnapshot lookupSnapshot = LookupSnapshot.EMPTY;
 
     public HashRangeAutoSplitStickyKeyConsumerSelector() {
         this(false);
@@ -100,6 +102,7 @@ public synchronized CompletableFuture<Optional<ImpactedConsumersResult>> addCons
                 return CompletableFuture.failedFuture(e);
             }
         }
+        updateLookupSnapshot();
         if (!addOrRemoveReturnsImpactedConsumersResult) {
             return CompletableFuture.completedFuture(Optional.empty());
         }
@@ -122,6 +125,7 @@ public synchronized Optional<ImpactedConsumersResult> removeConsumer(Consumer co
             } else {
                 rangeMap.remove(removeRange);
             }
+            updateLookupSnapshot();
         }
         if (!addOrRemoveReturnsImpactedConsumersResult) {
             return Optional.empty();
@@ -135,11 +139,33 @@ public synchronized Optional<ImpactedConsumersResult> removeConsumer(Consumer co
 
     @Override
     public Consumer select(int hash) {
-        if (!rangeMap.isEmpty()) {
-            return rangeMap.ceilingEntry(hash).getValue();
-        } else {
+        LookupSnapshot snapshot = lookupSnapshot;
+        if (snapshot.rangeEnds.length == 0) {
             return null;
         }
+        int low = 0;
+        int high = snapshot.rangeEnds.length - 1;
+        while (low < high) {
+            int mid = (low + high) >>> 1;
+            if (hash <= snapshot.rangeEnds[mid]) {
+                high = mid;
+            } else {
+                low = mid + 1;
+            }
+        }
+        return snapshot.consumers[low];
+    }
+
+    private void updateLookupSnapshot() {
+        int[] rangeEnds = new int[rangeMap.size()];
+        Consumer[] consumers = new Consumer[rangeEnds.length];
+        int index = 0;
+        for (Entry<Integer, Consumer> entry : rangeMap.entrySet()) {
+            rangeEnds[index] = entry.getKey();
+            consumers[index] = entry.getValue();
+            index++;
+        }
+        lookupSnapshot = new LookupSnapshot(rangeEnds, consumers);
     }
 
     @Override
@@ -199,4 +225,16 @@ private boolean is2Power(int num) {
         }
         return (num & num - 1) == 0;
     }
+
+    private static final class LookupSnapshot {
+        private static final LookupSnapshot EMPTY = new LookupSnapshot(new int[0], new Consumer[0]);
+
+        private final int[] rangeEnds;
+        private final Consumer[] consumers;
+
+        private LookupSnapshot(int[] rangeEnds, Consumer[] consumers) {
+            this.rangeEnds = rangeEnds;
+            this.consumers = consumers;
+        }
+    }
 }

pulsar-broker/src/test/java/org/apache/pulsar/broker/service/HashRangeAutoSplitStickyKeyConsumerSelectorTest.java

@@ -26,6 +26,10 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import org.apache.pulsar.client.api.Range;
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -81,4 +85,70 @@ public void testGetConsumerKeyHashRangesWithSameConsumerName() throws Exception
             prev = ranges;
         }
     }
+
+    @Test
+    public void testSelectionMatchesPublishedRangesAcrossMembershipChanges() throws Exception {
+        HashRangeAutoSplitStickyKeyConsumerSelector selector =
+                new HashRangeAutoSplitStickyKeyConsumerSelector(2 << 5, false);
+        List<Consumer> consumers = new ArrayList<>();
+        for (int i = 0; i < 8; i++) {
+            Consumer consumer = mock(Consumer.class);
+            selector.addConsumer(consumer).join();
+            consumers.add(consumer);
+            assertSelectionMatchesRanges(selector, 64);
+        }
+        for (Consumer consumer : consumers) {
+            selector.removeConsumer(consumer);
+            assertSelectionMatchesRanges(selector, 64);
+        }
+        Assert.assertNull(selector.select(0));
+    }
+
+    @Test
+    public void testConcurrentSelectionDuringMembershipChanges() throws Exception {
+        HashRangeAutoSplitStickyKeyConsumerSelector selector =
+                new HashRangeAutoSplitStickyKeyConsumerSelector(2 << 10, false);
+        Consumer stableConsumer = mock(Consumer.class);
+        selector.addConsumer(stableConsumer).join();
+        Set<Consumer> observedConsumers = java.util.concurrent.ConcurrentHashMap.newKeySet();
+        observedConsumers.add(stableConsumer);
+        ExecutorService executor = Executors.newFixedThreadPool(4);
+        try {
+            List<CompletableFuture<Void>> readers = new ArrayList<>();
+            for (int reader = 0; reader < 3; reader++) {
+                final int offset = reader;
+                readers.add(CompletableFuture.runAsync(() -> {
+                    for (int hash = offset; hash < 4096; hash += 3) {
+                        Consumer selected = selector.select(hash);
+                        Assert.assertNotNull(selected);
+                        Assert.assertTrue(observedConsumers.contains(selected));
+                    }
+                }, executor));
+            }
+            for (int i = 0; i < 100; i++) {
+                Consumer transientConsumer = mock(Consumer.class);
+                observedConsumers.add(transientConsumer);
+                selector.addConsumer(transientConsumer).join();
+                selector.removeConsumer(transientConsumer);
+            }
+            CompletableFuture.allOf(readers.toArray(CompletableFuture[]::new)).join();
+        } finally {
+            executor.shutdownNow();
+        }
+    }
+
+    private static void assertSelectionMatchesRanges(HashRangeAutoSplitStickyKeyConsumerSelector selector,
+                                                     int rangeSize) {
+        Map<Consumer, List<Range>> ranges = selector.getConsumerKeyHashRanges();
+        for (int hash = 0; hash < rangeSize; hash++) {
+            Consumer selected = selector.select(hash);
+            int currentHash = hash;
+            Consumer expected = ranges.entrySet().stream()
+                    .filter(entry -> entry.getValue().stream().anyMatch(range -> range.contains(currentHash)))
+                    .map(Map.Entry::getKey)
+                    .findFirst()
+                    .orElse(null);
+            Assert.assertSame(selected, expected, "hash " + hash);
+        }
+    }
 }