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

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -1755,6 +1755,11 @@ public class EntityGraphMapper {
                 // 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,12 +1775,8 @@ public class EntityGraphMapper {
                         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);
                         }
-                        }
                     } else {
                         if (LOG.isDebugEnabled()) {
                             LOG.debug(" --> Not propagating classification: [{}][{}] - no entities found to propagate to.", getTypeName(classificationVertex), entityType.getTypeName());
@@ -1801,7 +1802,7 @@ public class EntityGraphMapper {
                 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);