From e8de7313a855b2d9e8935218003573e0d0d02cca Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Fri, 20 Sep 2024 10:59:48 +0300 Subject: [PATCH 01/21] [fix][broker] Fix ordering issue with ConsistentHashingStickyKeyConsumerSelector --- ...stentHashingStickyKeyConsumerSelector.java | 172 ++++++++++++++---- ...tHashingStickyKeyConsumerSelectorTest.java | 98 +++++++++- 2 files changed, 227 insertions(+), 43 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index b2b2b512c8cfc..3a5397722fed1 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -18,17 +18,19 @@ */ package org.apache.pulsar.broker.service; -import com.google.common.collect.Lists; import java.util.ArrayList; +import java.util.Collections; import java.util.Comparator; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.NavigableMap; import java.util.TreeMap; +import java.util.WeakHashMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; +import org.apache.commons.lang3.mutable.MutableInt; import org.apache.pulsar.client.api.Range; import org.apache.pulsar.common.util.Murmur3_32Hash; @@ -44,15 +46,118 @@ public class ConsistentHashingStickyKeyConsumerSelector implements StickyKeyCons private final ReadWriteLock rwLock = new ReentrantReadWriteLock(); // Consistent-Hash ring - private final NavigableMap> hashRing; + private final NavigableMap hashRing; + // used for distributing consumer instance selections evenly in the hash ring when there + // are multiple instances of consumer with the same consumer name or when there are hash collisions + private final Map consumerSelectionCounters; private final int numberOfPoints; public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) { this.hashRing = new TreeMap<>(); + this.consumerSelectionCounters = new WeakHashMap<>(); this.numberOfPoints = numberOfPoints; } + /** + * This class is used to store the consumers and the selected consumer for a hash value in the hash ring. + * This attempts to distribute the consumers evenly in the hash ring for consumers with the same + * consumer name and priority level. These entries collide in the hash ring. + * The selected consumer is the consumer that is selected to serve the hash value. + * It is not changed unless a consumer is removed or a colliding consumer with higher priority or + * lower selection count is added. + */ + private static class HashRingEntry { + // This class is used to store the consumer which it was added to the hash ring + // sorting will be by priority, consumer name and usage count of the consumer instance + record ConsumerEntry(Consumer consumer, MutableInt consumerSelectionCounter) + implements Comparable { + private static final Comparator + BASE_CONSUMER_ENTRY_COMPARATOR = Comparator. + comparing(entry -> entry.consumer().getPriorityLevel()).reversed() + .thenComparing(entry -> entry.consumer().consumerName()); + + + private static final Comparator + CONSUMER_ENTRY_COMPARATOR = BASE_CONSUMER_ENTRY_COMPARATOR + // prefer the consumer instance with lowest selection count so that consumers get + // evenly distributed + .thenComparing(ConsumerEntry::consumerSelectionCounter); + + @Override + public int compareTo(ConsumerEntry o) { + return CONSUMER_ENTRY_COMPARATOR.compare(this, o); + } + + // comparison without the usage count so that the consumer doesn't get changed too eagerly + // when entries are removed + public int baseCompareTo(ConsumerEntry o) { + return BASE_CONSUMER_ENTRY_COMPARATOR.compare(this, o); + } + } + + private final List consumers; + ConsumerEntry selectedConsumerEntry; + + public HashRingEntry() { + this.consumers = new ArrayList<>(); + } + + public void addConsumer(Consumer consumer, MutableInt selectedCounter) { + consumers.add(new ConsumerEntry(consumer, selectedCounter)); + selectConsumer(null); + } + + public boolean removeConsumer(Consumer consumer) { + boolean removed = consumers.removeIf(consumerEntry -> consumerEntry.consumer().equals(consumer)); + selectConsumer(consumer); + return removed; + } + + public Consumer getSelectedConsumer() { + return selectedConsumerEntry != null ? selectedConsumerEntry.consumer() : null; + } + + private void selectConsumer(Consumer removedConsumer) { + if (consumers.size() > 1) { + boolean addOperation = removedConsumer == null; + if (addOperation) { + Collections.sort(consumers); + } + ConsumerEntry newSelectedConsumer = consumers.get(0); + // change the selected consumer only if the newer has higher priority, + // or the same priority and an earlier name in sorting order + if (selectedConsumerEntry == null || addOperation + || selectedConsumerEntry.consumer.equals(removedConsumer) + || selectedConsumerEntry.baseCompareTo(newSelectedConsumer) > 0) { + changeSelectedConsumerEntry(newSelectedConsumer); + } + } else if (consumers.size() == 1) { + changeSelectedConsumerEntry(consumers.get(0)); + } else { + changeSelectedConsumerEntry(null); + } + } + + private void changeSelectedConsumerEntry(ConsumerEntry newSelectedConsumer) { + beforeChangingSelectedConsumerEntry(); + selectedConsumerEntry = newSelectedConsumer; + afterChangingSelectedConsumerEntry(); + } + + private void beforeChangingSelectedConsumerEntry() { + if (selectedConsumerEntry != null) { + selectedConsumerEntry.consumerSelectionCounter.decrement(); + } + } + + private void afterChangingSelectedConsumerEntry() { + if (selectedConsumerEntry != null) { + selectedConsumerEntry.consumerSelectionCounter.increment(); + } + } + } + @Override public CompletableFuture addConsumer(Consumer consumer) { rwLock.writeLock().lock(); @@ -61,17 +166,9 @@ public CompletableFuture addConsumer(Consumer consumer) { // The points are deterministically added based on the hash of the consumer name for (int i = 0; i < numberOfPoints; i++) { int hash = calculateHashForConsumerAndIndex(consumer, i); - hashRing.compute(hash, (k, v) -> { - if (v == null) { - return Lists.newArrayList(consumer); - } else { - if (!v.contains(consumer)) { - v.add(consumer); - v.sort(Comparator.comparing(Consumer::consumerName, String::compareTo)); - } - return v; - } - }); + HashRingEntry hashRingEntry = hashRing.computeIfAbsent(hash, k -> new HashRingEntry()); + // Add the consumer to the hash ring entry + hashRingEntry.addConsumer(consumer, getConsumerSelectedCount(consumer)); } return CompletableFuture.completedFuture(null); } finally { @@ -79,6 +176,10 @@ public CompletableFuture addConsumer(Consumer consumer) { } } + private MutableInt getConsumerSelectedCount(Consumer consumer) { + return consumerSelectionCounters.computeIfAbsent(consumer, k -> new MutableInt()); + } + private static int calculateHashForConsumerAndIndex(Consumer consumer, int index) { String key = consumer.consumerName() + KEY_SEPARATOR + index; return Murmur3_32Hash.getInstance().makeHash(key.getBytes()); @@ -92,15 +193,11 @@ public void removeConsumer(Consumer consumer) { for (int i = 0; i < numberOfPoints; i++) { int hash = calculateHashForConsumerAndIndex(consumer, i); hashRing.compute(hash, (k, v) -> { - if (v == null) { - return null; - } else { - v.removeIf(c -> c.equals(consumer)); - if (v.isEmpty()) { - v = null; - } - return v; + v.removeConsumer(consumer); + if (v.getSelectedConsumer() == null) { + v = null; } + return v; }); } } finally { @@ -115,16 +212,14 @@ public Consumer select(int hash) { if (hashRing.isEmpty()) { return null; } - - List consumerList; - Map.Entry> ceilingEntry = hashRing.ceilingEntry(hash); + HashRingEntry hashRingEntry; + Map.Entry ceilingEntry = hashRing.ceilingEntry(hash); if (ceilingEntry != null) { - consumerList = ceilingEntry.getValue(); + hashRingEntry = ceilingEntry.getValue(); } else { - consumerList = hashRing.firstEntry().getValue(); + hashRingEntry = hashRing.firstEntry().getValue(); } - - return consumerList.get(hash % consumerList.size()); + return hashRingEntry.getSelectedConsumer(); } finally { rwLock.readLock().unlock(); } @@ -135,13 +230,24 @@ public Map> getConsumerKeyHashRanges() { Map> result = new LinkedHashMap<>(); rwLock.readLock().lock(); try { + if (hashRing.isEmpty()) { + return result; + } int start = 0; - for (Map.Entry> entry: hashRing.entrySet()) { - for (Consumer consumer: entry.getValue()) { - result.computeIfAbsent(consumer, key -> new ArrayList<>()) + int lastKey = 0; + for (Map.Entry entry: hashRing.entrySet()) { + Consumer consumer = entry.getValue().getSelectedConsumer(); + result.computeIfAbsent(consumer, key -> new ArrayList<>()) .add(Range.of(start, entry.getKey())); - } - start = entry.getKey() + 1; + lastKey = entry.getKey(); + start = lastKey + 1; + } + // Handle wrap-around + HashRingEntry firstHashRingEntry = hashRing.firstEntry().getValue(); + Consumer firstSelectedConsumer = firstHashRingEntry.getSelectedConsumer(); + List ranges = result.get(firstSelectedConsumer); + if (lastKey != Integer.MAX_VALUE - 1) { + ranges.add(Range.of(lastKey + 1, Integer.MAX_VALUE - 1)); } } finally { rwLock.readLock().unlock(); diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 48311c57338b5..7e21579741f0b 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -18,9 +18,9 @@ */ package org.apache.pulsar.broker.service; +import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; - import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -31,6 +31,7 @@ import java.util.stream.IntStream; import org.apache.pulsar.broker.service.BrokerServiceException.ConsumerAssignException; import org.apache.pulsar.client.api.Range; +import org.mockito.Mockito; import org.testng.Assert; import org.testng.annotations.Test; @@ -54,7 +55,7 @@ public void testConsumerSelect() throws ConsumerAssignException { selector.addConsumer(consumer2); final int N = 1000; - final double PERCENT_ERROR = 0.20; // 20 % + final double PERCENT_ERROR = 0.25; // 25 % Map selectionMap = new HashMap<>(); for (int i = 0; i < N; i++) { @@ -146,12 +147,17 @@ public void testGetConsumerKeyHashRanges() throws BrokerServiceException.Consume ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(3); List consumerName = Arrays.asList("consumer1", "consumer2", "consumer3"); List consumers = new ArrayList<>(); + long id=0; for (String s : consumerName) { - Consumer consumer = mock(Consumer.class); - when(consumer.consumerName()).thenReturn(s); + Consumer consumer = createMockConsumer(s, s, id++); selector.addConsumer(consumer); consumers.add(consumer); } + + // check that results are the same when called multiple times + assertThat(selector.getConsumerKeyHashRanges()) + .containsExactlyEntriesOf(selector.getConsumerKeyHashRanges()); + Map> expectedResult = new HashMap<>(); expectedResult.put(consumers.get(0), Arrays.asList( Range.of(119056335, 242013991), @@ -160,17 +166,47 @@ public void testGetConsumerKeyHashRanges() throws BrokerServiceException.Consume expectedResult.put(consumers.get(1), Arrays.asList( Range.of(0, 90164503), Range.of(90164504, 119056334), - Range.of(382436668, 722195656))); + Range.of(382436668, 722195656), + Range.of(1914695767, 2147483646))); expectedResult.put(consumers.get(2), Arrays.asList( Range.of(242013992, 242377547), Range.of(242377548, 382436667), Range.of(1656011843, 1707482097))); - for (Map.Entry> entry : selector.getConsumerKeyHashRanges().entrySet()) { - System.out.println(entry.getValue()); - Assert.assertEquals(entry.getValue(), expectedResult.get(entry.getKey())); - expectedResult.remove(entry.getKey()); + assertThat(selector.getConsumerKeyHashRanges()).containsExactlyInAnyOrderEntriesOf(expectedResult); + } + + @Test + public void testConsumersGetEvenlyMappedWhenThereAreCollisions() + throws BrokerServiceException.ConsumerAssignException { + ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(5); + List consumers = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + // use the same name for all consumers + Consumer consumer = createMockConsumer("consumer", "index " + i, i); + selector.addConsumer(consumer); + consumers.add(consumer); } - Assert.assertEquals(expectedResult.size(), 0); + // check that results are the same when called multiple times + assertThat(selector.getConsumerKeyHashRanges()) + .containsExactlyEntriesOf(selector.getConsumerKeyHashRanges()); + + Map> expectedResult = new HashMap<>(); + expectedResult.put(consumers.get(0), List.of(Range.of(306176209, 365902830))); + expectedResult.put(consumers.get(1), List.of(Range.of(216056714, 306176208))); + expectedResult.put(consumers.get(2), List.of(Range.of(365902831, 1240826377))); + expectedResult.put(consumers.get(3), List.of(Range.of(1240826378, 1862045174))); + expectedResult.put(consumers.get(4), List.of(Range.of(0, 216056713), Range.of(1862045175, 2147483646))); + assertThat(selector.getConsumerKeyHashRanges()).containsExactlyInAnyOrderEntriesOf(expectedResult); + } + + private static Consumer createMockConsumer(String consumerName, String toString, long id) { + // without stubOnly, the mock will record method invocations and run into OOME + Consumer consumer = mock(Consumer.class, Mockito.withSettings().stubOnly()); + when(consumer.consumerName()).thenReturn(consumerName); + when(consumer.getPriorityLevel()).thenReturn(0); + when(consumer.toString()).thenReturn(toString); + when(consumer.consumerId()).thenReturn(id); + return consumer; } // reproduces https://github.com/apache/pulsar/issues/22050 @@ -216,4 +252,46 @@ public void shouldRemoveConsumersFromConsumerKeyHashRanges() { // then there should be no mapping remaining Assert.assertEquals(selector.getConsumerKeyHashRanges().size(), 0); } + + @Test + public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { + final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(25); + final String consumerName = "consumer"; + final int numOfInitialConsumers = 25; + List consumers = new ArrayList<>(); + for (int i = 0; i < numOfInitialConsumers; i++) { + final Consumer consumer = createMockConsumer(consumerName, "index " + i, i); + consumers.add(consumer); + selector.addConsumer(consumer); + } + + int hashRangeSize = Integer.MAX_VALUE; + int validationPointCount = 100; + int increment = hashRangeSize / validationPointCount; + List selectedConsumerBeforeRemoval = new ArrayList<>(); + + for (int i = 0; i < validationPointCount; i++) { + selectedConsumerBeforeRemoval.add(selector.select(i * increment)); + } + + for (int i = 0; i < validationPointCount; i++) { + Consumer selected = selector.select(i * increment); + Consumer expected = selectedConsumerBeforeRemoval.get(i); + assertThat(selected.consumerId()).as("validationPoint %d", i).isEqualTo(expected.consumerId()); + } + + /* + TODO: failing test case + for (Consumer removedConsumer : consumers) { + selector.removeConsumer(removedConsumer); + for (int i = 0; i < validationPointCount; i++) { + Consumer selected = selector.select(i * increment); + Consumer expected = selectedConsumerBeforeRemoval.get(i); + if (expected != removedConsumer) { + assertThat(selected.consumerId()).as("validationPoint %d", i).isEqualTo(expected.consumerId()); + } + } + } + */ + } } From c0f101d5507458411730d4e5ef59cf5580f98275 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Fri, 20 Sep 2024 23:37:12 +0300 Subject: [PATCH 02/21] Improve test --- ...tHashingStickyKeyConsumerSelectorTest.java | 49 ++++++++++++++++--- 1 file changed, 43 insertions(+), 6 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 7e21579741f0b..25fe9170b76b2 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -255,9 +255,9 @@ public void shouldRemoveConsumersFromConsumerKeyHashRanges() { @Test public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { - final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(25); + final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(100); final String consumerName = "consumer"; - final int numOfInitialConsumers = 25; + final int numOfInitialConsumers = 100; List consumers = new ArrayList<>(); for (int i = 0; i < numOfInitialConsumers; i++) { final Consumer consumer = createMockConsumer(consumerName, "index " + i, i); @@ -266,7 +266,7 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { } int hashRangeSize = Integer.MAX_VALUE; - int validationPointCount = 100; + int validationPointCount = 200; int increment = hashRangeSize / validationPointCount; List selectedConsumerBeforeRemoval = new ArrayList<>(); @@ -280,8 +280,6 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { assertThat(selected.consumerId()).as("validationPoint %d", i).isEqualTo(expected.consumerId()); } - /* - TODO: failing test case for (Consumer removedConsumer : consumers) { selector.removeConsumer(removedConsumer); for (int i = 0; i < validationPointCount; i++) { @@ -292,6 +290,45 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { } } } - */ + } + + @Test + public void testShouldNotChangeSelectedConsumerWhenConsumerIsAdded() { + final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(100); + final String consumerName = "consumer"; + final int numOfInitialConsumers = 50; + List consumers = new ArrayList<>(); + for (int i = 0; i < numOfInitialConsumers; i++) { + final Consumer consumer = createMockConsumer(consumerName, "index " + i, i); + consumers.add(consumer); + selector.addConsumer(consumer); + } + + int hashRangeSize = Integer.MAX_VALUE; + int validationPointCount = 200; + int increment = hashRangeSize / validationPointCount; + List selectedConsumerBeforeRemoval = new ArrayList<>(); + + for (int i = 0; i < validationPointCount; i++) { + selectedConsumerBeforeRemoval.add(selector.select(i * increment)); + } + + for (int i = 0; i < validationPointCount; i++) { + Consumer selected = selector.select(i * increment); + Consumer expected = selectedConsumerBeforeRemoval.get(i); + assertThat(selected.consumerId()).as("validationPoint %d", i).isEqualTo(expected.consumerId()); + } + + for (int i = numOfInitialConsumers; i < numOfInitialConsumers * 2; i++) { + final Consumer addedConsumer = createMockConsumer(consumerName, "index " + i, i); + selector.addConsumer(addedConsumer); + for (int j = 0; j < validationPointCount; j++) { + Consumer selected = selector.select(j * increment); + Consumer expected = selectedConsumerBeforeRemoval.get(j); + if (expected != addedConsumer) { + assertThat(selected.consumerId()).as("validationPoint %d", j).isEqualTo(expected.consumerId()); + } + } + } } } From d9544ddf0f41175be6b51f94450ee6e863d34bb8 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 00:04:59 +0300 Subject: [PATCH 03/21] Use insertion sorting for more stability --- ...stentHashingStickyKeyConsumerSelector.java | 45 +++++++------------ ...tHashingStickyKeyConsumerSelectorTest.java | 3 +- 2 files changed, 18 insertions(+), 30 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index 3a5397722fed1..957755f4b7021 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -19,7 +19,6 @@ package org.apache.pulsar.broker.service; import java.util.ArrayList; -import java.util.Collections; import java.util.Comparator; import java.util.LinkedHashMap; import java.util.List; @@ -88,12 +87,6 @@ record ConsumerEntry(Consumer consumer, MutableInt consumerSelectionCounter) public int compareTo(ConsumerEntry o) { return CONSUMER_ENTRY_COMPARATOR.compare(this, o); } - - // comparison without the usage count so that the consumer doesn't get changed too eagerly - // when entries are removed - public int baseCompareTo(ConsumerEntry o) { - return BASE_CONSUMER_ENTRY_COMPARATOR.compare(this, o); - } } private final List consumers; @@ -104,13 +97,24 @@ public HashRingEntry() { } public void addConsumer(Consumer consumer, MutableInt selectedCounter) { - consumers.add(new ConsumerEntry(consumer, selectedCounter)); - selectConsumer(null); + ConsumerEntry consumerEntry = new ConsumerEntry(consumer, selectedCounter); + insertConsumer(consumerEntry); + selectConsumer(); + } + + private void insertConsumer(ConsumerEntry consumerEntry) { + for (int i = 0; i < consumers.size(); i++) { + if (consumers.get(i).compareTo(consumerEntry) > 0) { + consumers.add(i, consumerEntry); + return; + } + } + consumers.add(consumerEntry); } public boolean removeConsumer(Consumer consumer) { boolean removed = consumers.removeIf(consumerEntry -> consumerEntry.consumer().equals(consumer)); - selectConsumer(consumer); + selectConsumer(); return removed; } @@ -118,25 +122,8 @@ public Consumer getSelectedConsumer() { return selectedConsumerEntry != null ? selectedConsumerEntry.consumer() : null; } - private void selectConsumer(Consumer removedConsumer) { - if (consumers.size() > 1) { - boolean addOperation = removedConsumer == null; - if (addOperation) { - Collections.sort(consumers); - } - ConsumerEntry newSelectedConsumer = consumers.get(0); - // change the selected consumer only if the newer has higher priority, - // or the same priority and an earlier name in sorting order - if (selectedConsumerEntry == null || addOperation - || selectedConsumerEntry.consumer.equals(removedConsumer) - || selectedConsumerEntry.baseCompareTo(newSelectedConsumer) > 0) { - changeSelectedConsumerEntry(newSelectedConsumer); - } - } else if (consumers.size() == 1) { - changeSelectedConsumerEntry(consumers.get(0)); - } else { - changeSelectedConsumerEntry(null); - } + private void selectConsumer() { + changeSelectedConsumerEntry(consumers.isEmpty() ? null : consumers.get(0)); } private void changeSelectedConsumerEntry(ConsumerEntry newSelectedConsumer) { diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 25fe9170b76b2..51b2f38c647f5 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -286,7 +286,8 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { Consumer selected = selector.select(i * increment); Consumer expected = selectedConsumerBeforeRemoval.get(i); if (expected != removedConsumer) { - assertThat(selected.consumerId()).as("validationPoint %d", i).isEqualTo(expected.consumerId()); + assertThat(selected.consumerId()).as("validationPoint %d, removed %s", i, + removedConsumer.toString()).isEqualTo(expected.consumerId()); } } } From 8313f07b20eda7112919dc3fd852172a512ee7ad Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 00:15:38 +0300 Subject: [PATCH 04/21] Update test to match the impl --- .../ConsistentHashingStickyKeyConsumerSelectorTest.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 51b2f38c647f5..af3e6c3c83acf 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -191,10 +191,10 @@ public void testConsumersGetEvenlyMappedWhenThereAreCollisions() .containsExactlyEntriesOf(selector.getConsumerKeyHashRanges()); Map> expectedResult = new HashMap<>(); - expectedResult.put(consumers.get(0), List.of(Range.of(306176209, 365902830))); - expectedResult.put(consumers.get(1), List.of(Range.of(216056714, 306176208))); - expectedResult.put(consumers.get(2), List.of(Range.of(365902831, 1240826377))); - expectedResult.put(consumers.get(3), List.of(Range.of(1240826378, 1862045174))); + expectedResult.put(consumers.get(0), List.of(Range.of(216056714, 306176208))); + expectedResult.put(consumers.get(1), List.of(Range.of(365902831, 1240826377))); + expectedResult.put(consumers.get(2), List.of(Range.of(1240826378, 1862045174))); + expectedResult.put(consumers.get(3), List.of(Range.of(306176209, 365902830))); expectedResult.put(consumers.get(4), List.of(Range.of(0, 216056713), Range.of(1862045175, 2147483646))); assertThat(selector.getConsumerKeyHashRanges()).containsExactlyInAnyOrderEntriesOf(expectedResult); } From 9c8365c708af02d13b42759276f7f92273281bac Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 00:21:00 +0300 Subject: [PATCH 05/21] Improve --- .../service/ConsistentHashingStickyKeyConsumerSelector.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index 957755f4b7021..f0b74f6fa4c4a 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -127,6 +127,9 @@ private void selectConsumer() { } private void changeSelectedConsumerEntry(ConsumerEntry newSelectedConsumer) { + if (newSelectedConsumer == selectedConsumerEntry) { + return; + } beforeChangingSelectedConsumerEntry(); selectedConsumerEntry = newSelectedConsumer; afterChangingSelectedConsumerEntry(); From ce0e0c57c1d59b6abbfb0c4e0092ac93e9b1fdd4 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 00:27:19 +0300 Subject: [PATCH 06/21] Add test that checks hash ranges to get a visual clue of the problem --- ...tHashingStickyKeyConsumerSelectorTest.java | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index af3e6c3c83acf..e3b90d4262a80 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -293,6 +293,32 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { } } + @Test + public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemovedCheckHashRanges() { + final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(5); + final String consumerName = "consumer"; + final int numOfInitialConsumers = 10; + List consumers = new ArrayList<>(); + for (int i = 0; i < numOfInitialConsumers; i++) { + final Consumer consumer = createMockConsumer(consumerName, "index " + i, i); + consumers.add(consumer); + selector.addConsumer(consumer); + } + + int hashRangeSize = Integer.MAX_VALUE; + + Map> expected = selector.getConsumerKeyHashRanges(); + assertThat(selector.getConsumerKeyHashRanges()).as("sanity check").isEqualTo(expected); + + for (Consumer removedConsumer : consumers) { + selector.removeConsumer(removedConsumer); + Map> actual = selector.getConsumerKeyHashRanges(); + expected.remove(removedConsumer); + assertThat(actual).as("removed %s", removedConsumer.toString()) + .containsExactlyInAnyOrderEntriesOf(expected); + } + } + @Test public void testShouldNotChangeSelectedConsumerWhenConsumerIsAdded() { final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(100); From d39a9f2fd1fb93502b005af761b8089c87199275 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 00:43:50 +0300 Subject: [PATCH 07/21] Improve tests --- ...tHashingStickyKeyConsumerSelectorTest.java | 62 ++++++++++++++++--- 1 file changed, 54 insertions(+), 8 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index e3b90d4262a80..68737ed3060e2 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -295,9 +295,9 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { @Test public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemovedCheckHashRanges() { - final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(5); + final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(100); final String consumerName = "consumer"; - final int numOfInitialConsumers = 10; + final int numOfInitialConsumers = 25; List consumers = new ArrayList<>(); for (int i = 0; i < numOfInitialConsumers; i++) { final Consumer consumer = createMockConsumer(consumerName, "index " + i, i); @@ -305,17 +305,63 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemovedCheckHashRan selector.addConsumer(consumer); } - int hashRangeSize = Integer.MAX_VALUE; - Map> expected = selector.getConsumerKeyHashRanges(); assertThat(selector.getConsumerKeyHashRanges()).as("sanity check").isEqualTo(expected); + System.out.println(expected); for (Consumer removedConsumer : consumers) { selector.removeConsumer(removedConsumer); - Map> actual = selector.getConsumerKeyHashRanges(); - expected.remove(removedConsumer); - assertThat(actual).as("removed %s", removedConsumer.toString()) - .containsExactlyInAnyOrderEntriesOf(expected); + for (Map.Entry> entry : expected.entrySet()) { + if (entry.getKey() == removedConsumer) { + continue; + } + for (Range range : entry.getValue()) { + assertThat(selector.select(range.getStart())).as("removed %s, range %s", removedConsumer, range) + .isEqualTo(entry.getKey()); + assertThat(selector.select(range.getEnd())).as("removed %s, range %s", removedConsumer, range) + .isEqualTo(entry.getKey()); + } + } + expected = selector.getConsumerKeyHashRanges(); + } + } + + @Test + public void testShouldNotChangeSelectedConsumerUnnecessarilyWhenConsumerIsAddedCheckHashRanges() { + final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(100); + final String consumerName = "consumer"; + final int numOfInitialConsumers = 25; + List consumers = new ArrayList<>(); + for (int i = 0; i < numOfInitialConsumers; i++) { + final Consumer consumer = createMockConsumer(consumerName, "index " + i, i); + consumers.add(consumer); + selector.addConsumer(consumer); + } + + Map> expected = selector.getConsumerKeyHashRanges(); + assertThat(selector.getConsumerKeyHashRanges()).as("sanity check").isEqualTo(expected); + + for (int i = numOfInitialConsumers; i < numOfInitialConsumers * 2; i++) { + final Consumer addedConsumer = createMockConsumer(consumerName, "index " + i, i); + selector.addConsumer(addedConsumer); + for (Map.Entry> entry : expected.entrySet()) { + if (entry.getKey() == addedConsumer) { + continue; + } + for (Range range : entry.getValue()) { + Consumer rangeStartConsumer = selector.select(range.getStart()); + if (rangeStartConsumer != addedConsumer) { + assertThat(rangeStartConsumer).as("added %s, range start %s", addedConsumer, range) + .isEqualTo(entry.getKey()); + } + Consumer rangeEndConsumer = selector.select(range.getStart()); + if (rangeEndConsumer != addedConsumer) { + assertThat(rangeEndConsumer).as("added %s, range end %s", addedConsumer, range) + .isEqualTo(entry.getKey()); + } + } + } + expected = selector.getConsumerKeyHashRanges(); } } From 1c5967c98ca895f6fd9c8d3f43ed5be7cb9a3fb7 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 00:47:26 +0300 Subject: [PATCH 08/21] Additional check --- .../ConsistentHashingStickyKeyConsumerSelectorTest.java | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 68737ed3060e2..4c867ddc9a10e 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -316,10 +316,13 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemovedCheckHashRan continue; } for (Range range : entry.getValue()) { - assertThat(selector.select(range.getStart())).as("removed %s, range %s", removedConsumer, range) + Consumer rangeStartConsumer = selector.select(range.getStart()); + assertThat(rangeStartConsumer).as("removed %s, range %s", removedConsumer, range) .isEqualTo(entry.getKey()); - assertThat(selector.select(range.getEnd())).as("removed %s, range %s", removedConsumer, range) + Consumer rangeEndConsumer = selector.select(range.getEnd()); + assertThat(rangeEndConsumer).as("removed %s, range %s", removedConsumer, range) .isEqualTo(entry.getKey()); + assertThat(rangeStartConsumer).isSameAs(rangeEndConsumer); } } expected = selector.getConsumerKeyHashRanges(); From 366fb88fd53287bb130fad9364a6a6280a104ad7 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 01:02:06 +0300 Subject: [PATCH 09/21] Improve test --- ...nsistentHashingStickyKeyConsumerSelectorTest.java | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 4c867ddc9a10e..78206a7cc61c2 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -283,11 +283,12 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { for (Consumer removedConsumer : consumers) { selector.removeConsumer(removedConsumer); for (int i = 0; i < validationPointCount; i++) { - Consumer selected = selector.select(i * increment); + int hash = i * increment; + Consumer selected = selector.select(hash); Consumer expected = selectedConsumerBeforeRemoval.get(i); if (expected != removedConsumer) { - assertThat(selected.consumerId()).as("validationPoint %d, removed %s", i, - removedConsumer.toString()).isEqualTo(expected.consumerId()); + assertThat(selected.consumerId()).as("validationPoint %d, removed %s, hash %d", i, + removedConsumer.toString(), hash).isEqualTo(expected.consumerId()); } } } @@ -399,10 +400,11 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsAdded() { final Consumer addedConsumer = createMockConsumer(consumerName, "index " + i, i); selector.addConsumer(addedConsumer); for (int j = 0; j < validationPointCount; j++) { - Consumer selected = selector.select(j * increment); + int hash = j * increment; + Consumer selected = selector.select(hash); Consumer expected = selectedConsumerBeforeRemoval.get(j); if (expected != addedConsumer) { - assertThat(selected.consumerId()).as("validationPoint %d", j).isEqualTo(expected.consumerId()); + assertThat(selected.consumerId()).as("validationPoint %d, hash %d", j, hash).isEqualTo(expected.consumerId()); } } } From 1b6c6cc0facdb6efcc89b4111b05f9471bfc25e0 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 01:32:06 +0300 Subject: [PATCH 10/21] Fix test --- ...ntHashingStickyKeyConsumerSelectorTest.java | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 78206a7cc61c2..82e566dd11f8b 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -24,8 +24,10 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -267,7 +269,7 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { int hashRangeSize = Integer.MAX_VALUE; int validationPointCount = 200; - int increment = hashRangeSize / validationPointCount; + int increment = hashRangeSize / (validationPointCount + 1); List selectedConsumerBeforeRemoval = new ArrayList<>(); for (int i = 0; i < validationPointCount; i++) { @@ -280,15 +282,17 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemoved() { assertThat(selected.consumerId()).as("validationPoint %d", i).isEqualTo(expected.consumerId()); } + Set removedConsumers = new HashSet<>(); for (Consumer removedConsumer : consumers) { selector.removeConsumer(removedConsumer); + removedConsumers.add(removedConsumer); for (int i = 0; i < validationPointCount; i++) { int hash = i * increment; Consumer selected = selector.select(hash); Consumer expected = selectedConsumerBeforeRemoval.get(i); - if (expected != removedConsumer) { - assertThat(selected.consumerId()).as("validationPoint %d, removed %s, hash %d", i, - removedConsumer.toString(), hash).isEqualTo(expected.consumerId()); + if (!removedConsumers.contains(expected)) { + assertThat(selected.consumerId()).as("validationPoint %d, removed %s, hash %d ranges %s", i, + removedConsumer.toString(), hash, selector.getConsumerKeyHashRanges()).isEqualTo(expected.consumerId()); } } } @@ -383,7 +387,7 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsAdded() { int hashRangeSize = Integer.MAX_VALUE; int validationPointCount = 200; - int increment = hashRangeSize / validationPointCount; + int increment = hashRangeSize / (validationPointCount + 1); List selectedConsumerBeforeRemoval = new ArrayList<>(); for (int i = 0; i < validationPointCount; i++) { @@ -396,14 +400,16 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsAdded() { assertThat(selected.consumerId()).as("validationPoint %d", i).isEqualTo(expected.consumerId()); } + Set addedConsumers = new HashSet<>(); for (int i = numOfInitialConsumers; i < numOfInitialConsumers * 2; i++) { final Consumer addedConsumer = createMockConsumer(consumerName, "index " + i, i); selector.addConsumer(addedConsumer); + addedConsumers.add(addedConsumer); for (int j = 0; j < validationPointCount; j++) { int hash = j * increment; Consumer selected = selector.select(hash); Consumer expected = selectedConsumerBeforeRemoval.get(j); - if (expected != addedConsumer) { + if (!addedConsumers.contains(addedConsumer)) { assertThat(selected.consumerId()).as("validationPoint %d, hash %d", j, hash).isEqualTo(expected.consumerId()); } } From 01d77a2d08eab44033ca9da65d265002e6c278c5 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Sat, 21 Sep 2024 01:32:15 +0300 Subject: [PATCH 11/21] Polish --- ...stentHashingStickyKeyConsumerSelector.java | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index f0b74f6fa4c4a..43dbbd59164fc 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -67,21 +67,20 @@ public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) { * lower selection count is added. */ private static class HashRingEntry { - // This class is used to store the consumer which it was added to the hash ring - // sorting will be by priority, consumer name and usage count of the consumer instance + // This class is used to store the consumer added to the hash ring + // sorting will be by priority, consumer name and active "selection" count of the consumer instance + // so that consumers get evenly distributed record ConsumerEntry(Consumer consumer, MutableInt consumerSelectionCounter) implements Comparable { - private static final Comparator - BASE_CONSUMER_ENTRY_COMPARATOR = Comparator. - comparing(entry -> entry.consumer().getPriorityLevel()).reversed() - .thenComparing(entry -> entry.consumer().consumerName()); - - - private static final Comparator - CONSUMER_ENTRY_COMPARATOR = BASE_CONSUMER_ENTRY_COMPARATOR - // prefer the consumer instance with lowest selection count so that consumers get - // evenly distributed - .thenComparing(ConsumerEntry::consumerSelectionCounter); + private static final Comparator CONSUMER_ENTRY_COMPARATOR = + Comparator. + // priority level is the primary sorting key + comparing(entry -> entry.consumer().getPriorityLevel()).reversed() + // consumer name is the secondary sorting key + .thenComparing(entry -> entry.consumer().consumerName()) + // then prefer the consumer instance with lowest selection count + // so that consumers get evenly distributed + .thenComparing(ConsumerEntry::consumerSelectionCounter); @Override public int compareTo(ConsumerEntry o) { @@ -102,6 +101,7 @@ public void addConsumer(Consumer consumer, MutableInt selectedCounter) { selectConsumer(); } + // use insertion sort so that the consumers don't get unnecessarily switched in removal or addition private void insertConsumer(ConsumerEntry consumerEntry) { for (int i = 0; i < consumers.size(); i++) { if (consumers.get(i).compareTo(consumerEntry) > 0) { From 54d989f5a1279050df343693444b0717240d8d81 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Mon, 23 Sep 2024 07:50:54 +0300 Subject: [PATCH 12/21] Sort if selected consumer was removed --- ...stentHashingStickyKeyConsumerSelector.java | 41 +++++++++---------- 1 file changed, 19 insertions(+), 22 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index 43dbbd59164fc..b7701ce351e0c 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -19,6 +19,7 @@ package org.apache.pulsar.broker.service; import java.util.ArrayList; +import java.util.Collections; import java.util.Comparator; import java.util.LinkedHashMap; import java.util.List; @@ -95,37 +96,33 @@ public HashRingEntry() { this.consumers = new ArrayList<>(); } - public void addConsumer(Consumer consumer, MutableInt selectedCounter) { - ConsumerEntry consumerEntry = new ConsumerEntry(consumer, selectedCounter); - insertConsumer(consumerEntry); - selectConsumer(); + public Consumer getSelectedConsumer() { + return selectedConsumerEntry != null ? selectedConsumerEntry.consumer() : null; } - // use insertion sort so that the consumers don't get unnecessarily switched in removal or addition - private void insertConsumer(ConsumerEntry consumerEntry) { - for (int i = 0; i < consumers.size(); i++) { - if (consumers.get(i).compareTo(consumerEntry) > 0) { - consumers.add(i, consumerEntry); - return; - } - } + public void addConsumer(Consumer consumer, MutableInt selectedCounter) { + ConsumerEntry consumerEntry = new ConsumerEntry(consumer, selectedCounter); consumers.add(consumerEntry); + if (selectedConsumerEntry == null || consumerEntry.compareTo(selectedConsumerEntry) < 0) { + // if the new consumer has a higher priority or lower selection count + // than the currently selected consumer, select the new consumer + changeSelectedConsumerEntry(consumerEntry); + } } public boolean removeConsumer(Consumer consumer) { - boolean removed = consumers.removeIf(consumerEntry -> consumerEntry.consumer().equals(consumer)); - selectConsumer(); + boolean removed = consumers.removeIf(consumerEntry -> consumerEntry.consumer() == consumer); + if (removed && consumer == getSelectedConsumer()) { + // if the selected consumer was removed, a new consumer will be selected. + // The consumers are sorted here to ensure that the consumer with the + // lowest selection count is selected + Collections.sort(consumers); + // select the first consumer in sorting order + changeSelectedConsumerEntry(consumers.isEmpty() ? null : consumers.get(0)); + } return removed; } - public Consumer getSelectedConsumer() { - return selectedConsumerEntry != null ? selectedConsumerEntry.consumer() : null; - } - - private void selectConsumer() { - changeSelectedConsumerEntry(consumers.isEmpty() ? null : consumers.get(0)); - } - private void changeSelectedConsumerEntry(ConsumerEntry newSelectedConsumer) { if (newSelectedConsumer == selectedConsumerEntry) { return; From 079fb85cd7f0d864476c4fc68d984f053ab31f26 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Wed, 25 Sep 2024 19:13:47 +0300 Subject: [PATCH 13/21] Fix NPE when non-existing consumers were removed --- .../ConsistentHashingStickyKeyConsumerSelector.java | 8 +++++--- .../ConsistentHashingStickyKeyConsumerSelectorTest.java | 2 ++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index b7701ce351e0c..abea49a670341 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -180,9 +180,11 @@ public void removeConsumer(Consumer consumer) { for (int i = 0; i < numberOfPoints; i++) { int hash = calculateHashForConsumerAndIndex(consumer, i); hashRing.compute(hash, (k, v) -> { - v.removeConsumer(consumer); - if (v.getSelectedConsumer() == null) { - v = null; + if (v != null) { + v.removeConsumer(consumer); + if (v.getSelectedConsumer() == null) { + v = null; + } } return v; }); diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 82e566dd11f8b..d0f6282617680 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -253,6 +253,8 @@ public void shouldRemoveConsumersFromConsumerKeyHashRanges() { consumers.forEach(selector::removeConsumer); // then there should be no mapping remaining Assert.assertEquals(selector.getConsumerKeyHashRanges().size(), 0); + // when consumers are removed again, should not fail + consumers.forEach(selector::removeConsumer); } @Test From 8ac47e49a70a32a86b7c8a081ba43b8df54b4d52 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Wed, 25 Sep 2024 23:16:01 +0300 Subject: [PATCH 14/21] Remove ordering by name and priority level --- ...stentHashingStickyKeyConsumerSelector.java | 26 +++++++------------ 1 file changed, 9 insertions(+), 17 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index abea49a670341..13c0a3fe9173e 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -48,7 +48,7 @@ public class ConsistentHashingStickyKeyConsumerSelector implements StickyKeyCons // Consistent-Hash ring private final NavigableMap hashRing; // used for distributing consumer instance selections evenly in the hash ring when there - // are multiple instances of consumer with the same consumer name or when there are hash collisions + // are multiple instances of consumer with the same consumer name or when there are other hash collisions private final Map consumerSelectionCounters; private final int numberOfPoints; @@ -61,31 +61,23 @@ public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) { /** * This class is used to store the consumers and the selected consumer for a hash value in the hash ring. - * This attempts to distribute the consumers evenly in the hash ring for consumers with the same - * consumer name and priority level. These entries collide in the hash ring. + * This distributes the consumers evenly in the hash ring for consumers that collide. Consumers with the same + * consumer name in all cases. * The selected consumer is the consumer that is selected to serve the hash value. - * It is not changed unless a consumer is removed or a colliding consumer with higher priority or + * It is not changed unless a consumer is removed or a colliding consumer with a * lower selection count is added. */ private static class HashRingEntry { // This class is used to store the consumer added to the hash ring - // sorting will be by priority, consumer name and active "selection" count of the consumer instance - // so that consumers get evenly distributed + // sorting will be by active "selection" count of the consumer instance so that consumers get evenly distributed record ConsumerEntry(Consumer consumer, MutableInt consumerSelectionCounter) implements Comparable { - private static final Comparator CONSUMER_ENTRY_COMPARATOR = - Comparator. - // priority level is the primary sorting key - comparing(entry -> entry.consumer().getPriorityLevel()).reversed() - // consumer name is the secondary sorting key - .thenComparing(entry -> entry.consumer().consumerName()) - // then prefer the consumer instance with lowest selection count - // so that consumers get evenly distributed - .thenComparing(ConsumerEntry::consumerSelectionCounter); + private static final Comparator COMPARE_BY_SELECTION_COUNT = + Comparator.comparing(ConsumerEntry::consumerSelectionCounter); @Override public int compareTo(ConsumerEntry o) { - return CONSUMER_ENTRY_COMPARATOR.compare(this, o); + return COMPARE_BY_SELECTION_COUNT.compare(this, o); } } @@ -104,7 +96,7 @@ public void addConsumer(Consumer consumer, MutableInt selectedCounter) { ConsumerEntry consumerEntry = new ConsumerEntry(consumer, selectedCounter); consumers.add(consumerEntry); if (selectedConsumerEntry == null || consumerEntry.compareTo(selectedConsumerEntry) < 0) { - // if the new consumer has a higher priority or lower selection count + // if the new consumer lower selection count // than the currently selected consumer, select the new consumer changeSelectedConsumerEntry(consumerEntry); } From 346b344dd84997fa9fb15fe8a7e04a12f3f7700a Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Thu, 26 Sep 2024 01:14:20 +0300 Subject: [PATCH 15/21] Use a weights based solution where the size of the range have an impact --- ...stentHashingStickyKeyConsumerSelector.java | 56 +++++++++++----- ...tHashingStickyKeyConsumerSelectorTest.java | 66 ++++++++++++++----- 2 files changed, 92 insertions(+), 30 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index 13c0a3fe9173e..9735ad87d2659 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -30,7 +30,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; -import org.apache.commons.lang3.mutable.MutableInt; +import org.apache.commons.lang3.mutable.MutableLong; import org.apache.pulsar.client.api.Range; import org.apache.pulsar.common.util.Murmur3_32Hash; @@ -49,13 +49,13 @@ public class ConsistentHashingStickyKeyConsumerSelector implements StickyKeyCons private final NavigableMap hashRing; // used for distributing consumer instance selections evenly in the hash ring when there // are multiple instances of consumer with the same consumer name or when there are other hash collisions - private final Map consumerSelectionCounters; + private final Map consumerSelectionWeightCounters; private final int numberOfPoints; public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) { this.hashRing = new TreeMap<>(); - this.consumerSelectionCounters = new WeakHashMap<>(); + this.consumerSelectionWeightCounters = new WeakHashMap<>(); this.numberOfPoints = numberOfPoints; } @@ -68,12 +68,14 @@ public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) { * lower selection count is added. */ private static class HashRingEntry { + private final int rangeSize; + // This class is used to store the consumer added to the hash ring // sorting will be by active "selection" count of the consumer instance so that consumers get evenly distributed - record ConsumerEntry(Consumer consumer, MutableInt consumerSelectionCounter) + record ConsumerEntry(Consumer consumer, MutableLong selectionWeight) implements Comparable { private static final Comparator COMPARE_BY_SELECTION_COUNT = - Comparator.comparing(ConsumerEntry::consumerSelectionCounter); + Comparator.comparing(ConsumerEntry::selectionWeight); @Override public int compareTo(ConsumerEntry o) { @@ -84,7 +86,8 @@ public int compareTo(ConsumerEntry o) { private final List consumers; ConsumerEntry selectedConsumerEntry; - public HashRingEntry() { + public HashRingEntry(int rangeSize) { + this.rangeSize = rangeSize; this.consumers = new ArrayList<>(); } @@ -92,10 +95,11 @@ public Consumer getSelectedConsumer() { return selectedConsumerEntry != null ? selectedConsumerEntry.consumer() : null; } - public void addConsumer(Consumer consumer, MutableInt selectedCounter) { - ConsumerEntry consumerEntry = new ConsumerEntry(consumer, selectedCounter); + public void addConsumer(Consumer consumer, MutableLong selectionWeight) { + ConsumerEntry consumerEntry = new ConsumerEntry(consumer, selectionWeight); consumers.add(consumerEntry); - if (selectedConsumerEntry == null || consumerEntry.compareTo(selectedConsumerEntry) < 0) { + if (selectedConsumerEntry == null || consumerEntry.selectionWeight().longValue() + rangeSize + <= selectedConsumerEntry.selectionWeight().longValue()) { // if the new consumer lower selection count // than the currently selected consumer, select the new consumer changeSelectedConsumerEntry(consumerEntry); @@ -126,13 +130,13 @@ private void changeSelectedConsumerEntry(ConsumerEntry newSelectedConsumer) { private void beforeChangingSelectedConsumerEntry() { if (selectedConsumerEntry != null) { - selectedConsumerEntry.consumerSelectionCounter.decrement(); + selectedConsumerEntry.selectionWeight.add(-rangeSize); } } private void afterChangingSelectedConsumerEntry() { if (selectedConsumerEntry != null) { - selectedConsumerEntry.consumerSelectionCounter.increment(); + selectedConsumerEntry.selectionWeight.add(rangeSize); } } } @@ -143,11 +147,33 @@ public CompletableFuture addConsumer(Consumer consumer) { try { // Insert multiple points on the hash ring for every consumer // The points are deterministically added based on the hash of the consumer name + + // since there might be hash collisions, we need to ensure that the consumers are evenly distributed + // in the hash ring. This is done by calculating the range size for each consumer so that a + // weight can be calculated for the selection of a consumer in a specific entry + + // first calculate the hash for each point and sort them + List hashKeys = new ArrayList<>(); for (int i = 0; i < numberOfPoints; i++) { int hash = calculateHashForConsumerAndIndex(consumer, i); - HashRingEntry hashRingEntry = hashRing.computeIfAbsent(hash, k -> new HashRingEntry()); + hashKeys.add(hash); + } + Collections.sort(hashKeys); + + // start from the last hash key, mapping it to a negative value so that range size can be calculated + // for the first entry + int maxHashValue = Integer.MAX_VALUE - 1; + int lastHashValue = hashKeys.get(hashKeys.size() - 1); + int remainingSlotSize = maxHashValue - (lastHashValue + 1); + int start = -remainingSlotSize; + MutableLong consumerSelectionWeight = getConsumerSelectionWeight(consumer); + for (int hash : hashKeys) { + // Calculate the range size for the hash + int rangeSize = hash - start; + HashRingEntry hashRingEntry = hashRing.computeIfAbsent(hash, k -> new HashRingEntry(rangeSize)); // Add the consumer to the hash ring entry - hashRingEntry.addConsumer(consumer, getConsumerSelectedCount(consumer)); + hashRingEntry.addConsumer(consumer, consumerSelectionWeight); + start = hash + 1; } return CompletableFuture.completedFuture(null); } finally { @@ -155,8 +181,8 @@ public CompletableFuture addConsumer(Consumer consumer) { } } - private MutableInt getConsumerSelectedCount(Consumer consumer) { - return consumerSelectionCounters.computeIfAbsent(consumer, k -> new MutableInt()); + private MutableLong getConsumerSelectionWeight(Consumer consumer) { + return consumerSelectionWeightCounters.computeIfAbsent(consumer, k -> new MutableLong()); } private static int calculateHashForConsumerAndIndex(Consumer consumer, int index) { diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index d0f6282617680..57ec0584f73f5 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -21,8 +21,10 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; +import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -31,8 +33,10 @@ import java.util.UUID; import java.util.stream.Collectors; import java.util.stream.IntStream; +import org.apache.commons.lang3.mutable.MutableInt; import org.apache.pulsar.broker.service.BrokerServiceException.ConsumerAssignException; import org.apache.pulsar.client.api.Range; +import org.assertj.core.data.Offset; import org.mockito.Mockito; import org.testng.Assert; import org.testng.annotations.Test; @@ -178,27 +182,58 @@ public void testGetConsumerKeyHashRanges() throws BrokerServiceException.Consume } @Test - public void testConsumersGetEvenlyMappedWhenThereAreCollisions() + public void testConsumersGetSufficientlyAccuratelyEvenlyMapped() throws BrokerServiceException.ConsumerAssignException { - ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(5); + ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(200); List consumers = new ArrayList<>(); - for (int i = 0; i < 5; i++) { - // use the same name for all consumers - Consumer consumer = createMockConsumer("consumer", "index " + i, i); + for (int i = 0; i < 20; i++) { + // use the same name for all consumers, use toString to distinguish them + Consumer consumer = createMockConsumer("consumer", String.format("index %02d", i), i); selector.addConsumer(consumer); consumers.add(consumer); } - // check that results are the same when called multiple times - assertThat(selector.getConsumerKeyHashRanges()) - .containsExactlyEntriesOf(selector.getConsumerKeyHashRanges()); + printConsumerRangesStats(selector); - Map> expectedResult = new HashMap<>(); - expectedResult.put(consumers.get(0), List.of(Range.of(216056714, 306176208))); - expectedResult.put(consumers.get(1), List.of(Range.of(365902831, 1240826377))); - expectedResult.put(consumers.get(2), List.of(Range.of(1240826378, 1862045174))); - expectedResult.put(consumers.get(3), List.of(Range.of(306176209, 365902830))); - expectedResult.put(consumers.get(4), List.of(Range.of(0, 216056713), Range.of(1862045175, 2147483646))); - assertThat(selector.getConsumerKeyHashRanges()).containsExactlyInAnyOrderEntriesOf(expectedResult); + int totalSelections = 10000; + + Map consumerSelectionCount = new HashMap<>(); + for (int i = 0; i < totalSelections; i++) { + Consumer selectedConsumer = selector.select(("key " + i).getBytes(StandardCharsets.UTF_8)); + consumerSelectionCount.computeIfAbsent(selectedConsumer, c -> new MutableInt()).increment(); + } + + printSelectionCountStats(consumerSelectionCount); + + int averageCount = totalSelections / consumers.size(); + int allowedVariance = (int) (0.5d * averageCount); + System.out.println("averageCount: " + averageCount + " allowedVariance: " + allowedVariance); + + for (Map.Entry entry : consumerSelectionCount.entrySet()) { + assertThat(entry.getValue().intValue()).describedAs("consumer: %s", entry.getKey()) + .isCloseTo(averageCount, Offset.offset(allowedVariance)); + } + + } + + private static void printSelectionCountStats(Map consumerSelectionCount) { + int totalSelections = consumerSelectionCount.values().stream().mapToInt(MutableInt::intValue).sum(); + consumerSelectionCount.entrySet().stream() + .sorted(Map.Entry.comparingByKey(Comparator.comparing(Consumer::toString))) + .forEach(entry -> System.out.println( + String.format("consumer: %s got selected %d times. ratio: %.2f%%", entry.getKey(), + entry.getValue().intValue(), + ((double) entry.getValue().intValue() / totalSelections) * 100.0d))); + } + + private static void printConsumerRangesStats(ConsistentHashingStickyKeyConsumerSelector selector) { + selector.getConsumerKeyHashRanges().entrySet().stream() + .map(entry -> Map.entry(entry.getKey(), + entry.getValue().stream().mapToInt(r -> r.getEnd() - r.getStart() + 1).sum())) + .sorted(Map.Entry.comparingByKey(Comparator.comparing(Consumer::toString))) + .forEach(entry -> System.out.println( + String.format("consumer: %s total ranges size: %d ratio: %.2f%%", entry.getKey(), + entry.getValue(), + ((double) entry.getValue() / (Integer.MAX_VALUE - 1)) * 100.0d))); } private static Consumer createMockConsumer(String consumerName, String toString, long id) { @@ -417,4 +452,5 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsAdded() { } } } + } From 586d7453250ed4f12f748af68b232f4f4c91ecb2 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Thu, 26 Sep 2024 09:01:33 +0300 Subject: [PATCH 16/21] Use approach where each duplicate consumer name will have different name index which is sufficiently consistent --- ...stentHashingStickyKeyConsumerSelector.java | 219 ++++++++---------- ...tHashingStickyKeyConsumerSelectorTest.java | 31 +-- 2 files changed, 112 insertions(+), 138 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index 9735ad87d2659..8068995af6339 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -19,20 +19,19 @@ package org.apache.pulsar.broker.service; import java.util.ArrayList; -import java.util.Collections; -import java.util.Comparator; -import java.util.LinkedHashMap; +import java.util.HashMap; +import java.util.IdentityHashMap; import java.util.List; import java.util.Map; import java.util.NavigableMap; import java.util.TreeMap; -import java.util.WeakHashMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; -import org.apache.commons.lang3.mutable.MutableLong; +import org.apache.commons.lang3.mutable.MutableInt; import org.apache.pulsar.client.api.Range; import org.apache.pulsar.common.util.Murmur3_32Hash; +import org.roaringbitmap.RoaringBitmap; /** * This is a consumer selector based fixed hash range. @@ -46,98 +45,95 @@ public class ConsistentHashingStickyKeyConsumerSelector implements StickyKeyCons private final ReadWriteLock rwLock = new ReentrantReadWriteLock(); // Consistent-Hash ring - private final NavigableMap hashRing; - // used for distributing consumer instance selections evenly in the hash ring when there - // are multiple instances of consumer with the same consumer name or when there are other hash collisions - private final Map consumerSelectionWeightCounters; + private final NavigableMap hashRing; + private final ConsumerNameIndexTracker consumerNameIndexTracker = new ConsumerNameIndexTracker(); + private final int numberOfPoints; public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) { this.hashRing = new TreeMap<>(); - this.consumerSelectionWeightCounters = new WeakHashMap<>(); this.numberOfPoints = numberOfPoints; } - /** - * This class is used to store the consumers and the selected consumer for a hash value in the hash ring. - * This distributes the consumers evenly in the hash ring for consumers that collide. Consumers with the same - * consumer name in all cases. - * The selected consumer is the consumer that is selected to serve the hash value. - * It is not changed unless a consumer is removed or a colliding consumer with a - * lower selection count is added. - */ - private static class HashRingEntry { - private final int rangeSize; + private static class ConsumerIdentityWrapper { + final Consumer consumer; - // This class is used to store the consumer added to the hash ring - // sorting will be by active "selection" count of the consumer instance so that consumers get evenly distributed - record ConsumerEntry(Consumer consumer, MutableLong selectionWeight) - implements Comparable { - private static final Comparator COMPARE_BY_SELECTION_COUNT = - Comparator.comparing(ConsumerEntry::selectionWeight); + public ConsumerIdentityWrapper(Consumer consumer) { + this.consumer = consumer; + } - @Override - public int compareTo(ConsumerEntry o) { - return COMPARE_BY_SELECTION_COUNT.compare(this, o); + @Override + public boolean equals(Object obj) { + if (obj instanceof ConsumerIdentityWrapper) { + ConsumerIdentityWrapper other = (ConsumerIdentityWrapper) obj; + return consumer == other.consumer; } + return false; + } + + @Override + public int hashCode() { + return consumer.hashCode(); } - private final List consumers; - ConsumerEntry selectedConsumerEntry; + @Override + public String toString() { + return consumer.toString(); + } + } - public HashRingEntry(int rangeSize) { - this.rangeSize = rangeSize; - this.consumers = new ArrayList<>(); + private static class ConsumerNameIndexTracker { + private final Map consumerNameCounters = new HashMap<>(); + private final Map consumerEntries = new HashMap<>(); + record ConsumerEntry(String consumerName, int nameIndex, MutableInt refCount) { } - public Consumer getSelectedConsumer() { - return selectedConsumerEntry != null ? selectedConsumerEntry.consumer() : null; + private RoaringBitmap getConsumerNameIndexBitmap(String consumerName) { + return consumerNameCounters.computeIfAbsent(consumerName, + k -> new RoaringBitmap()); } - public void addConsumer(Consumer consumer, MutableLong selectionWeight) { - ConsumerEntry consumerEntry = new ConsumerEntry(consumer, selectionWeight); - consumers.add(consumerEntry); - if (selectedConsumerEntry == null || consumerEntry.selectionWeight().longValue() + rangeSize - <= selectedConsumerEntry.selectionWeight().longValue()) { - // if the new consumer lower selection count - // than the currently selected consumer, select the new consumer - changeSelectedConsumerEntry(consumerEntry); + private int allocateConsumerNameIndex(String consumerName) { + RoaringBitmap bitmap = getConsumerNameIndexBitmap(consumerName); + // find the first index that is not set, if there is no such index, add a new one + int index = (int) bitmap.nextAbsentValue(0); + if (index == -1) { + index = bitmap.getCardinality(); } + bitmap.add(index); + return index; } - public boolean removeConsumer(Consumer consumer) { - boolean removed = consumers.removeIf(consumerEntry -> consumerEntry.consumer() == consumer); - if (removed && consumer == getSelectedConsumer()) { - // if the selected consumer was removed, a new consumer will be selected. - // The consumers are sorted here to ensure that the consumer with the - // lowest selection count is selected - Collections.sort(consumers); - // select the first consumer in sorting order - changeSelectedConsumerEntry(consumers.isEmpty() ? null : consumers.get(0)); + private void deallocateConsumerNameIndex(String consumerName, int index) { + RoaringBitmap bitmap = getConsumerNameIndexBitmap(consumerName); + bitmap.remove(index); + if (bitmap.isEmpty()) { + consumerNameCounters.remove(consumerName); } - return removed; } - private void changeSelectedConsumerEntry(ConsumerEntry newSelectedConsumer) { - if (newSelectedConsumer == selectedConsumerEntry) { - return; + public void removeHashRingReference(ConsumerIdentityWrapper removed) { + ConsumerEntry consumerEntry = consumerEntries.get(removed); + int refCount = consumerEntry.refCount.decrementAndGet(); + if (refCount == 0) { + deallocateConsumerNameIndex(consumerEntry.consumerName, consumerEntry.nameIndex); + consumerEntries.remove(removed, consumerEntry); } - beforeChangingSelectedConsumerEntry(); - selectedConsumerEntry = newSelectedConsumer; - afterChangingSelectedConsumerEntry(); } - private void beforeChangingSelectedConsumerEntry() { - if (selectedConsumerEntry != null) { - selectedConsumerEntry.selectionWeight.add(-rangeSize); - } + public int addHashRingReference(ConsumerIdentityWrapper wrapper) { + String consumerName = wrapper.consumer.consumerName(); + ConsumerEntry entry = consumerEntries.computeIfAbsent(wrapper, + k -> new ConsumerEntry(consumerName, allocateConsumerNameIndex(consumerName), + new MutableInt(0))); + entry.refCount.increment(); + return entry.nameIndex; } - private void afterChangingSelectedConsumerEntry() { - if (selectedConsumerEntry != null) { - selectedConsumerEntry.selectionWeight.add(rangeSize); - } + public int getTrackedConsumerNameIndex(ConsumerIdentityWrapper wrapper) { + ConsumerEntry consumerEntry = consumerEntries.get(wrapper); + return consumerEntry != null ? consumerEntry.nameIndex : -1; } } @@ -145,35 +141,16 @@ private void afterChangingSelectedConsumerEntry() { public CompletableFuture addConsumer(Consumer consumer) { rwLock.writeLock().lock(); try { + ConsumerIdentityWrapper consumerIdentityWrapper = new ConsumerIdentityWrapper(consumer); // Insert multiple points on the hash ring for every consumer // The points are deterministically added based on the hash of the consumer name - - // since there might be hash collisions, we need to ensure that the consumers are evenly distributed - // in the hash ring. This is done by calculating the range size for each consumer so that a - // weight can be calculated for the selection of a consumer in a specific entry - - // first calculate the hash for each point and sort them - List hashKeys = new ArrayList<>(); for (int i = 0; i < numberOfPoints; i++) { - int hash = calculateHashForConsumerAndIndex(consumer, i); - hashKeys.add(hash); - } - Collections.sort(hashKeys); - - // start from the last hash key, mapping it to a negative value so that range size can be calculated - // for the first entry - int maxHashValue = Integer.MAX_VALUE - 1; - int lastHashValue = hashKeys.get(hashKeys.size() - 1); - int remainingSlotSize = maxHashValue - (lastHashValue + 1); - int start = -remainingSlotSize; - MutableLong consumerSelectionWeight = getConsumerSelectionWeight(consumer); - for (int hash : hashKeys) { - // Calculate the range size for the hash - int rangeSize = hash - start; - HashRingEntry hashRingEntry = hashRing.computeIfAbsent(hash, k -> new HashRingEntry(rangeSize)); - // Add the consumer to the hash ring entry - hashRingEntry.addConsumer(consumer, consumerSelectionWeight); - start = hash + 1; + int consumerNameIndex = consumerNameIndexTracker.addHashRingReference(consumerIdentityWrapper); + int hash = calculateHashForConsumerAndIndex(consumer, consumerNameIndex, i); + ConsumerIdentityWrapper removed = hashRing.put(hash, consumerIdentityWrapper); + if (removed != null) { + consumerNameIndexTracker.removeHashRingReference(removed); + } } return CompletableFuture.completedFuture(null); } finally { @@ -181,12 +158,8 @@ public CompletableFuture addConsumer(Consumer consumer) { } } - private MutableLong getConsumerSelectionWeight(Consumer consumer) { - return consumerSelectionWeightCounters.computeIfAbsent(consumer, k -> new MutableLong()); - } - - private static int calculateHashForConsumerAndIndex(Consumer consumer, int index) { - String key = consumer.consumerName() + KEY_SEPARATOR + index; + private static int calculateHashForConsumerAndIndex(Consumer consumer, int consumerNameIndex, int index) { + String key = consumer.consumerName() + KEY_SEPARATOR + consumerNameIndex + KEY_SEPARATOR + index; return Murmur3_32Hash.getInstance().makeHash(key.getBytes()); } @@ -194,18 +167,16 @@ private static int calculateHashForConsumerAndIndex(Consumer consumer, int index public void removeConsumer(Consumer consumer) { rwLock.writeLock().lock(); try { - // Remove all the points that were added for this consumer - for (int i = 0; i < numberOfPoints; i++) { - int hash = calculateHashForConsumerAndIndex(consumer, i); - hashRing.compute(hash, (k, v) -> { - if (v != null) { - v.removeConsumer(consumer); - if (v.getSelectedConsumer() == null) { - v = null; - } + ConsumerIdentityWrapper consumerIdentityWrapper = new ConsumerIdentityWrapper(consumer); + int consumerNameIndex = consumerNameIndexTracker.getTrackedConsumerNameIndex(consumerIdentityWrapper); + if (consumerNameIndex > -1) { + // Remove all the points that were added for this consumer + for (int i = 0; i < numberOfPoints; i++) { + int hash = calculateHashForConsumerAndIndex(consumer, consumerNameIndex, i); + if (hashRing.remove(hash, consumerIdentityWrapper)) { + consumerNameIndexTracker.removeHashRingReference(consumerIdentityWrapper); } - return v; - }); + } } } finally { rwLock.writeLock().unlock(); @@ -219,14 +190,13 @@ public Consumer select(int hash) { if (hashRing.isEmpty()) { return null; } - HashRingEntry hashRingEntry; - Map.Entry ceilingEntry = hashRing.ceilingEntry(hash); + + Map.Entry ceilingEntry = hashRing.ceilingEntry(hash); if (ceilingEntry != null) { - hashRingEntry = ceilingEntry.getValue(); + return ceilingEntry.getValue().consumer; } else { - hashRingEntry = hashRing.firstEntry().getValue(); + return hashRing.firstEntry().getValue().consumer; } - return hashRingEntry.getSelectedConsumer(); } finally { rwLock.readLock().unlock(); } @@ -234,7 +204,7 @@ public Consumer select(int hash) { @Override public Map> getConsumerKeyHashRanges() { - Map> result = new LinkedHashMap<>(); + Map> result = new IdentityHashMap<>(); rwLock.readLock().lock(); try { if (hashRing.isEmpty()) { @@ -242,17 +212,16 @@ public Map> getConsumerKeyHashRanges() { } int start = 0; int lastKey = 0; - for (Map.Entry entry: hashRing.entrySet()) { - Consumer consumer = entry.getValue().getSelectedConsumer(); + for (Map.Entry entry: hashRing.entrySet()) { + Consumer consumer = entry.getValue().consumer; result.computeIfAbsent(consumer, key -> new ArrayList<>()) - .add(Range.of(start, entry.getKey())); - lastKey = entry.getKey(); - start = lastKey + 1; + .add(Range.of(start, entry.getKey())); + lastKey = entry.getKey() + 1; + start = lastKey; } // Handle wrap-around - HashRingEntry firstHashRingEntry = hashRing.firstEntry().getValue(); - Consumer firstSelectedConsumer = firstHashRingEntry.getSelectedConsumer(); - List ranges = result.get(firstSelectedConsumer); + Consumer firstConsumer = hashRing.firstEntry().getValue().consumer; + List ranges = result.get(firstConsumer); if (lastKey != Integer.MAX_VALUE - 1) { ranges.add(Range.of(lastKey + 1, Integer.MAX_VALUE - 1)); } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 57ec0584f73f5..0fb1250b808a7 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -165,19 +165,22 @@ public void testGetConsumerKeyHashRanges() throws BrokerServiceException.Consume .containsExactlyEntriesOf(selector.getConsumerKeyHashRanges()); Map> expectedResult = new HashMap<>(); + assertThat(consumers.get(0).consumerName()).isEqualTo("consumer1"); expectedResult.put(consumers.get(0), Arrays.asList( - Range.of(119056335, 242013991), - Range.of(722195657, 1656011842), - Range.of(1707482098, 1914695766))); + Range.of(95615213, 440020355), + Range.of(440020356, 455987436), + Range.of(1189794593, 1264144431))); + assertThat(consumers.get(1).consumerName()).isEqualTo("consumer2"); expectedResult.put(consumers.get(1), Arrays.asList( - Range.of(0, 90164503), - Range.of(90164504, 119056334), - Range.of(382436668, 722195656), - Range.of(1914695767, 2147483646))); + Range.of(939655188, 1189794592), + Range.of(1314727625, 1977451233), + Range.of(1977451234, 2016237253))); + assertThat(consumers.get(2).consumerName()).isEqualTo("consumer3"); expectedResult.put(consumers.get(2), Arrays.asList( - Range.of(242013992, 242377547), - Range.of(242377548, 382436667), - Range.of(1656011843, 1707482097))); + Range.of(0, 95615212), + Range.of(455987437, 939655187), + Range.of(1264144432, 1314727624), + Range.of(2016237255, 2147483646))); assertThat(selector.getConsumerKeyHashRanges()).containsExactlyInAnyOrderEntriesOf(expectedResult); } @@ -205,7 +208,7 @@ public void testConsumersGetSufficientlyAccuratelyEvenlyMapped() printSelectionCountStats(consumerSelectionCount); int averageCount = totalSelections / consumers.size(); - int allowedVariance = (int) (0.5d * averageCount); + int allowedVariance = (int) (0.2d * averageCount); System.out.println("averageCount: " + averageCount + " allowedVariance: " + allowedVariance); for (Map.Entry entry : consumerSelectionCount.entrySet()) { @@ -213,6 +216,8 @@ public void testConsumersGetSufficientlyAccuratelyEvenlyMapped() .isCloseTo(averageCount, Offset.offset(allowedVariance)); } + consumers.forEach(selector::removeConsumer); + assertThat(selector.getConsumerKeyHashRanges()).isEmpty(); } private static void printSelectionCountStats(Map consumerSelectionCount) { @@ -348,7 +353,7 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsRemovedCheckHashRan } Map> expected = selector.getConsumerKeyHashRanges(); - assertThat(selector.getConsumerKeyHashRanges()).as("sanity check").isEqualTo(expected); + assertThat(selector.getConsumerKeyHashRanges()).as("sanity check").containsExactlyInAnyOrderEntriesOf(expected); System.out.println(expected); for (Consumer removedConsumer : consumers) { @@ -384,7 +389,7 @@ public void testShouldNotChangeSelectedConsumerUnnecessarilyWhenConsumerIsAddedC } Map> expected = selector.getConsumerKeyHashRanges(); - assertThat(selector.getConsumerKeyHashRanges()).as("sanity check").isEqualTo(expected); + assertThat(selector.getConsumerKeyHashRanges()).as("sanity check").containsExactlyInAnyOrderEntriesOf(expected); for (int i = numOfInitialConsumers; i < numOfInitialConsumers * 2; i++) { final Consumer addedConsumer = createMockConsumer(consumerName, "index " + i, i); From 45d56fde0eaf619e3b84acc74ec12fd875e6fb8e Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Thu, 26 Sep 2024 11:39:52 +0300 Subject: [PATCH 17/21] Add test to validate that range didn't change --- ...tHashingStickyKeyConsumerSelectorTest.java | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 0fb1250b808a7..1d072dd472148 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -458,4 +458,26 @@ public void testShouldNotChangeSelectedConsumerWhenConsumerIsAdded() { } } + @Test + public void testShouldNotChangeMappingWhenConsumerLeavesAndRejoins() { + final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(100); + final String consumerName = "consumer"; + final int numOfInitialConsumers = 25; + List consumers = new ArrayList<>(); + for (int i = 0; i < numOfInitialConsumers; i++) { + final Consumer consumer = createMockConsumer(consumerName, "index " + i, i); + consumers.add(consumer); + selector.addConsumer(consumer); + } + + Map> expected = selector.getConsumerKeyHashRanges(); + assertThat(selector.getConsumerKeyHashRanges()).as("sanity check").containsExactlyInAnyOrderEntriesOf(expected); + + selector.removeConsumer(consumers.get(0)); + selector.removeConsumer(consumers.get(numOfInitialConsumers / 2)); + selector.addConsumer(consumers.get(0)); + selector.addConsumer(consumers.get(numOfInitialConsumers / 2)); + + assertThat(selector.getConsumerKeyHashRanges()).as("ranges shouldn't change").containsExactlyInAnyOrderEntriesOf(expected); + } } From c2f1c592e6c81bee671eaabe905b2a6a925eaeab Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Thu, 26 Sep 2024 15:41:44 +0300 Subject: [PATCH 18/21] Increase number of points to distribute more evenly --- .../ConsistentHashingStickyKeyConsumerSelectorTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 1d072dd472148..01e15895986a2 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -47,7 +47,7 @@ public class ConsistentHashingStickyKeyConsumerSelectorTest { @Test public void testConsumerSelect() throws ConsumerAssignException { - ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(100); + ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(200); String key1 = "anyKey"; Assert.assertNull(selector.select(key1.getBytes())); @@ -61,7 +61,7 @@ public void testConsumerSelect() throws ConsumerAssignException { selector.addConsumer(consumer2); final int N = 1000; - final double PERCENT_ERROR = 0.25; // 25 % + final double PERCENT_ERROR = 0.20; // 20 % Map selectionMap = new HashMap<>(); for (int i = 0; i < N; i++) { From 83e24795cebc5508a89e73df5e5d9412d90545b0 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Tue, 1 Oct 2024 07:24:32 +0300 Subject: [PATCH 19/21] Move to separate files --- ...stentHashingStickyKeyConsumerSelector.java | 84 ------------------- .../service/ConsumerIdentityWrapper.java | 51 +++++++++++ .../service/ConsumerNameIndexTracker.java | 79 +++++++++++++++++ 3 files changed, 130 insertions(+), 84 deletions(-) create mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapper.java create mode 100644 pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerNameIndexTracker.java diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index 8068995af6339..c04ad1f41f84a 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -19,7 +19,6 @@ package org.apache.pulsar.broker.service; import java.util.ArrayList; -import java.util.HashMap; import java.util.IdentityHashMap; import java.util.List; import java.util.Map; @@ -28,10 +27,8 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; -import org.apache.commons.lang3.mutable.MutableInt; import org.apache.pulsar.client.api.Range; import org.apache.pulsar.common.util.Murmur3_32Hash; -import org.roaringbitmap.RoaringBitmap; /** * This is a consumer selector based fixed hash range. @@ -56,87 +53,6 @@ public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) { this.numberOfPoints = numberOfPoints; } - private static class ConsumerIdentityWrapper { - final Consumer consumer; - - public ConsumerIdentityWrapper(Consumer consumer) { - this.consumer = consumer; - } - - @Override - public boolean equals(Object obj) { - if (obj instanceof ConsumerIdentityWrapper) { - ConsumerIdentityWrapper other = (ConsumerIdentityWrapper) obj; - return consumer == other.consumer; - } - return false; - } - - @Override - public int hashCode() { - return consumer.hashCode(); - } - - @Override - public String toString() { - return consumer.toString(); - } - } - - private static class ConsumerNameIndexTracker { - private final Map consumerNameCounters = new HashMap<>(); - private final Map consumerEntries = new HashMap<>(); - record ConsumerEntry(String consumerName, int nameIndex, MutableInt refCount) { - } - - private RoaringBitmap getConsumerNameIndexBitmap(String consumerName) { - return consumerNameCounters.computeIfAbsent(consumerName, - k -> new RoaringBitmap()); - } - - private int allocateConsumerNameIndex(String consumerName) { - RoaringBitmap bitmap = getConsumerNameIndexBitmap(consumerName); - // find the first index that is not set, if there is no such index, add a new one - int index = (int) bitmap.nextAbsentValue(0); - if (index == -1) { - index = bitmap.getCardinality(); - } - bitmap.add(index); - return index; - } - - private void deallocateConsumerNameIndex(String consumerName, int index) { - RoaringBitmap bitmap = getConsumerNameIndexBitmap(consumerName); - bitmap.remove(index); - if (bitmap.isEmpty()) { - consumerNameCounters.remove(consumerName); - } - } - - public void removeHashRingReference(ConsumerIdentityWrapper removed) { - ConsumerEntry consumerEntry = consumerEntries.get(removed); - int refCount = consumerEntry.refCount.decrementAndGet(); - if (refCount == 0) { - deallocateConsumerNameIndex(consumerEntry.consumerName, consumerEntry.nameIndex); - consumerEntries.remove(removed, consumerEntry); - } - } - - public int addHashRingReference(ConsumerIdentityWrapper wrapper) { - String consumerName = wrapper.consumer.consumerName(); - ConsumerEntry entry = consumerEntries.computeIfAbsent(wrapper, - k -> new ConsumerEntry(consumerName, allocateConsumerNameIndex(consumerName), - new MutableInt(0))); - entry.refCount.increment(); - return entry.nameIndex; - } - - public int getTrackedConsumerNameIndex(ConsumerIdentityWrapper wrapper) { - ConsumerEntry consumerEntry = consumerEntries.get(wrapper); - return consumerEntry != null ? consumerEntry.nameIndex : -1; - } - } - @Override public CompletableFuture addConsumer(Consumer consumer) { rwLock.writeLock().lock(); diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapper.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapper.java new file mode 100644 index 0000000000000..fe71c6a510410 --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapper.java @@ -0,0 +1,51 @@ +/* + * 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; + +/** + * A wrapper class for a Consumer instance that provides custom implementations + * of equals and hashCode methods. The equals method returns true if and only if + * the compared instance is the same instance. + */ +class ConsumerIdentityWrapper { + final Consumer consumer; + + public ConsumerIdentityWrapper(Consumer consumer) { + this.consumer = consumer; + } + + @Override + public boolean equals(Object obj) { + if (obj instanceof ConsumerIdentityWrapper) { + ConsumerIdentityWrapper other = (ConsumerIdentityWrapper) obj; + return consumer == other.consumer; + } + return false; + } + + @Override + public int hashCode() { + return consumer.hashCode(); + } + + @Override + public String toString() { + return consumer.toString(); + } +} diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerNameIndexTracker.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerNameIndexTracker.java new file mode 100644 index 0000000000000..ff424ff175eac --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerNameIndexTracker.java @@ -0,0 +1,79 @@ +/* + * 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.HashMap; +import java.util.Map; +import org.apache.commons.lang3.mutable.MutableInt; +import org.roaringbitmap.RoaringBitmap; + +class ConsumerNameIndexTracker { + private final Map consumerNameCounters = new HashMap<>(); + private final Map consumerEntries = new HashMap<>(); + + record ConsumerEntry(String consumerName, int nameIndex, MutableInt refCount) { + } + + private RoaringBitmap getConsumerNameIndexBitmap(String consumerName) { + return consumerNameCounters.computeIfAbsent(consumerName, + k -> new RoaringBitmap()); + } + + private int allocateConsumerNameIndex(String consumerName) { + RoaringBitmap bitmap = getConsumerNameIndexBitmap(consumerName); + // find the first index that is not set, if there is no such index, add a new one + int index = (int) bitmap.nextAbsentValue(0); + if (index == -1) { + index = bitmap.getCardinality(); + } + bitmap.add(index); + return index; + } + + private void deallocateConsumerNameIndex(String consumerName, int index) { + RoaringBitmap bitmap = getConsumerNameIndexBitmap(consumerName); + bitmap.remove(index); + if (bitmap.isEmpty()) { + consumerNameCounters.remove(consumerName); + } + } + + public void removeHashRingReference(ConsumerIdentityWrapper removed) { + ConsumerEntry consumerEntry = consumerEntries.get(removed); + int refCount = consumerEntry.refCount.decrementAndGet(); + if (refCount == 0) { + deallocateConsumerNameIndex(consumerEntry.consumerName, consumerEntry.nameIndex); + consumerEntries.remove(removed, consumerEntry); + } + } + + public int addHashRingReference(ConsumerIdentityWrapper wrapper) { + String consumerName = wrapper.consumer.consumerName(); + ConsumerEntry entry = consumerEntries.computeIfAbsent(wrapper, + k -> new ConsumerEntry(consumerName, allocateConsumerNameIndex(consumerName), + new MutableInt(0))); + entry.refCount.increment(); + return entry.nameIndex; + } + + public int getTrackedConsumerNameIndex(ConsumerIdentityWrapper wrapper) { + ConsumerEntry consumerEntry = consumerEntries.get(wrapper); + return consumerEntry != null ? consumerEntry.nameIndex : -1; + } +} From 1a11b4eb762db77ce233420e8559398cdba04488 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Tue, 1 Oct 2024 08:26:41 +0300 Subject: [PATCH 20/21] Refactor --- ...stentHashingStickyKeyConsumerSelector.java | 38 +++-- .../service/ConsumerIdentityWrapper.java | 21 ++- .../service/ConsumerNameIndexTracker.java | 111 ++++++++++--- ...tHashingStickyKeyConsumerSelectorTest.java | 17 +- .../service/ConsumerIdentityWrapperTest.java | 68 ++++++++ .../service/ConsumerNameIndexTrackerTest.java | 157 ++++++++++++++++++ .../org/apache/pulsar/client/api/Range.java | 11 +- 7 files changed, 381 insertions(+), 42 deletions(-) create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapperTest.java create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerNameIndexTrackerTest.java diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java index c04ad1f41f84a..1ae9a6ff96b7d 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelector.java @@ -43,9 +43,9 @@ public class ConsistentHashingStickyKeyConsumerSelector implements StickyKeyCons // Consistent-Hash ring private final NavigableMap hashRing; + // Tracks the used consumer name indexes for each consumer name private final ConsumerNameIndexTracker consumerNameIndexTracker = new ConsumerNameIndexTracker(); - private final int numberOfPoints; public ConsistentHashingStickyKeyConsumerSelector(int numberOfPoints) { @@ -61,11 +61,15 @@ public CompletableFuture addConsumer(Consumer consumer) { // Insert multiple points on the hash ring for every consumer // The points are deterministically added based on the hash of the consumer name for (int i = 0; i < numberOfPoints; i++) { - int consumerNameIndex = consumerNameIndexTracker.addHashRingReference(consumerIdentityWrapper); + int consumerNameIndex = + consumerNameIndexTracker.increaseConsumerRefCountAndReturnIndex(consumerIdentityWrapper); int hash = calculateHashForConsumerAndIndex(consumer, consumerNameIndex, i); + // When there's a collision, the new consumer will replace the old one. + // This is a rare case, and it is acceptable to replace the old consumer since there + // are multiple points for each consumer. This won't affect the overall distribution significantly. ConsumerIdentityWrapper removed = hashRing.put(hash, consumerIdentityWrapper); if (removed != null) { - consumerNameIndexTracker.removeHashRingReference(removed); + consumerNameIndexTracker.decreaseConsumerRefCount(removed); } } return CompletableFuture.completedFuture(null); @@ -74,8 +78,19 @@ public CompletableFuture addConsumer(Consumer consumer) { } } - private static int calculateHashForConsumerAndIndex(Consumer consumer, int consumerNameIndex, int index) { - String key = consumer.consumerName() + KEY_SEPARATOR + consumerNameIndex + KEY_SEPARATOR + index; + /** + * Calculate the hash for a consumer and hash ring point. + * The hash is calculated based on the consumer name, consumer name index, and hash ring point index. + * The resulting hash is used as the key to insert the consumer into the hash ring. + * + * @param consumer the consumer + * @param consumerNameIndex the index of the consumer name + * @param hashRingPointIndex the index of the hash ring point + * @return the hash value + */ + private static int calculateHashForConsumerAndIndex(Consumer consumer, int consumerNameIndex, + int hashRingPointIndex) { + String key = consumer.consumerName() + KEY_SEPARATOR + consumerNameIndex + KEY_SEPARATOR + hashRingPointIndex; return Murmur3_32Hash.getInstance().makeHash(key.getBytes()); } @@ -84,13 +99,13 @@ public void removeConsumer(Consumer consumer) { rwLock.writeLock().lock(); try { ConsumerIdentityWrapper consumerIdentityWrapper = new ConsumerIdentityWrapper(consumer); - int consumerNameIndex = consumerNameIndexTracker.getTrackedConsumerNameIndex(consumerIdentityWrapper); + int consumerNameIndex = consumerNameIndexTracker.getTrackedIndex(consumerIdentityWrapper); if (consumerNameIndex > -1) { // Remove all the points that were added for this consumer for (int i = 0; i < numberOfPoints; i++) { int hash = calculateHashForConsumerAndIndex(consumer, consumerNameIndex, i); if (hashRing.remove(hash, consumerIdentityWrapper)) { - consumerNameIndexTracker.removeHashRingReference(consumerIdentityWrapper); + consumerNameIndexTracker.decreaseConsumerRefCount(consumerIdentityWrapper); } } } @@ -106,11 +121,11 @@ public Consumer select(int hash) { if (hashRing.isEmpty()) { return null; } - Map.Entry ceilingEntry = hashRing.ceilingEntry(hash); if (ceilingEntry != null) { return ceilingEntry.getValue().consumer; } else { + // Handle wrap-around in the hash ring, return the first consumer return hashRing.firstEntry().getValue().consumer; } } finally { @@ -132,10 +147,11 @@ public Map> getConsumerKeyHashRanges() { Consumer consumer = entry.getValue().consumer; result.computeIfAbsent(consumer, key -> new ArrayList<>()) .add(Range.of(start, entry.getKey())); - lastKey = entry.getKey() + 1; - start = lastKey; + lastKey = entry.getKey(); + start = lastKey + 1; } - // Handle wrap-around + // Handle wrap-around in the hash ring, the first consumer will also contain the range from the last key + // to the maximum value of the hash range Consumer firstConsumer = hashRing.firstEntry().getValue().consumer; List ranges = result.get(firstConsumer); if (lastKey != Integer.MAX_VALUE - 1) { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapper.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapper.java index fe71c6a510410..2aae1d9b0622e 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapper.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapper.java @@ -22,6 +22,10 @@ * A wrapper class for a Consumer instance that provides custom implementations * of equals and hashCode methods. The equals method returns true if and only if * the compared instance is the same instance. + * + *

