Fixed type checking for Publisher<Non-Message>; fixes gh-2079

benchmarks/src/main/java/org/springframework/cloud/sleuth/benchmarks/app/stream/SleuthBenchmarkingStreamApplication.java

@@ -134,7 +134,7 @@ public class SleuthBenchmarkingStreamApplication {
 	@Bean(name = "myFlux")
 	@ConditionalOnProperty(value = "spring.sleuth.function.type", havingValue = "DECORATE_ON_EACH",
 			matchIfMissing = true)
-	public Function<Flux<Message<String>>, Flux<Message<String>>> onEachFunction() {
+	public Function<Flux<String>, Flux<String>> onEachFunction() {
 		log.info("on each function");
 		return new SleuthFunction();
 	}
@@ -142,14 +142,14 @@ public class SleuthBenchmarkingStreamApplication {
 	@Bean(name = "myFlux")
 	@ConditionalOnProperty(value = "spring.sleuth.function.type", havingValue = "DECORATE_QUEUES",
 			matchIfMissing = true)
-	public Function<Flux<Message<String>>, Flux<Message<String>>> decorateQueuesFunction() {
+	public Function<Flux<String>, Flux<String>> decorateQueuesFunction() {
 		log.info("decorate queues function");
 		return new SleuthFunction();
 	}
 
 	@Bean(name = "myFlux")
 	@ConditionalOnProperty(value = "spring.sleuth.function.type", havingValue = "DECORATE_ON_LAST")
-	public Function<Flux<Message<String>>, Flux<Message<String>>> onLastFunction() {
+	public Function<Flux<String>, Flux<String>> onLastFunction() {
 		log.info("on last function");
 		return new SleuthFunction();
 	}
@@ -269,18 +269,18 @@ class SleuthNonReactiveFunction implements Function<String, String> {
 
 }
 
-class SleuthFunction implements Function<Flux<Message<String>>, Flux<Message<String>>> {
+class SleuthFunction implements Function<Flux<String>, Flux<String>> {
 
 	private static final Logger log = LoggerFactory.getLogger(SleuthFunction.class);
 
 	static final Scheduler SCHEDULER = Schedulers.newParallel("sleuthFunction");
 
 	@Override
-	public Flux<Message<String>> apply(Flux<Message<String>> input) {
+	public Flux<String> apply(Flux<String> input) {
 		return input.doOnEach(signal -> log.info("Got a message"))
 				.flatMap(s -> Mono.delay(Duration.ofMillis(1), SCHEDULER).map(aLong -> {
 					log.info("Logging [{}] from flat map", s);
-					return MessageBuilder.withPayload(s.getPayload().toUpperCase()).build();
+					return s.toUpperCase();
 				}));
 	}
 

benchmarks/src/test/java/org/springframework/cloud/sleuth/benchmarks/jmh/stream/MicroBenchmarkStreamTests.java

@@ -56,9 +56,9 @@ import org.springframework.messaging.support.MessageBuilder;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
-@Measurement(iterations = 1, time = 1)
-@Warmup(iterations = 1, time = 1)
-@Fork(1)
+@Measurement(iterations = 10, time = 1)
+@Warmup(iterations = 10, time = 1)
+@Fork(4)
 @BenchmarkMode(Mode.SampleTime)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 @Microbenchmark

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

@@ -94,19 +94,27 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 			return targetFunction.apply(message); // no instrumentation
 		}
 		else if (targetFunction.isInputTypePublisher() || targetFunction.isOutputTypePublisher()) {
+			if (message != null && !(message instanceof Publisher)) {
+				logDebugAboutMessageTypes(message);
+				return targetFunction.apply(message); // no instrumentation
+			}
 			return reactorStream((Publisher) message, targetFunction);
 		}
 		else if (message != null && !(message instanceof Message)) {
-			if (log.isDebugEnabled()) {
-				String messageClass = message.getClass().getName();
-				log.debug("We only support tracing for Message types. You need to wrap your function type ["
-						+ messageClass + "] into [Message<" + messageClass + ">]");
-			}
+			logDebugAboutMessageTypes(message);
 			return targetFunction.apply(message); // no instrumentation
 		}
 		return nonReactorStream((Message<byte[]>) message, targetFunction);
 	}
 
+	private void logDebugAboutMessageTypes(Object message) {
+		if (log.isDebugEnabled()) {
+			String messageClass = message.getClass().getName();
+			log.debug("We only support tracing for Message types. You need to wrap your function type [" + messageClass
+					+ "] into [Message<" + messageClass + ">]");
+		}
+	}
+
 	private Object reactorStream(Publisher messageStream,
 			SimpleFunctionRegistry.FunctionInvocationWrapper targetFunction) {
 		if (messageStream == null && targetFunction.isSupplier()) { // Supplier

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

@@ -25,8 +25,12 @@ import java.util.function.Supplier;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Scheduler;
+import reactor.core.scheduler.Schedulers;
 
 import org.springframework.cloud.function.context.FunctionRegistration;
 import org.springframework.cloud.function.context.FunctionType;
@@ -134,6 +138,19 @@ class TraceFunctionAroundWrapperTests {
 		assertThat(tracer.spans).isEmpty();
 	}
 
+	@Test
+	void test_tracing_with_flux_to_flux_without_message() {
+		FunctionRegistration<FluxToFluxFunction> registration = new FunctionRegistration<>(new FluxToFluxFunction(),
+				"greeter").type(FunctionType.of(FluxToFluxFunction.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		String result = (String) ((Flux) wrapper.apply("hello", function)).blockFirst();
+
+		assertThat(result).isEqualTo("HELLO");
+		assertThat(tracer.spans).isEmpty();
+	}
+
 	@Test
 	void test_tracing_with_consumer() {
 		GreeterConsumer consumer = new GreeterConsumer();
@@ -375,4 +392,21 @@ class TraceFunctionAroundWrapperTests {
 
 	}
 
+	static class FluxToFluxFunction implements Function<Flux<String>, Flux<String>> {
+
+		private static final Logger log = LoggerFactory.getLogger(FluxToFluxFunction.class);
+
+		static final Scheduler SCHEDULER = Schedulers.newParallel("sleuthFunction");
+
+		@Override
+		public Flux<String> apply(Flux<String> input) {
+			return input.doOnEach(signal -> log.info("Got a message"))
+					.flatMap(s -> Mono.delay(Duration.ofMillis(1), SCHEDULER).map(aLong -> {
+						log.info("Logging [{}] from flat map", s);
+						return s.toUpperCase();
+					}));
+		}
+
+	}
+
 }