Skip to content

[fix][test] Make KeyShared draining-hashes delivery test deterministic #26602

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Dream95 wants to merge 4 commits into
base: master
Choose a base branch
from
Open

[fix][test] Make KeyShared draining-hashes delivery test deterministic #26602

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
75df0ef
[fix][test] Make KeyShared draining-hashes delivery test deterministic
Dream95 Sep 16, 2026
2cb47cd
[fix][test] Drop invocationCount from draining-hashes delivery test
Dream95 Sep 16, 2026
a4ac78b
[fix][test] Use fixed initial consumer name in draining-hashes test
Dream95 Sep 25, 2026
e73de21
[fix][test] Bound the takeover-hash search in draining-hashes test to…
Dream95 Sep 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
151 changes: 92 additions & 59 deletions pulsar-broker/src/test/java/org/apache/pulsar/client/api/KeySharedSubscriptionTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pulsar.broker.BrokerTestUtil;
import org.apache.pulsar.broker.ServiceConfiguration;
import org.apache.pulsar.broker.service.ConsistentHashingStickyKeyConsumerSelector;
import org.apache.pulsar.broker.service.Dispatcher;
import org.apache.pulsar.broker.service.DrainingHashesTracker;
import org.apache.pulsar.broker.service.PendingAcksMap;
Expand Down Expand Up @@ -936,14 +937,14 @@ public void testWithMessageCompression(KeySharedImplementationType impl) throws
2. "consumer" receives BUT DO NOT acknowledge message with key "testMessageKey"
Example: "consumer" hash range: [0-256]. "testMessage" hash range: [0-256]

3. Add more consumers until hash range of message with key "testMessageKey" removed from "consumer" hash ranges
Example: "consumer" hash range: [0-123]. "consumer N" hash range: [124-136]. "testMessage" hash range: [124-136]
3. Add a consumer whose name is chosen so that it takes over the hash of "testMessageKey"
Example: "consumer" hash range: [0-123]. "takeover-0" hash range: [124-136]. "testMessage" hash range: [124-136]

4. Stop the added consumers until message key "testMessageKey" hash range moved to "consumer" hash ranges
4. Stop the added consumer so that the hash of "testMessageKey" moves back to "consumer"
Example: "consumer" hash range: [0-128]. "testMessage" hash range: [0-128]

5. Add more consumers until hash range of message with key "testMessageKey" removed from "consumer" hash ranges
Example: "consumer" hash range: [0-96]. "consumer N" hash range: [121-154]. "testMessage" hash range: [121-154]
5. Add such a consumer again so that it takes over the hash of "testMessageKey"
Example: "consumer" hash range: [0-96]. "takeover-0" hash range: [121-154]. "testMessage" hash range: [121-154]

6. Stop "consumer"

Expand All @@ -952,18 +953,18 @@ public void testWithMessageCompression(KeySharedImplementationType impl) throws

Expected result:
The message with key "testMessageKey" received by consumer who owns message's hash range
Example: "consumer N" with hash range [121-154] received the message
Example: "takeover-0" with hash range [121-154] received the message

