ATLAS-3568: Performance improvements in writing audit logs

intg/src/main/java/org/apache/atlas/listener/EntityChangeListenerV2.java

@@ -74,6 +74,15 @@ public interface EntityChangeListenerV2 {
     void onClassificationsAdded(AtlasEntity entity, List<AtlasClassification> classifications) throws AtlasBaseException;
 
     /**
+     * This is upon adding new classifications to entities.
+     *
+     * @param entities              list of entities
+     * @param classifications classifications that are to be added to entities
+     * @throws AtlasBaseException if the listener notification fails
+     */
+    void onClassificationsAdded(List<AtlasEntity> entities, List<AtlasClassification> classifications) throws AtlasBaseException;
+
+    /**
      * This is upon updating classifications to an entity.
      *
      * @param entity          the entity
@@ -92,6 +101,15 @@ public interface EntityChangeListenerV2 {
     void onClassificationsDeleted(AtlasEntity entity, List<AtlasClassification> classifications) throws AtlasBaseException;
 
     /**
+     * This is upon deleting classifications from entities.
+     *
+     * @param entities              list of entities
+     * @param classifications classifications that needs to be deleted from entities
+     * @throws AtlasBaseException if the listener notification fails
+     */
+    void onClassificationsDeleted(List<AtlasEntity> entities, List<AtlasClassification> classifications) throws AtlasBaseException;
+
+    /**
      * This is upon adding a new term to an entity.
      *
      * @param term     the term

repository/src/main/java/org/apache/atlas/repository/audit/EntityAuditListenerV2.java

@@ -44,6 +44,7 @@ import org.springframework.stereotype.Component;
 import javax.inject.Inject;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -172,6 +173,28 @@ public class EntityAuditListenerV2 implements EntityChangeListenerV2 {
     }
 
     @Override
+    public void onClassificationsAdded(List<AtlasEntity> entities, List<AtlasClassification> classifications) throws AtlasBaseException {
+        if (CollectionUtils.isNotEmpty(classifications)) {
+            MetricRecorder           metric = RequestContext.get().startMetricRecord("entityAudit");
+            List<EntityAuditEventV2> events = Collections.synchronizedList(new ArrayList<>());
+
+            for (AtlasClassification classification : classifications) {
+                for (AtlasEntity entity : entities) {
+                    if (entity.getGuid().equals(classification.getEntityGuid())) {
+                        events.add(createEvent(entity, CLASSIFICATION_ADD, "Added classification: " + AtlasType.toJson(classification)));
+                    } else {
+                        events.add(createEvent(entity, PROPAGATED_CLASSIFICATION_ADD, "Added propagated classification: " + AtlasType.toJson(classification)));
+                    }
+                }
+            }
+
+            auditRepository.putEventsV2(events);
+
+            RequestContext.get().endMetricRecord(metric);
+        }
+    }
+
+    @Override
     public void onClassificationsUpdated(AtlasEntity entity, List<AtlasClassification> classifications) throws AtlasBaseException {
         if (CollectionUtils.isNotEmpty(classifications)) {
             MetricRecorder metric = RequestContext.get().startMetricRecord("entityAudit");
@@ -221,6 +244,28 @@ public class EntityAuditListenerV2 implements EntityChangeListenerV2 {
     }
 
     @Override
+    public void onClassificationsDeleted(List<AtlasEntity> entities, List<AtlasClassification> classifications) throws AtlasBaseException {
+        if (CollectionUtils.isNotEmpty(classifications) && CollectionUtils.isNotEmpty(entities)) {
+            MetricRecorder           metric = RequestContext.get().startMetricRecord("entityAudit");
+            List<EntityAuditEventV2> events = Collections.synchronizedList(new ArrayList<>());
+
+            for (AtlasClassification classification : classifications) {
+                for (AtlasEntity entity : entities) {
+                    if (StringUtils.equals(entity.getGuid(), classification.getEntityGuid())) {
+                        events.add(createEvent(entity, CLASSIFICATION_DELETE, "Deleted classification: " + classification.getTypeName()));
+                    } else {
+                        events.add(createEvent(entity, PROPAGATED_CLASSIFICATION_DELETE, "Deleted propagated classification: " + classification.getTypeName()));
+                    }
+                }
+            }
+
+            auditRepository.putEventsV2(events);
+
+            RequestContext.get().endMetricRecord(metric);
+        }
+    }
+
+    @Override
     public void onTermAdded(AtlasGlossaryTerm term, List<AtlasRelatedObjectId> entities) throws AtlasBaseException {
         if (term != null && CollectionUtils.isNotEmpty(entities)) {
             MetricRecorder metric = RequestContext.get().startMetricRecord("entityAudit");

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

@@ -151,7 +151,7 @@ public class AtlasEntityChangeNotifier {
                 Referenceable entityRef = toReferenceable(entity.getGuid());
                 List<Struct>  traits    = toStruct(addedClassifications);
 
-                if (entity == null || CollectionUtils.isEmpty(traits)) {
+                if (entityRef == null || CollectionUtils.isEmpty(traits)) {
                     return;
                 }
 
@@ -166,6 +166,41 @@ public class AtlasEntityChangeNotifier {
         }
     }
 
+    public void onClassificationsAddedToEntities(List<AtlasEntity> entities, List<AtlasClassification> addedClassifications) throws AtlasBaseException {
+        if (isV2EntityNotificationEnabled) {
+            doFullTextMappingHelper(entities);
+
+            for (EntityChangeListenerV2 listener : entityChangeListenersV2) {
+                listener.onClassificationsAdded(entities, addedClassifications);
+            }
+        } else {
+            updateFullTextMapping(entities, addedClassifications);
+
+            if (instanceConverter != null) {
+                List<Struct> traits = toStruct(addedClassifications);
+
+                if (!CollectionUtils.isEmpty(traits)) {
+                    for(AtlasEntity entity : entities) {
+                        Referenceable entityRef = toReferenceable(entity.getGuid());
+
+                        if (entityRef == null) {
+                            LOG.warn("EntityRef with guid {} not found while adding classifications {} ", entity.getGuid(), addedClassifications);
+                            continue;
+                        }
+
+                        for (EntityChangeListener listener : entityChangeListeners) {
+                            try {
+                                listener.onTraitsAdded(entityRef, traits);
+                            } catch (AtlasException e) {
+                                throw new AtlasBaseException(AtlasErrorCode.NOTIFICATION_FAILED, e, getListenerName(listener), "TraitAdd");
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     public void onClassificationUpdatedToEntity(AtlasEntity entity, List<AtlasClassification> updatedClassifications) throws AtlasBaseException {
         doFullTextMapping(entity.getGuid());
 
@@ -220,6 +255,39 @@ public class AtlasEntityChangeNotifier {
         }
     }
 
+    public void onClassificationsDeletedFromEntities(List<AtlasEntity> entities, List<AtlasClassification> deletedClassifications) throws AtlasBaseException {
+        doFullTextMappingHelper(entities);
+
+        if (isV2EntityNotificationEnabled) {
+            for (EntityChangeListenerV2 listener : entityChangeListenersV2) {
+                listener.onClassificationsDeleted(entities, deletedClassifications);
+            }
+        } else {
+            if (instanceConverter != null) {
+                List<Struct> traits = toStruct(deletedClassifications);
+
+                if(!CollectionUtils.isEmpty(deletedClassifications)) {
+                    for(AtlasEntity entity : entities) {
+                        Referenceable entityRef = toReferenceable(entity.getGuid());
+
+                        if (entityRef == null) {
+                            LOG.warn("EntityRef with guid {} not found while deleting classifications {} ", entity.getGuid(), deletedClassifications);
+                            continue;
+                        }
+
+                        for (EntityChangeListener listener : entityChangeListeners) {
+                            try {
+                                listener.onTraitsDeleted(entityRef, traits);
+                            } catch (AtlasException e) {
+                                throw new AtlasBaseException(AtlasErrorCode.NOTIFICATION_FAILED, e, getListenerName(listener), "TraitDelete");
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     public void onTermAddedToEntities(AtlasGlossaryTerm term, List<AtlasRelatedObjectId> entityIds) throws AtlasBaseException {
         // listeners notified on term-entity association only if v2 notifications are enabled
         if (isV2EntityNotificationEnabled) {
@@ -575,6 +643,12 @@ public class AtlasEntityChangeNotifier {
         RequestContext.get().endMetricRecord(metric);
     }
 
+    private void updateFullTextMapping(List<AtlasEntity> entities, List<AtlasClassification> classifications) {
+        for (AtlasEntity entity : entities) {
+            updateFullTextMapping(entity.getGuid(), classifications);
+        }
+    }
+
     private void doFullTextMapping(String guid) {
         if(AtlasRepositoryConfiguration.isFreeTextSearchEnabled() || !AtlasRepositoryConfiguration.isFullTextSearchEnabled()) {
             return;
@@ -586,6 +660,12 @@ public class AtlasEntityChangeNotifier {
         doFullTextMapping(Collections.singletonList(entityHeader));
     }
 
+    private void doFullTextMappingHelper(List<AtlasEntity> entities) {
+        for (AtlasEntity entity : entities) {
+            doFullTextMapping(entity.getGuid());
+        }
+    }
+
     private void pruneResponse(EntityMutationResponse resp) {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> pruneResponse()");

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

@@ -631,6 +631,12 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             LOG.debug("Updating classifications={} for entity={}", classifications, guid);
         }
 
+        AtlasPerfTracer perf = null;
+
+        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
+            AtlasPerfTracer.getPerfTracer(PERF_LOG, "AtlasEntityStoreV2.updateClassification()");
+        }
+
         if (StringUtils.isEmpty(guid)) {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "Guid not specified");
         }
@@ -663,6 +669,8 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
         }
 
         entityGraphMapper.updateClassifications(context, guid, classifications);
+
+        AtlasPerfTracer.log(perf);
     }
 
     @Override

webapp/src/main/java/org/apache/atlas/notification/EntityNotificationListenerV2.java

@@ -107,6 +107,11 @@ public class EntityNotificationListenerV2 implements EntityChangeListenerV2 {
     }
 
     @Override
+    public void onClassificationsAdded(List<AtlasEntity> entities, List<AtlasClassification> classifications) throws AtlasBaseException {
+        notifyEntityEvents(entities, CLASSIFICATION_ADD);
+    }
+
+    @Override
     public void onClassificationsUpdated(AtlasEntity entity, List<AtlasClassification> classifications) throws AtlasBaseException {
         Map<String, List<AtlasClassification>> addedPropagations   = RequestContext.get().getAddedPropagations();
         Map<String, List<AtlasClassification>> removedPropagations = RequestContext.get().getRemovedPropagations();
@@ -124,6 +129,11 @@ public class EntityNotificationListenerV2 implements EntityChangeListenerV2 {
     }
 
     @Override
+    public void onClassificationsDeleted(List<AtlasEntity> entities, List<AtlasClassification> classifications) throws AtlasBaseException {
+        notifyEntityEvents(entities, CLASSIFICATION_DELETE);
+    }
+
+    @Override
     public void onTermAdded(AtlasGlossaryTerm term, List<AtlasRelatedObjectId> entities) {
         // do nothing -> notification not sent out for term assignment to entities
     }