From 99c5d57168b8741afd25277a752e08c1ec5347b4 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 19 Sep 2026 14:37:22 +0300 Subject: [PATCH] [improve][broker] Speed up auto-split Key Shared selection Publish an immutable lookup snapshot after membership changes so the hot selection path uses primitive-bound binary search instead of a concurrent skip-list lookup. Assisted-by: Codex --- .../service/AutoSplitSelectorBenchmark.java | 85 +++++++++++++++++++ ...ngeAutoSplitStickyKeyConsumerSelector.java | 46 +++++++++- ...utoSplitStickyKeyConsumerSelectorTest.java | 70 +++++++++++++++ 3 files changed, 197 insertions(+), 4 deletions(-) create mode 100644 microbench/src/main/java/org/apache/pulsar/broker/service/AutoSplitSelectorBenchmark.java diff --git a/microbench/src/main/java/org/apache/pulsar/broker/service/AutoSplitSelectorBenchmark.java b/microbench/src/main/java/org/apache/pulsar/broker/service/AutoSplitSelectorBenchmark.java new file mode 100644 index 0000000000000..5df5f7258365b --- /dev/null +++ b/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 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> 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; + } + } +} diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/HashRangeAutoSplitStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/HashRangeAutoSplitStickyKeyConsumerSelector.java index 48d5491d119b2..743cf07d4bda0 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/HashRangeAutoSplitStickyKeyConsumerSelector.java +++ b/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 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> addCons return CompletableFuture.failedFuture(e); } } + updateLookupSnapshot(); if (!addOrRemoveReturnsImpactedConsumersResult) { return CompletableFuture.completedFuture(Optional.empty()); } @@ -122,6 +125,7 @@ public synchronized Optional removeConsumer(Consumer co } else { rangeMap.remove(removeRange); } + updateLookupSnapshot(); } if (!addOrRemoveReturnsImpactedConsumersResult) { return Optional.empty(); @@ -135,11 +139,33 @@ public synchronized Optional 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 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; + } + } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/HashRangeAutoSplitStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/HashRangeAutoSplitStickyKeyConsumerSelectorTest.java index b3c9bc43d255a..71bdc7c86d65a 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/HashRangeAutoSplitStickyKeyConsumerSelectorTest.java +++ b/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 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 observedConsumers = java.util.concurrent.ConcurrentHashMap.newKeySet(); + observedConsumers.add(stableConsumer); + ExecutorService executor = Executors.newFixedThreadPool(4); + try { + List> 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> 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); + } + } }