spring-projectsGH-3879: Add cache to optimize header match performance.

Fixes: spring-projects#3879
Issue link: spring-projects#3879

What
Add a LRU cache for pattern match of KafkaHeaderMapper.

Why?
To improve CPU usage used by pattern match of KafkaHeaderMapper.
Commonly, many Kafka records in the same topic will have the same header name.
Currently, Pattern Match has O(M*N) time complexity, where M is pattern
length, N is String length.

If results of patterns match are cached and KafkaHeaderMapper uses it,
KafkaHeaderMapper can expect improvement in terms of CPU usage.

Signed-off-by: Sanghyeok An <[email protected]>

Author: chickenchickenlove

spring-kafka/src/main/java/org/springframework/kafka/support/AbstractKafkaHeaderMapper.java

@@ -36,6 +36,7 @@
 import org.springframework.core.log.LogAccessor;
 import org.springframework.messaging.MessageHeaders;
 import org.springframework.util.Assert;
+import org.springframework.util.ConcurrentLruCache;
 import org.springframework.util.ObjectUtils;
 import org.springframework.util.PatternMatchUtils;
 
@@ -67,6 +68,8 @@ public abstract class AbstractKafkaHeaderMapper implements KafkaHeaderMapper {
 
 	private final List<HeaderMatcher> multiValueHeaderMatchers = new ArrayList<>();
 
+	private final HeaderPatternMatchCache headerMatchedCache = new HeaderPatternMatchCache();
+
 	private final Map<String, Boolean> rawMappedHeaders = new HashMap<>();
 
 	{
@@ -272,11 +275,22 @@ protected Object headerValueToAddOut(String key, Object value) {
 	 * @since 4.0
 	 */
 	protected boolean doesMatchMultiValueHeader(String headerName) {
+		if (this.headerMatchedCache.isMultiValuePattern(headerName)) {
+			return true;
+		}
+
+		if (this.headerMatchedCache.isSingleValuePattern(headerName)) {
+			return false;
+		}
+
 		for (HeaderMatcher headerMatcher : this.multiValueHeaderMatchers) {
 			if (headerMatcher.matchHeader(headerName)) {
+				this.headerMatchedCache.cacheAsMultiValueHeader(headerName);
 				return true;
 			}
 		}
+
+		this.headerMatchedCache.cacheAsSingleValueHeader(headerName);
 		return false;
 	}
 
@@ -427,4 +441,33 @@ public boolean isNegated() {
 
 	}
 
+	/**
+	 * A Cache that remembers whether a header name matches the multi-value pattern.
+	 */
+	class HeaderPatternMatchCache {
+
+		private static final int MAX_SIZE = 1000;
+
+		private final ConcurrentLruCache<String, Boolean> multiValueHeaderPatternMatchCache = new ConcurrentLruCache<>(MAX_SIZE, key -> Boolean.TRUE);
+
+		private final ConcurrentLruCache<String, Boolean> singleValueHeaderPatternMatchCache = new ConcurrentLruCache<>(MAX_SIZE, key -> Boolean.TRUE);
+
+		public boolean isMultiValuePattern(String headerName) {
+			return this.multiValueHeaderPatternMatchCache.contains(headerName);
+		}
+
+		public boolean isSingleValuePattern(String headerName) {
+			return this.singleValueHeaderPatternMatchCache.contains(headerName);
+		}
+
+		public void cacheAsSingleValueHeader(String headerName) {
+			this.singleValueHeaderPatternMatchCache.get(headerName);
+		}
+
+		public void cacheAsMultiValueHeader(String headerName) {
+			this.multiValueHeaderPatternMatchCache.get(headerName);
+		}
+
+	}
+
 }

spring-kafka/src/test/java/org/springframework/kafka/support/DefaultKafkaHeaderMapperTests.java

@@ -31,6 +31,8 @@
 import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 import org.springframework.core.log.LogAccessor;
 import org.springframework.kafka.retrytopic.RetryTopicHeaders;
@@ -413,6 +415,103 @@ void multiValueHeaderToTest() {
 				.containsExactly(multiValueWildCardHeader2Value1, multiValueWildCardHeader2Value2);
 	}
 
+	@ParameterizedTest
+	@ValueSource(ints = {2000})
+//	@ValueSource(ints = {500, 1000, 2000})
+	void hugeNumberOfSingleValueHeaderToTest(int numberOfSingleValueHeaderCount) {
+		// GIVEN
+		Headers rawHeaders = new RecordHeaders();
+
+		String multiValueHeader1 = "test-multi-value1";
+		byte[] multiValueHeader1Value1 = { 0, 0, 0, 0 };
+		byte[] multiValueHeader1Value2 = { 0, 0, 0, 1 };
+
+		rawHeaders.add(multiValueHeader1, multiValueHeader1Value1);
+		rawHeaders.add(multiValueHeader1, multiValueHeader1Value2);
+
+		byte[] deliveryAttemptHeaderValue = { 0, 0, 0, 1 };
+		byte[] originalOffsetHeaderValue = { 0, 0, 0, 2 };
+		byte[] defaultHeaderAttemptsValues = { 0, 0, 0, 5 };
+
+		rawHeaders.add(KafkaHeaders.DELIVERY_ATTEMPT, deliveryAttemptHeaderValue);
+		rawHeaders.add(KafkaHeaders.ORIGINAL_OFFSET, originalOffsetHeaderValue);
+		rawHeaders.add(RetryTopicHeaders.DEFAULT_HEADER_ATTEMPTS, defaultHeaderAttemptsValues);
+
+		byte[] singleValueHeaderValue = { 0, 0, 0, 6 };
+		for (int i = 0; i < numberOfSingleValueHeaderCount; i++) {
+			String singleValueHeader = "test-single-value" + i;
+			rawHeaders.add(singleValueHeader, singleValueHeaderValue);
+		}
+
+		DefaultKafkaHeaderMapper mapper = new DefaultKafkaHeaderMapper();
+		mapper.setMultiValueHeaderPatterns(multiValueHeader1);
+
+		// WHEN
+		Map<String, Object> mappedHeaders = new HashMap<>();
+		mapper.toHeaders(rawHeaders, mappedHeaders);
+
+		// THEN
+		assertThat(mappedHeaders.get(KafkaHeaders.DELIVERY_ATTEMPT)).isEqualTo(1);
+		assertThat(mappedHeaders.get(KafkaHeaders.ORIGINAL_OFFSET)).isEqualTo(originalOffsetHeaderValue);
+		assertThat(mappedHeaders.get(RetryTopicHeaders.DEFAULT_HEADER_ATTEMPTS)).isEqualTo(defaultHeaderAttemptsValues);
+
+		for (int i = 0; i < numberOfSingleValueHeaderCount; i++) {
+			String singleValueHeader = "test-single-value" + i;
+			assertThat(mappedHeaders.get(singleValueHeader)).isEqualTo(singleValueHeaderValue);
+		}
+
+		assertThat(mappedHeaders)
+				.extractingByKey(multiValueHeader1, InstanceOfAssertFactories.list(byte[].class))
+				.containsExactly(multiValueHeader1Value1, multiValueHeader1Value2);
+	}
+
+	@ParameterizedTest
+	@ValueSource(ints = {500, 1000, 2000})
+	void hugeNumberOfMultiValueHeaderToTest(int numberOfMultiValueHeaderCount) {
+		// GIVEN
+		DefaultKafkaHeaderMapper mapper = new DefaultKafkaHeaderMapper();
+		Headers rawHeaders = new RecordHeaders();
+
+		byte[] multiValueHeader1Value1 = { 0, 0, 0, 0 };
+		byte[] multiValueHeader1Value2 = { 0, 0, 0, 1 };
+
+		for (int i = 0; i < numberOfMultiValueHeaderCount; i++) {
+			String multiValueHeader = "test-multi-value" + i;
+			mapper.setMultiValueHeaderPatterns(multiValueHeader);
+			rawHeaders.add(multiValueHeader, multiValueHeader1Value1);
+			rawHeaders.add(multiValueHeader, multiValueHeader1Value2);
+		}
+
+		byte[] deliveryAttemptHeaderValue = { 0, 0, 0, 1 };
+		byte[] originalOffsetHeaderValue = { 0, 0, 0, 2 };
+		byte[] defaultHeaderAttemptsValues = { 0, 0, 0, 5 };
+
+		rawHeaders.add(KafkaHeaders.DELIVERY_ATTEMPT, deliveryAttemptHeaderValue);
+		rawHeaders.add(KafkaHeaders.ORIGINAL_OFFSET, originalOffsetHeaderValue);
+		rawHeaders.add(RetryTopicHeaders.DEFAULT_HEADER_ATTEMPTS, defaultHeaderAttemptsValues);
+
+		String singleValueHeader = "test-single-value";
+		byte[] singleValueHeaderValue = { 0, 0, 0, 6 };
+		rawHeaders.add(singleValueHeader, singleValueHeaderValue);
+
+		// WHEN
+		Map<String, Object> mappedHeaders = new HashMap<>();
+		mapper.toHeaders(rawHeaders, mappedHeaders);
+
+		// THEN
+		assertThat(mappedHeaders.get(KafkaHeaders.DELIVERY_ATTEMPT)).isEqualTo(1);
+		assertThat(mappedHeaders.get(KafkaHeaders.ORIGINAL_OFFSET)).isEqualTo(originalOffsetHeaderValue);
+		assertThat(mappedHeaders.get(RetryTopicHeaders.DEFAULT_HEADER_ATTEMPTS)).isEqualTo(defaultHeaderAttemptsValues);
+		assertThat(mappedHeaders.get(singleValueHeader)).isEqualTo(singleValueHeaderValue);
+
+		for (int i = 0; i < numberOfMultiValueHeaderCount; i++) {
+			String multiValueHeader = "test-multi-value" + i;
+			assertThat(mappedHeaders)
+					.extractingByKey(multiValueHeader, InstanceOfAssertFactories.list(byte[].class))
+					.containsExactly(multiValueHeader1Value1, multiValueHeader1Value2);
+		}
+	}
+
 	@Test
 	void multiValueHeaderFromTest() {
 		// GIVEN