From f680892dfc277deabf44be57758de5b18e67e9ad Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Mon, 29 Aug 2022 21:14:55 +0100 Subject: [PATCH] Switch to io.micrometer:context-propagation library See gh-459 --- build.gradle | 3 +- spring-graphql-test/build.gradle | 1 + spring-graphql/build.gradle | 2 + .../method/InvocableHandlerMethodSupport.java | 4 +- .../CompositeThreadLocalAccessor.java | 1 + .../ContextDataFetcherDecorator.java | 18 +-- .../DataFetcherExceptionResolverAdapter.java | 22 ++- .../execution/DefaultBatchLoaderRegistry.java | 33 ++-- .../DefaultExecutionGraphQlService.java | 3 +- .../ExceptionResolversExceptionHandler.java | 8 +- .../execution/GraphQlContextAccessor.java | 64 ++++++++ .../execution/ReactorContextManager.java | 141 ------------------ .../SecurityContextThreadLocalAccessor.java | 112 +++++++++++++- .../SubscriptionExceptionResolverAdapter.java | 33 ++-- .../execution/ThreadLocalAccessor.java | 4 + .../DefaultWebGraphQlHandlerBuilder.java | 32 ++-- .../graphql/server/WebGraphQlHandler.java | 6 + .../webmvc/GraphQlWebSocketHandler.java | 40 ++--- .../io.micrometer.context.ContextAccessor | 1 + .../io.micrometer.context.ThreadLocalAccessor | 1 + .../graphql/TestThreadLocalAccessor.java | 64 ++++---- ...gPrincipalMethodArgumentResolverTests.java | 5 +- .../DataFetcherHandlerMethodTests.java | 8 +- ...gPrincipalMethodArgumentResolverTests.java | 7 +- ...iteSubscriptionExceptionResolverTests.java | 26 ++-- .../ContextDataFetcherDecoratorTests.java | 23 ++- .../DefaultBatchLoaderRegistryTests.java | 18 +-- ...ceptionResolversExceptionHandlerTests.java | 19 ++- .../execution/ReactorContextManagerTests.java | 73 --------- .../server/WebGraphQlHandlerTests.java | 24 ++- .../server/WebSocketHandlerTestSupport.java | 9 -- .../webmvc/GraphQlWebSocketHandlerTests.java | 79 +++++----- 32 files changed, 413 insertions(+), 471 deletions(-) create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/execution/GraphQlContextAccessor.java delete mode 100644 spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java create mode 100644 spring-graphql/src/main/resources/META-INF/services/io.micrometer.context.ContextAccessor create mode 100644 spring-graphql/src/main/resources/META-INF/services/io.micrometer.context.ThreadLocalAccessor delete mode 100644 spring-graphql/src/test/java/org/springframework/graphql/execution/ReactorContextManagerTests.java diff --git a/build.gradle b/build.gradle index a4420ee0..76430941 100644 --- a/build.gradle +++ b/build.gradle @@ -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" diff --git a/spring-graphql-test/build.gradle b/spring-graphql-test/build.gradle index 9bf491ac..7c02613a 100644 --- a/spring-graphql-test/build.gradle +++ b/spring-graphql-test/build.gradle @@ -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' diff --git a/spring-graphql/build.gradle b/spring-graphql/build.gradle index 66311ff6..050bbb11 100644 --- a/spring-graphql/build.gradle +++ b/spring-graphql/build.gradle @@ -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' diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java index 9ef9f8e5..315e9aae 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java @@ -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( diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/CompositeThreadLocalAccessor.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/CompositeThreadLocalAccessor.java index 55ddc3bb..a946968f 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/CompositeThreadLocalAccessor.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/CompositeThreadLocalAccessor.java @@ -28,6 +28,7 @@ import java.util.Map; * @author Rossen Stoyanchev * @since 1.0.0 */ +@SuppressWarnings("deprecation") class CompositeThreadLocalAccessor implements ThreadLocalAccessor { private final List accessors; diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java index 1f2d3d9b..21a7a026 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java @@ -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 { @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; diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java index 81c48911..b8c6c471 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java @@ -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 resolveInternal(Throwable ex, DataFetchingEnvironment env) { + private List 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; } } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java index 7dbee83d..ad20e6da 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java @@ -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> load(List 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> load(Set 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); } } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java index 27f96409..0b1352fd 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java @@ -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)); diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java index 1d35d4df..372003ba 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java @@ -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 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) { diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/GraphQlContextAccessor.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/GraphQlContextAccessor.java new file mode 100644 index 00000000..55890196 --- /dev/null +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/GraphQlContextAccessor.java @@ -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 { + + @Override + public boolean canReadFrom(Class contextType) { + return GraphQLContext.class.equals(contextType); + } + + @Override + public void readValues(GraphQLContext context, Predicate keyPredicate, Map readValues) { + context.stream().forEach(entry -> { + if (keyPredicate.test(entry.getKey())) { + readValues.put(entry.getKey(), entry.getValue()); + } + }); + } + + @Override + public 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 valuesToWrite, GraphQLContext targetContext) { + return targetContext.putAll(valuesToWrite); + } + +} diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java deleted file mode 100644 index 14480a8a..00000000 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java +++ /dev/null @@ -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. - * - *

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 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 invokeCallable(Callable 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); - } - -} diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/SecurityContextThreadLocalAccessor.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/SecurityContextThreadLocalAccessor.java index 84134680..ade37e06 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/SecurityContextThreadLocalAccessor.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/SecurityContextThreadLocalAccessor.java @@ -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, + 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 void setValueInternal(Object value) { + ((ThreadLocalAccessor) 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 container) { - container.put(KEY, SecurityContextHolder.getContext()); + container.put((String) key(), SecurityContextHolder.getContext()); } @Override public void restoreValues(Map 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 { + + @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 { + + @Override + public Object key() { + return getClass().getName(); + } + + @Override + public Object getValue() { + return null; + } + + @Override + public void setValue(Object value) { + } + + @Override + public void reset() { + } + + } + } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java index 7c3c8cb3..836c00be 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java @@ -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> resolveException(Throwable exception) { - if (!this.threadLocalContextAware) { + if (this.threadLocalContextAware) { + return Mono.deferContextual(contextView -> { + ContextSnapshot snapshot = ContextSnapshot.captureFrom(contextView); + try { + List 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 errors; - try { - ReactorContextManager.restoreThreadLocalValues(contextView); - errors = resolveToMultipleErrors(exception); - } - finally { - ReactorContextManager.resetThreadLocalValues(contextView); - } - return Mono.justOrEmpty(errors); - }); } /** diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ThreadLocalAccessor.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ThreadLocalAccessor.java index f6617949..8a66a864 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ThreadLocalAccessor.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ThreadLocalAccessor.java @@ -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 { /** diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java b/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java index 774ddae9..68bfd146 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java @@ -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 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 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); } - }; } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java index 02471343..a3a7e9e1 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java @@ -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 accessors); /** diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java index 052fdd04..4553811e 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java @@ -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 attributes) { - if (this.accessor != null) { - Map 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 valuesMap = (Map) 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; } - } diff --git a/spring-graphql/src/main/resources/META-INF/services/io.micrometer.context.ContextAccessor b/spring-graphql/src/main/resources/META-INF/services/io.micrometer.context.ContextAccessor new file mode 100644 index 00000000..f238cecb --- /dev/null +++ b/spring-graphql/src/main/resources/META-INF/services/io.micrometer.context.ContextAccessor @@ -0,0 +1 @@ +org.springframework.graphql.execution.GraphQlContextAccessor \ No newline at end of file diff --git a/spring-graphql/src/main/resources/META-INF/services/io.micrometer.context.ThreadLocalAccessor b/spring-graphql/src/main/resources/META-INF/services/io.micrometer.context.ThreadLocalAccessor new file mode 100644 index 00000000..8cdc96eb --- /dev/null +++ b/spring-graphql/src/main/resources/META-INF/services/io.micrometer.context.ThreadLocalAccessor @@ -0,0 +1 @@ +org.springframework.graphql.execution.SecurityContextThreadLocalAccessor \ No newline at end of file diff --git a/spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java b/spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java index 80bebbd6..75137aa7 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/TestThreadLocalAccessor.java @@ -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 implements ThreadLocalAccessor { +public class TestThreadLocalAccessor implements ThreadLocalAccessor { private final ThreadLocal threadLocal; @Nullable private Long threadId; - private final boolean suppressThreadIdCheck; public TestThreadLocalAccessor(ThreadLocal threadLocal) { - this(threadLocal, false); - } - - public TestThreadLocalAccessor(ThreadLocal threadLocal, boolean suppressThreadIdCheck) { this.threadLocal = threadLocal; - this.suppressThreadIdCheck = suppressThreadIdCheck; + } + + + @Override + public Object key() { + return getClass().getName(); } @Override - public void extractValues(Map 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 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 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(); - } - } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingPrincipalMethodArgumentResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingPrincipalMethodArgumentResolverTests.java index 0e12d127..591e1304 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingPrincipalMethodArgumentResolverTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingPrincipalMethodArgumentResolverTests.java @@ -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 threadLocalContextWriter = context -> - ReactorContextManager.extractThreadLocalValues(new SecurityContextThreadLocalAccessor(), context); + ContextSnapshot.capture().updateContext(context); private static Stream controllers() { diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java index 7ccade17..2838e890 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java @@ -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); diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingPrincipalMethodArgumentResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingPrincipalMethodArgumentResolverTests.java index 0490fb9c..3ad30fb2 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingPrincipalMethodArgumentResolverTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingPrincipalMethodArgumentResolverTests.java @@ -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 threadLocalContextWriter = context -> - ReactorContextManager.extractThreadLocalValues(new SecurityContextThreadLocalAccessor(), context); + ContextSnapshot.capture().updateContext(context); private final GreetingController greetingController = new GreetingController(); diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/CompositeSubscriptionExceptionResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/CompositeSubscriptionExceptionResolverTests.java index 95f36f2d..b9b46e43 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/execution/CompositeSubscriptionExceptionResolverTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/CompositeSubscriptionExceptionResolverTests.java @@ -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 nameThreadLocal = new ThreadLocal<>(); - nameThreadLocal.set("007"); - TestThreadLocalAccessor accessor = new TestThreadLocalAccessor<>(nameThreadLocal); + ThreadLocal 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 flux = Mono.delay(Duration.ofMillis(10)) - .flatMap((aLong) -> Mono.fromFuture(graphQL.executeAsync(input)).map(ResponseHelper::forSubscription)) - .block(TIMEOUT); + Flux 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(); } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextDataFetcherDecoratorTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextDataFetcherDecoratorTests.java index 56befead..b3857329 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextDataFetcherDecoratorTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/ContextDataFetcherDecoratorTests.java @@ -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 nameThreadLocal = new ThreadLocal<>(); - nameThreadLocal.set("007"); - TestThreadLocalAccessor accessor = new TestThreadLocalAccessor<>(nameThreadLocal); + ThreadLocal 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 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(); } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java index 9d4eb347..26f5196e 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java @@ -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> 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> 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(); - } - } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandlerTests.java index 87ed8dc3..575dffa6 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandlerTests.java @@ -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 nameThreadLocal = new ThreadLocal<>(); - nameThreadLocal.set("007"); - TestThreadLocalAccessor accessor = new TestThreadLocalAccessor<>(nameThreadLocal); + ThreadLocal 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 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(); } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/execution/ReactorContextManagerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/execution/ReactorContextManagerTests.java deleted file mode 100644 index dbea7df7..00000000 --- a/spring-graphql/src/test/java/org/springframework/graphql/execution/ReactorContextManagerTests.java +++ /dev/null @@ -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 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 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"); - } - -} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java index dfe33c37..2ffe9596 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java @@ -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 nameThreadLocal = new ThreadLocal<>(); - nameThreadLocal.set("007"); - TestThreadLocalAccessor threadLocalAccessor = new TestThreadLocalAccessor<>(nameThreadLocal); + ThreadLocal threadLocal = new ThreadLocal<>(); + threadLocal.set("007"); + ContextRegistry.getInstance().registerThreadLocalAccessor(new TestThreadLocalAccessor<>(threadLocal)); try { - Mono 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 nameThreadLocal = new ThreadLocal<>(); - nameThreadLocal.set("007"); - TestThreadLocalAccessor threadLocalAccessor = new TestThreadLocalAccessor<>(nameThreadLocal); + ThreadLocal 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 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(); } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java b/spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java index b84078ee..d152a8a7 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java @@ -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(); } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java index 03aeaa9a..48f53d13 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java @@ -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 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);