Introduced a Reactor operator that always puts a span to the context

without this change for different reactor operators (e.g. flatMap) we were losing trace context.
with this change we're introducing a reactor operator that for each reactor operator ensures that operations get executed within a trace context

fixes #850

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ReactorSleuth.java

@@ -24,6 +24,7 @@ import reactor.core.Fuseable;
 import reactor.core.Scannable;
 import reactor.core.publisher.Operators;
 import org.reactivestreams.Publisher;
+import reactor.util.context.Context;
 
 /**
  * Reactive Span pointcuts factories
@@ -60,6 +61,33 @@ public abstract class ReactorSleuth {
 		}));
 	}
 
+	/**
+	 * Return a span operator pointcut given a {@link Tracing}. This can be used in reactor
+	 * via {@link reactor.core.publisher.Flux#transform(Function)}, {@link
+	 * reactor.core.publisher.Mono#transform(Function)}, {@link
+	 * reactor.core.publisher.Hooks#onEachOperator(Function)} or {@link
+	 * reactor.core.publisher.Hooks#onLastOperator(Function)}. The Span operator
+	 * pointcut will pass the Scope of the Span without ever creating any new spans.
+	 *
+	 * @param tracing the {@link Tracing} instance to use in this span operator
+	 * @param <T> an arbitrary type that is left unchanged by the span operator
+	 *
+	 * @return a new Span operator pointcut
+	 */
+	public static <T> Function<? super Publisher<T>, ? extends Publisher<T>> scopePassingSpanOperator(
+			Tracing tracing) {
+		return Operators.lift(POINTCUT_FILTER, ((scannable, sub) -> {
+			//do not trace fused flows
+			if(scannable instanceof Fuseable && sub instanceof Fuseable.QueueSubscription){
+				return sub;
+			}
+			return new ScopePassingSpanSubscriber<>(
+					sub,
+					sub != null ? sub.currentContext() : Context.empty(),
+					tracing);
+		}));
+	}
+
 	private static final Predicate<Scannable> POINTCUT_FILTER =
 			s ->  !(s instanceof Fuseable.ScalarCallable);
 

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/ScopePassingSpanSubscriber.java

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2013-2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.instrument.reactor;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import brave.Span;
+import brave.Tracer;
+import brave.Tracing;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.util.Logger;
+import reactor.util.Loggers;
+import reactor.util.context.Context;
+
+/**
+ * A trace representation of the {@link Subscriber} that always
+ * continues a span
+ *
+ * @author Marcin Grzejszczak
+ * @since 2.0.0
+ */
+final class ScopePassingSpanSubscriber<T> extends AtomicBoolean implements Subscription,
+		CoreSubscriber<T> {
+
+	private static final Logger log = Loggers.getLogger(
+			ScopePassingSpanSubscriber.class);
+
+	private final Span span;
+	private final Subscriber<? super T> subscriber;
+	private final Context context;
+	private final Tracer tracer;
+	private Subscription s;
+
+	ScopePassingSpanSubscriber(Subscriber<? super T> subscriber, Context ctx, Tracing tracing) {
+		this.subscriber = subscriber;
+		this.tracer = tracing.tracer();
+		Span root = ctx != null ?
+				ctx.getOrDefault(Span.class, this.tracer.currentSpan()) : null;
+		this.span = root;
+		this.context = root != null ?
+				ctx.put(Span.class, root): Context.empty();
+	}
+
+	@Override public void onSubscribe(Subscription subscription) {
+		this.s = subscription;
+		try (Tracer.SpanInScope inScope = this.tracer.withSpanInScope(this.span)) {
+			this.subscriber.onSubscribe(this);
+		}
+	}
+
+	@Override public void request(long n) {
+		try (Tracer.SpanInScope inScope = this.tracer.withSpanInScope(this.span)) {
+			this.s.request(n);
+		}
+	}
+
+	@Override public void cancel() {
+		try (Tracer.SpanInScope inScope = this.tracer.withSpanInScope(this.span)) {
+			this.s.cancel();
+		}
+	}
+
+	@Override public void onNext(T o) {
+		try (Tracer.SpanInScope inScope = this.tracer.withSpanInScope(this.span)) {
+			this.subscriber.onNext(o);
+		}
+	}
+
+	@Override public void onError(Throwable throwable) {
+		this.subscriber.onError(throwable);
+	}
+
+	@Override public void onComplete() {
+		this.subscriber.onComplete();
+	}
+
+	@Override public Context currentContext() {
+		return this.context;
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/SpanSubscriber.java

@@ -22,12 +22,12 @@ import brave.Span;
 import brave.Tracer;
 import brave.Tracing;
 import brave.propagation.TraceContextOrSamplingFlags;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
 import reactor.util.Logger;
 import reactor.util.Loggers;
 import reactor.util.context.Context;
-import org.reactivestreams.Subscriber;
-import org.reactivestreams.Subscription;
 
 /**
  * A trace representation of the {@link Subscriber}
@@ -139,23 +139,7 @@ final class SpanSubscriber<T> extends AtomicBoolean implements Subscription,
 			if (log.isTraceEnabled()) {
 				log.trace("Cleaning up");
 			}
-			Tracer.SpanInScope ws = null;
-			if (this.tracer.currentSpan() != this.span) {
-				if (log.isTraceEnabled()) {
-					log.trace("Detaching span");
-				}
-				ws = this.tracer.withSpanInScope(this.span);
-				if (log.isTraceEnabled()) {
-					log.trace("Continuing span");
-				}
-			}
-			if (log.isTraceEnabled()) {
-				log.trace("Closing span");
-			}
 			this.span.finish();
-			if (ws != null) {
-				ws.close();
-			}
 			if (log.isTraceEnabled()) {
 				log.trace("Span closed");
 			}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/reactor/TraceReactorAutoConfiguration.java

@@ -75,6 +75,7 @@ public class TraceReactorAutoConfiguration {
 		@PostConstruct
 		public void setupHooks() {
 			this.lastOperatorWrapper.wrapLastOperator(this.tracing);
+			Hooks.onEachOperator(SLEUTH_TRACE_REACTOR_KEY, ReactorSleuth.scopePassingSpanOperator(this.tracing));
 			Schedulers.setFactory(new Schedulers.Factory() {
 				@Override public ScheduledExecutorService decorateExecutorService(String schedulerType,
 						Supplier<? extends ScheduledExecutorService> actual) {
@@ -87,7 +88,8 @@ public class TraceReactorAutoConfiguration {
 
 		@PreDestroy
 		public void cleanupHooks() {
-			Hooks.resetOnLastOperator();
+			Hooks.resetOnLastOperator(SLEUTH_TRACE_REACTOR_KEY);
+			Hooks.resetOnEachOperator(SLEUTH_TRACE_REACTOR_KEY);
 			Schedulers.resetFactory();
 		}
 	}

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

@@ -40,6 +40,7 @@ import org.springframework.web.reactive.HandlerMapping;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilter;
 import org.springframework.web.server.WebFilterChain;
+import reactor.util.context.Context;
 
 /**
  * A {@link WebFilter} that creates / continues / closes and detaches spans
@@ -131,6 +132,10 @@ public final class TraceWebFilter implements WebFilter, Ordered {
 	}
 
 	@Override public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
+		if (tracer().currentSpan() != null) {
+			// clear any previous trace
+			tracer().withSpanInScope(null);
+		}
 		ServerHttpRequest request = exchange.getRequest();
 		ServerHttpResponse response = exchange.getResponse();
 		String uri = request.getPath().pathWithinApplication().value();
@@ -149,8 +154,7 @@ public final class TraceWebFilter implements WebFilter, Ordered {
 								.map(c -> c.put(CONTEXT_ERROR, t)))
 						.flatMap(c -> {
 							//reactivate span from context
-							SpanAndScope spanAndScope = c.getOrDefault(SpanAndScope.class, defaultSpanAndScope());
-							Span span = spanAndScope.span;
+							Span span = spanFromContext(c);
 							Mono<Void> continuation;
 							Throwable t = null;
 							if (c.hasKey(CONTEXT_ERROR)) {
@@ -168,27 +172,41 @@ public final class TraceWebFilter implements WebFilter, Ordered {
 							}
 							addResponseTagsForSpanWithoutParent(exchange, response, span);
 							handler().handleSend(response, t, span);
-							spanAndScope.scope.close();
+							if (log.isDebugEnabled()) {
+								log.debug("Handled send of " + span);
+							}
 							return continuation;
 						})
 						.subscriberContext(c -> {
 							Span span;
-							if (c.hasKey(SpanAndScope.class)) {
-								SpanAndScope spanAndScope = c.get(SpanAndScope.class);
-								Span parent = spanAndScope.span;
+							if (c.hasKey(Span.class)) {
+								Span parent = c.get(Span.class);
 								span = tracer()
 										.nextSpan(TraceContextOrSamplingFlags.create(parent.context()))
 										.start();
+								if (log.isDebugEnabled()) {
+									log.debug("Found span in reactor context" + span);
+								}
 							} else {
 								try {
-									if (skip) {
+									boolean hasTracingContextInHeaders = extractor()
+											.extract(request.getHeaders()) != TraceContextOrSamplingFlags.EMPTY;
+									// if there was a span received then we must not change
+									// the sampling decision
+									if (skip && !hasTracingContextInHeaders) {
 										span = unsampledSpan(name);
 									} else {
 										if (spanFromAttribute != null) {
 											span = spanFromAttribute;
+											if (log.isDebugEnabled()) {
+												log.debug("Found span in attribute " + span);
+											}
 										} else {
 											span = handler().handleReceive(extractor(),
 													request.getHeaders(), request);
+											if (log.isDebugEnabled()) {
+												log.debug("Handled receive of span " + span);
+											}
 										}
 									}
 									exchange.getAttributes().put(TRACE_REQUEST_ATTR, span);
@@ -200,16 +218,33 @@ public final class TraceWebFilter implements WebFilter, Ordered {
 									} else {
 										span = tracer().nextSpan().name(name).start();
 										exchange.getAttributes().put(TRACE_SPAN_WITHOUT_PARENT, span);
+										if (log.isDebugEnabled()) {
+											log.debug("Created a new 'fallback' span " + span);
+										}
 									}
 								}
 							}
-							return c.put(SpanAndScope.class, new SpanAndScope(span, tracer().withSpanInScope(span)));
+							return c.put(Span.class, span);
 						}));
 	}
 
-	private SpanAndScope defaultSpanAndScope() {
-		Span defaultSpan = tracer().nextSpan().start();
-		return new SpanAndScope(defaultSpan, tracer().withSpanInScope(defaultSpan));
+	private Span spanFromContext(Context c) {
+		if (c.hasKey(Span.class)) {
+			Span span = c.get(Span.class);
+			if (log.isDebugEnabled()) {
+				log.debug("Found span in context " + span);
+			}
+			return span;
+		}
+		Span span = defaultSpan();
+		if (log.isDebugEnabled()) {
+			log.debug("No span found in context. Creating a new one " + span);
+		}
+		return span;
+	}
+
+	private Span defaultSpan() {
+		return tracer().nextSpan().start();
 	}
 
 	private void addResponseTagsForSpanWithoutParent(ServerWebExchange exchange,
@@ -222,9 +257,13 @@ public final class TraceWebFilter implements WebFilter, Ordered {
 	}
 
 	private Span unsampledSpan(String name) {
-		return tracer().nextSpan(TraceContextOrSamplingFlags.create(
+		Span span = tracer().nextSpan(TraceContextOrSamplingFlags.create(
 				SamplingFlags.NOT_SAMPLED)).name(name)
 				.kind(Span.Kind.SERVER).start();
+		if (log.isDebugEnabled()) {
+			log.debug("Created a new unsampled span " + span);
+		}
+		return span;
 	}
 
 	private Span getSpanFromAttribute(ServerWebExchange exchange) {
@@ -245,22 +284,6 @@ public final 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) {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java

@@ -22,7 +22,6 @@ import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
 import brave.propagation.Propagation;
 import brave.propagation.TraceContext;
-import reactor.core.publisher.Mono;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.beans.BeansException;
@@ -34,6 +33,7 @@ import org.springframework.web.reactive.function.client.ClientResponse;
 import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
 import org.springframework.web.reactive.function.client.ExchangeFunction;
 import org.springframework.web.reactive.function.client.WebClient;
+import reactor.core.publisher.Mono;
 
 /**
  * {@link BeanPostProcessor} to wrap a {@link WebClient} instance into
@@ -113,41 +113,42 @@ class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 					Object any = anyAndContext.getT1();
 					Span clientSpan = anyAndContext.getT2().get(CLIENT_SPAN_KEY);
 					Mono<ClientResponse> continuation;
-					Throwable throwable = null;
-					ClientResponse response = null;
-					try (Tracer.SpanInScope ws = tracer().withSpanInScope(clientSpan)) {
+					final Tracer.SpanInScope ws = tracer().withSpanInScope(clientSpan);
 						if (any instanceof Throwable) {
-							throwable = (Throwable) any;
-							continuation = Mono.error(throwable);
+							continuation = Mono.error((Throwable) any);
 						} else {
-							response = (ClientResponse) any;
-							boolean error = response.statusCode().is4xxClientError() ||
-									response.statusCode().is5xxServerError();
-							if (error) {
-								if (log.isDebugEnabled()) {
-									log.debug(
-											"Non positive status code was returned from the call. Will close the span ["
-													+ clientSpan + "]");
-								}
-								throwable = new RestClientException(
-										"Status code of the response is [" + response.statusCode()
-												.value() + "] and the reason is [" + response
-												.statusCode().getReasonPhrase() + "]");
-							}
-							continuation = Mono.just(response);
+							continuation = Mono.just((ClientResponse) any);
 						}
-					} finally {
-						handler().handleReceive(response, throwable, clientSpan);
-					}
-					return continuation;
+					return continuation.doAfterSuccessOrError(
+							(clientResponse, throwable1) -> {
+								Throwable throwable = throwable1;
+								boolean error = clientResponse.statusCode().is4xxClientError() ||
+										clientResponse.statusCode().is5xxServerError();
+								if (error) {
+									if (log.isDebugEnabled()) {
+										log.debug(
+												"Non positive status code was returned from the call. Will close the span ["
+														+ clientSpan + "]");
+									}
+									throwable = new RestClientException(
+											"Status code of the response is [" + clientResponse.statusCode()
+													.value() + "] and the reason is [" + clientResponse
+													.statusCode().getReasonPhrase() + "]");
+								}
+								handler().handleReceive(clientResponse, throwable, clientSpan);
+								ws.close();
+							});
 				})
 				.subscriberContext(c -> {
 					if (log.isDebugEnabled()) {
-						log.debug("Creating a client span for the WebClient");
+						log.debug("Instrumenting WebClient call");
 					}
 					Span parent = c.getOrDefault(Span.class, null);
-					Span clientSpan = handler().handleSend(injector(), builder, request,
-							parent != null ? parent : tracer().nextSpan());
+					Span clientSpan = handler().handleSend(injector(), builder,
+							request, tracer().nextSpan());
+					if (log.isDebugEnabled()) {
+						log.debug("Created a client span for the WebClient " + clientSpan);
+					}
 					if (parent == null) {
 						c = c.put(Span.class, clientSpan);
 						if (log.isDebugEnabled()) {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jCurrentTraceContext.java

@@ -90,7 +90,7 @@ public final class Slf4jCurrentTraceContext extends CurrentTraceContext {
 			String sampled = String.valueOf(currentSpan.sampled());
 			MDC.put("spanExportable", sampled);
 			MDC.put(LEGACY_EXPORTABLE_NAME, sampled);
-			log("Starting span: {}", currentSpan);
+			log("Starting scope for span: {}", currentSpan);
 			if (currentSpan.parentId() != null) {
 				if (log.isTraceEnabled()) {
 					log.trace("With parent: {}", currentSpan.parentId());
@@ -112,7 +112,7 @@ public final class Slf4jCurrentTraceContext extends CurrentTraceContext {
 
 		class ThreadContextCurrentTraceContextScope implements Scope {
 			@Override public void close() {
-				log("Closing span: {}", currentSpan);
+				log("Closing scope for span: {}", currentSpan);
 				scope.close();
 				replace("traceId", previousTraceId);
 				replace("parentId", previousParentId);
@@ -128,7 +128,7 @@ public final class Slf4jCurrentTraceContext extends CurrentTraceContext {
 	}
 
 	private void log(String text, TraceContext span) {
-		if (span != null) {
+		if (span == null) {
 			return;
 		}
 		if (log.isTraceEnabled()) {

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/sample/FlatMapTests.java

@@ -0,0 +1,185 @@
+/*
+ * Copyright 2013-2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.instrument.reactor.sample;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import brave.Tracer;
+import brave.sampler.Sampler;
+import org.awaitility.Awaitility;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.boot.WebApplicationType;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.autoconfigure.security.reactive.ReactiveSecurityAutoConfiguration;
+import org.springframework.boot.autoconfigure.security.reactive.ReactiveUserDetailsServiceAutoConfiguration;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.boot.test.rule.OutputCapture;
+import org.springframework.cloud.sleuth.util.ArrayListSpanReporter;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.env.Environment;
+import org.springframework.web.reactive.function.client.ClientResponse;
+import org.springframework.web.reactive.function.client.WebClient;
+import org.springframework.web.reactive.function.server.RouterFunction;
+import org.springframework.web.reactive.function.server.ServerResponse;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Hooks;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
+import zipkin2.Span;
+
+import static org.assertj.core.api.BDDAssertions.then;
+import static org.springframework.web.reactive.function.server.RequestPredicates.GET;
+import static org.springframework.web.reactive.function.server.RouterFunctions.route;
+
+// https://github.com/spring-cloud/spring-cloud-sleuth/issues/850
+public class FlatMapTests {
+
+	private static final Logger LOGGER = LoggerFactory.getLogger(FlatMapTests.class);
+
+	@BeforeClass
+	public static void setup() {
+		Hooks.resetOnLastOperator();
+		Schedulers.resetFactory();
+	}
+
+	@Rule public OutputCapture capture = new OutputCapture();
+
+	@Test public void should_work_with_flat_maps() {
+		//given
+		ConfigurableApplicationContext context = new SpringApplicationBuilder(
+				FlatMapTests.TestConfiguration.class).web(WebApplicationType.REACTIVE)
+				.properties("server.port=0", "spring.jmx.enabled=false",
+						"spring.application.name=TraceWebFluxTests", "security.basic.enabled=false",
+						"management.security.enabled=false").run();
+		ArrayListSpanReporter accumulator = context.getBean(ArrayListSpanReporter.class);
+		int port = context.getBean(Environment.class).getProperty("local.server.port", Integer.class);
+		RequestSender sender = context.getBean(RequestSender.class);
+		TestConfiguration config = context.getBean(TestConfiguration.class);
+		sender.port = port;
+		accumulator.clear();
+
+		Awaitility.await().untilAsserted(() -> {
+			//when
+			accumulator.clear();
+			String firstTraceId = flatMapTraceId(accumulator, callFlatMap(port).block());
+			//then
+			thenAllWebClientCallsHaveSameTraceId(firstTraceId, sender);
+			thenSpanInFooHasSameTraceId(firstTraceId, config);
+			accumulator.clear();
+
+			//when
+			String secondTraceId = flatMapTraceId(accumulator, callFlatMap(port).block());
+			//then
+			then(firstTraceId)
+					.as("Id will not be reused between calls")
+					.isNotEqualTo(secondTraceId);
+			thenSpanInFooHasSameTraceId(secondTraceId, config);
+			//and
+			then(Arrays.stream(capture.toString().split("\n"))
+					.filter(s -> s.contains("Received a request to uri"))
+					.map(s -> s.split(",")[1])
+					.collect(Collectors.toList()))
+					.as("TraceFilter should not have any trace when receiving a request")
+					.containsOnly("");
+		});
+	}
+
+	private void thenAllWebClientCallsHaveSameTraceId(String traceId,
+			RequestSender sender) {
+		then(sender.span.context().traceIdString()).isEqualTo(traceId);
+	}
+
+	private void thenSpanInFooHasSameTraceId(String traceId,
+			TestConfiguration config) {
+		then(config.spanInFoo.context().traceIdString()).isEqualTo(traceId);
+	}
+
+	private Mono<ClientResponse> callFlatMap(int port) {
+		return WebClient.create().get()
+				.uri("http://localhost:" + port + "/withFlatMap").exchange();
+	}
+
+	private String flatMapTraceId(ArrayListSpanReporter accumulator,
+			ClientResponse response) {
+		then(response.statusCode().value()).isEqualTo(200);
+		then(accumulator.getSpans()).isNotEmpty();
+		LOGGER.info("Accumulated spans: " + accumulator.getSpans());
+		List<String> traceIdOfFlatMap = accumulator.getSpans().stream()
+				.filter(span -> span.tags().containsKey("http.path") && span.tags()
+						.get("http.path").equals("/withFlatMap")).map(Span::traceId)
+				.collect(Collectors.toList());
+		then(traceIdOfFlatMap).hasSize(1);
+		return traceIdOfFlatMap.get(0);
+	}
+
+	@Configuration
+	@EnableAutoConfiguration(
+			exclude = { ReactiveUserDetailsServiceAutoConfiguration.class,
+					ReactiveSecurityAutoConfiguration.class })
+	static class TestConfiguration {
+
+		brave.Span spanInFoo;
+
+		@Bean RouterFunction<ServerResponse> handlers(Tracer tracer, RequestSender requestSender) {
+			return route(GET("/noFlatMap"), request -> {
+				LOGGER.info("noFlatMap");
+				Flux<Integer> one = requestSender.getAll().map(string -> string.length());
+				return ServerResponse.ok().body(one, Integer.class);
+			}).andRoute(GET("/withFlatMap"), request -> {
+				LOGGER.info("withFlatMap");
+				Flux<Integer> one = requestSender.getAll().map(string -> string.length());
+				Flux<Integer> response = one.flatMap(size -> requestSender.getAll()
+						.doOnEach(sig -> LOGGER.info(sig.getContext().toString())))
+						.map(string -> {
+							LOGGER.info("WHATEVER YEAH");
+							return string.length();
+						});
+				return ServerResponse.ok().body(response, Integer.class);
+			}).andRoute(GET("/foo"), request -> {
+				LOGGER.info("foo");
+				spanInFoo = tracer.currentSpan();
+				return ServerResponse.ok().body(Flux.just(1), Integer.class);
+			});
+		}
+
+		@Bean WebClient webClient() {
+			return WebClient.create();
+		}
+
+		@Bean ArrayListSpanReporter reporter() {
+			return new ArrayListSpanReporter();
+		}
+
+		@Bean Sampler sampler() {
+			return Sampler.ALWAYS_SAMPLE;
+		}
+
+		@Bean
+		RequestSender sender(WebClient client, Tracer tracer) {
+			return new RequestSender(client, tracer);
+		}
+
+	}
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/reactor/sample/RequestSender.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2013-2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.instrument.reactor.sample;
+
+import brave.Span;
+import brave.Tracer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.http.HttpMethod;
+import org.springframework.web.reactive.function.client.WebClient;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+class RequestSender {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(RequestSender.class);
+
+    private final WebClient webClient;
+    private final Tracer tracer;
+    int port;
+    Span span;
+
+    public RequestSender(WebClient webClient, Tracer tracer) {
+        this.webClient = webClient;
+        this.tracer = tracer;
+    }
+
+    public Mono<String> get(Integer someParameterNotUsedNow){
+        LOGGER.info("getting for parameter {}", someParameterNotUsedNow);
+        this.span = this.tracer.currentSpan();
+        return webClient
+                .method(HttpMethod.GET)
+                .uri("http://localhost:" + this.port + "/foo")
+                .retrieve().bodyToMono(String.class);
+    }
+
+    public Flux<String> getAll(){
+        LOGGER.info("Before merge");
+        Flux<String> merge = Flux.merge(get(1), get(2), get(3));
+        LOGGER.info("after merge");
+        return merge;
+    }
+
+
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterWebIntegrationTests.java

@@ -23,7 +23,6 @@ import java.util.stream.Collectors;
 
 import brave.Tracing;
 import brave.sampler.Sampler;
-import zipkin2.Span;
 import org.assertj.core.api.BDDAssertions;
 import org.junit.After;
 import org.junit.Before;
@@ -40,13 +39,14 @@ import org.springframework.context.annotation.Configuration;
 import org.springframework.core.env.Environment;
 import org.springframework.http.ResponseEntity;
 import org.springframework.http.client.ClientHttpResponse;
-import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.context.junit4.SpringRunner;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestMethod;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.client.DefaultResponseErrorHandler;
 import org.springframework.web.client.HttpClientErrorException;
 import org.springframework.web.client.RestTemplate;
+import zipkin2.Span;
 
 import static org.assertj.core.api.Assertions.fail;
 import static org.assertj.core.api.BDDAssertions.then;
@@ -54,7 +54,7 @@ import static org.assertj.core.api.BDDAssertions.then;
 /**
  * @author Marcin Grzejszczak
  */
-@RunWith(SpringJUnit4ClassRunner.class)
+@RunWith(SpringRunner.class)
 @SpringBootTest(classes = TraceFilterWebIntegrationTests.Config.class,
 		webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
 		properties = "spring.sleuth.http.legacy.enabled=true")

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java

@@ -16,6 +16,10 @@
 
 package org.springframework.cloud.sleuth.instrument.web;
 
+import java.util.Random;
+
+import brave.Span;
+import brave.Tracer;
 import brave.sampler.Sampler;
 import org.awaitility.Awaitility;
 import org.junit.BeforeClass;
@@ -46,6 +50,8 @@ import static org.assertj.core.api.BDDAssertions.then;
 
 public class TraceWebFluxTests {
 
+	public static final String EXPECTED_TRACE_ID = "b919095138aa4c6e";
+
 	@BeforeClass
 	public static void setup() {
 		Hooks.resetOnLastOperator();
@@ -53,6 +59,7 @@ public class TraceWebFluxTests {
 	}
 
 	@Test public void should_instrument_web_filter() throws Exception {
+		// setup
 		ConfigurableApplicationContext context = new SpringApplicationBuilder(
 				TraceWebFluxTests.Config.class).web(WebApplicationType.REACTIVE)
 				.properties("server.port=0", "spring.jmx.enabled=false",
@@ -60,12 +67,31 @@ public class TraceWebFluxTests {
 								"management.security.enabled=false").run();
 		ArrayListSpanReporter accumulator = context.getBean(ArrayListSpanReporter.class);
 		int port = context.getBean(Environment.class).getProperty("local.server.port", Integer.class);
+		Controller2 controller2 = context.getBean(Controller2.class);
+		clean(accumulator, controller2);
+
+		// when
+		ClientResponse response = whenRequestIsSent(port);
+		//then
+		thenSpanWasReportedWithTags(accumulator, response);
+		clean(accumulator, controller2);
+
+		// when
+		ClientResponse nonSampledResponse = whenNonSampledRequestIsSent(port);
+		// then
+		thenNoSpanWasReported(accumulator, nonSampledResponse, controller2);
+
+		// cleanup
+		context.close();
+	}
+
+	private void clean(ArrayListSpanReporter accumulator, Controller2 controller2) {
 		accumulator.clear();
+		controller2.span = null;
+	}
 
-		Mono<ClientResponse> exchange = WebClient.create().get()
-				.uri("http://localhost:" + port + "/api/c2/10").exchange();
-		ClientResponse response = exchange.block();
-
+	private void thenSpanWasReportedWithTags(ArrayListSpanReporter accumulator,
+			ClientResponse response) {
 		Awaitility.await().untilAsserted(() -> {
 			then(response.statusCode().value()).isEqualTo(200);
 			then(accumulator.getSpans()).hasSize(1);
@@ -75,6 +101,32 @@ public class TraceWebFluxTests {
 				.containsEntry("mvc.controller.class", "Controller2");
 	}
 
+	private void thenNoSpanWasReported(ArrayListSpanReporter accumulator,
+			ClientResponse response, Controller2 controller2) {
+		Awaitility.await().untilAsserted(() -> {
+			then(response.statusCode().value()).isEqualTo(200);
+			then(accumulator.getSpans()).isEmpty();
+		});
+		then(controller2.span).isNotNull();
+		then(controller2.span.context().traceIdString()).isEqualTo(EXPECTED_TRACE_ID);
+	}
+
+	private ClientResponse whenRequestIsSent(int port) {
+		Mono<ClientResponse> exchange = WebClient.create().get()
+				.uri("http://localhost:" + port + "/api/c2/10").exchange();
+		return exchange.block();
+	}
+
+	private ClientResponse whenNonSampledRequestIsSent(int port) {
+		Mono<ClientResponse> exchange = WebClient.create().get()
+				.uri("http://localhost:" + port + "/api/c2/10")
+				.header("X-B3-SpanId", EXPECTED_TRACE_ID)
+				.header("X-B3-TraceId", EXPECTED_TRACE_ID)
+				.header("X-B3-Sampled", "0")
+				.exchange();
+		return exchange.block();
+	}
+
 	@Configuration
 	@EnableAutoConfiguration(
 			exclude = { TraceWebClientAutoConfiguration.class,
@@ -94,18 +146,27 @@ public class TraceWebFluxTests {
 			return new ArrayListSpanReporter();
 		}
 
-		@Bean Controller2 controller2() {
-			return new Controller2();
+		@Bean Controller2 controller2(Tracer tracer) {
+			return new Controller2(tracer);
 		}
 	}
 
 	@RestController
 	static class Controller2 {
 
+		Span span;
+
+		private final Tracer tracer;
+
+		Controller2(Tracer tracer) {
+			this.tracer = tracer;
+		}
+
 		@GetMapping("/api/c2/{id}")
 		public Flux<String> successful(@PathVariable Long id) {
 			// #786
 			then(MDC.get("X-B3-TraceId")).isNotEmpty();
+			this.span = this.tracer.currentSpan();
 			return Flux.just(id.toString());
 		}
 	}