Skip to content

Commit 2652a57

Browse files
rstoyanchevrstoyanchev
committed
Enable use of single, specific ContextSnapshotFactory instance
Closes gh-919
1 parent d9a583a commit 2652a57

12 files changed

+246
-39
lines changed

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,12 @@
2626
import java.util.concurrent.Executor;
2727

2828
import graphql.GraphQLContext;
29-
import io.micrometer.context.ContextSnapshotFactory;
3029
import reactor.core.publisher.Mono;
3130

3231
import org.springframework.core.CoroutinesUtils;
3332
import org.springframework.core.KotlinDetector;
3433
import org.springframework.data.util.KotlinReflectionUtils;
34+
import org.springframework.graphql.execution.ContextSnapshotFactoryHelper;
3535
import org.springframework.lang.Nullable;
3636
import org.springframework.util.Assert;
3737

@@ -46,8 +46,6 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
4646

4747
private static final Object NO_VALUE = new Object();
4848

49-
private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build();
50-
5149

5250
private final boolean hasCallableReturnValue;
5351

@@ -131,7 +129,7 @@ private Object handleReturnValue(GraphQLContext graphQLContext, @Nullable Object
131129
return CompletableFuture.supplyAsync(
132130
() -> {
133131
try {
134-
return SNAPSHOT_FACTORY.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
132+
return ContextSnapshotFactoryHelper.captureFrom(graphQLContext).wrap((Callable<?>) result).call();
135133
}
136134
catch (Exception ex) {
137135
throw new IllegalStateException(

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,6 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
5959

6060
private final SubscriptionExceptionResolver subscriptionExceptionResolver;
6161

62-
private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
6362

6463
private ContextDataFetcherDecorator(
6564
DataFetcher<?> delegate, boolean subscription,
@@ -72,17 +71,18 @@ private ContextDataFetcherDecorator(
7271
this.subscriptionExceptionResolver = subscriptionExceptionResolver;
7372
}
7473

74+
7575
@Override
76-
public Object get(DataFetchingEnvironment environment) throws Exception {
76+
public Object get(DataFetchingEnvironment env) throws Exception {
7777

78-
ContextSnapshot snapshot;
79-
if (environment.getLocalContext() instanceof GraphQLContext localContext) {
80-
snapshot = this.snapshotFactory.captureFrom(environment.getGraphQlContext(), localContext);
81-
}
82-
else {
83-
snapshot = this.snapshotFactory.captureFrom(environment.getGraphQlContext());
84-
}
85-
Object value = snapshot.wrap(() -> this.delegate.get(environment)).call();
78+
GraphQLContext graphQlContext = env.getGraphQlContext();
79+
ContextSnapshotFactory snapshotFactory = ContextSnapshotFactoryHelper.getInstance(graphQlContext);
80+
81+
ContextSnapshot snapshot = (env.getLocalContext() instanceof GraphQLContext localContext) ?
82+
snapshotFactory.captureFrom(graphQlContext, localContext) :
83+
snapshotFactory.captureFrom(graphQlContext);
84+
85+
Object value = snapshot.wrap(() -> this.delegate.get(env)).call();
8686

8787
if (this.subscription) {
8888
Assert.state(value instanceof Publisher, "Expected Publisher for a subscription");

‎spring-graphql/src/main/java/org/springframework/graphql/execution/ContextSnapshotFactoryHelper.java‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
/*
2+
* Copyright 2002-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.graphql.execution;
18+
19+
import graphql.GraphQLContext;
20+
import io.micrometer.context.ContextSnapshot;
21+
import io.micrometer.context.ContextSnapshotFactory;
22+
import reactor.util.context.Context;
23+
import reactor.util.context.ContextView;
24+
25+
import org.springframework.lang.Nullable;
26+
27+
/**
28+
* Helper to use a single {@link ContextSnapshotFactory} instance by saving and
29+
* obtaining it to and from Reactor and GraphQL contexts.
30+
*
31+
* @author Rossen Stoyanchev
32+
* @since 1.3
33+
*/
34+
public abstract class ContextSnapshotFactoryHelper {
35+
36+
private static final ContextSnapshotFactory sharedInstance = ContextSnapshotFactory.builder().build();
37+
38+
private static final String CONTEXT_SNAPSHOT_FACTORY_KEY = ContextSnapshotFactoryHelper.class.getName() + ".KEY";
39+
40+
41+
/**
42+
* Select a {@code ContextSnapshotFactory} instance to use, either the one
43+
* passed in if it is not {@code null}, or a shared, static instance.
44+
* @param factory the candidate factory instance to use if not {@code null}
45+
* @return the instance to use
46+
*/
47+
public static ContextSnapshotFactory selectInstance(@Nullable ContextSnapshotFactory factory) {
48+
if (factory != null) {
49+
return factory;
50+
}
51+
return sharedInstance;
52+
}
53+
54+
/**
55+
* Save the {@code ContextSnapshotFactory} in the given {@link Context}.
56+
* @param factory the instance to save
57+
* @param context the context to save the instance to
58+
* @return a new context with the saved instance
59+
*/
60+
public static Context saveInstance(ContextSnapshotFactory factory, Context context) {
61+
return context.put(CONTEXT_SNAPSHOT_FACTORY_KEY, factory);
62+
}
63+
64+
/**
65+
* Save the {@code ContextSnapshotFactory} in the given {@link Context}.
66+
* @param factory the instance to save
67+
* @param context the context to save the instance to
68+
*/
69+
public static void saveInstance(ContextSnapshotFactory factory, GraphQLContext context) {
70+
context.put(CONTEXT_SNAPSHOT_FACTORY_KEY, factory);
71+
}
72+
73+
/**
74+
* Access the {@code ContextSnapshotFactory} from the given {@link ContextView}
75+
* or return a shared, static instance.
76+
* @param contextView the context where the instance is saved
77+
* @return the instance to use
78+
*/
79+
public static ContextSnapshotFactory getInstance(ContextView contextView) {
80+
ContextSnapshotFactory factory = contextView.getOrDefault(CONTEXT_SNAPSHOT_FACTORY_KEY, null);
81+
return selectInstance(factory);
82+
}
83+
84+
/**
85+
* Access the {@code ContextSnapshotFactory} from the given {@link GraphQLContext}
86+
* or return a shared, static instance.
87+
* @param context the context where the instance is saved
88+
* @return the instance to use
89+
*/
90+
public static ContextSnapshotFactory getInstance(GraphQLContext context) {
91+
ContextSnapshotFactory factory = context.get(CONTEXT_SNAPSHOT_FACTORY_KEY);
92+
return selectInstance(factory);
93+
}
94+
95+
/**
96+
* Shortcut to obtain the {@code ContextSnapshotFactory} instance, and to
97+
* capture from the given {@link ContextView}.
98+
* @param contextView the context to capture from
99+
* @return a snapshot from the capture
100+
*/
101+
public static ContextSnapshot captureFrom(ContextView contextView) {
102+
ContextSnapshotFactory factory = getInstance(contextView);
103+
return selectInstance(factory).captureFrom(contextView);
104+
}
105+
106+
/**
107+
* Shortcut to obtain the {@code ContextSnapshotFactory} instance, and to
108+
* capture from the given {@link GraphQLContext}.
109+
* @param context the context to capture from
110+
* @return a snapshot from the capture
111+
*/
112+
public static ContextSnapshot captureFrom(GraphQLContext context) {
113+
ContextSnapshotFactory factory = getInstance(context);
114+
return selectInstance(factory).captureFrom(context);
115+
}
116+
117+
}

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222

2323
import graphql.GraphQLError;
2424
import graphql.schema.DataFetchingEnvironment;
25-
import io.micrometer.context.ContextSnapshotFactory;
2625
import io.micrometer.context.ThreadLocalAccessor;
2726
import org.apache.commons.logging.Log;
2827
import org.apache.commons.logging.LogFactory;
@@ -53,8 +52,6 @@ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcher
5352

5453
protected final Log logger = LogFactory.getLog(getClass());
5554

56-
protected final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
57-
5855
private boolean threadLocalContextAware;
5956

6057

@@ -101,7 +98,7 @@ private List<GraphQLError> resolveInternal(Throwable exception, DataFetchingEnvi
10198
return resolveToMultipleErrors(exception, env);
10299
}
103100
try {
104-
return this.snapshotFactory.captureFrom(env.getGraphQlContext())
101+
return ContextSnapshotFactoryHelper.captureFrom(env.getGraphQlContext())
105102
.wrap(() -> resolveToMultipleErrors(exception, env))
106103
.call();
107104
}

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828

2929
import graphql.GraphQLContext;
3030
import io.micrometer.context.ContextSnapshot;
31-
import io.micrometer.context.ContextSnapshotFactory;
3231
import org.dataloader.BatchLoaderContextProvider;
3332
import org.dataloader.BatchLoaderEnvironment;
3433
import org.dataloader.BatchLoaderWithContext;
@@ -54,8 +53,6 @@
5453
*/
5554
public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
5655

57-
private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build();
58-
5956
private final List<ReactorBatchLoader<?, ?>> loaders = new ArrayList<>();
6057

6158
private final List<ReactorMappedBatchLoader<?, ?>> mappedLoaders = new ArrayList<>();
@@ -231,7 +228,7 @@ DataLoaderOptions getOptions() {
231228
@Override
232229
public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
233230
GraphQLContext graphQLContext = environment.getContext();
234-
ContextSnapshot snapshot = SNAPSHOT_FACTORY.captureFrom(graphQLContext);
231+
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(graphQLContext);
235232
try {
236233
return snapshot.wrap(() ->
237234
this.loader.apply(keys, environment)
@@ -279,7 +276,7 @@ DataLoaderOptions getOptions() {
279276
@Override
280277
public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
281278
GraphQLContext graphQLContext = environment.getContext();
282-
ContextSnapshot snapshot = SNAPSHOT_FACTORY.captureFrom(graphQLContext);
279+
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(graphQLContext);
283280
try {
284281
return snapshot.wrap(() ->
285282
this.loader.apply(keys, environment)

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,6 @@ public class DefaultExecutionGraphQlService implements ExecutionGraphQlService {
4646
private static final BiFunction<ExecutionInput, ExecutionInput.Builder, ExecutionInput> RESET_EXECUTION_ID_CONFIGURER =
4747
(executionInput, builder) -> builder.executionId(null).build();
4848

49-
private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
5049

5150
private final GraphQlSource graphQlSource;
5251

@@ -90,7 +89,10 @@ public final Mono<ExecutionGraphQlResponse> execute(ExecutionGraphQlRequest requ
9089

9190
ExecutionInput executionInput = request.toExecutionInput();
9291

93-
this.snapshotFactory.captureFrom(contextView).updateContext(executionInput.getGraphQLContext());
92+
ContextSnapshotFactory factory = ContextSnapshotFactoryHelper.getInstance(contextView);
93+
GraphQLContext graphQLContext = executionInput.getGraphQLContext();
94+
ContextSnapshotFactoryHelper.saveInstance(factory, graphQLContext);
95+
factory.captureFrom(contextView).updateContext(graphQLContext);
9496

9597
ExecutionInput updatedExecutionInput =
9698
(this.hasDataLoaderRegistrations ? registerDataLoaders(executionInput) : executionInput);

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
import graphql.execution.ExecutionId;
3030
import graphql.schema.DataFetchingEnvironment;
3131
import io.micrometer.context.ContextSnapshot;
32-
import io.micrometer.context.ContextSnapshotFactory;
3332
import org.apache.commons.logging.Log;
3433
import org.apache.commons.logging.LogFactory;
3534
import reactor.core.publisher.Flux;
@@ -47,8 +46,6 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler
4746

4847
private static final Log logger = LogFactory.getLog(ExceptionResolversExceptionHandler.class);
4948

50-
private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
51-
5249
private final List<DataFetcherExceptionResolver> resolvers;
5350

5451
/**
@@ -65,7 +62,7 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler
6562
public CompletableFuture<DataFetcherExceptionHandlerResult> handleException(DataFetcherExceptionHandlerParameters params) {
6663
Throwable exception = unwrapException(params);
6764
DataFetchingEnvironment env = params.getDataFetchingEnvironment();
68-
ContextSnapshot snapshot = this.snapshotFactory.captureFrom(env.getGraphQlContext());
65+
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(env.getGraphQlContext());
6966
try {
7067
return Flux.fromIterable(this.resolvers)
7168
.flatMap((resolver) -> resolver.resolveException(exception, env))

‎spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222

2323
import graphql.GraphQLError;
2424
import io.micrometer.context.ContextSnapshot;
25-
import io.micrometer.context.ContextSnapshotFactory;
2625
import io.micrometer.context.ThreadLocalAccessor;
2726
import org.apache.commons.logging.Log;
2827
import org.apache.commons.logging.LogFactory;
@@ -51,8 +50,6 @@ public abstract class SubscriptionExceptionResolverAdapter implements Subscripti
5150

5251
protected final Log logger = LogFactory.getLog(getClass());
5352

54-
protected final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
55-
5653
private boolean threadLocalContextAware;
5754

5855

@@ -86,7 +83,7 @@ public boolean isThreadLocalContextAware() {
8683
public final Mono<List<GraphQLError>> resolveException(Throwable exception) {
8784
if (this.threadLocalContextAware) {
8885
return Mono.deferContextual((contextView) -> {
89-
ContextSnapshot snapshot = this.snapshotFactory.captureFrom(contextView);
86+
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(contextView);
9087
try {
9188
List<GraphQLError> errors = snapshot.wrap(() -> resolveToMultipleErrors(exception)).call();
9289
return Mono.justOrEmpty(errors);

‎spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import reactor.core.publisher.Mono;
2626

2727
import org.springframework.graphql.ExecutionGraphQlService;
28+
import org.springframework.graphql.execution.ContextSnapshotFactoryHelper;
2829
import org.springframework.graphql.server.WebGraphQlInterceptor.Chain;
2930
import org.springframework.lang.Nullable;
3031
import org.springframework.util.Assert;
@@ -41,6 +42,9 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder {
4142

4243
private final List<WebGraphQlInterceptor> interceptors = new ArrayList<>();
4344

45+
@Nullable
46+
private ContextSnapshotFactory snapshotFactory;
47+
4448
@Nullable
4549
private WebSocketGraphQlInterceptor webSocketInterceptor;
4650

@@ -68,10 +72,16 @@ public WebGraphQlHandler.Builder interceptors(List<WebGraphQlInterceptor> interc
6872
return this;
6973
}
7074

75+
@Override
76+
public WebGraphQlHandler.Builder contextSnapshotFactory(ContextSnapshotFactory snapshotFactory) {
77+
this.snapshotFactory = snapshotFactory;
78+
return this;
79+
}
80+
7181
@Override
7282
public WebGraphQlHandler build() {
7383

74-
ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build();
84+
ContextSnapshotFactory snapshotFactory = ContextSnapshotFactoryHelper.selectInstance(this.snapshotFactory);
7585

7686
Chain endOfChain = (request) -> this.service.execute(request).map(WebGraphQlResponse::new);
7787

@@ -88,10 +98,18 @@ public WebSocketGraphQlInterceptor getWebSocketInterceptor() {
8898
DefaultWebGraphQlHandlerBuilder.this.webSocketInterceptor : new WebSocketGraphQlInterceptor() { };
8999
}
90100

101+
@Override
102+
public ContextSnapshotFactory contextSnapshotFactory() {
103+
return snapshotFactory;
104+
}
105+
91106
@Override
92107
public Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request) {
93108
ContextSnapshot snapshot = snapshotFactory.captureAll();
94-
return executionChain.next(request).contextWrite(snapshot::updateContext);
109+
return executionChain.next(request).contextWrite((context) -> {
110+
context = ContextSnapshotFactoryHelper.saveInstance(snapshotFactory, context);
111+
return snapshot.updateContext(context);
112+
});
95113
}
96114
};
97115
}

0 commit comments

Comments
 (0)