GH-1327 Fix wrapped flag not reset on exception in FunctionAroundWrapper

Ensure the wrapped flag on FunctionInvocationWrapper is reset in a
finally block so that an exception from doApply does not leave it stuck
at true, which would cause subsequent invocations to bypass the
FunctionAroundWrapper (breaking tracing context propagation).

Resolves #1327
Resolves #1332
Signed-off-by: Agustino Alim <agustino.alim@gmail.com>

Author: agustino-lim

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

@@ -41,9 +41,12 @@ public final Object apply(Object input, FunctionInvocationWrapper targetFunction
 		boolean functionalTracingEnabled = !StringUtils.hasText(functionalTracingEnabledStr)
 				|| Boolean.parseBoolean(functionalTracingEnabledStr);
 		if (functionalTracingEnabled && !(input instanceof Publisher) && input instanceof Message && !FunctionTypeUtils.isCollectionOfMessage(targetFunction.getOutputType())) {
-			Object result = this.doApply(input, targetFunction);
-			targetFunction.wrapped = false;
-			return result;
+			try {
+				return this.doApply(input, targetFunction);
+			}
+			finally {
+				targetFunction.wrapped = false;
+			}
 		}
 		else {
 			return targetFunction.apply(input);

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

@@ -708,6 +708,33 @@ public void testWrappedWithAroundAdviseConfiguration() {
 		assertThat(result.getHeaders().get("after")).isEqualTo("bar");
 	}
 
+	@SuppressWarnings({ "rawtypes", "unchecked" })
+	@Test
+	public void testAroundWrapperAppliedOnEveryInvocation() {
+		FunctionCatalog catalog = this.configureCatalog(AroundWrapperExceptionResetConfiguration.class);
+		FunctionInvocationWrapper f = catalog.lookup("uppercase");
+		AtomicInteger wrapperCallCount = (AtomicInteger) this.context.getBean("wrapperCallCount");
+
+		// successful invocation
+		Message result = (Message) f.apply(MessageBuilder.withPayload("hello").build());
+		assertThat(result.getPayload()).isEqualTo("HELLO");
+		assertThat(wrapperCallCount.get()).isEqualTo(1);
+
+		// failed invocation
+		try {
+			f.apply(MessageBuilder.withPayload("exception").build());
+		}
+		catch (RuntimeException e) {
+			// expected
+		}
+		assertThat(wrapperCallCount.get()).isEqualTo(2);
+
+		// subsequent invocation must still go through the wrapper
+		result = (Message) f.apply(MessageBuilder.withPayload("world").build());
+		assertThat(result.getPayload()).isEqualTo("WORLD");
+		assertThat(wrapperCallCount.get()).isEqualTo(3);
+	}
+
 	@SuppressWarnings({ "rawtypes", "unchecked" })
 	@Test
 	public void testEachElementInFluxIsProcessed() {
@@ -1607,4 +1634,35 @@ public Function<Message<?>, Message<?>> myFunction() {
 			return msg -> msg;
 		}
 	}
+
+	@EnableAutoConfiguration
+	@Configuration
+	protected static class AroundWrapperExceptionResetConfiguration {
+
+		@Bean
+		public Function<Message<String>, Message<String>> uppercase() {
+			return v -> {
+				if ("exception".equals(v.getPayload())) {
+					throw new RuntimeException("Expected exception");
+				}
+				return MessageBuilder.withPayload(v.getPayload().toUpperCase(Locale.ROOT)).copyHeaders(v.getHeaders()).build();
+			};
+		}
+
+		@Bean
+		public AtomicInteger wrapperCallCount() {
+			return new AtomicInteger();
+		}
+
+		@Bean
+		public FunctionAroundWrapper wrapper(AtomicInteger wrapperCallCount) {
+			return new FunctionAroundWrapper() {
+				@Override
+				protected Object doApply(Object input, FunctionInvocationWrapper targetFunction) {
+					wrapperCallCount.incrementAndGet();
+					return targetFunction.apply(input);
+				}
+			};
+		}
+	}
 }