Skip to content

Commit 774e43a

Browse files
agustino-limolegz
authored andcommitted
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>
1 parent d79ec64 commit 774e43a

2 files changed

Lines changed: 64 additions & 3 deletions

File tree

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
@@ -708,6 +708,33 @@ public void testWrappedWithAroundAdviseConfiguration() {
708708
assertThat(result.getHeaders().get("after")).isEqualTo("bar");
709709
}
710710

711+
@SuppressWarnings({ "rawtypes", "unchecked" })
712+
@Test
713+
public void testAroundWrapperAppliedOnEveryInvocation() {
714+
FunctionCatalog catalog = this.configureCatalog(AroundWrapperExceptionResetConfiguration.class);
715+
FunctionInvocationWrapper f = catalog.lookup("uppercase");
716+
AtomicInteger wrapperCallCount = (AtomicInteger) this.context.getBean("wrapperCallCount");
717+
718+
// successful invocation
719+
Message result = (Message) f.apply(MessageBuilder.withPayload("hello").build());
720+
assertThat(result.getPayload()).isEqualTo("HELLO");
721+
assertThat(wrapperCallCount.get()).isEqualTo(1);
722+
723+
// failed invocation
724+
try {
725+
f.apply(MessageBuilder.withPayload("exception").build());
726+
}
727+
catch (RuntimeException e) {
728+
// expected
729+
}
730+
assertThat(wrapperCallCount.get()).isEqualTo(2);
731+
732+
// subsequent invocation must still go through the wrapper
733+
result = (Message) f.apply(MessageBuilder.withPayload("world").build());
734+
assertThat(result.getPayload()).isEqualTo("WORLD");
735+
assertThat(wrapperCallCount.get()).isEqualTo(3);
736+
}
737+
711738
@SuppressWarnings({ "rawtypes", "unchecked" })
712739
@Test
713740
public void testEachElementInFluxIsProcessed() {
@@ -1607,4 +1634,35 @@ public Function<Message<?>, Message<?>> myFunction() {
16071634
return msg -> msg;
16081635
}
16091636
}
1637+
1638+
@EnableAutoConfiguration
1639+
@Configuration
1640+
protected static class AroundWrapperExceptionResetConfiguration {
1641+
1642+
@Bean
1643+
public Function<Message<String>, Message<String>> uppercase() {
1644+
return v -> {
1645+
if ("exception".equals(v.getPayload())) {
1646+
throw new RuntimeException("Expected exception");
1647+
}
1648+
return MessageBuilder.withPayload(v.getPayload().toUpperCase(Locale.ROOT)).copyHeaders(v.getHeaders()).build();
1649+
};
1650+
}
1651+
1652+
@Bean
1653+
public AtomicInteger wrapperCallCount() {
1654+
return new AtomicInteger();
1655+
}
1656+
1657+
@Bean
1658+
public FunctionAroundWrapper wrapper(AtomicInteger wrapperCallCount) {
1659+
return new FunctionAroundWrapper() {
1660+
@Override
1661+
protected Object doApply(Object input, FunctionInvocationWrapper targetFunction) {
1662+
wrapperCallCount.incrementAndGet();
1663+
return targetFunction.apply(input);
1664+
}
1665+
};
1666+
}
1667+
}
16101668
}

0 commit comments

Comments
 (0)