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

Fixing invalid scoping in TraceHandlerFunction; fixes gh-2026

Browse Source
This commit is contained in:
Marcin Grzejszczak
2022-06-15 10:39:26 +02:00
parent 871fc9fb11
commit b3338ff739
2 changed files with 15 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
spring-cloud-sleuth-brave/src/test/java/org/springframework/cloud/sleuth/brave/bridge/W3CBaggagePropagatorTest.java
View File

@@ -110,8 +110,8 @@ class W3CBaggagePropagatorTest {
}
/**
* We need to use {@link HttpServletRequestWrapper} for the carrier for this test, since it is what combines the
* multiple baggage headers into one.
* We need to use {@link HttpServletRequestWrapper} for the carrier for this test,
* since it is what combines the multiple baggage headers into one.
*/
@Test
void extract_multipleBaggageHeaders() {
@@ -120,14 +120,11 @@ class W3CBaggagePropagatorTest {
mockRequest.addHeader("baggage", "key2=value2,key3=value3");
HttpServletRequestWrapper carrier = (HttpServletRequestWrapper) HttpServletRequestWrapper.create(mockRequest);
TraceContextOrSamplingFlags contextWithBaggage = propagator
.contextWithBaggage(carrier, context(), HttpServletRequestWrapper::header);
TraceContextOrSamplingFlags contextWithBaggage = propagator.contextWithBaggage(carrier, context(),
HttpServletRequestWrapper::header);
Map<String, String> baggageEntries = BaggageField.getAllValues(contextWithBaggage);
assertThat(baggageEntries)
.hasSize(3)
.containsEntry("key1", "value1")
.containsEntry("key2", "value2")
assertThat(baggageEntries).hasSize(3).containsEntry("key1", "value1").containsEntry("key2", "value2")
.containsEntry("key3", "value3");
}

21
spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHandlerFunction.java
View File

@@ -16,7 +16,7 @@
package org.springframework.cloud.sleuth.instrument.web;
import java.util.concurrent.atomic.AtomicReference;
import java.util.Optional;
import reactor.core.publisher.Mono;
@@ -47,16 +47,15 @@ public class TraceHandlerFunction implements HandlerFunction {
@Override
public Mono<?> handle(ServerRequest serverRequest) {
AtomicReference<CurrentTraceContext.Scope> scope = new AtomicReference<>();
return Mono.just(scope)
.doFirst(() -> serverRequest.attribute(TraceWebFilter.TRACE_REQUEST_ATTR)
.ifPresent(span -> scope.set(currentTraceContext().maybeScope(((Span) span).context()))))
.flatMap(r -> this.delegate.handle(serverRequest)).doFinally(signalType -> {
CurrentTraceContext.Scope spanInScope = scope.get();
if (spanInScope != null) {
spanInScope.close();
}
});
Optional<Object> spanOptional = serverRequest.attribute(TraceWebFilter.TRACE_REQUEST_ATTR);
if (!spanOptional.isPresent()) {
return this.delegate.handle(serverRequest);
}
return Mono.justOrEmpty(spanOptional).cast(Span.class).flatMap((Span span) -> {
try (CurrentTraceContext.Scope scope = currentTraceContext().maybeScope(span.context())) {
return this.delegate.handle(serverRequest);
}
});
}
private CurrentTraceContext currentTraceContext() {