Fixed closing of Zuul span, added local component span for zuul

fixes #242

Author: marcingrzejszczak

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/zuul/TracePostZuulFilter.java

@@ -17,7 +17,7 @@
 package org.springframework.cloud.sleuth.instrument.zuul;
 
 import org.springframework.cloud.sleuth.Span;
-import org.springframework.cloud.sleuth.SpanAccessor;
+import org.springframework.cloud.sleuth.Tracer;
 
 import com.netflix.zuul.ZuulFilter;
 
@@ -29,10 +29,10 @@
  */
 public class TracePostZuulFilter extends ZuulFilter {
 
-	private final SpanAccessor spanAccessor;
+	private final Tracer tracer;
 
-	public TracePostZuulFilter(SpanAccessor spanAccessor) {
-		this.spanAccessor = spanAccessor;
+	public TracePostZuulFilter(Tracer tracer) {
+		this.tracer = tracer;
 	}
 
 	@Override
@@ -44,6 +44,7 @@ public boolean shouldFilter() {
 	public Object run() {
 		// TODO: the client sent event should come from the client not the filter!
 		getCurrentSpan().logEvent(Span.CLIENT_RECV);
+		this.tracer.close(getCurrentSpan());
 		return null;
 	}
 
@@ -58,6 +59,6 @@ public int filterOrder() {
 	}
 
 	private Span getCurrentSpan() {
-		return this.spanAccessor.getCurrentSpan();
+		return this.tracer.getCurrentSpan();
 	}
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/zuul/TracePreZuulFilter.java

@@ -20,7 +20,9 @@
 import org.springframework.cloud.sleuth.SpanInjector;
 import org.springframework.cloud.sleuth.Tracer;
 
+import com.netflix.zuul.ExecutionStatus;
 import com.netflix.zuul.ZuulFilter;
+import com.netflix.zuul.ZuulFilterResult;
 import com.netflix.zuul.context.RequestContext;
 
 /**
@@ -32,6 +34,8 @@
  */
 public class TracePreZuulFilter extends ZuulFilter {
 
+	private static final String ZUUL_COMPONENT = "zuul";
+
 	private final Tracer tracer;
 	private final SpanInjector<RequestContext> spanInjector;
 
@@ -47,12 +51,22 @@ public boolean shouldFilter() {
 
 	@Override
 	public Object run() {
+		getCurrentSpan().logEvent(Span.CLIENT_SEND);
+		return null;
+	}
+
+	@Override
+	public ZuulFilterResult runFilter() {
 		RequestContext ctx = RequestContext.getCurrentContext();
 		Span span = getCurrentSpan();
-		this.spanInjector.inject(span, ctx);
-		// TODO: the client sent event should come from the client not the filter!
-		span.logEvent(Span.CLIENT_SEND);
-		return null;
+		Span newSpan = this.tracer.createSpan(span.getName(), span);
+		newSpan.tag(Span.SPAN_LOCAL_COMPONENT_TAG_NAME, ZUUL_COMPONENT);
+		this.spanInjector.inject(newSpan, ctx);
+		ZuulFilterResult result = super.runFilter();
+		if (ExecutionStatus.SUCCESS != result.getStatus()) {
+			this.tracer.close(newSpan);
+		}
+		return result;
 	}
 
 	private Span getCurrentSpan() {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/zuul/TraceZuulAutoConfiguration.java

@@ -22,7 +22,6 @@
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
 import org.springframework.cloud.netflix.ribbon.SpringClientFactory;
-import org.springframework.cloud.sleuth.SpanAccessor;
 import org.springframework.cloud.sleuth.SpanInjector;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration;
@@ -63,8 +62,8 @@ public TraceRestClientRibbonCommandFactory traceRestClientRibbonCommandFactory(S
 
 	@Bean
 	@ConditionalOnMissingBean
-	public TracePostZuulFilter tracePostZuulFilter(SpanAccessor accessor) {
-		return new TracePostZuulFilter(accessor);
+	public TracePostZuulFilter tracePostZuulFilter(Tracer tracer) {
+		return new TracePostZuulFilter(tracer);
 	}
 
 	@Bean

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/zuul/TracePostZuulFilterTests.java

@@ -52,10 +52,11 @@ public void clean() {
 	}
 
 	@Test
-	public void filterPublishesEvent() throws Exception {
+	public void filterPublishesEventAndClosesSpan() throws Exception {
 		Span span = this.tracer.createSpan("http:start");
 		this.filter.run();
 
 		then(span).hasLoggedAnEvent(Span.CLIENT_RECV);
+		then(this.tracer.getCurrentSpan()).isNull();
 	}
 }

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/zuul/TracePreZuulFilterTests.java

@@ -17,6 +17,7 @@
 package org.springframework.cloud.sleuth.instrument.zuul;
 
 import java.util.Random;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.After;
 import org.junit.Before;
@@ -31,6 +32,7 @@
 import org.springframework.cloud.sleuth.trace.TestSpanContextHolder;
 
 import com.netflix.zuul.context.RequestContext;
+import com.netflix.zuul.monitoring.MonitoringHelper;
 
 import static org.assertj.core.api.BDDAssertions.then;
 
@@ -45,6 +47,11 @@ public class TracePreZuulFilterTests {
 
 	private TracePreZuulFilter filter = new TracePreZuulFilter(this.tracer, new RequestContextInjector());
 
+	@Before
+	public void setup() {
+		MonitoringHelper.initMocks();
+	}
+
 	@After
 	@Before
 	public void clean() {
@@ -56,7 +63,7 @@ public void clean() {
 	public void filterAddsHeaders() throws Exception {
 		this.tracer.createSpan("http:start");
 
-		this.filter.run();
+		this.filter.runFilter();
 
 		RequestContext ctx = RequestContext.getCurrentContext();
 		then(ctx.getZuulRequestHeaders().get(Span.TRACE_ID_NAME))
@@ -69,7 +76,7 @@ public void filterAddsHeaders() throws Exception {
 	public void notSampledIfNotExportable() throws Exception {
 		this.tracer.createSpan("http:start", NeverSampler.INSTANCE);
 
-		this.filter.run();
+		this.filter.runFilter();
 
 		RequestContext ctx = RequestContext.getCurrentContext();
 		then(ctx.getZuulRequestHeaders().get(Span.TRACE_ID_NAME))
@@ -78,4 +85,43 @@ public void notSampledIfNotExportable() throws Exception {
 				.isEqualTo(Span.SPAN_NOT_SAMPLED);
 	}
 
+
+	@Test
+	public void shouldCloseSpanWhenExceptionIsThrown() throws Exception {
+		Span startedSpan = this.tracer.createSpan("http:start", NeverSampler.INSTANCE);
+		final AtomicReference<Span> span = new AtomicReference<>();
+
+		new TracePreZuulFilter(this.tracer, new RequestContextInjector()) {
+			@Override
+			public Object run() {
+				super.run();
+				span.set(TracePreZuulFilterTests.this.tracer.getCurrentSpan());
+				throw new RuntimeException();
+			}
+		}.runFilter();
+
+		then(startedSpan).isNotEqualTo(span.get());
+		then(span.get().logs()).extracting("event").contains(Span.CLIENT_SEND);
+		then(this.tracer.getCurrentSpan()).isEqualTo(startedSpan);
+	}
+
+	@Test
+	public void shouldNotCloseSpanWhenNoExceptionIsThrown() throws Exception {
+		Span startedSpan = this.tracer.createSpan("http:start", NeverSampler.INSTANCE);
+		final AtomicReference<Span> span = new AtomicReference<>();
+
+		new TracePreZuulFilter(this.tracer, new RequestContextInjector()) {
+			@Override
+			public Object run() {
+				span.set(TracePreZuulFilterTests.this.tracer.getCurrentSpan());
+				return super.run();
+			}
+		}.runFilter();
+
+		then(startedSpan).isNotEqualTo(span.get());
+		then(span.get().logs()).extracting("event").contains(Span.CLIENT_SEND);
+		then(span.get().tags()).containsKey(Span.SPAN_LOCAL_COMPONENT_TAG_NAME);
+		then(this.tracer.getCurrentSpan()).isEqualTo(span.get());
+	}
+
 }