diff --git a/pom.xml b/pom.xml index 99e37bb07..ac04794c1 100644 --- a/pom.xml +++ b/pom.xml @@ -67,7 +67,7 @@ 3.0.4-SNAPSHOT 2.0.3-SNAPSHOT 3.1.3 - 3.1.3 + 3.1.4-SNAPSHOT 3.0.4-SNAPSHOT 3.0.4-SNAPSHOT 5.13.2 diff --git a/spring-cloud-sleuth-instrumentation/pom.xml b/spring-cloud-sleuth-instrumentation/pom.xml index f7e966f9f..28bfe02aa 100644 --- a/spring-cloud-sleuth-instrumentation/pom.xml +++ b/spring-cloud-sleuth-instrumentation/pom.xml @@ -75,6 +75,7 @@ org.springframework.cloud spring-cloud-function-context + 3.1.4-SNAPSHOT true diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapper.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapper.java index 95ec169f2..644aa5031 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapper.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapper.java @@ -25,6 +25,7 @@ import org.apache.commons.logging.LogFactory; import org.springframework.cloud.context.scope.refresh.RefreshScopeRefreshedEvent; import org.springframework.cloud.function.context.catalog.FunctionAroundWrapper; import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry; +import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.propagation.Propagator; import org.springframework.context.ApplicationListener; @@ -55,6 +56,8 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper private final Propagator.Getter extractor; + private final TraceMessageHandler traceMessageHandler; + final Map functionToDestinationCache = new ConcurrentHashMap<>(); public TraceFunctionAroundWrapper(Environment environment, Tracer tracer, Propagator propagator, @@ -64,31 +67,40 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper this.propagator = propagator; this.injector = injector; this.extractor = extractor; + this.traceMessageHandler = TraceMessageHandler.forNonSpringIntegration(this.tracer, this.propagator, + this.injector, this.extractor); } @Override protected Object doApply(Message message, SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction) { - TraceMessageHandler traceMessageHandler = TraceMessageHandler.forNonSpringIntegration(this.tracer, - this.propagator, this.injector, this.extractor); - if (log.isDebugEnabled()) { - log.debug("Will retrieve the tracing headers from the message"); + MessageAndSpans invocationMessage = null; + Span span; + if (message == null && targetFunction.isSupplier()) { // Supplier + span = traceMessageHandler.tracer.nextSpan().name(targetFunction.getFunctionDefinition()); } - MessageAndSpans wrappedInputMessage = traceMessageHandler.wrapInputMessage(message, - inputDestination(targetFunction.getFunctionDefinition())); - if (log.isDebugEnabled()) { - log.debug("Wrapped input msg " + wrappedInputMessage); + else { + if (log.isDebugEnabled()) { + log.debug("Will retrieve the tracing headers from the message"); + } + invocationMessage = traceMessageHandler.wrapInputMessage(message, + inputDestination(targetFunction.getFunctionDefinition())); + if (log.isDebugEnabled()) { + log.debug("Wrapped input msg " + invocationMessage); + } + span = invocationMessage.childSpan; } + Object result; Throwable throwable = null; - try (Tracer.SpanInScope ws = tracer.withSpan(wrappedInputMessage.childSpan.start())) { - result = targetFunction.apply(wrappedInputMessage.msg); + try (Tracer.SpanInScope ws = tracer.withSpan(span.start())) { + result = invocationMessage == null ? targetFunction.get() : targetFunction.apply(invocationMessage.msg); } catch (Exception e) { throwable = e; throw e; } finally { - traceMessageHandler.afterMessageHandled(wrappedInputMessage.childSpan, throwable); + traceMessageHandler.afterMessageHandled(span, throwable); } if (result == null) { if (log.isDebugEnabled()) { @@ -96,21 +108,73 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper } return null; } - Message msgResult = toMessage(result); - MessageAndSpan wrappedOutputMessage = traceMessageHandler.wrapOutputMessage(msgResult, - wrappedInputMessage.parentSpan, outputDestination(targetFunction.getFunctionDefinition())); + Message msgResult = toMessage(result); + + MessageAndSpan wrappedOutputMessage; + if (invocationMessage != null) { + wrappedOutputMessage = traceMessageHandler.wrapOutputMessage(msgResult, invocationMessage.parentSpan, + outputDestination(targetFunction.getFunctionDefinition())); + } + else { + wrappedOutputMessage = this.getMessageAndSpans(msgResult, targetFunction.getFunctionDefinition(), span); + } if (log.isDebugEnabled()) { log.debug("Wrapped output msg " + wrappedOutputMessage); } traceMessageHandler.afterMessageHandled(wrappedOutputMessage.span, null); return wrappedOutputMessage.msg; + + // if (log.isDebugEnabled()) { + // log.debug("Will retrieve the tracing headers from the message"); + // } + // MessageAndSpans wrappedInputMessage = + // traceMessageHandler.wrapInputMessage(message, + // inputDestination(targetFunction.getFunctionDefinition())); + // if (log.isDebugEnabled()) { + // log.debug("Wrapped input msg " + wrappedInputMessage); + // } + // Object result; + // Throwable throwable = null; + // try (Tracer.SpanInScope ws = + // tracer.withSpan(wrappedInputMessage.childSpan.start())) { + // result = targetFunction.apply(wrappedInputMessage.msg); + // } + // catch (Exception e) { + // throwable = e; + // throw e; + // } + // finally { + // traceMessageHandler.afterMessageHandled(wrappedInputMessage.childSpan, + // throwable); + // } + // if (result == null) { + // if (log.isDebugEnabled()) { + // log.debug("Returned message is null - we have a consumer"); + // } + // return null; + // } + // Message msgResult = toMessage(result); + // MessageAndSpan wrappedOutputMessage = + // traceMessageHandler.wrapOutputMessage(msgResult, + // wrappedInputMessage.parentSpan, + // outputDestination(targetFunction.getFunctionDefinition())); + // if (log.isDebugEnabled()) { + // log.debug("Wrapped output msg " + wrappedOutputMessage); + // } + // traceMessageHandler.afterMessageHandled(wrappedOutputMessage.span, null); + // return wrappedOutputMessage.msg; + } - private Message toMessage(Object result) { + MessageAndSpan getMessageAndSpans(Message resultMessage, String name, Span spanFromMessage) { + return traceMessageHandler.wrapOutputMessage(resultMessage, spanFromMessage, outputDestination(name)); + } + + private Message toMessage(Object result) { if (!(result instanceof Message)) { return MessageBuilder.withPayload(result).build(); } - return (Message) result; + return (Message) result; } String inputDestination(String functionDefinition) { diff --git a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java index 45973688c..b1c8d993f 100644 --- a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java +++ b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java @@ -16,8 +16,23 @@ package org.springframework.cloud.sleuth.instrument.messaging; +import java.util.Collections; +import java.util.function.Supplier; + +import com.fasterxml.jackson.databind.ObjectMapper; import org.junit.jupiter.api.Test; +import org.springframework.cloud.function.context.FunctionRegistration; +import org.springframework.cloud.function.context.FunctionType; +import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry; +import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper; +import org.springframework.cloud.function.context.config.JsonMessageConverter; +import org.springframework.cloud.function.json.JacksonMapper; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.tracer.SimpleTracer; +import org.springframework.core.convert.support.DefaultConversionService; +import org.springframework.messaging.Message; +import org.springframework.messaging.converter.CompositeMessageConverter; import org.springframework.mock.env.MockEnvironment; import static org.assertj.core.api.Assertions.assertThat; @@ -25,6 +40,30 @@ import static org.assertj.core.api.BDDAssertions.then; class TraceFunctionAroundWrapperTests { + @Test + void test_tracing_with_supplier() { + CompositeMessageConverter messageConverter = new CompositeMessageConverter( + Collections.singletonList(new JsonMessageConverter(new JacksonMapper(new ObjectMapper())))); + + SimpleTracer tracer = new SimpleTracer(); + TraceFunctionAroundWrapper wrapper = new TraceFunctionAroundWrapper(null, tracer, null, null, null) { + @Override + MessageAndSpan getMessageAndSpans(Message resultMessage, String name, Span spanFromMessage) { + return new MessageAndSpan(resultMessage, spanFromMessage); + } + }; + + FunctionRegistration registration = new FunctionRegistration<>(new Greeter(), "greeter") + .type(FunctionType.of(Greeter.class)); + SimpleFunctionRegistry catalog = new SimpleFunctionRegistry(new DefaultConversionService(), messageConverter, + new JacksonMapper(new ObjectMapper())); + catalog.register(registration); + FunctionInvocationWrapper function = catalog.lookup("greeter"); + Message result = (Message) wrapper.apply(null, function); + assertThat(result.getPayload()).isEqualTo("hello"); + assertThat(tracer.getOnlySpan().name).isEqualTo("greeter"); + } + @Test void should_clear_cache_on_refresh() { TraceFunctionAroundWrapper wrapper = new TraceFunctionAroundWrapper(null, null, null, null, null); @@ -66,4 +105,13 @@ class TraceFunctionAroundWrapperTests { assertThat(wrapper.outputDestination("marcin")).isEqualTo("bob"); } + private static class Greeter implements Supplier { + + @Override + public String get() { + return "hello"; + } + + } + } diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-messaging-tests/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/BraveTraceFunctionAwareWrapperTests.java b/tests/brave/spring-cloud-sleuth-instrumentation-messaging-tests/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/BraveTraceFunctionAwareWrapperTests.java new file mode 100644 index 000000000..f94705c42 --- /dev/null +++ b/tests/brave/spring-cloud-sleuth-instrumentation-messaging-tests/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/BraveTraceFunctionAwareWrapperTests.java @@ -0,0 +1,56 @@ +/* + * Copyright 2021-2021 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.messaging; + +import brave.sampler.Sampler; + +import org.springframework.cloud.sleuth.brave.BraveTestSpanHandler; +import org.springframework.cloud.sleuth.test.TestSpanHandler; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +/** + * @author Oleg Zhurakousky + */ +public class BraveTraceFunctionAwareWrapperTests extends TraceFunctionAroundWrapperTests { + + @Override + protected Class configuration() { + return App.class; + } + + @Configuration(proxyBeanMethods = false) + static class App { + + @Bean + TestSpanHandler testSpanHandlerSupplier(brave.test.TestSpanHandler testSpanHandler) { + return new BraveTestSpanHandler(testSpanHandler); + } + + @Bean + Sampler alwaysSampler() { + return Sampler.ALWAYS_SAMPLE; + } + + @Bean + brave.test.TestSpanHandler braveTestSpanHandler() { + return new brave.test.TestSpanHandler(); + } + + } + +} diff --git a/tests/common/pom.xml b/tests/common/pom.xml index b318e7bde..edd2d0aaf 100644 --- a/tests/common/pom.xml +++ b/tests/common/pom.xml @@ -39,6 +39,12 @@ org.springframework.cloud spring-cloud-sleuth-instrumentation + + org.springframework.cloud + spring-cloud-function-context + + true + org.springframework.boot spring-boot-starter-test diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java new file mode 100644 index 000000000..c9b802c07 --- /dev/null +++ b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java @@ -0,0 +1,93 @@ +/* + * Copyright 2021-2021 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.messaging; + +import java.util.function.Function; +import java.util.function.Supplier; + +import org.junit.jupiter.api.Test; + +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.builder.SpringApplicationBuilder; +import org.springframework.cloud.function.context.FunctionCatalog; +import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper; +import org.springframework.cloud.sleuth.test.TestSpanHandler; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.messaging.Message; +import org.springframework.messaging.support.MessageBuilder; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Oleg Zhurakousky + * + */ +public abstract class TraceFunctionAroundWrapperTests { + + @Test + public void test_tracing_with_supplier() { + try (ConfigurableApplicationContext context = new SpringApplicationBuilder(configuration(), + SampleConfiguration.class).run("--logging.level.org.springframework.cloud.function=DEBUG", + "--spring.main.lazy-initialization=true");) { + TestSpanHandler spanHandler = context.getBean(TestSpanHandler.class); + assertThat(spanHandler.reportedSpans()).isEmpty(); + FunctionCatalog catalog = context.getBean(FunctionCatalog.class); + FunctionInvocationWrapper function = catalog.lookup("greeter"); + function.setSkipOutputConversion(true); + Message result = (Message) function.get(); + assertThat(result.getPayload()).isEqualTo("hello"); + assertThat(spanHandler.reportedSpans().size()).isEqualTo(2); + assertThat(((String) result.getHeaders().get("b3"))).contains(spanHandler.get(0).getTraceId()); + } + } + + @Test + public void test_tracing_with_function() { + try (ConfigurableApplicationContext context = new SpringApplicationBuilder(configuration(), + SampleConfiguration.class).run("--logging.level.org.springframework.cloud.function=DEBUG", + "--spring.main.lazy-initialization=true");) { + TestSpanHandler spanHandler = context.getBean(TestSpanHandler.class); + assertThat(spanHandler.reportedSpans()).isEmpty(); + FunctionCatalog catalog = context.getBean(FunctionCatalog.class); + FunctionInvocationWrapper function = catalog.lookup("uppercase"); + function.setSkipOutputConversion(true); + Message result = (Message) function.apply(MessageBuilder.withPayload("hello").build()); + assertThat(result.getPayload()).isEqualTo("HELLO"); + assertThat(spanHandler.reportedSpans().size()).isEqualTo(3); + assertThat(((String) result.getHeaders().get("b3"))).contains(spanHandler.get(0).getTraceId()); + } + } + + protected abstract Class configuration(); + + @EnableAutoConfiguration + public static class SampleConfiguration { + + @Bean + public Supplier greeter() { + return () -> "hello"; + } + + @Bean + public Function uppercase() { + return v -> v.toUpperCase(); + } + + } + +};