From 0a07ba1bcc4190e53ae9333083f733ed66c0ecc2 Mon Sep 17 00:00:00 2001 From: 98001yash Date: Fri, 23 Oct 2026 20:31:48 +0530 Subject: [PATCH] Improve default partition distribution for UUID keys Signed-off-by: 98001yash --- .../cloud/stream/binder/PartitionHandler.java | 60 +++++++++++++-- .../stream/binder/PartitionHandlerTests.java | 77 +++++++++++++++++++ 2 files changed, 131 insertions(+), 6 deletions(-) create mode 100644 core/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/binder/PartitionHandlerTests.java diff --git a/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/PartitionHandler.java b/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/PartitionHandler.java index c420998a11..59c02723d2 100644 --- a/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/PartitionHandler.java +++ b/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/PartitionHandler.java @@ -16,6 +16,7 @@ package org.springframework.cloud.stream.binder; +import java.nio.charset.StandardCharsets; import java.util.Map; import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; @@ -183,19 +184,66 @@ private PartitionSelectorStrategy getPartitionSelectorStrategy( } /** - * Default partition strategy; only works on keys with "real" hash codes, such as - * String. Caller now always applies modulo so no need to do so here. + * Default partition strategy using Murmur2 hashing to provide a more even + * distribution of keys across partitions. */ private static class DefaultPartitionSelector implements PartitionSelectorStrategy { + private static final int SEED = 0x9747b28c; + + private static final int M = 0x5bd1e995; + @Override public int selectPartition(Object key, int partitionCount) { - int hashCode = key.hashCode(); - if (hashCode == Integer.MIN_VALUE) { - hashCode = 0; + byte[] bytes = key.toString().getBytes(StandardCharsets.UTF_8); + return murmur2(bytes) & 0x7fffffff; + } + + private int murmur2(byte[] data) { + int length = data.length; + int h = SEED ^ length; + int length4 = length / 4; + + for (int i = 0; i < length4; i++) { + int i4 = i * 4; + + int k = (data[i4] & 0xff) + | ((data[i4 + 1] & 0xff) << 8) + | ((data[i4 + 2] & 0xff) << 16) + | ((data[i4 + 3] & 0xff) << 24); + + k *= M; + k ^= k >>> 24; + k *= M; + + h *= M; + h ^= k; + } + + int remaining = length % 4; + int offset = length4 * 4; + + switch (remaining) { + case 3: + h ^= (data[offset + 2] & 0xff) << 16; + // fall through + case 2: + h ^= (data[offset + 1] & 0xff) << 8; + // fall through + case 1: + h ^= data[offset] & 0xff; + h *= M; + default: + break; } - return Math.abs(hashCode); + + h ^= h >>> 13; + h *= M; + h ^= h >>> 15; + + return h; } + } diff --git a/core/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/binder/PartitionHandlerTests.java b/core/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/binder/PartitionHandlerTests.java new file mode 100644 index 0000000000..f0dfeff8ab --- /dev/null +++ b/core/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/binder/PartitionHandlerTests.java @@ -0,0 +1,77 @@ +/* + * Copyright 2026-present the original author or authors. + * + * Licensed 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 + * + * https://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.springframework.cloud.stream.binder; + +import java.util.UUID; + +import org.junit.jupiter.api.Test; + +import org.springframework.beans.factory.support.DefaultListableBeanFactory; +import org.springframework.expression.spel.standard.SpelExpressionParser; +import org.springframework.expression.spel.support.StandardEvaluationContext; +import org.springframework.messaging.Message; +import org.springframework.messaging.support.MessageBuilder; + +import static org.assertj.core.api.Assertions.assertThat; + +class PartitionHandlerTests { + + @Test + void defaultPartitionSelectorDistributesUuidKeysAcrossPartitions() { + int partitionCount = 16; + int numberOfKeys = 10_000; + int[] distribution = new int[partitionCount]; + + PartitionHandler partitionHandler = createPartitionHandler(partitionCount); + + for (int i = 0; i < numberOfKeys; i++) { + String key = new UUID(0L, i).toString(); + Message message = MessageBuilder.withPayload(key).build(); + + int partition = partitionHandler.determinePartition(message); + distribution[partition]++; + } + + int minimum = numberOfKeys / partitionCount; + int maximum = 0; + int actualMinimum = numberOfKeys; + + for (int count : distribution) { + assertThat(count) + .as("Each partition should receive a reasonable share of keys") + .isGreaterThan(0); + maximum = Math.max(maximum, count); + actualMinimum = Math.min(actualMinimum, count); + } + + assertThat(maximum - actualMinimum) + .as("Difference between most- and least-used partitions: %s", + distribution) + .isLessThan(minimum / 2); + } + + private PartitionHandler createPartitionHandler(int partitionCount) { + ProducerProperties properties = new ProducerProperties(); + properties.setPartitionCount(partitionCount); + properties.setPartitionKeyExpression( + new SpelExpressionParser().parseExpression("payload")); + + return new PartitionHandler(new StandardEvaluationContext(), properties, + new DefaultListableBeanFactory()); + } + +}