Switch to io.micrometer:context-propagation library

See gh-459
This commit is contained in:
rstoyanchev
2022-08-29 21:14:55 +01:00
parent f583eb302c
commit f680892dfc
32 changed files with 413 additions and 471 deletions

View File

@@ -59,7 +59,7 @@ configure(moduleProjects) {
dependencyManagement {
imports {
mavenBom "com.fasterxml.jackson:jackson-bom:2.13.3"
mavenBom "io.projectreactor:reactor-bom:2022.0.0-M4"
mavenBom "io.projectreactor:reactor-bom:2022.0.0-M5"
mavenBom "org.springframework:spring-framework-bom:6.0.0-M5"
mavenBom "org.springframework.data:spring-data-bom:2022.0.0-M5"
mavenBom "org.springframework.security:spring-security-bom:6.0.0-M6"
@@ -76,6 +76,7 @@ configure(moduleProjects) {
dependency "jakarta.persistence:jakarta.persistence-api:3.0.0"
dependency "jakarta.servlet:jakarta.servlet-api:5.0.0"
dependency "com.google.code.findbugs:jsr305:3.0.2"
dependency "io.micrometer:context-propagation:1.0.0-SNAPSHOT"
dependency "org.assertj:assertj-core:3.23.1"
dependency "com.jayway.jsonpath:json-path:2.7.0"
dependency "org.skyscreamer:jsonassert:1.5.0"

View File

@@ -31,6 +31,7 @@ dependencies {
testImplementation 'io.projectreactor:reactor-test'
testImplementation 'io.projectreactor.netty:reactor-netty'
testImplementation 'io.rsocket:rsocket-transport-local'
testImplementation 'io.micrometer:context-propagation'
testImplementation 'com.squareup.okhttp3:mockwebserver:3.14.9'
testImplementation 'com.fasterxml.jackson.core:jackson-databind'

View File

@@ -12,6 +12,7 @@ dependencies {
compileOnly 'org.springframework:spring-webmvc'
compileOnly 'org.springframework:spring-websocket'
compileOnly 'org.springframework:spring-messaging'
compileOnly 'io.micrometer:context-propagation'
compileOnly 'jakarta.servlet:jakarta.servlet-api'
compileOnly 'jakarta.validation:jakarta.validation-api'
@@ -41,6 +42,7 @@ dependencies {
testImplementation 'org.springframework.data:spring-data-commons'
testImplementation 'org.springframework.data:spring-data-keyvalue'
testImplementation 'org.springframework.data:spring-data-jpa'
testImplementation 'io.micrometer:context-propagation'
testImplementation 'com.h2database:h2'
testImplementation 'org.hibernate:hibernate-core-jakarta'
testImplementation 'org.hibernate.validator:hibernate-validator'

View File

@@ -26,11 +26,11 @@ import java.util.stream.Collectors;
import java.util.stream.Stream;
import graphql.GraphQLContext;
import io.micrometer.context.ContextSnapshot;
import reactor.core.publisher.Mono;
import org.springframework.core.CoroutinesUtils;
import org.springframework.core.KotlinDetector;
import org.springframework.graphql.execution.ReactorContextManager;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -114,7 +114,7 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
return CompletableFuture.supplyAsync(
() -> {
try {
return ReactorContextManager.invokeCallable((Callable<?>) result, graphQLContext);
return ContextSnapshot.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
}
catch (Exception ex) {
throw new IllegalStateException(

View File

@@ -28,6 +28,7 @@ import java.util.Map;
* @author Rossen Stoyanchev
* @since 1.0.0
*/
@SuppressWarnings("deprecation")
class CompositeThreadLocalAccessor implements ThreadLocalAccessor {
private final List<ThreadLocalAccessor> accessors;

View File

@@ -29,10 +29,10 @@ import graphql.schema.GraphQLTypeVisitor;
import graphql.schema.GraphQLTypeVisitorStub;
import graphql.util.TraversalControl;
import graphql.util.TraverserContext;
import io.micrometer.context.ContextSnapshot;
import org.reactivestreams.Publisher;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.context.ContextView;
import org.springframework.util.Assert;
@@ -70,29 +70,23 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
@Override
public Object get(DataFetchingEnvironment environment) throws Exception {
Object value = ReactorContextManager.invokeCallable(() ->
this.delegate.get(environment), environment.getGraphQlContext());
ContextView contextView = ReactorContextManager.getReactorContext(environment.getGraphQlContext());
ContextSnapshot snapshot = ContextSnapshot.captureFrom(environment.getGraphQlContext());
Object value = snapshot.wrap(() -> this.delegate.get(environment)).call();
if (this.subscription) {
Assert.state(value instanceof Publisher, "Expected Publisher for a subscription");
Flux<?> flux = Flux.from((Publisher<?>) value).onErrorResume(exception ->
this.subscriptionExceptionResolver.resolveException(exception)
.flatMap(errors -> Mono.error(new SubscriptionPublisherException(errors, exception))));
return (!contextView.isEmpty() ? flux.contextWrite(contextView) : flux);
return flux.contextWrite(snapshot::updateContext);
}
if (value instanceof Flux) {
value = ((Flux<?>) value).collectList();
}
if (value instanceof Mono) {
Mono<?> valueMono = (Mono<?>) value;
if (!contextView.isEmpty()) {
valueMono = valueMono.contextWrite(contextView);
}
value = valueMono.toFuture();
if (value instanceof Mono<?> valueMono) {
value = valueMono.contextWrite(snapshot::updateContext).toFuture();
}
return value;

View File

@@ -21,8 +21,11 @@ import java.util.function.BiFunction;
import graphql.GraphQLError;
import graphql.schema.DataFetchingEnvironment;
import io.micrometer.context.ContextSnapshot;
import io.micrometer.context.ThreadLocalAccessor;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import reactor.core.publisher.Mono;
import reactor.util.context.ContextView;
import org.springframework.lang.Nullable;
@@ -47,6 +50,8 @@ import org.springframework.lang.Nullable;
*/
public abstract class DataFetcherExceptionResolverAdapter implements DataFetcherExceptionResolver {
protected final Log logger = LogFactory.getLog(getClass());
private boolean threadLocalContextAware;
@@ -88,17 +93,18 @@ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcher
}
@Nullable
private List<GraphQLError> resolveInternal(Throwable ex, DataFetchingEnvironment env) {
private List<GraphQLError> resolveInternal(Throwable exception, DataFetchingEnvironment env) {
if (!this.threadLocalContextAware) {
return resolveToMultipleErrors(ex, env);
return resolveToMultipleErrors(exception, env);
}
ContextView contextView = ReactorContextManager.getReactorContext(env.getGraphQlContext());
try {
ReactorContextManager.restoreThreadLocalValues(contextView);
return resolveToMultipleErrors(ex, env);
return ContextSnapshot.captureFrom(env.getGraphQlContext())
.wrap(() -> resolveToMultipleErrors(exception, env))
.call();
}
finally {
ReactorContextManager.resetThreadLocalValues(contextView);
catch (Exception ex2) {
logger.warn("Failed to resolve " + exception, ex2);
return null;
}
}

View File

@@ -19,11 +19,13 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import graphql.GraphQLContext;
import io.micrometer.context.ContextSnapshot;
import org.dataloader.BatchLoaderContextProvider;
import org.dataloader.BatchLoaderEnvironment;
import org.dataloader.BatchLoaderWithContext;
@@ -34,7 +36,6 @@ import org.dataloader.DataLoaderRegistry;
import org.dataloader.MappedBatchLoaderWithContext;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.context.ContextView;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -188,16 +189,20 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
@Override
public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
ContextView contextView = ReactorContextManager.getReactorContext(environment.getContext());
GraphQLContext graphQLContext = environment.getContext();
ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
try {
ReactorContextManager.restoreThreadLocalValues(contextView);
return this.loader.apply(keys, environment).collectList().contextWrite(contextView).toFuture();
return snapshot.wrap(() ->
this.loader.apply(keys, environment)
.collectList()
.contextWrite(snapshot::updateContext)
.toFuture())
.call();
}
finally {
ReactorContextManager.resetThreadLocalValues(contextView);
catch (Exception ex) {
return CompletableFuture.failedFuture(ex);
}
}
}
@@ -239,13 +244,17 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
@Override
public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
ContextView contextView = ReactorContextManager.getReactorContext(environment.getContext());
GraphQLContext graphQLContext = environment.getContext();
ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
try {
ReactorContextManager.restoreThreadLocalValues(contextView);
return this.loader.apply(keys, environment).contextWrite(contextView).toFuture();
return snapshot.wrap(() ->
this.loader.apply(keys, environment)
.contextWrite(snapshot::updateContext)
.toFuture())
.call();
}
finally {
ReactorContextManager.resetThreadLocalValues(contextView);
catch (Exception ex) {
return CompletableFuture.failedFuture(ex);
}
}

View File

@@ -24,6 +24,7 @@ import graphql.ExecutionInput;
import graphql.GraphQL;
import graphql.GraphQLContext;
import graphql.execution.ExecutionIdProvider;
import io.micrometer.context.ContextSnapshot;
import org.dataloader.DataLoaderRegistry;
import reactor.core.publisher.Mono;
@@ -76,7 +77,7 @@ public class DefaultExecutionGraphQlService implements ExecutionGraphQlService {
request.configureExecutionInput(RESET_EXECUTION_ID_CONFIGURER);
}
ExecutionInput executionInput = request.toExecutionInput();
ReactorContextManager.setReactorContext(contextView, executionInput.getGraphQLContext());
ContextSnapshot.captureFrom(contextView).updateContext(executionInput.getGraphQLContext());
ExecutionInput updatedExecutionInput = registerDataLoaders(executionInput);
return Mono.fromFuture(this.graphQlSource.graphQl().executeAsync(updatedExecutionInput))
.map(result -> new DefaultExecutionGraphQlResponse(updatedExecutionInput, result));

View File

@@ -28,11 +28,11 @@ import graphql.execution.DataFetcherExceptionHandlerParameters;
import graphql.execution.DataFetcherExceptionHandlerResult;
import graphql.execution.ExecutionId;
import graphql.schema.DataFetchingEnvironment;
import io.micrometer.context.ContextSnapshot;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.context.ContextView;
import org.springframework.util.Assert;
@@ -70,6 +70,7 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler
public CompletableFuture<DataFetcherExceptionHandlerResult> handleException(DataFetcherExceptionHandlerParameters params) {
Throwable exception = unwrapException(params);
DataFetchingEnvironment env = params.getDataFetchingEnvironment();
ContextSnapshot snapshot = ContextSnapshot.captureFrom(env.getGraphQlContext());
try {
return Flux.fromIterable(this.resolvers)
.flatMap(resolver -> resolver.resolveException(exception, env))
@@ -78,10 +79,7 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler
.doOnNext(result -> logResolvedException(exception, result))
.onErrorResume(resolverEx -> Mono.just(handleResolverError(resolverEx, exception, env)))
.switchIfEmpty(Mono.fromCallable(() -> createInternalError(exception, env)))
.contextWrite((context) -> {
ContextView contextView = ReactorContextManager.getReactorContext(env.getGraphQlContext());
return (contextView.isEmpty() ? context : context.putAll(contextView));
})
.contextWrite(snapshot::updateContext)
.toFuture();
}
catch (Exception resolverEx) {

View File

@@ -0,0 +1,64 @@
/*
* Copyright 2002-2022 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.graphql.execution;
import java.util.Map;
import java.util.function.Predicate;
import graphql.GraphQLContext;
import io.micrometer.context.ContextAccessor;
/**
* {@code ContextAccessor} that enables support for reading and writing values
* to and from a {@link GraphQLContext}. This accessor is automatically
* registered via {@link java.util.ServiceLoader}.
*
* @author Rossen Stoyanchev
* @since 1.1.0
*/
public class GraphQlContextAccessor implements ContextAccessor<GraphQLContext, GraphQLContext> {
@Override
public boolean canReadFrom(Class<?> contextType) {
return GraphQLContext.class.equals(contextType);
}
@Override
public void readValues(GraphQLContext context, Predicate<Object> keyPredicate, Map<Object, Object> readValues) {
context.stream().forEach(entry -> {
if (keyPredicate.test(entry.getKey())) {
readValues.put(entry.getKey(), entry.getValue());
}
});
}
@Override
public <T> T readValue(GraphQLContext context, Object key) {
return context.get(key);
}
@Override
public boolean canWriteTo(Class<?> contextType) {
return GraphQLContext.class.equals(contextType);
}
@Override
public GraphQLContext writeValues(Map<Object, Object> valuesToWrite, GraphQLContext targetContext) {
return targetContext.putAll(valuesToWrite);
}
}

View File

@@ -1,141 +0,0 @@
/*
* Copyright 2002-2022 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.graphql.execution;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.Callable;
import graphql.GraphQLContext;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
/**
* Provides helper methods to save Reactor context in the {@link GraphQLContext}
* so it can be subsequently obtained and propagated to data fetchers, exception
* handlers, and others.
*
* <p>The Reactor context is also used to carry ThreadLocal values that are also
* restored around the execution of data fetchers and exceptions handlers.
*
* @author Rossen Stoyanchev
* @since 1.0.0
*/
public abstract class ReactorContextManager {
private static final String CONTEXT_VIEW_KEY = ReactorContextManager.class.getName() + ".CONTEXT_VIEW";
private static final String THREAD_ID = ReactorContextManager.class.getName() + ".THREAD_ID";
private static final String THREAD_LOCAL_VALUES_KEY = ReactorContextManager.class.getName() + ".THREAD_VALUES_ACCESSOR";
private static final String THREAD_LOCAL_ACCESSOR_KEY = ReactorContextManager.class.getName() + ".THREAD_LOCAL_ACCESSOR";
/**
* Save the given Reactor {@link ContextView} in the given {@link GraphQLContext}.
* @param contextView the reactor {@code ContextView} to save
* @param graphQLContext the {@code GraphQLContext} where to save
*/
static void setReactorContext(ContextView contextView, GraphQLContext graphQLContext) {
graphQLContext.put(CONTEXT_VIEW_KEY, contextView);
}
/**
* Return the Reactor {@link ContextView} saved in the given {@link GraphQLContext}.
* @param graphQlContext the DataFetchingEnvironment
* @return the reactor {@link ContextView}
*/
static ContextView getReactorContext(GraphQLContext graphQlContext) {
Assert.notNull(graphQlContext, "GraphQLContext is required");
return graphQlContext.getOrDefault(CONTEXT_VIEW_KEY, Context.empty());
}
/**
* Use the given accessor to extract ThreadLocal values and save them in a
* sub-map in the given {@link Context}, so those can be restored later
* around the execution of data fetchers and exception resolvers. The accessor
* instance is also saved in the Reactor Context, so it can be used to
* actually restore and reset ThreadLocal values.
* @param accessor the accessor to use
* @param context the context to write to if there are ThreadLocal values
* @return a new Reactor {@link ContextView} or the {@code Context} instance
* that was passed in, if there were no ThreadLocal values to extract.
*/
public static Context extractThreadLocalValues(ThreadLocalAccessor accessor, Context context) {
Map<String, Object> valuesMap = new LinkedHashMap<>();
accessor.extractValues(valuesMap);
if (valuesMap.isEmpty()) {
return context;
}
return context.putAll((ContextView) Context.of(
THREAD_LOCAL_VALUES_KEY, valuesMap,
THREAD_LOCAL_ACCESSOR_KEY, accessor,
THREAD_ID, Thread.currentThread().getId()));
}
/**
* Restore {@code ThreadLocal} values, invoke the given {@code Callable},
* and reset the {@code ThreadLocal} values.
* @param callable the callable to invoke
* @param graphQlContext the current {@code GraphQLContext}
* @return the return value from the invocation
*/
public static <T> T invokeCallable(Callable<T> callable, GraphQLContext graphQlContext) throws Exception {
ContextView contextView = getReactorContext(graphQlContext);
try {
ReactorContextManager.restoreThreadLocalValues(contextView);
return callable.call();
}
finally {
ReactorContextManager.resetThreadLocalValues(contextView);
}
}
/**
* Look up saved ThreadLocal values and restore them if any are found.
* This is a no-op if invoked on the thread that values were extracted on.
* @param contextView the reactor {@link ContextView}
*/
static void restoreThreadLocalValues(ContextView contextView) {
ThreadLocalAccessor accessor = getThreadLocalAccessor(contextView);
if (accessor != null) {
accessor.restoreValues(contextView.get(THREAD_LOCAL_VALUES_KEY));
}
}
/**
* Look up saved ThreadLocal values and remove the ThreadLocal values.
* This is a no-op if invoked on the thread that values were extracted on.
* @param contextView the reactor {@link ContextView}
*/
static void resetThreadLocalValues(ContextView contextView) {
ThreadLocalAccessor accessor = getThreadLocalAccessor(contextView);
if (accessor != null) {
accessor.resetValues(contextView.get(THREAD_LOCAL_VALUES_KEY));
}
}
@Nullable
private static ThreadLocalAccessor getThreadLocalAccessor(ContextView view) {
Long id = view.getOrDefault(THREAD_ID, null);
return (id != null && id != Thread.currentThread().getId() ? view.get(THREAD_LOCAL_ACCESSOR_KEY) : null);
}
}

View File

@@ -17,30 +17,82 @@ package org.springframework.graphql.execution;
import java.util.Map;
import io.micrometer.context.ThreadLocalAccessor;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.util.ClassUtils;
/**
* {@link ThreadLocalAccessor} to extract and restore security context through
* {@link SecurityContextHolder}.
* {@link SecurityContextHolder}. This accessor is automatically registered via
* {@link java.util.ServiceLoader} but applies if Spring Security is present on
* the classpath.
*
* @author Rob Winch
* @author Rossen Stoyanchev
* @since 1.0.0
*/
public class SecurityContextThreadLocalAccessor implements ThreadLocalAccessor {
@SuppressWarnings("deprecation")
public class SecurityContextThreadLocalAccessor implements ThreadLocalAccessor<Object>,
org.springframework.graphql.execution.ThreadLocalAccessor {
private static final String KEY = SecurityContext.class.getName();
private final static boolean springSecurityPresent = ClassUtils.isPresent(
"org.springframework.security.core.context.SecurityContext",
SecurityContextThreadLocalAccessor.class.getClassLoader());
private final ThreadLocalAccessor<?> delegate;
public SecurityContextThreadLocalAccessor() {
if (springSecurityPresent) {
this.delegate = new DelegateAccessor();
}
else {
this.delegate = new NoOpAccessor();
}
}
@Override
public Object key() {
return this.delegate.key();
}
@Override
public Object getValue() {
return this.delegate.getValue();
}
@Override
public void setValue(Object value) {
setValueInternal(value);
}
@SuppressWarnings("unchecked")
private <V> void setValueInternal(Object value) {
((ThreadLocalAccessor<V>) this.delegate).setValue((V) value);
}
@Override
public void reset() {
this.delegate.reset();
}
// Temporary implementation of deprecated ThreadLocalAccessor while it is still used
// in the Boot starter. If registered as such, it is ignored.
@Override
public void extractValues(Map<String, Object> container) {
container.put(KEY, SecurityContextHolder.getContext());
container.put((String) key(), SecurityContextHolder.getContext());
}
@Override
public void restoreValues(Map<String, Object> values) {
if (values.containsKey(KEY)) {
SecurityContextHolder.setContext((SecurityContext) values.get(KEY));
if (values.containsKey((String) key())) {
SecurityContextHolder.setContext((SecurityContext) values.get((String) key()));
}
}
@@ -49,4 +101,52 @@ public class SecurityContextThreadLocalAccessor implements ThreadLocalAccessor {
SecurityContextHolder.clearContext();
}
private static class DelegateAccessor implements ThreadLocalAccessor<Object> {
@Override
public Object key() {
return SecurityContext.class.getName();
}
@Override
public Object getValue() {
return SecurityContextHolder.getContext();
}
@Override
public void setValue(Object value) {
SecurityContextHolder.setContext((SecurityContext) value);
}
@Override
public void reset() {
SecurityContextHolder.clearContext();
}
}
private static class NoOpAccessor implements ThreadLocalAccessor<Object> {
@Override
public Object key() {
return getClass().getName();
}
@Override
public Object getValue() {
return null;
}
@Override
public void setValue(Object value) {
}
@Override
public void reset() {
}
}
}

View File

@@ -21,6 +21,10 @@ import java.util.List;
import java.util.function.Function;
import graphql.GraphQLError;
import io.micrometer.context.ContextSnapshot;
import io.micrometer.context.ThreadLocalAccessor;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import reactor.core.publisher.Mono;
import org.springframework.lang.Nullable;
@@ -44,6 +48,8 @@ import org.springframework.lang.Nullable;
*/
public abstract class SubscriptionExceptionResolverAdapter implements SubscriptionExceptionResolver {
protected final Log logger = LogFactory.getLog(getClass());
private boolean threadLocalContextAware;
@@ -72,22 +78,25 @@ public abstract class SubscriptionExceptionResolverAdapter implements Subscripti
}
@SuppressWarnings({"unused", "try"})
@Override
public final Mono<List<GraphQLError>> resolveException(Throwable exception) {
if (!this.threadLocalContextAware) {
if (this.threadLocalContextAware) {
return Mono.deferContextual(contextView -> {
ContextSnapshot snapshot = ContextSnapshot.captureFrom(contextView);
try {
List<GraphQLError> errors = snapshot.wrap(() -> resolveToMultipleErrors(exception)).call();
return Mono.justOrEmpty(errors);
}
catch (Exception ex2) {
logger.warn("Failed to resolve " + exception, ex2);
return Mono.empty();
}
});
}
else {
return Mono.justOrEmpty(resolveToMultipleErrors(exception));
}
return Mono.deferContextual(contextView -> {
List<GraphQLError> errors;
try {
ReactorContextManager.restoreThreadLocalValues(contextView);
errors = resolveToMultipleErrors(exception);
}
finally {
ReactorContextManager.resetThreadLocalValues(contextView);
}
return Mono.justOrEmpty(errors);
});
}
/**

View File

@@ -37,8 +37,12 @@ import org.springframework.beans.factory.ObjectProvider;
*
* @author Rossen Stoyanchev
* @since 1.0.0
* @deprecated as of 1.1.0 in favor of using
* {@link io.micrometer.context.ThreadLocalAccessor} from the
* {@code "io.micrometer:context-propagation"} library.
* @see org.springframework.graphql.server.WebGraphQlHandler.Builder#threadLocalAccessor(ThreadLocalAccessor...)
*/
@Deprecated
public interface ThreadLocalAccessor {
/**

View File

@@ -20,11 +20,13 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import io.micrometer.context.ContextSnapshot;
import org.springframework.graphql.execution.SecurityContextThreadLocalAccessor;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import reactor.core.publisher.Mono;
import org.springframework.graphql.ExecutionGraphQlService;
import org.springframework.graphql.execution.ReactorContextManager;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.graphql.server.WebGraphQlInterceptor.Chain;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -36,6 +38,7 @@ import org.springframework.util.CollectionUtils;
*
* @author Rossen Stoyanchev
*/
@SuppressWarnings("deprecation")
class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder {
private final ExecutionGraphQlService service;
@@ -72,13 +75,23 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder {
return this;
}
@SuppressWarnings("deprecation")
@Override
public WebGraphQlHandler.Builder threadLocalAccessor(ThreadLocalAccessor... accessors) {
return threadLocalAccessors(Arrays.asList(accessors));
}
@SuppressWarnings("deprecation")
@Override
public WebGraphQlHandler.Builder threadLocalAccessors(List<ThreadLocalAccessor> accessors) {
// Filter out SecurityContextThreadLocalAccessor which is registered as a ThreadLocalAccessor
// from the micrometer-metrics/context-propagatation library. This code can be removed when
// SecurityContextThreadLocalAccessor no longer implements the deprecated ThreadLocalAccessor.
accessors = accessors.stream()
.filter(a -> !(a instanceof SecurityContextThreadLocalAccessor))
.toList();
if (!CollectionUtils.isEmpty(accessors)) {
this.accessors = (this.accessors != null) ? this.accessors : new ArrayList<>();
this.accessors.addAll(accessors);
@@ -96,9 +109,6 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder {
.map(interceptor -> interceptor.apply(endOfChain))
.orElse(endOfChain);
ThreadLocalAccessor accessor = (CollectionUtils.isEmpty(this.accessors) ? null :
ThreadLocalAccessor.composite(this.accessors));
return new WebGraphQlHandler() {
@Override
@@ -110,20 +120,14 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder {
@Nullable
@Override
public ThreadLocalAccessor getThreadLocalAccessor() {
return accessor;
return (CollectionUtils.isEmpty(accessors) ? null : ThreadLocalAccessor.composite(accessors));
}
@Override
public Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request) {
return executionChain.next(request)
.contextWrite(context -> {
if (accessor != null) {
return ReactorContextManager.extractThreadLocalValues(accessor, context);
}
return context;
});
ContextSnapshot snapshot = ContextSnapshot.capture();
return executionChain.next(request).contextWrite(snapshot::updateContext);
}
};
}

View File

@@ -44,7 +44,9 @@ public interface WebGraphQlHandler {
/**
* Return the composite {@link ThreadLocalAccessor} that the handler is
* configured with.
* @deprecated as of 1.1.0, together with {@link ThreadLocalAccessor}.
*/
@Deprecated
@Nullable
ThreadLocalAccessor getThreadLocalAccessor();
@@ -102,7 +104,9 @@ public interface WebGraphQlHandler {
* fetchers and exception resolvers.
* @param accessors the accessors to add
* @return this builder
* @deprecated as of 1.1.0 together with {@link ThreadLocalAccessor}.
*/
@Deprecated
Builder threadLocalAccessor(ThreadLocalAccessor... accessors);
/**
@@ -110,7 +114,9 @@ public interface WebGraphQlHandler {
* List.
* @param accessors the list of accessors to add
* @return this builder
* @deprecated as of 1.1.0 together with {@link ThreadLocalAccessor}.
*/
@Deprecated
Builder threadLocalAccessors(List<ThreadLocalAccessor> accessors);
/**

View File

@@ -18,7 +18,6 @@ package org.springframework.graphql.server.webmvc;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
@@ -28,7 +27,6 @@ import java.security.Principal;
import java.time.Duration;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@@ -37,6 +35,7 @@ import java.util.concurrent.atomic.AtomicReference;
import graphql.ExecutionResult;
import graphql.GraphQLError;
import graphql.GraphqlErrorBuilder;
import io.micrometer.context.ContextSnapshot;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.reactivestreams.Publisher;
@@ -49,7 +48,6 @@ import reactor.core.scheduler.Schedulers;
import org.springframework.graphql.execution.ErrorType;
import org.springframework.graphql.execution.SubscriptionPublisherException;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.graphql.server.WebGraphQlHandler;
import org.springframework.graphql.server.WebGraphQlResponse;
import org.springframework.graphql.server.WebSocketGraphQlInterceptor;
@@ -119,7 +117,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
Assert.notNull(converter, "HttpMessageConverter for JSON is required");
this.graphQlHandler = graphQlHandler;
this.contextHandshakeInterceptor = new ContextHandshakeInterceptor(graphQlHandler.getThreadLocalAccessor());
this.contextHandshakeInterceptor = new ContextHandshakeInterceptor();
this.webSocketGraphQlInterceptor = this.graphQlHandler.getWebSocketInterceptor();
this.initTimeoutDuration = connectionInitTimeout;
this.converter = converter;
@@ -134,7 +132,9 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
* Return a {@link WebSocketHttpRequestHandler} that uses this instance as
* its {@link WebGraphQlHandler} and adds a {@link HandshakeInterceptor} to
* propagate context.
* @deprecated as of 1.1.0 without a replacement, there should be no need for it
*/
@Deprecated
public WebSocketHttpRequestHandler asWebSocketHttpRequestHandler(HandshakeHandler handshakeHandler) {
WebSocketHttpRequestHandler handler = new WebSocketHttpRequestHandler(this, handshakeHandler);
handler.setHandshakeInterceptors(Collections.singletonList(this.contextHandshakeInterceptor));
@@ -169,7 +169,8 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
@SuppressWarnings({"unused", "try"})
@Override
protected void handleTextMessage(WebSocketSession session, TextMessage webSocketMessage) throws Exception {
try (Closeable closeable = this.contextHandshakeInterceptor.restoreThreadLocalValue(session)) {
ContextSnapshot snapshot = this.contextHandshakeInterceptor.getContextSnapshot(session);
try (AutoCloseable closeable = snapshot.setThreadLocalValues()) {
handleInternal(session, webSocketMessage);
}
}
@@ -344,25 +345,12 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
*/
private static class ContextHandshakeInterceptor implements HandshakeInterceptor {
private static final String SAVED_CONTEXT_KEY = ContextHandshakeInterceptor.class.getName();
@Nullable
private final ThreadLocalAccessor accessor;
ContextHandshakeInterceptor(@Nullable ThreadLocalAccessor accessor) {
this.accessor = accessor;
}
@Override
public boolean beforeHandshake(
ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
Map<String, Object> attributes) {
if (this.accessor != null) {
Map<String, Object> valuesMap = new LinkedHashMap<>();
this.accessor.extractValues(valuesMap);
attributes.put(SAVED_CONTEXT_KEY, valuesMap);
}
attributes.put(ContextSnapshot.class.getName(), ContextSnapshot.capture());
return true;
}
@@ -372,17 +360,11 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
@Nullable Exception exception) {
}
@SuppressWarnings("unchecked")
public Closeable restoreThreadLocalValue(WebSocketSession session) {
if (this.accessor != null) {
Map<String, Object> valuesMap = (Map<String, Object>) session.getAttributes().get(SAVED_CONTEXT_KEY);
Assert.state(valuesMap != null, "No ThreadLocal context in WebSocketSession attributes");
this.accessor.restoreValues(valuesMap);
return () -> this.accessor.resetValues(valuesMap);
}
return () -> {};
public ContextSnapshot getContextSnapshot(WebSocketSession session) {
ContextSnapshot snapshot = (ContextSnapshot) session.getAttributes().get(ContextSnapshot.class.getName());
Assert.notNull(snapshot, "No ContextSnapshot in WebSocketSession attributes");
return snapshot;
}
}

View File

@@ -0,0 +1 @@
org.springframework.graphql.execution.GraphQlContextAccessor

View File

@@ -0,0 +1 @@
org.springframework.graphql.execution.SecurityContextThreadLocalAccessor

View File

@@ -16,72 +16,58 @@
package org.springframework.graphql;
import java.util.Map;
import io.micrometer.context.ThreadLocalAccessor;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import static org.assertj.core.api.Assertions.assertThat;
/**
* {@link ThreadLocalAccessor} that operates on the ThreadLocal it is given.
*/
public class TestThreadLocalAccessor<T> implements ThreadLocalAccessor {
public class TestThreadLocalAccessor<T> implements ThreadLocalAccessor<T> {
private final ThreadLocal<T> threadLocal;
@Nullable
private Long threadId;
private final boolean suppressThreadIdCheck;
public TestThreadLocalAccessor(ThreadLocal<T> threadLocal) {
this(threadLocal, false);
}
public TestThreadLocalAccessor(ThreadLocal<T> threadLocal, boolean suppressThreadIdCheck) {
this.threadLocal = threadLocal;
this.suppressThreadIdCheck = suppressThreadIdCheck;
}
@Override
public Object key() {
return getClass().getName();
}
@Override
public void extractValues(Map<String, Object> container) {
saveThreadId();
T name = this.threadLocal.get();
Assert.notNull(name, "No ThreadLocal value");
container.put("name", name);
public T getValue() {
T value = this.threadLocal.get();
// Only save thread id on initial call (restore looks up previous value) and if there is a value.
if (value != null && this.threadId == null) {
this.threadId = Thread.currentThread().getId();
}
return value;
}
@Override
@SuppressWarnings("unchecked")
public void restoreValues(Map<String, Object> values) {
checkThreadId();
T name = (T) values.get("name");
Assert.notNull(name, "No value to set");
this.threadLocal.set(name);
public void setValue(T value) {
if (this.threadId != null) {
assertThat(Thread.currentThread().getId() != this.threadId)
.as("ThreadLocal restored on the same thread. Propagation not tested effectively.")
.isTrue();
}
this.threadLocal.set(value);
}
@Override
public void resetValues(Map<String, Object> values) {
public void reset() {
this.threadLocal.remove();
}
private void saveThreadId() {
if (this.suppressThreadIdCheck) {
return;
}
this.threadId = Thread.currentThread().getId();
}
private void checkThreadId() {
if (this.suppressThreadIdCheck) {
return;
}
assertThat(this.threadId).as("No threadId to check. Was extractValues not called?").isNotNull();
assertThat(Thread.currentThread().getId() != this.threadId)
.as("ThreadLocal value extracted and restored on the same thread. Propagation not tested effectively.")
.isTrue();
}
}

View File

@@ -23,6 +23,7 @@ import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import io.micrometer.context.ContextSnapshot;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
@@ -34,8 +35,6 @@ import org.springframework.graphql.ExecutionGraphQlResponse;
import org.springframework.graphql.ResponseHelper;
import org.springframework.graphql.TestExecutionRequest;
import org.springframework.graphql.data.method.annotation.BatchMapping;
import org.springframework.graphql.execution.ReactorContextManager;
import org.springframework.graphql.execution.SecurityContextThreadLocalAccessor;
import org.springframework.lang.Nullable;
import org.springframework.security.authentication.TestingAuthenticationToken;
import org.springframework.security.core.Authentication;
@@ -61,7 +60,7 @@ public class BatchMappingPrincipalMethodArgumentResolverTests extends BatchMappi
ReactiveSecurityContextHolder.withAuthentication(this.authentication);
private final Function<Context, Context> threadLocalContextWriter = context ->
ReactorContextManager.extractThreadLocalValues(new SecurityContextThreadLocalAccessor(), context);
ContextSnapshot.capture().updateContext(context);
private static Stream<Arguments> controllers() {

View File

@@ -22,6 +22,7 @@ import java.util.concurrent.CompletableFuture;
import graphql.GraphQLContext;
import graphql.schema.DataFetchingEnvironment;
import graphql.schema.DataFetchingEnvironmentImpl;
import io.micrometer.context.ContextSnapshot;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
@@ -51,10 +52,9 @@ public class DataFetcherHandlerMethodTests {
new HandlerMethod(new TestController(), TestController.class.getMethod("handleAndReturnCallable")),
resolvers, null, new SimpleAsyncTaskExecutor(), false);
GraphQLContext graphQLContext = new GraphQLContext.Builder().build();
DataFetchingEnvironment environment = DataFetchingEnvironmentImpl.newDataFetchingEnvironment()
.graphQLContext(graphQLContext)
DataFetchingEnvironment environment = DataFetchingEnvironmentImpl
.newDataFetchingEnvironment()
.graphQLContext(GraphQLContext.newContext().build())
.build();
Object result = handlerMethod.invoke(environment);

View File

@@ -20,6 +20,7 @@ import java.security.Principal;
import java.time.Duration;
import java.util.function.Function;
import io.micrometer.context.ContextSnapshot;
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
@@ -33,13 +34,11 @@ import org.springframework.context.annotation.AnnotationConfigApplicationContext
import org.springframework.core.MethodParameter;
import org.springframework.graphql.ExecutionGraphQlResponse;
import org.springframework.graphql.ExecutionGraphQlService;
import org.springframework.graphql.ResponseHelper;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.ResponseHelper;
import org.springframework.graphql.TestExecutionRequest;
import org.springframework.graphql.data.method.annotation.QueryMapping;
import org.springframework.graphql.data.method.annotation.SubscriptionMapping;
import org.springframework.graphql.execution.ReactorContextManager;
import org.springframework.graphql.execution.SecurityContextThreadLocalAccessor;
import org.springframework.lang.Nullable;
import org.springframework.security.authentication.TestingAuthenticationToken;
import org.springframework.security.core.Authentication;
@@ -66,7 +65,7 @@ public class SchemaMappingPrincipalMethodArgumentResolverTests {
ReactiveSecurityContextHolder.withAuthentication(this.authentication);
private final Function<Context, Context> threadLocalContextWriter = context ->
ReactorContextManager.extractThreadLocalValues(new SecurityContextThreadLocalAccessor(), context);
ContextSnapshot.capture().updateContext(context);
private final GreetingController greetingController = new GreetingController();

View File

@@ -22,12 +22,13 @@ import graphql.ExecutionInput;
import graphql.GraphQL;
import graphql.GraphQLError;
import graphql.GraphqlErrorBuilder;
import io.micrometer.context.ContextRegistry;
import io.micrometer.context.ContextSnapshot;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import reactor.test.StepVerifier;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.ResponseHelper;
@@ -85,14 +86,14 @@ public class CompositeSubscriptionExceptionResolverTests {
String query = "subscription { greetings }";
String schema = "type Subscription { greetings: String! } type Query { greeting: String! }";
ThreadLocal<String> nameThreadLocal = new ThreadLocal<>();
nameThreadLocal.set("007");
TestThreadLocalAccessor<String> accessor = new TestThreadLocalAccessor<>(nameThreadLocal);
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("007");
ContextRegistry.getInstance().registerThreadLocalAccessor(new TestThreadLocalAccessor<>(threadLocal));
try {
SubscriptionExceptionResolverAdapter resolver = SubscriptionExceptionResolver.forSingleError(exception ->
GraphqlErrorBuilder.newError()
.message("Error: " + exception.getMessage() + ", name=" + nameThreadLocal.get())
.message("Error: " + exception.getMessage() + ", name=" + threadLocal.get())
.errorType(ErrorType.BAD_REQUEST)
.build());
resolver.setThreadLocalContextAware(true);
@@ -106,13 +107,14 @@ public class CompositeSubscriptionExceptionResolverTests {
.subscriptionExceptionResolvers(resolver)
.toGraphQl();
ContextView view = ReactorContextManager.extractThreadLocalValues(accessor, Context.empty());
ExecutionInput input = ExecutionInput.newExecutionInput(query).build();
ReactorContextManager.setReactorContext(view, input.getGraphQLContext());
ContextSnapshot.capture().updateContext(input.getGraphQLContext());
Flux<ResponseHelper> flux = Mono.delay(Duration.ofMillis(10))
.flatMap((aLong) -> Mono.fromFuture(graphQL.executeAsync(input)).map(ResponseHelper::forSubscription))
.block(TIMEOUT);
Flux<ResponseHelper> flux = Mono.defer(() -> Mono.fromFuture(graphQL.executeAsync(input)))
.map(ResponseHelper::forSubscription)
.subscribeOn(Schedulers.boundedElastic()) // restore on different thread for DataFetcher
.block(TIMEOUT)
.subscribeOn(Schedulers.boundedElastic()); // restore on different thread for SubscriptionExceptionResolver
StepVerifier.create(flux)
.consumeNextWith((helper) -> assertThat(helper.toEntity("greetings", String.class)).isEqualTo("a"))
@@ -126,7 +128,7 @@ public class CompositeSubscriptionExceptionResolverTests {
.verify(TIMEOUT);
}
finally {
nameThreadLocal.remove();
threadLocal.remove();
}
}

View File

@@ -25,12 +25,12 @@ import graphql.ExecutionResult;
import graphql.GraphQL;
import graphql.GraphQLError;
import graphql.GraphqlErrorBuilder;
import io.micrometer.context.ContextRegistry;
import io.micrometer.context.ContextSnapshot;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.ResponseHelper;
@@ -59,7 +59,7 @@ public class ContextDataFetcherDecoratorTests {
.toGraphQl();
ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
ReactorContextManager.setReactorContext(Context.of("name", "007"), input.getGraphQLContext());
input.getGraphQLContext().put("name", "007");
ExecutionResult executionResult = graphQl.executeAsync(input).get();
@@ -79,7 +79,7 @@ public class ContextDataFetcherDecoratorTests {
.toGraphQl();
ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greetings }").build();
ReactorContextManager.setReactorContext(Context.of("name", "007"), input.getGraphQLContext());
input.getGraphQLContext().put("name", "007");
ExecutionResult result = graphQl.executeAsync(input).get();
@@ -99,7 +99,7 @@ public class ContextDataFetcherDecoratorTests {
.toGraphQl();
ExecutionInput input = ExecutionInput.newExecutionInput().query("subscription { greetings }").build();
ReactorContextManager.setReactorContext(Context.of("name", "007"), input.getGraphQLContext());
input.getGraphQLContext().put("name", "007");
ExecutionResult executionResult = graphQl.executeAsync(input).get();
@@ -154,17 +154,16 @@ public class ContextDataFetcherDecoratorTests {
@Test
void dataFetcherWithThreadLocalContext() {
ThreadLocal<String> nameThreadLocal = new ThreadLocal<>();
nameThreadLocal.set("007");
TestThreadLocalAccessor<String> accessor = new TestThreadLocalAccessor<>(nameThreadLocal);
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("007");
ContextRegistry.getInstance().registerThreadLocalAccessor(new TestThreadLocalAccessor<>(threadLocal));
try {
GraphQL graphQl = GraphQlSetup.schemaContent(SCHEMA_CONTENT)
.queryFetcher("greeting", (env) -> "Hello " + nameThreadLocal.get())
.queryFetcher("greeting", (env) -> "Hello " + threadLocal.get())
.toGraphQl();
ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
ContextView view = ReactorContextManager.extractThreadLocalValues(accessor, Context.empty());
ReactorContextManager.setReactorContext(view, input.getGraphQLContext());
ContextSnapshot.capture().updateContext(input.getGraphQLContext());
Mono<ExecutionResult> resultMono = Mono.delay(Duration.ofMillis(10))
.flatMap((aLong) -> Mono.fromFuture(graphQl.executeAsync(input)));
@@ -173,7 +172,7 @@ public class ContextDataFetcherDecoratorTests {
assertThat(greeting).isEqualTo("Hello 007");
}
finally {
nameThreadLocal.remove();
threadLocal.remove();
}
}

View File

@@ -28,8 +28,6 @@ import org.dataloader.stats.StatisticsCollector;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.graphql.Book;
import org.springframework.graphql.BookSource;
@@ -59,7 +57,10 @@ public class DefaultBatchLoaderRegistryTests {
return Flux.fromIterable(ids).map(BookSource::getBook);
}));
GraphQLContext graphQLContext = initGraphQLContext(Context.of("key", "value"));
ExecutionInput input = ExecutionInput.newExecutionInput().query("").build();
GraphQLContext graphQLContext = input.getGraphQLContext();
graphQLContext.put("key", "value");
this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry, graphQLContext);
Map<String, DataLoader<?, ?>> map = this.dataLoaderRegistry.getDataLoadersMap();
@@ -82,7 +83,10 @@ public class DefaultBatchLoaderRegistryTests {
return Flux.fromIterable(ids).map(BookSource::getBook).collectMap(Book::getId, Function.identity());
}));
GraphQLContext graphQLContext = initGraphQLContext(Context.of("key", "value"));
ExecutionInput input = ExecutionInput.newExecutionInput().query("").build();
GraphQLContext graphQLContext = input.getGraphQLContext();
graphQLContext.put("key", "value");
this.batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry, graphQLContext);
Map<String, DataLoader<?, ?>> map = this.dataLoaderRegistry.getDataLoadersMap();
@@ -109,10 +113,4 @@ public class DefaultBatchLoaderRegistryTests {
assertThat(map.get(name).getStatistics()).isSameAs(collector.getStatistics());
}
private GraphQLContext initGraphQLContext(ContextView context) {
ExecutionInput executionInput = ExecutionInput.newExecutionInput().query("").build();
ReactorContextManager.setReactorContext(context, executionInput.getGraphQLContext());
return executionInput.getGraphQLContext();
}
}

View File

@@ -22,10 +22,11 @@ import java.util.Collections;
import graphql.ExecutionInput;
import graphql.ExecutionResult;
import graphql.GraphqlErrorBuilder;
import io.micrometer.context.ContextRegistry;
import io.micrometer.context.ContextSnapshot;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.ResponseHelper;
@@ -76,7 +77,7 @@ public class ExceptionResolversExceptionHandlerTests {
.message("Resolved error: " + ex.getMessage() + ", name=" + view.get("name"))
.errorType(ErrorType.BAD_REQUEST).build())));
ReactorContextManager.setReactorContext(Context.of("name", "007"), input.getGraphQLContext());
this.input.getGraphQLContext().put("name", "007");
ExecutionResult result = this.graphQlSetup.exceptionResolver(resolver).toGraphQl()
.executeAsync(this.input).get();
@@ -88,21 +89,19 @@ public class ExceptionResolversExceptionHandlerTests {
@Test
void resolveExceptionWithThreadLocal() {
ThreadLocal<String> nameThreadLocal = new ThreadLocal<>();
nameThreadLocal.set("007");
TestThreadLocalAccessor<String> accessor = new TestThreadLocalAccessor<>(nameThreadLocal);
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("007");
ContextRegistry.getInstance().registerThreadLocalAccessor(new TestThreadLocalAccessor<>(threadLocal));
try {
DataFetcherExceptionResolverAdapter resolver =
DataFetcherExceptionResolver.forSingleError((ex, env) ->
GraphqlErrorBuilder.newError(env)
.message("Resolved error: " + ex.getMessage() + ", name=" + nameThreadLocal.get())
.message("Resolved error: " + ex.getMessage() + ", name=" + threadLocal.get())
.errorType(ErrorType.BAD_REQUEST)
.build());
resolver.setThreadLocalContextAware(true);
ContextView view = ReactorContextManager.extractThreadLocalValues(accessor, Context.empty());
ReactorContextManager.setReactorContext(view, input.getGraphQLContext());
ContextSnapshot.capture().updateContext(this.input.getGraphQLContext());
Mono<ExecutionResult> result = Mono.delay(Duration.ofMillis(10)).flatMap((aLong) ->
Mono.fromFuture(this.graphQlSetup.exceptionResolver(resolver).toGraphQl().executeAsync(this.input)));
@@ -112,7 +111,7 @@ public class ExceptionResolversExceptionHandlerTests {
assertThat(response.error(0).message()).isEqualTo("Resolved error: Invalid greeting, name=007");
}
finally {
nameThreadLocal.remove();
threadLocal.remove();
}
}

View File

@@ -1,73 +0,0 @@
/*
* Copyright 2002-2021 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.graphql.execution;
import java.time.Duration;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;
import org.springframework.graphql.TestThreadLocalAccessor;
import static org.assertj.core.api.Assertions.assertThat;
/**
* Unit tests for {@link ReactorContextManager}.
* @author Rossen Stoyanchev
*/
public class ReactorContextManagerTests {
@Test
void restoreThreadLocaValues() {
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("myValue");
Context context = ReactorContextManager.extractThreadLocalValues(
new TestThreadLocalAccessor<>(threadLocal), Context.empty());
try {
Mono.delay(Duration.ofMillis(10))
.doOnNext(aLong -> {
assertThat(threadLocal.get()).isNull();
ReactorContextManager.restoreThreadLocalValues(context);
assertThat(threadLocal.get()).isEqualTo("myValue");
ReactorContextManager.resetThreadLocalValues(context);
})
.block();
}
finally {
threadLocal.remove();
}
}
@Test
void restoreThreadLocaValuesOnSameThreadIsNoOp() {
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("myValue");
Context context = ReactorContextManager.extractThreadLocalValues(
new TestThreadLocalAccessor<>(threadLocal, true), Context.empty());
threadLocal.remove();
ReactorContextManager.restoreThreadLocalValues(context);
assertThat(threadLocal.get()).isNull();
threadLocal.set("anotherValue");
ReactorContextManager.resetThreadLocalValues(context);
assertThat(threadLocal.get()).isEqualTo("anotherValue");
}
}

View File

@@ -22,6 +22,7 @@ import java.util.Collections;
import graphql.GraphqlErrorBuilder;
import graphql.schema.DataFetcher;
import io.micrometer.context.ContextRegistry;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Mono;
@@ -92,15 +93,13 @@ public class WebGraphQlHandlerTests {
@Test
void threadLocalContextPropagation() {
ThreadLocal<String> nameThreadLocal = new ThreadLocal<>();
nameThreadLocal.set("007");
TestThreadLocalAccessor<String> threadLocalAccessor = new TestThreadLocalAccessor<>(nameThreadLocal);
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("007");
ContextRegistry.getInstance().registerThreadLocalAccessor(new TestThreadLocalAccessor<>(threadLocal));
try {
Mono<WebGraphQlResponse> responseMono = this.graphQlSetup
.queryFetcher("greeting", env -> "Hello " + nameThreadLocal.get())
.queryFetcher("greeting", env -> "Hello " + threadLocal.get())
.interceptor((input, next) -> Mono.delay(Duration.ofMillis(10)).flatMap((aLong) -> next.next(input)))
.threadLocalAccessor(threadLocalAccessor)
.toWebGraphQlHandler()
.handleRequest(webInput);
@@ -108,27 +107,26 @@ public class WebGraphQlHandlerTests {
assertThat(greeting).isEqualTo("Hello 007");
}
finally {
nameThreadLocal.remove();
threadLocal.remove();
}
}
@Test
void threadLocalContextPropagationToExceptionResolver() {
ThreadLocal<String> nameThreadLocal = new ThreadLocal<>();
nameThreadLocal.set("007");
TestThreadLocalAccessor<String> threadLocalAccessor = new TestThreadLocalAccessor<>(nameThreadLocal);
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("007");
ContextRegistry.getInstance().registerThreadLocalAccessor(new TestThreadLocalAccessor<>(threadLocal));
try {
DataFetcherExceptionResolverAdapter exceptionResolver =
DataFetcherExceptionResolver.forSingleError((ex, env) ->
GraphqlErrorBuilder.newError(env)
.message("Resolved error: " + ex.getMessage() + ", name=" + nameThreadLocal.get())
.message("Resolved error: " + ex.getMessage() + ", name=" + threadLocal.get())
.errorType(ErrorType.BAD_REQUEST).build());
exceptionResolver.setThreadLocalContextAware(true);
Mono<WebGraphQlResponse> responseMono = this.graphQlSetup.queryFetcher("greeting", this.errorDataFetcher)
.exceptionResolver(exceptionResolver)
.interceptor((input, next) -> Mono.delay(Duration.ofMillis(10)).flatMap((aLong) -> next.next(input)))
.threadLocalAccessor(threadLocalAccessor)
.toWebGraphQlHandler()
.handleRequest(webInput);
@@ -137,7 +135,7 @@ public class WebGraphQlHandlerTests {
assertThat(response.error(0).message()).isEqualTo("Resolved error: Invalid greeting, name=007");
}
finally {
nameThreadLocal.remove();
threadLocal.remove();
}
}

View File

@@ -20,8 +20,6 @@ import reactor.core.publisher.Flux;
import org.springframework.graphql.BookSource;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.lang.Nullable;
public abstract class WebSocketHandlerTestSupport {
@@ -67,12 +65,6 @@ public abstract class WebSocketHandlerTestSupport {
protected WebGraphQlHandler initHandler(WebGraphQlInterceptor... interceptors) {
return initHandler(null, interceptors);
}
protected WebGraphQlHandler initHandler(
@Nullable ThreadLocalAccessor accessor, WebGraphQlInterceptor... interceptors) {
return GraphQlSetup.schemaResource(BookSource.schema)
.queryFetcher("bookById", environment -> {
Long id = Long.parseLong(environment.getArgument("id"));
@@ -83,7 +75,6 @@ public abstract class WebSocketHandlerTestSupport {
return Flux.fromIterable(BookSource.books())
.filter((book) -> book.getAuthor().getFullName().contains(author));
})
.threadLocalAccessor(accessor)
.interceptor(interceptors)
.toWebGraphQlHandler();
}

View File

@@ -28,6 +28,8 @@ import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import io.micrometer.context.ContextRegistry;
import io.micrometer.context.ContextSnapshot;
import org.assertj.core.api.InstanceOfAssertFactories;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
@@ -37,7 +39,6 @@ import reactor.test.StepVerifier;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.TestThreadLocalAccessor;
import org.springframework.graphql.execution.ErrorType;
import org.springframework.graphql.execution.ThreadLocalAccessor;
import org.springframework.graphql.server.ConsumeOneAndNeverCompleteInterceptor;
import org.springframework.graphql.server.WebGraphQlHandler;
import org.springframework.graphql.server.WebGraphQlInterceptor;
@@ -51,7 +52,6 @@ import org.springframework.http.HttpInputMessage;
import org.springframework.http.converter.GenericHttpMessageConverter;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter;
import org.springframework.lang.Nullable;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
@@ -369,58 +369,59 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
void contextPropagation() throws Exception {
ThreadLocal<String> threadLocal = new ThreadLocal<>();
threadLocal.set("foo");
ContextRegistry.getInstance().registerThreadLocalAccessor(new TestThreadLocalAccessor<>(threadLocal));
try {
WebGraphQlInterceptor threadLocalInterceptor = (request, chain) -> {
assertThat(threadLocal.get()).isEqualTo("foo");
return chain.next(request);
};
WebGraphQlInterceptor threadLocalInterceptor = (request, chain) -> {
assertThat(threadLocal.get()).isEqualTo("foo");
return chain.next(request);
};
GraphQlWebSocketHandler handler = initWebSocketHandler(threadLocalInterceptor);
GraphQlWebSocketHandler handler = initWebSocketHandler(
new TestThreadLocalAccessor<>(threadLocal), threadLocalInterceptor);
// Ensure ContextSnapshot is present in WebSocketSession attributes
this.session.getAttributes().put(ContextSnapshot.class.getName(), ContextSnapshot.capture());
// Use HandshakeInterceptor to capture ThreadLocal context
handler.asWebSocketHttpRequestHandler((request, response, wsHandler, attributes) -> false)
.getHandshakeInterceptors().get(0)
.beforeHandshake(null, null, null, this.session.getAttributes());
// Context should propagate, if message is handled on different thread
Thread thread = new Thread(() -> {
try {
handle(handler,
new TextMessage("{\"type\":\"connection_init\"}"),
new TextMessage(BOOK_QUERY));
}
catch (Exception ex) {
throw new IllegalStateException(ex);
}
});
thread.start();
// Context should propagate, if message is handled on different thread
Thread thread = new Thread(() -> {
try {
handle(handler,
new TextMessage("{\"type\":\"connection_init\"}"),
new TextMessage(BOOK_QUERY));
}
catch (Exception ex) {
throw new IllegalStateException(ex);
}
});
thread.start();
StepVerifier.create(this.session.getOutput())
.expectNextCount(2)
.consumeNextWith((message) -> assertMessageType(message, GraphQlWebSocketMessageType.COMPLETE))
.then(this.session::close) // Complete output Flux
.expectComplete()
.verify(TIMEOUT);
StepVerifier.create(this.session.getOutput())
.expectNextCount(2)
.consumeNextWith((message) -> assertMessageType(message, GraphQlWebSocketMessageType.COMPLETE))
.then(this.session::close) // Complete output Flux
.expectComplete()
.verify(TIMEOUT);
}
finally {
threadLocal.remove();
}
}
private void handle(GraphQlWebSocketHandler handler, TextMessage... textMessages) throws Exception {
handler.afterConnectionEstablished(this.session);
if (!this.session.getAttributes().containsKey(ContextSnapshot.class.getName())) {
// Ensure ContextSnapshot is present in WebSocketSession attributes
this.session.getAttributes().put(ContextSnapshot.class.getName(), ContextSnapshot.capture());
}
for (TextMessage message : textMessages) {
handler.handleTextMessage(this.session, message);
}
}
private GraphQlWebSocketHandler initWebSocketHandler(WebGraphQlInterceptor... interceptors) {
return initWebSocketHandler(null, interceptors);
}
private GraphQlWebSocketHandler initWebSocketHandler(
@Nullable ThreadLocalAccessor accessor, WebGraphQlInterceptor... interceptors) {
try {
return new GraphQlWebSocketHandler(
initHandler(accessor, interceptors), converter, Duration.ofSeconds(60));
return new GraphQlWebSocketHandler(initHandler(interceptors), converter, Duration.ofSeconds(60));
}
catch (Exception ex) {
throw new IllegalStateException(ex);