diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java index 70c75c177..c02dbccdd 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptor.java @@ -25,8 +25,10 @@ import brave.propagation.TraceContext; import brave.propagation.TraceContextOrSamplingFlags; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; +import org.springframework.aop.support.AopUtils; import org.springframework.cloud.sleuth.util.SpanNameUtil; import org.springframework.integration.channel.AbstractMessageChannel; +import org.springframework.integration.channel.DirectChannel; import org.springframework.integration.context.IntegrationObjectSupport; import org.springframework.messaging.Message; import org.springframework.messaging.MessageChannel; @@ -113,11 +115,23 @@ public final class TracingChannelInterceptor extends ChannelInterceptorAdapter log.debug("Created a new span in pre send" + span); } headers.setImmutable(); - return new GenericMessage<>(message.getPayload(), headers.getMessageHeaders()); + Message outputMessage = new GenericMessage<>(message.getPayload(), headers.getMessageHeaders()); + if (isDirectChannel(channel)) { + beforeHandle(outputMessage, channel, null); + } + return outputMessage; + } + + private boolean isDirectChannel(MessageChannel channel) { + return DirectChannel.class + .isAssignableFrom(AopUtils.getTargetClass(channel)); } @Override public void afterSendCompletion(Message message, MessageChannel channel, boolean sent, Exception ex) { + if (isDirectChannel(channel)) { + afterMessageHandled(message, channel, null, ex); + } if (log.isDebugEnabled()) { log.debug("Will finish the current span after completion " + this.tracer.currentSpan()); } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptorTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptorTest.java index 7af1c87db..5ba281dc1 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptorTest.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TracingChannelInterceptorTest.java @@ -23,6 +23,8 @@ import java.util.Map; import brave.Tracing; import brave.propagation.StrictCurrentTraceContext; +import org.springframework.integration.channel.DirectChannel; +import org.springframework.messaging.MessagingException; import zipkin2.Span; import org.junit.After; import org.junit.Test; @@ -48,6 +50,14 @@ public class TracingChannelInterceptorTest { .build()); QueueChannel channel = new QueueChannel(); + DirectChannel directChannel = new DirectChannel(); + Message message; + MessageHandler handler = new MessageHandler() { + @Override + public void handleMessage(Message msg) throws MessagingException { + message = msg; + } + }; @Test public void injectsProducerSpan() { channel.addInterceptor(producerSideOnly(interceptor)); @@ -61,6 +71,19 @@ public class TracingChannelInterceptorTest { .containsExactly(Span.Kind.PRODUCER); } + @Test public void injectsProducerAndConsumerSpan() { + directChannel.addInterceptor(interceptor); + directChannel.subscribe(this.handler); + directChannel.send(MessageBuilder.withPayload("foo").build()); + + assertThat(message).isNotNull(); + assertThat(message.getHeaders()) + .containsKeys("X-B3-TraceId", "X-B3-SpanId", "X-B3-Sampled", + "nativeHeaders"); + assertThat(spans).flatExtracting(Span::kind) + .contains(Span.Kind.CONSUMER, Span.Kind.PRODUCER); + } + @Test public void injectsProducerSpan_nativeHeaders() { channel.addInterceptor(producerSideOnly(interceptor)); diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/multiple/MultipleHopsIntegrationTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/multiple/MultipleHopsIntegrationTests.java index 95ebaee6a..2dafd11f9 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/multiple/MultipleHopsIntegrationTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/multiple/MultipleHopsIntegrationTests.java @@ -79,10 +79,14 @@ public class MultipleHopsIntegrationTests { this.restTemplate.getForObject("http://localhost:" + this.config.port + "/greeting", String.class); await().atMost(5, SECONDS).untilAsserted(() -> { - then(this.reporter.getSpans()).hasSize(5); + then(this.reporter.getSpans()).hasSize(13); }); then(this.reporter.getSpans().stream().map(zipkin2.Span::name) .collect(toList())).containsAll(asList("http:/greeting", "send")); + then(this.reporter.getSpans().stream().map(zipkin2.Span::kind) + // no server kind due to test constraints + .collect(toList())).containsAll(asList(zipkin2.Span.Kind.CONSUMER, + zipkin2.Span.Kind.PRODUCER, zipkin2.Span.Kind.SERVER)); then(this.reporter.getSpans().stream() .map(span -> span.tags().get("channel")) .filter(Objects::nonNull) diff --git a/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java b/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java index 352270700..b966a104b 100644 --- a/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java +++ b/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java @@ -128,7 +128,7 @@ public class MessagingApplicationTests extends AbstractIntegrationTest { Optional lastHttpSpansParent = findLastHttpSpansParent(); // "http:/parent/" -> "message:messages" -> "http:/foo" (CS + CR) -> "http:/foo" (SS) thenAllSpansArePresent(firstHttpSpan, eventSpans, lastHttpSpansParent, eventSentSpan, producerSpan); - then(this.integrationTestSpanCollector.hashedSpans).as("There were 3 spans").hasSize(3); + then(this.integrationTestSpanCollector.hashedSpans).as("There were 5 spans").hasSize(5); log.info("Checking the parent child structure"); List> parentChild = this.integrationTestSpanCollector.hashedSpans.stream() .filter(span -> span.parentId() != null)