feat(graphql): custom complexity calculator and separate configurable thread pool for graphQL

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -22,6 +22,7 @@
 import com.linkedin.datahub.graphql.analytics.resolver.GetMetadataAnalyticsResolver;
 import com.linkedin.datahub.graphql.analytics.resolver.IsAnalyticsEnabledResolver;
 import com.linkedin.datahub.graphql.analytics.service.AnalyticsService;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
 import com.linkedin.datahub.graphql.generated.AccessToken;
 import com.linkedin.datahub.graphql.generated.AccessTokenMetadata;
@@ -394,7 +395,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -2900,7 +2900,7 @@ private <T, K> DataLoader<K, DataFetcherResult<T>> createDataLoader(
         DataLoaderOptions.newOptions().setBatchLoaderContextProvider(contextProvider);
     return DataLoader.newDataLoader(
         (keys, context) ->
-            CompletableFuture.supplyAsync(
+            GraphQLConcurrencyUtils.supplyAsync(
                 () -> {
                   try {
                     log.debug(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java

@@ -3,6 +3,7 @@
 import static graphql.schema.idl.RuntimeWiring.*;
 
 import com.linkedin.datahub.graphql.exception.DataHubDataFetcherExceptionHandler;
+import com.linkedin.datahub.graphql.instrumentation.DataHubFieldComplexityCalculator;
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import graphql.GraphQL;
@@ -80,7 +81,9 @@ private GraphQLEngine(
     List<Instrumentation> instrumentations = new ArrayList<>(3);
     instrumentations.add(new TracingInstrumentation());
     instrumentations.add(new MaxQueryDepthInstrumentation(graphQLQueryDepthLimit));
-    instrumentations.add(new MaxQueryComplexityInstrumentation(graphQLQueryComplexityLimit));
+    instrumentations.add(
+        new MaxQueryComplexityInstrumentation(
+            graphQLQueryComplexityLimit, new DataHubFieldComplexityCalculator()));
     ChainedInstrumentation chainedInstrumentation = new ChainedInstrumentation(instrumentations);
     _graphQL =
         new GraphQL.Builder(graphQLSchema)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/SubTypesResolver.java

@@ -2,6 +2,7 @@
 
 import com.linkedin.common.SubTypes;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.generated.Entity;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.client.EntityClient;
@@ -26,7 +27,7 @@ public class SubTypesResolver implements DataFetcher<CompletableFuture<SubTypes>
   @Override
   @Nullable
   public CompletableFuture<SubTypes> get(DataFetchingEnvironment environment) throws Exception {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
           SubTypes subType = null;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/WeaklyTypedAspectsResolver.java

@@ -5,6 +5,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.DataMap;
 import com.linkedin.data.codec.JacksonDataCodec;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.generated.AspectParams;
 import com.linkedin.datahub.graphql.generated.AspectRenderSpec;
 import com.linkedin.datahub.graphql.generated.Entity;
@@ -48,7 +49,7 @@ private boolean shouldReturnAspect(AspectSpec aspectSpec, AspectParams params) {
   @Override
   public CompletableFuture<List<RawAspect>> get(DataFetchingEnvironment environment)
       throws Exception {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           List<RawAspect> results = new ArrayList<>();
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/concurrency/GraphQLConcurrencyUtils.java

@@ -0,0 +1,28 @@
+package com.linkedin.datahub.graphql.concurrency;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.function.Supplier;
+
+public class GraphQLConcurrencyUtils {
+  private GraphQLConcurrencyUtils() {}
+
+  private static ExecutorService graphQLExecutorService = null;
+
+  public static ExecutorService getExecutorService() {
+    return GraphQLConcurrencyUtils.graphQLExecutorService;
+  }
+
+  public static void setExecutorService(ExecutorService executorService) {
+    GraphQLConcurrencyUtils.graphQLExecutorService = executorService;
+  }
+
+  public static <T> CompletableFuture<T> supplyAsync(Supplier<T> supplier) {
+    if (GraphQLConcurrencyUtils.graphQLExecutorService == null) {
+      return CompletableFuture.supplyAsync(supplier);
+    } else {
+      return CompletableFuture.supplyAsync(
+          supplier, GraphQLConcurrencyUtils.graphQLExecutorService);
+    }
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/concurrency/GraphQLWorkerPoolThreadFactory.java

@@ -0,0 +1,25 @@
+package com.linkedin.datahub.graphql.concurrency;
+
+import java.util.concurrent.ThreadFactory;
+
+public class GraphQLWorkerPoolThreadFactory implements ThreadFactory {
+
+  private static int threadInitNumber;
+
+  private static synchronized int nextThreadNum() {
+    return threadInitNumber++;
+  }
+
+  private long stackSize;
+
+  public GraphQLWorkerPoolThreadFactory(long stackSize) {
+    this.stackSize = stackSize;
+  }
+
+  @Override
+  public final Thread newThread(Runnable runnable) {
+    Thread thread = new Thread(null, runnable, "GraphQLWorkerThread-" + nextThreadNum(), stackSize);
+
+    return thread;
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/instrumentation/DataHubFieldComplexityCalculator.java

@@ -0,0 +1,69 @@
+package com.linkedin.datahub.graphql.instrumentation;
+
+import graphql.analysis.FieldComplexityCalculator;
+import graphql.analysis.FieldComplexityEnvironment;
+import graphql.language.Field;
+import graphql.language.FragmentSpread;
+import graphql.language.Selection;
+import graphql.language.SelectionSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class DataHubFieldComplexityCalculator implements FieldComplexityCalculator {
+
+  private static final String COUNT_ARG = "count";
+  private static final String INPUT_ARG = "input";
+  private static final String SEARCH_RESULTS_FIELD = "searchResults";
+  private static final String ENTITY_FIELD = "entity";
+  private static final String SEARCH_RESULT_FIELDS_FIELD = "searchResultFields";
+  private static final String GRAPHQL_QUERY_TYPE = "Query";
+
+  @SuppressWarnings("rawtypes")
+  @Override
+  public int calculate(FieldComplexityEnvironment environment, int childComplexity) {
+    int complexity = 1;
+    Map<String, Object> args = environment.getArguments();
+    if (args.containsKey(INPUT_ARG)) {
+      Map<String, Object> input = (Map<String, Object>) args.get(INPUT_ARG);
+      if (input.containsKey(COUNT_ARG) && (Integer) input.get(COUNT_ARG) > 1) {
+        Integer count = (Integer) input.get(COUNT_ARG);
+        Field field = environment.getField();
+        List<Selection> subFields = field.getSelectionSet().getSelections();
+        Optional<FragmentSpread> searchResultsFieldsField =
+            subFields.stream()
+                .filter(selection -> selection instanceof Field)
+                .map(selection -> (Field) selection)
+                .filter(subField -> SEARCH_RESULTS_FIELD.equals(subField.getName()))
+                .map(Field::getSelectionSet)
+                .map(SelectionSet::getSelections)
+                .flatMap(List::stream)
+                .filter(selection -> selection instanceof Field)
+                .map(selection -> (Field) selection)
+                .filter(subField -> ENTITY_FIELD.equals(subField.getName()))
+                .map(Field::getSelectionSet)
+                .map(SelectionSet::getSelections)
+                .flatMap(List::stream)
+                .filter(selection -> selection instanceof FragmentSpread)
+                .map(selection -> (FragmentSpread) selection)
+                .filter(subField -> SEARCH_RESULT_FIELDS_FIELD.equals(subField.getName()))
+                .findFirst();
+        if (searchResultsFieldsField.isPresent()) {
+          // This fragment includes 2 lineage queries, we account for this additional complexity by
+          // multiplying
+          // by the count of entities attempting to be returned
+          complexity += 2 * count;
+        }
+      }
+    }
+    if (GRAPHQL_QUERY_TYPE.equals(environment.getParentType().getName())) {
+      log.info(
+          "Query complexity for query: {} is {}",
+          environment.getField().getName(),
+          complexity + childComplexity);
+    }
+    return complexity + childComplexity;
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/MeResolver.java

@@ -11,6 +11,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
 import com.linkedin.datahub.graphql.generated.AuthenticatedUser;
 import com.linkedin.datahub.graphql.generated.CorpUser;
@@ -49,7 +50,7 @@ public MeResolver(final EntityClient entityClient, final FeatureFlags featureFla
   @Override
   public CompletableFuture<AuthenticatedUser> get(DataFetchingEnvironment environment) {
     final QueryContext context = environment.getContext();
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           try {
             // 1. Get currently logged in user profile.

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/assertion/AssertionRunEventResolver.java

@@ -4,6 +4,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.generated.Assertion;
 import com.linkedin.datahub.graphql.generated.AssertionResultType;
 import com.linkedin.datahub.graphql.generated.AssertionRunEvent;
@@ -40,7 +41,7 @@ public AssertionRunEventResolver(final EntityClient client) {
 
   @Override
   public CompletableFuture<AssertionRunEventsResult> get(DataFetchingEnvironment environment) {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/assertion/DeleteAssertionResolver.java

@@ -9,6 +9,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
@@ -38,7 +39,7 @@ public CompletableFuture<Boolean> get(final DataFetchingEnvironment environment)
       throws Exception {
     final QueryContext context = environment.getContext();
     final Urn assertionUrn = Urn.createFromString(environment.getArgument("urn"));
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
 
           // 1. check the entity exists. If not, return false.

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/assertion/EntityAssertionsResolver.java

@@ -6,6 +6,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.generated.Assertion;
 import com.linkedin.datahub.graphql.generated.Entity;
 import com.linkedin.datahub.graphql.generated.EntityAssertionsResult;
@@ -45,7 +46,7 @@ public EntityAssertionsResolver(final EntityClient entityClient, final GraphClie
 
   @Override
   public CompletableFuture<EntityAssertionsResult> get(DataFetchingEnvironment environment) {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/CreateAccessTokenResolver.java

@@ -10,6 +10,7 @@
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.AccessToken;
 import com.linkedin.datahub.graphql.generated.AccessTokenMetadata;
@@ -37,7 +38,7 @@ public CreateAccessTokenResolver(final StatefulTokenService statefulTokenService
   @Override
   public CompletableFuture<AccessToken> get(final DataFetchingEnvironment environment)
       throws Exception {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
           final CreateAccessTokenInput input =

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/DebugAccessResolver.java

@@ -9,6 +9,7 @@
 import com.linkedin.data.template.StringArray;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.DebugAccessResult;
 import com.linkedin.entity.EntityResponse;
@@ -54,7 +55,7 @@ public DebugAccessResolver(EntityClient entityClient, GraphClient graphClient) {
   @Override
   public CompletableFuture<DebugAccessResult> get(DataFetchingEnvironment environment)
       throws Exception {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/GetAccessTokenMetadataResolver.java

@@ -4,6 +4,7 @@
 import com.google.common.collect.ImmutableList;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.AccessTokenMetadata;
 import com.linkedin.datahub.graphql.types.auth.AccessTokenMetadataType;
@@ -31,7 +32,7 @@ public GetAccessTokenMetadataResolver(
   @Override
   public CompletableFuture<AccessTokenMetadata> get(final DataFetchingEnvironment environment)
       throws Exception {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
           final String token = environment.getArgument("token");

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/GetAccessTokenResolver.java

@@ -9,6 +9,7 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.AccessToken;
 import com.linkedin.datahub.graphql.generated.AccessTokenType;
@@ -33,7 +34,7 @@ public GetAccessTokenResolver(final StatelessTokenService tokenService) {
   @Override
   public CompletableFuture<AccessToken> get(final DataFetchingEnvironment environment)
       throws Exception {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
           final GetAccessTokenInput input =

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolver.java

@@ -5,6 +5,7 @@
 import com.google.common.collect.ImmutableList;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.AccessTokenMetadata;
 import com.linkedin.datahub.graphql.generated.EntityType;
@@ -40,7 +41,7 @@ public ListAccessTokensResolver(final EntityClient entityClient) {
   @Override
   public CompletableFuture<ListAccessTokenResult> get(DataFetchingEnvironment environment)
       throws Exception {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
           final ListAccessTokenInput input =

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/RevokeAccessTokenResolver.java

@@ -9,6 +9,7 @@
 import com.linkedin.data.DataMap;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.client.EntityClient;
@@ -35,7 +36,7 @@ public RevokeAccessTokenResolver(
 
   @Override
   public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throws Exception {
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           final QueryContext context = environment.getContext();
           final String tokenId = bindArgument(environment.getArgument("tokenId"), String.class);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/browse/BrowsePathsResolver.java

@@ -2,6 +2,7 @@
 
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
 
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.generated.BrowsePath;
 import com.linkedin.datahub.graphql.generated.BrowsePathsInput;
 import com.linkedin.datahub.graphql.generated.EntityType;
@@ -35,7 +36,7 @@ public CompletableFuture<List<BrowsePath>> get(DataFetchingEnvironment environme
     final BrowsePathsInput input =
         bindArgument(environment.getArgument("input"), BrowsePathsInput.class);
 
-    return CompletableFuture.supplyAsync(
+    return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           try {
             _logger.debug(