Stream reactive (#2038)

Added native support for Reactive Spring Cloud Stream.

docs/.asciidoctorconfig.adoc

@@ -0,0 +1,12 @@
+// +++++++++++++++++++++++++++++++++++++++++++++++++++++++
+// +  Initial AsciiDoc editor configuration file - V1.0  +
+// ++++++++++++++++++++++++++++++++++++++++++++++++++++++
+// 
+// Did not found any configuration files, so create this at project root level.
+// If you do not like those files to be generated - you can turn it off inside Asciidoctor Editor preferences.
+// 
+// You can define editor specific parts here.
+// For example: with next line you could set imagesdir attribute to subfolder "images" relative to the folder where this config file is located.
+// :imagesdir: {asciidoctorconfigdir}/images
+// 
+// For more information please take a look at https://github.com/de-jcup/eclipse-asciidoctor-editor/wiki/Asciidoctor-configfiles

docs/src/main/asciidoc/integrations.adoc

@@ -315,16 +315,39 @@ include::{common_tests_path}/src/main/java/org/springframework/cloud/sleuth/inst
 
 This feature is available for all tracer implementations.
 
-Spring Cloud Sleuth can instrument Spring Cloud Function.
+Spring Cloud Sleuth can instrument Spring Cloud Function. Since Spring Cloud Stream uses Spring Cloud Function
+you will get the messaging instrumentation out of the box.
+
 The way to achieve it is to provide a `Function` or `Consumer` or `Supplier` that takes in a `Message` as a parameter e.g. `Function<Message<String>, Message<Integer>>`.
-If the type is not `Message` then instrumentation will not take place.
-Out of the box instrumentation will not take place when dealing with Reactor based streams - e.g. `Function<Flux<Message<String>>, Flux<Message<Integer>>>`.
+If the type **is not** `Message` then instrumentation **will not** take place.
 
-Since Spring Cloud Stream reuses Spring Cloud Function, you'll get the instrumentation out of the box.
+For a reactive `Consumer<Flux<Message<?>>>` remember to manually close the span and clear the context before you call `.subscribe()`. Example:
 
-You can disable this behavior by setting the value of `spring.sleuth.function.enabled` to `false`.
+[source,java,indent=0]
+----
+@Bean
+	Consumer<Flux<Message<String>>> channel(Tracer tracer) {
+		// For the reactive consumer remember to call "subscribe()" at the end, otherwise
+		// you'll get the "Dispatcher has no subscribers" error
+		return i -> i
+					.doOnNext(s -> log.info("HELLO"))
+					// You must finish the span yourself and clear the tracing context like presented below. 
+					// Otherwise you will be missing out the span that wraps the function execution.
+					.doOnNext(s -> {
+						tracer.currentSpan().end();
+						tracer.withSpan(null);
+					})
+					.subscribe();
+	}
+}
+----
+
+You can disable Spring Cloud Stream integration by setting the value of `spring.sleuth.function.enabled` to `false`.
+
+If you want to fully control the life cycle of spans within the reactive messaging context of Spring Cloud Stream
+remember to disable the Spring Cloud Stream integration and leverage the `MessagingSleuthOperators` utility 
+class that allows you to manipulate the input and output messages in order to continue the tracing context and to execute custom code within the tracing context. 
 
-In order to work with reactive Stream functions you can leverage the `MessagingSleuthOperators` utility class that allows you to manipulate the input and output messages in order to continue the tracing context and to execute custom code within the tracing context.
 
 [source,java,indent=0]
 -----

spring-cloud-sleuth-api/src/main/java/org/springframework/cloud/sleuth/Tracer.java

