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

Clearing thread locals whenever an RSocket request arrives

Browse Source
This commit is contained in:
Marcin Grzejszczak
2022-12-13 19:04:45 +01:00
parent e953f2fa48
commit 6d55869656
2 changed files with 16 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/rsocket/TracingRequesterRSocketProxy.java
View File

@@ -71,13 +71,19 @@ public class TracingRequesterRSocketProxy extends RSocketProxy {
this.isZipkinPropagationEnabled = isZipkinPropagationEnabled;
}
private void clearThreadLocal() {
this.tracer.withSpan(null);
}
@Override
public Mono<Void> fireAndForget(Payload payload) {
clearThreadLocal();
return setSpan(super::fireAndForget, payload, FrameType.REQUEST_FNF);
}
@Override
public Mono<Payload> requestResponse(Payload payload) {
clearThreadLocal();
return setSpan(super::requestResponse, payload, FrameType.REQUEST_RESPONSE);
}
@@ -138,11 +144,13 @@ public class TracingRequesterRSocketProxy extends RSocketProxy {
@Override
public Flux<Payload> requestStream(Payload payload) {
clearThreadLocal();
return Flux.deferContextual(contextView -> setSpan(super::requestStream, payload, contextView));
}
@Override
public Flux<Payload> requestChannel(Publisher<Payload> inbound) {
clearThreadLocal();
return Flux.from(inbound).switchOnFirst((firstSignal, flux) -> {
final Payload firstPayload = firstSignal.get();
if (firstPayload != null) {

8
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/rsocket/TracingResponderRSocketProxy.java
View File

@@ -76,6 +76,7 @@ public class TracingResponderRSocketProxy extends RSocketProxy {
@Override
public Mono<Void> fireAndForget(Payload payload) {
clearThreadLocal();
// called on Netty EventLoop
// there can't be trace context in thread local here
Span handle = consumerSpanBuilder(payload.sliceMetadata(), FrameType.REQUEST_FNF);
@@ -86,8 +87,13 @@ public class TracingResponderRSocketProxy extends RSocketProxy {
return ReactorSleuth.tracedMono(this.tracer, handle, () -> super.fireAndForget(newPayload));
}
private void clearThreadLocal() {
this.tracer.withSpan(null);
}
@Override
public Mono<Payload> requestResponse(Payload payload) {
clearThreadLocal();
Span handle = consumerSpanBuilder(payload.sliceMetadata(), FrameType.REQUEST_RESPONSE);
if (log.isDebugEnabled()) {
log.debug("Created consumer span " + handle);
@@ -98,6 +104,7 @@ public class TracingResponderRSocketProxy extends RSocketProxy {
@Override
public Flux<Payload> requestStream(Payload payload) {
clearThreadLocal();
Span handle = consumerSpanBuilder(payload.sliceMetadata(), FrameType.REQUEST_STREAM);
if (log.isDebugEnabled()) {
log.debug("Created consumer span " + handle);
@@ -108,6 +115,7 @@ public class TracingResponderRSocketProxy extends RSocketProxy {
@Override
public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
clearThreadLocal();
return Flux.from(payloads).switchOnFirst((firstSignal, flux) -> {
final Payload firstPayload = firstSignal.get();
if (firstPayload != null) {