Merge branch 'master' of https://github.com/michaelhyatt/elastic-apm-…

…mule3-agent
michaelhyatt · Jul 8, 2020 · 161c858 · 161c858
2 parents 07385ef + 0c202e6
commit 161c858
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 8 deletions.
diff --git a/apm-mule3-agent/src/main/java/co/elastic/apm/mule/listeners/DomainApplicationListener.java b/apm-mule3-agent/src/main/java/co/elastic/apm/mule/listeners/DomainApplicationListener.java
@@ -64,13 +64,10 @@ public void onNotification(PipelineMessageNotification notification) {
 					transactionUtils.startTransactionIfNone(notification);
 					break;
 
+				case PipelineMessageNotification.PROCESS_COMPLETE:
 				case PipelineMessageNotification.PROCESS_END:
 					transactionUtils.endTransactionIfNeeded(notification);
 					break;
-
-				case PipelineMessageNotification.PROCESS_COMPLETE:
-					// Ignored, as it is skipped when flow exception is thrown
-					break;
 				}
 			}
 		});

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/src/main/java/co/elastic/apm/mule/listeners/DomainApplicationListener.java b/apm-mule3.8-agent/src/main/java/co/elastic/apm/mule/listeners/DomainApplicationListener.java
@@ -64,13 +64,10 @@ public void onNotification(PipelineMessageNotification notification) {
 					transactionUtils.startTransactionIfNone(notification);
 					break;
 
+				case PipelineMessageNotification.PROCESS_COMPLETE:					
 				case PipelineMessageNotification.PROCESS_END:
 					transactionUtils.endTransactionIfNeeded(notification);
 					break;
-
-				case PipelineMessageNotification.PROCESS_COMPLETE:
-					// Ignored, as it is skipped when flow exception is thrown
-					break;
 				}
 			}
 		});

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();
 	}