Polish async support

An @Async method should start a new trace. Still not finished that bit
because if the user also customizes the async post processor there will
be a clash.

Biggest change here is better support for async web requests.
Each async request goes through the TraceFilter multiple times so you
have to re-attach to the span in the request if there is one.

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/MilliSpan.java

@@ -24,6 +24,8 @@ public class MilliSpan implements Span {
 	@Singular
 	private List<String> parents;
 	private final String spanId;
+	@NonFinal
+	private boolean remote = false;
 	private Map<String, String> kVAnnotations = new LinkedHashMap<>();
 	private final String processId;
 	@Singular
@@ -63,7 +65,8 @@ public class MilliSpan implements Span {
 
 	@Override
 	public void addTimelineAnnotation(String msg) {
-		this.timelineAnnotations.add(new TimelineAnnotation(System.currentTimeMillis(), msg));
+		this.timelineAnnotations.add(new TimelineAnnotation(System.currentTimeMillis(),
+				msg));
 	}
 
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/Span.java

@@ -4,48 +4,56 @@ import java.util.List;
 import java.util.Map;
 
 /**
- * Base interface for gathering and reporting statistics about a block of
- * execution.
+ * Base interface for gathering and reporting statistics about a block of execution.
+ * <p/>
+ * Spans should form a directed acyclic graph structure. It should be possible to keep
+ * following the parents of a span until you arrive at a span with no parents.
  * <p/>
- * Spans should form a directed acyclic graph structure.  It should be possible
- * to keep following the parents of a span until you arrive at a span with no
- * parents.<p/>
  */
 public interface Span {
 
 	/**
-	 * A human-readable name assigned to this span instance.<p/>
+	 * A human-readable name assigned to this span instance.
+	 * <p/>
 	 */
 	String getName();
 
 	/**
-	 * A pseudo-unique (random) number assigned to this span instance.<p/>
+	 * A pseudo-unique (random) number assigned to this span instance.
 	 * <p/>
-	 * The spanId is immutable and cannot be changed.  It is safe to access this
-	 * from multiple threads.
+	 * <p/>
+	 * The spanId is immutable and cannot be changed. It is safe to access this from
+	 * multiple threads.
 	 */
 	String getSpanId();
 
 	/**
-	 * A pseudo-unique (random) number assigned to the trace associated with this
-	 * span
+	 * A pseudo-unique (random) number assigned to the trace associated with this span
 	 */
 	String getTraceId();
 
 	/**
-	 * Return a unique id for the process from which this Span originated.<p/>
+	 * Return a unique id for the process from which this Span originated.
+	 * <p/>
 	 * <p/>
 	 * Will never be null.
 	 */
 	String getProcessId();
 
 	/**
-	 * Returns the parent IDs of the span.<p/>
+	 * Returns the parent IDs of the span.
+	 * <p/>
 	 * <p/>
 	 * The collection will be empty if there are no parents.
 	 */
 	List<String> getParents();
 
+	/**
+	 * Flag that tells us whether the span was started in another process. Useful in RPC
+	 * tracing when the receiver actually has to add annotations to the senders span.
+	 */
+	boolean isRemote();
+
 	/**
 	 * The block has completed, stop the clock
 	 */
@@ -62,8 +70,8 @@ public interface Span {
 	long getEnd();
 
 	/**
-	 * Return the total amount of time elapsed since start was called, if running,
-	 * or difference between stop and start
+	 * Return the total amount of time elapsed since start was called, if running, or
+	 * difference between stop and start
 	 */
 	long getAccumulatedMillis();
 
@@ -83,14 +91,16 @@ public interface Span {
 	void addTimelineAnnotation(String msg);
 
 	/**
-	 * Get data associated with this span (read only)<p/>
+	 * Get data associated with this span (read only)
+	 * <p/>
 	 * <p/>
 	 * Will never be null.
 	 */
 	Map<String, String> getKVAnnotations();
 
 	/**
-	 * Get any timeline annotations (read only)<p/>
+	 * Get any timeline annotations (read only)
+	 * <p/>
 	 * <p/>
 	 * Will never be null.
 	 */

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceScope.java

@@ -47,17 +47,17 @@ public class TraceScope implements Closeable {
 	 */
 	public Span detach() {
 		if (this.detached) {
-			ExceptionUtils.error("Tried to detach trace span " + this.span + " but "
-					+ "it has already been detached.");
+			ExceptionUtils.error("Tried to detach trace span but "
+					+ "it has already been detached: " + this.span);
 		}
 		this.detached = true;
 
 		Span cur = TraceContextHolder.getCurrentSpan();
 		if (cur != this.span) {
-			ExceptionUtils.error("Tried to detach trace span " + this.span + " but "
-					+ "it is not the current span for the "
-					+ Thread.currentThread().getName() + " thread.  You have "
-					+ "probably forgotten to close or detach " + cur);
+			ExceptionUtils.error("Tried to detach trace span but "
+					+ "it is not the current span for the '"
+					+ Thread.currentThread().getName() + "' thread: " + this.span
+					+ ". You have " + "probably forgotten to close or detach " + cur);
 		}
 		else {
 			TraceContextHolder.setCurrentSpan(this.savedSpan);
@@ -74,15 +74,17 @@ public class TraceScope implements Closeable {
 		this.detached = true;
 		Span cur = TraceContextHolder.getCurrentSpan();
 		if (cur != this.span) {
-			ExceptionUtils.error("Tried to close trace span " + this.span + " but "
-					+ "it is not the current span for the "
-					+ Thread.currentThread().getName() + " thread.  You have "
-					+ "probably forgotten to close or detach " + cur);
+			ExceptionUtils.error("Tried to close trace span but "
+					+ "it is not the current span for the '"
+					+ Thread.currentThread().getName() + "' thread" + this.span
+					+ ".  You have " + "probably forgotten to close or detach " + cur);
 		}
 		else {
 			this.span.stop();
-			if (this.savedSpan != null && this.span.getParents().contains(this.savedSpan.getSpanId())) {
-				this.publisher.publishEvent(new SpanStoppedEvent(this, this.savedSpan, this.span));
+			if (this.savedSpan != null
+					&& this.span.getParents().contains(this.savedSpan.getSpanId())) {
+				this.publisher.publishEvent(new SpanStoppedEvent(this, this.savedSpan,
+						this.span));
 			}
 			else {
 				this.publisher.publishEvent(new SpanStoppedEvent(this, this.span));

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/TraceExecutor.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2015 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
+ *
+ *      http://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.scheduling;
+
+import java.util.concurrent.Executor;
+
+import lombok.RequiredArgsConstructor;
+
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.cloud.sleuth.instrument.TraceRunnable;
+
+/**
+ * @author Dave Syer
+ *
+ */
+@RequiredArgsConstructor
+public class TraceExecutor implements Executor {
+
+	private final Trace trace;
+	private final Executor delegate;
+
+	@Override
+	public void execute(Runnable command) {
+		this.delegate.execute(new TraceRunnable(this.trace, command));
+	}
+
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/TraceSchedulingAspect.java

@@ -28,8 +28,8 @@ public class TraceSchedulingAspect {
 	}
 
 	@Around("execution (@org.springframework.scheduling.annotation.Scheduled  * *.*(..))")
-	public Object traceSceduledThread(final ProceedingJoinPoint pjp) throws Throwable {
-		TraceScope scope = trace.startSpan(pjp.toShortString());
+	public Object traceBackgroundThread(final ProceedingJoinPoint pjp) throws Throwable {
+		TraceScope scope = this.trace.startSpan(pjp.toShortString());
 		try {
 			return pjp.proceed();
 		} finally {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/TraceSchedulingAutoConfiguration.java

@@ -4,13 +4,18 @@ package org.springframework.cloud.sleuth.instrument.scheduling;
  * @author Spencer Gibb
  */
 
+import java.util.concurrent.Executor;
+
 import org.aspectj.lang.ProceedingJoinPoint;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.cloud.sleuth.Trace;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.EnableAspectJAutoProxy;
-import org.springframework.scheduling.annotation.EnableScheduling;
+import org.springframework.core.task.SimpleAsyncTaskExecutor;
+import org.springframework.scheduling.annotation.AsyncConfigurerSupport;
+import org.springframework.scheduling.annotation.EnableAsync;
 
 /**
  * Registers beans related to task scheduling.
@@ -21,13 +26,28 @@ import org.springframework.scheduling.annotation.EnableScheduling;
  * @author Spencer Gibb
  */
 @Configuration
-@EnableScheduling
 @EnableAspectJAutoProxy
-@ConditionalOnClass(ProceedingJoinPoint.class)
 public class TraceSchedulingAutoConfiguration {
 
+	@ConditionalOnClass(ProceedingJoinPoint.class)
 	@Bean
 	public TraceSchedulingAspect traceSchedulingAspect(Trace trace) {
 		return new TraceSchedulingAspect(trace);
 	}
+
+	@EnableAsync
+	@Configuration
+	protected static class AsyncConfiguration extends AsyncConfigurerSupport {
+
+		@Autowired
+		private Trace trace;
+
+		// TODO: look for an existing AsyncConfigurer and steal its Executor
+		@Override
+		public Executor getAsyncExecutor() {
+			return new TraceExecutor(this.trace, new SimpleAsyncTaskExecutor());
+		}
+
+	}
+
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java

@@ -54,6 +54,9 @@ import org.springframework.web.util.UrlPathHelper;
 @Order(Ordered.HIGHEST_PRECEDENCE + 5)
 public class TraceFilter extends OncePerRequestFilter {
 
+	private static final String TRACE_REQUEST_ATTR = TraceFilter.class.getName()
+			+ ".TRACE";
+
 	public static final Pattern DEFAULT_SKIP_PATTERN = Pattern
 			.compile("/api-docs.*|/autoconfig|/configprops|/dump|/info|/metrics.*|/mappings|/trace|/swagger.*|.*\\.png|.*\\.css|.*\\.js|.*\\.html|/favicon.ico");
 
@@ -79,29 +82,35 @@ public class TraceFilter extends OncePerRequestFilter {
 		String uri = this.urlPathHelper.getPathWithinApplication(request);
 		boolean skip = this.skipPattern.matcher(uri).matches();
 
-		TraceScope traceScope = null;
-		if (!skip) {
+		TraceScope traceScope = (TraceScope) request.getAttribute(TRACE_REQUEST_ATTR);
+		if (traceScope != null) {
+			this.trace.continueSpan(traceScope.getSpan());
+		}
+		else if (!skip) {
 			String spanId = getHeader(request, response, SPAN_ID_NAME);
 			String traceId = getHeader(request, response, TRACE_ID_NAME);
 			String name = "http" + uri;
 			if (hasText(spanId) && hasText(traceId)) {
 
-				MilliSpanBuilder traceInfo = MilliSpan.builder().traceId(traceId).spanId(spanId);
+				MilliSpanBuilder traceInfo = MilliSpan.builder().traceId(traceId)
+						.spanId(spanId);
 				String parentId = getHeader(request, response, PARENT_ID_NAME);
 				String processId = getHeader(request, response, PROCESS_ID_NAME);
 				String parentName = getHeader(request, response, SPAN_NAME_NAME);
-				if (parentName!=null) {
+				if (parentName != null) {
 					traceInfo.name(parentName);
 				}
-				if (processId!=null) {
+				if (processId != null) {
 					traceInfo.processId(processId);
 				}
-				if (parentId!=null) {
+				if (parentId != null) {
 					traceInfo.parent(parentId);
 				}
+				traceInfo.remote(true);
 
 				// TODO: trace description?
 				traceScope = this.trace.startSpan(name, traceInfo.build());
+				request.setAttribute(TRACE_REQUEST_ATTR, traceScope);
 				// Send new span id back
 				addToResponseIfNotPresent(response, SPAN_ID_NAME, traceScope.getSpan()
 						.getSpanId());
@@ -115,6 +124,9 @@ public class TraceFilter extends OncePerRequestFilter {
 			filterChain.doFilter(request, response);
 		}
 		finally {
+			if (request.isAsyncSupported() && request.isAsyncStarted()) {
+				return;
+			}
 			if (traceScope != null) {
 				traceScope.close();
 			}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceHttpResponse.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2015 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
+ *
+ *      http://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.web.client;
+
+import java.io.IOException;
+import java.io.InputStream;
+
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.client.ClientHttpResponse;
+
+/**
+ * @author Dave Syer
+ *
+ */
+public class TraceHttpResponse implements ClientHttpResponse {
+
+	private final ClientHttpResponse delegate;
+	private final TraceRestTemplateInterceptor interceptor;
+
+	public TraceHttpResponse(TraceRestTemplateInterceptor interceptor,
+			ClientHttpResponse delegate) {
+		this.interceptor = interceptor;
+		this.delegate = delegate;
+	}
+
+	@Override
+	public HttpHeaders getHeaders() {
+		return this.delegate.getHeaders();
+	}
+
+	@Override
+	public InputStream getBody() throws IOException {
+		return this.delegate.getBody();
+	}
+
+	@Override
+	public HttpStatus getStatusCode() throws IOException {
+		return this.delegate.getStatusCode();
+	}
+
+	@Override
+	public int getRawStatusCode() throws IOException {
+		return this.delegate.getRawStatusCode();
+	}
+
+	@Override
+	public String getStatusText() throws IOException {
+		return this.delegate.getStatusText();
+	}
+
+	@Override
+	public void close() {
+		try {
+			this.delegate.close();
+		}
+		finally {
+			this.interceptor.close();
+		}
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java

@@ -72,11 +72,11 @@ ApplicationEventPublisherAware {
 			setHeader(request, PROCESS_ID_NAME, processId);
 		}
 		publish(new ClientSentEvent(this, getCurrentSpan()));
-		try {
-			return execution.execute(request, body);
-		} finally {
-			publish(new ClientReceivedEvent(this, getCurrentSpan()));
-		}
+		return new TraceHttpResponse(this, execution.execute(request, body));
+	}
+
+	public void close() {
+		publish(new ClientReceivedEvent(this, getCurrentSpan()));
 	}
 
 	private void publish(ApplicationEvent event) {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/slf4j/Slf4jSpanListener.java

@@ -28,7 +28,7 @@ public class Slf4jSpanListener {
 		//TODO: what log level?
 		log.info("Starting span: {}", span);
 		if (event.getParent()!=null) {
-			log.info("Starting parent: {}", event.getParent());
+			log.info("With parent: {}", event.getParent());
 		}
 	}
 
@@ -37,7 +37,7 @@ public class Slf4jSpanListener {
 		//TODO: what should this log level be?
 		log.info("Stopped span: {}", event.getSpan());
 		if (event.getParent()!=null) {
-			log.info("Stopped parent: {}", event.getParent());
+			log.info("With parent: {}", event.getParent());
 		}
 		MDC.remove(SPAN_ID_NAME);
 		MDC.remove(TRACE_ID_NAME);

spring-cloud-sleuth-core/src/main/resources/META-INF/spring.factories

@@ -2,5 +2,6 @@
 org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
 org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration,\
 org.springframework.cloud.sleuth.slf4j.SleuthSlf4jAutoConfiguration,\
+org.springframework.cloud.sleuth.instrument.scheduling.TraceSchedulingAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.web.TraceWebAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.web.client.TraceWebClientAutoConfiguration

spring-cloud-sleuth-sample/src/main/java/org/springframework/cloud/sleuth/sample/SampleApplication.java

@@ -1,19 +1,17 @@
 package org.springframework.cloud.sleuth.sample;
 
 import org.springframework.boot.SpringApplication;
-import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
 import org.springframework.cloud.sleuth.Sampler;
 import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
 import org.springframework.context.annotation.Bean;
-import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.EnableAspectJAutoProxy;
 import org.springframework.scheduling.annotation.EnableAsync;
 
 /**
  * @author Spencer Gibb
  */
-@Configuration
-@EnableAutoConfiguration
+@SpringBootApplication
 @EnableAspectJAutoProxy(proxyTargetClass = true)
 @EnableAsync
 public class SampleApplication {
@@ -35,9 +33,9 @@ public class SampleApplication {
 	}
 
 	// Use this for debugging (or if there is no Zipkin collector running on port 9410)
-	// @Bean
-	// public SpanCollector spanCollector() {
-	// return new LoggingSpanCollectorImpl();
-	// }
+	//	@Bean
+	//	public SpanCollector spanCollector() {
+	//		return new LoggingSpanCollectorImpl();
+	//	}
 
 }

spring-cloud-sleuth-sample/src/main/java/org/springframework/cloud/sleuth/sample/SampleBackground.java

@@ -0,0 +1,30 @@
+package org.springframework.cloud.sleuth.sample;
+
+import java.util.Random;
+
+import lombok.SneakyThrows;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.scheduling.annotation.Async;
+import org.springframework.stereotype.Component;
+
+/**
+ * @author Spencer Gibb
+ */
+@Component
+public class SampleBackground {
+
+	@Autowired
+	private Trace trace;
+
+	@SneakyThrows
+	@Async
+	public void background() {
+		final Random random = new Random();
+		int millis = random.nextInt(1000);
+		Thread.sleep(millis);
+		this.trace.addKVAnnotation("background-sleep-millis", String.valueOf(millis));
+	}
+
+}

spring-cloud-sleuth-sample/src/main/java/org/springframework/cloud/sleuth/sample/SampleController.java

@@ -29,6 +29,8 @@ ApplicationListener<EmbeddedServletContainerInitializedEvent> {
 	private RestTemplate restTemplate;
 	@Autowired
 	private Trace trace;
+	@Autowired
+	private SampleBackground controller;
 	private int port;
 
 	@SneakyThrows
@@ -37,8 +39,8 @@ ApplicationListener<EmbeddedServletContainerInitializedEvent> {
 		final Random random = new Random();
 		Thread.sleep(random.nextInt(1000));
 
-		String s = this.restTemplate.getForObject("http://localhost:" + this.port + "/hi2",
-				String.class);
+		String s = this.restTemplate.getForObject("http://localhost:" + this.port
+				+ "/hi2", String.class);
 		return "hi/" + s;
 	}
 
@@ -47,12 +49,21 @@ ApplicationListener<EmbeddedServletContainerInitializedEvent> {
 		return new Callable<String>() {
 			@Override
 			public String call() throws Exception {
+				final Random random = new Random();
+				int millis = random.nextInt(1000);
+				Thread.sleep(millis);
+				SampleController.this.trace.addKVAnnotation("callable-sleep-millis", String.valueOf(millis));
 				Span currentSpan = TraceContextHolder.getCurrentSpan();
-				return "async hi: "+currentSpan;
+				return "async hi: " + currentSpan;
 			}
 		};
 	}
 
+	@RequestMapping("/async")
+	public String async() {
+		this.controller.background();
+		return "ho";
+	}
 
 	@SneakyThrows
 	@RequestMapping("/hi2")
@@ -67,15 +78,32 @@ ApplicationListener<EmbeddedServletContainerInitializedEvent> {
 	@SneakyThrows
 	@RequestMapping("/traced")
 	public String traced() {
-		TraceScope scope = this.trace.startSpan("customTraceEndpoint", new AlwaysSampler(), null);
+		TraceScope scope = this.trace.startSpan("customTraceEndpoint",
+				new AlwaysSampler(), null);
 		final Random random = new Random();
 		int millis = random.nextInt(1000);
 		log.info("Sleeping for {} millis", millis);
 		Thread.sleep(millis);
+		this.trace.addKVAnnotation("random-sleep-millis", String.valueOf(millis));
 
-		String s = this.restTemplate.getForObject("http://localhost:" + this.port + "/hi2", String.class);
+		String s = this.restTemplate.getForObject("http://localhost:" + this.port
+				+ "/call", String.class);
 		scope.close();
-		return "hi/" + s;
+		return "traced/" + s;
+	}
+
+	@SneakyThrows
+	@RequestMapping("/start")
+	public String start() {
+		final Random random = new Random();
+		int millis = random.nextInt(1000);
+		log.info("Sleeping for {} millis", millis);
+		Thread.sleep(millis);
+		this.trace.addKVAnnotation("random-sleep-millis", String.valueOf(millis));
+
+		String s = this.restTemplate.getForObject("http://localhost:" + this.port
+				+ "/call", String.class);
+		return "start/" + s;
 	}
 
 	@Override

spring-cloud-sleuth-zipkin/src/main/java/org/springframework/cloud/sleuth/zipkin/ZipkinSpanListener.java

@@ -48,9 +48,10 @@ public class ZipkinSpanListener {
 
 	@EventListener
 	public void start(SpanStartedEvent event) {
-		if (event.getParent()!=null) {
+		if (event.getParent()!=null && event.getParent().isRemote()) {
 			event.getParent().addTimelineAnnotation(zipkinCoreConstants.SERVER_RECV);
 		}
+		event.getSpan().addTimelineAnnotation("start");
 	}
 
 	@EventListener
@@ -65,10 +66,11 @@ public class ZipkinSpanListener {
 
 	@EventListener
 	public void start(SpanStoppedEvent event) {
-		if (event.getParent()!=null) {
+		if (event.getParent()!=null && event.getParent().isRemote()) {
 			event.getParent().addTimelineAnnotation(zipkinCoreConstants.SERVER_SEND);
 			this.spanCollector.collect(convert(event.getParent()));
 		}
+		event.getSpan().addTimelineAnnotation("stop");
 		this.spanCollector.collect(convert(event.getSpan()));
 	}