*/
@Test
void testMessageDeliveredFromDrainingHashes() throws PulsarClientException {
String messageKey = "testMessageKey";
String topic = "testMessageDeliveredFromDrainingHashes" + UUID.randomUUID();

List<Consumer<Integer>> consumers = new ArrayList<>();

// Fixed name verified not to collide with the hash of "testMessageKey" under the current
// consistent-hashing settings, so a "takeover-N" consumer can always take over that hash.
@Cleanup
Consumer<Integer> consumer = createConsumer(topic);
Consumer<Integer> consumer = createConsumerWithName(topic,"initial-owner");
String initialOwnerName = consumer.getConsumerName();

@Cleanup
Expand All @@ -978,75 +979,98 @@ void testMessageDeliveredFromDrainingHashes() throws PulsarClientException {
StickyKeyConsumerSelector selector = dispatcher.getSelector();
int messageKeyHash = selector.makeStickyKeyHash(messageKey.getBytes(StandardCharsets.UTF_8));

Consumer<Integer> newOwner = null;
try {
// add new consumers until hash range of the initial message owner moved
Pair<String, List<Consumer<Integer>>> addedConsumers = addConsumersUntilOwnerChanged(
topic, initialOwnerName, messageKeyHash, selector
);
String currentOwnerName = addedConsumers.getKey();
consumers.addAll(addedConsumers.getValue());
// move the hash of the message key away from the initial owner
newOwner = addConsumerThatTakesOverHash(topic, initialOwnerName, messageKeyHash, selector);

// returning hash range to the initial owner
for (int i = consumers.size() - 1; i > -1 && !initialOwnerName.equals(currentOwnerName); i--) {
consumers.remove(i).close();
currentOwnerName = findOwnerName(selector, messageKeyHash);
}

// add new consumers until hash range of the initial message owner moved
Pair<String, List<Consumer<Integer>>> addedConsumersAfter = addConsumersUntilOwnerChanged(
topic, initialOwnerName, messageKeyHash, selector
);
newOwner.close();
assertThat(findOwnerName(selector, messageKeyHash)).isEqualTo(initialOwnerName);

String currentOwnerNameAfter = addedConsumersAfter.getKey();
consumers.addAll(addedConsumersAfter.getValue());
// move the hash of the message key away from the initial owner again
newOwner = addConsumerThatTakesOverHash(topic, initialOwnerName, messageKeyHash, selector);

// remove the initial owner
consumer.close();

// verify the message sent to new consumer which owns hash range
Optional<Consumer<Integer>> theNewOwnerMaybe = consumers
.stream()
.filter(c -> c.getConsumerName().equals(currentOwnerNameAfter))
.findAny();

assertThat(theNewOwnerMaybe).isPresent();

Consumer<Integer> theNewOwner = theNewOwnerMaybe.get();

Message<Integer> message = theNewOwner.receive(5, TimeUnit.SECONDS);
Message<Integer> message = newOwner.receive(5, TimeUnit.SECONDS);
assertThat(message).describedAs("Message with key " + messageKey
+ " expected to be delivered to consumer " + theNewOwner.getConsumerName()).isNotNull();
+ " expected to be delivered to consumer " + newOwner.getConsumerName()).isNotNull();
assertThat(message.getKey()).isEqualTo(messageKey);
} finally {
for (Consumer<Integer> c : consumers) {
c.close();
if (newOwner != null) {
newOwner.close();
}
}
}

private Pair<String, List<Consumer<Integer>>> addConsumersUntilOwnerChanged(
private Consumer<Integer> addConsumerThatTakesOverHash(
String topic,
String initialOwnerName,
int messageKeyHash,
StickyKeyConsumerSelector selector
) throws PulsarClientException {
String currentOwnerName;
List<Consumer<Integer>> consumers = new ArrayList<>();
do {
Consumer<Integer> addedC = createConsumer(topic);
consumers.add(addedC);
currentOwnerName = findOwnerName(selector, messageKeyHash);
} while (initialOwnerName.equals(currentOwnerName));
return Pair.of(currentOwnerName, consumers);
String takeoverName = findConsumerNameThatTakesHash(selector, initialOwnerName, messageKeyHash);
Consumer<Integer> added = createConsumerWithName(topic, takeoverName);
String currentOwnerName = findOwnerName(selector, messageKeyHash);
if (!takeoverName.equals(currentOwnerName)) {
added.close();
fail("Consumer " + takeoverName + " was expected to take over the hash " + messageKeyHash
+ ", but the owner is " + currentOwnerName);
}
return added;
}

/**
* Predicts which consumer name will take over the given hash from its current owner once a
* consumer with that name joins the subscription.
*/
private String findConsumerNameThatTakesHash(StickyKeyConsumerSelector selector, String initialOwnerName,
int messageKeyHash) {
Range range = selector.getKeyHashRange();
ConsistentHashingStickyKeyConsumerSelector probe = new ConsistentHashingStickyKeyConsumerSelector(
conf.getSubscriptionKeySharedConsistentHashingReplicaPoints(), false, range.getEnd());
Set<String> existingNames = new HashSet<>();
for (var existing : selector.getConsumerKeyHashRanges().keySet()) {
existingNames.add(existing.consumerName());
probe.addConsumer(mockBrokerConsumer(existing.consumerName())).join();
}
// Search the whole hash-key space. If no name takes over the hash once it's exhausted, an exact
// hash-point collision means takeover can't happen, so fail fast instead of looping.
int maxAttempts = selector.getKeyHashRange().getEnd();
for (int i = 0; i <maxAttempts ; i++) {
String name = "takeover-" + i;
if (!existingNames.add(name)) {
continue;
}

@lhotari lhotari Sep 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[INTENT MISMATCH] Choose a suitable initial name and bound the takeover search

The initial consumer still gets a random five-character name. With the valid name 00807, the modern selector's replica point 78 (00807 + NUL + 0 + NUL + 78) and testMessageKey both hash to slot 21672 with the default 100 points and range 1..65535. The existing ring point always wins: putIfAbsent retains its owner, and select uses ceilingEntry. No takeover-N can therefore succeed. This loop never terminates, and existingNames keeps every attempted name.

The collision also affected the old loop, but moving the search into a probe does not resolve this case of the flaky test. Please give the initial consumer a fixed suitable name (or choose a compatible pair), bound the search, and cover the exact-point case.

var candidate = mockBrokerConsumer(name);
probe.addConsumer(candidate).join();
var selected = probe.select(messageKeyHash);
if (selected != null && !initialOwnerName.equals(selected.consumerName())) {
return name;
}
probe.removeConsumer(candidate);
}
fail("Could not find a consumer name to take over hash " + messageKeyHash
+ " after exhausting the hash-key space (" + maxAttempts + "slots)");
return null;
}

private org.apache.pulsar.broker.service.Consumer mockBrokerConsumer(String consumerName) {
org.apache.pulsar.broker.service.Consumer consumer =
Mockito.mock(org.apache.pulsar.broker.service.Consumer.class);
Mockito.when(consumer.consumerName()).thenReturn(consumerName);
return consumer;
}

private String findOwnerName(StickyKeyConsumerSelector selector, int hash) {
return selector.getConsumerKeyHashRanges().entrySet().stream()
.filter(entry ->
entry.getValue().stream().anyMatch(range -> range.contains(hash))
)
.findAny().orElseThrow(() -> new IllegalArgumentException("No owner for the hash " + hash))
.getKey().consumerName();
var owner = selector.select(hash);
if (owner == null) {
throw new IllegalArgumentException("No owner for the hash " + hash);
}
return owner.consumerName();
}

@Test(dataProvider = "currentImplementationType")
Expand Down Expand Up @@ -1539,17 +1563,26 @@ private Consumer<Integer> createConsumer(String topic) throws PulsarClientExcept

private Consumer<Integer> createConsumer(String topic, KeySharedPolicy keySharedPolicy)
throws PulsarClientException {
ConsumerBuilder<Integer> builder = pulsarClient.newConsumer(Schema.INT32);
builder.topic(topic)
.subscriptionName(SUBSCRIPTION_NAME)
.subscriptionType(SubscriptionType.Key_Shared)
.ackTimeout(3, TimeUnit.SECONDS);
ConsumerBuilder<Integer> builder = newConsumerBuilder(topic);
if (keySharedPolicy != null) {
builder.keySharedPolicy(keySharedPolicy);
}
return builder.subscribe();
}

private Consumer<Integer> createConsumerWithName(String topic, String consumerName)
throws PulsarClientException {
return newConsumerBuilder(topic).consumerName(consumerName).subscribe();
}

private ConsumerBuilder<Integer> newConsumerBuilder(String topic) {
return pulsarClient.newConsumer(Schema.INT32)
.topic(topic)
.subscriptionName(SUBSCRIPTION_NAME)
.subscriptionType(SubscriptionType.Key_Shared)
.ackTimeout(3, TimeUnit.SECONDS);
}

private void receive(List<Consumer<?>> consumers) throws PulsarClientException {
// Add a key so that we know this key was already assigned to one consumer
Map<String, Consumer<?>> keyToConsumer = new HashMap<>();
Expand Down