diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/SleuthAdvisorConfig.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/SleuthAdvisorConfig.java index dbc2cbde8..bd9dd414a 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/SleuthAdvisorConfig.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/annotation/SleuthAdvisorConfig.java @@ -45,6 +45,7 @@ import org.springframework.util.ReflectionUtils; import org.springframework.util.StringUtils; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.publisher.SignalType; /** * Custom pointcut advisor that picks all classes / interfaces that @@ -259,14 +260,14 @@ class SleuthInterceptor implements IntroductionInterceptor, BeanFactoryAware { if(publisher instanceof Mono){ return startSpan.flatMap(spanStarted -> ((Mono)publisher) .doOnError(onFailureReactor(log, hasLog, spanStarted)) - .doOnTerminate(afterReactor(startNewSpan, log, hasLog, spanStarted))) + .doFinally(afterReactor(startNewSpan, log, hasLog, spanStarted))) //put span in context so it can be used by ScopePassingSpanSubscriber .subscriberContext(context -> context.put(Span.class, span)); } else if(publisher instanceof Flux){ return startSpan.flatMapMany(spanStarted -> ((Flux)publisher) .doOnError(onFailureReactor(log, hasLog, spanStarted)) - .doOnTerminate(afterReactor(startNewSpan, log, hasLog, spanStarted))) + .doFinally(afterReactor(startNewSpan, log, hasLog, spanStarted))) //put span in context so it can be used by ScopePassingSpanSubscriber .subscriberContext(context -> context.put(Span.class, span)); } @@ -282,8 +283,8 @@ class SleuthInterceptor implements IntroductionInterceptor, BeanFactoryAware { } } - private Runnable afterReactor(boolean isNewSpan, String log, boolean hasLog, Span span) { - return () -> { + private Consumer afterReactor(boolean isNewSpan, String log, boolean hasLog, Span span) { + return signalType -> { try(Tracer.SpanInScope ws = tracer().withSpanInScope(span)) { after(span, isNewSpan, log, hasLog); } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java index 63ee1ea7b..741432d71 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java @@ -83,14 +83,19 @@ final class ScopePassingSpanSubscriber extends AtomicBoolean implements SpanS } @Override public void onError(Throwable throwable) { - this.subscriber.onError(throwable); + try (Tracer.SpanInScope inScope = this.tracer.withSpanInScope(this.span)) { + this.subscriber.onError(throwable); + } } @Override public void onComplete() { - this.subscriber.onComplete(); + try (Tracer.SpanInScope inScope = this.tracer.withSpanInScope(this.span)) { + this.subscriber.onComplete(); + } } @Override public Context currentContext() { return this.context; } + } \ No newline at end of file diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectWebFluxTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectWebFluxTests.java index b09f50fb9..41c73a067 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectWebFluxTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectWebFluxTests.java @@ -16,6 +16,9 @@ package org.springframework.cloud.sleuth.annotation; +import java.util.List; +import java.util.concurrent.ConcurrentLinkedQueue; + import brave.Tracer; import brave.sampler.Sampler; import org.apache.commons.lang3.StringUtils; @@ -23,7 +26,10 @@ import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.actuate.trace.http.HttpTrace; +import org.springframework.boot.actuate.trace.http.HttpTraceRepository; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.autoconfigure.security.reactive.ReactiveUserDetailsServiceAutoConfiguration; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.boot.web.server.LocalServerPort; import org.springframework.cloud.sleuth.util.ArrayListSpanReporter; @@ -38,21 +44,20 @@ import reactor.core.publisher.Mono; import zipkin2.Span; import zipkin2.reporter.Reporter; -import java.util.List; - import static org.assertj.core.api.BDDAssertions.then; @RunWith(SpringJUnit4ClassRunner.class) -@SpringBootTest(classes = { +@SpringBootTest( + properties = {"spring.main.web-application-type=reactive"}, + classes = { SleuthSpanCreatorAspectWebFluxTests.TestEndpoint.class, SleuthSpanCreatorAspectWebFluxTests.TestConfiguration.class} , webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT) -@EnableAutoConfiguration public class SleuthSpanCreatorAspectWebFluxTests { - + @Autowired Tracer tracer; @Autowired ArrayListSpanReporter reporter; - + @Before public void setup() { this.reporter.clear(); @@ -60,11 +65,15 @@ public class SleuthSpanCreatorAspectWebFluxTests { @LocalServerPort private int port; + + private final WebClient webClient = WebClient.create(); + + private static final ConcurrentLinkedQueue spanIdsInHttpTrace = new ConcurrentLinkedQueue<>(); @Test public void shouldReturnSpanFromWebFluxTraceContext() { - Mono mono = WebClient.create().get().uri("http://localhost:"+port+"/test/ping") + Mono mono = webClient.get().uri("http://localhost:"+port+"/test/ping") .retrieve().bodyToMono(Long.class); then(this.reporter.getSpans()).isEmpty(); @@ -82,7 +91,7 @@ public class SleuthSpanCreatorAspectWebFluxTests { @Test public void shouldReturnSpanFromWebFluxSubscriptionContext() { - Mono mono = WebClient.create().get().uri("http://localhost:"+port+"/test/pingFromContext") + Mono mono = webClient.get().uri("http://localhost:"+port+"/test/pingFromContext") .retrieve().bodyToMono(Long.class); then(this.reporter.getSpans()).isEmpty(); @@ -100,7 +109,7 @@ public class SleuthSpanCreatorAspectWebFluxTests { @Test public void shouldContinueSpanInWebFlux() { - Mono mono = WebClient.create().get().uri("http://localhost:"+port+"/test/continueSpan") + Mono mono = webClient.get().uri("http://localhost:"+port+"/test/continueSpan") .retrieve().bodyToMono(Long.class); then(this.reporter.getSpans()).isEmpty(); @@ -118,7 +127,7 @@ public class SleuthSpanCreatorAspectWebFluxTests { @Test public void shouldCreateNewSpanInWebFlux() { - Mono mono = WebClient.create().get().uri("http://localhost:"+port+"/test/newSpan1") + Mono mono = webClient.get().uri("http://localhost:"+port+"/test/newSpan1") .retrieve().bodyToMono(Long.class); then(this.reporter.getSpans()).isEmpty(); @@ -137,7 +146,7 @@ public class SleuthSpanCreatorAspectWebFluxTests { @Test public void shouldCreateNewSpanInWebFluxInSubscriberContext() { - Mono mono = WebClient.create().get().uri("http://localhost:"+port+"/test/newSpan2") + Mono mono = webClient.get().uri("http://localhost:"+port+"/test/newSpan2") .retrieve().bodyToMono(Long.class); then(this.reporter.getSpans()).isEmpty(); @@ -153,6 +162,28 @@ public class SleuthSpanCreatorAspectWebFluxTests { then(this.tracer.currentSpan()).isNull(); } + @Test + public void shouldSetupCorrectSpanInHttpTrace() { + + spanIdsInHttpTrace.clear(); + + Mono mono = webClient.get().uri("http://localhost:"+port+"/test/ping") + .retrieve().bodyToMono(Long.class); + + then(this.reporter.getSpans()).isEmpty(); + + Long newSpanId = mono.block(); + + List spans = this.reporter.getSpans(); + then(spans).hasSize(1); + then(spans.get(0).kind()).isEqualTo(Span.Kind.SERVER); + then(spans.get(0).name()).isEqualTo("get /test/ping"); + then(spans.get(0).id()) + .isEqualTo(toHexString(newSpanId)) + .isEqualTo(toHexString(spanIdsInHttpTrace.poll())); + then(this.tracer.currentSpan()).isNull(); + } + private static String toHexString(long value){ return StringUtils.leftPad(Long.toHexString(value), 16, '0'); @@ -160,7 +191,7 @@ public class SleuthSpanCreatorAspectWebFluxTests { @Configuration - @EnableAutoConfiguration + @EnableAutoConfiguration(exclude = {org.springframework.boot.autoconfigure.security.reactive.ReactiveSecurityAutoConfiguration.class, ReactiveUserDetailsServiceAutoConfiguration.class}) protected static class TestConfiguration { @Bean @@ -175,6 +206,25 @@ public class SleuthSpanCreatorAspectWebFluxTests { @Bean Sampler alwaysSampler() { return Sampler.ALWAYS_SAMPLE; } + + @Bean AccessLoggingHttpTraceRepository accessLoggingHttpTraceRepository(){ + return new AccessLoggingHttpTraceRepository(); + } + } + + static class AccessLoggingHttpTraceRepository implements HttpTraceRepository { + + @Autowired Tracer tracer; + + @Override + public List findAll() { + return null; + } + + @Override + public void add(HttpTrace trace) { + spanIdsInHttpTrace.add(tracer.currentSpan().context().spanId()); + } } @RestController