Upgrade to Context Propagation 1.0.3

bclozel · bclozel · commit 48c16267e8da · 2023-06-07T16:13:42.000+02:00
Closes gh-716 Closes gh-718
diff --git a/platform/build.gradle b/platform/build.gradle
@@ -25,7 +25,7 @@ dependencies {
 
 	constraints {
 		api("com.graphql-java:graphql-java:${graphQlJavaVersion}")
-		api("io.micrometer:context-propagation:1.0.2")
+		api("io.micrometer:context-propagation:1.0.3")
 
 		api("jakarta.annotation:jakarta.annotation-api:2.1.1")
 		api("jakarta.servlet:jakarta.servlet-api:6.0.0")
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
@@ -27,6 +27,7 @@
 
 import graphql.GraphQLContext;
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import reactor.core.publisher.Mono;
 
 import org.springframework.core.CoroutinesUtils;
@@ -45,6 +46,8 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
 
 	private static final Object NO_VALUE = new Object();
 
+	private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build();
+
 
 	private final boolean hasCallableReturnValue;
 
@@ -114,7 +117,7 @@ private Object handleReturnValue(GraphQLContext graphQLContext, @Nullable Object
 			return CompletableFuture.supplyAsync(
 					() -> {
 						try {
-							return ContextSnapshot.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
+							return SNAPSHOT_FACTORY.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
 						}
 						catch (Exception ex) {
 							throw new IllegalStateException(
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
@@ -33,6 +33,7 @@
 import graphql.util.TraversalControl;
 import graphql.util.TraverserContext;
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -59,6 +60,8 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
 
 	private final SubscriptionExceptionResolver subscriptionExceptionResolver;
 
+	private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
+
 	private ContextDataFetcherDecorator(
 			DataFetcher<?> delegate, boolean subscription,
 			SubscriptionExceptionResolver subscriptionExceptionResolver) {
@@ -73,18 +76,13 @@ private ContextDataFetcherDecorator(
 	@Override
 	public Object get(DataFetchingEnvironment environment) throws Exception {
 
-		GraphQLContext context;
-		// temporarily merge global and local graphql context until https://github.com/micrometer-metrics/context-propagation/pull/98
+		ContextSnapshot snapshot;
 		if (environment.getLocalContext() instanceof GraphQLContext localContext) {
-			context = GraphQLContext.newContext()
-					.of(environment.getGraphQlContext())
-					.of(localContext)
-					.build();
+			snapshot = snapshotFactory.captureFrom(environment.getGraphQlContext(), localContext);
 		}
 		else {
-			context = environment.getGraphQlContext();
+			snapshot = snapshotFactory.captureFrom(environment.getGraphQlContext());
 		}
-		ContextSnapshot snapshot = ContextSnapshot.captureFrom(context);
 		Object value = snapshot.wrap(() -> this.delegate.get(environment)).call();
 
 		if (this.subscription) {
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
@@ -22,6 +22,7 @@
 import graphql.GraphQLError;
 import graphql.schema.DataFetchingEnvironment;
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import io.micrometer.context.ThreadLocalAccessor;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -52,6 +53,8 @@ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcher
 
 	protected final Log logger = LogFactory.getLog(getClass());
 
+	protected final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
+
 	private boolean threadLocalContextAware;
 
 
@@ -98,7 +101,7 @@ private List<GraphQLError> resolveInternal(Throwable exception, DataFetchingEnvi
 			return resolveToMultipleErrors(exception, env);
 		}
 		try {
-			return ContextSnapshot.captureFrom(env.getGraphQlContext())
+			return snapshotFactory.captureFrom(env.getGraphQlContext())
 					.wrap(() -> resolveToMultipleErrors(exception, env))
 					.call();
 		}
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
@@ -27,6 +27,7 @@
 
 import graphql.GraphQLContext;
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import org.dataloader.BatchLoaderContextProvider;
 import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.BatchLoaderWithContext;
@@ -52,13 +53,16 @@
  */
 public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
 
+	private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build();
+
 	private final List<ReactorBatchLoader<?,?>> loaders = new ArrayList<>();
 
 	private final List<ReactorMappedBatchLoader<?,?>> mappedLoaders = new ArrayList<>();
 
 	private final Supplier<DataLoaderOptions> defaultOptionsSupplier;
 
 
+
 	/**
 	 * Default constructor
 	 */
@@ -230,7 +234,7 @@ public DataLoaderOptions getOptions() {
 		@Override
 		public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
 			GraphQLContext graphQLContext = environment.getContext();
-			ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
+			ContextSnapshot snapshot = SNAPSHOT_FACTORY.captureFrom(graphQLContext);
 			try {
 				return snapshot.wrap(() ->
 								this.loader.apply(keys, environment)
@@ -280,7 +284,7 @@ public DataLoaderOptions getOptions() {
 		@Override
 		public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
 			GraphQLContext graphQLContext = environment.getContext();
-			ContextSnapshot snapshot = ContextSnapshot.captureFrom(graphQLContext);
+			ContextSnapshot snapshot = SNAPSHOT_FACTORY.captureFrom(graphQLContext);
 			try {
 				return snapshot.wrap(() ->
 								this.loader.apply(keys, environment)
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
@@ -25,6 +25,7 @@
 import graphql.GraphQLContext;
 import graphql.execution.ExecutionIdProvider;
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import org.dataloader.DataLoaderRegistry;
 import reactor.core.publisher.Mono;
 
@@ -45,6 +46,7 @@ public class DefaultExecutionGraphQlService implements ExecutionGraphQlService {
 	private static final BiFunction<ExecutionInput, ExecutionInput.Builder, ExecutionInput> RESET_EXECUTION_ID_CONFIGURER =
 			(executionInput, builder) -> builder.executionId(null).build();
 
+	private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
 
 	private final GraphQlSource graphQlSource;
 
@@ -77,7 +79,7 @@ public final Mono<ExecutionGraphQlResponse> execute(ExecutionGraphQlRequest requ
 				request.configureExecutionInput(RESET_EXECUTION_ID_CONFIGURER);
 			}
 			ExecutionInput executionInput = request.toExecutionInput();
-			ContextSnapshot.captureFrom(contextView).updateContext(executionInput.getGraphQLContext());
+			snapshotFactory.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
@@ -29,6 +29,7 @@
 import graphql.execution.ExecutionId;
 import graphql.schema.DataFetchingEnvironment;
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import reactor.core.publisher.Flux;
@@ -47,6 +48,8 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler
 
 	private static final Log logger = LogFactory.getLog(ExceptionResolversExceptionHandler.class);
 
+	private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
+
 	private final List<DataFetcherExceptionResolver> resolvers;
 
 	/**
@@ -70,7 +73,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());
+		ContextSnapshot snapshot = snapshotFactory.captureFrom(env.getGraphQlContext());
 		try {
 			return Flux.fromIterable(this.resolvers)
 					.flatMap(resolver -> resolver.resolveException(exception, env))
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
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -72,10 +72,30 @@ private <V> void setValueInternal(Object value) {
 	}
 
 	@Override
+	public void setValue() {
+		this.delegate.setValue();
+	}
+
+	@Override
+	@Deprecated
 	public void reset() {
 		this.delegate.reset();
 	}
 
+	@Override
+	public void restore(Object previousValue) {
+		restoreInternal(previousValue);
+	}
+
+	@SuppressWarnings("unchecked")
+	public <V> void restoreInternal(Object previousValue) {
+		((ThreadLocalAccessor<V>) this.delegate).restore((V) previousValue);
+	}
+
+	@Override
+	public void restore() {
+		this.delegate.restore();
+	}
 
 	private static class DelegateAccessor implements ThreadLocalAccessor<Object> {
 
@@ -95,6 +115,22 @@ public void setValue(Object value) {
 		}
 
 		@Override
+		public void setValue() {
+			SecurityContextHolder.clearContext();
+		}
+
+		@Override
+		public void restore(Object previousValue) {
+			SecurityContextHolder.setContext((SecurityContext) previousValue);
+		}
+
+		@Override
+		public void restore() {
+			SecurityContextHolder.clearContext();
+		}
+
+		@Override
+		@Deprecated
 		public void reset() {
 			SecurityContextHolder.clearContext();
 		}
@@ -119,6 +155,19 @@ public void setValue(Object value) {
 		}
 
 		@Override
+		public void setValue() {
+		}
+
+		@Override
+		public void restore(Object previousValue) {
+		}
+
+		@Override
+		public void restore() {
+		}
+
+		@Override
+		@Deprecated
 		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
@@ -22,6 +22,7 @@
 
 import graphql.GraphQLError;
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import io.micrometer.context.ThreadLocalAccessor;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -50,6 +51,8 @@ public abstract class SubscriptionExceptionResolverAdapter implements Subscripti
 
     protected final Log logger = LogFactory.getLog(getClass());
 
+    protected final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
+
     private boolean threadLocalContextAware;
 
 
@@ -83,7 +86,7 @@ public boolean isThreadLocalContextAware() {
     public final Mono<List<GraphQLError>> resolveException(Throwable exception) {
         if (this.threadLocalContextAware) {
             return Mono.deferContextual(contextView -> {
-                ContextSnapshot snapshot = ContextSnapshot.captureFrom(contextView);
+                ContextSnapshot snapshot = snapshotFactory.captureFrom(contextView);
                 try {
                     List<GraphQLError> errors = snapshot.wrap(() -> resolveToMultipleErrors(exception)).call();
                     return Mono.justOrEmpty(errors);
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
@@ -21,6 +21,7 @@
 import java.util.List;
 
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import reactor.core.publisher.Mono;
 
 import org.springframework.graphql.ExecutionGraphQlService;
@@ -70,6 +71,8 @@ public WebGraphQlHandler.Builder interceptors(List<WebGraphQlInterceptor> interc
 	@Override
 	public WebGraphQlHandler build() {
 
+		ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
+
 		Chain endOfChain = request -> this.service.execute(request).map(WebGraphQlResponse::new);
 
 		Chain executionChain = this.interceptors.stream()
@@ -87,7 +90,7 @@ public WebSocketGraphQlInterceptor getWebSocketInterceptor() {
 
 			@Override
 			public Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request) {
-				ContextSnapshot snapshot = ContextSnapshot.captureAll();
+				ContextSnapshot snapshot = snapshotFactory.captureAll();
 				return executionChain.next(request).contextWrite(snapshot::updateContext);
 			}
 		};
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
@@ -36,6 +36,7 @@
 import graphql.GraphQLError;
 import graphql.GraphqlErrorBuilder;
 import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshotFactory;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Publisher;
@@ -357,12 +358,14 @@ private static class ContextHandshakeInterceptor implements HandshakeInterceptor
 
 		private static final String KEY = ContextSnapshot.class.getName();
 
+		private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build();
+
 		@Override
 		public boolean beforeHandshake(
 				ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
 				Map<String, Object> attributes) {
 
-			attributes.put(KEY, ContextSnapshot.captureAll());
+			attributes.put(KEY, SNAPSHOT_FACTORY.captureAll());
 			return true;
 		}
 