@@ -140,6 +140,14 @@ public interface Tracer extends BaggageManager {
 	 */
 	TraceContext.Builder traceContextBuilder();
 
+	/**
+	 * Returns the {@link CurrentTraceContext}. Can be {@code null} so that we don't break
+	 * backward compatibility.
+	 * @return current trace context
+	 */
+	@Nullable
+	CurrentTraceContext currentTraceContext();
+
 	/**
 	 * Allows to customize the current span in scope.
 	 * @return current span customizer

spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/NoOpTracer.java

@@ -20,6 +20,7 @@ import java.util.HashMap;
 import java.util.Map;
 
 import org.springframework.cloud.sleuth.BaggageInScope;
+import org.springframework.cloud.sleuth.CurrentTraceContext;
 import org.springframework.cloud.sleuth.ScopedSpan;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanCustomizer;
@@ -99,4 +100,9 @@ class NoOpTracer implements Tracer {
 		return new NoOpBaggageInScope();
 	}
 
+	@Override
+	public CurrentTraceContext currentTraceContext() {
+		return new NoOpCurrentTraceContext();
+	}
+
 }

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveSpanBuilder.java

@@ -128,4 +128,10 @@ class BraveSpanBuilder implements Span.Builder {
 		return new BraveSpanBuilder(tracer, context);
 	}
 
+	@Override
+	public String toString() {
+		return "{" + " delegate='" + this.delegate + "'" + ", parentContext='" + this.parentContext + "'"
+				+ ", startTimestamp='" + this.startTimestamp + "'" + "}";
+	}
+
 }

spring-cloud-sleuth-brave/src/main/java/org/springframework/cloud/sleuth/brave/bridge/BraveTracer.java

@@ -21,6 +21,7 @@ import java.util.Map;
 import brave.propagation.TraceContextOrSamplingFlags;
 
 import org.springframework.cloud.sleuth.BaggageInScope;
+import org.springframework.cloud.sleuth.CurrentTraceContext;
 import org.springframework.cloud.sleuth.ScopedSpan;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanCustomizer;
@@ -40,9 +41,18 @@ public class BraveTracer implements Tracer {
 
 	private final BraveBaggageManager braveBaggageManager;
 
+	private final CurrentTraceContext currentTraceContext;
+
 	public BraveTracer(brave.Tracer tracer, BraveBaggageManager braveBaggageManager) {
 		this.tracer = tracer;
 		this.braveBaggageManager = braveBaggageManager;
+		this.currentTraceContext = null;
+	}
+
+	public BraveTracer(brave.Tracer tracer, CurrentTraceContext context, BraveBaggageManager braveBaggageManager) {
+		this.tracer = tracer;
+		this.braveBaggageManager = braveBaggageManager;
+		this.currentTraceContext = context;
 	}
 
 	@Override
@@ -122,6 +132,11 @@ public class BraveTracer implements Tracer {
 		return this.braveBaggageManager.createBaggage(name).set(value);
 	}
 
+	@Override
+	public CurrentTraceContext currentTraceContext() {
+		return this.currentTraceContext;
+	}
+
 }
 
 class BraveSpanInScope implements Tracer.SpanInScope {

spring-cloud-sleuth-instrumentation/pom.xml

@@ -31,6 +31,10 @@
 		<version>3.1.0-SNAPSHOT</version>
 		<relativePath>..</relativePath>
 	</parent>
+	
+	<properties>
+		<!-- <kotlin.compiler.incremental>true</kotlin.compiler.incremental> -->
+	</properties>
 
 	<dependencies>
 		<dependency>

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/SleuthMessagingSpan.java

@@ -48,6 +48,16 @@ enum SleuthMessagingSpan implements DocumentedSpan {
 			}
 		},
 
+		/**
+		 * Name of the Spring Cloud Function function name.
+		 */
+		FUNCTION_NAME {
+			@Override
+			public String getKey() {
+				return "function.name";
+			}
+		},
+
 		/**
 		 * User provided keys via customization options.
 		 */

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapper.java

@@ -16,6 +16,7 @@
 
 package org.springframework.cloud.sleuth.instrument.messaging;
 
+import java.lang.reflect.Type;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -23,6 +24,9 @@ import java.util.concurrent.ConcurrentHashMap;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import org.springframework.cloud.context.scope.refresh.RefreshScopeRefreshedEvent;
 import org.springframework.cloud.function.context.catalog.FunctionAroundWrapper;
@@ -30,6 +34,7 @@ import org.springframework.cloud.function.context.catalog.FunctionTypeUtils;
 import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth;
 import org.springframework.cloud.sleuth.propagation.Propagator;
 import org.springframework.context.ApplicationListener;
 import org.springframework.core.env.Environment;
@@ -84,28 +89,201 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 	}
 
 	@Override
