diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java index 0e9889403..4773fab11 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java @@ -364,6 +364,7 @@ public class TraceFilter extends GenericFilterBean { final Span span; final Tracer.SpanInScope scope; + SpanAndScope(Span span, Tracer.SpanInScope scope) { this.span = span; this.scope = scope; diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java index 57f0320d6..6768bd452 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java @@ -133,7 +133,8 @@ public class TraceWebFilter implements WebFilter, Ordered { .map(c -> c.put(CONTEXT_ERROR, t))) .flatMap(c -> { //reactivate span from context - Span span = c.getOrDefault(Span.class, tracer().nextSpan().start()); + SpanAndScope spanAndScope = c.getOrDefault(SpanAndScope.class, defaultSpanAndScope()); + Span span = spanAndScope.span; Mono continuation; Throwable t = null; if (c.hasKey(CONTEXT_ERROR)) { @@ -151,12 +152,14 @@ public class TraceWebFilter implements WebFilter, Ordered { } addResponseTagsForSpanWithoutParent(exchange, response, span); handler().handleSend(response, t, span); + spanAndScope.scope.close(); return continuation; }) .subscriberContext(c -> { Span span; - if (c.hasKey(Span.class)) { - Span parent = c.get(Span.class); + if (c.hasKey(SpanAndScope.class)) { + SpanAndScope spanAndScope = c.get(SpanAndScope.class); + Span parent = spanAndScope.span; span = tracer() .nextSpan(TraceContextOrSamplingFlags.create(parent.context())) .start(); @@ -184,10 +187,15 @@ public class TraceWebFilter implements WebFilter, Ordered { } } } - return c.put(Span.class, span); + return c.put(SpanAndScope.class, new SpanAndScope(span, tracer().withSpanInScope(span))); })); } + private SpanAndScope defaultSpanAndScope() { + Span defaultSpan = tracer().nextSpan().start(); + return new SpanAndScope(defaultSpan, tracer().withSpanInScope(defaultSpan)); + } + private void addResponseTagsForSpanWithoutParent(ServerWebExchange exchange, ServerHttpResponse response, Span span) { if (spanWithoutParent(exchange) && response.getStatusCode() != null @@ -221,6 +229,22 @@ public class TraceWebFilter implements WebFilter, Ordered { } } + class SpanAndScope { + + final Span span; + final Tracer.SpanInScope scope; + + SpanAndScope(Span span, Tracer.SpanInScope scope) { + this.span = span; + this.scope = scope; + } + + SpanAndScope() { + this.span = null; + this.scope = null; + } + } + private void addClassNameTag(Object handler, Span span) { String className; if (handler instanceof HandlerMethod) { diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java index 942cbcdf8..127b373f2 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java @@ -1,14 +1,10 @@ package org.springframework.cloud.sleuth.instrument.web; import brave.sampler.Sampler; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Hooks; -import reactor.core.publisher.Mono; -import reactor.core.scheduler.Schedulers; -import org.assertj.core.api.BDDAssertions; import org.awaitility.Awaitility; import org.junit.BeforeClass; import org.junit.Test; +import org.slf4j.MDC; import org.springframework.boot.WebApplicationType; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.autoconfigure.security.reactive.ReactiveSecurityAutoConfiguration; @@ -24,6 +20,12 @@ import org.springframework.web.bind.annotation.PathVariable; import org.springframework.web.bind.annotation.RestController; import org.springframework.web.reactive.function.client.ClientResponse; import org.springframework.web.reactive.function.client.WebClient; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Hooks; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +import static org.assertj.core.api.BDDAssertions.then; public class TraceWebFluxTests { @@ -47,10 +49,10 @@ public class TraceWebFluxTests { Awaitility.await().untilAsserted(() -> { ClientResponse response = exchange.block(); - BDDAssertions.then(response.statusCode().value()).isEqualTo(200); + then(response.statusCode().value()).isEqualTo(200); }); - BDDAssertions.then(accumulator.getSpans()).hasSize(1); - BDDAssertions.then(accumulator.getSpans().get(0).tags()) + then(accumulator.getSpans()).hasSize(1); + then(accumulator.getSpans().get(0).tags()) .containsEntry("mvc.controller.method", "successful") .containsEntry("mvc.controller.class", "Controller2"); } @@ -80,8 +82,11 @@ public class TraceWebFluxTests { @RestController static class Controller2 { + @GetMapping("/api/c2/{id}") public Flux successful(@PathVariable Long id) { + // #786 + then(MDC.get("X-B3-TraceId")).isNotEmpty(); return Flux.just(id.toString()); } }