PLT-191: Add more metrics for recording

atlanhq · Nov 16, 2023 · 258acc4 · 258acc4
1 parent 1cb28a0
commit 258acc4
Showing 1 changed file with 32 additions and 27 deletions.
diff --git a/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -243,6 +243,7 @@ public Object getEntityAttribute(AtlasVertex entityVertex, AtlasAttribute attrib
     }
 
     public AtlasObjectId toAtlasObjectId(AtlasVertex entityVertex) throws AtlasBaseException {
+        AtlasPerfMetrics.MetricRecorder toAtlasObjectId = RequestContext.get().startMetricRecord("toAtlasObjectId");
         AtlasObjectId   ret        = null;
         String          typeName   = entityVertex.getProperty(Constants.TYPE_NAME_PROPERTY_KEY, String.class);
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
@@ -274,7 +275,7 @@ public AtlasObjectId toAtlasObjectId(AtlasVertex entityVertex) throws AtlasBaseE
             }
             ret = new AtlasObjectId(entityVertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class), typeName, uniqueAttributes, attributes);
         }
-
+        RequestContext.get().endMetricRecord(toAtlasObjectId);
         return ret;
     }
 
@@ -1058,7 +1059,6 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeader2(AtlasVertex entityVertex
                 attributes.stream().parallel().forEach(att->{
                     try {
                         RequestContext.setCurrentContext(context);
-                        Set<String> params = RequestContext.get().getRelationAttrsForSearch();
                         enrichAttributes(entityVertex, entityType, att, ret);
                     } catch (AtlasBaseException e) {
                         throw new RuntimeException(e);
@@ -1305,6 +1305,7 @@ private Object mapVertexToAttribute(AtlasVertex entityVertex, AtlasAttribute att
     }
 
     private Object mapVertexToAttribute(AtlasVertex entityVertex, AtlasAttribute attribute, AtlasEntityExtInfo entityExtInfo, final boolean isMinExtInfo, boolean includeReferences, boolean ignoreInactive) throws AtlasBaseException {
+        AtlasPerfMetrics.MetricRecorder mapVertexToAttribute = RequestContext.get().startMetricRecord("mapVertexToAttribute");
         Object    ret                = null;
         AtlasType attrType           = attribute.getAttributeType();
         String    edgeLabel          = attribute.getRelationshipEdgeLabel();
@@ -1386,7 +1387,7 @@ private Object mapVertexToAttribute(AtlasVertex entityVertex, AtlasAttribute att
                 // do nothing
                 break;
         }
-
+        RequestContext.get().endMetricRecord(mapVertexToAttribute);
         return ret;
     }
 
@@ -1626,42 +1627,46 @@ public static Object mapVertexToPrimitive(AtlasElement entityVertex, final Strin
     private AtlasObjectId mapVertexToObjectId(AtlasVertex entityVertex, String edgeLabel, AtlasEdge edge,
                                               AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute,
                                               AtlasRelationshipEdgeDirection edgeDirection, final boolean isMinExtInfo) throws AtlasBaseException {
+        AtlasPerfMetrics.MetricRecorder mapVertexToObjectId = RequestContext.get().startMetricRecord("mapVertexToObjectId");
         AtlasObjectId ret = null;
 
-        if (edge == null) {
-            edge = graphHelper.getEdgeForLabel(entityVertex, edgeLabel, edgeDirection);
-        }
-
-        if (GraphHelper.elementExists(edge)) {
-            if (!RequestContext.get().isAllowDeletedRelationsIndexsearch() && getState(edge) == Id.EntityState.DELETED ) {
-                return null;
+        try {
+            if (edge == null) {
+                edge = graphHelper.getEdgeForLabel(entityVertex, edgeLabel, edgeDirection);
             }
 
-            AtlasVertex referenceVertex = edge.getInVertex();
+            if (GraphHelper.elementExists(edge)) {
+                if (!RequestContext.get().isAllowDeletedRelationsIndexsearch() && getState(edge) == Id.EntityState.DELETED) {
+                    return null;
+                }
 
-            if (StringUtils.equals(getIdFromVertex(referenceVertex), getIdFromVertex(entityVertex))) {
-                referenceVertex = edge.getOutVertex();
-            }
+                AtlasVertex referenceVertex = edge.getInVertex();
 
-            if (referenceVertex != null) {
-                if (entityExtInfo != null && isOwnedAttribute) {
-                    final AtlasEntity entity;
+                if (StringUtils.equals(getIdFromVertex(referenceVertex), getIdFromVertex(entityVertex))) {
+                    referenceVertex = edge.getOutVertex();
+                }
 
-                    if (isMinExtInfo) {
-                        entity = mapVertexToAtlasEntityMin(referenceVertex, entityExtInfo);
-                    } else {
-                        entity = mapVertexToAtlasEntity(referenceVertex, entityExtInfo);
-                    }
+                if (referenceVertex != null) {
+                    if (entityExtInfo != null && isOwnedAttribute) {
+                        final AtlasEntity entity;
 
-                    if (entity != null) {
-                        ret = AtlasTypeUtil.getAtlasObjectId(entity);
+                        if (isMinExtInfo) {
+                            entity = mapVertexToAtlasEntityMin(referenceVertex, entityExtInfo);
+                        } else {
+                            entity = mapVertexToAtlasEntity(referenceVertex, entityExtInfo);
+                        }
+
+                        if (entity != null) {
+                            ret = AtlasTypeUtil.getAtlasObjectId(entity);
+                        }
+                    } else {
+                        ret = toAtlasObjectId(referenceVertex);
                     }
-                } else {
-                    ret = toAtlasObjectId(referenceVertex);
                 }
             }
+        } finally {
+            RequestContext.get().endMetricRecord(mapVertexToObjectId);
         }
-
         return ret;
     }