-	protected Object doApply(Message<byte[]> message, SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction) {
-
-		/*
-		 * This code is temporary to disable conditions for which this interceptor is not
-		 * ready. For example, - it does not handle properly input or output of type
-		 * Publisher - it wraps output in Message when function returns a
-		 * Collection<Message> which it should not do.
-		 *
-		 */
-		if ((FunctionTypeUtils.isCollectionOfMessage(targetFunction.getOutputType())
-				|| targetFunction.isOutputTypePublisher())
-				|| (targetFunction.isSupplier() && targetFunction.isOutputTypePublisher())) {
+	protected Object doApply(Object message, SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction) {
+		if (FunctionTypeUtils.isCollectionOfMessage(targetFunction.getOutputType())) {
 			return targetFunction.apply(message); // no instrumentation
 		}
+		else if (targetFunction.isInputTypePublisher() || targetFunction.isOutputTypePublisher()) {
+			return reactorStream((Publisher) message, targetFunction);
+		}
+		return nonReactorStream((Message<byte[]>) message, targetFunction);
+	}
 
+	private Object reactorStream(Publisher messageStream,
+			SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction) {
+		if (messageStream == null && targetFunction.isSupplier()) { // Supplier
+			return reactorStreamSupplier(messageStream, targetFunction);
+		}
+		Type itemType = FunctionTypeUtils.getGenericType(targetFunction.getInputType());
+		Class<?> itemTypeClass = FunctionTypeUtils.getRawType(itemType);
+		if (!itemTypeClass.equals(Message.class)) {
+			if (log.isDebugEnabled()) {
+				log.debug("Target function [" + targetFunction.getFunctionDefinition() + "] has raw input type ["
+						+ itemType + "] and should be [" + Message.class + "]. Will not wrap it.");
+				return targetFunction.apply(messageStream);
+			}
+		}
+		Publisher<Message> messagePublisher = messageStream;
+		if (FunctionTypeUtils.isMono(targetFunction.getInputType())) {
+			return reactorMonoStream(targetFunction, messagePublisher);
+		}
+		return reactorFluxStream(targetFunction, messagePublisher);
+	}
+
+	private Object reactorMonoStream(SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction,
+			Publisher<Message> messagePublisher) {
+		if (log.isDebugEnabled()) {
+			log.debug("Will instrument a stream Mono function");
+		}
+		Mono<Message> mono = Mono.from(messagePublisher)
+				// ensure there are no previous spans
+				.doOnNext(m -> tracer.withSpan(null))
+				.map(msg -> this.traceMessageHandler.wrapInputMessage(msg,
+						inputDestination(targetFunction.getFunctionDefinition())))
+				.flatMap(msg -> Mono.deferContextual(ctx -> {
+					MessageAndSpansAndScope messageAndSpansAndScope = ctx.get(MessageAndSpansAndScope.class);
+					messageAndSpansAndScope.messageAndSpans = msg;
+					messageAndSpansAndScope.span = msg.childSpan;
+					setNameAndTag(targetFunction, msg.childSpan);
+					messageAndSpansAndScope.scope = tracer.withSpan(msg.childSpan);
+					return Mono.just(msg.msg);
+				}));
+		if (targetFunction.isConsumer()) {
+			return targetFunction.apply(reactorStreamConsumer(mono));
+		}
+		final Mono<Message> function = ((Mono<Message>) targetFunction.apply(mono));
+		return Mono.deferContextual(contextView -> {
+			MessageAndSpansAndScope msg = contextView.get(MessageAndSpansAndScope.class);
+			return function.doOnNext(message -> {
+				msg.end();
+				msg.handle();
+			}).map(msgResult -> {
+				MessageAndSpan messageAndSpan = traceMessageHandler.wrapOutputMessage(msgResult,
+						msg.messageAndSpans.parentSpan, outputDestination(targetFunction.getFunctionDefinition()));
+				traceMessageHandler.afterMessageHandled(messageAndSpan.span, null);
+				return messageAndSpan.msg;
+			})
+					// TODO: Fix me when this is resolved in Reactor
+					// .doOnSubscribe(__ -> scope.close())
+					.doOnError(msg::error).doFinally(signalType -> {
+						if (!msg.isHandled()) {
+							msg.end();
+						}
+					});
+		}).contextWrite(contextView -> contextView.put(MessageAndSpansAndScope.class, new MessageAndSpansAndScope()));
+	}
+
+	private Object reactorFluxStream(SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction,
+			Publisher<Message> messagePublisher) {
+		if (log.isDebugEnabled()) {
+			log.debug("Will instrument a stream Flux function");
+		}
+		Flux<Message> flux = Flux.from(messagePublisher)
+				// ensure there are no previous spans
+				.doOnNext(m -> tracer.withSpan(null))
+				.map(msg -> this.traceMessageHandler.wrapInputMessage(msg,
+						inputDestination(targetFunction.getFunctionDefinition())))
+				.flatMap(msg -> Flux.deferContextual(ctx -> {
+					MessageAndSpansAndScope messageAndSpansAndScope = ctx.get(MessageAndSpansAndScope.class);
+					messageAndSpansAndScope.messageAndSpans = msg;
+					messageAndSpansAndScope.span = msg.childSpan;
+					setNameAndTag(targetFunction, msg.childSpan);
+					messageAndSpansAndScope.scope = tracer.withSpan(msg.childSpan);
+					return Mono.just(msg.msg);
+				}));
+		if (targetFunction.isConsumer()) {
+			return targetFunction.apply(reactorStreamConsumer(flux));
+		}
+		final Flux<Message> function = ((Flux<Message>) targetFunction.apply(flux));
+		return Flux.deferContextual(contextView -> {
+			MessageAndSpansAndScope msg = contextView.get(MessageAndSpansAndScope.class);
+			return function.doOnNext(message -> {
+				msg.end();
+				msg.handle();
+			}).map(msgResult -> {
+				MessageAndSpan messageAndSpan = traceMessageHandler.wrapOutputMessage(msgResult,
+						msg.messageAndSpans.parentSpan, outputDestination(targetFunction.getFunctionDefinition()));
+				traceMessageHandler.afterMessageHandled(messageAndSpan.span, null);
+				return messageAndSpan.msg;
+			})
+					// TODO: Fix me when this is resolved in Reactor
+					// .doOnSubscribe(__ -> scope.close())
+					.doOnError(msg::error).doFinally(signalType -> {
+						if (!msg.isHandled()) {
+							msg.end();
+						}
+					});
+		}).contextWrite(contextView -> contextView.put(MessageAndSpansAndScope.class, new MessageAndSpansAndScope()));
+	}
+
+	private Object reactorStreamConsumer(Object result) {
+		if (result instanceof Mono) {
+			return Mono.deferContextual(contextView -> {
+				MessageAndSpansAndScope msg = contextView.get(MessageAndSpansAndScope.class);
+				return ((Mono<Message>) result)
+						// TODO: Fix me when this is resolved in Reactor
+						// .doOnSubscribe(__ -> scope.close())
+						.doOnError(msg::error).doFinally(signalType -> {
+							msg.end();
+						});
+			}).contextWrite(
+					contextView -> contextView.put(MessageAndSpansAndScope.class, new MessageAndSpansAndScope()));
+		}
+		return Flux.deferContextual(contextView -> {
+			MessageAndSpansAndScope msg = contextView.get(MessageAndSpansAndScope.class);
+			return ((Flux<Message>) result)
+					// TODO: Fix me when this is resolved in Reactor
+					// .doOnSubscribe(__ -> scope.close())
+					.doOnError(msg::error).doFinally(signalType -> {
+						msg.end();
+					});
+		}).contextWrite(contextView -> contextView.put(MessageAndSpansAndScope.class, new MessageAndSpansAndScope()));
+	}
+
+	private Object reactorStreamSupplier(Publisher<?> message,
+			SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction) {
+		Publisher<?> publisher = (Publisher<?>) targetFunction.get();
+		if (publisher instanceof Mono) {
+			if (log.isDebugEnabled()) {
+				log.debug("Will instrument a stream Mono supplier");
+			}
+			Mono mono = (Mono) publisher;
+			publisher = ReactorSleuth.tracedMono(tracer, tracer.currentTraceContext(),
+					targetFunction.getFunctionDefinition(), () -> mono, (msg, s) -> {
+						customizedInputMessageSpan(s, msg instanceof Message ? (Message) msg : null);
+					}).map(object -> toMessage(object))
+					.map(object -> this.getMessageAndSpans((Message) object, targetFunction.getFunctionDefinition(),
+							setNameAndTag(targetFunction, tracer.currentSpan())))
+					.doOnNext(wrappedOutputMessage -> customizedOutputMessageSpan(
+							((MessageAndSpan) wrappedOutputMessage).span, ((MessageAndSpan) wrappedOutputMessage).msg))
+					.doOnNext(wrappedOutputMessage -> traceMessageHandler
+							.afterMessageHandled(((MessageAndSpan) wrappedOutputMessage).span, null))
+					.map(wrappedOutputMessage -> ((MessageAndSpan) wrappedOutputMessage).msg);
+		}
+		else {
+			if (log.isDebugEnabled()) {
+				log.debug("Will instrument a stream Flux supplier");
+			}
+			Flux flux = (Flux) publisher;
+			publisher = ReactorSleuth.tracedFlux(tracer, tracer.currentTraceContext(),
+					targetFunction.getFunctionDefinition(), () -> flux, (msg, s) -> {
+						customizedInputMessageSpan(s, msg instanceof Message ? (Message) msg : null);
+					}).map(object -> toMessage(object))
+					.map(object -> this.getMessageAndSpans((Message) object, targetFunction.getFunctionDefinition(),
+							setNameAndTag(targetFunction, tracer.currentSpan())))
+					.doOnNext(wrappedOutputMessage -> customizedOutputMessageSpan(
+							((MessageAndSpan) wrappedOutputMessage).span, ((MessageAndSpan) wrappedOutputMessage).msg))
+					.doOnNext(wrappedOutputMessage -> traceMessageHandler
+							.afterMessageHandled(((MessageAndSpan) wrappedOutputMessage).span, null))
+					.map(wrappedOutputMessage -> ((MessageAndSpan) wrappedOutputMessage).msg);
+		}
+		return publisher;
+	}
+
+	private Span setNameAndTag(SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction, Span span) {
+		return span.name(targetFunction.getFunctionDefinition()).tag(SleuthMessagingSpan.Tags.FUNCTION_NAME.getKey(),
+				targetFunction.getFunctionDefinition());
+	}
+
+	private Object nonReactorStream(Message<byte[]> message,
+			SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction) {
 		MessageAndSpans invocationMessage = null;
 		Span span;
 		if (message == null && targetFunction.isSupplier()) { // Supplier
 			if (log.isDebugEnabled()) {
 				log.debug("Creating a span for a supplier");
 			}
-			span = this.tracer.nextSpan().name(targetFunction.getFunctionDefinition());
+			span = setNameAndTag(targetFunction, this.tracer.nextSpan());
 			customizedInputMessageSpan(span, null);
 		}
 		else {
@@ -117,7 +295,7 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 			if (log.isDebugEnabled()) {
 				log.debug("Wrapped input msg " + invocationMessage);
 			}
-			span = invocationMessage.childSpan;
+			span = setNameAndTag(targetFunction, invocationMessage.childSpan);
 		}
 		Object result;
 		Throwable throwable = null;
@@ -164,6 +342,10 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 		this.customizers.forEach(cust -> cust.customizeInputMessageSpan(spanToCustomize, msg));
 	}
 
+	private void customizedOutputMessageSpan(Span spanToCustomize, Message<?> msg) {
+		this.customizers.forEach(cust -> cust.customizeOutputMessageSpan(spanToCustomize, msg));
+	}
+
 	private Message<?> toMessage(Object result) {
 		if (!(result instanceof Message)) {
 			return MessageBuilder.withPayload(result).build();
@@ -197,4 +379,39 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 		this.functionToDestinationCache.clear();
 	}
 
+	static class MessageAndSpansAndScope {
+
+		MessageAndSpans messageAndSpans;
+
+		Span span;
+
+		Tracer.SpanInScope scope;
+
+		boolean handled;
+
+		void error(Throwable throwable) {
+			if (this.span != null) {
+				this.span.error(throwable);
+			}
+		}
+
+		void handle() {
+			this.handled = true;
+		}
+
+		boolean isHandled() {
+			return this.handled;
+		}
+
+		void end() {
+			if (this.span != null) {
+				this.span.end();
+			}
+			if (this.scope != null) {
+				this.scope.close();
+			}
+		}
+
+	}
+
 }

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceMessageHandler.java

@@ -152,7 +152,6 @@ class TraceMessageHandler {
 		Span.Builder consumerSpanBuilder = SleuthMessagingSpan.MESSAGING_SPAN
 				.wrap(this.propagator.extract(headers, this.extractor));
 		Span consumerSpan = consumerSpan(destinationName, consumerSpanBuilder, message);
-
 		if (log.isDebugEnabled()) {
 			log.debug("Built a consumer span " + consumerSpan);
 		}

spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java

@@ -16,11 +16,17 @@
 
 package org.springframework.cloud.sleuth.instrument.messaging;
 
+import java.time.Duration;
 import java.util.Collections;
+import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import org.springframework.cloud.function.context.FunctionRegistration;
 import org.springframework.cloud.function.context.FunctionType;
@@ -29,10 +35,15 @@ import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry
 import org.springframework.cloud.function.context.config.JsonMessageConverter;
 import org.springframework.cloud.function.json.JacksonMapper;
 import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.Span.Builder;
+import org.springframework.cloud.sleuth.TraceContext;
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.propagation.Propagator;
 import org.springframework.cloud.sleuth.tracer.SimpleTracer;
 import org.springframework.core.convert.support.DefaultConversionService;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.converter.CompositeMessageConverter;
+import org.springframework.messaging.support.MessageBuilder;
 import org.springframework.mock.env.MockEnvironment;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -40,28 +51,175 @@ import static org.assertj.core.api.BDDAssertions.then;
 
 class TraceFunctionAroundWrapperTests {
 
-	@Test
-	void test_tracing_with_supplier() {
-		CompositeMessageConverter messageConverter = new CompositeMessageConverter(
-				Collections.singletonList(new JsonMessageConverter(new JacksonMapper(new ObjectMapper()))));
+	CompositeMessageConverter messageConverter = new CompositeMessageConverter(
+			Collections.singletonList(new JsonMessageConverter(new JacksonMapper(new ObjectMapper()))));
+
+	SimpleFunctionRegistry catalog = new SimpleFunctionRegistry(new DefaultConversionService(), messageConverter,
+			new JacksonMapper(new ObjectMapper()));
+
+	SimpleTracer tracer = new SimpleTracer();
+
+	MockEnvironment mockEnvironment = mockEnvironment();
+
+	TraceFunctionAroundWrapper wrapper = new TraceFunctionAroundWrapper(mockEnvironment, tracer, testPropagator(),
+			new MessageHeaderPropagatorSetter(), new MessageHeaderPropagatorGetter()) {
+		@Override
+		MessageAndSpan getMessageAndSpans(Message<?> resultMessage, String name, Span spanFromMessage) {
+			return new MessageAndSpan(resultMessage, spanFromMessage);
+		}
+	};
+
+	private Propagator testPropagator() {
+		return new Propagator() {
 
-		SimpleTracer tracer = new SimpleTracer();
-		TraceFunctionAroundWrapper wrapper = new TraceFunctionAroundWrapper(null, tracer, null, null, null) {
 			@Override
-			MessageAndSpan getMessageAndSpans(Message<?> resultMessage, String name, Span spanFromMessage) {
-				return new MessageAndSpan(resultMessage, spanFromMessage);
+			public <C> void inject(TraceContext context, C carrier, Setter<C> setter) {
+				setter.set(carrier, "superHeader", "test");
+			}
+
+			@Override
+			public List<String> fields() {
+				return Collections.singletonList("superHeader");
+			}
+
+			@Override
+			public <C> Builder extract(C carrier, Getter<C> getter) {
+				return tracer.spanBuilder();
 			}
 		};
+	}
 
+	@Test
+	void test_tracing_with_supplier() {
 		FunctionRegistration<Greeter> registration = new FunctionRegistration<>(new Greeter(), "greeter")
 				.type(FunctionType.of(Greeter.class));
-		SimpleFunctionRegistry catalog = new SimpleFunctionRegistry(new DefaultConversionService(), messageConverter,
-				new JacksonMapper(new ObjectMapper()));
 		catalog.register(registration);
 		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
 		Message<?> result = (Message<?>) wrapper.apply(null, function);
+
 		assertThat(result.getPayload()).isEqualTo("hello");
 		assertThat(tracer.getOnlySpan().name).isEqualTo("greeter");
+		assertThatAllSpansAreStartedAndStopped();
+	}
+
+	@Test
+	void test_tracing_with_function() {
+		FunctionRegistration<GreeterFunction> registration = new FunctionRegistration<>(new GreeterFunction(),
+				"greeter").type(FunctionType.of(GreeterFunction.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		Message<?> result = (Message<?>) wrapper
+				.apply(MessageBuilder.withPayload("hello").setHeader("superHeader", "someValue").build(), function);
+
+		assertThat(result.getPayload()).isEqualTo("HELLO");
+		assertThat(tracer.spans).hasSize(3);
+		assertThat(tracer.spans.get(0).name).isEqualTo("handle");
+		assertThat(tracer.spans.get(1).name).isEqualTo("greeter");
+		assertThat(tracer.spans.get(2).name).isEqualTo("send");
+		assertThatAllSpansAreStartedAndStopped();
+	}
+
+	@Test
+	void test_tracing_with_consumer() {
+		GreeterConsumer consumer = new GreeterConsumer();
+		FunctionRegistration<GreeterConsumer> registration = new FunctionRegistration<>(consumer, "greeter")
+				.type(FunctionType.of(GreeterConsumer.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		wrapper.apply(MessageBuilder.withPayload("hello").setHeader("superHeader", "someValue").build(), function);
+
+		assertThat(consumer.result).isEqualTo("HELLO");
+		assertThat(tracer.spans).hasSize(2);
+		assertThat(tracer.spans.get(0).name).isEqualTo("handle");
+		assertThat(tracer.spans.get(1).name).isEqualTo("greeter");
+		assertThatAllSpansAreStartedAndStopped();
+	}
+
+	@Test
+	void should_trace_when_reactive_mono_supplier() {
+		FunctionRegistration<ReactiveMonoGreeter> registration = new FunctionRegistration<>(new ReactiveMonoGreeter(),
+				"greeter").type(FunctionType.of(ReactiveMonoGreeter.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		Message<?> result = ((Mono<Message<?>>) wrapper.apply(null, function)).block(Duration.ofSeconds(5));
+
+		assertThat(result.getPayload()).isEqualTo("hello");
+		assertThat(tracer.getOnlySpan().name).isEqualTo("greeter");
+		assertThatAllSpansAreStartedAndStopped();
+	}
+	
+	@Test
+	void should_trace_when_reactive_mono_function() {
+		FunctionRegistration<ReactiveMonoGreeterFunction> registration = new FunctionRegistration<>(
+				new ReactiveMonoGreeterFunction(), "greeter").type(FunctionType.of(ReactiveMonoGreeterFunction.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		Message<?> result = ((Mono<Message<?>>) wrapper.apply(
+				Mono.just(MessageBuilder.withPayload("hello").setHeader("superHeader", "someValue").build()), function))
+						.block(Duration.ofSeconds(5));
+
+		assertThat(result.getPayload()).isEqualTo("HELLO");
+		assertThat(tracer.spans).hasSize(3);
+		assertThat(tracer.spans.get(0).name).isEqualTo("handle");
+		assertThat(tracer.spans.get(1).name).isEqualTo("greeter");
+		assertThat(tracer.spans.get(2).name).isEqualTo("send");
+		assertThatAllSpansAreStartedAndStopped();
+	}
+
+	@Test
+	void should_trace_when_reactive_flux_supplier() {
+		FunctionRegistration<ReactiveFluxGreeter> registration = new FunctionRegistration<>(new ReactiveFluxGreeter(),
+				"greeter").type(FunctionType.of(ReactiveFluxGreeter.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		Message<?> result = ((Flux<Message<?>>) wrapper.apply(null, function)).blockFirst(Duration.ofSeconds(5));
+
+		assertThat(result.getPayload()).isEqualTo("hello");
+		assertThat(tracer.getOnlySpan().name).isEqualTo("greeter");
+		assertThatAllSpansAreStartedAndStopped();
+	}
+
+	@Test
+	void should_trace_when_reactive_flux_function() {
+		FunctionRegistration<ReactiveFluxGreeterFunction> registration = new FunctionRegistration<>(
+				new ReactiveFluxGreeterFunction(), "greeter").type(FunctionType.of(ReactiveFluxGreeterFunction.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		Message<?> result = ((Flux<Message<?>>) wrapper.apply(
+				Flux.just(MessageBuilder.withPayload("hello").setHeader("superHeader", "someValue").build()), function))
+						.blockFirst(Duration.ofSeconds(5));
+
+		assertThat(result.getPayload()).isEqualTo("HELLO");
+		assertThat(tracer.spans).hasSize(3);
+		assertThat(tracer.spans.get(0).name).isEqualTo("handle");
+		assertThat(tracer.spans.get(1).name).isEqualTo("greeter");
+		assertThat(tracer.spans.get(2).name).isEqualTo("send");
+		assertThatAllSpansAreStartedAndStopped();
+	}
+
+	@Test
+	void should_trace_when_reactive_flux_consumer() {
+		ReactiveFluxGreeterConsumer consumer = new ReactiveFluxGreeterConsumer(this.tracer);
+		FunctionRegistration<ReactiveFluxGreeterConsumer> registration = new FunctionRegistration<>(consumer, "greeter")
+				.type(FunctionType.of(ReactiveFluxGreeterConsumer.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		wrapper.apply(Flux.just(MessageBuilder.withPayload("hello").setHeader("superHeader", "someValue").build()),
+				function);
+
+		assertThat(consumer.result).isEqualTo("HELLO");
+		assertThat(tracer.spans).hasSize(2);
+		assertThat(tracer.spans.get(0).name).isEqualTo("handle");
+		assertThat(tracer.spans.get(1).name).isEqualTo("greeter");
+		assertThatAllSpansAreStartedAndStopped();
 	}
 
 	@Test
@@ -105,6 +263,18 @@ class TraceFunctionAroundWrapperTests {
 		assertThat(wrapper.outputDestination("marcin")).isEqualTo("bob");
 	}
 
+	private MockEnvironment mockEnvironment() {
+		MockEnvironment mockEnvironment = new MockEnvironment();
+		mockEnvironment.setProperty("spring.cloud.stream.bindings.greeter-in-0.destination", "oleg");
+		mockEnvironment.setProperty("spring.cloud.stream.bindings.greeter-out-0.destination", "bob");
+		return mockEnvironment;
+	}
+	
+	private void assertThatAllSpansAreStartedAndStopped() {
+		assertThat(tracer.spans.stream()
+				.allMatch(s -> s.started && s.ended)).as("All spans must be started and stopped").isTrue();
+	}
+
 	private static class Greeter implements Supplier<String> {
 
 		@Override
@@ -114,4 +284,84 @@ class TraceFunctionAroundWrapperTests {
 
 	}
 
+	private static class GreeterFunction implements Function<String, String> {
+
+		@Override
+		public String apply(String in) {
+			return in.toUpperCase();
+		}
+
+	}
+
+	private static class GreeterConsumer implements Consumer<String> {
+
+		String result;
+
+		@Override
+		public void accept(String in) {
+			this.result = in.toUpperCase();
+		}
+
+	}
+
+	private static class ReactiveMonoGreeter implements Supplier<Mono<Message<String>>> {
+
+		@Override
+		public Mono<Message<String>> get() {
+			return Mono.just(MessageBuilder.withPayload("hello").build());
+		}
+
+	}
+
+	private static class ReactiveMonoGreeterFunction implements Function<Mono<Message<String>>, Mono<Message<String>>> {
+
+		@Override
+		public Mono<Message<String>> apply(Mono<Message<String>> in) {
+			return in.map(s -> MessageBuilder.fromMessage(s).withPayload(s.getPayload().toUpperCase()).build());
+		}
+
+	}
+
+	private static class ReactiveFluxGreeter implements Supplier<Flux<Message<String>>> {
+
+		@Override
+		public Flux<Message<String>> get() {
+			return Flux.just(MessageBuilder.withPayload("hello").build());
+		}
+
+	}
+
+	private static class ReactiveFluxGreeterFunction implements Function<Flux<Message<String>>, Flux<Message<String>>> {
+
+		@Override
+		public Flux<Message<String>> apply(Flux<Message<String>> in) {
+			return in.map(s -> MessageBuilder.fromMessage(s).withPayload(s.getPayload().toUpperCase()).build());
+		}
+
+	}
+
+	private static class ReactiveFluxGreeterConsumer implements Consumer<Flux<Message<String>>> {
+
+		String result;
+		
+		private final Tracer tracer;
+		
+		ReactiveFluxGreeterConsumer(Tracer tracer) {
+			this.tracer = tracer;
+		}
+
+		@Override
+		public void accept(Flux<Message<String>> in) {
+			in.map(s -> s.getPayload().toUpperCase()).doOnNext(s -> {
+				result = s;
+			})
+			.doOnNext(s -> {
+				tracer.currentSpan().end();
+				tracer.withSpan(null);
+			})
+			.subscribe();
+		}
+
+	}
+
 }

spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/tracer/SimpleSpan.java

@@ -54,6 +54,8 @@ public class SimpleSpan implements Span {
 
 	public boolean noOp;
 
+	public Span parent;
+
 	@Override
 	public boolean isNoop() {
 		return this.noOp;

spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/tracer/SimpleTracer.java

@@ -24,6 +24,7 @@ import java.util.Map;
 import org.assertj.core.api.BDDAssertions;
 
 import org.springframework.cloud.sleuth.BaggageInScope;
+import org.springframework.cloud.sleuth.CurrentTraceContext;
 import org.springframework.cloud.sleuth.ScopedSpan;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanCustomizer;
@@ -42,7 +43,9 @@ public class SimpleTracer implements Tracer {
 
 	@Override
 	public Span nextSpan(Span parent) {
-		return new SimpleSpan();
+		SimpleSpan span = nextSpan();
+		span.parent = parent;
+		return span;
 	}
 
 	public SimpleSpan getOnlySpan() {
@@ -125,4 +128,9 @@ public class SimpleTracer implements Tracer {
 		return null;
 	}
 
+	@Override
+	public CurrentTraceContext currentTraceContext() {
+		return null;
+	}
+
 }

tests/brave/spring-cloud-sleuth-instrumentation-messaging-tests/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/StreamMessageOperatorsTests.java

@@ -43,7 +43,8 @@ import org.springframework.messaging.support.MessageBuilder;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
-@SpringBootTest(classes = StreamMessageOperatorsTests.Config.class, webEnvironment = SpringBootTest.WebEnvironment.NONE)
+@SpringBootTest(classes = StreamMessageOperatorsTests.Config.class, webEnvironment = SpringBootTest.WebEnvironment.NONE,
+		properties = "spring.sleuth.function.enabled=false")
 public class StreamMessageOperatorsTests {
 
 	@Autowired(required = false)
@@ -60,8 +61,7 @@ public class StreamMessageOperatorsTests {
 
 	@Test
 	void should_instrument_a_simple_message_to_message_function() {
-		assertThat(tracingChannelInterceptor).as("Ensure that we're doing instrumentation via function wrapper")
-				.isNull();
+		assertThat(this.tracingChannelInterceptor).as("Ensure that we're doing instrumentation manually").isNull();
 
 		this.inputDestination.send(MessageBuilder.withPayload("hello".getBytes())
 				.setHeader("b3", "4883117762eb9420-4883117762eb9420-1").build());

tests/common/src/main/java/org/springframework/cloud/sleuth/test/TestTracer.java

@@ -21,6 +21,7 @@ import java.util.Map;
 import java.util.Queue;
 
 import org.springframework.cloud.sleuth.BaggageInScope;
+import org.springframework.cloud.sleuth.CurrentTraceContext;
 import org.springframework.cloud.sleuth.ScopedSpan;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanCustomizer;
@@ -118,4 +119,9 @@ public class TestTracer implements Tracer, AutoCloseable {
 		return createdSpans;
 	}
 
+	@Override
+	public CurrentTraceContext currentTraceContext() {
+		return this.delegate.currentTraceContext();
+	}
+
 }