move trace intercepting from aop to spring mvc interceptors

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

@@ -0,0 +1,45 @@
+package org.springframework.cloud.sleuth.instrument.web;
+
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.cloud.sleuth.TraceScope;
+import org.springframework.web.servlet.HandlerInterceptor;
+import org.springframework.web.servlet.ModelAndView;
+
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+/**
+ * @author Spencer Gibb
+ */
+public class TraceHandlerInterceptor implements HandlerInterceptor {
+
+	private static final String ATTR_NAME = "__CURRENT_TRACE_HANDLER_TRACE_SCOPE_ATTR___";
+
+	private final Trace trace;
+
+	public TraceHandlerInterceptor(Trace trace) {
+		this.trace = trace;
+	}
+
+	@Override
+	public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
+		//TODO: get trace data from request?
+		//TODO: what is the description?
+		TraceScope scope = trace.startSpan("");
+		request.setAttribute(ATTR_NAME, scope);
+		return true;
+	}
+
+	@Override
+	public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
+
+	}
+
+	@Override
+	public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
+		TraceScope scope = TraceScope.class.cast(request.getAttribute(ATTR_NAME));
+		if (scope != null) {
+			scope.close();
+		}
+	}
+}

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

@@ -1,115 +0,0 @@
-package org.springframework.cloud.sleuth.instrument.web;
-
-import static org.springframework.cloud.sleuth.Trace.SPAN_ID_NAME;
-import static org.springframework.cloud.sleuth.Trace.TRACE_ID_NAME;
-
-import lombok.extern.apachecommons.CommonsLog;
-
-import org.aspectj.lang.ProceedingJoinPoint;
-import org.aspectj.lang.annotation.Around;
-import org.aspectj.lang.annotation.Pointcut;
-import org.springframework.cloud.sleuth.Span;
-import org.springframework.cloud.sleuth.Trace;
-import org.springframework.cloud.sleuth.TraceScope;
-import org.springframework.http.HttpEntity;
-import org.springframework.http.HttpHeaders;
-import org.springframework.stereotype.Controller;
-import org.springframework.web.bind.annotation.RestController;
-import org.springframework.web.client.RestOperations;
-
-/**
- * Aspect that adds trace information to
- * <p/>
- * <ul>
- * <li>{@link RestController} annotated classes</li>
- * <li>{@link Controller} annotated classes</li>
- * <li>explicit {@link RestOperations}.exchange(..) method calls</li>
- * </ul>
- * <p/>
- * For controllers an around aspect is created that create a {@link Span} for each call.
- * <p/>
- * For {@link RestOperations} we are wrapping all executions of the <b>exchange</b>
- * methods and we are extracting {@link HttpHeaders} from the passed {@link HttpEntity}.
- * Next we are adding span id header * {@link Trace#SPAN_ID_NAME} with the value taken
- * from the current Span. Finally the method execution proceeds.
- *
- * @see RestController
- * @see Controller
- * @see RestOperations
- * @see Trace
- *
- * @author Tomasz Nurkewicz, 4financeIT
- * @author Marcin Grzejszczak, 4financeIT
- * @author Michal Chmielarz, 4financeIT
- * @author Spencer Gibb
- */
-@CommonsLog
-public class TraceWebAspect {
-
-	private final Trace trace;
-
-	public TraceWebAspect(Trace trace) {
-		this.trace = trace;
-	}
-
-	@Pointcut("@target(org.springframework.web.bind.annotation.RestController)")
-	private void anyRestControllerAnnotated() {
-	}
-
-	@Pointcut("@target(org.springframework.stereotype.Controller)")
-	private void anyControllerAnnotated() {
-	}
-
-	@Pointcut("anyRestControllerAnnotated() || anyControllerAnnotated()")
-	private void anyControllerOrRestController() {
-	}
-
-	@Around("anyControllerOrRestController()")
-	public Object wrapWithCorrelationId(ProceedingJoinPoint pjp) throws Throwable {
-		//TODO: get trace data from request?
-		TraceScope scope = trace.startSpan(pjp.toShortString());
-		try {
-			return pjp.proceed();
-		}
-		finally {
-			scope.close();
-		}
-	}
-
-	@Pointcut("execution(public * org.springframework.web.client.RestOperations.exchange(..))")
-	private void anyExchangeRestOperationsMethod() {
-	}
-
-	@Around("anyExchangeRestOperationsMethod()")
-	public Object traceRestOperations(ProceedingJoinPoint pjp)
-			throws Throwable {
-		TraceScope scope = trace.startSpan(pjp.toShortString());
-		try {
-			String spanId = scope.getSpan().getSpanId();
-			String traceId = scope.getSpan().getTraceId();
-			log.debug("Wrapping RestTemplate call with trace id [" + traceId + "] and span id [" + spanId + "]");
-			HttpEntity httpEntity = findHttpEntity(pjp.getArgs());
-			if (httpEntity != null) {
-				httpEntity.getHeaders().set(SPAN_ID_NAME, spanId);
-				httpEntity.getHeaders().set(TRACE_ID_NAME, traceId);
-			}
-			return pjp.proceed();
-		}
-		finally {
-			scope.close();
-		}
-	}
-
-	protected HttpEntity findHttpEntity(Object[] args) {
-		for (Object arg : args) {
-			if (isHttpEntity(arg)) {
-				return HttpEntity.class.cast(arg);
-			}
-		}
-		return null;
-	}
-
-	protected boolean isHttpEntity(Object arg) {
-		return HttpEntity.class.isAssignableFrom(arg.getClass());
-	}
-}

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

