GH-653 Retain original 'lookupDestination' while adding 'reconsiledLookupDestination'

This is necessary to retain user provided information while also prpviding a way to reconsile the actual destination/function
This commit is contained in:
Oleg Zhurakousky
2021-02-22 14:08:02 +01:00
parent 0abce5a2b8
commit 3821e5745c
3 changed files with 79 additions and 2 deletions

View File

@@ -18,6 +18,7 @@ package org.springframework.cloud.function.rsocket;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
@@ -48,6 +49,7 @@ import org.springframework.messaging.MessageHeaders;
import org.springframework.messaging.MessagingException;
import org.springframework.messaging.handler.CompositeMessageCondition;
import org.springframework.messaging.handler.DestinationPatternsMessageCondition;
import org.springframework.messaging.handler.MessageCondition;
import org.springframework.messaging.handler.invocation.reactive.HandlerMethodArgumentResolver;
import org.springframework.messaging.handler.invocation.reactive.HandlerMethodReturnValueHandler;
import org.springframework.messaging.handler.invocation.reactive.SyncHandlerMethodArgumentResolver;
@@ -59,6 +61,7 @@ import org.springframework.messaging.rsocket.annotation.support.RSocketPayloadRe
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.util.MimeTypeUtils;
import org.springframework.util.ReflectionUtils;
import org.springframework.util.RouteMatcher;
import org.springframework.util.RouteMatcher.Route;
import org.springframework.util.StringUtils;
import org.springframework.web.util.pattern.PathPatternRouteMatcher;
@@ -73,6 +76,8 @@ import org.springframework.web.util.pattern.PathPatternRouteMatcher;
*/
class FunctionRSocketMessageHandler extends RSocketMessageHandler {
public static final String RECONSILED_LOOKUP_DESTINATION_HEADER = "reconsiledLookupDestination";
private final FunctionCatalog functionCatalog;
private final FunctionProperties functionProperties;
@@ -130,6 +135,27 @@ class FunctionRSocketMessageHandler extends RSocketMessageHandler {
return super.handleMessage(message);
}
@Override
protected RouteMatcher.Route getDestination(Message<?> message) {
RouteMatcher.Route reconsiledDestination = (RouteMatcher.Route) message.getHeaders().get(RECONSILED_LOOKUP_DESTINATION_HEADER);
return reconsiledDestination == null ? super.getDestination(message) : reconsiledDestination;
}
@Override
protected CompositeMessageCondition getMatchingMapping(CompositeMessageCondition mapping, Message<?> message) {
List<MessageCondition<?>> result = new ArrayList<>(mapping.getMessageConditions().size());
for (MessageCondition<?> condition : mapping.getMessageConditions()) {
MessageCondition<?> matchingCondition = condition instanceof DestinationPatternsMessageCondition
? condition
: (MessageCondition<?>) condition.getMatchingCondition(message);
if (matchingCondition == null) {
return null;
}
result.add(matchingCondition);
}
return new CompositeMessageCondition(result.toArray(new MessageCondition[0]));
}
void registerFunctionHandler(Function<?, ?> function, String route) {
CompositeMessageCondition condition =
new CompositeMessageCondition(REQUEST_CONDITION,
@@ -180,7 +206,7 @@ class FunctionRSocketMessageHandler extends RSocketMessageHandler {
Map<String, Object> headersMap = (Map<String, Object>) ReflectionUtils
.getField(this.headersField, message.getHeaders());
PathPatternRouteMatcher matcher = new PathPatternRouteMatcher();
headersMap.put(DestinationPatternsMessageCondition.LOOKUP_DESTINATION_HEADER, matcher.parseRoute(destination));
headersMap.put(RECONSILED_LOOKUP_DESTINATION_HEADER, matcher.parseRoute(destination));
}
protected static final class MessageHandlerMethodArgumentResolver implements SyncHandlerMethodArgumentResolver {

View File

@@ -96,7 +96,13 @@ class RSocketListenerFunction implements Function<Message<Flux<byte[]>>, Publish
}
else {
dataFlux = dataFlux.flatMap((data) -> {
Object result = this.targetFunction.isSupplier() ? this.targetFunction.apply(null) : this.targetFunction.apply(data);
Message<?> incoming = (Message<?>) data;
Message sanitizedMessage = MessageBuilder.withPayload(incoming.getPayload()).copyHeaders(incoming.getHeaders())
.removeHeader("dataBufferFactory")
.removeHeader("rsocketRequester")
.removeHeader("rsocketResponse")
.build();
Object result = this.targetFunction.isSupplier() ? this.targetFunction.apply(null) : this.targetFunction.apply(sanitizedMessage);
return result instanceof Publisher<?>
? (Publisher<Message<byte[]>>) result
: Mono.just((Message<byte[]>) result);