ATLAS-3143: PatchFx: Memory usage and performance improvement.

graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasGraphQuery.java

@@ -96,6 +96,27 @@ public interface AtlasGraphQuery<V, E> {
      */
     Iterable<AtlasVertex<V, E>> vertices(int offset, int limit);
 
+    /**
+     * Executes the query and returns IDs of matching vertices.
+     * @return
+     */
+    Iterable<Object> vertexIds();
+
+    /**
+     * Executes the query and returns IDs of the matching vertices from given offset till the max limit
+     * @param limit max number of vertices
+     * @return
+     */
+    Iterable<Object> vertexIds(int limit);
+
+    /**
+     * Executes the query and returns IDs of the matching vertices from given offset till the max limit
+     * @param offset starting offset
+     * @param limit max number of vertices
+     * @return
+     */
+    Iterable<Object> vertexIds(int offset, int limit);
+
 
     /**
      * Adds a predicate that the returned vertices must have the specified

graphdb/common/src/main/java/org/apache/atlas/repository/graphdb/tinkerpop/query/NativeTinkerpopGraphQuery.java

@@ -77,6 +77,27 @@ public interface NativeTinkerpopGraphQuery<V, E> {
      */
     Iterable<AtlasVertex<V, E>> vertices(int offset, int limit);
 
+    /**
+     * Executes the graph query.
+     * @return
+     */
+    Iterable<Object> vertexIds();
+
+    /**
+     * Executes graph query
+     * @param limit Max vertices to return
+     * @return
+     */
+    Iterable<Object> vertexIds(int limit);
+
+    /**
+     * Executes graph query
+     * @param offset Starting offset
+     * @param limit Max vertices to return
+     * @return
+     */
+    Iterable<Object> vertexIds(int offset, int limit);
+
 
     /**
      * Adds an in condition to the query.

graphdb/common/src/main/java/org/apache/atlas/repository/graphdb/tinkerpop/query/TinkerpopGraphQuery.java

@@ -237,6 +237,63 @@ public abstract class TinkerpopGraphQuery<V, E> implements AtlasGraphQuery<V, E>
     }
 
     @Override
+    public Iterable<Object> vertexIds() {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Executing: " + queryCondition);
+        }
+
+        // Compute the overall result by combining the results of all the AndConditions (nested within OR) together.
+        Set<Object> result = new HashSet<>();
+
+        for(AndCondition andExpr : queryCondition.getAndTerms()) {
+            NativeTinkerpopGraphQuery<V, E> andQuery = andExpr.create(getQueryFactory());
+            for(Object vertexId : andQuery.vertexIds()) {
+                result.add(vertexId);
+            }
+        }
+        return result;
+    }
+
+    @Override
+    public Iterable<Object> vertexIds(int limit) {
+        return vertexIds(0, limit);
+    }
+
+    @Override
+    public Iterable<Object> vertexIds(int offset, int limit) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Executing: " + queryCondition);
+        }
+
+        Preconditions.checkArgument(offset >= 0, "Offset must be non-negative");
+        Preconditions.checkArgument(limit >= 0, "Limit must be non-negative");
+
+        // Compute the overall result by combining the results of all the AndConditions (nested within OR) together.
+        Set<Object> result = new HashSet<>();
+        long resultIdx = 0;
+        for(AndCondition andExpr : queryCondition.getAndTerms()) {
+            if (result.size() == limit) {
+                break;
+            }
+
+            NativeTinkerpopGraphQuery<V, E> andQuery = andExpr.create(getQueryFactory());
+            for(Object vertexId : andQuery.vertexIds(offset + limit)) {
+                if (resultIdx >= offset) {
+                    result.add(vertexId);
+
+                    if (result.size() == limit) {
+                        break;
+                    }
+                }
+
+                resultIdx++;
+            }
+        }
+
+        return result;
+    }
+
+    @Override
     public AtlasGraphQuery<V, E> has(String propertyKey, QueryOperator operator,
             Object value) {
         queryCondition.andWith(new HasPredicate(propertyKey, operator, value));

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/query/NativeJanusGraphQuery.java

@@ -158,6 +158,64 @@ public class NativeJanusGraphQuery implements NativeTinkerpopGraphQuery<AtlasJan
     }
 
     @Override
+    public Iterable<Object> vertexIds() {
+        Set<Object>                result = new HashSet<>();
+        Iterable<JanusGraphVertex> it     = query.vertices();
+
+        for (Iterator<? extends Vertex> iter = it.iterator(); iter.hasNext(); ) {
+            result.add(iter.next().id());
+        }
+
+        return result;
+    }
+
+    @Override
+    public Iterable<Object> vertexIds(int limit) {
+        Set<Object>                result = new HashSet<>(limit);
+        Iterable<JanusGraphVertex> it     = query.limit(limit).vertices();
+
+        if (LOG.isDebugEnabled()) {
+            if (query instanceof GraphCentricQueryBuilder) {
+                LOG.debug("NativeJanusGraphQuery.vertices({}): resultSize={}, {}", limit, getCountForDebugLog(it), ((GraphCentricQueryBuilder) query).constructQuery(ElementCategory.VERTEX));
+            } else {
+                LOG.debug("NativeJanusGraphQuery.vertices({}): resultSize={}, {}", limit, getCountForDebugLog(it), query);
+            }
+        }
+
+        for (Iterator<? extends Vertex> iter = it.iterator(); iter.hasNext(); ) {
+            result.add(iter.next().id());
+        }
+
+        return result;
+    }
+
+    @Override
+    public Iterable<Object> vertexIds(int offset, int limit) {
+        Set<Object>                result = new HashSet<>(limit);
+        Iterable<JanusGraphVertex> it     = query.limit(offset + limit).vertices();
+
+        if (LOG.isDebugEnabled()) {
+            if (query instanceof GraphCentricQueryBuilder) {
+                LOG.debug("NativeJanusGraphQuery.vertices({}, {}): resultSize={}, {}", offset, limit, getCountForDebugLog(it), ((GraphCentricQueryBuilder) query).constructQuery(ElementCategory.VERTEX));
+            } else {
+                LOG.debug("NativeJanusGraphQuery.vertices({}, {}): resultSize={}, {}", offset, limit, getCountForDebugLog(it), query);
+            }
+        }
+
+        Iterator<? extends Vertex> iter = it.iterator();
+
+        for (long resultIdx = 0; iter.hasNext() && result.size() < limit; resultIdx++) {
+            if (resultIdx < offset) {
+                continue;
+            }
+
+            result.add(iter.next().id());
+        }
+
+        return result;
+    }
+
+    @Override
     public void in(String propertyName, Collection<? extends Object> values) {
         query.has(propertyName, Contain.IN, values);
 

intg/src/main/java/org/apache/atlas/pc/WorkItemConsumer.java

@@ -30,7 +30,7 @@ import java.util.concurrent.atomic.AtomicLong;
 public abstract class WorkItemConsumer<T> implements Runnable {
     private static final Logger LOG = LoggerFactory.getLogger(WorkItemConsumer.class);
 
-    private static final int POLLING_DURATION_SECONDS  = 30;
+    private static final int POLLING_DURATION_SECONDS  = 5;
     private static final int DEFAULT_COMMIT_TIME_IN_MS = 15000;
 
     private final BlockingQueue<T> queue;
@@ -50,7 +50,7 @@ public abstract class WorkItemConsumer<T> implements Runnable {
                 T item = queue.poll(POLLING_DURATION_SECONDS, TimeUnit.SECONDS);
 
                 if (item == null) {
-                    LOG.warn("WorkItemConsumer.run(): no more items found in the queue. Will exit after committing");
+                    LOG.debug("WorkItemConsumer.run(): no more items found in the queue. Will exit after committing");
 
                     commitDirty();
 

repository/src/main/java/org/apache/atlas/repository/patches/AtlasPatchHandler.java

@@ -41,8 +41,8 @@ public abstract class AtlasPatchHandler {
     private void register() {
         PatchStatus patchStatus = getStatus();
 
-        if (patchStatus == UNKNOWN) {
-            patchRegistry.register(patchId, patchDescription, JAVA_PATCH_TYPE, getStatus());
+        if (patchStatus == null || patchStatus == UNKNOWN) {
+            patchRegistry.register(patchId, patchDescription, JAVA_PATCH_TYPE, "apply", UNKNOWN);
         }
     }
 

repository/src/main/java/org/apache/atlas/repository/patches/AtlasPatchRegistry.java

@@ -59,8 +59,16 @@ public class AtlasPatchRegistry {
     private final AtlasGraph               graph;
 
     public AtlasPatchRegistry(AtlasGraph graph) {
+        LOG.info("AtlasPatchRegistry: initializing..");
+
         this.graph              = graph;
         this.patchNameStatusMap = getPatchNameStatusForAllRegistered(graph);
+
+        LOG.info("AtlasPatchRegistry: found {} patches", patchNameStatusMap.size());
+
+        for (Map.Entry<String, PatchStatus> entry : patchNameStatusMap.entrySet()) {
+            LOG.info("AtlasPatchRegistry: patchId={}, status={}", entry.getKey(), entry.getValue());
+        }
     }
 
     public boolean isApplicable(String incomingId, String patchFile, int index) {
@@ -83,8 +91,8 @@ public class AtlasPatchRegistry {
         return patchNameStatusMap.get(id);
     }
 
-    public void register(String patchId, String description, String action, PatchStatus patchStatus) {
-        createOrUpdatePatchVertex(graph, patchId, description, action, patchStatus);
+    public void register(String patchId, String description, String patchType, String action, PatchStatus patchStatus) {
+        createOrUpdatePatchVertex(graph, patchId, description, patchType, action, patchStatus);
     }
 
     public void updateStatus(String patchId, PatchStatus patchStatus) {
@@ -118,14 +126,14 @@ public class AtlasPatchRegistry {
         return getAllPatches(graph);
     }
 
-    private void createOrUpdatePatchVertex(AtlasGraph graph, String patchId,
-                                           String description, String action, PatchStatus patchStatus) {
+    private void createOrUpdatePatchVertex(AtlasGraph graph, String patchId, String description,
+                                           String patchType, String action, PatchStatus patchStatus) {
         boolean     isPatchRegistered = MapUtils.isNotEmpty(patchNameStatusMap) && patchNameStatusMap.containsKey(patchId);
         AtlasVertex patchVertex       = isPatchRegistered ? findByPatchId(patchId) : graph.addVertex();
 
         setEncodedProperty(patchVertex, PATCH_ID_PROPERTY_KEY, patchId);
         setEncodedProperty(patchVertex, PATCH_DESCRIPTION_PROPERTY_KEY, description);
-        setEncodedProperty(patchVertex, PATCH_TYPE_PROPERTY_KEY, TYPEDEF_PATCH_TYPE);
+        setEncodedProperty(patchVertex, PATCH_TYPE_PROPERTY_KEY, patchType);
         setEncodedProperty(patchVertex, PATCH_ACTION_PROPERTY_KEY, action);
         setEncodedProperty(patchVertex, PATCH_STATE_PROPERTY_KEY, patchStatus.toString());
         setEncodedProperty(patchVertex, TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());

repository/src/main/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializer.java

@@ -76,6 +76,7 @@ import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.PUBLIC_
 import static org.apache.atlas.model.patches.AtlasPatch.PatchStatus.APPLIED;
 import static org.apache.atlas.model.patches.AtlasPatch.PatchStatus.FAILED;
 import static org.apache.atlas.model.patches.AtlasPatch.PatchStatus.SKIPPED;
+import static org.apache.atlas.model.patches.AtlasPatch.PatchStatus.UNKNOWN;
 
 /**
  * Class that handles initial loading of models and patches into typedef store
@@ -475,7 +476,7 @@ public class AtlasTypeDefStoreInitializer implements ActiveStateChangeHandler {
                                                patch.getId(), status.toString(), patch.getAction(), patchFile);
                                 }
 
-                                patchRegistry.register(patch.id, patch.description, patch.action, status);
+                                patchRegistry.register(patch.id, patch.description, TYPEDEF_PATCH_TYPE, patch.action, status);
                                 LOG.info("{} (status: {}; action: {}) in file: {}", patch.getId(), status.toString(), patch.getAction(), patchFile);
                             } else {
                                 LOG.info("{} in file: {} already {}. Ignoring.", patch.getId(), patchFile, patchRegistry.getStatus(patch.getId()).toString());
@@ -783,7 +784,7 @@ public class AtlasTypeDefStoreInitializer implements ActiveStateChangeHandler {
         public PatchStatus applyPatch(TypeDefPatch patch) throws AtlasBaseException {
             String           typeName       = patch.getTypeName();
             AtlasBaseTypeDef typeDef        = typeRegistry.getTypeDefByName(typeName);
-            PatchStatus      ret            = null;
+            PatchStatus      ret            = UNKNOWN;
 
             if (typeDef == null) {
                 throw new AtlasBaseException(AtlasErrorCode.PATCH_FOR_UNKNOWN_TYPE, patch.getAction(), typeName);

repository/src/test/java/org/apache/atlas/patches/AtlasPatchRegistryTest.java

@@ -27,6 +27,7 @@ import org.testng.annotations.Test;
 
 import javax.inject.Inject;
 
+import static org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer.TYPEDEF_PATCH_TYPE;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 
@@ -46,7 +47,7 @@ public class AtlasPatchRegistryTest {
     public void registerPatch() {
         AtlasPatchRegistry registry = new AtlasPatchRegistry(graph);
 
-        registry.register("1", "test patch", "apply", AtlasPatch.PatchStatus.UNKNOWN);
+        registry.register("1", "test patch", TYPEDEF_PATCH_TYPE, "apply", AtlasPatch.PatchStatus.UNKNOWN);
 
         assertPatches(registry, 1);
     }