From b0533f7a60bcf535a51699601c4e3baa7cbeb457 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin@grzejszczak.pl>
Date: Tue, 20 Dec 2016 18:13:57 +0100
Subject: [PATCH]     Adding parent id to MDC

    without this change we couldn't reference parent id in the logs
    with this change we add the parent id to MDC context

    fixes #480
---
 .../cloud/sleuth/log/Slf4jSpanLogger.java     | 10 +++++++
 .../cloud/sleuth/log/Slf4JSpanLoggerTest.java | 26 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jSpanLogger.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jSpanLogger.java
index 08468ffbf..ecc1e41bd 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jSpanLogger.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/log/Slf4jSpanLogger.java
@@ -51,6 +51,7 @@ public class Slf4jSpanLogger implements SpanLogger {
 		log("Starting span: {}", span);
 		if (parent != null) {
 			log("With parent: {}", parent);
+			MDC.put(Span.PARENT_ID_NAME, Span.idToHex(parent.getSpanId()));
 		}
 	}
 
@@ -59,9 +60,16 @@ public class Slf4jSpanLogger implements SpanLogger {
 		MDC.put(Span.SPAN_ID_NAME, Span.idToHex(span.getSpanId()));
 		MDC.put(Span.TRACE_ID_NAME, span.traceIdString());
 		MDC.put(Span.SPAN_EXPORT_NAME, String.valueOf(span.isExportable()));
+		setParentIdIfPresent(span);
 		log("Continued span: {}", span);
 	}
 
+	private void setParentIdIfPresent(Span span) {
+		if (!span.getParents().isEmpty()) {
+			MDC.put(Span.PARENT_ID_NAME, Span.idToHex(span.getParents().get(0)));
+		}
+	}
+
 	@Override
 	public void logStoppedSpan(Span parent, Span span) {
 		if (span != null) {
@@ -71,11 +79,13 @@ public class Slf4jSpanLogger implements SpanLogger {
 			log("With parent: {}", parent);
 			MDC.put(Span.SPAN_ID_NAME, Span.idToHex(parent.getSpanId()));
 			MDC.put(Span.SPAN_EXPORT_NAME, String.valueOf(parent.isExportable()));
+			setParentIdIfPresent(parent);
 		}
 		else {
 			MDC.remove(Span.SPAN_ID_NAME);
 			MDC.remove(Span.SPAN_EXPORT_NAME);
 			MDC.remove(Span.TRACE_ID_NAME);
+			MDC.remove(Span.PARENT_ID_NAME);
 		}
 	}
 
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/log/Slf4JSpanLoggerTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/log/Slf4JSpanLoggerTest.java
index 649ea1a9d..76c84d70d 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/log/Slf4JSpanLoggerTest.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/log/Slf4JSpanLoggerTest.java
@@ -37,7 +37,7 @@ import static org.mockito.Mockito.times;
 public class Slf4JSpanLoggerTest {
 
 	Span spanWithNameToBeExcluded = Span.builder().name("Hystrix").build();
-	Span spanWithNameNotToBeExcluded = Span.builder().name("Aspect").build();
+	Span spanWithNameNotToBeExcluded = Span.builder().name("Aspect").parent(3L).build();
 	String nameExcludingPattern = "^.*Hystrix.*$";
 	Logger log = Mockito.mock(Logger.class);
 	Slf4jSpanLogger slf4JSpanLogger = new Slf4jSpanLogger(this.nameExcludingPattern, this.log);
@@ -152,4 +152,28 @@ public class Slf4JSpanLoggerTest {
 
 		then(this.log).should().trace(anyString(), anyList());
 	}
+
+	@Test
+	public void should_set_mdc_entry_for_parent_when_starting() throws Exception {
+		this.slf4JSpanLogger.logStartedSpan(this.spanWithNameNotToBeExcluded,
+				this.spanWithNameNotToBeExcluded);
+
+		assertThat(MDC.get(Span.PARENT_ID_NAME)).isEqualTo(
+				Span.idToHex(this.spanWithNameNotToBeExcluded.getSpanId()));
+	}
+
+	@Test
+	public void should_set_mdc_entry_for_parent_when_continuing() throws Exception {
+  		this.slf4JSpanLogger.logContinuedSpan(this.spanWithNameNotToBeExcluded);
+
+		assertThat(MDC.get(Span.PARENT_ID_NAME)).isEqualTo(Span.idToHex(3L));
+	}
+
+	@Test
+	public void should_set_mdc_entry_for_parent_when_stopping() throws Exception {
+		this.slf4JSpanLogger.logStoppedSpan(this.spanWithNameNotToBeExcluded,
+				this.spanWithNameNotToBeExcluded);
+
+		assertThat(MDC.get(Span.PARENT_ID_NAME)).isEqualTo(Span.idToHex(3L));
+	}
 }
\ No newline at end of file