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

Polish

Browse Source
This commit is contained in:
Marcin Grzejszczak
2022-01-17 11:17:22 +01:00
parent 49e5372b7c
commit e4ff2dd051
1 changed files with 22 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -24,7 +24,6 @@ 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;
@@ -166,20 +165,19 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
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 -> {
msg.end();
msg.handle();
}).map(msgResult -> {
MessageAndSpan messageAndSpan = traceMessageHandler.wrapOutputMessage(msgResult,
msg.messageAndSpans.parentSpan, outputDestination(targetFunction.getFunctionDefinition()));
traceMessageHandler.afterMessageHandled(messageAndSpan.span, null);
return messageAndSpan.msg;
})
msg.end();
msg.handle();
}).map(msgResult -> {
MessageAndSpan messageAndSpan = traceMessageHandler.wrapOutputMessage(msgResult,
msg.messageAndSpans.parentSpan, outputDestination(targetFunction.getFunctionDefinition()));
traceMessageHandler.afterMessageHandled(messageAndSpan.span, null);
return messageAndSpan.msg;
})
// TODO: Fix me when this is resolved in Reactor
// .doOnSubscribe(__ -> scope.close())
.doOnError(msg::error).doFinally(signalType -> {
@@ -218,20 +216,19 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
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 -> {
msg.end();
msg.handle();
}).map(msgResult -> {
MessageAndSpan messageAndSpan = traceMessageHandler.wrapOutputMessage(msgResult,
msg.messageAndSpans.parentSpan, outputDestination(targetFunction.getFunctionDefinition()));
traceMessageHandler.afterMessageHandled(messageAndSpan.span, null);
return messageAndSpan.msg;
})
msg.end();
msg.handle();
}).map(msgResult -> {
MessageAndSpan messageAndSpan = traceMessageHandler.wrapOutputMessage(msgResult,
msg.messageAndSpans.parentSpan, outputDestination(targetFunction.getFunctionDefinition()));
traceMessageHandler.afterMessageHandled(messageAndSpan.span, null);
return messageAndSpan.msg;
})
// TODO: Fix me when this is resolved in Reactor
// .doOnSubscribe(__ -> scope.close())
.doOnError(msg::error).doFinally(signalType -> {
@@ -275,9 +272,9 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
}
Mono mono = (Mono) publisher;
publisher = ReactorSleuth.tracedMono(tracer, tracer.currentTraceContext(),
targetFunction.getFunctionDefinition(), () -> mono, (msg, s) -> {
customizedInputMessageSpan(s, msg instanceof Message ? (Message) msg : null);
}).map(object -> toMessage(object))
targetFunction.getFunctionDefinition(), () -> mono, (msg, s) -> {
customizedInputMessageSpan(s, msg instanceof Message ? (Message) msg : null);
}).map(object -> toMessage(object))
.map(object -> this.getMessageAndSpans((Message) object, targetFunction.getFunctionDefinition(),
setNameAndTag(targetFunction, tracer.currentSpan())))
.doOnNext(wrappedOutputMessage -> customizedOutputMessageSpan(
@@ -292,9 +289,9 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
}
Flux flux = (Flux) publisher;
publisher = ReactorSleuth.tracedFlux(tracer, tracer.currentTraceContext(),
targetFunction.getFunctionDefinition(), () -> flux, (msg, s) -> {
customizedInputMessageSpan(s, msg instanceof Message ? (Message) msg : null);
}).map(object -> toMessage(object))
targetFunction.getFunctionDefinition(), () -> flux, (msg, s) -> {
customizedInputMessageSpan(s, msg instanceof Message ? (Message) msg : null);
}).map(object -> toMessage(object))
.map(object -> this.getMessageAndSpans((Message) object, targetFunction.getFunctionDefinition(),
setNameAndTag(targetFunction, tracer.currentSpan())))
.doOnNext(wrappedOutputMessage -> customizedOutputMessageSpan(