GH-2245 Ensure DefaultPartitionInterceptor is added to partitioned destinations managed by StreamBridge

Resolves #2245
This commit is contained in:
Oleg Zhurakousky
2021-11-16 14:06:44 +01:00
parent 4e011c26f3
commit d1358789b9
2 changed files with 75 additions and 2 deletions

View File

@@ -0,0 +1,68 @@
/*
* Copyright 2021-2021 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.binding;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.cloud.stream.binder.BinderHeaders;
import org.springframework.cloud.stream.binder.PartitionHandler;
import org.springframework.cloud.stream.config.BindingProperties;
import org.springframework.integration.expression.ExpressionUtils;
import org.springframework.integration.support.MessageBuilder;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.support.ChannelInterceptor;
/**
*
* @author Oleg Zhurakousky
* @since 3.1
*
*/
public class DefaultPartitioningInterceptor implements ChannelInterceptor {
private final PartitionHandler partitionHandler;
public DefaultPartitioningInterceptor(BindingProperties bindingProperties, ConfigurableListableBeanFactory beanFactory) {
this.partitionHandler = new PartitionHandler(
ExpressionUtils.createStandardEvaluationContext(beanFactory),
bindingProperties.getProducer(), beanFactory);
}
public void setPartitionCount(int partitionCount) {
this.partitionHandler.setPartitionCount(partitionCount);
}
@Override
public Message<?> preSend(Message<?> message, MessageChannel channel) {
if (!message.getHeaders().containsKey(BinderHeaders.PARTITION_OVERRIDE)) {
int partition = this.partitionHandler.determinePartition(message);
return MessageBuilder
.fromMessage(message)
.setHeader(BinderHeaders.PARTITION_HEADER, partition).build();
}
else {
return MessageBuilder
.fromMessage(message)
.setHeader(BinderHeaders.PARTITION_HEADER,
message.getHeaders()
.get(BinderHeaders.PARTITION_OVERRIDE))
.removeHeader(BinderHeaders.PARTITION_OVERRIDE).build();
}
}
}

View File

@@ -37,6 +37,8 @@ import org.springframework.cloud.stream.binder.BinderFactory;
import org.springframework.cloud.stream.binder.ProducerProperties;
import org.springframework.cloud.stream.binding.BinderAwareChannelResolver.NewDestinationBindingCallback;
import org.springframework.cloud.stream.binding.BindingService;
import org.springframework.cloud.stream.binding.DefaultPartitioningInterceptor;
import org.springframework.cloud.stream.config.BindingProperties;
import org.springframework.cloud.stream.config.BindingServiceProperties;
import org.springframework.cloud.stream.messaging.DirectWithAttributesChannel;
import org.springframework.context.ConfigurableApplicationContext;
@@ -205,8 +207,6 @@ public final class StreamBridge implements SmartInitializingSingleton {
ProducerProperties producerProperties = this.bindingServiceProperties.getProducerProperties(bindingName);
SubscribableChannel messageChannel = this.resolveDestination(bindingName, producerProperties, binderName);
// Function functionToInvoke = this.functionCatalog.lookup(STREAM_BRIDGE_FUNC_NAME, outputContentType.toString());
// ((FunctionInvocationWrapper) functionToInvoke).setSkipOutputConversion(producerProperties.isUseNativeEncoding());
Function functionToInvoke = this.getStreamBridgeFunction(outputContentType.toString(), producerProperties);
if (producerProperties != null && producerProperties.isPartitioned()) {
@@ -273,6 +273,11 @@ public final class StreamBridge implements SmartInitializingSingleton {
this.bindingService.bindProducer(messageChannel, destinationName, false, binder);
this.channelCache.put(destinationName, messageChannel);
if (producerProperties.isPartitioned()) {
BindingProperties bindingProperties = this.bindingServiceProperties.getBindingProperties(destinationName);
((AbstractMessageChannel) messageChannel)
.addInterceptor(new DefaultPartitioningInterceptor(bindingProperties, this.applicationContext.getBeanFactory()));
}
this.addInterceptors((AbstractMessageChannel) messageChannel, destinationName);
}