[#132] Fixed closing of spans in messaging

fixes #132

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/AbstractTraceChannelInterceptor.java

@@ -18,6 +18,8 @@ import org.springframework.messaging.support.ChannelInterceptorAdapter;
  */
 abstract class AbstractTraceChannelInterceptor extends ChannelInterceptorAdapter {
 
+	protected static final String MESSAGE_NAME_PREFIX = "message/";
+
 	private final Tracer tracer;
 
 	private final Random random;
@@ -96,7 +98,7 @@ abstract class AbstractTraceChannelInterceptor extends ChannelInterceptorAdapter
 	}
 
 	String getMessageChannelName(MessageChannel channel) {
-		return "message/" + getChannelName(channel);
+		return MESSAGE_NAME_PREFIX + getChannelName(channel);
 	}
 
 }

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

@@ -22,6 +22,7 @@ import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.instrument.TraceKeys;
 import org.springframework.cloud.sleuth.sampler.IsTracingSampler;
+import org.springframework.cloud.sleuth.trace.SpanContextHolder;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 
@@ -31,7 +32,7 @@ import org.springframework.messaging.MessageChannel;
  */
 public class TraceChannelInterceptor extends AbstractTraceChannelInterceptor {
 
-	private ThreadLocal<Span> traceHolder = new ThreadLocal<>();
+	private final ThreadLocal<Span> traceHolder = new ThreadLocal<>();
 
 	public TraceChannelInterceptor(Tracer tracer, TraceKeys traceKeys, Random random) {
 		super(tracer, traceKeys, random);
@@ -39,9 +40,20 @@ public class TraceChannelInterceptor extends AbstractTraceChannelInterceptor {
 
 	@Override
 	public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
-		Span trace = this.traceHolder.get();
-		getTracer().close(trace);
+		Span trace = closeAllParentMessageRelatedSpans(this.traceHolder.get());
 		this.traceHolder.remove();
+		if (trace != null) {
+			SpanContextHolder.setCurrentSpan(trace);
+		}
+	}
+
+	private Span closeAllParentMessageRelatedSpans(Span trace) {
+		Span traceToClose = trace;
+		while (traceToClose != null && traceToClose.getName() != null &&
+				traceToClose.getName().startsWith(MESSAGE_NAME_PREFIX)) {
+			traceToClose = getTracer().close(traceToClose);
+		}
+		return traceToClose;
 	}
 
 	@Override

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/multiple/DemoApplication.java

@@ -0,0 +1,75 @@
+package org.springframework.cloud.sleuth.instrument.web.multiple;
+
+import java.util.Arrays;
+import java.util.List;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.integration.annotation.Aggregator;
+import org.springframework.integration.annotation.Gateway;
+import org.springframework.integration.annotation.IntegrationComponentScan;
+import org.springframework.integration.annotation.MessageEndpoint;
+import org.springframework.integration.annotation.MessagingGateway;
+import org.springframework.integration.annotation.ServiceActivator;
+import org.springframework.integration.annotation.Splitter;
+import org.springframework.util.StringUtils;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestParam;
+import org.springframework.web.bind.annotation.RestController;
+
+@RestController
+@MessageEndpoint
+@IntegrationComponentScan
+public class DemoApplication {
+
+	@Autowired
+	Sender sender;
+
+	@RequestMapping("/greeting")
+	public Greeting greeting(@RequestParam(defaultValue="Hello World!") String message) {
+		this.sender.send(message);
+		return new Greeting(message);
+	}
+
+	@Splitter(inputChannel="greetings", outputChannel="words")
+	public List<String> words(String greeting) {
+		return Arrays.asList(StringUtils.delimitedListToStringArray(greeting, " "));
+	}
+
+	@Aggregator(inputChannel="words", outputChannel="counts")
+	public int count(List<String> greeting) {
+		return greeting.size();
+	}
+
+	@ServiceActivator(inputChannel="counts")
+	public void report(int count) {
+		System.err.println("Count: " + count);
+	}
+
+}
+
+@MessagingGateway(name = "greeter")
+interface Sender {
+	@Gateway(requestChannel = "greetings")
+	void send(String message);
+}
+
+class Greeting {
+	private String message;
+
+	Greeting() {
+	}
+
+	public Greeting(String message) {
+		super();
+		this.message = message;
+	}
+
+	public String getMessage() {
+		return this.message;
+	}
+
+	public void setMessage(String message) {
+		this.message = message;
+	}
+
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/multiple/MultipleHopsIntegrationTests.java

@@ -0,0 +1,68 @@
+package org.springframework.cloud.sleuth.instrument.web.multiple;
+
+import java.util.Arrays;
+import java.util.stream.Collectors;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.cloud.sleuth.Sampler;
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.event.ArrayListSpanAccumulator;
+import org.springframework.cloud.sleuth.instrument.TraceKeys;
+import org.springframework.cloud.sleuth.instrument.web.TraceFilter;
+import org.springframework.cloud.sleuth.instrument.web.common.AbstractMvcIntegrationTest;
+import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.test.web.servlet.result.MockMvcResultMatchers;
+import org.springframework.test.web.servlet.setup.DefaultMockMvcBuilder;
+
+import static com.jayway.awaitility.Awaitility.await;
+import static java.util.Arrays.asList;
+import static java.util.stream.Collectors.toList;
+import static org.assertj.core.api.BDDAssertions.then;
+import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
+
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration(classes = MultipleHopsIntegrationTests.Config.class)
+public class MultipleHopsIntegrationTests extends AbstractMvcIntegrationTest {
+
+	@Autowired Tracer tracer;
+	@Autowired TraceKeys traceKeys;
+	@Autowired ArrayListSpanAccumulator arrayListSpanAccumulator;
+
+	@Override
+	protected void configureMockMvcBuilder(DefaultMockMvcBuilder mockMvcBuilder) {
+		mockMvcBuilder.addFilters(new TraceFilter(this.tracer, this.traceKeys));
+	}
+
+	@Test
+	public void should_prepare_spans_for_export() throws Exception {
+		this.mockMvc.perform(get("/greeting")).andExpect(
+				MockMvcResultMatchers.status().isOk());
+
+		await().until(() -> {
+			then(this.arrayListSpanAccumulator.getSpans().stream().map(Span::getName).collect(
+					toList())).containsAll(asList("http/greeting", "message/greetings",
+													"message/words", "message/counts"));
+		});
+	}
+
+	@Configuration
+	@SpringBootApplication
+	public static class Config {
+
+		@Bean ArrayListSpanAccumulator arrayListSpanAccumulator() {
+			return new ArrayListSpanAccumulator();
+		}
+
+		@Bean Sampler defaultTraceSampler() {
+			return new AlwaysSampler();
+		}
+	}
+}