Skip to content

Commit 36b1d11

Browse files
committed
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 Signed-off-by: Agustino Alim <agustino.alim@gmail.com>
1 parent a92217a commit 36b1d11

File tree

2 files changed

+64
-3
lines changed

2 files changed

+64
-3
lines changed

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,12 @@ public final Object apply(Object input, FunctionInvocationWrapper targetFunction
4141
boolean functionalTracingEnabled = !StringUtils.hasText(functionalTracingEnabledStr)
4242
|| Boolean.parseBoolean(functionalTracingEnabledStr);
4343
if (functionalTracingEnabled && !(input instanceof Publisher) && input instanceof Message && !FunctionTypeUtils.isCollectionOfMessage(targetFunction.getOutputType())) {
44-
Object result = this.doApply(input, targetFunction);
45-
targetFunction.wrapped = false;
46-
return result;
44+
try {
45+
return this.doApply(input, targetFunction);
46+
}
47+
finally {
48+
targetFunction.wrapped = false;
49+
}
4750
}
4851
else {
4952
return targetFunction.apply(input);

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -717,6 +717,33 @@ public void testWrappedWithAroundAdviseConfiguration() {
717717
assertThat(result.getHeaders().get("after")).isEqualTo("bar");
718718
}
719719

720+
@SuppressWarnings({ "rawtypes", "unchecked" })
721+
@Test
722+
public void testAroundWrapperAppliedOnEveryInvocation() {
723+
FunctionCatalog catalog = this.configureCatalog(AroundWrapperExceptionResetConfiguration.class);
724+
FunctionInvocationWrapper f = catalog.lookup("uppercase");
725+
AtomicInteger wrapperCallCount = (AtomicInteger) this.context.getBean("wrapperCallCount");
726+
727+
// successful invocation
728+
Message result = (Message) f.apply(MessageBuilder.withPayload("hello").build());
729+
assertThat(result.getPayload()).isEqualTo("HELLO");
730+
assertThat(wrapperCallCount.get()).isEqualTo(1);
731+
732+
// failed invocation
733+
try {
734+
f.apply(MessageBuilder.withPayload("exception").build());
735+
}
736+
catch (RuntimeException e) {
737+
// expected
738+
}
739+
assertThat(wrapperCallCount.get()).isEqualTo(2);
740+
741+
// subsequent invocation must still go through the wrapper
742+
result = (Message) f.apply(MessageBuilder.withPayload("world").build());
743+
assertThat(result.getPayload()).isEqualTo("WORLD");
744+
assertThat(wrapperCallCount.get()).isEqualTo(3);
745+
}
746+
720747
@SuppressWarnings({ "rawtypes", "unchecked" })
721748
@Test
722749
public void testEachElementInFluxIsProcessed() {
@@ -1616,4 +1643,35 @@ public Function<Message<?>, Message<?>> myFunction() {
16161643
return msg -> msg;
16171644
}
16181645
}
1646+
1647+
@EnableAutoConfiguration
1648+
@Configuration
1649+
protected static class AroundWrapperExceptionResetConfiguration {
1650+
1651+
@Bean
1652+
public Function<Message<String>, Message<String>> uppercase() {
1653+
return v -> {
1654+
if ("exception".equals(v.getPayload())) {
1655+
throw new RuntimeException("Expected exception");
1656+
}
1657+
return MessageBuilder.withPayload(v.getPayload().toUpperCase(Locale.ROOT)).copyHeaders(v.getHeaders()).build();
1658+
};
1659+
}
1660+
1661+
@Bean
1662+
public AtomicInteger wrapperCallCount() {
1663+
return new AtomicInteger();
1664+
}
1665+
1666+
@Bean
1667+
public FunctionAroundWrapper wrapper(AtomicInteger wrapperCallCount) {
1668+
return new FunctionAroundWrapper() {
1669+
@Override
1670+
protected Object doApply(Object input, FunctionInvocationWrapper targetFunction) {
1671+
wrapperCallCount.incrementAndGet();
1672+
return targetFunction.apply(input);
1673+
}
1674+
};
1675+
}
1676+
}
16191677
}

0 commit comments

Comments
 (0)