Fixed tests for TraceWebAspect, added WebTask aspect

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

@@ -16,6 +16,7 @@
 
 package org.springframework.cloud.sleuth.instrument.web;
 
+import java.lang.reflect.Field;
 import java.util.concurrent.Callable;
 
 import org.aspectj.lang.ProceedingJoinPoint;
@@ -28,6 +29,7 @@ import org.springframework.cloud.sleuth.instrument.TraceCallable;
 import org.springframework.stereotype.Controller;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.client.RestOperations;
+import org.springframework.web.context.request.async.WebAsyncTask;
 
 import lombok.extern.apachecommons.CommonsLog;
 
@@ -37,6 +39,7 @@ import lombok.extern.apachecommons.CommonsLog;
  * <ul>
  * <li>{@link RestController} annotated classes with public {@link Callable} methods</li>
  * <li>{@link Controller} annotated classes with public {@link Callable} methods</li>
+ * <li>{@link Controller} or {@link RestController} annotated classes with public {@link WebAsyncTask} methods</li>
  * </ul>
  * <p/>
  * For controllers an around aspect is created that wraps the {@link Callable#call()}
@@ -82,6 +85,14 @@ public class TraceWebAspect {
 	private void anyControllerOrRestControllerWithPublicAsyncMethod() {
 	}
 
+	@Pointcut("execution(public org.springframework.web.context.request.async.WebAsyncTask *(..))")
+	private void anyPublicMethodReturningWebAsyncTask() {
+	}
+
+	@Pointcut("(anyRestControllerAnnotated() || anyControllerAnnotated()) && anyPublicMethodReturningWebAsyncTask()")
+	private void anyControllerOrRestControllerWithPublicWebAsyncTaskMethod() {
+	}
+
 	@Around("anyControllerOrRestControllerWithPublicAsyncMethod()")
 	@SuppressWarnings("unchecked")
 	public Object wrapWithCorrelationId(ProceedingJoinPoint pjp) throws Throwable {
@@ -96,4 +107,22 @@ public class TraceWebAspect {
 		}
 	}
 
+	@SuppressWarnings("unchecked")
+	@Around("anyControllerOrRestControllerWithPublicWebAsyncTaskMethod()")
+	public Object wrapWebAsyncTaskWithCorrelationId(ProceedingJoinPoint pjp) throws Throwable {
+		final WebAsyncTask webAsyncTask = (WebAsyncTask) pjp.proceed();
+		if (this.accessor.isTracing()) {
+			try {
+				log.debug("Wrapping callable with span ["
+						+ this.accessor.getCurrentSpan() + "]");
+				Field callableField = WebAsyncTask.class.getDeclaredField("callable");
+				callableField.setAccessible(true);
+				callableField.set(webAsyncTask, new TraceCallable<>(this.traceManager, webAsyncTask.getCallable()));
+			} catch (NoSuchFieldException ex) {
+				log.warn("Cannot wrap webAsyncTask's callable with TraceCallable", ex);
+			}
+		}
+		return webAsyncTask;
+	}
+
 }

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

@@ -16,7 +16,7 @@
 
 package org.springframework.cloud.sleuth.trace;
 
-import static org.springframework.cloud.sleuth.util.ExceptionUtils.error;
+import static org.springframework.cloud.sleuth.util.ExceptionUtils.warn;
 
 import java.util.concurrent.Callable;
 
@@ -59,7 +59,7 @@ public class DefaultTraceManager implements TraceManager {
 		}
 		Span currentSpan = getCurrentSpan();
 		if (currentSpan != null && !parent.equals(currentSpan)) {
-			error("Trace client error: thread " + Thread.currentThread().getName()
+			warn("Trace client warn: thread " + Thread.currentThread().getName()
 					+ " tried to start a new Span " + "with parent " + parent.toString()
 					+ ", but there is already a " + "currentSpan " + currentSpan);
 		}
@@ -88,7 +88,7 @@ public class DefaultTraceManager implements TraceManager {
 		Span cur = TraceContextHolder.getCurrentSpan();
 		Span span = trace.getSpan();
 		if (cur != span) {
-			ExceptionUtils.error("Tried to detach trace span but "
+			ExceptionUtils.warn("Tried to detach trace span but "
 					+ "it is not the current span for the '"
 					+ Thread.currentThread().getName() + "' thread: " + span
 					+ ". You have " + "probably forgotten to close or detach " + cur);
@@ -113,7 +113,7 @@ public class DefaultTraceManager implements TraceManager {
 		Span span = trace.getSpan();
 		Trace savedTrace = trace.getSavedTrace();
 		if (cur != span) {
-			ExceptionUtils.error("Tried to close trace span but "
+			ExceptionUtils.warn("Tried to close trace span but "
 					+ "it is not the current span for the '"
 					+ Thread.currentThread().getName() + "' thread" + span
 					+ ".  You have " + "probably forgotten to close or detach " + cur);

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

@@ -23,8 +23,7 @@ import lombok.extern.apachecommons.CommonsLog;
  */
 @CommonsLog
 public abstract class ExceptionUtils {
-	public static void error(String msg) {
-		log.error(msg);
-		throw new RuntimeException(msg);
+	public static void warn(String msg) {
+		log.warn(msg);
 	}
 }

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

@@ -16,7 +16,6 @@ import java.util.concurrent.Callable;
 
 import org.junit.Before;
 import org.junit.ClassRule;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -43,7 +42,6 @@ import junitparams.JUnitParamsRunner;
 import junitparams.Parameters;
 
 @SpringApplicationConfiguration(classes = {RestTemplateTraceAspectITest.CorrelationIdAspectSpecConfiguration.class})
-@Ignore("Will fail due to not setting initial values for Trace and Span IDs")
 @RunWith(JUnitParamsRunner.class)
 public class RestTemplateTraceAspectITest extends MvcWiremockITest {
 
@@ -78,7 +76,7 @@ public class RestTemplateTraceAspectITest extends MvcWiremockITest {
 	}
 
 	private void thenTraceIdHasBeenSetOnARequestHeader() {
-		this.wireMock.verifyThat(getRequestedFor(urlMatching(".*")).withHeader(Trace.TRACE_ID_NAME, matching("^(?!\\s*$).+/))")));
+		this.wireMock.verifyThat(getRequestedFor(urlMatching(".*")).withHeader(Trace.TRACE_ID_NAME, matching("^(?!\\s*$).+")));
 	}
 
 	private void whenARequestIsSentToAnAsyncEndpoint(String url) throws Exception {
@@ -119,7 +117,7 @@ public class RestTemplateTraceAspectITest extends MvcWiremockITest {
 		}
 
 		@RequestMapping(value = "/webAsyncTaskPing", method = RequestMethod.GET, produces = MediaType.TEXT_PLAIN_VALUE)
-		WebAsyncTask<String> webAsyncTaskPing() {
+		public WebAsyncTask<String> webAsyncTaskPing() {
 			return new WebAsyncTask<>(new Callable<String>() {
 				@Override
 				public String call() throws Exception {