From 18dc5a41416606b47c63131894ceed47a7e69478 Mon Sep 17 00:00:00 2001 From: Oleh Dokuka <5380167+OlegDokuka@users.noreply.github.com> Date: Tue, 14 Mar 2023 21:47:28 +0200 Subject: [PATCH] ensures span is restored back after onComplete is propagated to the downstream (#2275) --- .../sleuth/instrument/web/TraceWebFilter.java | 14 ++--- .../instrument/web/TraceWebFluxTests.java | 51 +++++++++++++++++++ 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java index 54ead3555..5d0c02484 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java @@ -274,14 +274,16 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa @Override public void onError(Throwable t) { - terminateSpan(t); - this.actual.onError(t); + try (Tracer.SpanInScope ignored = terminateSpan(t)) { + this.actual.onError(t); + } } @Override public void onComplete() { - terminateSpan(null); - this.actual.onComplete(); + try (Tracer.SpanInScope ignored = terminateSpan(null)) { + this.actual.onComplete(); + } } @Override @@ -289,7 +291,7 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa return this.context; } - private void terminateSpan(@Nullable Throwable t) { + private Tracer.SpanInScope terminateSpan(@Nullable Throwable t) { Object attribute = this.exchange.getAttribute(HandlerMapping.BEST_MATCHING_HANDLER_ATTRIBUTE); addClassMethodTag(attribute, this.span); addClassNameTag(attribute, this.span); @@ -302,7 +304,7 @@ public class TraceWebFilter implements WebFilter, Ordered, ApplicationContextAwa if (log.isDebugEnabled()) { log.debug("Handled send of " + this.span); } - tracer.withSpan(null); + return tracer.withSpan(null); } private void addClassMethodTag(Object handler, Span span) { diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-webflux-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/TraceWebFluxTests.java b/tests/brave/spring-cloud-sleuth-instrumentation-webflux-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/TraceWebFluxTests.java index 0338b0359..3448f715b 100644 --- a/tests/brave/spring-cloud-sleuth-instrumentation-webflux-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/TraceWebFluxTests.java +++ b/tests/brave/spring-cloud-sleuth-instrumentation-webflux-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/web/TraceWebFluxTests.java @@ -70,6 +70,7 @@ public class TraceWebFluxTests { .run(); TestSpanHandler spans = context.getBean(TestSpanHandler.class); AssertingWebFilter assertingWebFilter = context.getBean(AssertingWebFilter.class); + InnerAssertingWebFilter innerAssertingWebFilter = context.getBean(InnerAssertingWebFilter.class); int port = context.getBean(Environment.class).getProperty("local.server.port", Integer.class); Controller2 controller2 = context.getBean(Controller2.class); clean(spans, controller2); @@ -78,6 +79,7 @@ public class TraceWebFluxTests { ClientResponse response = whenRequestIsSent(port, "/api/c2/10"); // then thenSpanWasReportedWithTags(spans, response); + thenTraceWasAvailableInDoAfter(innerAssertingWebFilter, spans.spans().get(0).id()); // then #2002 then(response.headers().header("mytraceid")).isNotEmpty(); clean(spans, controller2); @@ -86,35 +88,41 @@ public class TraceWebFluxTests { response = whenRequestIsSent(port, "/api/fn/20"); // then thenFunctionalSpanWasReportedWithTags(spans, response); + thenTraceWasAvailableInDoAfter(innerAssertingWebFilter, spans.spans().get(0).id()); spans.clear(); // when response = whenRequestIsSent(port, "/missing-endpoint"); // then thenSpanWith404StatusCodeWasReported(spans, response); + thenTraceWasAvailableInDoAfter(innerAssertingWebFilter, spans.spans().get(0).id()); spans.clear(); // when response = whenRequestIsSent(port, "/exception"); // then thenSpanWithExceptionWasReported(spans, response); + thenTraceWasAvailableInDoAfter(innerAssertingWebFilter, spans.spans().get(0).id()); spans.clear(); // when ClientResponse nonSampledResponse = whenNonSampledRequestIsSent(port); // then thenNoSpanWasReported(spans, nonSampledResponse, controller2); + thenSomeTraceWasAvailableInDoAfter(innerAssertingWebFilter); spans.clear(); // when ClientResponse skippedPatternResponse = whenRequestIsSentToSkippedPattern(port); // then thenNoSpanWasReported(spans, skippedPatternResponse, controller2); + thenSomeTraceWasAvailableInDoAfter(innerAssertingWebFilter); // when (issue #1683) response = whenRequestWithXForwardedForIsSent(port, "/api/fn/20"); // then thenSpanWasReportedWithRemoteIpTags(spans, response); + thenTraceWasAvailableInDoAfter(innerAssertingWebFilter, spans.spans().get(0).id()); thenNoTraceWasLeaked(assertingWebFilter); @@ -166,6 +174,17 @@ public class TraceWebFluxTests { then(assertingWebFilter.getSpans()).isEmpty(); } + + private void thenSomeTraceWasAvailableInDoAfter(InnerAssertingWebFilter assertingWebFilter) { + then(assertingWebFilter.getSpans()).hasSize(1); + assertingWebFilter.getSpans().clear(); + } + + private void thenTraceWasAvailableInDoAfter(InnerAssertingWebFilter assertingWebFilter, String spanId) { + then(assertingWebFilter.getSpans()).hasSize(1).first().matches(s -> s.context().traceIdString().equals(spanId)); + assertingWebFilter.getSpans().clear(); + } + private void thenNoSpanWasReported(TestSpanHandler spans, ClientResponse response, Controller2 controller2) { Awaitility.await().untilAsserted(() -> { then(response.statusCode().value()).isEqualTo(200); @@ -247,6 +266,38 @@ public class TraceWebFluxTests { return new AssertingWebFilter(tracer); } + + @Bean + @Order(Ordered.LOWEST_PRECEDENCE) + WebFilter innerFilterDoAfterTerminate(Tracer tracer) { + return new InnerAssertingWebFilter(tracer); + } + + } + + static class InnerAssertingWebFilter implements WebFilter { + final Queue spans = new ConcurrentLinkedQueue<>(); + + private final Tracer tracer; + + InnerAssertingWebFilter(Tracer tracer) { + this.tracer = tracer; + } + + @Override + public Mono filter(ServerWebExchange exchange, WebFilterChain webFilterChain) { + return webFilterChain.filter(exchange) + .doAfterTerminate(() -> { + Span currentSpan = tracer.currentSpan(); + if (currentSpan != null) { + spans.add(currentSpan); + } + }); + } + + Queue getSpans() { + return spans; + } } static class AssertingWebFilter implements WebFilter {