diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java index b470da3c7c74a..de77ff9444c6e 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java +++ b/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>> loaderSuppliers( + public static Map>> loaderSuppliers( final Collection> loadableTypes) { return loadableTypes.stream() .collect( @@ -1135,16 +1135,16 @@ private DataFetcher getEntityResolver() { }); } - private DataFetcher getResolver(LoadableType loadableType) { - return getResolver(loadableType, this::getUrnField); + private static DataFetcher getResolver(LoadableType loadableType) { + return getResolver(loadableType, GmsGraphQLEngine::getUrnField); } - private DataFetcher getResolver( + private static DataFetcher getResolver( LoadableType loadableType, Function 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 DataLoader> createDataLoader( + private static DataLoader> createDataLoader( final LoadableType graphType, final QueryContext queryContext) { BatchLoaderContextProvider contextProvider = () -> queryContext; DataLoaderOptions loaderOptions = diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java index dd8eabd3ce06f..97e282c106a9a 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java +++ b/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>> dataLoaderSuppliers, - final QueryContext context) { - final DataLoaderRegistry registry = new DataLoaderRegistry(); - for (String key : dataLoaderSuppliers.keySet()) { - registry.register(key, dataLoaderSuppliers.get(key).apply(context)); - } - return registry; - } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/LazyDataLoaderRegistry.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/LazyDataLoaderRegistry.java new file mode 100644 index 0000000000000..1a1d2d5f71f33 --- /dev/null +++ b/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>> dataLoaderSuppliers; + + public LazyDataLoaderRegistry( + QueryContext queryContext, + Map>> dataLoaderSuppliers) { + super(); + this.queryContext = queryContext; + this.dataLoaderSuppliers = new ConcurrentHashMap<>(dataLoaderSuppliers); + } + + @Override + public DataLoader getDataLoader(String key) { + return super.computeIfAbsent( + key, + k -> { + Function> supplier = dataLoaderSuppliers.get(key); + if (supplier == null) { + throw new IllegalArgumentException("No DataLoader registered for key: " + key); + } + return supplier.apply(queryContext); + }); + } + + @Override + public Set getKeys() { + return Stream.concat(dataLoaders.keySet().stream(), dataLoaderSuppliers.keySet().stream()) + .collect(Collectors.toSet()); + } + + @Override + public DataLoaderRegistry combine(DataLoaderRegistry registry) { + throw new UnsupportedOperationException(); + } +} diff --git a/metadata-service/configuration/src/main/resources/application.yaml b/metadata-service/configuration/src/main/resources/application.yaml index 3f1c6dd1a3d7d..37399466c50b0 100644 --- a/metadata-service/configuration/src/main/resources/application.yaml +++ b/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 diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java index 61c641e358f09..622cf20af9ff5 100644 --- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/config/SpringWebConfig.java +++ b/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 OPENLINEAGE_PACKAGES = Set.of("io.datahubproject.openapi.openlineage"); + @Value("${datahub.gms.async.request-timeout-ms}") + private long asyncTimeoutMilliseconds; + @Override public void configureMessageConverters(List> messageConverters) { messageConverters.add(new StringHttpMessageConverter()); @@ -158,4 +164,10 @@ private Map concat(Supplier> a, Supplier> b) { (v1, v2) -> v2, LinkedHashMap::new)); } + + @Override + public void configureAsyncSupport(@Nonnull AsyncSupportConfigurer configurer) { + WebMvcConfigurer.super.configureAsyncSupport(configurer); + configurer.setDefaultTimeout(asyncTimeoutMilliseconds); + } }