Fixing sc function (#1983)

- Removed an unnecessary child span; the parent span was never ended
- Added tests

spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/messaging/TraceFunctionAutoConfiguration.java

@@ -16,19 +16,28 @@
 
 package org.springframework.cloud.sleuth.autoconfig.instrument.messaging;
 
+import java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.cloud.context.scope.refresh.RefreshScopeRefreshedEvent;
 import org.springframework.cloud.function.context.catalog.FunctionAroundWrapper;
+import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.autoconfig.brave.BraveAutoConfiguration;
+import org.springframework.cloud.sleuth.instrument.messaging.FunctionMessageSpanCustomizer;
 import org.springframework.cloud.sleuth.instrument.messaging.TraceFunctionAroundWrapper;
 import org.springframework.cloud.sleuth.propagation.Propagator;
+import org.springframework.cloud.stream.messaging.DirectWithAttributesChannel;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.env.Environment;
+import org.springframework.messaging.Message;
 import org.springframework.messaging.support.MessageHeaderAccessor;
 
 /**
@@ -48,8 +57,61 @@ public class TraceFunctionAutoConfiguration {
 
 	@Bean
 	TraceFunctionAroundWrapper traceFunctionAroundWrapper(Environment environment, Tracer tracer, Propagator propagator,
-			Propagator.Setter<MessageHeaderAccessor> injector, Propagator.Getter<MessageHeaderAccessor> extractor) {
-		return new TraceFunctionAroundWrapper(environment, tracer, propagator, injector, extractor);
+			Propagator.Setter<MessageHeaderAccessor> injector, Propagator.Getter<MessageHeaderAccessor> extractor,
+			ObjectProvider<List<FunctionMessageSpanCustomizer>> customizers) {
+		return new TraceFunctionAroundWrapper(environment, tracer, propagator, injector, extractor,
+				customizers.getIfAvailable(ArrayList::new));
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnClass(DirectWithAttributesChannel.class)
+	static class TraceFunctionStreamConfiguration {
+
+		@Configuration(proxyBeanMethods = false)
+		@ConditionalOnClass(name = "org.springframework.cloud.stream.binder.kafka.config.KafkaBinderConfiguration")
+		@ConditionalOnMissingClass("org.springframework.cloud.stream.binder.rabbit.properties.RabbitBinderConfigurationProperties")
+		static class KafkaOnlyStreamConfiguration {
+
+			@Bean
+			FunctionMessageSpanCustomizer traceKafkaFunctionMessageSpanCustomizer() {
+				return new FunctionMessageSpanCustomizer() {
+					@Override
+					public void customizeInputMessageSpan(Span span, Message<?> message) {
+						span.remoteServiceName("kafka");
+					}
+
+					@Override
+					public void customizeOutputMessageSpan(Span span, Message<?> message) {
+						span.remoteServiceName("kafka");
+					}
+				};
+			}
+
+		}
+
+		@Configuration(proxyBeanMethods = false)
+		@ConditionalOnClass(
+				name = "org.springframework.cloud.stream.binder.rabbit.properties.RabbitBinderConfigurationProperties")
+		@ConditionalOnMissingClass("org.springframework.cloud.stream.binder.kafka.config.KafkaBinderConfiguration")
+		static class RabbitOnlyStreamConfiguration {
+
+			@Bean
+			FunctionMessageSpanCustomizer traceRabbitFunctionMessageSpanCustomizer() {
+				return new FunctionMessageSpanCustomizer() {
+					@Override
+					public void customizeInputMessageSpan(Span span, Message<?> message) {
+						span.remoteServiceName("rabbitmq");
+					}
+
+					@Override
+					public void customizeOutputMessageSpan(Span span, Message<?> message) {
+						span.remoteServiceName("rabbitmq");
+					}
+				};
+			}
+
+		}
+
 	}
 
 }

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

@@ -54,7 +54,6 @@ public class BravePropagator implements Propagator {
 	public <C> Span.Builder extract(C carrier, Getter<C> getter) {
 		TraceContextOrSamplingFlags extract = this.tracing.propagation().extractor(getter::get).extract(carrier);
 		if (extract.samplingFlags() == SamplingFlags.EMPTY) {
-			this.tracing.tracer().nextSpan();
 			return new BraveSpanBuilder(this.tracing.tracer());
 		}
 		return BraveSpanBuilder.toBuilder(this.tracing.tracer(), extract);

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

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2013-2021 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
+ *
+ *      https://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.messaging;
+
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.messaging.Message;
+
+/**
+ * Allows customization of messaging spans for Spring Cloud Function instrumentation.
+ *
+ * @author Marcin Grzejszczak
+ * @since 3.0.4
+ */
+public interface FunctionMessageSpanCustomizer {
+
+	/**
+	 * Customizes the span created after wrapping the input message in a span
+	 * representation.
+	 * @param span current span to customize
+	 * @param message received or sent message
+	 */
+	default void customizeInputMessageSpan(Span span, Message<?> message) {
+
+	}
+
+	/**
+	 * Customizes the span wrapping the function execution.
+	 * @param span current span to customize
+	 * @param message message to be sent
+	 */
+	default void customizeFunctionSpan(Span span, Message<?> message) {
+
+	}
+
+	/**
+	 * Customizes the span created for the output message.
+	 * @param span current span to customize
+	 * @param message message to be sent
+	 */
+	default void customizeOutputMessageSpan(Span span, Message<?> message) {
+
+	}
+
+}

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

@@ -188,6 +188,23 @@ public final class MessagingSleuthOperators {
 	 * @return instrumented message
 	 */
 	public static <T> Message<T> handleOutputMessage(BeanFactory beanFactory, Message<T> message, Throwable throwable) {
+		return handleOutputMessage(beanFactory, message, span -> {
+		}, throwable);
+	}
+
+	/**
+	 * Creates an output message with tracer headers and reports the corresponding
+	 * producer span. If the message contains a header called {@code destination} it will
+	 * be used to tag the span with destination name.
+	 * @param beanFactory - bean factory
+	 * @param message - message to which tracer headers should be injected
+	 * @param spanCustomizer - customizer of the output span
+	 * @param throwable - exception that took place while processing the message
+	 * @param <T> - message payload
+	 * @return instrumented message
+	 */
+	public static <T> Message<T> handleOutputMessage(BeanFactory beanFactory, Message<T> message,
+			Consumer<Span> spanCustomizer, Throwable throwable) {
 		TraceMessageHandler traceMessageHandler = TraceMessageHandler.forNonSpringIntegration(beanFactory);
 		Span span = traceMessageHandler.parentSpan(message);
 		span = span != null ? span : traceMessageHandler.consumerSpan(message);
@@ -198,6 +215,7 @@ public final class MessagingSleuthOperators {
 		}
 		MessageAndSpan messageAndSpan = traceMessageHandler.wrapOutputMessage(message, span,
 				String.valueOf(message.getHeaders().getOrDefault("destination", "")));
+		spanCustomizer.accept(messageAndSpan.span);
 		traceMessageHandler.afterMessageHandled(messageAndSpan.span, throwable);
 		return messageAndSpan.msg;
 	}

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

@@ -16,6 +16,8 @@
 
 package org.springframework.cloud.sleuth.instrument.messaging;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -58,17 +60,26 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 
 	private final TraceMessageHandler traceMessageHandler;
 
+	private final List<FunctionMessageSpanCustomizer> customizers;
+
 	final Map<String, String> functionToDestinationCache = new ConcurrentHashMap<>();
 
 	public TraceFunctionAroundWrapper(Environment environment, Tracer tracer, Propagator propagator,
 			Propagator.Setter<MessageHeaderAccessor> injector, Propagator.Getter<MessageHeaderAccessor> extractor) {
+		this(environment, tracer, propagator, injector, extractor, Collections.emptyList());
+	}
+
+	public TraceFunctionAroundWrapper(Environment environment, Tracer tracer, Propagator propagator,
+			Propagator.Setter<MessageHeaderAccessor> injector, Propagator.Getter<MessageHeaderAccessor> extractor,
+			List<FunctionMessageSpanCustomizer> customizers) {
 		this.environment = environment;
 		this.tracer = tracer;
 		this.propagator = propagator;
 		this.injector = injector;
 		this.extractor = extractor;
+		this.customizers = customizers;
 		this.traceMessageHandler = TraceMessageHandler.forNonSpringIntegration(this.tracer, this.propagator,
-				this.injector, this.extractor);
+				this.injector, this.extractor, this.customizers);
 	}
 
 	@Override
@@ -76,23 +87,26 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 		MessageAndSpans invocationMessage = null;
 		Span span;
 		if (message == null && targetFunction.isSupplier()) { // Supplier
-			span = traceMessageHandler.tracer.nextSpan().name(targetFunction.getFunctionDefinition());
+			if (log.isDebugEnabled()) {
+				log.debug("Creating a span for a supplier");
+			}
+			span = this.tracer.nextSpan().name(targetFunction.getFunctionDefinition());
+			customizedInputMessageSpan(span, null);
 		}
 		else {
 			if (log.isDebugEnabled()) {
 				log.debug("Will retrieve the tracing headers from the message");
 			}
-			invocationMessage = traceMessageHandler.wrapInputMessage(message,
+			invocationMessage = this.traceMessageHandler.wrapInputMessage(message,
 					inputDestination(targetFunction.getFunctionDefinition()));
 			if (log.isDebugEnabled()) {
 				log.debug("Wrapped input msg " + invocationMessage);
 			}
 			span = invocationMessage.childSpan;
 		}
-
 		Object result;
 		Throwable throwable = null;
-		try (Tracer.SpanInScope ws = tracer.withSpan(span.start())) {
+		try (Tracer.SpanInScope ws = this.tracer.withSpan(span.start())) {
 			result = invocationMessage == null ? targetFunction.get() : targetFunction.apply(invocationMessage.msg);
 		}
 		catch (Exception e) {
@@ -100,7 +114,7 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 			throw e;
 		}
 		finally {
-			traceMessageHandler.afterMessageHandled(span, throwable);
+			this.traceMessageHandler.afterMessageHandled(span, throwable);
 		}
 		if (result == null) {
 			if (log.isDebugEnabled()) {
@@ -109,10 +123,12 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 			return null;
 		}
 		Message<?> msgResult = toMessage(result);
-
 		MessageAndSpan wrappedOutputMessage;
+		if (log.isDebugEnabled()) {
+			log.debug("Will instrument the output message");
+		}
 		if (invocationMessage != null) {
-			wrappedOutputMessage = traceMessageHandler.wrapOutputMessage(msgResult, invocationMessage.parentSpan,
+			wrappedOutputMessage = this.traceMessageHandler.wrapOutputMessage(msgResult, invocationMessage.parentSpan,
 					outputDestination(targetFunction.getFunctionDefinition()));
 		}
 		else {
@@ -129,6 +145,10 @@ public class TraceFunctionAroundWrapper extends FunctionAroundWrapper
 		return traceMessageHandler.wrapOutputMessage(resultMessage, spanFromMessage, outputDestination(name));
 	}
 
+	private void customizedInputMessageSpan(Span spanToCustomize, Message<?> msg) {
+		this.customizers.forEach(cust -> cust.customizeInputMessageSpan(spanToCustomize, msg));
+	}
+
 	private Message<?> toMessage(Object result) {
 		if (!(result instanceof Message)) {
 			return MessageBuilder.withPayload(result).build();

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

@@ -81,10 +81,12 @@ class TraceMessageHandler {
 
 	private final Function<Span, Span.Builder> outputMessageSpanFunction;
 
+	private final List<FunctionMessageSpanCustomizer> customizers;
+
 	TraceMessageHandler(Tracer tracer, Propagator propagator, Propagator.Setter<MessageHeaderAccessor> injector,
 			Propagator.Getter<MessageHeaderAccessor> extractor, Function<Span, Span> preSendFunction,
 			TriConsumer<MessageHeaderAccessor, Span, Span> preSendMessageManipulator,
-			Function<Span, Span.Builder> outputMessageSpanFunction) {
+			Function<Span, Span.Builder> outputMessageSpanFunction, List<FunctionMessageSpanCustomizer> customizers) {
 		this.tracer = tracer;
 		this.propagator = propagator;
 		this.injector = injector;
@@ -93,18 +95,20 @@ class TraceMessageHandler {
 		this.preSendFunction = preSendFunction;
 		this.preSendMessageManipulator = preSendMessageManipulator;
 		this.outputMessageSpanFunction = outputMessageSpanFunction;
+		this.customizers = customizers;
 	}
 
 	static TraceMessageHandler forNonSpringIntegration(Tracer tracer, Propagator propagator,
-			Propagator.Setter<MessageHeaderAccessor> injector, Propagator.Getter<MessageHeaderAccessor> extractor) {
-		Function<Span, Span> preSendFunction = span -> tracer.nextSpan(span).name("handle").start();
+			Propagator.Setter<MessageHeaderAccessor> injector, Propagator.Getter<MessageHeaderAccessor> extractor,
+			List<FunctionMessageSpanCustomizer> customizers) {
+		Function<Span, Span> preSendFunction = span -> tracer.nextSpan(span).name("function").start();
 		TriConsumer<MessageHeaderAccessor, Span, Span> preSendMessageManipulator = (headers, parentSpan, childSpan) -> {
 			headers.setHeader("traceHandlerParentSpan", parentSpan);
 			headers.setHeader(Span.class.getName(), childSpan);
 		};
 		Function<Span, Span.Builder> postReceiveFunction = span -> tracer.spanBuilder().setParent(span.context());
 		return new TraceMessageHandler(tracer, propagator, injector, extractor, preSendFunction,
-				preSendMessageManipulator, postReceiveFunction);
+				preSendMessageManipulator, postReceiveFunction, customizers);
 	}
 
 	@SuppressWarnings("unchecked")
@@ -112,7 +116,7 @@ class TraceMessageHandler {
 		Propagator.Setter<MessageHeaderAccessor> setter = firstBeanOrException(beanFactory, Propagator.Setter.class);
 		Propagator.Getter<MessageHeaderAccessor> getter = firstBeanOrException(beanFactory, Propagator.Getter.class);
 		return forNonSpringIntegration(beanFactory.getBean(Tracer.class), beanFactory.getBean(Propagator.class), setter,
-				getter);
+				getter, customizers(beanFactory));
 	}
 
 	private static <T> T firstBeanOrException(BeanFactory beanFactory, Class<T> clazz) {
@@ -125,6 +129,16 @@ class TraceMessageHandler {
 		return object;
 	}
 
+	private static List<FunctionMessageSpanCustomizer> customizers(BeanFactory beanFactory) {
+		List<FunctionMessageSpanCustomizer> customizers = new ArrayList<>();
+		ObjectProvider<FunctionMessageSpanCustomizer> provider = beanFactory
+				.getBeanProvider(FunctionMessageSpanCustomizer.class);
+		for (FunctionMessageSpanCustomizer functionMessageSpanCustomizer : provider) {
+			customizers.add(functionMessageSpanCustomizer);
+		}
+		return customizers;
+	}
+
 	/**
 	 * Wraps the given input message with tracing headers and returns a corresponding
 	 * span.
@@ -134,41 +148,33 @@ class TraceMessageHandler {
 	 */
 	MessageAndSpans wrapInputMessage(Message<?> message, String destinationName) {
 		MessageHeaderAccessor headers = mutableHeaderAccessor(message);
-		Span extracted = this.propagator.extract(headers, this.extractor).start();
-		// Start and finish a consumer span as we will immediately process it.
-		Span.Builder consumerSpanBuilder = this.tracer.spanBuilder().setParent(extracted.context());
-		Span consumerSpan = consumerSpan(destinationName, extracted, consumerSpanBuilder);
-		// create and scope a span for the message processor
-		Span span = this.preSendFunction.apply(consumerSpan);
-		// remove any trace headers, but don't re-inject as we are synchronously
-		// processing the
-		// message and can rely on scoping to access this span later.
-		clearTracingHeaders(headers);
-		this.preSendMessageManipulator.accept(headers, consumerSpan, span);
+		Span.Builder consumerSpanBuilder = this.propagator.extract(headers, this.extractor);
+		Span consumerSpan = consumerSpan(destinationName, consumerSpanBuilder, message);
 		if (log.isDebugEnabled()) {
-			log.debug("Created a handle span after retrieving the message " + consumerSpanBuilder);
+			log.debug("Built a consumer span " + consumerSpan);
 		}
+		Span childSpan = this.preSendFunction.apply(consumerSpan);
+		clearTracingHeaders(headers);
+		this.preSendMessageManipulator.accept(headers, consumerSpan, childSpan);
+		this.customizers.forEach(customizer -> customizer.customizeFunctionSpan(childSpan, message));
 		if (message instanceof ErrorMessage) {
 			return new MessageAndSpans(new ErrorMessage((Throwable) message.getPayload(), headers.getMessageHeaders()),
-					consumerSpan, span);
+					consumerSpan, childSpan);
 		}
 		headers.setImmutable();
 		return new MessageAndSpans(new GenericMessage<>(message.getPayload(), headers.getMessageHeaders()),
-				consumerSpan, span);
+				consumerSpan, childSpan);
 	}
 
-	private Span consumerSpan(String destinationName, Span extracted, Span.Builder consumerSpanBuilder) {
-		Span consumerSpan;
-		if (!extracted.isNoop()) {
-			consumerSpanBuilder.kind(Span.Kind.CONSUMER).start();
-			addTags(consumerSpanBuilder, destinationName);
-			consumerSpanBuilder.remoteServiceName(REMOTE_SERVICE_NAME);
-			consumerSpan = consumerSpanBuilder.start();
-			consumerSpan.end();
-		}
-		else {
-			consumerSpan = consumerSpanBuilder.start();
-		}
+	private Span consumerSpan(String destinationName, Span.Builder consumerSpanBuilder, Message<?> message) {
+		consumerSpanBuilder.kind(Span.Kind.CONSUMER).name("handle");
+		addTags(consumerSpanBuilder, destinationName);
+		consumerSpanBuilder.remoteServiceName(REMOTE_SERVICE_NAME);
+		// this is the consumer part of the producer->consumer mechanism
+		Span consumerSpan = consumerSpanBuilder.start();
+		this.customizers.forEach(customizer -> customizer.customizeInputMessageSpan(consumerSpan, message));
+		// we're ending this immediately just to have a properly nested graph
+		consumerSpan.end();
 		return consumerSpan;
 	}
 
@@ -191,12 +197,6 @@ class TraceMessageHandler {
 		}
 	}
 
-	private void addTags(Span result, String destinationName) {
-		if (StringUtils.hasText(destinationName)) {
-			result.tag("channel", SpanNameUtil.shorten(destinationName));
-		}
-	}
-
 	/**
 	 * Called either when message got received and processed or message got sent.
 	 * @param span - span that corresponds to the given operation
@@ -233,7 +233,7 @@ class TraceMessageHandler {
 		MessageHeaderAccessor headers = mutableHeaderAccessor(retrievedMessage);
 		Span.Builder span = this.outputMessageSpanFunction.apply(parentSpan);
 		clearTracingHeaders(headers);
-		Span producerSpan = createProducerSpan(headers, span, destinationName);
+		Span producerSpan = createProducerSpan(headers, span, destinationName, message);
 		this.propagator.inject(producerSpan.context(), headers, this.injector);
 		if (log.isDebugEnabled()) {
 			log.debug("Created a new span output message " + span);
@@ -241,12 +241,14 @@ class TraceMessageHandler {
 		return new MessageAndSpan(outputMessage(message, retrievedMessage, headers), producerSpan);
 	}
 
-	private Span createProducerSpan(MessageHeaderAccessor headers, Span.Builder spanBuilder, String destinationName) {
+	private Span createProducerSpan(MessageHeaderAccessor headers, Span.Builder spanBuilder, String destinationName,
+			Message<?> message) {
 		spanBuilder.kind(Span.Kind.PRODUCER).name("send").remoteServiceName(toRemoteServiceName(headers));
 		Span span = spanBuilder.start();
 		if (!span.isNoop()) {
 			addTags(spanBuilder, destinationName);
 		}
+		this.customizers.forEach(customizer -> customizer.customizeOutputMessageSpan(span, message));
 		return span;
 	}
 

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

@@ -109,7 +109,6 @@ public final class TracingChannelInterceptor implements ExecutorChannelIntercept
 	public TracingChannelInterceptor(Tracer tracer, Propagator propagator,
 			Propagator.Setter<MessageHeaderAccessor> setter, Propagator.Getter<MessageHeaderAccessor> getter,
 			Function<String, String> remoteServiceNameMapper, MessageSpanCustomizer messageSpanCustomizer) {
-
 		this.tracer = tracer;
 		this.propagator = propagator;
 		this.injector = setter;

tests/common/src/main/java/org/springframework/cloud/sleuth/brave/BraveTestSpanHandler.java

@@ -18,6 +18,7 @@ package org.springframework.cloud.sleuth.brave;
 
 import java.util.Iterator;
 import java.util.List;
+import java.util.Queue;
 import java.util.stream.Collectors;
 
 import brave.test.IntegrationTestSpanHandler;
@@ -27,6 +28,8 @@ import org.springframework.cloud.sleuth.brave.bridge.BraveAccessor;
 import org.springframework.cloud.sleuth.exporter.FinishedSpan;
 import org.springframework.cloud.sleuth.test.TestSpanHandler;
 
+import static org.assertj.core.api.BDDAssertions.then;
+
 public class BraveTestSpanHandler implements TestSpanHandler {
 
 	final brave.test.TestSpanHandler spans;
@@ -81,6 +84,33 @@ public class BraveTestSpanHandler implements TestSpanHandler {
 		return BraveAccessor.finishedSpan(this.spans.get(index));
 	}
 
+	@Override
+	public void assertAllSpansWereFinishedOrAbandoned(Queue<Span> createdSpans) {
+		List<FinishedSpan> finishedSpans = reportedSpans();
+		then(finishedSpans).as("There should be that many finished spans as many created ones")
+				.hasSize(createdSpans.size());
+		// finished -> a,b,c ; created -> b,c,d => matchedFinished = b,c
+		List<FinishedSpan> matchedFinishedSpans = finishedSpans.stream()
+				.filter(f -> createdSpans.stream().anyMatch(cs -> f.getSpanId().equals(cs.context().spanId())))
+				.collect(Collectors.toList());
+		// finished -> a,b,c ; created -> b,c,d => matchedCreated = b,c
+		List<Span> matchedCreatedSpans = createdSpans.stream()
+				.filter(cs -> finishedSpans.stream().anyMatch(f -> cs.context().spanId().equals(f.getSpanId())))
+				.collect(Collectors.toList());
+		// finished -> a,b,c ; created -> b,c,d => missingFinished = a
+		List<FinishedSpan> missingFinishedSpans = finishedSpans.stream()
+				.filter(f -> matchedFinishedSpans.stream().noneMatch(m -> m.getSpanId().equals(f.getSpanId())))
+				.collect(Collectors.toList());
+		// finished -> a,b,c ; created -> b,c,d => missingCreated = d
+		List<Span> missingCreatedSpans = createdSpans.stream().filter(
+				f -> matchedCreatedSpans.stream().noneMatch(m -> m.context().spanId().equals(f.context().spanId())))
+				.collect(Collectors.toList());
+		if (!missingFinishedSpans.isEmpty() || !missingCreatedSpans.isEmpty()) {
+			throw new AssertionError("There were unmatched created spans " + missingCreatedSpans
+					+ " and/or finished span " + missingFinishedSpans);
+		}
+	}
+
 	@Override
 	public Iterator<FinishedSpan> iterator() {
 		return reportedSpans().iterator();

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceFunctionAroundWrapperTests.java

@@ -26,6 +26,8 @@ import org.springframework.boot.builder.SpringApplicationBuilder;
 import org.springframework.cloud.function.context.FunctionCatalog;
 import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
 import org.springframework.cloud.sleuth.test.TestSpanHandler;
+import org.springframework.cloud.sleuth.test.TestTracer;
+import org.springframework.cloud.sleuth.test.TestTracingBeanPostProcessor;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.messaging.Message;
@@ -49,10 +51,13 @@ public abstract class TraceFunctionAroundWrapperTests {
 			FunctionCatalog catalog = context.getBean(FunctionCatalog.class);
 			FunctionInvocationWrapper function = catalog.lookup("greeter");
 			function.setSkipOutputConversion(true);
+
 			Message<?> result = (Message<?>) function.get();
+
 			assertThat(result.getPayload()).isEqualTo("hello");
 			assertThat(spanHandler.reportedSpans().size()).isEqualTo(2);
 			assertThat(((String) result.getHeaders().get("b3"))).contains(spanHandler.get(0).getTraceId());
+			spanHandler.assertAllSpansWereFinishedOrAbandoned(context.getBean(TestTracer.class).createdSpans());
 		}
 	}
 
@@ -66,10 +71,13 @@ public abstract class TraceFunctionAroundWrapperTests {
 			FunctionCatalog catalog = context.getBean(FunctionCatalog.class);
 			FunctionInvocationWrapper function = catalog.lookup("uppercase");
 			function.setSkipOutputConversion(true);
+
 			Message<?> result = (Message<?>) function.apply(MessageBuilder.withPayload("hello").build());
+
 			assertThat(result.getPayload()).isEqualTo("HELLO");
 			assertThat(spanHandler.reportedSpans().size()).isEqualTo(3);
 			assertThat(((String) result.getHeaders().get("b3"))).contains(spanHandler.get(0).getTraceId());
+			spanHandler.assertAllSpansWereFinishedOrAbandoned(context.getBean(TestTracer.class).createdSpans());
 		}
 	}
 
@@ -88,6 +96,11 @@ public abstract class TraceFunctionAroundWrapperTests {
 			return v -> v.toUpperCase();
 		}
 
+		@Bean
+		static TestTracingBeanPostProcessor testTracerBeanPostProcessor() {
+			return new TestTracingBeanPostProcessor();
+		}
+
 	}
 
 };

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2013-2021 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
+ *
+ *      https://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.test;
+
+import java.util.List;
+
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.TraceContext;
+import org.springframework.cloud.sleuth.propagation.Propagator;
+
+/**
+ * {@link Propagator} that stores information about started spans.
+ */
+public class TestPropagator implements Propagator {
+
+	private final Propagator delegate;
+
+	private final TestTracer testTracer;
+
+	public TestPropagator(Propagator delegate, TestTracer testTracer) {
+		this.delegate = delegate;
+		this.testTracer = testTracer;
+	}
+
+	@Override
+	public List<String> fields() {
+		return this.delegate.fields();
+	}
+
+	@Override
+	public <C> void inject(TraceContext context, C carrier, Setter<C> setter) {
+		this.delegate.inject(context, carrier, setter);
+	}
+
+	@Override
+	public <C> Span.Builder extract(C carrier, Getter<C> getter) {
+		return new TestSpanBuilder(this.delegate.extract(carrier, getter), this.testTracer);
+	}
+
+}

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

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2013-2021 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
+ *
+ *      https://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.test;
+
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.TraceContext;
+
+class TestSpanBuilder implements Span.Builder {
+
+	private final Span.Builder delegate;
+
+	private final TestTracer testTracer;
+
+	TestSpanBuilder(Span.Builder delegate, TestTracer testTracer) {
+		this.delegate = delegate;
+		this.testTracer = testTracer;
+	}
+
+	@Override
+	public Span.Builder setParent(TraceContext context) {
+		delegate.setParent(context);
+		return this;
+	}
+
+	@Override
+	public Span.Builder setNoParent() {
+		delegate.setNoParent();
+		return this;
+	}
+
+	@Override
+	public Span.Builder name(String name) {
+		delegate.name(name);
+		return this;
+	}
+
+	@Override
+	public Span.Builder event(String value) {
+		delegate.event(value);
+		return this;
+	}
+
+	@Override
+	public Span.Builder tag(String key, String value) {
+		delegate.tag(key, value);
+		return this;
+	}
+
+	@Override
+	public Span.Builder error(Throwable throwable) {
+		delegate.error(throwable);
+		return this;
+	}
+
+	@Override
+	public Span.Builder kind(Span.Kind spanKind) {
+		delegate.kind(spanKind);
+		return this;
+	}
+
+	@Override
+	public Span.Builder remoteServiceName(String remoteServiceName) {
+		delegate.remoteServiceName(remoteServiceName);
+		return this;
+	}
+
+	@Override
+	public Span start() {
+		Span span = delegate.start();
+		this.testTracer.createdSpans.add(span);
+		return span;
+	}
+
+}

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

@@ -17,6 +17,7 @@
 package org.springframework.cloud.sleuth.test;
 
 import java.util.List;
+import java.util.Queue;
 
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.exporter.FinishedSpan;
@@ -35,4 +36,6 @@ public interface TestSpanHandler extends Iterable<FinishedSpan> {
 
 	FinishedSpan get(int index);
 
+	void assertAllSpansWereFinishedOrAbandoned(Queue<Span> createdSpans);
+
 }

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

@@ -0,0 +1,116 @@
+/*
+ * Copyright 2013-2021 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
+ *
+ *      https://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.test;
+
+import java.util.LinkedList;
+import java.util.Map;
+import java.util.Queue;
+
+import org.springframework.cloud.sleuth.BaggageInScope;
+import org.springframework.cloud.sleuth.ScopedSpan;
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.SpanCustomizer;
+import org.springframework.cloud.sleuth.TraceContext;
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.lang.Nullable;
+
+public class TestTracer implements Tracer, AutoCloseable {
+
+	private final Tracer delegate;
+
+	final Queue<Span> createdSpans = new LinkedList<>();
+
+	public TestTracer(Tracer delegate) {
+		this.delegate = delegate;
+	}
+
+	@Override
+	public Map<String, String> getAllBaggage() {
+		return delegate.getAllBaggage();
+	}
+
+	@Override
+	public BaggageInScope getBaggage(String name) {
+		return delegate.getBaggage(name);
+	}
+
+	@Override
+	public BaggageInScope getBaggage(TraceContext traceContext, String name) {
+		return delegate.getBaggage(traceContext, name);
+	}
+
+	@Override
+	public BaggageInScope createBaggage(String name) {
+		return delegate.createBaggage(name);
+	}
+
+	@Override
+	public BaggageInScope createBaggage(String name, String value) {
+		return delegate.createBaggage(name, value);
+	}
+
+	@Override
+	public Span nextSpan() {
+		Span span = delegate.nextSpan();
+		this.createdSpans.add(span);
+		return span;
+	}
+
+	@Override
+	public Span nextSpan(Span parent) {
+		Span span = delegate.nextSpan(parent);
+		this.createdSpans.add(span);
+		return span;
+	}
+
+	@Override
+	public SpanInScope withSpan(Span span) {
+		return delegate.withSpan(span);
+	}
+
+	@Override
+	public ScopedSpan startScopedSpan(String name) {
+		return delegate.startScopedSpan(name);
+	}
+
+	@Override
+	public Span.Builder spanBuilder() {
+		return new TestSpanBuilder(delegate.spanBuilder(), this);
+	}
+
+	@Override
+	@Nullable
+	public SpanCustomizer currentSpanCustomizer() {
+		return delegate.currentSpanCustomizer();
+	}
+
+	@Override
+	@Nullable
+	public Span currentSpan() {
+		return delegate.currentSpan();
+	}
+
+	@Override
+	public void close() throws Exception {
+		this.createdSpans.clear();
+	}
+
+	public Queue<Span> createdSpans() {
+		return createdSpans;
+	}
+
+}

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2013-2021 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
+ *
+ *      https://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.test;
+
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.propagation.Propagator;
+
+/**
+ * Wraps all tracing related components into test representations. That way additional
+ * assertions can take place.
+ */
+public class TestTracingBeanPostProcessor implements BeanPostProcessor {
+
+	TestTracer testTracer;
+
+	@Override
+	public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
+		if (bean instanceof Tracer && !(bean instanceof TestTracer)) {
+			this.testTracer = new TestTracer((Tracer) bean);
+			return this.testTracer;
+		}
+		else if (bean instanceof Propagator && !(bean instanceof TestPropagator)) {
+			return new TestPropagator((Propagator) bean, this.testTracer);
+		}
+		return bean;
+	}
+
+}