ALmost there

Author: marcingrzejszczak

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

@@ -25,6 +25,9 @@
 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;
 import org.springframework.cloud.function.context.catalog.FunctionTypeUtils;
@@ -39,9 +42,6 @@
 import org.springframework.messaging.support.MessageBuilder;
 import org.springframework.messaging.support.MessageHeaderAccessor;
 
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
-
 /**
  * Trace representation of a {@link FunctionAroundWrapper}.
  *

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

@@ -16,15 +16,18 @@
 
 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.Disabled;
 import org.junit.jupiter.api.Test;
-import org.mockito.Answers;
-import org.mockito.BDDMockito;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import org.springframework.cloud.function.context.FunctionRegistration;
 import org.springframework.cloud.function.context.FunctionType;
@@ -33,93 +36,204 @@
 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.TraceContext;
 import org.springframework.cloud.sleuth.Span.Builder;
+import org.springframework.cloud.sleuth.TraceContext;
 import org.springframework.cloud.sleuth.propagation.Propagator;
-import org.springframework.cloud.sleuth.propagation.Propagator.Getter;
-import org.springframework.cloud.sleuth.propagation.Propagator.Setter;
 import org.springframework.cloud.sleuth.tracer.SimpleTracer;
 import org.springframework.core.convert.support.DefaultConversionService;
-import org.springframework.integration.support.MessageBuilder;
 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;
 import static org.assertj.core.api.BDDAssertions.then;
 
 class TraceFunctionAroundWrapperTests {
-	
+
 	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(), 
+	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);
 		}
 	};
 
-	/**
-	 * @return
-	 */
 	private Propagator testPropagator() {
 		return new Propagator() {
-			
+
 			@Override
 			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) {
-				Builder builder = BDDMockito.mock(Builder.class, Answers.RETURNS_SELF);
-				BDDMockito.given(builder.start()).willReturn(tracer.nextSpan());
-				return builder;
+				return tracer.spanBuilder();
 			}
 		};
 	}
-	
+
 	@Test
 	void test_tracing_with_supplier() {
 		FunctionRegistration<Greeter> registration = new FunctionRegistration<>(new Greeter(), "greeter")
 				.type(FunctionType.of(Greeter.class));
 		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");
 	}
-	
+
 	@Test
 	void test_tracing_with_function() {
-		FunctionRegistration<GreeterFunction> registration = new FunctionRegistration<>(new GreeterFunction(), "greeter")
-				.type(FunctionType.of(GreeterFunction.class));
+		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);
-		
+
+		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");
+	}
+
+	@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");
+	}
+
+	@Disabled("TODO: don't we want to support this?")
+	@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");
 	}
 
+	@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");
+	}
+
+	@Disabled("TODO: don't we want to support this?")
+	@Test
+	void should_trace_when_reactive_mono_consumer() {
+		ReactiveMonoGreeterConsumer consumer = new ReactiveMonoGreeterConsumer();
+		FunctionRegistration<ReactiveMonoGreeterConsumer> registration = new FunctionRegistration<>(consumer, "greeter")
+				.type(FunctionType.of(ReactiveMonoGreeterConsumer.class));
+		catalog.register(registration);
+		FunctionInvocationWrapper function = catalog.lookup("greeter");
+
+		wrapper.apply(Mono.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");
+	}
+
+	@Disabled("TODO: don't we want to support this?")
+	@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");
+	}
+
+	@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");
+	}
+
+	@Test
+	void should_trace_when_reactive_flux_consumer() {
+		ReactiveFluxGreeterConsumer consumer = new ReactiveFluxGreeterConsumer();
+		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");
+	}
+
 	@Test
 	void should_clear_cache_on_refresh() {
 		TraceFunctionAroundWrapper wrapper = new TraceFunctionAroundWrapper(null, null, null, null, null);
@@ -160,7 +274,7 @@ void should_point_to_proper_destination_when_working_with_remapped_functions() {
 
 		assertThat(wrapper.outputDestination("marcin")).isEqualTo("bob");
 	}
-	
+
 	private MockEnvironment mockEnvironment() {
 		MockEnvironment mockEnvironment = new MockEnvironment();
 		mockEnvironment.setProperty("spring.cloud.stream.bindings.greeter-in-0.destination", "oleg");
@@ -176,14 +290,87 @@ public String get() {
 		}
 
 	}
-	
+
 	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 ReactiveMonoGreeterConsumer implements Consumer<Mono<Message<String>>> {
+
+		String result;
+
+		@Override
+		public void accept(Mono<Message<String>> in) {
+			in.map(s -> s.getPayload().toUpperCase()).doOnNext(s -> {
+				result = s;
+			}).subscribe();
+		}
+
+	}
+
+	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;
+
+		@Override
+		public void accept(Flux<Message<String>> in) {
+			in.map(s -> s.getPayload().toUpperCase()).doOnNext(s -> {
+				result = s;
+			}).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

@@ -43,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() {