Lin 1309 lineage timeout

intg/src/main/java/org/apache/atlas/AtlasConfiguration.java

@@ -95,6 +95,7 @@ public enum AtlasConfiguration {
     LINEAGE_ON_DEMAND_ENABLED("atlas.lineage.on.demand.enabled", true),
     LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT("atlas.lineage.on.demand.default.node.count", 3),
     LINEAGE_MAX_NODE_COUNT("atlas.lineage.max.node.count", 100),
+    LINEAGE_TIMEOUT_MS("atlas.lineage.max.timeout.ms", 15000),
 
     SUPPORTED_RELATIONSHIP_EVENTS("atlas.notification.relationships.filter", "asset_readme,asset_links"),
 
@@ -114,7 +115,8 @@ public enum AtlasConfiguration {
 
     INDEXSEARCH_ASYNC_SEARCH_KEEP_ALIVE_TIME_IN_SECONDS("atlas.indexsearch.async.search.keep.alive.time.in.seconds", 300),
     ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT("atlas.indexsearch.query.size.max.limit", 100000),
-    ATLAS_INDEXSEARCH_LIMIT_UTM_TAGS("atlas.indexsearch.limit.ignore.utm.tags", "project_sdk_python"),
+    ATLAS_INDEXSEARCH_LIMIT_UTM_TAGS("atlas.indexsearch.limit.ignore.utm.tags", ""),
+    ATLAS_INDEXSEARCH_ENABLE_API_LIMIT("atlas.indexsearch.enable.api.limit", false),
 
     ATLAS_MAINTENANCE_MODE("atlas.maintenance.mode", false),
 

intg/src/main/java/org/apache/atlas/model/lineage/AtlasLineageOnDemandInfo.java

@@ -33,6 +33,7 @@ public class AtlasLineageOnDemandInfo implements Serializable {
     private LineageOnDemandRequest                  lineageOnDemandPayload;
     private boolean                                 upstreamEntityLimitReached;
     private boolean                                 downstreamEntityLimitReached;
+    private boolean                                 timeoutOccurred;
 
     public AtlasLineageOnDemandInfo() {
     }
@@ -135,6 +136,14 @@ public void setDownstreamEntityLimitReached(boolean downstreamEntityLimitReached
         this.downstreamEntityLimitReached = downstreamEntityLimitReached;
     }
 
+    public void setTimeoutOccurred(boolean timeoutOccurred) {
+        this.timeoutOccurred = timeoutOccurred;
+    }
+
+    public boolean isTimeoutOccurred() {
+        return timeoutOccurred;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;

repository/src/main/java/org/apache/atlas/discovery/AtlasLineageOnDemandContext.java

@@ -23,6 +23,7 @@ public class AtlasLineageOnDemandContext {
     private Set<String>                             attributes;
     private Set<String>                             relationAttributes;
     private LineageOnDemandBaseParams               defaultParams;
+    private TimeoutChecker                          timeoutChecker;
 
     public AtlasLineageOnDemandContext(LineageOnDemandRequest lineageOnDemandRequest, AtlasTypeRegistry typeRegistry) {
         this.constraints        = lineageOnDemandRequest.getConstraints();
@@ -81,6 +82,14 @@ public void setDefaultParams(LineageOnDemandBaseParams defaultParams) {
         this.defaultParams = defaultParams;
     }
 
+    public TimeoutChecker getTimeoutChecker() {
+        return timeoutChecker;
+    }
+
+    public void setTimeoutChecker(TimeoutChecker timeoutChecker) {
+        this.timeoutChecker = timeoutChecker;
+    }
+
     protected Predicate constructInMemoryPredicate(AtlasTypeRegistry typeRegistry, SearchParameters.FilterCriteria filterCriteria) {
         LineageSearchProcessor lineageSearchProcessor = new LineageSearchProcessor();
         return lineageSearchProcessor.constructInMemoryPredicate(typeRegistry, filterCriteria);

repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java

@@ -88,6 +88,7 @@ public class EntityLineageService implements AtlasLineageService {
     private static final boolean LINEAGE_USING_GREMLIN = AtlasConfiguration.LINEAGE_USING_GREMLIN.getBoolean();
     private static final Integer DEFAULT_LINEAGE_MAX_NODE_COUNT       = 9000;
     private static final int     LINEAGE_ON_DEMAND_DEFAULT_DEPTH      = 3;
+    private static final long LINEAGE_TIMEOUT_MS = AtlasConfiguration.LINEAGE_TIMEOUT_MS.getLong();
     private static final String  SEPARATOR                            = "->";
 
     private final AtlasGraph graph;
@@ -298,6 +299,8 @@ private AtlasLineageOnDemandInfo getLineageInfoOnDemand(String guid, AtlasLineag
         AtomicInteger inputEntitiesTraversed = new AtomicInteger(0);
         AtomicInteger outputEntitiesTraversed = new AtomicInteger(0);
         AtomicInteger traversalOrder = new AtomicInteger(1);
+        TimeoutChecker timeoutChecker = new TimeoutChecker(LINEAGE_TIMEOUT_MS);
+        atlasLineageOnDemandContext.setTimeoutChecker(timeoutChecker);
         if (isDataSet) {
             AtlasVertex datasetVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
             if (direction == AtlasLineageOnDemandInfo.LineageDirection.INPUT || direction == AtlasLineageOnDemandInfo.LineageDirection.BOTH)
@@ -365,6 +368,14 @@ private void traverseEdgesOnDemand(Iterator<AtlasEdge> processEdges, boolean isI
     }
 
     private void traverseEdgesOnDemand(AtlasVertex datasetVertex, boolean isInput, int depth, int level, Set<String> visitedVertices, AtlasLineageOnDemandContext atlasLineageOnDemandContext, AtlasLineageOnDemandInfo ret, String baseGuid, AtomicInteger entitiesTraversed, AtomicInteger traversalOrder) throws AtlasBaseException {
+        // Get timeout checker from context or create new one
+        TimeoutChecker timeoutChecker = atlasLineageOnDemandContext.getTimeoutChecker();
+        // Check timeout before starting traversal
+        if (timeoutChecker.hasTimedOut()) {
+            handleTimeout(ret);
+            return;
+        }
+
         if (isEntityTraversalLimitReached(entitiesTraversed))
             return;
         if (depth != 0) { // base condition of recursion for depth
@@ -379,6 +390,12 @@ private void traverseEdgesOnDemand(AtlasVertex datasetVertex, boolean isInput, i
             RequestContext.get().endMetricRecord(traverseEdgesOnDemandGetEdgesIn);
 
             while (incomingEdges.hasNext()) {
+                // Check timeout periodically
+                if (timeoutChecker.hasTimedOut()) {
+                    handleTimeout(ret);
+                    return;
+                }
+
                 AtlasEdge incomingEdge = incomingEdges.next();
                 AtlasVertex processVertex = incomingEdge.getOutVertex();
 
@@ -407,6 +424,12 @@ private void traverseEdgesOnDemand(AtlasVertex datasetVertex, boolean isInput, i
                 RequestContext.get().endMetricRecord(traverseEdgesOnDemandGetEdgesOut);
 
                 while (outgoingEdges.hasNext()) {
+                    // Check timeout in inner loop as well
+                    if (timeoutChecker.hasTimedOut()) {
+                        handleTimeout(ret);
+                        return;
+                    }
+
                     AtlasEdge outgoingEdge = outgoingEdges.next();
                     AtlasVertex entityVertex = outgoingEdge.getInVertex();
 
@@ -436,7 +459,8 @@ private void traverseEdgesOnDemand(AtlasVertex datasetVertex, boolean isInput, i
                     if (entityVertex != null && !visitedVertices.contains(getId(entityVertex))) {
                         traverseEdgesOnDemand(entityVertex, isInput, depth - 1, nextLevel, visitedVertices, atlasLineageOnDemandContext, ret, baseGuid, entitiesTraversed, traversalOrder); // execute inner depth
                         AtlasEntityHeader traversedEntity = ret.getGuidEntityMap().get(AtlasGraphUtilsV2.getIdFromVertex(entityVertex));
-                        traversedEntity.setFinishTime(traversalOrder.get());
+                        if (traversedEntity != null)
+                            traversedEntity.setFinishTime(traversalOrder.get());
                     }
                 }
             }
@@ -445,6 +469,11 @@ private void traverseEdgesOnDemand(AtlasVertex datasetVertex, boolean isInput, i
         }
     }
 
+    private void handleTimeout(AtlasLineageOnDemandInfo ret) {
+        ret.setTimeoutOccurred(true);
+        LOG.warn("Lineage traversal timed out after {} ms", LINEAGE_TIMEOUT_MS);
+    }
+
     private static void setEntityLimitReachedFlag(boolean isInput, AtlasLineageOnDemandInfo ret) {
         if (isInput)
             ret.setUpstreamEntityLimitReached(true);

repository/src/main/java/org/apache/atlas/discovery/TimeoutChecker.java

@@ -0,0 +1,15 @@
+package org.apache.atlas.discovery;
+
+public class TimeoutChecker {
+    private final long startTime;
+    private final long timeoutMs;
+
+    public TimeoutChecker(long timeoutMs) {
+        this.startTime = System.currentTimeMillis();
+        this.timeoutMs = timeoutMs;
+    }
+
+    public boolean hasTimedOut() {
+        return System.currentTimeMillis() - startTime > timeoutMs;
+    }
+}

webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java

@@ -401,7 +401,7 @@ public AtlasSearchResult indexSearch(@Context HttpServletRequest servletRequest,
                 perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.indexSearch(" + parameters + ")");
             }
 
-            if (parameters.getQuerySize() > AtlasConfiguration.ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT.getLong()) {
+            if (AtlasConfiguration.ATLAS_INDEXSEARCH_ENABLE_API_LIMIT.getBoolean() && parameters.getQuerySize() > AtlasConfiguration.ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT.getLong()) {
                 if(CollectionUtils.isEmpty(parameters.getUtmTags())) {
                     throw new AtlasBaseException(AtlasErrorCode.INVALID_DSL_QUERY_SIZE, String.valueOf(AtlasConfiguration.ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT.getLong()));
                 }