ATLAS-3360: Duplicate audits in atlas when HMS and hive hook is enabled #2 - Add…

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

@@ -838,17 +838,13 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
                         }
 
                         entitiesToSkipUpdate.add(entity);
+                        RequestContext.get().recordEntityToSkip(entity.getGuid());
                     }
                 }
 
                 if (entitiesToSkipUpdate != null) {
                     // remove entitiesToSkipUpdate from EntityMutationContext
                     context.getUpdatedEntities().removeAll(entitiesToSkipUpdate);
-
-                    // remove entitiesToSkipUpdate from RequestContext
-                    for (AtlasEntity entity : entitiesToSkipUpdate) {
-                        RequestContext.get().removeEntityUpdate(entity);
-                    }
                 }
 
                 // Check if authorized to update entities

server-api/src/main/java/org/apache/atlas/RequestContext.java

@@ -54,6 +54,7 @@ public class RequestContext {
     private final Map<String, List<AtlasClassification>> removedPropagations  = new HashMap<>();
     private final AtlasPerfMetrics                       metrics              = isMetricsEnabled ? new AtlasPerfMetrics() : null;
     private       List<EntityGuidPair>                   entityGuidInRequest  = null;
+    private final Set<String>                            entitiesToSkipUpdate = new HashSet<>();
 
     private String       user;
     private Set<String>  userGroups;
@@ -108,6 +109,7 @@ public class RequestContext {
         this.entityExtInfoCache.clear();
         this.addedPropagations.clear();
         this.removedPropagations.clear();
+        this.entitiesToSkipUpdate.clear();
 
         if (metrics != null && !metrics.isEmpty()) {
             METRICS.debug(metrics.toString());
@@ -199,17 +201,18 @@ public class RequestContext {
     }
 
     public void recordEntityUpdate(AtlasEntityHeader entity) {
-        if (entity != null && entity.getGuid() != null) {
+        if (entity != null && entity.getGuid() != null && ! entitiesToSkipUpdate.contains(entity.getGuid())) {
             updatedEntities.put(entity.getGuid(), entity);
         }
     }
 
-    public void removeEntityUpdate(AtlasEntity entity) {
-        if (entity != null && entity.getGuid() != null) {
-            updatedEntities.remove(entity.getGuid());
+    public void recordEntityToSkip(String guid) {
+        if(! StringUtils.isEmpty(guid)) {
+            entitiesToSkipUpdate.add(guid);
         }
     }
 
+
     public void recordEntityDelete(AtlasEntityHeader entity) {
         if (entity != null && entity.getGuid() != null) {
             deletedEntities.put(entity.getGuid(), entity);