ensures span is restored back after onComplete is propagated to the downstream (#2275)

Author: OlegDokuka

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFilter.java

@@ -274,22 +274,24 @@ public void onNext(Void aVoid) {
 
 			@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
 			public Context currentContext() {
 				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 @@ private void terminateSpan(@Nullable Throwable t) {
 				if (log.isDebugEnabled()) {
 					log.debug("Handled send of " + this.span);
 				}
-				tracer.withSpan(null);
+				return tracer.withSpan(null);
 			}
 
 			private void addClassMethodTag(Object handler, Span span) {

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 void should_instrument_web_filter() throws Exception {
 				.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 void should_instrument_web_filter() throws Exception {
 		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 void should_instrument_web_filter() throws Exception {
 		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 @@ private void thenNoTraceWasLeaked(AssertingWebFilter assertingWebFilter) {
 		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 @@ AssertingWebFilter traceIdInResponseLastFilter(Tracer tracer) {
 			return new AssertingWebFilter(tracer);
 		}
 
+
+		@Bean
+		@Order(Ordered.LOWEST_PRECEDENCE)
+		WebFilter innerFilterDoAfterTerminate(Tracer tracer) {
+			return new InnerAssertingWebFilter(tracer);
+		}
+
+	}
+
+	static class InnerAssertingWebFilter implements WebFilter {
+		final Queue<Span> spans  = new ConcurrentLinkedQueue<>();
+
+		private final Tracer tracer;
+
+		InnerAssertingWebFilter(Tracer tracer) {
+			this.tracer = tracer;
+		}
+
+		@Override
+		public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain webFilterChain) {
+			return webFilterChain.filter(exchange)
+					.doAfterTerminate(() -> {
+						Span currentSpan = tracer.currentSpan();
+						if (currentSpan != null) {
+							spans.add(currentSpan);
+						}
+					});
+		}
+
+		Queue<Span> getSpans() {
+			return spans;
+		}
 	}
 
 	static class AssertingWebFilter implements WebFilter {