@@ -17,10 +17,8 @@ package org.springframework.cloud.sleuth.instrument.web;
 
 import java.util.regex.Pattern;
 
-import org.aspectj.lang.ProceedingJoinPoint;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
-import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
@@ -30,6 +28,8 @@ import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.EnableAspectJAutoProxy;
 import org.springframework.util.StringUtils;
+import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
+import org.springframework.web.servlet.config.annotation.WebMvcConfigurerAdapter;
 
 /**
  * Registers beans that add tracing to requests
@@ -43,7 +43,6 @@ import org.springframework.util.StringUtils;
 @EnableAspectJAutoProxy
 @ConditionalOnProperty(value = "spring.cloud.sleuth.trace.web.enabled", matchIfMissing = true)
 @ConditionalOnWebApplication
-@ConditionalOnClass(ProceedingJoinPoint.class)
 public class TraceWebAutoConfiguration {
 
 	/**
@@ -57,8 +56,14 @@ public class TraceWebAutoConfiguration {
 
 	@Bean
 	@ConditionalOnMissingBean
-	public TraceWebAspect traceWebAspect() {
+	public TraceHandlerInterceptor traceHandlerInterceptor() {
-		return new TraceWebAspect(trace);
+		return new TraceHandlerInterceptor(trace);
+	}
+
+	@Bean
+	public WebMvcConfigurerAdapter webMvcConfigurerAdapter(
+			TraceHandlerInterceptor handlerInterceptor) {
+		return new TraceWebConfigurer(handlerInterceptor);
 	}
 
 	@Bean
@@ -68,4 +73,17 @@ public class TraceWebAutoConfiguration {
 				: TraceFilter.DEFAULT_SKIP_PATTERN;
 		return new FilterRegistrationBean(new TraceFilter(trace, pattern));
 	}
+
+	protected static class TraceWebConfigurer extends WebMvcConfigurerAdapter {
+		private TraceHandlerInterceptor interceptor;
+
+		public TraceWebConfigurer(TraceHandlerInterceptor interceptor) {
+			this.interceptor = interceptor;
+		}
+
+		@Override
+		public void addInterceptors(InterceptorRegistry registry) {
+			registry.addInterceptor(interceptor).addPathPatterns("/**");
+		}
+	}
 }

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

@@ -48,9 +48,9 @@ public class TraceRestTemplateInterceptor implements ClientHttpRequestIntercepto
 		return execution.execute(request, body);
 	}
 
-	public void setHeader(HttpRequest request, String spanIdName, String spanId) {
+	public void setHeader(HttpRequest request, String name, String value) {
-		if (!request.getHeaders().containsKey(spanIdName) && isTracing()) {
+		if (!request.getHeaders().containsKey(name) && isTracing()) {
-			request.getHeaders().add(spanIdName, spanId);
+			request.getHeaders().add(name, value);
 		}
 	}
 

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

@@ -1,6 +1,7 @@
 package org.springframework.cloud.sleuth.slf4j;
 
 import static org.springframework.cloud.sleuth.Trace.SPAN_ID_NAME;
+import static org.springframework.cloud.sleuth.Trace.TRACE_ID_NAME;
 
 import lombok.extern.slf4j.Slf4j;
 import org.slf4j.MDC;
@@ -18,8 +19,9 @@ public class Slf4jSpanReceiver implements SpanReceiver {
 	@Override
 	public void receiveSpan(Span span) {
 		//TODO: what should this log level be?
-		log.info("Received span {}", span);
+		log.info("Received span: {}", span);
 		MDC.remove(SPAN_ID_NAME);
+		MDC.remove(TRACE_ID_NAME);
 	}
 
 	@PostConstruct