Merge pull request #238 from spring-cloud/fixing_dependencies_messaging

Messages don't log events even though they should

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceChannelInterceptor.java

@@ -16,6 +16,7 @@
 
 package org.springframework.cloud.sleuth.instrument.messaging;
 
+import org.springframework.cloud.sleuth.Log;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanExtractor;
 import org.springframework.cloud.sleuth.SpanInjector;
@@ -37,6 +38,7 @@ import org.springframework.messaging.support.MessageBuilder;
 public class TraceChannelInterceptor extends AbstractTraceChannelInterceptor {
 
 	private static final String SPAN_HEADER = "X-Current-Span";
+	private static final String MESSAGE_SENT_FROM_CLIENT = "X-Message-Sent";
 
 	public TraceChannelInterceptor(Tracer tracer, TraceKeys traceKeys,
 			SpanExtractor<Message<?>> spanExtractor,
@@ -46,7 +48,25 @@ public class TraceChannelInterceptor extends AbstractTraceChannelInterceptor {
 
 	@Override
 	public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
-		getTracer().close(getSpanFromHeader(message));
+		Span spanFromHeader = getSpanFromHeader(message);
+		if (containsServerReceived(spanFromHeader)) {
+			spanFromHeader.logEvent(Span.SERVER_SEND);
+		} else if (spanFromHeader != null) {
+			spanFromHeader.logEvent(Span.CLIENT_RECV);
+		}
+		getTracer().close(spanFromHeader);
+	}
+
+	private boolean containsServerReceived(Span span) {
+		if (span == null) {
+			return false;
+		}
+		for (Log log : span.logs()) {
+			if (Span.SERVER_RECV.equals(log.getEvent())) {
+				return true;
+			}
+		}
+		return false;
 	}
 
 	@Override
@@ -56,6 +76,12 @@ public class TraceChannelInterceptor extends AbstractTraceChannelInterceptor {
 		String name = getMessageChannelName(channel);
 		Span span = startSpan(parentSpan, name, message);
 		MessageBuilder<?> messageBuilder = MessageBuilder.fromMessage(message);
+		if (message.getHeaders().containsKey(MESSAGE_SENT_FROM_CLIENT)) {
+			span.logEvent(Span.SERVER_RECV);
+		} else {
+			span.logEvent(Span.CLIENT_SEND);
+			messageBuilder.setHeader(MESSAGE_SENT_FROM_CLIENT, true);
+		}
 		getSpanInjector().inject(span, messageBuilder);
 		return messageBuilder.build();
 	}
@@ -73,14 +99,22 @@ public class TraceChannelInterceptor extends AbstractTraceChannelInterceptor {
 	@Override
 	public Message<?> beforeHandle(Message<?> message, MessageChannel channel,
 			MessageHandler handler) {
-		getTracer().continueSpan(getSpanFromHeader(message));
+		Span spanFromHeader = getSpanFromHeader(message);
+		if (spanFromHeader!= null) {
+			spanFromHeader.logEvent(Span.SERVER_RECV);
+		}
+		getTracer().continueSpan(spanFromHeader);
 		return message;
 	}
 
 	@Override
 	public void afterMessageHandled(Message<?> message, MessageChannel channel,
 			MessageHandler handler, Exception ex) {
-		getTracer().detach(getSpanFromHeader(message));
+		Span spanFromHeader = getSpanFromHeader(message);
+		if (spanFromHeader!= null) {
+			spanFromHeader.logEvent(Span.SERVER_SEND);
+		}
+		getTracer().detach(spanFromHeader);
 	}
 
 	private Span getSpanFromHeader(Message<?> message) {

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceChannelInterceptorTests.java

@@ -27,10 +27,10 @@ import org.springframework.boot.test.IntegrationTest;
 import org.springframework.boot.test.SpringApplicationConfiguration;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.Tracer;
-import org.springframework.cloud.sleuth.util.ArrayListSpanAccumulator;
 import org.springframework.cloud.sleuth.instrument.messaging.TraceChannelInterceptorTests.App;
 import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
 import org.springframework.cloud.sleuth.trace.TestSpanContextHolder;
+import org.springframework.cloud.sleuth.util.ArrayListSpanAccumulator;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.integration.channel.DirectChannel;
@@ -44,7 +44,6 @@ import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 
 import static org.assertj.core.api.BDDAssertions.then;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 
@@ -98,9 +97,9 @@ public class TraceChannelInterceptorTests implements MessageHandler {
 		assertNotNull("message was null", this.message);
 
 		String spanId = this.message.getHeaders().get(Span.SPAN_ID_NAME, String.class);
-		assertNotNull("spanId was null", spanId);
-		assertNull(TestSpanContextHolder.getCurrentSpan());
-		assertFalse(this.span.isExportable());
+		then(spanId).isNotNull();
+		then(TestSpanContextHolder.getCurrentSpan()).isNull();
+		then(this.span.isExportable()).isFalse();
 	}
 
 	@Test
@@ -132,6 +131,20 @@ public class TraceChannelInterceptorTests implements MessageHandler {
 		assertNull(TestSpanContextHolder.getCurrentSpan());
 	}
 
+	@Test
+	public void shouldLogClientReceivedClientSentEventWhenTheMessageIsSentAndReceived() {
+		this.channel.send(MessageBuilder.withPayload("hi").build());
+
+		then(this.span.logs()).extracting("event").contains(Span.CLIENT_SEND, Span.CLIENT_RECV);
+	}
+
+	@Test
+	public void shouldLogServerReceivedServerSentEventWhenTheMessageIsPropagatedToTheNextListener() {
+		this.channel.send(MessageBuilder.withPayload("hi").setHeader("X-Message-Sent", true).build());
+
+		then(this.span.logs()).extracting("event").contains(Span.SERVER_RECV, Span.SERVER_SEND);
+	}
+
 	@Test
 	public void headerCreation() {
 		Span span = this.tracer.createSpan("http:testSendMessage", new AlwaysSampler());