ATLAS-2856: added utility methods to RequestContext, to find number of active…

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2Test.java

@@ -19,6 +19,7 @@ package org.apache.atlas.repository.store.graph.v2;
 
 import com.google.common.collect.ImmutableSet;
 import org.apache.atlas.AtlasErrorCode;
+import org.apache.atlas.RequestContext;
 import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtilsV2;
 import org.apache.atlas.exception.AtlasBaseException;
@@ -37,15 +38,17 @@ import org.apache.atlas.model.typedef.AtlasEntityDef;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.type.AtlasTypeUtil;
 import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
 import javax.inject.Inject;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -63,6 +66,8 @@ import static org.testng.Assert.fail;
 
 @Guice(modules = TestModules.TestOnlyModule.class)
 public class AtlasEntityStoreV2Test extends AtlasEntityTestBase {
+    private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityStoreV2Test.class);
+
     private AtlasEntitiesWithExtInfo deptEntity;
     private AtlasEntityWithExtInfo   dbEntity;
     private AtlasEntityWithExtInfo   tblEntity;
@@ -102,6 +107,14 @@ public class AtlasEntityStoreV2Test extends AtlasEntityTestBase {
 
         typeDefStore.createTypesDef(typesDef11);
     }
+    @BeforeTest
+    public void init() throws Exception {
+        entityStore = new AtlasEntityStoreV2(deleteHandler, typeRegistry, mockChangeNotifier, graphMapper);
+        RequestContext.clear();
+        RequestContext.get().setUser(TestUtilsV2.TEST_USER, null);
+
+        LOG.debug("RequestContext: activeCount={}, earliestActiveRequestTime={}", RequestContext.getActiveRequestsCount(), RequestContext.earliestActiveRequestTime());
+    }
 
     @Test
     public void testDefaultValueForPrimitiveTypes() throws Exception  {

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

@@ -32,6 +32,7 @@ public class RequestContext {
     private static final Logger LOG = LoggerFactory.getLogger(RequestContext.class);
 
     private static final ThreadLocal<RequestContext> CURRENT_CONTEXT = new ThreadLocal<>();
+    private static final Set<RequestContext>         ACTIVE_REQUESTS = new HashSet<>();
 
     private final Map<String, AtlasObjectId>             updatedEntities     = new HashMap<>();
     private final Map<String, AtlasObjectId>             deletedEntities     = new HashMap<>();
@@ -60,6 +61,10 @@ public class RequestContext {
         if (ret == null) {
             ret = new RequestContext();
             CURRENT_CONTEXT.set(ret);
+
+            synchronized (ACTIVE_REQUESTS) {
+                ACTIVE_REQUESTS.add(ret);
+            }
         }
 
         return ret;
@@ -79,6 +84,10 @@ public class RequestContext {
             if (instance.entityGuidInRequest != null) {
                 instance.entityGuidInRequest.clear();
             }
+
+            synchronized (ACTIVE_REQUESTS) {
+                ACTIVE_REQUESTS.remove(instance);
+            }
         }
 
         CURRENT_CONTEXT.remove();
@@ -148,6 +157,30 @@ public class RequestContext {
         }
     }
 
+    public static RequestContext createContext() {
+        clear();
+
+        return get();
+    }
+
+    public static int getActiveRequestsCount() {
+        return ACTIVE_REQUESTS.size();
+    }
+
+    public static long earliestActiveRequestTime() {
+        long ret = System.currentTimeMillis();
+
+        synchronized (ACTIVE_REQUESTS) {
+            for (RequestContext context : ACTIVE_REQUESTS) {
+                if (ret > context.getRequestTime()) {
+                    ret = context.getRequestTime();
+                }
+            }
+        }
+
+        return ret;
+    }
+
     public void recordRemovedPropagation(String guid, AtlasClassification classification) {
         if (StringUtils.isNotEmpty(guid) && classification != null) {
             List<AtlasClassification> classifications = removedPropagations.get(guid);