The reason for this class is the custom implementation of {@link Consumer#equals(Object)}. + * Using this wrapper class will be useful in use cases where it's necessary to match a key + * in a map by instance or a value in a set by instance.

*/ class ConsumerIdentityWrapper { final Consumer consumer; @@ -30,6 +34,15 @@ public ConsumerIdentityWrapper(Consumer consumer) { this.consumer = consumer; } + /** + * Compares this wrapper to the specified object. The result is true if and only if + * the argument is not null and is a ConsumerIdentityWrapper object that wraps + * the same Consumer instance. + * + * @param obj the object to compare this ConsumerIdentityWrapper against + * @return true if the given object represents a ConsumerIdentityWrapper + * equivalent to this wrapper, false otherwise + */ @Override public boolean equals(Object obj) { if (obj instanceof ConsumerIdentityWrapper) { @@ -39,6 +52,12 @@ public boolean equals(Object obj) { return false; } + /** + * Returns a hash code for this wrapper. The hash code is computed based on + * the wrapped Consumer instance. + * + * @return a hash code value for this object + */ @Override public int hashCode() { return consumer.hashCode(); @@ -48,4 +67,4 @@ public int hashCode() { public String toString() { return consumer.toString(); } -} +} \ No newline at end of file diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerNameIndexTracker.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerNameIndexTracker.java index ff424ff175eac..1f93313ab1b71 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerNameIndexTracker.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ConsumerNameIndexTracker.java @@ -20,41 +20,90 @@ import java.util.HashMap; import java.util.Map; +import javax.annotation.concurrent.NotThreadSafe; import org.apache.commons.lang3.mutable.MutableInt; import org.roaringbitmap.RoaringBitmap; +/** + * Tracks the used consumer name indexes for each consumer name. + * This is used by {@link ConsistentHashingStickyKeyConsumerSelector} to get a unique "consumer name index" + * for each consumer name. It is useful when there are multiple consumers with the same name, but they are + * different consumers. The purpose of the index is to prevent collisions in the hash ring. + * + * The consumer name index serves as an additional key for the hash ring assignment. The logic keeps track of + * used "index slots" for each consumer name and assigns the first unused index when a new consumer is added. + * This approach minimizes hash collisions due to using the same consumer name. + * + * An added benefit of this tracking approach is that a consumer that leaves and then rejoins immediately will get the + * same index and therefore the same assignments in the hash ring. This improves stability since the hash assignment + * changes are minimized over time, although a better solution would be to avoid reusing the same consumer name + * in the first place. + * + * When a consumer is removed, the index is deallocated. RoaringBitmap is used to keep track of the used indexes. + * The data structure to track a consumer name is removed when the reference count of the consumer name is zero. + * + * This class is not thread-safe and should be used in a synchronized context in the caller. + */ +@NotThreadSafe class ConsumerNameIndexTracker { - private final Map consumerNameCounters = new HashMap<>(); + // tracks the used index slots for each consumer name + private final Map consumerNameIndexSlotsMap = new HashMap<>(); + // tracks the active consumer entries private final Map consumerEntries = new HashMap<>(); + // Represents a consumer entry in the tracker, including the consumer name, index, and reference count. record ConsumerEntry(String consumerName, int nameIndex, MutableInt refCount) { } - private RoaringBitmap getConsumerNameIndexBitmap(String consumerName) { - return consumerNameCounters.computeIfAbsent(consumerName, - k -> new RoaringBitmap()); + /* + * Tracks the used indexes for a consumer name using a RoaringBitmap. + * A specific index slot is used when the bit is set. + * When all bits are cleared, the customer name can be removed from tracking. + */ + static class ConsumerNameIndexSlots { + private RoaringBitmap indexSlots = new RoaringBitmap(); + + public int allocateIndexSlot() { + // find the first index that is not set, if there is no such index, add a new one + int index = (int) indexSlots.nextAbsentValue(0); + if (index == -1) { + index = indexSlots.getCardinality(); + } + indexSlots.add(index); + return index; + } + + public boolean deallocateIndexSlot(int index) { + indexSlots.remove(index); + return indexSlots.isEmpty(); + } + } + + /* + * Adds a reference to the consumer and returns the index assigned to this consumer. + */ + public int increaseConsumerRefCountAndReturnIndex(ConsumerIdentityWrapper wrapper) { + ConsumerEntry entry = consumerEntries.computeIfAbsent(wrapper, k -> { + String consumerName = wrapper.consumer.consumerName(); + return new ConsumerEntry(consumerName, allocateConsumerNameIndex(consumerName), new MutableInt(0)); + }); + entry.refCount.increment(); + return entry.nameIndex; } private int allocateConsumerNameIndex(String consumerName) { - RoaringBitmap bitmap = getConsumerNameIndexBitmap(consumerName); - // find the first index that is not set, if there is no such index, add a new one - int index = (int) bitmap.nextAbsentValue(0); - if (index == -1) { - index = bitmap.getCardinality(); - } - bitmap.add(index); - return index; + return getConsumerNameIndexBitmap(consumerName).allocateIndexSlot(); } - private void deallocateConsumerNameIndex(String consumerName, int index) { - RoaringBitmap bitmap = getConsumerNameIndexBitmap(consumerName); - bitmap.remove(index); - if (bitmap.isEmpty()) { - consumerNameCounters.remove(consumerName); - } + private ConsumerNameIndexSlots getConsumerNameIndexBitmap(String consumerName) { + return consumerNameIndexSlotsMap.computeIfAbsent(consumerName, k -> new ConsumerNameIndexSlots()); } - public void removeHashRingReference(ConsumerIdentityWrapper removed) { + /* + * Decreases the reference count of the consumer and removes the consumer name from tracking if the ref count is + * zero. + */ + public void decreaseConsumerRefCount(ConsumerIdentityWrapper removed) { ConsumerEntry consumerEntry = consumerEntries.get(removed); int refCount = consumerEntry.refCount.decrementAndGet(); if (refCount == 0) { @@ -63,17 +112,25 @@ public void removeHashRingReference(ConsumerIdentityWrapper removed) { } } - public int addHashRingReference(ConsumerIdentityWrapper wrapper) { - String consumerName = wrapper.consumer.consumerName(); - ConsumerEntry entry = consumerEntries.computeIfAbsent(wrapper, - k -> new ConsumerEntry(consumerName, allocateConsumerNameIndex(consumerName), - new MutableInt(0))); - entry.refCount.increment(); - return entry.nameIndex; + private void deallocateConsumerNameIndex(String consumerName, int index) { + if (getConsumerNameIndexBitmap(consumerName).deallocateIndexSlot(index)) { + consumerNameIndexSlotsMap.remove(consumerName); + } } - public int getTrackedConsumerNameIndex(ConsumerIdentityWrapper wrapper) { + /* + * Returns the currently tracked index for the consumer. + */ + public int getTrackedIndex(ConsumerIdentityWrapper wrapper) { ConsumerEntry consumerEntry = consumerEntries.get(wrapper); return consumerEntry != null ? consumerEntry.nameIndex : -1; } + + int getTrackedConsumerNamesCount() { + return consumerNameIndexSlotsMap.size(); + } + + int getTrackedConsumersCount() { + return consumerEntries.size(); + } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 01e15895986a2..0d0d663502c4f 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -180,8 +180,21 @@ public void testGetConsumerKeyHashRanges() throws BrokerServiceException.Consume Range.of(0, 95615212), Range.of(455987437, 939655187), Range.of(1264144432, 1314727624), - Range.of(2016237255, 2147483646))); - assertThat(selector.getConsumerKeyHashRanges()).containsExactlyInAnyOrderEntriesOf(expectedResult); + Range.of(2016237254, 2147483646))); + Map> consumerKeyHashRanges = selector.getConsumerKeyHashRanges(); + assertThat(consumerKeyHashRanges).containsExactlyInAnyOrderEntriesOf(expectedResult); + + // check that ranges are continuous and cover the whole range + List allRanges = + consumerKeyHashRanges.values().stream().flatMap(List::stream).sorted().collect(Collectors.toList()); + Range previousRange = null; + for (Range range : allRanges) { + if (previousRange != null) { + assertThat(range.getStart()).isEqualTo(previousRange.getEnd() + 1); + } + previousRange = range; + } + assertThat(allRanges.stream().mapToInt(r -> r.getEnd() - r.getStart() + 1).sum()).isEqualTo(Integer.MAX_VALUE); } @Test diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapperTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapperTest.java new file mode 100644 index 0000000000000..75c8e6db5d2a0 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerIdentityWrapperTest.java @@ -0,0 +1,68 @@ +/* + * 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 static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotEquals; +import org.testng.annotations.Test; + +@Test(groups = "broker") +public class ConsumerIdentityWrapperTest { + private static Consumer mockConsumer() { + return mockConsumer("consumer"); + } + + private static Consumer mockConsumer(String consumerName) { + Consumer consumer = mock(Consumer.class); + when(consumer.consumerName()).thenReturn(consumerName); + return consumer; + } + + @Test + public void testEquals() { + Consumer consumer = mockConsumer(); + assertEquals(new ConsumerIdentityWrapper(consumer), new ConsumerIdentityWrapper(consumer)); + } + + @Test + public void testHashCode() { + Consumer consumer = mockConsumer(); + assertEquals(new ConsumerIdentityWrapper(consumer).hashCode(), + new ConsumerIdentityWrapper(consumer).hashCode()); + } + + @Test + public void testEqualsAndHashCode() { + Consumer consumer1 = mockConsumer(); + Consumer consumer2 = mockConsumer(); + ConsumerIdentityWrapper wrapper1 = new ConsumerIdentityWrapper(consumer1); + ConsumerIdentityWrapper wrapper2 = new ConsumerIdentityWrapper(consumer1); + ConsumerIdentityWrapper wrapper3 = new ConsumerIdentityWrapper(consumer2); + + // Test equality + assertEquals(wrapper1, wrapper2); + assertNotEquals(wrapper1, wrapper3); + + // Test hash code + assertEquals(wrapper1.hashCode(), wrapper2.hashCode()); + assertNotEquals(wrapper1.hashCode(), wrapper3.hashCode()); + } +} \ No newline at end of file diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerNameIndexTrackerTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerNameIndexTrackerTest.java new file mode 100644 index 0000000000000..0f18ecce2ffb4 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsumerNameIndexTrackerTest.java @@ -0,0 +1,157 @@ +/* + * 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.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotEquals; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +@Test(groups = "broker") +public class ConsumerNameIndexTrackerTest { + private ConsumerNameIndexTracker tracker; + + @BeforeMethod + public void setUp() { + tracker = new ConsumerNameIndexTracker(); + } + + private static Consumer mockConsumer() { + return mockConsumer("consumer"); + } + + + private static Consumer mockConsumer(String consumerName) { + Consumer consumer = mock(Consumer.class); + when(consumer.consumerName()).thenReturn(consumerName); + return consumer; + } + + @Test + public void testIncreaseConsumerRefCountAndReturnIndex() { + Consumer consumer1 = mockConsumer(); + Consumer consumer2 = mockConsumer(); + ConsumerIdentityWrapper wrapper1 = new ConsumerIdentityWrapper(consumer1); + ConsumerIdentityWrapper wrapper2 = new ConsumerIdentityWrapper(consumer2); + int index1 = tracker.increaseConsumerRefCountAndReturnIndex(wrapper1); + int index2 = tracker.increaseConsumerRefCountAndReturnIndex(wrapper2); + assertNotEquals(index1, index2); + assertEquals(index1, tracker.getTrackedIndex(wrapper1)); + assertEquals(index2, tracker.getTrackedIndex(wrapper2)); + } + + @Test + public void testTrackingReturnsStableIndexWhenRemovedAndAddedInSameOrder() { + List consumerIdentityWrappers = + IntStream.range(0, 100).mapToObj(i -> mockConsumer()).map(ConsumerIdentityWrapper::new).toList(); + Map trackedIndexes = + consumerIdentityWrappers.stream().collect(Collectors.toMap( + wrapper -> wrapper, wrapper -> tracker.increaseConsumerRefCountAndReturnIndex(wrapper))); + // stop tracking every other consumer + for (int i = 0; i < consumerIdentityWrappers.size(); i++) { + if (i % 2 == 0) { + tracker.decreaseConsumerRefCount(consumerIdentityWrappers.get(i)); + } + } + // check that others are tracked + for (int i = 0; i < consumerIdentityWrappers.size(); i++) { + ConsumerIdentityWrapper wrapper = consumerIdentityWrappers.get(i); + int trackedIndex = tracker.getTrackedIndex(wrapper); + assertEquals(trackedIndex, i % 2 == 0 ? -1 : trackedIndexes.get(wrapper)); + } + // check that new consumers are tracked with the same index + for (int i = 0; i < consumerIdentityWrappers.size(); i++) { + ConsumerIdentityWrapper wrapper = consumerIdentityWrappers.get(i); + if (i % 2 == 0) { + int trackedIndex = tracker.increaseConsumerRefCountAndReturnIndex(wrapper); + assertEquals(trackedIndex, trackedIndexes.get(wrapper)); + } + } + // check that all consumers are tracked with the original indexes + for (int i = 0; i < consumerIdentityWrappers.size(); i++) { + ConsumerIdentityWrapper wrapper = consumerIdentityWrappers.get(i); + int trackedIndex = tracker.getTrackedIndex(wrapper); + assertEquals(trackedIndex, trackedIndexes.get(wrapper)); + } + } + + @Test + public void testTrackingMultipleTimes() { + List consumerIdentityWrappers = + IntStream.range(0, 100).mapToObj(i -> mockConsumer()).map(ConsumerIdentityWrapper::new).toList(); + Map trackedIndexes = + consumerIdentityWrappers.stream().collect(Collectors.toMap( + wrapper -> wrapper, wrapper -> tracker.increaseConsumerRefCountAndReturnIndex(wrapper))); + Map trackedIndexes2 = + consumerIdentityWrappers.stream().collect(Collectors.toMap( + wrapper -> wrapper, wrapper -> tracker.increaseConsumerRefCountAndReturnIndex(wrapper))); + assertThat(tracker.getTrackedConsumerNamesCount()).isEqualTo(1); + assertThat(trackedIndexes).containsExactlyInAnyOrderEntriesOf(trackedIndexes2); + consumerIdentityWrappers.forEach(wrapper -> tracker.decreaseConsumerRefCount(wrapper)); + for (ConsumerIdentityWrapper wrapper : consumerIdentityWrappers) { + int trackedIndex = tracker.getTrackedIndex(wrapper); + assertEquals(trackedIndex, trackedIndexes.get(wrapper)); + } + consumerIdentityWrappers.forEach(wrapper -> tracker.decreaseConsumerRefCount(wrapper)); + assertThat(tracker.getTrackedConsumersCount()).isEqualTo(0); + assertThat(tracker.getTrackedConsumerNamesCount()).isEqualTo(0); + } + + @Test + public void testDecreaseConsumerRefCount() { + Consumer consumer1 = mockConsumer(); + ConsumerIdentityWrapper wrapper1 = new ConsumerIdentityWrapper(consumer1); + int index1 = tracker.increaseConsumerRefCountAndReturnIndex(wrapper1); + assertNotEquals(index1, -1); + tracker.decreaseConsumerRefCount(wrapper1); + assertEquals(tracker.getTrackedIndex(wrapper1), -1); + } + + @Test + public void testGetTrackedIndex() { + Consumer consumer1 = mockConsumer(); + Consumer consumer2 = mockConsumer(); + ConsumerIdentityWrapper wrapper1 = new ConsumerIdentityWrapper(consumer1); + ConsumerIdentityWrapper wrapper2 = new ConsumerIdentityWrapper(consumer2); + int index1 = tracker.increaseConsumerRefCountAndReturnIndex(wrapper1); + int index2 = tracker.increaseConsumerRefCountAndReturnIndex(wrapper2); + assertEquals(index1, tracker.getTrackedIndex(wrapper1)); + assertEquals(index2, tracker.getTrackedIndex(wrapper2)); + } + + @Test + public void testTrackingMultipleNames() { + List consumerIdentityWrappers = + IntStream.range(0, 100).mapToObj(i -> mockConsumer("consumer" + i)).map(ConsumerIdentityWrapper::new) + .toList(); + consumerIdentityWrappers.forEach(wrapper -> tracker.increaseConsumerRefCountAndReturnIndex(wrapper)); + assertThat(tracker.getTrackedConsumerNamesCount()).isEqualTo(100); + assertThat(tracker.getTrackedConsumersCount()).isEqualTo(100); + consumerIdentityWrappers.forEach(wrapper -> tracker.decreaseConsumerRefCount(wrapper)); + assertThat(tracker.getTrackedConsumersCount()).isEqualTo(0); + assertThat(tracker.getTrackedConsumerNamesCount()).isEqualTo(0); + } +} \ No newline at end of file diff --git a/pulsar-client-api/src/main/java/org/apache/pulsar/client/api/Range.java b/pulsar-client-api/src/main/java/org/apache/pulsar/client/api/Range.java index 4437ffc4ac6a2..488083f484b76 100644 --- a/pulsar-client-api/src/main/java/org/apache/pulsar/client/api/Range.java +++ b/pulsar-client-api/src/main/java/org/apache/pulsar/client/api/Range.java @@ -27,7 +27,7 @@ */ @InterfaceAudience.Public @InterfaceStability.Stable -public class Range { +public class Range implements Comparable { private final int start; private final int end; @@ -84,4 +84,13 @@ public int hashCode() { public String toString() { return "[" + start + ", " + end + "]"; } + + @Override + public int compareTo(Range o) { + int result = Integer.compare(start, o.start); + if (result == 0) { + result = Integer.compare(end, o.end); + } + return result; + } } From 22b95a7c6f5249610349fa9071dc78392538826e Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Tue, 1 Oct 2024 09:51:18 +0300 Subject: [PATCH 21/21] Add testConsumersReconnect test --- ...tHashingStickyKeyConsumerSelectorTest.java | 53 +++++++++++++++++++ ...ckyKeyDispatcherMultipleConsumersTest.java | 9 ++-- 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java index 0d0d663502c4f..04aafc49b47e6 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ConsistentHashingStickyKeyConsumerSelectorTest.java @@ -24,6 +24,7 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; @@ -493,4 +494,56 @@ public void testShouldNotChangeMappingWhenConsumerLeavesAndRejoins() { assertThat(selector.getConsumerKeyHashRanges()).as("ranges shouldn't change").containsExactlyInAnyOrderEntriesOf(expected); } + + @Test + public void testConsumersReconnect() { + final ConsistentHashingStickyKeyConsumerSelector selector = new ConsistentHashingStickyKeyConsumerSelector(100); + final String consumerName = "consumer"; + final int numOfInitialConsumers = 50; + final int validationPointCount = 200; + final List pointsToTest = pointsToTest(validationPointCount); + List consumers = new ArrayList<>(); + for (int i = 0; i < numOfInitialConsumers; i++) { + final Consumer consumer = createMockConsumer(consumerName, "index " + i, i); + consumers.add(consumer); + selector.addConsumer(consumer); + } + + // Mark original results. + List selectedConsumersBeforeRemove = new ArrayList<>(); + for (int i = 0; i < validationPointCount; i++) { + int point = pointsToTest.get(i); + selectedConsumersBeforeRemove.add(selector.select(point)); + } + + // All consumers leave (in any order) + List randomOrderConsumers = new ArrayList<>(consumers); + Collections.shuffle(randomOrderConsumers); + for (Consumer c : randomOrderConsumers) { + selector.removeConsumer(c); + } + + // All consumers reconnect in the same order as originally + for (Consumer c : consumers) { + selector.addConsumer(c); + } + + // Check that the same consumers are selected as before + for (int j = 0; j < validationPointCount; j++) { + int point = pointsToTest.get(j); + Consumer selected = selector.select(point); + Consumer expected = selectedConsumersBeforeRemove.get(j); + assertThat(selected.consumerId()).as("validationPoint %d, hash %d", j, point).isEqualTo(expected.consumerId()); + } + } + + private List pointsToTest(int validationPointCount) { + List res = new ArrayList<>(); + int hashRangeSize = Integer.MAX_VALUE; + final int increment = hashRangeSize / (validationPointCount + 1); + for (int i = 0; i < validationPointCount; i++) { + res.add(i * increment); + } + return res; + } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java index dcd852f409dbb..a0054f7e71425 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumersTest.java @@ -20,6 +20,7 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.apache.pulsar.common.protocol.Commands.serializeMetadataAndPayload; +import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyBoolean; import static org.mockito.Mockito.anyInt; @@ -326,7 +327,7 @@ public void testSkipRedeliverTemporally() { redeliverEntries.add(EntryImpl.create(1, 1, createMessage("message1", 1, "key1"))); final List readEntries = new ArrayList<>(); readEntries.add(EntryImpl.create(1, 2, createMessage("message2", 2, "key1"))); - readEntries.add(EntryImpl.create(1, 3, createMessage("message3", 3, "key22"))); + readEntries.add(EntryImpl.create(1, 3, createMessage("message3", 3, "key2"))); try { Field totalAvailablePermitsField = PersistentDispatcherMultipleConsumers.class.getDeclaredField("totalAvailablePermits"); @@ -417,7 +418,7 @@ public void testMessageRedelivery() throws Exception { // Messages with key1 are routed to consumer1 and messages with key2 are routed to consumer2 final List allEntries = new ArrayList<>(); - allEntries.add(EntryImpl.create(1, 1, createMessage("message1", 1, "key22"))); + allEntries.add(EntryImpl.create(1, 1, createMessage("message1", 1, "key2"))); allEntries.add(EntryImpl.create(1, 2, createMessage("message2", 2, "key1"))); allEntries.add(EntryImpl.create(1, 3, createMessage("message3", 3, "key1"))); allEntries.forEach(entry -> ((EntryImpl) entry).retain()); @@ -518,8 +519,8 @@ public void testMessageRedelivery() throws Exception { persistentDispatcher.readMoreEntries(); } - assertEquals(actualEntriesToConsumer1, expectedEntriesToConsumer1); - assertEquals(actualEntriesToConsumer2, expectedEntriesToConsumer2); + assertThat(actualEntriesToConsumer1).containsExactlyElementsOf(expectedEntriesToConsumer1); + assertThat(actualEntriesToConsumer2).containsExactlyElementsOf(expectedEntriesToConsumer2); allEntries.forEach(entry -> entry.release()); }