feat(graphql): Lazy dataLoaders (#11293)

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

@@ -724,7 +724,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
    * Returns a {@link Supplier} responsible for creating a new {@link DataLoader} from a {@link
    * LoadableType}.
    */
-  public Map<String, Function<QueryContext, DataLoader<?, ?>>> loaderSuppliers(
+  public static Map<String, Function<QueryContext, DataLoader<?, ?>>> loaderSuppliers(
       final Collection<? extends LoadableType<?, ?>> loadableTypes) {
     return loadableTypes.stream()
         .collect(
@@ -1135,16 +1135,16 @@ private DataFetcher getEntityResolver() {
         });
   }
 
-  private DataFetcher getResolver(LoadableType<?, String> loadableType) {
-    return getResolver(loadableType, this::getUrnField);
+  private static DataFetcher getResolver(LoadableType<?, String> loadableType) {
+    return getResolver(loadableType, GmsGraphQLEngine::getUrnField);
   }
 
-  private <T, K> DataFetcher getResolver(
+  private static <T, K> DataFetcher getResolver(
       LoadableType<T, K> loadableType, Function<DataFetchingEnvironment, K> keyProvider) {
     return new LoadableTypeResolver<>(loadableType, keyProvider);
   }
 
-  private String getUrnField(DataFetchingEnvironment env) {
+  private static String getUrnField(DataFetchingEnvironment env) {
     return env.getArgument(URN_FIELD_NAME);
   }
 
@@ -3025,7 +3025,7 @@ private void configureTestResultResolvers(final RuntimeWiring.Builder builder) {
                     })));
   }
 
-  private <T, K> DataLoader<K, DataFetcherResult<T>> createDataLoader(
+  private static <T, K> DataLoader<K, DataFetcherResult<T>> createDataLoader(
       final LoadableType<T, K> graphType, final QueryContext queryContext) {
     BatchLoaderContextProvider contextProvider = () -> queryContext;
     DataLoaderOptions loaderOptions =

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

@@ -28,7 +28,6 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.dataloader.DataLoader;
-import org.dataloader.DataLoaderRegistry;
 
 /**
  * Simple wrapper around a {@link GraphQL} instance providing APIs for building an engine and
@@ -100,7 +99,7 @@ public ExecutionResult execute(
     /*
      * Init DataLoaderRegistry - should be created for each request.
      */
-    DataLoaderRegistry register = createDataLoaderRegistry(_dataLoaderSuppliers, context);
+    LazyDataLoaderRegistry register = new LazyDataLoaderRegistry(context, _dataLoaderSuppliers);
 
     /*
      * Construct execution input
@@ -218,14 +217,4 @@ public GraphQLEngine build() {
           graphQLQueryIntrospectionEnabled);
     }
   }
-
-  private DataLoaderRegistry createDataLoaderRegistry(
-      final Map<String, Function<QueryContext, DataLoader<?, ?>>> dataLoaderSuppliers,
-      final QueryContext context) {
-    final DataLoaderRegistry registry = new DataLoaderRegistry();
-    for (String key : dataLoaderSuppliers.keySet()) {
-      registry.register(key, dataLoaderSuppliers.get(key).apply(context));
-    }
-    return registry;
-  }
 }

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

@@ -0,0 +1,53 @@
+package com.linkedin.datahub.graphql;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import lombok.extern.slf4j.Slf4j;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderRegistry;
+
+/**
+ * The purpose of this class is to avoid loading 42+ dataLoaders when many of the graphql queries do
+ * not use all of them.
+ */
+@Slf4j
+public class LazyDataLoaderRegistry extends DataLoaderRegistry {
+  private final QueryContext queryContext;
+  private final Map<String, Function<QueryContext, DataLoader<?, ?>>> dataLoaderSuppliers;
+
+  public LazyDataLoaderRegistry(
+      QueryContext queryContext,
+      Map<String, Function<QueryContext, DataLoader<?, ?>>> dataLoaderSuppliers) {
+    super();
+    this.queryContext = queryContext;
+    this.dataLoaderSuppliers = new ConcurrentHashMap<>(dataLoaderSuppliers);
+  }
+
+  @Override
+  public <K, V> DataLoader<K, V> getDataLoader(String key) {
+    return super.computeIfAbsent(
+        key,
+        k -> {
+          Function<QueryContext, DataLoader<?, ?>> supplier = dataLoaderSuppliers.get(key);
+          if (supplier == null) {
+            throw new IllegalArgumentException("No DataLoader registered for key: " + key);
+          }
+          return supplier.apply(queryContext);
+        });
+  }
+
+  @Override
+  public Set<String> getKeys() {
+    return Stream.concat(dataLoaders.keySet().stream(), dataLoaderSuppliers.keySet().stream())
+        .collect(Collectors.toSet());
+  }
+
+  @Override
+  public DataLoaderRegistry combine(DataLoaderRegistry registry) {
+    throw new UnsupportedOperationException();
+  }
+}

metadata-service/configuration/src/main/resources/application.yaml

@@ -72,6 +72,8 @@ datahub:
     host: ${DATAHUB_GMS_HOST:localhost}
     port: ${DATAHUB_GMS_PORT:8080}
     useSSL: ${DATAHUB_GMS_USE_SSL:${GMS_USE_SSL:false}}
+    async:
+      request-timeout-ms: ${DATAHUB_GMS_ASYNC_REQUEST_TIMEOUT_MS:55000}
 
     # URI instead of above host/port/ssl
     # Priority is given to the URI setting over separate host/port/useSSL parameters

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java

@@ -23,7 +23,9 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import javax.annotation.Nonnull;
 import org.springdoc.core.models.GroupedOpenApi;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.format.FormatterRegistry;
@@ -32,6 +34,7 @@
 import org.springframework.http.converter.HttpMessageConverter;
 import org.springframework.http.converter.StringHttpMessageConverter;
 import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter;
+import org.springframework.web.servlet.config.annotation.AsyncSupportConfigurer;
 import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
 
 @OpenAPIDefinition(
@@ -48,6 +51,9 @@ public class SpringWebConfig implements WebMvcConfigurer {
   private static final Set<String> OPENLINEAGE_PACKAGES =
       Set.of("io.datahubproject.openapi.openlineage");
 
+  @Value("${datahub.gms.async.request-timeout-ms}")
+  private long asyncTimeoutMilliseconds;
+
   @Override
   public void configureMessageConverters(List<HttpMessageConverter<?>> messageConverters) {
     messageConverters.add(new StringHttpMessageConverter());
@@ -158,4 +164,10 @@ private <K, V> Map<K, V> concat(Supplier<Map<K, V>> a, Supplier<Map<K, V>> b) {
                         (v1, v2) -> v2,
                         LinkedHashMap::new));
   }
+
+  @Override
+  public void configureAsyncSupport(@Nonnull AsyncSupportConfigurer configurer) {
+    WebMvcConfigurer.super.configureAsyncSupport(configurer);
+    configurer.setDefaultTimeout(asyncTimeoutMilliseconds);
+  }
 }