ATLAS-3568: Performance improvements in writing audit logs #2

Signed-off-by: Sarath Subramanian <sarath@apache.org>
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java
index 765ba36..e427a59 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java
@@ -1755,6 +1755,11 @@
                 // add the attributes for the trait instance
                 mapClassification(EntityOperation.CREATE, context, classification, entityType, entityVertex, classificationVertex);
                 updateModificationMetadata(entityVertex);
+                if(addedClassifications.get(classification) == null) {
+                    addedClassifications.put(classification, new HashSet<>());
+                }
+                //Add current Vertex to be notified
+                addedClassifications.get(classification).add(entityVertex);
 
                 if (propagateTags) {
                     // compute propagatedEntityVertices only once
@@ -1770,11 +1775,7 @@
                         List<AtlasVertex> entitiesPropagatedTo = deleteDelegate.getHandler().addTagPropagation(classificationVertex, entitiesToPropagateTo);
 
                         if (CollectionUtils.isNotEmpty(entitiesPropagatedTo)) {
-                            if(addedClassifications.get(classification) == null) {
-                                addedClassifications.put(classification, new HashSet<>(entitiesPropagatedTo));
-                            } else {
-                                addedClassifications.get(classification).addAll(entitiesPropagatedTo);
-                            }
+                            addedClassifications.get(classification).addAll(entitiesPropagatedTo);
                         }
                     } else {
                         if (LOG.isDebugEnabled()) {
@@ -1801,7 +1802,7 @@
                 Set<AtlasVertex>  vertices           = addedClassifications.get(classification);
                 List<AtlasEntity> propagatedEntities = updateClassificationText(classification, vertices);
 
-                entityChangeNotifier.onClassificationsAddedToEntities(propagatedEntities, classifications);
+                entityChangeNotifier.onClassificationsAddedToEntities(propagatedEntities, Collections.singletonList(classification));
             }
 
             RequestContext.get().endMetricRecord(metric);