Merge pull request #3822 from atlanhq/fix/LIN-1388

Added perf loggers
atlanhq · Dec 4, 2024 · f2b516f · f2b516f
2 parents 1672d81 + 343fef4
commit f2b516f
Showing 1 changed file with 38 additions and 21 deletions.
diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java
@@ -658,18 +658,23 @@ private static boolean isInputDirection(AtlasLineageListContext lineageListConte
     }
 
     private boolean checkForOffset(AtlasEdge atlasEdge, AtlasVertex entityVertex, AtlasLineageOnDemandContext atlasLineageOnDemandContext, AtlasLineageOnDemandInfo ret) {
-        String entityGuid = getGuid(entityVertex);
-        LineageOnDemandConstraints entityConstraints = getAndValidateLineageConstraintsByGuid(entityGuid, atlasLineageOnDemandContext);
-        LineageInfoOnDemand entityLineageInfo = ret.getRelationsOnDemand().containsKey(entityGuid) ? ret.getRelationsOnDemand().get(entityGuid) : new LineageInfoOnDemand(entityConstraints);
-
-        if (entityConstraints.getFrom() != 0 && entityLineageInfo.getFromCounter() < entityConstraints.getFrom()) {
-            if (! lineageContainsSkippedEdgeV2(ret, atlasEdge)) {
-                addEdgeToSkippedEdges(ret, atlasEdge);
-                entityLineageInfo.incrementFromCounter();
+        AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("checkForOffset");
+        try {
+            String entityGuid = getGuid(entityVertex);
+            LineageOnDemandConstraints entityConstraints = getAndValidateLineageConstraintsByGuid(entityGuid, atlasLineageOnDemandContext);
+            LineageInfoOnDemand entityLineageInfo = ret.getRelationsOnDemand().containsKey(entityGuid) ? ret.getRelationsOnDemand().get(entityGuid) : new LineageInfoOnDemand(entityConstraints);
+
+            if (entityConstraints.getFrom() != 0 && entityLineageInfo.getFromCounter() < entityConstraints.getFrom()) {
+                if (! lineageContainsSkippedEdgeV2(ret, atlasEdge)) {
+                    addEdgeToSkippedEdges(ret, atlasEdge);
+                    entityLineageInfo.incrementFromCounter();
+                }
+                return true;
             }
-            return true;
+            return false;
+        } finally {
+            RequestContext.get().endMetricRecord(metric);
         }
-        return false;
     }
 
     private boolean checkOffsetAndSkipEntity(AtlasLineageListContext atlasLineageListContext, AtlasLineageListInfo ret) {
@@ -996,17 +1001,22 @@ private int getLineageMaxNodeAllowedCount() {
     }
 
     private String getEdgeLabel(AtlasEdge edge) {
-        AtlasVertex inVertex     = edge.getInVertex();
-        AtlasVertex outVertex    = edge.getOutVertex();
-        String      inGuid       = AtlasGraphUtilsV2.getIdFromVertex(inVertex);
-        String      outGuid      = AtlasGraphUtilsV2.getIdFromVertex(outVertex);
-        String      relationGuid = AtlasGraphUtilsV2.getEncodedProperty(edge, RELATIONSHIP_GUID_PROPERTY_KEY, String.class);
-        boolean     isInputEdge  = edge.getLabel().equalsIgnoreCase(PROCESS_INPUTS_EDGE);
-
-        if (isLineageOnDemandEnabled()) {
-            return getEdgeLabelFromGuids(isInputEdge, inGuid, outGuid);
+        AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("getEdgeLabel");
+        try {
+            AtlasVertex inVertex     = edge.getInVertex();
+            AtlasVertex outVertex    = edge.getOutVertex();
+            String      inGuid       = AtlasGraphUtilsV2.getIdFromVertex(inVertex);
+            String      outGuid      = AtlasGraphUtilsV2.getIdFromVertex(outVertex);
+            String      relationGuid = AtlasGraphUtilsV2.getEncodedProperty(edge, RELATIONSHIP_GUID_PROPERTY_KEY, String.class);
+            boolean     isInputEdge  = edge.getLabel().equalsIgnoreCase(PROCESS_INPUTS_EDGE);
+
+            if (isLineageOnDemandEnabled()) {
+                return getEdgeLabelFromGuids(isInputEdge, inGuid, outGuid);
+            }
+            return relationGuid;
+        } finally {
+            RequestContext.get().endMetricRecord(metric);
         }
-        return relationGuid;
     }
 
     private String getEdgeLabelFromGuids(boolean isInputEdge, String inGuid, String outGuid) {
@@ -1309,7 +1319,12 @@ private List<AtlasEdge> getEdgesOfProcess(boolean isInput, AtlasLineageContext l
     }
 
     private boolean vertexMatchesEvaluation(AtlasVertex currentVertex, AtlasLineageContext lineageContext) {
-        return currentVertex.equals(lineageContext.getStartDatasetVertex()) || lineageContext.evaluate(currentVertex);
+        AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("vertexMatchesEvaluation");
+        try {
+            return currentVertex.equals(lineageContext.getStartDatasetVertex()) || lineageContext.evaluate(currentVertex);
+        } finally {
+            RequestContext.get().endMetricRecord(metric);
+        }
     }
 
     private boolean vertexMatchesEvaluation(AtlasVertex currentVertex, AtlasLineageOnDemandContext atlasLineageOnDemandContext) {
@@ -1384,9 +1399,11 @@ private boolean lineageContainsSkippedEdgeV2(AtlasLineageOnDemandInfo lineageInf
     }
 
     private void addEdgeToSkippedEdges(AtlasLineageOnDemandInfo lineageInfo, AtlasEdge edge) {
+        AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("addEdgeToSkippedEdges");
         if (lineageInfo.getSkippedEdges() != null) {
             lineageInfo.getSkippedEdges().add(getEdgeLabel(edge));
         }
+        RequestContext.get().endMetricRecord(metric);
     }
 
     private AtlasLineageInfo initializeLineageInfo(String guid, LineageDirection direction, int depth, int limit, int offset) {