From 604649d083d4f868a75e219025c4fd461384bdbe Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Tue, 18 Jun 2024 19:17:00 -0400 Subject: [PATCH] GH-2961: Concurrent writes and partition header - When concurrent threads publish to a binding, PartitionAwareFunctionWrapper resets to null between invocations. Addressing this issue by guarding this reset from occurring if the partiton header on the producer is found. Fixes https://github.com/spring-cloud/spring-cloud-stream/issues/2961 --- .../stream/function/PartitionAwareFunctionWrapper.java | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java b/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java index 882040ff9..11c5b8fbd 100644 --- a/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java +++ b/core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java @@ -1,5 +1,5 @@ /* - * Copyright 2020-2023 the original author or authors. + * Copyright 2020-2024 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. @@ -90,7 +90,13 @@ class PartitionAwareFunctionWrapper implements Function, Supplie this.setEnhancerIfNecessary(); } Object result = this.function.apply(input); - if (!((FunctionInvocationWrapper) this.function).isInputTypePublisher()) { + boolean messageContainsPartitionHeader = false; + if (result != null && Message.class.isAssignableFrom(result.getClass())) { + if (((Message) result).getHeaders().containsKey(BinderHeaders.PARTITION_HEADER)) { + messageContainsPartitionHeader = true; + } + } + if (!((FunctionInvocationWrapper) this.function).isInputTypePublisher() && !messageContainsPartitionHeader) { ((FunctionInvocationWrapper) this.function).setEnhancer(null); } return result;