Re-enable support for around advise

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/FunctionAroundWrapper.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2020-2020 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.
@@ -22,6 +22,12 @@ import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry
 import org.springframework.messaging.Message;
 
 /**
+ * Wrapper that acts as around advise over function invocation.
+ * If registered as bean it will be autowired into {@link FunctionInvocationWrapper}.
+ * Keep in mind that it only affects imperative invocations where input is {@link Message}
+ *
+ * NOTE: This API is experimental and and could change without notice. It is
+ * intended for internal use only (e.g., spring-cloud-sleuth)
  *
  * @author Oleg Zhurakousky
  * @since 3.1

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java

@@ -46,6 +46,7 @@ import reactor.core.publisher.Mono;
 import reactor.util.function.Tuples;
 
 import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.cloud.function.context.FunctionCatalog;
 import org.springframework.cloud.function.context.FunctionProperties;
 import org.springframework.cloud.function.context.FunctionRegistration;
@@ -96,6 +97,9 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 
 	private final JsonMapper jsonMapper;
 
+	@Autowired(required = false)
+	private FunctionAroundWrapper functionAroundWrapper;
+
 	public SimpleFunctionRegistry(ConversionService conversionService, CompositeMessageConverter messageConverter, JsonMapper jsonMapper) {
 		Assert.notNull(messageConverter, "'messageConverter' must not be null");
 		Assert.notNull(jsonMapper, "'jsonMapper' must not be null");
@@ -106,13 +110,6 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 		this.headersField.setAccessible(true);
 	}
 
-	@Override
-	public FunctionRegistration<?> getRegistration(Object function) {
-		throw new UnsupportedOperationException("FunctionInspector is deprecated. There is no need "
-				+ "to access FunctionRegistration directly since you can interogate the actual "
-				+ "looked-up function (see FunctionInvocationWrapper.");
-	}
-
 	@SuppressWarnings("unchecked")
 	@Override
 	public <T> T lookup(Class<?> type, String functionDefinition, String... expectedOutputMimeTypes) {
@@ -129,6 +126,13 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 		return (T) function;
 	}
 
+	@Override
+	public FunctionRegistration<?> getRegistration(Object function) {
+		throw new UnsupportedOperationException("FunctionInspector is deprecated. There is no need "
+				+ "to access FunctionRegistration directly since you can interogate the actual "
+				+ "looked-up function (see FunctionInvocationWrapper.");
+	}
+
 	@Override
 	public <T> void register(FunctionRegistration<T> registration) {
 		this.functionRegistrations.add(registration);
@@ -168,8 +172,11 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 			function.expectedOutputContentType = expectedOutputMimeTypes;
 		}
 		else {
-			logger.debug("Function '" + functionDefinition + "' is not found");
+			logger.debug("Function '" + functionDefinition + "' is not found in cache");
 		}
+
+		function = this.wrapInAroundAviceIfNecessary(function);
+
 		return (T) function;
 	}
 
@@ -198,6 +205,25 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 		return functionDefinition;
 	}
 
+	/**
+	 * This is primarily to support spring-cloud-sleauth.
+	 * There is no current use cases in functions where it is used.
+	 * The approach may change in the future.
+	 */
+	private FunctionInvocationWrapper wrapInAroundAviceIfNecessary(FunctionInvocationWrapper function) {
+		FunctionInvocationWrapper wrappedFunction = function;
+		if (function != null && this.functionAroundWrapper != null) {
+			wrappedFunction = new FunctionInvocationWrapper(function) {
+				@Override
+				Object doApply(Object input) {
+					logger.info("Executing around advise(s)");
+					return functionAroundWrapper.apply(input, function);
+				}
+			};
+		}
+		return wrappedFunction;
+	}
+
 	/*
 	 *
 	 */
@@ -259,7 +285,7 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 	 *
 	 */
 	@SuppressWarnings("rawtypes")
-	public final class FunctionInvocationWrapper implements Function<Object, Object>, Consumer<Object>, Supplier<Object>, Runnable {
+	public class FunctionInvocationWrapper implements Function<Object, Object>, Consumer<Object>, Supplier<Object>, Runnable {
 
 		private final Object target;
 
@@ -288,7 +314,15 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 		 */
 		private Function<Object, Message> enhancer;
 
-		private FunctionInvocationWrapper(String functionDefinition,  Object target, Type inputType, Type outputType) {
+		FunctionInvocationWrapper(FunctionInvocationWrapper function) {
+			this.target = function.target;
+			this.inputType = function.inputType;
+			this.outputType = function.outputType;
+			this.functionDefinition = function.functionDefinition;
+			this.message = this.inputType != null && FunctionTypeUtils.isMessage(this.inputType);
+		}
+
+		FunctionInvocationWrapper(String functionDefinition,  Object target, Type inputType, Type outputType) {
 			this.target = target;
 			this.inputType = this.normalizeType(inputType);
 			this.outputType = this.normalizeType(outputType);
@@ -558,7 +592,7 @@ public class SimpleFunctionRegistry implements FunctionRegistry, FunctionInspect
 		 *
 		 */
 		@SuppressWarnings("unchecked")
-		private Object doApply(Object input) {
+		Object doApply(Object input) {
 			Object result;
 
 			input = this.fluxifyInputIfNecessary(input);

spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java

@@ -507,6 +507,16 @@ public class BeanFactoryAwareFunctionRegistryTests {
 		f.run();
 	}
 
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	@Test
+	public void testWrappedWithAroundAdviseConfiguration() {
+		FunctionCatalog catalog = this.configureCatalog(WrappedWithAroundAdviseConfiguration.class);
+		Function f = catalog.lookup("uppercase");
+		Message result = (Message) f.apply(new GenericMessage<String>("hello"));
+		assertThat(result.getHeaders().get("before")).isEqualTo("foo");
+		assertThat(result.getHeaders().get("after")).isEqualTo("bar");
+	}
+
 	@EnableAutoConfiguration
 	public static class PojoToMessageFunctionCompositionConfiguration {
 
@@ -657,6 +667,29 @@ public class BeanFactoryAwareFunctionRegistryTests {
 		}
 	}
 
+	@EnableAutoConfiguration
+	@Configuration
+	protected static class WrappedWithAroundAdviseConfiguration {
+		@Bean
+		public Function<Message<String>, Message<String>> uppercase() {
+			return v -> MessageBuilder.withPayload(v.getPayload().toUpperCase()).copyHeaders(v.getHeaders()).build();
+		}
+
+		@Bean
+		public FunctionAroundWrapper wrapper() {
+			return new FunctionAroundWrapper() {
+
+				@SuppressWarnings("unchecked")
+				@Override
+				protected Object doApply(Message<byte[]> input, FunctionInvocationWrapper targetFunction) {
+					MessageBuilder.fromMessage(input).setHeader("before", "foo").build();
+					Message<Object> result = (Message<Object>) targetFunction.apply(MessageBuilder.fromMessage(input).setHeader("before", "foo").build());
+					return MessageBuilder.fromMessage(result).setHeader("after", "bar").build();
+				}
+			};
+		}
+	}
+
 	@EnableAutoConfiguration
 	@Configuration
 	protected static class SampleFunctionConfiguration {