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 377edae7ad6f9..f4008ab2e551f 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,16 +18,19 @@ */ package org.apache.pulsar.broker.service; +import com.google.common.collect.Lists; +import org.apache.pulsar.broker.service.BrokerServiceException.ConsumerAssignException; +import org.apache.pulsar.common.util.Murmur3_32Hash; + import java.util.Collections; +import java.util.Comparator; +import java.util.List; import java.util.Map; import java.util.NavigableMap; import java.util.TreeMap; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; -import org.apache.pulsar.broker.service.BrokerServiceException.ConsumerAssignException; -import org.apache.pulsar.common.util.Murmur3_32Hash; - /** * This is a consumer selector based fixed hash range. * @@ -39,7 +42,7 @@ public class ConsistentHashingStickyKeyConsumerSelector implements StickyKeyCons private final ReadWriteLock rwLock = new ReentrantReadWriteLock(); // Consistent-Hash ring - private final NavigableMap hashRing; + private final NavigableMap> hashRing; private final int numberOfPoints; @@ -57,7 +60,17 @@ public void addConsumer(Consumer consumer) throws ConsumerAssignException { for (int i = 0; i < numberOfPoints; i++) { String key = consumer.consumerName() + i; int hash = Murmur3_32Hash.getInstance().makeHash(key.getBytes()); - hashRing.put(hash, consumer); + 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; + } + }); } } finally { rwLock.writeLock().unlock(); @@ -72,7 +85,17 @@ public void removeConsumer(Consumer consumer) { for (int i = 0; i < numberOfPoints; i++) { String key = consumer.consumerName() + i; int hash = Murmur3_32Hash.getInstance().makeHash(key.getBytes()); - hashRing.remove(hash, consumer); + hashRing.compute(hash, (k, v) -> { + if (v == null) { + return null; + } else { + v.removeIf(c -> c.consumerName().equals(consumer.consumerName())); + if (v.isEmpty()) { + v = null; + } + return v; + } + }); } } finally { rwLock.writeLock().unlock(); @@ -89,18 +112,21 @@ public Consumer select(byte[] stickyKey) { return null; } - Map.Entry ceilingEntry = hashRing.ceilingEntry(hash); + List consumerList; + Map.Entry> ceilingEntry = hashRing.ceilingEntry(hash); if (ceilingEntry != null) { - return ceilingEntry.getValue(); + consumerList = ceilingEntry.getValue(); } else { - return hashRing.firstEntry().getValue(); + consumerList = hashRing.firstEntry().getValue(); } + + return consumerList.get(hash % consumerList.size()); } finally { rwLock.readLock().unlock(); } } - Map getRangeConsumer() { + Map> getRangeConsumer() { return Collections.unmodifiableMap(hashRing); } }