Switch to io.micrometer:context-propagation library

See gh-459

Author: rstoyanchev

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"

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'
 

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'

spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java

@@ -26,11 +26,11 @@
 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 @@ private Object handleReturnValue(GraphQLContext graphQLContext, @Nullable Object
 			return CompletableFuture.supplyAsync(
 					() -> {
 						try {
-							return ReactorContextManager.invokeCallable((Callable<?>) result, graphQLContext);
+							return ContextSnapshot.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
 						}
 						catch (Exception ex) {
 							throw new IllegalStateException(

spring-graphql/src/main/java/org/springframework/graphql/execution/CompositeThreadLocalAccessor.java

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

spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java

@@ -29,10 +29,10 @@
 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 @@ private ContextDataFetcherDecorator(
 	@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;

spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java

@@ -21,8 +21,11 @@
 
 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 @@
  */
 public abstract class DataFetcherExceptionResolverAdapter implements DataFetcherExceptionResolver {
 
+	protected final Log logger = LogFactory.getLog(getClass());
+
 	private boolean threadLocalContextAware;
 
 
@@ -88,17 +93,18 @@ public final Mono<List<GraphQLError>> resolveException(Throwable ex, DataFetchin
 	}
 
 	@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;
 		}
 	}
 

spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java

@@ -19,11 +19,13 @@
 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.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 DataLoaderOptions getOptionsOrDefault(
 
 		@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 DataLoaderOptions getOptionsOrDefault(
 
 		@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);
 			}
 		}
 

spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java

@@ -24,6 +24,7 @@
 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 final Mono<ExecutionGraphQlResponse> execute(ExecutionGraphQlRequest requ
 				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));

spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java

@@ -28,11 +28,11 @@
 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 @@ public DataFetcherExceptionHandlerResult onException(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 @@ public CompletableFuture<DataFetcherExceptionHandlerResult> handleException(Data
 					.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) {

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<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);
+	}
+
+}