From 3c2c43adff3581016bf0345722a7c5930dbe98e1 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <mgrzejszczak@pivotal.io>
Date: Mon, 23 Oct 2017 19:39:08 +0200
Subject: [PATCH] Fixed the invalid closing of spans; fixes #759

---
 .../sleuth/instrument/web/TraceFilter.java    | 11 +++++---
 .../web/TraceFilterWebIntegrationTests.java   | 28 +++++++++++++++++--
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java
index c9b0350c2..43e89f536 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceFilter.java
@@ -280,10 +280,9 @@ public class TraceFilter extends GenericFilterBean {
 					log.debug(
 							"Won't detach the span " + span + " since error has already been handled");
 				}
-			}  else if (shouldCloseSpan(request) && tracer().isTracing() && stillTracingCurrentSapn(span)) {
+			}  else if ((shouldCloseSpan(request) || isRootSpan(span)) && tracer().isTracing() && stillTracingCurrentSpan(span)) {
 				if (log.isDebugEnabled()) {
-					log.debug(
-							"Will close span " + span + " since some component marked it for closure");
+					log.debug("Will close span " + span + " since " + (shouldCloseSpan(request) ? "some component marked it for closure" : "response was unsuccessful for the root span"));
 				}
 				tracer().close(span);
 				clearTraceAttribute(request);
@@ -297,7 +296,11 @@ public class TraceFilter extends GenericFilterBean {
 		}
 	}
 
-	private boolean stillTracingCurrentSapn(Span span) {
+	private boolean isRootSpan(Span span) {
+		return span.getTraceId() == span.getSpanId();
+	}
+
+	private boolean stillTracingCurrentSpan(Span span) {
 		return tracer().getCurrentSpan().equals(span);
 	}
 
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterWebIntegrationTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterWebIntegrationTests.java
index 92b8740cd..5e676bb30 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterWebIntegrationTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterWebIntegrationTests.java
@@ -16,6 +16,9 @@
 
 package org.springframework.cloud.sleuth.instrument.web;
 
+import static org.assertj.core.api.Assertions.fail;
+import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then;
+
 import java.io.IOException;
 
 import org.junit.After;
@@ -36,15 +39,16 @@ import org.springframework.cloud.sleuth.util.ExceptionUtils;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.env.Environment;
+import org.springframework.http.ResponseEntity;
 import org.springframework.http.client.ClientHttpResponse;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestMethod;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.client.DefaultResponseErrorHandler;
+import org.springframework.web.client.HttpClientErrorException;
 import org.springframework.web.client.RestTemplate;
 
-import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then;
-
 /**
  * @author Marcin Grzejszczak
  */
@@ -63,6 +67,7 @@ public class TraceFilterWebIntegrationTests {
 	public void cleanup() {
 		ExceptionUtils.setFail(true);
 		TestSpanContextHolder.removeCurrentSpan();
+		this.accumulator.clear();
 	}
 
 	@Test
@@ -80,6 +85,20 @@ public class TraceFilterWebIntegrationTests {
 				.hasRpcTagsInProperOrder();
 	}
 
+	@Test
+	public void should_create_spans_for_endpoint_returning_unsuccessful_result() {
+		try {
+			new RestTemplate().getForObject("http://localhost:" + port() + "/test_bad_request", String.class);
+			fail("should throw exception");
+		} catch (HttpClientErrorException e) {
+		}
+
+		then(this.tracer.getCurrentSpan()).isNull();
+		then(ExceptionUtils.getLastException()).isNull();
+		then(new ListOfSpans(this.accumulator.getSpans()))
+				.hasServerSideSpansInProperOrder();
+	}
+
 	private int port() {
 		return this.environment.getProperty("local.server.port", Integer.class);
 	}
@@ -119,5 +138,10 @@ public class TraceFilterWebIntegrationTests {
 		public void throwException() {
 			throw new RuntimeException("Throwing exception");
 		}
+
+		@RequestMapping(path = "/test_bad_request", method = RequestMethod.GET)
+		public ResponseEntity<?> processFail() {
+			return ResponseEntity.badRequest().build();
+		}
 	}
 }