Merge pull request #26 from michaelhyatt/MDC_support

Support for log correlation
michaelhyatt · May 13, 2020 · eb191ee · eb191ee
2 parents 9ad90d2 + 979aa67
commit eb191ee
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 6 deletions.
diff --git a/README.md b/README.md
@@ -27,13 +27,13 @@ Go to the [Releases](https://github.com/michaelhyatt/elastic-apm-mule3-agent/rel
 Mule 3.9:
 ```
 mvn install:install-file -Dfile=<path-to-file> -DgroupId=co.elastic.apm \
-    -DartifactId=apm-mule3-agent -Dversion=1.15.0 -Dpackaging=jar
+    -DartifactId=apm-mule3-agent -Dversion=1.15.1 -Dpackaging=jar
 ```
 
 Mule 3.8:
 ```
 mvn install:install-file -Dfile=<path-to-file> -DgroupId=co.elastic.apm \
-    -DartifactId=apm-mule3.8-agent -Dversion=1.15.0 -Dpackaging=jar
+    -DartifactId=apm-mule3.8-agent -Dversion=1.15.1 -Dpackaging=jar
 ```
 
 ### Or, get the code and build it from scratch
@@ -46,7 +46,7 @@ Mule 3.9:
 <dependency>
     <groupId>co.elastic.apm</groupId>
     <artifactId>apm-mule3-agent</artifactId>
-    <version>1.15.0</version>
+    <version>1.15.1</version>
 </dependency>
 ```
 
@@ -55,7 +55,7 @@ Mule 3.8:
 <dependency>
     <groupId>co.elastic.apm</groupId>
     <artifactId>apm-mule3.8-agent</artifactId>
-    <version>1.15.0</version>
+    <version>1.15.1</version>
 </dependency>
 ```
 

diff --git a/apm-mule3-agent/pom.xml b/apm-mule3-agent/pom.xml
@@ -5,7 +5,7 @@
 	<modelVersion>4.0.0</modelVersion>
 	<groupId>co.elastic.apm</groupId>
 	<artifactId>apm-mule3-agent</artifactId>
-	<version>1.15.0</version>
+	<version>1.15.1</version>
 	<packaging>${packaging}</packaging>
 	<name>Mule apm-mule3-agent Application</name>
 

diff --git a/apm-mule3-agent/src/main/java/co/elastic/apm/mule/utils/TransactionUtils.java b/apm-mule3-agent/src/main/java/co/elastic/apm/mule/utils/TransactionUtils.java
@@ -1,6 +1,7 @@
 package co.elastic.apm.mule.utils;
 
 import org.apache.commons.lang3.tuple.ImmutablePair;
+import org.apache.log4j.MDC;
 import org.mule.api.MessagingException;
 import org.mule.api.MuleEvent;
 import org.mule.api.MuleMessage;
@@ -21,6 +22,9 @@
  *
  */
 public class TransactionUtils {
+
+	private static final String MDC_TRANSACTION_ID = "transaction.id";
+	private static final String MDC_TRACE_ID = "trace.id";
 
 	@Autowired
 	private SpanStore txMap;
@@ -56,6 +60,8 @@ public void startTransactionIfNone(PipelineMessageNotification notification) {
 
 		txMap.storeTransactionOrSpan(messageId, notification, transaction);
 
+		if (isLogCorrelationEnabled())
+			addMdcHeaders(transaction);
 	}
 
 	/**
@@ -82,9 +88,27 @@ public void endTransactionIfNeeded(PipelineMessageNotification notification) {
 				transaction.captureException(exceptionThrown);
 
 			transaction.end(notification.getTimestamp() * 1_000);
+
+			if (isLogCorrelationEnabled())
+				removeMdcHeaders();
 		}
 	}
 
+	private boolean isLogCorrelationEnabled() {
+		return "true".equals(System.getProperty("elastic.apm.enable_log_correlation"))
+				|| "true".equals(System.getenv("ELASTIC_APM_ENABLE_LOG_CORRELATION"));
+	}
+
+	private void addMdcHeaders(Transaction transaction) {
+		MDC.put(MDC_TRACE_ID, transaction.getTraceId());
+		MDC.put(MDC_TRANSACTION_ID, transaction.getId());
+	}
+
+	private void removeMdcHeaders() {
+		MDC.remove(MDC_TRACE_ID);
+		MDC.remove(MDC_TRANSACTION_ID);
+	}
+
 	private MuleMessage getMuleMessage(PipelineMessageNotification notification) {
 		return ((MuleEvent) notification.getSource()).getMessage();
 	}

diff --git a/apm-mule3.8-agent/pom.xml b/apm-mule3.8-agent/pom.xml
@@ -5,7 +5,7 @@
 	<modelVersion>4.0.0</modelVersion>
 	<groupId>co.elastic.apm</groupId>
 	<artifactId>apm-mule3.8-agent</artifactId>
-	<version>1.15.0</version>
+	<version>1.15.1</version>
 	<packaging>${packaging}</packaging>
 	<name>Mule apm-mule3-agent Application built for Mule 3.8</name>
 

diff --git a/apm-mule3.8-agent/src/main/java/co/elastic/apm/mule/utils/TransactionUtils.java b/apm-mule3.8-agent/src/main/java/co/elastic/apm/mule/utils/TransactionUtils.java
@@ -1,6 +1,7 @@
 package co.elastic.apm.mule.utils;
 
 import org.apache.commons.lang3.tuple.ImmutablePair;
+import org.apache.log4j.MDC;
 import org.mule.api.MessagingException;
 import org.mule.api.MuleEvent;
 import org.mule.api.MuleMessage;
@@ -22,6 +23,9 @@
  */
 public class TransactionUtils {
 
+	private static final String MDC_TRANSACTION_ID = "transaction.id";
+	private static final String MDC_TRACE_ID = "trace.id";
+
 	@Autowired
 	private SpanStore txMap;
 
@@ -56,6 +60,9 @@ public void startTransactionIfNone(PipelineMessageNotification notification) {
 
 		txMap.storeTransactionOrSpan(messageId, notification, transaction);
 
+		if (isLogCorrelationEnabled())
+			addMdcHeaders(transaction);
+
 	}
 
 	/**
@@ -82,9 +89,27 @@ public void endTransactionIfNeeded(PipelineMessageNotification notification) {
 				transaction.captureException(exceptionThrown);
 
 			transaction.end(notification.getTimestamp() * 1_000);
+
+			if (isLogCorrelationEnabled())
+				removeMdcHeaders();
 		}
 	}
 
+	private boolean isLogCorrelationEnabled() {
+		return "true".equals(System.getProperty("elastic.apm.enable_log_correlation"))
+				|| "true".equals(System.getenv("ELASTIC_APM_ENABLE_LOG_CORRELATION"));
+	}
+
+	private void addMdcHeaders(Transaction transaction) {
+		MDC.put(MDC_TRACE_ID, transaction.getTraceId());
+		MDC.put(MDC_TRANSACTION_ID, transaction.getId());
+	}
+
+	private void removeMdcHeaders() {
+		MDC.remove(MDC_TRACE_ID);
+		MDC.remove(MDC_TRANSACTION_ID);
+	}
+
 	private MuleMessage getMuleMessage(PipelineMessageNotification notification) {
 		return ((MuleEvent) notification.getSource()).getMessage();
 	}