Add MessageContext to the DataUnit

Author: gayaldassanayake

modules/core/src/main/java/org/apache/synapse/aspects/flow/statistics/collectors/CallbackStatisticCollector.java

@@ -46,6 +46,7 @@ public static void addCallback(MessageContext messageContext, String callbackId)
 			dataUnit.setCallbackId(callbackId);
 			dataUnit.setStatisticId(StatisticDataCollectionHelper.getStatisticTraceId(messageContext));
 			dataUnit.setCurrentIndex(StatisticDataCollectionHelper.getParentFlowPosition(messageContext, null));
+			dataUnit.setMessageContext(messageContext);
 
 			CallbackSentEvent callbackSentEvent = new CallbackSentEvent(dataUnit);
             addEventAndIncrementCallbackCount(messageContext, callbackSentEvent);
@@ -73,6 +74,7 @@ public static void callbackCompletionEvent(MessageContext oldMessageContext, Str
 			dataUnit.setStatisticId(StatisticDataCollectionHelper.getStatisticTraceId(oldMessageContext));
 			dataUnit.setCurrentIndex(StatisticDataCollectionHelper.getParentFlowPosition(oldMessageContext, null));
 			dataUnit.setIsOutOnlyFlow(StatisticDataCollectionHelper.isOutOnlyFlow(oldMessageContext));
+			dataUnit.setMessageContext(oldMessageContext);
 
 			CallbackCompletionEvent callbackCompletionEvent = new CallbackCompletionEvent(dataUnit);
             addEventAndDecrementCallbackCount(oldMessageContext, callbackCompletionEvent);
@@ -99,6 +101,7 @@ public static void updateParentsForCallback(MessageContext oldMessageContext, St
 			dataUnit.setStatisticId(StatisticDataCollectionHelper.getStatisticTraceId(oldMessageContext));
 			dataUnit.setCurrentIndex(StatisticDataCollectionHelper.getParentFlowPosition(oldMessageContext, null));
 			dataUnit.setIsOutOnlyFlow(StatisticDataCollectionHelper.isOutOnlyFlow(oldMessageContext));
+			dataUnit.setMessageContext(oldMessageContext);
 
 			CallbackReceivedEvent callbackReceivedEvent = new CallbackReceivedEvent(dataUnit);
             addEvent(oldMessageContext, callbackReceivedEvent);
@@ -125,6 +128,7 @@ public static void reportCallbackHandlingCompletion(MessageContext synapseOutMsg
 			dataUnit.setSynapseEnvironment(synapseOutMsgCtx.getEnvironment());
 			dataUnit.setStatisticId(StatisticDataCollectionHelper.getStatisticTraceId(synapseOutMsgCtx));
 			dataUnit.setCurrentIndex(StatisticDataCollectionHelper.getParentFlowPosition(synapseOutMsgCtx, null));
+			dataUnit.setMessageContext(synapseOutMsgCtx);
 
 			CallbackHandledEvent callbackHandledEvent = new CallbackHandledEvent(dataUnit);
             addEventAndDecrementCallbackCount(synapseOutMsgCtx, callbackHandledEvent);

modules/core/src/main/java/org/apache/synapse/aspects/flow/statistics/collectors/CloseEventCollector.java

@@ -105,6 +105,7 @@ public static void closeEntryEvent(MessageContext messageContext, String compone
 			}
 			StatisticDataCollectionHelper
 					.collectData(messageContext, isContentAltering, isCollectingTracing, statisticDataUnit);
+			statisticDataUnit.setMessageContext(messageContext);
 
 			StatisticsCloseEvent closeEvent = new StatisticsCloseEvent(statisticDataUnit);
 			if (currentIndex == null) {
@@ -139,6 +140,7 @@ public static void closeFlowForcefully(MessageContext messageContext, boolean er
 			dataUnit.setSynapseEnvironment(messageContext.getEnvironment());
 			dataUnit.setStatisticId(StatisticDataCollectionHelper.getStatisticTraceId(messageContext));
 			dataUnit.setCurrentIndex(StatisticDataCollectionHelper.getParentFlowPosition(messageContext, null));
+			dataUnit.setMessageContext(messageContext);
 
             EndFlowEvent endFlowEvent = new EndFlowEvent(dataUnit);
             if (!error) {

modules/core/src/main/java/org/apache/synapse/aspects/flow/statistics/collectors/FaultStatisticCollector.java

@@ -45,6 +45,7 @@ public static void reportFault(MessageContext messageContext) {
 				BasicStatisticDataUnit dataUnit = new BasicStatisticDataUnit();
 				dataUnit.setStatisticId(StatisticDataCollectionHelper.getStatisticTraceId(messageContext));
 				dataUnit.setCurrentIndex(StatisticDataCollectionHelper.getParentFlowPosition(messageContext, null));
+				dataUnit.setMessageContext(messageContext);
 
 				FaultEvent faultEvent = new FaultEvent(dataUnit);
 				faultEvent.getDataUnit().generateElasticMetadata(messageContext);

modules/core/src/main/java/org/apache/synapse/aspects/flow/statistics/collectors/OpenEventCollector.java

@@ -54,21 +54,15 @@ public class OpenEventCollector extends RuntimeStatisticCollector {
 	 */
 	public static Integer reportEntryEvent(MessageContext messageContext, String componentName,
 	                                       AspectConfiguration aspectConfiguration, ComponentType componentType) {
-		boolean isCollectingStatistics = (aspectConfiguration != null && aspectConfiguration.isStatisticsEnable());
-
 		// Enable statistics, if user enabled for all artifacts
-		if (!isCollectingStatistics) {
-			isCollectingStatistics = isCollectingStatistics || RuntimeStatisticCollector.isCollectingAllStatistics();
-		}
+		boolean isCollectingStatistics = (aspectConfiguration != null && aspectConfiguration.isStatisticsEnable())
+				|| RuntimeStatisticCollector.isCollectingAllStatistics();
 
-		boolean isCollectingTracing = false;
-		if (isCollectingProperties() || isCollectingPayloads() || isCollectingVariables()) {
-			isCollectingTracing = (aspectConfiguration != null && aspectConfiguration.isTracingEnabled());
-		}
+        boolean isCollectingTracing = (isCollectingProperties() || isCollectingPayloads() || isCollectingVariables())
+				&& aspectConfiguration != null && aspectConfiguration.isTracingEnabled();
 
 		Boolean isFlowStatisticEnabled =
 				(Boolean) messageContext.getProperty(StatisticsConstants.FLOW_STATISTICS_IS_COLLECTED);//todo try to use single object for "FLOW_TRACE_IS_COLLECTED"
-		Boolean isTracingEnabled;
 		if (isCollectingStatistics) {
 			messageContext.setProperty(StatisticsConstants.FLOW_STATISTICS_IS_COLLECTED, true);
 			setStatisticsTraceId(messageContext);
@@ -80,7 +74,7 @@ public static Integer reportEntryEvent(MessageContext messageContext, String com
 			messageContext.setProperty(StatisticsConstants.FLOW_STATISTICS_IS_COLLECTED, false);
 		}
 
-		isTracingEnabled = (Boolean) messageContext.getProperty(StatisticsConstants.FLOW_TRACE_IS_COLLECTED);
+		Boolean isTracingEnabled = (Boolean) messageContext.getProperty(StatisticsConstants.FLOW_TRACE_IS_COLLECTED);
 		if (shouldReportStatistic(messageContext)) {
 			StatisticDataUnit statisticDataUnit = new StatisticDataUnit();
 			statisticDataUnit.setComponentName(componentName);
@@ -103,6 +97,7 @@ public static Integer reportEntryEvent(MessageContext messageContext, String com
 				statisticDataUnit.setIsIndividualStatisticCollected(isCollectingStatistics);
 			}
 			StatisticDataCollectionHelper.collectData(messageContext, true, isTracingEnabled, statisticDataUnit);
+			statisticDataUnit.setMessageContext(messageContext);
 
 			StatisticsOpenEvent openEvent = new StatisticsOpenEvent(statisticDataUnit);
             addEventAndIncrementCount(messageContext, openEvent);
@@ -263,6 +258,7 @@ public static void reportFlowAsynchronousEvent(MessageContext messageContext) {
 			BasicStatisticDataUnit dataUnit = new BasicStatisticDataUnit();
 			dataUnit.setStatisticId(StatisticDataCollectionHelper.getStatisticTraceId(messageContext));
 			dataUnit.setCurrentIndex(StatisticDataCollectionHelper.getParentFlowPosition(messageContext, null));
+			dataUnit.setMessageContext(messageContext);
 			AsynchronousExecutionEvent asynchronousExecutionEvent = new AsynchronousExecutionEvent(dataUnit);
 
 			if (isOpenTelemetryEnabled()) {
@@ -291,6 +287,7 @@ private static void reportMediatorStatistics(MessageContext messageContext, Stri
 		statisticDataUnit.setParentIndex(parentIndex);
 		StatisticDataCollectionHelper
 				.collectData(messageContext, isContentAltering, isCollectingTracing, statisticDataUnit);
+		statisticDataUnit.setMessageContext(messageContext);
 
 		StatisticsOpenEvent openEvent = new StatisticsOpenEvent(statisticDataUnit);
         addEventAndIncrementCount(messageContext, openEvent);
@@ -308,6 +305,7 @@ public static void openContinuationEvents(MessageContext synCtx) {
 
             basicStatisticDataUnit.setCurrentIndex(StatisticDataCollectionHelper.getParentFlowPosition(synCtx, null));
             basicStatisticDataUnit.setStatisticId(StatisticDataCollectionHelper.getStatisticTraceId(synCtx));
+			basicStatisticDataUnit.setMessageContext(synCtx);
 
             ParentReopenEvent parentReopenEvent = new ParentReopenEvent(basicStatisticDataUnit);
 			addEvent(synCtx, parentReopenEvent);

modules/core/src/main/java/org/apache/synapse/aspects/flow/statistics/data/raw/BasicStatisticDataUnit.java

@@ -62,6 +62,11 @@ public class BasicStatisticDataUnit {
 	 */
 	private ElasticMetadata elasticMetadata;
 
+	/**
+	 * Message context of the message flow.
+	 */
+	private MessageContext messageContext;
+
 	public String getStatisticId() {
 		return statisticId;
 	}
@@ -117,4 +122,12 @@ public void generateElasticMetadata(MessageContext messageContext) {
 	public ElasticMetadata getElasticMetadata() {
 		return elasticMetadata;
 	}
+
+	public void setMessageContext(MessageContext messageContext) {
+		this.messageContext = messageContext;
+	}
+
+	public MessageContext getMessageContext() {
+		return messageContext;
+	}
 }