SpringCloud/spring-cloud-sleuth
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Added support for Flux->Mono, Mono->Flux; fixes gh-2077

Browse Source
This commit is contained in:
Marcin Grzejszczak
2022-01-17 11:14:40 +01:00
parent e7850379e3
commit 49e5372b7c
2 changed files with 77 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapper.java
View File

@@ -24,6 +24,7 @@ import java.util.concurrent.ConcurrentHashMap;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.jetbrains.annotations.NotNull;
import org.reactivestreams.Publisher;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
@@ -158,7 +159,16 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
if (targetFunction.isConsumer()) {
return targetFunction.apply(reactorStreamConsumer(mono));
}
final Mono<Message> function = ((Mono<Message>) targetFunction.apply(mono));
final Publisher<Message> function = ((Publisher<Message>) targetFunction.apply(mono));
if (function instanceof Mono) {
return messageMono(targetFunction, (Mono<Message>) function);
}
return messageFlux(targetFunction, (Flux<Message>) function);
}
@NotNull
private Mono<Message> messageMono(SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction,
Mono<Message> function) {
return Mono.deferContextual(contextView -> {
MessageAndSpansAndScope msg = contextView.get(MessageAndSpansAndScope.class);
return function.doOnNext(message -> {
@@ -201,7 +211,16 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
if (targetFunction.isConsumer()) {
return targetFunction.apply(reactorStreamConsumer(flux));
}
final Flux<Message> function = ((Flux<Message>) targetFunction.apply(flux));
final Publisher<Message> function = ((Publisher<Message>) targetFunction.apply(flux));
if (function instanceof Mono) {
return messageMono(targetFunction, (Mono<Message>) function);
}
return messageFlux(targetFunction, (Flux<Message>) function);
}
@NotNull
private Flux<Message> messageFlux(SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction,
Flux<Message> function) {
return Flux.deferContextual(contextView -> {
MessageAndSpansAndScope msg = contextView.get(MessageAndSpansAndScope.class);
return function.doOnNext(message -> {

56
spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java
View File

@@ -201,6 +201,25 @@ class TraceFunctionAroundWrapperTests {
assertThatAllSpansAreStartedAndStopped();
}
@Test
void should_trace_when_reactive_mono_to_flux_function() {
FunctionRegistration<ReactiveMonoToFluxFunction> registration = new FunctionRegistration<>(
new ReactiveMonoToFluxFunction(), "greeter").type(FunctionType.of(ReactiveMonoToFluxFunction.class));
catalog.register(registration);
FunctionInvocationWrapper function = catalog.lookup("greeter");
Message result = ((Flux<Message>) wrapper.apply(
Mono.just(MessageBuilder.withPayload("hello").setHeader("superHeader", "someValue").build()), function))
.blockFirst(Duration.ofSeconds(5));
assertThat(result.getPayload()).isEqualTo("HELLO");
assertThat(tracer.spans).hasSize(3);
assertThat(tracer.spans.get(0).name).isEqualTo("handle");
assertThat(tracer.spans.get(1).name).isEqualTo("greeter");
assertThat(tracer.spans.get(2).name).isEqualTo("send");
assertThatAllSpansAreStartedAndStopped();
}
@Test
void should_trace_when_reactive_flux_supplier() {
FunctionRegistration<ReactiveFluxGreeter> registration = new FunctionRegistration<>(new ReactiveFluxGreeter(),
@@ -234,6 +253,23 @@ class TraceFunctionAroundWrapperTests {
assertThatAllSpansAreStartedAndStopped();
}
@Test
void should_trace_when_reactive_flux_function_returns_mono() {
FunctionRegistration<ReactiveFluxToMonoFunction> registration = new FunctionRegistration<>(
new ReactiveFluxToMonoFunction(), "greeter").type(FunctionType.of(ReactiveFluxToMonoFunction.class));
catalog.register(registration);
FunctionInvocationWrapper function = catalog.lookup("greeter");
((Mono<Void>) wrapper.apply(
Flux.just(MessageBuilder.withPayload("hello").setHeader("superHeader", "someValue").build()), function))
.block(Duration.ofSeconds(5));
assertThat(tracer.spans).hasSize(2);
assertThat(tracer.spans.get(0).name).isEqualTo("handle");
assertThat(tracer.spans.get(1).name).isEqualTo("greeter");
assertThatAllSpansAreStartedAndStopped();
}
@Test
void should_trace_when_reactive_flux_consumer() {
ReactiveFluxGreeterConsumer consumer = new ReactiveFluxGreeterConsumer(this.tracer);
@@ -352,6 +388,16 @@ class TraceFunctionAroundWrapperTests {
}
private static class ReactiveMonoToFluxFunction implements Function<Mono<Message<String>>, Flux<Message<String>>> {
@Override
public Flux<Message<String>> apply(Mono<Message<String>> in) {
return Flux
.from(in.map(s -> MessageBuilder.fromMessage(s).withPayload(s.getPayload().toUpperCase()).build()));
}
}
private static class ReactiveFluxGreeter implements Supplier<Flux<Message<String>>> {
@Override
@@ -370,6 +416,16 @@ class TraceFunctionAroundWrapperTests {
}
private static class ReactiveFluxToMonoFunction implements Function<Flux<Message<String>>, Mono<Void>> {
@Override
public Mono<Void> apply(Flux<Message<String>> in) {
return in.map(s -> MessageBuilder.fromMessage(s).withPayload(s.getPayload().toUpperCase()).build())
.then(Mono.empty());
}
}
private static class ReactiveFluxGreeterConsumer implements Consumer<Flux<Message<String>>> {
String result;