From 747ff21237b8e00b0906597a95598cd3cdba3970 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Wed, 15 Jun 2022 10:39:26 +0200 Subject: [PATCH] Fixing invalid scoping in TraceHandlerFunction; fixes gh-2026 --- .../instrument/web/TraceHandlerFunction.java | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHandlerFunction.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHandlerFunction.java index 8df043c86..4005a505f 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHandlerFunction.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHandlerFunction.java @@ -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 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 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() {