From bcc0af4f97781068b6cf200faad42529670036f1 Mon Sep 17 00:00:00 2001
From: aarshi <aarshi.arora@atlan.com>
Date: Thu, 28 Nov 2024 13:11:17 +0530
Subject: [PATCH] add processing time

---
 .../repository/store/graph/v2/EntityGraphRetriever.java  | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
index e2c767e816..41386da015 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -1014,6 +1014,8 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeader(AtlasVertex entityVertex)
 
     private AtlasEntityHeader mapVertexToAtlasEntityHeader(AtlasVertex entityVertex, Set<String> attributes) throws AtlasBaseException {
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("mapVertexToAtlasEntityHeader");
+        AtlasPerfMetrics.MetricRecorder attributeProcessingRecorder = null;
+        boolean metricEnded = false;
         AtlasEntityHeader ret = new AtlasEntityHeader();
         try {
             //pre-fetching the properties
@@ -1068,6 +1070,7 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeader(AtlasVertex entityVertex,
                 }
 
                 if (CollectionUtils.isNotEmpty(attributes)) {
+                    attributeProcessingRecorder = RequestContext.get().startMetricRecord("attributeProcessingTime");
                     for (String attrName : attributes) {
                         AtlasAttribute attribute = entityType.getAttribute(attrName);
 
@@ -1090,11 +1093,17 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeader(AtlasVertex entityVertex,
                         if (attrValue != null) {
                             ret.setAttribute(attrName, attrValue);
                         }
+
                     }
+                    RequestContext.get().endMetricRecord(attributeProcessingRecorder);
+                    metricEnded = true;
                 }
             }
         }
         finally {
+            if (!metricEnded) {
+                RequestContext.get().endMetricRecord(attributeProcessingRecorder);
+            }
             RequestContext.get().endMetricRecord(metricRecorder);
         }
         return ret;